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 7e12819eb..43b80a629 100644 --- a/linopy/config.py +++ b/linopy/config.py @@ -84,4 +84,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 0c887db9a..eebd04b72 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -31,7 +31,7 @@ from xarray.core.utils import Frozen from linopy import expressions, variables -from linopy.alignment import broadcast_to_coords +from linopy.alignment import as_dataarray, broadcast_to_coords from linopy.common import ( ConstraintLabelIndex, LabelPositionIndex, @@ -55,8 +55,10 @@ get_label_position, get_printout_labels, has_optimized_model, + is_constant, iterate_slices, maybe_group_terms_polars, + maybe_replace_sign, maybe_replace_signs, replace_by_map, save_join, @@ -74,6 +76,7 @@ SIGNS_pretty, ) from linopy.scaling import ensure_scaling, validate_scaling +from linopy.semantics import check_user_nan from linopy.types import ( ConstantLike, ConstraintLike, @@ -85,6 +88,7 @@ if TYPE_CHECKING: from linopy.model import Model + from linopy.sparse_expression import CSRPayload FILL_VALUE = { @@ -1193,6 +1197,109 @@ def from_mutable( scaling=scaling, ) + @classmethod + def from_payload( + cls, + model: Model, + payload: CSRPayload, + sign: str, + rhs: Any, + name: str, + ) -> CSRConstraint: + """ + Staple sign and rhs onto a CSR-backed lhs to form a CSRConstraint. + + The sparse counterpart of :meth:`from_mutable`: instead of converting a + dense :class:`Constraint`, it realizes a + :class:`~linopy.sparse_expression.CSRPayload` directly. 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. + """ + 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), + ) + ) + has_terms = np.diff(csr.indptr) > 0 + csr.eliminate_zeros() + + rhs_flat = _rhs_grid_values(payload, rhs) - payload.const + + cindex = model._cCounter + model._cCounter += full_size + active = has_terms & ~np.isnan(rhs_flat) + + return cls( + 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, + ) + + +def extract_csr_pending( + lhs: Any, sign: Any, rhs: Any +) -> tuple[CSRPayload, str, Any] | None: + """Return (payload, sign, rhs) if lhs is a realizable CSR constraint.""" + 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, expressions.LinearExpression) and lhs._payload 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._payload.grid_dims): + return None + return lhs._payload, sign, rhs + + +def _as_rhs_dataarray(rhs: Any) -> DataArray | None: + 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: + """ + Broadcast the rhs onto the payload grid 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. + """ + rhs_da = _as_rhs_dataarray(rhs) + assert rhs_da is not None + 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)." + ) + 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) + class Constraint(ConstraintBase): """ @@ -1201,7 +1308,9 @@ 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") + + _pending: tuple[expressions.LinearExpression, str, Any] | None def __init__( self, @@ -1232,9 +1341,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.sparse_expression).""" + obj = cls.__new__(cls) + obj._model = model + obj._data = None # type: ignore[assignment] + 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 + dense = expressions.LinearExpression(lhs.data, lhs.model) + self._data = dense.to_constraint(sign, rhs).data + self._assigned = "labels" in self._data + self._pending = None return self._data @property diff --git a/linopy/expressions.py b/linopy/expressions.py index f4b938236..de42536aa 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -537,7 +537,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. @@ -551,6 +554,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.sparse_expression`. observed : bool Only applies when grouping by a list of coordinate names. If True, keep the result stacked over the observed key combinations (a @@ -575,6 +586,38 @@ def sum( multikey_frame = ( None if use_fallback else _multikey_value_frame(group, self.data) ) + + explicit_sparse = sparse is True + 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: + series = group.to_pandas() if isinstance(group, DataArray) else group + supported = ( + not use_fallback + and not observed + and multikey_frame is None + and isinstance(series, pd.Series) + and series.index.name in self.data.dims + ) + if supported: + from linopy.sparse_expression import CSRPayload + + expr = LinearExpression(self.data, self.model) + group_name = str(series.name or "group") + payload = CSRPayload.from_grouper(expr, series, group_name) + return LinearExpression._from_payload(payload, self.model) + if explicit_sparse: + raise ValueError( + "sparse=True supports only a single-key grouper (pandas " + "Series or 1-D DataArray) over an existing dimension, " + "without use_fallback or observed." + ) + if multikey_frame is not None: group = multikey_frame @@ -771,7 +814,7 @@ def sum(self, **kwargs: Any) -> LinearExpression: class BaseExpression(ABC): - __slots__ = ("_data", "_model") + __slots__ = ("_data", "_model", "_payload") __array_ufunc__ = None __array_priority__ = 10000 __pandas_priority__ = 10000 @@ -846,6 +889,7 @@ def __init__(self, data: Dataset | Any | None, model: Model) -> None: data = data.assign_attrs(name=None) self._model = model self._data = cast(Dataset, data) + self._payload = None def __repr__(self) -> str: """ @@ -946,6 +990,8 @@ def __neg__(self) -> Self: """ Get the negative of the expression. """ + if self._payload is not None: + return self._from_payload(self._payload.scaled(-1.0), self._model) return self.assign_multiindex_safe(coeffs=-self.coeffs, const=-self.const) def _multiply_by_linear_expression( @@ -1531,8 +1577,20 @@ def name(self) -> str: @property def data(self) -> Dataset: + if self._data is None and self._payload is not None: + self._data = self._payload.materialize().data + self._payload = None return self._data + @classmethod + def _from_payload(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._payload = payload + return obj + @property def model(self) -> Model: return self._model @@ -1544,6 +1602,8 @@ def dims(self) -> tuple[Hashable, ...]: @property def coord_dims(self) -> tuple[Hashable, ...]: + if self._data is None and self._payload is not None: + return tuple(self._payload.grid_dims) return tuple(k for k in self.dims if k not in HELPER_DIMS) @property @@ -1740,6 +1800,9 @@ def to_constraint( Legacy instead keeps a NaN RHS as that auto-mask, restoring the mask after the subtraction filled it with 0. """ + if self._payload 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( @@ -2344,6 +2407,10 @@ def __mul__( """ Multiply the expr by a factor. """ + if self._payload is not None and isinstance(other, int | float | np.number): + return type(self)._from_payload( + self._payload.scaled(float(other)), self._model + ) other = as_constant(other) if isinstance(other, QuadraticExpression): return other.__rmul__(self) @@ -3138,6 +3205,13 @@ def merge( model = exprs[0].model + if issubclass(cls, LinearExpression) and not has_quad_expression: + from linopy.sparse_expression 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 769ec1a50..600051e5f 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -1200,6 +1200,23 @@ def add_constraints( scaling: ConstantLike = ..., ) -> CSRConstraint: ... + @overload + def add_constraints( + self, + lhs: VariableLike + | ExpressionLike + | ConstraintLike + | Sequence[tuple[ConstantLike, VariableLike | str]] + | Callable, + sign: SignLike | None = ..., + rhs: ConstantLike | VariableLike | ExpressionLike | None = ..., + name: str | None = ..., + coords: Sequence[Sequence | pd.Index] | Mapping | None = ..., + mask: MaskLike | None = ..., + freeze: bool | None = ..., + scaling: ConstantLike = ..., + ) -> ConstraintBase: ... + def add_constraints( self, lhs: VariableLike @@ -1264,6 +1281,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.constraints import CSRConstraint, extract_csr_pending + + extracted = extract_csr_pending(lhs, sign, rhs) + if extracted is not None: + payload, csr_sign, csr_rhs = extracted + con = CSRConstraint.from_payload(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/linopy/sparse_expression.py b/linopy/sparse_expression.py new file mode 100644 index 000000000..3c6acb5d7 --- /dev/null +++ b/linopy/sparse_expression.py @@ -0,0 +1,248 @@ +""" +The sparse payload behind a LinearExpression: ``A @ x + c`` in CSR form. + +``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. 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. + +This module covers the expression layer only. Stapling sign and rhs onto a +payload to form a :class:`~linopy.constraints.CSRConstraint` lives in +:meth:`linopy.constraints.CSRConstraint.from_payload`. +""" + +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 Dataset + +from linopy.constants import TERM_DIM + +if TYPE_CHECKING: + 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: dict[str, pd.Index] = { + str(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 cls._from_scatter(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 cls._from_scatter( + expr, template.grid_dims, template.indexes, first, first, codes + ) + + @classmethod + def _from_scatter( + cls, + 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) & ~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 cls(scipy.sparse.csr_array(coo), const, grid_dims, indexes, expr.model) + + 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. 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, as on the dense path. + """ + 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]) + coo = scipy.sparse.coo_array((data, (rows, cols)), shape=shape) + return replace( + self, csr=scipy.sparse.csr_array(coo), 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.sort_indices() + lengths = np.diff(csr.indptr) + nterm = max(int(lengths.max(initial=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 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._payload for e in exprs if e._payload 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._payload 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_payload(combined, exprs[0].model) diff --git a/test/test_sparse_groupby.py b/test/test_sparse_groupby.py new file mode 100644 index 000000000..1691d3e4c --- /dev/null +++ b/test/test_sparse_groupby.py @@ -0,0 +1,259 @@ +""" +Tests for sparse groupby-sum (linopy.sparse_expression): type stability, transparent +materialization, and direct CSR realization under freeze. v1-only feature. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import pandas as pd +import polars as pl +import pytest +import xarray as xr + +import linopy +from linopy import LinearExpression, Model, Variable +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") + + +@dataclass +class Case: + """Model with gen_p and flow on a ring; load ordered like the sorted groups (v1).""" + + m: Model + gen_p: Variable + flow: Variable + eff: xr.DataArray + gbus: pd.Series + bus0: pd.Series + bus1: pd.Series + load: xr.DataArray + + def balance_lhs(self, sparse: bool | None) -> LinearExpression: + return ( + (self.eff * self.gen_p).groupby(self.gbus).sum(sparse=sparse) + + (1.0 * self.flow).groupby(self.bus0).sum(sparse=sparse) + - (1.0 * self.flow).groupby(self.bus1).sum(sparse=sparse) + ) + + +def base_model( + gens_per_bus: tuple[int, ...] = (7, 1, 3, 1, 2), n_snap: int = 3, seed: int = 0 +) -> Case: + 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 Case(m, gen_p, flow, eff, gbus, bus0, bus1, load) + + +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() -> None: + c = base_model() + if is_v1(): + res = (c.eff * c.gen_p).groupby(c.gbus).sum(sparse=True) + assert type(res) is LinearExpression + return + with pytest.raises(ValueError, match="requires v1 semantics"): + (c.eff * c.gen_p).groupby(c.gbus).sum(sparse=True) + linopy.options["sparse_groupby"] = True + try: + res = (c.eff * c.gen_p).groupby(c.gbus).sum() + finally: + linopy.options["sparse_groupby"] = False + assert res._payload is None + + +def test_csr_is_plain_linear_expression_and_materializes_equivalently() -> None: + require_v1() + c = base_model() + sparse = (c.eff * c.gen_p).groupby(c.gbus).sum(sparse=True) + eager = (c.eff * c.gen_p).groupby(c.gbus).sum() + assert type(sparse) is LinearExpression + assert_linequal(sparse, eager) + + +def test_csr_composition_materializes_equivalently() -> None: + require_v1() + c = base_model() + sparse = c.balance_lhs(sparse=True) + assert type(sparse) is LinearExpression + assert_linequal(sparse, c.balance_lhs(sparse=False)) + + +def test_scalar_ops_stay_csr() -> None: + require_v1() + c = base_model() + sparse = -2.0 * (c.eff * c.gen_p).groupby(c.gbus).sum(sparse=True) + assert sparse._payload is not None + assert_linequal(sparse, -2.0 * (c.eff * c.gen_p).groupby(c.gbus).sum()) + + +@pytest.mark.parametrize("sparse", [True, False], ids=["sparse", "dense"]) +def test_zero_coefficient_rows_stay_active(sparse: bool) -> None: + require_v1() + c = base_model() + lhs = (0.0 * c.gen_p).groupby(c.gbus).sum(sparse=sparse) + lhs = lhs + (0.0 * c.flow).groupby(c.bus0).sum(sparse=sparse) + con = c.m.add_constraints(lhs == c.load, name="bal", freeze=True) + assert len(con.active_labels()) == c.load.size + + +@pytest.mark.parametrize( + "grouper, kwargs", + [ + (["gen", "snapshot"], {}), + ("gen", {"use_fallback": True}), + (["gen"], {"observed": True}), + ], +) +def test_explicit_sparse_raises_on_unsupported_grouper( + grouper: str | list[str], kwargs: dict +) -> None: + require_v1() + c = base_model() + with pytest.raises(ValueError, match="single-key grouper"): + (1.0 * c.gen_p).groupby(grouper).sum(sparse=True, **kwargs) + + +def test_freeze_realizes_csr_without_dense_rectangle() -> None: + require_v1() + c1, c2 = base_model(), base_model() + con1 = c1.m.add_constraints(c1.balance_lhs(sparse=False) == c1.load, name="bal") + con2 = c2.m.add_constraints( + c2.balance_lhs(sparse=True) == c2.load, name="bal", freeze=True + ) + + assert isinstance(con2, CSRConstraint) + d1, d2 = canon(con1.to_polars()), canon(con2.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(con1.labels.values.ravel()), np.sort(con2.active_labels()) + ) + + +def test_freeze_false_falls_back_to_identical_dense_constraint() -> None: + """The fallback is canonical-form, so compare mathematically (strict=False).""" + require_v1() + c1, c2 = base_model(), base_model() + con1 = c1.m.add_constraints(c1.balance_lhs(sparse=False) == c1.load, name="bal") + con2 = c2.m.add_constraints(c2.balance_lhs(sparse=True) == c2.load, name="bal") + assert isinstance(con2, Constraint) + assert_conequal(con1, con2, strict=False) + assert np.array_equal(con1.labels.values, con2.labels.values) + + +def test_option_gates_csr_and_freeze_model_default() -> None: + require_v1() + c = base_model() + c.m.freeze_constraints = True + linopy.options["sparse_groupby"] = True + try: + con = c.m.add_constraints(c.balance_lhs(sparse=None) == c.load, name="bal") + finally: + linopy.options["sparse_groupby"] = False + assert isinstance(con, CSRConstraint) + + +def test_materialized_csr_still_freezes_via_dense_path() -> None: + require_v1() + c = base_model() + lhs = c.balance_lhs(sparse=True) + _ = lhs.nterm + con = c.m.add_constraints(lhs == c.load, name="bal", freeze=True) + assert isinstance(con, CSRConstraint) + + +@pytest.mark.parametrize("sparse", [True, False], ids=["sparse", "dense"]) +def test_nan_rhs_raises(sparse: bool) -> None: + require_v1() + c = base_model() + load = c.load.copy() + load[0, 0] = np.nan + with pytest.raises(ValueError, match="NaN"): + c.m.add_constraints(c.balance_lhs(sparse) == load, name="bal", freeze=True) + + +@pytest.mark.parametrize("sparse", [True, False], ids=["sparse", "dense"]) +def test_reordered_rhs_raises(sparse: bool) -> None: + require_v1() + c = base_model() + load = c.load.isel(bus=slice(None, None, -1)) + with pytest.raises(ValueError, match="[Cc]oordinate"): + c.m.add_constraints(c.balance_lhs(sparse) == load, name="bal", freeze=True) + + +def test_nan_grouper_raises_eagerly() -> None: + require_v1() + c = base_model() + gbus = c.gbus.copy() + gbus.iloc[0] = np.nan + with pytest.raises(ValueError, match="NaN values"): + (1.0 * c.gen_p).groupby(gbus).sum(sparse=True) + + +def test_lp_files_identical(tmp_path: Path) -> None: + require_v1() + sizes = (7, 1, 3, 1, 2, 1, 1, 4, 1, 2, 1, 1) + c1, c2 = base_model(gens_per_bus=sizes), base_model(gens_per_bus=sizes) + c1.m.add_constraints(c1.balance_lhs(sparse=False) == c1.load, name="bal") + c1.m.add_objective((1.0 * c1.gen_p).sum()) + c2.m.add_constraints( + c2.balance_lhs(sparse=True) == c2.load, name="bal", freeze=True + ) + c2.m.add_objective((1.0 * c2.gen_p).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" + c1.m.to_file(f1) + c2.m.to_file(f2) + assert canon_lp(f1.read_text()) == canon_lp(f2.read_text())