How to best structure data where one coordinate's values vary with another dimension?

[mr483]

I have created a dummy dataset below, in which I would like the "ti" dimension to have coordinates which depend on the selected "regime". It seems the way I define the DataArray actually accomplishes this (as evidenced when I print the ti values after selecting "free" regime, I notice that it returns only those specific to this regime. However, the line which is commented out fails due to no index being assigned to "ti". Is there a way I can set up this index correctly so that I could select a regime and ti value at once?

I could imagine that maybe this isn't possible (because it begs the question: which index should we use when you don't do a selection specific to regime?).

I thought also to perhaps create separate data arrays and combine into a dataset, but this has undesired behavior for my application, because when I do a selection that is specific to a regime, I notice that all ti values across all regimes are present. This is undesired, because in the actual real-world application for example:

• I will see in another dataset a value of 0.043, under conditions corresponding to the "free" regime
• I want to find in my data what the nearest ti value is in the dataset, specific to the "free" regime.
• For other variables I have no problem because the index values remain consistent across the regimes.

Am I basically stuck with three different DataArrays, one for each regime? Or can someone help me fix this indexing problem?

import numpy as np
import xarray as xr

regimes = ['free','restricted','closed']
wd_index = np.arange(40,85,3)
ti_indices = {'free':[0.013651005,0.02097709,0.028776308,0.041105852,0.083883279],
             'restricted':[0.0243316245,0.037467915,0.04940437,0.055962806,0.065437968],
             'closed':[0.042476781,0.055642339,0.064246267,0.078399973,0.120953288]}

values = np.random.rand(3,5,15)

all_one_array = xr.DataArray(values, dims=['regime','ti','wd'],
                                     coords={'regime':regimes,
                                         'ti':(['regime', 'ti'],np.array([ti_indices['free'],ti_indices['restricted'],ti_indices['closed']],dtype=object)),#ti_indices['closed']],dtype=object)),
                                         'wd':wd_index})

# desired behavior:
# all_one_array.sel({'regime':'free',
#                    'ti':0.013651005})

# But the following behavior is good!
print(all_one_array.sel({'regime':'free'}).ti)

# Another approach might be dataset combining three arrays
free_array = xr.DataArray(values[0,:,:],
                          dims=['ti','wd'],
                          coords={'ti':ti_indices['free'],
                                  'wd':wd_index})

restricted_array = xr.DataArray(values[1,:,:],
                          dims=['ti','wd'],
                          coords={'ti':ti_indices['restricted'],
                                  'wd':wd_index})

closed_array = xr.DataArray(values[2,:,:],
                          dims=['ti','wd'],
                          coords={'ti':ti_indices['closed'],
                                  'wd':wd_index})

all_one_dataset = xr.Dataset({'free':free_array,
                       'restricted':restricted_array,
                       'closed':closed_array})

# but the following behavior is not desired:
print(all_one_array.sel({'regime':'free'}).ti)

[benbovy]

If you don't mind a couple of extra steps, one simple way to achieve your desired behavior is:

(
    all_one_array
    .sel(regime="free")
    .set_xindex("ti")
    .sel(ti=0.013, method="nearest")
)

The .set_xindex() step allows to build a (pandas) index from the ti coordinate once it has been reduced to a 1-dimension coordinate, which in turn allows using it with .sel().

---

Alternatively, it is possible to create a custom Xarray Index to be associated directly with both the regime and ti coordinates:

(
    all_one_array
    .set_xindex(["regime", "ti"], CustomIndexSubclass)
    .sel(regime="free", ti=0.013)
)

Where the regime and ti label values passed to .sel() will be handled together by the logic implemented in the CustomIndexSubclass.sel() method.

Unfortunately Xarray custom indexes are not much documented yet, apart from a rudimentary guide and some examples gathered here or scattered elsewhere. Hopefully this will improve in the near future.