perf(csr): sparse shift/roll/diff and merge(dim=...); densify notice on by default in sparse models

doc/release_notes.rst

@@ -44,6 +44,10 @@ Upcoming Version
 * ``linopy.merge`` of several sparse expressions on the same coordinates now adds them in one pass instead of one pairwise sum per operand, and a sparse ``reindex`` that only reorders the existing cells gathers the rows directly instead of rebuilding the matrix. Merging three sparse operands is about 30% faster and such a reindex about 4x faster. (`#1010 <https://github.com/PyPSA/linopy/issues/1010>`__)
 * Building ``Model.matrices`` no longer allocates label-sized scaling lookups and skips ``eliminate_zeros`` for frozen constraints. On a model with 2M variables this makes each build of a sparse or frozen model about 20 to 35 ms faster. (`#1008 <https://github.com/PyPSA/linopy/issues/1008>`__)
 * Freezing a dense constraint builds the CSR matrix directly from the term rectangle and sorts it only when a row has more than one term. On a model with 2M variables this cuts the freeze of the nodal balance from about 465 to 185 ms and its peak memory from 1.74 GB to 0.25 GB. The export of a mutable constraint also no longer copies a strided term array whose terms are mostly zero. (`#1009 <https://github.com/PyPSA/linopy/issues/1009>`__)
+* ``shift``, ``roll`` and ``diff`` keep a CSR-backed ``LinearExpression`` sparse: the shift is evaluated on the grid's row numbers and the rows are gathered from the sparse backing, so the shifted-in cells are absent as on the dense path, ``roll`` wraps around (``roll_coords`` included) and the constant and auxiliary coordinates follow. ``shift`` with a ``fill_value`` still densifies. On a ``groupby`` result with uneven group sizes (2000 x 720 cells, 800 terms wide) ``shift`` takes 14 ms and 54 MB instead of 1 s and 4.4 GB. (`#1011 <https://github.com/PyPSA/linopy/issues/1011>`__)
+* ``linopy.merge(..., dim=<grid dimension>)`` keeps CSR-backed operands sparse: the rows are stacked in input order and the coordinates are what the dense path's ``xr.concat`` yields, so overlapping and duplicate labels along ``dim``, the ``join`` of the other dimensions (cells the join creates are absent), the v1 label check of the auto-detected join and auxiliary coordinates behave as on the dense path. The operands are left sparse, too. A merge along a new dimension, with extra ``xr.concat`` arguments or over non-unique labels that need aligning still densifies, with a notice. (`#1011 <https://github.com/PyPSA/linopy/issues/1011>`__)
+* ``linopy.options["warn_on_densify"]`` now defaults to ``None``, which means: on in a sparse model (``Model(sparse=True)``), off otherwise. Every operation that implicitly drops a sparse backing in a sparse model, e.g. ``.data`` or an operation without a sparse path, emits a ``PerformanceWarning`` naming the reason. Explicit conversions (``CSRConstraint.mutable()``/``to_dense()``, ``add_constraints(..., freeze=False)``) only warn when the option is ``True``. Set the option to ``False`` to silence it, or to ``True`` to see it in a dense model and on explicit conversions as well. (`#969 <https://github.com/PyPSA/linopy/issues/969>`__)
+* A merge of operands over different dimensions, e.g. ``bal + s`` with ``s`` over a subset of the dimensions of ``bal``, still densifies and now warns in a sparse model. (`#969 <https://github.com/PyPSA/linopy/issues/969>`__)
 
 *Other*
 
@@ -68,7 +72,7 @@ Upcoming Version
 * ``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>`__)
 * Reading metadata of a CSR-backed ``LinearExpression`` (``repr()``, ``shape``, ``sizes``, ``dims``, ``coords``, ``indexes``, ``isnull()``) no longer converts it to dense; it is served from the sparse backing, auxiliary coordinates included. (`#962 <https://github.com/PyPSA/linopy/issues/962>`__)
