Skip to content

Performance: chunked MarsInterp, lazy MarsFormat, vinterp/find_n without per-level copies - #199

Merged
rurata merged 4 commits into
aeolisfrom
aeolis-perf
Aug 28, 2026
Merged

Performance: chunked MarsInterp, lazy MarsFormat, vinterp/find_n without per-level copies#199
rurata merged 4 commits into
aeolisfrom
aeolis-perf

Conversation

@mirfjc

@mirfjc mirfjc commented Aug 27, 2026

Copy link
Copy Markdown

Summary

Performance and memory work on the interpolation path and on MarsFormat. No change in results: every output was checked bit-for-bit (all variables, NaNs included) against the previous code on four files, with the default and with forced small chunk sizes. Builds on #198 (branch aeolis); the three commits here are the ones after it.

Measured on a one-year planetWRF file (7.4 GB, 2676 frames, 60x36x52) on a laptop:

step before after
MarsFormat -gcm marswrf 28 s, 20 GB peak 19 s, 2.0 GB peak
MarsInterp -t pstd 53 min, 21 GB peak 76 s, 5.4 GB peak

On a 16-frame file MarsInterp goes from 13 s to 1.5 s (pstd), 6.4 s to 1.3 s (zagl), 9.5 s to 1.5 s (zstd).

What changed

  1. vinterp (amescap/FV3_utils.py): Lfull.flatten() and varIN.flatten() (full copies) were called up to six times per output level inside the level loop. They are now taken once as contiguous views before the loop. Same arithmetic, same order of operations.
  2. find_n: the index search looped in Python over every column with np.argmin (millions of calls on a year of data). For monotonically increasing input along the first axis, the normal case after reverse_input, the index of the largest X_IN <= X_OUT is now a vectorised count, chunked to bound memory. It gives the same index as the nearest-then-step-down search. Non-monotonic input still takes the original path.
  3. MarsInterp time chunking: interpolation is column-local, so the level field, the indices and each variable are computed and written per time slice (Ncdf.log_variable_slice, new). Chunk size is automatic (about 2e7 input elements) or -chunk N. The copy-through of non-interpolated variables is unchanged.
  4. MarsFormat lazy processing: the file is opened with dask chunks along time when dask is importable, and the marswrf block was rewritten to stay lazy (no 4D .values), so to_netcdf streams. dask runs with the synchronous scheduler: netCDF4/HDF5 is not thread safe and the threaded scheduler deadlocked on the -ba binning write in the test suite; chunking alone gives the memory benefit. Without dask the behaviour is exactly as before, with a one-line hint. dask is added as an optional extra [large] in pyproject.toml, not a hard dependency. The eta fit for reduced files now uses the first 16 frames instead of all (eta is constant; the 8-frame test file is unchanged).

Verification

  • MarsInterp outputs (pstd, zagl, zstd on the full-wrfout bite; pstd on the reduced bite) identical to the pre-change outputs, both with automatic chunking and with -chunk 3 / -chunk 5.
  • MarsFormat outputs (default and -stag, both bites) identical to the pre-change outputs.
  • pytest tests/test_marsformat.py tests/test_marsinterp.py tests/test_marsfiles.py tests/test_marsvars.py: 60 passed, 4 skipped (with test_marsplot.py).

Not done

  • Ames-format (atmos_daily) files were not re-timed here; the same code paths apply, and the find_n fast path triggers whenever the level field is monotonic, which it is for fms_press_calc output.
  • Diurn files (time_of_day dimension) go through the same chunked loop (time is still the first axis); tested only via the unit tests.

mirfjc added 3 commits August 26, 2026 18:44
…on loop

vinterp called Lfull.flatten() and varIN.flatten() (full copies) up to
six times per output level inside the level loop; they are now taken
once as contiguous views before the loop. find_n looped in Python over
every column with np.argmin (millions of calls on a full-year file);
for monotonically increasing input, the normal case after
reverse_input, the index of the largest X_IN <= X_OUT is now a
vectorised count, chunked to bound memory. The original search is kept
for non-monotonic input.

MarsInterp on a 16-frame planetWRF file: pstd 13.0 s -> 1.5 s, zagl
6.4 s -> 1.3 s, zstd 9.5 s -> 1.5 s; reduced file pstd 4.8 s -> 0.7 s.
Output bit-identical to the previous code on all four cases (every
variable, NaNs included).

tests: test_marsinterp, test_marsvars, test_marsfiles 36 passed, 4 skipped.
Interpolation is column-local, so the 3D level field, the indices and
every variable are now computed and written per time slice instead of
holding the whole file (plus float64 pressure and int64 index arrays,
about 3x the file size) in memory. Chunk size is automatic (about 5e7
input elements per chunk) or set with -chunk N. Results are independent
of the chunk size: output bit-identical to the previous code with the
default chunking and with forced 3- and 5-step chunks, on pstd, zagl
and zstd.

Ncdf_wrapper: new log_variable_slice() writes a slice along the first
dimension, defining the variable on first use.

tests: test_marsinterp, test_marsfiles, test_marsvars 36 passed, 4 skipped.
Open the input with dask chunks along time (about 2e7 elements of a
3D field per chunk) and keep the marswrf block lazy (no 4D .values:
zfull, the hydrostatic surface estimate and the PH/PHB path are now
xarray expressions), so that to_netcdf streams instead of holding the
whole dataset plus temporaries in memory. The eta fit for reduced
files uses the first 16 frames (eta is constant).

dask runs with the synchronous scheduler: netCDF4/HDF5 is not thread
safe and the threaded scheduler deadlocked on the -ba binning write
(caught in test_marsformat, all threads waiting on locks). The memory
benefit comes from chunking, not threads.

dask is optional: without it a one-line hint is printed and the
behaviour is exactly as before. Added as the [large] extra in
pyproject.toml.

One-year planetWRF file (7.4 GB): 28 s and 20 GB peak before, 19 s and
2.0 GB peak after. Outputs bit-identical to the previous code on the
full-wrfout bite (default and -stag) and the reduced bite.

tests: test_marsformat, test_marsinterp, test_marsfiles, test_marsvars,
test_marsplot 60 passed, 4 skipped.
@mirfjc
mirfjc requested review from falconstryker and rurata August 27, 2026 02:23
Lowers the automatic chunk size (was 5e7), which cut the full-year
peak memory from 9.6 GB to 5.4 GB at the same run time. This is the
setting the PR numbers were measured with.
@rurata
rurata merged commit 99d3a9f into aeolis Aug 28, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants