feat(expressions): sparse @/dot for CSR-backed expressions

doc/release_notes.rst

@@ -53,6 +53,7 @@ Upcoming Version
 
 * ``@``/``dot`` against a constant matrix that holds zeros no longer densifies the result to one term per contracted member. The zero-coefficient terms are dropped, so the term dimension shrinks to the widest non-zero cell. On PyPSA's Kirchhoff Voltage Law constraint (a cycle matrix with ~3 branches per cycle) this cuts the expression from 852 to 3 terms — 284x fewer cells — which in turn shrinks the downstream ``merge``. A constant without zeros is unaffected. (`#748 <https://github.com/PyPSA/linopy/issues/748>`__)
 * ``densify_terms`` (used by ``sum(drop_zeros=True)`` and the sparse ``@`` path) is now fully vectorised. It previously counted the non-zero positions with a Python loop that scaled quadratically in the number of non-zero terms — 127 s for a (2000 x 60) expression, now 3 ms — and allocated the compacted output at the full original term width. It now allocates only the compacted width and returns the expression unchanged when it holds no zeros.
+* Under v1, ``@``/``dot`` against a constant now runs as sparse linear algebra instead of building the dense broadcast intermediate (``self * other`` then ``.sum()``): peak memory scales with ``nnz(C) x nterm`` rather than with the full broadcast shape. A ``CSRLinearExpression``-backed operand (from ``groupby(...).sum(sparse=True)``) stays CSR-backed through ``@``, and the result is the compact canonical form (duplicate variables summed, terms label-ordered, explicit zeros pruned — cell activeness is carried by ``const`` alone). (`#748 <https://github.com/PyPSA/linopy/issues/748>`__, `#756 <https://github.com/PyPSA/linopy/issues/756>`__, `#925 <https://github.com/PyPSA/linopy/issues/925>`__)
 * Persistent snapshots of tz-aware ``DatetimeIndex`` coordinates no longer materialise an object array of ``Timestamp`` per container per capture and diff. Coordinates are stored as UTC-ns arrays with the timezone identity carried alongside, making snapshot capture ~24x and warm-start diffs ~33x faster on tz-aware models, while naive and tz-aware coordinates — and differing timezones — stay correctly unequal. (`#960 <https://github.com/PyPSA/linopy/pull/960>`__)
 
 **Bug fixes**

linopy/alignment.py

@@ -29,6 +29,7 @@
 import numpy as np
 import pandas as pd
 import polars as pl
+import scipy.sparse
 from numpy import arange
 from xarray import Coordinates, DataArray, Dataset, broadcast
 from xarray import align as xr_align
@@ -43,6 +44,12 @@
     CoordinateValidationError = ValueError  # type: ignore[assignment, misc]
 
 from linopy.constants import HELPER_DIMS
+from linopy.semantics import (
+    _shared_dim_mismatch_message,
+    check_user_nan,
+    enforce_aux_conflict,
+    first_mismatched_dim,
+)
 from linopy.types import UNLABELED_TYPES, CoordsLike, DimsLike
 
 
