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

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions