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.
Performance: chunked MarsInterp, lazy MarsFormat, vinterp/find_n without per-level copies
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
2 participants
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.
This is just the last bit to get aeolis into devel. I think my last pull request, I messed up, and left at aeolis-perf requesting to pull into aeolis, but after aeolis was merged into devel.