"""
Unit handling components, based on the `Pint <https://github.com/hgrecco/pint>`__
library.
.. note::
By default,
`Pint's application registry <https://pint.readthedocs.io/en/stable/getting/pint-in-your-projects.html#having-a-shared-registry>`__
is used.
"""
from __future__ import annotations
from typing import Any
import xarray as xr
import pint
# Internal unit registry. If None, use application registry
_ureg: pint.UnitRegistry | None = None
[docs]
def set_unit_registry(ureg: pint.UnitRegistry | None) -> None:
"""
Set internal unit registry.
Parameters
----------
ureg : pint.UnitRegistry or None
New default unit registry. Set to ``None`` to reset to the default
(the application registry).
"""
global _ureg
_ureg = ureg
[docs]
def get_unit_registry() -> pint.UnitRegistry:
"""
Access the internal unit registry. By default, the Pint application registry
is returned.
"""
global _ureg
return _ureg if _ureg is not None else pint.get_application_registry()
[docs]
def ensure_units(
value: Any, default_units: pint.Unit, convert: bool = False
) -> pint.Quantity:
"""
Ensure that a value is wrapped in a Pint quantity container.
Parameters
----------
value
Checked value.
default_units : pint.Unit
Units to use to initialize the :class:`pint.Quantity` if ``value`` is
not a :class:`pint.Quantity`.
convert : bool, default: False
If ``True``, ``value`` will also be converted to ``default_units`` if it
is a :class:`pint.Quantity`.
Returns
-------
Converted ``value``.
"""
if isinstance(value, pint.Quantity):
if convert:
return value.to(default_units)
else:
return value
else:
return value * default_units
[docs]
def xarray_to_quantity(da: xr.DataArray) -> pint.Quantity:
"""
Converts a :class:`~xarray.DataArray` to a :class:`~pint.Quantity`.
The array's ``attrs`` metadata mapping must contain a ``units`` field.
Parameters
----------
da : DataArray
:class:`~xarray.DataArray` instance which will be converted.
Returns
-------
quantity
The corresponding Pint quantity.
Raises
------
ValueError
If array attributes do not contain a ``units`` field.
Notes
-----
This function can also be used on coordinate variables.
"""
try:
units = da.attrs["units"]
except KeyError as e:
raise ValueError("this DataArray has no 'units' attribute field") from e
ureg = get_unit_registry()
return ureg.Quantity(da.values, units)
