swesphere 0.3.0__py3-none-any.whl

swesphere/__init__.py

@@ -0,0 +1,13 @@
+"""swesphere: shallow-water equations on the sphere for data-assimilation testbeds.
+
+    from swesphere import presets, climatology
+    model, x0 = presets.two_jets()
+    S = climatology.build_climatology(model, x0, spinup_days=40, n_snapshots=200, every_days=2)
+"""
+from .model import SWEModel
+from .grid import SphereGrid
+from .forcing import ZonalRelaxation
+from . import presets, climatology, initial_conditions, integrators, dynamics, diagnostics
+
+__version__ = "0.3.0"
+__all__ = ["SWEModel", "SphereGrid", "ZonalRelaxation", "presets", "climatology", "initial_conditions", "integrators", "dynamics", "diagnostics", "__version__"]

swesphere/climatology.py

@@ -0,0 +1,38 @@
+# -*- coding: utf-8 -*-
+"""Climatology builder: spin-up, snapshots, decorrelation check, normalization."""
+from __future__ import annotations
+
+import numpy as np
+
+DAY = 86400.0
+
+
+def build_climatology(model, x0, spinup_days, n_snapshots, every_days, verbose=False, label=""):
+    """Spin up ``spinup_days`` from ``x0`` and keep ``n_snapshots`` states ``every_days`` apart (float32).
+
+    For a forced preset this samples a stationary regime. For the unforced ``waves`` preset the record is a
+    slowly decaying flow (the filter is the only sink), so its early and late states differ in amplitude.
+    """
+    if verbose:
+        print(f"  {label} spin-up of {spinup_days:g} days", flush=True)
+    x = model.propagate(x0, [0.0, spinup_days * DAY])
+    if verbose:
+        print(f"  {label} spin-up done", flush=True)
+    snaps = [x.astype(np.float32)]
+    for k in range(n_snapshots - 1):
+        x = model.propagate(x, [0.0, every_days * DAY]); snaps.append(x.astype(np.float32))
+        if verbose and (k + 1) % 10 == 0:
+            print(f"  {label} {k + 1}/{n_snapshots} snapshots", flush=True)
+    return np.array(snaps)
+
+
+def anomaly_scales(model, snapshots):
+    """Per-field standard deviation of the climatological anomalies, repeated over the state."""
+    A = snapshots - snapshots.mean(axis=0); fs = model.field_size
+    sd = np.array([A[:, k * fs:(k + 1) * fs].std() for k in range(3)])
+    return np.repeat(sd, fs), sd
+
+
+def lag_correlation(model, snapshots, lag=1, field="h"):
+    A = snapshots - snapshots.mean(axis=0); blk = model.var_blocks[field]; a = A[:, blk]
+    return float(np.mean([np.corrcoef(a[i], a[i + lag])[0, 1] for i in range(len(a) - lag)]))

swesphere/diagnostics.py

@@ -0,0 +1,39 @@
+# -*- coding: utf-8 -*-
+"""Global invariants of the shallow-water equations and error norms.
+
+All integrals are area weighted (cos φ) over the whole grid, or over the rows
+outside the polar sponge with ``interior=True``. The continuous equations
+conserve mass, total energy and potential enstrophy; the discrete model loses a
+little of each through the spectral filter and the polar sponge, and the forced
+presets exchange them with the relaxation. EXP-00 reports the measured drifts.
+"""
+from __future__ import annotations
+
+import numpy as np
+
+from .dynamics import relative_vorticity
+
+
+def area_weights(model, interior: bool = False) -> np.ndarray:
+    w = np.clip(np.asarray(model.grid.cosL), 0.0, None).copy()
+    if interior:
+        k = model.n_pole_rows; w[:k] = 0.0; w[model.Nlat - k:] = 0.0
+    return w / w.sum()
+
+
+def invariants(model, x, interior: bool = False) -> dict:
+    """Mean depth (m), total energy per unit area (J/kg m) and potential enstrophy."""
+    u, v, h = model.unpack(np.asarray(x, dtype=float)); G = model.grid; w = area_weights(model, interior)
+    eta = relative_vorticity(u, v, G) + G.f_cor
+    return dict(mass=float((w * h).sum()),
+                energy=float((w * (0.5 * h * (u ** 2 + v ** 2) + 0.5 * G.g * h ** 2)).sum()),
+                enstrophy=float((w * 0.5 * eta ** 2 / h).sum()))
+
+
+def error_norms(model, x, x_ref, field: str = "h", interior: bool = True) -> dict:
+    """Normalized l1, l2 and l-infinity errors of Williamson et al. (1992)."""
+    i = "uvh".index(field); q = model.unpack(np.asarray(x, dtype=float))[i]; r = model.unpack(np.asarray(x_ref, dtype=float))[i]
+    w = area_weights(model, interior); m = w > 0
+    return dict(l1=float((w * np.abs(q - r)).sum() / (w * np.abs(r)).sum()),
+                l2=float(np.sqrt((w * (q - r) ** 2).sum() / (w * r ** 2).sum())),
+                linf=float(np.abs(q - r)[m].max() / np.abs(r)[m].max()))

swesphere/dynamics.py

