diff --git a/.gitignore b/.gitignore index 7e6d63e2..70b76c39 100644 --- a/.gitignore +++ b/.gitignore @@ -58,3 +58,6 @@ benchmarks/.ipynb_checkpoints/ # direnv .envrc coverage.xml + +# graft's local graph cache — regenerable, not committed (run `graft build`). +/graft/ diff --git a/doc/release_notes.rst b/doc/release_notes.rst index ff01781b..dae3a6dd 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -33,6 +33,10 @@ Upcoming Version * ``add_piecewise_formulation`` gained a ``mask`` parameter declaring which breakpoint slots hold a real breakpoint. It is needed for **ragged** curves — entities with different numbers of breakpoints — which are stored densely with the surplus slots left absent. Under v1 that absence must be declared (``mask=x_pts.notnull()``) rather than read off the NaN padding. (https://github.com/PyPSA/linopy/issues/884) +*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. + *Documentation* * The example notebooks now opt into the v1 arithmetic convention (``linopy.options["semantics"] = "v1"``). The coordinate-alignment and expression tutorials were reworked to teach strict label-based alignment: a mismatch on a shared dimension raises rather than silently filling or pairing by position, and is resolved explicitly with ``.sel`` / ``.reindex`` / ``.assign_coords`` or an explicit ``join=`` on the named ``.add`` / ``.mul`` / ``.le`` / … methods. diff --git a/linopy/constraints.py b/linopy/constraints.py index 20ec381f..f3b301cc 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -76,6 +76,7 @@ PerformanceWarning, SIGNS_pretty, ) +from linopy.csr import Grid, 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 ( @@ -88,8 +89,8 @@ ) if TYPE_CHECKING: + from linopy.csr import CSRLinearExpression from linopy.model import Model - from linopy.sparse_expression import CSRExpression FILL_VALUE = { @@ -599,8 +600,8 @@ class CSRConstraint(ConstraintBase): sign : str or np.ndarray Constraint sign. Either a single str ('=', '<=', '>=') for uniform signs, or a per-row np.ndarray of sign strings for mixed signs. - coords : list of pd.Index - One index per coordinate dimension defining the constraint grid. + grid : Grid + The coordinate grid the constraint rows live on. model : Model The linopy model this constraint belongs to. name : str @@ -617,7 +618,7 @@ class CSRConstraint(ConstraintBase): "_rhs", "_sign", "_scaling", - "_coords", + "_grid", "_model", "_name", "_cindex", @@ -633,7 +634,7 @@ def __init__( active_positions: np.ndarray, rhs: np.ndarray, sign: str | np.ndarray, - coords: list[pd.Index], + grid: Grid, model: Model, name: str = "", cindex: int | None = None, @@ -651,7 +652,7 @@ def __init__( if scaling is not None else np.ones_like(rhs, dtype=float) ) - self._coords = coords + self._grid = grid self._model = model self._name = name self._cindex = cindex @@ -674,11 +675,11 @@ def is_assigned(self) -> bool: @property def shape(self) -> tuple[int, ...]: - return tuple(len(c) for c in self._coords) + return self._grid.shape @property def full_size(self) -> int: - return int(np.prod(shape)) if (shape := self.shape) else 1 + return self._grid.size @property def range(self) -> tuple[int, int]: @@ -700,11 +701,11 @@ def attrs(self) -> dict[str, Any]: @property def coords(self) -> DatasetCoordinates: - return Dataset(coords={c.name: c for c in self._coords}).coords + return Dataset(coords=self._grid.indexes).coords @property def dims(self) -> Frozen[Hashable, int]: - d: dict[Hashable, int] = {c.name: len(c) for c in self._coords} + d: dict[Hashable, int] = dict(zip(self._grid.dims, self._grid.shape)) d[TERM_DIM] = self.nterm return Frozen(d) @@ -719,15 +720,15 @@ def sizes(self) -> Frozen[Hashable, int]: @property def indexes(self) -> Indexes: - return Dataset(coords={c.name: c for c in self._coords}).indexes + return Dataset(coords=self._grid.indexes).indexes @property def nterm(self) -> int: - return int(np.diff(self._csr.indptr).max()) if self._csr.nnz > 0 else 1 + return csr_nterm(self._csr) @property def coord_names(self) -> list[str]: - return [str(c.name) for c in self._coords] + return list(self._grid.dims) def __getstate__(self) -> dict[str, Any]: return self._init_kwargs() @@ -741,7 +742,7 @@ def _init_kwargs(self) -> dict[str, Any]: active_positions=self._active_positions, rhs=self._rhs, sign=self._sign, - coords=self._coords, + grid=self._grid, model=self._model, name=self._name, cindex=self._cindex, @@ -791,7 +792,7 @@ def _active_to_dataarray( ) -> DataArray: full = np.full(self.full_size, fill, dtype=active_values.dtype) full[self.active_positions] = active_values - return DataArray(full.reshape(self.shape), coords=self._coords) + return DataArray(full.reshape(self.shape), coords=self._grid.coords) @property def labels(self) -> DataArray: @@ -826,7 +827,7 @@ def vars(self) -> DataArray: def sign(self) -> DataArray: """Get sign DataArray.""" if isinstance(self._sign, str): - return DataArray(np.full(self.shape, self._sign), coords=self._coords) + return DataArray(np.full(self.shape, self._sign), coords=self._grid.coords) return self._active_to_dataarray(self._sign, fill="") @property @@ -906,23 +907,20 @@ def _to_dataset(self, nterm: int) -> Dataset: ------- Dataset with variables ``labels``, ``coeffs``, ``vars``. """ - csr = self._csr - counts = np.diff(csr.indptr) shape = self.shape full_size = self.full_size - - # Map active row i -> flat position in full shape via con_labels active_positions = self.active_positions - coeffs_2d = np.full((full_size, nterm), np.nan, dtype=csr.dtype) - vars_2d = np.full((full_size, nterm), -1, dtype=self._model._dtypes["labels"]) - if csr.nnz > 0: - row_indices = np.repeat(active_positions, counts) - term_cols = np.arange(csr.nnz) - np.repeat(csr.indptr[:-1], counts) - vars_2d[row_indices, term_cols] = csr.indices - coeffs_2d[row_indices, term_cols] = csr.data + vars_2d, coeffs_2d = csr_to_term_arrays( + self._csr, + nterm, + self._model._dtypes["labels"], + full_size, + active_positions, + coeff_dtype=self._csr.dtype, + ) dim_names = self.coord_names - xr_coords = {c.name: c for c in self._coords} + xr_coords = self._grid.indexes dims_with_term = dim_names + [TERM_DIM] coeffs_da = DataArray( coeffs_2d.reshape(shape + (nterm,)), @@ -940,7 +938,7 @@ def _to_dataset(self, nterm: int) -> Dataset: labels_flat[active_positions] = self.active_labels() ds = assign_multiindex_safe( ds, - labels=DataArray(labels_flat.reshape(shape), coords=self._coords), + labels=DataArray(labels_flat.reshape(shape), coords=self._grid.coords), ) return ds @@ -969,7 +967,7 @@ def data(self) -> Dataset: def __repr__(self) -> str: """Print the constraint without reconstructing the full Dataset.""" max_lines = options["display_max_rows"] - coords = self._coords + coords = self._grid.coords shape = self.shape dim_names = self.coord_names size = self.full_size @@ -1071,10 +1069,10 @@ def to_netcdf_ds(self) -> Dataset: } if isinstance(self._sign, np.ndarray): data_vars["_sign"] = DataArray(self._sign, dims=["_flat"]) - data_vars.update(coords_to_dataset_vars(self._coords)) + data_vars.update(coords_to_dataset_vars(self._grid.coords)) if self._dual is not None: data_vars["dual"] = DataArray(self._dual, dims=["_flat"]) - dim_names = [c.name for c in self._coords] + dim_names = list(self._grid.dims) attrs: dict[str, Any] = { "_linopy_format": "csr", "_csr_columns": "labels", @@ -1127,7 +1125,7 @@ def from_netcdf_ds(cls, ds: Dataset, model: Model, name: str) -> CSRConstraint: coord_dims = attrs["coord_dims"] if isinstance(coord_dims, str): coord_dims = [coord_dims] - coords = coords_from_dataset(ds, coord_dims) + grid = Grid.from_coords(coords_from_dataset(ds, coord_dims)) dual = ds["dual"].values if "dual" in ds else None if "_active_positions" in ds: active_positions = ds["_active_positions"].values @@ -1145,7 +1143,7 @@ def from_netcdf_ds(cls, ds: Dataset, model: Model, name: str) -> CSRConstraint: active_positions, rhs, sign, - coords, + grid, model, name, cindex=cindex, @@ -1231,10 +1229,14 @@ def freeze(self) -> CSRConstraint: """Return self (already immutable).""" return self - def mutable(self) -> Constraint: + def to_dense(self) -> Constraint: """Convert to a Constraint.""" return Constraint(self.data, self._model, self._name) + def mutable(self) -> Constraint: + """Convert to a Constraint.""" + return self.to_dense() + def to_polars(self) -> pl.DataFrame: """Convert frozen constraint to polars DataFrame directly from CSR.""" csr = self._csr @@ -1275,9 +1277,9 @@ def iterate_slices( Yield row-batched sub-Constraints without Dataset reconstruction. Batches are raw CSR slices suitable only for ``to_polars()``. They are - yielded with ``coords=[]`` because batches cover contiguous active rows, - not a contiguous slice of the coordinate grid, so the original coords - would be misleading. Do not call ``.data``, ``.mutable()``, or any + yielded with an empty ``grid`` because batches cover contiguous active + rows, not a contiguous slice of the coordinate grid, so the original + grid would be misleading. Do not call ``.data``, ``.mutable()``, or any coord-dependent property on batch slices. """ nnz = self._csr.nnz @@ -1293,14 +1295,14 @@ def iterate_slices( rhs=self._rhs[rows], sign=sign, scaling=self._scaling[rows], - coords=[], + grid=Grid({}), model=self._model, name=self._name, cindex=self._cindex, ) @classmethod - def from_mutable( + def from_dense( cls, con: Constraint, cindex: int | None = None, @@ -1324,7 +1326,7 @@ def from_mutable( ) csr.sum_duplicates() csr.eliminate_zeros() - coords = [con.indexes[d] for d in con.coord_dims] + grid = Grid.from_coords(con.indexes[d] for d in con.coord_dims) rhs = con.rhs.values.ravel()[active_mask] scaling = con.scaling.values.ravel()[active_mask] sign_vals = con.sign.values.ravel() @@ -1354,7 +1356,7 @@ def from_mutable( active_positions, rhs, sign, - coords, + grid, con.model, con.name, cindex=cindex, @@ -1365,13 +1367,15 @@ def from_mutable( ) @classmethod - def from_csr(cls, expr: CSRExpression, sign: str, rhs: DataArray) -> CSRConstraint: + def from_csr( + cls, expr: CSRLinearExpression, sign: str, rhs: DataArray + ) -> CSRConstraint: """ Staple sign and rhs onto a CSR-backed lhs to form an unassigned CSRConstraint. - The sparse counterpart of :meth:`from_mutable`: instead of converting a + The sparse counterpart of :meth:`from_dense`: instead of converting a dense :class:`Constraint`, it realizes a - :class:`~linopy.sparse_expression.CSRExpression` directly. The expression's + :class:`~linopy.csr.CSRLinearExpression` directly. The expression's label columns are kept as they are, its constant moves to the rhs, and rows with a NaN rhs are inactive, as on the dense path. ``rhs`` must come from :func:`csr_rhs`. @@ -1384,12 +1388,12 @@ def from_csr(cls, expr: CSRExpression, sign: str, rhs: DataArray) -> CSRConstrai active, rhs_flat[active], sign, - coords=[expr.indexes[d] for d in expr.grid_dims], + grid=expr.grid, model=expr.model, ) -def csr_rhs(expr: CSRExpression, rhs: Any) -> DataArray | None: +def csr_rhs(expr: CSRLinearExpression, rhs: Any) -> DataArray | None: """ 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 @@ -1401,12 +1405,12 @@ def csr_rhs(expr: CSRExpression, rhs: Any) -> DataArray | None: da = as_dataarray(rhs) except (TypeError, ValueError): return None - if set(da.dims) & set(HELPER_DIMS) or not set(da.dims) <= set(expr.grid_dims): + if set(da.dims) & set(HELPER_DIMS) or not set(da.dims) <= set(expr.grid.dims): return None return da -def _rhs_grid_values(expr: CSRExpression, rhs: DataArray) -> np.ndarray: +def _rhs_grid_values(expr: CSRLinearExpression, rhs: DataArray) -> np.ndarray: """ Broadcast the rhs onto the expression grid and flatten it, with v1 parity: NaN in the rhs raises (§5) and a reordered or differing index on a @@ -1415,16 +1419,16 @@ def _rhs_grid_values(expr: CSRExpression, rhs: DataArray) -> np.ndarray: if bool(rhs.isnull().any()): check_user_nan() for d in rhs.dims: - if not rhs.get_index(d).equals(expr.indexes[str(d)]): + if not rhs.get_index(d).equals(expr.grid.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: expr.indexes[d] for d in expr.grid_dims if d not in rhs.dims} + missing = {d: i for d, i in expr.grid.indexes.items() if d not in rhs.dims} if missing: rhs = rhs.expand_dims(missing) - return rhs.transpose(*expr.grid_dims).to_numpy().reshape(-1) + return rhs.transpose(*expr.grid.dims).to_numpy().reshape(-1) class Constraint(ConstraintBase): @@ -1938,7 +1942,7 @@ def sanitize_infinities(self) -> Constraint: def freeze(self) -> CSRConstraint: """Convert to an immutable Constraint.""" - return CSRConstraint.from_mutable(self) + return CSRConstraint.from_dense(self) def mutable(self) -> Constraint: """Return self (already mutable).""" diff --git a/linopy/csr.py b/linopy/csr.py new file mode 100644 index 00000000..424ecf71 --- /dev/null +++ b/linopy/csr.py @@ -0,0 +1,476 @@ +""" +The sparse backing of 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:`CSRLinearExpression` 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 documents the CSR structure only, working on plain datasets. The +one bridge back to a dense type is :meth:`CSRLinearExpression.to_dense`, which +wraps the expanded dataset in a :class:`~linopy.expressions.LinearExpression`. +The reverse bridges live at the dense call sites, in +:class:`~linopy.expressions.LinearExpression` and +:meth:`linopy.constraints.CSRConstraint.from_csr`. +""" + +from __future__ import annotations + +from collections.abc import Hashable, Iterable, Mapping +from dataclasses import dataclass, field, 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 HELPER_DIMS, TERM_DIM +from linopy.semantics import absorb_absence, enforce_aux_conflict + +if TYPE_CHECKING: + from linopy.expressions import LinearExpression + from linopy.model import Model + + +@dataclass(frozen=True, eq=False) +class Grid: + """ + The flat coordinate grid a CSR row layout is defined on. + + ``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. + """ + + indexes: dict[str, pd.Index] + + @classmethod + def from_coords(cls, coords: Iterable[pd.Index]) -> Grid: + """Build from one index per dimension, each named after its dim.""" + 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}) + + @property + def dims(self) -> tuple[str, ...]: + return tuple(self.indexes) + + @property + def coords(self) -> list[pd.Index]: + return list(self.indexes.values()) + + @property + def shape(self) -> tuple[int, ...]: + return tuple(len(i) for i in self.indexes.values()) + + @property + def size(self) -> int: + """Number of flat cells; one for the zero-dimensional grid.""" + shape = self.shape + return int(np.prod(shape, dtype=np.int64)) if shape else 1 + + @property + def strides(self) -> dict[str, int]: + """C-order row stride per dimension.""" + shape = self.shape + return { + d: int(np.prod(shape[i + 1 :], dtype=np.int64)) + for i, d in enumerate(self.dims) + } + + @property + def is_unique(self) -> bool: + """Whether every dimension's labels are unique.""" + return all(i.is_unique for i in self.indexes.values()) + + def indexer(self, other: Grid) -> tuple[np.ndarray, np.ndarray]: + """ + Map ``other``'s cells onto this grid: the flat target row of each + source cell, and a mask of the cells whose labels all survive. + """ + strides = self.strides + positions = { + d: self.indexes[d].get_indexer(i) for d, i in other.indexes.items() + } + valid = _outer_sum([pos == -1 for pos in positions.values()]) == 0 + row_map = _outer_sum([pos * strides[d] for d, pos in positions.items()]) + return row_map, valid + + 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() + } + ) + + 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}) + + def with_indexes(self, indexers: Mapping[Hashable, 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() + } + ) + + 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. + """ + combine = pd.Index.union if how == "outer" else pd.Index.intersection + others = list(others) + indexes = {} + for d in self.dims: + index = self.indexes[d] + for other in others: + index = combine(index, other.indexes[d]) + indexes[d] = pd.Index(index, name=d) + return Grid(indexes) + + 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() + ) + + +@dataclass(frozen=True) +class CSRLinearExpression: + """ + An expression as ``A @ x + c`` over a fixed coordinate grid. + + ``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. + """ + + 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, ...]: + return self.grid.shape + + @property + def n_cells(self) -> int: + return self.csr.shape[0] + + @property + def nterm(self) -> int: + return csr_nterm(self.csr) + + @classmethod + def from_grouper( + cls, + ds: Dataset, + model: Model, + grouper: pd.Series | pd.DataFrame, + group_dim: str, + stacked: bool, + coord_dims: tuple[str, ...], + ) -> CSRLinearExpression: + """ + 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. A + DataFrame grouper (one column per key) yields one grid dim per key + -- the cartesian grid, absent combinations being empty cells -- or, + ``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. + """ + member_dim = str(grouper.index.name) + if member_dim in ds.indexes: + grouper = grouper.reindex(ds.indexes[member_dim]) + elif len(grouper) != ds.sizes[member_dim]: + raise ValueError(f"grouper length does not match dimension {member_dim!r}") + if grouper.isna().to_numpy().any(): + raise ValueError( + "Cannot group by a pandas object containing NaN values. " + "Drop or fill the corresponding entries before grouping." + ) + frame = grouper if isinstance(grouper, pd.DataFrame) else grouper.to_frame() + 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]] = {} + if len(keys) == 1: + codes, uniques = pd.factorize(frame.iloc[:, 0], sort=True) + scatter_codes[group_dim] = codes + indexes[group_dim] = pd.Index(uniques, name=group_dim) + elif stacked: + 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 = { + k: (group_dim, uniques.get_level_values(i).to_numpy()) + for i, k in enumerate(keys) + } + else: + for col, k in zip(frame.columns, keys): + codes, uniques = pd.factorize(frame[col], sort=True) + scatter_codes[k] = codes + indexes[k] = pd.Index(uniques, name=k) + new_dims = tuple(scatter_codes) + grid_dims = tuple( + d for dim in coord_dims for d in (new_dims if dim == member_dim else (dim,)) + ) + 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 + ) + + @classmethod + def from_dense(cls, ds: Dataset, model: Model) -> CSRLinearExpression: + """Convert a dense expression to CSR form on its own coordinate grid.""" + grid_dims = tuple(str(d) for d in ds.coeffs.dims if d not in HELPER_DIMS) + 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) + + @classmethod + def _from_scatter( + cls, + ds: Dataset, + model: Model, + grid: Grid, + 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``): + ``member_dim`` lands in the contiguous block of grid dims named by + ``scatter_codes`` (one row-position array per dim), every other grid + dim maps one-to-one, and the COO→CSR conversion sums duplicates -- + which is the group sum. Cells no member lands in stay absent (NaN + const). With ``skipna`` the constant is reduced as by the dense group + kernel (NaN members count as 0); without it an absent cell (NaN const) + 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 + ] + ) + + 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 = grid.size + coo = scipy.sparse.coo_array( + (coeffs[keep], (rows[keep], vars_[keep])), + shape=(full_size, model._xCounter), + ) + + const_vals = ds.const.transpose(*transposed).to_numpy().reshape(-1) + if skipna: + const_vals = np.where(np.isnan(const_vals), 0.0, const_vals) + const = np.full(full_size, np.nan) + const[cell_rows] = 0.0 + np.add.at(const, cell_rows, const_vals) + + return cls(scipy.sparse.csr_array(coo), const, grid, model, coords) + + def scaled(self, factor: float) -> CSRLinearExpression: + return replace(self, csr=self.csr * factor, const=self.const * factor) + + 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). + """ + row_map, valid = grid.indexer(self.grid) + + coo = self.csr.tocoo() + keep = valid[coo.coords[0]] + rows = row_map[coo.coords[0][keep]] + cols = coo.coords[1][keep] + n_cells = grid.size + coo = scipy.sparse.coo_array( + (coo.data[keep], (rows, cols)), shape=(n_cells, self.csr.shape[1]) + ) + 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, + ) + + def filled(self, value: float) -> CSRLinearExpression: + """Resolve absent cells (NaN const) to a constant; terms untouched.""" + const = np.where(np.isnan(self.const), value, self.const) + return replace(self, const=const) + + 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) + + def same_grid(self, other: CSRLinearExpression) -> bool: + return self.grid == other.grid + + def added(self, other: CSRLinearExpression) -> CSRLinearExpression: + """ + 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 cell absent in either operand + is absent in the sum and carries no terms (v1 dead-term invariant). + Auxiliary coordinates propagate and conflicting ones raise (§11), as + on the dense path. + """ + const = self.const + other.const + 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]) + present = ~np.isnan(const)[rows] + 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 + ) + + def to_dense(self) -> LinearExpression: + """ + Expand to the dense equivalent in canonical form: terms label-ordered, + duplicates summed, padded to the widest cell with the usual fill. + Absent cells (NaN const) carry no terms, per the v1 dead-term invariant. + The expanded dataset is wrapped in a :class:`LinearExpression`. + """ + from linopy.expressions import LinearExpression + + csr = self.csr.copy() + csr.sort_indices() + nterm = self.nterm + vars_flat, coeffs_flat = csr_to_term_arrays( + csr, nterm, self.model._dtypes["labels"] + ) + + shape = self.grid.shape + dims = (*self.grid.dims, TERM_DIM) + ds = Dataset( + { + "coeffs": (dims, coeffs_flat.reshape(*shape, nterm)), + "vars": (dims, vars_flat.reshape(*shape, nterm)), + "const": (self.grid.dims, self.const.reshape(shape)), + }, + coords=self.grid.indexes | self.coords, + ) + return LinearExpression(absorb_absence(ds), self.model) + + +def csr_nterm(csr: scipy.sparse.csr_array) -> int: + """Widest CSR row, floored at one term.""" + return max(int(np.diff(csr.indptr).max(initial=0)), 1) + + +def csr_to_term_arrays( + csr: scipy.sparse.csr_array, + nterm: int, + label_dtype: Any, + n_rows: int | None = None, + row_positions: np.ndarray | None = None, + coeff_dtype: Any = float, +) -> tuple[np.ndarray, np.ndarray]: + """ + Expand CSR rows into padded ``(n_rows, nterm)`` term arrays. + + Returns variable labels (absent terms filled with ``-1``) and coefficients + (absent terms filled with NaN). ``row_positions`` places CSR row ``i`` at + output row ``row_positions[i]`` of an ``n_rows``-row output, for a CSR that + holds only the active rows of a larger grid. + """ + counts = np.diff(csr.indptr) + n_rows = csr.shape[0] if n_rows is None else n_rows + vars_ = np.full((n_rows, nterm), -1, dtype=label_dtype) + coeffs = np.full((n_rows, nterm), np.nan, dtype=coeff_dtype) + if csr.nnz: + positions = np.arange(csr.shape[0]) if row_positions is None else row_positions + rows = np.repeat(positions, counts) + cols = np.arange(csr.nnz) - np.repeat(csr.indptr[:-1], counts) + vars_[rows, cols] = csr.indices + coeffs[rows, cols] = csr.data + 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 _outer_sum(axis_positions: list[np.ndarray]) -> np.ndarray: + """Outer sum of per-axis offsets, flattened in C order.""" + cells = np.zeros((), dtype=np.int64) + for pos in axis_positions: + cells = cells[..., None] + pos + return cells.reshape(-1) diff --git a/linopy/expressions.py b/linopy/expressions.py index 9817ba50..65472232 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -20,7 +20,7 @@ Mapping, Sequence, ) -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from itertools import product, zip_longest from typing import ( TYPE_CHECKING, @@ -106,6 +106,7 @@ STACKED_TERM_DIM, TERM_DIM, ) +from linopy.csr import CSRLinearExpression, _aux_coords from linopy.semantics import ( AbsentType, FillValueLike, @@ -560,7 +561,7 @@ def sum( over an existing dimension; with a name list and ``observed=True`` the CSR result stays compact over the observed key combinations. Requires v1 semantics. Defaults to - ``linopy.options["sparse_groupby"]``. See :mod:`linopy.sparse_expression`. + ``linopy.options["sparse_groupby"]``. See :mod:`linopy.csr`. observed : bool Only applies when grouping by a list of coordinate names. If True, keep the result stacked over the observed key combinations (a @@ -610,9 +611,6 @@ def sum( and grouper.index.name in self.data.dims ) if supported: - from linopy.sparse_expression import CSRExpression - - expr = LinearExpression(self.data, self.model) stacked = observed or multikey_frame is None group_name = ( "group" @@ -622,8 +620,8 @@ def sum( coord_dims = tuple( str(d) for d in self.data.coeffs.dims if d != TERM_DIM ) - csr = CSRExpression.from_grouper( - expr, grouper, group_name, stacked, coord_dims + csr = CSRLinearExpression.from_grouper( + self.data, self.model, grouper, group_name, stacked, coord_dims ) return LinearExpression._from_csr(csr, self.model) if explicit_sparse: @@ -831,7 +829,7 @@ def sum(self, **kwargs: Any) -> LinearExpression: class BaseExpression(ABC): - __slots__ = ("_data", "_model", "_csr") + __slots__ = ("_data", "_model") __array_ufunc__ = None __array_priority__ = 10000 __pandas_priority__ = 10000 @@ -906,7 +904,6 @@ 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._csr = None def __repr__(self) -> str: """ @@ -1007,8 +1004,6 @@ def __neg__(self) -> Self: """ Get the negative of the expression. """ - if self._csr is not None: - return self._from_csr(self._csr.scaled(-1.0), self._model) return self.assign_multiindex_safe(coeffs=-self.coeffs, const=-self.const) def _multiply_by_linear_expression( @@ -1594,20 +1589,8 @@ def name(self) -> str: @property def data(self) -> Dataset: - if self._data is None and self._csr is not None: - self._data = self._csr.materialize().data - self._csr = None return self._data - @classmethod - def _from_csr(cls, csr: Any, model: Model) -> Self: - """Construct an expression backed by a CSRExpression.""" - obj = cls.__new__(cls) - obj._model = model - obj._data = None # type: ignore[assignment] - obj._csr = csr - return obj - @property def model(self) -> Model: return self._model @@ -1619,8 +1602,6 @@ def dims(self) -> tuple[Hashable, ...]: @property def coord_dims(self) -> tuple[Hashable, ...]: - if self._data is None and self._csr is not None: - return tuple(self._csr.grid_dims) return tuple(k for k in self.dims if k not in HELPER_DIMS) @property @@ -1817,11 +1798,6 @@ 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._csr is not None and isinstance(sign, str): - rhs_da = constraints.csr_rhs(self._csr, rhs) - if rhs_da is not None: - return constraints.CSRConstraint.from_csr(self._csr, sign, rhs_da) - rhs = as_constant(rhs) if self.is_constant and is_constant(rhs): raise ValueError( @@ -1985,13 +1961,6 @@ def fillna( ``to_linexpr``), which still holds the absence labels. """ value = _expr_unwrap(value) - csr = self._csr - if ( - csr is not None - and isinstance(value, np.floating | np.integer | int | float) - and not isinstance(value, bool) - ): - return type(self)._from_csr(csr.filled(float(value)), self._model) if isinstance(value, DataArray | np.floating | np.integer | int | float): value = {"const": value} return self.__class__(self.data.fillna(value), self.model) @@ -2089,8 +2058,6 @@ def nterm(self) -> int: """ Get the number of terms in the linear expression. """ - if self._csr is not None: - return self._csr.nterm return len(self.data._term) @property @@ -2355,6 +2322,79 @@ class LinearExpression(BaseExpression): """ + __slots__ = ("_csr",) + + def __init__(self, data: Dataset | Any | None, model: Model) -> None: + super().__init__(data, model) + self._csr: CSRLinearExpression | None = None + + @classmethod + def _from_csr(cls, csr: CSRLinearExpression, model: Model) -> Self: + """Construct an expression backed by a CSRLinearExpression.""" + obj = cls.__new__(cls) + obj._model = model + obj._data = None # type: ignore[assignment] + obj._csr = csr + return obj + + @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 + return self._data + + @property + def coord_dims(self) -> tuple[Hashable, ...]: + if self._data is None and self._csr is not None: + return self._csr.grid.dims + return super().coord_dims + + @property + def nterm(self) -> int: + """ + Get the number of terms in the linear expression. + """ + if self._csr is not None: + return self._csr.nterm + return super().nterm + + def __neg__(self) -> Self: + """ + Get the negative of the expression. + """ + if self._csr is not None: + return self._from_csr(self._csr.scaled(-1.0), self._model) + return super().__neg__() + + def fillna( + self, + value: int + | float + | DataArray + | Dataset + | LinearExpression + | dict[str, float | int | DataArray], + ) -> Self: + csr = self._csr + unwrapped = _expr_unwrap(value) + if ( + csr is not None + and isinstance(unwrapped, np.floating | np.integer | int | float) + and not isinstance(unwrapped, bool) + ): + return type(self)._from_csr(csr.filled(float(unwrapped)), self._model) + return super().fillna(value) + + def to_constraint( + self, sign: SignLike, rhs: SideLike, join: JoinOptions | None = None + ) -> 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: + return constraints.CSRConstraint.from_csr(self._csr, sign, rhs_da) + return super().to_constraint(sign, rhs, join) + @overload def __add__( self, @@ -2540,17 +2580,14 @@ def reindex( csr = self._csr if ( csr is not None - and set(indexers) <= set(csr.grid_dims) + and set(indexers) <= set(csr.grid.dims) and method is None and tolerance is None and copy and fill_value is self._fill_value ): - indexes = { - d: pd.Index(indexers.get(d, csr.indexes[d]), name=d) - for d in csr.grid_dims - } - return type(self)._from_csr(csr.reindexed(indexes), self._model) + grid = csr.grid.with_indexes(indexers) + return type(self)._from_csr(csr.reindexed(grid), self._model) return super().reindex( indexers, method=method, @@ -2570,7 +2607,7 @@ def rename( """ name_dict = either_dict_or_kwargs(name_dict, names, "rename") csr = self._csr - if csr is not None and set(name_dict) <= set(csr.grid_dims): + if csr is not None and set(name_dict) <= set(csr.grid.dims): relabel = {str(k): str(v) for k, v in name_dict.items()} return type(self)._from_csr(csr.renamed(relabel), self._model) return super().rename(name_dict) @@ -3177,6 +3214,87 @@ def as_expression( return LinearExpression(obj, model) +def _aligned( + csrs: list[CSRLinearExpression], join: JoinOptions | None, fill: float +) -> list[CSRLinearExpression] | None: + """ + 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. + """ + template = csrs[0].grid + dims = template.dims + if any(not p.grid.is_unique for p in csrs): + return None + if join == "override": + if any(p.grid.dims != dims or p.grid.shape != template.shape for p in csrs): + return None + return [replace(p, grid=template) for p in csrs] + if join in ("left", "right"): + source = csrs[0] if join == "left" else csrs[-1] + grid = source.grid.reordered(dims) + elif join in ("outer", "inner"): + grid = template.combined([p.grid for p in csrs[1:]], join) + else: + return None + return [p.reindexed(grid, fill) for p in csrs] + + +def _try_csr_merge( + exprs: Any, + dim: str, + join: JoinOptions | None, + fill_value: FillValueLike, + kwargs: dict[str, Any], +) -> LinearExpression | None: + """ + Sparse branch of :func:`merge`: combine plain LinearExpressions over one + set of grid dimensions (CSR-backed or dense-convertible) as sparse matrix + 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 + None to fall through to the dense path. + """ + if dim != TERM_DIM or kwargs: + 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): + return None + dims = set(exprs[0].coord_dims) + if any(set(e.coord_dims) != dims for e in exprs[1:]): + return None + for e in exprs: + if e._csr is None and set(e.data.coords) - dims != set( + _aux_coords(e.data, dims) + ): + return None + + csrs = [e._csr or CSRLinearExpression.from_dense(e.data, e.model) for e in exprs] + template = csrs[0] + order = template.grid.dims + csrs = [ + p.reindexed(p.grid.reordered(order)) if p.grid.dims != order else p + 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 + aligned = _aligned(csrs, join, join_fill(fill_value, 0.0)) + if aligned is None: + return None + csrs = aligned + + combined = csrs[0] + for csr in csrs[1:]: + combined = combined.added(csr) + return LinearExpression._from_csr(combined, exprs[0].model) + + Mergeable: TypeAlias = BaseExpression | variables.Variable | Dataset @@ -3293,9 +3411,7 @@ 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( + csr_result = _try_csr_merge( exprs, dim=dim, join=join, fill_value=fill_value, kwargs=kwargs ) if csr_result is not None: diff --git a/linopy/sparse_expression.py b/linopy/sparse_expression.py deleted file mode 100644 index bdd2697d..00000000 --- a/linopy/sparse_expression.py +++ /dev/null @@ -1,446 +0,0 @@ -""" -The sparse backing of 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:`CSRExpression` 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 -CSR expression to form a :class:`~linopy.constraints.CSRConstraint` lives in -:meth:`linopy.constraints.CSRConstraint.from_csr`. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field, replace -from typing import TYPE_CHECKING, Any - -import numpy as np -import pandas as pd -import scipy.sparse -from xarray import Dataset -from xarray.core.types import JoinOptions - -from linopy.constants import TERM_DIM -from linopy.semantics import FillValueLike, enforce_aux_conflict, join_fill - -if TYPE_CHECKING: - from linopy.expressions import LinearExpression - from linopy.model import Model - - -@dataclass(frozen=True) -class CSRExpression: - """ - 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, 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. - """ - - csr: scipy.sparse.csr_array - const: np.ndarray - grid_dims: tuple[str, ...] - indexes: dict[str, pd.Index] - model: Model - coords: dict[str, tuple[str, np.ndarray]] = field(default_factory=dict) - - @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] - - @property - def nterm(self) -> int: - return max(int(np.diff(self.csr.indptr).max(initial=0)), 1) - - @classmethod - def from_grouper( - cls, - expr: LinearExpression, - grouper: pd.Series | pd.DataFrame, - group_dim: str, - stacked: bool, - coord_dims: tuple[str, ...], - ) -> CSRExpression: - """ - 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. A - DataFrame grouper (one column per key) yields one grid dim per key - -- the cartesian grid, absent combinations being empty cells -- or, - ``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. - """ - 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}") - if grouper.isna().to_numpy().any(): - raise ValueError( - "Cannot group by a pandas object containing NaN values. " - "Drop or fill the corresponding entries before grouping." - ) - frame = grouper if isinstance(grouper, pd.DataFrame) else grouper.to_frame() - 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]] = {} - if len(keys) == 1: - codes, uniques = pd.factorize(frame.iloc[:, 0], sort=True) - scatter_codes[group_dim] = codes - indexes[group_dim] = pd.Index(uniques, name=group_dim) - elif stacked: - 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 = { - k: (group_dim, uniques.get_level_values(i).to_numpy()) - for i, k in enumerate(keys) - } - else: - for col, k in zip(frame.columns, keys): - codes, uniques = pd.factorize(frame[col], sort=True) - scatter_codes[k] = codes - indexes[k] = pd.Index(uniques, name=k) - new_dims = tuple(scatter_codes) - grid_dims = tuple( - d for dim in coord_dims for d in (new_dims if dim == member_dim else (dim,)) - ) - for d in coord_dims: - if d != member_dim: - indexes[d] = expr.data.get_index(d).rename(d) - coords |= _aux_coords(expr, set(coord_dims) - {member_dim}) - return cls._from_scatter( - expr, grid_dims, indexes, member_dim, scatter_codes, True, coords - ) - - @classmethod - def from_expression(cls, expr: LinearExpression) -> CSRExpression: - """Convert a dense expression to CSR form on its own coordinate grid.""" - grid_dims = tuple(str(d) for d in expr.coord_dims) - indexes = {d: expr.data.get_index(d).rename(d) for d in grid_dims} - first = grid_dims[0] - codes = {first: np.arange(len(indexes[first]))} - coords = _aux_coords(expr, set(grid_dims)) - return cls._from_scatter(expr, grid_dims, indexes, first, codes, False, coords) - - @classmethod - def _from_scatter( - cls, - expr: LinearExpression, - grid_dims: tuple[str, ...], - indexes: dict[str, pd.Index], - member_dim: str, - scatter_codes: dict[str, np.ndarray], - skipna: bool, - coords: dict[str, tuple[str, np.ndarray]], - ) -> CSRExpression: - """ - Scatter an expression's terms into grid rows (conceptually ``G @ A``): - ``member_dim`` lands in the contiguous block of grid dims named by - ``scatter_codes`` (one row-position array per dim), every other grid - dim maps one-to-one, and the COO→CSR conversion sums duplicates -- - which is the group sum. Cells no member lands in stay absent (NaN - const). With ``skipna`` the constant is reduced as by the dense group - kernel (NaN members count as 0); without it an absent cell (NaN const) - stays absent, as on the dense v1 merge path. - """ - ds = expr.data - shape, strides = _grid_layout(grid_dims, indexes) - stride = dict(zip(grid_dims, 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 = _flat_cells( - [ - member_rows - if d == member_dim - else np.arange(len(indexes[d])) * stride[d] - for d in transposed - ] - ) - - 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) - if skipna: - const_vals = np.where(np.isnan(const_vals), 0.0, const_vals) - const = np.full(full_size, np.nan) - const[cell_rows] = 0.0 - np.add.at(const, cell_rows, const_vals) - - return cls( - scipy.sparse.csr_array(coo), const, grid_dims, indexes, expr.model, coords - ) - - def scaled(self, factor: float) -> CSRExpression: - return replace(self, csr=self.csr * factor, const=self.const * factor) - - def reindexed( - self, - indexes: dict[str, pd.Index], - grid_dims: tuple[str, ...] | None = None, - fill: float = np.nan, - ) -> CSRExpression: - """ - Remap rows onto new per-dim indexes, optionally in a new dim order, - without the dense rectangle: dropped labels vanish, new labels get - ``fill`` as their constant (NaN: absent cells). - """ - grid_dims = grid_dims or self.grid_dims - shape, strides = _grid_layout(grid_dims, indexes) - stride_of = dict(zip(grid_dims, strides)) - positions = [indexes[d].get_indexer(self.indexes[d]) for d in self.grid_dims] - valid = _flat_cells([pos == -1 for pos in positions]) == 0 - row_map = _flat_cells( - [pos * stride_of[d] for pos, d in zip(positions, self.grid_dims)] - ) - - coo = self.csr.tocoo() - keep = valid[coo.coords[0]] - rows = row_map[coo.coords[0][keep]] - cols = coo.coords[1][keep] - n_cells = int(np.prod(shape, dtype=np.int64)) - coo = scipy.sparse.coo_array( - (coo.data[keep], (rows, cols)), shape=(n_cells, self.csr.shape[1]) - ) - const = np.full(n_cells, fill) - const[row_map[valid]] = self.const[valid] - coords = { - name: ( - d, - pd.Series(v, index=self.indexes[d]).reindex(indexes[d]).to_numpy(), - ) - for name, (d, v) in self.coords.items() - } - return replace( - self, - csr=scipy.sparse.csr_array(coo), - const=const, - grid_dims=grid_dims, - indexes=indexes, - coords=coords, - ) - - def filled(self, value: float) -> CSRExpression: - """Resolve absent cells (NaN const) to a constant; terms untouched.""" - const = np.where(np.isnan(self.const), value, self.const) - return replace(self, const=const) - - def renamed(self, names: dict[str, str]) -> CSRExpression: - """Relabel grid dims; the CSR row layout is unchanged.""" - grid_dims = tuple(names.get(d, d) for d in self.grid_dims) - indexes = { - names.get(d, d): self.indexes[d].rename(names.get(d, d)) - for d in self.grid_dims - } - coords = {n: (names.get(d, d), v) for n, (d, v) in self.coords.items()} - return replace(self, grid_dims=grid_dims, indexes=indexes, coords=coords) - - def same_grid(self, other: CSRExpression) -> 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: CSRExpression) -> CSRExpression: - """ - 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 cell absent in either operand - is absent in the sum and carries no terms (v1 dead-term invariant). - Auxiliary coordinates propagate and conflicting ones raise (§11), as - on the dense path. - """ - const = self.const + other.const - 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]) - present = ~np.isnan(const)[rows] - 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 - ) - - 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. - Absent cells (NaN const) carry no terms, per the v1 dead-term invariant. - """ - from linopy.expressions import LinearExpression - from linopy.semantics import absorb_absence - - 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} | self.coords, - ) - return LinearExpression(absorb_absence(ds), self.model) - - -def _aux_coords( - expr: LinearExpression, dims: set[str] -) -> dict[str, tuple[str, np.ndarray]]: - """One-dimensional auxiliary coordinates of ``expr`` lying on ``dims``.""" - return { - str(n): (str(c.dims[0]), c.to_numpy()) - for n, c in expr.data.coords.items() - if n not in expr.data.dims and len(c.dims) == 1 and str(c.dims[0]) in dims - } - - -def _grid_layout( - grid_dims: tuple[str, ...], indexes: dict[str, pd.Index] -) -> tuple[tuple[int, ...], list[int]]: - """C-order shape and row strides of the grid.""" - 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))] - return shape, strides - - -def _flat_cells(axis_positions: list[np.ndarray]) -> np.ndarray: - """Outer sum of per-axis offsets, flattened in C order.""" - cells = np.zeros((), dtype=np.int64) - for pos in axis_positions: - cells = cells[..., None] + pos - return cells.reshape(-1) - - -def _aligned( - csrs: list[CSRExpression], join: JoinOptions | None, fill: float -) -> list[CSRExpression] | None: - """ - 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. - """ - template = csrs[0] - dims = template.grid_dims - if any(not p.indexes[d].is_unique for p in csrs for d in dims): - return None - if join == "override": - if any(p.grid_dims != dims or p.shape != template.shape for p in csrs): - return None - return [replace(p, indexes=template.indexes) for p in csrs] - if join == "left": - indexes = template.indexes - elif join == "right": - indexes = csrs[-1].indexes - elif join in ("outer", "inner"): - combine = pd.Index.union if join == "outer" else pd.Index.intersection - indexes = {} - for d in dims: - index = template.indexes[d] - for p in csrs[1:]: - index = combine(index, p.indexes[d]) - indexes[d] = pd.Index(index, name=d) - else: - return None - return [p.reindexed(indexes, dims, fill) for p in csrs] - - -def try_csr_merge( - exprs: Any, - dim: str, - join: JoinOptions | None, - fill_value: FillValueLike, - kwargs: dict[str, Any], -) -> LinearExpression | None: - """ - Sparse branch of :func:`linopy.expressions.merge`: combine plain - LinearExpressions over one set of grid dimensions (CSR-backed or - dense-convertible) as sparse matrix addition. 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 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 - if all(e._csr is None for e in exprs): - return None - dims = set(exprs[0].coord_dims) - if any(set(e.coord_dims) != dims for e in exprs[1:]): - return None - for e in exprs: - if e._csr is None and set(e.data.coords) - dims != set(_aux_coords(e, dims)): - return None - - csrs = [e._csr or CSRExpression.from_expression(e) for e in exprs] - template = csrs[0] - if not all(template.same_grid(p) for p in csrs[1:]): - if any(p.coords for p in csrs): - return None - aligned = _aligned(csrs, join, join_fill(fill_value, 0.0)) - if aligned is None: - return None - csrs = aligned - - combined = csrs[0] - for csr in csrs[1:]: - combined = combined.add(csr) - return LinearExpression._from_csr(combined, exprs[0].model) diff --git a/test/test_constraint.py b/test/test_constraint.py index cc56bdca..24ae7f93 100644 --- a/test/test_constraint.py +++ b/test/test_constraint.py @@ -899,7 +899,7 @@ def test_freeze_mutable_roundtrip(m: Model) -> None: assert isinstance(frozen, linopy.constraints.CSRConstraint) mc = frozen.mutable() assert isinstance(mc, Constraint) - refrozen = linopy.constraints.CSRConstraint.from_mutable(mc, frozen._cindex) + refrozen = linopy.constraints.CSRConstraint.from_dense(mc, frozen._cindex) assert_equal(frozen.labels, refrozen.labels) assert_equal(frozen.rhs, refrozen.rhs) assert_equal(frozen.sign, refrozen.sign) @@ -907,6 +907,17 @@ def test_freeze_mutable_roundtrip(m: Model) -> None: np.testing.assert_array_equal(frozen.active_labels(), refrozen.active_labels()) +def test_frozen_coeff_dtype_preserved() -> None: + m = Model() + i = pd.RangeIndex(4, name="i") + x = m.add_variables(coords=[i], name="x") + coeff = xr.DataArray(np.arange(1, 5, dtype=np.float32), coords=[i]) + frozen = m.add_constraints(coeff * x >= 1, name="c", freeze=True) + assert frozen._csr.dtype == np.float32 + assert frozen.coeffs.dtype == np.float32 + assert frozen.mutable().coeffs.dtype == np.float32 + + def test_frozen_csr_stores_variable_labels(m: Model, x: linopy.Variable) -> None: frozen = m.constraints["c"] assert isinstance(frozen, linopy.constraints.CSRConstraint) @@ -959,20 +970,20 @@ def test_freeze_mutable_roundtrip_with_masking() -> None: frozen = m.constraints["c"] assert isinstance(frozen, linopy.constraints.CSRConstraint) mc = frozen.mutable() - refrozen = linopy.constraints.CSRConstraint.from_mutable(mc, frozen._cindex) + refrozen = linopy.constraints.CSRConstraint.from_dense(mc, frozen._cindex) assert_equal(frozen.labels, refrozen.labels) assert_equal(frozen.rhs, refrozen.rhs) assert frozen.ncons == refrozen.ncons == 3 -def test_from_mutable_mixed_signs() -> None: +def test_from_dense_mixed_signs() -> None: m = Model() x = m.add_variables(coords=[pd.RangeIndex(3, name="i")], name="x") m.add_constraints(x >= 0, name="mixed", freeze=False) mc = m.constraints["mixed"] assert isinstance(mc, Constraint) mc._data["sign"] = xr.DataArray(["<=", ">=", "<="], dims=["i"]) - frozen = linopy.constraints.CSRConstraint.from_mutable(mc) + frozen = linopy.constraints.CSRConstraint.from_dense(mc) assert isinstance(frozen._sign, np.ndarray) assert list(frozen._sign) == ["<=", ">=", "<="] assert_equal(frozen.sign, mc.sign) diff --git a/test/test_sparse_groupby.py b/test/test_csr.py similarity index 93% rename from test/test_sparse_groupby.py rename to test/test_csr.py index 88a7b60f..8b498e56 100644 --- a/test/test_sparse_groupby.py +++ b/test/test_csr.py @@ -1,5 +1,5 @@ """ -Tests for sparse groupby-sum (linopy.sparse_expression): type stability, transparent +Tests for sparse groupby-sum (linopy.csr): type stability, transparent materialization, and direct CSR realization under freeze. v1-only feature. """ @@ -167,9 +167,11 @@ def test_merge_keeps_absent_cell_absent() -> None: flow = flow.where(mask) sparse = (c.eff * c.gen_p).groupby(c.gbus).sum(sparse=True) - tot = linopy.merge([sparse, flow], join="outer") + tot = linopy.merge([sparse, flow], join="outer", cls=LinearExpression) assert tot._csr is not None - assert_linequal(tot, linopy.merge([dense, flow], join="outer")) + assert_linequal( + tot, linopy.merge([dense, flow], join="outer", cls=LinearExpression) + ) con = c.m.add_constraints(tot >= c.load, name="bal", freeze=True) assert isinstance(con, CSRConstraint) @@ -230,7 +232,7 @@ def test_namelist_sparse_matches_dense(observed: bool, member_first: bool) -> No assert set(csr.coords) == {"period", "season"} else: assert np.isnan(csr.const).sum() == 2 * 2 - assert csr.grid_dims == dense.coord_dims + assert csr.grid.dims == dense.coord_dims assert_linequal(sparse, dense) @@ -662,10 +664,26 @@ def test_cross_grid_merge_stays_csr_and_matches_dense( sparse, dense = cross_grid_parts(c, True), cross_grid_parts(c, False) if order == "flow-gen": sparse, dense = sparse[::-1], dense[::-1] - res = linopy.merge(sparse, join=join) + res = linopy.merge(sparse, join=join, cls=LinearExpression) assert res._csr is not None assert res.coord_dims == dense[0].coord_dims - assert_terms_equal(res, linopy.merge(dense, join=join)) + assert_terms_equal(res, linopy.merge(dense, join=join, cls=LinearExpression)) + + +def test_transposed_grid_exact_merge_stays_csr_and_matches_dense() -> None: + """Same labels in a transposed dim order stay sparse under the default join.""" + require_v1() + c = base_model() + a = (1.0 * c.flow).groupby(c.bus0).sum(sparse=True) + b = (1.0 * c.flow_t).groupby(c.bus0).sum(sparse=True) + assert a.coord_dims == b.coord_dims[::-1] != b.coord_dims + res = linopy.merge([a, b], cls=LinearExpression) + assert res._csr is not None + dense = [ + (1.0 * c.flow).groupby(c.bus0).sum(), + (1.0 * c.flow_t).groupby(c.bus0).sum(), + ] + assert_terms_equal(res, linopy.merge(dense, cls=LinearExpression)) @pytest.mark.parametrize("join", ["outer", "inner", "left", "right"]) @@ -675,17 +693,21 @@ def test_three_operand_cross_grid_merge_matches_dense(join: JoinOptions) -> None third_lines = ("line3", "line4") sparse = cross_grid_parts(c, True) + cross_grid_parts(c, True, third_lines)[1:] dense = cross_grid_parts(c, False) + cross_grid_parts(c, False, third_lines)[1:] - res = linopy.merge(sparse, join=join) + res = linopy.merge(sparse, join=join, cls=LinearExpression) assert res._csr is not None - assert_terms_equal(res, linopy.merge(dense, join=join)) + assert_terms_equal(res, linopy.merge(dense, join=join, cls=LinearExpression)) def test_cross_grid_merge_absent_fill_matches_dense() -> None: require_v1() c = base_model() sparse, dense = cross_grid_parts(c, True), cross_grid_parts(c, False) - res = linopy.merge(sparse, join="outer", fill_value=linopy.ABSENT) - expected = linopy.merge(dense, join="outer", fill_value=linopy.ABSENT) + res = linopy.merge( + sparse, join="outer", fill_value=linopy.ABSENT, cls=LinearExpression + ) + expected = linopy.merge( + dense, join="outer", fill_value=linopy.ABSENT, cls=LinearExpression + ) assert res._csr is not None filled = res.fillna(0) assert filled._csr is not None @@ -700,9 +722,9 @@ def test_cross_grid_merge_keeps_absent_cell_absent() -> None: sparse, dense = cross_grid_parts(c, True), cross_grid_parts(c, False) mask = xr.DataArray([True, False], coords=[dense[1].indexes["bus"]]) dense[1] = dense[1].where(mask) - res = linopy.merge([sparse[0], dense[1]], join="outer") + res = linopy.merge([sparse[0], dense[1]], join="outer", cls=LinearExpression) assert res._csr is not None - assert_terms_equal(res, linopy.merge(dense, join="outer")) + assert_terms_equal(res, linopy.merge(dense, join="outer", cls=LinearExpression)) def test_cross_grid_merge_mixed_dense_operand_stays_csr() -> None: @@ -710,8 +732,11 @@ def test_cross_grid_merge_mixed_dense_operand_stays_csr() -> None: c = base_model() sparse, dense = cross_grid_parts(c, True), cross_grid_parts(c, False) res = sparse[0].add(dense[1], join="outer") + assert isinstance(res, LinearExpression) assert res._csr is not None - assert_terms_equal(res, dense[0].add(dense[1], join="outer")) + expected = dense[0].add(dense[1], join="outer") + assert isinstance(expected, LinearExpression) + assert_terms_equal(res, expected) @pytest.mark.parametrize( @@ -764,13 +789,13 @@ def test_override_merge_same_shape_stays_csr() -> None: c = base_model() gen = (c.eff * c.gen_p).groupby(c.gbus).sum(sparse=True) flow = (1.0 * c.flow).groupby(c.bus1.str.upper()).sum(sparse=True) - res = linopy.merge([gen, flow], join="override") + res = linopy.merge([gen, flow], join="override", cls=LinearExpression) assert res._csr is not None dense = [ (c.eff * c.gen_p).groupby(c.gbus).sum(), (1.0 * c.flow).groupby(c.bus1.str.upper()).sum(), ] - assert_terms_equal(res, linopy.merge(dense, join="override")) + assert_terms_equal(res, linopy.merge(dense, join="override", cls=LinearExpression)) def test_cross_grid_balance_freezes_csr() -> None: @@ -778,7 +803,7 @@ def test_cross_grid_balance_freezes_csr() -> None: c1, c2 = base_model(), base_model() lhs1 = linopy.merge(cross_grid_parts(c1, False), join="outer") con1 = c1.m.add_constraints(lhs1 == c1.load, name="bal") - lhs2 = linopy.merge(cross_grid_parts(c2, True), join="outer") + lhs2 = linopy.merge(cross_grid_parts(c2, True), join="outer", cls=LinearExpression) assert lhs2._csr is not None con2 = c2.m.add_constraints(lhs2 == c2.load, name="bal", freeze=True) assert isinstance(con2, CSRConstraint)