"""TRACMASS-JAX: Lagrangian particle tracking on JAX-accelerated GCM grids.
This package provides the core data structures, integration routines, and
GCM interfaces for running Lagrangian particle trajectories on ocean model
velocity fields, using JAX for hardware-accelerated computation.
The grid is split into two objects so that flux data can be loaded lazily:
- :class:`~tracmass.gcm.GridFields` — static geometry (``kmt``, ``dxdy``)
and metadata. Closed over by compiled kernels as XLA constants.
- :class:`~tracmass.gcm.FluxFields` — time-varying flux data for a two-slot
window. Passed as a dynamic argument; load it just before calling the kernel.
Typical usage
-------------
>>> import tracmass
>>> grid, fluxes = tracmass.new_flux(tracmass.gcm.theoretical.fluxes)
>>> state = tracmass.gcm.theoretical.initial_state()
>>> result = tracmass.steady_state(state, grid, fluxes)
``new_flux`` builds the grid and the flux buffer together. To reuse one
grid across several flux buffers, build it once with
:func:`~tracmass.gcm.new_grid` and pass it back in::
>>> grid = tracmass.new_grid(tracmass.gcm.theoretical.fluxes)
>>> fluxes = tracmass.new_flux(grid, tracmass.gcm.theoretical.fluxes)
"""
import os
try:
import pynvml
except ImportError:
pynvml = None
if pynvml is not None:
try:
pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
info = pynvml.nvmlDeviceGetMemoryInfo(handle)
print(
f"GPU: {info.free // 1024**2} MiB free / {info.total // 1024**2} MiB total"
)
except pynvml.NVMLError as err:
print(f"GPU: NVML unavailable ({err})")
os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
import gc
import jax
from tracmass import config
settings = config.settings
from tracmass import advection_scheme, gcm
from tracmass.gcm import (
FluxFields,
GridFields,
LazyFluxFields,
new_flux,
new_grid,
swap_fields,
update_flux,
)
from tracmass.integrate import (
steady_state,
time_dependent,
time_dependent_patched,
time_dependent_streaming,
)
from tracmass.driver import run
from tracmass.verbose import set_verbose
from tracmass.output import load
from tracmass.monitoring import RunMonitor
from tracmass.particle import (
REASON_ACTIVE,
REASON_BOUNDS_H,
REASON_BOUNDS_V,
REASON_DXYZ,
REASON_ITERLIMIT,
REASON_LAND,
REASON_PROGRESS,
REASON_SLOW,
REASON_TIME,
ParticleRecorder,
ParticleState,
write_trajectories_parquet,
)
[docs]
def free_device_memory() -> None:
"""Release GPU/device memory held by tracmass between runs.
Clears the JAX JIT compilation cache (which pins ``kmt`` and ``dxdy``
as device constants), runs Python garbage collection to drop any
remaining JAX array references, and waits for all asynchronous device
operations to complete so memory is actually reclaimed before the next
run starts.
Call this in an interactive session (IPython / Jupyter) after
:func:`time_dependent` returns to avoid GPU memory accumulating across
successive runs::
result = tracmass.time_dependent(state, grid, lazy, ...)
tracmass.free_device_memory() # GPU back to baseline
Notes
-----
This also evicts compiled kernels, so the *next* call to
:func:`time_dependent` will retrigger JIT compilation (typically a
few seconds). It has no effect on CPU-only runs.
"""
jax.clear_caches()
gc.collect()
try:
jax.effects_barrier()
except Exception:
pass # no-op on CPU-only or if no device is initialised
__all__ = [
# Core data structures
"ParticleState",
"GridFields",
"FluxFields",
"LazyFluxFields",
# Flux buffer management
"new_grid",
"new_flux",
"update_flux",
"swap_fields",
# Integration
"steady_state",
"time_dependent",
"time_dependent_patched",
"time_dependent_streaming",
"run",
"set_verbose",
# Deactivation reason codes
"REASON_ACTIVE",
"REASON_TIME",
"REASON_LAND",
"REASON_BOUNDS_H",
"REASON_BOUNDS_V",
"REASON_PROGRESS",
"REASON_SLOW",
"REASON_ITERLIMIT",
"REASON_DXYZ",
# Monitoring
"RunMonitor",
# I/O utilities
"ParticleRecorder",
"load",
"write_trajectories_parquet",
# GPU memory management
"free_device_memory",
# Submodules
"gcm",
"advection_scheme",
]
