Wrapping and checking functions#
In some cases you might want to use pint with a pre-existing web service or library which is not units aware. Or you might want to write a fast implementation of a numerical algorithm that requires the input values in some specific units.
For example, consider a function to return the period of the pendulum within a hypothetical physics library. The library does not use units, but instead requires you to provide numerical values in certain units:
>>> from simple_physics import pendulum_period
>>> help(pendulum_period)
Help on function pendulum_period in module simple_physics:
pendulum_period(length)
Return the pendulum period in seconds. The length of the pendulum
must be provided in meters.
>>> pendulum_period(1)
2.0064092925890407
This behaviour is very error prone, in particular when combining multiple libraries. You could wrap this function to use Quantities instead:
>>> from pint import UnitRegistry
>>> ureg = UnitRegistry()
>>> Q_ = ureg.Quantity
>>> def mypp_caveman(length):
... return pendulum_period(length.to(ureg.meter).magnitude) * ureg.second
and:
>>> mypp_caveman(100 * ureg.centimeter)
<Quantity(2.00640929, 'second')>
Pint provides a more convenient way to do this:
>>> mypp = ureg.wraps(ureg.second, ureg.meter)(pendulum_period)
Or in the decorator format:
>>> @ureg.wraps(ureg.second, ureg.meter)
... def mypp(length):
... return pendulum_period(length)
>>> mypp(100 * ureg.centimeter)
<Quantity(2.00640929, 'second')>
wraps takes 3 input arguments:
- ret: the return units.
Use None to skip conversion.
- args: the inputs units for each argument, as an iterable.
Use None to skip conversion of any given element.
- strict: if True all convertible arguments must be a Quantity
and others will raise a ValueError (True by default)
Strict Mode#
By default, the function is wrapped in strict mode. In this mode, the input arguments assigned to units must be a Quantities.
>>> mypp(1. * ureg.meter)
<Quantity(2.00640929, 'second')>
>>> mypp(1.)
Traceback (most recent call last):
...
ValueError: A wrapped function using strict=True requires quantity for all arguments with not None units. (error found for meter, 1.0)
To enable using non-Quantity numerical values, set strict to False`.
>>> mypp_ns = ureg.wraps(ureg.second, ureg.meter, False)(pendulum_period)
>>> mypp_ns(1. * ureg.meter)
<Quantity(2.00640929, 'second')>
>>> mypp_ns(1.)
<Quantity(2.00640929, 'second')>
In this mode, the value is assumed to have the correct units.
Multiple arguments or return values#
For a function with more arguments, use a tuple:
>>> from simple_physics import pendulum_period2
>>> help(pendulum_period2)
Help on function pendulum_period2 in module simple_physics:
pendulum_period2(length, swing_amplitude)
Return the pendulum period in seconds. The length of the pendulum
must be provided in meters. The swing_amplitude must be in radians.
>>> mypp2 = ureg.wraps(ureg.second, (ureg.meter, ureg.radians))(pendulum_period2)
...
Or if the function has multiple outputs:
>>> mypp3 = ureg.wraps((ureg.second, ureg.meter / ureg.second),
... (ureg.meter, ureg.radians))(pendulum_period_maxspeed)
...
If there are more return values than specified units, None
is assumed for
the extra outputs. For example, given the NREL SOLPOS calculator that outputs
solar zenith, azimuth and air mass, the following wrapper assumes no units for
airmass
@ureg.wraps(('deg', 'deg'), ('deg', 'deg', 'millibar', 'degC'))
def solar_position(lat, lon, press, tamb, timestamp):
return zenith, azimuth, airmass
Optional arguments#
For a function with named keywords with optional values, use a tuple for all arguments:
>>> @ureg.wraps(ureg.second, (ureg.meters, ureg.meters/ureg.second**2, None))
... def calculate_time_to_fall(height, gravity=Q_(9.8, 'm/s^2'), verbose=False):
... """Calculate time to fall from a height h.
...
... By default, the gravity is assumed to be earth gravity,
... but it can be modified.
...
... d = .5 * g * t**2
... t = sqrt(2 * d / g)
... """
... t = math.sqrt(2 * height / gravity)
... if verbose: print(str(t) + " seconds to fall")
... return t
...
>>> lunar_module_height = Q_(22, 'feet') + Q_(11, 'inches')
>>> calculate_time_to_fall(lunar_module_height, verbose=True)
1.1939473204801092 seconds to fall
<Quantity(1.19394732, 'second')>
>>> moon_gravity = Q_(1.625, 'm/s^2')
>>> calculate_time_to_fall(lunar_module_height, gravity=moon_gravity)
<Quantity(2.932051, 'second')>
Specifying relations between arguments#
In certain cases, you may not be concerned with the actual units and only care about the unit relations among arguments.
This is done using a string starting with the equal sign =:
>>> @ureg.wraps('=A**2', ('=A', '=A'))
... def sqsum(x, y):
... return x * x + 2 * x * y + y * y
which can be read as the first argument (x) has certain units (we labeled them A), the second argument (y) has the same units as the first (A again). The return value has the unit of x squared (A**2)
You can use more than one label:
>>> @ureg.wraps('=A**2*B', ('=A', '=A*B', '=B'))
... def some_function(x, y, z):
... pass
With optional arguments
>>> @ureg.wraps('=A*B', ('=A', '=B'))
... def get_displacement(time, rate=Q_(1, 'm/s')):
... return time * rate
...
>>> get_displacement(Q_(2, 's'))
<Quantity(2, 'meter')>
>>> get_displacement(Q_(2, 's'), Q_(1, 'deg/s'))
<Quantity(2, 'degree')>
Ignoring an argument or return value#
To avoid the conversion of an argument or return value, use None
>>> mypp3 = ureg.wraps((ureg.second, None), ureg.meter)(pendulum_period_error)
Checking dimensionality#
When you want pint quantities to be used as inputs to your functions, pint provides a wrapper to ensure units are of correct type - or more precisely, they match the expected dimensionality of the physical quantity.
Similar to wraps(), you can pass None to skip checking of some parameters, but the return parameter type is not checked.
>>> mypp = ureg.check('[length]')(pendulum_period)
In the decorator format:
>>> @ureg.check('[length]')
... def pendulum_period(length):
... return 2*math.pi*math.sqrt(length/G)
If you just want to check the dimensionality of a quantity, you can do so with the built-in ‘check’ function.
>>> distance = 1 * ureg.m
>>> distance.check('[length]')
True
>>> distance.check('[time]')
False