perf: sparse-aware dot/matmul against a mostly-zero constant

doc/release_notes.rst

@@ -25,6 +25,7 @@ Upcoming Version
 
 **Performance**
 
+* ``LinearExpression.__matmul__``/``dot`` against a mostly-zero constant operand (e.g. a cycle-incidence matrix) no longer materializes one term per element of the contracted dimension. A sparse-aware kernel pairs each nonzero entry of the operand with its slice of the expression and group-sums the entries, so the result only holds contributing terms — on a Kirchhoff-voltage-law shaped contraction this shrinks the result by orders of magnitude. The represented expression is unchanged; only the term layout (count, order, padding) differs. The kernel engages when the operand's nonzero density is at most ``linopy.options["sparse_dot_max_density"]`` (default ``0.5``, ``0`` disables it) and the contracted labels match the expression exactly; all other cases keep the dense path unchanged. (`#748 <https://github.com/PyPSA/linopy/issues/748>`__)
 * ``LinearExpression.groupby(...).sum()`` now scatters terms directly into the padded result arrays via ``xarray.apply_ufunc``, avoiding intermediate copies and speeding up the grouping. A single kernel covers both numpy and chunked (dask) data, the latter staying lazy. On representative models this lowers build and export peak memory by up to ~3x. The kernel emits the grouped result in its final axis order in one contiguous allocation; on dask inputs the reduction now runs over a single chunk (it no longer parallelises over the surviving dimensions). (`#802 <https://github.com/PyPSA/linopy/pull/802>`__)
 
 **Deprecations**

linopy/config.py

@@ -62,4 +62,5 @@ def __repr__(self) -> str:
 options = OptionSettings(
     display_max_rows=14,
     display_max_terms=6,
+    sparse_dot_max_density=0.5,
 )

linopy/expressions.py

@@ -101,6 +101,8 @@
 
 FILL_VALUE = {"vars": -1, "coeffs": np.nan, "const": np.nan}
 
