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

Issue:
Agree on API for how we want to use this package
How will users interact with the gcm-filters package? We can answer this a bit by looking at existing code by @NoraLoose.

I image something like this (pseudocode)

```python
from gcm_filters import Filter

filter = Filter(scale=100_000., other_param=3, model="mom6", grid=grid_object)
scalar_filtered = filter.apply_scalar(scalar)
u_filt, v_filt = filter.apply_vector(u, v)
```

Answering this requires us to determine what are the input parameters to the filter design algorithm.


Discussion / hints:
What about "bring your own kernel"? BYOK.
What about

```python
Filter(kernel="mom6")
````

or 
```python
Filter(kernel=my_custom_kernel)
```
I like BYOK. The package can provide whatever kernels we write (e.g. POP for the paper?). If a user wants to use this for something new then they just have to write a Laplacian kernel for their grid, and we can encourage them to submit it as a pull request so we can add it to the package.

Fix patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index c2f84b6..889e378 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1,12 +1,5 @@
 repos:
 
-- repo: https://github.com/econchick/interrogate
-  rev: 1.2.0
-  hooks:
-    - id: interrogate
-      exclude: ^(docs|setup.py|tests)
-      args: [--config=pyproject.toml]
-
 - repo: https://github.com/pre-commit/pre-commit-hooks
   rev: v3.1.0
   hooks:
@@ -52,10 +45,3 @@ repos:
   - id: mypy
     exclude: docs/source/conf.py
     args: [--ignore-missing-imports]
-
-- repo: https://github.com/codespell-project/codespell
-  rev: v1.16.0
-  hooks:
-    - id: codespell
-      args:
-        - --quiet-level=2
diff --git a/docs/environment.yml b/docs/environment.yml
index 7211231..6ac8485 100644
--- a/docs/environment.yml
+++ b/docs/environment.yml
@@ -2,6 +2,7 @@ name: gcm-filters_docs
 dependencies:
   - python=3.8
   - numpy
+  - scipy
   - xarray
   - numpydoc
   - sphinx
diff --git a/gcm_filters/__init__.py b/gcm_filters/__init__.py
index 3dac500..e2ef750 100644
--- a/gcm_filters/__init__.py
+++ b/gcm_filters/__init__.py
@@ -6,3 +6,10 @@
     from ._version import __version__
 except ImportError:
     __version__ = "unknown"
+
+
+from .filter import Filter, FilterShape
+from .kernels import GridType
+
+
+__all__ = ["Filter", "FilterShape", "GridType"]
diff --git a/gcm_filters/filter.py b/gcm_filters/filter.py
new file mode 100644
index 0000000..cd54f15
--- /dev/null
+++ b/gcm_filters/filter.py
@@ -0,0 +1,228 @@
+"""Main Filter class."""
+import enum
+
+from dataclasses import dataclass, field
+from typing import Iterable, NamedTuple
+
+import numpy as np
+import xarray as xr
+
+from scipy import interpolate
+
+from .gpu_compat import get_array_module
+from .kernels import ALL_KERNELS, BaseLaplacian, GridType
+
+
+FilterShape = enum.Enum("FilterShape", ["GAUSSIAN", "TAPER"])
+
+
+class TargetSpec(NamedTuple):
+    s_max: float
+    filter_scale: float
+    transition_width: int
+
+
+# these functions return functions
+def _gaussian_target(target_spec: TargetSpec):
+    return lambda t: np.exp(
+        -(target_spec.s_max * (t + 1) / 2) * (target_spec.filter_scale / 2) ** 2
+    )
+
+
+def _taper_target(target_spec: TargetSpec):
+    return interpolate.PchipInterpolator(
+        np.array(
+            [
+                -1,
+                (2 / target_spec.s_max)
+                * (np.pi / (target_spec.transition_width * target_spec.filter_scale))
+                ** 2
+                - 1,
+                (2 / target_spec.s_max) * (np.pi / target_spec.filter_scale) ** 2 - 1,
+                2,
+            ]
+        ),
+        np.array([1, 1, 0, 0]),
+    )
+
+
+_target_function = {
+    FilterShape.GAUSSIAN: _gaussian_target,
+    FilterShape.TAPER: _taper_target,
+}
+
+
+class FilterSpec(NamedTuple):
+    n_lap_steps: int
+    s_l: Iterable[float]
+    n_bih_steps: int
+    s_b: Iterable[complex]
+
+
+def _compute_filter_spec(
+    filter_scale, dx_min, n_steps, filter_shape, transition_width, root_tolerance=1e-12
+):
+    # First set up the mass matrix for the Galerkin basis from Shen (SISC95)
+    M = (np.pi / 2) * (
+        2 * np.eye(n_steps - 1)
+        - np.diag(np.ones(n_steps - 3), 2)
+        - np.diag(np.ones(n_steps - 3), -2)
+    )
+    M[0, 0] = 3 * np.pi / 2
+
+    # The range of wavenumbers is 0<=|k|<=sqrt(2)*pi/dxMin. Nyquist here is for a 2D grid.
+    # Per the notes, define s=k^2.
+    # Need to rescale to t in [-1,1]: t = (2/sMax)*s -1; s = sMax*(t+1)/2
+    s_max = 2 * (np.pi / dx_min) ** 2
+
+    target_spec = TargetSpec(s_max, filter_scale, transition_width)
+    F = _target_function[filter_shape](target_spec)
+
+    # Compute inner products of Galerkin basis with target
+    b = np.zeros(n_steps - 1)
+    points, weights = np.polynomial.chebyshev.chebgauss(n_steps + 1)
+    for i in range(n_steps - 1):
+        tmp = np.zeros(n_steps + 1)
+        tmp[i] = 1
+        tmp[i + 2] = -1
+        phi = np.polynomial.chebyshev.chebval(points, tmp)
+        b[i] = np.sum(
+            weights * phi * (F(points) - ((1 - points) / 2 + F(1) * (points + 1) / 2))
+        )
+
+    # Get polynomial coefficients in Galerkin basis
+    c_hat = np.linalg.solve(M, b)
+    # Convert back to Chebyshev basis coefficients
+    p = np.zeros(n_steps + 1)
+    p[0] = c_hat[0] + (1 + F(1)) / 2
+    p[1] = c_hat[1] - (1 - F(1)) / 2
+    for i in range(2, n_steps - 1):
+        p[i] = c_hat[i] - c_hat[i - 2]
+    p[n_steps - 1] = -c_hat[n_steps - 3]
+    p[n_steps] = -c_hat[n_steps - 2]
+
+    # Get roots of the polynomial
+    r = np.polynomial.chebyshev.chebroots(p)
+
+    # convert back to s in [0,sMax]
+    s = s_max / 2 * (r + 1)
+    # Separate out the real and complex roots
+    n_lap_steps = np.size(s[np.where(np.abs(np.imag(r)) < root_tolerance)])
+    s_l = np.real(s[np.where(np.abs(np.imag(r)) < root_tolerance)])
+    n_bih_steps = (n_steps - n_lap_steps) // 2
+    s_b_re, indices = np.unique(
+        np.real(s[np.where(np.abs(np.imag(r)) > root_tolerance)]), return_index=True
+    )
+    s_b_im = np.imag(s[np.where(np.abs(np.imag(r)) > root_tolerance)])[indices]
+    s_b = s_b_re + s_b_im * 1j
+
+    return FilterSpec(n_lap_steps, s_l, n_bih_steps, s_b)
+
+
+def _create_filter_func(
+    filter_spec: FilterSpec,
+    Laplacian: BaseLaplacian,
+):
+    """Returns a function whose first argument is the field to be filtered
+    and whose subsequent arguments are the require grid variables
+    """
+
+    def filter_func(field, *args):
+        # these next steps are a kind of hack we have to turn keyword arugments into regular arguments
+        # the reason for doing this is that Xarray's apply_ufunc machinery works a lot better
+        # with regular arguments
+        assert len(args) == len(Laplacian.required_grid_args())
+        grid_vars = {k: v for k, v in zip(Laplacian.required_grid_args(), args)}
+        laplacian = Laplacian(**grid_vars)
+        np = get_array_module(field)
+        field_bar = field.copy()  # Initalize the filtering process
+        for i in range(filter_spec.n_lap_steps):
+            s_l = filter_spec.s_l[i]
+            tendency = laplacian(field_bar)  # Compute Laplacian
+            field_bar += (1 / s_l) * tendency  # Update filtered field
+        for i in range(filter_spec.n_bih_steps):
+            s_b = filter_spec.s_b[i]
+            temp_l = laplacian(field_bar)  # Compute Laplacian
+            temp_b = laplacian(temp_l)  # Compute Biharmonic (apply Laplacian twice)
+            field_bar += (
+                temp_l * 2 * np.real(s_b) / np.abs(s_b) ** 2
+                + temp_b * 1 / np.abs(s_b) ** 2
+            )
+        return field_bar
+
+    return filter_func
+
+
+@dataclass
+class Filter:
+    """A class for applying diffusion-based smoothing filters to gridded data.
+
+    ̦Parameters
+    -----------
+    filter_scale : float
+        The filter scale, which has different meaning depending on filter shape
+    dx_min : float
+        The smallest grid spacing. Should have same units as ``filter_scale``
+    n_steps : int, optional
+        Number of total steps in the filter
+    filter_shape : {"Gaussian", "Taper"}
+        - Gaussian: The target filter has kernel $e^{-|x/Lf|^2}$
+        - 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.
+    grid_type: what sort of grid we are dealing with
+    grid_vars: dictionary of extra parameters used to initialize the grid laplacian
+
+    Attributes
+    ----------
+    filter_spec: FilterSpec
+    """
+
+    filter_scale: float
+    dx_min: float
+    n_steps: int = 40
+    filter_shape: FilterShape = FilterShape.GAUSSIAN
+    transition_width: float = np.pi
+    grid_type: GridType = GridType.CARTESIAN
+    grid_vars: dict = field(default_factory=dict)
+
+    def __post_init__(self):
+
+        if self.n_steps <= 2:
+            raise ValueError("Filter requires N>2")
+
+        self.filter_spec = _compute_filter_spec(
+            self.filter_scale,
+            self.dx_min,
+            self.n_steps,
+            self.filter_shape,
+            self.transition_width,
+        )
+
+        # check that we have all the required grid aguments
+        self.Laplacian = ALL_KERNELS[self.grid_type]
+        if not set(self.Laplacian.required_grid_args()) == set(self.grid_vars):
+            raise ValueError(
+                f"Provided `grid_vars` {list(self.grid_vars)} do not match expected "
+                f"{list(self.Laplacian.required_grid_args())}"
+            )
+        self.grid_ds = xr.Dataset({name: da for name, da in self.grid_vars.items()})
+
+    def apply(self, field, dims):
+        """Filter a field across the dimensions specified by dims."""
+
+        filter_func = _create_filter_func(self.filter_spec, self.Laplacian)
+        grid_args = [self.grid_ds[name] for name in self.Laplacian.required_grid_args()]
+        assert len(dims) == 2
+        n_args = 1 + len(grid_args)
+        field_smooth = xr.apply_ufunc(
+            filter_func,
+            field,
+            *grid_args,
+            input_core_dims=n_args * [dims],
+            output_core_dims=[dims],
+            output_dtypes=[field.dtype],
+            dask="parallelized",
+        )
+        return field_smooth
diff --git a/gcm_filters/gpu_compat.py b/gcm_filters/gpu_compat.py
new file mode 100644
index 0000000..8049c10
--- /dev/null
+++ b/gcm_filters/gpu_compat.py
@@ -0,0 +1,15 @@
+"""GPU compatibility stuff."""
+import numpy as np
+
+
+try:
+    from cupy import get_array_module
+except ImportError:
+
+    def get_array_module(*args):
+        return np
+
+
+# this will work once we require numpy >= 1.20
+# ArrayType = np.typing.ArrayLike
+ArrayType = np.ndarray
diff --git a/gcm_filters/kernels.py b/gcm_filters/kernels.py
index 372046b..07dfcaf 100644
--- a/gcm_filters/kernels.py
+++ b/gcm_filters/kernels.py
@@ -1,36 +1,89 @@
 """
 Core smoothing routines that operate on 2D arrays.
 """
+import enum
 
-import functools
+from abc import ABC
+from dataclasses import dataclass
+from typing import Any, Dict
 
+from .gpu_compat import ArrayType, get_array_module
 
-try:
-    from cupy import get_array_module as _get_array_module
-except ImportError:
-    import numpy as np
 
-    def _get_array_module(*args):
-        return np
+# not married to the term "Cartesian"
+GridType = enum.Enum("GridType", ["CARTESIAN", "CARTESIAN_WITH_LAND"])
 
+ALL_KERNELS = {}  # type: Dict[GridType, Any]
 
-def simple_diffusion_kernel(phi):
-    """Classic diffusion stencil for regular grid.
 
-    Parameters
-    ----------
-    phi : array_like
+@dataclass
+class BaseLaplacian(ABC):
+    def __call__(self, field):
+        pass  # pragma: no cover
+
+    # change to property when we are using python 3.9
+    # https://stackoverflow.com/questions/128573/using-property-on-classmethods
+    @classmethod
+    def required_grid_args(self):
+        try:
+            return list(self.__annotations__)
+        except AttributeError:
+            return []
+
+
+@dataclass
+class CartesianLaplacian(BaseLaplacian):
+    """̵Laplacian for regularly spaced Cartesian grids."""
+
+    def __call__(self, field: ArrayType):
+        np = get_array_module(field)
+        return (
+            -4 * field
+            + np.roll(field, -1, axis=-1)
+            + np.roll(field, 1, axis=-1)
+            + np.roll(field, -1, axis=-2)
+            + np.roll(field, 1, axis=-2)
+        )
+
+
+ALL_KERNELS[GridType.CARTESIAN] = CartesianLaplacian
+
+
+@dataclass
+class CartesianLaplacianWithLandMask(BaseLaplacian):
+    """̵Laplacian for regularly spaced Cartesian grids with land mask.
 
-    Returns
-    -------
-    array_like
-        The diffusive tendency for `phi`
+    Attributes
+    ----------
+    wet_mask: Mask array, 1 for ocean, 0 for land
     """
-    np = _get_array_module(phi)
-    return (
-        -4 * phi
-        + np.roll(phi, -1, axis=-1)
-        + np.roll(phi, 1, axis=-1)
-        + np.roll(phi, -1, axis=-2)
-        + np.roll(phi, 1, axis=-2)
-    )
+
+    wet_mask: ArrayType
+
+    def __call__(self, field: ArrayType):
+        np = get_array_module(field)
+
+        out = field.copy()  # is this necessary?
+        out = np.nan_to_num(field)  # is this necessary?
+        out = self.wet_mask * out
+
+        fac = (
+            np.roll(self.wet_mask, -1, axis=-1)
+            + np.roll(self.wet_mask, 1, axis=-1)
+            + np.roll(self.wet_mask, -1, axis=-2)
+            + np.roll(self.wet_mask, 1, axis=-2)
+        )
+
+        out = (
+            -fac * out
+            + np.roll(out, -1, axis=-1)
+            + np.roll(out, 1, axis=-1)
+            + np.roll(out, -1, axis=-2)
+            + np.roll(out, 1, axis=-2)
+        )
+
+        out = self.wet_mask * out
+        return out
+
+
+ALL_KERNELS[GridType.CARTESIAN_WITH_LAND] = CartesianLaplacianWithLandMask
diff --git a/pyproject.toml b/pyproject.toml
index 85f2732..25e2582 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -2,19 +2,5 @@
 requires = ["setuptools>=42", "wheel", "setuptools_scm[toml]>=3.4"]
 build-backend = "setuptools.build_meta"
 
-[tool.interrogate]
-ignore-init-method = true
-ignore-init-module = false
-ignore-magic = false
-ignore-semiprivate = true
-ignore-private = true
-ignore-property-decorators = true
-ignore-module = false
-fail-under = 95
-exclude = ["setup.py", "docs", "tests"]
-verbose = 1
-quiet = false
-color = true
-
 [tool.isort]
-known_third_party = ["numpy", "pytest", "setuptools"]
+known_third_party = ["numpy", "pytest", "scipy", "setuptools", "xarray"]
diff --git a/requirements.txt b/requirements.txt
index 755e313..c9d3fd4 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,2 +1,3 @@
 dask[array]
+scipy
 xarray


Test patch:
diff --git a/tests/test_filter.py b/tests/test_filter.py
new file mode 100644
index 0000000..73ef8c1
--- /dev/null
+++ b/tests/test_filter.py
@@ -0,0 +1,99 @@
+import numpy as np
+import pytest
+import xarray as xr
+
+from gcm_filters import Filter, FilterShape, GridType
+from gcm_filters.filter import FilterSpec
+
+
+def _check_equal_filter_spec(spec1, spec2):
+    assert spec1.n_lap_steps == spec2.n_lap_steps
+    assert spec1.n_bih_steps == spec2.n_bih_steps
+    np.testing.assert_allclose(spec1.s_l, spec2.s_l)
+    np.testing.assert_allclose(spec1.s_b, spec2.s_b)
+
+
+# These values were just hard copied from my dev environment.
+# All they do is check that the results match what I got when I ran the code.
+# They do NOT assure that the filter spec is correct.
+@pytest.mark.parametrize(
+    "filter_args, expected_filter_spec",
+    [
+        (
+            dict(filter_scale=10.0, dx_min=1.0, n_steps=4),
+            FilterSpec(
+                n_lap_steps=4,
+                s_l=[2.36715983, 8.36124821, 15.17931706, 19.7392088],
+                n_bih_steps=0,
+                s_b=[],
+            ),
+        ),
+        (
+            dict(
+                filter_scale=1.0, dx_min=1.0, n_steps=10, filter_shape=FilterShape.TAPER
+            ),
+            FilterSpec(
+                n_lap_steps=6,
+                s_l=[
+                    9.87341331,
+                    12.66526236,
+                    15.23856752,
+                    17.38217753,
+                    18.91899844,
+                    19.7392088,
+                ],
+                n_bih_steps=2,
+                s_b=[-1.83974928 - 2.24294603j, 1.21758518 - 8.29775049j],
+            ),
+        ),
+    ],
+)
+def test_filter_spec(filter_args, expected_filter_spec):
+    """This test just verifies that the filter specification looks as expected."""
+    filter = Filter(**filter_args)
+    _check_equal_filter_spec(filter.filter_spec, expected_filter_spec)
+    # TODO: check other properties of filter_spec?
+
+
+@pytest.fixture(scope="module", params=list(GridType))
+def grid_type_and_input_ds(request):
+    grid_type = request.param
+
+    ny, nx = (128, 256)
+    data = np.random.rand(ny, nx)
+    da = xr.DataArray(data, dims=["y", "x"])
+
+    grid_vars = {}
+
+    if grid_type == GridType.CARTESIAN_WITH_LAND:
+        mask_data = np.ones_like(data)
+        mask_data[: (ny // 2), : (nx // 2)] = 0
+        da_mask = xr.DataArray(mask_data, dims=["y", "x"])
+        grid_vars = {"wet_mask": da_mask}
+
+    return grid_type, da, grid_vars
+
+
+@pytest.mark.parametrize(
+    "filter_args",
+    [dict(filter_scale=1.0, dx_min=1.0, n_steps=10, filter_shape=FilterShape.TAPER)],
+)
+def test_filter(grid_type_and_input_ds, filter_args):
+    grid_type, da, grid_vars = grid_type_and_input_ds
+    filter = Filter(grid_type=grid_type, grid_vars=grid_vars, **filter_args)
+    filtered = filter.apply(da, dims=["y", "x"])
+
+    # check conservation
+    # this would need to be replaced by a proper area-weighted integral
+    xr.testing.assert_allclose(da.sum(), filtered.sum())
+
+    # check variance reduction
+    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:
+        grid_vars_missing = {k: v for k, v in grid_vars.items() if k != gv}
+        with pytest.raises(ValueError, match=r"Provided `grid_vars` .*"):
+            filter = Filter(
+                grid_type=grid_type, grid_vars=grid_vars_missing, **filter_args
+            )
diff --git a/tests/test_kernels.py b/tests/test_kernels.py
index be8f09f..c43b4b5 100644
--- a/tests/test_kernels.py
+++ b/tests/test_kernels.py
@@ -1,17 +1,27 @@
 import numpy as np
 import pytest
 
-from gcm_filters.kernels import simple_diffusion_kernel
+from gcm_filters.kernels import ALL_KERNELS, GridType
 
 
-@pytest.fixture
-def numpy_array_2d():
+@pytest.fixture(scope="module", params=list(GridType))
+def grid_type_field_and_extra_kwargs(request):
+    grid_type = request.param
+
     ny, nx = (128, 256)
     data = np.random.rand(ny, nx)
-    return data
+
+    extra_kwargs = {}
+    if grid_type == GridType.CARTESIAN_WITH_LAND:
+        mask_data = np.ones_like(data)
+        mask_data[: (ny // 2), : (nx // 2)] = 0
+        extra_kwargs["wet_mask"] = mask_data
+    return grid_type, data, extra_kwargs
 
 
-@pytest.mark.parametrize("kernel_function", [simple_diffusion_kernel])
-def test_conservation(numpy_array_2d, kernel_function):
-    res = kernel_function(numpy_array_2d)
+def test_conservation(grid_type_field_and_extra_kwargs):
+    grid_type, data, extra_kwargs = grid_type_field_and_extra_kwargs
+    LaplacianClass = ALL_KERNELS[grid_type]
+    laplacian = LaplacianClass(**extra_kwargs)
+    res = laplacian(data)
     np.testing.assert_allclose(res.sum(), 0.0, atol=1e-12)