Impossible to create a Dataset with varying DataArrays dimensions dependence?

[doronbehar]

One of the greatest features of xarray is the ability to create a complex Dataset with some data variables depending on some dimensions and some data variables depending on other dimensions. Am I correct to understand this is not supported by scipp.Dataset? Here's an example written by a LLM:

import scipp as sc
import numpy as np

# Create coordinates for 4 dimensions
x_coord = sc.array(dims=['x'], values=np.linspace(0, 10, 11))  # 11 points
y_coord = sc.array(dims=['y'], values=np.linspace(0, 5, 6))    # 6 points  
z_coord = sc.array(dims=['z'], values=np.linspace(0, 2, 3))    # 3 points
time_coord = sc.array(dims=['time'], values=np.arange(0, 100, 10), unit='s')  # 10 time points

# Create first data variable that depends on 3 dimensions (x, y, z)
# Shape: (11, 6, 3)
temperature_data = sc.array(
    dims=['x', 'y', 'z'], 
    values=np.random.normal(20, 5, (11, 6, 3)),  # Random temperature data around 20°C
    unit='K'
)

# Create second data variable that depends on all 4 dimensions (x, y, z, time)
# Shape: (11, 6, 3, 10)
pressure_data = sc.array(
    dims=['x', 'y', 'z', 'time'],
    values=np.random.exponential(101325, (11, 6, 3, 10)),  # Random pressure data around 1 atm
    unit='Pa'
)

# Create the Dataset
dataset = sc.Dataset(
    coords={
        'x': x_coord,
        'y': y_coord, 
        'z': z_coord,
        'time': time_coord
    },
    data={
        'temperature': temperature_data,  # depends on x, y, z
        'pressure': pressure_data         # depends on x, y, z, time
    }
)

Where the following xarray equivalent works:

import xarray as xr
import numpy as np

# Create coordinates
x = np.linspace(0, 10, 11)               # 11 points
y = np.linspace(0, 5, 6)                 # 6 points
z = np.linspace(0, 2, 3)                 # 3 points
time_coord = np.arange(0, 100, 10)       # 10 time points

# Create temperature data (x, y, z)
temperature = xr.DataArray(
    data=np.random.normal(20, 5, (11, 6, 3)),
    dims=('x', 'y', 'z'),
    coords={'x': x, 'y': y, 'z': z},
    attrs={'units': 'K'}
)

# Create pressure data (x, y, z, time)
pressure = xr.DataArray(
    data=np.random.exponential(101325, (11, 6, 3, 10)),
    dims=('x', 'y', 'z', 'time'),
    coords={'x': x, 'y': y, 'z': z, 'time': time_coord},
    attrs={'units': 'Pa'}
)

# Combine into a Dataset
ds = xr.Dataset(
    data_vars={
        'temperature': temperature,
        'pressure': pressure
    }
)

Is this limitation by design? Or am I missing something?

[SimonHeybrock]

Hi @doronbehar!
Yes you are correct. This was in fact supported in earlier versions of Scipp, but as explained in ADR17 it was removed due to conceptual problems.

I would recommend to have a look at scipp.DataGroup to see if it fulfills you needs. It is conceptually much simpler, but also way more flexible, in particular since it supports nesting/tree-like structures. The closest equivalent would be xarray.DataTree, but our DataGroup is simpler.

[SimonHeybrock]

Note you DataGroup does support positional as well as label-based indexing, and this works even if the axis lengths and coords do not match between items:

import numpy as np
import scipp as sc


x1 = sc.linspace('angle', -1.0, 1.0, 50, unit='rad')
y1 = sc.sin(x1)
x2 = sc.linspace('angle', -1.0, 2.0, 70, unit='rad')
y2 = sc.cos(x2)
da1 = sc.DataArray(y1, coords={'angle': x1})
da2 = sc.DataArray(y2, coords={'angle': x2})

dg = sc.DataGroup(a=da1, b=da2, c=np.arange(100))
dg['angle', 30:70]  # positional, works with items without coords
dg['angle', sc.scalar(0.5, unit='rad'):]  # label-based, all items must have an `angle` coord

The latter (passing a 0-D scipp.Variable as start or stop in a slice) is the equivalent of sel in Xarray.

[doronbehar]

The latter (passing a 0-D scipp.Variable as start or stop in a slice) is the equivalent of sel in Xarray.

Hmm that's still not satisfying - .sel has the ability to select a single value of the array, as well as multiple values via a k: [ v1, v2] like dictionary arguments. It is true that I can get a single value with e.g:

dg['angle', 4]

But I'd have to know the index, and be sure that all DataArrays in the DataGroup use the same order?

[nvaytet]

dg['angle', sc.scalar(1.0, unit='rad')] should also work, selecting a single value.

Note that in this case, we require an exact match, so angle coordinates must contain a value = 1rad.
If your angle coordinate is bin edges, it is a bit friendlier in the sense that it will select the bin that contains the requested value.

[doronbehar]

Hmm OK, Thanks! the error message I got when I used sc.scalar(0.5, unit='rad') was confusing:

Traceback (most recent call last):
  File "/home/doron/repos/lab-ion-trap-simulations/./t.py", line 16, in <module>
    selected = dg['angle', sc.scalar(0.5, unit='rad')]  # label-based, all items must have an `angle` coord
               ~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nix/store/3ynldpj6dqq0h7g4dirl7hhg42nxc191-python3-3.12.10-env/lib/python3.12/site-packages/scipp/core/data_group.py", line 182, in __getitem__
    key: var[dim, index]  # type: ignore[index]
         ~~~^^^^^^^^^^^^
IndexError: Coord angle does not contain unique point with value <scipp.Variable> ()    float64            [rad]  0.5

Also seeing the index variable there in that traceback suggests that only an index is allowed.. Note also that xarray's sel method has an optional argument method='nearest', which is mentioned in error messages raised when a value that wasn't found exactly is used.

[nvaytet]

I guess it's always possible to create your own little function that finds the nearest point and returns that?
This has not been requested by anyone so far, so we did not implement it ;-)