Skip to content

Streamfunction memory allocation scales poorly with trajectory count (online and offline) #19

Description

@hdrake

Summary

Computing Lagrangian streamfunctions is memory-intensive in a way that scales with the number of trajectories, which makes large runs (high-resolution fields and/or large particle counts) prone to out-of-memory failures. This affects both calculation modes, for different reasons:

  • Online accumulation allocates flux arrays whose size grows linearly with the number of trajectories.
  • Offline accumulation keeps the flux arrays small, but requires writing and re-reading the full trajectory dataset, so its disk and post-processing memory grow with trajectory count × output frequency.

Online: flux arrays scale linearly with the number of trajectories

A streamfunction is partitioned by the zone in which each trajectory terminates (lbas). But that destination zone is only known after a trajectory ends — during accumulation we don't yet know which zone a given crossing belongs to. The current code works around this by giving every trajectory its own slot in the flux arrays (the third dimension is 0:ntracmax) and only sorting the contributions by zone at the very end.

The consequence is that flux storage scales as grid_cells × n_trajectories:

fluxes_xy : imt × jmt × n_trajectories × 8 bytes

Illustrative example (not a real configuration): a 360×180 horizontal grid with 1,000,000 trajectories needs roughly 360 × 180 × 10⁶ × 8 B ≈ 0.5 TB for the barotropic streamfunction alone. Ten times the trajectories ⇒ ten times the memory. Tracer-coordinate streamfunctions are far worse, because they additionally multiply by the tracer-bin resolution (hundreds) and the number of tracers — pushing the requirement into the many-TB range.

The root inefficiency: we allocate per-trajectory storage only because we don't know a priori where each contribution belongs, even though the final result only has as many "bins" as there are killing zones (tens).

Offline: small flux arrays, but it needs the full trajectory dataset

Offline accumulation sidesteps the per-trajectory flux dimension (its flux arrays are sized by the number of zones, not trajectories). But it pays for that elsewhere:

  1. It must write every trajectory's positions to disk (e.g. at each grid-crossing) so the fluxes can be reconstructed afterward, and
  2. the post-processing step loads those trajectory records back into memory to do the reconstruction.

Both costs scale with n_trajectories × records_per_trajectory.

Illustrative example: 1,000,000 trajectories each written ~100 times is ~10⁸ position records — order gigabytes on disk, and a comparable in-memory buffer during post-processing. Higher output resolution or frequency makes both larger.

So offline trades a memory problem for an I/O-and-post-processing-memory problem; it raises the ceiling but does not remove the scaling.

How the linked PRs address this

Together, #16#17#18 make online streamfunction memory scale with the number of zones rather than the number of trajectories.

Remaining / follow-up work

These PRs address the online flux scaling and offer a route that avoids full-trajectory I/O. One related scaling cost is not yet addressed: the offline post-processing step still loads the entire trajectory dataset into memory at once (rather than streaming it), so pure-offline runs retain a trajectory-count-dependent memory footprint in post-processing. Converting that to a streaming pass would be a natural follow-up.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions