xarray icon indicating copy to clipboard operation
xarray copied to clipboard

Should Xarray stop doing automatic index-based alignment?

Open shoyer opened this issue 1 year ago • 8 comments

What is your issue?

I am increasingly thinking that automatic index-based alignment in Xarray (copied from pandas) may have been a design mistake. Almost every time I work with datasets with different indexes, I find myself writing code to explicitly align them:

  1. Automatic alignment is hard to predict. The implementation is complicated, and the exact mode of automatic alignment (outer vs inner vs left join) depends on the specific operation. It's also no longer possible to predict the shape (or even the dtype) resulting from most Xarray operations purely from input shape/dtype.
  2. Automatic alignment brings unexpected performance penalty. In some domains (analytics) this is OK, but in others (e.g,. numerical modeling or deep learning) this is a complete deal-breaker.
  3. Automatic alignment is not useful for float indexes, because exact matches are rare. In practice, this makes it less useful in Xarray's usual domains than it for pandas.

Would it be insane to consider changing Xarray's behavior to stop doing automatic alignment? I imagine we could roll this out slowly, first with warnings and then with an option for disabling it.

If you think this is a good or bad idea, consider responding to this issue with a 👍 or 👎 reaction.

shoyer avatar Sep 16 '22 15:09 shoyer

I think I agree here but a lot of things are going to break.

IMO we could first align (hah) these choices to be the same:

the exact mode of automatic alignment (outer vs inner vs left join) depends on the specific operation.

so that they're all controlled by OPTIONS["arithmetic_join"] (rename to "default_join"?) and then change the default after a long period of warnings.

Automatic alignment is not useful for float indexes, because exact matches are rare. In practice, this makes it less useful in Xarray's usual domains than it for pandas.

What do you think of making the default FloatIndex use a reasonable (hard to define!) rtol for comparisons?

dcherian avatar Sep 16 '22 16:09 dcherian

IMO we could first align (hah) these choices to be the same:

the exact mode of automatic alignment (outer vs inner vs left join) depends on the specific operation.

The problem is that user expectations are actually rather different for different options:

  • With data movement operations like xarray.merge, you expect to keep around all existing data -- so you want an outer join.
  • With inplace operations that modify an existing Dataset, e.g., by adding new variables, you don't expect the existing coordinates to change -- so you want a left join.
  • With computate based operations (like arithmatic), you don't have an expectation that all existing data is unmodified, so keeping around a bunch of NaN values felt very wasteful -- hence the inner join.

What do you think of making the default FloatIndex use a reasonable (hard to define!) rtol for comparisons?

This would definitely be a step forward! However, it's a tricky nut to crack. We would both need a heuristic for defining rtol (some fraction of coordinate spacing?) and a method for deciding what the resulting coordinates should be (use values from the first object?).

Even then, automatic alignment is often problematic, e.g., imagine cases where a coordinate is defined in separate units.

shoyer avatar Sep 16 '22 17:09 shoyer

@shoyer could you maybe provide a code example of the current index aligned behaviour and a future not index aligned behaviour?

I am a bit worried about transitioning previous code bases to such new xarray releases

aaronspring avatar Sep 16 '22 20:09 aaronspring

As a concrete example, suppose we have two datasets:

  1. Hourly predictions for 10 days
  2. Daily observations for a month.
import numpy as np
import pandas as pd
import xarray

predictions = xarray.DataArray(
    np.random.RandomState(0).randn(24*10),
    {'time': pd.date_range('2022-01-01', '2022-01-11', freq='1h', closed='left')},
)
observations = xarray.DataArray(
    np.random.RandomState(1).randn(31),
    {'time': pd.date_range('2022-01-01', '2022-01-31', freq='24h')},
)

Today, if you compare these datasets, they automatically align:

>>> predictions - observations
<xarray.DataArray (time: 10)>
array([ 0.13970698,  2.88151104, -1.0857261 ,  2.21236931, -0.85490761,
        2.67796423,  0.63833301,  1.94923669, -0.35832191,  0.23234996])
Coordinates:
  * time     (time) datetime64[ns] 2022-01-01 2022-01-02 ... 2022-01-10

With this proposed change, you would get an error, e.g., something like:

>>> predictions - observations
ValueError: xarray objects are not aligned along dimension 'time':  
array(['2022-01-01T00:00:00.000000000', '2022-01-02T00:00:00.000000000',
       '2022-01-03T00:00:00.000000000', '2022-01-04T00:00:00.000000000',
       '2022-01-05T00:00:00.000000000', '2022-01-06T00:00:00.000000000',
       '2022-01-07T00:00:00.000000000', '2022-01-08T00:00:00.000000000',
       '2022-01-09T00:00:00.000000000', '2022-01-10T00:00:00.000000000',
       '2022-01-11T00:00:00.000000000', '2022-01-12T00:00:00.000000000',
       '2022-01-13T00:00:00.000000000', '2022-01-14T00:00:00.000000000',
       '2022-01-15T00:00:00.000000000', '2022-01-16T00:00:00.000000000',
       '2022-01-17T00:00:00.000000000', '2022-01-18T00:00:00.000000000',
       '2022-01-19T00:00:00.000000000', '2022-01-20T00:00:00.000000000',
       '2022-01-21T00:00:00.000000000', '2022-01-22T00:00:00.000000000',
       '2022-01-23T00:00:00.000000000', '2022-01-24T00:00:00.000000000',
       '2022-01-25T00:00:00.000000000', '2022-01-26T00:00:00.000000000',
       '2022-01-27T00:00:00.000000000', '2022-01-28T00:00:00.000000000',
       '2022-01-29T00:00:00.000000000', '2022-01-30T00:00:00.000000000',
       '2022-01-31T00:00:00.000000000'], dtype='datetime64[ns]')
