diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 7c3506333..53541eff0 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -44,6 +44,7 @@ Upcoming Version *Internal* * The sparse backing of a ``LinearExpression`` moved from ``linopy.sparse_expression`` to ``linopy.csr`` and the class ``CSRExpression`` was renamed to ``CSRLinearExpression``. Dense/sparse conversion is now spelled the same way on both CSR types: ``CSRLinearExpression.from_dense`` / ``.to_dense`` and ``CSRConstraint.from_dense`` (previously ``CSRConstraint.from_mutable``) / ``.to_dense``. ``Constraint.freeze()`` and ``CSRConstraint.mutable()`` are unchanged. +* The first constructor argument of ``LinearExpressionGroupby`` is now named ``obj`` (previously ``data``), as it holds a ``Dataset`` or an expression; ``LinearExpressionGroupby(data=...)`` must be spelled ``LinearExpressionGroupby(obj=...)``. *Documentation* @@ -54,6 +55,12 @@ Upcoming Version * ``@``/``dot`` against a constant matrix that holds zeros no longer densifies the result to one term per contracted member. The zero-coefficient terms are dropped, so the term dimension shrinks to the widest non-zero cell. On PyPSA's Kirchhoff Voltage Law constraint (a cycle matrix with ~3 branches per cycle) this cuts the expression from 852 to 3 terms — 284x fewer cells — which in turn shrinks the downstream ``merge``. A constant without zeros is unaffected. (`#748 `__) * ``densify_terms`` (used by ``sum(drop_zeros=True)`` and the sparse ``@`` path) is now fully vectorised. It previously counted the non-zero positions with a Python loop that scaled quadratically in the number of non-zero terms — 127 s for a (2000 x 60) expression, now 3 ms — and allocated the compacted output at the full original term width. It now allocates only the compacted width and returns the expression unchanged when it holds no zeros. * Under v1, ``@``/``dot`` against a constant now runs as sparse linear algebra instead of building the dense broadcast intermediate (``self * other`` then ``.sum()``): peak memory scales with ``nnz(C) x nterm`` rather than with the full broadcast shape. A ``CSRLinearExpression``-backed operand (from ``groupby(...).sum(sparse=True)``) stays CSR-backed through ``@``, and the result is the compact canonical form (duplicate variables summed, terms label-ordered, explicit zeros pruned — cell activeness is carried by ``const`` alone). (`#748 `__, `#756 `__, `#925 `__) +* Reading metadata of a CSR-backed ``LinearExpression`` (``repr()``, ``shape``, ``sizes``, ``dims``, ``coords``, ``indexes``, ``isnull()``) no longer converts it to dense; it is served from the sparse backing, auxiliary coordinates included. (`#962 `__) +* New ``LinearExpression.is_sparse`` tells whether an expression is CSR-backed, and its repr header reads ``LinearExpression (sparse)``. The opt-in option ``linopy.options["warn_on_densify"]`` (default ``False``) emits a ``PerformanceWarning`` naming the reason whenever a sparse backing is dropped, e.g. on ``.data``, an operation without a sparse path, or ``CSRConstraint.mutable()``. (`#969 `__) +* Multiplying, dividing, adding or subtracting a non-scalar constant (``DataArray``, ``Series``, ``ndarray``, …) keeps a CSR-backed ``LinearExpression`` sparse: the coefficients are scaled per cell and only the constant changes, instead of expanding the dense rectangle. The constant is aligned exactly as on the dense path (NaN, label mismatch, joins and auxiliary coordinates); an operand that adds dimensions still densifies. ``expr.const`` is now served from the sparse backing as well. (`#965 `__) +* ``sum(dim=...)`` and a further ``groupby(...).sum()`` keep a CSR-backed ``LinearExpression`` sparse: summed rows are merged in the sparse backing instead of stacking the summed dimensions into the term dimension of the dense rectangle. Auxiliary coordinates, absent cells and the constant follow the dense path; ``groupby(...).sum(sparse=False)`` still densifies. (`#964 `__) +* ``where``, ``sel``, ``isel``, ``loc`` and ``[]`` keep a CSR-backed ``LinearExpression`` sparse: the selection is evaluated on the grid's row numbers with the dense (xarray) semantics, including ``drop=``, ``method=``, scalar coordinates left by a scalar selection and auxiliary coordinates, and the selected rows are gathered from the sparse backing; masked cells become absent. A condition or indexer that introduces dimensions still densifies. A merge of differing grids now stays sparse when operands carry auxiliary coordinates, which are checked for conflicts as on the dense path. (`#966 `__) +* ``Model.add_constraints(..., mask=..., freeze=True)`` with a CSR-backed left-hand side applies the mask sparsely and returns a ``CSRConstraint`` instead of expanding the dense rectangle. (`#970 `__) * Persistent snapshots of tz-aware ``DatetimeIndex`` coordinates no longer materialise an object array of ``Timestamp`` per container per capture and diff. Coordinates are stored as UTC-ns arrays with the timezone identity carried alongside, making snapshot capture ~24x and warm-start diffs ~33x faster on tz-aware models, while naive and tz-aware coordinates — and differing timezones — stay correctly unequal. (`#960 `__) **Bug fixes** @@ -67,6 +74,7 @@ Upcoming Version * ``linopy.read_netcdf`` now restores variables, expressions and constraints in their original insertion order instead of alphabetically, so ``matrices.A`` of a round-tripped model is no longer a row/column permutation of the original. The order is stored in the file; files written by earlier versions still load in sorted order. ``linopy.testing.assert_model_equal`` now also compares container order. (`#934 `__) * Adding or removing variables after a constraint was added with ``freeze=True`` no longer breaks ``model.matrices``. The frozen constraint stored dense variable positions as its matrix columns, so blocks frozen at different times disagreed on their width and stacking them raised ``ValueError: inconsistent shapes``. Raw variable labels are stored instead and mapped to positions when the matrix is assembled. Frozen constraints in netCDF files written by earlier versions are read as before. (`#926 `__) +* ``linopy.merge`` with ``join="outer"``, ``"left"`` or ``"right"`` no longer raises an xarray ``MergeError`` when only one operand carries an auxiliary coordinate on a joined dimension. The constant was reindexed with a fill value of ``0`` that was also written into the auxiliary coordinate, so it no longer matched the coefficients' copy. The coordinate now follows its rows onto the joined grid and is NaN where no operand provides it, under both semantics. The same holds for ``.add`` / ``.sub`` / ``.mul`` / ``.div`` with a constant and a reindexing ``join=``. * A frozen constraint caches the label-to-position mapping of its matrix columns by weak reference and only rebuilds it when the constraint or the set of variables changes. While a persistent snapshot holds the arrays, repeated matrix assembly on an unchanged model returns the same objects, so the snapshot diff can again skip the comparison of untouched frozen constraints by object identity; one-off exports such as ``to_file`` retain no extra memory. (`#933 `__) * ``Solver.close()`` no longer leaves dangling native handles behind. The solver model is now dropped before the environment that owns it, instead of after. And the COPT and MindOpt file interfaces no longer hand back a model they already disposed: after a file-based COPT or MindOpt solve, ``model.solver_model`` is ``None`` rather than a handle into freed memory. (`#899 `__) diff --git a/linopy/alignment.py b/linopy/alignment.py index c566a21c5..df85bd081 100644 --- a/linopy/alignment.py +++ b/linopy/alignment.py @@ -77,6 +77,7 @@ def as_constant(other: Any) -> Any: if TYPE_CHECKING: + from linopy.csr import AuxCoords from linopy.expressions import LinearExpression, QuadraticExpression from linopy.variables import Variable @@ -696,7 +697,7 @@ def _matmul_operand_to_matrix( contracted: Sequence[str], new_dims: Sequence[str], indexes: Mapping[str, pd.Index], - aux_coords: Mapping[str, tuple[str, np.ndarray]], + aux_coords: AuxCoords, ) -> scipy.sparse.csr_array: """ Flatten a ``@`` constant to the ``(contracted, new)`` matrix of the sparse diff --git a/linopy/config.py b/linopy/config.py index 43b80a629..53541c829 100644 --- a/linopy/config.py +++ b/linopy/config.py @@ -85,4 +85,5 @@ def __repr__(self) -> str: display_max_terms=6, semantics=LEGACY_SEMANTICS, sparse_groupby=False, + warn_on_densify=False, ) diff --git a/linopy/constraints.py b/linopy/constraints.py index c2172fedc..fd88ab577 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -79,7 +79,7 @@ PerformanceWarning, SIGNS_pretty, ) -from linopy.csr import Grid, csr_nterm, csr_to_term_arrays +from linopy.csr import Grid, _densify_notice, csr_nterm, csr_to_term_arrays from linopy.scaling import ensure_scaling, validate_scaling from linopy.semantics import check_user_nan from linopy.types import ( @@ -1262,6 +1262,7 @@ def freeze(self) -> CSRConstraint: def to_dense(self) -> Constraint: """Convert to a Constraint.""" + _densify_notice("frozen constraint converted by `to_dense()`/`mutable()`") return Constraint(self.data, self._model, self._name) def mutable(self) -> Constraint: @@ -1419,25 +1420,26 @@ def from_csr( active, rhs_flat[active], sign, - grid=expr.grid, + grid=Grid(expr.grid.indexes), model=expr.model, ) -def csr_rhs(expr: CSRLinearExpression, rhs: Any) -> DataArray | None: +def csr_rhs(expr: CSRLinearExpression, rhs: Any) -> DataArray | str: """ - Return ``rhs`` as a DataArray on the expression grid, or None if the sparse - path cannot take it: a non-constant rhs, one that is no DataArray-like, or - one with helper dims or dims outside the grid falls back to the dense path. + Return ``rhs`` as a DataArray on the expression grid, or the reason the + sparse path cannot take it: a non-constant rhs, one that is no + DataArray-like, or one with helper dims or dims outside the grid falls + back to the dense path. """ if not is_constant(rhs): - return None + return "constraint with a non-constant rhs" try: da = as_dataarray(rhs) except (TypeError, ValueError): - return None + return "constraint with an rhs that is not array-like" if set(da.dims) & set(HELPER_DIMS) or not set(da.dims) <= set(expr.grid.dims): - return None + return "constraint with an rhs over dimensions outside the grid" return da diff --git a/linopy/csr.py b/linopy/csr.py index 90d4a8817..1a24b2dad 100644 --- a/linopy/csr.py +++ b/linopy/csr.py @@ -7,10 +7,13 @@ 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``/``+``/``-``, scaling and ``@``/``dot`` (:meth:`contracted`) -become sparse linear algebra. Unlike :meth:`added`, which preserves explicit -zeros through COO, ``@``/``dot`` prunes them: cell activeness is carried by -``const`` alone (issue #925). +grouping, ``sum``, ``merge``/``+``/``-``, scaling and ``@``/``dot`` +(:meth:`contracted`) become sparse linear algebra. Zero policy: the structural +operations (grouping and ``sum`` via :meth:`aggregated`, merge via +:meth:`added`, scaling, reindexing) go through COO and keep explicit zero +coefficients, like the dense path; only the product with a constant matrix, +``@``/``dot``, prunes them. Either way cell activeness is carried by ``const`` +alone (issue #925), so the two differ in term layout only. 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. @@ -25,17 +28,23 @@ from __future__ import annotations -from collections.abc import Iterable, Mapping +import operator +from collections.abc import Callable, Iterable, Mapping from dataclasses import dataclass, field, replace -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, TypeAlias import numpy as np import pandas as pd import scipy.sparse -from xarray import Dataset +from xarray import DataArray, Dataset -from linopy.constants import HELPER_DIMS, TERM_DIM -from linopy.semantics import absorb_absence, enforce_aux_conflict +from linopy.config import options +from linopy.constants import HELPER_DIMS, TERM_DIM, PerformanceWarning +from linopy.semantics import ( + absorb_absence, + enforce_aux_conflict, + warn_outside_linopy, +) if TYPE_CHECKING: from linopy.expressions import LinearExpression @@ -44,6 +53,9 @@ CONTRACTION_CHUNK = 64 """Kept-axis block size of the chunked Kronecker product in ``contracted``.""" +AuxCoords: TypeAlias = dict[str, tuple[str | tuple[()], np.ndarray]] +"""Auxiliary coordinates as ``name -> (grid dim, values)``, dim ``()`` for a scalar.""" + @dataclass(frozen=True, eq=False) class Grid: @@ -52,10 +64,14 @@ class Grid: ``indexes`` maps each grid dimension to its labels, ordered as the dimensions are, so row ``i`` of a CSR matrix is cell ``i`` of the - C-order flattening of that grid. + C-order flattening of that grid. ``aux`` holds the auxiliary + coordinates as ``name -> (grid dim, values)``, e.g. the key levels of a + grouped result kept stacked over the observed key combinations; a + scalar coordinate, e.g. left by a scalar selection, has dim ``()``. """ indexes: dict[str, pd.Index] + aux: AuxCoords = field(default_factory=dict) @classmethod def from_coords(cls, coords: Iterable[pd.Index]) -> Grid: @@ -63,9 +79,11 @@ def from_coords(cls, coords: Iterable[pd.Index]) -> Grid: return cls({str(c.name): c for c in coords}) @classmethod - def from_dataset(cls, ds: Dataset, dims: Iterable[str]) -> Grid: - """Build from the indexes ``ds`` carries on ``dims``.""" - return cls({d: ds.get_index(d).rename(d) for d in dims}) + def from_dataset(cls, ds: Dataset | DataArray, dims: Iterable[str]) -> Grid: + """Build from the indexes and auxiliary coordinates ``ds`` carries on ``dims``.""" + dims = tuple(dims) + indexes = {d: ds.get_index(d).rename(d) for d in dims} + return cls(indexes, _aux_coords(ds, set(dims))) @property def dims(self) -> tuple[str, ...]: @@ -79,6 +97,10 @@ def coords(self) -> list[pd.Index]: def shape(self) -> tuple[int, ...]: return tuple(len(i) for i in self.indexes.values()) + def to_dataset(self) -> Dataset: + """The grid's indexes and auxiliary coordinates as a coordinate-only Dataset.""" + return Dataset(coords=self.indexes | self.aux) + @property def size(self) -> int: """Number of flat cells; one for the zero-dimensional grid.""" @@ -114,31 +136,53 @@ def indexer(self, other: Grid) -> tuple[np.ndarray, np.ndarray]: def renamed(self, names: Mapping[str, str]) -> Grid: """Relabel dimensions; the cell layout is unchanged.""" - return Grid( - { - names.get(d, d): i.rename(names.get(d, d)) - for d, i in self.indexes.items() - } - ) + indexes = { + names.get(d, d): i.rename(names.get(d, d)) for d, i in self.indexes.items() + } + aux = { + n: (names.get(d, d) if isinstance(d, str) else d, v) + for n, (d, v) in self.aux.items() + } + return Grid(indexes, aux) def reordered(self, dims: Iterable[str]) -> Grid: - """Select and order the given dimensions; labels unchanged.""" - return Grid({d: self.indexes[d] for d in dims}) + """ + Select and order the given dimensions; labels unchanged, auxiliary + coordinates on dropped dimensions dropped, scalar ones kept. + """ + dims = tuple(dims) + aux = { + n: (d, v) + for n, (d, v) in self.aux.items() + if not isinstance(d, str) or d in dims + } + return Grid({d: self.indexes[d] for d in dims}, aux) def with_indexes(self, indexers: Mapping[Any, Any]) -> Grid: - """Replace the labels of the named dimensions; the rest unchanged.""" - return Grid( - { - d: pd.Index(indexers[d], name=d) if d in indexers else i - for d, i in self.indexes.items() - } - ) + """ + Replace the labels of the named dimensions; the rest, and the + auxiliary coordinates, unchanged. + """ + indexes = { + d: pd.Index(indexers[d], name=d) if d in indexers else i + for d, i in self.indexes.items() + } + return replace(self, indexes=indexes) + + def conformed(self, target: Grid) -> Grid: + """``target`` carrying this grid's auxiliary coordinates, reindexed onto its labels.""" + aux = dict(self.aux) + for name, (d, values) in self.aux.items(): + if isinstance(d, str): + series = pd.Series(values, index=self.indexes[d]) + aux[name] = (d, series.reindex(target.indexes[d]).to_numpy()) + return Grid(target.indexes, aux) def combined(self, others: Iterable[Grid], how: str) -> Grid: """ Join with ``others`` along shared dimensions: per dimension the union (``how="outer"``) or intersection (``how="inner"``) of labels, kept in - this grid's dimension order. + this grid's dimension order. Auxiliary coordinates are not carried. """ combine = pd.Index.union if how == "outer" else pd.Index.intersection others = list(others) @@ -150,11 +194,32 @@ def combined(self, others: Iterable[Grid], how: str) -> Grid: indexes[d] = pd.Index(index, name=d) return Grid(indexes) + def same_layout(self, other: Grid) -> bool: + """Whether both grids have the same dims and labels, auxiliary coordinates aside.""" + return self.dims == other.dims and all( + i.equals(other.indexes[d]) for d, i in self.indexes.items() + ) + def __eq__(self, other: object) -> bool: if not isinstance(other, Grid): return NotImplemented - return self.dims == other.dims and all( - i.equals(other.indexes[d]) for d, i in self.indexes.items() + if not self.same_layout(other) or self.aux.keys() != other.aux.keys(): + return False + for name, (d, values) in self.aux.items(): + other_d, other_values = other.aux[name] + if d != other_d: + return False + if not pd.Index(np.atleast_1d(values)).equals( + pd.Index(np.atleast_1d(other_values)) + ): + return False + return True + + def dataarray(self, values: np.ndarray, name: str | None = None) -> DataArray: + """Wrap one value per flat cell as a DataArray on the grid's coordinates.""" + coords = self.to_dataset().coords + return DataArray( + values.reshape(self.shape), coords=coords, dims=self.dims, name=name ) @@ -166,16 +231,14 @@ class CSRLinearExpression: ``csr`` has one row per flat cell of ``grid`` 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, NaN for an absent cell. ``coords`` holds auxiliary - coordinates as ``name -> (grid dim, values)``, e.g. the key levels of a - grouped result kept stacked over the observed key combinations. + the per-cell constant, NaN for an absent cell. Auxiliary coordinates + live on the ``grid``. """ csr: scipy.sparse.csr_array const: np.ndarray grid: Grid model: Model - coords: dict[str, tuple[str, np.ndarray]] = field(default_factory=dict) @property def shape(self) -> tuple[int, ...]: @@ -189,10 +252,24 @@ def n_cells(self) -> int: def nterm(self) -> int: return csr_nterm(self.csr) + def cell(self, indices: tuple[Any, ...]) -> tuple[np.ndarray, np.ndarray, float]: + """ + Coefficients, label-ordered variable labels and constant of the grid + cell at ``indices``, as in the dense form: an absent cell has no terms. + """ + row = int(np.ravel_multi_index(indices, self.grid.shape)) if indices else 0 + const = float(self.const[row]) + start, end = self.csr.indptr[row], self.csr.indptr[row + 1] + if np.isnan(const): + end = start + vars_, coeffs = self.csr.indices[start:end], self.csr.data[start:end] + order = np.argsort(vars_, kind="stable") + return coeffs[order], vars_[order], const + @classmethod def from_grouper( cls, - ds: Dataset, + source: Dataset | CSRLinearExpression, model: Model, grouper: pd.Series | pd.DataFrame, group_dim: str, @@ -210,7 +287,10 @@ def from_grouper( ``stacked``, a single ``group_dim`` over the observed key combinations only, the key values attached as auxiliary coordinates. The new dims take the member dim's slot in ``coord_dims``, as on the dense path. + ``source`` is a dense expression dataset or an already CSR-backed + expression, which is regrouped through :meth:`aggregated`. """ + ds = source if isinstance(source, Dataset) else source.grid.to_dataset() member_dim = str(grouper.index.name) if member_dim in ds.indexes: grouper = grouper.reindex(ds.indexes[member_dim]) @@ -225,7 +305,7 @@ def from_grouper( keys = [str(k) for k in frame.columns] scatter_codes: dict[str, np.ndarray] = {} indexes: dict[str, pd.Index] = {} - coords: dict[str, tuple[str, np.ndarray]] = {} + aux: AuxCoords = {} if len(keys) == 1: codes, uniques = pd.factorize(frame.iloc[:, 0], sort=True) scatter_codes[group_dim] = codes @@ -234,7 +314,7 @@ def from_grouper( codes, uniques = pd.factorize(pd.MultiIndex.from_frame(frame), sort=True) scatter_codes[group_dim] = codes indexes[group_dim] = pd.RangeIndex(len(uniques), name=group_dim) - coords = { + aux = { k: (group_dim, uniques.get_level_values(i).to_numpy()) for i, k in enumerate(keys) } @@ -250,11 +330,13 @@ def from_grouper( for d in coord_dims: if d != member_dim: indexes[d] = ds.get_index(d).rename(d) - coords |= _aux_coords(ds, set(coord_dims) - {member_dim}) - grid = Grid({d: indexes[d] for d in grid_dims}) - return cls._from_scatter( - ds, model, grid, member_dim, scatter_codes, True, coords - ) + aux |= _aux_coords(ds, set(coord_dims) - {member_dim}) + grid = Grid({d: indexes[d] for d in grid_dims}, aux) + if isinstance(source, Dataset): + return cls._from_scatter(ds, model, grid, member_dim, scatter_codes, True) + member_rows = _member_rows(grid, scatter_codes, ds.sizes[member_dim]) + rows = _cell_rows(grid, member_rows, member_dim, source.grid.dims) + return source.aggregated(grid, rows) @classmethod def from_dense(cls, ds: Dataset, model: Model) -> CSRLinearExpression: @@ -263,8 +345,7 @@ def from_dense(cls, ds: Dataset, model: Model) -> CSRLinearExpression: grid = Grid.from_dataset(ds, grid_dims) first = grid_dims[0] codes = {first: np.arange(len(grid.indexes[first]))} - coords = _aux_coords(ds, set(grid_dims)) - return cls._from_scatter(ds, model, grid, first, codes, False, coords) + return cls._from_scatter(ds, model, grid, first, codes, False) @classmethod def _from_scatter( @@ -275,7 +356,6 @@ def _from_scatter( member_dim: str, scatter_codes: dict[str, np.ndarray], skipna: bool, - coords: dict[str, tuple[str, np.ndarray]], ) -> CSRLinearExpression: """ Scatter an expression's terms into grid rows (conceptually ``G @ A``): @@ -288,22 +368,11 @@ def _from_scatter( stays absent, as on the dense v1 merge path. """ grid_dims = grid.dims - stride = grid.strides - slot = min(grid_dims.index(d) for d in scatter_codes) transposed = [d for d in grid_dims if d not in scatter_codes] transposed.insert(slot, member_dim) - member_rows = np.zeros(ds.sizes[member_dim], dtype=np.int64) - for d, codes in scatter_codes.items(): - member_rows += codes * stride[d] - cell_rows = _outer_sum( - [ - member_rows - if d == member_dim - else np.arange(len(grid.indexes[d])) * stride[d] - for d in transposed - ] - ) + member_rows = _member_rows(grid, scatter_codes, ds.sizes[member_dim]) + cell_rows = _cell_rows(grid, member_rows, member_dim, transposed) coeffs = ds.coeffs.transpose(*transposed, TERM_DIM).to_numpy().reshape(-1) vars_ = ds.vars.transpose(*transposed, TERM_DIM).to_numpy().reshape(-1) @@ -323,16 +392,101 @@ def _from_scatter( const[cell_rows] = 0.0 np.add.at(const, cell_rows, const_vals) - return cls(scipy.sparse.csr_array(coo), const, grid, model, coords) + return cls(scipy.sparse.csr_array(coo), const, grid, model) - def scaled(self, factor: float) -> CSRLinearExpression: - return replace(self, csr=self.csr * factor, const=self.const * factor) + def aggregated(self, grid: Grid, rows: np.ndarray) -> CSRLinearExpression: + """ + Sum source rows into the cells of ``grid``, row ``i`` landing in cell + ``rows[i]`` (conceptually ``G @ A``). Goes through COO, so duplicate + variables are summed and explicit zeros kept, as by :meth:`added`. The + constant is reduced as by the dense group kernel (NaN counts as 0); + cells no row lands in are absent. Auxiliary coordinates are ``grid``'s. + """ + coo = self.csr.tocoo() + coo = scipy.sparse.coo_array( + (coo.data, (rows[coo.coords[0]], coo.coords[1])), + shape=(grid.size, self.csr.shape[1]), + ) + weights = np.nan_to_num(self.const) + const = np.bincount(rows, weights=weights, minlength=grid.size).astype(float) + const[np.bincount(rows, minlength=grid.size) == 0] = np.nan + return replace(self, csr=scipy.sparse.csr_array(coo), const=const, grid=grid) + + def summed(self, dims: Iterable[str]) -> CSRLinearExpression: + """ + Sum over grid dimensions, as the dense ``sum``: the kept dims stay in + grid order with their auxiliary coordinates, and every kept cell is + present, its constant the NaN-skipping sum of its members. + """ + dims = set(dims) + grid = self.grid.reordered(d for d in self.grid.dims if d not in dims) + stride = grid.strides + rows = _outer_sum( + [ + np.zeros(n, dtype=np.int64) if d in dims else np.arange(n) * stride[d] + for d, n in zip(self.grid.dims, self.grid.shape) + ] + ) + return self.aggregated(grid, rows).filled(0.0) + + def pruned(self) -> CSRLinearExpression: + """Drop explicit zero coefficients; cell activeness stays with ``const``.""" + csr = self.csr.copy() + csr.eliminate_zeros() + return replace(self, csr=csr) + + def scaled( + self, + factor: float | np.ndarray, + op: Callable[[Any, Any], Any] = operator.mul, + ) -> CSRLinearExpression: + """ + Apply ``op`` with ``factor`` -- a scalar or one value per cell -- to + every coefficient and to the constant. Explicit zeros are kept, as by + :meth:`added`. + """ + factor = np.broadcast_to(factor, (self.n_cells,)) + per_term = np.repeat(factor, np.diff(self.csr.indptr)) + csr = scipy.sparse.csr_array( + (op(self.csr.data, per_term), self.csr.indices, self.csr.indptr), + shape=self.csr.shape, + ) + return replace(self, csr=csr).with_const(op(self.const, factor)) + + def with_const(self, const: np.ndarray) -> CSRLinearExpression: + """ + Replace the per-cell constant, leaving the terms alone. A cell it makes + absent (NaN) drops its terms (v1 dead-term invariant). + """ + absent = np.isnan(const) + counts = np.diff(self.csr.indptr) + if not counts[absent].any(): + return replace(self, const=const) + keep = np.repeat(~absent, counts) + indptr = np.concatenate([[0], np.cumsum(np.where(absent, 0, counts))]) + csr = scipy.sparse.csr_array( + (self.csr.data[keep], self.csr.indices[keep], indptr), + shape=self.csr.shape, + ) + return replace(self, csr=csr, const=const) + + def taken(self, rows: np.ndarray, grid: Grid) -> CSRLinearExpression: + """ + Gather rows into the cells of ``grid``, cell ``i`` taking row + ``rows[i]`` with its terms, explicit zeros included, and its constant; + ``-1`` leaves the cell absent. Auxiliary coordinates are ``grid``'s. + """ + present = rows >= 0 + const = np.where(present, self.const[rows], np.nan) + csr = self.csr[np.where(present, rows, 0)] + return replace(self, csr=csr, grid=grid).with_const(const) def reindexed(self, grid: Grid, fill: float = np.nan) -> CSRLinearExpression: """ Remap rows onto a new grid, possibly in a new dim order, without the dense rectangle: dropped labels vanish, new labels get ``fill`` as - their constant (NaN: absent cells). + their constant (NaN: absent cells). Auxiliary coordinates follow the + rows; those of ``grid`` are ignored. """ row_map, valid = grid.indexer(self.grid) @@ -346,21 +500,11 @@ def reindexed(self, grid: Grid, fill: float = np.nan) -> CSRLinearExpression: ) const = np.full(n_cells, fill) const[row_map[valid]] = self.const[valid] - coords = { - name: ( - d, - pd.Series(v, index=self.grid.indexes[d]) - .reindex(grid.indexes[d]) - .to_numpy(), - ) - for name, (d, v) in self.coords.items() - } return replace( self, csr=scipy.sparse.csr_array(coo), const=const, - grid=grid, - coords=coords, + grid=self.grid.conformed(grid), ) def filled(self, value: float) -> CSRLinearExpression: @@ -370,11 +514,11 @@ def filled(self, value: float) -> CSRLinearExpression: def renamed(self, names: dict[str, str]) -> CSRLinearExpression: """Relabel grid dims; the CSR row layout is unchanged.""" - coords = {n: (names.get(d, d), v) for n, (d, v) in self.coords.items()} - return replace(self, grid=self.grid.renamed(names), coords=coords) + return replace(self, grid=self.grid.renamed(names)) def same_grid(self, other: CSRLinearExpression) -> bool: - return self.grid == other.grid + """Whether both live on the same cells, auxiliary coordinates aside.""" + return self.grid.same_layout(other.grid) def added(self, other: CSRLinearExpression) -> CSRLinearExpression: """ @@ -396,11 +540,9 @@ def added(self, other: CSRLinearExpression) -> CSRLinearExpression: coo = scipy.sparse.coo_array( (data[present], (rows[present], cols[present])), shape=shape ) - enforce_aux_conflict([Dataset(coords=p.coords) for p in (self, other)]) - coords = other.coords | self.coords - return replace( - self, csr=scipy.sparse.csr_array(coo), const=const, coords=coords - ) + enforce_aux_conflict([Dataset(coords=p.grid.aux) for p in (self, other)]) + grid = replace(self.grid, aux=other.grid.aux | self.grid.aux) + return replace(self, csr=scipy.sparse.csr_array(coo), const=const, grid=grid) def contracted( self, @@ -459,14 +601,12 @@ def contracted( blocks.append(block @ source.csr[rows]) const_blocks.append(block @ const[rows]) - indexes = kept_grid.indexes - indexes |= {str(i.name): i for i in new_indexes} + indexes = kept_grid.indexes | {str(i.name): i for i in new_indexes} return replace( source, csr=scipy.sparse.csr_array(scipy.sparse.vstack(blocks, format="csr")), const=np.concatenate(const_blocks), - grid=Grid(indexes), - coords={n: (d, v) for n, (d, v) in source.coords.items() if d in kept}, + grid=Grid(indexes, kept_grid.aux), ) def to_dense(self) -> LinearExpression: @@ -493,7 +633,7 @@ def to_dense(self) -> LinearExpression: "vars": (dims, vars_flat.reshape(*shape, nterm)), "const": (self.grid.dims, self.const.reshape(shape)), }, - coords=self.grid.indexes | self.coords, + coords=self.grid.indexes | self.grid.aux, ) return LinearExpression(absorb_absence(ds), self.model) @@ -532,13 +672,57 @@ def csr_to_term_arrays( return vars_, coeffs -def _aux_coords(ds: Dataset, dims: set[str]) -> dict[str, tuple[str, np.ndarray]]: - """One-dimensional auxiliary coordinates of ``ds`` lying on ``dims``.""" - return { - str(n): (str(c.dims[0]), c.to_numpy()) - for n, c in ds.coords.items() - if n not in ds.dims and len(c.dims) == 1 and str(c.dims[0]) in dims - } +def _densify_notice(reason: str) -> None: + """ + Emit a :class:`~linopy.constants.PerformanceWarning` naming why a sparse + (CSR) backing is dropped, if ``options["warn_on_densify"]`` is set. + """ + if options["warn_on_densify"]: + warn_outside_linopy( + f"Sparse (CSR) backing densified: {reason}.", PerformanceWarning + ) + + +def _aux_coords(ds: Dataset | DataArray, dims: set[str]) -> AuxCoords: + """Scalar and one-dimensional auxiliary coordinates of ``ds`` lying on ``dims``.""" + aux: AuxCoords = {} + for n, c in ds.coords.items(): + if n in ds.dims: + continue + if c.ndim == 0: + aux[str(n)] = ((), c.to_numpy()) + elif c.ndim == 1 and str(c.dims[0]) in dims: + aux[str(n)] = (str(c.dims[0]), c.to_numpy()) + return aux + + +def _member_rows( + grid: Grid, scatter_codes: dict[str, np.ndarray], n_members: int +) -> np.ndarray: + """Row offset in ``grid`` of each member, from its per-dim group codes.""" + stride = grid.strides + member_rows = np.zeros(n_members, dtype=np.int64) + for d, codes in scatter_codes.items(): + member_rows += codes * stride[d] + return member_rows + + +def _cell_rows( + grid: Grid, member_rows: np.ndarray, member_dim: str, dims: Iterable[str] +) -> np.ndarray: + """ + Grid row of every cell over ``dims`` in C order, ``member_dim`` landing on + ``member_rows`` and every other dim mapping one-to-one onto ``grid``. + """ + stride = grid.strides + return _outer_sum( + [ + member_rows + if d == member_dim + else np.arange(len(grid.indexes[d])) * stride[d] + for d in dims + ] + ) def _outer_sum(axis_positions: list[np.ndarray]) -> np.ndarray: diff --git a/linopy/expressions.py b/linopy/expressions.py index 5385c9494..0726afc25 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -108,7 +108,7 @@ STACKED_TERM_DIM, TERM_DIM, ) -from linopy.csr import CSRLinearExpression, _aux_coords +from linopy.csr import CSRLinearExpression, Grid, _aux_coords, _densify_notice from linopy.semantics import ( AbsentType, FillValueLike, @@ -464,13 +464,30 @@ def _restore_multikey_index( class LinearExpressionGroupby: """ GroupBy object specialized to grouping LinearExpression objects. + + ``obj`` is the grouped expression, or its dense dataset. A CSR-backed + expression is only densified by operations without a sparse path. """ - data: xr.Dataset + obj: xr.Dataset | BaseExpression group: Hashable | DataArray | IndexVariable | pd.Series | pd.DataFrame model: Any kwargs: Mapping[str, Any] = field(default_factory=dict) + @property + def data(self) -> Dataset: + return self.obj if isinstance(self.obj, Dataset) else self.obj.data + + @property + def _csr(self) -> CSRLinearExpression | None: + return self.obj._csr if isinstance(self.obj, LinearExpression) else None + + @property + def _labels(self) -> Dataset: + """The coordinates to resolve groupers against, without densifying.""" + csr = self._csr + return self.data if csr is None else csr.grid.to_dataset() + @property def groupby(self) -> xarray.core.groupby.DatasetGroupBy: """ @@ -584,16 +601,16 @@ def sum( "`observed=True` is not supported with `use_fallback=True`." ) - group = _resolve_group(self.group, self.data) - _check_grouper_alignment(group, self.data) + labels = self._labels + group = _resolve_group(self.group, labels) + _check_grouper_alignment(group, labels) - multikey_frame = ( - None if use_fallback else _multikey_value_frame(group, self.data) - ) + multikey_frame = None if use_fallback else _multikey_value_frame(group, labels) + csr = self._csr explicit_sparse = sparse is True if sparse is None: - sparse = is_v1() and options["sparse_groupby"] + sparse = is_v1() and (options["sparse_groupby"] or csr is not None) elif sparse and not is_v1(): raise ValueError( "sparse groupby-sum requires v1 semantics; opt in with " @@ -607,10 +624,17 @@ def sum( if grouper is None: is_1d = isinstance(group, DataArray) and group.ndim == 1 grouper = group.to_pandas() if is_1d else group + if csr is None: + source: Dataset | CSRLinearExpression = self.data + coord_dims = tuple( + str(d) for d in self.data.coeffs.dims if d != TERM_DIM + ) + else: + source, coord_dims = csr, csr.grid.dims supported = ( not use_fallback and isinstance(grouper, (pd.Series, pd.DataFrame)) - and grouper.index.name in self.data.dims + and grouper.index.name in coord_dims ) if supported: stacked = observed or multikey_frame is None @@ -619,13 +643,10 @@ def sum( if isinstance(grouper, pd.DataFrame) else str(grouper.name or "group") ) - coord_dims = tuple( - str(d) for d in self.data.coeffs.dims if d != TERM_DIM - ) - csr = CSRLinearExpression.from_grouper( - self.data, self.model, grouper, group_name, stacked, coord_dims + res = CSRLinearExpression.from_grouper( + source, self.model, grouper, group_name, stacked, coord_dims ) - return LinearExpression._from_csr(csr, self.model) + return LinearExpression._from_csr(res, self.model) if explicit_sparse: raise ValueError( "sparse=True supports only a pandas Series or DataFrame, 1-D " @@ -920,10 +941,10 @@ def __repr__(self) -> str: masked_entries = 0 # (~self.mask).sum().values.item() if self.mask lines = [] - header_string = self.type + header_string = f"{self.type} (sparse)" if self.is_sparse else self.type if size > 1 or ndim > 0: - row_labels = get_printout_labels(self.data, dims) + row_labels = get_printout_labels(self.coords.to_dataset(), dims) for indices in generate_indices_for_printout(dim_sizes, max_lines): if indices is None: lines.append("\t\t...") @@ -931,10 +952,7 @@ def __repr__(self) -> str: coord = [row_labels[i][ind] for i, ind in enumerate(indices)] if self.mask is None or self.mask.values[indices]: expr = format_single_expression( - self.coeffs.values[indices], - self.vars.values[indices], - self.const.values[indices], - self.model, + *self._printout_cell(indices), self.model ) line = format_coord(coord) + f": {expr}" @@ -947,15 +965,28 @@ def __repr__(self) -> str: underscore = "-" * (len(shape_str) + len(mask_str) + len(header_string) + 4) lines.insert(0, f"{header_string} [{shape_str}]{mask_str}:\n{underscore}") elif size == 1: - expr = format_single_expression( - self.coeffs.values, self.vars.values, self.const.item(), self.model - ) + expr = format_single_expression(*self._printout_cell(()), self.model) lines.append(f"{header_string}\n{'-' * len(header_string)}\n{expr}") else: lines.append(f"{header_string}\n{'-' * len(header_string)}\n") return "\n".join(lines) + def _printout_cell( + self, indices: tuple[Any, ...] + ) -> tuple[np.ndarray, np.ndarray, Any]: + """Coefficients, variable labels and constant of one cell of the repr.""" + return ( + self.coeffs.values[indices], + self.vars.values[indices], + self.const.values[indices], + ) + + @property + def is_sparse(self) -> bool: + """Whether the expression is backed by a sparse CSR matrix (:mod:`linopy.csr`).""" + return False + @property def is_constant(self) -> bool: """True if the expression contains no variables.""" @@ -1085,7 +1116,7 @@ def _align_constant( API-stable across xarray releases. """ enforce_aux_conflict([self.const, other], stacklevel=4) - other_fill = join_fill(fill_value, 0) + other_fill = {other.name: join_fill(fill_value, 0)} if join is None: if is_v1(): join = "exact" @@ -1150,6 +1181,25 @@ def _reindexed_to(self, const: DataArray, needs_data_reindex: bool) -> Self: data = self.data.reindex_like(const, fill_value=self._fill_value) return self.__class__(data, self.model) + def _combined_with_constant( + self, + self_const: DataArray, + operand: Any, + op: Callable[[Any, Any], Any], + *, + needs_data_reindex: bool, + scale: bool, + ) -> Self: + """ + Combine the aligned constant ``operand`` into ``const`` with ``op``, + and into the coefficients too if ``scale``. + """ + expr = self._reindexed_to(self_const, needs_data_reindex) + const = op(self_const, operand) + if not scale: + return expr.assign(const=const) + return expr.assign(coeffs=op(expr.coeffs, operand), const=const) + def _add_constant( self, other: ConstantLike, @@ -1170,12 +1220,16 @@ def _add_constant_v1( if np.isscalar(other) and join is None: if is_nan_scalar(other): check_user_nan() - return self.assign(const=self.const + other) + return self._combined_with_constant( + self.const, other, operator.add, needs_data_reindex=False, scale=False + ) self_const, da, needs_reindex = self._broadcast_and_align( other, fill_value, join ) - expr = self._reindexed_to(self_const, needs_reindex) - return expr.assign(const=self_const + da)._absorb_join_absence(fill_value) + expr = self._combined_with_constant( + self_const, da, operator.add, needs_data_reindex=needs_reindex, scale=False + ) + return expr._absorb_join_absence(fill_value) # LEGACY: remove at 1.0 — see doc/design/legacy-removal.rst. def _add_constant_legacy( @@ -1226,11 +1280,10 @@ def _apply_constant_op_v1( self_const, factor, needs_reindex = self._broadcast_and_align( other, fill_value, join, op_kind ) - expr = self._reindexed_to(self_const, needs_reindex) - result = expr.assign( - coeffs=op(expr.coeffs, factor), const=op(self_const, factor) + expr = self._combined_with_constant( + self_const, factor, op, needs_data_reindex=needs_reindex, scale=True ) - return result._absorb_join_absence(fill_value) + return expr._absorb_join_absence(fill_value) # LEGACY: remove at 1.0 — see doc/design/legacy-removal.rst. def _apply_constant_op_legacy( @@ -1616,7 +1669,7 @@ def coord_sizes(self) -> dict[Hashable, int]: @property def coord_names(self) -> list[str]: - return get_dims_with_index_levels(self.data, self.coord_dims) + return get_dims_with_index_levels(self.coords.to_dataset(), self.coord_dims) @property def vars(self) -> DataArray: @@ -1891,9 +1944,10 @@ def isnull(self) -> DataArray: def where( self, - cond: DataArray, + cond: DataArray | Callable[[Any], DataArray], other: LinearExpression | int + | float | DataArray | dict[str, float | int | DataArray] | None = None, @@ -2030,9 +2084,8 @@ def groupby( A `LinearExpressionGroupBy` containing the xarray groups and ensuring the correct return type. """ - ds = self.data kwargs = dict(restore_coord_dims=restore_coord_dims, **kwargs) - return LinearExpressionGroupby(ds, group, model=self.model, kwargs=kwargs) + return LinearExpressionGroupby(self, group, model=self.model, kwargs=kwargs) def rolling( self, @@ -2357,17 +2410,225 @@ def _from_csr(cls, csr: CSRLinearExpression, model: Model) -> Self: @property def data(self) -> Dataset: - if self._data is None and self._csr is not None: - self._data = self._csr.to_dense()._data - self._csr = None + self._densify("`.data` read, e.g. by an operation without a sparse path") return self._data + def _densify(self, reason: str) -> None: + """Convert the CSR backing to dense for good, with a notice naming ``reason``.""" + _densify_all([self], reason) + + @property + def is_sparse(self) -> bool: + """ + Whether the expression is backed by a sparse CSR matrix + (:mod:`linopy.csr`) instead of a dense Dataset. Reading ``.data``, or + an operation without a sparse path, converts it to dense for good. + """ + return self._csr is not None + @property def coord_dims(self) -> tuple[Hashable, ...]: - if self._data is None and self._csr is not None: + if self._csr is not None: return self._csr.grid.dims return super().coord_dims + @property + def dims(self) -> tuple[Hashable, ...]: + if self._csr is not None: + return (*self._csr.grid.dims, TERM_DIM) + return super().dims + + @property + def shape(self) -> tuple[int, ...]: + if self._csr is not None: + return (*self._csr.grid.shape, self._csr.nterm) + return super().shape + + @property + def size(self) -> int: + if self._csr is not None: + return int(np.prod(self.shape, dtype=np.int64)) + return super().size + + @property + def sizes(self) -> Frozen: + if self._csr is not None: + return Frozen(dict(zip(self.dims, self.shape))) + return super().sizes + + @property + def coords(self) -> DatasetCoordinates: + if self._csr is not None: + return self._csr.grid.to_dataset().coords + return super().coords + + @property + def indexes(self) -> Indexes: + if self._csr is not None: + return self._csr.grid.to_dataset().indexes + return super().indexes + + @property + def const(self) -> DataArray: + csr = self._csr + if csr is None: + return super().const + return csr.grid.dataarray(csr.const, name="const") + + @const.setter + def const(self, value: DataArray) -> None: + self._data = assign_multiindex_safe(self.data, const=value) + + def _combined_with_constant( + self, + self_const: DataArray, + operand: Any, + op: Callable[[Any, Any], Any], + *, + needs_data_reindex: bool, + scale: bool, + ) -> Self: + csr = self._csr + flags = {"needs_data_reindex": needs_data_reindex, "scale": scale} + if csr is None: + return super()._combined_with_constant(self_const, operand, op, **flags) + dims = csr.grid.dims + const = op(self_const, operand) + if set(const.dims) != set(dims) or _has_multiindex(const.indexes.values()): + self._densify( + "elementwise operation with a constant over new dimensions or " + "MultiIndex labels" + ) + return super()._combined_with_constant(self_const, operand, op, **flags) + const = const.transpose(*dims) + grid = Grid.from_dataset(const, dims) + if needs_data_reindex: + csr = csr.reindexed(grid) + if scale: + factor = operand.broadcast_like(const).transpose(*dims) + csr = csr.scaled(factor.to_numpy().reshape(-1), op) + csr = replace(csr.with_const(const.to_numpy().reshape(-1)), grid=grid) + return type(self)._from_csr(csr, self._model) + + def _absorb_join_absence(self, fill_value: FillValueLike) -> Self: + if self._csr is not None: + return self + return super()._absorb_join_absence(fill_value) + + def _printout_cell( + self, indices: tuple[Any, ...] + ) -> tuple[np.ndarray, np.ndarray, Any]: + if self._csr is not None: + return self._csr.cell(indices) + return super()._printout_cell(indices) + + def _selected( + self, select: Callable[[DataArray], DataArray], name: str + ) -> CSRLinearExpression | None: + """ + Apply the dense selection ``select`` to the grid's row numbers and + gather the selected rows, a row number of ``-1`` or NaN leaving the + cell absent. None, after a densify notice, where the selection leaves + the grid: new or unlabelled dimensions, MultiIndex labels, or + coordinates that are no auxiliary coordinates of the grid. + """ + csr = self._csr + if csr is None: + return None + reason = f"`{name}` over MultiIndex labels, or introducing dimensions" + if _has_multiindex(csr.grid.indexes.values()): + self._densify(reason) + return None + rows = select(csr.grid.dataarray(np.arange(csr.n_cells))) + dims = tuple(str(d) for d in rows.dims) + leaves_grid = not set(dims) <= set(csr.grid.dims) & set(rows.indexes) + if leaves_grid or _has_multiindex(rows.indexes.values()): + self._densify(reason) + return None + grid = Grid.from_dataset(rows, dims) + if set(rows.coords) != set(dims) | set(grid.aux): + self._densify(reason) + return None + flat = rows.fillna(-1).to_numpy().reshape(-1).astype(np.int64) + return csr.taken(flat, grid) + + def sel(self, *args: Any, **kwargs: Any) -> LinearExpression: + """Select by label as ``Dataset.sel``; a CSR-backed expression stays sparse.""" + csr = self._selected(lambda rows: rows.sel(*args, **kwargs), "sel") + if csr is None: + return super().sel(*args, **kwargs) + return type(self)._from_csr(csr, self._model) + + def isel(self, *args: Any, **kwargs: Any) -> LinearExpression: + """Select by position as ``Dataset.isel``; a CSR-backed expression stays sparse.""" + csr = self._selected(lambda rows: rows.isel(*args, **kwargs), "isel") + if csr is None: + return super().isel(*args, **kwargs) + return type(self)._from_csr(csr, self._model) + + def __getitem__(self, selector: int | tuple[slice, list[int]] | slice) -> Self: + csr = self._selected(lambda rows: rows[selector], "__getitem__") + if csr is None: + return super().__getitem__(selector) + return type(self)._from_csr(csr, self._model) + + def where( + self, + cond: DataArray | Callable[[Any], DataArray], + other: LinearExpression + | int + | float + | DataArray + | dict[str, float | int | DataArray] + | None = None, + **kwargs: Any, + ) -> Self: + cond_da = _expr_unwrap(cond) + drop = kwargs.get("drop", False) + absent = other is None or other is np.nan + supported = ( + isinstance(cond_da, DataArray) + and set(kwargs) <= {"drop"} + and (absent or (isinstance(other, int | float) and not drop)) + ) + if not supported: + self._densify( + "`where` with a condition that is no DataArray, an `other` that " + "is no scalar, or extra arguments" + ) + return super().where(cond, other, **kwargs) + csr = self._selected( + lambda rows: ( + rows.where(cond_da, drop=True) if drop else rows.where(cond_da, -1) + ), + "where", + ) + if csr is None: + return super().where(cond, other, **kwargs) + if not absent: + const = self.const.where(cond_da, other).transpose(*csr.grid.dims) + csr = csr.with_const(const.to_numpy().reshape(-1)) + return type(self)._from_csr(csr, self._model) + + def sum( + self, + dim: DimsLike | None = None, + drop_zeros: bool = False, + **kwargs: Any, + ) -> Self: + csr = self._csr + if csr is None or kwargs: + return super().sum(dim, drop_zeros, **kwargs) + if dim is None or isinstance(dim, EllipsisType): + dims: list[Hashable] = list(csr.grid.dims) + else: + dims = [dim] if isinstance(dim, str) else list(dim) + summed = [str(d) for d in dims if d != TERM_DIM] + if not set(summed) <= set(csr.grid.dims): + return super().sum(dim, drop_zeros) + res = csr.summed(summed) if summed else csr + return type(self)._from_csr(res.pruned() if drop_zeros else res, self._model) + @property def nterm(self) -> int: """ @@ -2409,8 +2670,9 @@ def to_constraint( ) -> ConstraintBase: if self._csr is not None and isinstance(sign, str): rhs_da = constraints.csr_rhs(self._csr, rhs) - if rhs_da is not None: + if isinstance(rhs_da, DataArray): return constraints.CSRConstraint.from_csr(self._csr, sign, rhs_da) + self._densify(rhs_da) return super().to_constraint(sign, rhs, join) @overload @@ -2497,7 +2759,11 @@ def __mul__( """ Multiply the expr by a factor. """ - if self._csr is not None and isinstance(other, int | float | np.number): + if ( + self._csr is not None + and isinstance(other, int | float | np.number) + and not is_nan_scalar(other) + ): return type(self)._from_csr(self._csr.scaled(float(other)), self._model) other = as_constant(other) if isinstance(other, QuadraticExpression): @@ -2585,22 +2851,29 @@ def _sparse_matmul(self, other: ConstantLike) -> LinearExpression | None: if self._csr is None and _has_multiindex(self.data.indexes.values()): return None if not self.coord_dims: + self._densify("`@` on a zero-dimensional expression") return None csr = self._csr or CSRLinearExpression.from_dense(self.data, self.model) if not csr.grid.is_unique or _has_multiindex(csr.grid.indexes.values()): + self._densify("`@` over non-unique or MultiIndex labels") return None if csr.grid.size == 0: + self._densify("`@` over a zero-size grid") return None - coords = Dataset(coords=dict(csr.grid.indexes) | dict(csr.coords)).coords + coords = csr.grid.to_dataset().coords da = _matmul_operand_to_dataarray(other, coords, csr.grid.dims) if _has_multiindex(da.indexes.values()): + self._densify("`@` with a MultiIndex operand") return None contracted = [d for d in csr.grid.dims if d in da.dims] new_dims = [str(d) for d in da.dims if d not in csr.grid.dims] if not contracted or any(d not in da.indexes for d in new_dims): + self._densify( + "`@` with an operand sharing no dimension or with an unlabelled one" + ) return None matrix = _matmul_operand_to_matrix( - da, contracted, new_dims, csr.grid.indexes, csr.coords + da, contracted, new_dims, csr.grid.indexes, csr.grid.aux ) new_indexes = [da.indexes[d].rename(d) for d in new_dims] res = csr.contracted(matrix, contracted, new_indexes) @@ -3288,23 +3561,37 @@ def _has_multiindex(indexes: Iterable[pd.Index]) -> bool: return any(isinstance(i, pd.MultiIndex) for i in indexes) +def _densify_all(exprs: Iterable[Any], reason: str) -> None: + """Convert the CSR-backed LinearExpressions among ``exprs`` to dense, with one notice.""" + sparse = [ + (e, csr) + for e in exprs + if isinstance(e, LinearExpression) and (csr := e._csr) is not None + ] + if sparse: + _densify_notice(reason) + for e, csr in sparse: + e._data = csr.to_dense()._data + e._csr = None + + def _aligned( csrs: list[CSRLinearExpression], join: JoinOptions | None, fill: float -) -> list[CSRLinearExpression] | None: +) -> list[CSRLinearExpression] | str: """ Conform the CSR expressions to the grid an explicit join produces, the cells the join creates carrying ``fill`` as constant. None where the dense path owns the semantics: ``exact`` and the auto-detected join raise there on differing grids, ``override`` on differing shapes, any join on - non-unique labels. + non-unique labels; the reason is returned instead. """ template = csrs[0].grid dims = template.dims if any(not p.grid.is_unique for p in csrs): - return None + return "merge over non-unique labels" if join == "override": if any(p.grid.dims != dims or p.grid.shape != template.shape for p in csrs): - return None + return '`join="override"` merge of differently shaped grids' return [replace(p, grid=template) for p in csrs] if join in ("left", "right"): source = csrs[0] if join == "left" else csrs[-1] @@ -3312,7 +3599,7 @@ def _aligned( elif join in ("outer", "inner"): grid = template.combined([p.grid for p in csrs[1:]], join) else: - return None + return f"merge of differing grids with join={join!r}" return [p.reindexed(grid, fill) for p in csrs] @@ -3329,23 +3616,32 @@ def _try_csr_merge( addition. Grids that share dims in a different order are transposed onto the template order first. Grids that differ in their labels are aligned row-wise onto the joined grid, the cells the join creates carrying the - fill of the dense path (zero, or NaN for ``fill_value=ABSENT``); auxiliary - coordinates across differing grids are left to the dense path. Returns + fill of the dense path (zero, or NaN for ``fill_value=ABSENT``). Auxiliary + coordinates are checked for conflicts on the operands as given (§11, as + on the dense path) and follow their rows onto the joined grid. Returns None to fall through to the dense path. """ + if not any(type(e) is LinearExpression and e._csr is not None for e in exprs): + return None if dim != TERM_DIM or kwargs: + _densify_all( + exprs, "merge along a coordinate dimension or with extra arguments" + ) return None if not all(type(e) is LinearExpression for e in exprs): - return None - if all(e._csr is None for e in exprs): + _densify_all(exprs, "merge with an operand that is not a LinearExpression") return None dims = set(exprs[0].coord_dims) if any(set(e.coord_dims) != dims for e in exprs[1:]): + _densify_all(exprs, "merge of operands over different dimensions") return None for e in exprs: if e._csr is None and set(e.data.coords) - dims != set( _aux_coords(e.data, dims) ): + _densify_all( + exprs, "merge with a dense operand carrying non-grid coordinates" + ) return None csrs = [e._csr or CSRLinearExpression.from_dense(e.data, e.model) for e in exprs] @@ -3356,10 +3652,10 @@ def _try_csr_merge( for p in csrs ] if not all(template.same_grid(p) for p in csrs[1:]): - if any(p.coords for p in csrs): - return None + enforce_aux_conflict([Dataset(coords=p.grid.aux) for p in csrs]) aligned = _aligned(csrs, join, join_fill(fill_value, 0.0)) - if aligned is None: + if isinstance(aligned, str): + _densify_all(exprs, aligned) return None csrs = aligned @@ -3528,10 +3824,9 @@ def merge( skipna = not is_v1() if dim == TERM_DIM: ds = xr.concat([d[["coeffs", "vars"]] for d in data], dim, **kwargs) - subkwargs = {**kwargs, "fill_value": join_fill(fill_value, 0)} - const = xr.concat([d["const"] for d in data], dim, **subkwargs).sum( - TERM_DIM, skipna=skipna - ) + subkwargs = {**kwargs, "fill_value": {"const": join_fill(fill_value, 0)}} + const = xr.concat([d[["const"]] for d in data], dim, **subkwargs)["const"] + const = const.sum(TERM_DIM, skipna=skipna) ds = assign_multiindex_safe(ds, const=const) elif dim == FACTOR_DIM: ds = xr.concat([d[["vars"]] for d in data], dim, **kwargs) diff --git a/linopy/model.py b/linopy/model.py index 6258c4ffe..fd4066c97 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -1329,6 +1329,15 @@ def add_constraints( rhs_da = as_dataarray(rhs) original_rhs_mask = (rhs_da.coords, rhs_da.dims, ~np.isnan(rhs_da.values)) + if ( + isinstance(lhs, LinearExpression) + and lhs.is_sparse + and freeze + and mask is not None + ): + mask = broadcast_to_coords(mask, lhs.coords, label="mask").astype(bool) + lhs = lhs.where(mask) + mask = None con = self._constraint_from_lhs(lhs, sign, rhs, coords) if isinstance(con, CSRConstraint) and freeze and mask is None: _check_infinities(con._sign, con._rhs, name) diff --git a/linopy/objective.py b/linopy/objective.py index dc2e8ad41..6cc2cd6fd 100644 --- a/linopy/objective.py +++ b/linopy/objective.py @@ -270,7 +270,7 @@ def to_matrix(self, *args: Any, **kwargs: Any) -> csc_matrix: assign = objwrap(expressions.LinearExpression.assign) - sel = objwrap(expressions.LinearExpression.sel) + sel = objwrap(expressions.BaseExpression.sel) def __add__( self, expr: ConstantLike | QuadraticExpression | LinearExpression | Objective diff --git a/linopy/semantics.py b/linopy/semantics.py index e8e584d7f..ea44af51f 100644 --- a/linopy/semantics.py +++ b/linopy/semantics.py @@ -303,29 +303,36 @@ def _legacy_multiindex_message(dim: str, context: str) -> str: _LINOPY_ROOT = os.path.dirname(os.path.abspath(__file__)) -def warn_legacy(message: str, *, stacklevel: int | None = None) -> None: +def warn_outside_linopy(message: str, category: type[Warning]) -> None: + """ + Emit a warning whose source-frame points at the first call-stack frame + *outside* the linopy package, as the call-chain depth varies per site. + Python 3.12+ uses the stdlib ``skip_file_prefixes``; older versions walk + the stack to compute the ``stacklevel``. """ - Emit a `LinopySemanticsWarning` whose source-frame points at the - first call-stack frame *outside* the linopy package. + if sys.version_info >= (3, 12): + warn(message, category, skip_file_prefixes=(_LINOPY_ROOT,)) + return + frame = sys._getframe() + stacklevel = 1 + while frame.f_back is not None and frame.f_code.co_filename.startswith( + _LINOPY_ROOT + ): + frame = frame.f_back + stacklevel += 1 + warn(message, category, stacklevel=stacklevel) - Static ``stacklevel`` doesn't fit here — the call-chain depth from - ``warn_legacy`` to the user's code varies per site (e.g. masked-var - via ``__add__`` is 5 frames deep, via ``Variable.fillna`` is 4). On - Python 3.12+ we use the stdlib ``skip_file_prefixes`` argument - (implemented and tested in CPython); on 3.11 we fall back to a - static ``stacklevel=5``, good enough for the common merge chain. - Pass an explicit ``stacklevel`` to override (e.g. for tests). + +def warn_legacy(message: str, *, stacklevel: int | None = None) -> None: """ - if stacklevel is not None: - warn(message, LinopySemanticsWarning, stacklevel=stacklevel) - elif sys.version_info >= (3, 12): - warn( - message, - LinopySemanticsWarning, - skip_file_prefixes=(_LINOPY_ROOT,), - ) + Emit a `LinopySemanticsWarning` pointing at the first call-stack frame + outside the linopy package (see :func:`warn_outside_linopy`). Pass an + explicit ``stacklevel`` to override (e.g. for tests). + """ + if stacklevel is None: + warn_outside_linopy(message, LinopySemanticsWarning) else: - warn(message, LinopySemanticsWarning, stacklevel=5) + warn(message, LinopySemanticsWarning, stacklevel=stacklevel) def _short_repr(values: Any, limit: int = 6) -> str: @@ -491,6 +498,7 @@ def reindex_like_if_needed( ) -> DataArray: """ ``arr.reindex_like(ref, fill_value=...)`` without the copy when already aligned. + Only the values take ``fill_value``; auxiliary coordinates are filled with NaN. ``broadcast_to_coords`` leaves an operand spanning ``ref``'s dims but with fresh index objects (equal, not identical), so a plain ``reindex_like`` would @@ -501,7 +509,7 @@ def reindex_like_if_needed( """ if set(ref.dims) <= set(arr.dims) and first_mismatched_dim(ref, arr) is None: return arr - return arr.reindex_like(ref, fill_value=fill_value) + return arr.reindex_like(ref, fill_value={arr.name: fill_value}) def shared_dim_mismatches( diff --git a/test/test_csr.py b/test/test_csr.py index 6ea20f5a3..bb30e86a4 100644 --- a/test/test_csr.py +++ b/test/test_csr.py @@ -8,8 +8,9 @@ import re import tracemalloc import warnings -from collections.abc import Callable -from dataclasses import dataclass +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from dataclasses import dataclass, replace from pathlib import Path from typing import Any @@ -22,10 +23,10 @@ from xarray.core.types import JoinOptions import linopy -from linopy import LinearExpression, Model, Variable +from linopy import LinearExpression, Model, QuadraticExpression, Variable from linopy.constants import TERM_DIM from linopy.constraints import Constraint, ConstraintBase, CSRConstraint -from linopy.csr import CSRLinearExpression +from linopy.csr import CSRLinearExpression, Grid from linopy.semantics import is_v1 from linopy.testing import assert_conequal, assert_linequal, assert_quadequal @@ -233,7 +234,7 @@ def test_namelist_sparse_matches_dense(observed: bool, member_first: bool) -> No csr = sparse._csr assert csr is not None if observed: - assert set(csr.coords) == {"period", "season"} + assert set(csr.grid.aux) == {"period", "season"} else: assert np.isnan(csr.const).sum() == 2 * 2 assert csr.grid.dims == dense.coord_dims @@ -316,13 +317,13 @@ def test_namelist_sparse_observed_keeps_aux_coords_through_merge() -> None: dense = expr.groupby(keys).sum(observed=True) tot = sparse + dense assert tot._csr is not None - assert set(tot._csr.coords) == {"period", "season"} + assert set(tot._csr.grid.aux) == {"period", "season"} assert_linequal(tot, 2.0 * dense) other = dense.assign_coords(region=("group", list("abcd"))) tot = sparse + other assert tot._csr is not None - assert set(tot._csr.coords) == {"period", "season", "region"} + assert set(tot._csr.grid.aux) == {"period", "season", "region"} xr.testing.assert_equal( tot.data.coords.to_dataset(), (dense + other).data.coords.to_dataset() ) @@ -761,16 +762,33 @@ def test_cross_grid_merge_raises_like_dense(kwargs: dict, error: type) -> None: linopy.merge(cross_grid_parts(c, True), **kwargs) -def test_merge_with_aux_coord_operand_raises_like_dense() -> None: +@pytest.mark.parametrize("join", ["outer", "inner", "left", "right"]) +def test_cross_grid_merge_with_aux_coord_operand_stays_csr(join: JoinOptions) -> None: + """The aux coord follows its rows onto the joined grid, like dense.""" require_v1() c = base_model() sparse, dense = cross_grid_parts(c, True), cross_grid_parts(c, False) tag = xr.DataArray(["x", "y"], coords=[dense[1].indexes["bus"]]) tagged = LinearExpression(dense[1].data.assign_coords(tag=tag), c.m) - with pytest.raises(xr.MergeError, match="conflicting values for variable 'tag'"): - linopy.merge([dense[0], tagged], join="outer") - with pytest.raises(xr.MergeError, match="conflicting values for variable 'tag'"): - linopy.merge([sparse[0], tagged], join="outer") + with no_densify(): + res = linopy.merge([sparse[0], tagged], join=join, cls=LinearExpression) + assert res.is_sparse + labels = res.indexes["bus"] + want = tag.to_series().reindex(labels).to_numpy() + assert pd.Series(res.coords["tag"].values).equals(pd.Series(want)) + want_expr = linopy.merge([dense[0], tagged], join=join, cls=LinearExpression) + assert_sparse_matches(res, want_expr) + + +def test_cross_grid_merge_aux_coord_conflict_raises_like_dense() -> None: + require_v1() + sparse, dense = sparse_and_dense("aux") + parts = [sparse, sparse.isel(group=[2, 1])] + with pytest.raises(ValueError) as want: + linopy.merge([dense, dense.isel(group=[2, 1])], join="outer") + with pytest.raises(ValueError) as got: + linopy.merge(parts, join="outer") + assert str(got.value) == str(want.value) def test_cross_grid_merge_with_duplicate_labels_raises_like_dense() -> None: @@ -835,9 +853,9 @@ def cell_matrix( ) -> tuple[np.ndarray, np.ndarray]: """Per-cell coefficient row over raw variable labels, plus the flat constant.""" ds = expr.data - nterm = ds.sizes[TERM_DIM] - vars_ = ds.vars.transpose(*dims, TERM_DIM).to_numpy().reshape(-1, nterm) - coeffs = ds.coeffs.transpose(*dims, TERM_DIM).to_numpy().reshape(-1, nterm) + shape = (int(np.prod([ds.sizes[d] for d in dims], dtype=int)), ds.sizes[TERM_DIM]) + vars_ = ds.vars.transpose(*dims, TERM_DIM).to_numpy().reshape(shape) + coeffs = ds.coeffs.transpose(*dims, TERM_DIM).to_numpy().reshape(shape) active = (vars_ != -1) & ~np.isnan(coeffs) rows = np.broadcast_to(np.arange(len(vars_))[:, None], vars_.shape) matrix = np.zeros((len(vars_), expr.model._xCounter)) @@ -991,7 +1009,7 @@ def test_contracted_keeps_aux_coords_on_kept_dims_only() -> None: c = base_model() csr = tagged_group(c)._csr assert csr is not None - assert set(csr.coords) == {"bus", "tag"} + assert set(csr.grid.aux) == {"bus", "tag"} kept = csr.contracted( flat_operand( loc_operand(csr.grid.indexes["snapshot"]), ("snapshot",), ("loc",) @@ -999,13 +1017,53 @@ def test_contracted_keeps_aux_coords_on_kept_dims_only() -> None: ["snapshot"], [LOC], ) - assert set(kept.coords) == {"bus", "tag"} + assert set(kept.grid.aux) == {"bus", "tag"} dropped = csr.contracted( flat_operand(loc_operand(csr.grid.indexes["group"]), ("group",), ("loc",)), ["group"], [LOC], ) - assert dropped.coords == {} + assert dropped.grid.aux == {} + + +@pytest.mark.parametrize( + ("op", "aux_dims"), + [ + (lambda g: g.renamed({"group": "g"}), {"bus": "g", "tag": "g"}), + ( + lambda g: g.reordered(["snapshot", "group"]), + {"bus": "group", "tag": "group"}, + ), + (lambda g: g.reordered(["snapshot"]), {}), + ( + lambda g: g.with_indexes({"snapshot": [5, 6, 7]}), + {"bus": "group", "tag": "group"}, + ), + (lambda g: g.combined([g], "outer"), {}), + ], + ids=["renamed", "transposed", "dropped", "relabelled", "combined"], +) +def test_grid_ops_carry_aux_coords( + op: Callable[[Grid], Grid], aux_dims: dict[str, str] +) -> None: + require_v1() + csr = tagged_group(base_model())._csr + assert csr is not None + assert {n: d for n, (d, _) in op(csr.grid).aux.items()} == aux_dims + + +def test_grid_conformed_reindexes_aux_coords_and_equality_sees_them() -> None: + require_v1() + csr = tagged_group(base_model())._csr + assert csr is not None + grid = csr.grid + labels = grid.indexes["group"][::-1] + conformed = grid.conformed(grid.with_indexes({"group": labels})) + tags = pd.Series(grid.aux["tag"][1], index=grid.indexes["group"]) + assert np.array_equal(conformed.aux["tag"][1], tags.loc[labels].to_numpy()) + assert grid == replace(grid) + assert grid != replace(grid, aux={}) + assert grid.same_layout(replace(grid, aux={})) def test_matmul_absent_cells_have_const_zero_and_no_terms() -> None: @@ -1219,3 +1277,479 @@ def test_matmul_peak_memory() -> None: tracemalloc.stop() assert res.coord_dims == ("snapshot", "cycle") assert peak < dense_rectangle_bytes / 4 + + +def full_contraction(c: Case) -> LinearExpression: + grouped = (1.0 * c.gen_p).groupby(c.gbus).sum(sparse=True) + ones = xr.DataArray(np.ones(grouped.shape[:2]), coords=grouped.coords) + return grouped @ ones + + +SPARSE_BUILDS: dict[str, Callable[[Case], LinearExpression]] = { + "grouped": lambda c: (c.eff * c.gen_p).groupby(c.gbus).sum(sparse=True), + "aux": tagged_group, + "absent": lambda c: keyed_model()[1].groupby(["period", "season"]).sum(sparse=True), + "zero-dim": full_contraction, +} + + +def sparse_and_dense( + build: str, c: Case | None = None +) -> tuple[LinearExpression, LinearExpression]: + """A sparse build and its dense conversion, the sparse one left sparse.""" + sparse = SPARSE_BUILDS[build](base_model() if c is None else c) + csr = sparse._csr + assert sparse.is_sparse and csr is not None + return sparse, csr.to_dense() + + +METADATA: dict[str, Callable[[LinearExpression], Any]] = { + "shape": lambda e: e.shape, + "size": lambda e: e.size, + "ndim": lambda e: e.ndim, + "sizes": lambda e: dict(e.sizes), + "dims": lambda e: e.dims, + "coord_dims": lambda e: e.coord_dims, + "coord_sizes": lambda e: e.coord_sizes, + "coord_names": lambda e: e.coord_names, + "coords": lambda e: xr.Dataset(coords=e.coords), + "indexes": lambda e: {k: list(v) for k, v in e.indexes.items()}, + "isnull": lambda e: e.isnull(), + "repr": lambda e: repr(e).splitlines()[2:], +} + + +@pytest.mark.parametrize("attr", list(METADATA)) +@pytest.mark.parametrize("build", list(SPARSE_BUILDS)) +def test_metadata_is_served_without_densifying(build: str, attr: str) -> None: + require_v1() + sparse, dense = sparse_and_dense(build) + got, want = METADATA[attr](sparse), METADATA[attr](dense) + assert sparse.is_sparse + if isinstance(got, xr.Dataset | xr.DataArray): + xr.testing.assert_identical(got, want) + else: + assert got == want + + +def test_is_sparse_tracks_backing_and_repr_marks_it() -> None: + require_v1() + c = base_model() + sparse = SPARSE_BUILDS["grouped"](c) + assert sparse.is_sparse + assert repr(sparse).startswith("LinearExpression (sparse) [") + sparse.data + assert not sparse.is_sparse + assert repr(sparse).startswith("LinearExpression [") + assert not (1.0 * c.gen_p).is_sparse + + +DENSIFY_OPS: dict[str, tuple[Callable[[LinearExpression, Case], Any], str]] = { + "data": (lambda e, c: e.data, "`.data` read"), + "merge": (lambda e, c: e + 1.0 * c.flow, "over different dimensions"), + "matmul": (lambda e, c: e @ loc_operand(c.gbus.index), "sharing no dimension"), + "rhs": (lambda e, c: e <= 1.0 * c.gen_p.sum("gen"), "non-constant rhs"), + "mutable": (lambda e, c: (e == c.load).mutable(), "`mutable\\(\\)`"), + "new-dim": (lambda e, c: e * xr.DataArray([1.0, 2.0], coords=[LOC]), "new dim"), + "sum-kwargs": (lambda e, c: e.sum(dims="bus"), "`.data` read"), + "groupby-fallback": ( + lambda e, c: e.groupby(halves(e)).sum(use_fallback=True), + "`.data` read", + ), + "isel-pointwise": ( + lambda e, c: e.isel(bus=xr.DataArray([0, 1], dims="pt")), + "`isel` over MultiIndex labels, or introducing dimensions", + ), + "where-new-dim": ( + lambda e, c: e.where(xr.DataArray([True, False], coords=[LOC])), + "`where` over MultiIndex labels, or introducing dimensions", + ), + "where-callable": ( + lambda e, c: e.where(lambda ds: ds.const > 0), + "`where` with a condition that is no DataArray", + ), +} + + +def halves(e: LinearExpression) -> pd.Series: + """Alternating ``a``/``b`` grouper over the expression's first dimension.""" + index = e.indexes[e.coord_dims[0]] + return pd.Series(np.arange(len(index)) % 2, index=index).map({0: "a", 1: "b"}) + + +@pytest.mark.parametrize("enabled", [True, False], ids=["enabled", "default"]) +@pytest.mark.parametrize("op", list(DENSIFY_OPS)) +def test_warn_on_densify_names_the_reason(op: str, enabled: bool) -> None: + require_v1() + c = base_model() + sparse = SPARSE_BUILDS["grouped"](c) + func, reason = DENSIFY_OPS[op] + with linopy.options as opts, warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always", linopy.PerformanceWarning) + opts.set_value(warn_on_densify=enabled) + func(sparse, c) + notices = [w for w in caught if issubclass(w.category, linopy.PerformanceWarning)] + if not enabled: + assert notices == [] + return + assert len(notices) == 1 + assert re.search(reason, str(notices[0].message)) + assert notices[0].filename == __file__ + + +def grid_operand(e: LinearExpression) -> xr.DataArray: + """Positive values over the expression's full grid.""" + values = np.random.default_rng(1).uniform(1, 2, e.shape[:-1]) + return xr.DataArray(values, coords=[e.indexes[d] for d in e.coord_dims]) + + +def first_dim_operand(e: LinearExpression) -> xr.DataArray: + return grid_operand(e).isel({e.coord_dims[-1]: 0}, drop=True) + + +OPERANDS: dict[str, Callable[[LinearExpression], Any]] = { + "dataarray": grid_operand, + "subset-dims": first_dim_operand, + "series": lambda e: first_dim_operand(e).to_series(), + "ndarray": lambda e: grid_operand(e).to_numpy(), + "zeros": lambda e: xr.zeros_like(grid_operand(e)), + "operand-aux": lambda e: first_dim_operand(e).assign_coords( + extra=(e.coord_dims[0], np.arange(e.shape[0])) + ), +} + +ELEMENTWISE_OPS: dict[ + str, Callable[[LinearExpression, Any], LinearExpression | QuadraticExpression] +] = { + "mul": lambda e, x: e * x, + "rmul": lambda e, x: x * e, + "truediv": lambda e, x: e / x, + "div": lambda e, x: e.div(x, join="left"), + "add": lambda e, x: e + x, + "radd": lambda e, x: x + e, + "sub": lambda e, x: e - x, + "rsub": lambda e, x: x - e, +} + + +@pytest.mark.parametrize( + ("build", "operand", "op"), + [ + (b, o, op) + for b in ["grouped", "aux", "absent"] + for o in OPERANDS + for op in ELEMENTWISE_OPS + if (b, o) != ("absent", "ndarray") and not (o == "zeros" and "div" in op) + ], +) +def test_elementwise_constant_ops_stay_csr_and_match_dense( + build: str, operand: str, op: str +) -> None: + require_v1() + sparse, dense = sparse_and_dense(build) + x = OPERANDS[operand](dense) + func = ELEMENTWISE_OPS[op] + res = func(sparse, x) + assert res.is_sparse + assert_linequal(res, func(dense, x)) + + +@pytest.mark.parametrize("join", ["inner", "outer", "left", "right"]) +@pytest.mark.parametrize("fill_value", [None, linopy.ABSENT], ids=["fill", "absent"]) +@pytest.mark.parametrize("method", ["add", "sub", "mul", "div"]) +def test_elementwise_join_on_mismatched_labels_stays_csr_and_matches_dense( + method: str, fill_value: Any, join: JoinOptions +) -> None: + require_v1() + sparse, dense = sparse_and_dense("grouped") + x = first_dim_operand(dense).isel(bus=[1, 3]).reindex(bus=["bus3", "bus1", "x"]) + x = x.fillna(2.0) + res = getattr(sparse, method)(x, join=join, fill_value=fill_value) + want = getattr(dense, method)(x, join=join, fill_value=fill_value) + assert res.is_sparse + xr.testing.assert_identical(res.isnull(), want.isnull()) + assert_cells_equal(res, want, ("bus", "snapshot")) + + +@pytest.mark.parametrize("op", ["mul", "add"]) +@pytest.mark.parametrize("kind", ["nan", "nan-scalar", "mismatch", "reorder", "aux"]) +def test_elementwise_operand_errors_match_dense(kind: str, op: str) -> None: + """§5, §8 and §11 on the constant are raised as on the dense path.""" + require_v1() + sparse, dense = sparse_and_dense("aux") + x: Any = grid_operand(dense) + if kind == "nan": + x[0, 0] = np.nan + elif kind == "nan-scalar": + x = np.float32("nan") + elif kind == "mismatch": + x = x.isel(group=slice(1, None)) + elif kind == "reorder": + x = x.isel(group=slice(None, None, -1)) + else: + x = x.assign_coords(tag=("group", ["X"] * x.sizes["group"])) + func = ELEMENTWISE_OPS[op] + with pytest.raises(ValueError) as want: + func(dense, x) + with pytest.raises(ValueError) as got: + func(sparse, x) + assert str(got.value) == str(want.value) + assert sparse.is_sparse + + +@pytest.mark.parametrize("op", ["mul", "add"]) +def test_elementwise_operand_over_new_dim_falls_back_to_dense(op: str) -> None: + require_v1() + sparse, dense = sparse_and_dense("grouped") + x = xr.DataArray([1.0, 2.0], coords=[LOC]) + func = ELEMENTWISE_OPS[op] + res = func(sparse, x) + assert not res.is_sparse + assert_linequal(res, func(dense, x)) + + +@pytest.mark.parametrize("join", ["outer", "inner", "left", "right"]) +@pytest.mark.parametrize("op", ["add", "sub", "mul", "div"]) +def test_join_with_constant_keeps_aux_coords_like_dense( + op: str, join: JoinOptions +) -> None: + require_v1() + sparse, dense = sparse_and_dense("aux") + n = sparse.sizes["group"] + x = xr.DataArray( + np.arange(1.0, n + 2), coords=[pd.RangeIndex(1, n + 2, name="group")] + ) + with no_densify(): + res = getattr(sparse, op)(x, join=join) + assert_sparse_matches(res, getattr(dense, op)(x, join=join)) + + +def assert_sparse_matches(res: LinearExpression, want: LinearExpression) -> None: + """CSR backing kept, and equal to the dense reference up to term layout.""" + assert res.is_sparse + assert res.coord_dims == want.coord_dims + xr.testing.assert_identical( + xr.Dataset(coords=res.coords), xr.Dataset(coords=want.coords) + ) + xr.testing.assert_identical(res.isnull(), want.isnull()) + assert_cells_equal(res, want, tuple(map(str, want.coord_dims))) + + +@contextmanager +def no_densify() -> Iterator[None]: + """Fail on any densify notice inside the block.""" + with linopy.options as opts, warnings.catch_warnings(): + warnings.simplefilter("error", linopy.PerformanceWarning) + opts.set_value(warn_on_densify=True) + yield + + +SUM_DIMS: dict[str, Callable[[tuple[str, ...]], Any]] = { + "first": lambda d: d[0], + "last-list": lambda d: [d[-1]], + "leading": lambda d: list(d[:-1]), + "all": lambda d: None, + "all-reversed": lambda d: list(d)[::-1], + "ellipsis": lambda d: ..., + "empty": lambda d: [], + "term": lambda d: [TERM_DIM], +} + + +@pytest.mark.parametrize("dims", list(SUM_DIMS)) +@pytest.mark.parametrize("build", ["grouped", "aux", "absent"]) +def test_sum_stays_csr_and_matches_dense(build: str, dims: str) -> None: + require_v1() + sparse, dense = sparse_and_dense(build) + dim = SUM_DIMS[dims](tuple(map(str, sparse.coord_dims))) + with no_densify(): + res = sparse.sum(dim) + assert_sparse_matches(res, dense.sum(dim)) + + +@pytest.mark.parametrize("drop_zeros", [False, True]) +def test_sum_keeps_explicit_zeros_unless_dropped(drop_zeros: bool) -> None: + require_v1() + c = base_model() + sparse = (0.0 * c.gen_p).groupby(c.gbus).sum(sparse=True) + assert sparse._csr is not None + dense = sparse._csr.to_dense() + res = sparse.sum("snapshot", drop_zeros=drop_zeros) + assert res._csr is not None + assert (res._csr.csr.nnz == 0) == drop_zeros + assert_sparse_matches(res, dense.sum("snapshot", drop_zeros=drop_zeros)) + + +def test_sum_unknown_dim_raises_like_dense() -> None: + require_v1() + sparse, dense = sparse_and_dense("grouped") + with pytest.raises(KeyError) as want: + dense.sum("nodim") + with pytest.raises(KeyError) as got: + sparse.sum("nodim") + assert str(got.value) == str(want.value) + + +CHAINS: dict[str, Callable[[LinearExpression], LinearExpression]] = { + "group-group": lambda e: e.groupby(halves(e)).sum(), + "group-sum": lambda e: e.groupby(halves(e)).sum().sum("snapshot"), + "sum-group": lambda e: e.sum("snapshot").groupby(halves(e)).sum(), + "empty-sum": lambda e: e.isel({dim0(e): []}).sum(dim0(e)), + "empty-group": lambda e: ( + (empty := e.isel({dim0(e): []})).groupby(halves(empty)).sum() + ), +} + + +@pytest.mark.parametrize("chain", list(CHAINS)) +@pytest.mark.parametrize("build", ["grouped", "aux", "absent"]) +def test_chained_groupby_stays_csr_and_matches_dense(build: str, chain: str) -> None: + require_v1() + sparse, dense = sparse_and_dense(build) + func = CHAINS[chain] + with no_densify(): + res = func(sparse) + assert_sparse_matches(res, func(dense)) + + +@pytest.mark.parametrize("observed", [False, True]) +def test_chained_namelist_groupby_on_aux_coords_matches_dense(observed: bool) -> None: + require_v1() + sparse, dense = sparse_and_dense("aux") + res = sparse.groupby(["bus", "tag"]).sum(observed=observed) + assert_sparse_matches(res, dense.groupby(["bus", "tag"]).sum(observed=observed)) + + +def test_chained_groupby_sparse_false_densifies() -> None: + require_v1() + sparse, dense = sparse_and_dense("grouped") + res = sparse.groupby(halves(sparse)).sum(sparse=False) + assert not res.is_sparse + assert_linequal(res, dense.groupby(halves(dense)).sum()) + + +def dim0(e: LinearExpression) -> str: + return str(e.coord_dims[0]) + + +def labels0(e: LinearExpression) -> pd.Index: + return e.indexes[dim0(e)] + + +def alternating(e: LinearExpression) -> xr.DataArray: + """True on every other label of the expression's first dimension.""" + labels = labels0(e) + return xr.DataArray(np.arange(len(labels)) % 2 == 0, coords=[labels]) + + +SELECTIONS: dict[str, Callable[[LinearExpression], LinearExpression]] = { + "sel-scalar": lambda e: e.sel({dim0(e): labels0(e)[1]}), + "sel-scalar-drop": lambda e: e.sel({dim0(e): labels0(e)[1]}, drop=True), + "sel-list": lambda e: e.sel({dim0(e): list(labels0(e)[[2, 0]])}), + "sel-slice": lambda e: e.sel({dim0(e): slice(labels0(e)[1], None)}), + "isel-int": lambda e: e.isel({str(e.coord_dims[-1]): 0}), + "isel-list": lambda e: e.isel({dim0(e): [2, 1, 1]}), + "isel-slice": lambda e: e.isel({dim0(e): slice(None, None, -2)}), + "isel-empty": lambda e: e.isel({dim0(e): []}), + "loc-dict": lambda e: e.loc[{dim0(e): list(labels0(e)[:2])}], + "loc-scalar": lambda e: e.loc[labels0(e)[1]], + "getitem-int": lambda e: e[0], + "getitem-tuple": lambda e: e[1:, [1, 0]], + "where-dim": lambda e: e.where(alternating(e)), + "where-grid": lambda e: e.where(grid_operand(e) > 1.5), + "where-other": lambda e: e.where(alternating(e), 3.0), + "where-drop": lambda e: e.where(alternating(e), drop=True), + "where-drop-mismatch": lambda e: e.where( + alternating(e).isel({dim0(e): slice(1, None)}), drop=True + ), + "sel-scalar-sum": lambda e: e.sel({dim0(e): labels0(e)[1]}).sum(), +} + + +@pytest.mark.parametrize("select", list(SELECTIONS)) +@pytest.mark.parametrize("build", ["grouped", "aux", "absent"]) +def test_selection_stays_csr_and_matches_dense(build: str, select: str) -> None: + require_v1() + sparse, dense = sparse_and_dense(build) + func = SELECTIONS[select] + with no_densify(): + res = func(sparse) + assert_sparse_matches(res, func(dense)) + + +ABSENT_KEEPING_OPS: dict[str, Callable[[LinearExpression], LinearExpression]] = { + "elementwise": lambda e: e * grid_operand(e) + 1.0, + "selection": lambda e: e.where(alternating(e), 3.0).isel(season=[1, 0, 1]), +} + + +@pytest.mark.parametrize("op", list(ABSENT_KEEPING_OPS)) +def test_absent_cells_stay_termless(op: str) -> None: + require_v1() + sparse = SPARSE_BUILDS["absent"](base_model()) + csr = ABSENT_KEEPING_OPS[op](sparse)._csr + assert csr is not None + absent = np.isnan(csr.const) + assert absent.any() + assert (np.diff(csr.csr.indptr)[absent] == 0).all() + + +def test_scalar_selection_coords_merge_like_dense() -> None: + require_v1() + c = base_model() + sparse, dense = sparse_and_dense("grouped", c) + flow = (1.0 * c.flow).groupby(c.bus0).sum() + with no_densify(): + res = sparse.sel(snapshot=1) + flow.sel(snapshot=1) + assert_sparse_matches(res, dense.sel(snapshot=1) + flow.sel(snapshot=1)) + with pytest.raises(ValueError) as want: + dense.sel(snapshot=1) + flow.sel(snapshot=2) + with pytest.raises(ValueError) as got: + sparse.sel(snapshot=1) + flow.sel(snapshot=2) + assert str(got.value) == str(want.value) + + +def test_where_on_mismatched_labels_raises_like_dense() -> None: + require_v1() + sparse, dense = sparse_and_dense("grouped") + cond = alternating(dense).isel(bus=slice(1, None)) + with pytest.raises(ValueError) as want: + dense.where(cond) + with pytest.raises(ValueError) as got: + sparse.where(cond) + assert str(got.value) == str(want.value) + assert sparse.is_sparse + + +SELECTION_FALLBACKS = ["isel-pointwise", "where-new-dim", "where-callable"] + + +@pytest.mark.parametrize("op", SELECTION_FALLBACKS) +def test_selection_fallbacks_match_dense(op: str) -> None: + require_v1() + c = base_model() + sparse, dense = sparse_and_dense("grouped", c) + func = DENSIFY_OPS[op][0] + res = func(sparse, c) + assert not res.is_sparse + assert_linequal(res, func(dense, c)) + + +MASKS: dict[str, Callable[[LinearExpression], Any]] = { + "dim": alternating, + "grid": lambda e: grid_operand(e) > 1.5, + "ndarray": lambda e: (grid_operand(e) > 1.5).to_numpy(), +} + + +@pytest.mark.parametrize("mask", list(MASKS)) +def test_add_constraints_mask_freezes_sparse_and_matches_dense(mask: str) -> None: + require_v1() + c1, c2 = base_model(), base_model() + lhs = c2.balance_lhs(sparse=True) + m = MASKS[mask](lhs) + con1 = c1.m.add_constraints(c1.balance_lhs(False), ">=", c1.load, "bal", mask=m) + with no_densify(): + con2 = c2.m.add_constraints(lhs, ">=", c2.load, "bal", mask=m, freeze=True) + assert isinstance(con2, CSRConstraint) + assert_frozen_equal(con1, con2) diff --git a/test/test_legacy_violations.py b/test/test_legacy_violations.py index 645519ea7..db6b805d4 100644 --- a/test/test_legacy_violations.py +++ b/test/test_legacy_violations.py @@ -2467,6 +2467,29 @@ def test_aux_coord_only_on_one_side_propagates( result = v + w assert "B" in result.coords + @pytest.mark.parametrize("join", ["outer", "inner", "left", "right"]) + @pytest.mark.parametrize("tagged_first", [False, True]) + def test_aux_coord_only_on_one_side_follows_join( + self, m: Model, v: Variable, join: Any, tagged_first: bool + ) -> None: + """The one-sided aux coord follows its rows, NaN where no operand has it.""" + w = m.add_variables(lower=0, coords=[pd.Index([0, 1, 2], name="A")], name="w") + parts = [1 * v, 1 * w][:: -1 if tagged_first else 1] + result = merge(parts, join=join) + want = v.coords["B"].to_series().reindex(result.indexes["A"]) + assert result.coords["B"].to_series().equals(want) + + @pytest.mark.parametrize("join", ["outer", "inner", "left", "right"]) + @pytest.mark.parametrize("op", ["add", "sub", "mul", "div"]) + def test_aux_coord_follows_join_with_constant( + self, v: Variable, op: str, join: Any + ) -> None: + """Labels the join creates get a NaN aux coord, not the constant's fill.""" + const = xr.DataArray([1.0, 2.0, 3.0], coords=[pd.Index([2, 3, 4], name="A")]) + result = getattr(1 * v, op)(const, join=join) + want = v.coords["B"].to_series().reindex(result.indexes["A"]) + assert result.coords["B"].to_series().equals(want) + def test_aux_coord_object_dtype_with_nan_compares_equal( self, m: Model, A: pd.Index ) -> None: