One variable in xr.DataArray has an index depending on another variable, how to correctly implement this?

Question

I am trying to get started using Xarray, but having an issue with a specific task: the labels I want to use in one dimension are dependent on another dimension.

Specifically, in the example below, the labels/indices which should be used in the "ti" dimension are dependent on the category in the "regime" dimension.

import numpy as np
import xarray as xr

wd_index = np.arange(40,85,3)
ws_index = np.arange(0,100,5)
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,20)

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

check = array.sel({'regime':'free','ws':5,'ti':ti_indices['free'][0]})

This last line returns:

   raise KeyError(f"no index found for coordinate {key!r}")
KeyError: "no index found for coordinate 'ti'"

But I haven't yet figured out how to correctly assign the index. I also tried:

xr.DataArray(values, dims=['regime','ti','wd','ws'],
                                     coords={'regime':['free','restricted','closed'],
                                         'ti':(['regime', 'ti'],{'free':ti_indices['free'],
                                                                 'restricted':ti_indices['restricted'],
                                                                 'closed':ti_indices['closed']}),
                                         'wd':wd_index,
                                         'ws':ws_index})

but this is apparently not a correct syntax, and returns:

ValueError: Variable 'ti': Could not convert tuple of form (dims, data[, attrs, encoding]): (['regime', 'ti'], {'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]}) to Variable.

I'm wondering if perhaps this is not possible, I read the following in the docs: "The distinction between whether a variable falls in data or coordinates (borrowed from CF conventions) is mostly semantic, and you can probably get away with ignoring it if you like: dictionary like access on a dataset will supply variables found in either category. However, xarray does make use of the distinction for indexing and computations. Coordinates indicate constant/fixed/independent quantities, unlike the varying/measured/dependent quantities that belong in data." — user13132640
– user13132640, Commented Apr 23 at 15:03

Divyansh Gupta · Accepted Answer · 2025-05-28 15:23:41Z

I found a not-so-elegant solution which would be to add each of the three ti vectors as separate (non-indexed) coordinate to the same ti dimension. It won't scale so well beyond a handful of regimes though.

import numpy as np
import xarray as xr

wd_index = np.arange(40, 85, 3)
ws_index = np.arange(0, 100, 5)
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, 20)

array = xr.DataArray(
    values,
    dims=["regime", "ti", "wd", "ws"],
    coords={
        "regime": ["free", "restricted", "closed"],
        "ti_free": ("ti", ti_indices["free"]),
        "ti_restricted": ("ti", ti_indices["restricted"]),
        "ti_closed": ("ti", ti_indices["closed"]),
        "wd": wd_index,
        "ws": ws_index,
    },
)

check = array.sel(
    {"regime": "free", "ws": 5, "ti": array.ti_free == ti_indices["free"][0]}
)

The trick to select from a non-indexed coord comes from this answer

Collectives™ on Stack Overflow

One variable in xr.DataArray has an index depending on another variable, how to correctly implement this?

1 Answer 1

Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related