@@ -684,6 +691,33 @@ def _matmul_operand_to_dataarray(
     return as_dataarray(other, coords=coords, dims=dims)
 
 
+def _matmul_operand_to_matrix(
+    other: DataArray,
+    contracted: Sequence[str],
+    new_dims: Sequence[str],
+    indexes: Mapping[str, pd.Index],
+    aux_coords: Mapping[str, tuple[str, np.ndarray]],
+) -> scipy.sparse.csr_array:
+    """
+    Flatten a ``@`` constant to the ``(contracted, new)`` matrix of the sparse
+    contraction, enforcing the same rules ``*`` enforces on a constant: §5 on
+    NaN, §8 on the labels of the shared (contracted) dims and §11 on auxiliary
+    coordinates. The expression side is represented by its grid labels alone
+    and never broadcasts.
+    """
+    if other.isnull().any():
+        check_user_nan(op_kind="mul")
+    reference = Dataset(coords={d: indexes[d] for d in contracted} | dict(aux_coords))
+    mismatch = first_mismatched_dim(reference, other)
+    if mismatch is not None:
+        raise ValueError(_shared_dim_mismatch_message(*mismatch))
+    enforce_aux_conflict([reference, other])
+    values = other.transpose(*contracted, *new_dims).values
+    n_contracted = int(np.prod([len(indexes[d]) for d in contracted], dtype=np.int64))
+    n_new = int(np.prod([other.sizes[d] for d in new_dims], dtype=np.int64))
+    return scipy.sparse.csr_array(values.reshape(n_contracted, n_new))
+
+
 def _dims_for_positional_input(
     arr: Any, expected: dict[Hashable, Any], dims: DimsLike | None
 ) -> DimsLike | None:

linopy/csr.py

@@ -7,7 +7,10 @@
 public type, different backing, akin to dask-backed xarray objects. The CSR
 form is canonical (duplicate variables summed, terms label-ordered) and ragged
 along ``_term``, so the group-size padding of issue #745 has no analog;
-grouping, ``merge``/``+``/``-`` and scaling become sparse linear algebra.
+grouping, ``merge``/``+``/``-``, scaling and ``@``/``dot`` (:meth:`contracted`)
+become sparse linear algebra. Unlike :meth:`added`, which preserves explicit
+zeros through COO, ``@``/``dot`` prunes them: cell activeness is carried by
+``const`` alone (issue #925).
 Anything without a sparse branch expands through ``.data`` to the
 mathematically identical dense rectangle in canonical term layout — the reason
 the feature is v1-gated, where term layout is non-contractual.
@@ -38,6 +41,9 @@
     from linopy.expressions import LinearExpression
     from linopy.model import Model
 
+CONTRACTION_CHUNK = 64
+"""Kept-axis block size of the chunked Kronecker product in ``contracted``."""
+
 
 @dataclass(frozen=True, eq=False)
 class Grid:
@@ -396,6 +402,73 @@ def added(self, other: CSRLinearExpression) -> CSRLinearExpression:
             self, csr=scipy.sparse.csr_array(coo), const=const, coords=coords
         )
 
+    def contracted(
+        self,
+        matrix: scipy.sparse.csr_array,
+        contracted_dims: Iterable[str],
+        new_indexes: Iterable[pd.Index],
+    ) -> CSRLinearExpression:
+        """
+        Contract grid dimensions against a sparse constant (``expr @ C``).
+
+        ``matrix`` is the constant flattened to
+        ``(prod(contracted shape), prod(new shape))`` in C order over
+        ``contracted_dims``, which are given in grid order. Each entry of
+        ``new_indexes`` must be named after the dim it creates, which is read
+        off its ``name``. The result lives on the kept grid dims followed by
+        ``new_indexes`` and is
+        ``kron(I_kept, matrix.T) @ csr``, evaluated in chunks of the kept axis
+        so the operator never grows with the kept size.
+
+        The result is compact canonical form: duplicate variables summed, terms
+        label-ordered and explicit zeros pruned -- unlike :meth:`added`, the
+        sparse product drops them, so cell activeness is carried by ``const``
+        alone (see issue #925). Auxiliary coordinates on kept dims propagate,
+        those on contracted dims drop.
+        """
+        contracted_dims = tuple(contracted_dims)
+        kept = tuple(d for d in self.grid.dims if d not in contracted_dims)
+        target = kept + contracted_dims
+        source = (
+            self
+            if self.grid.dims == target
+            else self.reindexed(self.grid.reordered(target))
+        )
+
+        n_contracted = matrix.shape[0]
+        kept_grid = source.grid.reordered(kept)
+        n_kept = kept_grid.size
+        const = np.nan_to_num(source.const)
+        chunk = min(CONTRACTION_CHUNK, n_kept)
+        operator = scipy.sparse.kron(
+            scipy.sparse.eye_array(chunk), matrix.T, format="csr"
+        )
+
+        blocks = []
+        const_blocks = []
+        for start in range(0, n_kept, chunk):
+            size = min(chunk, n_kept - start)
+            block = (
+                operator
+                if size == chunk
+                else scipy.sparse.kron(
+                    scipy.sparse.eye_array(size), matrix.T, format="csr"
+                )
+            )
+            rows = slice(start * n_contracted, (start + size) * n_contracted)
+            blocks.append(block @ source.csr[rows])
+            const_blocks.append(block @ const[rows])
+
+        indexes = kept_grid.indexes
+        indexes |= {str(i.name): i for i in new_indexes}
+        return replace(
+            source,
+            csr=scipy.sparse.csr_array(scipy.sparse.vstack(blocks, format="csr")),
+            const=np.concatenate(const_blocks),
+            grid=Grid(indexes),
+            coords={n: (d, v) for n, (d, v) in source.coords.items() if d in kept},
+        )
+
     def to_dense(self) -> LinearExpression:
         """
         Expand to the dense equivalent in canonical form: terms label-ordered,

linopy/expressions.py

@@ -39,7 +39,7 @@
 import scipy
 import xarray as xr
 import xarray.core.groupby
-from numpy import array, nan, ndarray
+from numpy import array, nan
 from pandas.core.frame import DataFrame
 from pandas.core.series import Series
 from scipy.sparse import csc_matrix
@@ -62,6 +62,7 @@
 from linopy import constraints, variables
 from linopy.alignment import (
     _matmul_operand_to_dataarray,
+    _matmul_operand_to_matrix,
     as_constant,
     as_dataarray,
     broadcast_to_coords,
@@ -1523,9 +1524,13 @@ def pow(self, other: int) -> QuadraticExpression:
         """
         return self.__pow__(other)
 
-    def dot(self, other: ndarray) -> Self | QuadraticExpression:
+    def dot(self, other: SideLike) -> Self | QuadraticExpression:
         """
         Matrix multiplication with other, similar to xarray dot.
+
+        Identical to ``@``. For a :class:`LinearExpression` that includes the
+        sparse contraction under v1; :class:`QuadraticExpression` always takes
+        the dense path. There is no per-call ``sparse=`` flag.
         """
         return self.__matmul__(other)
 
@@ -2539,9 +2544,20 @@ def __matmul__(
     ) -> LinearExpression | QuadraticExpression:
         """
         Matrix multiplication with other, similar to xarray dot.
+
+        Under v1, a constant ``other`` is contracted as one sparse matrix
+        product (:meth:`_sparse_matmul`) instead of the dense broadcast
+        ``(self * other).sum(dim)``. The result is then the compact canonical
+        form -- duplicate variables summed, terms label-ordered, explicit
+        zeros pruned -- so its term count may differ from the dense path's
+        while the values agree. A CSR-backed expression stays CSR-backed.
         """
         other = as_constant(other)
         other_is_const = not isinstance(other, LinearExpression | variables.Variable)
+        if other_is_const and is_v1() and type(self) is LinearExpression:
+            sparse = self._sparse_matmul(other)
+            if sparse is not None:
+                return sparse
         if other_is_const:
             other = _matmul_operand_to_dataarray(other, self.coords, self.coord_dims)
 
@@ -2551,6 +2567,47 @@ def __matmul__(
             res = res.densify_terms()
         return res
 
+    def _sparse_matmul(self, other: ConstantLike) -> LinearExpression | None:
+        """
+        Contract against a constant as one sparse matrix product, skipping the
+        dense broadcast intermediate of ``(self * other).sum(dim)``.
+
+        Returns None where the dense path owns the semantics: a MultiIndex on
+        the expression or on the operand, non-unique or MultiIndex grid
+        labels, a zero-size grid, an operand sharing no dimension with the
+        grid, and an unlabelled output dimension. The result is the compact
+        canonical form of :meth:`CSRLinearExpression.contracted`, so its term
+        count may differ from the dense path's while the values agree; it
+        stays CSR-backed when the input was.
+        """
+        if is_nan_scalar(other):
+            check_user_nan(op_kind="mul")
+        if self._csr is None and _has_multiindex(self.data.indexes.values()):
+            return None
+        if not self.coord_dims:
+            return None
+        csr = self._csr or CSRLinearExpression.from_dense(self.data, self.model)
+        if not csr.grid.is_unique or _has_multiindex(csr.grid.indexes.values()):
+            return None
+        if csr.grid.size == 0:
+            return None
+        coords = Dataset(coords=dict(csr.grid.indexes) | dict(csr.coords)).coords
+        da = _matmul_operand_to_dataarray(other, coords, csr.grid.dims)
+        if _has_multiindex(da.indexes.values()):
+            return None
+        contracted = [d for d in csr.grid.dims if d in da.dims]
+        new_dims = [str(d) for d in da.dims if d not in csr.grid.dims]
+        if not contracted or any(d not in da.indexes for d in new_dims):
+            return None
+        matrix = _matmul_operand_to_matrix(
+            da, contracted, new_dims, csr.grid.indexes, csr.coords
+        )
+        new_indexes = [da.indexes[d].rename(d) for d in new_dims]
+        res = csr.contracted(matrix, contracted, new_indexes)
+        if self._csr is not None or options["sparse_groupby"]:
+            return type(self)._from_csr(res, self.model)
+        return res.to_dense()
+
     @property
     def flat(self) -> pd.DataFrame:
         """
@@ -3227,6 +3284,10 @@ def as_expression(
         return LinearExpression(obj, model)
 
 
+def _has_multiindex(indexes: Iterable[pd.Index]) -> bool:
+    return any(isinstance(i, pd.MultiIndex) for i in indexes)
+
+
 def _aligned(
     csrs: list[CSRLinearExpression], join: JoinOptions | None, fill: float
 ) -> list[CSRLinearExpression] | None:

linopy/semantics.py

@@ -462,7 +462,9 @@ def enforce_no_multiindex(
     warn_legacy(_legacy_multiindex_message(str(dim), context), stacklevel=stacklevel)
 
 
-def first_mismatched_dim(a: DataArray, b: DataArray) -> tuple[str, Any, Any] | None:
+def first_mismatched_dim(
+    a: DataArray | Dataset, b: DataArray | Dataset
+) -> tuple[str, Any, Any] | None:
     """
     Return ``(dim, a_labels, b_labels)`` for the first shared dim that
     disagrees on coordinate labels OR size, or ``None`` if all agree.