Reindex on a coordinate dependent on two dimensions

[maresb]

Overview

Is there a clean way in Xarray to reindex or a DataArray based on a 2D coordinate like local_time below, without resorting to manual index manipulation, potentially leveraging Xarray's built-in functions?

I provide code for mock input data and show my desired result after outlining my motivation.

Motivation

I'm fairly new to Xarray, and I feel like the following problem should have an elegant solution, but unfortunately I'm not finding it, even after spending significant time thinking and experimenting. I hope there is a simple answer that I just haven't discovered yet.

I'm working with some gridded climatological data; essentially a 3D DataArray of temperature values with dimensions lat, lon, and utc_time. In terms of day/night, it would be very convenient to switch the dimension from utc_time to local_time so that I could easily perform aggregations like max/min temperature for the local-time day.

Concrete example: input/desired output

For my purposes, I can use the approximation

local_time = utc_time + round(lon / 15) * HOUR

Note that local_time depends not just on the utc_time dimension, but also on lon, and this is the source of the complications that I'm really struggling to overcome. The reindexing machinery I'd like to use seems to assume 1-d coordinates.

Input

For simplicity I'm going to ignore lat, but eventually I need to broadcast along that dimension.

import xarray as xr
import numpy as np
import pandas as pd

utc_time = pd.to_datetime(["2020-01-01 12:00", "2020-01-01 15:00"])
lon = [-45, 0, 45]

da = xr.DataArray(
    273 + np.arange(6).reshape(2, 3),
    coords=[("utc_time", utc_time), ("lon", lon)],
    name="temperature",
)
da = da.assign_coords(local_time=da.utc_time + pd.Timedelta(hours=1) * da.lon / 15)
"""
<xarray.DataArray 'temperature' (utc_time: 2, lon: 3)>
array([[273, 274, 275],
       [276, 277, 278]])
Coordinates:
  * utc_time    (utc_time) datetime64[ns] 2020-01-01T12:00:00 2020-01-01T15:0...
  * lon         (lon) int64 -45 0 45
    local_time  (utc_time, lon) datetime64[ns] 2020-01-01T09:00:00 ... 2020-0...
"""

Desired output:

"""
<xarray.DataArray 'temperature' (local_time: 4, lon: 3)>
array([[273.,  nan,  nan],
       [276., 274.,  nan],
       [ nan, 277., 275.],
       [ nan,  nan, 278.]])
Coordinates:
  * local_time  (local_time) datetime64[ns] 2020-01-01T09:00:00 ... 2020-01-0...
  * lon         (lon) int64 -45 0 45
    utc_time    (local_time, lon) datetime64[ns] 2020-01-01T12:00:00 ... 2020...
"""

Note how the values are shifted vertically in the time dimension according to an offset determined by longitude dimension.

My question is not about how to generate the NumPy array of the result. It's about if I can use Xarray's coordinate machinery to generate the above result in a relatively declarative manner, without performing error-prone and confusing index manipulations by hand.

---

OPTIONAL: The approaches I've already attempted

I've actually gone substantially further in the following approaches than shown here. Apologies if this isn't so coherent, since it's tricky to judge what's useful or not from a failed approach.

groupby()

One way to approach this would be to simply group by ["local_time", "lon"], but unfortunately .groupby() simply doesn't work with multiple indices. (I've also attempted without success doing things with multiindices, similar to my "Multiindices" section below.

set_xindex()

I think this would be my ideal solution; something like:

res = da.set_xindex(["local_time", "lon"])
"""ValueError: those coordinates already have an index: {'lon'}"""

res = da.set_xindex(["local_time"])
"""ValueError: PandasIndex only accepts a 1-dimensional variable, variable 'local_time' has 2 dimensions"""

Multiindices

Maybe I could use set_xindex or groupby() after stacking to get a 1-d multiindex:

