perf(matrices): keep broadcast constraint fields lazy when assembling matrices

doc/release_notes.rst

@@ -107,6 +107,7 @@ Version 1.0.0
 * ``Model.from_spec`` builds the PyPSA spec about 4 to 5 times faster with about a third less peak memory: evaluated nodes are reused within a build, empty constraints are skipped before evaluation, and ``sum(..., by=...)`` on a dense expression groups in numpy.
 * ``Store.terms(active=None)`` returns the number of live terms per cell instead of a cell index per term, and drops the cells outside ``active``. ``Constraint.to_polars`` spreads labels with ``np.repeat`` and masks rows inside the store, so it no longer allocates an int64 index per nonzero or copies its frame to filter masked rows. (`#1058 <https://github.com/PyPSA/linopy/issues/1058>`__)
 * ``Model.assign_coords`` no longer relabels containers through a Python dict. It locates each container's labels with ``master.get_indexer(index)`` and takes the new values with ``new.take(positions)``: on a tz-aware ``DatetimeIndex``, the rolling-horizon case it is built for, the previous ``Index.map`` over a dict fell into pandas' object-to-datetime inference and cost milliseconds per container. Reassigning a dataset's coordinates also builds it once instead of twice. A model with ~30 containers on a 4380-label tz-aware axis reassigns in ~9 ms, down from ~810 ms.
+* ``Model.matrices`` keeps broadcast scalar constraint fields (sign, rhs, scaling) as scalars and skips the row-mask copy when all rows are active, speeding up matrix assembly; output is unchanged.
 
 **Bug fixes**
 

linopy/matrices.py

@@ -29,6 +29,7 @@ class _RowBlock(NamedTuple):
     """
     Active rows of one constraint. ``positional`` returns their positional
     CSR, built only while the rows are gathered; ``nnz`` bounds its nonzeros.
+    ``rhs``, ``sign`` and ``scaling`` are 0-d when the field is a broadcast scalar.
     """
 
     positional: Callable[[], scipy.sparse.csr_array]
@@ -48,17 +49,26 @@ def _concat(arrays: list[ndarray]) -> ndarray:
     return np.concatenate(arrays) if arrays else np.empty(0, dtype=float)
 
 
+def _active(values: ndarray, rows: ndarray | slice) -> ndarray:
+    """``values`` raveled at ``rows``; a broadcast scalar stays a 0-d view."""
+    flat = values.reshape(-1)
+    if flat.size and not flat.strides[0]:
+        return flat[:1].reshape(())
+    return flat[rows]
+
+
 def _row_block(
     con: Constraint, active: ndarray, label_index: VariableLabelIndex
 ) -> _RowBlock:
     store = con.lhs._store
+    rows = slice(None) if active.all() else active
     return _RowBlock(
         partial(store.positional_csr, label_index, active, con.name),
-        int(active.sum()),
+        int(np.count_nonzero(active)),
         store.nnz(active),
-        con.rhs.values.ravel()[active],
-        con.sign.values.ravel()[active].astype("U1"),
-        con.scaling.values.ravel()[active],
+        _active(con.rhs.values, rows),
+        _active(con.sign.values, rows).astype("U1"),
+        _active(con.scaling.values, rows),
     )
 
 
@@ -112,7 +122,11 @@ def _stack(
             data[entries] = block_data
             b[rows] = block.rhs
         else:
-            per_row = repeat_per_row(row_scaling, block_indptr)
+            per_row = (
+                repeat_per_row(row_scaling, block_indptr)
+                if row_scaling.ndim
+                else row_scaling
+            )
             np.multiply(block_data, per_row, out=data[entries])
             np.multiply(block.rhs, row_scaling, out=b[rows])
         if col_scaling is not None:
@@ -293,7 +307,9 @@ def _build_cons(m: Model, col_scaling: ndarray | None) -> _ConArrays:
     label_dtype = m._dtypes["labels"]
     return _ConArrays(
         clabels=m.constraints.label_index.clabels,
-        con_scaling=_concat([block.scaling for block in regular]),
+        con_scaling=_concat(
+            [np.broadcast_to(block.scaling, block.n_rows) for block in regular]
+        ),
         A=A,
         b=b,
         sense=sense,

test/test_matrices.py

@@ -11,7 +11,7 @@
 import xarray as xr
 
 from linopy import EQUAL, GREATER_EQUAL, Model
-from linopy.matrices import Matrices
+from linopy.matrices import Matrices, _row_block
 
 
 def test_basic_matrices() -> None:
@@ -238,3 +238,51 @@ def test_matrices_take_rejects_indicator_constraints() -> None:
 
     with pytest.raises(NotImplementedError, match="indicator"):
         m.matrices.take(rows, cols)
+
+
+def test_row_block_keeps_broadcast_fields_lazy() -> None:
+    m = Model()
+    x = m.add_variables(coords=[pd.RangeIndex(4, name="i")], name="x")
+    c = m.add_constraints(2 * x >= 3, name="c", scaling=2.0)
+    block = _row_block(c, np.array([True, False, True, True]), m.variables.label_index)
+    assert block.n_rows == 3
+    assert block.sign.shape == block.scaling.shape == ()
+    assert block.sign.dtype == np.dtype("U1")
+
+
+@pytest.mark.parametrize("sparse", [False, True], ids=["dense", "sparse"])
+def test_matrices_scalar_fields(sparse: bool) -> None:
+    m = Model()
+    i = pd.RangeIndex(3, name="i")
+    x = m.add_variables(coords=[i], name="x")
+    cons = [
+        m.add_constraints(2 * x >= 1, name="a", scaling=2.0),
+        m.add_constraints(
+            x == 4,
+            name="d",
+            scaling=4.0,
+            mask=xr.DataArray([True, False, True], coords=[i]),
+        ),
+        # fully masked: empty block with a 0-d rhs and scaling
+        m.add_constraints(
+            x >= 0, name="off", scaling=3.0, mask=xr.DataArray([False] * 3, coords=[i])
+        ),
+        # dimensionless constraint
+        m.add_constraints(x.sum() <= 7, name="tot", scaling=2.0),
+    ]
+    if sparse:
+        for con in cons:
+            con.to_sparse()
+    M = m.matrices
+
+    assert M.A is not None
+    np.testing.assert_array_equal(
+        M.A.toarray(),
+        [[4, 0, 0], [0, 4, 0], [0, 0, 4], [4, 0, 0], [0, 0, 4], [2, 2, 2]],
+    )
+    np.testing.assert_array_equal(M.b, [2, 2, 2, 16, 16, 14])
+    assert M.sense.dtype == np.dtype("U1")
+    np.testing.assert_array_equal(M.sense, [">", ">", ">", "=", "=", "<"])
+    assert M.con_scaling.dtype == float
+    np.testing.assert_array_equal(M.con_scaling, [2, 2, 2, 4, 4, 2])
+    np.testing.assert_array_equal(M.clabels, [0, 1, 2, 3, 5, 9])