About#
The TRACMASS method#
The original idea behind TRACMASS is that the trajectory path through each grid cell is an analytical solution of a differential equation that depends on the velocities on the walls of the grid box. The scheme was first developed by Döös (1995) and Blanke and Raynaud (1997) for stationary velocity fields. Vries and Döös (2001) extended it to time-dependent fields by interpolating the velocity linearly in time as well as in space over each grid box. Time-stepping schemes such as forward Euler or fourth-order Runge–Kutta (Butcher, 2008; Fabbroni, 2009) work differently: they integrate trajectories forward with time steps that are as short as possible.
Lagrangian trajectories and Eulerian tracers have a lot in common, and the two are often confused. Both are carried passively by the velocity fields of the GCM, so both can trace water or air masses, or substances such as pollutants, as currents and winds move them. A tracer equation usually has to be integrated online, together with the GCM. Trajectories can be computed online or offline. Offline is by far the faster option, because it only needs the velocity fields the GCM has already produced.
TRACMASS has been developed over many years and used in studies of the global ocean (Döös and Coward, 1997; Drijfhout et al., 2003; Döös et al., 2008; Berglund et al., 2017), of regional seas such as the Mediterranean and the Baltic (Döös et al., 2004; Jönsson et al., 2004; Engqvist et al., 2006; Soomere et al., 2011), and of the large-scale atmospheric circulation (Kjellsson and Döös, 2012).
The code was first written in Fortran 77 in the early 1990s, for the FRAM ocean model at the Institute of Oceanographic Sciences Deacon Laboratory in Wormley, UK. The name comes from the EU project TRACMASS, in which it was used together with the similar trajectory code Ariane (Blanke and Raynaud, 1997). The current Fortran code, TRACMASS 7, is written in Fortran 90 and can be driven by velocity fields from most finite-difference GCMs.
The JAX port#
TRACMASS-JAX reimplements the TRACMASS trajectory core in Python with
JAX. It is a port, not a new scheme. The equations
for the crossing time and the new position in each grid box are those of
mod_pos_tstep.F90 and mod_pos_tanalytical.F90 (see
Advection schemes). tests/test_fortran_comparison.py
reruns the Fortran Theoretical project and checks the JAX trajectory
against the Fortran output (TRACMASS_run.csv).
The port changes how the computation is organised:
Fortran TRACMASS |
TRACMASS-JAX |
|---|---|
|
|
|
|
|
|
Project folder + |
A GCM adapter (Python module) chosen at run time |
Namelist ( |
|
Text output ( |
Parquet files with embedded run metadata |
Because the kernel is compiled once and then reused for every GCM time step,
and the flux data is passed in as an argument, a run reads each velocity
snapshot from disk once and keeps the expensive part on the accelerator.
GCM data is opened lazily with xarray and Dask, and is only read when the
kernel needs it. Two integrators read only the part of the grid near the
particles (see Memory-efficient integration).
Status#
TRACMASS-JAX is under active development at the same time as it is used for research runs. The core is in place: both advection schemes, forward and backward runs, multi-window time stepping, seeding, recording, three reanalysis adapters and two synthetic test cases. Tracers, stream functions, diffusion, kill zones and the ORCA north fold are not yet ported. See what is supported and the changelog.