-* New ``LinearExpression.is_sparse`` tells whether an expression is CSR-backed, and its repr header reads ``LinearExpression (sparse)``. The opt-in option ``linopy.options["warn_on_densify"]`` (default ``False``) emits a ``PerformanceWarning`` naming the reason whenever a sparse backing is dropped, e.g. on ``.data``, an operation without a sparse path, or ``CSRConstraint.mutable()``. (`#969 <https://github.com/PyPSA/linopy/issues/969>`__)
+* New ``LinearExpression.is_sparse`` tells whether an expression is CSR-backed, and its repr header reads ``LinearExpression (sparse)``. The option ``linopy.options["warn_on_densify"]`` emits a ``PerformanceWarning`` naming the reason whenever a sparse backing is dropped, e.g. on ``.data``, an operation without a sparse path, or ``CSRConstraint.mutable()``. (`#969 <https://github.com/PyPSA/linopy/issues/969>`__)
 * Multiplying, dividing, adding or subtracting a non-scalar constant (``DataArray``, ``Series``, ``ndarray``, …) keeps a CSR-backed ``LinearExpression`` sparse: the coefficients are scaled per cell and only the constant changes, instead of expanding the dense rectangle. The constant is aligned exactly as on the dense path (NaN, label mismatch, joins and auxiliary coordinates); an operand that adds dimensions still densifies. ``expr.const`` is now served from the sparse backing as well. (`#965 <https://github.com/PyPSA/linopy/issues/965>`__)
 * ``sum(dim=...)`` and a further ``groupby(...).sum()`` keep a CSR-backed ``LinearExpression`` sparse: summed rows are merged in the sparse backing instead of stacking the summed dimensions into the term dimension of the dense rectangle. Auxiliary coordinates, absent cells and the constant follow the dense path; ``groupby(...).sum(sparse=False)`` still densifies. (`#964 <https://github.com/PyPSA/linopy/issues/964>`__)
 * ``where``, ``sel``, ``isel``, ``loc`` and ``[]`` keep a CSR-backed ``LinearExpression`` sparse: the selection is evaluated on the grid's row numbers with the dense (xarray) semantics, including ``drop=``, ``method=``, scalar coordinates left by a scalar selection and auxiliary coordinates, and the selected rows are gathered from the sparse backing; masked cells become absent. A condition or indexer that introduces dimensions still densifies. A merge of differing grids now stays sparse when operands carry auxiliary coordinates, which are checked for conflicts as on the dense path. (`#966 <https://github.com/PyPSA/linopy/issues/966>`__)

linopy/config.py

@@ -96,5 +96,5 @@ def __repr__(self) -> str:
     display_max_terms=6,
     semantics=LEGACY_SEMANTICS,
     sparse_groupby=False,
-    warn_on_densify=False,
+    warn_on_densify=None,
 )

linopy/constraints.py

@@ -1520,7 +1520,11 @@ def freeze(self) -> CSRConstraint:
 
     def to_dense(self) -> Constraint:
         """Convert to a Constraint."""
-        _densify_notice("frozen constraint converted by `to_dense()`/`mutable()`")
+        _densify_notice(
+            "frozen constraint converted by `to_dense()`/`mutable()`",
+            self._model,
+            explicit=True,
+        )
         return Constraint(self.data, self._model, self._name)
 
     def mutable(self) -> Constraint:
@@ -2787,7 +2791,11 @@ def set_blocks(self, block_map: np.ndarray) -> None:
 
         for name, constraint in self.items():
             if not isinstance(constraint, Constraint):
-                self.data[name] = constraint = constraint.mutable()
+                _densify_notice(
+                    "frozen constraint converted by `set_blocks`", self.model
+                )
+                constraint = Constraint(constraint.data, self.model, name)
+                self.data[name] = constraint
             res = xr.full_like(constraint.labels, N + 1, dtype=block_map.dtype)
             entries = replace_by_map(constraint.vars, block_map)
 

