Add utility function to efficiently unstack xarrays

docs/index.rst

@@ -83,6 +83,7 @@ Documentation contents
    reading_files
    agipd_lpd_data
    streaming
+   misc_api
    validation
    cli
    data_format
@@ -105,4 +106,3 @@ Indices and tables
 
 * :ref:`genindex`
 * :ref:`search`
-

docs/misc_api.rst

@@ -0,0 +1,9 @@
+Other functions
+===============
+
+.. module:: extra_data.utils
+
+There is one function to help working with `xarray <https://docs.xarray.dev/en/stable/index.html>`_
+DataArrays:
+
+.. autofunction:: unstack_regular

extra_data/components.py

@@ -13,7 +13,7 @@
 from .exceptions import SourceNameError
 from .reader import DataCollection, by_id, by_index
 from .read_machinery import DataChunk, roi_shape, split_trains
-from .utils import default_num_threads
+from .utils import default_num_threads, unstack_regular
 from .writer import FileWriter
 from .write_cxi import XtdfCXIWriter, JUNGFRAUCXIWriter
 
@@ -1342,10 +1342,7 @@ def xarray(self, *, pulses=None, fill_value=None, roi=(), astype=None,
         out = self._wrap_xarray(arr, subtrain_index)
 
         if unstack_pulses:
-            # Separate train & pulse dimensions, and arrange dimensions
-            # so that the data is contiguous in memory.
-            dim_order = ['module'] + out.indexes['train_pulse'].names + self.dimensions[2:]
-            return out.unstack('train_pulse').transpose(*dim_order)
+            return unstack_regular(out, 'train_pulse', fallback=True)
 
         return out
 
@@ -1600,8 +1597,7 @@ def _get_pulse_data(self, source, key, tid):
 
         # Separate train & pulse dimensions, and arrange dimensions
         # so that the data is contiguous in memory.
-        dim_order = train_pulse_ids.names + dims[1:]
-        return arr.unstack('train_pulse').transpose(*dim_order)
+        return unstack_regular(arr, 'train_pulse', fallback=True)
 
     def _select_pulse_ids(self, pulse_ids):
         """Select pulses by ID

extra_data/tests/test_utils.py

@@ -0,0 +1,33 @@
+import numpy as np
+import pandas as pd
+import pytest
+import xarray as xr
+
+from extra_data.utils import unstack_regular
+
+
+def test_unstack_regular():
+    stacked = xr.DataArray(
+        np.arange(4*5*6*2).reshape((1, 4*5*6, 2)),
+        dims=('a', 'combined', 'e'),
+        coords={'combined': pd.MultiIndex.from_product(
+            [range(4), range(5), range(6)], names=['b', 'c', 'd']
+        ), 'e': ['x', 'y']}
+    )
+
+    res = unstack_regular(stacked, 'combined', fallback=False)
+    assert res.shape == (1, 4, 5, 6, 2)
+    assert res.dims == ('a', 'b', 'c', 'd', 'e')
+    xr.testing.assert_equal(
+        res, stacked.unstack('combined').transpose('a', 'b', 'c', 'd', 'e')
+    )
+
+    with pytest.raises(ValueError, match="cartesian"):
+        # Not a cartesian product
+        unstack_regular(stacked[:, :-1], 'combined', fallback=False)
+
+    # Test the fallback to arr.unstack()
+    fallback_res = unstack_regular(stacked[:, :-1], fallback=True)
+    assert fallback_res.shape == (1, 4, 5, 6, 2)
+    assert fallback_res.dims == ('a', 'b', 'c', 'd', 'e')
+    assert fallback_res.dtype.kind == 'f'

extra_data/utils.py

@@ -11,6 +11,11 @@
 import os
 import sys
 from shutil import get_terminal_size
+from typing import TYPE_CHECKING
+
+import numpy as np
+if TYPE_CHECKING:
+    import pandas
 
 
 def available_cpu_cores():
@@ -49,3 +54,115 @@ def isinstance_no_import(obj, mod: str, cls: str):
         return False
 
     return isinstance(obj, getattr(m, cls))
+
+
+def _multiindex_regular_labels(mix: "pandas.MultiIndex"):
+    """Return a tuple of indexes if mix is a cartesian product, else None"""
+    import pandas as pd
+    if mix.has_duplicates:
+        return None
+
+    if len(mix) == 0:  # Empty dimension counts as regular
+        return tuple([mix.get_level_values(i) for i in range(mix.nlevels)])
+
+    k1_sel, k1_subix = mix.get_loc_level(mix[0][0])
+    rpt_len = len(k1_subix)
+    rpt, rem = divmod(len(mix), rpt_len)
+    if rem != 0:
+        return None
+
+    if isinstance(k1_subix, pd.MultiIndex):
+        inner_labels = _multiindex_regular_labels(k1_subix)
+        if inner_labels is None:
+            return None
+    else:
+        inner_labels = (k1_subix,)
+
+    # Check that the outermost level has each value rpt_len times
+    outer_codes = mix.codes[0].reshape(-1, rpt_len)
+    if (outer_codes != outer_codes[:, 0, np.newaxis]).any():
+        return None
+
+    # Check that each inner level is repeating regularly
+    for level in range(1, mix.nlevels):
+        codes = mix.codes[level].reshape(-1, rpt_len)
+        if (codes != codes[0]).any():
+            return None
+
+    outer_labels = mix.levels[0][outer_codes[:, 0]]
+    return (outer_labels,) + inner_labels
+
+
+def _unstack_regular_once(arr, dim, fallback=False, fill_value=None):
+    import pandas as pd
+    import xarray as xr
+
+    mix = arr.indexes[dim]
+    assert isinstance(mix, pd.MultiIndex)
+    if (mix_labels := _multiindex_regular_labels(mix)) is None:
+        # Not a cartesian product -> cannot reshape
+        if fallback:
+            return _unstack_fallback(arr, dim, fill_value)
+        raise ValueError(f"MultiIndex for {dim!r} is not a cartesian product")
+
+    mix_shape = tuple(len(l) for l in mix_labels)
+    dim_ix = arr.dims.index(dim)
+    new_shape = arr.shape[:dim_ix] + mix_shape + arr.shape[dim_ix + 1:]
+    data = arr.values.reshape(new_shape)
+
+    coords = {
+        k: v for (k, v) in arr.coords.items() if dim not in v.dims  # Unchanged
+    } | dict(
+        zip(mix.names, mix_labels)  # Unstacked coordinates
+    ) | {
+        # Other coordinates along unstacked dimension
+        k: (mix.names, v.values.reshape(mix_shape)) for (k, v) in arr.coords.items()
+        if (dim in v.dims and k != dim and k not in mix.names)
+    }
+
+    return xr.DataArray(
+        data,
+        dims=arr.dims[:dim_ix] + mix.names + arr.dims[dim_ix + 1:],
+        coords=coords,
+    )
+
+
+def _unstack_fallback(arr, dim, fill_value=None):
+    from xarray.core.dtypes import NA
+    if fill_value is None:
+        fill_value = NA
+
+    res = arr.unstack(dim, fill_value=fill_value)
+
+    # Restore the obvious axis order
+    dim_ix = arr.dims.index(dim)
+    new_dims = res.dims[arr.ndim - 1:]
+    dim_order = arr.dims[:dim_ix] + new_dims + arr.dims[dim_ix + 1:]
+    return res.transpose(*dim_order)
+
+
+def unstack_regular(arr, dim=None, *, fallback=False, fill_value=None):
+    """Unstack an xarray.DataArray efficiently when no fill values are needed.
+
+    Where the stacked index is a full cartesian product, we can make a view of
+    the original data instead of copying it, which is much more efficient. In
+    this case, we also don't have to convert integers to floats to allow for
+    NaN values.
+
+    If ``fallback=True``, this also accepts arrays where fill values are needed,
+    and uses xarray's implementation. Otherwise, it raises ValueError if
+    unstacking would require inserting fill values.
+
+    The unstacked dimensions are expanded in-place in the dimension order,
+    rather than being moved to the end.
+    """
+    import pandas as pd
+
+    if dim is None:
+        dim = [d for d in arr.dims if isinstance(arr.indexes.get(d), pd.MultiIndex)]
+    if isinstance(dim, str):
+        dim = [dim]
+    for d in dim:
+        arr = _unstack_regular_once(arr, d, fallback, fill_value)
+
+    return arr