Units with pint_xarray.PintIndex

Units with pint_xarray.PintIndex#

See also

Learn more at the pint-xarray documentation page.

Highlights#

pint-xarray provides an index that wraps other indexes and attaches units to the indexed coordinates. This allows operations like integrate or sel to take the units into account.

Example#

First we open the dataset, fill in missing units attributes, and calculate the length of the vectors for later:

%xmode minimal
import numpy as np
import xarray as xr

xr.set_options(
    display_expand_indexes=True,
    display_expand_attrs=False,
    display_expand_data=False,
)

ds = (
    xr.tutorial.open_dataset("eraint_uvz")
    .load()
    .assign_coords(
        month=lambda ds: ds["month"].assign_attrs({"units": "months"})
    )
    .assign(windspeed=lambda ds: np.hypot(ds["u"], ds["v"]))
)
ds
Exception reporting mode: Minimal
/home/docs/checkouts/readthedocs.org/user_builds/xarray-indexes/envs/39/lib/python3.13/site-packages/xarray/conventions.py:204: SerializationWarning: variable 'z' has non-conforming '_FillValue' np.float64(nan) defined, dropping '_FillValue' entirely.
  var = coder.decode(var, name=name)
/home/docs/checkouts/readthedocs.org/user_builds/xarray-indexes/envs/39/lib/python3.13/site-packages/xarray/conventions.py:204: SerializationWarning: variable 'u' has non-conforming '_FillValue' np.float64(nan) defined, dropping '_FillValue' entirely.
  var = coder.decode(var, name=name)
/home/docs/checkouts/readthedocs.org/user_builds/xarray-indexes/envs/39/lib/python3.13/site-packages/xarray/conventions.py:204: SerializationWarning: variable 'v' has non-conforming '_FillValue' np.float64(nan) defined, dropping '_FillValue' entirely.
  var = coder.decode(var, name=name)
<xarray.Dataset> Size: 22MB
Dimensions:    (month: 2, level: 3, latitude: 241, longitude: 480)
Coordinates:
  * month      (month) int32 8B 1 7
  * level      (level) int32 12B 200 500 850
  * latitude   (latitude) float32 964B 90.0 89.25 88.5 ... -88.5 -89.25 -90.0
  * longitude  (longitude) float32 2kB -180.0 -179.2 -178.5 ... 178.5 179.2
Data variables:
    z          (month, level, latitude, longitude) float64 6MB 1.068e+05 ... ...
    u          (month, level, latitude, longitude) float64 6MB 1.282 ... 3.539
    v          (month, level, latitude, longitude) float64 6MB -0.04676 ... 3...
    windspeed  (month, level, latitude, longitude) float64 6MB 1.283 ... 4.896
Attributes: (2)

Quantifying#

Now we can quantify to convert arrays with a "units" attribute to quantity arrays:

import cf_xarray.units
import pint_xarray

quantified = ds.pint.quantify()
quantified
<xarray.Dataset> Size: 22MB
Dimensions:    (month: 2, level: 3, latitude: 241, longitude: 480)
Coordinates:
  * month      (month) int32 8B [month] 1 7
  * level      (level) int32 12B [mbar] 200 500 850
  * latitude   (latitude) float32 964B [degrees_north] 90.0 89.25 ... -90.0
  * longitude  (longitude) float32 2kB [degrees_east] -180.0 -179.2 ... 179.2
Data variables:
    z          (month, level, latitude, longitude) float64 6MB [m²/s²] 1.068e...
    u          (month, level, latitude, longitude) float64 6MB [m/s] 1.282 .....
    v          (month, level, latitude, longitude) float64 6MB [m/s] -0.04676...
    windspeed  (month, level, latitude, longitude) float64 6MB [m/s] 1.283 .....
Indexes:
    longitude  PintIndex(PandasIndex, units={'longitude': 'degrees_east'})
    latitude   PintIndex(PandasIndex, units={'latitude': 'degrees_north'})
    level      PintIndex(PandasIndex, units={'level': 'mbar'})
    month      PintIndex(PandasIndex, units={'month': 'month'})
Attributes: (2)

Note how all variables are associated with a pint.Quantity array, and how all coordinate variables are associated with a pint_xarray.PintIndex wrapping a PandasIndex.

Selection#

With the PintIndex, selecting with quantities will convert the indexers to the index’ units:

ureg = pint_xarray.unit_registry

quantified.sel(
    latitude=slice(
        ureg.Quantity(4800, "arcmin"), ureg.Quantity(600, "arcmin")
    ),
    longitude=slice(
        ureg.Quantity(-10, "degree"), ureg.Quantity(np.pi, "radians")
    ),
)
<xarray.Dataset> Size: 5MB
Dimensions:    (month: 2, level: 3, latitude: 93, longitude: 253)
Coordinates:
  * month      (month) int32 8B [month] 1 7
  * level      (level) int32 12B [mbar] 200 500 850
  * latitude   (latitude) float32 372B [degrees_north] 79.5 78.75 ... 11.25 10.5
  * longitude  (longitude) float32 1kB [degrees_east] -9.75 -9.0 ... 178.5 179.2
Data variables:
    z          (month, level, latitude, longitude) float64 1MB [m²/s²] 1.075e...
    u          (month, level, latitude, longitude) float64 1MB [m/s] 9.438 .....
    v          (month, level, latitude, longitude) float64 1MB [m/s] 1.211 .....
    windspeed  (month, level, latitude, longitude) float64 1MB [m/s] 9.515 .....
Indexes:
    longitude  PintIndex(PandasIndex, units={'longitude': 'degrees_east'})
    latitude   PintIndex(PandasIndex, units={'latitude': 'degrees_north'})
    level      PintIndex(PandasIndex, units={'level': 'mbar'})
    month      PintIndex(PandasIndex, units={'month': 'month'})
