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"},