Repository: ocean-eddy-cpt/gcm-filters (language: python)

Issue:
Add vector-based filtering for B-grid models
First off, thanks to the developers for this really great package with excellent documentation!

I would like to apply a vector-based filtering scheme to a B grid model (POP). It looks like the only current implementation of vector filtering is `VECTOR_C_GRID` for models with an Arakawa C grid. 

Would a corresponding `VECTOR_B_GRID` be a desirable feature to add to `gcm_filters`? If so, I'd be happy to take a stab at implementing it!
Add vector-based filtering for B-grid models
First off, thanks to the developers for this really great package with excellent documentation!

I would like to apply a vector-based filtering scheme to a B grid model (POP). It looks like the only current implementation of vector filtering is `VECTOR_C_GRID` for models with an Arakawa C grid. 

Would a corresponding `VECTOR_B_GRID` be a desirable feature to add to `gcm_filters`? If so, I'd be happy to take a stab at implementing it!
Add MOM6 kernel
Right now we just have a trivial laplacian kernel:

https://github.com/ocean-eddy-cpt/gcm-filters/blob/7b8d2906d0bdacb98846f4e39c3883ed3a3fc44a/gcm_filters/kernels.py#L17-L22

We need to add the kernel for MOM6. This should be based on whatever @NoraLoose is actively using in her notebooks.


Discussion / hints:
For context, this would enable filtering of velocity fields on MOM5 and POP grids. I support it 💯 .
I agree - this would be useful.
Hi Paige, nice to see you here!

Having a vector Laplacian on a B-grid would be super useful - thanks for taking this on! :rocket: 
(Griffies' book (2004), the POP manual, and the POP code itself will probably be useful references.)

> For context, this would enable filtering of velocity fields on MOM5 and POP grids.

If we are planning to implement a tripole boundary condition for the B-grid Laplacian, I can probably help with that (if needed). At some point I thought about how to implement the tripole exchanges on a vector grid, and will dig out the relevant code snippet. But I imagine that the boundary condition will be something we will tackle further down the road.


Great, it sounds like this is a desired feature so I'll get started! Thanks @NoraLoose for the useful references - I'll start by looking at those references and the vector Laplacian on a  C-grid code to get started. It might take a little time to spin up on all of this, but very keen to contribute! 

Your help with the triple boundary condition, and with the implementation in general, would be greatly appreciated @NoraLoose!
For context, this would enable filtering of velocity fields on MOM5 and POP grids. I support it 💯 .
I agree - this would be useful.
Hi Paige, nice to see you here!

Having a vector Laplacian on a B-grid would be super useful - thanks for taking this on! :rocket: 
(Griffies' book (2004), the POP manual, and the POP code itself will probably be useful references.)

> For context, this would enable filtering of velocity fields on MOM5 and POP grids.

If we are planning to implement a tripole boundary condition for the B-grid Laplacian, I can probably help with that (if needed). At some point I thought about how to implement the tripole exchanges on a vector grid, and will dig out the relevant code snippet. But I imagine that the boundary condition will be something we will tackle further down the road.


Great, it sounds like this is a desired feature so I'll get started! Thanks @NoraLoose for the useful references - I'll start by looking at those references and the vector Laplacian on a  C-grid code to get started. It might take a little time to spin up on all of this, but very keen to contribute! 

Your help with the triple boundary condition, and with the implementation in general, would be greatly appreciated @NoraLoose!
@gustavo-marques, @sdbachman, and all other MOM6 pros - I am happy to hand this task over to you if you'd like to contribute to gcm-filters (which I heard you do :smiley: ).

The only reason why I am assigned to this issue is that I have filtering code for the NW2 (MOM6) domain. But I haven't done the global domain yet. Would be great to have your expertise here. I think I have some ideas for how to do the global Laplacian for MOM6, so I am happy to discuss these ideas (or leave it completely up to you if you prefer).  Just wanted to mention this, because I didn't want to give the false impression that I want to do all kernel implementations for all models. :nerd_face:
Sounds good! Thanks, @NoraLoose. I will talk to @sdbachman and we will get back to you. 
This issue was accidentally closed. 
We have both scalar and vector Laplacians for MOM6 now, so I'm closing this issue.

Fix patch:
diff --git a/docs/basic_filtering.rst b/docs/basic_filtering.rst
index 379c3b5..243538b 100644
--- a/docs/basic_filtering.rst
+++ b/docs/basic_filtering.rst
@@ -46,6 +46,8 @@ The following table provides an overview of these different grid type options: w
 +--------------------------------+-----------------------------------------------------------------+--------------+--------------------+------------------+------------------------------------------+
 | ``VECTOR_C_GRID``              | `Arakawa C-Grid <https://en.wikipedia.org/wiki/Arakawa_grids>`_ | yes          | periodic           | Vector Laplacian | :doc:`examples/example_vector_laplacian` |
 +--------------------------------+-----------------------------------------------------------------+--------------+--------------------+------------------+------------------------------------------+
+| ``VECTOR_B_GRID``              | `Arakawa B-Grid <https://en.wikipedia.org/wiki/Arakawa_grids>`_ | no           | periodic           | Vector Laplacian |                                          |
++--------------------------------+-----------------------------------------------------------------+--------------+--------------------+------------------+------------------------------------------+
 
 Grid types with scalar Laplacians can be used for filtering scalar fields (such as temperature), and grid types with vector Laplacians can be used for filtering vector fields (such as velocity).
 
diff --git a/gcm_filters/kernels.py b/gcm_filters/kernels.py
index f6598e6..6922157 100644
--- a/gcm_filters/kernels.py
+++ b/gcm_filters/kernels.py
@@ -23,6 +23,7 @@
         "TRIPOLAR_REGULAR_WITH_LAND_AREA_WEIGHTED",
         "TRIPOLAR_POP_WITH_LAND",
         "VECTOR_C_GRID",
+        "VECTOR_B_GRID",
     ],
 )
 
