Integrators#

An integrator takes a particle state, a grid and a flux buffer, and moves the particles through one or more GCM windows. There are four:

Function	Windows	Flux data per window	Kernel loop
`steady_state`	1, frozen field	whole grid	`lax.while_loop`
`time_dependent`	many	whole grid (or the `subset`)	`lax.while_loop`
`time_dependent_patched`	many	bounding box of the active particles	`lax.while_loop`
`time_dependent_streaming`	many	the six faces of each particle’s cell	Python `for` loop

All of them use the same advection equations (see Advection schemes). They differ in how much flux data they read, which decides memory use and speed.

Windows and sub-steps#

A window is one GCM timestep, the interval between two snapshots (grid.timestep, the Fortran ngcm_step). Inside a window:

Particle time is kept in seconds (tt) and in normalised form (ts = tt / timestep), so window n runs from ts = n to ts = n + 1.
The window is split into iter_sub sub-steps. A sub-cycle advances each active particle until it hits a wall or reaches its next sub-step boundary, whichever comes first. A particle can need many sub-cycles per sub-step if it crosses many cells.
The loop runs until no active particle is left inside the window. Particles that reach the window end are paused (active=False, reason == 0) and resumed at the start of the next window.

The number of windows is n_windows if given, else lazy.n_windows (from the dtm_end passed to new_flux), else 1.

Safety limits#

Some particles never reach the end of a window: they are trapped in a convergence zone, oscillate at a wall, or move so slowly that they hold everyone else up. The integrators stop them and record a reason:

Check	Reason
A sub-cycle made no progress in time	`no_progress`
A particle falls behind a pace of one window per `(max_sub_iter_per_window - 2) / pace_factor` sub-cycles	`too_slow`
The window used `max_sub_iter_per_window` sub-cycles (default 10 000)	`iter_limit`
`tt - t0` exceeded the trajectory lifetime	`time_limit`

`time_dependent`#

The basic multi-window integrator. Each window it reads both snapshots of the whole grid (or the [env.gcm] subset) with lazy.compute(), runs the compiled loop, and moves on with update_flux. Use it for the synthetic test cases, for regional subsets, and whenever particles cover a large part of the grid.

Options of its own:

scheme: "tstep" or "analytical".
n_particle_batches: run the particles in this many groups, one compiled loop per group, to limit GPU memory. Between windows the particle state is kept on the host.
rampup_windows, rampup_substeps: extra saves in the first windows (see Output).

`time_dependent_patched`#

Reads only a rectangular patch: the bounding box of the active particles plus patch_buffer cells on each side. The patch is snapped to the chunk boundaries of the source Zarr store, so no chunk is read partially. At each window the patch is re-centred on the particles, which costs nothing until it is read. It only ever grows, never shrinks, so the kernel (piecewise_local, which takes the patch’s kmt and dxdy as arguments) is recompiled only when the particle cloud has spread beyond the current patch size.

This is several times faster than time_dependent for clustered particles on a large grid, such as a regional release in a global reanalysis. It is not a good fit when the particles spread over most of the globe, because the patch then grows to cover it.

`time_dependent_streaming`#

Reads only what each particle needs: the six face fluxes and the geometry of its current cell (a FluxStencil), gathered with one batched Dask read per sub-cycle (gcm.extract_stencil). Because the data a particle needs depends on where it is, the sub-cycle loop runs in Python instead of lax.while_loop. Two optimisations keep the number of reads down:

only active particles are read;
a stencil cache keeps each particle’s stencil until it moves to another cell, and is cleared when a new window starts.

Use it for widely scattered particles and very long runs, where any bounding box would cover most of the grid. Its cost grows with the number of sub-cycles, not with the size of the grid. It always uses the time-stepping scheme and does not support ramp-up saves.

Choosing an integrator#

Situation	Integrator
Synthetic tests, small or subset grids	`time_dependent`
A regional release in a global model, run for weeks to months	`time_dependent_patched`
Particles all over the globe, or runs of many years	`time_dependent_streaming`
Comparing with Fortran time-analytical runs	`time_dependent(scheme="analytical")`
Many particles on a GPU that runs out of memory	`time_dependent` with `n_particle_batches`, or a `subset`

See Memory-efficient integration for examples and timing advice.

Common arguments#

All three multi-window integrators accept:

integrator(state, grid, lazy, fluxes_fn=None, *,
           iter_sub=1, max_iter=100_000, n_windows=None,
           max_sub_iter_per_window=10_000,
           timax_seconds=..., traj_max_time=...,
           recorder=..., monitor=..., no_wflux=...,
           pace_factor=5, n_particle_batches=...,
           rampup_windows=..., rampup_substeps=..., seed_freq=...)

Arguments shown as ... default to the matching setting (see Settings). fluxes_fn defaults to lazy.fluxes_fn. monitor and recorder take an instance (which you keep ownership of), a string ("rich", "traj.parquet", built and closed for you), None to turn them off, or nothing to use the settings.