Pint support for DataArray

[keewis]

This is part of the effort to add support for pint (see #3594) on DataArray. This depends on #3611, so we should merge that first.

• Tests added
• Passes black . && mypy . && flake8
• Fully documented, including whats-new.rst for all changes and api.rst for new API

list of failing tests from #3594:

• __init__: Needs unit support in IndexVariable, blocked by Explicit indexes in xarray's data-model (Future of MultiIndex) #1603
• aggregation:
  • pint currently does not implement prod
  • calling np.median returns a numpy scalar instead of a DataArray because DataArray does not implement __array_function__
  • everything else works
• bivariate ufuncs: fixed by ufunc calls with mixed args and upcast types hgrecco/pint#951 (but the unit won't be correct if the array is not passed first. See Make commutative operations commutative in both magnitude and units hgrecco/pint#1019)
• numpy methods: rank only works with numpy.ndarray (uses bottleneck)
• ffill, bfill: uses bottleneck, which does not support NEP-18?
• where: works (but the choice of treating array(nan) and array(0) as special needs to be discussed)
• interpolate_na: enforces the use of numpy.vectorize which does not support NEP18
• combine_first: works, but does not test indexing (which should fail)
• equals, identical: works (right now identical returns the same as equals)
• drop_sel, sel, loc: indexes strip units (Explicit indexes in xarray's data-model (Future of MultiIndex) #1603)
• interp: uses scipy.interpolate.interp1d which strips units
• interp_like: same as interp
• to_unstacked_dataset: blocked by IndexVariable
• quantile: works (but needs a new pint version since it uses numpy.nanquantile
• rolling: uses numpy.lib.stride_tricks.as_strided which is not supported by NEP18
• rolling_exp: numbagg

[keewis]

While trying to get the tests to pass, I encountered a few issues:

• the implementation of __array_ufunc__:

  xarray/xarray/core/arithmetic.py

  Lines 37 to 39 in 3cbc459

  	for x in inputs + out:
  	if not isinstance(x, self._HANDLED_TYPES + (SupportsArithmetic,)):
  	return NotImplemented

  checks if it handles all types. However, the way it is implemented will cause it to return NotImplemented for duck arrays, even if it would be fine to treat them the same as ndarray. Adding and not hasattr(x, "__array_function__") will make my tests pass, but I thought I'd ask first.
• I've been testing if the aggregation functions work, both as method and as function. This works fine for mean but not for median: calling np.median(da) returns a numpy scalar instead of the data array wrapping a 0d array returned by da.median(). I can't seem to find the reason why this happens, mostly because I don't know where to look.

[keewis]

The aggregation tests fail because numpy.median needs __array_function__ to be implemented, but I think that's out of scope for this PR. I marked it as xfailing.

The bivariate ufunc test is fixed by hgrecco/pint#951.

shift could be fixed by using

        filler = np.full_like(np.resize(trimmed_data, shape), fill_value, dtype=dtype)

instead of

xarray/xarray/core/variable.py

Lines 1105 to 1112 in 651f27f

	if isinstance(trimmed_data, dask_array_type):
	chunks = list(trimmed_data.chunks)
	chunks[axis] = (shape[axis],)
	full = functools.partial(da.full, chunks=chunks)
	else:
	full = np.full
	filler = full(shape, fill_value, dtype=dtype)

but I didn't check if that would be a bad for performance (and pint would have to implement np.resize). However, that would not work for dask arrays so I'm not sure what to do.

rolling uses bottleneck which does not play nice with pint, yet. I'd suggest using

       	if bottleneck_move_func	is not None and not isinstance(
	    self.obj.data, dask_array_type
	) and not hasattr(self.obj.data, "__array_function__"):

instead of

xarray/xarray/core/rolling.py

Lines 363 to 365 in 651f27f

	if bottleneck_move_func is not None and not isinstance(
	self.obj.data, dask_array_type
	):

If we go with that, we'd need pint to implement np.pad.

Edit: I added it, but this still requires np.pad. However, it makes a lot of other, unrelated tests fail, so I removed the commit again.

Unless I missed something, these should be the only issues left for pint support in DataArray (not counting the IndexVariable issue).

[keewis]

@jthielen, it looks like this needs numpy.pad and numpy.resize in pint. I tried to implement it myself, but while numpy.resize is easy (add an entry to the implement_consistent_units_by_argument loop), that is not the case for numpy.pad. Regardless of whether my proposed fixes are actually the best way, numpy.pad is used in other places of the code base so we'd probably need both, anyway. What do you think?

[jthielen]

@jthielen, it looks like this needs numpy.pad and numpy.resize in pint. I tried to implement it myself, but while numpy.resize is easy (add an entry to the implement_consistent_units_by_argument loop), that is not the case for numpy.pad. Regardless of whether my proposed fixes are actually the best way, numpy.pad is used in other places of the code base so we'd probably need both, anyway. What do you think?

No problem at all to add these! See hgrecco/pint#956.

[crusaderky]

if you merge from master you should get all green on the upstream-dev test now

[keewis]

most of these are failing tests related to this PR but it is on hold until #3706 and #3611 (in that order) have been merged.

[keewis]

A few updates on the progress here:

The implementation of interpolate_na uses apply_ufunc with vectorize=True which in turn uses numpy.vectorize. numpy.vectorize is explicitly mentioned as deferred in NEP18 (i.e. not supported). I guess we can't support interpolate_na unless someone knows a way around that? Passing vectorize=False to apply_ufunc does make it work, though.

The failure of test_bivariate_ufunc is due to the way pint infers units, but whether or not that should matter was discussed in #3706 (but we didn't really reach a consensus there, I think?).

interp does work up until it calls scipy.interpolate.interp1d, which seems to strip the units?

reindex works for data, but the tests need a rewrite so they also test coords / dims with units (the same for the interp tests).

[max-sixty]

LGTM! Thanks as ever @keewis

[max-sixty]

not sure what the bot is checking, xarray/core/arithmetic.py doesn't have 1934 lines (the last line is 107) and black doesn't complain about indentation for any line

I put in an issue

[keewis]

I put in an issue

thanks, @max-sixty.

[keewis]

it seems we can't use __array_function__ to check if the type should be handled or not (using __array_function__ works, but that's by accident, I think, and stops working once/if we implement __array_function__on xarray objects).

I'd say it is pretty rare the other type knows what to do with a xarray object, so I think we should retry using the wrapped object. I don't know the reason for the current implementation, though, so I could be missing something. For reference, this only is a problem with functions that accept multiple values, like numpy.maximum. Any thoughts on this, @shoyer?

Edit: actually, that smells like a type hierarchy issue, as discussed in hgrecco/pint#845 and dask/dask#5329

[keewis]

Ugh. Trying to add pint.Quantityto the list of handled types does not work since pint tries to add xarray to its known non-handled types – that means we have a circular import. I guess I'll xfail the bivariate_ufuncs tests, too.

To resolve this, the type registration mentioned in #1617 and

xarray/xarray/core/arithmetic.py

Line 23 in ba9f822

# TODO: allow extending this with some sort of registration system

might help?

[dcherian]

pint tries to add xarray to its known non-handled types

I don't know much about this stuff can't pint have just a list of handled types and return NotImplementedfor anything else?

[keewis]

see dask/dask#5329 (comment). Basically, pint is high enough in the type hierarchy that there are too many types below it and it is easier for it to list those above it.

Actually, pint and xarray using different (conflicting) ways to check the support was the reason why I asked for advice in that issue.

[jthielen]

@keewis Is part of the specific problem that Pint has a guarded import of xarray when defining its upcast types? Would it help if it checked for the fully qualified class name instead?

[jthielen]

Also, after glancing through all this, it seems like xarray is dealing inconsistently with the type casting hierarchy:

• Construction/wrapping allowing __array_function__-defining duck arrays among other explictly listed additions (pandas.Index, dask arrays, etc., see

  xarray/xarray/core/variable.py

  Line 166 in 9b5140e

  def as_compatible_data(data, fastpath=False):

  )
• Binary ops only denying based on xarray's internal hierarchy (Dataset, DataArray, Variable), otherwise deferring to underlying data
• __array_ufunc__ allowing a list of supported types (

  xarray/xarray/core/arithmetic.py

  Lines 24 to 30 in 9b5140e

  	_HANDLED_TYPES = (
  	np.ndarray,
  	np.generic,
  	numbers.Number,
  	bytes,
  	str,
  	) + dask_array_type

  along with SupportsArithmetic)
• (__array_function__ yet to be defined)

Should this be its own issue?

[keewis]

Is part of the specific problem that Pint has a guarded import of xarray when defining its upcast types? Would it help if it checked for the fully qualified class name instead?

SupportsArithmetic is a base class of Variable, DataArray and Dataset, so if I try to import pint for type checking, pint can't use these classes, otherwise we'd have a circular import. Comparing fully qualified class names would work, but it might make sense to wait for a recommendation regarding my question.

Should this be its own issue?

Definitely, even if it's just to clarify why we accept different types in those functions: I think __init__ may convert its arguments so it supports more types than __array_ufunc__ (I don't know about the others, though). For the missing __array_function__ see #3917.

[dcherian]

IMO @keewis you should merge this if you think its ready.

even if it's just to clarify why we accept different types in those functions: I think init may convert its arguments so it supports more types than array_ufunc (I don't know about the others, though).

open an issue for this?

[keewis]

see #3950

Merging: sure, I just wanted to avoid merging my own PRs without a final review