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.

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

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, ...]:

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]))

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