perf(sparse): merge CSR expressions across differing grids without densifying

doc/release_notes.rst

@@ -73,6 +73,8 @@ Version 0.9.1
 
 * ``LinearExpression.reindex`` keeps a sparse (CSR-backed) expression sparse for plain label changes — reorder, add, or drop coordinates — instead of expanding to the dense rectangle. New coordinates become absent cells and the result matches the dense reindex, so a ``groupby(sparse=True) → reindex → merge`` chain stays sparse and the build peak stays low (v1 only; other arguments fall back to dense). (`#932 <https://github.com/PyPSA/linopy/issues/932>`__)
 
+* ``linopy.merge`` (and ``+`` / ``-`` / ``.add`` / ``.sub`` with an explicit ``join=``) keeps sparse (CSR-backed) expressions sparse when the operands live on different label subsets of the same dimensions, e.g. a nodal balance summing grouped generator, line and load terms. The payloads are aligned row-wise onto the joined grid (``outer`` / ``inner`` / ``left`` / ``right`` / ``override``) instead of falling back to the dense rectangle; the positions the join creates carry the same fill as the dense path (zero, or absent with ``fill_value=linopy.ABSENT``) and the result equals the dense one. Dense operands on a different grid are converted on the fly. (`#749 <https://github.com/PyPSA/linopy/issues/749>`__)
+
 **Bug fixes**
 
 * An SOS set is now ordered by declaration, not by the values of its coordinates. Labels are names, but they were handed to the solver as SOS weights, so element-for-element identical models could reach different optima — silently changing who is adjacent in a ``sos_type=2`` set — just because their coordinates were named differently. **Behaviour change:** ascending coordinates are unaffected (this covers every piecewise formulation), descending ones now run in reverse with the same adjacency, and only a ``sos_type=2`` set whose numeric coordinates neither ascend nor descend changes meaning — that case now warns, and sorting the index restores the old order. (`#892 <https://github.com/PyPSA/linopy/issues/892>`__, `#893 <https://github.com/PyPSA/linopy/pull/893>`__)

linopy/expressions.py

