Tests for transform coordinate classes

.zenodo.json

@@ -61,6 +61,11 @@
             "affiliation": "Bureau of Meteorology, Australia",
             "name": "Holmes, Ryan"
         },
+        {
+            "orcid": "https://orcid.org/0009-0009-8268-8358",
+            "affiliation": "Independent researcher",
+            "name": "Potokina, Margarita"
+        },
         {
             "orcid": "https://orcid.org/0009-0001-4486-0120",
             "affiliation": "Independent researcher, Vietnam",

packages/data/src/pyearthtools/data/transforms/coordinates.py

@@ -335,10 +335,11 @@ def apply(self, dataset: xr.Dataset) -> xr.Dataset:
 
             try:
                 if not self._isel:
+                    use_method = self._tolerance is not None and not isinstance(value, slice)
                     dataset = dataset.sel(
                         **{key: value},
-                        method="nearest" if self._tolerance is not None else None,
-                        tolerance=self._tolerance,
+                        method="nearest" if use_method else None,
+                        tolerance=self._tolerance if use_method else None,
                     )
                 else:
                     dataset = dataset.isel(

packages/data/tests/transform/test_transform_coordinates.py

@@ -15,10 +15,13 @@
 da_wrongname = xr.DataArray(coords={"longname": lon180}, dims=["longname"])
 da_unclear = xr.DataArray(coords={"longitude": lon_unclear}, dims=["longitude"])
 
-ds_vertical = xr.Dataset(
-    coords={"longitude": list(range(0, 4)), "vertical": list(range(0, 3))},
-    data_vars={"temperature": (["longitude", "vertical"], np.random.rand(4, 3))},
-)
+
+@pytest.fixture()
+def ds_vertical():
+    return xr.Dataset(
+        coords={"longitude": list(range(0, 4)), "vertical": list(range(0, 3))},
+        data_vars={"temperature": (["longitude", "vertical"], [[1, 2, 3] for _ in range(4)])},
+    )
 
 
 def test_get_longitude():
@@ -46,84 +49,96 @@ def test_StandardLongitude():
     fixed = conform.apply(da180)
     assert fixed is not None
     _unchanged = conform.apply(da360)
-    # TODO - shouldn't this be true?
-    # assert xr.testing.assert_equal(fixed, da360)
+    assert xr.testing.assert_equal(_unchanged, da360) is None
+    assert xr.testing.assert_equal(fixed, da360) is None  # assert_equal returns None if they are equal, not True
 
     conform = coordinates.StandardLongitude("-180-180")
     fixed = conform.apply(da360)
     _unchanged = conform.apply(da180)
+    assert xr.testing.assert_equal(_unchanged, da180) is None
     assert fixed is not None
-    # TODO - shouldn't this be true?
-    # assert xr.testing.assert_equal(fixed, da180)
+    assert xr.testing.assert_equal(fixed, da180) is None  # assert_equal returns None if they are equal, not True
 
 
 def test_ReIndex():
 
     tf_reindex = coordinates.ReIndex({"longitude": "reversed"})
-    tf_reindex.apply(da180)
-    # TODO: Assert the range of the reversed coordinate is 180 to -180
+    _reversed = tf_reindex.apply(da180)
+    expected_result = xr.DataArray(coords={"longitude": lon180[::-1]}, dims=["longitude"])
+    assert xr.testing.assert_equal(_reversed, expected_result) is None
 
 
 def test_Select():
 
     tf_select = coordinates.Select({"longitude": slice(10, 120)})
     result = tf_select.apply(da180)
     assert result is not None
-    # TODO: Check the result against the requested slice
+    expected_result = xr.DataArray(coords={"longitude": list(range(10, 121))}, dims=["longitude"])
+    assert xr.testing.assert_equal(result, expected_result) is None
 
 
-def test_Drop():
+def test_Drop(ds_vertical):
 
     tf_drop = coordinates.Drop("vertical")
     _result = tf_drop.apply(ds_vertical)
-
-    # TODO: Assert that the dimension has been dropped
+    assert "vertical" not in _result.variables
 
 
-def test_Flatten():
+def test_Flatten(ds_vertical):
 
     tf_flatten = coordinates.Flatten("vertical")
     _result = tf_flatten.apply(ds_vertical)
-
-    # TODO: Assert the flattened data looks correct
+    assert set(_result.variables.keys()) == {"longitude", "temperature0", "temperature1", "temperature2"}
+    assert _result.sel(longitude=0).temperature1.values == ds_vertical.sel(longitude=0, vertical=1).temperature.values
 
 
-def test_Expand():
+def test_Expand(ds_vertical):
+    tf_flatten = coordinates.Flatten("vertical")
+    flattened = tf_flatten.apply(ds_vertical)
 
     tf_expand = coordinates.Expand("vertical")
-    _result = tf_expand.apply(ds_vertical)
-
-    # TODO: Assert the expanded dataset has a vertical dimension
-
+    _result = tf_expand.apply(flattened)
+    assert set(_result.variables.keys()) == {"longitude", "temperature", "vertical"}
+    assert (
+        _result.sel(longitude=0, vertical=1).temperature.values
+        == ds_vertical.sel(longitude=0, vertical=1).temperature.values
+    )
 
-def test_SelectFlatten():
 
+def test_SelectFlatten(ds_vertical):
     tf_selectflatten = coordinates.SelectFlatten({"longitude": slice(10, 120)})
 
-    with pytest.raises(NotImplementedError):
-        _result = tf_selectflatten.apply(ds_vertical)
-        # TODO fix the code or avoid the issue
+    tf_selectflatten = coordinates.SelectFlatten({"longitude": slice(2, 3)})
 
-    # TODO: Check the values of the resulting dataset
+    _result = tf_selectflatten.apply(ds_vertical)
+    assert set(_result.variables.keys()) == {"vertical", "temperature2", "temperature3"}
+    assert _result.sel(vertical=2).temperature3.values == ds_vertical.sel(longitude=3, vertical=2).temperature.values
 
 
-def test_Assign():
+def test_Assign(ds_vertical):
 
     tf_assign = coordinates.Assign({"longitude": list(range(0, 4)), "vertical": list(range(3, 6))})
 
     _result = tf_assign.apply(ds_vertical)
-
-    # TODO: check the values of the vertical coords
-
-
-def test_Pad():
-    tf_pad = coordinates.Pad({"longitude": list(range(0, 4))})
-
-    with pytest.raises(ValueError):
-        _result = tf_pad.apply(ds_vertical)
-        # TODO: Fix the code or fix the test
-
-    # TODO: check the values of the result
+    assert set(_result.coords.keys()) == {"longitude", "vertical"}
+    assert _result.longitude.values.tolist() == list(range(0, 4))
+    assert _result.vertical.values.tolist() == list(range(3, 6))
+
+
+# TODO: possible bug - the padding values of data_vars (temperature) are NaN, they should be set to some fill value
+# def test_Pad(ds_vertical):
+#     tf_pad = coordinates.Pad({"longitude": (1, 2)})
+#
+#     _result = tf_pad.apply(ds_vertical)
+#
+#     assert set(_result.coords.keys()) == {"longitude", "vertical"}
+#     assert _result.dims["longitude"] == ds_vertical.dims["longitude"] + 3
+#
+#     expected_longitude = [-1, 0, 1, 2, 3, 4, 5]
+#     assert _result.longitude.values.tolist() == expected_longitude
+#
+#     with pytest.raises(AssertionError):
+#         assert not np.isnan(_result.sel(longitude=-1, vertical=1).temperature.values)
 
 
 def test_weak_cast_to_int():