TRACMASS-JAX#

TRACMASS-JAX is a JAX port of TRACMASS, the Lagrangian trajectory code for ocean and atmosphere general circulation models (GCMs). Like the Fortran code, it computes trajectories offline from GCM velocity fields. Each particle crosses each grid box along an analytical solution of the velocities on the box walls. The time stepping is written in JAX, so all particles are advanced at once with jax.vmap and the inner loop is compiled with XLA. It runs on CPUs and on NVIDIA GPUs.

Get started

Install the package and run the synthetic double-gyre case in a few lines.

Quick start

Configure a run

Every tracmass.toml key, and how the Fortran namelist maps onto it.

Settings reference

Components

The driver, the three-layer grid split, the advection schemes and the GCM adapters.

Components

Tutorials

Seeding, stopping trajectories, backward runs, subsets, adding a new GCM.

Tutorials

Output

The Parquet trajectory files, their metadata, and tracmass.load.

Output

API reference

Generated from the NumPy-style docstrings in the source.

API reference

Why a JAX port?#

TRACMASS has been used for more than three decades, in studies of the global ocean, the Baltic and Mediterranean Seas, coastal regions and the atmosphere (see About). The Fortran code loops over trajectories one at a time. TRACMASS-JAX keeps the same scheme and the same box-wall equations. It changes how the work is organised:

All particles at once. The per-particle kernel is vectorised with jax.vmap and the sub-step loop runs inside lax.while_loop, so millions of particles share one compiled kernel, on the CPU or the GPU.
Lazy GCM data. Velocity fields come from Dask-backed xarray datasets (Zarr, NetCDF, OPeNDAP). Nothing is read from disk until the kernel needs it, and only the part of the grid near the particles has to be read.
Settings, not recompilation. Fortran TRACMASS is configured with a project folder, a Makefile and a namelist, and is recompiled for each case. A TRACMASS-JAX run is one Python call, tracmass.run(env=...), with the options in a tracmass.toml file (see Settings).
Analysis-ready output. Trajectories go to Parquet files with the run metadata embedded, and tracmass.load reads them back as pandas DataFrames, optionally with longitude and latitude.

What is supported#

	Fortran TRACMASS 7	TRACMASS-JAX
Time-stepping scheme (`mod_pos_tstep`)	✓	✓ `scheme="tstep"` (default)
Time-analytical scheme (`mod_pos_tanalytical`)	✓	✓ `scheme="analytical"`
Forward and backward trajectories	✓	✓ direction from `dtm_start` / `dtm_end`
Zonal periodicity (`iperio`)	✓	✓
North fold (`jperio = 1`, ORCA grids)	✓	—
Subdomains (`l_subdom`)	✓	✓ `[env.gcm] subset`
Seeding from boxes, masks, files	✓	✓ plus lat/lon boxes, points and lists
Repeated seeding (`isec`, `tst1`/`tst2`)	✓	✓ `seed_freq`
Maximum trajectory time (`timax`)	✓	✓ `traj_max_time`
Kill zones (`exitType`)	✓	domain edges and land only
Tracers, stream functions, diffusion	✓	not yet
GCMs	NEMO, MOM5, ROMS, IFS, AVISO, …	GLORYS12, ECCO2, Scripps GoG MITgcm, synthetic test cases

The Fortran-to-JAX map lists each Fortran routine and namelist variable with its TRACMASS-JAX equivalent.

Citing#

If you use TRACMASS-JAX, please cite the TRACMASS reference paper and the software:

Döös, K., Jönsson, B., and Kjellsson, J. (2017). Evaluation of oceanic and atmospheric trajectory schemes in the TRACMASS trajectory model v6.0. Geoscientific Model Development, 10, 1733–1749. https://doi.org/10.5194/gmd-10-1733-2017

Jönsson, B., Döös, K., and Kjellsson, J. (2015). TRACMASS: Lagrangian trajectory code. Zenodo. https://doi.org/10.5281/zenodo.34157