@@ -695,6 +696,147 @@ def __call__(self, ufield: ArrayType, vfield: ArrayType):
 ALL_KERNELS[GridType.VECTOR_C_GRID] = CgridVectorLaplacian
 
 
+@dataclass
+class BgridVectorLaplacian(BaseVectorLaplacian):
+    """̵Vector Laplacian on B-Grid. Implemented for viscosity operators on B-grids based on POP code. Assumes periodic boundary conditions.
+
+    Attributes
+    ----------
+
+    DXU: x-spacing centered at U points
+    DYU: y-spacing centered at U points
+    HUS: cell widths on south side of U cell
+    HUW: cell widths on west side of U cell
+    HTE: cell widths on east side of T cell
+    HTN: cell widths on north side of T cell
+    UAREA: U-cell area
+    TAREA: T-cell area
+    """
+
+    is_dimensional = True
+
+    DXU: ArrayType
+    DYU: ArrayType
+    HUS: ArrayType
+    HUW: ArrayType
+    HTE: ArrayType
+    HTN: ArrayType
+    UAREA: ArrayType
+    TAREA: ArrayType
+
+    def __post_init__(self):
+        np = get_array_module(self.DXU)
+
+        # Derived quantities
+        self.UAREA_R = 1 / self.UAREA
+        self.TAREA_R = 1 / self.TAREA
+
+        self.DXUR = 1 / self.DXU
+        self.DYUR = 1 / self.DYU
+
+    def __call__(self, ufield: ArrayType, vfield: ArrayType):
+        np = get_array_module(ufield)
+
+        ufield = np.nan_to_num(ufield)
+        vfield = np.nan_to_num(vfield)
+
+        # Constants
+        c2 = 2
+        p5 = 0.5
+
+        # Calculate coefficients for the stencil without metric terms
+        WORK1 = self.HUS / self.HTE
+
+        DUS = (
+            WORK1 * self.UAREA_R
+        )  # South coefficient of 5-point stencil for the Del**2 operator acting at U points
+        DUN = (
+            np.roll(WORK1, 1, axis=-1) * self.UAREA_R
+        )  # North coefficient of 5-point stencil
+
+        WORK1 = self.HUW / self.HTN
+
+        DUW = WORK1 * self.UAREA_R  # West coefficient of 5-point stencil
+        DUE = (
+            np.roll(WORK1, 1, axis=-2) * self.UAREA_R
+        )  # East coefficient of 5-point stencil
+
+        # Calculate coefficients for metric terms and for metric advection terms (KXU,KYU)
+        KXU = (
+            np.roll(self.HUW, 1, axis=-2) - self.HUW
+        ) * self.UAREA_R  # defined in (3.24) of POP manual
+        KYU = (np.roll(self.HUS, 1, axis=-1) - self.HUS) * self.UAREA_R
+
+        WORK1 = (self.HTE - np.roll(self.HTE, -1, axis=-2)) * self.TAREA_R  # KXT
+        WORK2 = p5 * (WORK1 + np.roll(WORK1, 1, axis=-1))
+        DXKX = (np.roll(WORK2, 1, axis=-2) - WORK2) * self.DXUR
+
+        WORK2 = p5 * (WORK1 + np.roll(WORK1, 1, axis=-2))
+        DYKX = (np.roll(WORK2, 1, axis=-1) - WORK2) * self.DYUR
+
+        WORK1 = (self.HTN - np.roll(self.HTN, -1, axis=-1)) * self.TAREA_R  # KYT
+        WORK2 = p5 * (WORK1 + np.roll(WORK1, 1, axis=-2))
+        DYKY = (np.roll(WORK2, 1, axis=-1) - WORK2) * self.DYUR
+
+        WORK2 = p5 * (WORK1 + np.roll(WORK1, 1, axis=-1))
+        DXKY = (np.roll(WORK2, axis=-2, shift=1) - WORK2) * self.DXUR
+
+        DUM = -(
+            DXKX + DYKY + c2 * (KXU * KXU + KYU * KYU)
+        )  # central coefficient for metric terms that do not mix U,V
+        DMC = (
+            DXKY - DYKX
+        )  # central coefficient of 5-point stencil for the metric terms that mix U,V
+
+        # Calculate the central coefficient for metric mixing terms that mix U,V
+        DME = (c2 * KYU) / (
+            self.HTN + np.roll(self.HTN, axis=-2, shift=1)
+        )  # East coefficient of 5-point stencil for the metric terms that mix U,V
+
+        DMN = -(c2 * KXU) / (
+            self.HTE + np.roll(self.HTE, axis=-1, shift=1)
+        )  # North coefficient of 5-point stencil
+
+        DUC = -(DUN + DUS + DUE + DUW)  # central coefficient of 5-point stencil
+        DMW = -DME  # West coefficient of 5-point stencil
+        DMS = -DMN  # East coefficient of 5-point stencil
+
+        # Compute the horizontal diffusion of momentum
+        am = 1
+        cc = DUC + DUM
+
+        u_component = am * (
+            cc * ufield
+            + DUN * np.roll(ufield, -1, axis=-2)
+            + DUS * np.roll(ufield, 1, axis=-2)
+            + DUE * np.roll(ufield, -1, axis=-1)
+            + DUW * np.roll(ufield, 1, axis=-1)
+            + DMC * vfield
+            + DMN * np.roll(vfield, -1, axis=-2)
+            + DMS * np.roll(vfield, 1, axis=-2)
+            + DME * np.roll(vfield, -1, axis=-1)
+            + DMW * np.roll(vfield, 1, axis=-1)
+        )
+
+        v_component = am * (
+            cc * vfield
+            + DUN * np.roll(vfield, -1, axis=-2)
+            + DUS * np.roll(vfield, 1, axis=-2)
+            + DUE * np.roll(vfield, -1, axis=-1)
+            + DUW * np.roll(vfield, 1, axis=-1)
+            + DMC * ufield
+            + DMN * np.roll(ufield, -1, axis=-2)
+            + DMS * np.roll(ufield, 1, axis=-2)
+            + DME * np.roll(ufield, -1, axis=-1)
+            + DMW * np.roll(ufield, 1, axis=-1)
+        )
+
+        return (u_component, v_component)
+
+
+ALL_KERNELS[GridType.VECTOR_B_GRID] = BgridVectorLaplacian
+
+
 def required_grid_vars(grid_type: GridType):
     """Utility function for figuring out the required grid variables needed by each grid type.
 
diff --git a/requirements-dev.txt b/requirements-dev.txt
index f6164ae..e57ebf3 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -9,6 +9,7 @@ flake8-print
 interrogate
 isort
 nbsphinx
+numcodecs==0.10.1
 pre-commit
 pylint
 pytest


Test patch:
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/conftest.py b/tests/conftest.py
index d201233..a1ab515 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -30,6 +30,34 @@
     GridType.TRIPOLAR_POP_WITH_LAND: ["wet_mask", "dxe", "dye", "dxn", "dyn", "tarea"],
 }
 
+_vector_grid_kwargs = {
+    GridType.VECTOR_C_GRID: [
+        "wet_mask_t",
+        "wet_mask_q",
+        "dxT",
+        "dyT",
+        "dxCu",
+        "dyCu",
+        "dxCv",
+        "dyCv",
+        "dxBu",
+        "dyBu",
+        "area_u",
+        "area_v",
+        "kappa_iso",
+        "kappa_aniso",
+    ],
+    GridType.VECTOR_B_GRID: [
+        "DXU",
+        "DYU",
+        "HUS",
+        "HUW",
+        "HTE",
+        "HTN",
+        "UAREA",
+        "TAREA",
+    ],
+}
 
 scalar_grids = [
     GridType.REGULAR,
@@ -45,7 +73,7 @@
     GridType.TRIPOLAR_REGULAR_WITH_LAND_AREA_WEIGHTED,
     GridType.TRIPOLAR_POP_WITH_LAND,
 ]
-vector_grids = [GridType.VECTOR_C_GRID]
+vector_grids = [GridType.VECTOR_C_GRID, GridType.VECTOR_B_GRID]
 
 
 def _make_random_data(shape: Tuple[int, int], seed: int) -> np.ndarray:
@@ -150,72 +178,90 @@ def grid_type_and_input_ds(request, scalar_grid_type_data_and_extra_kwargs):
 
 
 @pytest.fixture(scope="session")
+# compute latitudes and longitudes for a spherical C-grid geometry
 def spherical_geometry():
+
     ny, nx = (128, 256)
 
     # construct spherical coordinate system similar to MOM6 NeverWorld2 grid
     # define latitudes and longitudes
     lat_min = -70
     lat_max = 70
-    lat_u = np.linspace(
+
+    # compute latitude of u-points on C-grid
+    latCu = np.linspace(
         lat_min + 0.5 * (lat_max - lat_min) / ny,
         lat_max - 0.5 * (lat_max - lat_min) / ny,
         ny,
     )
-    lat_v = np.linspace(lat_min + (lat_max - lat_min) / ny, lat_max, ny)
+    # compute latitude of v-points on C-grid
+    latCv = np.linspace(lat_min + (lat_max - lat_min) / ny, lat_max, ny)
+
     lon_min = 0
     lon_max = 60
-    lon_u = np.linspace(lon_min + (lon_max - lon_min) / nx, lon_max, nx)
-    lon_v = np.linspace(
+
+    lonCu = np.linspace(lon_min + (lon_max - lon_min) / nx, lon_max, nx)
+    lonCv = np.linspace(
         lon_min + 0.5 * (lon_max - lon_min) / nx,
         lon_max - 0.5 * (lon_max - lon_min) / nx,
         nx,
     )
-    (geolon_u, geolat_u) = np.meshgrid(lon_u, lat_u)
-    (geolon_v, geolat_v) = np.meshgrid(lon_v, lat_v)
 
-    return geolon_u, geolat_u, geolon_v, geolat_v
+    (geolonCu, geolatCu) = np.meshgrid(lonCu, latCu)
+    (geolonCv, geolatCv) = np.meshgrid(lonCv, latCv)
+
+    return geolonCu, geolatCu, geolonCv, geolatCv
 
 
 @pytest.fixture(scope="session", params=vector_grids)
 def vector_grid_type_data_and_extra_kwargs(request, spherical_geometry):
     grid_type = request.param
-    geolon_u, geolat_u, geolon_v, geolat_v = spherical_geometry
-    ny, nx = geolon_u.shape
+    geolonCu, geolatCu, geolonCv, geolatCv = spherical_geometry
+    ny, nx = geolonCu.shape
 
     extra_kwargs = {}
 
-    # for now, we assume that the only implemented vector grid is VECTOR_C_GRID
-    # we can relax this if we implement other vector grids
-    assert grid_type == GridType.VECTOR_C_GRID
-
     R = 6378000
+
     # dx varies spatially
-    extra_kwargs["dxCu"] = R * np.cos(geolat_u / 360 * 2 * np.pi)
-    extra_kwargs["dxCv"] = R * np.cos(geolat_v / 360 * 2 * np.pi)
-    extra_kwargs["dxBu"] = extra_kwargs["dxCv"] + np.roll(
-        extra_kwargs["dxCv"], -1, axis=1
-    )
-    extra_kwargs["dxT"] = extra_kwargs["dxCu"] + np.roll(
-        extra_kwargs["dxCu"], 1, axis=1
-    )
-    # dy is set constant, equal to dx at the equator
-    dy = np.max(extra_kwargs["dxCu"]) * np.ones((ny, nx))
-    extra_kwargs["dyCu"] = dy
-    extra_kwargs["dyCv"] = dy
-    extra_kwargs["dyBu"] = dy
-    extra_kwargs["dyT"] = dy
+    for name in _vector_grid_kwargs[grid_type]:
+        if name in ["dxCu", "dxT", "HUS", "HTE"]:
+            extra_kwargs[name] = R * np.cos(geolatCu / 360 * 2 * np.pi)
+            # compute dy for later
+            # dy is set constant, equal to dx at the equator
+            dy = np.max(extra_kwargs[name]) * np.ones((ny, nx))
+        if name in ["dxCv", "dxBu", "DXU", "HUW", "HTN"]:
+            extra_kwargs[name] = R * np.cos(geolatCv / 360 * 2 * np.pi)
+
+    # dy
+    for name in _vector_grid_kwargs[grid_type]:
+        if name in ["dyCu", "dyCv", "dyBu", "dyT", "DYU"]:
+            extra_kwargs[name] = dy
+
     # compute grid cell areas
-    extra_kwargs["area_u"] = extra_kwargs["dxCu"] * extra_kwargs["dyCu"]
-    extra_kwargs["area_v"] = extra_kwargs["dxCv"] * extra_kwargs["dyCv"]
+    for name in _vector_grid_kwargs[grid_type]:
+        if name == "area_u":
+            extra_kwargs[name] = extra_kwargs["dxCu"] * extra_kwargs["dyCu"]
+        elif name == "area_v":
+            extra_kwargs[name] = extra_kwargs["dxCv"] * extra_kwargs["dyCv"]
+        elif name == "UAREA":
+            extra_kwargs[name] = extra_kwargs["DXU"] * extra_kwargs["DYU"]
+        elif name == "TAREA":
+            # on a spherical grid, HTE is the same as x-spacing centered at tracer point
+            # on a spherical grid, DYU is the same as y-spacing centered at tracer point
+            extra_kwargs[name] = extra_kwargs["HTE"] * extra_kwargs["DYU"]
+
     # set isotropic and anisotropic kappas
-    extra_kwargs["kappa_iso"] = np.ones((ny, nx))
-    extra_kwargs["kappa_aniso"] = np.ones((ny, nx))
+    for name in _vector_grid_kwargs[grid_type]:
+        if name in ["kappa_iso", "kappa_aniso"]:
+            extra_kwargs[name] = np.ones((ny, nx))
+
     # put a big island in the middle
     mask_data = np.ones((ny, nx))
     mask_data[: (ny // 2), : (nx // 2)] = 0
-    extra_kwargs["wet_mask_t"] = mask_data
-    extra_kwargs["wet_mask_q"] = mask_data
+    for name in _vector_grid_kwargs[grid_type]:
+        if name in ["wet_mask_t", "wet_mask_q"]:
+            extra_kwargs[name] = mask_data
 
     data_u = _make_random_data((ny, nx), 42)
     data_v = _make_random_data((ny, nx), 43)
diff --git a/tests/test_data_filter/VECTOR_B_GRID.zarr/.zarray b/tests/test_data_filter/VECTOR_B_GRID.zarr/.zarray
new file mode 100644
index 0000000..e6651bd
--- /dev/null
+++ b/tests/test_data_filter/VECTOR_B_GRID.zarr/.zarray
@@ -0,0 +1,24 @@
+{
+    "chunks": [
+        2,
+        128,
+        256
+    ],
+    "compressor": {
+        "blocksize": 0,
+        "clevel": 5,
+        "cname": "lz4",
+        "id": "blosc",
+        "shuffle": 1
+    },
+    "dtype": "<f4",
+    "fill_value": 0.0,
+    "filters": null,
+    "order": "C",
+    "shape": [
+        2,
+        128,
+        256
+    ],
+    "zarr_format": 2
+}
\ No newline at end of file
diff --git a/tests/test_data_filter/VECTOR_B_GRID.zarr/0.0.0 b/tests/test_data_filter/VECTOR_B_GRID.zarr/0.0.0
new file mode 100644
index 0000000..4af7da9
Binary files /dev/null and b/tests/test_data_filter/VECTOR_B_GRID.zarr/0.0.0 differ
diff --git a/tests/test_data_filter/VECTOR_C_GRID.zarr/0.0.0 b/tests/test_data_filter/VECTOR_C_GRID.zarr/0.0.0
index 9a0fd93..4af7da9 100644
Binary files a/tests/test_data_filter/VECTOR_C_GRID.zarr/0.0.0 and b/tests/test_data_filter/VECTOR_C_GRID.zarr/0.0.0 differ
diff --git a/tests/test_data_kernels/VECTOR_B_GRID.zarr/.zarray b/tests/test_data_kernels/VECTOR_B_GRID.zarr/.zarray
new file mode 100644
index 0000000..e6651bd
--- /dev/null
+++ b/tests/test_data_kernels/VECTOR_B_GRID.zarr/.zarray
@@ -0,0 +1,24 @@
+{
+    "chunks": [
+        2,
+        128,
+        256
+    ],
+    "compressor": {
+        "blocksize": 0,
+        "clevel": 5,
+        "cname": "lz4",
+        "id": "blosc",
+        "shuffle": 1
+    },
+    "dtype": "<f4",
+    "fill_value": 0.0,
+    "filters": null,
+    "order": "C",
+    "shape": [
+        2,
+        128,
+        256
+    ],
+    "zarr_format": 2
+}
\ No newline at end of file
diff --git a/tests/test_data_kernels/VECTOR_B_GRID.zarr/0.0.0 b/tests/test_data_kernels/VECTOR_B_GRID.zarr/0.0.0
new file mode 100644
index 0000000..f14879d
Binary files /dev/null and b/tests/test_data_kernels/VECTOR_B_GRID.zarr/0.0.0 differ
diff --git a/tests/test_data_kernels/VECTOR_C_GRID.zarr/0.0.0 b/tests/test_data_kernels/VECTOR_C_GRID.zarr/0.0.0
index 66c6fb4..c61ca7c 100644
Binary files a/tests/test_data_kernels/VECTOR_C_GRID.zarr/0.0.0 and b/tests/test_data_kernels/VECTOR_C_GRID.zarr/0.0.0 differ
diff --git a/tests/test_kernels.py b/tests/test_kernels.py
index e86bd5e..d457eb2 100644
--- a/tests/test_kernels.py
+++ b/tests/test_kernels.py
@@ -274,16 +274,8 @@ def test_zero_area(vector_grid_type_data_and_extra_kwargs):
 
     grid_type, (data_u, data_v), extra_kwargs = vector_grid_type_data_and_extra_kwargs
 
-    test_kwargs = copy.deepcopy(extra_kwargs)
-    # fill area_u, area_v with zeros over land; e.g., you will find that in MOM6 model output
-    test_kwargs["area_u"] = np.where(
-        extra_kwargs["wet_mask_t"] > 0, test_kwargs["area_u"], 0
-    )
-    test_kwargs["area_v"] = np.where(
-        extra_kwargs["wet_mask_t"] > 0, test_kwargs["area_v"], 0
-    )
     LaplacianClass = ALL_KERNELS[grid_type]
-    laplacian = LaplacianClass(**test_kwargs)
+    laplacian = LaplacianClass(**extra_kwargs)
     res_u, res_v = laplacian(data_u, data_v)
     assert not np.any(np.isinf(res_u))
     assert not np.any(np.isnan(res_u))