Main program

In Fortran TRACMASS the main program TRACMASS.F90 reads the namelist, sets up the grid, seeds, and calls the main loop in mod_loop.F90. TRACMASS-JAX has no single main program. tracmass.run plays the part of TRACMASS.F90, and the integrators (time_dependent and its variants) play the part of mod_loop.F90.

tracmass.run

tracmass.run(dtm_start=None, dtm_end=None, state=None, *, env=None, gcm=None, integrator=None, verbose=None, **kwargs) works like this:

1. Activate the settings environment (env), the counterpart of init_namelist. Every argument left out below is then read from that environment (see Settings).

2. Resolve the run: the GCM adapter ([env.gcm] name), the integrator, and the start and end dates (dtm_start, then dtm_end or dtm_start + runtime).

3. Build the grid and the flux buffer with new_flux(gcm, dtm_start, dtm_end), the counterpart of init_subdomain, init_alloc and setup_grid. The [env.gcm] subset is applied to the adapter first. This returns a static GridFields and a lazy LazyFluxFields.

4. Seed (init_seed): state can be a ParticleState, a function seed(grid) -> ParticleState, or left out to seed from the [env.seed] table.

5. Integrate by calling the integrator with (state, grid, lazy, **kwargs). It returns the final ParticleState.

The integrator loop

time_dependent is the counterpart of loop in mod_loop.F90. Its structure is:

resolve options from arguments and settings (no_wflux, n_windows, timax, ...)
compile the kernel:   step_all = piecewise(grid) | analytical(grid, tseas)
open monitor / recorder
for each window (GCM time step) win = 0 .. n_windows-1:      # intsTimeLoop
    release new particles if a seed_freq release falls in this window
    record positions at the window start
    concrete = lazy.compute()                 # read the two snapshots
    while particles are active and inside the window:        # niterLoop
        intrpg  = frac(mean ts)               # time-interpolation weight
        cap     = time to each particle's next sub-step boundary
        state   = step_all(concrete, state, intrpg, cap)     # ntracLoop (vmap)
        stop particles that exceed timax, make no progress or fall behind
    tally and record the particles that stopped in this window
    free the flux arrays, resume particles paused at the window end
    lazy = update_flux(lazy)                  # swap slots, open next snapshot
record final positions, close monitor / recorder

The three Fortran loops become three different things:

Fortran loop	What it runs over	TRACMASS-JAX
intsTimeLoop	GCM time steps	Python for loop over windows
ntracLoop	trajectories	jax.vmap inside the compiled step_all
niterLoop	sub-steps within a GCM step	lax.while_loop around step_all

Fortran runs the trajectory loop outside the sub-step loop: one trajectory goes through all its sub-steps before the next starts. TRACMASS-JAX turns this around. Every sub-step advances all active particles, and the loop stops when no particle is left in the window. A particle that reaches the end of the window is paused, not stopped (active=False with reason == 0), and resumes in the next window.

Steady state

steady_state(state, grid, fluxes, ...) integrates in a single, frozen velocity field. It is the counterpart of running Fortran TRACMASS with one GCM field and is mostly used for tests. It takes the same integration options as time_dependent, without windows or I/O.

Running from the command line

There is no runtracmass executable. Use a short script, or call tracmass.run from the shell:

uv run python -c "import tracmass; tracmass.run(env='omi')"

The arguments of Fortran’s runtracmass map like this:

./runtracmass argument	TRACMASS-JAX
rerun	Not built in. Select the particles you want with tracmass.load(...) and seed their positions again with seed.from_df_xyz
norun	Not needed: analyse output with tracmass.load and pandas
summary	The monitor prints a per-window and final tally of stopped particles