From 63ea76e5358d7955b23c8593d01fd6c2164e5769 Mon Sep 17 00:00:00 2001 From: Dmitrii Kochkov Date: Tue, 24 Feb 2026 15:44:02 -0800 Subject: [PATCH] Added experimental `isel`, `sel` methods to Coordinate and Field classes. PiperOrigin-RevId: 874824835 --- coordax/__init__.py | 2 +- coordax/coordinate_systems.py | 452 ++++++++++++++++++++++++++++- coordax/coordinate_systems_test.py | 170 +++++++++++ coordax/experimental.py | 4 + coordax/fields.py | 151 +++++++++- coordax/fields_test.py | 227 +++++++++++++-- pyproject.toml | 2 +- 7 files changed, 979 insertions(+), 29 deletions(-) diff --git a/coordax/__init__.py b/coordax/__init__.py index e2476e0..4c4d64f 100644 --- a/coordax/__init__.py +++ b/coordax/__init__.py @@ -56,4 +56,4 @@ ) import coordax.testing # pylint: disable=unused-import -__version__ = '0.2.5' # keep sync with pyproject.toml +__version__ = '0.2.6' # keep sync with pyproject.toml diff --git a/coordax/coordinate_systems.py b/coordax/coordinate_systems.py index 6155ef4..ef37243 100644 --- a/coordax/coordinate_systems.py +++ b/coordax/coordinate_systems.py @@ -16,6 +16,7 @@ ``Coordinate`` objects define a discretization schema, dimension names and provide methods & coordinate field values to facilitate computations. """ + from __future__ import annotations import abc @@ -25,7 +26,7 @@ import functools import itertools import typing -from typing import Any, Self, TYPE_CHECKING, Type, TypeAlias, TypeGuard, TypeVar +from typing import Any, Literal, Self, TYPE_CHECKING, Type, TypeAlias, TypeGuard, TypeVar import warnings from coordax import utils @@ -42,6 +43,8 @@ Pytree: TypeAlias = Any Sequence = collections.abc.Sequence +SelMethod: TypeAlias = Literal['nearest'] | None + @functools.partial(utils.export, module='coordax.coords') @dataclasses.dataclass(frozen=True) @@ -51,6 +54,26 @@ class NoCoordinateMatch: reason: str +def normalize_indexers( + indexers: dict[str | Coordinate, Any] | None, + **indexers_kwargs, +) -> dict[str | Coordinate, Any]: + """Returns indexers replacing sequence with arrays, checks spec mode.""" + if indexers is None: # only kwargs is allowed. + normalized_indexers = dict(indexers_kwargs) + else: + if not isinstance(indexers, dict): + raise ValueError(f'Indexers must be a dict, got {type(indexers)}.') + if indexers_kwargs: + raise ValueError( + 'Using dict and kwarg indexers simultaneously is dangerous and not' + f' supported, got {indexers=}, {indexers_kwargs=}.' + ) + normalized_indexers = indexers + seq_to_array = lambda x: np.asarray(x) if isinstance(x, Sequence) else x + return {k: seq_to_array(v) for k, v in normalized_indexers.items()} + + @utils.export class Coordinate(abc.ABC): """Abstract class for coordinate objects. @@ -108,6 +131,158 @@ def axes(self) -> tuple[Coordinate, ...]: else: return tuple(SelectedAxis(self, i) for i in range(self.ndim)) + def isel( + self, + indexers: dict[str | Coordinate, Any] | None = None, + **indexers_kwargs, + ) -> Coordinate: + """Returns a new coordinate with the given integer indexers applied. + + Note: This is an experimental feature, and may be changed or completely + removed in the future. + + Note that ``isel`` should only be used with integer or slice indexers, or + array-like objects with integer or slice values. All keys provided to `isel` + are expected to be present in the dimension names or axes of a the + coordinate. By default, empty and ``slice(None)`` indexers return the + original coordinate. Otherwise a subclass-specific ``_isel`` is called. + + For label-based selection, use ``sel`` instead that is implemented using + ``map_indexers`` in subclasses that allows mapping from label indexers to + integer indexers used by ``isel``. + + Args: + indexers: A mapping from dimensions to indices, slices, or arrays. + **indexers_kwargs: The keyword arguments form of ``indexers``. + + Returns: + A new coordinate with the selection applied. + + Examples: + >>> import coordax as cx + >>> import numpy as np + >>> x = cx.SizedAxis('x', 5) + >>> x.isel(x=0) + Scalar() + >>> x.isel(x=slice(1, 4)) + coordax.SizedAxis('x', size=3) + >>> x.isel(x=[0, 2, 4]) + coordax.SizedAxis('x', size=3) + + >>> y = cx.LabeledAxis('y', np.arange(5)) + >>> y.isel(y=slice(0, 2)) + coordax.LabeledAxis('y', ticks=array([0, 1])) + """ + indexers = normalize_indexers(indexers, **indexers_kwargs) + self._validate_indexers(indexers) + + indexer = indexers.get(self) + if not indexers or (isinstance(indexer, slice) and indexer == slice(None)): + return self + + return self._isel(indexers) + + def _isel(self, indexers: dict[str | Coordinate, Any]) -> Coordinate: + """Returns a new coordinate with the given indexers applied. + + Args: + indexers: A mapping from dimensions to indices, slices, or arrays. + Guaranteed to be non-empty and contain only valid keys for this + coordinate. + """ + raise NotImplementedError(f'{type(self).__name__} does not implement _isel') + + def _validate_indexers(self, indexers: dict[str | Coordinate, Any]) -> None: + """Validates that the integer indexers are valid for this coordinate.""" + unknown_dims = set(indexers.keys()) - set(self.axes) - set(self.dims) + if unknown_dims: + for k in unknown_dims: + if is_coord(k) and k.ndim > 1: + raise ValueError( + f'Indexing with coordinate {k=} with ndim > 1 is not supported.' + ) + raise ValueError( + f'Dimensions {unknown_dims} do not exist in coordinate' + f' {type(self).__name__}' + ) + + def sel( + self, + indexers: dict[str | Coordinate, Any] | None = None, + method: Literal['nearest'] | None = None, + **indexers_kwargs, + ) -> Coordinate: + """Returns a new coordinate with the given selection applied. + + Note: This is an experimental feature, and may be changed or completely + removed in the future. + + ``sel`` is designed to work with label-based indexers, which may include + indexers that are not in the coordinate dimensions. The selection is + accomplished by mapping label-based indexers to integer-based indexers via + ``map_indexers``, which can be customized in subclasses. Indexers that are + not processed by ``map_indexers`` are considered "unused" and will raise an + error. + + Args: + indexers: A mapping specifying labels to be selected from the coordinate. + method: Optional method to use for inexact matches. Cannot be used when + ``indexers`` contain slices. Default is `None`. + **indexers_kwargs: The keyword arguments form of ``indexers``. + + Returns: + A new coordinate with the selection applied. + + Examples: + >>> import coordax as cx + >>> import numpy as np + >>> x = cx.LabeledAxis('x', np.array([10, 20, 30])) + >>> x.sel(x=20) + Scalar() + >>> x.sel(x=slice(10, 20)) + coordax.LabeledAxis('x', ticks=array([10, 20])) + >>> x.sel(x=12, method='nearest') + Scalar() + """ + sel_indexers = normalize_indexers(indexers, **indexers_kwargs) + if not sel_indexers: + return self + + unpacked_indexers, unpacked_c = unpack_and_validate_indexers(sel_indexers) + mapped_indexers, consumed = self.map_indexers( + unpacked_indexers, method=method + ) + + final_consumed = set() + for c in consumed: + if c in unpacked_c: + final_consumed.add(unpacked_c[c]) + else: + final_consumed.add(c) + + unused_sel_indexers = set(sel_indexers.keys()) - final_consumed + if unused_sel_indexers: + raise ValueError( + f'Indexers {unused_sel_indexers} were not processed by any component' + f' in {self}' + ) + + return self.isel(mapped_indexers) + + def map_indexers( + self, + indexers: dict[str | Coordinate, Any], + method: Literal['nearest'] | None = None, + ) -> tuple[dict[str | Coordinate, Any], set[str | Coordinate]]: + """Maps label-based indexers to integer-based indexers (indices/slices).""" + del method # unused. + # support `sel` on composite coords as long as they are not being indexed. + if not indexers or not any(dim in indexers for dim in self.dims + (self,)): + return {}, set() + raise NotImplementedError( + f'{type(self).__name__} does not implement map_indexers' + ) + def to_xarray(self) -> dict[str, xarray.Variable]: """Convert this coordinate into xarray variables.""" import xarray # pylint: disable=g-import-not-at-top @@ -145,6 +320,42 @@ def from_xarray( raise NotImplementedError('from_xarray not implemented') +def unpack_and_validate_indexers( + indexers: dict[str | Coordinate, Any], +) -> tuple[dict[str | Coordinate, Any], dict[str | Coordinate, Coordinate]]: + """Unpacks multidimensional indexers and raises if slice.step is not None.""" + unpacked_indexers = {} + unpacked_coords = {} + # pytype: disable=attribute-error + for k, v in indexers.items(): + if is_coord(k) and k.ndim > 1: + key_coord = k + if is_coord(v): + if key_coord.dims != v.dims: + raise ValueError( + f'{key_coord.dims=} do not match indexer dimensions {v.dims=}.' + ) + for ax, v_ax in zip(key_coord.axes, v.axes): + unpacked_indexers[ax] = v_ax + assert isinstance(ax, Coordinate) + assert isinstance(key_coord, Coordinate) + unpacked_coords[ax] = key_coord + else: + raise ValueError( + f'Indexer for {key_coord=} with {key_coord.ndim=} > 0 must be a' + f' coordinate, got {v}.' + ) + else: + if isinstance(v, slice) and v.step is not None: + raise ValueError( + f'Indexer for {k=} uses slice with {v.step=} != None, which is ' + 'not supported.' + ) + unpacked_indexers[k] = v + # pytype: enable=attribute-error + return unpacked_indexers, unpacked_coords + + @functools.partial(utils.export, module='coordax.coords') @dataclasses.dataclass(frozen=True) class ArrayKey: @@ -197,6 +408,14 @@ def dims(self) -> tuple[str, ...]: def shape(self) -> tuple[int, ...]: return () + def map_indexers( + self, + indexers: dict[str | Coordinate, Any], + method: Literal['nearest'] | None = None, + ) -> tuple[dict[str | Coordinate, Any], set[str | Coordinate]]: + del indexers, method # unused. + return {}, set() + @utils.export @jax.tree_util.register_static @@ -217,6 +436,20 @@ def __post_init__(self): f'dimension {self.axis=} of {self.coordinate=} is not named' ) + def map_indexers( + self, + indexers: dict[str | Coordinate, Any], + method: Literal['nearest'] | None = None, + ) -> tuple[dict[str | Coordinate, Any], set[str | Coordinate]]: + return self.coordinate.map_indexers(indexers, method=method) + + def _isel(self, indexers: dict[str | Coordinate, Any]) -> Coordinate: + start_ndim = self.coordinate.ndim + new_coord = self.coordinate.isel(indexers) + if new_coord.ndim == start_ndim: + return SelectedAxis(new_coord, self.axis) + return Scalar() + @property def dims(self) -> tuple[str, ...]: """Dimension names of the coordinate.""" @@ -392,6 +625,26 @@ def axes(self) -> tuple[Coordinate, ...]: """Returns a tuple of Axis objects for each dimension.""" return _concat_tuples(c.axes for c in self.coordinates) + def _isel(self, indexers: dict[str | Coordinate, Any]) -> Coordinate: + new_coords = [] + for c in self.coordinates: # already canonicalized. + c_indexers = {k: v for k, v in indexers.items() if contains_dims(c, k)} + new_coords.append(c.isel(c_indexers)) + return compose(*new_coords) + + def map_indexers( + self, + indexers: dict[str | Coordinate, Any], + method: Literal['nearest'] | None = None, + ) -> tuple[dict[str | Coordinate, Any], set[str | Coordinate]]: + mapped = {} + consumed = set() + for c in self.coordinates: + sub_mapped, sub_consumed = c.map_indexers(indexers, method=method) + mapped.update(sub_mapped) + consumed.update(sub_consumed) + return mapped, consumed + @utils.export @jax.tree_util.register_static @@ -410,6 +663,35 @@ def dims(self) -> tuple[str, ...]: def shape(self) -> tuple[int, ...]: return (self.size,) + def map_indexers( + self, + indexers: dict[str | Coordinate, Any], + method: Literal['nearest'] | None = None, + ) -> tuple[dict[str | Coordinate, Any], set[str | Coordinate]]: + if indexers and (self.name in indexers or self in indexers): + key = self.name if self.name in indexers else self + assert isinstance(key, (str, Coordinate)) + if indexers[key] == self: + return {key: slice(None)}, {key} + raise ValueError(f'{type(self).__name__} does not support `sel`.') + return {}, set() + + def _isel(self, indexers: dict[str | Coordinate, Any]) -> Coordinate: + key = self.name if self.name in indexers else self + indexer = indexers[key] + if isinstance(indexer, int): + return Scalar() + + if isinstance(indexer, slice): + start, stop, step = indexer.indices(self.size) + new_size = len(range(start, stop, step)) + return SizedAxis(self.name, new_size) + + if hasattr(indexer, '__len__'): # array or list. + return SizedAxis(self.name, len(indexer)) + + raise ValueError(f'Unsupported indexer type {type(indexer)}') + def __repr__(self): return f'coordax.SizedAxis({self.name!r}, size={self.size})' @@ -456,6 +738,35 @@ def dims(self) -> tuple[str | None, ...]: def shape(self) -> tuple[int, ...]: return (self.size,) + def map_indexers( + self, + indexers: dict[str | Coordinate, Any], + method: Literal['nearest'] | None = None, + ) -> tuple[dict[str | Coordinate, Any], set[str | Coordinate]]: + if indexers and (self.name in indexers or self in indexers): + key = self.name if self.name in indexers else self + assert isinstance(key, (str, Coordinate)) + if indexers[key] == self: + return {key: slice(None)}, {key} + raise ValueError(f'{type(self).__name__} does not support `sel`.') + return {}, set() + + def _isel(self, indexers: dict[str | Coordinate, Any]) -> Coordinate: + key = self.name if self.name in indexers else self + indexer = indexers[key] + if isinstance(indexer, int): + return Scalar() + + if isinstance(indexer, slice): + start, stop, step = indexer.indices(self.size) + new_size = len(range(start, stop, step)) + return DummyAxis(self.name, new_size) + + if hasattr(indexer, '__len__'): + return DummyAxis(self.name, len(indexer)) + + raise ValueError(f'Unsupported indexer type {type(indexer)}') + def __repr__(self): return f'coordax.DummyAxis({self.name!r}, size={self.size})' @@ -473,8 +784,118 @@ def from_xarray( return cls(name=dim, size=coords.sizes[dim]) -# TODO(dkochkov): consider storing tuple values instead of np.ndarray (which -# could be exposed as a property). +@functools.partial(utils.export, module='coordax.experimental') +def map_indexers_using_ticks( + axis: Coordinate, + indexers: dict[str | Coordinate, Any], + ticks: np.ndarray | None = None, + ticks_are_sorted: bool | None = None, + method: Literal['nearest'] | None = None, +) -> tuple[dict[str | Coordinate, Any], set[str | Coordinate]]: + """Maps indexers to ticks using either fancy indexing or slices. + + Args: + axis: The coordinate for which to map indexers. + indexers: A mapping from dimensions to slices. + ticks: The tick values to use for mapping indexers. If None, ticks are taken + from the `axis`. Default is None. Can be used to map custom ticks. + ticks_are_sorted: Whether the `ticks` are sorted. + method: Method to use for inexact matches. Either None (default) for exact + matches or 'nearest'. Cannot be used when indexers contain slices. + + Returns: + A tuple of (mapped_indexers, consumed_keys). + + Raises: + KeyError: If any values in the indexer are not found in the ticks. + ValueError: If an unsupported method is provided or if ticks are not unique. + """ + if axis.ndim != 1: + raise ValueError( + f'Mapping indexers using ticks requires 1D axis, got {axis.ndim=}' + ) + [dim] = axis.dims + if axis not in indexers and dim not in indexers: + return {}, set() + + ticks = ticks if ticks is not None else axis.fields[dim].data + if np.unique(ticks).size != ticks.size: + raise ValueError(f'Ticks must be unique, got {ticks}') + if ticks_are_sorted is None: + ticks_are_sorted = np.all(ticks[:-1] < ticks[1:]) + + key = axis if axis in indexers else dim + assert isinstance(key, (str, Coordinate)) # make pytype happy + value = indexers[key] + if is_coord(value) and value == axis: + return {key: slice(None)}, {key} + if is_coord(value): + if not isinstance(value, type(axis)) or value.dims != axis.dims: + raise ValueError( + 'Indexing with axis requires same type and dims, got index axis' + f' {value} for slicing {axis=}.' + ) + value = value.fields[dim].data + + if isinstance(value, slice): + if method is not None: + raise NotImplementedError('Method argument not supported for slices') + + start, stop = value.start, value.stop + if ticks_are_sorted: + start_idx, stop_idx = 0, ticks.size + if start is not None: + start_idx = np.searchsorted(ticks, start, side='left') + if stop is not None: + stop_idx = np.searchsorted(ticks, stop, side='right') + + return {key: slice(start_idx, stop_idx)}, {key} + else: + mask = np.ones(ticks.size, dtype=bool) + if start is not None: + mask &= ticks >= start + if stop is not None: + mask &= ticks <= stop + return {key: np.where(mask)[0]}, {key} + + if method == 'nearest': + if ticks_are_sorted: + candidates = np.searchsorted(ticks, value, side='left') + left = np.maximum(candidates - 1, 0) + right = np.minimum(candidates, len(ticks) - 1) + d_left = np.abs(value - ticks[left]) + d_right = np.abs(value - ticks[right]) + # In case of ties, prefer the left (smaller) index. + idx = np.where(d_left <= d_right, left, right) + elif np.ndim(value) == 0: + idx = np.abs(ticks - value).argmin() + else: + ticks_view = ticks.reshape((-1,) + (1,) * np.ndim(value)) + idx = np.abs(ticks_view - value).argmin(axis=0) + + if np.ndim(idx) == 0: + idx = int(idx) + return {key: idx}, {key} + + if method is None: + sort_indices = None if ticks_are_sorted else np.argsort(ticks) + sorted_ticks = ticks if ticks_are_sorted else ticks[sort_indices] + indices = np.searchsorted(sorted_ticks, value) + if sort_indices is not None: + indices = sort_indices[indices] + unique_retrieved = np.sort(np.unique(ticks[indices])) + unique_value = np.sort(np.unique(value)) + if unique_retrieved.size != unique_value.size or np.any( + unique_retrieved != unique_value + ): + raise KeyError(f'Not all values in {value} were found in {ticks}') + if np.ndim(value) == 0: # if value is not an array, index must be an int. + indices = indices.item() + return {key: indices}, {key} + + raise ValueError(f'Unknown method {method}') + + @utils.export @jax.tree_util.register_static @dataclasses.dataclass(frozen=True) @@ -504,6 +925,31 @@ def fields(self) -> dict[str, 'fields.Field']: return {self.name: fields.field(self.ticks, self)} + @functools.cached_property + def _sorted_ticks(self) -> bool: + return np.all(self.ticks[:-1] <= self.ticks[1:]) + + def map_indexers( + self, + indexers: dict[str | Coordinate, Any], + method: Literal['nearest'] | None = None, + ) -> tuple[dict[str | Coordinate, Any], set[str | Coordinate]]: + return map_indexers_using_ticks( + axis=self, + indexers=indexers, + ticks=self.ticks, + ticks_are_sorted=self._sorted_ticks, + method=method, + ) + + def _isel(self, indexers: dict[str | Coordinate, Any]) -> Coordinate: + key = self.name if self.name in indexers else self + indexer = indexers[key] + if isinstance(indexer, int): + return Scalar() + + return LabeledAxis(self.name, self.ticks[indexer]) + def _components(self): return (self.name, ArrayKey(self.ticks)) diff --git a/coordax/coordinate_systems_test.py b/coordax/coordinate_systems_test.py index e5c0cf2..851b49b 100644 --- a/coordax/coordinate_systems_test.py +++ b/coordax/coordinate_systems_test.py @@ -480,6 +480,176 @@ def test_extract(self): with self.assertRaisesRegex(ValueError, 'Expected exactly one instance'): cx.coords.extract(x, cx.LabeledAxis) + def test_isel_sized_axis(self): + axis = cx.SizedAxis('x', 5) + with self.subTest('integer_indexing'): + self.assertEqual(axis.isel({'x': 0}), cx.Scalar()) + with self.subTest('slice_indexing'): + self.assertEqual(axis.isel({'x': slice(1, 4)}), cx.SizedAxis('x', 3)) + with self.subTest('none_slice'): + sliced = axis.isel(x=slice(None)) + self.assertEqual(sliced, axis) + with self.subTest('array_indexing'): + self.assertEqual(axis.isel({'x': [0, 2, 4]}), cx.SizedAxis('x', 3)) + with self.subTest('array_indexing_kwargs'): + self.assertEqual(axis.isel(x=[2, 4]), cx.SizedAxis('x', 2)) + with self.subTest('slice_out_of_bounds'): # xarray semantics. + self.assertEqual(axis.isel(x=slice(3, 10)), cx.SizedAxis('x', 2)) + with self.subTest('axis_as_key'): + self.assertEqual(axis.isel({axis: 0}), cx.Scalar()) + + def test_isel_labeled_axis(self): + axis = cx.LabeledAxis('x', np.arange(5)) + with self.subTest('integer_indexing'): + self.assertEqual(axis.isel(x=0), cx.Scalar()) + with self.subTest('slice_indexing'): + sliced = axis.isel(x=slice(1, 4)) + expected = cx.LabeledAxis('x', np.arange(1, 4)) + self.assertEqual(sliced, expected) + with self.subTest('none_slice'): + sliced = axis.isel(x=slice(None)) + self.assertEqual(sliced, axis) + with self.subTest('array_indexing'): + sliced = axis.isel(x=[0, 2, 4]) + expected = cx.LabeledAxis('x', np.array([0, 2, 4])) + self.assertEqual(sliced, expected) + with self.subTest('axis_as_key'): + sliced = axis.isel({axis: [0, 1]}) + expected = cx.LabeledAxis('x', np.array([0, 1])) + self.assertEqual(sliced, expected) + with self.subTest('supports_repeated_indices'): + sliced = axis.isel({axis: [0, 0]}) + expected = cx.LabeledAxis('x', np.array([0, 0])) + self.assertEqual(sliced, expected) + + def test_isel_unknown_dim_raises(self): + x = cx.SizedAxis('x', 2) + with self.assertRaisesRegex(ValueError, 'Dimensions .* do not exist'): + x.isel({'y': 0}) + with self.subTest('for_scalar'): + with self.assertRaisesRegex(ValueError, 'Dimensions .* do not exist'): + cx.Scalar().isel({'x': 0}) + + def test_sel_labeled_axis(self): + axis = cx.LabeledAxis('x', np.array([10, 20, 30, 40, 50])) + with self.subTest('exact_match'): + self.assertEqual(axis.sel({'x': 20}), cx.Scalar()) + with self.subTest('slice_match'): + # xarray semantics: [10, 40] -> 10, 20, 30, 40 + sel_slice = axis.sel({'x': slice(10, 40)}) + expected = cx.LabeledAxis('x', np.array([10, 20, 30, 40])) + self.assertEqual(sel_slice, expected) + with self.subTest('multiple_matches'): + sel_values = axis.sel({'x': [10, 40]}) + expected = cx.LabeledAxis('x', np.array([10, 40])) + self.assertEqual(sel_values, expected) + with self.subTest('nearest_match'): + sel_nearest = axis.sel({'x': 22}, method='nearest') + self.assertEqual(sel_nearest, cx.Scalar()) + with self.subTest('nearest_match_sequence'): + sel_nearest = axis.sel({'x': [22, 29]}, method='nearest') + self.assertEqual(sel_nearest, cx.LabeledAxis('x', np.array([20, 30]))) + with self.subTest('nearest_match_sequence_same_index'): + sel_nearest = axis.sel({'x': [22, 23]}, method='nearest') + self.assertEqual(sel_nearest, cx.LabeledAxis('x', np.array([20, 20]))) + with self.subTest('nearest_match_unsorted'): + axis_unsorted = cx.LabeledAxis('x', np.array([30, 10, 50, 20, 40])) + sel_nearest = axis_unsorted.sel({'x': [22, 31]}, method='nearest') + self.assertEqual(sel_nearest, cx.LabeledAxis('x', np.array([20, 30]))) + with self.subTest('preserves_sel_order_ordered_ticks'): + sel_ordered = axis.sel({'x': [40, 20]}) + expected = cx.LabeledAxis('x', np.array([40, 20])) + self.assertEqual(sel_ordered, expected) + with self.subTest('preserves_sel_order_unsorted_ticks'): + axis_unsorted = cx.LabeledAxis('x', np.array([30, 10, 50, 20, 40])) + sel_ordered = axis_unsorted.sel({'x': [40, 20]}) + expected = cx.LabeledAxis('x', np.array([40, 20])) + self.assertEqual(sel_ordered, expected) + with self.subTest('supports_repeated_labels'): + sel_repeated = axis.sel({'x': [20, 20]}) + expected = cx.LabeledAxis('x', np.array([20, 20])) + self.assertEqual(sel_repeated, expected) + + def test_sel_labeled_axis_raises_on_no_match(self): + axis = cx.LabeledAxis('x', np.array([-1, 2, 3])) + with self.subTest('value_not_found'): + with self.assertRaises(KeyError): + axis.sel({'x': 0}) + with self.subTest('not_all_values_found'): + with self.assertRaisesRegex(KeyError, 'Not all values in .* were found'): + axis.sel({'x': [2, 2.5, 3]}) + + def test_sel_labeled_axis_raises_on_slice_with_step(self): + axis = cx.LabeledAxis('x', np.arange(10)) + with self.assertRaisesRegex( + ValueError, + "Indexer for k='x' uses slice with v.step=2 != None, which is not" + ' supported.', + ): + axis.sel({'x': slice(0, 5, 2)}) + + def test_isel_composed(self): + x, y = cx.SizedAxis('x', 3), cx.SizedAxis('y', 4) + xy = cx.coords.compose(x, y) + with self.subTest('index_one_axis'): + self.assertEqual(xy.isel(x=0), y) + with self.subTest('slice_one_axis'): + expected = cx.coords.compose(x, cx.SizedAxis('y', 2)) + self.assertEqual(xy.isel(y=slice(0, 2)), expected) + with self.subTest('index_both'): + self.assertEqual(xy.isel(x=0, y=1), cx.Scalar()) + with self.subTest('index_using_axis'): + expected = cx.coords.compose(cx.SizedAxis('x', 1), y) + self.assertEqual(xy.isel({x: slice(0, 1)}), expected) + with self.subTest('none_slice'): + sliced = xy.isel(x=slice(None)) + self.assertEqual(sliced, xy) + + def test_sel_composed(self): + x = cx.SizedAxis('x', 2) + y = cx.LabeledAxis('y', np.array([10, 20, 30])) + coord = cx.coords.compose(x, y) + + with self.subTest('single_value'): + # SizedAxis 'x' ignored (no sel support), LabeledAxis 'y' selected. + selected = coord.sel(y=20) + # x remains, y becomes scalar (dropped from composition) -> x + self.assertEqual(selected, x) + + with self.subTest('slice'): + selected = coord.sel(y=slice(10, 20)) + expected_y = cx.LabeledAxis('y', np.array([10, 20])) + expected = cx.coords.compose(x, expected_y) + self.assertEqual(selected, expected) + + with self.subTest('axis_as_key'): + selected = coord.sel({y: 20}) + self.assertEqual(selected, x) + + with self.subTest('axis_as_value'): + target_y = cx.LabeledAxis('y', np.array([10, 20])) + selected = coord.sel(y=target_y) + expected = cx.coords.compose(x, target_y) + self.assertEqual(selected, expected) + + with self.subTest('none_slice'): + selected = coord.sel(y=slice(None)) + self.assertEqual(selected, coord) + + z = cx.LabeledAxis('z', np.linspace(0, np.pi, 10)) + xyz = cx.coords.compose(x, y, z) + with self.subTest('multidim_axis_as_value'): + sub_y = cx.LabeledAxis('y', np.array([10, 20])) + sub_z = cx.LabeledAxis('z', np.linspace(0, np.pi, 10)[::2]) + yz = cx.coords.compose(y, z) + sub_yz = cx.coords.compose(sub_y, sub_z) + selected = xyz.sel({yz: sub_yz}) + self.assertEqual(selected, cx.coords.compose(x, sub_y, sub_z)) + + with self.subTest('full_coordinate'): + selected = coord.sel({coord: coord}) + self.assertEqual(selected, coord) + def test_deprecated_aliases(self): with self.assertWarnsRegex( DeprecationWarning, diff --git a/coordax/experimental.py b/coordax/experimental.py index a1f2f48..a2b8a56 100644 --- a/coordax/experimental.py +++ b/coordax/experimental.py @@ -16,6 +16,10 @@ # Note: import as is required for names to be exported. # See PEP 484 & https://github.com/jax-ml/jax/issues/7570 # pylint: disable=g-multiple-import,useless-import-alias,g-importing-member,unused-import +from coordax.coordinate_systems import ( + map_indexers_using_ticks as map_indexers_using_ticks, + SelMethod as SelMethod, +) from coordax.ndarrays import ( NDArray as NDArray, register_ndarray as register_ndarray, diff --git a/coordax/fields.py b/coordax/fields.py index a29297a..294ee3d 100644 --- a/coordax/fields.py +++ b/coordax/fields.py @@ -17,6 +17,7 @@ Named dimensions of a ``Field`` are associated with coordinates that describe their discretization. """ + from __future__ import annotations import collections @@ -41,6 +42,7 @@ import xarray +# pylint: disable=redefined-outer-name Pytree: TypeAlias = Any Sequence = collections.abc.Sequence @@ -876,6 +878,147 @@ def broadcast_like(self, other: Self | Coordinate) -> Self: self.named_array.broadcast_like(other.named_array), other.axes ) + def isel( + self, + indexers: dict[str | Coordinate, Any] | None = None, + **indexers_kwargs, + ) -> Field: + """Returns a new Field with the given indexers applied. + + Note: This is an experimental feature, and may be changed or completely + removed in the future. + + ``isel`` mimics the behavior of ``xarray.DataArray.isel`` and expects + integer, slice, or array-like objects with integer or slice values. All keys + provided to `isel` are expected to be present in `self.dims`. For + label-based selection, use ``sel`` instead. + + Note: ``isel`` might not be supported for some coordinate types or result in + outputs where coordinate components have different type than that of the + original field. This is a reflection of the limitations of the associated + coordinate class. + + Args: + indexers: A mapping from dimensions to indices, slices, or arrays. + **indexers_kwargs: The keyword arguments form of ``indexers``. + + Returns: + A new Field with the selection applied. + + Examples: + >>> import coordax as cx + >>> import jax.numpy as jnp + >>> field = cx.field(jnp.arange(6).reshape(2, 3), 'x', 'y') + >>> field.isel(x=0) + + >>> field.isel(y=slice(0, 2)) + + >>> field.isel(x=[0, 1]) + + """ + indexers = coordinate_systems.normalize_indexers( + indexers, **indexers_kwargs + ) + if not indexers: + return self + + for dim in indexers: + if coordinate_systems.is_coord(dim): + if dim not in self.coordinate.axes: + raise ValueError( + f'Dimension {dim!r} not found in field with {self.coordinate=}' + ) + else: + if dim not in self.named_axes: + raise ValueError( + f'Dimension {dim!r} not found in field with {self.dims=}' + ) + + dim_names = _dimension_names(*indexers.keys()) + f = self + n_positional = len(f.positional_shape) + tmp_axes = [] + if n_positional > 0: + tmp_axes.append(new_axis_name(f)) + f = f.tag(tmp_axes[-1]) + + for dim, indexer in zip(dim_names, indexers.values(), strict=True): + post_slice_coord = f.coordinate.isel({dim: indexer}) + data_slice = [slice(None)] * f.ndim + data_slice[f.named_axes[dim]] = indexer + f = field(f.data[tuple(data_slice)], post_slice_coord) + return f.untag(*tmp_axes) + + def sel( + self, + indexers: dict[str | Coordinate, Any] | None = None, + method: Literal['nearest'] | None = None, + **indexers_kwargs, + ) -> Field: + """Returns a new Field with the given selection applied. + + Note: This is an experimental feature, and may be changed or completely + removed in the future. + + ``sel`` mimics the behavior of ``xarray.DataArray.sel`` and expects + label values. The selection mechanics is delegated to the underlying + coordinate, via mapping to integer-based indexers, slicing and retagging + with the appropriate coordinate. + + Note: ``sel`` might not be supported for some coordinate types or result in + outputs where coordinate components have different type than that of the + original field. This is a reflection of the limitations of the associated + coordinate class. + + Args: + indexers: A mapping from dimension names to values or slices. + method: Optional method to use for inexact matches. Cannot be used when + ``indexers` contain slices. Default is `None`. + **indexers_kwargs: The keyword arguments form of ``indexers``. + + Returns: + A new Field with the selection applied. + + Examples: + >>> import coordax as cx + >>> import jax.numpy as jnp + >>> import numpy as np + >>> x = cx.LabeledAxis('x', np.array([10, 20])) + >>> field = cx.field(jnp.arange(6).reshape(2, 3), x, 'y') + >>> field.sel(x=20) + + >>> field.sel(x=slice(10, 20)) + + """ + sel_indexers = coordinate_systems.normalize_indexers( + indexers, **indexers_kwargs + ) + if not sel_indexers: + return self + + unpacked_indexers, unpacked_c = ( + coordinate_systems.unpack_and_validate_indexers(sel_indexers) + ) + mapped_indexers, consumed = self.coordinate.map_indexers( + unpacked_indexers, method=method + ) + + final_consumed = set() + for c in consumed: + if c in unpacked_c: + final_consumed.add(unpacked_c[c]) + else: + final_consumed.add(c) + + unused_sel_indexers = set(sel_indexers.keys()) - final_consumed + if unused_sel_indexers: + raise ValueError( + f'Indexers {unused_sel_indexers} were not processed by any of the ' + f'coordinates in {self.coordinate}' + ) + + return self.isel(mapped_indexers) + def __repr__(self): if _in_treescope_abbreviation_mode(): return treescope.render_to_text(self) @@ -1236,9 +1379,9 @@ def get_coordinate_part( c = field_or_coord.coordinate if is_field(field_or_coord) else field_or_coord dim_to_axes = {d: ax for d, ax in zip(c.dims, c.axes)} - return coordinate_systems.compose(*[ - d if coordinate_systems.is_coord(d) else dim_to_axes[d] for d in dims - ]) + return coordinate_systems.compose( + *[d if coordinate_systems.is_coord(d) else dim_to_axes[d] for d in dims] + ) PyTree = Any @@ -1300,10 +1443,12 @@ def untag( :meth:`coordax.Field.untag` """ if allow_missing: + def untag_fn(x): if not is_field(x) or not any([contains_dims(x, d) for d in dims]): return x return x.untag(*[d for d in dims if contains_dims(x, d)]) + else: untag_fn = lambda x: x.untag(*dims) if is_field(x) else x return jax.tree.map(untag_fn, tree, is_leaf=is_field) diff --git a/coordax/fields_test.py b/coordax/fields_test.py index 5ecc084..b9eed2a 100644 --- a/coordax/fields_test.py +++ b/coordax/fields_test.py @@ -27,6 +27,7 @@ import jax import jax.numpy as jnp import numpy as np +import pytest class FieldTest(parameterized.TestCase): @@ -135,7 +136,7 @@ def test_field_constructor_invalid(self): def test_field_coordinate_property(self): x = coordax.LabeledAxis('x', np.arange(2)) field = coordax.field(np.zeros((2, 3)), x, 'y') - expected_coord = coordax.compose_coordinates(x, coordax.DummyAxis('y', 3)) + expected_coord = coordax.coords.compose(x, coordax.DummyAxis('y', 3)) self.assertEqual(field.coordinate, expected_coord) def test_field_treedef_independent_of_tag_order(self): @@ -215,13 +216,13 @@ def test_field_repr(self): testcase_name='product_coord_&_product_coord', array=np.arange(2 * 3).reshape((2, 3)), tags=( - coordax.compose_coordinates( + coordax.coords.compose( coordax.SizedAxis('x', 2), coordax.SizedAxis('y', 3), ), ), untags=( - coordax.compose_coordinates( + coordax.coords.compose( coordax.SizedAxis('x', 2), coordax.SizedAxis('y', 3), ), @@ -232,7 +233,7 @@ def test_field_repr(self): testcase_name='product_coord_&_names', array=np.arange(2 * 3).reshape((2, 3)), tags=( - coordax.compose_coordinates( + coordax.coords.compose( coordax.SizedAxis('x', 2), coordax.SizedAxis('y', 3), ), @@ -317,7 +318,7 @@ def test_broadcast_like(self): x = coordax.LabeledAxis('x', np.linspace(0, 1, 4)) y = coordax.LabeledAxis('y', np.linspace(5, 10, 5)) z = coordax.LabeledAxis('z', np.linspace(0, np.pi, 7)) - yxz = coordax.compose_coordinates(y, x, z) + yxz = coordax.coords.compose(y, x, z) field = coordax.field(np.arange(4), x) other = coordax.field(np.ones((5, 4, 7)), yxz) expected_data = np.tile(np.arange(4)[np.newaxis, :, np.newaxis], (5, 1, 7)) @@ -339,7 +340,7 @@ def test_broadcast_like_invalid_coords(self): re.escape( 'cannot broadcast field because axes corresponding to dimension ' f"'x' do not match: {x} vs {x_mismatch}" - ) + ), ): field.broadcast_like(other) @@ -347,7 +348,7 @@ def test_broadcast_to_coordinate(self): x, y = coordax.SizedAxis('x', 4), coordax.SizedAxis('y', 5) z = coordax.LabeledAxis('z', np.linspace(0, np.pi, 7)) field = coordax.field(np.arange(4), x) - yxz = coordax.compose_coordinates(y, x, z) + yxz = coordax.coords.compose(y, x, z) expected_data = np.tile(np.arange(4)[np.newaxis, :, np.newaxis], (5, 1, 7)) actual = field.broadcast_like(yxz) expected = coordax.field(expected_data, yxz) @@ -678,20 +679,16 @@ def is_duck_identity(x): x = coordax.LabeledAxis('x', np.array([np.e, np.pi])) y = coordax.LabeledAxis('y', np.linspace(0, 1, 2)) - other = coordax.field( - Duck(a=jnp.zeros((2, 2)), b=jnp.zeros((2, 2))), x, y - ) + other = coordax.field(Duck(a=jnp.zeros((2, 2)), b=jnp.zeros((2, 2))), x, y) actual = field.tag(x).broadcast_like(other) expected = coordax.field( - Duck(a=jnp.array([[1, 1], [2, 2]]), b=jnp.array([[3, 3], [4, 4]])), - x, y + Duck(a=jnp.array([[1, 1], [2, 2]]), b=jnp.array([[3, 3], [4, 4]])), x, y ) testing.assert_fields_equal(actual, expected) actual = field.tag(y).broadcast_like(other) expected = coordax.field( - Duck(a=jnp.array([[1, 2], [1, 2]]), b=jnp.array([[3, 4], [3, 4]])), - x, y + Duck(a=jnp.array([[1, 2], [1, 2]]), b=jnp.array([[3, 4], [3, 4]])), x, y ) testing.assert_fields_equal(actual, expected) @@ -707,6 +704,7 @@ def mapped_fun(*args) -> jax.Array: leaves = [x.a if isinstance(x, NonPytree) else x for x in leaves] args = jax.tree.unflatten(argdef, leaves) return jax.vmap(fun, in_axes, out_axes, **kwargs)(*args) + return mapped_fun def foo(x, y): @@ -733,21 +731,21 @@ def test_get_coordinate(self): ) with self.subTest('default'): actual = coordax.get_coordinate(field) - expected = coordax.compose_coordinates( + expected = coordax.coords.compose( *[axes[d] for d in ['x', 'y', 'z']] ) self.assertEqual(actual, expected) with self.subTest('with_positional_dims'): actual = coordax.get_coordinate(field.untag('y')) - expected = coordax.compose_coordinates( + expected = coordax.coords.compose( axes['x'], coordax.DummyAxis(None, 3), axes['z'] ) self.assertEqual(actual, expected) with self.subTest('with_name_only_dims'): actual = coordax.get_coordinate(field.untag('y', 'z').tag('g', 'h')) - expected = coordax.compose_coordinates( + expected = coordax.coords.compose( axes['x'], coordax.DummyAxis('g', 3), coordax.DummyAxis('h', 4), @@ -797,9 +795,7 @@ def test_untag_allow_missing(self): f1_c = coordax.field(jnp.zeros((2,)), x_axis) tree_c = {'a': f1_c, 'b': f2} untagged_c = coordax.untag(tree_c, x_axis, allow_missing=True) - testing.assert_fields_equal( - untagged_c['a'], f1_c.untag(x_axis) - ) + testing.assert_fields_equal(untagged_c['a'], f1_c.untag(x_axis)) testing.assert_fields_equal(untagged_c['b'], f2) def test_contains_dims(self): @@ -838,7 +834,7 @@ def test_get_coordinate_part(self): self.assertEqual(coordax.get_coordinate_part(field, x), x) with self.subTest('by_composite_coordinate'): - xy = coordax.compose_coordinates(x, y) + xy = coordax.coords.compose(x, y) self.assertEqual(coordax.get_coordinate_part(field, xy), xy) with self.subTest('dims_not_in_inputs_raises'): @@ -851,6 +847,195 @@ def test_get_coordinate_part(self): with self.assertRaisesRegex(ValueError, 'is not a part of'): coordax.get_coordinate_part(field, z) + def test_isel(self): + y = coordax.SizedAxis('y', 3) + data = jnp.arange(6).reshape(2, 3) + field = coordax.field(data, 'x', y) + + with self.subTest('single_value'): + f_x0 = field.isel(x=0) + self.assertEqual(f_x0.dims, ('y',)) + self.assertEqual(f_x0.shape, (3,)) + np.testing.assert_array_equal(f_x0.data, data[0]) + + with self.subTest('negative_indices'): + f_xm2 = field.isel(x=-2) + self.assertEqual(f_xm2.dims, ('y',)) + self.assertEqual(f_xm2.shape, (3,)) + np.testing.assert_array_equal(f_xm2.data, data[-2]) + + with self.subTest('slice_one_axis'): + f_slice = field.isel(y=slice(0, 2)) + self.assertEqual(f_slice.dims, ('x', 'y')) + self.assertEqual(f_slice.shape, (2, 2)) + np.testing.assert_array_equal(f_slice.data, data[:, 0:2]) + + with self.subTest('multiple_axes'): + f_mixed = field.isel(x=1, y=slice(3, 0, -1)) # reverse slice. + self.assertEqual(f_mixed.dims, ('y',)) + self.assertEqual(f_mixed.shape, (2,)) + np.testing.assert_array_equal(f_mixed.data, data[1, 1:3][::-1]) + + with self.subTest('axis_as_key'): + f_axis = field.isel({y: slice(0, 2)}) + self.assertEqual(f_axis.dims, ('x', 'y')) + self.assertEqual(f_axis.shape, (2, 2)) + np.testing.assert_array_equal(f_axis.data, data[:, 0:2]) + + def test_isel_raises_on_unknown_dim(self): + field = coordax.field(jnp.zeros((2, 3)), 'x', 'y') + with self.assertRaisesRegex(ValueError, 'Dimension .* not found in field'): + field.isel(z=0) + + def test_sel(self): + data = jnp.arange(6).reshape(2, 3) + x = coordax.LabeledAxis('x', np.array([10, 20])) + y = coordax.LabeledAxis('y', np.array([100, 200, 300])) + field = coordax.Field(data, dims=('x', 'y'), axes={'x': x, 'y': y}) + + with self.subTest('exact_value'): + f_val = field.sel(x=20) + self.assertEqual(f_val.dims, ('y',)) + np.testing.assert_array_equal(f_val.data, data[1]) + self.assertEqual(f_val.axes['y'], y) + with self.subTest('nearest_match'): + f_val = field.sel(x=90, method='nearest') # 20 is the nearest value. + self.assertEqual(f_val.dims, ('y',)) + np.testing.assert_array_equal(f_val.data, data[1]) + self.assertEqual(f_val.axes['y'], y) + with self.subTest('nearest_match_multiple'): + f_val = field.sel(y=[105, 290], method='nearest') # 100, 300 + self.assertEqual(f_val.dims, ('x', 'y')) + np.testing.assert_array_equal(f_val.data, data[:, [0, 2]]) + expected_new_y = coordax.LabeledAxis('y', np.array([100, 300])) + self.assertEqual(f_val.axes['y'], expected_new_y) + with self.subTest('slice'): + f_slice = field.sel(y=slice(100, 200)) + self.assertEqual(f_slice.dims, ('x', 'y')) + self.assertEqual(f_slice.shape, (2, 2)) + np.testing.assert_array_equal(f_slice.data, data[:, 0:2]) + expected_y_ticks = np.array([100, 200]) + np.testing.assert_array_equal(f_slice.axes['y'].ticks, expected_y_ticks) + with self.subTest('using_axes'): + y_sel = coordax.LabeledAxis('y', np.array([100, 200])) + f_axis = field.sel({y: y_sel}) + self.assertEqual(f_axis.dims, ('x', 'y')) + np.testing.assert_array_equal(f_axis.data, data[:, :2]) + self.assertEqual(f_axis.axes['y'], y_sel) + with self.subTest('repeated_indices'): + f_sel = field.sel({y: [100, 100]}) + self.assertEqual(f_sel.dims, ('x', 'y')) + np.testing.assert_array_equal(f_sel.data, data[:, [0, 0]]) + expected_y_ticks = np.array([100, 100]) + np.testing.assert_array_equal(f_sel.axes['y'].ticks, expected_y_ticks) + with self.subTest('preserves_sel_order'): + f_sel = field.sel({y: [200, 100]}) + self.assertEqual(f_sel.dims, ('x', 'y')) + np.testing.assert_array_equal(f_sel.data, data[:, [1, 0]]) + expected_y_ticks = np.array([200, 100]) + np.testing.assert_array_equal(f_sel.axes['y'].ticks, expected_y_ticks) + + def test_sel_supports_multidim_coords_in_indexers(self): + x = coordax.LabeledAxis('x', np.arange(10)) + p = coordax.LabeledAxis('p', np.array([100, 200, 300, 500, 600])) + p_to_sel = coordax.LabeledAxis('p', np.array([100, 300, 600])) + px = coordax.coords.compose(p, x) + field = coordax.field(jnp.zeros((5, 10)), px) + px_sel = coordax.coords.compose(p_to_sel, x) + f_sel = field.sel({px: px_sel}) + self.assertEqual(f_sel.dims, ('p', 'x')) + self.assertEqual(f_sel.axes['x'], x) + self.assertEqual(f_sel.axes['p'], p_to_sel) + + def test_sel_raises_on_unused_indexer(self): + x = coordax.LabeledAxis('x', np.array([10, 20])) + y = coordax.LabeledAxis('y', np.array([100, 200, 300])) + field = coordax.field(jnp.zeros((2, 3)), x, y) + with self.assertRaisesRegex( + ValueError, + re.escape("Indexers {'z'} were not processed") + ): + field.sel(z=slice(0, 20), x=10) + + def test_isel_matches_xarray(self): + pytest.importorskip('xarray') + rng = np.random.RandomState(seed=0) + data = rng.randn(2, 3, 4) + x = coordax.LabeledAxis('x', np.arange(2)) + y = coordax.LabeledAxis('y', np.arange(3)) + z = coordax.LabeledAxis('z', np.arange(4)) + + field = coordax.field(data, x, y, z) + da = field.to_xarray() + + with self.subTest('with_indices'): + isel_args = {'x': 0, 'y': 2} + actual = field.isel(isel_args) + expected = coordax.from_xarray(da.isel(isel_args)) + testing.assert_fields_equal(actual, expected) + with self.subTest('with_mixed_indexers'): + isel_args = {'x': 0, 'y': slice(1, 2), 'z': 0} + actual = field.isel(isel_args) + expected = coordax.from_xarray(da.isel(isel_args)) + testing.assert_fields_equal(actual, expected) + with self.subTest('with_repeated_and_unordered_indices'): + isel_args = {'x': -1, 'y': [1, 0], 'z': [1, 1, 2]} + actual = field.isel(isel_args) + expected = coordax.from_xarray(da.isel(isel_args)) + testing.assert_fields_equal(actual, expected) + + def test_sel_same_as_xarray(self): + pytest.importorskip('xarray') + data = np.arange(20).reshape(2, 10) + x = coordax.LabeledAxis('x', np.arange(2)) + y = coordax.LabeledAxis('y', np.arange(10)) + + field = coordax.field(data, x, y) + da = field.to_xarray() + + with self.subTest('with_indices'): + sel_args = {'y': 8, 'x': 0} + actual = field.sel(sel_args) + expected = coordax.from_xarray(da.sel(sel_args)) + testing.assert_fields_equal(actual, expected) + + with self.subTest('with_multiple_indices'): + sel_args = {'y': [0, 4, 6, 8], 'x': 0} + actual = field.sel(sel_args) + expected = coordax.from_xarray(da.sel(sel_args)) + testing.assert_fields_equal(actual, expected) + + with self.subTest('with_slice'): + isel_args = {'x': 0, 'y': slice(1, 2)} + actual = field.isel(isel_args) + expected = coordax.from_xarray(da.isel(isel_args)) + testing.assert_fields_equal(actual, expected) + + with self.subTest('with_unique_nearest_match'): + sel_args = {'x': 0, 'y': [1.2, 5.05]} + actual = field.sel(sel_args, method='nearest') + expected = coordax.from_xarray(da.sel(sel_args, method='nearest')) + testing.assert_fields_equal(actual, expected) + + with self.subTest('with_non_unique_nearest_match'): + sel_args = {'y': [1.2, 1.4]} + actual = field.sel(sel_args, method='nearest') + expected = coordax.from_xarray(da.sel(sel_args, method='nearest')) + testing.assert_fields_equal(actual, expected) + + def test_isel_preserves_positional_shape(self): + data = jnp.zeros((2, 3, 4)) + field = coordax.field(data, 'x', None, 'y') + # x: 0, None: 1, y: 2 + + f_sel = field.isel(x=0) + self.assertEqual(f_sel.dims, (None, 'y')) + self.assertEqual(f_sel.shape, (3, 4)) + + f_sel_y = field.isel(y=slice(0, 2)) + self.assertEqual(f_sel_y.dims, ('x', None, 'y')) + self.assertEqual(f_sel_y.shape, (2, 3, 2)) + def test_deprecated_tmp_axis_name(self): with self.assertWarnsRegex( DeprecationWarning, diff --git a/pyproject.toml b/pyproject.toml index 665d89e..6950da7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ packages = ["coordax"] [project] name = "coordax" -version = "0.2.5" # keep sync with __init__.py +version = "0.2.6" # keep sync with __init__.py description = "Coordinate axes for scientific computing in JAX" authors = [ {name = "Google LLC", email = "noreply@google.com"},