+ENTRY_DIM = "_entry"
+
 
 def exprwrap(
     method: Callable, *default_args: Any, **new_default_kwargs: Any
@@ -2057,13 +2059,114 @@ def __matmul__(
     ) -> LinearExpression | QuadraticExpression:
         """
         Matrix multiplication with other, similar to xarray dot.
+
+        A constant operand that is mostly zeros (e.g. an incidence matrix) is
+        contracted by :meth:`_matmul_sparse_constant`, which only materializes
+        terms for its nonzero entries.
         """
         if not isinstance(other, LinearExpression | variables.Variable):
             other = as_dataarray(other, coords=self.coords, dims=self.coord_dims)
+            res = self._matmul_sparse_constant(other)
+            if res is not None:
+                return res
 
         common_dims = list(set(self.coord_dims).intersection(other.dims))
         return (self * other).sum(dim=common_dims)
 
+    def _matmul_sparse_constant(self, other: DataArray) -> LinearExpression | None:
+        """
+        Sparse-aware kernel for ``self @ other`` with a mostly-zero constant.
+
+        The dense path ``(self * other).sum(common_dims)`` stacks one term per
+        element of the contracted dimensions regardless of the values in
+        ``other``, so a sparse operand densifies the result. This kernel
+        instead pairs each nonzero entry of ``other`` with its slice of the
+        expression and group-sums the entries into the result, so only
+        contributing terms are materialized. The represented expression is
+        identical to the dense path's; only the positional term layout (term
+        count, order and padding) differs.
+
+        Returns None whenever the shortcut does not apply and the caller must
+        take the dense path. The kernel only engages when the contracted
+        labels already match the expression exactly, so the dense path's
+        alignment semantics never need to be reproduced here. The operand
+        density at which the kernel engages is set by the
+        ``sparse_dot_max_density`` option (0 disables the sparse path).
+        """
+        common_dims = [d for d in self.coord_dims if d in other.dims]
+        free_dims = [d for d in other.dims if d not in common_dims]
+        if not common_dims or len(free_dims) > 1:
+            return None
+        free_dim = free_dims[0] if free_dims else None
+        if free_dim in HELPER_DIMS or ENTRY_DIM in (*self.data.dims, *other.dims):
+            return None
+        if self.data.chunks or other.chunks is not None:
+            return None
+        if not np.issubdtype(other.dtype, np.number):
+            return None
+        for dim in common_dims:
+            if other.sizes[dim] != self.data.sizes[dim] or not other.get_index(
+                dim
+            ).equals(self.data.get_index(dim)):
+                return None
+        if free_dim is not None and not other.get_index(free_dim).is_unique:
+            return None
+
+        values = other.transpose(*common_dims, *free_dims).to_numpy()
+        nonzero = values != 0  # keeps NaN entries: NaN coefficients must stay terms
+        nnz = int(nonzero.sum())
+        if nnz == 0 or nnz > options["sparse_dot_max_density"] * values.size:
+            return None
+
+        entries = np.nonzero(nonzero)
+        indexers = {
+            dim: DataArray(positions, dims=ENTRY_DIM)
+            for dim, positions in zip(common_dims, entries)
+        }
+        ds = self.data[["coeffs", "vars"]].isel(indexers)
+        ds = ds.drop_vars(
+            [name for name, coord in ds.coords.items() if ENTRY_DIM in coord.dims]
+        )
+        ds = ds.assign(coeffs=ds.coeffs * DataArray(values[entries], dims=ENTRY_DIM))
+
+        if (self.const == 0).all():
+            const: DataArray | float = 0.0
+        else:
+            # the dense path reduces the constant with xarray's skipna sum, so
+            # NaNs must not propagate through the contraction
+            const = xr.dot(self.const.fillna(0.0), other.fillna(0.0), dim=common_dims)
+
+        expr = LinearExpression(ds, self.model)
+        if free_dim is None:
+            res = expr.sum(ENTRY_DIM).data
+        else:
+            free_index = other.get_index(free_dim)
+            grouper = pd.Series(
+                np.asarray(free_index)[entries[-1]],
+                index=pd.RangeIndex(nnz, name=ENTRY_DIM),
+                name=free_dim,
+            )
+            res = expr.groupby(grouper).sum().data
+            res = res.reindex({free_dim: free_index}, fill_value=self._fill_value)
+
+        allowed_dims = set() if free_dim is None else {free_dim}
+        carried = {
+            name: coord
+            for name, coord in other.coords.items()
+            if name not in res.coords and set(coord.dims) <= allowed_dims
+        }
+        if carried:
+            res = res.assign_coords(carried)
+
+        res = assign_multiindex_safe(res, const=const)
+        order = [
+            d for d in self.data.coeffs.dims if d not in common_dims and d != TERM_DIM
+        ]
+        if free_dim is not None:
+            order.append(free_dim)
+        res = res.transpose(*order, TERM_DIM)
+        return LinearExpression(res, self.model)
+
     @property
     def flat(self) -> pd.DataFrame:
         """

test/test_linear_expression.py

@@ -27,6 +27,7 @@
     Variable,
     merge,
 )
+from linopy.config import options
 from linopy.constants import HELPER_DIMS, TERM_DIM
 from linopy.expressions import ScalarLinearExpression
 from linopy.testing import assert_linequal, assert_quadequal
@@ -647,6 +648,177 @@ def test_matmul_wrong_input(x: Variable, y: Variable, z: Variable) -> None:
         expr @ expr
 
 
+def canonical_terms(expr: LinearExpression, dims: list[str]) -> pd.Series:
+    """
+    Cell-wise variable/coefficient map of an expression, ignoring term layout.
+
+    Absent terms and zero coefficients carry no information, so two
+    expressions with equal canonical terms represent the same linear function
+    regardless of term count, order or padding.
+    """
+    ds = xr.Dataset({"vars": expr.vars, "coeffs": expr.coeffs})
+    df = ds.to_dataframe().reset_index()
+    df = df[(df["vars"] != -1) & (df["coeffs"] != 0) & df["coeffs"].notna()]
+    return df.groupby(dims + ["vars"])["coeffs"].sum().sort_index()
+
+
+def assert_same_linear_function(a: LinearExpression, b: LinearExpression) -> None:
+    """Dimension order is not semantic, so it is deliberately not compared."""
+    assert set(a.coord_dims) == set(b.coord_dims)
+    dims = sorted(map(str, a.coord_dims))
+    pd.testing.assert_series_equal(canonical_terms(a, dims), canonical_terms(b, dims))
+    xr.testing.assert_allclose(a.const.transpose(*dims), b.const.transpose(*dims))
+
+
+@pytest.fixture
+def flow_and_incidence() -> tuple[LinearExpression, xr.DataArray]:
+    m = Model()
+    snapshots = pd.Index(range(3), name="snapshot")
+    branches = pd.Index([f"b{i}" for i in range(8)], name="branch")
+    cycles = pd.Index(list("dbac"), name="cycle")  # deliberately unsorted
+    flow = m.add_variables(coords=[snapshots, branches], name="flow")
+
+    C = xr.DataArray(
+        np.zeros((8, 4)),
+        coords={"branch": branches, "cycle": cycles},
+        dims=["branch", "cycle"],
+    )
+    C.loc["b0", "d"] = 1.0
+    C.loc["b1", "d"] = -1.0
+    C.loc["b2", "b"] = 1.0
+    C.loc["b3", "b"] = 1.0
+    C.loc["b0", "b"] = -1.0  # b0 belongs to two cycles
+    C.loc["b4", "a"] = 2.5
+    # cycle "c" has no branches
+    return 1 * flow, C
+
+
+def test_matmul_sparse_const_incidence(
+    flow_and_incidence: tuple[LinearExpression, xr.DataArray],
+) -> None:
+    """A sparse incidence matrix contracts to observed terms only (#748)."""
+    expr, C = flow_and_incidence
+    res = expr @ C
+    dense = (expr * C).sum("branch")
+
+    assert_same_linear_function(res, dense)
+    assert res.nterm == 3  # largest cycle, not the branch count
+    assert dense.nterm == len(expr.indexes["branch"])
+    assert res.data.coeffs.dims == dense.data.coeffs.dims
+    assert res.indexes["cycle"].equals(C.indexes["cycle"])
+
+
+def test_matmul_sparse_const_carries_free_dim_aux_coords(
+    flow_and_incidence: tuple[LinearExpression, xr.DataArray],
+) -> None:
+    expr, C = flow_and_incidence
+    C = C.assign_coords(voltage=("cycle", [380, 220, 380, 110]))
+    res = expr @ C
+    dense = (expr * C).sum("branch")
+    assert_same_linear_function(res, dense)
+    assert_equal(res.data.voltage, dense.data.voltage)
+
+
+def test_matmul_sparse_const_vector(
+    flow_and_incidence: tuple[LinearExpression, xr.DataArray],
+) -> None:
+    """Contraction without a free dimension keeps only nonzero entries."""
+    expr, C = flow_and_incidence
+    w = xr.DataArray(
+        np.array([0.0, 2.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0]),
+        coords=[expr.indexes["branch"]],
+    )
+    res = expr @ w
+    dense = (expr * w).sum("branch")
+    assert_same_linear_function(res, dense)
+    assert res.nterm == 2
+
+
+def test_matmul_sparse_const_with_expression_const(
+    flow_and_incidence: tuple[LinearExpression, xr.DataArray],
+) -> None:
+    expr, C = flow_and_incidence
+    res = (expr + 5.0) @ C
+    dense = ((expr + 5.0) * C).sum("branch")
+    assert_same_linear_function(res, dense)
+
+
+def test_matmul_sparse_const_keeps_nan_entries(
+    flow_and_incidence: tuple[LinearExpression, xr.DataArray],
+) -> None:
+    """A NaN entry in the operand must stay a term, like on the dense path."""
+    expr, C = flow_and_incidence
+    C.loc["b5", "a"] = np.nan
+    res = expr @ C
+    cell = res.data.sel(cycle="a", snapshot=0)
+    nan_terms = cell.vars.values[np.isnan(cell.coeffs.values)]
+    assert expr.vars.sel(branch="b5", snapshot=0).item() in nan_terms
+
+
+def test_matmul_sparse_const_falls_back_on_reordered_labels(
+    flow_and_incidence: tuple[LinearExpression, xr.DataArray],
+) -> None:
+    """Mismatching contracted labels take the dense path unchanged."""
+    expr, C = flow_and_incidence
+    C_reordered = C.reindex(branch=C.indexes["branch"][::-1])
+    res = expr @ C_reordered
+    assert_linequal(res, (expr * C_reordered).sum("branch"))
+    assert res.nterm == len(expr.indexes["branch"])
+
+
+def test_matmul_sparse_const_falls_back_on_dense_operand(
+    flow_and_incidence: tuple[LinearExpression, xr.DataArray],
+) -> None:
+    """A mostly-nonzero operand keeps the dense path's exact layout."""
+    expr, C = flow_and_incidence
+    dense_C = xr.ones_like(C)
+    res = expr @ dense_C
+    assert_linequal(res, (expr * dense_C).sum("branch"))
+    assert res.nterm == len(expr.indexes["branch"])
+
+
+def test_matmul_sparse_const_falls_back_on_duplicate_free_labels(
+    flow_and_incidence: tuple[LinearExpression, xr.DataArray],
+) -> None:
+    expr, C = flow_and_incidence
+    C_dup = C.assign_coords(cycle=["d", "d", "a", "c"])
+    res = expr @ C_dup
+    assert_linequal(res, (expr * C_dup).sum("branch"))
+
+
+def test_matmul_sparse_const_falls_back_on_multiple_free_dims(
+    flow_and_incidence: tuple[LinearExpression, xr.DataArray],
+) -> None:
+    expr, C = flow_and_incidence
+    C_3d = C.expand_dims(scenario=pd.Index([0, 1], name="scenario"))
+    res = expr @ C_3d
+    common = [d for d in expr.coord_dims if d in C_3d.dims]
+    assert_linequal(res, (expr * C_3d).sum(common))
+
+
+def test_matmul_sparse_const_disabled_by_option(
+    flow_and_incidence: tuple[LinearExpression, xr.DataArray],
+) -> None:
+    expr, C = flow_and_incidence
+    with options:
+        options(sparse_dot_max_density=0.0)
+        res = expr @ C
+    assert_linequal(res, (expr * C).sum("branch"))
+    assert res.nterm == len(expr.indexes["branch"])
+
+
+def test_matmul_sparse_const_free_dim_without_coords(
+    flow_and_incidence: tuple[LinearExpression, xr.DataArray],
+) -> None:
+    expr, C = flow_and_incidence
+    C_nocoords = C.drop_vars("cycle")
+    res = expr @ C_nocoords
+    dense = (expr * C_nocoords).sum("branch")
+    # the constructor fills integer coords for coordinate-less dims either way
+    assert res.indexes["cycle"].equals(dense.indexes["cycle"])
+    assert_same_linear_function(res, dense)
+
+
 def test_linear_expression_multiplication_invalid(
     x: Variable, y: Variable, z: Variable
 ) -> None: