"""Fabboni analytical test-case GCM for TRACMASS JAX.
Provides spatially-uniform, analytically-defined velocity fluxes and the
canonical single-particle initial state matching the Fabboni test case from
the TRACMASS Fortran distribution (``projects/Theoretical``).
The velocity field decays exponentially from initial conditions
``u = U0 = 0.5 m/s``, ``v = 0``, with inertial oscillation governed by
the Coriolis parameter at 45° N and two damping time-scales ``GAMMA`` and
``GAMMAG``. Because the field is spatially uniform, all particles experience
identical forcing — making it ideal for comparison with the Fortran reference
trajectory.
This module mirrors the interface of :mod:`tracmass.gcm.theoretical`:
- :func:`fluxes` — returns an ``xr.Dataset`` compatible with
:func:`~tracmass.gcm.new_flux` and :func:`~tracmass.gcm.update_flux`.
- :func:`initial_state` — returns a single-particle :class:`ParticleState`
that matches the Fortran seeding (``isec=1, idir=1``).
"""
from __future__ import annotations
import math
import jax.numpy as jnp
import numpy as np
import pandas as pd
import xarray as xr
from tracmass.particle import ParticleState
# ---------------------------------------------------------------------------
# Grid constants
# ---------------------------------------------------------------------------
IMT: int = 300 # number of x-cells
JMT: int = 300 # number of y-cells
KM: int = 1 # number of vertical levels
DXV: float = 250.0 # m — zonal cell width
DYU: float = 250.0 # m — meridional cell height
DZT: float = 10.0 # m — layer thickness
DXDY: float = DXV * DYU # m²
DTM0: pd.Timestamp = pd.Timestamp("2000-01-01")
TIMESTEP: pd.Timedelta = pd.Timedelta(1, "h")
# ---------------------------------------------------------------------------
# Fabboni velocity parameters
# ---------------------------------------------------------------------------
_PI = math.pi
_RADIAN = _PI / 180.0
UG: float = 0.04
U0: float = 0.5
FCOR: float = 2.0 * 2.0 * _PI / (24.0 * 3600.0) * math.cos(45.0 * _RADIAN)
GAMMA: float = 1.0 / (2.89 * 24.0 * 3600.0)
GAMMAG: float = 1.0 / (28.9 * 24.0 * 3600.0)
# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------
def _velocity(ints: int) -> tuple[float, float]:
"""Compute Fabboni zonal and meridional velocities at GCM step *ints*.
Parameters
----------
ints : int
GCM step index (0 = ``DTM0``).
Returns
-------
tuple[float, float]
``(u_vel, v_vel)`` in m/s.
"""
omtime = float(ints) * TIMESTEP.total_seconds()
u_vel = (
UG * math.exp(-GAMMAG * omtime)
+ (U0 - UG) * math.exp(-GAMMA * omtime) * math.cos(FCOR * omtime)
)
v_vel = -(U0 - UG) * math.exp(-GAMMA * omtime) * math.sin(FCOR * omtime)
return u_vel, v_vel
# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------
[docs]
def fluxes(
dtm: pd.Timestamp | str = "2000-01-01",
) -> xr.Dataset:
"""Generate Fabboni spatially-uniform flux fields at time *dtm*.
The step index is derived from the offset of *dtm* from
:data:`DTM0 = 2000-01-01 <DTM0>`, counted in units of
:data:`TIMESTEP` (1 hour).
Backward-in-time tracking is handled at the :func:`~tracmass.gcm.new_grid`
/ :func:`~tracmass.gcm.new_flux` level by passing ``dtm_end < dtm_start``; this function always returns
forward (positive) fluxes at the requested time.
Parameters
----------
dtm : pd.Timestamp or str, optional
Target time for the flux snapshot. Default ``"2000-01-01"``.
Returns
-------
xr.Dataset
Dataset containing ``uflux``, ``vflux``, ``wflux``, ``dzt``,
``zstot``, ``kmt``, ``dxdy``, ``timestep``, ``dtm``, ``reverse``,
``imt``, ``jmt``, ``km``, ``iperio``, ``jperio``.
"""
dtm = pd.to_datetime(dtm)
ints = int(round((dtm - DTM0) / TIMESTEP))
u_vel, v_vel = _velocity(ints)
uf = DYU * DZT * u_vel # volume flux (m³/s)
vf = DXV * DZT * v_vel
uflux = np.full((IMT, JMT, KM, 1), uf, dtype=np.float64)
vflux = np.full((IMT, JMT + 1, KM, 1), vf, dtype=np.float64)
vflux[:, 0, :, :] = 0.0 # no-flux at southern boundary
vflux[:, JMT, :, :] = 0.0 # no-flux at northern boundary
wflux = np.zeros((IMT, JMT, KM + 1, 1), dtype=np.float64)
dzt = np.full((IMT, JMT, KM, 1), DZT, dtype=np.float64)
dxdy = np.full((IMT, JMT), DXDY, dtype=np.float64)
kmt = np.ones((IMT, JMT), dtype=np.int32)
zstot = np.ones((IMT, JMT, 1), dtype=np.float64)
return xr.Dataset({
"uflux": (("x", "y", "z", "t"), uflux),
"vflux": (("x", "y_v", "z", "t"), vflux),
"wflux": (("x", "y", "z_w", "t"), wflux),
"dzt": (("x", "y", "z", "t"), dzt),
"zstot": (("x", "y", "t"), zstot),
"kmt": (("x", "y"), kmt),
"dxdy": (("x", "y"), dxdy),
"timestep": TIMESTEP,
"dtm": dtm,
"reverse": False,
"imt": IMT,
"jmt": JMT,
"km": KM,
"iperio": 1,
"jperio": 1,
})
[docs]
def initial_state() -> ParticleState:
"""Create the Fabboni reference initial particle state.
Seeds a single particle at the eastern wall of cell (150, 100, 1),
matching the Fortran seeding ``isec=1, idir=1``:
- ``x = 150.0`` — eastern wall of i-cell 150 → ``ia = 151``
- ``y = 99.5`` — centre of j-cell 100 → ``ja = 100``
- ``z = 0.5`` — centre of k-cell 1 → ``ka = 1``
Returns
-------
ParticleState
Single-particle state with all time counters at zero and
``subvol`` set to one :data:`TIMESTEP` in seconds.
"""
n = 1
tseas = TIMESTEP.total_seconds()
return ParticleState(
x=jnp.array([150.0], dtype=jnp.float64),
y=jnp.array([99.5], dtype=jnp.float64),
z=jnp.array([0.5], dtype=jnp.float64),
ia=jnp.array([151], dtype=jnp.int32),
ja=jnp.array([100], dtype=jnp.int32),
ka=jnp.array([1], dtype=jnp.int32),
tt=jnp.zeros(n, dtype=jnp.float64),
ts=jnp.zeros(n, dtype=jnp.float64),
t0=jnp.zeros(n, dtype=jnp.float64),
subvol=jnp.array([tseas], dtype=jnp.float64),
niter=jnp.zeros(n, dtype=jnp.int32),
active=jnp.ones(n, dtype=jnp.bool_),
reason=jnp.zeros(n, dtype=jnp.int16),
)