@@ -0,0 +1,88 @@
+# -*- coding: utf-8 -*-
+"""
+Vector-invariant form of the shallow-water equations on the sphere.
+
+State variables
+---------------
+* ``u(λ, φ)``, zonal velocity (m/s)
+* ``v(λ, φ)``, meridional velocity (m/s)
+* ``h(λ, φ)``, layer thickness (m)
+
+Equations (Hack & Jakob 1992 form)
+----------------------------------
+    ∂u/∂t = η v − (1 / a cosφ) ∂B/∂λ
+    ∂v/∂t = −η u − (1 / a) ∂B/∂φ
+    ∂h/∂t = −∇·(h V) = −(1 / a cosφ) [∂(h u)/∂λ + ∂(h v cosφ)/∂φ]
+
+where:
+    ζ = (1 / a cosφ) ∂v/∂λ − (1/a) ∂u/∂φ + (tanφ/a) u    (relative vorticity)
+    η = ζ + f                                            (absolute vorticity)
+    B = g h + ½(u² + v²)                                 (Bernoulli potential)
+    δ = (1 / a cosφ) ∂u/∂λ + (1 / a cosφ) ∂(v cosφ)/∂φ   (divergence, diagnostic only)
+
+The continuity equation is discretized in flux form, so the zonal FFT derivative and
+the centered meridional difference telescope and the global mass changes only through
+the spectral filter and the polar rows. (Versions before 0.2.0 used ``dh = -h δ`` and
+omitted the advection of ``h``; see CHANGELOG.md.)
+"""
+
+from __future__ import annotations
+
+from typing import Tuple
+
+import numpy as np
+
+from .grid import SphereGrid, d_lon, d_lat
+
+
+def rhs(
+    u: np.ndarray,
+    v: np.ndarray,
+    h: np.ndarray,
+    grid: SphereGrid,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Return ``(du/dt, dv/dt, dh/dt)`` for the SWE system."""
+    cos = grid.cos_safe
+    a = grid.a
+
+    # Relative vorticity ζ
+    zeta = (d_lon(v, grid) / (a * cos)
+            - d_lat(u, grid) / a
+            + grid.sinL / cos * u / a)
+
+    # Absolute vorticity η
+    eta = zeta + grid.f_cor
+
+    # Bernoulli potential
+    B = grid.g * h + 0.5 * (u ** 2 + v ** 2)
+
+    # Tendencies: vector-invariant momentum, flux-form continuity
+    du = eta * v - d_lon(B, grid) / (a * cos)
+    dv = -eta * u - d_lat(B, grid) / a
+    dh = -(d_lon(h * u, grid) + d_lat(h * v * grid.cosL, grid)) / (a * cos)
+
+    return du, dv, dh
+
+
+def relative_vorticity(u: np.ndarray, v: np.ndarray,
+                       grid: SphereGrid) -> np.ndarray:
+    """Diagnostic: relative vorticity ζ from (u, v)."""
+    cos = grid.cos_safe
+    a = grid.a
+    return (d_lon(v, grid) / (a * cos)
+            - d_lat(u, grid) / a
+            + grid.sinL / cos * u / a)
+
+
+def divergence(u: np.ndarray, v: np.ndarray,
+               grid: SphereGrid) -> np.ndarray:
+    """Diagnostic: horizontal divergence δ from (u, v)."""
+    cos = grid.cos_safe
+    a = grid.a
+    return (d_lon(u, grid) / (a * cos)
+            + d_lat(v * grid.cosL, grid) / (a * cos))
+
+
+def divergence(u: np.ndarray, v: np.ndarray, grid: SphereGrid) -> np.ndarray:
+    """Diagnostic: horizontal divergence δ from (u, v)."""
+    return (d_lon(u, grid) + d_lat(v * grid.cosL, grid)) / (grid.a * grid.cos_safe)

swesphere/forcing.py

@@ -0,0 +1,26 @@
+# -*- coding: utf-8 -*-
+"""Forcing that gives the model a stationary turbulent regime.
+
+Without forcing the shallow-water equations on the sphere do not sustain eddies: the filter dissipates
+whatever the initial condition put in and the flow relaxes toward a zonal state with weak, slowly
+decaying waves (the ``waves`` preset). ``ZonalRelaxation`` relaxes only the zonal mean of (u, v, h)
+toward a reference jet and leaves the eddies alone; with a barotropically unstable reference
+(Galewsky) the instability regenerates eddies against the dissipation and the flow reaches a
+stationary turbulent state (EXP-02 measures its statistics). Relaxing the full fields toward the
+zonal state kills the eddies instead.
+"""
+from __future__ import annotations
+
+import numpy as np
+
+
+class ZonalRelaxation:
+    def __init__(self, u_ref, v_ref, h_ref, tau):
+        self.k = 1.0 / float(tau)
+        self.ref = tuple(np.repeat(np.asarray(f).mean(axis=1, keepdims=True), np.asarray(f).shape[1], axis=1) for f in (u_ref, v_ref, h_ref))
+
+    def __call__(self, u, v, h, du, dv, dh):
+        for f, r_, d_ in ((u, self.ref[0], du), (v, self.ref[1], dv), (h, self.ref[2], dh)):
+            zm = np.repeat(f.mean(axis=1, keepdims=True), f.shape[1], axis=1)
+            d_ -= self.k * (zm - r_)
+        return du, dv, dh

swesphere/grid.py

@@ -0,0 +1,161 @@
+# -*- coding: utf-8 -*-
+"""
+Driscoll-Healy spherical grid and numerical operators.
+
+Geometry
+--------
+* Grid is stored North-to-South with `Nlat = 2*(LMAX+1)` rows and
+  `Nlon = 4*(LMAX+1)` columns.
+* Latitudes range from +90° (row 0) down to -90° (row Nlat-1), exclusive
+  of the south pole (DH convention).
+* Longitudes are equally spaced 0 to 360°, exclusive of 360.
+
+Operators
+---------
+* ``d_lon(q)``, exact zonal derivative via FFT.
+* ``d_lat(q)``, centered finite-difference meridional derivative
+  (with sign flip for N→S storage).
+* ``sh_filter(q, alpha, p)``, exponential filter in spherical harmonics:
+  attenuates degree n by ``exp(-alpha * (n/Lmax)^p)``. Removes aliasing
+  and polar noise.
+* ``make_sponge(...)``, multiplicative damping near the poles.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Tuple
+
+import numpy as np
+
+import pyshtools as pysh
+
+
+# ----------------------------------------------------------------------
+# Grid
+# ----------------------------------------------------------------------
+@dataclass
+class SphereGrid:
+    """Driscoll-Healy spherical grid metadata.
+
+    Attributes
+    ----------
+    LMAX : int
+        Spectral truncation (T-LMAX).
+    Nlat, Nlon : int
+        Grid dimensions in latitude and longitude.
+    a : float
+        Earth radius (m).
+    Omega : float
+        Earth angular velocity (rad/s).
+    g : float
+        Gravitational acceleration (m/s^2).
+    LON_ns, LAT_ns : ndarray
+        Meshgrids of longitude and latitude in radians, stored N→S.
+    cosL, sinL, f_cor : ndarray
+        Geometric quantities (cos φ, sin φ, Coriolis parameter).
+    cos_safe : ndarray
+        cos φ with a small minimum value to avoid polar singularity.
+    kx : ndarray
+        Zonal wavenumbers for FFT.
+    dlat : float
+        Meridional grid spacing in radians.
+    """
+    LMAX: int
+    Nlat: int
+    Nlon: int
+    a: float
+    Omega: float
+    g: float
+    LON_ns: np.ndarray
+    LAT_ns: np.ndarray
+    cosL: np.ndarray
+    sinL: np.ndarray
+    f_cor: np.ndarray
+    cos_safe: np.ndarray
+    kx: np.ndarray
+    dlat: float
+
+    @classmethod
+    def make(
+        cls,
+        LMAX: int = 32,
+        a: float = 6.371e6,
+        Omega: float = 7.292e-5,
+        g: float = 9.81,
+        cos_floor: float = 0.08,
+    ) -> "SphereGrid":
+        Nlat = 2 * (LMAX + 1)
+        Nlon = 4 * (LMAX + 1)
+        lats = 90 - np.arange(Nlat) * 180 / Nlat            # degrees, N→S
+        lons = np.arange(Nlon) * 360 / Nlon
+        LON_ns, LAT_ns = np.meshgrid(np.radians(lons), np.radians(lats))
+        cosL = np.cos(LAT_ns)
+        sinL = np.sin(LAT_ns)
+        f_cor = 2 * Omega * sinL
+        cos_safe = np.where(np.abs(cosL) < cos_floor,
+                            np.sign(cosL + 1e-10) * cos_floor, cosL)
+        kx = np.fft.rfftfreq(Nlon, d=1.0 / Nlon)
+        dlat = np.radians(180 / Nlat)
+        return cls(LMAX=LMAX, Nlat=Nlat, Nlon=Nlon, a=a, Omega=Omega, g=g,
+                   LON_ns=LON_ns, LAT_ns=LAT_ns, cosL=cosL, sinL=sinL,
+                   f_cor=f_cor, cos_safe=cos_safe, kx=kx, dlat=dlat)
+
+
+# ----------------------------------------------------------------------
+# Operators
+# ----------------------------------------------------------------------
+def d_lon(q: np.ndarray, grid: SphereGrid) -> np.ndarray:
+    """Exact zonal derivative ∂q/∂λ via FFT."""
+    return np.fft.irfft(1j * grid.kx * np.fft.rfft(q, axis=1),
+                        n=grid.Nlon, axis=1)
+
+
+def d_lat(q: np.ndarray, grid: SphereGrid) -> np.ndarray:
+    """Centered finite-difference meridional derivative ∂q/∂φ.
+
+    The grid is stored N→S, so dq/d(row index) has opposite sign to
+    dq/d(latitude), we flip the sign.
+    """
+    dq = np.zeros_like(q)
+    dq[1:-1] = (q[2:] - q[:-2]) / (2 * grid.dlat)
+    dq[0] = (q[1] - q[0]) / grid.dlat
+    dq[-1] = (q[-1] - q[-2]) / grid.dlat
+    return -dq
+
+
+def sh_filter(
+    q: np.ndarray,
+    grid: SphereGrid,
+    alpha: float = 36.0,
+    p: int = 16,
+) -> np.ndarray:
+    """Exponential filter in spherical harmonics.
+
+    Attenuates spherical-harmonic degree ``n`` by ``exp(-alpha*(n/LMAX)^p)``.
+    With ``p=16`` this is essentially a brick-wall filter that passes
+    ``n < LMAX`` and cuts ``n ~ LMAX``. Removes aliasing and polar noise
+    that builds up under explicit time integration.
+    """
+    L = grid.LMAX
+    c = pysh.expand.SHExpandDH(q, norm=1, sampling=2, lmax_calc=L)
+    n_arr = np.arange(L + 1)
+    decay = np.exp(-alpha * (n_arr / max(L, 1)) ** p)
+    c = c * decay[None, :, None]
+    return pysh.expand.MakeGridDH(c, norm=1, sampling=2, lmax=L)
+
+
+def make_sponge(Nlat: int, n_pole_rows: int = 10,
+                strength: float = 0.98) -> np.ndarray:
+    """Build a multiplicative damping mask near the poles.
+
+    Returns an array of shape ``(Nlat, 1)`` that is < 1 within
+    ``n_pole_rows`` of each pole, ramping linearly from full damping
+    (``1 - strength``) at the pole to 1 at the edge of the sponge band.
+    """
+    sponge = np.ones((Nlat, 1))
+    for i in range(n_pole_rows):
+        w = (n_pole_rows - i) / n_pole_rows * strength
+        sponge[i] = 1 - w
+        sponge[-1 - i] = 1 - w
+    return sponge

swesphere/initial_conditions.py

@@ -0,0 +1,160 @@
+# -*- coding: utf-8 -*-
+"""
+Initial conditions for the SWE on the sphere.
+
+Two ingredients:
+
+* **Solid-body rotation** (Williamson et al. 1992, Test Case 2). An
+  exact stationary solution of the SWE on the sphere::
+
+       u = U0 · cos φ
+       v = 0
+       h = H0 − (a Ω U0 + U0² / 2) / g · sin² φ
+
+  The default base state of every IC.
+
+* **Rossby-wave perturbations** in each hemisphere, gaussian-enveloped
+  in latitude and zonally sinusoidal::
+
+       δh = A_h · cos(k λ) · exp(−(φ − φ₀)² / (2 σ²))
+       δu = A_uv · cos(k λ) · exp(...)
+       δv = A_uv · sin(k λ) · exp(...)
+
+  Default amplitudes and shapes match the reference setup from the
+  upstream visualization script.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Tuple
+
+import numpy as np
+
+from .grid import SphereGrid
+
+
+@dataclass
+class RossbyPert:
+    """Rossby-wave-like perturbation envelope and amplitudes."""
+    amp_h: float = 220.0      # Amplitude in h (m)
+    amp_uv: float = 6.0       # Amplitude in (u, v) (m/s)
+    wavenumber: int = 4       # Zonal wavenumber
+    lat_center: float = 42.0  # Latitude of the gaussian center (deg)
+    width: float = 12.0       # Gaussian width (deg)
+    phase: float = 0.0        # Phase offset (rad)
+
+
+def williamson_tc2(grid: SphereGrid, U0: float = 38.0,
+                   H0: float = 2800.0
+                   ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Williamson 1992 Test Case 2: solid-body rotation."""
+    u = U0 * grid.cosL
+    v = np.zeros_like(u)
+    h = H0 - (grid.a * grid.Omega * U0 + 0.5 * U0 ** 2) / grid.g * grid.sinL ** 2
+    return u, v, h
+
+
+def add_rossby_perturbation(
+    u: np.ndarray, v: np.ndarray, h: np.ndarray,
+    grid: SphereGrid, pert: RossbyPert,
+):
+    """Add a single Rossby-wave perturbation in-place."""
+    phi0 = np.radians(pert.lat_center)
+    sigma = np.radians(pert.width)
+    env = np.exp(-((grid.LAT_ns - phi0) ** 2) / (2 * sigma ** 2))
+    cos_arg = pert.wavenumber * grid.LON_ns + pert.phase
+    sin_arg = pert.wavenumber * grid.LON_ns + pert.phase
+    h += pert.amp_h * np.cos(cos_arg) * env
+    u += pert.amp_uv * np.cos(cos_arg) * env
+    v += pert.amp_uv * np.sin(sin_arg) * env
+
+
+def default_ic(grid: SphereGrid, U0: float = 38.0,
+               H0: float = 2800.0
+               ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Default IC: TC2 + NH wavenumber-4 + SH wavenumber-3 perturbations."""
+    u, v, h = williamson_tc2(grid, U0=U0, H0=H0)
+    add_rossby_perturbation(u, v, h, grid, RossbyPert(
+        amp_h=220.0, amp_uv=6.0, wavenumber=4,
+        lat_center=42.0, width=12.0, phase=0.0,
+    ))
+    add_rossby_perturbation(u, v, h, grid, RossbyPert(
+        amp_h=150.0, amp_uv=0.0, wavenumber=3,
+        lat_center=-38.0, width=11.0, phase=1.1,
+    ))
+    # Small extra meridional component on the SH band, matching the
+    # reference setup (proportional to amp_h / 50 in the original).
+    phi2 = np.radians(-38.0)
+    env2 = np.exp(-((grid.LAT_ns - phi2) ** 2) / (2 * np.radians(11.0) ** 2))
+    v += (150.0 / 50) * np.sin(3 * grid.LON_ns + 1.1) * env2
+    return u, v, h
+
+
+def perturbed_ic(
+    grid: SphereGrid,
+    seed: int = 0,
+    U0: float = 38.0,
+    H0: float = 2800.0,
+    pert_strength: float = 0.05,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Default IC plus a small random perturbation of ``u`` and ``h`` at zonal wavenumbers 1-8,
+    confined to low and middle latitudes, to start runs from different states.
+
+    The perturbations of ``u`` and ``h`` have independent phases, so they are not in geostrophic
+    balance: they excite weak gravity waves that the model radiates and filters within a few days.
+    Use a spin-up before sampling.
+    """
+    u, v, h = default_ic(grid, U0=U0, H0=H0)
+    rng = np.random.default_rng(seed)
+
+    # Low zonal wavenumbers, gaussian envelope in latitude.
+    Nlon = grid.Nlon
+    n_modes = 8
+    for k in range(1, n_modes + 1):
+        amp_u = pert_strength * U0 * rng.standard_normal()
+        amp_h = pert_strength * H0 * 0.05 * rng.standard_normal()
+        phase_u = 2 * np.pi * rng.random()
+        phase_h = 2 * np.pi * rng.random()
+        # Latitude-banded modulation
+        lat_mod = np.exp(-(grid.LAT_ns / np.radians(45.0)) ** 2)
+        u += amp_u * np.cos(k * grid.LON_ns + phase_u) * lat_mod
+        h += amp_h * np.cos(k * grid.LON_ns + phase_h) * lat_mod
+    return u, v, h
+
+
+# ---------------------------------------------------------------------------
+# Galewsky, Scott & Polvani (2004): a balanced, barotropically unstable jet
+# ---------------------------------------------------------------------------
+def galewsky_jet(grid: SphereGrid, umax: float = 80.0, phi0_deg: float = 25.7, phi1_deg: float = 64.6,
+                 h0: float = 10000.0, bump: bool = True, both_hemispheres: bool = True):
+    """Balanced zonal jet(s) with the small height bump that triggers the instability.
+
+    Returns ``(u, v, h)`` on the grid. With ``both_hemispheres`` a mirror jet
+    (also westerly) is placed in the south so the whole sphere is active; the
+    two hemispheres perturb each other across the equator and the resulting
+    turbulence is irregular (with the zonal relaxation of :mod:`swesphere.forcing`).
+    """
+    Nlat, Nlon = grid.Nlat, grid.Nlon
+    lats = np.radians(90 - np.arange(Nlat) * 180 / Nlat); lons = np.radians(np.arange(Nlon) * 360 / Nlon)
+    phi0, phi1 = np.radians(phi0_deg), np.radians(phi1_deg); en = np.exp(-4.0 / (phi1 - phi0) ** 2)
+
+    def uprof(phi):
+        out = np.zeros_like(phi); m = (phi > phi0) & (phi < phi1)
+        out[m] = umax / en * np.exp(1.0 / ((phi[m] - phi0) * (phi[m] - phi1)))
+        if both_hemispheres:
+            ms = (-phi > phi0) & (-phi < phi1)
+            out[ms] = umax / en * np.exp(1.0 / ((-phi[ms] - phi0) * (-phi[ms] - phi1)))
+        return out
+
+    fine = np.linspace(-np.pi / 2, np.pi / 2, 4001); uf = uprof(fine); f = 2 * grid.Omega * np.sin(fine)
+    dh = -grid.a * uf * (f + uf * np.tan(fine) / grid.a) / grid.g
+    hf = np.concatenate([[0.0], np.cumsum(0.5 * (dh[1:] + dh[:-1]) * np.diff(fine))]); hf = h0 + hf - hf.mean()
+    u = np.repeat(uprof(lats)[:, None], Nlon, axis=1); v = np.zeros_like(u)
+    h = np.repeat(np.interp(lats, fine, hf)[:, None], Nlon, axis=1)
+    if bump:
+        LON, LAT = np.meshgrid(lons, lats); phi2 = np.pi / 4; alpha, beta = 1 / 3, 1 / 15
+        h = h + 120.0 * np.cos(LAT) * np.exp(-((LON - np.pi) / alpha) ** 2) * np.exp(-((phi2 - LAT) / beta) ** 2)
+        if both_hemispheres:
+            h = h + 120.0 * np.cos(LAT) * np.exp(-((LON - 0.4 * np.pi) / alpha) ** 2) * np.exp(-((-phi2 - LAT) / beta) ** 2)
+    return u, v, h

swesphere/integrators.py

@@ -0,0 +1,22 @@
+# -*- coding: utf-8 -*-
+"""Time stepper for the shallow-water right-hand side: fourth-order Runge-Kutta.
+
+RK4 with the spectral filter and the polar sponge integrates every preset for hundreds of days at
+dt = 120 s. EXP-01 in ``experiments/`` measures the stable range of dt (a property of this scheme
+against the fastest gravity waves; another scheme has another range) and how the eddy amplitude
+depends on the filter cadence.
+"""
+from __future__ import annotations
+
+
+def rk4(rhs, u, v, h, dt):
+    k1 = rhs(u, v, h)
+    k2 = rhs(u + 0.5 * dt * k1[0], v + 0.5 * dt * k1[1], h + 0.5 * dt * k1[2])
+    k3 = rhs(u + 0.5 * dt * k2[0], v + 0.5 * dt * k2[1], h + 0.5 * dt * k2[2])
+    k4 = rhs(u + dt * k3[0], v + dt * k3[1], h + dt * k3[2])
+    return (u + dt / 6 * (k1[0] + 2 * k2[0] + 2 * k3[0] + k4[0]),
+            v + dt / 6 * (k1[1] + 2 * k2[1] + 2 * k3[1] + k4[1]),
+            h + dt / 6 * (k1[2] + 2 * k2[2] + 2 * k3[2] + k4[2]))
+
+
+INTEGRATORS = {"rk4": rk4}

swesphere/model.py

@@ -0,0 +1,89 @@
+# -*- coding: utf-8 -*-
+"""The shallow-water model on the sphere with a state-vector API.
+
+State layout: ``x = [u.ravel(), v.ravel(), h.ravel()]`` on the ``(Nlat, Nlon)``
+grid, row-major, rows from north to south. ``field_size = Nlat * Nlon``.
+"""
+from __future__ import annotations
+
+import numpy as np
+
+from .grid import SphereGrid, sh_filter, make_sponge
+from .dynamics import rhs as _rhs
+from .integrators import INTEGRATORS
+from . import initial_conditions as ics
+
+
+class SWEModel:
+    def __init__(self, LMAX: int = 32, dt: float = 120.0, scheme: str = "rk4",
+                 filter_every: int = 3, filter_alpha: float = 36.0, filter_p: int = 16,
+                 n_pole_rows: int = 10, forcing=None, a: float = 6.371e6, Omega: float = 7.292e-5, g: float = 9.81):
+        self.grid = SphereGrid.make(LMAX=LMAX, a=a, Omega=Omega, g=g)
+        self.dt, self.scheme = float(dt), scheme
+        self.filter_every, self.filter_alpha, self.filter_p = int(filter_every), float(filter_alpha), int(filter_p)
+        self.n_pole_rows = int(n_pole_rows)
+        self._sponge = make_sponge(self.grid.Nlat, n_pole_rows=n_pole_rows)
+        self.forcing = forcing
+        self._step = INTEGRATORS[scheme]
+
+    # ---- layout
+    @property
+    def Nlat(self): return self.grid.Nlat
+    @property
+    def Nlon(self): return self.grid.Nlon
+    @property
+    def field_size(self): return self.grid.Nlat * self.grid.Nlon
+    @property
+    def dim(self): return 3 * self.field_size
+    @property
+    def var_blocks(self):
+        fs = self.field_size
+        return {"u": slice(0, fs), "v": slice(fs, 2 * fs), "h": slice(2 * fs, 3 * fs)}
+
+    def pack(self, u, v, h):
+        return np.concatenate([np.asarray(u).ravel(), np.asarray(v).ravel(), np.asarray(h).ravel()])
+
+    def unpack(self, x):
+        fs = self.field_size; s = (self.Nlat, self.Nlon)
+        return x[:fs].reshape(s), x[fs:2 * fs].reshape(s), x[2 * fs:3 * fs].reshape(s)
+
+    def interior_mask(self):
+        """Boolean over one field: True outside the polar sponge (the estimated points)."""
+        rows = np.repeat(np.arange(self.Nlat), self.Nlon)
+        return (rows >= self.n_pole_rows) & (rows < self.Nlat - self.n_pole_rows)
+
+    # ---- dynamics
+    def rhs(self, u, v, h):
+        du, dv, dh = _rhs(u, v, h, self.grid)
+        if self.forcing is not None:
+            du, dv, dh = self.forcing(u, v, h, du, dv, dh)
+        return du, dv, dh
+
+    def propagate(self, x0, T, just_final_state=True):
+        if not just_final_state:
+            raise NotImplementedError
+        u, v, h = (f.copy() for f in self.unpack(np.asarray(x0, dtype=float)))
+        n_steps = max(1, int(round((float(T[-1]) - float(T[0])) / self.dt)))
+        for step in range(1, n_steps + 1):
+            u, v, h = self._step(self.rhs, u, v, h, self.dt)
+            u = u * self._sponge; v = v * self._sponge
+            if step % self.filter_every == 0:
+                u = sh_filter(u, self.grid, alpha=self.filter_alpha, p=self.filter_p)
+                v = sh_filter(v, self.grid, alpha=self.filter_alpha, p=self.filter_p)
+                h = sh_filter(h, self.grid, alpha=self.filter_alpha, p=self.filter_p)
+            if not np.isfinite(h).all():
+                raise RuntimeError(f"SWE integration blew up at step {step}/{n_steps} (dt={self.dt}); reduce dt or filter more often")
+        return self.pack(u, v, h)
+
+    # ---- initial conditions
+    def initial_condition(self, kind: str = "rossby", seed: int = 0, **kw):
+        if kind == "rossby":
+            perturbed = kw.pop("perturbed", False)
+            u, v, h = ics.perturbed_ic(self.grid, seed=seed, **kw) if perturbed else ics.default_ic(self.grid, **kw)
+        elif kind == "tc2":
+            u, v, h = ics.williamson_tc2(self.grid, **kw)
+        elif kind == "galewsky":
+            u, v, h = ics.galewsky_jet(self.grid, **kw)
+        else:
+            raise ValueError(kind)
+        return self.pack(u, v, h)

swesphere/presets.py

@@ -0,0 +1,64 @@
+# -*- coding: utf-8 -*-
+"""Named configurations. The numbers that characterize each one are measured by the experiments
+(``results/EXP-01/stability.csv``, ``results/EXP-02/regimes.csv``, ``results/EXP-03/footprint.csv``)
+for the default ``dt = 120 s`` with the filter every 3 steps; they change with the filter cadence.
+
+``waves``
+    Solid-body rotation (Williamson test case 2) plus a wavenumber-4 wave packet in the north and a
+    wavenumber-3 packet in the south, RK4, no forcing. The initial eddies decay in ~40 days and zonal
+    flow with weak Rossby waves remains. With no source of energy the waves keep decaying slowly (the
+    spectral filter is the only sink): the regime is quasi-linear and slowly decaying, not stationary.
+    Spin-up 40 days; snapshots every 3 days.
+
+``one_jet``
+    Galewsky jet at 45 N, RK4, relaxation of the zonal mean (tau = 5 d) toward the jet. Sustained and
+    fairly regular eddies in the north, the south quiescent. Spin-up 40 days; snapshots every 2 days.
+
+``two_jets``
+    Galewsky jets at +-45 deg, same relaxation. Barotropic instability sustains irregular eddies in both
+    hemispheres: a stationary turbulent regime in which balanced perturbations travel with the mean wind
+    (EXP-03). Spin-up 40 days; snapshots every 2 days.
+"""
+from __future__ import annotations
+
+from .model import SWEModel
+from .forcing import ZonalRelaxation
+
+DAY = 86400.0
+
+
+def waves(**overrides) -> SWEModel:
+    kw = dict(dt=120.0, scheme="rk4"); kw.update(overrides)
+    return SWEModel(**kw)
+
+
+def one_jet(tau_days: float = 5.0, h0: float = 10000.0, **overrides):
+    """Galewsky jet in the north only: sustained, fairly regular eddies; the southern hemisphere quiescent.
+    Returns ``(model, x0)``."""
+    return _jets(False, tau_days, h0, **overrides)
+
+
+def two_jets(tau_days: float = 5.0, h0: float = 10000.0, **overrides):
+    """Galewsky jets at +-45 deg: irregular sustained turbulence in both hemispheres. Returns ``(model, x0)``."""
+    return _jets(True, tau_days, h0, **overrides)
+
+
+def _jets(both, tau_days, h0, **overrides):
+    kw = dict(dt=120.0, scheme="rk4"); kw.update(overrides)
+    model = SWEModel(**kw)
+    x0 = model.initial_condition("galewsky", h0=h0, both_hemispheres=both)
+    u, v, h = model.unpack(x0)
+    model.forcing = ZonalRelaxation(u, v, h, tau_days * DAY)
+    return model, x0
+
+
+PRESETS = {"waves": lambda **kw: (waves(**kw), waves(**kw).initial_condition("rossby")), "one_jet": one_jet, "two_jets": two_jets}
+SPINUP_DAYS = {"waves": 40.0, "one_jet": 40.0, "two_jets": 40.0}
+SNAPSHOT_DAYS = {"waves": 3.0, "one_jet": 2.0, "two_jets": 2.0}
+
+
+def make(name: str, **overrides):
+    """``(model, x0)`` for a preset name."""
+    if name == "waves":
+        m = waves(**overrides); return m, m.initial_condition("rossby")
+    return PRESETS[name](**overrides)

swesphere-0.3.0.dist-info/METADATA

@@ -0,0 +1,194 @@
+Metadata-Version: 2.4
+Name: swesphere
+Version: 0.3.0
+Summary: Shallow-water equations on the rotating sphere: a verified model with documented, reproducible flow regimes
+Author-email: "Elias D. Nino-Ruiz" <enino@uninorte.edu.co>
+License: MIT
+Project-URL: Homepage, https://aml-cs.github.io/software/swesphere/
+Project-URL: Documentation, https://aml-cs.github.io/software/swesphere/
+Project-URL: Repository, https://github.com/enino84/swesphere
+Project-URL: Issues, https://github.com/enino84/swesphere/issues
+Project-URL: Changelog, https://github.com/enino84/swesphere/blob/main/CHANGELOG.md
+Keywords: shallow water,sphere,geophysical fluid dynamics,numerical methods,data assimilation,testbed
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Atmospheric Science
+Classifier: Topic :: Scientific/Engineering :: Physics
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.24
+Requires-Dist: scipy>=1.10
+Requires-Dist: pyshtools>=4.10
+Provides-Extra: test
+Requires-Dist: pytest; extra == "test"
+Provides-Extra: plots
+Requires-Dist: matplotlib; extra == "plots"
+Provides-Extra: experiments
+Requires-Dist: pandas; extra == "experiments"
+Requires-Dist: matplotlib; extra == "experiments"
+Dynamic: license-file
+
+# swesphere
+
+[![tests](https://github.com/enino84/swesphere/actions/workflows/tests.yml/badge.svg)](https://github.com/enino84/swesphere/actions/workflows/tests.yml)
+[![PyPI](https://img.shields.io/pypi/v/swesphere.svg)](https://pypi.org/project/swesphere/)
+[![Python](https://img.shields.io/pypi/pyversions/swesphere.svg)](https://pypi.org/project/swesphere/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](https://github.com/enino84/swesphere/blob/main/LICENSE)
+
+Shallow-water equations on the sphere with three documented flow regimes: a small geophysical
+flow for numerical experiments, verified against the standard test cases and reproducible from a
+container. This package is the model only; assimilation filters live in `pyteda`.
+
+![The two_jets preset on the sphere](https://raw.githubusercontent.com/enino84/swesphere/main/docs/img/hero_two_jets.png)
+
+*A state of the `two_jets` record: zonal wind, meridional wind, and the departure of `h` from its
+zonal mean drawn as relief (exaggerated). Polar caps under the sponge are grey. Rendered with
+`python figures/make_globe3d.py two_jets`.*
+
+## Install
+
+```
+pip install swesphere
+```
+
+Documentation and API: https://aml-cs.github.io/software/swesphere/
+
+## Quick start
+
+```python
+from swesphere import presets, climatology
+model, x0 = presets.two_jets()
+S = climatology.build_climatology(model, x0, spinup_days=40, n_snapshots=200, every_days=2)
+scales, per_field = climatology.anomaly_scales(model, S)
+```
+
+State vector `[u, v, h]` (row-major, north to south), `model.var_blocks`,
+`model.interior_mask()` (outside the polar sponge), `model.propagate(x, [t0, t1])`.
+
+## The model
+
+Equations: vector-invariant momentum and flux-form continuity,
+
+    du/dt = eta v - (1 / a cos(lat)) dB/dlon        eta = zeta + f,   B = g h + (u^2 + v^2) / 2
+    dv/dt = -eta u - (1 / a) dB/dlat
+    dh/dt = -div(h V)
+
+Discretization: regular latitude-longitude (Driscoll-Healy) grid, `Nlat = 2 (LMAX + 1)`,
+`Nlon = 4 (LMAX + 1)` (66 x 132 at the default `LMAX = 32`); exact zonal derivative by FFT,
+second-order centered meridional difference; RK4 in time (`dt = 120 s`). Dissipation: an
+exponential filter in spherical harmonics (`pyshtools`) applied every `filter_every = 3` steps,
+and a multiplicative sponge on the wind over the `n_pole_rows = 10` rows next to each pole
+(poleward of about 63 deg), which is how the model avoids the polar singularity of the grid.
+The rows under the sponge are not meant to be used (`model.interior_mask()`).
+
+![Verification](https://raw.githubusercontent.com/enino84/swesphere/main/docs/img/verification.png)
+
+*Left: normalized l2 error of `h` in Williamson test case 2, with and without the polar sponge.
+Right: relative vorticity of the Galewsky instability at day 6, LMAX 64 (compare with the
+published reference solution).*
+
+Two consequences that the experiments quantify:
+
+- The scheme is second order (EXP-00: Williamson test case 2 without the sponge, l2 error of h
+  2.2e-4 at LMAX 32 and 5.5e-5 at LMAX 64 after 5 days). With the sponge the steady state of
+  that test is disturbed from the polar caps and the error in the interior is about 1e-2.
+- The dissipation acts per step, not per second, so the eddy amplitude of a regime depends on
+  `dt * filter_every`. The statistics quoted for the presets hold for the defaults (EXP-01).
+
+## Presets
+
+| preset | what it is | regime |
+|---|---|---|
+| `waves` | solid-body rotation (Williamson TC2) + wave packets (wavenumber 4 north, 3 south), no forcing | quasi-linear Rossby waves on a zonal flow, **slowly decaying** (no energy source; the filter is the only sink) |
+| `one_jet` | Galewsky jet at 45 N + relaxation of the zonal mean (5 d) | stationary, fairly regular eddies in the north |
+| `two_jets` | Galewsky jets at +-45 deg + the same relaxation | stationary irregular turbulence in both hemispheres; errors advected by the flow |
+
+![The three regimes](https://raw.githubusercontent.com/enino84/swesphere/main/docs/img/regimes.png)
+
+The measured statistics of each regime (eddy std of h, wind, anomaly scales, lag correlations,
+first-half against second-half amplitude) are written by EXP-02 to `results/EXP-02/regimes.csv`.
+Measured with the defaults: eddy std of `h` 8.8 m (`waves`, decaying), 124 m (`one_jet`) and
+156 m (`two_jets`); `|U|max` about 42, 76 and 76 m/s. Note the scales in the figure above: the
+meridional wind and the depth anomaly of `waves` are an order of magnitude weaker than those of
+the forced presets.
+
+![Nine days of the two_jets regime](https://raw.githubusercontent.com/enino84/swesphere/main/docs/img/chaos_vorticity.png)
+
+*Relative vorticity of the `two_jets` regime over nine days, starting from a spun-up state. The jet
+meanders, sheds eddies and reforms: the flow is chaotic but statistically stationary, which is what
+makes it usable as a reference.*
+
+## The article's experiments (Docker)
+
+```
+docker compose build test
+docker compose run --rm test              # unit tests
+QUICK=1 docker compose run --rm all       # smoke
+docker compose run --rm all               # everything, one core
+```
+
+or detached and in parallel (`verification`, `stability` and `regimes` are independent;
+`footprint` needs `regimes`; `figures` needs all):
+
+```
+docker compose run -d verification; docker compose run -d stability; docker compose run -d regimes
+docker compose run -d footprint           # when regimes has finished
+docker compose run --rm figures           # when everything has finished
+```
+
+| experiment | what it measures | output |
+|---|---|---|
+| `exp00_verification.py` | Williamson TC2 error norms (LMAX 32, 64; with and without sponge), Galewsky instability: invariants and vorticity | `results/EXP-00/` |
+| `exp01_stability.py` | the stable range of `dt` for the shipped RK4 (pushed until it blows up) and the eddy amplitude against the filter interval | `results/EXP-01/` |
+| `exp02_regimes.py` | 200-state record per preset, per-snapshot series, statistics, stationarity, decorrelation | `results/EXP-02/` |
+| `exp03_footprint.py` | balanced against unbalanced perturbations: advection against gravity waves | `results/EXP-03/` |
+| `figures/make_figures.py` | figures and tables of the article, one file per panel | `results/figures/` |
+
+Besides those, `figures/make_globe3d.py <preset> [state.npy]` renders the illustrative
+three-dimensional views above (it spins the preset up itself if no state is given).
+
+Progress of a detached run: `docker ps`, `docker logs -f <name>`; every script logs the preset
+and the step, and rewrites its CSV after every case.
+
+## Using your own integrator
+
+The package publishes the right-hand side and ships RK4 as the tested stepper. Any other
+scheme can be plugged in and runs with the same filter, sponge and forcing (its stable range
+of `dt` is its own: rerun EXP-01):
+
+```python
+from swesphere import presets
+from swesphere.integrators import INTEGRATORS
+
+def my_step(rhs, u, v, h, dt):          # rhs(u, v, h) -> (du, dv, dh)
+    du, dv, dh = rhs(u, v, h)
+    return u + dt * du, v + dt * dv, h + dt * dh
+
+INTEGRATORS["mine"] = my_step
+model, x0 = presets.two_jets(scheme="mine")
+```
+
+`swesphere.dynamics.rhs(u, v, h, grid)` is the physics alone; `model.rhs(u, v, h)` adds the
+preset's forcing. `swesphere.diagnostics` gives mass, energy, potential enstrophy and the
+Williamson error norms.
+
+Changes between versions are listed in `CHANGELOG.md` (0.2.0 fixes the continuity equation of 0.1.0).
+
+## Citing, contributing and license
+
+If you use swesphere in your work, please cite it: the `CITATION.cff` file gives the reference, and
+GitHub shows it under "Cite this repository". Bug reports and pull requests are welcome; see
+`CONTRIBUTING.md`. Questions: enino@uninorte.edu.co.
+
+swesphere is released under the MIT license (`LICENSE`). Its dependencies are NumPy and SciPy
+(BSD-3-Clause) and pyshtools (BSD-3-Clause), all compatible with it.
+

swesphere-0.3.0.dist-info/RECORD

@@ -0,0 +1,15 @@
+swesphere/__init__.py,sha256=vDaLzIseojqBTVUfkfFVjAWZ0a5Fl2Fg9rudc1aAh4Y,636
+swesphere/climatology.py,sha256=DNfPW5IQfRSL34mQPXn-RtzCs33riNC1oStFvBVKRKg,1675
+swesphere/diagnostics.py,sha256=Y6LmfbvHdRwHBxJCuM1zA4sR6PKBfxoVHV4R6CzdJXo,1901
+swesphere/dynamics.py,sha256=6hG6xM0GD07oZ2psldd46q8U9dlYmcAkxglg0Via2nI,2897
+swesphere/forcing.py,sha256=l3ptF9w9dGoMwuALx1dAo8h9sMScirFMW2ywTcWmCuA,1300
+swesphere/grid.py,sha256=LPUS4Pn9HaegkubcYMJbRQo0vCIFqNLXaAB91hdore4,5194
+swesphere/initial_conditions.py,sha256=qfQojCTBnLgjAXrPjveRmUmL4_uqj7qQMB5e4XJN9Xw,6687
+swesphere/integrators.py,sha256=-NnjmGNxOffWtfJT9RJM_Z0dApIfBt2HxJTqpqTXnBk,965
+swesphere/model.py,sha256=9wXh0ZijWEu4nCYaNVooanfBOCWgeNhajGYBA9ihVT0,3921
+swesphere/presets.py,sha256=cOksXQxeJOp5bzX04Cunmv2Z0hyWnCtq73tLBpwIZJU,2821
+swesphere-0.3.0.dist-info/licenses/LICENSE,sha256=11w9-5xx_a4OvX5f1CA7ngXqNoKoRgHRq6NxYoeW2Hk,1075
+swesphere-0.3.0.dist-info/METADATA,sha256=WaEB3FqzJlMRUXA4abn3g9neHkGxnOIg_PzMVSgHTvM,9649
+swesphere-0.3.0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+swesphere-0.3.0.dist-info/top_level.txt,sha256=GryoOsOGt5MZI7nQJw1NYdvyClatMMYBq6cwBkyzQgI,10
+swesphere-0.3.0.dist-info/RECORD,,

swesphere-0.3.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

swesphere-0.3.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Elias D. Nino-Ruiz
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

swesphere-0.3.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+swesphere