linopy/csr.py

@@ -536,6 +536,44 @@ def taken(self, rows: np.ndarray, grid: Grid) -> CSRLinearExpression:
         csr = self.csr[np.where(present, rows, 0)]
         return replace(self, csr=csr, grid=grid).with_const(const)
 
+    def concatenated(
+        self, others: Iterable[CSRLinearExpression], dim: str, grid: Grid
+    ) -> CSRLinearExpression:
+        """
+        Stack this and ``others``, in that order, along ``dim`` into the
+        cells of ``grid``, all operands sharing the labels of the other dims
+        in this grid's dim order. Rows are gathered with their terms,
+        explicit zeros included, and their constants. Auxiliary coordinates
+        are ``grid``'s.
+        """
+        parts = (self, *others)
+        n_cols = max(p.csr.shape[1] for p in parts)
+        blocks = [
+            scipy.sparse.csr_array(
+                (p.csr.data, p.csr.indices, p.csr.indptr), shape=(p.n_cells, n_cols)
+            )
+            for p in parts
+        ]
+        csr = scipy.sparse.vstack(blocks, format="csr")
+        if csr.shape[0] != grid.size:
+            raise ValueError(
+                f"Stacked {csr.shape[0]} rows into a grid of {grid.size} cells."
+            )
+        const = np.concatenate([p.const for p in parts])
+        stacked = replace(self, csr=csr, const=const, grid=grid)
+        axis = self.grid.dims.index(dim)
+        if not axis:
+            return stacked
+        offsets = np.cumsum([0, *(p.n_cells for p in parts[:-1])])
+        order = np.concatenate(
+            [
+                (o + np.arange(p.n_cells)).reshape(p.shape)
+                for o, p in zip(offsets, parts)
+            ],
+            axis=axis,
+        ).reshape(-1)
+        return stacked.taken(order, grid)
+
     def reindexed(self, grid: Grid, fill: float = np.nan) -> CSRLinearExpression:
         """
         Remap rows onto a new grid, possibly in a new dim order: dropped
@@ -549,12 +587,7 @@ def reindexed(self, grid: Grid, fill: float = np.nan) -> CSRLinearExpression:
             source = np.full(grid.size, -1, dtype=row_map.dtype)
             source[row_map] = np.arange(grid.size)
             if (source >= 0).all():
-                return replace(
-                    self,
-                    csr=self.csr[source],
-                    const=self.const[source],
-                    grid=self.grid.conformed(grid),
-                )
+                return self.taken(source, self.grid.conformed(grid))
 
         coo = self.csr.tocoo()
         keep = valid[coo.coords[0]]
@@ -781,12 +814,15 @@ def csr_to_term_arrays(
     return vars_, coeffs
 
 
-def _densify_notice(reason: str) -> None:
+def _densify_notice(reason: str, model: Model, explicit: bool = False) -> None:
     """
     Emit a :class:`~linopy.constants.PerformanceWarning` naming why a sparse
-    (CSR) backing is dropped, if ``options["warn_on_densify"]`` is set.
+    (CSR) backing is dropped, if ``options["warn_on_densify"]`` is set, or
+    left at its default ``None``, ``model`` is sparse and the conversion is
+    an implicit fallback rather than an ``explicit`` user request.
     """
-    if options["warn_on_densify"]:
+    enabled = options["warn_on_densify"]
+    if model.sparse and not explicit if enabled is None else enabled:
         warn_outside_linopy(
             f"Sparse (CSR) backing densified: {reason}.", PerformanceWarning
         )

linopy/expressions.py

@@ -666,7 +666,9 @@ def sum(
                     "existing dimension, without use_fallback."
                 )
             if csr is None:
-                _densify_notice("groupby-sum with a grouper without a sparse path")
+                _densify_notice(
+                    "groupby-sum with a grouper without a sparse path", self.model
+                )
 
         if multikey_frame is not None:
             group = multikey_frame
@@ -2609,31 +2611,51 @@ def _selected(
         flat = rows.fillna(-1).to_numpy().reshape(-1).astype(np.int64)
         return csr.taken(flat, grid)
 
-    def sel(self, *args: Any, **kwargs: Any) -> LinearExpression:
+    def _gathered(
+        self, name: str, dense: Callable[..., Self], /, *args: Any, **kwargs: Any
+    ) -> Self:
         """
