Repository: xgcm/xgcm
Pull request #350 (2021-06-28 17:26:10)

Add periodic and boundary conditions to interp_like
 - [X] Periodic condition added to test_grid/test_interp_like
 - [x] Implement changes to interp_like (include boundary conditions as input)
 - [x] Boundary conditions  and fill value added to test_grid/test_interp_like
 - [x] Include a test for metics.py/iterate_axis_combinations




Patch:
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index c9ce5ab..a925816 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -10,6 +10,8 @@ v0.6.0 (unreleased)
 
 New Features
 ~~~~~~~~~~~~
+- :py:meth:`~xgcm.grid.Grid.interp_like` enables users to interpolate arrays onto the grid positions of another array (:issue:`234` , :issue:`343`, :pull:`350`). By `Dianne Deauna <https://github.com/jdldeauna>`_ under the `SIParCS internship <https://www2.cisl.ucar.edu/siparcs-2021-projects#8>`_.
+
 - :py:meth:`~xgcm.grid.Grid.set_metrics` now enables adding metrics to a grid object (:pull:`336`, :issue:`199`).
   By `Dianne Deauna <https://github.com/jdldeauna>`_ under the `SIParCS internship <https://www2.cisl.ucar.edu/siparcs-2021-projects#8>`_.
 
diff --git a/xgcm/grid.py b/xgcm/grid.py
index 5985f97..45ed989 100644
--- a/xgcm/grid.py
+++ b/xgcm/grid.py
@@ -108,7 +108,7 @@ class Axis:
               points nearest to the edge
             This sets the default value. It can be overriden by specifying the
             boundary kwarg when calling specific methods.
-        fill_value : {float}, optional
+        fill_value : float, optional
             The value to use in the boundary condition when `boundary='fill'`.
 
         REFERENCES
@@ -1319,8 +1319,7 @@ class Grid:
         array_dims = set(array.dims)
 
         possible_metric_vars = set(tuple(k) for k in self._metrics.keys())
-        input_axes = tuple(axes)
-        possible_combos = set(itertools.permutations(input_axes))
+        possible_combos = set(itertools.permutations(tuple(axes)))
         overlap_metrics = possible_metric_vars.intersection(possible_combos)
 
         if len(overlap_metrics) > 0:
@@ -1333,10 +1332,10 @@ class Grid:
                     metric_vars = mv
                     break
             # if metric_vars is None:
-            #     # Condition 2: interpolate metric with matching axis to desired dimensions
+            #     # Condition 4: interpolate metric with matching axis to desired dimensions
             #     metric_vars = self.interp_like(mv, array)
         else:
-            # Condition 3: use provided metrics to calculate for required metric
+            # Condition 2: use provided metrics to calculate for required metric
             for axis_combinations in iterate_axis_combinations(axes):
                 try:
                     # will raise KeyError if the axis combination is not in metrics
@@ -1354,7 +1353,7 @@ class Grid:
                             metric_vars = possible_combinations
                             break
                         # else:
-                        #     # Condition 4: metrics in the wrong position (must interpolate before multiplying)
+                        #     # Condition 5: metrics in the wrong position (must interpolate before multiplying)
                         #     metric_vars = tuple(
                         #         self.interp_like(pc, array)
                         #         for pc in possible_combinations
@@ -1372,7 +1371,7 @@ class Grid:
         return metric_vars
 
     @docstrings.dedent
