diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 53541eff0..67399d44d 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -60,7 +60,9 @@ Upcoming Version * 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 `__) +* ``Model.add_constraints(..., mask=..., freeze=True)`` with a CSR-backed left-hand side, or with a ``CSRConstraint`` built beforehand, applies the mask sparsely and returns a ``CSRConstraint`` instead of expanding the dense rectangle. An expression on the right-hand side (``lhs <= other_expr`` or ``add_constraints(lhs, "<=", other_expr)``) no longer densifies a CSR-backed left-hand side: it is moved to the left-hand side, its constant becoming the right-hand side, as on the dense path. ``penalty=`` stays unsupported together with ``freeze=True`` and raises. (`#970 `__) +* ``LinearExpression.flat`` and ``LinearExpression.to_polars`` on a CSR-backed expression are emitted directly from the sparse backing instead of expanding the dense rectangle; the rows equal those of the dense path, with absent cells and zero coefficients dropped. (`#968 `__) +* A CSR-backed linear objective stays sparse when set with ``Model.add_objective`` and when the model is exported or solved: the objective vector ``matrices.c``, LP/MPS files, netcdf output, ``Model.copy``, persistent snapshots and direct solver APIs all read it without expanding the dense rectangle, with results equal to the dense objective. The objective name is now owned by the ``Objective`` itself (``Objective.name`` and ``Objective.attrs``) instead of being written into the expression's attributes. (`#967 `__) * 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** @@ -73,7 +75,9 @@ Upcoming Version * ``Model.to_netcdf``/``linopy.read_netcdf`` downgraded a quadratic objective the same way; the expression type is now stored alongside the objective and restored on read. Files written by earlier versions are read as before. (`#903 `__) * ``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 `__) - +* A ``CSRConstraint`` built from a CSR-backed expression now keeps its auxiliary coordinates, e.g. the key levels of a multi-key ``groupby(..., observed=True).sum(sparse=True)`` or a scalar coordinate left by a scalar selection, so its ``coords`` equal those of the dense constraint. They are kept by ``to_dense()``/``mutable()`` and written to and read from netcdf, and ``Constraint.freeze()`` keeps them as well. (`#941 `__) +* ``linopy.read_netcdf`` no longer fails for a whole model when a frozen constraint has a dimension or MultiIndex level whose name contains ``-``: the coordinate metadata of a ``CSRConstraint`` is now stored under positional variable names with the names as attributes. Files written by earlier versions are read as before. +* ``CSRConstraint.loc``, ``update``, ``from_rule`` and ``soften``, and assigning to ``coeffs``, ``vars`` or ``sign``, now raise an ``AttributeError`` saying the operation is not supported on a frozen constraint and naming ``.mutable()`` as the way out (for ``from_rule``: build with ``Constraint.from_rule`` and ``.freeze()`` the result; for ``soften``: add the constraint with ``freeze=False``), like the read-only ``rhs``, ``lhs`` and ``scaling`` setters, instead of a bare ``AttributeError``. (`#963 `__) * ``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/common.py b/linopy/common.py index ba055859f..792910502 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -7,6 +7,7 @@ from __future__ import annotations +import json import operator from collections.abc import Callable, Generator, Hashable, Iterable, Mapping, Sequence from functools import cached_property, reduce, wraps @@ -1318,24 +1319,25 @@ def coords_to_dataset_vars(coords: list[pd.Index]) -> dict[str, DataArray]: Suitable for embedding coordinate metadata as plain data variables in a Dataset that has its own unrelated dimensions (e.g. CSR netcdf format). + The variables are named by position, never by the (possibly dashed) + dimension or level names, since the netcdf reader splits variable names + on ``-``; the level names of a MultiIndex are stored as a JSON attribute. Reconstruct with :func:`coords_from_dataset`. """ data_vars: dict[str, DataArray] = {} - for c in coords: + for i, c in enumerate(coords): if isinstance(c, pd.MultiIndex): - for level_name, level_values in zip(c.names, c.levels): - data_vars[f"_coord_{c.name}_level_{level_name}"] = DataArray( - np.array(level_values), - dims=[f"_coorddim_{c.name}_level_{level_name}"], + for j, level_values in enumerate(c.levels): + data_vars[f"_index{i}_level{j}"] = DataArray( + np.array(level_values), dims=[f"_indexdim{i}_level{j}"] ) - data_vars[f"_coord_{c.name}_codes"] = DataArray( + data_vars[f"_index{i}_codes"] = DataArray( np.array(c.codes).T, - dims=[f"_coorddim_{c.name}", f"_coorddim_{c.name}_nlevels"], + dims=[f"_indexdim{i}", f"_indexdim{i}_nlevels"], + attrs={"level_names": json.dumps([str(n) for n in c.names])}, ) else: - data_vars[f"_coord_{c.name}"] = DataArray( - np.array(c), dims=[f"_coorddim_{c.name}"] - ) + data_vars[f"_index{i}"] = DataArray(np.array(c), dims=[f"_indexdim{i}"]) return data_vars @@ -1343,29 +1345,44 @@ def coords_from_dataset(ds: Dataset, coord_dims: list[str]) -> list[pd.Index]: """ Deserialize a list of pd.Index (including MultiIndex) from a Dataset. - Reconstructs coordinates previously serialized by :func:`coords_to_dataset_vars`. + Reconstructs coordinates previously serialized by :func:`coords_to_dataset_vars`, + or by its earlier name-keyed format (``_coord_``). """ coords = [] - for d in coord_dims: - if f"_coord_{d}_codes" in ds: - codes_2d = ds[f"_coord_{d}_codes"].values.T + for i, d in enumerate(coord_dims): + if f"_index{i}_codes" in ds: + codes = ds[f"_index{i}_codes"] + level_names = json.loads(codes.attrs["level_names"]) + level_keys = [f"_index{i}_level{j}" for j in range(len(level_names))] + coords.append(_multiindex(ds, codes.values.T, level_keys, level_names, d)) + elif f"_index{i}" in ds: + coords.append(pd.Index(ds[f"_index{i}"].values, name=d)) + elif f"_coord_{d}_codes" in ds: + prefix = f"_coord_{d}_level_" level_names = [ - str(k)[len(f"_coord_{d}_level_") :] - for k in ds - if str(k).startswith(f"_coord_{d}_level_") + str(k)[len(prefix) :] for k in ds if str(k).startswith(prefix) ] - arrays = [ - ds[f"_coord_{d}_level_{ln}"].values[codes_2d[i]] - for i, ln in enumerate(level_names) - ] - mi = pd.MultiIndex.from_arrays(arrays, names=level_names) - mi.name = d - coords.append(mi) + level_keys = [prefix + ln for ln in level_names] + codes_2d = ds[f"_coord_{d}_codes"].values.T + coords.append(_multiindex(ds, codes_2d, level_keys, level_names, d)) else: coords.append(pd.Index(ds[f"_coord_{d}"].values, name=d)) return coords +def _multiindex( + ds: Dataset, + codes_2d: np.ndarray, + level_keys: list[str], + level_names: list[str], + name: str, +) -> pd.MultiIndex: + arrays = [ds[k].values[codes_2d[j]] for j, k in enumerate(level_keys)] + mi = pd.MultiIndex.from_arrays(arrays, names=level_names) + mi.name = name + return mi + + def is_constant(x: SideLike) -> bool: """ Check if the given object is a constant type or an expression type without diff --git a/linopy/constraints.py b/linopy/constraints.py index fd88ab577..1a61d8bbc 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -17,6 +17,7 @@ TYPE_CHECKING, Any, NamedTuple, + NoReturn, overload, ) from warnings import warn @@ -46,8 +47,6 @@ check_has_nulls, check_has_nulls_polars, contains_labels, - coords_from_dataset, - coords_to_dataset_vars, filter_nulls_polars, format_coord, format_single_constraint, @@ -58,7 +57,6 @@ get_label_position, get_printout_labels, has_optimized_model, - is_constant, iterate_slices, maybe_group_terms_polars, maybe_replace_sign, @@ -79,7 +77,13 @@ PerformanceWarning, SIGNS_pretty, ) -from linopy.csr import Grid, _densify_notice, csr_nterm, csr_to_term_arrays +from linopy.csr import ( + Grid, + _densify_notice, + csr_nterm, + csr_to_term_arrays, + index_dtype, +) from linopy.scaling import ensure_scaling, validate_scaling from linopy.semantics import check_user_nan from linopy.types import ( @@ -592,6 +596,17 @@ def _positional_csr( return csr +_UNSUPPORTED = "is not supported on a frozen constraint" + + +def _frozen_error( + attr: str, + what: str = "is read-only", + remedy: str = "call .mutable() to modify", +) -> AttributeError: + return AttributeError(f"CSRConstraint.{attr} {what}; {remedy}.") + + _PositionalCache = tuple[scipy.sparse.csr_array, np.ndarray, "weakref.ref[np.ndarray]"] @@ -715,7 +730,7 @@ def attrs(self) -> dict[str, Any]: @property def coords(self) -> DatasetCoordinates: - return Dataset(coords=self._grid.indexes).coords + return self._grid.to_dataset().coords @property def dims(self) -> Frozen[Hashable, int]: @@ -785,21 +800,38 @@ def assign_labels( a zero coefficient counts as a term. ``scaling`` is a row scaling over the full flat grid. """ - keep = np.diff(self._csr.indptr) > 0 - positions = self._active_positions[keep] - csr = self._csr[keep] - csr.eliminate_zeros() - changes: dict[str, Any] = dict( - csr=csr, - active_positions=positions, + kept = self._kept(np.diff(self._csr.indptr) > 0) + kept._csr.eliminate_zeros() + changes: dict[str, Any] = dict(cindex=cindex, name=name) + if scaling is not None: + changes["scaling"] = scaling[kept._active_positions] + return kept._replace(**changes) + + def _kept(self, keep: np.ndarray) -> CSRConstraint: + """Copy holding only the active rows where ``keep`` is True.""" + + def rows(values: Any) -> Any: + is_rows = isinstance(values, np.ndarray) and values.ndim + return values[keep] if is_rows else values + + return self._replace( + csr=self._csr[keep], + active_positions=self._active_positions[keep], rhs=self._rhs[keep], - sign=self._sign if isinstance(self._sign, str) else self._sign[keep], - cindex=cindex, - name=name, + sign=rows(self._sign), + scaling=self._scaling[keep], + dual=rows(self._dual), + binvar_labels=rows(self._binvar_labels), + binval=rows(self._binval), ) - if scaling is not None: - changes["scaling"] = scaling[positions] - return self._replace(**changes) + + def masked(self, mask: DataArray) -> CSRConstraint: + """ + Copy with the cells where the boolean ``mask`` is False made inactive, + without the dense rectangle. ``mask`` must lie on the constraint grid. + """ + flat = mask.transpose(*self._grid.dims).to_numpy().reshape(-1) + return self._kept(flat[self._active_positions]) def _assign_coords(self, **coords: Any) -> CSRConstraint: """ @@ -823,7 +855,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._grid.coords) + return self._grid.dataarray(full) @property def labels(self) -> DataArray: @@ -843,6 +875,10 @@ def coeffs(self) -> DataArray: ) return self.data.coeffs + @coeffs.setter + def coeffs(self, value: ConstantLike) -> None: + raise _frozen_error("coeffs") + @property def vars(self) -> DataArray: """Get variable labels DataArray, shape (*coord_dims, _term).""" @@ -854,13 +890,21 @@ def vars(self) -> DataArray: ) return self.data.vars + @vars.setter + def vars(self, value: variables.Variable | DataArray) -> None: + raise _frozen_error("vars") + @property def sign(self) -> DataArray: """Get sign DataArray.""" if isinstance(self._sign, str): - return DataArray(np.full(self.shape, self._sign), coords=self._grid.coords) + return self._grid.dataarray(np.full(self.full_size, self._sign)) return self._active_to_dataarray(self._sign, fill="") + @sign.setter + def sign(self, value: SignLike) -> None: + raise _frozen_error("sign") + @property def rhs(self) -> DataArray: """Get RHS DataArray, shape (*coord_dims).""" @@ -868,9 +912,7 @@ def rhs(self) -> DataArray: @rhs.setter def rhs(self, value: ConstantLike) -> None: - raise AttributeError( - "CSRConstraint.rhs is read-only; call .mutable() to modify." - ) + raise _frozen_error("rhs") @property def scaling(self) -> DataArray: @@ -879,9 +921,7 @@ def scaling(self) -> DataArray: @scaling.setter def scaling(self, value: ConstantLike) -> None: - raise AttributeError( - "CSRConstraint.scaling is read-only; call .mutable() to modify." - ) + raise _frozen_error("scaling") @property def lhs(self) -> expressions.LinearExpression: @@ -891,8 +931,28 @@ def lhs(self) -> expressions.LinearExpression: @lhs.setter def lhs(self, value: ExpressionLike | VariableLike | ConstantLike) -> None: - raise AttributeError( - "CSRConstraint.lhs is read-only; call .mutable() to modify term structure." + raise _frozen_error("lhs") + + @property + def loc(self) -> LocIndexer: + raise _frozen_error("loc", _UNSUPPORTED) + + def update(self, *args: Any, **kwargs: Any) -> NoReturn: + raise _frozen_error("update", _UNSUPPORTED) + + def soften(self, *args: Any, **kwargs: Any) -> NoReturn: + raise _frozen_error( + "soften", + _UNSUPPORTED, + "add the constraint with freeze=False to soften it", + ) + + @classmethod + def from_rule(cls, *args: Any, **kwargs: Any) -> NoReturn: + raise _frozen_error( + "from_rule", + "is not supported", + "build with Constraint.from_rule and call .freeze() on the result", ) @property @@ -951,7 +1011,7 @@ def _to_dataset(self, nterm: int) -> Dataset: ) dim_names = self.coord_names - xr_coords = self._grid.indexes + xr_coords = self._grid.indexes | self._grid.aux dims_with_term = dim_names + [TERM_DIM] coeffs_da = DataArray( coeffs_2d.reshape(shape + (nterm,)), @@ -969,7 +1029,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._grid.coords), + labels=self._grid.dataarray(labels_flat), ) return ds @@ -1100,7 +1160,7 @@ 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._grid.coords)) + data_vars.update(self._grid.to_netcdf_vars()) if self._dual is not None: data_vars["dual"] = DataArray(self._dual, dims=["_flat"]) dim_names = list(self._grid.dims) @@ -1129,8 +1189,8 @@ def from_netcdf_ds(cls, ds: Dataset, model: Model, name: str) -> CSRConstraint: """ Reconstruct a Constraint from a netcdf Dataset (CSR format). - Files without the ``_csr_columns`` attribute were written before #926 - and hold dense variable positions instead of labels. + Files without the ``_csr_columns`` attribute hold dense variable + positions instead of labels. """ attrs = ds.attrs shape = tuple(attrs["shape"]) @@ -1156,7 +1216,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] - grid = Grid.from_coords(coords_from_dataset(ds, coord_dims)) + grid = Grid.from_netcdf_vars(ds, coord_dims) dual = ds["dual"].values if "dual" in ds else None if "_active_positions" in ds: active_positions = ds["_active_positions"].values @@ -1358,7 +1418,7 @@ def from_dense( ) csr.sum_duplicates() csr.eliminate_zeros() - grid = Grid.from_coords(con.indexes[d] for d in con.coord_dims) + grid = Grid.from_dataset(con.data, map(str, con.coord_dims)) rhs = con.rhs.values.ravel()[active_mask] scaling = con.scaling.values.ravel()[active_mask] sign_vals = con.sign.values.ravel() @@ -1405,12 +1465,10 @@ def from_csr( """ Staple sign and rhs onto a CSR-backed lhs to form an unassigned CSRConstraint. - The sparse counterpart of :meth:`from_dense`: instead of converting a - dense :class:`Constraint`, it realizes a - :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`. + Builds directly from a :class:`~linopy.csr.CSRLinearExpression`. The + expression's label columns are kept as they are, its constant moves to + the rhs, and rows with a NaN rhs are inactive. ``rhs`` must come from + :func:`csr_rhs`. """ sign = maybe_replace_sign(sign) rhs_flat = _rhs_grid_values(expr, rhs) - expr.const @@ -1420,20 +1478,19 @@ def from_csr( active, rhs_flat[active], sign, - grid=Grid(expr.grid.indexes), + grid=expr.grid, model=expr.model, ) def csr_rhs(expr: CSRLinearExpression, rhs: Any) -> DataArray | str: """ - 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. + Return the constant ``rhs`` as a DataArray on the expression grid, or the + reason the sparse path cannot take it: an rhs that is no DataArray-like, + or one with helper dims or dims outside the grid falls back to the dense + path. An expression rhs is moved to the lhs before, see + :meth:`LinearExpression.to_constraint`. """ - if not is_constant(rhs): - return "constraint with a non-constant rhs" try: da = as_dataarray(rhs) except (TypeError, ValueError): @@ -1925,7 +1982,8 @@ def _matrix_export_data( data = coeffs_final[valid_final] counts = valid_final.sum(axis=1) - indptr = np.empty(len(con_labels) + 1, dtype=np.int32) + dtype = index_dtype(len(data), (len(con_labels),), self.model) + indptr = np.empty(len(con_labels) + 1, dtype=dtype) indptr[0] = 0 np.cumsum(counts, out=indptr[1:]) return con_labels, row_mask, vlabel_cols, data, indptr diff --git a/linopy/csr.py b/linopy/csr.py index 1a24b2dad..ea5d0da7d 100644 --- a/linopy/csr.py +++ b/linopy/csr.py @@ -1,22 +1,22 @@ """ 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, ``sum``, ``merge``/``+``/``-``, scaling and ``@``/``dot`` -(:meth:`contracted`) become sparse linear algebra. Zero policy: the structural -operations (grouping and ``sum`` via :meth:`aggregated`, merge via +Under v1 semantics, ``expr.groupby(g).sum(sparse=True)`` (or +``linopy.options["sparse_groupby"]``) returns an ordinary +:class:`~linopy.expressions.LinearExpression` backed by a +:class:`CSRLinearExpression`: 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``, with no fixed term +count per row; grouping, ``sum``, ``merge``/``+``/``-``, scaling and +``@``/``dot`` (:meth:`contracted`) are 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. +coefficients; only the product with a constant matrix, ``@``/``dot``, prunes +them. Either way cell activeness is carried by ``const`` alone, independent of +term layout. +Any operation without a sparse branch expands the expression through +``.data`` to the mathematically identical dense rectangle in canonical term +layout; this is valid because v1 semantics do not fix the term layout. 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 @@ -38,6 +38,7 @@ import scipy.sparse from xarray import DataArray, Dataset +from linopy.common import coords_from_dataset, coords_to_dataset_vars from linopy.config import options from linopy.constants import HELPER_DIMS, TERM_DIM, PerformanceWarning from linopy.semantics import ( @@ -56,6 +57,9 @@ AuxCoords: TypeAlias = dict[str, tuple[str | tuple[()], np.ndarray]] """Auxiliary coordinates as ``name -> (grid dim, values)``, dim ``()`` for a scalar.""" +_AUX_PREFIX = "_aux" +_SCALAR_DIM = "_scalar" + @dataclass(frozen=True, eq=False) class Grid: @@ -85,6 +89,31 @@ def from_dataset(cls, ds: Dataset | DataArray, dims: Iterable[str]) -> Grid: indexes = {d: ds.get_index(d).rename(d) for d in dims} return cls(indexes, _aux_coords(ds, set(dims))) + @classmethod + def from_netcdf_vars(cls, ds: Dataset, dims: Iterable[str]) -> Grid: + """Read back a grid written by :meth:`to_netcdf_vars`.""" + aux: AuxCoords = { + da.attrs["name"]: (da.attrs.get("dim", ()), da.to_numpy()) + for k, da in ds.data_vars.items() + if str(k).startswith(_AUX_PREFIX) + } + return cls(cls.from_coords(coords_from_dataset(ds, list(dims))).indexes, aux) + + def to_netcdf_vars(self) -> dict[str, DataArray]: + """ + The indexes and auxiliary coordinates as plain data variables for + netcdf, named by position with the coordinate names as attributes. + """ + aux = { + f"{_AUX_PREFIX}{j}": DataArray( + v, + dims=[f"{_AUX_PREFIX}dim{j}"] if isinstance(d, str) else [], + attrs={"name": n} | ({"dim": d} if isinstance(d, str) else {}), + ) + for j, (n, (d, v)) in enumerate(self.aux.items()) + } + return coords_to_dataset_vars(self.coords) | aux + @property def dims(self) -> tuple[str, ...]: return tuple(self.indexes) @@ -240,6 +269,14 @@ class CSRLinearExpression: grid: Grid model: Model + def __post_init__(self) -> None: + csr = self.csr + dtype = index_dtype(csr.nnz, csr.shape, self.model) + if csr.indices.dtype != dtype or csr.indptr.dtype != dtype: + indices, indptr = csr.indices.astype(dtype), csr.indptr.astype(dtype) + csr = scipy.sparse.csr_array((csr.data, indices, indptr), shape=csr.shape) + object.__setattr__(self, "csr", csr) + @property def shape(self) -> tuple[int, ...]: return self.grid.shape @@ -255,7 +292,8 @@ def nterm(self) -> int: 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. + cell at ``indices``. An absent cell returns empty coefficient and + label arrays. """ row = int(np.ravel_multi_index(indices, self.grid.shape)) if indices else 0 const = float(self.const[row]) @@ -277,18 +315,18 @@ def from_grouper( coord_dims: tuple[str, ...], ) -> CSRLinearExpression: """ - Build the grouped sum directly in CSR form (no padded rectangle). + Build the grouped sum in CSR form. 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. - ``source`` is a dense expression dataset or an already CSR-backed - expression, which is regrouped through :meth:`aggregated`. + labels are sorted. 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``. ``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) @@ -343,6 +381,9 @@ 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) + if not grid_dims: + scalar = ds.expand_dims(_SCALAR_DIM) + return cls._from_scatter(scalar, model, grid, _SCALAR_DIM, {}, False) first = grid_dims[0] codes = {first: np.arange(len(grid.indexes[first]))} return cls._from_scatter(ds, model, grid, first, codes, False) @@ -361,14 +402,14 @@ def _from_scatter( 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. + dim maps one-to-one, and the COO to CSR conversion sums duplicate + variables, giving the group sum. Cells no member lands in stay absent + (NaN const). With ``skipna``, NaN member constants count as 0 in the + sum; without it, a NaN member constant propagates and leaves the cell + absent (NaN const). """ grid_dims = grid.dims - slot = min(grid_dims.index(d) for d in scatter_codes) + slot = min((grid_dims.index(d) for d in scatter_codes), default=0) transposed = [d for d in grid_dims if d not in scatter_codes] transposed.insert(slot, member_dim) member_rows = _member_rows(grid, scatter_codes, ds.sizes[member_dim]) @@ -380,9 +421,8 @@ def _from_scatter( 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), + csr = coo_to_csr( + coeffs[keep], rows[keep], vars_[keep], (full_size, model._xCounter), model ) const_vals = ds.const.transpose(*transposed).to_numpy().reshape(-1) @@ -392,31 +432,31 @@ 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) + return cls(csr, const, grid, model) 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. + constant is the NaN-skipping sum of its rows' constants (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]), - ) + shape = (grid.size, self.csr.shape[1]) + rows_ = rows[coo.coords[0]] + csr = coo_to_csr(coo.data, rows_, coo.coords[1], shape, self.model) 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) + return replace(self, csr=csr, 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. + Sum over grid dimensions. 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) @@ -429,6 +469,14 @@ def summed(self, dims: Iterable[str]) -> CSRLinearExpression: ) return self.aggregated(grid, rows).filled(0.0) + def live_terms(self) -> np.ndarray: + """ + Mask over the stored terms: the cell is present and the coefficient + is nonzero. + """ + present = np.repeat(~np.isnan(self.const), np.diff(self.csr.indptr)) + return present & (self.csr.data != 0) + def pruned(self) -> CSRLinearExpression: """Drop explicit zero coefficients; cell activeness stays with ``const``.""" csr = self.csr.copy() @@ -455,8 +503,8 @@ def scaled( 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). + Replace the per-cell constant, leaving the terms alone. A cell made + absent (NaN) has its terms dropped: an absent cell carries no terms. """ absent = np.isnan(const) counts = np.diff(self.csr.indptr) @@ -483,10 +531,10 @@ def taken(self, rows: np.ndarray, grid: Grid) -> CSRLinearExpression: 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). Auxiliary coordinates follow the - rows; those of ``grid`` are ignored. + Remap rows onto a new grid, possibly in a new dim order: dropped + labels vanish, new labels get ``fill`` as their constant (NaN: absent + cells). Auxiliary coordinates follow the rows; those of ``grid`` are + ignored. """ row_map, valid = grid.indexer(self.grid) @@ -495,14 +543,13 @@ def reindexed(self, grid: Grid, fill: float = np.nan) -> CSRLinearExpression: 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]) - ) + shape = (n_cells, self.csr.shape[1]) + csr = coo_to_csr(coo.data[keep], rows, cols, shape, self.model) const = np.full(n_cells, fill) const[row_map[valid]] = self.const[valid] return replace( self, - csr=scipy.sparse.csr_array(coo), + csr=csr, const=const, grid=self.grid.conformed(grid), ) @@ -525,10 +572,9 @@ 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. + an empty cell. A cell absent in either operand is absent in the sum + and carries no terms. Auxiliary coordinates propagate and conflicting + ones raise (§11). """ const = self.const + other.const a, b = self.csr.tocoo(), other.csr.tocoo() @@ -537,12 +583,10 @@ def added(self, other: CSRLinearExpression) -> CSRLinearExpression: 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 - ) + csr = coo_to_csr(data[present], rows[present], cols[present], shape, self.model) 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) + return replace(self, csr=csr, const=const, grid=grid) def contracted( self, @@ -562,11 +606,11 @@ def contracted( ``kron(I_kept, matrix.T) @ csr``, evaluated in chunks of the kept axis so the operator never grows with the kept size. - The result is compact canonical form: duplicate variables summed, terms - label-ordered and explicit zeros pruned -- unlike :meth:`added`, the - sparse product drops them, so cell activeness is carried by ``const`` - alone (see issue #925). Auxiliary coordinates on kept dims propagate, - those on contracted dims drop. + The result is in compact canonical form: duplicate variables summed, + terms label-ordered and explicit zeros pruned -- unlike :meth:`added`, + the sparse product drops them, so cell activeness is carried by + ``const`` alone. Auxiliary coordinates on kept dims propagate, those + on contracted dims drop. """ contracted_dims = tuple(contracted_dims) kept = tuple(d for d in self.grid.dims if d not in contracted_dims) @@ -613,8 +657,8 @@ 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`. + Absent cells (NaN const) carry no terms. The expanded dataset is + wrapped in a :class:`LinearExpression`. """ from linopy.expressions import LinearExpression @@ -638,6 +682,28 @@ def to_dense(self) -> LinearExpression: return LinearExpression(absorb_absence(ds), self.model) +def index_dtype(nnz: int, shape: tuple[int, ...], model: Model) -> np.dtype: + """ + Index dtype of a CSR store: the model's label dtype, widened to int64 + only when the nonzeros or the shape outgrow it. + """ + dtype = np.dtype(model._dtypes["labels"]) + return dtype if max(nnz, *shape) <= np.iinfo(dtype).max else np.dtype(np.int64) + + +def coo_to_csr( + data: np.ndarray, + rows: np.ndarray, + cols: np.ndarray, + shape: tuple[int, int], + model: Model, +) -> scipy.sparse.csr_array: + """Build a CSR store from COO triplets in the model's index dtype.""" + dtype = index_dtype(len(data), shape, model) + coords = (rows.astype(dtype, copy=False), cols.astype(dtype, copy=False)) + return scipy.sparse.csr_array(scipy.sparse.coo_array((data, coords), shape=shape)) + + 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) @@ -687,7 +753,7 @@ 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: + if n in ds.xindexes: continue if c.ndim == 0: aux[str(n)] = ((), c.to_numpy()) diff --git a/linopy/expressions.py b/linopy/expressions.py index 0726afc25..b4b19bc28 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -567,9 +567,9 @@ def sum( Parameters ---------- use_fallback : bool - Fall back to the previous, slower groupby-sum implementation, kept - as an escape hatch. Leave at False unless the default misbehaves. - Defaults to False. + Use the slower fallback groupby-sum implementation instead of the + default one. Kept as an escape hatch. Leave at False unless the + default misbehaves. Defaults to False. sparse : bool, optional Build the grouped sum in CSR form behind the ordinary LinearExpression type — no group-size padding; a still-sparse @@ -867,6 +867,9 @@ def flat(self) -> pd.DataFrame: ... @abstractmethod def to_polars(self) -> pl.DataFrame: ... + @abstractmethod + def linear_terms(self) -> tuple[np.ndarray, np.ndarray]: ... + def __init__(self, data: Dataset | Any | None, model: Model) -> None: from linopy.model import Model @@ -2553,14 +2556,20 @@ def _selected( 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.""" + """ + Select by label as ``Dataset.sel``. For a CSR-backed expression, + returns a CSR-backed result when the selection stays on the grid. + """ 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.""" + """ + Select by position as ``Dataset.isel``. For a CSR-backed expression, + returns a CSR-backed result when the selection stays on the grid. + """ csr = self._selected(lambda rows: rows.isel(*args, **kwargs), "isel") if csr is None: return super().isel(*args, **kwargs) @@ -2669,6 +2678,9 @@ def to_constraint( self, sign: SignLike, rhs: SideLike, join: JoinOptions | None = None ) -> ConstraintBase: if self._csr is not None and isinstance(sign, str): + rhs = as_constant(rhs) + if join is not None or not isinstance(rhs, CONSTANT_TYPES): + return self.sub(rhs, join=join).to_constraint(sign, 0) rhs_da = constraints.csr_rhs(self._csr, rhs) if isinstance(rhs_da, DataArray): return constraints.CSRConstraint.from_csr(self._csr, sign, rhs_da) @@ -2816,7 +2828,8 @@ def __matmul__( ``(self * other).sum(dim)``. The result is then the compact canonical form -- duplicate variables summed, terms label-ordered, explicit zeros pruned -- so its term count may differ from the dense path's - while the values agree. A CSR-backed expression stays CSR-backed. + while the values agree. Returns a CSR-backed result when ``self`` is + CSR-backed. """ other = as_constant(other) other_is_const = not isinstance(other, LinearExpression | variables.Variable) @@ -2843,8 +2856,8 @@ def _sparse_matmul(self, other: ConstantLike) -> LinearExpression | None: labels, a zero-size grid, an operand sharing no dimension with the grid, and an unlabelled output dimension. The result is the compact canonical form of :meth:`CSRLinearExpression.contracted`, so its term - count may differ from the dense path's while the values agree; it - stays CSR-backed when the input was. + count may differ from the dense path's while the values agree; the + result is CSR-backed when the input was. """ if is_nan_scalar(other): check_user_nan(op_kind="mul") @@ -2894,13 +2907,15 @@ def flat(self) -> pd.DataFrame: ------- df : pandas.DataFrame """ - ds = self.data + if self._csr is not None: + df = pd.DataFrame(self._csr_terms(self._csr)) + else: - def mask_func(data: dict) -> pd.Series: - mask = (data["vars"] != -1) & (data["coeffs"] != 0) - return mask + def mask_func(data: dict) -> pd.Series: + mask = (data["vars"] != -1) & (data["coeffs"] != 0) + return mask - df = to_dataframe(ds, mask_func=mask_func) + df = to_dataframe(self.data, mask_func=mask_func) df = df.groupby("vars", as_index=False).sum() check_has_nulls(df, name=self.type) return df @@ -2916,8 +2931,9 @@ def reindex( **indexers_kwargs: Any, ) -> LinearExpression: """ - Conform to new coordinates as ``Dataset.reindex``; a CSR-backed - expression stays sparse when only labels change. + Conform to new coordinates as ``Dataset.reindex``. For a CSR-backed + expression, returns a CSR-backed result when only grid labels change + and no other keyword argument is given. """ indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "reindex") csr = self._csr @@ -2945,8 +2961,8 @@ def rename( **names: Any, ) -> LinearExpression: """ - Rename dimensions as ``Dataset.rename``; a CSR-backed expression - stays sparse when only grid dims are relabelled. + Rename dimensions as ``Dataset.rename``. For a CSR-backed expression, + returns a CSR-backed result when only grid dims are relabelled. """ name_dict = either_dict_or_kwargs(name_dict, names, "rename") csr = self._csr @@ -2975,18 +2991,51 @@ def to_polars(self) -> pl.DataFrame: ------- df : polars.DataFrame """ - if self.is_constant: + if self._csr is not None: + df = pl.DataFrame(self._csr_terms(self._csr)) + elif self.is_constant: df = pl.DataFrame( {"const": self.data["const"].values.reshape(-1)} ).with_columns(pl.lit(None).alias("coeffs"), pl.lit(None).alias("vars")) return df.select(["vars", "coeffs", "const"]) - - df = to_polars(self.data) - df = filter_nulls_polars(df) + else: + df = filter_nulls_polars(to_polars(self.data)) df = maybe_group_terms_polars(df) check_has_nulls_polars(df, name=self.type) return df + def _csr_terms(self, csr: CSRLinearExpression) -> dict[str, np.ndarray]: + """ + Stored terms of a CSR backing as ``const``, ``coeffs`` and ``vars`` + columns, dropping absent cells and zero coefficients like the dense + long format. + """ + keep = csr.live_terms() + const = np.repeat(csr.const, np.diff(csr.csr.indptr)) + return { + "const": const[keep], + "coeffs": csr.csr.data[keep].astype(float), + "vars": csr.csr.indices[keep].astype(self.model._dtypes["labels"]), + } + + def linear_terms(self) -> tuple[np.ndarray, np.ndarray]: + """ + Variable labels and coefficients of the stored terms. For a CSR-backed + expression, reads the CSR arrays directly; otherwise reads the dense + term arrays. Absent terms and zero coefficients are dropped, duplicate + labels are not summed. + """ + if self._csr is not None: + csr = self._csr.csr + keep = self._csr.live_terms() + if keep.all(): + return csr.indices, csr.data + return csr.indices[keep], csr.data[keep] + labels = self.data.vars.values.ravel() + coeffs = self.data.coeffs.values.ravel() + keep = (labels != -1) & (coeffs != 0) + return labels[keep], coeffs[keep] + def simplify(self) -> LinearExpression: """ Simplify the linear expression by combining terms with the same variable. @@ -3461,6 +3510,19 @@ def mask_func(data: dict) -> pd.Series: check_has_nulls(df, name=self.type) return df + def linear_terms(self) -> tuple[np.ndarray, np.ndarray]: + """ + Variable labels and coefficients of the terms with a single variable; + zero coefficients are dropped, duplicate labels are not summed. + """ + coeffs = self.data.coeffs + factors = self.data.vars.transpose(FACTOR_DIM, *coeffs.dims).values + vars1, vars2 = factors[0].ravel(), factors[1].ravel() + labels = np.where(vars1 != -1, vars1, vars2) + values = coeffs.values.ravel() + keep = ((vars1 == -1) | (vars2 == -1)) & (labels != -1) & (values != 0) + return labels[keep], values[keep] + def to_polars(self, **kwargs: Any) -> pl.DataFrame: """ Convert the expression to a polars DataFrame. @@ -3580,10 +3642,10 @@ def _aligned( ) -> 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; the reason is returned instead. + cells the join creates carrying ``fill`` as constant. Returns the reason + as a string where the dense path owns the semantics instead: ``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 @@ -3617,9 +3679,9 @@ def _try_csr_merge( 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 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. + coordinates are checked for conflicts on the operands as given 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 diff --git a/linopy/io.py b/linopy/io.py index 08701ed03..718757e51 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -12,6 +12,7 @@ import time import warnings from collections.abc import Callable, Iterable +from dataclasses import replace from importlib.metadata import version from io import BufferedWriter from pathlib import Path @@ -30,7 +31,7 @@ to_polars, ) from linopy.constants import CONCAT_DIM, FACTOR_DIM, SOS_DIM_ATTR, SOS_TYPE_ATTR -from linopy.objective import Objective +from linopy.objective import Objective, linear_part from linopy.scaling import constraint_scaling_lookup, variable_scaling_lookup if TYPE_CHECKING: @@ -301,14 +302,7 @@ def objective_to_file( elif m.is_quadratic: df = _scale_objective_dataframe(df, variable_scaling, m.objective.scaling) - linear_terms = df.filter(pl.col("vars1").eq(-1) | pl.col("vars2").eq(-1)) - linear_terms = linear_terms.with_columns( - pl.when(pl.col("vars1").eq(-1)) - .then(pl.col("vars2")) - .otherwise(pl.col("vars1")) - .alias("vars") - ) - objective_write_linear_terms(f, linear_terms, print_variable) + objective_write_linear_terms(f, linear_part(df), print_variable) quads = df.filter(pl.col("vars1").ne(-1) & pl.col("vars2").ne(-1)) objective_write_quadratic_terms(f, quads, print_variable) @@ -1091,8 +1085,7 @@ def with_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: ) for name, expr in m.expressions.items() ] - objective = m.objective.data - objective = objective.assign_attrs( + objective = m.objective.to_netcdf_ds().assign_attrs( sense=m.objective.sense, scaling=m.objective.scaling, **{EXPR_TYPE_ATTR: m.objective.expression.type}, @@ -1399,8 +1392,12 @@ def _copy_con_data(con: ConstraintBase) -> xr.Dataset: new_model, ) - obj_expr = type(m.objective.expression)( - m.objective.expression.data.copy(deep=deep), new_model + expr = m.objective.expression + csr = expr._csr if isinstance(expr, LinearExpression) else None + obj_expr = ( + type(expr)(expr.data.copy(deep=deep), new_model) + if csr is None + else LinearExpression._from_csr(replace(csr, model=new_model), new_model) ) new_model._objective = Objective( obj_expr, new_model, m.objective.sense, m.objective.scaling diff --git a/linopy/matrices.py b/linopy/matrices.py index 9aec46fa9..bca66d70f 100644 --- a/linopy/matrices.py +++ b/linopy/matrices.py @@ -116,6 +116,7 @@ def _build_cons(self) -> None: label_index = m.variables.label_index label_to_pos = label_index.label_to_pos con_scaling_by_label = constraint_scaling_lookup(m) + unit_cols = bool((self.var_scaling == 1).all()) def scale_rows_and_cols( csr: scipy.sparse.csr_array, con_labels: np.ndarray, b: np.ndarray @@ -123,18 +124,20 @@ def scale_rows_and_cols( if csr.shape[0] == 0: return csr, b row_scaling = con_scaling_by_label[con_labels] + unit_rows = bool((row_scaling == 1).all()) + if unit_rows and unit_cols: + return csr, b # With solver variables y = Scol * x, constraints A x = b become # Srow * A * Scol^-1 * y = Srow * b. - csr = cast( - scipy.sparse.csr_array, - csr.multiply(row_scaling[:, np.newaxis]).tocsr(), + data = csr.data + if not unit_rows: + data = data * np.repeat(row_scaling, np.diff(csr.indptr)) + if not unit_cols: + data = data / self.var_scaling[csr.indices] + scaled = scipy.sparse.csr_array( + (data, csr.indices, csr.indptr), shape=csr.shape ) - if csr.shape[1] and len(self.var_scaling): - csr = cast( - scipy.sparse.csr_array, - csr.multiply(1 / self.var_scaling[np.newaxis, :]).tocsr(), - ) - return csr, b * row_scaling + return scaled, b * row_scaling reg_csrs, reg_b, reg_sense = [], [], [] ind_csrs, ind_b, ind_sense, ind_binvar, ind_binval = [], [], [], [], [] @@ -174,22 +177,9 @@ def c(self) -> ndarray: label_index = m.variables.label_index label_to_pos = label_index.label_to_pos - expr = m.objective.expression - if isinstance(expr, expressions.QuadraticExpression): - # vars has shape (_factor=2, _term); linear terms have one factor == -1 - vars_2d = expr.data.vars.values # shape (2, n_term) - coeffs_all = expr.data.coeffs.values.ravel() - vars1, vars2 = vars_2d[0], vars_2d[1] - linear = (vars1 == -1) | (vars2 == -1) - var_labels = np.where(vars1[linear] != -1, vars1[linear], vars2[linear]) - coeffs = coeffs_all[linear] - else: - var_labels = expr.data.vars.values.ravel() - coeffs = expr.data.coeffs.values.ravel() - - mask = var_labels != -1 - positions = label_to_pos[var_labels[mask]] - scaled_coeffs = coeffs[mask] / self.var_scaling[positions] + var_labels, coeffs = m.objective.linear_terms() + positions = label_to_pos[var_labels] + scaled_coeffs = coeffs / self.var_scaling[positions] scaled_coeffs = scaled_coeffs * m.objective.scaling np.add.at(result, positions, scaled_coeffs) return result diff --git a/linopy/model.py b/linopy/model.py index fd4066c97..93d80b191 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -61,6 +61,7 @@ Constraints, CSRConstraint, ) +from linopy.csr import _densify_notice from linopy.dualization import dualize from linopy.expressions import ( Expressions, @@ -1308,11 +1309,12 @@ def add_constraints( name = self._resolve_constraint_name(name) if freeze is None: freeze = self.freeze_constraints + chunked = bool(freeze and self.chunk) freeze = freeze and not self.chunk if penalty is not None and freeze: raise ValueError( - "`penalty` cannot be combined with `freeze=True` (or a model default of `freeze_constraints=True`)," + "`penalty` cannot be combined with `freeze=True` (or a model default of `freeze_constraints=True`), " "since `soften` is not supported on frozen constraints." ) @@ -1329,17 +1331,11 @@ 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: + if isinstance(con, CSRConstraint) and freeze: + if mask is not None: + mask = broadcast_to_coords(mask, con.coords, label="mask") + con = con.masked(mask.astype(bool)) _check_infinities(con._sign, con._rhs, name) self.check_force_dim_names(con.coords.to_dataset()) enforce_no_multiindex(con, context=f"constraint {name!r}") @@ -1351,6 +1347,14 @@ def add_constraints( self._cCounter += con.full_size con = con.assign_labels(cindex, name, scaling_grid.values.ravel()) return self.constraints.add(con) + if isinstance(con, CSRConstraint): + if penalty is not None: + reason = "`penalty` given, softening needs a mutable constraint" + elif chunked: + reason = "chunked model, `Model.chunk` adds constraints unfrozen" + else: + reason = "constraint added unfrozen, `freeze=False`" + _densify_notice(reason) data = con.data _check_infinities(data.sign, data.rhs, name) diff --git a/linopy/objective.py b/linopy/objective.py index 6cc2cd6fd..f3801ce91 100644 --- a/linopy/objective.py +++ b/linopy/objective.py @@ -53,6 +53,19 @@ def _objwrap(obj: Objective, *args: Any, **kwargs: Any) -> Objective: return _objwrap +def linear_part(df: pl.DataFrame) -> pl.DataFrame: + """ + Linear terms of a quadratic ``to_polars`` frame, their variable in ``vars``. + """ + linear = df.filter(pl.col("vars1").eq(-1) | pl.col("vars2").eq(-1)) + return linear.with_columns( + pl.when(pl.col("vars1").eq(-1)) + .then(pl.col("vars2")) + .otherwise(pl.col("vars1")) + .alias("vars") + ) + + class Objective: """ An objective expression containing all relevant information. @@ -92,12 +105,19 @@ def __repr__(self) -> str: return f"Objective:\n----------\n{expr_string}\n{sense_string}\n{value_string}" + @property + def name(self) -> str: + """ + Returns the name of the objective, owned by the objective itself. + """ + return "objective" + @property def attrs(self) -> dict[str, Any]: """ Returns the attributes of the objective. """ - return self.expression.attrs + return {"name": self.name} @property def coords(self) -> DatasetCoordinates: @@ -133,6 +153,26 @@ def to_polars(self, **kwargs: Any) -> pl.DataFrame: """ return self.expression.to_polars(**kwargs) + def linear_terms(self) -> tuple[np.ndarray, np.ndarray]: + """ + Returns the variable labels and coefficients of the linear objective + terms, see ``LinearExpression.linear_terms``. + """ + return self.expression.linear_terms() + + def to_netcdf_ds(self) -> Dataset: + """ + Returns the objective in the dense layout for netcdf serialization. + + The expression setter reduces the objective to a single cell, so a + CSR-backed expression is densified into that one cell here, without + mutating the stored CSR data. + """ + expr = self.expression + csr = expr._csr if isinstance(expr, expressions.LinearExpression) else None + ds = expr.data if csr is None else csr.to_dense().data + return ds.assign_attrs(name=self.name) + @property def coeffs(self) -> DataArray: """ @@ -197,7 +237,8 @@ def expression( if (expr.const != 0.0) and not np.isnan(expr.const): raise ValueError("Constant values in objective function not supported.") - expr.attrs["name"] = "objective" + if not expr.is_sparse: + expr.attrs["name"] = self.name self._expression = expr @property diff --git a/linopy/persistent/snapshot.py b/linopy/persistent/snapshot.py index 76396d5ad..ee09390cf 100644 --- a/linopy/persistent/snapshot.py +++ b/linopy/persistent/snapshot.py @@ -7,7 +7,6 @@ import numpy as np import pandas as pd -from linopy import expressions from linopy.constraints import Constraint if TYPE_CHECKING: @@ -38,19 +37,8 @@ def _objective_linear_vector(model: Model) -> np.ndarray: vlabels = model.variables.label_index.vlabels label_to_pos = model.variables.label_index.label_to_pos result = np.zeros(len(vlabels), dtype=np.float64) - expr = model.objective.expression - if isinstance(expr, expressions.QuadraticExpression): - vars_2d = expr.data.vars.values - coeffs_all = expr.data.coeffs.values.ravel() - vars1, vars2 = vars_2d[0], vars_2d[1] - linear = (vars1 == -1) | (vars2 == -1) - var_labels = np.where(vars1[linear] != -1, vars1[linear], vars2[linear]) - coeffs = coeffs_all[linear] - else: - var_labels = expr.data.vars.values.ravel() - coeffs = expr.data.coeffs.values.ravel() - mask = var_labels != -1 - np.add.at(result, label_to_pos[var_labels[mask]], coeffs[mask]) + var_labels, coeffs = model.objective.linear_terms() + np.add.at(result, label_to_pos[var_labels], coeffs) return result diff --git a/test/test_common.py b/test/test_common.py index 7109314ef..ead6a5410 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -79,20 +79,39 @@ def test_assign_multiindex_safe() -> None: def test_coords_dataset_vars_roundtrip_multiindex() -> None: """MultiIndex and plain coords survive serialization to Dataset vars and back.""" mi = pd.MultiIndex.from_product( - [[2020, 2030], ["t1", "t2"]], names=("period", "timestep") + [[2020, 2030], ["t1", "t2"]], names=("period", "time-step") ) - mi.name = "snapshot" - plain = pd.Index([1, 2, 3], name="simple") + mi.name = "snap-shot" + plain = pd.Index([1, 2, 3], name="my-simple") ds = xr.Dataset(coords_to_dataset_vars([mi, plain])) - restored = coords_from_dataset(ds, ["snapshot", "simple"]) + assert not any("-" in str(k) for k in ds) + restored = coords_from_dataset(ds, ["snap-shot", "my-simple"]) assert isinstance(restored[0], pd.MultiIndex) assert restored[0].equals(mi) - assert list(restored[0].names) == ["period", "timestep"] - assert restored[0].name == "snapshot" + assert list(restored[0].names) == ["period", "time-step"] + assert restored[0].name == "snap-shot" assert restored[1].equals(plain) - assert restored[1].name == "simple" + assert restored[1].name == "my-simple" + + +def test_coords_from_dataset_reads_name_keyed_format() -> None: + """Files written before the positional format keyed the variables by name.""" + ds = xr.Dataset( + { + "_coord_snapshot_level_period": ("a", [2020, 2030]), + "_coord_snapshot_level_timestep": ("b", ["t1", "t2"]), + "_coord_snapshot_codes": (("c", "d"), [[0, 0], [0, 1], [1, 0], [1, 1]]), + "_coord_simple": ("e", [1, 2, 3]), + } + ) + mi, plain = coords_from_dataset(ds, ["snapshot", "simple"]) + want = pd.MultiIndex.from_product( + [[2020, 2030], ["t1", "t2"]], names=("period", "timestep") + ) + assert mi.equals(want) and mi.name == "snapshot" + assert plain.equals(pd.Index([1, 2, 3])) and plain.name == "simple" def test_iterate_slices_basic() -> None: diff --git a/test/test_csr.py b/test/test_csr.py index bb30e86a4..b75f49de7 100644 --- a/test/test_csr.py +++ b/test/test_csr.py @@ -92,7 +92,7 @@ def canon(df: pl.DataFrame) -> pl.DataFrame: ) -def assert_frozen_equal(con1: Constraint, con2: CSRConstraint) -> None: +def assert_frozen_equal(con1: ConstraintBase, con2: ConstraintBase) -> None: d1, d2 = canon(con1.to_polars()), canon(con2.to_polars()) assert d1["labels"].equals(d2["labels"]) assert d1["vars"].equals(d2["vars"]) @@ -499,6 +499,22 @@ def test_nan_grouper_raises_eagerly() -> None: (1.0 * c.gen_p).groupby(gbus).sum(sparse=True) +TERM_LINE = re.compile(r"^[+-][0-9.e+-]+ x[0-9]+$") + + +def canon_lp(text: str) -> list[str]: + """LP file lines with each block of term lines sorted.""" + out: list[str] = [] + buf: list[str] = [] + for line in text.splitlines(): + if TERM_LINE.match(line): + buf.append(line) + else: + out += sorted(buf) + [line] + buf = [] + return out + sorted(buf) + + def test_lp_files_identical(tmp_path: Path) -> None: require_v1() sizes = (7, 1, 3, 1, 2, 1, 1, 4, 1, 2, 1, 1) @@ -510,19 +526,6 @@ def test_lp_files_identical(tmp_path: Path) -> None: ) c2.m.add_objective((1.0 * c2.gen_p).sum()) - term_line = re.compile(r"^[+-][0-9.e+-]+ x[0-9]+$") - - def canon_lp(text: str) -> list[str]: - out: list[str] = [] - buf: list[str] = [] - for line in text.splitlines(): - if term_line.match(line): - buf.append(line) - else: - out += sorted(buf) + [line] - buf = [] - return out + sorted(buf) - f1, f2 = tmp_path / "eager.lp", tmp_path / "sparse.lp" c1.m.to_file(f1) c2.m.to_file(f2) @@ -1344,12 +1347,32 @@ def test_is_sparse_tracks_backing_and_repr_marks_it() -> None: assert not (1.0 * c.gen_p).is_sparse +def add_chunked(e: LinearExpression, c: Case) -> Any: + c.m.chunk = {"bus": 2} + return c.m.add_constraints(e >= 1, freeze=True) + + +def add_with_penalty(e: LinearExpression, c: Case) -> Any: + c.m.add_objective(1.0 * c.gen_p.sum()) + return c.m.add_constraints(e >= 1, penalty=1.0) + + 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"), + "rhs": ( + lambda e, c: e <= xr.DataArray([1.0, 2.0], coords=[LOC]), + "rhs over dimensions outside the grid", + ), + "rhs-expr": (lambda e, c: e <= 1.0 * c.gen_p.sum("gen"), "over different dim"), "mutable": (lambda e, c: (e == c.load).mutable(), "`mutable\\(\\)`"), + "add-unfrozen": ( + lambda e, c: c.m.add_constraints(e >= 1), + "constraint added unfrozen, `freeze=False`", + ), + "add-chunked": (add_chunked, "chunked model"), + "add-penalty": (add_with_penalty, "`penalty` given"), "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": ( @@ -1397,6 +1420,109 @@ def test_warn_on_densify_names_the_reason(op: str, enabled: bool) -> None: assert notices[0].filename == __file__ +@pytest.mark.parametrize("scale", [1.0, 0.0], ids=["plain", "zeros"]) +@pytest.mark.parametrize("build", list(SPARSE_BUILDS)) +def test_flat_and_to_polars_served_without_densifying(build: str, scale: float) -> None: + require_v1() + sparse, _ = sparse_and_dense(build) + sparse = scale * sparse + assert sparse._csr is not None + dense = sparse._csr.to_dense() + with no_densify(): + got_pl, got_flat = sparse.to_polars(), sparse.flat + assert sparse.is_sparse + assert got_pl.sort("vars").equals(dense.to_polars().sort("vars")) + pd.testing.assert_frame_equal(got_flat, dense.flat) + + +@pytest.mark.parametrize("build", list(SPARSE_BUILDS)) +def test_sparse_store_indices_follow_model_label_dtype(build: str) -> None: + require_v1() + sparse, _ = sparse_and_dense(build) + m = sparse.model + con = m.add_constraints(sparse >= 1, name="con", freeze=True) + A = m.matrices.A + assert isinstance(con, CSRConstraint) and sparse._csr is not None + assert A is not None + dtypes = {sparse._csr.csr.indices.dtype, con._csr.indices.dtype} + assert dtypes == {np.dtype(m.dtypes["labels"]), A.indices.dtype} + + +@pytest.mark.parametrize("label_dtype", [np.int32, np.int64]) +def test_sparse_group_sum_indices_widen_with_model(label_dtype: type) -> None: + require_v1() + m = Model(dtypes={"labels": label_dtype}) + x = m.add_variables(coords=[pd.RangeIndex(4, name="i")], name="x") + group = xr.DataArray([0, 0, 1, 1], coords=[x.indexes["i"]], name="g") + expr = x.groupby(group).sum(sparse=True) + assert expr._csr is not None + assert expr._csr.csr.indices.dtype == expr._csr.csr.indptr.dtype == label_dtype + + +def objective_twins(build: str, scale: float) -> tuple[Model, Model]: + """Twin models, the first with the sparse build as objective, the second dense.""" + sparse, _ = sparse_and_dense(build) + _, dense = sparse_and_dense(build) + with no_densify(): + sparse.model.add_objective(scale * (sparse - sparse.const.fillna(0))) + dense.model.add_objective(scale * (dense - dense.const.fillna(0))) + return sparse.model, dense.model + + +def test_sparse_objective_rejects_constant_without_densifying() -> None: + require_v1() + sparse, _ = sparse_and_dense("absent") + with no_densify(), pytest.raises(ValueError, match="Constant values"): + sparse.model.add_objective(sparse) + + +@pytest.mark.parametrize("scale", [1.0, 0.0], ids=["plain", "zeros"]) +@pytest.mark.parametrize("build", list(SPARSE_BUILDS)) +def test_objective_stays_csr_and_exports_like_dense( + build: str, scale: float, tmp_path: Path +) -> None: + require_v1() + ms, md = objective_twins(build, scale) + with no_densify(): + c = ms.matrices.c + terms = ms.objective.linear_terms() + ms.to_file(tmp_path / "sparse.lp") + ms.to_netcdf(tmp_path / "sparse.nc") + copied = ms.copy() + assert ms.objective.attrs == {"name": "objective"} + repr(ms.objective) + assert ms.objective.expression.is_sparse and copied.objective.expression.is_sparse + md.to_file(tmp_path / "dense.lp") + md.to_netcdf(tmp_path / "dense.nc") + assert np.array_equal(c, md.matrices.c) + dense_terms = md.objective.linear_terms() + assert sorted(zip(*map(list, terms))) == sorted(zip(*map(list, dense_terms))) + assert (dense_terms[1] != 0).all() + lp_sparse, lp_dense = (tmp_path / f"{k}.lp" for k in ("sparse", "dense")) + assert canon_lp(lp_sparse.read_text()) == canon_lp(lp_dense.read_text()) + rs, rd = (linopy.read_netcdf(tmp_path / f"{k}.nc") for k in ("sparse", "dense")) + es, ed = rs.objective.expression, rd.objective.expression + assert isinstance(es, LinearExpression) and isinstance(ed, LinearExpression) + assert_cells_equal(es, ed, ()) + assert rs.objective.expression.attrs["name"] == "objective" + + +@pytest.mark.skipif("highs" not in linopy.available_solvers, reason="needs highs") +@pytest.mark.parametrize("io_api", ["lp", "direct"]) +@pytest.mark.parametrize("build", list(SPARSE_BUILDS)) +def test_sparse_objective_solves_like_dense(build: str, io_api: str) -> None: + require_v1() + values = [] + for m in objective_twins(build, 1.0): + for var in m.variables.data.values(): + var.update(lower=0, upper=1) + m.objective.sense = "max" + with no_densify(): + m.solve("highs", io_api=io_api) + values.append(m.objective.value) + assert values[0] == pytest.approx(values[1]) + + 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]) @@ -1454,6 +1580,28 @@ def test_elementwise_constant_ops_stay_csr_and_match_dense( assert_linequal(res, func(dense, x)) +DENSE_ADDENDS: dict[str, Callable[[Case, LinearExpression], LinearExpression]] = { + "same-grid": lambda c, dense: 2 * dense, + "scalar-var": lambda c, dense: 2 * c.gen_p.isel(gen=0, snapshot=0), +} + + +@pytest.mark.parametrize( + ("build", "addend"), + [(b, "same-grid") for b in SPARSE_BUILDS] + [("zero-dim", "scalar-var")], +) +def test_merge_with_dense_expression_stays_csr_and_matches_dense( + build: str, addend: str +) -> None: + require_v1() + c = base_model() + sparse, dense = sparse_and_dense(build, c) + other = DENSE_ADDENDS[addend](c, dense) + with no_densify(): + res = sparse + other + assert_sparse_matches(res, dense + other) + + @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"]) @@ -1524,6 +1672,24 @@ def test_join_with_constant_keeps_aux_coords_like_dense( assert_sparse_matches(res, getattr(dense, op)(x, join=join)) +@pytest.mark.parametrize("join", ["outer", "inner", "left", "right"]) +def test_to_constraint_join_with_constant_matches_dense(join: JoinOptions) -> None: + require_v1() + cons = [] + for sparse in (False, True): + c = base_model() + lhs = (c.eff * c.gen_p).groupby(c.gbus).sum(sparse=sparse) + with no_densify(): + con = lhs.to_constraint("<=", c.load.isel(bus=[0, 2]), join=join) + cons.append(c.m.add_constraints(con, name="c", freeze=sparse)) + dense, frozen = cons + assert isinstance(frozen, CSRConstraint) + xr.testing.assert_identical( + xr.Dataset(coords=frozen.coords), xr.Dataset(coords=dense.coords) + ) + assert_frozen_equal(dense, frozen) + + 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 @@ -1742,14 +1908,155 @@ def test_selection_fallbacks_match_dense(op: str) -> None: } +@pytest.mark.parametrize("prebuilt", [False, True], ids=["lhs", "constraint"]) @pytest.mark.parametrize("mask", list(MASKS)) -def test_add_constraints_mask_freezes_sparse_and_matches_dense(mask: str) -> None: +def test_add_constraints_mask_freezes_sparse_and_matches_dense( + mask: str, prebuilt: bool +) -> 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) + if prebuilt: + con2 = c2.m.add_constraints(lhs >= c2.load, name="bal", mask=m, freeze=True) + else: + con2 = c2.m.add_constraints(lhs, ">=", c2.load, "bal", mask=m, freeze=True) + assert isinstance(con2, CSRConstraint) + assert_frozen_equal(con1, con2) + + +def observed_keys_group() -> LinearExpression: + """The #941 reproducer: a multi-key observed grouping with aux coords on ``group``.""" + m = Model() + x = m.add_variables(coords=[pd.RangeIndex(4, name="s")], name="x") + expr = x.to_linexpr().assign_coords( + period=("s", [1, 1, 2, 2]), region=("s", ["n", "s", "n", "s"]) + ) + return expr.groupby(["period", "region"]).sum(observed=True, sparse=True) + + +def dashed_group() -> LinearExpression: + """An observed grouping whose dims and aux coords carry dashes, the netcdf name separator.""" + m = Model() + s = pd.RangeIndex(4, name="my-s") + x = m.add_variables(coords=[s], name="x") + grouper = pd.DataFrame( + {"my-bus": ["a", "a", "b", "b"], "tag": [1, 1, 2, 2]}, index=s + ) + grouped = x.to_linexpr().groupby(grouper).sum(sparse=True, observed=True) + return grouped.rename({"group": "my-group"}) + + +AUX_BUILDS: dict[str, Callable[[], LinearExpression]] = { + "observed-keys": observed_keys_group, + "dashed": dashed_group, + "aux": lambda: SPARSE_BUILDS["aux"](base_model()), + "scalar": lambda: SPARSE_BUILDS["grouped"](base_model()).sel(snapshot=1), +} + + +@pytest.mark.parametrize("build", list(AUX_BUILDS)) +def test_frozen_constraint_keeps_aux_coords_like_dense( + build: str, tmp_path: Path +) -> None: + require_v1() + sparse = AUX_BUILDS[build]() + m = sparse.model + assert sparse._csr is not None + ref = m.add_constraints(sparse._csr.to_dense() >= 1, name="dense") + with no_densify(): + con = m.add_constraints(sparse >= 1, name="sparse", freeze=True) + assert isinstance(con, CSRConstraint) + want = xr.Dataset(coords=ref.coords) + assert set(want.coords) > set(con.coord_names) + for got in (con, con.mutable(), ref.freeze()): + xr.testing.assert_identical(xr.Dataset(coords=got.coords), want) + m.to_netcdf(tmp_path / "m.nc") + read = linopy.read_netcdf(tmp_path / "m.nc").constraints["sparse"] + assert isinstance(read, CSRConstraint) + xr.testing.assert_identical(xr.Dataset(coords=read.coords), want) + + +def setter(attr: str) -> Callable[[CSRConstraint], None]: + return lambda con: setattr(con, attr, 1.0) + + +FROZEN_MUTATIONS: dict[str, Callable[[CSRConstraint], Any]] = { + "loc": lambda con: con.loc[{"bus": "bus0"}], + "update": lambda con: con.update(rhs=2.0), + **{ + attr: setter(attr) + for attr in ["coeffs", "vars", "sign", "rhs", "lhs", "scaling"] + }, +} + + +@pytest.mark.parametrize("op", list(FROZEN_MUTATIONS)) +def test_frozen_constraint_mutation_names_mutable(op: str) -> None: + require_v1() + c = base_model() + con = c.m.add_constraints(c.balance_lhs(sparse=True) >= c.load, freeze=True) + assert isinstance(con, CSRConstraint) + with pytest.raises(AttributeError, match=rf"CSRConstraint\.{op} .*\.mutable\(\)"): + FROZEN_MUTATIONS[op](con) + + +FROZEN_UNSUPPORTED: dict[str, tuple[Callable[[CSRConstraint], Any], str]] = { + "soften": (lambda con: con.soften(penalty=1.0), r"freeze=False"), + "from_rule": ( + lambda con: CSRConstraint.from_rule(con.model, lambda m, i: None, [[0]]), + r"Constraint\.from_rule .*\.freeze\(\)", + ), +} + + +@pytest.mark.parametrize("op", list(FROZEN_UNSUPPORTED)) +def test_frozen_unsupported_names_the_working_route(op: str) -> None: + require_v1() + c = base_model() + con = c.m.add_constraints(c.balance_lhs(sparse=True) >= c.load, freeze=True) + assert isinstance(con, CSRConstraint) + call, remedy = FROZEN_UNSUPPORTED[op] + with pytest.raises(AttributeError, match=rf"CSRConstraint\.{op} .*{remedy}"): + call(con) + + +EXPR_RHS: dict[str, Callable[[Case, bool], LinearExpression]] = { + "expr": lambda c, sparse: (1.0 * c.flow).groupby(c.bus1).sum(sparse=sparse), + "expr-const": lambda c, sparse: ( + (1.0 * c.flow).groupby(c.bus1).sum(sparse=sparse) + c.load + ), + "dense-expr-const": lambda c, sparse: (1.0 * c.flow).groupby(c.bus1).sum() - 2.0, +} + + +@pytest.mark.parametrize("form", ["operator", "add_constraints"]) +@pytest.mark.parametrize("rhs", list(EXPR_RHS)) +def test_expression_rhs_freezes_sparse_and_matches_dense(rhs: str, form: str) -> None: + require_v1() + c1, c2 = base_model(), base_model() + lhs1 = (c1.eff * c1.gen_p).groupby(c1.gbus).sum() + con1 = c1.m.add_constraints(lhs1 <= EXPR_RHS[rhs](c1, False), name="c", freeze=True) + lhs2 = (c2.eff * c2.gen_p).groupby(c2.gbus).sum(sparse=True) + rhs2 = EXPR_RHS[rhs](c2, True) + with no_densify(): + if form == "operator": + con2 = c2.m.add_constraints(lhs2 <= rhs2, name="c", freeze=True) + else: + con2 = c2.m.add_constraints(lhs2, "<=", rhs2, name="c", freeze=True) assert isinstance(con2, CSRConstraint) assert_frozen_equal(con1, con2) + + +@pytest.mark.parametrize("default", [False, True], ids=["argument", "model-default"]) +def test_penalty_with_freeze_raises(default: bool) -> None: + require_v1() + c = base_model() + c.m.freeze_constraints = default + lhs = c.balance_lhs(sparse=True) + with pytest.raises(ValueError, match="`penalty` cannot be combined with `freeze"): + c.m.add_constraints( + lhs >= c.load, penalty=1.0, freeze=None if default else True + ) diff --git a/test/test_quadratic_expression.py b/test/test_quadratic_expression.py index 40ab8c41a..8355c5b3a 100644 --- a/test/test_quadratic_expression.py +++ b/test/test_quadratic_expression.py @@ -296,6 +296,18 @@ def test_quadratic_expression_to_polars(x: Variable, y: Variable) -> None: assert len(df) == expr.nterm * 2 +@pytest.mark.parametrize("factor_last", [False, True]) +def test_quadratic_expression_linear_terms( + x: Variable, y: Variable, factor_last: bool +) -> None: + expr = x * y + 3 * x + 0 * y + assert isinstance(expr, QuadraticExpression) + if factor_last: + expr = QuadraticExpression(expr.data.transpose(..., FACTOR_DIM), expr.model) + labels, coeffs = expr.linear_terms() + assert sorted(zip(labels.tolist(), coeffs.tolist())) == [(0, 3.0), (1, 3.0)] + + def test_quadratic_expression_constant_to_polars() -> None: m = Model() arr = pd.Series(index=pd.Index([0, 1], name="t"), data=[10, 20])