-        Select by label as ``Dataset.sel``. For a CSR-backed expression,
-        returns a CSR-backed result when the selection stays on the grid.
+        Apply the dense method ``name`` to the grid's row numbers via
+        :meth:`_selected`, falling back to ``dense`` off the grid.
         """
-        csr = self._selected(lambda rows: rows.sel(*args, **kwargs), "sel")
+        csr = self._selected(operator.methodcaller(name, *args, **kwargs), name)
         if csr is None:
-            return super().sel(*args, **kwargs)
+            return dense(*args, **kwargs)
         return type(self)._from_csr(csr, self._model)
 
-    def isel(self, *args: Any, **kwargs: Any) -> LinearExpression:
+    def sel(self, *args: Any, **kwargs: Any) -> Self:
+        """
+        Select by label as ``Dataset.sel``. For a CSR-backed expression,
+        returns a CSR-backed result when the selection stays on the grid.
+        """
+        return self._gathered("sel", super().sel, *args, **kwargs)
+
+    def isel(self, *args: Any, **kwargs: Any) -> Self:
         """
         Select by position as ``Dataset.isel``. For a CSR-backed expression,
         returns a CSR-backed result when the selection stays on the grid.
         """
-        csr = self._selected(lambda rows: rows.isel(*args, **kwargs), "isel")
-        if csr is None:
-            return super().isel(*args, **kwargs)
-        return type(self)._from_csr(csr, self._model)
+        return self._gathered("isel", super().isel, *args, **kwargs)
 
     def __getitem__(self, selector: int | tuple[slice, list[int]] | slice) -> Self:
-        csr = self._selected(lambda rows: rows[selector], "__getitem__")
-        if csr is None:
-            return super().__getitem__(selector)
-        return type(self)._from_csr(csr, self._model)
+        return self._gathered("__getitem__", super().__getitem__, selector)
+
+    def shift(self, *args: Any, **kwargs: Any) -> Self:
+        """
+        Shift along dimensions as ``Dataset.shift``. For a CSR-backed
+        expression, returns a CSR-backed result with the shifted-in cells
+        absent, unless a ``fill_value`` is passed.
+        """
+        if len(args) > 1 or "fill_value" in kwargs:
+            self._densify("`shift` with a fill_value")
+        return self._gathered("shift", super().shift, *args, **kwargs)
+
+    def roll(self, *args: Any, **kwargs: Any) -> Self:
+        """
+        Roll along dimensions as ``Dataset.roll``. For a CSR-backed
+        expression, returns a CSR-backed result.
+        """
+        return self._gathered("roll", super().roll, *args, **kwargs)
 
     def where(
         self,
@@ -3703,7 +3725,7 @@ def _densify_all(exprs: Iterable[Any], reason: str) -> None:
         if isinstance(e, LinearExpression) and (csr := e._csr) is not None
     ]
     if sparse:
-        _densify_notice(reason)
+        _densify_notice(reason, sparse[0][0].model)
     for e, csr in sparse:
         e._data = csr.to_dense()._data
         e._csr = None
@@ -3737,6 +3759,40 @@ def _aligned(
     return [p.reindexed(grid, fill) for p in csrs]
 
 
+def _concatenated(
+    csrs: list[CSRLinearExpression], dim: str, join: JoinOptions | None
+) -> CSRLinearExpression | str:
+    """
+    Stack CSR expressions sharing one dim order along the grid dim ``dim``,
+    in order. The result's coordinates are what the dense path's
+    ``xr.concat`` yields on the grids alone: labels and auxiliary
+    coordinates concatenated along ``dim``, the other dims joined as
+    ``join`` says (``outer`` by default, after the v1 label check for the
+    auto-detected join), with the cells the join creates absent. Returns the
+    reason as a string where the dense path owns the semantics instead: a
+    join over non-unique labels.
+    """
+    dims = csrs[0].grid.dims
+    metadata = [p.grid.to_dataset() for p in csrs]
+    if join is None:
+        enforce_merge_dims(metadata, concat_dim=dim, context=f"merge along dim {dim!r}")
+    enforce_aux_conflict(metadata, concat_dim=dim)
+    combined = xr.concat(
+        metadata, dim, join=join or "outer", coords="minimal", compat="override"
+    )
+    grid = Grid.from_dataset(combined, dims)
+    aligned = []
+    for p in csrs:
+        target = grid.with_indexes({dim: p.grid.indexes[dim]})
+        if join == "override" or p.grid.same_layout(target):
+            aligned.append(replace(p, grid=target))
+        elif p.grid.is_unique:
+            aligned.append(p.reindexed(target))
+        else:
+            return "merge over non-unique labels"
+    return aligned[0].concatenated(aligned[1:], dim, grid)
+
+
 def _try_csr_merge(
     exprs: Any,
     dim: str,
@@ -3747,20 +3803,19 @@ def _try_csr_merge(
     """
     Sparse branch of :func:`merge`: combine plain LinearExpressions over one
     set of grid dimensions (CSR-backed or dense-convertible) as sparse matrix