-    def interp_like(self, array, like):
+    def interp_like(self, array, like, boundary=None, fill_value=None):
         """Compares positions between two data arrays and interpolates array to the position of like if necessary
 
         Parameters
@@ -1381,6 +1380,21 @@ class Grid:
             DataArray to interpolate to the position of like
         like : DataArray
             DataArray with desired grid positions for source array
+        boundary : str or dict, optional,
+            boundary can either be one of {None, 'fill', 'extend', 'extrapolate'}
+            * None:  Do not apply any boundary conditions. Raise an error if
+              boundary conditions are required for the operation.
+            * 'fill':  Set values outside the array boundary to fill_value
+              (i.e. a Dirichlet boundary condition.)
+            * 'extend': Set values outside the array to the nearest array
+              value. (i.e. a limited form of Neumann boundary condition where
+              the difference at the boundary will be zero.)
+            * 'extrapolate': Set values by extrapolating linearly from the two
+              points nearest to the edge
+            This sets the default value. It can be overriden by specifying the
+            boundary kwarg when calling specific methods.
+        fill_value : float, optional
+            The value to use in the boundary condition when `boundary='fill'`.
 
         Returns
         -------
@@ -1388,21 +1402,30 @@ class Grid:
             Source data array with updated positions along axes matching with target array
         """
 
+        interp_axes = []
         for axname, axis in self.axes.items():
-            # This will raise a KeyError since this for-loop goes through all axes contained in self,
-            # but it is possible to apply the method for only 1 axis at a time
             try:
                 position_array, _ = axis._get_axis_coord(array)
                 position_like, _ = axis._get_axis_coord(like)
             # This will raise a KeyError if you have multiple axes contained in self,
             # since the for-loop will go through all axes, but the method is applied for only 1 axis at a time
+            # This is for cases where an axis is present in self that is not available for either array or like.
+            # For the axis you are interested in interpolating, there should be data for it in grid, array, and like.
             except KeyError:
                 continue
             if position_like != position_array:
-                array = self.interp(array, axname)
+                interp_axes.append(axname)
+
+        array = self.interp(
+            array,
+            interp_axes,
+            fill_value=fill_value,
+            boundary=boundary,
+        )
         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


Test patch:
diff --git a/xgcm/test/test_grid.py b/xgcm/test/test_grid.py
index 4a1f44f..a015764 100644
--- a/xgcm/test/test_grid.py
+++ b/xgcm/test/test_grid.py
@@ -784,34 +784,52 @@ def test_boundary_kwarg_same_as_grid_constructor_kwarg():
 
 @pytest.mark.parametrize(
     "metric_axes,metric_name",
-    [("X", "dx_t"), ("Y", "dy_ne"), (["Y", "X"], "dy_n")],
+    [
+        (["Y", "X"], "area_n"),
+        ("X", "dx_t"),
+        ("Y", "dy_ne"),
+        (["Y", "X"], "dy_n"),
+        (["X"], "tracer"),
+    ],
 )
-def test_interp_like(metric_axes, metric_name):
-
-    ds, coords, _ = datasets_grid_metric("C")
-    grid = Grid(ds, coords=coords)
-    grid.set_metrics(metric_axes, metric_name)
-    metric_available = grid._metrics.get(frozenset(metric_axes), None)
-    metric_available = metric_available[0]
-    interp_metric = grid.interp_like(metric_available, ds.u)
-    test_metric = grid.interp(ds[metric_name], metric_axes)
-
-    xr.testing.assert_equal(interp_metric, test_metric)
-    xr.testing.assert_allclose(interp_metric, test_metric)
-
-
+@pytest.mark.parametrize("periodic", [True, False])
 @pytest.mark.parametrize(
-    "var_name,like_name,var_axes",
+    "boundary, boundary_expected",
     [
-        ("tracer", "u", "X"),
+        ({"X": "fill", "Y": "fill"}, {"X": "fill", "Y": "fill"}),
+        ({"X": "extend", "Y": "extend"}, {"X": "extend", "Y": "extend"}),
+        (
+            {"X": "extrapolate", "Y": "extrapolate"},
+            {"X": "extrapolate", "Y": "extrapolate"},
+        ),
+        ("fill", {"X": "fill", "Y": "fill"}),
+        ("extend", {"X": "extend", "Y": "extend"}),
+        ("extrapolate", {"X": "extrapolate", "Y": "extrapolate"}),
+        ({"X": "extend", "Y": "fill"}, {"X": "extend", "Y": "fill"}),
+        ({"X": "extrapolate", "Y": "fill"}, {"X": "extrapolate", "Y": "fill"}),
+        pytest.param(
+            "fill",
+            {"X": "fill", "Y": "extend"},
+            marks=pytest.mark.xfail,
+            id="boundary not equal to boundary_expected",
+        ),
     ],
 )