Attributes: (2)

or raise on incompatible units:

quantified.sel(month=ureg.Quantity(10, "m"))
  + Exception Group Traceback (most recent call last):
  |   File "/home/docs/checkouts/readthedocs.org/user_builds/xarray-indexes/envs/39/lib/python3.13/site-packages/IPython/core/interactiveshell.py", line 3701, in run_code
  |     exec(code_obj, self.user_global_ns, self.user_ns)
  |     ~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  |   File "/tmp/ipykernel_735/1181962863.py", line 1, in <module>
  |     quantified.sel(month=ureg.Quantity(10, "m"))
  |     ~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  |   File "/home/docs/checkouts/readthedocs.org/user_builds/xarray-indexes/envs/39/lib/python3.13/site-packages/xarray/core/dataset.py", line 3001, in sel
  |     query_results = map_index_queries(
  |         self, indexers=indexers, method=method, tolerance=tolerance
  |     )
  |   File "/home/docs/checkouts/readthedocs.org/user_builds/xarray-indexes/envs/39/lib/python3.13/site-packages/xarray/core/indexing.py", line 214, in map_index_queries
  |     results.append(index.sel(labels, **options))
  |                    ~~~~~~~~~^^^^^^^^^^^^^^^^^^^
  |   File "/home/docs/checkouts/readthedocs.org/user_builds/xarray-indexes/envs/39/lib/python3.13/site-packages/pint_xarray/index.py", line 63, in sel
  |     converted_labels = conversion.convert_indexer_units(labels, self.units)
  |   File "/home/docs/checkouts/readthedocs.org/user_builds/xarray-indexes/envs/39/lib/python3.13/site-packages/pint_xarray/conversion.py", line 485, in convert_indexer_units
  |     raise create_exception_group(invalid, "convert_indexers")
  | pint_xarray.errors.PintExceptionGroup: Cannot convert indexers (1 sub-exception)
  +-+---------------- 1 ----------------
    | Traceback (most recent call last):
    |   File "/home/docs/checkouts/readthedocs.org/user_builds/xarray-indexes/envs/39/lib/python3.13/site-packages/pint_xarray/conversion.py", line 480, in convert_indexer_units
    |     converted[name] = convert(indexer, indexer_units)
    |                       ~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^
    |   File "/home/docs/checkouts/readthedocs.org/user_builds/xarray-indexes/envs/39/lib/python3.13/site-packages/pint_xarray/conversion.py", line 473, in convert
    |     return array_convert_units(indexer, units)
    |   File "/home/docs/checkouts/readthedocs.org/user_builds/xarray-indexes/envs/39/lib/python3.13/site-packages/pint_xarray/conversion.py", line 89, in array_convert_units
    |     return data.to(unit)
    |            ~~~~~~~^^^^^^
    |   File "/home/docs/checkouts/readthedocs.org/user_builds/xarray-indexes/envs/39/lib/python3.13/site-packages/pint/facets/plain/quantity.py", line 537, in to
    |     magnitude = self._convert_magnitude_not_inplace(other, *contexts, **ctx_kwargs)
    |   File "/home/docs/checkouts/readthedocs.org/user_builds/xarray-indexes/envs/39/lib/python3.13/site-packages/pint/facets/plain/quantity.py", line 480, in _convert_magnitude_not_inplace
    |     return self._REGISTRY.convert(self._magnitude, self._units, other, **ctx_kwargs)
    |            ~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    |   File "/home/docs/checkouts/readthedocs.org/user_builds/xarray-indexes/envs/39/lib/python3.13/site-packages/pint/facets/plain/registry.py", line 1115, in convert
    |     return self._convert(value, src, dst, inplace, **ctx_kwargs)
    |            ~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    |   File "/home/docs/checkouts/readthedocs.org/user_builds/xarray-indexes/envs/39/lib/python3.13/site-packages/pint/facets/context/registry.py", line 410, in _convert
    |     return super()._convert(value, src, dst, inplace, **ctx_kwargs)
    |            ~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    |   File "/home/docs/checkouts/readthedocs.org/user_builds/xarray-indexes/envs/39/lib/python3.13/site-packages/pint/facets/nonmultiplicative/registry.py", line 264, in _convert
    |     return super()._convert(value, src, dst, inplace)
    |            ~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^
    |   File "/home/docs/checkouts/readthedocs.org/user_builds/xarray-indexes/envs/39/lib/python3.13/site-packages/pint/facets/plain/registry.py", line 1150, in _convert
    |     raise factor
    | pint.errors.DimensionalityError: Cannot convert from 'meter' ([length]) to 'month' ([time])
    | incompatible units for indexer for 'month'
    +------------------------------------

Numerical operations#

We can also perform numerical operations, like integration:

quantified["windspeed"].integrate("month")
<xarray.DataArray 'windspeed' (level: 3, latitude: 241, longitude: 480)> Size: 3MB
[m·month/s] 7.107 7.108 7.11 7.09 7.071 7.101 ... 25.59 25.57 25.6 25.56 25.6
Coordinates:
  * level      (level) int32 12B [mbar] 200 500 850
  * latitude   (latitude) float32 964B [degrees_north] 90.0 89.25 ... -90.0
  * longitude  (longitude) float32 2kB [degrees_east] -180.0 -179.2 ... 179.2
Indexes:
    longitude  PintIndex(PandasIndex, units={'longitude': 'degrees_east'})
    latitude   PintIndex(PandasIndex, units={'latitude': 'degrees_north'})
    level      PintIndex(PandasIndex, units={'level': 'mbar'})

Note how the units are displayed as "meter * months / second" and not the expected "meter"? This is caused by pint trying avoid implicit conversions as much as possible, which can substantially reduce the amount of computations.