From a4795d5bf97b2a111010848414ab246f602dad11 Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 25 Sep 2026 20:25:43 +0200 Subject: [PATCH 1/3] perf(csr): run each contraction chunk on its used columns only --- doc/release_notes.rst | 1 + linopy/csr.py | 22 +++++++++++++++++++++- 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index ea9416add..ddfaa4d46 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -38,6 +38,7 @@ Upcoming Version *Sparse model* * New read-only model key ``Model(sparse=True)`` turns on the whole sparse (CSR) path for one model: ``groupby(...).sum()`` and ``@``/``dot`` against a constant return CSR-backed expressions, and ``add_constraints`` freezes every constraint unless ``freeze=False`` is passed. It requires the v1 semantics: ``Model(sparse=True)`` raises under legacy, and so do ``groupby(...).sum()`` and ``@`` on a sparse model, and ``read_netcdf`` of one, once the semantics are switched back to legacy. It rejects ``chunk``. The key is kept by ``Model.copy`` and the netcdf round trip. ``@`` no longer reads ``linopy.options["sparse_groupby"]``. (`#976 `__) +* ``@``/``dot`` against a sparse constant no longer scales with the total number of variables in the model: each chunk of the sparse product now runs on only the variables it uses, which removes several seconds of allocation overhead on models with millions of variables. * Deprecated in favour of ``Model(sparse=True)``, each with a ``FutureWarning`` and to be removed with the legacy semantics: ``Model(freeze_constraints=...)`` and the ``Model.freeze_constraints`` setter, ``groupby(...).sum(sparse=...)`` and ``linopy.options["sparse_groupby"]``. They keep their current behaviour until then, except that ``@`` ignores ``sparse_groupby``, and netcdf files that store ``freeze_constraints`` still load. (`#976 `__) * Adding a frozen constraint from a sparse expression no longer copies the lhs matrix when every row stays active, and picks the mask and the row scaling at the active rows without expanding them over the full coordinate grid. This roughly halves the peak memory of ``add_constraints`` on a sparse model. (`#977 `__) diff --git a/linopy/csr.py b/linopy/csr.py index e20ac1c4b..e76a921b6 100644 --- a/linopy/csr.py +++ b/linopy/csr.py @@ -647,7 +647,7 @@ def contracted( ) ) rows = slice(start * n_contracted, (start + size) * n_contracted) - blocks.append(block @ source.csr[rows]) + blocks.append(column_compacted_matmul(block, source.csr[rows])) const_blocks.append(block @ const[rows]) indexes = kept_grid.indexes | {str(i.name): i for i in new_indexes} @@ -696,6 +696,26 @@ def index_dtype(nnz: int, shape: tuple[int, ...], model: Model) -> np.dtype: return dtype if max(nnz, *shape) <= np.iinfo(dtype).max else np.dtype(np.int64) +def column_compacted_matmul( + left: scipy.sparse.csr_array, right: scipy.sparse.csr_array +) -> scipy.sparse.csr_array: + """ + ``left @ right`` on the columns ``right`` uses only. + + scipy sizes its product scratch to the column count, which for a model + CSR is every variable label; compacting keeps it to the used columns. + """ + used, compact = np.unique(right.indices, return_inverse=True) + product = left @ scipy.sparse.csr_array( + (right.data, compact.reshape(-1), right.indptr), + shape=(right.shape[0], used.size), + ) + return scipy.sparse.csr_array( + (product.data, used[product.indices], product.indptr), + shape=(left.shape[0], right.shape[1]), + ) + + def coo_to_csr( data: np.ndarray, rows: np.ndarray, From 2e54c27c22b546a5eef321cfb17082bb7fa4bfcd Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 25 Sep 2026 20:25:52 +0200 Subject: [PATCH 2/3] docs: link #990 in release notes --- doc/release_notes.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index ddfaa4d46..24d5add07 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -38,7 +38,7 @@ Upcoming Version *Sparse model* * New read-only model key ``Model(sparse=True)`` turns on the whole sparse (CSR) path for one model: ``groupby(...).sum()`` and ``@``/``dot`` against a constant return CSR-backed expressions, and ``add_constraints`` freezes every constraint unless ``freeze=False`` is passed. It requires the v1 semantics: ``Model(sparse=True)`` raises under legacy, and so do ``groupby(...).sum()`` and ``@`` on a sparse model, and ``read_netcdf`` of one, once the semantics are switched back to legacy. It rejects ``chunk``. The key is kept by ``Model.copy`` and the netcdf round trip. ``@`` no longer reads ``linopy.options["sparse_groupby"]``. (`#976 `__) -* ``@``/``dot`` against a sparse constant no longer scales with the total number of variables in the model: each chunk of the sparse product now runs on only the variables it uses, which removes several seconds of allocation overhead on models with millions of variables. +* ``@``/``dot`` against a sparse constant no longer scales with the total number of variables in the model: each chunk of the sparse product now runs on only the variables it uses, which removes several seconds of allocation overhead on models with millions of variables. (`#990 `__) * Deprecated in favour of ``Model(sparse=True)``, each with a ``FutureWarning`` and to be removed with the legacy semantics: ``Model(freeze_constraints=...)`` and the ``Model.freeze_constraints`` setter, ``groupby(...).sum(sparse=...)`` and ``linopy.options["sparse_groupby"]``. They keep their current behaviour until then, except that ``@`` ignores ``sparse_groupby``, and netcdf files that store ``freeze_constraints`` still load. (`#976 `__) * Adding a frozen constraint from a sparse expression no longer copies the lhs matrix when every row stays active, and picks the mask and the row scaling at the active rows without expanding them over the full coordinate grid. This roughly halves the peak memory of ``add_constraints`` on a sparse model. (`#977 `__) From 7578d857951551a37876f7d26851c72482a52f44 Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 25 Sep 2026 20:39:32 +0200 Subject: [PATCH 3/3] perf(csr): sort compacted product and test column_compacted_matmul --- linopy/csr.py | 6 +++--- test/test_csr.py | 49 +++++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 51 insertions(+), 4 deletions(-) diff --git a/linopy/csr.py b/linopy/csr.py index e76a921b6..2222bd503 100644 --- a/linopy/csr.py +++ b/linopy/csr.py @@ -700,16 +700,16 @@ def column_compacted_matmul( left: scipy.sparse.csr_array, right: scipy.sparse.csr_array ) -> scipy.sparse.csr_array: """ - ``left @ right`` on the columns ``right`` uses only. + ``left @ right`` on the columns ``right`` uses only, label-ordered. scipy sizes its product scratch to the column count, which for a model CSR is every variable label; compacting keeps it to the used columns. """ used, compact = np.unique(right.indices, return_inverse=True) product = left @ scipy.sparse.csr_array( - (right.data, compact.reshape(-1), right.indptr), - shape=(right.shape[0], used.size), + (right.data, compact, right.indptr), shape=(right.shape[0], used.size) ) + product.sort_indices() return scipy.sparse.csr_array( (product.data, used[product.indices], product.indptr), shape=(left.shape[0], right.shape[1]), diff --git a/test/test_csr.py b/test/test_csr.py index 1a371ec56..a3349105e 100644 --- a/test/test_csr.py +++ b/test/test_csr.py @@ -26,7 +26,7 @@ from linopy import LinearExpression, Model, QuadraticExpression, Variable from linopy.constants import TERM_DIM from linopy.constraints import Constraint, ConstraintBase, CSRConstraint -from linopy.csr import CSRLinearExpression, Grid +from linopy.csr import CSRLinearExpression, Grid, column_compacted_matmul from linopy.semantics import is_v1 from linopy.testing import ( assert_conequal, @@ -1090,6 +1090,53 @@ def gen_csr(c: Case) -> CSRLinearExpression: return CSRLinearExpression.from_dense((c.eff * c.gen_p).data, c.m) +BIG = 2**31 + + +@pytest.mark.parametrize( + "n_col, right_rows, expected", + [ + ( + 1000, + [{900: 1.0, 3: 2.0, 500: 3.0}, {3: 4.0, 999: 5.0}], + [ + {3: 10.0, 500: 3.0, 900: 1.0, 999: 10.0}, + {3: 12.0, 999: 15.0}, + {3: 8.0, 500: 12.0, 900: 4.0}, + ], + ), + (1000, [{}, {}], [{}, {}, {}]), + ( + BIG + 10, + [{BIG + 5: 1.0, 7: 2.0}, {BIG: 3.0}], + [{7: 2.0, BIG: 6.0, BIG + 5: 1.0}, {BIG: 9.0}, {7: 8.0, BIG + 5: 4.0}], + ), + ], + ids=["wide", "empty", "beyond-int32"], +) +def test_column_compacted_matmul( + n_col: int, right_rows: list[dict[int, float]], expected: list[dict[int, float]] +) -> None: + right = scipy.sparse.csr_array( + ( + [v for row in right_rows for v in row.values()], + np.array([c for row in right_rows for c in row], dtype=np.int64), + np.cumsum([0, *map(len, right_rows)]), + ), + shape=(2, n_col), + ) + left = scipy.sparse.csr_array([[1.0, 2.0], [0.0, 3.0], [4.0, 0.0]]) + + res = column_compacted_matmul(left, right) + + assert res.shape == (3, n_col) + rows = [ + dict(zip(res.indices[a:b].tolist(), res.data[a:b].tolist())) + for a, b in zip(res.indptr[:-1], res.indptr[1:]) + ] + assert [list(r.items()) for r in rows] == [sorted(r.items()) for r in expected] + + @pytest.mark.parametrize("n_snap", [3, 70], ids=["one-chunk", "chunk-boundary"]) @pytest.mark.parametrize("zeros", [False, True], ids=["dense-C", "sparse-C"]) def test_contracted_partial_matches_dense(n_snap: int, zeros: bool) -> None: