perf: prune explicit zeros at matrix assembly (cut build peak memory)

linopy/constraints.py

@@ -243,6 +243,10 @@ def to_matrix_with_rhs(
     def active_labels(self) -> np.ndarray:
         """Active constraint labels in build order, without building the CSR."""
 
+    @abstractmethod
+    def active_row_mask(self) -> np.ndarray:
+        """Boolean mask over raveled rows selecting active constraint rows."""
+
     def __getitem__(
         self, selector: str | int | slice | list | tuple | dict
     ) -> Constraint:
@@ -955,6 +959,9 @@ def to_matrix_with_rhs(
     def active_labels(self) -> np.ndarray:
         return self._con_labels
 
+    def active_row_mask(self) -> np.ndarray:
+        return np.ones(self._csr.shape[0], dtype=bool)
+
     def sanitize_zeros(self) -> CSRConstraint:
         """
         Remove terms with zero or near-zero coefficients.
@@ -1531,9 +1538,8 @@ def _matrix_export_data(
         row_mask = (labels_flat != -1) & (vars_2d != -1).any(axis=1)
         con_labels = labels_flat[row_mask]
         vars_final = vars_2d[row_mask]
-        valid_final = vars_final != -1
-
         coeffs_final = self.coeffs.values.ravel().reshape(vars_2d.shape)[row_mask]
+        valid_final = (vars_final != -1) & (coeffs_final != 0)
         cols = label_to_pos[vars_final[valid_final]]
         data = coeffs_final[valid_final]
 
@@ -1566,7 +1572,8 @@ def to_matrix(
         csr.sum_duplicates()
         return csr, con_labels
 
-    def active_labels(self) -> np.ndarray:
+    def active_row_mask(self) -> np.ndarray:
+        """Boolean mask over raveled rows: label set and at least one variable present."""
         labels_flat = self.labels.values.ravel()
         vars_vals = self.vars.values
         n_rows = len(labels_flat)
@@ -1575,8 +1582,10 @@ def active_labels(self) -> np.ndarray:
             if n_rows > 0
             else vars_vals.reshape(0, max(1, vars_vals.size))
         )
-        row_mask = (labels_flat != -1) & (vars_2d != -1).any(axis=1)
-        return labels_flat[row_mask]
+        return (labels_flat != -1) & (vars_2d != -1).any(axis=1)
+
+    def active_labels(self) -> np.ndarray:
+        return self.labels.values.ravel()[self.active_row_mask()]
 
     def to_matrix_with_rhs(
         self, label_index: VariableLabelIndex

linopy/matrices.py

@@ -22,10 +22,20 @@
 
 
 def _stack(csrs: list) -> scipy.sparse.csr_array | None:
-    """Vertically stack CSR blocks, or None when there are none."""
+    """
+    Vertically stack CSR blocks, or None when there are none.
+
+    Explicit zeros are dropped: expressions that broadcast against a dense
+    coordinate store one coefficient per pair, most of them zero, and a zero
+    coefficient never changes a constraint. Keeping them only inflates the
+    stored nnz handed to the solvers/writers (e.g. ``highspy.addRows`` scales
+    with stored nnz), so we prune them once, centrally, for every backend.
+    """
     if not csrs:
         return None
-    return cast(scipy.sparse.csr_array, scipy.sparse.vstack(csrs, format="csr"))
+    stacked = cast(scipy.sparse.csr_array, scipy.sparse.vstack(csrs, format="csr"))
+    stacked.eliminate_zeros()
+    return stacked
 
 
 def _concat(arrays: list, dtype: type | None = None) -> ndarray:
@@ -178,7 +188,6 @@ def dual(self) -> ndarray:
         if not self._parent.status == "ok":
             raise ValueError("Model is not optimized.")
         m = self._parent
-        label_index = m.variables.label_index
         dual_list = []
         has_dual = False
         for c in m.constraints.data.values():
@@ -192,9 +201,7 @@ def dual(self) -> ndarray:
                 else:
                     dual_list.append(np.full(len(c._con_labels), np.nan))
             else:
-                csr, _ = c.to_matrix(label_index)
-                nonempty = np.diff(csr.indptr).astype(bool)
-                active_rows = np.flatnonzero(nonempty)
+                active_rows = np.flatnonzero(c.active_row_mask())
                 if "dual" in c.data:
                     dual_list.append(c.dual.values.ravel()[active_rows])
                     has_dual = True

test/test_matrices.py

@@ -68,6 +68,27 @@ def test_matrices_duplicated_variables() -> None:
     assert np.isin(np.unique(np.array(A)), [0.0, 2.0]).all()
 
 
+def test_matrices_drops_explicit_zeros() -> None:
+    # https://github.com/PyPSA/linopy/issues/814
+    # Expressions that broadcast against a dense coordinate store one coefficient
+    # per pair, most of them structurally zero. Those must not reach A, whose
+    # stored nnz drives the direct-solver handoff (e.g. highspy.addRows).
+    m = Model()
+
+    x = m.add_variables(coords=[range(4)], name="x")
+    coeff = xr.DataArray(
+        np.eye(4), dims=["j", "i"], coords={"j": range(4), "i": range(4)}
+    )
+    m.add_constraints((coeff * x.rename(dim_0="i")).sum("i") <= 1, name="c")
+
+    A = m.matrices.A
+    assert A is not None
+    # 4 structural nonzeros on the diagonal; the 12 broadcast zeros are dropped.
+    assert A.nnz == 4
+    assert not (A.data == 0).any()
+    assert np.array_equal(A.todense(), np.eye(4))
+
+
 def test_matrices_float_c() -> None:
     # https://github.com/PyPSA/linopy/issues/200
     m = Model()

test/test_persistent_snapshot_buffers.py

@@ -58,14 +58,27 @@ def test_shape_mismatch_triggers_sparsity_rebuild(baseline_model: Model) -> None
     snap = ModelSnapshot.capture(baseline_model)
     x = baseline_model.variables["x"]
     y = baseline_model.variables["y"]
-    baseline_model.constraints["c1"].lhs = 2 * x + 0 * y.sum()
+    baseline_model.constraints["c1"].lhs = 2 * x + 1 * y.sum()
     diff = ModelDiff.from_snapshot(snap, baseline_model)
     assert diff in {
         RebuildReason.SPARSITY,
         RebuildReason.STRUCTURAL_LABELS,
     }
 
 
+def test_zero_coefficient_term_needs_no_rebuild(baseline_model: Model) -> None:
+    """
+    An explicit zero coefficient is pruned from the matrix, so adding one
+    leaves the sparsity pattern unchanged and needs no rebuild.
+    """
+    snap = ModelSnapshot.capture(baseline_model)
+    x = baseline_model.variables["x"]
+    y = baseline_model.variables["y"]
+    baseline_model.constraints["c1"].lhs = 2 * x + 0 * y.sum()
+    diff = ModelDiff.from_snapshot(snap, baseline_model)
+    assert not isinstance(diff, RebuildReason)
+
+
 def test_zero_row_container_capture() -> None:
     m = Model()
     m.add_variables(0, 10, coords=[range(2)], name="x")