perf(csr): run each contraction chunk on its used columns only

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 <https://github.com/PyPSA/linopy/issues/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. (`#990 <https://github.com/PyPSA/linopy/pull/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 <https://github.com/PyPSA/linopy/issues/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 <https://github.com/PyPSA/linopy/issues/977>`__)
 

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, 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, 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]),
+    )
+
+
 def coo_to_csr(
     data: np.ndarray,
     rows: np.ndarray,

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: