"""Synthetic (theoretical) GCM interface for TRACMASS JAX.
Provides functions that generate idealised, analytically-defined
velocity flux fields and initial particle states for testing and
validation of the TRACMASS particle-tracking framework without
requiring real ocean reanalysis data.
"""
from typing import NewType
import jax
import jax.numpy as jnp
import numpy as np
import pandas as pd
import xarray as xr
from tracmass.gcm.flux_grid import FluxFields, GridFields
from tracmass.particle import ParticleState
FluxDataset = NewType("FluxDataset", xr.Dataset)
# ---------------------------------------------------------------------------
# Grid constants
# ---------------------------------------------------------------------------
DXV: float = 250.0 # m — zonal cell width
DYU: float = 250.0 # m — meridional cell height
DZT: float = 10.0 # m — nominal layer thickness
DXDY: float = DXV * DYU # m²
# Target max horizontal speed ~0.1 m/s; flux = speed × face_area
_FLUX_AMP: float = 0.1 * DYU * DZT # 250.0 m³/s
# ===========================================================================
# Convenience: create test data and run
# ===========================================================================
#
def _make_perturbed_field(grid, fluxes, seed_offset):
"""Create a slightly perturbed velocity field for multi-window testing.
Adds a small Gaussian random perturbation to the present-time
(slot 1) velocity fluxes stored in *fluxes*.
Parameters
----------
grid : GridFields
Static grid with dimension metadata.
fluxes : FluxFields
Source flux data from which the baseline velocity field (slot 1) is taken.
seed_offset : int
Offset applied to the JAX random seed to ensure independent
perturbations across different calls.
Returns
-------
dict
Dictionary with keys ``"uflux"`` and ``"vflux"`` containing
JAX arrays of the perturbed fluxes.
"""
key = jax.random.PRNGKey(seed_offset)
imt, jmt, km = grid.imt, grid.jmt, grid.km
return {
"uflux": fluxes.uflux[:, :, :, 1]
+ 0.002 * jax.random.normal(key, (imt, jmt, km)),
"vflux": fluxes.vflux[:, :, :, 1]
+ 0.002
* jax.random.normal(jax.random.PRNGKey(seed_offset + 1), (imt, jmt + 1, km)),
}
[docs]
def generate_keys(random_seed: int = 42):
"""Generate a set of JAX random keys from a single seed.
Splits a PRNGKey derived from *random_seed* into 10 independent
sub-keys for use in stochastic field generation.
Parameters
----------
random_seed : int, optional
Integer seed used to initialise the JAX PRNG. Default is ``42``.
Returns
-------
jnp.ndarray
Array of shape ``(10, 2)`` containing 10 independent JAX random keys.
"""
key = jax.random.PRNGKey(random_seed)
keys = jax.random.split(key, 10)
return keys
[docs]
def fluxes(
dtm="2007-06-01",
imt: int = 100,
jmt: int = 80,
km: int = 30,
random_seed: int = 42,
) -> FluxDataset:
"""Generate synthetic volume-flux fields on an idealised rectangular grid.
Constructs analytically-defined zonal and meridional volume fluxes
following a double-gyre pattern with small random perturbations, and
random vertical fluxes. When *dtm* is later than the reference date
``"2007-06-01"``, the fluxes are time-evolved by accumulating
day-by-day random increments.
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.
Parameters
----------
dtm : pd.Timestamp or str, optional
Target date for the flux snapshot. Default is ``"2007-06-01"``.
imt : int, optional
Number of grid cells in the x-direction. Default is ``100``.
jmt : int, optional
Number of grid cells in the y-direction. Default is ``80``.
km : int, optional
Number of vertical levels. Default is ``30``.
random_seed : int, optional
Seed for JAX random number generation. Default is ``42``.
Returns
-------
FluxDataset
xarray Dataset containing the variables required by
:class:`GridFields`: ``uflux``, ``vflux``, ``wflux``, ``kmt``,
``dzt``, ``dxdy``, ``timestep``, ``dtm``, ``reverse``,
``imt``, ``jmt``, ``km``, ``iperio``, ``jperio``, and
``zstot``.
"""
dtm = pd.to_datetime(dtm)
dtm0 = pd.to_datetime("2007-06-01")
timestep = pd.Timedelta(1, "h")
keys = generate_keys(random_seed=random_seed)
xi = jnp.arange(imt)[:, None, None, None] / imt
yj = jnp.arange(jmt)[None, :, None, None] / jmt
uflux = (
-_FLUX_AMP
* jnp.cos(jnp.pi * yj)
* jnp.sin(2 * jnp.pi * xi)
* jnp.ones((1, 1, km, 1))
)
uflux = uflux + 0.1 * _FLUX_AMP * jax.random.normal(keys[0], (imt, jmt, km, 1))
yj_v = jnp.arange(jmt + 1)[None, :, None, None] / jmt
xi_v = jnp.arange(imt)[:, None, None, None] / imt
vflux = (
_FLUX_AMP
* jnp.sin(jnp.pi * yj_v)
* jnp.cos(2 * jnp.pi * xi_v)
* jnp.ones((1, 1, km, 1))
)
vflux = vflux + 0.1 * _FLUX_AMP * jax.random.normal(keys[1], (imt, jmt + 1, km, 1))
n_days = (dtm - dtm0).days
if n_days > 0:
# Use fold_in to derive a fixed-shape field keyed by n_days.
# Scaling by sqrt(n_days) preserves random-walk variance growth.
# This avoids vmapping over jnp.arange(n_days) which changes shape
# each window and triggers JAX recompilation.
amp = 0.04 * _FLUX_AMP * float(np.sqrt(n_days))
uflux = uflux + amp * jax.random.normal(
jax.random.fold_in(jax.random.PRNGKey(0), n_days), (imt, jmt, km, 1)
)
vflux = vflux + amp * jax.random.normal(
jax.random.fold_in(jax.random.PRNGKey(1), n_days), (imt, jmt + 1, km, 1)
)
# Vertical flux: w ~ 1e-4 m/s → wflux = w × DXDY
wflux = 1e-4 * DXDY * jax.random.normal(keys[2], (imt, jmt, km + 1, 1))
kmt = jnp.full((imt, jmt), km, dtype=jnp.int32)
kmt = kmt.at[:, 0].set(0)
kmt = kmt.at[:, -1].set(0)
dsdict = {
"uflux": (("x", "y", "z", "t"), uflux),
"vflux": (("x", "y_v", "z", "t"), vflux),
"wflux": (("x", "y", "z_w", "t"), wflux),
"kmt": (("x", "y"), kmt),
"dzt": (("x", "y", "z", "t"), jnp.ones((imt, jmt, km, 1)) * DZT),
"dxdy": (("x", "y"), jnp.ones((imt, jmt)) * DXDY),
"timestep": timestep,
"dtm": dtm,
"reverse": False,
}
dsdict["km"], dsdict["jmt"], dsdict["imt"] = km, jmt, imt
dsdict["iperio"], dsdict["jperio"] = 1, 0
dsdict["zstot"] = (("x", "y", "t"), jnp.ones((imt, jmt, 1)))
return xr.Dataset(dsdict)
[docs]
def initial_state(
n_particles: int = 10_000,
imt: int = 100,
jmt: int = 80,
km: int = 30,
random_seed: int = 42,
) -> ParticleState:
"""Generate a random initial particle state on the theoretical grid.
Draws particle positions uniformly at random within the interior of
the grid (avoiding boundary cells) and computes the corresponding
integer cell indices.
Parameters
----------
n_particles : int, optional
Number of particles to initialise. Default is ``10_000``.
imt : int, optional
Number of grid cells in the x-direction. Default is ``100``.
jmt : int, optional
Number of grid cells in the y-direction. Default is ``80``.
km : int, optional
Number of vertical levels. Default is ``30``.
random_seed : int, optional
Seed for JAX random number generation. Default is ``42``.
Returns
-------
ParticleState
Named tuple containing particle positions (``x``, ``y``, ``z``),
integer cell indices (``ia``, ``ja``, ``ka``), time trackers
(``tt``, ``ts``, ``t0``), sub-volume (``subvol``), iteration
counter (``niter``), and active mask (``active``).
"""
keys = generate_keys(random_seed=random_seed)
x0 = jax.random.uniform(keys[3], (n_particles,), minval=1.0, maxval=imt)
y0 = jax.random.uniform(keys[4], (n_particles,), minval=2.0, maxval=jmt - 1.0)
z0 = jax.random.uniform(keys[5], (n_particles,), minval=1.0, maxval=km - 1.0)
ia = jnp.clip(jnp.floor(x0).astype(jnp.int32) + 1, 1, imt)
ja = jnp.clip(jnp.floor(y0).astype(jnp.int32) + 1, 1, jmt)
ka = jnp.clip(jnp.floor(z0).astype(jnp.int32) + 1, 1, km)
state = ParticleState(
x=x0,
y=y0,
z=z0,
ia=ia,
ja=ja,
ka=ka,
tt=jnp.zeros(n_particles),
ts=jnp.zeros(n_particles),
t0=jnp.zeros(n_particles),
subvol=jnp.ones(n_particles) * 86400.0, # tseas as subvol for ts tracking
niter=jnp.zeros(n_particles, dtype=jnp.int32),
active=jnp.ones(n_particles, dtype=jnp.bool_),
reason=jnp.zeros(n_particles, dtype=jnp.int16),
)
return state
