diff --git a/AUTHORS.md b/AUTHORS.md index 5df024637..49ecd1e24 100644 --- a/AUTHORS.md +++ b/AUTHORS.md @@ -38,3 +38,4 @@ * Valentin Gebhart * Dahyann Araya * Giovanni Cozzolongo +* Thomas Struys diff --git a/CHANGELOG.md b/CHANGELOG.md index 47c78a827..62d159fa0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,7 @@ Code freeze date: YYYY-MM-DD - Updated Impact Calculation Tutorial (`doc.climada_engine_Impact.ipynb`) [#1095](https://github.com/CLIMADA-project/climada_python/pull/1095). - Makes current `measure` module a legacy module, moving it to `_legacy_measure`, to retain compatibility with `CostBenefit` class and various tests. [#1274](https://github.com/CLIMADA-project/climada_python/pull/1274) +- `HazardForecast.quantile` and `ImpactForecast.quantile` compute quantiles block-wise instead of densifying the entire sparse matrix, greatly reducing peak memory. Results are unchanged. [#1203](https://github.com/CLIMADA-project/climada_python/issues/1203) ### Fixed diff --git a/climada/engine/impact_forecast.py b/climada/engine/impact_forecast.py index e45ae5b2a..c21700f87 100644 --- a/climada/engine/impact_forecast.py +++ b/climada/engine/impact_forecast.py @@ -28,7 +28,11 @@ from ..util import log_level from ..util.checker import size -from ..util.forecast import ForecastMixin, reduce_unique_selection +from ..util.forecast import ( + ForecastMixin, + reduce_unique_selection, + sparse_quantile_axis0, +) from .impact import Impact LOGGER = logging.getLogger(__name__) @@ -445,7 +449,7 @@ def _quantile( concat_kws={"reset_event_ids": True}, ) - red_imp_mat = sparse.csr_matrix(np.quantile(self.imp_mat.toarray(), q, axis=0)) + red_imp_mat = sparse.csr_matrix(sparse_quantile_axis0(self.imp_mat, q)) red_at_event = np.array([red_imp_mat.sum()]) if event_name is None: event_name = f"quantile_{q}" diff --git a/climada/engine/test/test_impact_forecast.py b/climada/engine/test/test_impact_forecast.py index 284e7d910..729ea0cf7 100644 --- a/climada/engine/test/test_impact_forecast.py +++ b/climada/engine/test/test_impact_forecast.py @@ -19,6 +19,8 @@ Tests for Impact Forecast. """ +from unittest.mock import patch + import numpy as np import numpy.testing as npt import pandas as pd @@ -26,6 +28,7 @@ from scipy.sparse import csr_matrix from climada.engine import Impact, ImpactForecast +from climada.util.forecast import sparse_quantile_axis0 from .test_impact import impact_kwargs as imp_kwargs @@ -467,3 +470,21 @@ def test_reduce_dim( imp_fc_reduced.at_event, reduction_results_dim[dim][attr]["at_event"], ) + + +def test_quantile_uses_block_wise_helper(impact_forecast): + """quantile passes the impact matrix itself to the block-wise helper""" + expected = np.quantile(impact_forecast.imp_mat.toarray(), 0.5, axis=0) + + # wraps, so the real helper still runs and the result stays checkable + with patch( + "climada.engine.impact_forecast.sparse_quantile_axis0", + wraps=sparse_quantile_axis0, + ) as helper: + reduced = impact_forecast.quantile(0.5) + + helper.assert_called_once() + matrix, q = helper.call_args.args + assert matrix is impact_forecast.imp_mat + assert q == 0.5 + npt.assert_array_equal(reduced.imp_mat.toarray().squeeze(), expected) diff --git a/climada/hazard/forecast.py b/climada/hazard/forecast.py index 6e8d4e51d..836b4499a 100644 --- a/climada/hazard/forecast.py +++ b/climada/hazard/forecast.py @@ -32,7 +32,11 @@ from climada.hazard.base import Hazard from climada.hazard.xarray import HazardXarrayReader from climada.util.checker import size -from climada.util.forecast import ForecastMixin, reduce_unique_selection +from climada.util.forecast import ( + ForecastMixin, + reduce_unique_selection, + sparse_quantile_axis0, +) LOGGER = logging.getLogger(__name__) @@ -236,7 +240,6 @@ def mean(self, dim: Literal["member", "lead_time"] | None = None): **self._reduce_attrs("mean"), ) - # TODO: Do not densify the entire matrix but compute quantiles column-wise! def _quantile( self, q: float, @@ -253,12 +256,8 @@ def _quantile( q=q, ) - red_intensity = sparse.csr_matrix( - np.quantile(self.intensity.toarray(), q, axis=0) - ) - red_fraction = sparse.csr_matrix( - np.quantile(self.fraction.toarray(), q, axis=0) - ) + red_intensity = sparse.csr_matrix(sparse_quantile_axis0(self.intensity, q)) + red_fraction = sparse.csr_matrix(sparse_quantile_axis0(self.fraction, q)) if event_name is None: event_name = f"quantile_{q}" return HazardForecast( diff --git a/climada/hazard/test/test_forecast.py b/climada/hazard/test/test_forecast.py index cdb8952fd..30e6c3516 100644 --- a/climada/hazard/test/test_forecast.py +++ b/climada/hazard/test/test_forecast.py @@ -20,6 +20,7 @@ """ import datetime as dt +from unittest.mock import patch import numpy as np import numpy.testing as npt @@ -31,6 +32,7 @@ from climada.hazard.base import Hazard from climada.hazard.forecast import HazardForecast, xarray_has_timedelta_bug from climada.hazard.test.test_base import hazard_kwargs +from climada.util.forecast import sparse_quantile_axis0 # See https://docs.xarray.dev/en/stable/whats-new.html#id80 xarray_leadtime = pytest.mark.skipif( @@ -533,6 +535,25 @@ def test_median_quantile(self, haz_fc): np.median(haz_fc.intensity.todense(), axis=0), ) + def test_quantile_uses_block_wise_helper(self, haz_fc): + """quantile passes intensity and fraction themselves to the block-wise helper""" + # wraps, so the real helper still runs and the result stays checkable + with patch( + "climada.hazard.forecast.sparse_quantile_axis0", + wraps=sparse_quantile_axis0, + ) as helper: + reduced = haz_fc.quantile(0.5) + + assert helper.call_count == 2 + intensity_call, fraction_call = helper.call_args_list + assert intensity_call.args[0] is haz_fc.intensity + assert fraction_call.args[0] is haz_fc.fraction + assert intensity_call.args[1] == fraction_call.args[1] == 0.5 + npt.assert_array_equal( + reduced.intensity.toarray().squeeze(), + np.quantile(haz_fc.intensity.toarray(), 0.5, axis=0), + ) + @pytest.mark.parametrize("attr", ["min", "mean", "max", "median", "quantile"]) @pytest.mark.parametrize("dim", ["lead_time", "member", "single"]) def test_reduce_dim_unique_or_single(self, haz_fc, q, attr, dim): diff --git a/climada/util/forecast.py b/climada/util/forecast.py index fb98c6113..993c9c88b 100644 --- a/climada/util/forecast.py +++ b/climada/util/forecast.py @@ -22,6 +22,7 @@ from typing import Any, Literal, Mapping import numpy as np +from scipy import sparse class ForecastMixin: @@ -200,3 +201,54 @@ def reduce_unique_selection( ], **concat_kws, ) + + +def sparse_quantile_axis0( + matrix: sparse.spmatrix, + q: np.typing.ArrayLike, + max_memory_mb: float = 64.0, +) -> np.ndarray: + """Quantile along axis 0 of a sparse matrix, without densifying it whole. + + Equivalent to ``np.quantile(matrix.toarray(), q, axis=0)``, but densifies + only as many columns at a time as fit into ``max_memory_mb``, so peak + memory no longer scales with the number of columns. + + Implicit zeros count as values: a column of ``[0, 0, 0, -1, 9]`` has median + ``0.0``, not the ``4.0`` given by its two stored values alone. + + Parameters + ---------- + matrix : scipy.sparse.spmatrix + Matrix to reduce, of shape (n_rows, n_cols). + q : float or array_like of float + Quantile or sequence of quantiles, each between 0 and 1. + max_memory_mb : float, optional + Approximate memory budget for a single densified block, in megabytes. + At least one column is always densified, so a budget too small for a + single column still works. Default: 64.0. + + Returns + ------- + np.ndarray + Quantiles along axis 0, shaped as ``np.quantile`` would return them. + """ + csc = matrix.tocsc() + n_rows, n_cols = csc.shape + if n_cols == 0: + return np.quantile(csc.toarray(), q, axis=0) + column_bytes = max(n_rows * csc.dtype.itemsize, 1) + block = max(int(max_memory_mb * 1e6) // column_bytes, 1) + + # np.quantile sorts its input, so it copies unless allowed not to. Each block is + # a throwaway nothing else references, so it can be sorted in place: skipping the + # per-block copy measured about twice as fast, at no cost in peak memory. + return np.concatenate( + [ + np.quantile( + csc[:, start : start + block].toarray(), q, axis=0, overwrite_input=True + ) + for start in range(0, n_cols, block) + ], + axis=-1, + ) diff --git a/climada/util/test/test_forecast.py b/climada/util/test/test_forecast.py index 8d0a107e6..596cc1981 100644 --- a/climada/util/test/test_forecast.py +++ b/climada/util/test/test_forecast.py @@ -19,11 +19,15 @@ Tests for ForecastMixin class. """ +from unittest.mock import patch + import numpy as np import numpy.testing as npt import pandas as pd +import pytest +from scipy.sparse import csc_matrix, csr_matrix -from climada.util.forecast import ForecastMixin +from climada.util.forecast import ForecastMixin, sparse_quantile_axis0 def test_forecast_init(): @@ -93,3 +97,101 @@ def test_idx_lead_time(): idx = forecast.idx_lead_time(None) npt.assert_array_equal(idx, np.array([False, False, False, False]), strict=True) + + +class TestSparseQuantile: + """Block-wise sparse quantile helper (issue #1203)""" + + @pytest.mark.parametrize("q", [0.0, 0.25, 0.5, 0.9, 1.0]) + # 1e-9 forces one column per block, 5e-4 gives 3 columns so the last block + # is a partial remainder, 10.0 takes the whole matrix in one block + @pytest.mark.parametrize("max_memory_mb", [1e-9, 5e-4, 10.0]) + def test_matches_dense_quantile(self, q, max_memory_mb): + """Block-wise result must equal the dense reference exactly""" + rng = np.random.default_rng(0) + dense = rng.random((20, 37)) + dense[dense < 0.7] = 0.0 # sparse, with implicit zeros dominating + mat = csr_matrix(dense) + npt.assert_array_equal( + sparse_quantile_axis0(mat, q, max_memory_mb=max_memory_mb), + np.quantile(dense, q, axis=0), + ) + + def test_counts_implicit_zeros(self): + """A column of [0,0,0,-1,9] has median 0, not 4 (the stored-only answer)""" + mat = csr_matrix(np.array([[0.0], [0.0], [0.0], [-1.0], [9.0]])) + npt.assert_array_equal(sparse_quantile_axis0(mat, 0.5), [0.0]) + + def test_negative_values(self): + """Negative stored values must not be confused with implicit zeros""" + dense = np.array([[-3.0, 0.0], [0.0, -1.0], [2.0, 0.0]]) + npt.assert_array_equal( + sparse_quantile_axis0(csr_matrix(dense), 0.5), + np.quantile(dense, 0.5, axis=0), + ) + + def test_all_zero_column(self): + """An entirely empty column yields zero, not NaN""" + npt.assert_array_equal( + sparse_quantile_axis0(csr_matrix((4, 3)), 0.5), np.zeros(3) + ) + + def test_single_row(self): + """One event: the quantile is that event's values for any q""" + dense = np.array([[5.0, 0.0, -2.0]]) + npt.assert_array_equal(sparse_quantile_axis0(csr_matrix(dense), 0.3), dense[0]) + + @pytest.mark.parametrize("q", [-0.1, 1.5]) + def test_quantile_out_of_range(self, q): + """An out-of-range quantile must raise, not return nonsense""" + mat = csr_matrix(np.array([[1.0, 0.0], [0.0, 2.0]])) + with pytest.raises(ValueError, match="Quantiles must be in the range"): + sparse_quantile_axis0(mat, q) + + def test_no_columns(self): + """A matrix with no columns yields an empty result, not an error""" + npt.assert_array_equal( + sparse_quantile_axis0(csr_matrix((5, 0)), 0.5), np.empty(0) + ) + + @pytest.mark.parametrize("dtype", [np.float32, np.float64]) + def test_preserves_dtype(self, dtype): + """The result dtype must match the dense reference, not be forced to float64""" + dense = np.array([[1, 0], [0, 2], [3, 0]], dtype=dtype) + result = sparse_quantile_axis0(csr_matrix(dense), 0.5) + assert result.dtype == np.quantile(dense, 0.5, axis=0).dtype + + def test_multiple_quantiles(self): + """A sequence of quantiles returns one row per quantile, as np.quantile does""" + dense = np.array([[1.0, 0.0], [0.0, 2.0]]) + npt.assert_array_equal( + sparse_quantile_axis0(csr_matrix(dense), [0.25, 0.75]), + np.quantile(dense, [0.25, 0.75], axis=0), + ) + + @pytest.mark.parametrize("max_memory_mb", [0.0, -1.0, 1e-12]) + def test_degenerate_memory_budget(self, max_memory_mb): + """A useless budget must still give the right answer, never uninitialised memory""" + dense = np.array([[1.0, 0.0], [0.0, 2.0], [4.0, 3.0]]) + npt.assert_array_equal( + sparse_quantile_axis0(csr_matrix(dense), 0.5, max_memory_mb=max_memory_mb), + np.quantile(dense, 0.5, axis=0), + ) + + def test_densifies_blocks_not_whole_matrix(self): + """Only column blocks are densified, never the full matrix (issue #1203)""" + dense = np.zeros((6, 12)) + dense[1, 3] = 4.0 + dense[4, 9] = -2.0 + mat = csr_matrix(dense) + # the helper converts to CSC first, so csc_matrix is what densifies + with patch.object( + csc_matrix, "toarray", autospec=True, side_effect=csc_matrix.toarray + ) as spy: + sparse_quantile_axis0(mat, 0.5, max_memory_mb=1e-9) + + densified = [call.args[0].shape for call in spy.call_args_list] + assert densified, "no densification seen at all: the spy is not wired up" + assert ( + max(shape[1] for shape in densified) < mat.shape[1] + ), f"a block covered every column: {densified}"