From 0832141d8b6967570a1597851fbb2c7f4cab994a Mon Sep 17 00:00:00 2001 From: Dmitrii Kochkov Date: Thu, 26 Feb 2026 07:58:25 -0800 Subject: [PATCH] Updated SelectedAxis to only contain coordinate fields that are associated with the sliced axis. This resolves an unnecessary conflict when a coordinate slice is combined with another coordinate that has a field with the name corresponding to other coordinate types. To ensure that multidimensional coordinate fields are not lost due to the order of appearance, CartesiaProduct canonicalize sorted components in fields property. PiperOrigin-RevId: 875724189 --- coordax/__init__.py | 2 +- coordax/coordinate_systems.py | 15 +++++++++-- coordax/xarray_test.py | 50 ++++++++++++++++++++++++++++++++++- pyproject.toml | 2 +- 4 files changed, 64 insertions(+), 5 deletions(-) diff --git a/coordax/__init__.py b/coordax/__init__.py index 4c4d64f..23c7fe7 100644 --- a/coordax/__init__.py +++ b/coordax/__init__.py @@ -56,4 +56,4 @@ ) import coordax.testing # pylint: disable=unused-import -__version__ = '0.2.6' # keep sync with pyproject.toml +__version__ = '0.2.7' # keep sync with pyproject.toml diff --git a/coordax/coordinate_systems.py b/coordax/coordinate_systems.py index ef37243..9fafc99 100644 --- a/coordax/coordinate_systems.py +++ b/coordax/coordinate_systems.py @@ -465,7 +465,10 @@ def shape(self) -> tuple[int, ...]: @property def fields(self) -> dict[str, 'fields.Field']: """A maps from field names to their values.""" - return self.coordinate.fields + coord_fields = self.coordinate.fields + return { + k: v for k, v in coord_fields.items() if set(self.dims).issubset(v.dims) + } def __repr__(self): return f'coordax.SelectedAxis({self.coordinate!r}, axis={self.axis})' @@ -618,7 +621,15 @@ def shape(self) -> tuple[int, ...]: @property def fields(self) -> dict[str, 'fields.Field']: """Returns a mapping from field names to their values.""" - return _merge_dicts(c.fields for c in self.coordinates) + # To extract all valid fields from multidimensional coordinates, we need to + # ensure that SelectedAxes are merged back into the original coordinates. + # To ensure that merging happens we reorder SelectedAxis by name and type. + selected = [x for x in self.coordinates if isinstance(x, SelectedAxis)] + key_by_class_and_axis = lambda x: (x.coordinate.__class__.__name__, x.axis) + ordered_selected = sorted(selected, key=key_by_class_and_axis) + non_selected = [x for x in self.coordinates if x not in selected] + reordered_canonized = canonicalize(*ordered_selected, *non_selected) + return _merge_dicts(c.fields for c in reordered_canonized) @property def axes(self) -> tuple[Coordinate, ...]: diff --git a/coordax/xarray_test.py b/coordax/xarray_test.py index 5a16b65..3cee7a0 100644 --- a/coordax/xarray_test.py +++ b/coordax/xarray_test.py @@ -108,6 +108,53 @@ def test_field_to_data_array_custom_coord(self): ) xarray.testing.assert_identical(actual, expected) + def test_field_to_data_array_custom_coord_selected_axis(self): + data = np.arange(2 * 2).reshape((2, 2)) + custom_coord = AdhocCoordinate( + dims=('x', 'y'), + shape=(2, 3), + fields=lambda c: { + 'custom': coordax.field(np.zeros(c.shape), c), + 'x': coordax.field(np.ones(c.shape[0]), coordax.SelectedAxis(c, 0)), + 'y': coordax.field(np.ones(c.shape[1]), coordax.SelectedAxis(c, 1)), + }, + ) + field = coordax.field(data, coordax.SelectedAxis(custom_coord, 0), 'z') + # Verify that coord_fields that are not in selected slice are excluded. + actual = field.to_xarray() + expected = xarray.DataArray( + data=np.arange(2 * 2).reshape((2, 2)), + dims=['x', 'z'], + coords={'x': (('x',), np.ones(2))}, + ) + xarray.testing.assert_identical(actual, expected) + + def test_field_to_data_array_disjoint_slice_combine_fields(self): + data = np.arange(3 * 5 * 2).reshape((3, 5, 2)) + custom_coord = AdhocCoordinate( + dims=('x', 'y'), + shape=(2, 3), + fields=lambda c: { + 'custom': coordax.field(np.zeros(c.shape), c), + 'x': coordax.field(np.ones(c.shape[0]), coordax.SelectedAxis(c, 0)), + 'y': coordax.field(np.ones(c.shape[1]), coordax.SelectedAxis(c, 1)), + }, + ) + x, y = custom_coord.axes + field = coordax.field(data, y, coordax.SizedAxis('z', 5), x) + # Both slices presernt, so all coord_fields should be included. + actual = field.to_xarray() + expected = xarray.DataArray( + data=np.arange(3 * 5 * 2).reshape((3, 5, 2)), + dims=['y', 'z', 'x'], + coords={ + 'x': (('x',), np.ones(2)), + 'y': (('y',), np.ones(3)), + 'custom': (('x', 'y'), np.zeros((2, 3))), + }, + ) + xarray.testing.assert_identical(actual, expected) + def test_field_to_data_array_missing_dimension_names(self): data = np.arange(2 * 3).reshape((2, 3)) field = coordax.field(data) @@ -240,7 +287,8 @@ def test_field_from_xarray_deprecated(self): ) with self.assertWarnsRegex( DeprecationWarning, - r'cx\.Field\.from_xarray\(\) is deprecated, use cx\.from_xarray\(\) instead', + r'cx\.Field\.from_xarray\(\) is deprecated, use cx\.from_xarray\(\)' + r' instead', ): actual = coordax.Field.from_xarray(data_array) expected = coordax.field(data, 'x', coordax.LabeledAxis('y', [1, 2, 3])) diff --git a/pyproject.toml b/pyproject.toml index 6950da7..f9de79e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ packages = ["coordax"] [project] name = "coordax" -version = "0.2.6" # keep sync with __init__.py +version = "0.2.7" # keep sync with __init__.py description = "Coordinate axes for scientific computing in JAX" authors = [ {name = "Google LLC", email = "noreply@google.com"},