vs
array(['2022-01-01T00:00:00.000000000', '2022-01-01T01:00:00.000000000',
       '2022-01-01T02:00:00.000000000', ..., '2022-01-10T21:00:00.000000000',
       '2022-01-10T22:00:00.000000000', '2022-01-10T23:00:00.000000000'],
      dtype='datetime64[ns]')

Instead, you would need to manually align these objects, e.g., with xarray.align, reindex_like() or interp_like(), e.g.,

>>> predictions, observations = xarray.align(predictions, observations)

or

>>> observations = observations.reindex_like(predictions)

or

>>> predictions = predictions.interp_like(observations)

To (partially) simulate the effect of this change on a codebase today, you could write xarray.set_options(arithmetic_join='exact') -- but presmably it would also make sense to change Xarray's other alignment code (e.g., in concat and merge).

shoyer avatar Sep 16 '22 22:09 shoyer

I think I really empathize with the pain here. There's a very real explicitness vs "helpfulness" tradeoff, often depending on whether people are doing exploratory research vs hardened production (a bit like Ask vs Guess culture!).

But from the perspective of someone who works with lots of people who use Xarray for their daily research, I think this would be a big hurdle, even without considering the change costs.

One analogy is xarray vs. pandas for 2D data — among my colleagues xarray is known to be a smaller, more reliable API surface, while pandas is more fully featured but also a maze of surprising methods and behavior (df['a'] * df!). Forcing explicit alignment would strengthen that case. But it could take it too far — operations that you expect to just work would now need nannying.

"Make another mode" can seem like an easy decision — "who doesn't want another mode" — but it could make development more difficult, since we'd need calls to check which mode we're in & tests for those. It's not insurmountable though, and maybe it would only be required in a couple of methods, so testing those would be sufficient to ensure the resulting behavior would be correct?

(FWIW we don't use float indexes, so it could be fine to dispense with those)

max-sixty avatar Sep 16 '22 23:09 max-sixty

I still find myself struggling to understand which of those options are needed for my use cases (inner, outer etc.). Default is working in many cases, but in other cases it is trial and error.

In that sense this proposal would make me have to really understand what's going on.

The suggestion of another mode by @max-sixty just made me think, if this automatic alignment machinery could be moved to another package. If that package is installed the current behaviour is preserved, if not then the new behaviour proposed by @shoyer comes into play.

kmuehlbauer avatar Sep 17 '22 05:09 kmuehlbauer

This suggestion looks roughly like what we are discussing in https://github.com/pydata/xarray/discussions/7041#discussioncomment-3662179, i.e., using a custom index that avoids this? So maybe the question here is whether such an ArrayIndex should be the default?

Aside from that, with my outside perspective (having used Xarray extremely little, looking at the docs and code occasionally, but developing a similar library that does not have indexes):

Indexes (including alignment behavior) feel like a massive complication of Xarray, both conceptually (which includes documentation and teaching efforts) as well as code. If all you require is the ArrayIndex behavior (i.e., exact coord comparison in operations) then the entire concept of indexes is just ballast, distraction in the documentation, and confusion. Example: Why can't we use loc/sel with a non-dimension (non-index) coord? --- without index we would just search the coord with no need to limit this to index-coords, and this is often fast enough?

SimonHeybrock avatar Sep 20 '22 04:09 SimonHeybrock

So maybe the question here is whether such an ArrayIndex should be the default?

Another solution for more flexibility or a smooth transition may be to add a build option to the Index base class API, so that it would be possible for the current default PandasIndex or any custom index to easily (and explicitly) deactivate automatic alignment while keeping it around for label-based selection.

Indexes (including alignment behavior) feel like a massive complication of Xarray, both conceptually (which includes documentation and teaching efforts) as well as code.

I agree, although this is getting addressed slowly but surely. In Xarray internals, most of the indexes logic is now in the core.indexes module. For the public API #4366, #6849 and #6971 will ultimately make things better. Object reprs are important too (#6795). There is still a good amount of work in order to improve the documentation, some of it is discussed in #6975.

IMO nearly all the complication and confusion emerge from the mixed concept of a dimension coordinate in the Xarray data model. Once the concept of an index is clearly decoupled from the concept of a coordinate and both concepts are represented as 1st-class citizens, it will help users focusing on the parts of the API and/or documentation that are relevant to their needs. It will also help "selling" Xarray to users who don't need much of the index capabilities (this has been discussed several times, either as external feedback or between Xarray devs, e.g., proposal of a "xarray-lite" package). Finally it will make more affordable major changes such as the one proposed here by @shoyer.

benbovy avatar Sep 20 '22 07:09 benbovy