Add B-grid vector Laplacian

gcm_filters/kernels.py

@@ -23,6 +23,7 @@
         "TRIPOLAR_REGULAR_WITH_LAND_AREA_WEIGHTED",
         "TRIPOLAR_POP_WITH_LAND",
         "VECTOR_C_GRID",
+        "VECTOR_B_GRID",
     ],
 )
 
@@ -675,6 +676,172 @@ def __call__(self, ufield: ArrayType, vfield: ArrayType):
 ALL_KERNELS[GridType.VECTOR_C_GRID] = CgridVectorLaplacian
 
 
+@dataclass
+class BgridVectorLaplacian(BaseVectorLaplacian):
+    """̵Vector Laplacian on B-Grid. To be implemented for viscosity operators on B-grids based on POP code.
+
+    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
+    """
+
+    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
+
+        radius = 6370.0e5
+        ny, nx = self.TAREA.shape
+        self.AMF = np.sqrt(self.UAREA / (2 * np.pi * radius / nx) ** 2)
+
+    def __call__(self, ufield: ArrayType, vfield: ArrayType):
+        np = get_array_module(ufield)
+
+        # Constants
+        c2 = 2
+        p5 = 0.5
+
+        # Calculate coefficients for the stencil without metric terms
+        WORK1 = (self.HUS / self.HTE) * p5 * (self.AMF + np.roll(self.AMF, -1, axis=1))
+
+        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) * p5 * (self.AMF + np.roll(self.AMF, -1, axis=0))
+
+        DUW = WORK1 * self.UAREA_R  # West coefficient of 5-point stencil
+        DUE = (
+            np.roll(WORK1, 1, axis=0) * 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=0) - 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=0)) * self.TAREA_R  # KXT
+        WORK2 = (
+            p5
+            * (WORK1 + np.roll(WORK1, 1, axis=1))
+            * p5
+            * (np.roll(self.AMF, -1, axis=0) + self.AMF)
+        )
+        DXKX = (np.roll(WORK2, 1, axis=0) - WORK2) * self.DXUR
+
+        WORK2 = (
+            p5
+            * (WORK1 + np.roll(WORK1, 1, axis=0))
+            * p5
+            * (np.roll(self.AMF, -1, axis=1) + self.AMF)
+        )
+        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=0))
+            * p5
+            * (np.roll(self.AMF, -1, axis=1) + self.AMF)
+        )
+        DYKY = (np.roll(WORK2, 1, axis=1) - WORK2) * self.DYUR
+
+        WORK2 = (
+            p5
+            * (WORK1 + np.roll(WORK1, 1, axis=1))
+            * p5
+            * (np.roll(self.AMF, -1, axis=0) + self.AMF)
+        )
+        DXKY = (np.roll(WORK2, axis=0, shift=1) - WORK2) * self.DXUR
+
+        DUM = -(
+            DXKX + DYKY + c2 * self.AMF * (KXU ** 2 + KYU ** 2)
+        )  # 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
+        WORK1 = (
+            np.roll(self.AMF, axis=1, shift=1) - np.roll(self.AMF, axis=1, shift=-1)
+        ) / (self.HTE + np.roll(self.HTE, axis=1, shift=1))
+        DME = (c2 * self.AMF * KYU + WORK1) / (
+            self.HTN + np.roll(self.HTN, axis=0, shift=1)
+        )  # East coefficient of 5-point stencil for the metric terms that mix U,V
+
+        WORK1 = (
+            np.roll(self.AMF, axis=0, shift=1) - np.roll(self.AMF, axis=0, shift=-1)
+        ) / (self.HTN + np.roll(self.HTN, axis=0, shift=1))
+        DMN = -(c2 * self.AMF * KXU + WORK1) / (
+            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=0)
+            + DUS * np.roll(ufield, 1, axis=0)
+            + DUE * np.roll(ufield, -1, axis=1)
+            + DUW * np.roll(ufield, 1, axis=1)
+            + DMC * vfield
+            + DMN * np.roll(vfield, -1, axis=0)
+            + DMS * np.roll(vfield, 1, axis=0)
+            + 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=0)
+            + DUS * np.roll(vfield, 1, axis=0)
+            + DUE * np.roll(vfield, -1, axis=1)
+            + DUW * np.roll(vfield, 1, axis=1)
+            + DMC * ufield
+            + DMN * np.roll(ufield, -1, axis=0)
+            + DMS * np.roll(ufield, 1, axis=0)
+            + 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.
 

requirements-dev.txt

@@ -9,6 +9,8 @@ flake8-print
 interrogate
 isort
 nbsphinx
+netcdf4
+pooch
 pre-commit
 pylint
 pytest

tests/conftest.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import pytest
+import xarray as xr
 
 from numpy.random import PCG64, Generator
 
