perf(sparse): trim the stages around the CSR solver hand-off

doc/release_notes.rst

@@ -73,6 +73,12 @@ Upcoming Version
 * ``LinearExpression.flat`` and ``LinearExpression.to_polars`` on a CSR-backed expression are emitted directly from the sparse backing instead of expanding the dense rectangle; the rows equal those of the dense path, with absent cells and zero coefficients dropped. (`#968 <https://github.com/PyPSA/linopy/issues/968>`__)
 * A CSR-backed linear objective stays sparse when set with ``Model.add_objective`` and when the model is exported or solved: the objective vector ``matrices.c``, LP/MPS files, netcdf output, ``Model.copy``, persistent snapshots and direct solver APIs all read it without expanding the dense rectangle, with results equal to the dense objective. The objective name is now owned by the ``Objective`` itself (``Objective.name`` and ``Objective.attrs``) instead of being written into the expression's attributes. (`#967 <https://github.com/PyPSA/linopy/issues/967>`__)
 * Persistent snapshots of tz-aware ``DatetimeIndex`` coordinates no longer materialise an object array of ``Timestamp`` per container per capture and diff. Coordinates are stored as UTC-ns arrays with the timezone identity carried alongside, making snapshot capture ~24x and warm-start diffs ~33x faster on tz-aware models, while naive and tz-aware coordinates — and differing timezones — stay correctly unequal. (`#960 <https://github.com/PyPSA/linopy/pull/960>`__)
+* The MOSEK direct API builds bound keys, bounds and constraint names vectorised instead of looping over every row and column in Python.
+* The LP writer computes the scaling lookups once per file, skips scaling when every factor is 1, does not re-sort rows that are already grouped (as frozen constraints deliver them) and builds each constraint line from fewer string branches. Constraint writing is about 2x faster; the output is byte-identical.
+* ``linopy.merge`` of CSR-backed expressions concatenates all operands once instead of folding them pairwise, so an N-way merge scales with the total number of nonzeros instead of quadratically with the operand count.
+* ``LinearExpression.solution`` and ``QuadraticExpression.solution`` no longer rebuild the full constraint matrix; the solution is mapped by variable label, and a CSR-backed expression is evaluated on its sparse backing without densifying. Roughly 10x faster and 6x less memory on a large model. ``Model.dualize`` reads ``model.matrices`` once instead of four times.
+* ``model.matrices`` assembles the constraint blocks into preallocated CSR buffers in one pass, without ``vstack``, ``eliminate_zeros`` or a full-grid scaling round trip; scaling is skipped when it is identity. Label lookups use the model's label dtype, so ``matrices.clabels`` and ``matrices.indicator_binvar`` are now ``int32`` by default. On an 8M-row model this is 45% faster with 200 MB less peak memory, with bit-identical output.
+* Duals of frozen constraints are read back on the active rows directly instead of through three full-grid arrays, and freezing a masked constraint gathers rows once instead of up to three times. Mixed-sign constraints convert their sense vectorised instead of in a Python loop.
 
 **Bug fixes**
 

linopy/common.py

@@ -930,13 +930,16 @@ def label_to_pos(self) -> np.ndarray:
         """
         Mapping from variable label to dense position, shape (_xCounter,).
 
+        Positions share the model's label dtype, since they never exceed a label.
+
         Position i in the active variable array corresponds to label vlabels[i].
         Masked or unused labels map to -1.
         """
         vlabels = self.vlabels
-        n = self._variables.model._xCounter
-        label_to_pos = np.full(n, -1, dtype=np.intp)
-        label_to_pos[vlabels] = np.arange(len(vlabels), dtype=np.intp)
+        model = self._variables.model
+        dtype = model._dtypes["labels"]
+        label_to_pos = np.full(model._xCounter, -1, dtype=dtype)
+        label_to_pos[vlabels] = np.arange(len(vlabels), dtype=dtype)
         return label_to_pos
 
     @property
