Skip to content

clean-up and implement the conversion methods #11

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 33 commits into from
Jul 10, 2020
Merged

clean-up and implement the conversion methods #11

Show file tree
Hide file tree
Changes from 30 commits
Commits
Show all changes
33 commits
Select commit Hold shift + click to select a range
e9f4c91
add a function to attach units to xarray objects
keewis Jun 16, 2020
d047f0d
remove the old and commented out assert_equal_with_units function
keewis Jun 16, 2020
d2c4a04
add utility functions to attach, convert, extract, and strip units
keewis Jun 27, 2020
f8567c2
add the missing assert functions
keewis Jun 27, 2020
33ddd39
reorder the array functions
keewis Jun 27, 2020
b1d9268
implement a extract_units function
keewis Jun 28, 2020
00ea3fd
implement convert_units
keewis Jun 28, 2020
54f385b
implement attach_units
keewis Jun 28, 2020
3101db2
implement strip_units
keewis Jun 28, 2020
b37bdd4
isort and black
keewis Jun 28, 2020
013b056
implement DataArray.pint.to
keewis Jun 28, 2020
860aad2
implement Dataset.pint.to
keewis Jun 28, 2020
0f076a6
set units as optional
keewis Jun 28, 2020
5b1d334
remove a leftover comment about mypy
keewis Jun 29, 2020
e3d3e72
update the minimum required pint version
keewis Jul 1, 2020
6855b31
Merge branch 'master' into to
keewis Jul 8, 2020
3fa00aa
Merge branch 'master' into to
keewis Jul 8, 2020
5052f76
document Dataset.pint.to
keewis Jul 8, 2020
9a17b22
add docstrings to the conversion methods
keewis Jul 8, 2020
50a7c8e
don't require the variable names to be str
keewis Jul 8, 2020
cc999e4
don't try to always convert the DataArray's unit
keewis Jul 8, 2020
cf47518
add the example output
keewis Jul 8, 2020
e56aa30
make sure conversion of only the coords works
keewis Jul 8, 2020
e970515
don't raise if the default value for units is passed
keewis Jul 8, 2020
30b75bc
undo the test changes
keewis Jul 8, 2020
f2c17c3
properly separate the conversions
keewis Jul 8, 2020
2e07ad8
try separating the example descriptions
keewis Jul 8, 2020
1e20fda
properly refer to pint.Unit
keewis Jul 8, 2020
3077b60
show how to use a dict-like to convert the data
keewis Jul 8, 2020
c875299
fix the docstrings
keewis Jul 8, 2020
682ef84
add a comment on the source of the utility function
keewis Jul 9, 2020
426c607
link to the source and add xarray's license
keewis Jul 10, 2020
7ac1c56
add a copyright notice to the xarray license
keewis Jul 10, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions docs/api.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ Dataset

Dataset.pint.quantify
Dataset.pint.dequantify
Dataset.pint.to
Dataset.pint.to_base_units
Dataset.pint.to_system

Expand Down
2 changes: 1 addition & 1 deletion docs/conf.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
import sphinx_autosummary_accessors

# need to import so accessors get registered
import pintxarray # noqa: F401
import pint_xarray # noqa: F401

# -- Project information -----------------------------------------------------

Expand Down
236 changes: 227 additions & 9 deletions pint_xarray/accessors.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@
)
from xarray.core.npcompat import IS_NEP18_ACTIVE

from . import conversion

