diff --git a/extra_data/keydata.py b/extra_data/keydata.py index 94bcf87c..4590bf41 100644 --- a/extra_data/keydata.py +++ b/extra_data/keydata.py @@ -1,3 +1,4 @@ +from collections.abc import Sequence from warnings import warn import h5py @@ -11,6 +12,25 @@ ) +def _expand_ellipsis(ndim, indexing): + """Expand numpy indexing to one slice/index per dimension""" + # expanding Ellipsis + if any(x is Ellipsis for x in indexing): + ellipsis_idx = indexing.index(Ellipsis) + # Count non-Ellipsis + n_before = len(indexing[:ellipsis_idx]) + n_after = len(indexing[ellipsis_idx + 1:]) + n_ellipsis = max(0, ndim - n_before - n_after) + + # Replace Ellipsis with appropriate number of colons + indexing = (indexing[:ellipsis_idx] + + (slice(None),) * n_ellipsis + + indexing[ellipsis_idx + 1:]) + + # Pad with slice(None) if indexing is shorter than ndim + return indexing + (slice(None),) * (ndim - len(indexing)) + + def expand_indexing(shape, indexing): """ Expand numpy indexing into explicit coordinate arrays for each dimension. @@ -32,23 +52,7 @@ def expand_indexing(shape, indexing): if not isinstance(indexing, tuple): indexing = (indexing,) - ndim = len(shape) - - # expanding Ellipsis - if any(x is Ellipsis for x in indexing): - ellipsis_idx = indexing.index(Ellipsis) - # Count non-Ellipsis - n_before = len(indexing[:ellipsis_idx]) - n_after = len(indexing[ellipsis_idx + 1:]) - n_ellipsis = max(0, ndim - n_before - n_after) - - # Replace Ellipsis with appropriate number of colons - indexing = (indexing[:ellipsis_idx] + - (slice(None),) * n_ellipsis + - indexing[ellipsis_idx + 1:]) - - # Pad with slice(None) if indexing is shorter than ndim - indexing += (slice(None),) * (ndim - len(indexing)) + indexing = _expand_ellipsis(len(shape), indexing) # Process each index and expand coordinates result = [] @@ -104,7 +108,7 @@ class KeyData: """ def __init__( self, source, key, *, train_ids, files, section, dtype, eshape, - inc_suspect_trains=True, + edims=None, roi=(), inc_suspect_trains=True, ): self.source = source self.key = key @@ -114,6 +118,10 @@ def __init__( self.dtype = dtype self.entry_shape = eshape self.ndim = len(eshape) + 1 + if edims is None: + edims = {f"dim_{i}": None for i in range(len(eshape))} + self._entry_dims = edims + self._roi = roi self.inc_suspect_trains = inc_suspect_trains def _find_chunks(self): @@ -274,6 +282,8 @@ def _without_virtual_overview(self): section=self.section, dtype=self.dtype, eshape=self.entry_shape, + edims=self._entry_dims, + roi=self._roi, inc_suspect_trains=self.inc_suspect_trains, ) @@ -342,6 +352,8 @@ def _only_tids(self, tids, files=None): section=self.section, dtype=self.dtype, eshape=self.entry_shape, + edims=self._entry_dims, + roi=self._roi, inc_suspect_trains=self.inc_suspect_trains, ) @@ -597,39 +609,43 @@ def xarray(self, extra_dims=None, roi=(), name=None, extra_coords=None): given, it should map dimension names to coordinate arrays. If True, default coordinate arrays will be generated. """ - import xarray + import xarray # Fail before loading data if xarray is missing ndarr = self.ndarray(roi=roi) - # Train ID index - coords = {'trainId': self.train_id_coordinates()} - dims = ['trainId'] - - def _dim_name(idx): - if extra_dims is not None: - return extra_dims[idx] - else: - return f'dim_{idx}' - - # Dimension labels after the train dimension - if extra_dims is not None and isinstance(extra_coords, dict): - dims += extra_dims - coords |= extra_coords + return self._wrap_xarray( + ndarr, extra_dims=extra_dims, roi=roi, name=name, extra_coords=extra_coords + ) - elif isinstance(extra_coords, dict): - coords |= extra_coords - dims += ['dim_%d' % i for i in range(ndarr.ndim - 1)] + def _wrap_xarray(self, arr, extra_dims=None, roi=(), name=None, extra_coords=None): + import xarray - elif extra_coords or extra_dims is not None: - # add default coordinates if extra_coords is True or extra_dims given. - for idx, coord in enumerate(expand_indexing(self.entry_shape, roi)): - dim = _dim_name(idx) - coords[dim] = coord + if (extra_dims is not None) or (extra_coords is not None): + if extra_coords is None: + dim_info = extra_dims + elif (extra_dims is None) and (extra_coords is True): + # Pass a sequence of names so that with_entry_dims will generate + # integer coordinates. + dim_info = list(self._entry_dims.keys()) + elif extra_dims is None: + # Allowed to attach coordinates to a subset of existing dimensions + if missing := set(extra_coords) - set(self._entry_dims): + raise ValueError(f"No dimensions named {missing}") + dim_info = {k: extra_coords.get(k, v) + for (k, v) in self._entry_dims.items()} + else: # Both specified + dim_info = {n: extra_coords.get(n, None) for n in extra_dims} + + return self.with_entry_dims(dim_info, roi=roi)._wrap_xarray(arr, name=name) + + roi = _expand_ellipsis(len(self.entry_shape), self._roi) + dims = ['trainId'] + [ + n for n, sel in zip(self._entry_dims, roi) if not isinstance(sel, int) + ] - if not isinstance(coord, int): - dims.append(dim) - else: - dims += ['dim_%d' % i for i in range(ndarr.ndim - 1)] + coords = {'trainId': self.train_id_coordinates()} | { + k: v for (k, v) in self._entry_dims.items() if (v is not None) + } # xarray attributes attrs = {} @@ -639,11 +655,12 @@ def _dim_name(idx): except Exception as e: warn(f"Exception fetching units: {e}") - if ndarr.dtype.names is not None: + if arr.dtype.names is not None: # Structured dtype. return xarray.Dataset( - {field: (dims, ndarr[field]) for field in ndarr.dtype.names}, - coords=coords, attrs=attrs) + {field: (dims, arr[field]) for field in arr.dtype.names}, + coords=coords, attrs=attrs + ) else: if name is None: name = f'{self.source}.{self.key}' @@ -653,7 +670,8 @@ def _dim_name(idx): # Primitive dtype. return xarray.DataArray( - ndarr, dims=dims, coords=coords, name=name, attrs=attrs) + arr, dims=dims, coords=coords, name=name, attrs=attrs + ) def series(self): """Load this data as a pandas Series. Only for 1D data. @@ -671,7 +689,8 @@ def series(self): data = self.ndarray() return pd.Series(data, name=name, index=index) - def dask_array(self, labelled=False): + def dask_array(self, labelled=False, *, roi=(), name=None, extra_dims=None, + extra_coords=None): """Make a Dask array for this data. Dask is a system for lazy parallel computation. This method doesn't @@ -691,6 +710,26 @@ def dask_array(self, labelled=False): labelled: bool If True, label the train IDs for the data, returning an xarray.DataArray object wrapping a Dask array. + roi: numpy.s_[], slice, or tuple of slices + The region of interest. This expression selects data in all + dimensions apart from the first (trains) dimension. If the data + holds a 1D array for each entry, roi=np.s_[:8] would get the first 8 + values from every train. If the data is 2D or more at each entry, + selection looks like roi=np.s_[:8, 5:10] . + name: str + Name the array itself. The default is the source and key joined by a + dot. Ignored if labelled is False, and for structured data when a + dataset is returned. + extra_dims: list of str + Name extra dimensions in the array. The first dimension is + automatically called 'train'. The default for extra dimensions is + dim_0, dim_1, ... Ignored if labelled is False. + extra_coords: bool or dict + Add coordinates to the returned DataArray. If roi is used, the + coordinates will match the selected region of interest. If a dict is + given, it should map dimension names to coordinate arrays. If True, + default coordinate arrays will be generated. Ignored if labelled is + False. """ import dask.array as da @@ -725,17 +764,79 @@ def dask_array(self, labelled=False): shape = (0,) + self.entry_shape dask_arr = da.zeros(shape=shape, dtype=self.dtype, chunks=shape) + dask_arr = dask_arr[(slice(None),) + roi] + if labelled: - # Dimension labels - dims = ['trainId'] + ['dim_%d' % i for i in range(dask_arr.ndim - 1)] + return self._wrap_xarray( + dask_arr, + extra_dims=extra_dims, + roi=roi, name=name, + extra_coords=extra_coords + ) + else: + return dask_arr + + def with_entry_dims(self, dims=None, roi=None, **kwargs): + """Attach dimension names, ROI and optional coordinate labels - # Train ID index - coords = {'trainId': self.train_id_coordinates()} + Parameters + ---------- - import xarray - return xarray.DataArray(dask_arr, dims=dims, coords=coords) + dims: list or dict + Either a list of dimension names, or a dictionary mapping dimension + names to coordinate labels. Every dimension except for the trains/ + entries dimension must be named, but coordinate labels are optional. + Uses None in the dict format to omit coordinate labels. If a ROI + is specified, coordinate labels should be only for that selection. + roi: numpy.s_[], slice, or tuple of slices + The region of interest. This expression selects data in all + dimensions apart from the first (trains) dimension. If the data + holds a 1D array for each entry, roi=np.s_[:8] would get the first 8 + values from every train. If the data is 2D or more at each entry, + selection looks like roi=np.s_[:8, 5:10] . + kwargs: + The dictionary format for dims can be passed as keyword arguments + instead. Don't mix keyword arguments with a dims list/dict. + """ + if not isinstance(roi, tuple): + roi = roi, + + if dims is None: + if not kwargs: + raise TypeError("No dimension information specified") + coords = kwargs + elif kwargs: + raise TypeError("Mixing positional and keyword arguments is not supported") + elif isinstance(dims, str): + coords = {dims: None} + elif isinstance(dims, Sequence): + # If we are given names, generate integer indexes + coords = {name: coord for (name, coord) in + zip(dims, expand_indexing(self.entry_shape, roi))} + elif isinstance(dims, dict): + coords = dims else: - return dask_arr + raise TypeError( + f"Unexpected type for dimension/coordinate info: {type(dims)}" + ) + + if len(coords) != len(self.entry_shape): + raise TypeError( + f"Expected names for {len(self.entry_shape)} dimensions, got {len(coords)}" + ) + + return KeyData( + self.source, self.key, + train_ids=self.train_ids, + files=self.files, + section=self.section, + dtype=self.dtype, + eshape=self.entry_shape, + edims=coords, + roi=roi, + inc_suspect_trains=self.inc_suspect_trains, + ) + # Getting data by train: -------------------------------------------------- diff --git a/extra_data/tests/test_keydata.py b/extra_data/tests/test_keydata.py index 52727727..7e9b8b9c 100644 --- a/extra_data/tests/test_keydata.py +++ b/extra_data/tests/test_keydata.py @@ -3,6 +3,7 @@ import pandas as pd import xarray as xr import pytest +from functools import partial import h5py @@ -566,7 +567,7 @@ def test_xarray_extra_dims_and_coords(mock_spb_raw_run): assert 'b' not in da8.coords np.testing.assert_allclose(da8.coords['fs'], custom_fs) - # Passing only extra_coords (no extra_dims) + # Passing only extra_coords (no extra_dims), name not already in dims with pytest.raises(ValueError): am0.xarray(roi=np.s_[:, 10:20, 5:15], extra_coords={'fs': np.arange(5, 15)}) @@ -586,6 +587,100 @@ def test_xarray_extra_dims_and_coords(mock_spb_raw_run): np.testing.assert_array_equal(da11.coords['dim_2'], np.arange(5, 15)) +def test_dask_array_extra_dims_and_coords(mock_spb_raw_run): + run = RunDirectory(mock_spb_raw_run) + am0 = run['SPB_DET_AGIPD1M-1/DET/0CH0:xtdf', 'image.data'][0] + assert am0.entry_shape == (2, 512, 128) + + lda = partial(am0.dask_array, labelled=True) + + # No extra dims and extra_coords==False -> only trainId coordinate + da_default = lda() + assert da_default.shape == (64, 2, 512, 128) + assert da_default.dims == ('trainId', 'dim_0', 'dim_1', 'dim_2') + assert set(da_default.coords) == {'trainId'} + + # ROI with slices, with extra_dims -> generate coords + da1 = lda(roi=np.s_[:, 10:20, 5:15], extra_dims=['raw_proc', 'ss', 'fs']) + assert da1.shape == (64, 2, 10, 10) + assert da1.dims == ('trainId', 'raw_proc', 'ss', 'fs') + np.testing.assert_array_equal(da1.coords['raw_proc'], np.arange(2)) + np.testing.assert_array_equal(da1.coords['ss'], np.arange(10, 20)) + np.testing.assert_array_equal(da1.coords['fs'], np.arange(5, 15)) + + # Short ROI: explicit non-full slice on first entry dim + da2 = lda(roi=np.s_[:1, 30:40, ...], extra_dims=['a', 'b', 'c']) + assert da2.shape == (64, 1, 10, 128) + assert da2.dims == ('trainId', 'a', 'b', 'c') + np.testing.assert_array_equal(da2.coords['a'], np.arange(0, 1)) + np.testing.assert_array_equal(da2.coords['b'], np.arange(30, 40)) + np.testing.assert_array_equal(da2.coords['c'], np.arange(128)) + + # Ellipsis at the start + da3 = lda(roi=np.s_[..., 30:40], extra_dims=['a', 'b', 'c']) + assert da3.shape == (64, 2, 512, 10) + assert da3.dims == ('trainId', 'a', 'b', 'c') + np.testing.assert_array_equal(da3.coords['a'], np.arange(2)) + np.testing.assert_array_equal(da3.coords['b'], np.arange(512)) + np.testing.assert_array_equal(da3.coords['c'], np.arange(30, 40)) + + # Integer-only index on first entry dim with extra_dims -> keep dim with size 1 + da4 = lda(roi=np.s_[0, :, :], extra_dims=['a', 'b', 'c']) + # The data is a single index in the first entry, dim is dropped + assert da4.shape == (64, 512, 128) + assert da4.dims == ('trainId', 'b', 'c') + # Coordinate for the dropped dimension is the selected index + a_coord = np.asarray(da4.coords['a']) + np.testing.assert_array_equal(a_coord, np.array(0)) + + # Fancy indexing with list of indices + da5 = lda(roi=np.s_[:, [30, 33], :], extra_dims=['a', 'b', 'c']) + assert da5.shape == (64, 2, 2, 128) + np.testing.assert_array_equal(da5.coords['b'], np.array([30, 33])) + + # Boolean indexing + mask = np.zeros((512,), dtype=bool) + mask[::2] = True + da6 = lda(roi=np.s_[:, mask, 5:15], extra_dims=['a', 'ss', 'fs']) + assert da6.shape == (64, 2, 256, 10) + np.testing.assert_array_equal(da6.coords['ss'], np.arange(0, 512, 2)) + np.testing.assert_array_equal(da6.coords['fs'], np.arange(5, 15)) + + # Negative indexing on last two dims + da7 = lda(roi=np.s_[:, -10:, [-10, -5, -1]], extra_dims=['a', 'b', 'c']) + assert da7.shape == (64, 2, 10, 3) + np.testing.assert_array_equal(da7.coords['b'], np.arange(502, 512)) + np.testing.assert_array_equal(da7.coords['c'], np.array([118, 123, 127])) + + # Custom coordinates via extra_coords + custom_fs = np.linspace(0.0, 9.0, 10) + da8 = lda(roi=np.s_[:, 10:20, 5:15], extra_dims=['a', 'b', 'fs'], extra_coords={'fs': custom_fs}) + assert da8.shape == (64, 2, 10, 10) + assert da8.dims == ('trainId', 'a', 'b', 'fs') + assert 'a' not in da8.coords + assert 'b' not in da8.coords + np.testing.assert_allclose(da8.coords['fs'], custom_fs) + + # Passing only extra_coords (no extra_dims), name not already in dims + with pytest.raises(ValueError): + lda(roi=np.s_[:, 10:20, 5:15], extra_coords={'fs': np.arange(5, 15)}) + + da9 = lda(roi=np.s_[:, 10:20, 5:15], extra_dims=['r/p', 'ss', 'fs'], extra_coords={'fs': np.arange(5, 15)}) + assert list(da9.coords) >= ['trainId', 'fs'] + + da10 = lda(roi=np.s_[:, 15, 5:15], extra_coords={'dim_0': [100, 200]}) + assert list(da10.coords) >= ['trainId', 'dim_0'] + + # extra_coords == True + da11 = lda(roi=np.s_[:, 15, 5:15], extra_coords=True) + assert da11.shape == (64, 2, 10) + assert da11.dims == ('trainId', 'dim_0', 'dim_2') + assert set(da11.coords) == {'trainId', 'dim_0', 'dim_1', 'dim_2'} + np.testing.assert_array_equal(da11.coords['dim_0'], np.arange(2)) + np.testing.assert_array_equal(da11.coords['dim_1'], np.arange(15, 16)) + np.testing.assert_array_equal(da11.coords['dim_2'], np.arange(5, 15)) + + @pytest.mark.skipif(not os.path.isdir("/pnfs/xfel.eu"), reason="xfel.eu dCache not available") def test_drop_empty_trains_out_of_order(): run = open_run(900174, 354, data='raw', include='*-AGIPD00-*.h5')