Repository: xgcm/xgcm
Pull request #351 (2021-07-07 20:54:53)

Allow for overwriting metrics in set_metrics method
 - [X] Include overwrite kwarg in set_metrics
 - [X] Tests added in test_metrics/test_set_metric_overwrite_true and test_set_metric_overwrite_false




Patch:
diff --git a/.github/workflows/stale.yaml b/.github/workflows/stale.yaml
index 1ff499b..1e0ee71 100644
--- a/.github/workflows/stale.yaml
+++ b/.github/workflows/stale.yaml
@@ -10,7 +10,7 @@ jobs:
     - uses: actions/stale@v3
       with:
         repo-token: ${{ secrets.GITHUB_TOKEN }}
-        exempt-issue-labels: 'keepOpen,In Progress'
+        exempt-issue-labels: 'keepOpen,In Progress,bug'
         exempt-pr-labels: 'keepOpen,In Progress'
         stale-issue-message: "This issue has been marked 'stale' due to lack of recent activity. If there is no further activity, the issue will be closed in another 30 days. Thank you for your contribution!"
         close-issue-message: "This issue has been closed due to inactivity. If you feel this is in error, please reopen the issue or file a new issue with the relevant details."
diff --git a/.gitignore b/.gitignore
index 34094a2..b657e9c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -74,5 +74,6 @@ dask-worker-space/dask-worker-space/purge.lock
 .mypy_cache
 xgcm/.vscode/settings.json
 02_Testing xgcm metrics.ipynb
+Testing xgcm.ipynb
 .ipynb_checkpoints/02_Testing xgcm metrics-checkpoint.ipynb
 xgcm/test/sample.py