-def test_interp_like_var(var_name, like_name, var_axes):
+@pytest.mark.parametrize("fill_value", [None, 0.1])
+def test_interp_like(
+    metric_axes, metric_name, periodic, boundary, boundary_expected, fill_value
+):
 
-    ds, coords, metrics = datasets_grid_metric("C")
-    grid = Grid(ds, coords=coords, metrics=metrics)
-    interp_var = grid.interp_like(ds[var_name], ds[like_name])
-    test_var = grid.interp(ds[var_name], var_axes)
+    ds, coords, _ = datasets_grid_metric("C")
+    grid = Grid(ds, coords=coords, periodic=periodic)
+    grid.set_metrics(metric_axes, metric_name)
+    metric_available = grid._metrics.get(frozenset(metric_axes), None)
+    metric_available = metric_available[0]
+    interp_metric = grid.interp_like(
+        metric_available, ds.u, boundary=boundary, fill_value=fill_value
+    )
+    expected_metric = grid.interp(
+        ds[metric_name], metric_axes, boundary=boundary_expected, fill_value=fill_value
+    )
 
-    xr.testing.assert_equal(interp_var, test_var)
-    xr.testing.assert_allclose(interp_var, test_var)
+    xr.testing.assert_allclose(interp_metric, expected_metric)
diff --git a/xgcm/test/test_metrics.py b/xgcm/test/test_metrics.py
index 01a8730..01d27bd 100644
--- a/xgcm/test/test_metrics.py
+++ b/xgcm/test/test_metrics.py
@@ -6,6 +6,7 @@ import pytest
 import xarray as xr
 
 from xgcm.grid import Grid
+from xgcm.metrics import iterate_axis_combinations
 from xgcm.test.datasets import datasets_grid_metric
 
 
@@ -96,6 +97,48 @@ def test_assign_metric(key, metric_vars):
     _ = Grid(ds, coords=coords, metrics={key: metric_vars})
 
 
