Assigning Function Units

In scientific computing, a substantial number of existing scientific computing functions are designed based on dimensionless data. brainunit provides interface that is applicable to these dimensionless functions without modifying existing frameworks or underlying implementations. The core idea is:

• Dimensionless processing before function calls: Prior to invoking scientific computing functions, input data undergoes dimensionless processing to ensure that the functions internally handle only unitless numerical operations. For example, using b = a.to_decimal(UNIT) method can normalize the quantity a into the dimensionless data b according to the given physical unit UNIT.

• Restoring physical units after computation: Once the computation is complete and results are returned, we can restore the appropriate physical units to the solutions.

Specifically, brainunit provides the assign_units function, which facilitates the automatic assignment and restoration of physical units at the input and output stages of functions. This method is, in principle, applicable to any Python-based scientific computing library, preserving the physical semantics and interpretability at the input and output levels without altering their existing structures.

First, we need to import the necessary libraries and modules.

import brainunit
from brainunit import volt, mV, meter, second, check_dims, check_units, assign_units, DimensionMismatchError, UnitMismatchError

assign_units Decorator

The assign_units decorator is used to automatically assign units to the input arguments or return values of a function. It ensures that the values are converted to the specified units, simplifying unit handling in scientific computations.

Basic Usage

We can use the assign_units decorator to automatically assign units to the input arguments of a function.

@assign_units(v=volt)
def a_function(v, x):
    """
    v will be assigned units of volt, and x can have any (or no) unit.
    """
    return v

Correct Units

The following calls are correct because the v argument is automatically converted to volts.

assert a_function(3 * mV, 5 * second) == (3 * mV).to_decimal(volt)
assert a_function(3 * volt, 5 * second) == (3 * volt).to_decimal(volt)
assert a_function(5 * volt, "something") == (5 * volt).to_decimal(volt)

Incorrect Units

The following calls will raise a UnitMismatchError or TypeError because the v argument cannot be converted to volts.

try:
    a_function(5 * second, None)
except UnitMismatchError as e:
    print(e)

try:
    a_function(5, None)
except TypeError as e:
    print(e)

try:
    a_function(object(), None)
except TypeError as e:
    print(e)

Cannot convert to the decimal number using a unit with different dimensions. (units are s and V).
Function 'a_function' expected a Quantity object for argument 'v' but got '5'
Function 'a_function' expected a Quantity object for argument 'v' but got '<object object at 0x7f495024b240>'

Assigning Units to Return Values

The assign_units decorator can also be used to automatically assign units to the return value of a function.

@assign_units(result=second)
def b_function():
    """
    The return value will be assigned units of seconds.
    """
    return 5

Correct Return Value

The following call is correct because the return value is automatically converted to seconds.

assert b_function() == 5 * second

Assigning Units to Multiple Return Values

The assign_units decorator can also assign units to multiple return values.

@assign_units(result=(second, volt))
def d_function():
    """
    The return values will be assigned units of seconds and volts, respectively.
    """
    return 5, 3

Correct Return Values

The following call is correct because the return values are automatically converted to seconds and volts.

assert d_function()[0] == 5 * second
assert d_function()[1] == 3 * volt

Assigning Units to Dictionary Return Values

The assign_units decorator can also assign units to dictionary return values.

@assign_units(result={'u': second, 'v': (volt, meter)})
def d_function2(true_result):
    """
    The return values will be assigned units based on the dictionary specification.
    """
    if true_result == 0:
        return {'u': 5, 'v': (3, 2)}
    elif true_result == 1:
        return 3, 5
    else:
        return 3, 5

Correct Return Values

The following call is correct because the return values are automatically converted to the specified units.

d_function2(0)

{'u': Quantity(5, "s"), 'v': (Quantity(3, "V"), Quantity(2, "m"))}

Incorrect Return Values

The following call will raise a TypeError because the return values do not match the expected structure.

try:
    d_function2(1)
except TypeError as e:
    print(e)

Expected a return value of pytree PyTreeDef({'u': *, 'v': (*, *)}) with type {'u': Unit("s"), 'v': (Unit("V"), Unit("m"))}, but got the pytree PyTreeDef((*, *)) and the value (3, 5)

Through the examples above, we can see the utility of the assign_units decorator in automatically assigning units to input arguments and return values. It simplifies unit handling in scientific computations, ensuring consistency and reducing the likelihood of errors.