Add periodic boundary conditions

Author: jrmaddison

bt_ocean/finite_difference.py

@@ -10,7 +10,8 @@
 
 __all__ = \
     [
-        "diff"
+        "diff_bounded",
+        "diff_periodic"
     ]
 
 
@@ -51,11 +52,11 @@ def difference_coefficients(beta, order):
 
 
 @partial(jax.jit, static_argnames={"order", "N", "axis", "i0", "i1", "boundary_expansion"})
-def diff(u, dx, order, N, *, axis=-1, i0=None, i1=None, boundary_expansion=None):
+def diff_bounded(u, dx, order, N, *, axis=-1, i0=None, i1=None, boundary_expansion=None):
     """Compute a centred finite difference approximation to a derivative for
-    data stored on a uniform grid. Transitions to one-sided differencing as the
-    end-points are approached. Selects an additional right-sided point if
-    `N` is even.
+    data stored on a uniform grid. Result is defined on the same grid as the
+    input (i.e. without staggering). Transitions to one-sided differencing as
+    the end-points are approached.
 
     Parameters
     ----------
@@ -67,7 +68,8 @@ def diff(u, dx, order, N, *, axis=-1, i0=None, i1=None, boundary_expansion=None)
     order : Integral
         Derivative order.
     N : Integral
-        Number of grid points in the difference approximation.
+        Number of grid points in the difference approximation. Centered
+        differencing uses an additional right-sided point if `N` is even.
     axis : Integral
         Axis.
     i0 : Integral
@@ -115,7 +117,7 @@ def diff(u, dx, order, N, *, axis=-1, i0=None, i1=None, boundary_expansion=None)
     i1 = i0 + N
     parity = (-1) ** order
 
-    for i in range(max(-i0, i1 - 1)):
+    for i in range(max(0, min(i0_b, u.shape[-1] - i1_b)), max(-i0, i1 - 1)):
         beta = tuple(range(-i, -i + N + int(bool(boundary_expansion))))
         alpha = tuple(map(dtype, difference_coefficients(beta, order)))
         if i < -i0 and i >= i0_b:
@@ -130,7 +132,7 @@ def diff(u, dx, order, N, *, axis=-1, i0=None, i1=None, boundary_expansion=None)
                 v = v.at[..., u.shape[-1] - 1 - i].add(
                     parity * alpha_j * u[..., u.shape[-1] - 1 - i - beta_j])
 
-    # Center
+    # Center points
     beta = tuple(range(i0, i1))
     alpha = tuple(map(dtype, difference_coefficients(beta, order)))
     i0_c = max(-i0, i0_b)
@@ -142,3 +144,35 @@ def diff(u, dx, order, N, *, axis=-1, i0=None, i1=None, boundary_expansion=None)
 
     v = jnp.moveaxis(v, -1, axis)
     return v / (dx ** order)
+
+
+@partial(jax.jit, static_argnames={"order", "N", "axis"})
+def diff_periodic(u, dx, order, N, *, axis=-1):
+    """Compute a centred finite difference approximation to a derivative for
+    data stored on a uniform grid. Result is defined on the same grid as the
+    input (i.e. without staggering). Applies periodic boundary conditions.
+
+    Arguments and return value are as for :func:`.diff_bounded`.
+    """
+
+    if axis < 0:
+        axis = len(u.shape) + axis
+    if axis < 0 or axis >= len(u.shape):
+        raise ValueError("Invalid axis")
+    if u.shape[axis] < N:
+        raise ValueError("Insufficient points")
+
+    u = jnp.moveaxis(u, axis, -1)
+    i0 = -(N // 2)
+    i1 = i0 + N
+
+    # Periodic extension
+    u_e = jnp.zeros_like(u, shape=u.shape[:-1] + (u.shape[-1] + N,))
+    u_e = u_e.at[..., -i0:-i1].set(u)
+    u_e = u_e.at[..., :-i0].set(u[..., i0:])
+    u_e = u_e.at[..., -i1:].set(u[..., :i1])
+
+    v = diff_bounded(u_e, dx, order, N, axis=-1, i0=-i0, i1=-i1)[..., -i0:-i1]
+
+    v = jnp.moveaxis(v, -1, axis)
+    return v

bt_ocean/grid.py

@@ -7,7 +7,7 @@
 from functools import cached_property, partial
 from numbers import Real
 
-from .finite_difference import diff
+from .finite_difference import diff_bounded as diff
 from .precision import default_idtype, default_fdtype
 
 __all__ = \

tests/test_finite_difference.py

@@ -1,5 +1,5 @@
 from bt_ocean.finite_difference import (
-    difference_coefficients, diff as centered_difference_bounded)
+    difference_coefficients, diff_bounded, diff_periodic)
 from bt_ocean.precision import default_fdtype
 
 import jax
@@ -53,7 +53,7 @@ def test_centered_difference_monomials(alpha):
                 for i in range(order):
                     diff_exact *= p - i
             for N in range(max(p + 1, order + 1), x.shape[0]):
-                diff = centered_difference_bounded(u, dx, order=order, N=N)
+                diff = diff_bounded(u, dx, order=order, N=N)
                 error_norm = abs(diff - diff_exact).max()
                 diff_exact_norm = abs(diff_exact).max()
                 if diff_exact_norm > 0:
@@ -77,12 +77,12 @@ def test_centered_difference_convergence():
         W = W.at[-1].set(0.5 * dx)
         u = jnp.sin(jnp.pi * x)
 
-        D1_error = (centered_difference_bounded(u, dx, order=1, N=3)
+        D1_error = (diff_bounded(u, dx, order=1, N=3)
                     - jnp.pi * jnp.cos(jnp.pi * x))
         error_norms_1 = error_norms_1.at[i].set(jnp.sqrt((W * (D1_error ** 2)).sum()))  # noqa: E501
         print(f"{p=:d} {error_norms_1[i]=:.6g}")
 
-        D2_error = (centered_difference_bounded(u, dx, order=2, N=3)
+        D2_error = (diff_bounded(u, dx, order=2, N=3)
                     + (jnp.pi ** 2) * jnp.sin(jnp.pi * x))
         error_norms_2 = error_norms_2.at[i].set(jnp.sqrt((W * (D2_error ** 2)).sum()))  # noqa: E501
         print(f"{p=:d} {error_norms_2[i]=:.6g}")
@@ -92,3 +92,33 @@ def test_centered_difference_convergence():
     print(f"{orders_2=}")
     assert orders_1.min() > 2
     assert orders_2.min() > 2
+
+
+def test_centered_difference_convergence_periodic():
+    if default_fdtype() != np.float64 or not jax.config.x64_enabled:
+        pytest.skip("float64 not available")
+
+    P = jnp.arange(7, 12, dtype=int)
+    error_norms_1 = jnp.zeros_like(P, dtype=float)
+    error_norms_2 = jnp.zeros_like(P, dtype=float)
+    for i, p in enumerate(P):
+        x = jnp.linspace(0, 1, 2 ** p + 1, dtype=float)[:-1]
+        dx = x[1] - x[0]
+        W = jnp.full_like(x, dx)
+        u = jnp.sin(2 * jnp.pi * x)
+
+        D1_error = (diff_periodic(u, dx, order=1, N=3)
+                    - 2 * jnp.pi * jnp.cos(2 * jnp.pi * x))
+        error_norms_1 = error_norms_1.at[i].set(jnp.sqrt((W * (D1_error ** 2)).sum()))  # noqa: E501
+        print(f"{p=:d} {error_norms_1[i]=:.6g}")
+
+        D2_error = (diff_periodic(u, dx, order=2, N=3)
+                    + 4 * (jnp.pi ** 2) * jnp.sin(2 * jnp.pi * x))
+        error_norms_2 = error_norms_2.at[i].set(jnp.sqrt((W * (D2_error ** 2)).sum()))  # noqa: E501
+        print(f"{p=:d} {error_norms_2[i]=:.6g}")
+    orders_1 = jnp.log2(error_norms_1[:-1] / error_norms_1[1:])
+    orders_2 = jnp.log2(error_norms_2[:-1] / error_norms_2[1:])
+    print(f"{orders_1=}")
+    print(f"{orders_2=}")
+    assert orders_1.min() > 1.99
+    assert orders_2.min() > 1.99