Memory-efficient integration#

For large grids (GLORYS12 is 4318 × 3059 × 50 cells), reading the full flux field into memory every window is the main cost of a run, and can exceed GPU memory. This tutorial shows the tools for reducing it, roughly in the order to try them.

1. Read less: subsets and 2-D runs#

The cheapest data is data you never read:

[myrun]
move_vertically = false      # wflux is not read at all

[myrun.gcm]
name = "glorys"
subset = { i1 = 2600, i2 = 3700, j1 = 500, j2 = 1600, k1 = 1, k2 = 1 }

LazyFluxFields.compute() warns when a single window is larger than 8 GB. time_dependent reads the first window before the monitor starts, so a global run without a subset can show nothing for minutes. Use verbose=True to see what it is doing (see Monitoring).

2. Read near the particles: time_dependent_patched#

result = tracmass.time_dependent_patched(
    state, grid, lazy,
    patch_buffer=50,       # extra cells around the particle bounding box
)

Each window reads only the bounding box of the active particles plus patch_buffer cells on each side, lined up with the Zarr chunks of the source. For a clustered release in a global model this can be several times faster than time_dependent. A larger patch_buffer means fewer patch resizes and recompilations, at the cost of reading more data.

3. Read per particle: time_dependent_streaming#

result = tracmass.time_dependent_streaming(state, grid, lazy)

Each sub-cycle reads just the six face fluxes of each active particle’s cell, in one batched Dask call, and caches them until the particle moves. This suits widely scattered particles and runs of thousands of windows. The cost is a Python-level loop over sub-cycles, about 10–20 per window for daily GLORYS.

Ideas for making streaming faster that are not implemented yet: overlap I/O with computation in a background thread, use Dask’s threaded scheduler for the stencil reads, and store the Zarr data with smaller horizontal chunks (about 100 × 100 instead of 1000 × 1000).

4. Cap GPU memory: n_particle_batches#

With many millions of particles, the scratch memory of the compiled loop can exceed the GPU’s memory even when the fluxes fit. n_particle_batches splits the particles into equal groups, one compiled loop per group. The particle state is kept on the host between windows:

[default]
n_particle_batches = 8

This applies to steady_state and time_dependent. The other integrators accept the argument but do not batch.

5. Free memory between runs#

In an interactive session, compiled kernels keep kmt and dxdy on the device as constants. Release them between runs:

result = tracmass.time_dependent(state, grid, lazy)
tracmass.free_device_memory()      # clears JAX caches, collects garbage

The next run compiles its kernel again, which takes a few seconds.

Reading speed#

On the GLORYS Zarr stores, read speed depends mostly on how the request lines up with the chunks:

  • Daily velocity Zarrs are stored in (10, 1000, 1000) chunks (depth, y, x). The grid definition is opened with chunks={}, so its chunks match the stored ones instead of being rechunked by Dask.

  • wvel stores downloaded before 2026-10-08 have one 50-level chunk per tile, so reading a few levels decompresses all 50. seagrid.glorys.rechunk_wvel(dtm1, dtm2) rewrites them in place with 10-level chunks (20–45 s per day). Don’t rechunk days that a running job is about to read.

  • The patched integrator snaps its patches to chunk boundaries.

  • compute() keeps the present slot’s arrays, and after update_flux reuses them as the next window’s past slot, so each snapshot is read only once.

A typical daily-GLORYS run with time_dependent_patched on a regional release spends about 0.7 s per window on I/O.