stacked = da.stack(utc_time_lon=("utc_time", "lon"))
print(stacked.utc_time_lon)
"""
<xarray.DataArray 'utc_time_lon' (utc_time_lon: 6)>
array([(Timestamp('2020-01-01 12:00:00'), -45),
       (Timestamp('2020-01-01 12:00:00'), 0),
       (Timestamp('2020-01-01 12:00:00'), 45),
       (Timestamp('2020-01-01 15:00:00'), -45),
       (Timestamp('2020-01-01 15:00:00'), 0),
       (Timestamp('2020-01-01 15:00:00'), 45)], dtype=object)
Coordinates:
    local_time    (utc_time_lon) datetime64[ns] 2020-01-01T09:00:00 ... 2020-...
  * utc_time_lon  (utc_time_lon) object MultiIndex
  * utc_time      (utc_time_lon) datetime64[ns] 2020-01-01T12:00:00 ... 2020-...
  * lon           (utc_time_lon) int64 -45 0 45 -45 0 45
"""

But this seems like the wrong way around, since I need to group on the inverse multiindex.

But how do I get the new/inverse multiindex? Maybe this?

da.stack(local_time_lon=("local_time", "lon"))
"""KeyError: 'local_time'"""

I can go further with Pandas as in the .sel() section below, but this approach feels quite ugly.

sel()

I could come at it from the opposite direction, constructing local_time as a dimension, constructing utc_time as a coordinate, and using a .sel() operation to select the temperatures from the array where utc_time is a dimension.

df = da.local_time.to_dataframe().reset_index()

# The 'local_time' column is duplicated because it's both a coordinate
# and the data variable, so drop it with the trick from
# <https://stackoverflow.com/a/40435354/10155767>.
df = df.loc[:,~df.columns.duplicated()]

df_pivot = df.pivot(index='local_time', columns='lon', values='utc_time')
utc_from_local_da = xr.DataArray(df_pivot, dims=['local_time', 'lon'])
# utc_from_local_da = utc_from_local_da.fillna(np.datetime64('2020-01-01T12:00:00.000000000'))
utc_from_local_da = utc_from_local_da.assign_coords(utc_time=utc_from_local_da)

# The moment of truth:
da.sel(utc_time=utc_from_local_da.utc_time, lon=utc_from_local_da.lon)
"""KeyError: "not all values found in index 'utc_time'""""
# Or uncommenting the above fillna line simulate a hypothetically work around the KeyError:
"""
IndexError: dimension coordinate 'local_time' conflicts between indexed and indexing objects:
<xarray.DataArray 'local_time' (local_time: 4, lon: 3)>
array([['2020-01-01T09:00:00.000000000', '2020-01-01T12:00:00.000000000',
        '2020-01-01T15:00:00.000000000'],
       ['2020-01-01T12:00:00.000000000', '2020-01-01T12:00:00.000000000',
        '2020-01-01T15:00:00.000000000'],
       ['2020-01-01T09:00:00.000000000', '2020-01-01T15:00:00.000000000',
        '2020-01-01T15:00:00.000000000'],
       ['2020-01-01T09:00:00.000000000', '2020-01-01T12:00:00.000000000',
        '2020-01-01T18:00:00.000000000']], dtype='datetime64[ns]')
Coordinates:
  * utc_time    (local_time, lon) datetime64[ns] 2020-01-01T12:00:00 ... 2020...
  * lon         (lon) int64 -45 0 45
    local_time  (local_time, lon) datetime64[ns] 2020-01-01T09:00:00 ... 2020...
vs.
<xarray.IndexVariable 'local_time' (local_time: 4)>
array(['2020-01-01T09:00:00.000000000', '2020-01-01T12:00:00.000000000',
       '2020-01-01T15:00:00.000000000', '2020-01-01T18:00:00.000000000'],
      dtype='datetime64[ns]')
"""

[dcherian]

One way to approach this would be to simply group by ["local_time", "lon"], but unfortunately .groupby() simply doesn't work with multiple indices.

Your intuition is correct, you'll just need to use flox to execute though.

import flox.xarray

flox.xarray.xarray_reduce(
    da, da.local_time, da.lon, func="first", fill_value=np.nan
)

This simpler groupby works too:

import flox.xarray

flox.xarray.xarray_reduce(
    da, da.local_time, func="first", dim="utc_time", fill_value=np.nan
).transpose()

[dcherian]

@benbovy this does truly feel like a reindexing problem, do you think it could be made to work

[benbovy]

Hmm not sure how this would work using Xarray's reindexing machinery, which works along existing dimensions whereas in the example above the dimensions change from the input to the desired output (utc_time -> local_time).

If it is straightforward to generate the result with numpy, would it be possible to use xarray.apply_func() here?