From 841ecdd8f69df391a1c5955d8b3e4610d48c5e33 Mon Sep 17 00:00:00 2001 From: Thomas Kluyver Date: Mon, 27 Jul 2026 17:50:32 +0100 Subject: [PATCH 1/4] Add utility function to efficiently unstack xarrays --- extra_data/components.py | 10 ++-- extra_data/utils.py | 114 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 117 insertions(+), 7 deletions(-) diff --git a/extra_data/components.py b/extra_data/components.py index 4d3d56e9..72c67446 100644 --- a/extra_data/components.py +++ b/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 diff --git a/extra_data/utils.py b/extra_data/utils.py index 376704f0..81aa1c12 100644 --- a/extra_data/utils.py +++ b/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,112 @@ 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 + + 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 From b7b64c865dcd07f67cc13b1b89f5651e68584542 Mon Sep 17 00:00:00 2001 From: Thomas Kluyver Date: Tue, 28 Jul 2026 08:31:33 +0100 Subject: [PATCH 2/4] Handle empty dimension in unstack_regular() --- extra_data/utils.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/extra_data/utils.py b/extra_data/utils.py index 81aa1c12..4f27eaab 100644 --- a/extra_data/utils.py +++ b/extra_data/utils.py @@ -62,6 +62,9 @@ def _multiindex_regular_labels(mix: "pandas.MultiIndex"): 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) From bf5dbf2a81e01898c079cfd0b8bf14257ea8798f Mon Sep 17 00:00:00 2001 From: Thomas Kluyver Date: Tue, 28 Jul 2026 09:03:02 +0100 Subject: [PATCH 3/4] Add dedicated test for unstack_regular() --- extra_data/tests/test_utils.py | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) create mode 100644 extra_data/tests/test_utils.py diff --git a/extra_data/tests/test_utils.py b/extra_data/tests/test_utils.py new file mode 100644 index 00000000..1dfb2598 --- /dev/null +++ b/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' From 7f6c971df488f7a3a567dc99385d86ad17e0e43a Mon Sep 17 00:00:00 2001 From: Thomas Kluyver Date: Tue, 28 Jul 2026 09:20:13 +0100 Subject: [PATCH 4/4] Document unstack_regular() --- docs/index.rst | 2 +- docs/misc_api.rst | 9 +++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) create mode 100644 docs/misc_api.rst diff --git a/docs/index.rst b/docs/index.rst index 98245103..20b2d807 100644 --- a/docs/index.rst +++ b/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` - diff --git a/docs/misc_api.rst b/docs/misc_api.rst new file mode 100644 index 00000000..53f2b722 --- /dev/null +++ b/docs/misc_api.rst @@ -0,0 +1,9 @@ +Other functions +=============== + +.. module:: extra_data.utils + +There is one function to help working with `xarray `_ +DataArrays: + +.. autofunction:: unstack_regular