@@ -969,17 +972,17 @@ def clabels(self) -> np.ndarray:
             for c in self._constraints.data.values()
             if not c.is_indicator
         ]
-        return (
-            np.concatenate(label_lists) if label_lists else np.array([], dtype=np.intp)
-        )
+        dtype = self._constraints.model._dtypes["labels"]
+        return np.concatenate([np.array([], dtype=dtype), *label_lists], dtype=dtype)
 
     @cached_property
     def label_to_pos(self) -> np.ndarray:
         """Mapping from constraint label to dense position, shape (_cCounter,)."""
         clabels = self.clabels
-        n = self._constraints.model._cCounter
-        label_to_pos = np.full(n, -1, dtype=np.intp)
-        label_to_pos[clabels] = np.arange(len(clabels), dtype=np.intp)
+        model = self._constraints.model
+        dtype = model._dtypes["labels"]
+        label_to_pos = np.full(model._cCounter, -1, dtype=dtype)
+        label_to_pos[clabels] = np.arange(len(clabels), dtype=dtype)
         return label_to_pos
 
     @property

linopy/constraints.py

@@ -939,31 +939,39 @@ def _replace(self, **changes: Any) -> CSRConstraint:
         return new
 
     def assign_labels(
-        self, cindex: int, name: str, scaling: float | DataArray = 1.0
+        self,
+        cindex: int,
+        name: str,
+        scaling: float | DataArray = 1.0,
+        mask: np.ndarray | None = None,
     ) -> CSRConstraint:
         """
         Return a copy labelled from ``cindex`` and named ``name``.
 
         Rows without terms are dropped, as when freezing a dense constraint;
-        a zero coefficient counts as a term. ``scaling`` is a scalar or a row
-        scaling broadcast on the grid; its distinct values are validated
-        without expanding a broadcast view.
+        a zero coefficient counts as a term. Active rows where the boolean
+        ``mask`` is False are dropped in the same gather. ``scaling`` is a
+        scalar or a row scaling broadcast on the grid; its distinct values are
+        validated without expanding a broadcast view.
         """
         values = np.asarray(scaling)
         distinct = values[tuple(slice(None) if s else 0 for s in values.strides)]
         validate_scaling(distinct, "constraint scaling")
-        kept = self._kept(np.diff(self._csr.indptr) > 0)
+        keep = np.diff(self._csr.indptr) > 0
+        if mask is not None:
+            keep &= mask
+        kept = self._kept(keep)
         csr = kept._csr
         if not csr.data.all():
             csr = csr.copy() if csr is self._csr else csr
             csr.eliminate_zeros()
         if isinstance(scaling, DataArray):
-            row_scaling = kept._active_values(scaling)
+            row_scaling = kept.active_values(scaling)
         else:
             row_scaling = np.full(csr.shape[0], float(scaling))
         return kept._replace(csr=csr, cindex=cindex, name=name, scaling=row_scaling)
 
-    def _active_values(self, values: DataArray) -> np.ndarray:
+    def active_values(self, values: DataArray) -> np.ndarray:
         """
         Values of ``values``, broadcast on the grid, at the active rows.
 
@@ -1002,13 +1010,6 @@ def rows(values: Any) -> Any:
             binval=rows(self._binval),
         )
 
-    def masked(self, mask: DataArray) -> CSRConstraint:
-        """
-        Copy with the cells where the boolean ``mask`` is False made inactive,
-        without the dense rectangle. ``mask`` must lie on the constraint grid.
-        """
-        return self._kept(self._active_values(mask).astype(bool))
-
     def _assign_coords(self, **coords: Any) -> CSRConstraint:
         """
         Reassign coordinate values on the constraint, keeping the shape.
