swesphere 0.3.0__tar.gz

swesphere-0.3.0/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/PKG-INFO

@@ -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/README.md

@@ -0,0 +1,155 @@
+# 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/pyproject.toml

@@ -0,0 +1,46 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "swesphere"
+version = "0.3.0"
+description = "Shallow-water equations on the rotating sphere: a verified model with documented, reproducible flow regimes"
+readme = {file = "README.md", content-type = "text/markdown"}
+license = {text = "MIT"}
+requires-python = ">=3.9"
+authors = [{name = "Elias D. Nino-Ruiz", email = "enino@uninorte.edu.co"}]
+keywords = ["shallow water", "sphere", "geophysical fluid dynamics", "numerical methods", "data assimilation", "testbed"]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering :: Atmospheric Science",
+    "Topic :: Scientific/Engineering :: Physics",
+    "Topic :: Scientific/Engineering :: Mathematics",
+]
+dependencies = ["numpy>=1.24", "scipy>=1.10", "pyshtools>=4.10"]
+
+[project.optional-dependencies]
+test = ["pytest"]
+plots = ["matplotlib"]
+experiments = ["pandas", "matplotlib"]
+
+[project.urls]
+Homepage = "https://aml-cs.github.io/software/swesphere/"
+Documentation = "https://aml-cs.github.io/software/swesphere/"
+Repository = "https://github.com/enino84/swesphere"
+Issues = "https://github.com/enino84/swesphere/issues"
+Changelog = "https://github.com/enino84/swesphere/blob/main/CHANGELOG.md"
+
+[tool.setuptools]
+package-dir = {"" = "src"}
+
+[tool.setuptools.packages.find]
+where = ["src"]

swesphere-0.3.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

swesphere-0.3.0/src/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-0.3.0/src/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-0.3.0/src/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-0.3.0/src/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-0.3.0/src/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