perf(expressions): sparse @/dot drops zero-coeff terms; vectorise densify_terms (#748)

doc/release_notes.rst

@@ -37,8 +37,14 @@ Upcoming Version
 
 * The example notebooks now opt into the v1 arithmetic convention (``linopy.options["semantics"] = "v1"``). The coordinate-alignment and expression tutorials were reworked to teach strict label-based alignment: a mismatch on a shared dimension raises rather than silently filling or pairing by position, and is resolved explicitly with ``.sel`` / ``.reindex`` / ``.assign_coords`` or an explicit ``join=`` on the named ``.add`` / ``.mul`` / ``.le`` / … methods.
 
+**Performance**
+
+* ``@``/``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.
+
 **Bug fixes**
 
+* ``densify_terms`` no longer raises on expressions without coordinate dimensions (``expr.sum(drop_zeros=True)`` over all dimensions) and now works on ``QuadraticExpression``, where it previously indexed the ``_factor`` axis as the term axis.
 * ``sum()`` over a dimension no longer raises when another dimension of the expression has size 0; it returns an expression without terms over the kept coordinates, as summing over the empty dimension itself already did. (https://github.com/PyPSA/linopy/issues/906)
 * A multi-key ``groupby`` now returns its groups sorted by key tuple, like the single-key path. The key combinations were numbered by iterating a ``set``, so the group order was arbitrary and changed between processes with ``PYTHONHASHSEED``.
 * The ``linopy.options`` context manager now restores the values that were active on entry instead of resetting all options to their defaults.

linopy/expressions.py

@@ -2158,31 +2158,35 @@ def densify_terms(self) -> Self:
         Move all non-zero term entries to the front and cut off all-zero
         entries in the term-axis.
         """
-        data = self.data.transpose(..., TERM_DIM)
-
-        cdata = data.coeffs.data
-        axis = cdata.ndim - 1
-        nnz = np.nonzero(cdata)
-        nterm = (cdata != 0).sum(axis).max()
-
-        mod_nnz = list(nnz)
-        mod_nnz.pop(axis)
-
-        remaining_axes = np.vstack(mod_nnz).T
-        _, idx_ = np.unique(remaining_axes, axis=0, return_inverse=True)
-        idx = list(idx_)
-        new_index = np.array([idx[:i].count(j) for i, j in enumerate(idx)])
-        mod_nnz.insert(axis, new_index)
-
-        vdata = np.full_like(cdata, -1)
-        vdata[tuple(mod_nnz)] = data.vars.data[nnz]
-        data.vars.data = vdata
-
-        cdata = np.zeros_like(cdata)
-        cdata[tuple(mod_nnz)] = data.coeffs.data[nnz]
-        data.coeffs.data = cdata
+        coeffs = self.data.coeffs.transpose(..., TERM_DIM)
+        vars = self.data.vars.transpose(*coeffs.dims, ...)
+        cdata = coeffs.data
+        mask = cdata != 0
+        if mask.all():
+            return self
 
-        return self.__class__(data.sel({TERM_DIM: slice(0, nterm)}), self.model)
+        lead = cdata.shape[:-1]
+        old_nterm = cdata.shape[-1]
+        trailing = vars.shape[cdata.ndim :]
+        mask = mask.reshape(-1, old_nterm)
+        nrows = mask.shape[0]
+        counts = mask.sum(1)
+        nterm = int(counts.max()) if nrows else 0
+        rows = np.repeat(np.arange(nrows), counts)
+        pos = np.arange(rows.size) - np.repeat(np.cumsum(counts) - counts, counts)
+
+        new_coeffs = np.zeros((nrows, nterm), dtype=cdata.dtype)
+        new_coeffs[rows, pos] = cdata.reshape(nrows, old_nterm)[mask]
+        new_vars = np.full((nrows, nterm, *trailing), -1, dtype=vars.dtype)
+        new_vars[rows, pos] = vars.data.reshape(nrows, old_nterm, *trailing)[mask]
+
+        new: dict[Hashable, Any] = {
+            "coeffs": (coeffs.dims, new_coeffs.reshape(*lead, nterm)),
+            "vars": (vars.dims, new_vars.reshape(*lead, nterm, *trailing)),
+        }
+        data_vars = {k: new.get(k, self.data[k]) for k in self.data.data_vars}
+        data = Dataset(data_vars, coords=self.data.coords, attrs=self.data.attrs)
+        return self.__class__(data, self.model)
 
     def sanitize(self) -> Self:
         """
@@ -2469,11 +2473,15 @@ def __matmul__(
         Matrix multiplication with other, similar to xarray dot.
         """
         other = as_constant(other)
-        if not isinstance(other, LinearExpression | variables.Variable):
+        other_is_const = not isinstance(other, LinearExpression | variables.Variable)
+        if other_is_const:
             other = _matmul_operand_to_dataarray(other, self.coords, self.coord_dims)
 
         common_dims = list(set(self.coord_dims).intersection(other.dims))
-        return (self * other).sum(dim=common_dims)
+        res = (self * other).sum(dim=common_dims)
+        if other_is_const and common_dims and bool((other == 0).any()):
+            res = res.densify_terms()
+        return res
 
     @property
     def flat(self) -> pd.DataFrame:

test/test_linear_expression.py

@@ -525,6 +525,55 @@ def test_linear_expression_sum_drop_zeros(z: Variable) -> None:
     assert res.nterm == 2
 
 
+def test_densify_terms_compacts_rows_in_order(
+    x: Variable, y: Variable, z: Variable
+) -> None:
+    mask = xr.DataArray([[1.0, 0, 1], [0, 1, 0]], dims=["dim_0", "dim_1"])
+    expr = x + y * mask + z
+    res = expr.densify_terms()
+
+    assert res.nterm == 3
+    assert res.vars.dtype == x.labels.dtype
+    assert (res.vars.isel(_term=0) == x.labels).all()
+    assert (res.vars.isel(_term=1) == z.labels.where(mask == 0, y.labels)).all()
+    assert (res.vars.isel(_term=2) == z.labels.where(mask == 1, -1)).all()
+    assert (res.coeffs.isel(_term=[0, 1]) == 1).all()
+    assert (res.coeffs.isel(_term=2) == mask).all()
+    pd.testing.assert_frame_equal(res.flat, expr.flat)
+
+
+def test_densify_terms_no_zeros_is_noop(x: Variable, y: Variable) -> None:
+    expr = x + y
+    assert expr.densify_terms() is expr
+
+
+def test_densify_terms_scalar_and_all_zero(x: Variable) -> None:
+    res = (0 * x + x).sum(drop_zeros=True)
+    assert res.nterm == 2
+    assert_linequal(res, x.sum())
+
+    res = (0 * x).sum(drop_zeros=True)
+    assert res.nterm == 0
+    assert res.const.item() == 0
+
+
+def test_densify_terms_quadratic(x: Variable, y: Variable) -> None:
+    mask = xr.DataArray([1.0, 0], dims=["dim_0"])
+    expr = x * y + x * x * mask + y * y
+
+    res = expr.densify_terms()
+    assert res.nterm == 3
+    assert res.vars.dims == expr.vars.dims
+    assert res.sel(dim_0=1).nterm == 3
+    assert (res.sel(dim_0=1).coeffs.isel(_term=2) == 0).all()
+    assert (res.sel(dim_0=1).vars.isel(_term=2) == -1).all()
+    pd.testing.assert_frame_equal(res.flat, expr.flat)
+
+    res = expr.sum(drop_zeros=True)
+    assert res.nterm == 5
+    pd.testing.assert_frame_equal(res.flat, expr.sum().flat)
+
+
 class TestCollapseAuxCoords:
     # collapsing a dimension carrying an auxiliary coordinate must drop it, not
     # leak it onto the term dimension where it breaks later arithmetic with a
@@ -679,6 +728,67 @@ def test_matmul_contracts_all_dims_when_const_covers_them(z: Variable) -> None:
     assert_linequal(res, (expr * b).sum(["dim_0", "dim_1"]))
 
 
+def test_matmul_sparse_operand_drops_zero_terms(z: Variable) -> None:
+    """
+    ``@`` against a zero-containing constant compacts the term dimension to
+    the widest non-zero cell instead of one term per contracted member (#748).
+    """
+    expr = 1 * z  # dims (dim_0, dim_1); contracting dim_1 (size 3)
+    b = xr.DataArray(
+        [[1.0, 0.0], [2.0, 3.0], [0.0, 0.0]],
+        coords={"dim_1": expr.indexes["dim_1"], "location": ["L1", "L2"]},
+        dims=["dim_1", "location"],
+    )
+
+    res = expr @ b
+
+    assert res.nterm == 2 < b.sizes["dim_1"]
+    assert_linequal(res, (expr * b).sum("dim_1").densify_terms())
+
+
+def test_matmul_dense_operand_keeps_all_terms(z: Variable) -> None:
+    """A constant without zeros contracts to the full term dimension, untouched."""
+    expr = 1 * z
+    b = xr.DataArray(
+        np.arange(1, 7).reshape(3, 2),
+        coords={"dim_1": expr.indexes["dim_1"], "location": ["L1", "L2"]},
+        dims=["dim_1", "location"],
+    )
+
+    res = expr @ b
+
+    assert res.nterm == b.sizes["dim_1"]
+    assert_linequal(res, (expr * b).sum("dim_1"))
+
+
+def test_matmul_all_zero_operand_yields_no_terms(z: Variable) -> None:
+    """An all-zero constant contracts to the zero expression, not a crash (#748)."""
+    expr = 1 * z
+    b = xr.DataArray(
+        np.zeros((3, 2)),
+        coords={"dim_1": expr.indexes["dim_1"], "location": ["L1", "L2"]},
+        dims=["dim_1", "location"],
+    )
+
+    res = expr @ b
+
+    assert res.nterm == 0
+    assert (res.data.vars == -1).all()
+
+
+def test_matmul_full_contraction_with_zero_operand(x: Variable) -> None:
+    """
+    ``variable @ vector`` contracting away every coord dim compacts a
+    zero-containing operand without crashing on the term-only shape (#748).
+    """
+    b = xr.DataArray([2.0, 0.0], coords={"dim_0": x.indexes["dim_0"]})
+
+    res = x @ b
+
+    assert res.nterm == 1
+    assert_linequal(res, (x * b).sum("dim_0").densify_terms())
+
+
 def test_matmul_wrong_input(x: Variable, y: Variable, z: Variable) -> None:
     expr = 10 * x + y + z
     with pytest.raises(TypeError):