diff --git a/README.rst b/README.rst
index 9136c7d..fbd9182 100644
--- a/README.rst
+++ b/README.rst
@@ -50,7 +50,7 @@ To learn how to install and use xgcm for your dataset, visit the `xgcm documenta
 .. _Pangeo: http://pangeo.io
 .. _dask: http://dask.pydata.org
 .. _xarray: http://xarray.pydata.org
-.. _Arakawa Grids: https://en.wikipedia.org/wiki/Arakawa_grid
+.. _Arakawa Grids: https://en.wikipedia.org/wiki/Arakawa_grids
 .. _xgcm documentation: https://xgcm.readthedocs.io/
 
 .. |conda forge| image:: https://img.shields.io/conda/vn/conda-forge/xgcm
diff --git a/xgcm/grid.py b/xgcm/grid.py
index 45ed989..b4e446e 100644
--- a/xgcm/grid.py
+++ b/xgcm/grid.py
@@ -1259,28 +1259,49 @@ class Grid:
             self.axes[axis]._facedim = facedim
             self.axes[axis]._connections = axis_links
 
-    def set_metrics(self, key, value):
+    def set_metrics(self, key, value, overwrite=False):
         metric_axes = frozenset(_maybe_promote_str_to_list(key))
         axes_not_found = [ma for ma in metric_axes if ma not in self.axes]
-
         if len(axes_not_found) > 0:
             raise KeyError(
                 f"Metric axes {axes_not_found!r} not compatible with grid axes {tuple(self.axes)!r}"
             )
 
-        # initialize empty list
-        self._metrics[metric_axes] = []
-        for metric_varname in _maybe_promote_str_to_list(value):
-            if metric_varname not in self._ds:
+        metric_value = _maybe_promote_str_to_list(value)
+        for metric_varname in metric_value:
+            if metric_varname not in self._ds.variables:
                 raise KeyError(
                     f"Metric variable {metric_varname} not found in dataset."
                 )
-            # resetting coords avoids potential broadcasting / alignment issues
-            metric_var = self._ds[metric_varname].reset_coords(drop=True)
 
-            # TODO: check for consistency of metric_var dims with axis dims
-            # check for duplicate dimensions among each axis metric
-            self._metrics[metric_axes].append(metric_var)
+        existing_metric_axes = set(self._metrics.keys())
+        if metric_axes in existing_metric_axes:
+            value_exist = self._metrics.get(metric_axes)
+            # resetting coords avoids potential broadcasting / alignment issues
+            value_new = self._ds[metric_varname].reset_coords(drop=True)
+            did_overwrite = False
+            # go through each existing value until data array with matching dimensions is selected
+            for idx, ve in enumerate(value_exist):
+                # double check if dimensions match
+                if set(value_new.dims) == set(ve.dims):
+                    if overwrite:
+                        # replace existing data array with new data array input
+                        self._metrics[metric_axes][idx] = value_new
+                        did_overwrite = True
+                    else:
+                        raise ValueError(
+                            f"Metric variable {ve.name} with dimensions {ve.dims} already assigned in metrics."
+                            f" Overwrite {ve.name} with {metric_varname} by setting overwrite=True."
+                        )
+            # if no existing value matches new value dimension-wise, just append new value
+            if not did_overwrite:
+                self._metrics[metric_axes].append(value_new)
+        else:
+            # no existing metrics for metric_axes yet; initialize empty list
+            self._metrics[metric_axes] = []
+            for metric_varname in metric_value:
+                metric_var = self._ds[metric_varname].reset_coords(drop=True)
+                self._metrics[metric_axes].append(metric_var)
 
     def _get_dims_from_axis(self, da, axis):
         dim = []
@@ -1424,15 +1445,6 @@ class Grid:
         )
         return array
 
-    def _interp_metric(self, da, axes):
-        # include boundary and fill_value as inputs
-        metric_available = self._metrics.get(frozenset(axes), None)
-        if metric_available is not None:
-            # this function works with only one metric at a time
-            metric = metric_available[0]
-            metric_interp = self.interp_like(metric, da)
-        return metric_interp
-
     def __repr__(self):
         summary = ["<xgcm.Grid>"]
         for name, axis in self.axes.items():


Test patch:
diff --git a/xgcm/test/test_metrics.py b/xgcm/test/test_metrics.py
index 01d27bd..c6e5636 100644
--- a/xgcm/test/test_metrics.py
+++ b/xgcm/test/test_metrics.py
@@ -233,7 +233,6 @@ def test_get_metric_with_conditions_02a():
 
     expected_metric = grid.interp(ds["area_e"], ("X", "Y"))
 
-    xr.testing.assert_equal(get_metric, expected_metric)
     xr.testing.assert_allclose(get_metric, expected_metric)
 
 
@@ -251,7 +250,6 @@ def test_get_metric_with_conditions_02b():
 
     expected_metric = grid.interp(ds["area_e"], ("X", "Y"))
 
-    xr.testing.assert_equal(get_metric, expected_metric)
     xr.testing.assert_allclose(get_metric, expected_metric)
 
 
@@ -268,7 +266,6 @@ def test_get_metric_with_conditions_04a():
     interp_metric = grid.interp(ds.dx_t, "Y")
     expected_metric = (interp_metric * ds.dy_n).reset_coords(drop=True)
 
-    xr.testing.assert_equal(get_metric, expected_metric)
     xr.testing.assert_allclose(get_metric, expected_metric)
 
 
@@ -286,7 +283,6 @@ def test_get_metric_with_conditions_04b():
     interp_metric_2 = grid.interp(ds.dy_t, "Y")
     expected_metric = (interp_metric_1 * interp_metric_2).reset_coords(drop=True)
 
-    xr.testing.assert_equal(get_metric, expected_metric)
     xr.testing.assert_allclose(get_metric, expected_metric)
 
 
@@ -311,28 +307,92 @@ def test_set_metric():
         assert k in grid_manual._metrics.keys()
 
         for metric_expected, metric in zip(v, grid_manual._metrics[k]):
-            xr.testing.assert_equal(metric_expected.reset_coords(drop=True), metric)
+            xr.testing.assert_allclose(metric_expected.reset_coords(drop=True), metric)
 
         for metric_expected, metric in zip(v, grid._metrics[k]):
-            xr.testing.assert_equal(metric_expected.reset_coords(drop=True), metric)
+            xr.testing.assert_allclose(metric_expected.reset_coords(drop=True), metric)
 
 
-@pytest.mark.skip(reason="For the next PR")
 @pytest.mark.parametrize(
-    "metric_axes,metric_name",
+    "metric_axes, exist_metric_varname, add_metric_varname, expected_varname",
     [
-        ("X", "dx_t"),
-        ("Y", "dy_ne"),
+        (
+            "X",
+            ["dx_t", "dx_n", "dx_e", "dx_ne"],
+            ["dx_n_overwrite"],
+            ["dx_t", "dx_n_overwrite", "dx_e", "dx_ne"],
+        ),
+        (
+            ("Y", "X"),
+            ["area_t", "area_n", "area_e", "area_ne"],
+            ["area_n_overwrite"],
+            ["area_t", "area_n_overwrite", "area_e", "area_ne"],
+        ),
+        # overwrite 1 existing metric, append 1 new metric
+        (
+            "X",
+            ["dx_t", "dx_n", "dx_e"],
+            ["dx_n_overwrite", "dx_ne"],
+            ["dx_t", "dx_n_overwrite", "dx_e", "dx_ne"],
+        ),
     ],
 )
-def test_interp_metrics(metric_axes, metric_name):
-    # need to test with boundary and fill_value conditions
+def test_set_metric_overwrite_true(
+    metric_axes, exist_metric_varname, add_metric_varname, expected_varname
+):
     ds, coords, metrics = datasets_grid_metric("C")
-    grid = Grid(ds, coords=coords)
-    grid.set_metrics(metric_axes, metric_name)
-    interp_metric = grid._interp_metric(ds.u, metric_axes)
 
-    expected_metric = grid.interp(ds[metric_name], metric_axes)
+    ds = ds.assign_coords({add_metric_varname[0]: ds[exist_metric_varname[1]] * 10})
+
+    metrics = {
+        k: [m for m in v if m in exist_metric_varname] for k, v in metrics.items()
+    }
+    grid = Grid(ds, coords=coords, metrics=metrics)
+
+    for av in add_metric_varname:
+        grid.set_metrics(metric_axes, av, overwrite=True)
+
+    key = frozenset(list(metric_axes))
+    set_metric = grid._metrics.get(key)
+
+    expected_metric = []
+    for ev in expected_varname:
+        metric_var = ds[ev].reset_coords(drop=True)
+        expected_metric.append(metric_var)
+
+    assert len(set_metric) == len(expected_metric)
+
+    for i in range(len(set_metric)):
+        assert set_metric[i].equals(expected_metric[i])
 
-    if interp_metric.equals(expected_metric) is False:
-        xr.testing.assert_allclose(interp_metric, expected_metric)
+
+@pytest.mark.parametrize(
+    "metric_axes,overwrite_metric,add_metric",
+    [("X", "dx_t_overwrite", "dx_t"), ("X", "dx_e", None)],
+)
+def test_set_metric_value_errors(metric_axes, overwrite_metric, add_metric):
+    ds, coords, metrics = datasets_grid_metric("C")
+
+    if add_metric is not None:
+        ds = ds.assign_coords({overwrite_metric: ds[add_metric] * 10})
+
+    grid = Grid(ds, coords=coords, metrics=metrics)
+
+    with pytest.raises(ValueError, match="setting overwrite=True."):
+        grid.set_metrics(metric_axes, overwrite_metric)
+
+
+@pytest.mark.parametrize(
+    "metric_axes,add_metric",
+    [("X", "foo"), (("U", "V"), "area_n")],
+)
+def test_set_metric_key_errors(metric_axes, add_metric):
+    ds, coords, metrics = datasets_grid_metric("C")
+    grid = Grid(ds, coords=coords, metrics=metrics)
+
+    if len(metric_axes) == 1:
+        with pytest.raises(KeyError, match="not found in dataset."):
+            grid.set_metrics(metric_axes, add_metric)
+    else:
+        with pytest.raises(KeyError, match="not compatible with grid axes"):
+            grid.set_metrics(metric_axes, add_metric)