+@pytest.mark.parametrize(
+    "axes, expected",
+    [
+        (
+            ("X", "Y"),
+            (
+                (frozenset({"X", "Y"}),),
+                (frozenset({"X"}), frozenset({"Y"})),
+                (frozenset({"Y"}), frozenset({"X"})),
+            ),
+        ),
+        (
+            ("X", "Y", "Z"),
+            (
+                (frozenset({"Y", "X", "Z"}),),
+                (frozenset({"Z"}), frozenset({"X"}), frozenset({"Y"})),
+                (frozenset({"X"}), frozenset({"Y"}), frozenset({"Z"})),
+                (frozenset({"Y"}), frozenset({"X"}), frozenset({"Z"})),
+                (frozenset({"Y", "Z"}), frozenset({"X"})),
+                (frozenset({"Y", "X"}), frozenset({"Z"})),
+                (frozenset({"X", "Z"}), frozenset({"Y"})),
+            ),
+        ),
+    ],
+)
+def test_iterate_axis_combinations(axes, expected):
+
+    actual = list(iterate_axis_combinations(axes))
+
+    expected = [set(e) for e in expected]
+    actual = [set(a) for a in actual]
+
+    # comparing these is a pain since output order is not guaranteed.
+    # instead we ensure actual and expected have the same number of elements
+    # and that every element in actual is present in expected and vice-versa
+    assert len(actual) == len(expected)
+    for a in actual:
+        assert set(a) in expected
+    for e in expected:
+        assert set(e) in actual
+
+
 @pytest.mark.parametrize(
     "axes, data_var, drop_vars, metric_expected_list, expected_error",
     [
@@ -144,13 +187,12 @@ def test_get_metric_with_conditions_01():
 
     expected_metric = ds["area_n"].reset_coords(drop=True)
 
-    if xr.testing.assert_equal(get_metric, expected_metric) is False:
-        xr.testing.assert_allclose(get_metric, expected_metric)
+    xr.testing.assert_allclose(get_metric, expected_metric)
 
 
 def test_get_metric_with_conditions_02():
     # Condition 2: use provided metrics to calculate for required metric
-    ds, coords, _ = datasets_grid_metric("C")
+    ds, coords, metrics = datasets_grid_metric("C")
     grid = Grid(ds, coords=coords)
     grid.set_metrics(("X"), "dx_n")
     grid.set_metrics(("Y"), "dy_n")
@@ -161,13 +203,12 @@ def test_get_metric_with_conditions_02():
     metric_var_2 = ds.dy_n
     expected_metric = (metric_var_1 * metric_var_2).reset_coords(drop=True)
 
-    if xr.testing.assert_equal(get_metric, expected_metric) is False:
-        xr.testing.assert_allclose(get_metric, expected_metric)
+    xr.testing.assert_allclose(get_metric, expected_metric)
 
 
 def test_get_metric_with_conditions_03():
     # Condition 3: Adapted from original version of test_get_metric
-    ds, coords, _ = datasets_grid_metric("C")
+    ds, coords, metrics = datasets_grid_metric("C")
     grid = Grid(ds, coords=coords)
     grid.set_metrics(("X", "Y"), "area_t")
     grid.set_metrics(("Z"), "dz_t")
@@ -178,69 +219,75 @@ def test_get_metric_with_conditions_03():
     metric_var_2 = ds.dz_t
     expected_metric = (metric_var_1 * metric_var_2).reset_coords(drop=True)
 
-    if xr.testing.assert_equal(get_metric, expected_metric) is False:
-        xr.testing.assert_allclose(get_metric, expected_metric)
+    xr.testing.assert_allclose(get_metric, expected_metric)
+
+
+@pytest.mark.skip(reason="For the next PR")
+def test_get_metric_with_conditions_02a():
+    # Condition 4, case a: interpolate metric with matching axis to desired dimensions
+    ds, coords, _ = datasets_grid_metric("C")
+    grid = Grid(ds, coords=coords)
+    grid.set_metrics(("X", "Y"), "area_e")
 
+    get_metric = grid.get_metric(ds.v, ("X", "Y"))
 
-# def test_get_metric_with_conditions_02a():
-#     # Condition 4, case a: interpolate metric with matching axis to desired dimensions
-#     ds, coords, _ = datasets_grid_metric("C")
-#     grid = Grid(ds, coords=coords)
-#     grid.set_metrics(("X","Y"), "area_e")
+    expected_metric = grid.interp(ds["area_e"], ("X", "Y"))
 
-#     get_metric = grid.get_metric(ds.v, ("X","Y"))
+    xr.testing.assert_equal(get_metric, expected_metric)
+    xr.testing.assert_allclose(get_metric, expected_metric)
 
-#     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)
+@pytest.mark.skip(reason="For the next PR")
+def test_get_metric_with_conditions_02b():
+    # Condition 4, case b: get_metric should select for the metric with matching axes and interpolate from there,
+    # even if other metrics in the desired positions are available
+    ds, coords, _ = datasets_grid_metric("C")
+    grid = Grid(ds, coords=coords)
+    grid.set_metrics(("X", "Y"), "area_e")
+    grid.set_metrics(("X"), "dx_n")
+    grid.set_metrics(("Y"), "dx_n")
+
+    get_metric = grid.get_metric(ds.v, ("X", "Y"))
 
-# def test_get_metric_with_conditions_02b():
-#     # Condition 4, case b: get_metric should select for the metric with matching axes and interpolate from there,
-#     # even if other metrics in the desired positions are available
-#     ds, coords, _ = datasets_grid_metric("C")
-#     grid = Grid(ds, coords=coords)
-#     grid.set_metrics(("X","Y"), "area_e")
-#     grid.set_metrics(("X"), "dx_n")
-#     grid.set_metrics(("Y"), "dx_n")
+    expected_metric = grid.interp(ds["area_e"], ("X", "Y"))
 
-#     get_metric = grid.get_metric(ds.v, ("X","Y"))
+    xr.testing.assert_equal(get_metric, expected_metric)
+    xr.testing.assert_allclose(get_metric, expected_metric)
 
-#     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)
+@pytest.mark.skip(reason="For the next PR")
+def test_get_metric_with_conditions_04a():
+    # Condition 5, case a: 1 metric on the wrong position (must interpolate before multiplying)
+    ds, coords, _ = datasets_grid_metric("C")
+    grid = Grid(ds, coords=coords)
+    grid.set_metrics(("X"), "dx_t")
+    grid.set_metrics(("Y"), "dy_n")
 
-# def test_get_metric_with_conditions_04a():
-#      # Condition 4, case a: 1 metric on the wrong position (must interpolate before multiplying)
-#     ds, coords, _ = datasets_grid_metric("C")
-#     grid = Grid(ds, coords=coords)
-#     grid.set_metrics(("X"), "dx_t")
-#     grid.set_metrics(("Y"), "dy_n")
+    get_metric = grid.get_metric(ds.v, ("X", "Y"))
 
-#     get_metric = grid.get_metric(ds.v, ("X","Y"))
+    interp_metric = grid.interp(ds.dx_t, "Y")
+    expected_metric = (interp_metric * ds.dy_n).reset_coords(drop=True)
 
-#     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)
 
-#     xr.testing.assert_equal(get_metric, expected_metric)
-#     xr.testing.assert_allclose(get_metric, expected_metric)
 
-# def test_get_metric_with_conditions_04b():
-#      # Condition 5, case b: 2 metrics in the wrong position (must interpolate before multiplying)
-#     ds, coords, _ = datasets_grid_metric("C")
-#     grid = Grid(ds, coords=coords)
-#     grid.set_metrics(("X"), "dx_t")
-#     grid.set_metrics(("Y"), "dy_t")
+@pytest.mark.skip(reason="For the next PR")
+def test_get_metric_with_conditions_04b():
+    # Condition 5, case b: 2 metrics in the wrong position (must interpolate before multiplying)
+    ds, coords, _ = datasets_grid_metric("C")
+    grid = Grid(ds, coords=coords)
+    grid.set_metrics(("X"), "dx_t")
+    grid.set_metrics(("Y"), "dy_t")
 
-#     get_metric = grid.get_metric(ds.v, ("X","Y"))
+    get_metric = grid.get_metric(ds.v, ("X", "Y"))
 
-#     interp_metric_1 = grid.interp(ds.dx_t,"Y")
-#     interp_metric_2 = grid.interp(ds.dy_t,"Y")
-#     expected_metric = (interp_metric_1*interp_metric_2).reset_coords(drop=True)
+    interp_metric_1 = grid.interp(ds.dx_t, "Y")
+    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)
+    xr.testing.assert_equal(get_metric, expected_metric)
+    xr.testing.assert_allclose(get_metric, expected_metric)
 
 
 def test_set_metric():
@@ -270,6 +317,7 @@ def test_set_metric():
             xr.testing.assert_equal(metric_expected.reset_coords(drop=True), metric)
 
 
+@pytest.mark.skip(reason="For the next PR")
 @pytest.mark.parametrize(
     "metric_axes,metric_name",
     [
@@ -278,13 +326,13 @@ def test_set_metric():
     ],
 )
 def test_interp_metrics(metric_axes, metric_name):
-
-    ds, coords, _ = datasets_grid_metric("C")
+    # need to test with boundary and fill_value conditions
+    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)
 
-    test_metric = grid.interp(ds[metric_name], metric_axes)
+    expected_metric = grid.interp(ds[metric_name], metric_axes)
 
-    xr.testing.assert_equal(interp_metric, test_metric)
-    xr.testing.assert_allclose(interp_metric, test_metric)
+    if interp_metric.equals(expected_metric) is False:
+        xr.testing.assert_allclose(interp_metric, expected_metric)