Performance: chunked MarsInterp, lazy MarsFormat, vinterp/find_n without per-level copies - #199
Merged
Conversation
…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.
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 (branchaeolis); 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:
MarsFormat -gcm marswrfMarsInterp -t pstdOn a 16-frame file
MarsInterpgoes 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
vinterp(amescap/FV3_utils.py):Lfull.flatten()andvarIN.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.find_n: the index search looped in Python over every column withnp.argmin(millions of calls on a year of data). For monotonically increasing input along the first axis, the normal case afterreverse_input, the index of the largestX_IN <= X_OUTis 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.MarsInterptime 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.MarsFormatlazy processing: the file is opened with dask chunks along time whendaskis importable, and the marswrf block was rewritten to stay lazy (no 4D.values), soto_netcdfstreams. dask runs with the synchronous scheduler: netCDF4/HDF5 is not thread safe and the threaded scheduler deadlocked on the-babinning write in the test suite; chunking alone gives the memory benefit. Without dask the behaviour is exactly as before, with a one-line hint.daskis added as an optional extra[large]inpyproject.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
MarsInterpoutputs (pstd,zagl,zstdon the full-wrfout bite;pstdon the reduced bite) identical to the pre-change outputs, both with automatic chunking and with-chunk 3/-chunk 5.MarsFormatoutputs (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 (withtest_marsplot.py).Not done
atmos_daily) files were not re-timed here; the same code paths apply, and thefind_nfast path triggers whenever the level field is monotonic, which it is forfms_press_calcoutput.time_of_daydimension) go through the same chunked loop (time is still the first axis); tested only via the unit tests.