Fortran-to-JAX map#

For users coming from Fortran TRACMASS 7: where each piece of the Fortran code went. For the namelist variables see the namelist-to-settings map.

Routines#

Fortran (TRACMASS 7)	TRACMASS-JAX
`TRACMASS.F90` (main program)	`tracmass.run` (`driver.py`)
`mod_init.F90 :: init_namelist`	`tracmass.config.settings`, `settings.setenv(env)`
`mod_subdomain.F90 :: init_subdomain`	`[env.gcm] subset`, `config.apply_gcm_subset`
`setup_grid.F90` (project)	adapter `fluxes()` → `new_grid` / `_grid_from_ds`
`read_field.F90` (project)	adapter `fluxes()` (lazy) + `LazyFluxFields.compute()`
`mod_getfile.F90`	the `tracmass.seagrid` modules (`open_dataset`, `grid_definition`)
`mod_seed.F90 :: init_seed`, `seed`	`tracmass.seed.*`, `seed.from_settings`
`mod_loop.F90 :: loop`	`time_dependent` / `_patched` / `_streaming`
`intsTimeLoop` (GCM steps)	Python loop over windows
`ntracLoop` (trajectories)	`jax.vmap` in `step_all`
`niterLoop` (sub-steps)	`lax.while_loop` (`_run_window`), or a Python loop for streaming
`mod_pos_tstep.F90 :: cross_time`	`advection_scheme.full_grid._cross_time_1d`
`mod_pos_tstep.F90 :: calc_pos`	`advection_scheme.full_grid._calc_pos_1d`
`mod_pos_tstep.F90 :: update_traj`, `update_bounce`	`advection_scheme.full_grid._subcycle_one_particle`
`mod_pos_tanalytical.F90 :: cross_time` (`apos`, `aneg`, `azer`)	`_cross_time_analytical_1d` (`_apos_crossing`, `_aneg_crossing`, `_azer_crossing`)
`mod_pos_tanalytical.F90 :: calc_pos`	`_calc_pos_analytical_1d`
`mod_pos_tanalytical.F90 :: daw`, `erfexp`	`_dawson`, `_erfexp_val`
`mod_swap.F90 :: swap_time`	`gcm.swap_fields`, `update_flux`
`mod_swap.F90 :: swap_sign` (`nff = -1`)	`dtm_end < dtm_start` in `new_flux` (negates fluxes, `grid.reverse = True`)
`mod_vertvel.F90`	not ported: `wflux` comes from the GCM; `move_vertically = false` for 2-D
`mod_calendar.F90`	`pandas` timestamps (`dtm`, `dtm_start`, `dtm_end`)
`mod_clock.F90`	`ts` / `tt` in `ParticleState`, `intrpg = frac(ts)`
`mod_error.F90`	`ParticleState.reason` codes
`kill_zones.F90` (project)	land, domain edges and `traj_max_time` only (see Stopping trajectories)
`mod_print.F90`	`RunMonitor`
`mod_write.F90`	`ParticleRecorder`, `write_trajectories_parquet`
`mod_vars.F90` `trajectory` type	`ParticleState` (struct of arrays)
`mod_postprocess.F90`	`tracmass.load` + `pandas`
`mod_stream.F90`	not yet
`mod_tracers.F90`, `mod_tracerf.F90`	not yet
`mod_diffusion.F90`	not yet
`mod_divergence.F90`	not yet

Variables#

Fortran	TRACMASS-JAX
`x1`, `y1`, `z1`	`state.x`, `state.y`, `state.z`
`ia`, `ja`, `ka` / `ib`, `jb`, `kb`	`state.ia`, `state.ja`, `state.ka`
`tt`, `ts`, `t0`	`state.tt`, `state.ts`, `state.t0`
`subvol`	`state.subvol`
`niter`	`state.niter`
`nrj(6, ntrac)` / exit flags	`state.active`, `state.reason`
`uflux`, `vflux`, `wflux` (`imt, jmt, km, 2`)	`FluxFields.uflux`, … (`imt, jmt, km, 2`)
`kmt`, `dxdy`, `dzt`	`GridFields.kmt`, `GridFields.dxdy`, `FluxFields.dzt`
`intrpg`, `intrpr`	`intrpg`, `1 - intrpg`
`ds`, `dse`, `dsw`, `dsn`, `dss`, `dsu`, `dsd`, `dsmin`	same names in `_subcycle_one_particle`
`tseas`	`grid.timestep.total_seconds()`
`iter`	`iter_sub`
`UNDEF`	`advection_scheme.full_grid.UNDEF = 1e20`

Behavioural differences#

Particles in lock-step. All particles advance together, one sub-cycle at a time. Fortran finishes one trajectory’s sub-steps before moving to the next. The time-interpolation weight is shared across particles, as in Fortran (see Time interpolation). tests/test_fortran_comparison.py checks the result against the Fortran Theoretical run.
Single precision positions. Positions and fluxes are 32-bit by default. Set JAX_ENABLE_X64=True (the pixi environment does) for 64-bit.
Domain edges. On non-periodic axes, entering the outermost row or column stops the particle (out_of_bounds_h), because the flux through the outer face of those cells is not part of the grid.
Safety stops. Particles that make no progress, fall far behind or hit the per-window sub-cycle limit are stopped, each with its own reason code.