Merge pull request #128 from paigem/vector-b-grid

Add B-grid vector Laplacian

Author: NoraLoose

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).
 

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.
 

tests/__init__.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)

tests/conftest.py

@@ -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
+}

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
+}