perf(csr): frozen constraint build copies the sparse matrix 2-3 times and expands scaling over the full grid #977

Describe the feature you'd like to see

Adding a frozen (CSRConstraint) constraint from a sparse expression allocates about 2x the size of the lhs CSR matrix, although the matrix could mostly be reused. Found while working on #974 (sparse/CSR plan #972); measured on feat/csr-boundaries @ 0cf4a200.

Where the copies come from:

1. CSRConstraint.from_csr (linopy/constraints.py) always does expr.csr[active], a fancy row index that copies data/indices/indptr even when every row is active (the common case).
2. assign_labels calls _kept(np.diff(indptr) > 0), another self._csr[keep] row copy, then eliminate_zeros(). Again a full copy when no row is empty.
3. With mask=, masked goes through _kept too, so that is a third copy. It runs after step 1 has already copied the full, unmasked matrix.
4. Model.add_constraints (linopy/model.py, frozen branch) broadcasts scaling over the full coordinate grid (broadcast_to_coords(scaling, con.coords) then .values.ravel()), and assign_labels then indexes it down to the active rows. When the grid is much bigger than the active rows (sparse masks, grouped results), this dense array is mostly wasted.

Measured (tracemalloc peak of add_constraints(..., freeze=True), lhs CSR built beforehand)

The last row keeps 1 % of the rows (about 0.4 MB of CSR) but peaks at 82 MB, because the whole matrix is copied before the mask is applied. On its own, the full-grid scaling there costs 6.4 ms and 20.1 MB peak, where the 12,523 active rows need 0.1 MB.

Per step (1e6 nnz, all rows active, 16.8 MB CSR): from_csr peaks at 20.0 MB, masked (all True) at 19.3 MB, assign_labels at 20.8 MB. None of the results share memory with their input.

Suggestions

• Skip the row slice when all rows are active (active.size == n_rows), both in from_csr and _kept. Reuse the CSR, or copy it once if it must be owned.
• Fuse steps 1–3: compute one combined keep-mask (non-NaN rhs & mask & non-empty row) and slice once. Apply eliminate_zeros only if (data == 0).any(), on the already sliced matrix.
• Compute the scaling only over the active rows: index a broadcast scalar/DataArray at active_positions, don't materialise the full grid. A scalar scaling (the default 1) needs no array at all before the final np.full.

Minimal reproducible example

Needs the freeze=True sparse path from #974. Run with uv run python repro.py.

import time, tracemalloc
import numpy as np, pandas as pd, xarray as xr
import linopy
from linopy import Model

linopy.options["semantics"] = "v1"

def run(n_gen, n_snap, n_bus, active_frac=None):
    m = Model()
    gen = pd.Index(range(n_gen), name="gen")
    snap = pd.Index(range(n_snap), name="snapshot")
    rng = np.random.default_rng(0)
    bus = xr.DataArray(rng.integers(0, n_bus, n_gen), coords=[gen], name="bus")
    p = m.add_variables(0, 10, coords=[gen, snap], name="p")
    lhs = (1.0 * p).groupby(bus).sum(sparse=True)
    shape = (lhs.indexes["bus"].size, n_snap)
    rhs = xr.DataArray(rng.uniform(1, 5, shape), coords=lhs.coords)
    mask = None if active_frac is None else xr.DataArray(rng.random(shape) < active_frac, coords=lhs.coords)
    csr = lhs._csr.csr
    mb = (csr.data.nbytes + csr.indices.nbytes + csr.indptr.nbytes) / 1e6
    tracemalloc.start(); t = time.perf_counter()
    m.add_constraints(lhs, "==", rhs, name="balance", freeze=True, mask=mask)
    dt = time.perf_counter() - t; _, peak = tracemalloc.get_traced_memory(); tracemalloc.stop()
    print(f"nnz={csr.nnz:,} grid={rhs.size:,} active={m.constraints['balance']._csr.shape[0]:,} "
          f"csr={mb:.1f} MB | {dt*1e3:.1f} ms, peak {peak/1e6:.1f} MB = {peak/1e6/mb:.1f}x csr")

run(100_000, 10, 1_000)
run(1_000_000, 1, 100_000)
run(1_000_000, 1, 100_000, active_frac=0.9)
run(2_000, 1_000, 2_000, active_frac=0.01)

import time, tracemalloc
import numpy as np, pandas as pd, xarray as xr
import linopy
from linopy import Model
from linopy.constraints import CSRConstraint, csr_rhs
linopy.options["semantics"] = "v1"
m = Model()
n_gen, n_snap, n_bus = 1_000_000, 1, 100_000
gen = pd.Index(range(n_gen), name="gen"); snap = pd.Index(range(n_snap), name="snapshot")
rng = np.random.default_rng(0)
bus = xr.DataArray(rng.integers(0, n_bus, n_gen), coords=[gen], name="bus")
p = m.add_variables(0, 10, coords=[gen, snap], name="p")
lhs = (1.0 * p).groupby(bus).sum(sparse=True)
rhs = xr.DataArray(rng.uniform(1, 5, (lhs.indexes["bus"].size, n_snap)), coords=lhs.coords)
def meas(label, f):
    tracemalloc.start(); t = time.perf_counter(); r = f(); dt = time.perf_counter() - t
    _, peak = tracemalloc.get_traced_memory(); tracemalloc.stop()
    print(f"{label:45s} {dt*1e3:6.1f} ms  peak {peak/1e6:6.1f} MB"); return r
e = lhs._csr
print("csr MB", (e.csr.data.nbytes + e.csr.indices.nbytes + e.csr.indptr.nbytes) / 1e6)
con = meas("from_csr (csr[active], all rows active)", lambda: CSRConstraint.from_csr(e, "=", csr_rhs(e, rhs)))
print("from_csr shares data with expr:", np.shares_memory(con._csr.data, e.csr.data))
meas("  expr.csr[arange(n)] alone", lambda: e.csr[np.arange(e.csr.shape[0])])
mask = xr.DataArray(np.ones(rhs.shape, bool), coords=rhs.coords)
meas("masked (all True)", lambda: con.masked(mask))
lab = meas("assign_labels (_kept + eliminate_zeros)", lambda: con.assign_labels(0, "c", np.ones(rhs.size)))
print("assign_labels shares data with input:", np.shares_memory(lab._csr.data, con._csr.data))

csr MB 16.799984
from_csr (csr[active], all rows active)          3.6 ms  peak   20.0 MB
from_csr shares data with expr: False
  expr.csr[arange(n)] alone                      2.3 ms  peak   18.4 MB
masked (all True)                                3.2 ms  peak   19.3 MB
assign_labels (_kept + eliminate_zeros)          4.9 ms  peak   20.8 MB
assign_labels shares data with input: False

grid 1256000 active rows 12523
scaling over full grid: 6.4 ms, peak 20.1 MB (active rows need 0.10 MB)