-    addition. Grids that share dims in a different order are transposed onto
-    the template order first. 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``). Auxiliary
-    coordinates are checked for conflicts on the operands as given and follow
-    their rows onto the joined grid. Returns None to fall through to the
-    dense path.
+    addition along the term dimension, or as a row stack along one of the
+    grid dimensions (:func:`_concatenated`). Grids that share dims in a
+    different order are transposed onto the template order first. 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``). Auxiliary coordinates are checked for
+    conflicts on the operands as given and follow their rows onto the joined
+    grid. Returns None to fall through to the dense path.
     """
     if not any(type(e) is LinearExpression and e._csr is not None for e in exprs):
         return None
-    if dim != TERM_DIM or kwargs:
-        _densify_all(
-            exprs, "merge along a coordinate dimension or with extra arguments"
-        )
+    if kwargs:
+        _densify_all(exprs, "merge with extra arguments")
         return None
     if not all(type(e) is LinearExpression for e in exprs):
         _densify_all(exprs, "merge with an operand that is not a LinearExpression")
@@ -3769,6 +3824,9 @@ def _try_csr_merge(
     if any(set(e.coord_dims) != dims for e in exprs[1:]):
         _densify_all(exprs, "merge of operands over different dimensions")
         return None
+    if dim != TERM_DIM and dim not in dims:
+        _densify_all(exprs, "merge along a new dimension")
+        return None
     for e in exprs:
         if e._csr is None and set(e.data.coords) - dims != set(
             _aux_coords(e.data, dims)
@@ -3785,6 +3843,12 @@ def _try_csr_merge(
         p.reindexed(p.grid.reordered(order)) if p.grid.dims != order else p
         for p in csrs
     ]
+    if dim != TERM_DIM:
+        stacked = _concatenated(csrs, dim, join)
+        if isinstance(stacked, str):
+            _densify_all(exprs, stacked)
+            return None
+        return LinearExpression._from_csr(stacked, exprs[0].model)
     if not all(template.same_grid(p) for p in csrs[1:]):
         enforce_aux_conflict([Dataset(coords=p.grid.aux) for p in csrs])
         aligned = _aligned(csrs, join, join_fill(fill_value, 0.0))