@@ -3272,7 +3272,9 @@ def merge(
     if issubclass(cls, LinearExpression) and not has_quad_expression:
         from linopy.sparse_expression import try_csr_merge
 
-        csr_result = try_csr_merge(exprs, dim=dim, join=join, kwargs=kwargs)
+        csr_result = try_csr_merge(
+            exprs, dim=dim, join=join, fill_value=fill_value, kwargs=kwargs
+        )
         if csr_result is not None:
             return csr_result
 

linopy/sparse_expression.py

@@ -26,8 +26,10 @@
 import pandas as pd
 import scipy.sparse
 from xarray import Dataset
+from xarray.core.types import JoinOptions
 
 from linopy.constants import TERM_DIM
+from linopy.semantics import FillValueLike, join_fill
 
 if TYPE_CHECKING:
     from linopy.expressions import LinearExpression
@@ -95,20 +97,13 @@ def from_grouper(
         )
 
     @classmethod
-    def from_expression(
-        cls, expr: LinearExpression, template: CSRPayload
-    ) -> CSRPayload | None:
-        """Convert a dense expression on the template's grid, else None."""
-        if set(expr.coord_dims) != set(template.grid_dims):
-            return None
-        for d in expr.coord_dims:
-            if not expr.data.get_index(d).equals(template.indexes[str(d)]):
-                return None
-        first = template.grid_dims[0]
-        codes = np.arange(len(template.indexes[first]))
-        return cls._from_scatter(
-            expr, template.grid_dims, template.indexes, first, first, codes, False
-        )
+    def from_expression(cls, expr: LinearExpression) -> CSRPayload:
+        """Convert a dense expression to CSR form on its own coordinate grid."""
+        grid_dims = tuple(str(d) for d in expr.coord_dims)
+        indexes = {d: expr.data.get_index(d).rename(d) for d in grid_dims}
+        first = grid_dims[0]
+        codes = np.arange(len(indexes[first]))
+        return cls._from_scatter(expr, grid_dims, indexes, first, first, codes, False)
 
     @classmethod
     def _from_scatter(
@@ -163,15 +158,25 @@ def _from_scatter(
     def scaled(self, factor: float) -> CSRPayload:
         return replace(self, csr=self.csr * factor, const=self.const * factor)
 
-    def reindexed(self, indexes: dict[str, pd.Index]) -> CSRPayload:
+    def reindexed(
+        self,
+        indexes: dict[str, pd.Index],
+        grid_dims: tuple[str, ...] | None = None,
+        fill: float = np.nan,
+    ) -> CSRPayload:
         """
-        Remap rows onto new per-dim indexes without the dense rectangle:
-        dropped labels vanish, new labels are absent cells (NaN const).
+        Remap rows onto new per-dim indexes, optionally in a new dim order,
+        without the dense rectangle: dropped labels vanish, new labels get
+        ``fill`` as their constant (NaN: absent cells).
         """
-        shape, strides = _grid_layout(self.grid_dims, indexes)
+        grid_dims = grid_dims or self.grid_dims
+        shape, strides = _grid_layout(grid_dims, indexes)
+        stride_of = dict(zip(grid_dims, strides))
         positions = [indexes[d].get_indexer(self.indexes[d]) for d in self.grid_dims]
         valid = _flat_cells([pos == -1 for pos in positions]) == 0
-        row_map = _flat_cells([pos * s for pos, s in zip(positions, strides)])
+        row_map = _flat_cells(
+            [pos * stride_of[d] for pos, d in zip(positions, self.grid_dims)]
+        )
 
         coo = self.csr.tocoo()
         keep = valid[coo.coords[0]]
@@ -181,10 +186,14 @@ def reindexed(self, indexes: dict[str, pd.Index]) -> CSRPayload:
         coo = scipy.sparse.coo_array(
             (coo.data[keep], (rows, cols)), shape=(n_cells, self.csr.shape[1])
         )
-        const = np.full(n_cells, np.nan)
+        const = np.full(n_cells, fill)
         const[row_map[valid]] = self.const[valid]
         return replace(
-            self, csr=scipy.sparse.csr_array(coo), const=const, indexes=indexes
+            self,
+            csr=scipy.sparse.csr_array(coo),
+            const=const,
+            grid_dims=grid_dims,
+            indexes=indexes,
         )
 
     def filled(self, value: float) -> CSRPayload:
@@ -211,17 +220,20 @@ def add(self, other: CSRPayload) -> CSRPayload:
         Sparse matrix addition == merge along the term dimension. Goes through
         COO so explicit zero coefficients survive (scipy's ``+`` drops them),
         keeping a cell with only zero-coefficient terms distinguishable from
-        an empty cell, as on the dense path.
+        an empty cell, as on the dense path. A cell absent in either operand
+        is absent in the sum and carries no terms (v1 dead-term invariant).
         """
+        const = self.const + other.const
         a, b = self.csr.tocoo(), other.csr.tocoo()
         shape = (self.n_cells, max(a.shape[1], b.shape[1]))
         rows = np.concatenate([a.coords[0], b.coords[0]])
         cols = np.concatenate([a.coords[1], b.coords[1]])
         data = np.concatenate([a.data, b.data])
-        coo = scipy.sparse.coo_array((data, (rows, cols)), shape=shape)
-        return replace(
-            self, csr=scipy.sparse.csr_array(coo), const=self.const + other.const
+        present = ~np.isnan(const)[rows]
+        coo = scipy.sparse.coo_array(
+            (data[present], (rows[present], cols[present])), shape=shape
         )
+        return replace(self, csr=scipy.sparse.csr_array(coo), const=const)
 
     def materialize(self) -> LinearExpression:
         """
@@ -275,33 +287,79 @@ def _flat_cells(axis_positions: list[np.ndarray]) -> np.ndarray:
     return cells.reshape(-1)
 
 
+def _aligned(
+    payloads: list[CSRPayload], join: JoinOptions | None, fill: float
+) -> list[CSRPayload] | None:
+    """
+    Conform the payloads to the grid an explicit join produces, cells the
+    join creates carrying ``fill`` as constant. None where the dense path
+    owns the semantics: ``exact`` and the auto-detected join raise there on
+    differing grids, ``override`` on differing shapes, any join on
+    non-unique labels.
+    """
+    template = payloads[0]
+    dims = template.grid_dims
+    if any(not p.indexes[d].is_unique for p in payloads for d in dims):
+        return None
+    if join == "override":
+        if any(p.grid_dims != dims or p.shape != template.shape for p in payloads):
+            return None
+        return [replace(p, indexes=template.indexes) for p in payloads]
+    if join == "left":
+        indexes = template.indexes
+    elif join == "right":
+        indexes = payloads[-1].indexes
+    elif join in ("outer", "inner"):
+        combine = pd.Index.union if join == "outer" else pd.Index.intersection
+        indexes = {}
+        for d in dims:
+            index = template.indexes[d]
+            for p in payloads[1:]:
+                index = combine(index, p.indexes[d])
+            indexes[d] = pd.Index(index, name=d)
+    else:
+        return None
+    return [p.reindexed(indexes, dims, fill) for p in payloads]
+
+
 def try_csr_merge(
-    exprs: Any, dim: str, join: Any, kwargs: dict
+    exprs: Any,
+    dim: str,
+    join: JoinOptions | None,
+    fill_value: FillValueLike,
+    kwargs: dict[str, Any],
 ) -> LinearExpression | None:
     """
     Sparse branch of :func:`linopy.expressions.merge`: combine plain
-    LinearExpressions on one shared grid (CSR-backed or dense-convertible),
-    where any join produces the identical result. Returns None to fall
-    through to the dense path.
+    LinearExpressions over one set of grid dimensions (CSR-backed or
+    dense-convertible) as sparse matrix addition. Grids that differ in
+    their labels are aligned row-wise onto the joined grid, the cells the
+    join creates carrying the fill of the dense path (zero, or NaN for
+    ``fill_value=ABSENT``). Returns None to fall through to the dense path.
     """
     from linopy.expressions import LinearExpression
 
     if dim != TERM_DIM or kwargs:
         return None
     if not all(type(e) is LinearExpression for e in exprs):
         return None
-    payloads = [e._payload for e in exprs if e._payload is not None]
-    if not payloads:
+    if all(e._payload is None for e in exprs):
         return None
-    template = payloads[0]
-    if not all(template.same_grid(p) for p in payloads[1:]):
+    dims = set(exprs[0].coord_dims)
+    if any(set(e.coord_dims) != dims for e in exprs[1:]):
+        return None
+    if any(e._payload is None and set(e.data.coords) - dims for e in exprs):
         return None
 
-    combined: CSRPayload | None = None
-    for e in exprs:
-        payload = e._payload or CSRPayload.from_expression(e, template)
-        if payload is None:
+    payloads = [e._payload or CSRPayload.from_expression(e) for e in exprs]
+    template = payloads[0]
+    if not all(template.same_grid(p) for p in payloads[1:]):
+        aligned = _aligned(payloads, join, join_fill(fill_value, 0.0))
+        if aligned is None:
             return None
-        combined = payload if combined is None else combined.add(payload)
-    assert combined is not None
+        payloads = aligned
+
+    combined = payloads[0]
+    for payload in payloads[1:]:
+        combined = combined.add(payload)
     return LinearExpression._from_payload(combined, exprs[0].model)