# 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](../configuration/settings.md)).
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`](grid.md#gridfields) and a lazy
   [`LazyFluxFields`](grid.md#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:

```text
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:

```bash
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 |