openscm / openscm-units

Units for use with simple climate modelling
https://openscm-units.readthedocs.io/en/latest/
BSD 3-Clause "New" or "Revised" License
9 stars 5 forks source link

automatic set up of pint_xarray accessors #17

Open mikapfl opened 3 years ago

mikapfl commented 3 years ago

Is your feature request related to a problem? Please describe.

It's annoying that I always have to call pint_xarray.accessors.setup_registry(unit_registry) if I want to use openscm-units together with pint_xarray. Of course, this has to be done only once after importing both, like this:

import xarray as xr
import pint
import pint_xarray
from openscm_units import unit_registry

pint_xarray.accessors.setup_registry(unit_registry)

but still, it is a mouthful and not completely obvious.

Describe the solution you'd like

Ideally, openscm-units would somehow detect if used together with pint_xarray and call the accessor setup automatically. Looking into setup_registry, it enables force_ndarray_like, so that operations with the unit_registry always yield arrays, like this:

In [1]: import xarray as xr
   ...: import pint
   ...: import pint_xarray
   ...: from openscm_units import unit_registry

In [2]: 2 * unit_registry('kg CO2')
Out[2]: 2 <Unit('CO2 * kilogram')>

In [3]: pint_xarray.accessors.setup_registry(unit_registry)
Out[3]: <openscm_units.unit_registry.ScmUnitRegistry at 0x7f6f1f448eb0>

In [4]: 2 * unit_registry('kg CO2')
Out[4]: array(2) <Unit('CO2 * kilogram')>

As a default behaviour, this would probably be annoying, and admittedly, it would be weird to change behaviour just because pint_xarray is available at import time. So maybe we need something like a unit_registry_xarray which is the same as the unit_registry, just with force_ndarray_like=True. Or maybe pint_xarray is too much of a moving target at the moment to integrate so tightly?

Describe alternatives you've considered

Alternatively, the need to call setup_registry could be documented somewhere. I could provide patches for that.

znicholls commented 3 years ago

So maybe we need something like a unit_registry_xarray which is the same as the unit_registry, just with force_ndarray_like=True.

This seems like a good solution (from my very vague understanding of pint and xarray's implementation). Happy to look at a PR.

Or maybe pint_xarray is too much of a moving target at the moment to integrate so tightly?

Given how lightweight this repo is, I'm happy to play around with getting the xarray integration right. It would be interesting to see (and seems a much more promising route forward than our pandas integration which is currently stuck e.g. https://github.com/pandas-dev/pandas/pull/35127).

Alternatively, the need to call setup_registry could be documented somewhere.

We should definitely do this (and/or add a reference to the pint xarray docs so people know where to look for the latest information).