if not hasattr(Quantity, "__array_function__"):
raise ImportError(
"Imported version of pint does not implement " "__array_function__"
Expand All @@ -29,6 +31,26 @@
# TODO type hints


def is_dict_like(obj):
return hasattr(obj, "keys") and hasattr(obj, "__getitem__")


def either_dict_or_kwargs(positional, keywords, method_name):
if positional is not None:
if not is_dict_like(positional):
raise ValueError(
f"the first argument to .{method_name} must be a dictionary"
)
if keywords:
raise ValueError(
"cannot specify both keyword and positional "
f"arguments to .{method_name}"
)
return positional
else:
return keywords


def _array_attach_units(data, unit, convert_from=None):
"""
Internal utility function for attaching units to a numpy-like array,
Expand Down Expand Up @@ -230,15 +252,109 @@ def registry(self):
def registry(self, _):
raise AttributeError("Don't try to change the registry once created")

def to(self, units):
quantity = self.da.data.to(units)
return DataArray(
dim=self.da.dims,
data=quantity,
coords=self.da.coords,
attrs=self.da.attrs,
encoding=self.da.encoding,
)
def to(self, units=None, **unit_kwargs):
""" convert the quantities in a DataArray

Parameters
----------
units : str or pint.Unit or mapping of hashable to str or pint.Unit, optional
The units to convert to. If a unit name or
:py:class`pint.Unit` object, convert the DataArray's
data. If a dict-like, it has to map a variable name to a
unit name or :py:class:`pint.Unit` object.
**unit_kwargs
The kwargs form of ``units``. Can only be used for
variable names that are strings and valid python identifiers.

Returns
-------
object : DataArray
A new object with converted units.

Examples
--------
>>> da = xr.DataArray(
... data=np.linspace(0, 1, 5) * ureg.m,
... coords={"u": ("x", np.arange(5) * ureg.s)},
... dims="x",
... name="arr",
... )
>>> da
<xarray.DataArray 'arr' (x: 5)>
<Quantity([0. 0.25 0.5 0.75 1. ], 'meter')>
Coordinates:
u (x) int64 <Quantity([0 1 2 3 4], 'second')>
Dimensions without coordinates: x

Convert the data

>>> da.pint.to("mm")
<xarray.DataArray 'arr' (x: 5)>
<Quantity([ 0. 250. 500. 750. 1000.], 'millimeter')>
Coordinates:
u (x) int64 <Quantity([0 1 2 3 4], 'second')>
Dimensions without coordinates: x
>>> da.pint.to(ureg.mm)
<xarray.DataArray 'arr' (x: 5)>
<Quantity([ 0. 250. 500. 750. 1000.], 'millimeter')>
Coordinates:
u (x) int64 <Quantity([0 1 2 3 4], 'second')>
Dimensions without coordinates: x
>>> da.pint.to({da.name: "mm"})
<xarray.DataArray 'arr' (x: 5)>
<Quantity([ 0. 250. 500. 750. 1000.], 'millimeter')>
Coordinates:
u (x) int64 <Quantity([0 1 2 3 4], 'second')>
Dimensions without coordinates: x

Convert coordinates

>>> da.pint.to({"u": ureg.ms})
<xarray.DataArray 'arr' (x: 5)>
<Quantity([0. 0.25 0.5 0.75 1. ], 'meter')>
Coordinates:
u (x) float64 <Quantity([ 0. 1000. 2000. 3000. 4000.], 'millisec...
Dimensions without coordinates: x
>>> da.pint.to(u="ms")
<xarray.DataArray 'arr' (x: 5)>
<Quantity([0. 0.25 0.5 0.75 1. ], 'meter')>
Coordinates:
u (x) float64 <Quantity([ 0. 1000. 2000. 3000. 4000.], 'millisec...
Dimensions without coordinates: x

Convert both simultaneously

>>> da.pint.to("mm", u="ms")
<xarray.DataArray 'arr' (x: 5)>
<Quantity([ 0. 250. 500. 750. 1000.], 'millimeter')>
Coordinates:
u (x) float64 <Quantity([ 0. 1000. 2000. 3000. 4000.], 'millisec...
Dimensions without coordinates: x
>>> da.pint.to({"arr": ureg.mm, "u": ureg.ms})
<xarray.DataArray 'arr' (x: 5)>
<Quantity([ 0. 250. 500. 750. 1000.], 'millimeter')>
Coordinates:
u (x) float64 <Quantity([ 0. 1000. 2000. 3000. 4000.], 'millisec...
Dimensions without coordinates: x
>>> da.pint.to(arr="mm", u="ms")
<xarray.DataArray 'arr' (x: 5)>
<Quantity([ 0. 250. 500. 750. 1000.], 'millimeter')>
Coordinates:
u (x) float64 <Quantity([ 0. 1000. 2000. 3000. 4000.], 'millisec...
Dimensions without coordinates: x
"""
if isinstance(units, (str, pint.Unit)):
unit_kwargs[self.da.name] = units
units = None
elif units is not None and not is_dict_like(units):
raise ValueError(
"units must be either a string, a pint.Unit object or a dict-like,"
f" but got {units!r}"
)

units = either_dict_or_kwargs(units, unit_kwargs, "to")

return conversion.convert_units(self.da, units)

def to_base_units(self):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should to_base_units also be adjusted analogously to use a util from conversion or is it fine as-is?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

to_base_units uses Quantity.to_base_units, so I'm not quite sure how we'd be able to use a conversion function here. Did you want to add a function to conversion that calls .to_base_units? If so, then yes, we could use that to reduce duplicated code.

We should also allow passing in a dict / kwargs to control which of the DataArray's variables should be converted (defaulting to "convert all variables").

Both should be a new PR, though.

quantity = self.da.data.to_base_units()
Expand Down Expand Up @@ -346,6 +462,108 @@ def dequantify(self):
}
return Dataset(dequantified_vars, coords=self.ds.coords, attrs=self.ds.attrs)

def to(self, units=None, **unit_kwargs):
""" convert the quantities in a DataArray

Parameters
----------
units : mapping of hashable to str or pint.Unit, optional
Maps variable names to the unit to convert to.
**unit_kwargs
The kwargs form of ``units``. Can only be used for
variable names that are strings and valid python identifiers.

Returns
-------
object : DataArray
A new object with converted units.

Examples
--------
>>> ds = xr.Dataset(
... data_vars={
... "a": ("x", np.linspace(0, 1, 5) * ureg.m),
... "b": ("x", np.linspace(-1, 0, 5) * ureg.kg),
... },
... coords={"u": ("x", np.arange(5) * ureg.s)},
... )
>>> ds
<xarray.Dataset>
Dimensions: (x: 5)
Coordinates:
u (x) int64 <Quantity([0 1 2 3 4], 'second')>
Dimensions without coordinates: x
Data variables:
a (x) float64 <Quantity([0. 0.25 0.5 0.75 1. ], 'meter')>
b (x) float64 <Quantity([-1. -0.75 -0.5 -0.25 0. ], 'kilogram')>

Convert the data

>>> ds.pint.to({"a": "mm", "b": ureg.g})
<xarray.Dataset>
Dimensions: (x: 5)
Coordinates:
u (x) int64 <Quantity([0 1 2 3 4], 'second')>
Dimensions without coordinates: x
Data variables:
a (x) float64 <Quantity([ 0. 250. 500. 750. 1000.], 'millimet...
b (x) float64 <Quantity([-1000. -750. -500. -250. 0.], 'gra...
>>> ds.pint.to(a=ureg.mm, b="g")
<xarray.Dataset>
Dimensions: (x: 5)
Coordinates:
u (x) int64 <Quantity([0 1 2 3 4], 'second')>
Dimensions without coordinates: x
Data variables:
a (x) float64 <Quantity([ 0. 250. 500. 750. 1000.], 'millimet...
b (x) float64 <Quantity([-1000. -750. -500. -250. 0.], 'gra...

Convert coordinates

>>> ds.pint.to({"u": ureg.ms})
<xarray.Dataset>
Dimensions: (x: 5)
Coordinates:
u (x) float64 <Quantity([ 0. 1000. 2000. 3000. 4000.], 'millisec...
Dimensions without coordinates: x
Data variables:
a (x) float64 <Quantity([0. 0.25 0.5 0.75 1. ], 'meter')>
b (x) float64 <Quantity([-1. -0.75 -0.5 -0.25 0. ], 'kilogram')>
>>> ds.pint.to(u="ms")
<xarray.Dataset>
Dimensions: (x: 5)
Coordinates:
u (x) float64 <Quantity([ 0. 1000. 2000. 3000. 4000.], 'millisec...
Dimensions without coordinates: x
Data variables:
a (x) float64 <Quantity([0. 0.25 0.5 0.75 1. ], 'meter')>
b (x) float64 <Quantity([-1. -0.75 -0.5 -0.25 0. ], 'kilogram')>

Convert both simultaneously

>>> ds.pint.to(a=ureg.mm, b=ureg.g, u="ms")
<xarray.Dataset>
Dimensions: (x: 5)
Coordinates:
u (x) float64 <Quantity([ 0. 1000. 2000. 3000. 4000.], 'millisec...
Dimensions without coordinates: x
Data variables:
a (x) float64 <Quantity([ 0. 250. 500. 750. 1000.], 'millimet...
b (x) float64 <Quantity([-1000. -750. -500. -250. 0.], 'gra...
>>> ds.pint.to({"a": "mm", "b": "g", "u": ureg.ms})
<xarray.Dataset>
Dimensions: (x: 5)
Coordinates:
u (x) float64 <Quantity([ 0. 1000. 2000. 3000. 4000.], 'millisec...
Dimensions without coordinates: x
Data variables:
a (x) float64 <Quantity([ 0. 250. 500. 750. 1000.], 'millimet...
b (x) float64 <Quantity([-1000. -750. -500. -250. 0.], 'gra...
"""
units = either_dict_or_kwargs(units, unit_kwargs, "to")

return conversion.convert_units(self.ds, units)

def to_base_units(self):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Likewise here?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same here: not sure how to use conversion, but we should definitely add a possibility to only convert certain variables.

base_vars = {name: da.pint.to_base_units() for name, da in self.ds.items()}
return Dataset(base_vars, coords=self.ds.coords, attrs=self.ds.attrs)
Expand Down
Loading