xarray-contrib / pint-xarray

Interface for using pint with xarray, providing convenience accessors
https://pint-xarray.readthedocs.io/en/latest/
Apache License 2.0
101 stars 12 forks source link

@ureg.wraps decorator #141

Open TomNicholas opened 2 years ago

TomNicholas commented 2 years ago

Looks like we forgot to ensure that @ureg.wraps works - see this example (inspired by this story):

from pint_xarray import unit_registry as ureg

@ureg.wraps(ret=None, args=["Newtons / second^2"])
def jpl_trajectory_code(acceleration):
    # do some rocket science
    ...

lockheed_acceleration_value = xr.DataArray(5).pint.quantify("pounds / second^2")

jpl_trajectory_code(lockheed_acceleration_value)

which returns

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-25-c85d33b3d64a> in <module>
----> 1 jpl_trajectory_code(lockheed_acceleration_value)

~/miniconda3/envs/py38-mamba/lib/python3.8/site-packages/pint/registry_helpers.py in wrapper(*values, **kw)
    263             # In principle, the values are used as is
    264             # When then extract the magnitudes when needed.
--> 265             new_values, values_by_name = converter(ureg, values, strict)
    266 
    267             result = func(*new_values, **kw)

~/miniconda3/envs/py38-mamba/lib/python3.8/site-packages/pint/registry_helpers.py in _converter(ureg, values, strict)
    147                         )
    148                     else:
--> 149                         raise ValueError(
    150                             "A wrapped function using strict=True requires "
    151                             "quantity or a string for all arguments with not None units. "

ValueError: A wrapped function using strict=True requires quantity or a string for all arguments with not None units. (error found for Newtons / second ** 2, <xarray.DataArray ()>
<Quantity(5, 'pound / second ** 2')>)

I would expect some kind of DimensionalityError instead.

It does not surprise me that passing an xarray.DataArray toa pint function does not work, but what does make no sense to me is that this still doesn't work even if I pass a pint.Quantity:

lockheed_acceleration_value = pint.Quantity(5, units="pounds / second^2")

In [32]: lockheed_acceleration_value
Out[32]: array(5) <Unit('pound / second ** 2')>

In [33]: jpl_trajectory_code(lockheed_acceleration_value)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-33-c85d33b3d64a> in <module>
----> 1 jpl_trajectory_code(lockheed_acceleration_value)

~/miniconda3/envs/py38-mamba/lib/python3.8/site-packages/pint/registry_helpers.py in wrapper(*values, **kw)
    263             # In principle, the values are used as is
    264             # When then extract the magnitudes when needed.
--> 265             new_values, values_by_name = converter(ureg, values, strict)
    266 
    267             result = func(*new_values, **kw)

~/miniconda3/envs/py38-mamba/lib/python3.8/site-packages/pint/registry_helpers.py in _converter(ureg, values, strict)
    147                         )
    148                     else:
--> 149                         raise ValueError(
    150                             "A wrapped function using strict=True requires "
    151                             "quantity or a string for all arguments with not None units. "

ValueError: A wrapped function using strict=True requires quantity or a string for all arguments with not None units. (error found for Newtons / second ** 2, 5 pound / second ** 2)

Am I using pint wrongly somehow? I can't see how what I am doing is different from the example in the docs.


Either way this definitely isn't going to work as is because the code for pint.registry_helpers.wraps always returns a pint.Quantity.

(Also if we are going to make our own version of this checking decorator can we call it something more informative than wraps? It's like calling a class Klass... I suggest something like @check, @verify, or @expects.)

TomNicholas commented 2 years ago

I was wondering whether pint.wraps should be able to handle any object that defines pint-like methods and attributes, a question that's related to https://github.com/hgrecco/pint/issues/1099. That seems like a lot of work though, so we should probably just define our own function in pint-xarray.