@@ -44,7 +45,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:
@@ -104,6 +105,17 @@ def _make_scalar_grid_data(grid_type):
     return grid_type, data, extra_kwargs
 
 
+def pop_test_data():
+    import pooch
+
+    fname = pooch.retrieve(
+        url="doi:10.5281/zenodo.5947728/pop_hires_test_data.nc",
+        known_hash="md5:e88ee4d310f476abe5fa3aac72d2e51e",
+    )
+    ds = xr.open_dataset(fname)
+    return ds
+
+
 @pytest.fixture(scope="session", params=scalar_grids)
 def scalar_grid_type_data_and_extra_kwargs(request):
     return _make_scalar_grid_data(request.param)
@@ -133,10 +145,7 @@ def tripolar_grid_type_data_and_extra_kwargs(request):
     return grid_type, data, extra_kwargs
 
 
-@pytest.fixture(scope="session")
-def spherical_geometry():
-    ny, nx = (128, 256)
-
+def spherical_geometry(ny, nx):
     # construct spherical coordinate system similar to MOM6 NeverWorld2 grid
     # define latitudes and longitudes
     lat_min = -70
@@ -161,18 +170,14 @@ def spherical_geometry():
     return geolon_u, geolat_u, geolon_v, geolat_v
 
 
-@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
+def gen_mom_vector_data():
+    ny, nx = (128, 256)
+
+    geolon_u, geolat_u, geolon_v, geolat_v = spherical_geometry(ny, nx)
     ny, nx = geolon_u.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)
@@ -201,8 +206,40 @@ def vector_grid_type_data_and_extra_kwargs(request, spherical_geometry):
     extra_kwargs["wet_mask_t"] = mask_data
     extra_kwargs["wet_mask_q"] = mask_data
 
+    # fill area_u, area_v with zeros over land; e.g., you will find that in MOM6 model output
+    extra_kwargs["area_u"] = np.where(
+        extra_kwargs["wet_mask_t"] > 0, extra_kwargs["area_u"], 0
+    )
+    extra_kwargs["area_v"] = np.where(
+        extra_kwargs["wet_mask_t"] > 0, extra_kwargs["area_v"], 0
+    )
+
     data_u = _make_random_data((ny, nx), 42)
     data_v = _make_random_data((ny, nx), 43)
 
+    return (data_u, data_v), extra_kwargs
+
+
+def gen_pop_vector_data():
+    ds = pop_test_data()
+    grid_vars = ["DXU", "DYU", "HUS", "HUW", "HTE", "HTN", "UAREA", "TAREA"]
+    extra_kwargs = {name: ds[name].values for name in grid_vars}
+    u_data = ds.U1_1.values
+    v_data = ds.V1_1.values
+    return (u_data, v_data), extra_kwargs
+
+
+@pytest.fixture(
+    scope="session",
+    params=[
+        (GridType.VECTOR_C_GRID, gen_mom_vector_data),
+        (GridType.VECTOR_B_GRID, gen_pop_vector_data),
+    ],
+    ids=["MOM-C-GRID", "POP-B-GRID"],
+)
+def vector_grid_type_data_and_extra_kwargs(request):
+    grid_type, gen_kwargs = request.param
+    (data_u, data_v), extra_kwargs = gen_kwargs()
+
     # use same return signature as other kernel fixtures
     return grid_type, (data_u, data_v), extra_kwargs

tests/test_kernels.py

@@ -11,6 +11,8 @@
     required_grid_vars,
 )
 
+from .conftest import spherical_geometry
+
 
 def test_conservation(scalar_grid_type_data_and_extra_kwargs):
     """This test checks that scalar Laplacians preserve the area integral."""
@@ -229,15 +231,15 @@ def test_tripolar_exchanges(tripolar_grid_type_data_and_extra_kwargs):
 #################### Vector Laplacian tests ########################################
 
 
-def test_conservation_under_solid_body_rotation(
-    vector_grid_type_data_and_extra_kwargs, spherical_geometry
-):
+def test_conservation_under_solid_body_rotation(vector_grid_type_data_and_extra_kwargs):
     """This test checks that vector Laplacians are invariant under solid body rotations:
     a corollary of conserving angular momentum."""
 
-    grid_type, _, extra_kwargs = vector_grid_type_data_and_extra_kwargs
+    grid_type, (u, _), extra_kwargs = vector_grid_type_data_and_extra_kwargs
 
-    _, geolat_u, _, _ = spherical_geometry
+    ny, nx = u.shape
+
+    _, geolat_u, _, _ = spherical_geometry(ny, nx)
     # u = cos(lat), v=0 is solid body rotation
     data_u = np.cos(geolat_u / 360 * 2 * np.pi)
     data_v = np.zeros_like(data_u)
@@ -255,16 +257,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))