diff --git a/benchmarks/patterns/nodal_balance.py b/benchmarks/patterns/nodal_balance.py index 458df39a4..4cc2e6884 100644 --- a/benchmarks/patterns/nodal_balance.py +++ b/benchmarks/patterns/nodal_balance.py @@ -7,6 +7,12 @@ result's ``_term`` axis blows up — most of it fill. ``severity`` dials that skew; the build's peak memory is expected to climb steeply with it on the current (dense) kernel. + +``nodal_balance_sparse`` builds the identical constraint through the +CSR-backed path (``sum(sparse=True)`` + ``freeze=True`` under v1): the +grouped sum never materializes the padded rectangle and is realized directly +as a CSRConstraint, so its peak memory should stay flat across the severity +sweep — the pair makes the padding cost visible. """ from __future__ import annotations @@ -16,7 +22,14 @@ import xarray as xr import linopy -from benchmarks.registry import SEVERITIES, BenchSpec, register_pattern +from benchmarks.registry import ( + BUILD, + MATRICES, + SEVERITIES, + TO_LP, + BenchSpec, + register_pattern, +) N_GEN = 2000 N_BUS = 50 @@ -43,7 +56,7 @@ def _bus_of_gen(severity: int) -> np.ndarray: return bus -def build_nodal_balance(severity: int) -> linopy.Model: +def _build(severity: int, sparse: bool) -> linopy.Model: gens = pd.RangeIndex(N_GEN, name="gen") time = pd.RangeIndex(N_TIME, name="time") buses = pd.RangeIndex(N_BUS, name="bus") @@ -53,15 +66,29 @@ def build_nodal_balance(severity: int) -> linopy.Model: gen = m.add_variables(lower=0, coords=[gens, time], name="gen") bus_of_gen = pd.Series(_bus_of_gen(severity), index=gens, name="bus") - supply = (1 * gen).groupby(bus_of_gen).sum() + supply = (1 * gen).groupby(bus_of_gen).sum(sparse=sparse) demand = xr.DataArray( rng.uniform(10.0, 100.0, size=(N_BUS, N_TIME)), coords=[buses, time] ) - m.add_constraints(supply == demand, name="balance") + m.add_constraints(supply == demand, name="balance", freeze=sparse) m.add_objective(gen.sum()) return m +def build_nodal_balance(severity: int) -> linopy.Model: + return _build(severity, sparse=False) + + +def build_nodal_balance_sparse(severity: int) -> linopy.Model: + """The same balance via the sparse groupby + frozen CSR path (v1-only).""" + previous = linopy.options["semantics"] + linopy.options["semantics"] = "v1" + try: + return _build(severity, sparse=True) + finally: + linopy.options["semantics"] = previous + + SPEC = register_pattern( BenchSpec( name="nodal_balance", @@ -70,3 +97,13 @@ def build_nodal_balance(severity: int) -> linopy.Model: axis="severity", ) ) + +SPARSE_SPEC = register_pattern( + BenchSpec( + name="nodal_balance_sparse", + build=build_nodal_balance_sparse, + sweep=SEVERITIES, + axis="severity", + phases=frozenset({BUILD, MATRICES, TO_LP}), + ) +) diff --git a/linopy/config.py b/linopy/config.py index 92cd09458..590a4cabc 100644 --- a/linopy/config.py +++ b/linopy/config.py @@ -82,4 +82,5 @@ def __repr__(self) -> str: display_max_rows=14, display_max_terms=6, semantics=LEGACY_SEMANTICS, + sparse_groupby=False, ) diff --git a/linopy/constraints.py b/linopy/constraints.py index d323a4de4..8a3c625b0 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -1144,7 +1144,7 @@ class Constraint(ConstraintBase): Supports setters, xarray operations via conwrap, and from_rule construction. """ - __slots__ = ("_data", "_model", "_assigned", "_coef_dirty") + __slots__ = ("_data", "_model", "_assigned", "_coef_dirty", "_pending") def __init__( self, @@ -1174,9 +1174,29 @@ def __init__( self._data = data self._model = model self._coef_dirty = False + self._pending = None + + @classmethod + def _from_pending( + cls, lhs: expressions.LinearExpression, sign: str, rhs: Any, model: Model + ) -> Constraint: + """Anonymous constraint over a still-sparse lhs (see linopy.csr).""" + obj = cls.__new__(cls) + obj._model = model + obj._data = None + obj._assigned = False + obj._coef_dirty = False + obj._pending = (lhs, sign, rhs) + return obj @property def data(self) -> Dataset: + if self._data is None and self._pending is not None: + lhs, sign, rhs = self._pending + lhs.data # noqa: B018 + self._data = lhs.to_constraint(sign, rhs).data + self._assigned = "labels" in self._data + self._pending = None return self._data @property diff --git a/linopy/csr.py b/linopy/csr.py new file mode 100644 index 000000000..fecd7e399 --- /dev/null +++ b/linopy/csr.py @@ -0,0 +1,353 @@ +""" +CSR-backed expressions: a sparse internal representation for grouped sums. + +``expr.groupby(g).sum(sparse=True)`` (or ``linopy.options["sparse_groupby"]`` +under v1) returns an ordinary :class:`~linopy.expressions.LinearExpression` +backed by a :class:`CSRPayload` instead of the dense dataset — same public +type, different backing, akin to dask-backed xarray objects. The CSR form is +canonical (duplicate variables summed, terms label-ordered) and ragged along +``_term``, so the group-size padding of issue #745 has no analog; grouping, +``merge``/``+``/``-`` and scaling become sparse linear algebra, and +``Model.add_constraints(..., freeze=True)`` staples sign and rhs on to form a +:class:`~linopy.constraints.CSRConstraint` directly. Anything without a +sparse branch expands through ``.data`` to the mathematically identical dense +rectangle in canonical term layout — the reason the feature is v1-gated, +where term layout is non-contractual. +""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING, Any + +import numpy as np +import pandas as pd +import scipy.sparse +from xarray import DataArray, Dataset + +from linopy.constants import HELPER_DIMS, TERM_DIM + +if TYPE_CHECKING: + from linopy.constraints import CSRConstraint + from linopy.expressions import LinearExpression + from linopy.model import Model + + +@dataclass(frozen=True) +class CSRPayload: + """ + An expression as ``A @ x + c`` over a fixed coordinate grid. + + ``csr`` has one row per flat grid cell (C order over ``grid_dims``) and + one column per raw variable label — label columns stay valid when + variables are added to the model later; realization maps them to dense + positions. ``const`` is the per-cell constant. + """ + + csr: scipy.sparse.csr_array + const: np.ndarray + grid_dims: tuple[str, ...] + indexes: dict[str, pd.Index] + model: Model + + @property + def shape(self) -> tuple[int, ...]: + return tuple(len(self.indexes[d]) for d in self.grid_dims) + + @property + def n_cells(self) -> int: + return self.csr.shape[0] + + @classmethod + def from_grouper( + cls, expr: LinearExpression, grouper: pd.Series, group_dim: str + ) -> CSRPayload: + """ + Build the grouped sum directly in CSR form (no padded rectangle). + + The grouper is conformed to the expression's member index by label + (upstream alignment checks guarantee equal label sets) and group + labels are sorted, matching the dense kernel's output grid. + """ + member_dim = str(grouper.index.name) + if member_dim in expr.data.indexes: + grouper = grouper.reindex(expr.data.indexes[member_dim]) + elif len(grouper) != expr.data.sizes[member_dim]: + raise ValueError(f"grouper length does not match dimension {member_dim!r}") + codes, uniques = pd.factorize(grouper, sort=True) + if (codes == -1).any(): + raise ValueError( + "Cannot group by a pandas object containing NaN values. " + "Drop or fill the corresponding entries before grouping." + ) + grid_dims = tuple( + group_dim if d == member_dim else str(d) for d in expr.coord_dims + ) + indexes = { + d: expr.data.get_index(d).rename(d) + for d in expr.coord_dims + if d != member_dim + } + indexes[group_dim] = pd.Index(uniques, name=group_dim) + return _assemble(expr, grid_dims, indexes, group_dim, member_dim, codes) + + @classmethod + def from_expression( + cls, expr: LinearExpression, template: CSRPayload + ) -> CSRPayload | None: + """Convert a dense expression on the template's grid, else None.""" + if set(expr.coord_dims) != set(template.grid_dims): + return None + for d in expr.coord_dims: + if not expr.data.get_index(d).equals(template.indexes[str(d)]): + return None + first = template.grid_dims[0] + codes = np.arange(len(template.indexes[first])) + return _assemble( + expr, template.grid_dims, template.indexes, first, first, codes + ) + + def scaled(self, factor: float) -> CSRPayload: + return replace(self, csr=self.csr * factor, const=self.const * factor) + + def same_grid(self, other: CSRPayload) -> bool: + return self.grid_dims == other.grid_dims and all( + self.indexes[d].equals(other.indexes[d]) for d in self.grid_dims + ) + + def add(self, other: CSRPayload) -> CSRPayload: + """Sparse matrix addition == merge along the term dimension.""" + width = max(self.csr.shape[1], other.csr.shape[1]) + csr = _widen(self.csr, width) + _widen(other.csr, width) + return replace(self, csr=csr, const=self.const + other.const) + + def materialize(self) -> LinearExpression: + """ + Expand to the dense rectangle in canonical form: terms label-ordered, + duplicates summed, padded to the widest cell with the usual fill. + """ + from linopy.expressions import LinearExpression + + csr = self.csr.copy() + csr.eliminate_zeros() + csr.sort_indices() + lengths = np.diff(csr.indptr) + nterm = max(int(lengths.max()) if len(lengths) else 0, 1) + + vars_flat = np.full( + (self.n_cells, nterm), -1, dtype=self.model._dtypes["labels"] + ) + coeffs_flat = np.full((self.n_cells, nterm), np.nan) + rows = np.repeat(np.arange(self.n_cells), lengths) + pos = np.arange(csr.nnz) - np.repeat(csr.indptr[:-1], lengths) + vars_flat[rows, pos] = csr.indices + coeffs_flat[rows, pos] = csr.data + + dims = (*self.grid_dims, TERM_DIM) + ds = Dataset( + { + "coeffs": (dims, coeffs_flat.reshape(*self.shape, nterm)), + "vars": (dims, vars_flat.reshape(*self.shape, nterm)), + "const": (self.grid_dims, self.const.reshape(self.shape)), + }, + coords={d: self.indexes[d] for d in self.grid_dims}, + ) + return LinearExpression(ds, self.model) + + +def _widen(csr: scipy.sparse.csr_array, width: int) -> scipy.sparse.csr_array: + if csr.shape[1] == width: + return csr + csr = csr.copy() + csr.resize((csr.shape[0], width)) + return csr + + +def _assemble( + expr: LinearExpression, + grid_dims: tuple[str, ...], + indexes: dict[str, pd.Index], + scatter_dim: str, + member_dim: str, + codes: np.ndarray, +) -> CSRPayload: + """ + Scatter an expression's terms into grid rows (conceptually ``G @ A``): + ``member_dim`` lands in the grid dim ``scatter_dim`` at row positions + ``codes``, every other grid dim maps one-to-one, and the COO→CSR + conversion sums duplicates — which is the group sum. The constant is + reduced with the dense kernel's skipna semantics. + """ + ds = expr.data + shape = tuple(len(indexes[d]) for d in grid_dims) + strides = [int(np.prod(shape[i + 1 :], dtype=np.int64)) for i in range(len(shape))] + + transposed = [member_dim if d == scatter_dim else d for d in grid_dims] + axis_positions = [ + codes * stride if d == scatter_dim else np.arange(n) * stride + for d, n, stride in zip(grid_dims, shape, strides) + ] + cell_rows = axis_positions[0] + for pos in axis_positions[1:]: + cell_rows = cell_rows[..., None] + pos + cell_rows = cell_rows.reshape(-1) + + coeffs = ds.coeffs.transpose(*transposed, TERM_DIM).to_numpy().reshape(-1) + vars_ = ds.vars.transpose(*transposed, TERM_DIM).to_numpy().reshape(-1) + rows = np.repeat(cell_rows, ds.sizes[TERM_DIM]) + keep = (vars_ != -1) & (coeffs != 0) & ~np.isnan(coeffs) + + full_size = int(np.prod(shape, dtype=np.int64)) if shape else 1 + coo = scipy.sparse.coo_array( + (coeffs[keep], (rows[keep], vars_[keep])), + shape=(full_size, expr.model._xCounter), + ) + + const_vals = ds.const.transpose(*transposed).to_numpy().reshape(-1) + const = np.zeros(full_size) + np.add.at(const, cell_rows, np.where(np.isnan(const_vals), 0.0, const_vals)) + + return CSRPayload( + scipy.sparse.csr_array(coo), const, grid_dims, indexes, expr.model + ) + + +def try_csr_merge( + exprs: Any, dim: str, join: Any, kwargs: dict +) -> LinearExpression | None: + """ + Sparse branch of :func:`linopy.expressions.merge`: combine plain + LinearExpressions on one shared grid (CSR-backed or dense-convertible), + where any join produces the identical result. Returns None to fall + through to the dense path. + """ + from linopy.expressions import LinearExpression + + if dim != TERM_DIM or kwargs: + return None + if not all(type(e) is LinearExpression for e in exprs): + return None + payloads = [e._csr for e in exprs if e._csr is not None] + if not payloads: + return None + template = payloads[0] + if not all(template.same_grid(p) for p in payloads[1:]): + return None + + combined: CSRPayload | None = None + for e in exprs: + payload = e._csr or CSRPayload.from_expression(e, template) + if payload is None: + return None + combined = payload if combined is None else combined.add(payload) + assert combined is not None + return LinearExpression._from_csr(combined, exprs[0].model) + + +def extract_pending( + lhs: Any, sign: Any, rhs: Any +) -> tuple[CSRPayload, str, Any] | None: + """Return (payload, sign, rhs) if lhs is a realizable CSR constraint.""" + from linopy.common import is_constant + from linopy.constraints import Constraint + from linopy.expressions import LinearExpression + + if ( + isinstance(lhs, Constraint) + and lhs._pending is not None + and sign is None + and rhs is None + ): + lhs, sign, rhs = lhs._pending + if not (isinstance(lhs, LinearExpression) and lhs._csr is not None): + return None + if not isinstance(sign, str) or rhs is None or not is_constant(rhs): + return None + rhs_da = _as_rhs_dataarray(rhs) + if rhs_da is None or not set(rhs_da.dims) <= set(lhs._csr.grid_dims): + return None + return lhs._csr, sign, rhs + + +def _as_rhs_dataarray(rhs: Any) -> DataArray | None: + from linopy.alignment import as_dataarray + + try: + da = as_dataarray(rhs) + except (TypeError, ValueError): + return None + return None if set(da.dims) & set(HELPER_DIMS) else da + + +def _rhs_grid_values(payload: CSRPayload, rhs: Any) -> np.ndarray: + """ + Align the rhs to the payload grid by label and flatten it, with v1 + parity: NaN in the rhs raises (§5) and a reordered or differing index + on a shared dim raises (§8), as on the dense path. + """ + from linopy.semantics import check_user_nan, is_v1 + + rhs_da = _as_rhs_dataarray(rhs) + assert rhs_da is not None + if is_v1(): + if bool(rhs_da.isnull().any()): + check_user_nan() + for d in rhs_da.dims: + if not rhs_da.get_index(d).equals(payload.indexes[str(d)]): + raise ValueError( + f"Coordinate mismatch on shared dimension {d!r} between " + "the rhs and the grouped result. Align the rhs with " + ".sel(...) / .reindex(...) before combining (§8)." + ) + rhs_da = rhs_da.reindex( + {d: payload.indexes[str(d)] for d in rhs_da.dims}, fill_value=np.nan + ) + missing = {d: payload.indexes[d] for d in payload.grid_dims if d not in rhs_da.dims} + if missing: + rhs_da = rhs_da.expand_dims(missing) + return rhs_da.transpose(*payload.grid_dims).to_numpy().reshape(-1) + + +def realize_csr_constraint( + model: Model, payload: CSRPayload, sign: str, rhs: Any, name: str +) -> CSRConstraint: + """ + Staple sign and rhs onto a CSR-backed lhs to form a CSRConstraint. + + Label columns are mapped to dense variable positions, the payload's + constant moves to the rhs, labels are allocated as in + ``Model._allocate_constraint_labels``, and rows without terms or with a + NaN rhs are inactive — all as on the frozen dense path. + """ + from linopy.common import maybe_replace_sign + from linopy.constraints import CSRConstraint + + sign = maybe_replace_sign(sign) + full_size = payload.n_cells + + label_index = model.variables.label_index + coo = payload.csr.tocoo() + csr = scipy.sparse.csr_array( + scipy.sparse.coo_array( + (coo.data, (coo.coords[0], label_index.label_to_pos[coo.coords[1]])), + shape=(full_size, label_index.n_active_vars), + ) + ) + csr.eliminate_zeros() + + rhs_flat = _rhs_grid_values(payload, rhs) - payload.const + + cindex = model._cCounter + model._cCounter += full_size + active = (np.diff(csr.indptr) > 0) & ~np.isnan(rhs_flat) + + return CSRConstraint( + csr[active], + np.arange(cindex, cindex + full_size)[active], + rhs_flat[active], + sign, + coords=[payload.indexes[d] for d in payload.grid_dims], + model=model, + name=name, + cindex=cindex, + ) diff --git a/linopy/expressions.py b/linopy/expressions.py index 9c957d69c..484087322 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -477,7 +477,10 @@ def map( return LinearExpression(self.groupby.map(func, args=args, **kwargs), self.model) def sum( - self, use_fallback: bool = False, observed: bool = False + self, + use_fallback: bool = False, + observed: bool = False, + sparse: bool | None = None, ) -> LinearExpression: """ Sum the expression over each group. @@ -491,6 +494,14 @@ def sum( Fall back to the previous, slower groupby-sum implementation, kept as an escape hatch. Leave at False unless the default misbehaves. Defaults to False. + sparse : bool, optional + Build the grouped sum in CSR form behind the ordinary + LinearExpression type — no group-size padding; a still-sparse + lhs reaching ``Model.add_constraints`` with ``freeze=True`` + becomes a CSRConstraint directly, other operations expand to + the dense rectangle in canonical term layout. Single-key + groupers only; requires v1 semantics. Defaults to + ``linopy.options["sparse_groupby"]``. See :mod:`linopy.csr`. observed : bool Only applies when grouping by a list of coordinate names. If True, keep the result stacked over the observed key combinations (a @@ -515,6 +526,24 @@ def sum( multikey_frame = ( None if use_fallback else _multikey_value_frame(group, self.data) ) + + if sparse is None: + sparse = is_v1() and options["sparse_groupby"] + elif sparse and not is_v1(): + raise ValueError( + "sparse groupby-sum requires v1 semantics; opt in with " + "linopy.options['semantics'] = 'v1'." + ) + if sparse and not use_fallback and not observed and multikey_frame is None: + series = group.to_pandas() if isinstance(group, DataArray) else group + if isinstance(series, pd.Series) and series.index.name in self.data.dims: + from linopy.csr import CSRPayload + + expr = LinearExpression(self.data, self.model) + group_name = str(getattr(series, "name", None) or "group") + payload = CSRPayload.from_grouper(expr, series, group_name) + return LinearExpression._from_csr(payload, self.model) + if multikey_frame is not None: group = multikey_frame @@ -711,7 +740,7 @@ def sum(self, **kwargs: Any) -> LinearExpression: class BaseExpression(ABC): - __slots__ = ("_data", "_model") + __slots__ = ("_data", "_model", "_csr") __array_ufunc__ = None __array_priority__ = 10000 __pandas_priority__ = 10000 @@ -785,6 +814,7 @@ def __init__(self, data: Dataset | Any | None, model: Model) -> None: self._model = model self._data = cast(Dataset, data) + self._csr = None def __repr__(self) -> str: """ @@ -885,6 +915,8 @@ def __neg__(self) -> Self: """ Get the negative of the expression. """ + if self._csr is not None: + return self._from_csr(self._csr.scaled(-1.0), self._model) return self.assign_multiindex_safe(coeffs=-self.coeffs, const=-self.const) def _multiply_by_linear_expression( @@ -1460,8 +1492,20 @@ def type(self) -> str: @property def data(self) -> Dataset: + if self._data is None and self._csr is not None: + self._data = self._csr.materialize().data + self._csr = None return self._data + @classmethod + def _from_csr(cls, payload: Any, model: Model) -> Self: + """Construct an expression backed by a CSRPayload.""" + obj = cls.__new__(cls) + obj._model = model + obj._data = None # type: ignore[assignment] + obj._csr = payload + return obj + @property def model(self) -> Model: return self._model @@ -1473,6 +1517,8 @@ def dims(self) -> tuple[Hashable, ...]: @property def coord_dims(self) -> tuple[Hashable, ...]: + if self._data is None and self._csr is not None: + return tuple(self._csr.grid_dims) return tuple(k for k in self.dims if k not in HELPER_DIMS) @property @@ -1661,6 +1707,9 @@ def to_constraint( which are moved to the left-hand-side and constant values which are moved to the right-hand side. """ + if self._csr is not None and isinstance(sign, str) and is_constant(rhs): + return constraints.Constraint._from_pending(self, sign, rhs, self.model) + rhs = as_constant(rhs) if self.is_constant and is_constant(rhs): raise ValueError( @@ -2298,6 +2347,8 @@ def __mul__( """ Multiply the expr by a factor. """ + if self._csr is not None and isinstance(other, int | float | np.number): + return type(self)._from_csr(self._csr.scaled(float(other)), self._model) other = as_constant(other) if isinstance(other, QuadraticExpression): return other.__rmul__(self) @@ -3069,6 +3120,13 @@ def merge( model = exprs[0].model + if issubclass(cls, LinearExpression) and not has_quad_expression: + from linopy.csr import try_csr_merge + + csr_result = try_csr_merge(exprs, dim=dim, join=join, kwargs=kwargs) + if csr_result is not None: + return csr_result + data = [e.data if isinstance(e, linopy_types) else e for e in exprs] data = [fill_missing_coords(ds, fill_helper_dims=True) for ds in data] diff --git a/linopy/model.py b/linopy/model.py index f15c086b3..ac39a5f9e 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -1138,6 +1138,17 @@ def add_constraints( """ name = self._resolve_constraint_name(name) + + resolved_freeze = self.freeze_constraints if freeze is None else freeze + if resolved_freeze and mask is None and not self.chunk: + from linopy.csr import extract_pending, realize_csr_constraint + + extracted = extract_pending(lhs, sign, rhs) + if extracted is not None: + payload, csr_sign, csr_rhs = extracted + con = realize_csr_constraint(self, payload, csr_sign, csr_rhs, name) + return self.constraints.add(con) + if sign is not None: sign = maybe_replace_signs(as_dataarray(sign)) diff --git a/test/test_sparse_groupby.py b/test/test_sparse_groupby.py new file mode 100644 index 000000000..beddef1d3 --- /dev/null +++ b/test/test_sparse_groupby.py @@ -0,0 +1,254 @@ +""" +Tests for sparse groupby-sum (linopy.csr): type stability, transparent +materialization, and direct CSR realization under freeze. v1-only feature. +""" + +from __future__ import annotations + +import re + +import numpy as np +import pandas as pd +import polars as pl +import pytest +import xarray as xr + +import linopy +from linopy import LinearExpression +from linopy.constraints import Constraint, CSRConstraint +from linopy.semantics import is_v1 +from linopy.testing import assert_conequal, assert_linequal + + +def require_v1() -> None: + if not is_v1(): + pytest.skip("sparse groupby-sum is gated behind v1 semantics") + + +def base_model(gens_per_bus=(7, 1, 3, 1, 2), n_snap=3, seed=0): + """Model with gen_p and flow on a ring; load ordered like the sorted groups (v1).""" + rng = np.random.default_rng(seed) + n_bus = len(gens_per_bus) + buses = pd.Index([f"bus{i}" for i in range(n_bus)], name="bus") + gen_bus = np.repeat(np.arange(n_bus), gens_per_bus) + gens = pd.Index([f"gen{i}" for i in range(len(gen_bus))], name="gen") + lines = pd.Index([f"line{i}" for i in range(n_bus)], name="line") + snaps = pd.Index(range(n_snap), name="snapshot") + + m = linopy.Model() + gen_p = m.add_variables(coords=[gens, snaps], name="gen_p") + flow = m.add_variables(coords=[lines, snaps], name="flow") + + gbus = pd.Series(buses[gen_bus], index=gens, name="bus") + bus0 = pd.Series(buses[np.arange(n_bus)], index=lines, name="bus") + bus1 = pd.Series(buses[(np.arange(n_bus) + 1) % n_bus], index=lines, name="bus") + load = xr.DataArray( + rng.uniform(1, 10, (n_bus, n_snap)), coords=[buses, snaps], name="load" + ).sortby("bus") + eff = xr.DataArray(rng.uniform(0.5, 1.5, len(gens)), coords=[gens]) + return m, gen_p, flow, eff, gbus, bus0, bus1, load, buses + + +def balance_lhs(gen_p, flow, eff, gbus, bus0, bus1, sparse): + return ( + (eff * gen_p).groupby(gbus).sum(sparse=sparse) + + (1.0 * flow).groupby(bus0).sum(sparse=sparse) + - (1.0 * flow).groupby(bus1).sum(sparse=sparse) + ) + + +def canon(df: pl.DataFrame) -> pl.DataFrame: + return ( + df.group_by(["labels", "vars"]) + .agg(pl.col("coeffs").sum(), pl.col("sign").first(), pl.col("rhs").first()) + .sort(["labels", "vars"]) + ) + + +def test_csr_requires_v1(): + m, gen_p, flow, eff, gbus, *_ = base_model() + if is_v1(): + res = (eff * gen_p).groupby(gbus).sum(sparse=True) + assert type(res) is LinearExpression + return + with pytest.raises(ValueError, match="requires v1 semantics"): + (eff * gen_p).groupby(gbus).sum(sparse=True) + linopy.options["sparse_groupby"] = True + try: + res = (eff * gen_p).groupby(gbus).sum() + finally: + linopy.options["sparse_groupby"] = False + assert res._csr is None + + +def test_csr_is_plain_linear_expression_and_materializes_equivalently(): + require_v1() + m, gen_p, flow, eff, gbus, *_ = base_model() + sparse = (eff * gen_p).groupby(gbus).sum(sparse=True) + eager = (eff * gen_p).groupby(gbus).sum() + assert type(sparse) is LinearExpression + assert_linequal(sparse, eager) + + +def test_csr_composition_materializes_equivalently(): + require_v1() + m, gen_p, flow, eff, gbus, bus0, bus1, *_ = base_model() + sparse = balance_lhs(gen_p, flow, eff, gbus, bus0, bus1, sparse=True) + eager = balance_lhs(gen_p, flow, eff, gbus, bus0, bus1, sparse=False) + assert type(sparse) is LinearExpression + assert_linequal(sparse, eager) + + +def test_scalar_ops_stay_csr(): + require_v1() + m, gen_p, flow, eff, gbus, *_ = base_model() + sparse = -2.0 * (eff * gen_p).groupby(gbus).sum(sparse=True) + assert sparse._csr is not None # still unmaterialized after neg/scale + eager = -2.0 * (eff * gen_p).groupby(gbus).sum() + from linopy.testing import _sort_by_vars_along_term + + a, b = _sort_by_vars_along_term(sparse), _sort_by_vars_along_term(eager) + assert (a.vars == b.vars).all() + xr.testing.assert_allclose( + a.coeffs.where(a.vars != -1), b.coeffs.where(b.vars != -1) + ) + xr.testing.assert_equal(a.const, b.const) + + +def test_freeze_realizes_csr_without_dense_rectangle(): + require_v1() + m1, *rest = base_model() + lhs1 = balance_lhs(*rest[:6], sparse=False) + c1 = m1.add_constraints(lhs1 == rest[6], name="balance") + + m2, *rest2 = base_model() + lhs2 = balance_lhs(*rest2[:6], sparse=True) + c2 = m2.add_constraints(lhs2 == rest2[6], name="balance", freeze=True) + + assert isinstance(c2, CSRConstraint) + d1, d2 = canon(c1.to_polars()), canon(c2.to_polars()) + assert d1["labels"].equals(d2["labels"]) + assert d1["vars"].equals(d2["vars"]) + assert np.allclose(d1["coeffs"], d2["coeffs"]) + assert (d1["sign"] == d2["sign"]).all() + assert np.allclose(d1["rhs"], d2["rhs"]) + assert np.array_equal( + np.sort(c1.labels.values.ravel()), np.sort(c2.active_labels()) + ) + + +def test_freeze_false_falls_back_to_identical_dense_constraint(): + """The fallback is canonical-form, so compare mathematically (strict=False).""" + require_v1() + m1, *rest = base_model() + c1 = m1.add_constraints( + balance_lhs(*rest[:6], sparse=False) == rest[6], name="balance" + ) + + m2, *rest2 = base_model() + c2 = m2.add_constraints( + balance_lhs(*rest2[:6], sparse=True) == rest2[6], name="balance" + ) + assert isinstance(c2, Constraint) + assert_conequal(c1, c2, strict=False) + assert np.array_equal(c1.labels.values, c2.labels.values) + + +def test_option_gates_csr_and_freeze_model_default(): + require_v1() + m, *rest = base_model() + m.freeze_constraints = True + linopy.options["sparse_groupby"] = True + try: + lhs = balance_lhs(*rest[:6], sparse=None) + con = m.add_constraints(lhs == rest[6], name="balance") + finally: + linopy.options["sparse_groupby"] = False + assert isinstance(con, CSRConstraint) + + +def test_materialized_csr_still_freezes_via_dense_path(): + require_v1() + m, *rest = base_model() + lhs = balance_lhs(*rest[:6], sparse=True) + _ = lhs.nterm # force materialization before constraint creation + con = m.add_constraints(lhs == rest[6], name="balance", freeze=True) + assert isinstance(con, CSRConstraint) + + +def test_nan_rhs_raises_on_both_paths(): + require_v1() + m1, *rest = base_model() + load = rest[6].copy() + load[0, 0] = np.nan + with pytest.raises(ValueError, match="NaN"): + m1.add_constraints( + balance_lhs(*rest[:6], sparse=False) == load, name="bal", freeze=True + ) + + m2, *rest2 = base_model() + load2 = rest2[6].copy() + load2[0, 0] = np.nan + with pytest.raises(ValueError, match="NaN"): + m2.add_constraints( + balance_lhs(*rest2[:6], sparse=True) == load2, name="bal", freeze=True + ) + + +def test_reordered_rhs_raises_on_both_paths(): + require_v1() + m1, *rest = base_model() + load = rest[6].isel(bus=slice(None, None, -1)) # same labels, reversed + with pytest.raises(ValueError, match="[Cc]oordinate"): + m1.add_constraints( + balance_lhs(*rest[:6], sparse=False) == load, name="bal", freeze=True + ) + + m2, *rest2 = base_model() + load2 = rest2[6].isel(bus=slice(None, None, -1)) + with pytest.raises(ValueError, match="[Cc]oordinate"): + m2.add_constraints( + balance_lhs(*rest2[:6], sparse=True) == load2, name="bal", freeze=True + ) + + +def test_nan_grouper_raises_eagerly(): + require_v1() + m, gen_p, flow, eff, gbus, *_ = base_model() + gbus = gbus.copy() + gbus.iloc[0] = np.nan + with pytest.raises(ValueError, match="NaN values"): + (1.0 * gen_p).groupby(gbus).sum(sparse=True) + + +def test_lp_files_identical(tmp_path): + require_v1() + sizes = (7, 1, 3, 1, 2, 1, 1, 4, 1, 2, 1, 1) + + m1, *rest = base_model(gens_per_bus=sizes) + m1.add_constraints(balance_lhs(*rest[:6], sparse=False) == rest[6], name="balance") + m1.add_objective((1.0 * rest[0]).sum()) + + m2, *rest2 = base_model(gens_per_bus=sizes) + m2.add_constraints( + balance_lhs(*rest2[:6], sparse=True) == rest2[6], name="balance", freeze=True + ) + m2.add_objective((1.0 * rest2[0]).sum()) + + term_line = re.compile(r"^[+-][0-9.e+-]+ x[0-9]+$") + + def canon_lp(text: str) -> list[str]: + out: list[str] = [] + buf: list[str] = [] + for line in text.splitlines(): + if term_line.match(line): + buf.append(line) + else: + out += sorted(buf) + [line] + buf = [] + return out + sorted(buf) + + f1, f2 = tmp_path / "eager.lp", tmp_path / "sparse.lp" + m1.to_file(f1) + m2.to_file(f2) + assert canon_lp(f1.read_text()) == canon_lp(f2.read_text())