Merge pull request #146 from iangrooms/fix_dimensionality

Fix dimensionality

.pre-commit-config.yaml

@@ -34,11 +34,17 @@ repos:
   hooks:
     - id: seed-isort-config
 
+# - repo: https://github.com/psf/black
+#   rev: 21.5b0
+#   hooks:
+#   - id: black
+#     language_version: python3
+
 - repo: https://github.com/psf/black
-  rev: 21.5b0
+  rev: 22.1.0
   hooks:
-  - id: black
-    language_version: python3
+    - id: black
+      additional_dependencies: ['click==8.0.4']
 
 #- repo: https://github.com/pre-commit/mirrors-mypy
 #  rev: v0.770

docs/basic_filtering.rst

@@ -97,6 +97,11 @@ Each grid type from the above two tables has different *grid variables* that mus
 
 We have made a big island.
 
+
+.. note:: Some more complicated grid types require more grid variables.
+    The units for these variables should be *consistent*, but no specific system of units is required.
+    For example, if grid cell edge lengths are defined using kilometers, then the filter scale and ``dx_min`` should also be defined using kilometers, and the grid cell areas should be defined in square kilometers.
+
 Creating the Filter Object
 --------------------------
 

gcm_filters/filter.py

@@ -77,6 +77,7 @@ class FilterSpec(NamedTuple):
     s_max: float
     p: Iterable[float]
     n_iterations: int
