perf: compute forecast quantiles block-wise to bound memory - #1313
Conversation
HazardForecast.quantile and ImpactForecast.quantile densified the entire sparse matrix before calling np.quantile, so peak memory scaled with n_events * n_centroids regardless of sparsity. On large forecasts this exhausts memory and kills the process. Add climada.util.forecast.sparse_quantile_axis0, which densifies at most block_size columns at a time. Each column's quantile is independent, so the result is unchanged; only the peak allocation differs. Implicit zeros still take part in the quantile, as reducing over stored values alone would silently inflate the result on sparse data. The blocks are temporaries owned by the helper, so np.quantile may partition them in place. That spares a second full-size copy: np.quantile otherwise duplicates its input, which is why the old peak was twice the dense size and why the block-wise version is faster rather than slower. Measured on synthetic forecasts (density 0.02), results asserted bit-identical to the previous code path: 200 ev x 5k centr: 16.2 -> 8.3 MB peak, 2.41x faster 500 ev x 20k centr: 160.8 -> 35.9 MB peak, 1.67x faster 1000 ev x 50k centr: 802.1 -> 79.0 MB peak, 1.96x faster The existing test_reduce, test_quantile_min_max and test_median_quantile assert against the dense np.quantile reference and pass unmodified. This PR fixes CLIMADA-project#1203
peanutfun
left a comment
peanutfun
left a comment
There was a problem hiding this comment.
@EnvDroneSense Thank you for raising a PR! I am happy about your contribution and your code generally looks good, but your PR description is lengthy, has many unsupported claims, and overall looks AI-generated.
In your "Results" section you claim a speedup and reduction in memory usage with the new code. How did you conduct these measurements? What code did you execute? Later, you claim that setting overwrite_input=True is faster. I see that it is reasonable to save memory, but the entire addition of the new function already serves that purpose. Did you measure the additional speedup as well? Also, the entire "Speedup" section never shows speedup but discusses saving memory. Finally, you claim that results are identical, but the unit tests use assert_allclose instead of assert_array_equal.
The dim is not None branch of both _quantile methods delegates to reduce_unique_selection and is not changed here. It was not measured, and if it densifies by another route then #1203 is only partly addressed by this PR. Happy to extend the scope if you would prefer that handled in one go.
This is not necessary. The "dim is not None branch" calls reduce_unique_selection, which performs a subselection and then delegates back to the calling attribute, in this case _quantile, so there is no extension needed to apply your changes to that particular branch.
I would be happy if you could answer my questions regarding your benchmarks. Please also review my code comments, which I would like to see resolved before approving a merge.
Again, thank you for your contribution!
Replace the block_size column count with a max_memory_mb budget. - preserve the input dtype instead of forcing float64 - accept a sequence of quantiles, as np.quantile does - handle a matrix with no columns - fix the densification spy: it patched csr_matrix while the helper densifies csc_matrix, so it asserted nothing
|
@peanutfun sorry for the messy PR. (I have rewritten it) On the speedup: part of that was me misinterpreting my own benchmark results, so I have redone them from scratch. This time I measured against real data rather than synthetic icon_full.py, and the results look quite good: On overwrite_input=True: I measured it separately this time. At the default 8 MB budget it makes no difference to peak memory (280.5 MB either way) but is about 2x faster (2078 vs 4985 ms), because without it every block gets copied before sorting. So it is a speed optimization rather than a memory one. icon_overwrite.py I have also implemented the changes you recommended. One thing I am unsure about: the default for max_memory_mb is currently an arbitrary 8 MB. Is that acceptable, and if not, how would you suggest picking a better one? Should it go into climada.conf instead? Thank you for your review, I am trying to learn more about code optimization. |
|
@EnvDroneSense Excellent work! Using your benchmark code, I could replicate these numbers on my computer. The tests also run through on my end. A final request: Please update the test |
Mock sparse_quantile_axis0 where each forecast class uses it and assert it receives the matrix object itself, rather than inferring the call from a toarray() spy. Add the matching test for HazardForecast, which reduces intensity and fraction and so must call the helper twice.
|
@peanutfun Great. Both tests mock |
peanutfun
left a comment
peanutfun
left a comment
There was a problem hiding this comment.
Excellent work! Thanks again
2 participants
Changes proposed in this PR:
climada.util.forecast.sparse_quantile_axis0, which takes a quantile along axis 0 of a sparse matrix by densifying a block of columns at a time rather than the whole matrix. The block size is derived from amax_memory_mbbudget.HazardForecast._quantile(intensity and fraction) andImpactForecast._quantile(imp_mat), replacing the.toarray()on the full matrix.This PR fixes #1203
Peak memory no longer scales with the number of centroids. Measured on a full DWD ICON-EU-EPS run, 2560 events x 164,984 cells, float32, with [icon_full.py](https://github.com/user-attachments/files/31507435/icon_full.py) :
The script asserts the two paths return equal arrays on every run.
One constraint worth flagging: implicit zeros have to take part in the quantile. A column of
[0, 0, 0, -1, 9]has median 0.0, not the 4.0 you get from its two stored values. That is why the helper densifies blocks instead of reducing overmatrix.data, which would be faster but silently wrong.