@@ -1448,7 +1449,7 @@ def to_matrix_with_rhs(
         if isinstance(self._sign, str):
             sense = np.full(len(self._rhs), self._sign[0])
         else:
-            sense = np.array([s[0] for s in self._sign])
+            sense = self._sign.astype("U1")
         return (
             self._to_positional_csr(label_index),
             self.active_labels(),
@@ -1472,7 +1473,9 @@ def sanitize_zeros(self) -> CSRConstraint:
         external holders of the previous arrays (e.g. a ModelSnapshot
         sharing them) keep a valid baseline.
         """
-        zeros = np.abs(self._csr.data) <= 1e-10
+        data = self._csr.data
+        zeros = data <= 1e-10
+        zeros &= data >= -1e-10
         if zeros.any():
             csr = self._csr.copy()
             csr.data[zeros] = 0
@@ -1617,14 +1620,12 @@ def from_dense(
         scaling = con.scaling.values.ravel()[active_mask]
         sign_vals = con.sign.values.ravel()
         active_signs = sign_vals[active_mask]
-        unique_signs = np.unique(active_signs)
-        if len(unique_signs) == 0:
+        sign: str | np.ndarray
+        if not len(active_signs):
             full_unique_signs = np.unique(sign_vals)
-            sign: str | np.ndarray = (
-                str(full_unique_signs.item()) if len(full_unique_signs) == 1 else "="
-            )
-        elif len(unique_signs) == 1:
-            sign = str(unique_signs[0])
+            sign = str(full_unique_signs.item()) if len(full_unique_signs) == 1 else "="
+        elif (active_signs == active_signs[0]).all():
+            sign = str(active_signs[0])
         else:
             sign = active_signs
         dual = (
@@ -2221,12 +2222,7 @@ def to_matrix_with_rhs(
         csr.sum_duplicates()
 
         b = self.rhs.values.ravel()[row_mask]
-        sign_flat = self.sign.values.ravel()[row_mask]
-        unique_signs = np.unique(sign_flat)
-        if len(unique_signs) == 1:
-            sense = np.full(len(con_labels), str(unique_signs[0])[0], dtype="U1")
-        else:
-            sense = sign_flat.astype("U1")
+        sense = self.sign.values.ravel()[row_mask].astype("U1")
         return csr, con_labels, b, sense
 
     def sanitize_zeros(self) -> Constraint:

linopy/csr.py

@@ -28,6 +28,7 @@
 
 from __future__ import annotations
 
+import functools
 import operator
 from collections.abc import Callable, Iterable, Mapping
 from dataclasses import dataclass, field, replace
@@ -572,25 +573,30 @@ def same_grid(self, other: CSRLinearExpression) -> bool:
         """Whether both live on the same cells, auxiliary coordinates aside."""
         return self.grid.same_layout(other.grid)
 
-    def added(self, other: CSRLinearExpression) -> CSRLinearExpression:
+    def added(self, *others: CSRLinearExpression) -> CSRLinearExpression:
         """
-        Sparse matrix addition == merge along the term dimension. Goes through
-        COO so explicit zero coefficients survive (scipy's ``+`` drops them),
+        Sparse matrix addition == merge along the term dimension, over any
+        number of operands on the same grid in one COO pass. 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. A cell absent in either operand is absent in the sum
-        and carries no terms. Auxiliary coordinates propagate and conflicting
-        ones raise (§11).
-        """
-        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])
+        an empty cell. A cell absent in any operand is absent in the sum and
+        carries no terms. Auxiliary coordinates propagate, earlier operands
+        taking precedence, and conflicting ones raise (§11).
+        """
+        parts = (self, *others)
+        const = functools.reduce(np.add, (p.const for p in parts))
+        coos = [p.csr.tocoo() for p in parts]
+        shape = (self.n_cells, max(c.shape[1] for c in coos))
+        rows = np.concatenate([c.coords[0] for c in coos])
+        cols = np.concatenate([c.coords[1] for c in coos])
+        data = np.concatenate([c.data for c in coos])
         present = ~np.isnan(const)[rows]
         csr = coo_to_csr(data[present], rows[present], cols[present], shape, self.model)
-        enforce_aux_conflict([Dataset(coords=p.grid.aux) for p in (self, other)])
-        grid = replace(self.grid, aux=other.grid.aux | self.grid.aux)
+        enforce_aux_conflict([Dataset(coords=p.grid.aux) for p in parts])
+        aux: AuxCoords = {}
+        for p in reversed(parts):
+            aux |= p.grid.aux
+        grid = replace(self.grid, aux=aux)
         return replace(self, csr=csr, const=const, grid=grid)
 
     def contracted(

linopy/dualization.py

@@ -479,12 +479,13 @@ def _add_dual_feasibility_constraints(
     dual_vars : dict
         ``{constraint_name: dual_variable}`` as returned by ``_add_dual_variables()``.
     """
-    A = m.matrices.A
+    M = m.matrices
+    A = M.A
     if A is None:
         raise ValueError("Constraint matrix is None, model has no constraints.")
 
-    vlabels = np.asarray(m.matrices.vlabels, dtype=np.int64)
-    clabels = np.asarray(m.matrices.clabels, dtype=np.int64)
+    vlabels = np.asarray(M.vlabels, dtype=np.int64)
+    clabels = np.asarray(M.clabels, dtype=np.int64)
 
     flat_con_to_dual = _build_flat_con_to_dual_label_lookup(m, dual_vars)
     if not len(flat_con_to_dual):
@@ -496,7 +497,7 @@ def _add_dual_feasibility_constraints(
     flat_v, flat_d, nnz_data = _extract_dual_feas_entries(
         A, vlabels, clabels, flat_con_to_dual
     )
-    c_lookup = _build_obj_coeff_lookup(vlabels, m.matrices.c)
+    c_lookup = _build_obj_coeff_lookup(vlabels, M.c)
 
     logger.debug("Building dual feasibility constraints for each primal variable.")
     for var_name, var in m.variables.items():

linopy/expressions.py

@@ -1737,17 +1737,25 @@ def mask(self) -> None:
         return None
 
     @has_optimized_model
+    def _label_solution(self) -> np.ndarray:
+        """
+        Solution values indexed by variable label, with a trailing NaN that
+        label ``-1`` reads.
+        """
+        sol = np.full(self.model._xCounter + 1, np.nan)
+        for _, var in self.model.variables.items():
+            labels = var.labels.values.ravel()
+            mask = labels != -1
+            sol[labels[mask]] = var.solution.values.ravel()[mask]
+        return sol
+
     def _map_solution(self) -> DataArray:
         """
         Replace variable labels by solution values.
         """
-        m = self.model
-        M = m.matrices
-        sol = pd.Series(M.sol, M.vlabels)
-        sol[-1] = np.nan
-        idx = np.ravel(self.vars)
-        values = np.asarray(sol[idx]).reshape(self.vars.shape)
-        return xr.DataArray(values, dims=self.vars.dims, coords=self.vars.coords)
+        labels = self.vars
+        values = self._label_solution()[labels.values]
+        return xr.DataArray(values, dims=labels.dims, coords=labels.coords)
 
     @property
     def solution(self) -> DataArray:
@@ -2496,6 +2504,20 @@ def const(self) -> DataArray:
     def const(self, value: DataArray) -> None:
         self._data = assign_multiindex_safe(self.data, const=value)
 
+    @property
+    def solution(self) -> DataArray:
+        """
+        Get the optimal values of the expression.
+
+        The function raises an error in case no model is set as a
+        reference or the model is not optimized.
+        """
+        csr = self._csr
+        if csr is None:
+            return super().solution
+        sol = np.nan_to_num(self._label_solution()[: csr.csr.shape[1]])
+        return csr.grid.dataarray(csr.csr @ sol + csr.const, name="solution")
+
     def _combined_with_constant(
         self,
         self_const: DataArray,
@@ -3736,9 +3758,7 @@ def _try_csr_merge(
             return None
         csrs = aligned
 
-    combined = csrs[0]
-    for csr in csrs[1:]:
-        combined = combined.added(csr)
+    combined = csrs[0].added(*csrs[1:])
     return LinearExpression._from_csr(combined, exprs[0].model)