+    dx_min_sq: float
 
 
 def _compute_filter_spec(
@@ -161,7 +162,20 @@ def _compute_filter_spec(
     s = np.array([y for x in s for y in x])
     is_laplacian = np.abs(s.imag / s.real) < root_tolerance
 
-    return FilterSpec(n_steps_total, s, is_laplacian, s_max, p, n_iterations)
+    dx_min_sq = dx_min**2
+
+    return FilterSpec(n_steps_total, s, is_laplacian, s_max, p, n_iterations, dx_min_sq)
+
+
+def _maybe_dimensionalize(data, dx_min_sq):
+    """The point of this function is to avoid the expense of
+    dividing a large array by 1 when dealing with nondimensional
+    Laplacians
+    """
+    if dx_min_sq == 1:
+        return data
+    else:
+        return data / dx_min_sq
 
 
 def _create_filter_func(
@@ -191,13 +205,24 @@ def filter_func(field, *args):
                 if filter_spec.is_laplacian[i]:
                     s_l = np.real(filter_spec.s[i])
                     tendency = laplacian(field_bar)  # Compute Laplacian
+                    if not Laplacian.is_dimensional:
+                        tendency = _maybe_dimensionalize(
+                            tendency, filter_spec.dx_min_sq
+                        )  # dimensionalize
                     field_bar += (1 / s_l) * tendency  # Update filtered field
                 else:
                     s_b = filter_spec.s[i]
                     temp_l = laplacian(field_bar)  # Compute Laplacian
                     temp_b = laplacian(
                         temp_l
                     )  # Compute Biharmonic (apply Laplacian twice)
+                    if not Laplacian.is_dimensional:
+                        temp_l = _maybe_dimensionalize(
+                            temp_l, filter_spec.dx_min_sq
+                        )  # dimensionalize
+                        temp_b = _maybe_dimensionalize(
+                            temp_b, filter_spec.dx_min_sq**2
+                        )  # dimensionalize
                     field_bar += (
                         temp_l * 2 * np.real(s_b) / np.abs(s_b) ** 2
                         + temp_b * 1 / np.abs(s_b) ** 2
@@ -242,6 +267,13 @@ def filter_func_vec(ufield, vfield, *args):
                     (utendency, vtendency) = laplacian(
                         ufield_bar, vfield_bar
                     )  # Compute Laplacian
+                    if not Laplacian.is_dimensional:
+                        utendency = _maybe_dimensionalize(
+                            utendency, filter_spec.dx_min_sq
+                        )  # dimensionalize
+                        vtendency = _maybe_dimensionalize(
+                            vtendency, filter_spec.dx_min_sq
+                        )  # dimensionalize
                     ufield_bar += (1 / s_l) * utendency  # Update filtered ufield
                     vfield_bar += (1 / s_l) * vtendency  # Update filtered vfield
                 else:
@@ -252,6 +284,19 @@ def filter_func_vec(ufield, vfield, *args):
                     (utemp_b, vtemp_b) = laplacian(
                         utemp_l, vtemp_l
                     )  # Compute Biharmonic (apply Laplacian twice)
+                    if not Laplacian.is_dimensional:
+                        utemp_l = _maybe_dimensionalize(
+                            utemp_l, filter_spec.dx_min_sq
+                        )  # dimensionalize
+                        vtemp_l = _maybe_dimensionalize(
+                            vtemp_l, filter_spec.dx_min_sq
+                        )  # dimensionalize
+                        utemp_b = _maybe_dimensionalize(
+                            utemp_b, filter_spec.dx_min_sq**2
+                        )  # dimensionalize
+                        vtemp_b = _maybe_dimensionalize(
+                            vtemp_b, filter_spec.dx_min_sq**2
+                        )  # dimensionalize
                     ufield_bar += (
                         utemp_l * 2 * np.real(s_b) / np.abs(s_b) ** 2
                         + utemp_b * 1 / np.abs(s_b) ** 2
@@ -291,7 +336,8 @@ class Filter:
         - ``FilterShape.TAPER``: The target filter has target grid scale Lf. Smaller scales are zeroed out.
           Scales larger than ``pi * filter_scale / 2`` are left as-is. In between is a smooth transition.
     transition_width : float, optional
-        Width of the transition region in the "Taper" filter. Theoretical minimum is 1; not recommended.
+        Width of the transition region in the "Taper" filter.
+        This is a nondimensional parameter. Theoretical minimum is 1; not recommended.
     ndim : int, optional
          Laplacian is applied on a grid of dimension ndim
     grid_type : GridType

gcm_filters/kernels.py

@@ -41,7 +41,7 @@ def _prepare_tripolar_exchanges(field):
 
 @dataclass
 class BaseScalarLaplacian(ABC):
-    """̵Base class for scalar Laplacians."""
+    """Base class for scalar Laplacians."""
 
     def prepare(self, field):
         return field
@@ -107,6 +107,8 @@ def finalize(self, field):
 class RegularLaplacian(BaseScalarLaplacian):
     """̵Scalar Laplacian for regularly spaced Cartesian grids."""
 
+    is_dimensional = False
+
     def __call__(self, field: ArrayType):
         np = get_array_module(field)
         return (
@@ -134,6 +136,8 @@ class RegularLaplacianWithArea(AreaWeightedMixin, RegularLaplacian):
     area: cell area
     """
 
+    is_dimensional = False
+
     area: ArrayType
 
     pass
@@ -151,6 +155,8 @@ class RegularLaplacianWithLandMask(BaseScalarLaplacian):
     wet_mask: Mask array, 1 for ocean, 0 for land
     """
 
+    is_dimensional = False
+
     wet_mask: ArrayType
 
     def __post_init__(self):
@@ -199,6 +205,8 @@ class RegularLaplacianWithLandMaskAndArea(
     wet_mask: Mask array, 1 for ocean, 0 for land
     """
 
+    is_dimensional = False
+
     area: ArrayType
     wet_mask: ArrayType
 
@@ -236,6 +244,8 @@ class IrregularLaplacianWithLandMask(BaseScalarLaplacian):
              least one place in the domain must have kappa_s = 1 if kappa_w < 1.
     """
 
+    is_dimensional = True
+
     wet_mask: ArrayType
     dxw: ArrayType
     dyw: ArrayType
@@ -319,6 +329,8 @@ class MOM5LaplacianU(BaseScalarLaplacian):
     area_u: area of U-cell, dxu*dyu
     """
 
+    is_dimensional = True
+
     wet_mask: ArrayType
     dxt: ArrayType
     dyt: ArrayType
@@ -374,6 +386,8 @@ class MOM5LaplacianT(BaseScalarLaplacian):
     area_t: area of T-cell, dxt*dyt
     """
 
+    is_dimensional = True
+
     wet_mask: ArrayType
     dxt: ArrayType
     dyt: ArrayType
@@ -428,6 +442,8 @@ class TripolarRegularLaplacianTpoint(AreaWeightedMixin, BaseScalarLaplacian):
     wet_mask: Mask array, 1 for ocean, 0 for land
     """
 
+    is_dimensional = False
+
     area: ArrayType
     wet_mask: ArrayType
 
@@ -485,6 +501,8 @@ class POPTripolarLaplacianTpoint(BaseScalarLaplacian):
     tarea: cell area, provided by POP model diagnostic TAREA(nlat, nlon)
     """
 
+    is_dimensional = True
+
     wet_mask: ArrayType
     dxe: ArrayType
     dye: ArrayType
@@ -588,6 +606,8 @@ class CgridVectorLaplacian(BaseVectorLaplacian):
     kappa_aniso: additive anisotropic viscosity aligned with x-direction
     """
 
+    is_dimensional = True
+
     wet_mask_t: ArrayType
     wet_mask_q: ArrayType
     dxT: ArrayType

tests/test_filter.py

@@ -17,6 +17,7 @@ def _check_equal_filter_spec(spec1, spec2):
     assert spec1.s_max == spec2.s_max
     np.testing.assert_allclose(spec1.p, spec2.p, rtol=1e-07, atol=1e-07)
     assert spec1.n_iterations == spec2.n_iterations
+    np.testing.assert_allclose(spec1.dx_min_sq, spec2.dx_min_sq)
 
 
 # These values were just hard copied from my dev environment.
@@ -76,6 +77,7 @@ def _check_equal_filter_spec(spec1, spec2):
                     -0.00454758,
                 ],
                 n_iterations=1,
+                dx_min_sq=1.0,
             ),
         ),
         (
@@ -106,6 +108,7 @@ def _check_equal_filter_spec(spec1, spec2):
                     0.00168445,
                 ],
                 n_iterations=1,
+                dx_min_sq=1.0,
             ),
         ),
     ],
@@ -331,7 +334,7 @@ def test_diffusion_filter(grid_type_and_input_ds, filter_args):
         filtered_u, filtered_v = filter.apply_to_vector(da, da, dims=["y", "x"])
 
     # check variance reduction
-    assert (filtered ** 2).sum() < (da ** 2).sum()
+    assert (filtered**2).sum() < (da**2).sum()
 
     # check that we get an error if we leave out any required grid_vars
     for gv in grid_vars:
@@ -430,6 +433,40 @@ def test_application_to_dataset():
         filter.apply(dataset, ["yy", "x"])
 
 
+def test_nondimensional_invariance():
+    # Create a dataset with spatial variables, as above
+    dataset = xr.Dataset(
+        data_vars=dict(
+            spatial=(("y", "x"), np.random.normal(size=(100, 100))),
+        ),
+        coords=dict(
+            x=np.linspace(0, 1e6, 100),
+            y=np.linspace(0, 1e6, 100),
+        ),
+    )
+
+    # Filter it using a nondimenisional filter, dx_min = 1
+    filter = Filter(
+        filter_scale=4,
+        dx_min=1,
+        filter_shape=FilterShape.GAUSSIAN,
+        grid_type=GridType.REGULAR,
+    )
+    filtered_dataset = filter.apply(dataset, ["y", "x"])
+
+    # Filter it using a nondimensional filter, dx_min = 2
+    filter = Filter(
+        filter_scale=8,
+        dx_min=2,
+        filter_shape=FilterShape.GAUSSIAN,
+        grid_type=GridType.REGULAR,
+    )
+    filtered_dataset_v2 = filter.apply(dataset, ["y", "x"])
+
+    # Check if they are the same
+    xr.testing.assert_allclose(filtered_dataset.spatial, filtered_dataset_v2.spatial)
+
+
 @pytest.mark.parametrize(
     "filter_args",
     [
@@ -472,7 +509,7 @@ def test_iterated_filter(grid_type_and_input_ds, filter_args, n_iterations):
     # See the "Factoring the Gaussian Filter" section of the docs for details.
     assert (((filtered - iteratively_filtered) ** 2) * area).sum() < (
         (0.01 * (1 + n_iterations)) ** 2
-    ) * ((da ** 2) * area).sum()
+    ) * ((da**2) * area).sum()
 
 
 #################### Visosity-based filter tests ########################################
@@ -547,7 +584,7 @@ def test_iterated_viscosity_filter(
     difference = (filtered_u - iteratively_filtered_u) ** 2 + (
         filtered_v - iteratively_filtered_v
     ) ** 2
-    unfiltered = da_u ** 2 + da_v ** 2
+    unfiltered = da_u**2 + da_v**2
     assert (difference * area).sum() < ((0.01 * (1 + n_iterations)) ** 2) * (
         unfiltered * area
     ).sum()

tests/test_kernels.py

@@ -42,6 +42,25 @@ def test_required_grid_vars(scalar_grid_type_data_and_extra_kwargs):
     assert set(grid_vars) == set(extra_kwargs)
 
 
+def test_dimensionality_scalar(scalar_grid_type_data_and_extra_kwargs):
+    """This test checks that REGULAR Laplacians are marked as nondimensional"""
+    grid_type, data, extra_kwargs = scalar_grid_type_data_and_extra_kwargs
+
+    LaplacianClass = ALL_KERNELS[grid_type]
+    switch = {
+        GridType.REGULAR: True,
+        GridType.REGULAR_AREA_WEIGHTED: True,
+        GridType.REGULAR_WITH_LAND: True,
+        GridType.REGULAR_WITH_LAND_AREA_WEIGHTED: True,
+        GridType.IRREGULAR_WITH_LAND: False,
+        GridType.MOM5U: False,
+        GridType.MOM5T: False,
+        GridType.TRIPOLAR_REGULAR_WITH_LAND_AREA_WEIGHTED: True,
+        GridType.TRIPOLAR_POP_WITH_LAND: False,
+    }
+    assert switch.get(grid_type, "Invalid input") != LaplacianClass.is_dimensional
+
+
 ################## Irregular grid tests for scalar Laplacians ##############################################
 # Irregular grids are grids that allow spatially varying dx, dy
 
@@ -275,3 +294,14 @@ def test_required_vector_grid_vars(vector_grid_type_data_and_extra_kwargs):
     grid_type, _, extra_kwargs = vector_grid_type_data_and_extra_kwargs
     grid_vars = required_grid_vars(grid_type)
     assert set(grid_vars) == set(extra_kwargs)
+
+
+def test_dimensionality_vector(vector_grid_type_data_and_extra_kwargs):
+    """This test checks that REGULAR Laplacians are marked as nondimensional"""
+    grid_type, (data_u, data_v), extra_kwargs = vector_grid_type_data_and_extra_kwargs
+
+    LaplacianClass = ALL_KERNELS[grid_type]
+    switch = {
+        GridType.VECTOR_C_GRID: False,
+    }
+    assert switch.get(grid_type, "Invalid input") != LaplacianClass.is_dimensional