"""Grid field containers and utilities for TRACMASS JAX.
This module defines the core data structures and helper functions used
to manage GCM grid geometry and velocity flux fields across two-slot
time windows required by the TRACMASS Lagrangian particle tracker.
The grid is split into three objects:
- :class:`GridFields` — static geometry (``kmt``, ``dxdy``) and metadata
(dimensions, periodicity). These fields never change during a run and
are closed over by the compiled advection kernels as XLA constants.
- :class:`LazyFluxFields` — two time-slots of flux data held as raw
xarray DataArrays (possibly Dask-backed), with no I/O triggered yet.
Also carries ``dtm``, the timestamp of the past (slot 0) time level.
Call :meth:`~LazyFluxFields.compute` to load data into device memory
just before calling the advection kernel. :func:`update_flux` advances
the window by a pure Python dict reference swap — no Dask compute, no
JAX, no GC required until the next :meth:`~LazyFluxFields.compute`.
- :class:`FluxFields` — the computed form: JAX arrays backed by device
memory, registered as a pytree so the compiled kernel can receive it as
a dynamic argument.
"""
from __future__ import annotations
import dataclasses
import gc
import importlib
import logging
import math
import sys
import time
import warnings
from dataclasses import dataclass, field
from typing import Callable, NewType
import dask
import jax
import jax.numpy as jnp
import numpy as np
import pandas as pd
import xarray as xr
from tracmass.verbose import timed
__all__ = ["glorys", "theoretical"]
_FLUX_VARS = frozenset({"uflux", "vflux", "wflux", "dzt", "zstot"})
# Largest finite float32 value — used to clip float64 arrays before casting.
# GCM fill values (e.g. netCDF default 9.96921e+36) are within float32 range,
# but derived quantities (flux × scale factor, interpolation artefacts) can
# occasionally exceed float32 max (~3.4e38), causing jnp.stack to emit a
# RuntimeWarning when it implicitly casts float64 → float32.
_F32_MAX = float(np.finfo(np.float32).max)
# Threshold above which flux values are treated as fill / missing and zeroed.
# The GLORYS12 / NEMO fill value is 9.96921e+36. Any |flux| above this
# threshold is physically unrealistic and is zeroed so that NaN-free land-mask
# encoded as fill values doesn't propagate into the advection kernel.
_FLUX_FILL_THRESHOLD = 1e20
# LazyFluxFields.compute warns when one window reads more than this.
_LARGE_LOAD_BYTES = 8e9
log = logging.getLogger(__name__)
_STATIC_VARS = frozenset(
{"kmt", "dxdy", "imt", "jmt", "km", "timestep", "reverse", "iperio", "jperio"}
)
# Flux variables whose sign must be flipped for backward-in-time tracking.
# dzt and zstot are geometric / sea-surface scalars and are never negated.
_FLUX_NEGATION_VARS = frozenset({"uflux", "vflux", "wflux"})
FluxDataset = NewType("FluxDataset", xr.Dataset)
# ---------------------------------------------------------------------------
# Static grid geometry
# ---------------------------------------------------------------------------
[docs]
@dataclass(frozen=True, kw_only=True)
class GridFields:
"""Static grid geometry and metadata for a TRACMASS run.
These fields are fixed for the lifetime of a run. The advection
kernels close over a ``GridFields`` instance so that ``kmt`` and
``dxdy`` become XLA constants in the compiled graph.
Attributes
----------
kmt : jnp.ndarray
Bottom-level index (integer), shape ``(imt, jmt)``.
A value of 0 indicates a land cell.
dxdy : jnp.ndarray
Horizontal cell area, shape ``(imt, jmt)``.
imt : int
Number of grid cells in the zonal (x) direction.
jmt : int
Number of grid cells in the meridional (y) direction.
km : int
Number of vertical levels.
timestep : pd.Timedelta
Duration between consecutive time levels.
reverse : bool
If ``True``, fluxes have been negated for backward-in-time tracking.
iperio : int
East-west periodicity flag (1 = cyclic, 0 = closed).
jperio : int
North-fold periodicity flag (1 = north-fold, 0 = closed).
i_offset, j_offset : int
Position of this grid's cell ``(0, 0)`` in the adapter's full
(un-subset) grid — the ``subset`` ``i1`` / ``j1``, else 0. Particle
positions are relative to this grid; the trajectory recorder adds
the offsets so saved ``x`` / ``y`` / ``ia`` / ``ja`` are absolute.
"""
kmt: jnp.ndarray # (imt, jmt) int — bottom index
dxdy: jnp.ndarray # (imt, jmt) — horizontal cell area
imt: int
jmt: int
km: int
timestep: pd.Timedelta
reverse: bool
iperio: int # 1 = east-west cyclic
jperio: int # 1 = north-fold
i_offset: int = 0 # full-grid index of cell i=0 (subset i1)
j_offset: int = 0 # full-grid index of cell j=0 (subset j1)
def _replace(self, **kwargs):
return dataclasses.replace(self, **kwargs)
def __repr__(self):
lines = ["GridFields("]
for f in ("kmt", "dxdy"):
v = getattr(self, f)
lines.append(f" {f}: shape={v.shape}, dtype={v.dtype},")
for f in ("imt", "jmt", "km", "timestep", "reverse", "iperio", "jperio",
"i_offset", "j_offset"):
lines.append(f" {f}: {getattr(self, f)},")
lines.append(")")
return "\n".join(lines)
jax.tree_util.register_dataclass(
GridFields,
data_fields=["kmt", "dxdy"],
meta_fields=[
"imt", "jmt", "km", "timestep", "reverse", "iperio", "jperio",
"i_offset", "j_offset",
],
)
[docs]
def edge_mask(grid: GridFields) -> np.ndarray:
"""Cells on the edge of the domain, where particles leave it.
The outer ring of cells along every axis that does not wrap around: the
first and last column when ``grid.iperio == 0`` (e.g. a ``subset`` in
longitude) and the first and last row when ``grid.jperio == 0``. A
particle that enters one of these cells stops with
``REASON_BOUNDS_H``. The edge cells cannot be integrated correctly:
the flux through their outer face is not part of the grid (the kernels
treat the west / south one as a closed wall, and a particle crossing the
east / north one would read fluxes outside the arrays).
Parameters
----------
grid : GridFields
Grid geometry; only ``imt``, ``jmt``, ``iperio`` and ``jperio``
are used.
Returns
-------
np.ndarray
Boolean array, shape ``(imt, jmt)``.
Examples
--------
>>> edge = edge_mask(grid) # subset grid: iperio == 0
>>> edge[0].all() and edge[-1].all()
True
"""
edge = np.zeros((grid.imt, grid.jmt), dtype=bool)
if grid.iperio == 0:
edge[[0, -1], :] = True
if grid.jperio == 0:
edge[:, [0, -1]] = True
return edge
# ---------------------------------------------------------------------------
# Computed flux data (JAX arrays, kernel-ready)
# ---------------------------------------------------------------------------
[docs]
@dataclass(frozen=True, kw_only=True)
class FluxFields:
"""Two-slot time window of velocity flux data for the advection kernels.
Each array field has a trailing dimension of size 2, where slot 0 holds
the 'past' time level and slot 1 holds the 'present' time level.
``FluxFields`` is a JAX pytree and is passed as a *dynamic* argument to
the compiled step functions. Obtain one from a
:class:`LazyFluxFields` by calling :meth:`~LazyFluxFields.compute`.
Attributes
----------
dzt : jnp.ndarray
Layer thickness at tracer points, shape ``(imt, jmt, km, 2)``.
zstot : jnp.ndarray
Total sea-surface height anomaly, shape ``(imt, jmt, 2)``.
uflux : jnp.ndarray
Zonal volume flux, shape ``(imt, jmt, km, 2)``.
vflux : jnp.ndarray
Meridional volume flux, shape ``(imt, jmt+1, km, 2)``.
wflux : jnp.ndarray or None
Vertical volume flux in the kernel convention: ``wflux[..., k, :]``
is the flux through the bottom face of level ``k`` (level 0 at the
surface), positive downward. Adapters supply the model convention
(top face, positive upward); see :func:`_w_to_kernel`.
``None`` disables vertical particle motion.
"""
dzt: jnp.ndarray # (imt, jmt, km, 2)
zstot: jnp.ndarray # (imt, jmt, 2)
uflux: jnp.ndarray # (imt, jmt, km, 2)
vflux: jnp.ndarray # (imt, jmt+1, km, 2)
wflux: jnp.ndarray | None = None # (imt, jmt, km+1, 2)
def _replace(self, **kwargs):
return dataclasses.replace(self, **kwargs)
def __repr__(self):
lines = ["FluxFields("]
for f in ("dzt", "zstot", "uflux", "vflux"):
v = getattr(self, f)
lines.append(f" {f}: shape={v.shape}, dtype={v.dtype},")
if self.wflux is None:
lines.append(" wflux: None,")
else:
lines.append(
f" wflux: shape={self.wflux.shape}, dtype={self.wflux.dtype},"
)
lines.append(")")
return "\n".join(lines)
jax.tree_util.register_dataclass(
FluxFields,
data_fields=["dzt", "zstot", "uflux", "vflux", "wflux"],
meta_fields=[],
)
# ---------------------------------------------------------------------------
# Local flux fields — patch-based, includes geometry as dynamic arrays
# ---------------------------------------------------------------------------
[docs]
@dataclass(frozen=True, kw_only=True)
class LocalFluxFields:
"""Flux data for a spatial sub-patch, with geometry included dynamically.
Unlike :class:`FluxFields`, ``kmt`` and ``dxdy`` are carried as ordinary
JAX arrays rather than being closed over as XLA constants. This allows
the advection kernel to be compiled **once** for a fixed patch shape and
then reused across windows as the patch content changes — the kernel
receives new values via the normal dynamic-argument path.
Attributes
----------
uflux : jnp.ndarray
Zonal volume flux, shape ``(imt_p, jmt_p, km, 2)``.
vflux : jnp.ndarray
Meridional volume flux, shape ``(imt_p, jmt_p+1, km, 2)``.
wflux : jnp.ndarray or None
Vertical volume flux, shape ``(imt_p, jmt_p, km+1, 2)``.
dzt : jnp.ndarray
Layer thickness, shape ``(imt_p, jmt_p, km, 2)``.
zstot : jnp.ndarray
Sea-surface height anomaly, shape ``(imt_p, jmt_p, 2)``.
kmt : jnp.ndarray
Bottom-level index, shape ``(imt_p, jmt_p)``. Dynamic — values
differ per patch but shape is fixed for the lifetime of the run.
dxdy : jnp.ndarray
Horizontal cell area, shape ``(imt_p, jmt_p)``. Dynamic.
edge : jnp.ndarray or None
:func:`edge_mask` of the full grid sliced to the patch, shape
``(imt_p, jmt_p)``: the patch cells on the domain edge, where
particles stop as leaving the domain. ``None``: no domain edge
in the patch. Dynamic.
"""
uflux: jnp.ndarray # (imt_p, jmt_p, km, 2)
vflux: jnp.ndarray # (imt_p, jmt_p+1, km, 2)
wflux: jnp.ndarray | None # (imt_p, jmt_p, km+1, 2)
dzt: jnp.ndarray # (imt_p, jmt_p, km, 2)
zstot: jnp.ndarray # (imt_p, jmt_p, 2)
kmt: jnp.ndarray # (imt_p, jmt_p)
dxdy: jnp.ndarray # (imt_p, jmt_p)
edge: jnp.ndarray | None = None # (imt_p, jmt_p) bool
def _replace(self, **kwargs):
return dataclasses.replace(self, **kwargs)
def __repr__(self):
lines = ["LocalFluxFields("]
for f in ("uflux", "vflux", "dzt", "zstot", "kmt", "dxdy"):
v = getattr(self, f)
lines.append(f" {f}: shape={v.shape}, dtype={v.dtype},")
if self.wflux is None:
lines.append(" wflux: None,")
else:
lines.append(
f" wflux: shape={self.wflux.shape}, dtype={self.wflux.dtype},"
)
lines.append(")")
return "\n".join(lines)
jax.tree_util.register_dataclass(
LocalFluxFields,
data_fields=["uflux", "vflux", "wflux", "dzt", "zstot", "kmt", "dxdy", "edge"],
meta_fields=[],
)
# ---------------------------------------------------------------------------
# Per-particle stencil — point-extracted flux values for streaming integration
# ---------------------------------------------------------------------------
[docs]
@dataclass(frozen=True, kw_only=True)
class FluxStencil:
"""Per-particle flux stencil extracted at each particle's current cell.
Each array carries values for all *N* active particles, with a trailing
slot dimension of size 2 (slot 0 = past, slot 1 = present).
Obtained via :func:`extract_stencil` using vectorised xarray point
indexing, which tells Dask to load only the Zarr chunks that contain the
requested grid cells. This enables memory-efficient integration of
widely scattered particles without loading a rectangular spatial patch.
Attributes
----------
u_east : jnp.ndarray
Zonal flux at the east wall of each particle's cell, shape ``(N, 2)``.
u_west : jnp.ndarray
Zonal flux at the west wall, shape ``(N, 2)``.
v_north : jnp.ndarray
Meridional flux at the north wall, shape ``(N, 2)``.
v_south : jnp.ndarray
Meridional flux at the south wall, shape ``(N, 2)``.
w_top : jnp.ndarray or None
Vertical flux at the top wall, shape ``(N, 2)``. ``None`` when
vertical motion is disabled (``no_wflux=True``).
w_bot : jnp.ndarray or None
Vertical flux at the bottom wall, shape ``(N, 2)`` or ``None``.
dzt : jnp.ndarray
Layer thickness at the particle's cell, shape ``(N, 2)``.
zstot : jnp.ndarray
Total column depth at the particle's position, shape ``(N, 2)``.
kmt : jnp.ndarray
Bottom level index at the particle's cell, shape ``(N,)``.
dxdy : jnp.ndarray
Horizontal cell area at the particle's position, shape ``(N,)``.
"""
u_east: jnp.ndarray # (N, 2)
u_west: jnp.ndarray # (N, 2)
v_north: jnp.ndarray # (N, 2)
v_south: jnp.ndarray # (N, 2)
w_top: jnp.ndarray | None # (N, 2) or None
w_bot: jnp.ndarray | None # (N, 2) or None
dzt: jnp.ndarray # (N, 2)
zstot: jnp.ndarray # (N, 2)
kmt: jnp.ndarray # (N,)
dxdy: jnp.ndarray # (N,)
def _replace(self, **kwargs):
return dataclasses.replace(self, **kwargs)
jax.tree_util.register_dataclass(
FluxStencil,
data_fields=[
"u_east",
"u_west",
"v_north",
"v_south",
"w_top",
"w_bot",
"dzt",
"zstot",
"kmt",
"dxdy",
],
meta_fields=[],
)
# ---------------------------------------------------------------------------
# Lazy flux data (xarray DataArrays, not yet loaded)
# ---------------------------------------------------------------------------
[docs]
@dataclass
class LazyFluxFields:
"""Two time-slots of flux data held as raw xarray DataArrays.
The DataArrays may be backed by Dask (for GLORYS from Zarr) or by
in-memory arrays (for the theoretical adapter). In either case,
**no** ``numpy`` or device memory is allocated until
:meth:`compute` is called.
This object is *not* a JAX pytree. Pass the result of
:meth:`compute` to the advection kernel instead.
The two time slots are stored as separate DataArrays **without** a
trailing time dimension. The shape of each DataArray matches the
spatial shape of the variable (e.g. ``(imt, jmt, km)`` for
``uflux``). :func:`update_flux` advances the window by reassigning
Python dict references — no Dask compute, no JAX, no GC.
Attributes
----------
_past : dict
Mapping ``{name: xr.DataArray}`` for the 'past' time slot.
_present : dict
Mapping ``{name: xr.DataArray}`` for the 'present' time slot.
dtm : pd.Timestamp or None
Timestamp of the 'past' (slot 0) time level. Set by
:func:`new_flux` and :func:`update_flux`; ``None`` when
constructed via lower-level helpers like :func:`swap_fields`.
fluxes_fn : callable or None
The GCM adapter callable used to create this object, e.g.
``gcm.glorys.fluxes``. Stored by :func:`new_flux` and
propagated by :func:`update_flux` so that
:func:`~tracmass.integrate.time_dependent` can infer it
automatically without the caller re-supplying it.
timestep : pd.Timedelta or None
GCM output interval (always positive), copied from
``grid.timestep`` by :func:`new_flux`. :func:`update_flux` steps
``dtm`` by this amount.
reverse : bool
``True`` for a backward run (copied from ``grid.reverse``):
:func:`update_flux` then walks ``dtm`` backward and negates the
velocity fluxes it loads.
dtm_end : pd.Timestamp or None
The *dtm_end* passed to :func:`new_flux`, or ``None``. Used by
:attr:`n_windows` to size a run.
Notes
-----
:meth:`compute` keeps the present slot's arrays (already read and
cleaned) in ``_computed``, keyed by the slot's DataArray. After
:func:`update_flux` that DataArray is the new past slot, so the next
:meth:`compute` reuses the arrays instead of reading the same snapshot
again: every snapshot is read once, not twice. ``_computed`` is shared
by the copies :func:`update_flux` makes and holds one slot per variable.
"""
_past: dict # {name: xr.DataArray} — past slot, no trailing t-dim
_present: dict # {name: xr.DataArray} — present slot, no trailing t-dim
dtm: pd.Timestamp | None = None
fluxes_fn: Callable | None = None
zarr_chunk_xy: tuple[int, int] = (1, 1) # (chunk_x, chunk_y) of source Zarr
timestep: pd.Timedelta | None = None
reverse: bool = False
dtm_end: pd.Timestamp | None = None
# {name: (DataArray, np.ndarray)}: the last computed present slot.
_computed: dict = field(default_factory=dict, repr=False, compare=False)
@property
def n_windows(self) -> int | None:
"""Number of GCM windows left between ``dtm`` and ``dtm_end``.
``ceil(|dtm_end - dtm| / timestep)``, so a *dtm_end* that falls
inside a window still has that window integrated. Decreases by one
with every :func:`update_flux` call. ``None`` when ``dtm_end``
(or ``dtm`` / ``timestep``) is not set.
Returns
-------
int or None
"""
if self.dtm_end is None or self.dtm is None or self.timestep is None:
return None
span = self.dtm - self.dtm_end if self.reverse else self.dtm_end - self.dtm
return max(0, math.ceil(span / self.timestep))
[docs]
def compute_patch(
self,
i0: int,
i1: int,
j0: int,
j1: int,
grid: "GridFields",
include_wflux: bool = True,
) -> "LocalFluxFields":
"""Load only the sub-patch ``[i0:i1, j0:j1]`` into device memory.
Slices the lazy xarray DataArrays before calling ``.values``, so
Dask only computes (and reads from Zarr) the chunks that intersect
the requested region. ``kmt`` and ``dxdy`` are sliced from *grid*
and included in the returned :class:`LocalFluxFields` as dynamic JAX
arrays, so the advection kernel can be compiled once for a fixed
patch shape and reused without recompilation as patch content changes.
Parameters
----------
i0, i1 : int
Zonal slice bounds (0-based, exclusive end).
j0, j1 : int
Meridional slice bounds (0-based, exclusive end).
grid : GridFields
Full-grid static geometry; ``kmt`` and ``dxdy`` are sliced from
here to match the patch.
include_wflux : bool, optional
False skips reading ``wflux`` (returned as ``None``), for runs
without vertical motion. Default True.
Returns
-------
LocalFluxFields
"""
_fill_names = frozenset({"uflux", "vflux", "wflux", "dzt", "zstot"})
# Collect all sliced Dask arrays first, then compute in one scheduler
# call. With small chunks (e.g. 200×200) there are many chunk files
# in the Zarr store; issuing 10 separate dask.compute() calls would
# build and dispatch 10 task graphs. A single dask.compute(*tasks)
# builds one graph and lets the scheduler parallelise all reads.
tasks: list = [] # Dask arrays to compute together
meta: list = [] # (name, tasks_index_t0, tasks_index_t1) or None
for name in ("uflux", "vflux", "dzt", "zstot", "wflux"):
t0_da = self._past.get(name)
t1_da = self._present.get(name)
if t0_da is None or t1_da is None or (
name == "wflux" and not include_wflux
):
meta.append((name, None, None))
continue
x_dim = t0_da.dims[0]
y_dim = t0_da.dims[1]
j_end = j1 + 1 if name == "vflux" else j1
sel = {x_dim: slice(i0, i1), y_dim: slice(j0, j_end)}
idx0 = len(tasks)
tasks.append(t0_da.isel(sel).data)
tasks.append(t1_da.isel(sel).data)
meta.append((name, idx0, idx0 + 1))
computed = dask.compute(*tasks) if tasks else ()
kw: dict = {}
for name, idx0, idx1 in meta:
if idx0 is None:
kw[name] = None
continue
t0 = np.asarray(computed[idx0])
t1 = np.asarray(computed[idx1])
if t0.dtype != np.float32:
t0 = np.clip(t0, -_F32_MAX, _F32_MAX).astype(np.float32)
t1 = np.clip(t1, -_F32_MAX, _F32_MAX).astype(np.float32)
elif not (t0.flags.writeable and t1.flags.writeable):
# Same guard as compute(): in-memory float32 slices can be
# read-only views, which the in-place fill cleanup below needs.
t0 = t0.copy()
t1 = t1.copy()
if name in _fill_names:
np.nan_to_num(t0, copy=False, nan=0.0, posinf=0.0, neginf=0.0)
np.nan_to_num(t1, copy=False, nan=0.0, posinf=0.0, neginf=0.0)
t0[np.abs(t0) > _FLUX_FILL_THRESHOLD] = 0.0
t1[np.abs(t1) > _FLUX_FILL_THRESHOLD] = 0.0
kw[name] = jnp.stack([t0, t1], axis=-1)
del t0, t1
gc.collect()
kw["kmt"] = jnp.asarray(np.asarray(grid.kmt[i0:i1, j0:j1]))
kw["dxdy"] = jnp.asarray(np.asarray(grid.dxdy[i0:i1, j0:j1]))
kw["edge"] = jnp.asarray(edge_mask(grid)[i0:i1, j0:j1])
return LocalFluxFields(**kw)
[docs]
def preload(self, include_wflux: bool = True) -> "LazyFluxFields":
"""Return a copy with all Dask arrays preloaded to numpy.
Issues a **single** :func:`dask.compute` call for every variable in
``_past`` and ``_present``, replacing Dask-backed DataArrays with
in-memory numpy-backed DataArrays. The returned object is otherwise
identical and can be used anywhere a :class:`LazyFluxFields` is
expected.
Calling this once per GCM window in
:func:`~tracmass.integrate.time_dependent_streaming` changes the I/O
pattern from *O(n_sub)* small Dask computes (one per sub-cycle that
has cell crossings) to *O(1)* per window. Subsequent
:func:`extract_stencil` calls on the returned object use pure numpy
fancy indexing — no further Zarr I/O regardless of how many
sub-cycles or cell crossings occur.
Parameters
----------
include_wflux : bool, optional
False leaves ``wflux`` out of the copy instead of reading it,
for runs without vertical motion. Default True.
Returns
-------
LazyFluxFields
A copy with numpy-backed DataArrays and the same metadata
(``dtm``, ``fluxes_fn``, ``timestep``, ``reverse``, ...).
"""
keep = lambda k: include_wflux or k != "wflux" # noqa: E731
past_keys = [k for k in self._past if keep(k)]
pres_keys = [k for k in self._present if keep(k)]
all_arrs = [self._past[k].data for k in past_keys] + [
self._present[k].data for k in pres_keys
]
computed = dask.compute(*all_arrs)
n_past = len(past_keys)
new_past = {
k: xr.DataArray(computed[i], dims=self._past[k].dims)
for i, k in enumerate(past_keys)
}
new_pres = {
k: xr.DataArray(computed[n_past + i], dims=self._present[k].dims)
for i, k in enumerate(pres_keys)
}
return dataclasses.replace(self, _past=new_past, _present=new_pres)
[docs]
def compute(self, include_wflux: bool = True) -> FluxFields:
"""Trigger Dask computation and transfer flux data to JAX device memory.
Variables are loaded one at a time (past then present) and stacked
along a new trailing axis, so the peak extra allocation per variable
is one numpy array of size ``t0``, one of size ``t1``, and the
resulting JAX array of size ``t0 + t1``.
NaN values in flux arrays (``uflux``, ``vflux``, ``wflux``) are
replaced with 0.0 before stacking. Land cells in GLORYS/NEMO are
masked as NaN in the Zarr files; the physical no-flow boundary
condition is 0.0, so NaN → 0 is the correct substitution. Without
this replacement, NaN fluxes propagate through ``_calc_pos_1d`` to
produce a NaN particle position, which after ``astype(int32)``
conversion gives a garbage cell index that may point to a kmt=0
(land) cell, causing spurious ``REASON_LAND`` deactivations.
Parameters
----------
include_wflux : bool, optional
False skips reading ``wflux`` (returned as ``None``), for runs
without vertical motion. Default True.
Returns
-------
FluxFields
Kernel-ready flux data with trailing slot dimension of size 2.
"""
names = ("uflux", "vflux", "dzt", "zstot") + (
("wflux",) if include_wflux else ()
)
nbytes = sum(
self._past[n].nbytes + self._present[n].nbytes
for n in names
if n in self._past and n in self._present
)
ref = self._past.get("dzt", self._past.get("uflux"))
shape = () if ref is None else ref.shape
if nbytes > _LARGE_LOAD_BYTES:
warnings.warn(
f"reading {nbytes / 1e9:.1f} GB of flux data per window "
f"(grid {shape}); this can take minutes before the run "
f"starts. Restrict the region / levels with [env.gcm] "
f"subset (e.g. k1/k2 for a 2-D run).",
stacklevel=2,
)
log.debug(
"compute: reading flux window dtm=%s, %.2f GB, grid %s",
self.dtm, nbytes / 1e9, shape,
)
_t_all = time.perf_counter()
kw: dict = {"wflux": None}
for name in names:
t0_da = self._past.get(name)
t1_da = self._present.get(name)
if t0_da is None or t1_da is None:
kw[name] = None
continue
_t = time.perf_counter()
cached = self._computed.get(name)
reused = cached is not None and cached[0] is t0_da
# The past slot was the present slot of the previous window.
t0 = cached[1] if reused else self._read_slot(name, t0_da)
t1 = self._read_slot(name, t1_da)
log.debug(
"compute: %-5s %s x %s: %.1f s", name, 1 if reused else 2,
t1.shape, time.perf_counter() - _t,
)
self._computed[name] = (t1_da, t1)
kw[name] = np.stack([t0, t1], axis=-1)
del t0, t1
gc.collect()
log.debug("compute: done in %.1f s", time.perf_counter() - _t_all)
return FluxFields(**kw)
@staticmethod
def _read_slot(name: str, da) -> np.ndarray:
"""Read one slot of *name* as clean float32 (fill values → 0).
Flux arrays and dzt/zstot all need fill-value sanitisation. GLORYS
(and other NEMO-based) grids store land cells with the NEMO fill
value 9.96921e36. If these survive into the advection kernel,
dxyz_raw overflows float32 to Inf, which makes ds=0 and
dt_particles = 0 * Inf = NaN, corrupting particle positions.
"""
arr = da.values # Dask compute → numpy (or no-op if already numpy)
if arr.dtype != np.float32:
arr = np.clip(arr, -_F32_MAX, _F32_MAX).astype(np.float32)
elif not arr.flags.writeable:
arr = arr.copy()
if name in ("uflux", "vflux", "wflux", "dzt", "zstot"):
np.nan_to_num(arr, copy=False, nan=0.0, posinf=0.0, neginf=0.0)
arr[np.abs(arr) > _FLUX_FILL_THRESHOLD] = 0.0
return arr
def __repr__(self):
lines = ["LazyFluxFields("]
for name in ("uflux", "vflux", "dzt", "zstot", "wflux"):
da = self._past.get(name)
if da is None:
lines.append(f" {name}: absent,")
else:
backed = (
"dask"
if hasattr(da, "data") and hasattr(da.data, "dask")
else "numpy"
)
lines.append(f" {name}: shape={da.shape}+2 (lazy, {backed}),")
lines.append(")")
return "\n".join(lines)
# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------
def _squeeze_t(da: xr.DataArray) -> xr.DataArray:
"""Remove a trailing size-1 time dimension from a DataArray, if present.
GCM adapters that already attach a t-dimension (e.g. the theoretical
adapter) return DataArrays with shape ``(..., 1)``. This helper
normalises them to ``(...)`` so that all adapters produce the same
shape going into :class:`LazyFluxFields`.
"""
# Find and drop any dimension of size 1 that looks like a time slot dim.
# We only squeeze the trailing dim to avoid touching spatial dims of size 1.
dims = da.dims
if dims and da.sizes.get(dims[-1], 2) == 1:
return da.isel({dims[-1]: 0})
return da
def _w_to_kernel(da: xr.DataArray) -> xr.DataArray:
"""Convert an adapter's vertical flux to the kernels' face convention.
Adapters return ``wflux`` the way ocean models store it, with level 0 at
the surface: ``wflux[..., k]`` is the volume flux through the **top**
face of level ``k``, positive **upward** (NEMO ``depthw``, MITgcm
``WVEL``). The kernels read ``wflux[..., k]`` as the flux through the
face between levels ``k`` and ``k + 1`` — the **bottom** face of level
``k`` — positive towards ``k + 1``, i.e. **downward**. The bottom face
of level ``k`` is the top face of level ``k + 1``, so::
w_kernel[..., k] = -wflux[..., k + 1], w_kernel[..., -1] = 0
The shift is lazy (no data is read) and keeps the array shape.
Parameters
----------
da : xr.DataArray
Vertical volume flux with dimensions ``(x, y, z, ...)``.
Returns
-------
xr.DataArray
The flux in the kernel convention, same shape and dimensions.
"""
zdim = da.dims[2]
return -da.shift({zdim: -1}, fill_value=0.0)
def _slot_var(name: str, da: xr.DataArray, reverse: bool) -> xr.DataArray:
"""Prepare one adapter variable for a :class:`LazyFluxFields` slot.
Drops a trailing size-1 time dimension, converts ``wflux`` to the
kernel convention (:func:`_w_to_kernel`) and negates the velocity fluxes
for a backward run. Lazy: nothing is read.
"""
da = _squeeze_t(da)
if name == "wflux":
da = _w_to_kernel(da)
if reverse and name in _FLUX_NEGATION_VARS:
da = -da
return da
[docs]
def subset_offset_attrs(subset: dict) -> dict[str, int]:
"""Dataset attrs recording where an adapter's x/y subset starts.
Subsetting adapters put these on the dataset ``fluxes()`` returns;
:func:`_static_from_ds` turns them into ``GridFields.i_offset`` /
``j_offset`` so recorded positions can be written in full-grid indices.
Parameters
----------
subset : dict
The adapter's ``settings["subset"]`` (``i1``, ``j1``, ``istep``, ...).
Returns
-------
dict[str, int]
``{"_i_offset": i1 or 0, "_j_offset": j1 or 0}``.
Warns
-----
UserWarning
If ``istep`` / ``jstep`` is set: a strided subset's cells are not
full-grid cells, so the offset alone does not make saved positions
absolute.
"""
if (subset.get("istep") or 1) != 1 or (subset.get("jstep") or 1) != 1:
warnings.warn(
"GCM subset uses istep/jstep; recorded x/y are offset by i1/j1 "
"but remain in strided-cell units, not full-grid indices",
stacklevel=3,
)
return {
"_i_offset": int(subset.get("i1") or 0),
"_j_offset": int(subset.get("j1") or 0),
}
def _static_from_ds(ds: xr.Dataset) -> GridFields:
"""Compute static geometry fields from *ds* into a ``GridFields``.
Only variables whose names are in ``_STATIC_VARS`` are processed;
all others are ignored. Each variable is converted to a JAX array
(or Python scalar for 0-D variables).
Parameters
----------
ds : xr.Dataset
Dataset containing at least the static fields required by
``GridFields`` (``kmt``, ``dxdy``, ``imt``, etc.).
Returns
-------
GridFields
"""
static_dict: dict = {}
for name in _STATIC_VARS:
if name not in ds:
continue
da = ds[name]
if np.iterable(da):
arr = np.asarray(da.values, dtype=float)
if name == "kmt":
# Land cells are sometimes stored as the NEMO fill value
# (9.96921e36) rather than 0 when the grid definition netCDF
# was written via xarray's ds.where() without explicit
# _FillValue encoding. Replace any value above a physically
# impossible threshold (no ocean has > 10 000 depth levels)
# with 0 so that land cells are correctly identified.
arr = np.where(
np.isnan(arr) | (arr > _FLUX_FILL_THRESHOLD), 0.0, arr
)
static_dict[name] = jnp.asarray(arr)
else:
static_dict[name] = da.item()
# Subsetting adapters record where the subset starts in the full grid.
static_dict["i_offset"] = int(ds.attrs.get("_i_offset", 0))
static_dict["j_offset"] = int(ds.attrs.get("_j_offset", 0))
return GridFields(**static_dict)
# ---------------------------------------------------------------------------
# swap_fields (operates on LazyFluxFields)
# ---------------------------------------------------------------------------
[docs]
def swap_fields(lazy_fluxes: LazyFluxFields) -> LazyFluxFields:
"""Swap time levels: present (slot 1) becomes past (slot 0).
Mirrors ``mod_swap.F90::swap_time()``. After calling this function,
the caller should provide new present-slot data (e.g. by calling
:func:`update_flux`).
This is a pure Python reference swap — no Dask compute, no JAX.
Parameters
----------
lazy_fluxes : LazyFluxFields
Lazy flux data to swap.
Returns
-------
LazyFluxFields
New ``LazyFluxFields`` where the old present data is now the past
slot. The present slot is left pointing at the old present data
(stale) and should be overwritten by the caller.
"""
return LazyFluxFields(_past=lazy_fluxes._present, _present=lazy_fluxes._present)
# ---------------------------------------------------------------------------
# fluxes_to_grid (static fields eager, flux fields lazy)
# ---------------------------------------------------------------------------
[docs]
def fluxes_to_grid(ds: xr.Dataset) -> tuple[GridFields, LazyFluxFields]:
"""Convert a two-slot concatenated xarray Dataset to ``(GridFields, LazyFluxFields)``.
Static variables (``kmt``, ``dxdy``, dimensions, timestamps) are
computed eagerly into a ``GridFields``. Flux variables
(``uflux``, ``vflux``, ``wflux``, ``dzt``, ``zstot``) are kept as
xarray DataArrays inside a ``LazyFluxFields`` — no I/O is triggered for
them until :meth:`~LazyFluxFields.compute` is called.
The DataArrays stored in the ``LazyFluxFields`` are expected to already
have a trailing time slot dimension of size 2 (as produced by
``xr.concat``). :func:`new_flux` builds the same two-slot structure
directly from two separate ``fluxes()`` calls.
Parameters
----------
ds : xr.Dataset
Two-slot concatenated dataset.
Returns
-------
tuple[GridFields, LazyFluxFields]
"""
grid = _static_from_ds(ds)
# Split the concatenated t-dim=2 DataArrays into separate past/present
# DataArrays (without the t dim) so LazyFluxFields stores them uniformly.
past: dict = {}
present: dict = {}
for name in _FLUX_VARS:
if name not in ds:
continue
da = ds[name]
dims = da.dims
# The trailing dim should be the concat dim "t" of size 2.
# isel with a scalar drops the dim; isel with a slice keeps it.
if dims and da.sizes.get(dims[-1]) == 2:
past[name] = _slot_var(name, da.isel({dims[-1]: 0}), False)
present[name] = _slot_var(name, da.isel({dims[-1]: 1}), False)
else:
# Fallback: no trailing t-dim (shouldn't happen after concat,
# but handle gracefully by using the same DataArray for both).
da_sq = _slot_var(name, da, False)
past[name] = da_sq
present[name] = da_sq
dtm = pd.Timestamp(ds["dtm"].item()) if "dtm" in ds else None
lazy = LazyFluxFields(
_past=past,
_present=present,
dtm=dtm,
timestep=grid.timestep,
reverse=grid.reverse,
)
return grid, lazy
# ---------------------------------------------------------------------------
# new_grid / new_flux
# ---------------------------------------------------------------------------
def _resolve_fluxes(fluxes):
"""Return ``(callable, module)`` for a fluxes callable or adapter name.
A string is imported as ``tracmass.gcm.<name>`` and its ``fluxes``
attribute taken. The settings-driven ``[env.gcm] subset`` table from
the active ``tracmass.toml`` environment is pushed into the adapter
module before its first ``fluxes()`` call (no-op for adapters without a
``settings["subset"]`` slot — theoretical, fabboni — or when config
cannot be loaded).
"""
if isinstance(fluxes, str):
module = importlib.import_module(f"tracmass.gcm.{fluxes}")
fluxes = module.fluxes
else:
module = sys.modules.get(getattr(fluxes, "__module__", ""), None)
try:
from tracmass.config import apply_gcm_subset
if module is not None:
apply_gcm_subset(module)
except Exception: # pragma: no cover - config is optional
pass
return fluxes, module
def _infer_backward(dtm_start: pd.Timestamp, dtm_end) -> bool:
"""``True`` when *dtm_end* is earlier than *dtm_start* (backward run)."""
return dtm_end is not None and pd.Timestamp(dtm_end) < dtm_start
def _fetch_slot0(fn, dtm_start) -> "tuple[xr.Dataset, pd.Timestamp]":
"""Fetch the slot-0 dataset, honoring an explicit *dtm_start* or else
the adapter's own default.
When *dtm_start* is ``None``, *fn* is called with no arguments at all,
so its own default (which for the GLORYS/ECCO2 adapters falls through
to their ``seagrid`` layer's ``open_dataset`` default) determines the
date — rather than a date hard-coded here. The actual date used is
always read back from the returned dataset's ``dtm`` variable.
Returns
-------
tuple[xr.Dataset, pd.Timestamp]
The slot-0 dataset and its actual timestamp.
"""
with timed(log, "new_flux: opening slot 0 (%s)", dtm_start or "default date"):
if dtm_start is not None:
ds1 = fn(pd.Timestamp(dtm_start))
else:
ds1 = fn()
return ds1, pd.Timestamp(ds1.dtm.item())
def _grid_from_ds(ds1: xr.Dataset, *, backward: bool) -> GridFields:
"""Build a :class:`GridFields` from a slot-0 dataset.
``timestep`` is forced positive (the GCM output interval) and
``reverse`` is taken from *backward*, overriding whatever the adapter
wrote into the dataset (adapters always report ``reverse=False``).
"""
grid = _static_from_ds(ds1)
return grid._replace(timestep=abs(ds1.timestep.item()), reverse=backward)
def _lazy_from_ds(
grid: GridFields, fluxes, ds1: xr.Dataset, dtm_end=None
) -> LazyFluxFields:
"""Build the two-slot :class:`LazyFluxFields` for *grid* from *ds1*.
*ds1* is the already-fetched slot-0 (past) dataset. Slot 1 (present) is
loaded by calling *fluxes* one ``grid.timestep`` further forward — or
backward when ``grid.reverse`` — in physical time. ``uflux`` /
``vflux`` / ``wflux`` are negated lazily for backward runs so particles
trace upstream; *fluxes* is never passed ``reverse=``. ``timestep``
and ``reverse`` are copied from *grid*; *dtm_end* is stored as given.
"""
backward = grid.reverse
dtm_past = pd.Timestamp(ds1.dtm.item())
actual_dtm2 = dtm_past - grid.timestep if backward else dtm_past + grid.timestep
with timed(log, "new_flux: opening slot 1 (%s)", actual_dtm2):
ds2 = fluxes(actual_dtm2)
past: dict = {}
present: dict = {}
for name in _FLUX_VARS:
if name in ds1:
past[name] = _slot_var(name, ds1[name], backward)
if name in ds2:
present[name] = _slot_var(name, ds2[name], backward)
zarr_chunk_xy = (
int(ds1.attrs.get("_zarr_chunk_x", 1)),
int(ds1.attrs.get("_zarr_chunk_y", 1)),
)
return LazyFluxFields(
_past=past,
_present=present,
dtm=dtm_past,
fluxes_fn=fluxes,
zarr_chunk_xy=zarr_chunk_xy,
timestep=grid.timestep,
reverse=grid.reverse,
dtm_end=None if dtm_end is None else pd.Timestamp(dtm_end),
)
[docs]
def new_grid(
fluxes,
dtm_start: pd.Timestamp | str | None = None,
dtm_end: pd.Timestamp | str | None = None,
) -> GridFields:
"""Build the static grid geometry for a run.
Calls *fluxes* once (at *dtm_start*) purely to read the time-invariant
geometry — ``kmt``, ``dxdy``, ``imt`` / ``jmt`` / ``km``, ``iperio`` /
``jperio`` — and the GCM output interval (``timestep``). No flux data
is loaded. The returned :class:`GridFields` is immutable for the
lifetime of the run and can be reused across many :func:`new_flux`
buffers on the same grid (successive experiments, restart windows).
Integration direction is inferred from *dtm_start* / *dtm_end* exactly
as in :func:`new_flux`: ``dtm_end < dtm_start`` marks a backward run
(``grid.reverse=True``). ``grid.timestep`` is always positive.
Parameters
----------
fluxes : callable or str
``fluxes(dtm) -> xr.Dataset``, or the name of a ``tracmass.gcm``
sub-module (e.g. ``"glorys"``). When the adapter can be identified,
the ``[env.gcm] subset`` table from the active ``tracmass.toml``
environment is applied first via
:func:`tracmass.config.apply_gcm_subset`.
dtm_start : pd.Timestamp or str or None, optional
Any valid time for the adapter; only the geometry is read. Default
``None`` — *fluxes* is called with no arguments at all, so its own
default applies (for ``gcm.glorys``/``gcm.ecco2`` this falls
through to :func:`tracmass.seagrid.glorys.open_dataset` /
:func:`tracmass.seagrid.ecco2.open_dataset`'s own ``dtm`` default).
dtm_end : pd.Timestamp or str or None, optional
When earlier than the resolved *dtm_start*, marks the grid
backward (``reverse=True``). Default ``None`` (forward).
Returns
-------
GridFields
Static geometry with ``timestep`` (positive) and ``reverse`` set.
See Also
--------
new_flux : build the flux buffer, optionally reusing a grid from here.
"""
fn, _module = _resolve_fluxes(fluxes)
ds1, actual_dtm_start = _fetch_slot0(fn, dtm_start)
return _grid_from_ds(ds1, backward=_infer_backward(actual_dtm_start, dtm_end))
[docs]
def new_flux(
grid_or_fluxes,
fluxes=None,
dtm_start: pd.Timestamp | str | None = None,
dtm_end: pd.Timestamp | str | None = None,
) -> "tuple[GridFields, LazyFluxFields] | LazyFluxFields":
"""Create the lazy two-slot flux buffer for a run.
Two call forms:
``new_flux(fluxes, dtm_start=..., dtm_end=None) -> (GridFields, LazyFluxFields)``
Build the grid and the flux buffer together. Equivalent to
``g = new_grid(fluxes, ...); return g, new_flux(g, fluxes, ...)``
but fetching slot 0 only once. This is the original signature and
return type.
``new_flux(grid, fluxes, dtm_start=..., dtm_end=None) -> LazyFluxFields``
Reuse a :class:`GridFields` from :func:`new_grid` (or an earlier
``new_flux``); build only the flux buffer. ``grid.reverse`` and
``grid.timestep`` drive direction and step size; *dtm_end*, if
given, is only checked for consistency with ``grid.reverse``.
Flux data (``uflux``, ``vflux``, ``wflux``, ``dzt``, ``zstot``) is kept
as lazy xarray DataArrays and **not** loaded into device memory until
:meth:`~LazyFluxFields.compute` is called. For a backward run
(``grid.reverse``) slot 1 is one ``grid.timestep`` further into the
past, the velocity fluxes are negated lazily, and each
:func:`update_flux` call walks ``lazy.dtm`` backward through the
archive. The ``fluxes`` callable is never passed ``reverse=``.
Parameters
----------
grid_or_fluxes : GridFields or callable or str
A :class:`GridFields` selects the grid-reuse form; anything else is
the ``fluxes`` callable / adapter name of the combined form.
fluxes : callable or str, optional
The ``fluxes(dtm) -> xr.Dataset`` callable (or adapter name).
Required — and only used — in the grid-reuse form.
dtm_start : pd.Timestamp or str or None, optional
Physical time of slot 0 (past). Default ``None`` — *fluxes* is
called with no arguments at all, so its own default applies (for
``gcm.glorys``/``gcm.ecco2`` this falls through to
:func:`tracmass.seagrid.glorys.open_dataset` /
:func:`tracmass.seagrid.ecco2.open_dataset`'s own ``dtm`` default).
dtm_end : pd.Timestamp or str or None, optional
When earlier than the resolved *dtm_start*, activates backward
tracking (combined form) or is checked against ``grid.reverse``
(grid-reuse form). Stored as ``lazy.dtm_end``, so
``lazy.n_windows`` gives the number of windows to reach it and
the ``time_dependent*`` integrators need no ``n_windows``.
Returns
-------
tuple[GridFields, LazyFluxFields] or LazyFluxFields
A ``(grid, lazy)`` tuple for the combined form; just ``lazy`` for
the grid-reuse form. ``lazy.dtm`` holds the slot-0 timestamp;
``lazy.timestep`` / ``lazy.reverse`` mirror the grid.
See Also
--------
new_grid : build just the static grid.
"""
if isinstance(grid_or_fluxes, GridFields):
grid = grid_or_fluxes
if fluxes is None:
raise TypeError(
"new_flux(grid, fluxes, ...): the 'fluxes' argument is "
"required when the first argument is a GridFields"
)
fn, _module = _resolve_fluxes(fluxes)
ds1, actual_dtm_start = _fetch_slot0(fn, dtm_start)
implied_backward = _infer_backward(actual_dtm_start, dtm_end)
if dtm_end is not None and implied_backward != grid.reverse:
raise ValueError(
"dtm_start / dtm_end imply "
f"{'backward' if implied_backward else 'forward'} tracking but "
f"grid.reverse={grid.reverse}; rebuild the grid with "
"new_grid(dtm_start=..., dtm_end=...)"
)
return _lazy_from_ds(grid, fn, ds1, dtm_end)
if fluxes is not None:
raise TypeError(
"new_flux(fluxes, dtm_start=..., dtm_end=...) takes no second "
"positional argument; did you mean new_flux(grid, fluxes, ...)?"
)
fn, _module = _resolve_fluxes(grid_or_fluxes)
ds1, actual_dtm_start = _fetch_slot0(fn, dtm_start)
grid = _grid_from_ds(ds1, backward=_infer_backward(actual_dtm_start, dtm_end))
return grid, _lazy_from_ds(grid, fn, ds1, dtm_end)
# ---------------------------------------------------------------------------
# update_flux
# ---------------------------------------------------------------------------
[docs]
def update_flux(*args, fluxes_fn=None) -> LazyFluxFields:
"""Advance the flux buffer by one timestep using a rolling two-slot window.
Two call forms:
``update_flux(lazy, fluxes_fn=None) -> LazyFluxFields``
Step size and direction come from ``lazy.timestep`` /
``lazy.reverse`` (set by :func:`new_flux`), and *fluxes_fn*
defaults to ``lazy.fluxes_fn``.
``update_flux(grid, lazy, fluxes_fn) -> LazyFluxFields``
Older form: step size and direction come from ``grid.timestep`` /
``grid.reverse``. Still needed for a ``lazy`` built by lower-level
helpers that do not record ``timestep``.
The time-slot swap is a pure Python dict reference assignment — the
old present slot's DataArrays become the new past slot at zero cost.
Only a single call to *fluxes_fn* is made to load the next timestep's
DataArrays (still lazy for Dask-backed sources).
No Dask computation, no JAX operations, and no GC are needed here;
all I/O is deferred until the next :meth:`~LazyFluxFields.compute`
call.
The timestep is always positive; for backward runs (``reverse`` is
``True``) ``lazy.dtm`` decrements by one timestep per call. The
``fluxes_fn`` callable is invoked without ``reverse=`` — flux negation
is applied by this function after loading.
Parameters
----------
*args
``(lazy,)``, ``(lazy, fluxes_fn)`` or ``(grid, lazy, fluxes_fn)``;
see the call forms above. ``lazy.dtm`` is the physical timestamp
of the current past (slot 0) time level.
fluxes_fn : callable, optional
``fluxes_fn(dtm) -> xr.Dataset`` for the next window. Defaults to
``lazy.fluxes_fn``.
Returns
-------
LazyFluxFields
New lazy flux data with ``dtm`` advanced (forward) or decremented
(backward) by one timestep; all other metadata is carried over.
Raises
------
TypeError
On a wrong number of positional arguments.
ValueError
If no *fluxes_fn* is given and ``lazy.fluxes_fn`` is unset, or if
the timestep is unknown (a ``lazy`` without ``timestep`` needs the
``(grid, lazy, fluxes_fn)`` form).
Examples
--------
>>> grid, lazy = new_flux(gcm.theoretical.fluxes)
>>> lazy = update_flux(lazy)
"""
if args and isinstance(args[0], GridFields):
if len(args) not in (2, 3):
raise TypeError("update_flux(grid, lazy, fluxes_fn) takes 3 arguments")
grid, lazy_fluxes = args[:2]
if len(args) == 3:
fluxes_fn = args[2]
timestep, reverse = grid.timestep, grid.reverse
else:
if len(args) not in (1, 2):
raise TypeError("update_flux(lazy, fluxes_fn=None) takes 1 or 2 arguments")
lazy_fluxes = args[0]
if len(args) == 2:
fluxes_fn = args[1]
timestep, reverse = lazy_fluxes.timestep, lazy_fluxes.reverse
if fluxes_fn is None:
fluxes_fn = lazy_fluxes.fluxes_fn
if fluxes_fn is None:
raise ValueError(
"update_flux: no fluxes_fn given and lazy.fluxes_fn is not set"
)
if timestep is None:
raise ValueError(
"update_flux: lazy.timestep is not set; build the buffer with "
"new_flux() or use update_flux(grid, lazy, fluxes_fn)"
)
delta = -timestep if reverse else timestep
actual_next_present = lazy_fluxes.dtm + 2 * delta
with timed(log, "update_flux: opening %s", actual_next_present):
new_ds = fluxes_fn(actual_next_present)
# Build new present-slot dict (wflux in kernel convention, velocity
# fluxes negated for backward runs)
new_present: dict = {
name: _slot_var(name, new_ds[name], reverse)
for name in _FLUX_VARS
if name in new_ds
}
# Swap: old present → new past (zero-cost reference assignment).
return dataclasses.replace(
lazy_fluxes,
_past=lazy_fluxes._present,
_present=new_present,
dtm=lazy_fluxes.dtm + delta,
timestep=timestep,
reverse=reverse,
)
