From 50fc4f412142dbdaaf7b332cac9b6c49a6152c1d Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Fri, 24 Jul 2026 11:59:28 +0200 Subject: [PATCH] perf(matmul): sparse-aware dot against a mostly-zero constant (#748) (self * other).sum(common_dims) stacks one term per element of the contracted dimensions regardless of zeros in the operand, so a sparse matrix (e.g. a cycle incidence) densifies the result. Pair each nonzero entry of the operand with its slice of the expression instead and group-sum the entries, so only contributing terms are materialized. The kernel engages when the operand's nonzero density is at most linopy.options['sparse_dot_max_density'] (default 0.5, 0 disables) and the contracted labels already match the expression exactly; every other case falls back to the dense path unchanged, so no alignment semantics are re-implemented. KVL-shaped contraction (852 branches, 268 cycles, 24 snapshots, 3 branches/cycle): result term dimension 852 -> 3, coeff+var cells 10,960,128 -> 38,592, memray process peak 282 MB -> 97 MB. Co-Authored-By: Claude Fable 5 --- doc/release_notes.rst | 1 + linopy/config.py | 1 + linopy/expressions.py | 103 ++++++++++++++++++++ test/test_linear_expression.py | 172 +++++++++++++++++++++++++++++++++ 4 files changed, 277 insertions(+) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index cc14abf92..725ab1f67 100644 --- a/doc/release_notes.rst +++ b/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 `__) * ``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 `__) **Deprecations** diff --git a/linopy/config.py b/linopy/config.py index 5d269c4e7..426453d58 100644 --- a/linopy/config.py +++ b/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, ) diff --git a/linopy/expressions.py b/linopy/expressions.py index 21a4160e9..c0e065fb1 100644 --- a/linopy/expressions.py +++ b/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: """ diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index 54b76dd42..691bb05ed 100644 --- a/test/test_linear_expression.py +++ b/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: