From f7eb1e17ceec7449a3ebb8ec59068fbd964a087c Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 23 Oct 2026 14:40:44 +0200 Subject: [PATCH 1/8] refactor(constraints): Constraint holds lhs store plus field arrays, data derived (#1046) Constraint stores lhs, sign, rhs, labels, scaling, dual, binary fields and blocks; data is built from them. Store.positional_csr/canonical; CSRStore caches the positional CSR by weakref. add_constraints, to_constraint, netcdf reader and Model.copy construct from fields. Step 4a of #1019. --- doc/release_notes.rst | 12 + linopy/constraints.py | 475 +++++++++++++++++----------- linopy/csr.py | 91 ++++++ linopy/expressions.py | 10 +- linopy/io.py | 31 +- linopy/model.py | 121 +++---- linopy/store.py | 56 ++++ test/test_constraint.py | 100 +++++- test/test_constraint_label_index.py | 10 +- test/test_io.py | 4 +- test/test_store_conformance.py | 45 +++ 11 files changed, 671 insertions(+), 284 deletions(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 46ee3004b..169b5d00b 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -40,6 +40,18 @@ Version 1.0.0 * A constant operation that produces NaN in a live cell, e.g. ``0 * inf``, ``0 / 0`` or ``inf - inf`` from multiplying, dividing, adding or subtracting a constant, now raises ``ValueError`` on both backends. Before, the cell kept its terms under a NaN constant. linopy trusts NaN only as absence, so mark cells to drop with ``.where(...)`` or ``mask=`` instead. Cells that were absent before stay absent. A sparse division by zero no longer emits numpy ``RuntimeWarning``\ s. (`#1045 `__) * Cells with a NaN constant now drop their terms when a ``LinearExpression`` or ``QuadraticExpression`` is built from a dataset or its ``vars`` or ``coeffs`` are assigned, so neither a dataset with terms under a NaN constant nor a ``vars`` assignment can revive an absent cell. (`#1045 `__) + +**Breaking Changes: constraints on stores** + +.. note:: + + The following content was generated by AI. + +* ``Constraint`` holds its left-hand side as a ``LinearExpression`` with a zero constant and its store in canonical form, next to ``sign``, ``rhs``, ``labels``, ``scaling``, ``dual``, ``binary_var``, ``binary_val`` and ``blocks`` as DataArrays on the grid of the lhs. ``Constraint.data`` is derived on each access, so writing into it, e.g. ``con.data["sign"] = ...``, no longer changes the constraint; use ``Constraint.update(...)``. ``Constraint.lhs`` returns a new expression on the held store on each access. (`#1046 `__) +* The ``Constraint`` constructor takes the fields instead of a dataset: ``Constraint(lhs, sign, rhs, model, name="", labels=None, scaling=1.0, dual=None, binary_var=None, binary_val=None, blocks=None, attrs=None)``; an lhs with a nonzero constant raises. Build a constraint from its dataset layout with ``Constraint.from_dataset(ds, model, name)``; the ``skip_broadcast`` argument is removed. A field with dimensions outside the lhs or with coordinates differing from those of the lhs raises, instead of broadcasting the lhs or filling the field with NaN. (`#1046 `__) +* ``Constraint.assign`` and the other wrapped xarray methods reject data variables other than the constraint fields, and ``Constraint.broadcast_like`` no longer broadcasts the fields over ``_term``. After ``Constraint.update(coeffs=...)`` or ``update(variables=...)``, ``coeffs`` and ``vars`` keep the ``_term`` dimension, a value without it is spread over the terms. (`#1046 `__) +* ``Constraints.set_blocks`` stores the block numbers as the ``blocks`` field of each constraint, read as ``con.data.blocks``. (`#1046 `__) + **Deprecations** * The ``active_fill`` parameter of ``add_piecewise_formulation`` is deprecated and emits a ``FutureWarning``; it is removed in 1.1. Pad a partial gate explicitly instead, e.g. ``active=status.reindex(gen=gens).fillna(1)``, which gives the same model. A partial ``active`` without padding still raises. diff --git a/linopy/constraints.py b/linopy/constraints.py index 206692781..5ee492494 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -82,14 +82,16 @@ from linopy.containers import LabelContainer from linopy.csr import ( Grid, + _csr_from_label_columns, _densify_notice, + _positional_csr, canonical, csr_nterm, csr_to_term_arrays, - index_dtype, ) -from linopy.scaling import ensure_scaling, validate_scaling +from linopy.scaling import validate_scaling from linopy.semantics import check_user_nan +from linopy.store import DenseStore, Store, _take_flat from linopy.types import ( ConstantLike, ConstraintLike, @@ -121,7 +123,7 @@ def conwrap( def _conwrap(con: Constraint, *args: Any, **kwargs: Any) -> Constraint: for k, v in new_default_kwargs.items(): kwargs.setdefault(k, v) - return con.__class__( + return con.from_dataset( method(con.data, *default_args, *args, **kwargs), con.model, con.name ) @@ -435,10 +437,9 @@ def __getitem__( Get selection from the constraint. Returns a Constraint with the selected data. """ - data = Dataset( - {k: self.data[k][selector] for k in self.data}, attrs=self.data.attrs - ) - return Constraint(data, self.model, self.name) + ds = self.data + data = Dataset({k: ds[k][selector] for k in ds}, attrs=ds.attrs) + return Constraint.from_dataset(data, self.model, self.name) @property def attrs(self) -> dict[str, Any]: @@ -661,51 +662,6 @@ def _equal_nnz_slices( start = end -def _csr_from_label_columns( - data: np.ndarray, - vlabel_cols: np.ndarray, - indptr: np.ndarray, - n_rows: int, - label_index: VariableLabelIndex, - name: str, -) -> scipy.sparse.csr_array: - """ - Build a CSR matrix from raw variable labels as column indices. - - Labels are mapped to dense positions in the active variable array, so the - result is ``(n_rows, label_index.n_active_vars)`` wide regardless of when - the label columns were assembled. ``data`` and, where its dtype is wide - enough, ``indptr`` are the very arrays passed in, so callers can compare - buffers by identity. - """ - n_cols = label_index.n_active_vars - cols = label_index.label_to_pos[vlabel_cols] - if cols.size and min(cols.min(), vlabel_cols.min()) < 0: - missing = vlabel_cols[(cols < 0) | (vlabel_cols < 0)] - raise ValueError( - f"Constraint '{name}' references variable labels that are no longer " - f"part of the model, e.g. {missing[:5].tolist()}." - ) - index_dtype = ( - indptr.dtype if n_cols <= np.iinfo(indptr.dtype).max else np.dtype(np.int64) - ) - return _positional_csr( - data, cols.astype(index_dtype, copy=False), indptr, n_rows, n_cols - ) - - -def _positional_csr( - data: np.ndarray, cols: np.ndarray, indptr: np.ndarray, n_rows: int, n_cols: int -) -> scipy.sparse.csr_array: - """Wrap positional CSR components without copying ``data`` or ``cols``.""" - csr = scipy.sparse.csr_array( - (data, cols, indptr.astype(cols.dtype, copy=False)), shape=(n_rows, n_cols) - ) - csr.data = data - csr.indices = cols - return csr - - _UNSUPPORTED = "is not supported on a frozen constraint" @@ -713,13 +669,6 @@ def _is_uniform(values: np.ndarray) -> bool: return values.size > 0 and bool((values == values[0]).all()) -def _take_flat(values: np.ndarray, flat: np.ndarray) -> np.ndarray: - """Gather C-order flat positions, copying a non-contiguous array only if dense.""" - if values.flags.c_contiguous or 4 * flat.size > values.size: - return values.ravel()[flat] - return values.flat[flat] - - def _slack_names(attrs: Mapping[str, Any]) -> tuple[str, str] | None: """Names of the positive and negative slack variables recorded by ``soften``.""" positive = attrs.get("slack_positive") @@ -1314,7 +1263,7 @@ def _to_positional_csr( csr.data, cols, csr.indptr, csr.shape[0], label_index.n_active_vars ) positional = _csr_from_label_columns( - csr.data, csr.indices, csr.indptr, csr.shape[0], label_index, self._name + csr.data, csr.indices, csr.indptr, csr.shape[0], label_index, self.name ) self._positional_cache = (csr, label_to_pos, weakref.ref(positional.indices)) return positional @@ -1505,7 +1454,7 @@ def to_dense(self) -> Constraint: self._model, explicit=True, ) - return Constraint(self.data, self._model, self._name) + return Constraint.from_dataset(self.data, self._model, self._name) def mutable(self) -> Constraint: """Return a modifiable Constraint (same as to_dense).""" @@ -1590,9 +1539,11 @@ def from_dense( cindex : int or None Starting label index; defaults to the one recorded on ``con``. """ - if cindex is None and "label_range" in con.data.attrs: - cindex = int(con.data.attrs["label_range"][0]) - _, active_mask, vlabel_cols, data, indptr = con._matrix_export_data() + if cindex is None and "label_range" in con.attrs: + cindex = int(con.attrs["label_range"][0]) + active_mask = con.active_row_mask() + store = con.lhs._store.to_dense() + data, vlabel_cols, indptr = store.active_terms(active_mask) active_positions = np.flatnonzero(active_mask) csr = scipy.sparse.csr_array( (data, vlabel_cols, indptr), @@ -1714,48 +1665,215 @@ class Slack(NamedTuple): negative: variables.Variable | None +_FIELDS = ( + "sign", + "rhs", + "labels", + "scaling", + "dual", + "binary_var", + "binary_val", + "blocks", +) +_DATASET_FIELDS = frozenset(("coeffs", "vars", *_FIELDS)) + + +def _on_grid(value: Any, template: DataArray, label: str) -> DataArray: + """``value`` broadcast onto the dims of ``template`` and in their order.""" + da = DataArray(value) + if extra := set(da.dims) - set(template.dims): + raise ValueError( + f"Constraint {label} has dimensions {sorted(map(str, extra))} outside " + f"the dimensions {list(map(str, template.dims))} of its lhs." + ) + indexes = template.indexes + aligned = all(i.equals(indexes[k]) for k, i in da.indexes.items() if k in indexes) + if aligned and da.dims != template.dims: + da = da.broadcast_like(template).transpose(*template.dims) + if not aligned or da.sizes != template.sizes: + raise ValueError( + f"Constraint {label} does not align with the coordinates of its lhs." + ) + missing = {k: v for k, v in template.coords.items() if k not in da.coords} + return da.assign_coords(missing) if missing else da + + class Constraint(ConstraintBase): """ - Constraint backed by an xarray Dataset. - - Supports setters, xarray operations via conwrap, and from_rule construction. + Constraint ``lhs sign rhs`` on the grid of its left-hand side. + + The lhs is a :class:`~linopy.expressions.LinearExpression` with a zero + constant, its store in canonical form. ``sign``, ``rhs``, ``labels``, + ``scaling``, ``dual``, ``binary_var``, ``binary_val`` and ``blocks`` are + DataArrays on its grid; :attr:`data` is the Dataset derived from them. + Supports setters, xarray operations via conwrap, and from_rule + construction. """ - __slots__ = ("_data", "_model", "_assigned", "_coef_dirty") + __slots__ = ( + "_lhs", + "_sign", + "_rhs", + "_labels", + "_scaling", + "_dual", + "_binary_var", + "_binary_val", + "_blocks", + "_grid", + "_attrs", + "_model", + "_coef_dirty", + ) + + _lhs: expressions.LinearExpression + _sign: DataArray + _rhs: DataArray + _labels: DataArray | None + _scaling: DataArray + _dual: DataArray | None + _binary_var: DataArray | None + _binary_val: DataArray | None + _blocks: DataArray | None + _grid: Dataset + _attrs: dict[str, Any] def __init__( self, - data: Dataset, + lhs: expressions.LinearExpression, + sign: SignLike, + rhs: ConstantLike, model: Model, name: str = "", - skip_broadcast: bool = False, + labels: DataArray | None = None, + scaling: ConstantLike = 1.0, + dual: DataArray | None = None, + binary_var: DataArray | None = None, + binary_val: ConstantLike | None = None, + blocks: DataArray | None = None, + attrs: Mapping[str, Any] | None = None, ) -> None: from linopy.model import Model - if not isinstance(data, Dataset): - raise ValueError(f"data must be a Dataset, got {type(data)}") - + if not isinstance(lhs, expressions.LinearExpression): + raise TypeError(f"lhs must be a LinearExpression, got {type(lhs)}") if not isinstance(model, Model): raise ValueError(f"model must be a Model, got {type(model)}") + self._model = model + self._attrs = {**(attrs or {}), "name": name} + self._coef_dirty = False + fields = { + "sign": sign, + "rhs": rhs, + "labels": labels, + "scaling": scaling, + "dual": dual, + "binary_var": binary_var, + "binary_val": binary_val, + "blocks": blocks, + } + self._bind(lhs, {k: v for k, v in fields.items() if v is not None}) + + @classmethod + def from_dataset(cls, ds: Dataset, model: Model, name: str = "") -> Constraint: + """ + Build a constraint from its dense Dataset layout, see :attr:`data`. + + The Dataset holds ``coeffs`` and ``vars`` on the grid dims followed + by ``_term``, and ``sign``, ``rhs`` and optionally ``labels``, + ``scaling``, ``dual``, ``binary_var``, ``binary_val`` and ``blocks``. + """ + if not isinstance(ds, Dataset): + raise ValueError(f"data must be a Dataset, got {type(ds)}") for attr in ("coeffs", "vars", "sign", "rhs"): - if attr not in data: + if attr not in ds: raise ValueError(f"missing '{attr}' in data") + if unknown := set(map(str, ds.data_vars)) - _DATASET_FIELDS: + raise ValueError(f"unknown constraint fields {sorted(unknown)} in data") + + (ds,) = xr.broadcast(ds, exclude=[TERM_DIM]) + lhs = expressions.LinearExpression(ds[["coeffs", "vars"]], model) + fields = {k: ds[k] for k in _FIELDS if k in ds} + return cls(lhs, model=model, name=name, attrs=ds.attrs, **fields) + + def _bind( + self, lhs: expressions.LinearExpression, fields: Mapping[Any, Any] + ) -> None: + """ + Hold ``lhs`` in canonical form and ``fields`` on its grid. Callers + replacing the lhs terms flip ``_coef_dirty``. + """ + store = lhs._store.canonical() + const = store.const + if np.any(const[~np.isnan(const)]): + raise ValueError( + "The lhs of a constraint must have a zero constant; move it to the rhs." + ) + template = store.dataarray(const, "const") + self._grid = template.to_dataset() + self._lhs = expressions.LinearExpression._from_store(store, self._model) + self._sign = _on_grid(fields["sign"], template, "sign") + self._rhs = _on_grid(fields["rhs"], template, "rhs") + self._scaling = _on_grid(fields.get("scaling", 1.0), template, "scaling") + self._labels = self._optional_on_grid(fields, "labels", template) + self._dual = self._optional_on_grid(fields, "dual", template) + self._binary_var = self._optional_on_grid(fields, "binary_var", template) + self._binary_val = self._optional_on_grid(fields, "binary_val", template) + self._blocks = self._optional_on_grid(fields, "blocks", template) + + @staticmethod + def _optional_on_grid( + fields: Mapping[Any, Any], key: str, template: DataArray + ) -> DataArray | None: + value = fields.get(key) + return None if value is None else _on_grid(value, template, key) + + def _grid_template(self) -> DataArray: + """A DataArray on the lhs grid, to broadcast a field onto.""" + return self._grid["const"] + + def _fields(self) -> dict[str, DataArray]: + """The fields next to the lhs that are set, in their Dataset order.""" + fields = { + "sign": self._sign, + "rhs": self._rhs, + "labels": self._labels, + "scaling": self._scaling, + "dual": self._dual, + "binary_var": self._binary_var, + "binary_val": self._binary_val, + "blocks": self._blocks, + } + return {k: v for k, v in fields.items() if v is not None} - data = data.assign_attrs(name=name) + def _on_store( + self, store: Store, model: Model, fields: Mapping[str, DataArray] + ) -> Constraint: + """A constraint with this one's attrs, the lhs held by ``store`` and ``fields``.""" + lhs = expressions.LinearExpression._from_store(store, model) + return Constraint(lhs, model=model, name=self.name, attrs=self._attrs, **fields) - if not skip_broadcast: - (data,) = xr.broadcast(data, exclude=[TERM_DIM]) - data = ensure_scaling(data, data.rhs) + def _copy(self, model: Model, deep: bool, include_solution: bool) -> Constraint: + """A copy on ``model``, without the solve artifacts unless ``include_solution``.""" + keep = self._fields().keys() if include_solution else set(self.data_attrs) + fields = {k: v.copy(deep=deep) for k, v in self._fields().items() if k in keep} + return self._on_store(self._lhs._store.copy(model, deep), model, fields) - self._assigned = "labels" in data - self._data = data - self._model = model - self._coef_dirty = False + def __getitem__( + self, selector: str | int | slice | list | tuple | dict + ) -> Constraint: + terms = self._lhs._store.to_dataset() + selected = Dataset({k: terms[k][selector] for k in terms}) + fields = {k: v[selector] for k, v in self._fields().items()} + return self._on_store(DenseStore(selected, self._model), self._model, fields) @property def data(self) -> Dataset: - return self._data + """The dense Dataset derived from the lhs and the fields.""" + terms = self._lhs._store.to_dataset() + fields = {"coeffs": terms.coeffs, "vars": terms.vars, **self._fields()} + return Dataset(fields, attrs=self._attrs)[list(fields)] @property def model(self) -> Model: @@ -1763,27 +1881,49 @@ def model(self) -> Model: @property def name(self) -> str: - return self.attrs["name"] + return self._attrs["name"] + + @property + def attrs(self) -> dict[str, Any]: + return self._attrs + + @property + def coords(self) -> DatasetCoordinates: + return self._grid.coords + + @property + def indexes(self) -> Indexes: + return self._grid.indexes + + @property + def dims(self) -> Frozen[Hashable, int]: + return self.sizes + + @property + def sizes(self) -> Frozen[Hashable, int]: + return Frozen({**self._grid_template().sizes, TERM_DIM: self.nterm}) + + @property + def coord_names(self) -> list[str]: + return get_dims_with_index_levels(self._grid, self.coord_dims) @property def is_assigned(self) -> bool: - return self._assigned + return self._labels is not None + + @property + def nterm(self) -> int: + return self._lhs._store.nterm @property def ncons(self) -> int: """Get the number of active constraints (non-masked, with at least one valid variable).""" - labels = self.labels.values - vars_arr = self.vars.values - if labels.ndim == 0: - return int(labels != FILL_VALUE["labels"] and (vars_arr != -1).any()) - return int( - ((labels != FILL_VALUE["labels"]) & (vars_arr != -1).any(axis=-1)).sum() - ) + return int(self.active_row_mask().sum()) @property def range(self) -> tuple[int, int]: """Return the range of the constraint.""" - return self.data.attrs["label_range"] + return self._attrs["label_range"] @property def loc(self) -> LocIndexer: @@ -1791,27 +1931,27 @@ def loc(self) -> LocIndexer: @property def labels(self) -> DataArray: - return self.data.get("labels", DataArray([])) + return DataArray([]) if self._labels is None else self._labels @property def coeffs(self) -> DataArray: - return self.data.coeffs + return self.lhs.coeffs @property def vars(self) -> DataArray: - return self.data.vars + return self.lhs.vars @property def sign(self) -> DataArray: - return self.data.sign + return self._sign @property def rhs(self) -> DataArray: - return self.data.rhs + return self._rhs @property def scaling(self) -> DataArray: - return self.data.scaling + return self._scaling @scaling.setter def scaling(self, value: ConstantLike) -> None: @@ -1822,52 +1962,45 @@ def scaling(self, value: ConstantLike) -> None: label=f"scaling for constraint '{self.name}'", ) scaling = validate_scaling(scaling, f"scaling for constraint '{self.name}'") - self._data = assign_fields(self.data, scaling=scaling) + self._scaling = _on_grid(scaling, self._grid_template(), "scaling") @property def is_indicator(self) -> bool: - return "binary_var" in self._data + return self._binary_var is not None @property def binary_var(self) -> DataArray | None: - return self._data.get("binary_var") + return self._binary_var @property def binary_val(self) -> int | np.ndarray | None: - if "binary_val" in self._data: - return self._data["binary_val"].values - return None + return None if self._binary_val is None else self._binary_val.values @property def lhs(self) -> expressions.LinearExpression: - data = self.data[["coeffs", "vars"]].rename({self.term_dim: TERM_DIM}) - return expressions.LinearExpression(data, self.model) + """The lhs, a new expression on the store the constraint holds.""" + return expressions.LinearExpression._from_store(self._lhs._store, self._model) def _assign_lhs( self, expr: expressions.LinearExpression, rhs: DataArray | None = None ) -> None: """ - Internal: replace coeffs/vars from ``expr``, adjusting rhs for - the expression's constant part. Sets ``_coef_dirty``. + Internal: replace the lhs by ``expr``, moving its constant part to + the rhs. Sets ``_coef_dirty``. """ base_rhs = self.rhs if rhs is None else rhs - self._data = self.data.drop_vars(["coeffs", "vars"]).assign( - coeffs=expr.coeffs, - vars=expr.vars, - rhs=base_rhs - expr.const, - ) + fields = self._fields() | {"rhs": base_rhs - expr.const} + self._bind(expr.reset_const(), fields) self._coef_dirty = True - def _update_data(self, **fields: Any) -> None: + def _assign_terms(self, **terms: DataArray) -> None: """ - Internal: write ``fields`` into ``self._data`` and update dirty bookkeeping. - - Writes that touch the lhs structure (``coeffs``, ``vars``) flip - ``_coef_dirty``. Other fields (``rhs``, ``sign``, …) leave it alone. + Internal: replace ``coeffs`` and/or ``vars`` of the dense lhs + layout. Sets ``_coef_dirty``. """ - self._data = assign_fields(self.data, **fields) - if "coeffs" in fields or "vars" in fields: - self._coef_dirty = True + ds = assign_fields(self.lhs.data[["coeffs", "vars"]], **terms) + self._bind(expressions.LinearExpression(ds, self._model), self._fields()) + self._coef_dirty = True def update( self, @@ -1980,7 +2113,6 @@ def update( "replacement), not both." ) - # 1. lhs replacement first so subsequent rhs= rearrangement sees the new lhs. if lhs is not None: expr = expressions.as_expression( lhs, self.model, coords=self.coords, dims=self.coord_dims @@ -1991,7 +2123,6 @@ def update( ) self._assign_lhs(expr) - # 2. rhs (rearranges non-constant part onto lhs). if rhs is not None: expr = expressions.as_expression( rhs, self.model, coords=self.coords, dims=self.coord_dims @@ -2000,14 +2131,13 @@ def update( if residual.nterm != 0: self._assign_lhs(self.lhs - residual, rhs=expr.const) else: - self._update_data(rhs=expr.const) + self._rhs = _on_grid(expr.const, self._grid_template(), "rhs") - # 3. coeffs / variables partial updates (only valid without lhs=). if coeffs is not None: new_coeffs = DataArray(coeffs).broadcast_like( self.vars, exclude=[self.term_dim] ) - self._update_data(coeffs=new_coeffs) + self._assign_terms(coeffs=new_coeffs) if variables is not None: from linopy.variables import Variable as _Variable @@ -2019,12 +2149,11 @@ def update( new_vars = variables.labels.broadcast_like( self.coeffs, exclude=[self.term_dim] ) - self._update_data(vars=new_vars) + self._assign_terms(vars=new_vars) - # 4. sign last so it composes cleanly with the rest. if sign is not None: new_sign = maybe_replace_signs(DataArray(sign)).broadcast_like(self.sign) - self._update_data(sign=new_sign) + self._sign = _on_grid(new_sign, self._grid_template(), "sign") return self @@ -2037,48 +2166,25 @@ def _assign_coords(self, **coords: Any) -> Constraint: reindexing, no shape change, and the order of the underlying data is preserved. """ - self._data = reassign_coords(self.data, **coords) + fields = reassign_coords(Dataset(self._fields()), **coords) + self._bind(self.lhs._assign_coords(**coords), fields) return self @property @has_optimized_model def dual(self) -> DataArray: - if "dual" not in self.data: + if self._dual is None: raise AttributeError( "Underlying is optimized but does not have dual values stored." ) - return self.data["dual"] + return self._dual @dual.setter def dual(self, value: ConstantLike) -> None: - value = DataArray(value).broadcast_like(self.labels) - self._data = assign_fields(self.data, dual=value) + self._dual = _on_grid(value, self._grid_template(), "dual") def has_labels(self, labels: np.ndarray) -> bool: - return contains_labels(self.data["vars"].values.ravel(), labels) - - def _matrix_export_data( - self, - ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: - """Return (con_labels, row_mask, vlabel_cols, data, indptr) with raw labels.""" - labels_flat = self.labels.values.ravel() - vars_vals = self.vars.values - coeffs_vals = self.coeffs.values - valid = vars_vals != -1 - row_mask = (labels_flat != -1) & valid.any(axis=-1).ravel() - con_labels = labels_flat[row_mask] - valid &= coeffs_vals != 0 - valid &= (self.labels.values != -1)[..., None] - flat = np.flatnonzero(valid) - vlabel_cols = _take_flat(vars_vals, flat) - data = _take_flat(coeffs_vals, flat) - flat //= valid.shape[-1] - counts = np.bincount(flat, minlength=len(row_mask))[row_mask] - 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 + return contains_labels(self.vars.values.ravel(), labels) def to_matrix( self, label_index: VariableLabelIndex @@ -2096,17 +2202,13 @@ def to_matrix( con_labels : np.ndarray Active constraint labels in row order. """ - con_labels, _, vlabel_cols, data, indptr = self._matrix_export_data() - csr = _csr_from_label_columns( - data, vlabel_cols, indptr, len(con_labels), label_index, self.name - ) - csr.sum_duplicates() - return csr, con_labels + active = self.active_row_mask() + csr = self._lhs._store.positional_csr(label_index, active, self.name) + return csr, self.labels.values.ravel()[active] def active_row_mask(self) -> np.ndarray: """Boolean mask over raveled rows: label set and at least one variable present.""" - present = (self.vars.values != -1).any(axis=-1).ravel() - return (self.labels.values.ravel() != -1) & present + return (self.labels.values.ravel() != -1) & self._lhs._store.has_terms() def active_labels(self) -> np.ndarray: return self.labels.values.ravel()[self.active_row_mask()] @@ -2115,20 +2217,17 @@ def to_matrix_with_rhs( self, label_index: VariableLabelIndex ) -> tuple[scipy.sparse.csr_array, np.ndarray, np.ndarray, np.ndarray]: """Return (csr, con_labels, b, sense) in one pass.""" - con_labels, row_mask, vlabel_cols, data, indptr = self._matrix_export_data() - csr = _csr_from_label_columns( - data, vlabel_cols, indptr, len(con_labels), label_index, self.name - ) - csr.sum_duplicates() - - b = self.rhs.values.ravel()[row_mask] - sense = self.sign.values.ravel()[row_mask].astype("U1") + active = self.active_row_mask() + csr = self._lhs._store.positional_csr(label_index, active, self.name) + con_labels = self.labels.values.ravel()[active] + b = self.rhs.values.ravel()[active] + sense = self.sign.values.ravel()[active].astype("U1") return csr, con_labels, b, sense def sanitize_zeros(self) -> Constraint: """Remove terms with zero or near-zero coefficients.""" not_zero = abs(self.coeffs) > 1e-10 - self._update_data( + self._assign_terms( vars=self.vars.where(not_zero, -1), coeffs=self.coeffs.where(not_zero), ) @@ -2137,8 +2236,7 @@ def sanitize_zeros(self) -> Constraint: def sanitize_missings(self) -> Constraint: """Mask out rows where all variables are missing (-1).""" contains_non_missing = (self.vars != -1).any(self.term_dim) - labels = self.labels.where(contains_non_missing, -1) - self._data = assign_fields(self.data, labels=labels) + self._labels = self.labels.where(contains_non_missing, -1) return self def sanitize_infinities(self) -> Constraint: @@ -2146,8 +2244,7 @@ def sanitize_infinities(self) -> Constraint: valid_infinity_values = ((self.sign == LESS_EQUAL) & (self.rhs == np.inf)) | ( (self.sign == GREATER_EQUAL) & (self.rhs == -np.inf) ) - labels = self.labels.where(~valid_infinity_values, -1) - self._data = assign_fields(self.data, labels=labels) + self._labels = self.labels.where(~valid_infinity_values, -1) return self def freeze(self) -> CSRConstraint: @@ -2222,8 +2319,7 @@ def from_rule(cls, model: Model, rule: Callable, coords: CoordsLike) -> Constrai ) sign = DataArray(array([c.sign for c in cons]).reshape(shape), coords) rhs = DataArray(array([c.rhs for c in cons]).reshape(shape), coords) - data = lhs.data.assign(sign=sign, rhs=rhs) - return cls(data, model=model) + return cls(lhs, sign, rhs, model) def _attach_slack(self, slack: Slack, sign: str) -> None: positive = slack.positive @@ -2231,7 +2327,7 @@ def _attach_slack(self, slack: Slack, sign: str) -> None: if slack.negative is not None: lhs = lhs + slack.negative self.update(lhs=lhs) - self._data = self._data.assign_attrs( + self._attrs.update( slack_positive=slack.positive.name, slack_negative=slack.negative.name if slack.negative is not None else "", ) @@ -2290,7 +2386,7 @@ def to_polars(self) -> pl.DataFrame: assign = conwrap(assign_fields) assign_attrs = conwrap(Dataset.assign_attrs) assign_coords = conwrap(Dataset.assign_coords) - broadcast_like = conwrap(Dataset.broadcast_like) + broadcast_like = conwrap(Dataset.broadcast_like, exclude=[TERM_DIM]) chunk = conwrap(Dataset.chunk) drop_sel = conwrap(Dataset.drop_sel) drop_isel = conwrap(Dataset.drop_isel) @@ -2537,7 +2633,7 @@ def set_blocks(self, block_map: np.ndarray) -> None: _densify_notice( "frozen constraint converted by `set_blocks`", self.model ) - constraint = Constraint(constraint.data, self.model, name) + constraint = Constraint.from_dataset(constraint.data, self.model, name) self.data[name] = constraint res = xr.full_like(constraint.labels, N + 1, dtype=block_map.dtype) entries = replace_by_map(constraint.vars, block_map) @@ -2551,7 +2647,7 @@ def set_blocks(self, block_map: np.ndarray) -> None: res = res.where(not_missing.any(constraint.term_dim), -1) res = res.where(not_zero.any(constraint.term_dim), 0) - constraint._data = assign_fields(constraint.data, blocks=res) + constraint._blocks = res @property def flat(self) -> pd.DataFrame: @@ -2611,8 +2707,7 @@ def reset_dual(self) -> None: if c._dual is not None: self.data[k] = c._replace(dual=None) elif isinstance(c, Constraint): - if "dual" in c.data: - c._data = c.data.drop_vars("dual") + c._dual = None else: msg = f"reset_dual encountered an unknown constraint type: {type(c)}" raise NotImplementedError(msg) @@ -2676,5 +2771,5 @@ def rhs(self) -> int | float | np.floating | np.integer: return self._rhs def to_constraint(self) -> Constraint: - data = self.lhs.to_linexpr().data.assign(sign=self.sign, rhs=self.rhs) - return Constraint(data=data, model=self.lhs.model) + model = self.lhs.model + return Constraint(self.lhs.to_linexpr(), self.sign, self.rhs, model) diff --git a/linopy/csr.py b/linopy/csr.py index 830603917..7302a6094 100644 --- a/linopy/csr.py +++ b/linopy/csr.py @@ -41,6 +41,7 @@ import functools import math import operator +import weakref from collections.abc import Callable, Iterable, Mapping from dataclasses import dataclass, field, replace from typing import TYPE_CHECKING, Any, TypeAlias, final @@ -57,6 +58,7 @@ from linopy.store import DenseStore, Reason, Store, merged_grid if TYPE_CHECKING: + from linopy.common import VariableLabelIndex from linopy.model import Model CONTRACTION_CHUNK = 64 @@ -271,6 +273,9 @@ class CSRStore(Store): const: np.ndarray grid: Grid model: Model + _positional_cache: _PositionalCache | None = field( + default=None, init=False, repr=False, compare=False + ) def __post_init__(self) -> None: csr = self.csr @@ -551,6 +556,46 @@ def solution( def canonical(self) -> CSRStore: return replace(self, csr=canonical(self.csr)) + def positional_csr( + self, label_index: VariableLabelIndex, active_rows: np.ndarray, name: str + ) -> scipy.sparse.csr_array: + """ + The stored rows at ``active_rows`` with label columns replaced by + dense positions, ``data`` and ``indptr`` the stored arrays where all + rows are active. The result arrays are cached by weak reference, + keyed on ``label_index.label_to_pos``, which is rebuilt whenever + variables are added or removed, and on ``active_rows``: while a + caller such as a ``ModelSnapshot`` holds them, repeated calls return + the same arrays. The store is immutable, so a new store starts + without a cache. + """ + label_to_pos = label_index.label_to_pos + n_cols = label_index.n_active_vars + if self._positional_cache is not None: + cached_label_to_pos, cached_rows, refs = self._positional_cache + arrays = [ref() for ref in refs] + data, cols, indptr = arrays + if ( + cached_label_to_pos is label_to_pos + and data is not None + and cols is not None + and indptr is not None + and np.array_equal(cached_rows, active_rows) + ): + return _positional_csr(data, cols, indptr, len(indptr) - 1, n_cols) + rows = self.csr if active_rows.all() else self.csr[active_rows] + positional = _csr_from_label_columns( + rows.data, rows.indices, rows.indptr, rows.shape[0], label_index, name + ) + refs = ( + weakref.ref(positional.data), + weakref.ref(positional.indices), + weakref.ref(positional.indptr), + ) + cache = (label_to_pos, active_rows.copy(), refs) + object.__setattr__(self, "_positional_cache", cache) + return positional + def with_grid(self, grid: Grid) -> CSRStore: return replace(self, grid=grid) @@ -836,6 +881,52 @@ def to_dense(self, reason: Reason | None = None) -> DenseStore: return DenseStore.from_arrays(self.grid, coeffs, vars_, self.const, self.model) +_PositionalCache: TypeAlias = tuple[ + np.ndarray, np.ndarray, tuple["weakref.ref[np.ndarray]", ...] +] + + +def _csr_from_label_columns( + data: np.ndarray, + vlabel_cols: np.ndarray, + indptr: np.ndarray, + n_rows: int, + label_index: VariableLabelIndex, + name: str, +) -> scipy.sparse.csr_array: + """ + Build a CSR matrix from raw variable labels as column indices. + + Labels are mapped to dense positions in the active variable array, so the + result is ``(n_rows, label_index.n_active_vars)`` wide regardless of when + the label columns were assembled. ``data`` and, where its dtype is wide + enough, ``indptr`` are the very arrays passed in, so callers can compare + buffers by identity. + """ + n_cols = label_index.n_active_vars + cols = label_index.label_to_pos[vlabel_cols] + if cols.size and min(cols.min(), vlabel_cols.min()) < 0: + missing = vlabel_cols[(cols < 0) | (vlabel_cols < 0)] + raise ValueError( + f"Constraint '{name}' references variable labels that are no longer " + f"part of the model, e.g. {missing[:5].tolist()}." + ) + dtype = indptr.dtype if n_cols <= np.iinfo(indptr.dtype).max else np.int64 + return _positional_csr(data, cols.astype(dtype, copy=False), indptr, n_rows, n_cols) + + +def _positional_csr( + data: np.ndarray, cols: np.ndarray, indptr: np.ndarray, n_rows: int, n_cols: int +) -> scipy.sparse.csr_array: + """Wrap positional CSR components without copying ``data`` or ``cols``.""" + csr = scipy.sparse.csr_array( + (data, cols, indptr.astype(cols.dtype, copy=False)), shape=(n_rows, n_cols) + ) + csr.data = data + csr.indices = cols + return csr + + 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 diff --git a/linopy/expressions.py b/linopy/expressions.py index fd72660a4..cfaa3cb1d 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -1672,12 +1672,12 @@ def to_constraint( if rhs.isnull().any(): check_user_nan() - all_to_lhs = self.sub(rhs, join=join).data - computed_rhs = -all_to_lhs.const - data = assign_fields( - all_to_lhs[["coeffs", "vars"]], sign=sign, rhs=computed_rhs + all_to_lhs = self.sub(rhs, join=join) + if not isinstance(all_to_lhs, LinearExpression): + raise TypeError("The lhs of a constraint must be a linear expression.") + return constraints.Constraint( + all_to_lhs.reset_const(), sign, -all_to_lhs.const, self.model ) - return constraints.Constraint(data, model=self.model) def _broadcast_rhs(self, rhs: ConstantLike) -> DataArray: """Broadcast a constant RHS onto the expression's coordinates.""" diff --git a/linopy/io.py b/linopy/io.py index e478d5f04..fcb87151a 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -1289,7 +1289,7 @@ def owner(k: str) -> str | None: if con_ds.attrs.get("_linopy_format") == "csr": constraints[name] = CSRConstraint.from_netcdf_ds(con_ds, m, name) else: - constraints[name] = Constraint(con_ds, m, name) + constraints[name] = Constraint.from_dataset(con_ds, m, name) m._constraints = Constraints(constraints, m) objective = get_prefix(ds, "objective") @@ -1427,25 +1427,22 @@ def _copy_expr( new_model, ) - def _copy_con(name: str, con: ConstraintBase) -> ConstraintBase: - if isinstance(con, CSRConstraint): - buffer_types = (np.ndarray, scipy.sparse.sparray) - changes: dict[str, Any] = {"model": new_model} - if deep: - kwargs = con._init_kwargs().items() - changes |= { - k: v.copy() for k, v in kwargs if isinstance(v, buffer_types) - } - if not include_solution: - changes["dual"] = None - return con._replace(**changes) - d = con.data + def _copy_con(con: ConstraintBase) -> ConstraintBase: + if isinstance(con, Constraint): + return con._copy(new_model, deep, include_solution) + if not isinstance(con, CSRConstraint): + raise TypeError(f"Cannot copy constraint of type {type(con)}") + buffer_types = (np.ndarray, scipy.sparse.sparray) + changes: dict[str, Any] = {"model": new_model} + if deep: + kwargs = con._init_kwargs().items() + changes |= {k: v.copy() for k, v in kwargs if isinstance(v, buffer_types)} if not include_solution: - d = d[con.data_attrs] - return Constraint(d.copy(deep=deep), new_model, name) + changes["dual"] = None + return con._replace(**changes) new_model._constraints = Constraints( - {name: _copy_con(name, con) for name, con in m.constraints.items()}, + {name: _copy_con(con) for name, con in m.constraints.items()}, new_model, ) diff --git a/linopy/model.py b/linopy/model.py index c28a20f99..2add8e245 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -30,7 +30,6 @@ from linopy.alignment import as_dataarray, broadcast_to_coords from linopy.common import ( _as_renamed_index, - assign_fields, assigned_labels, best_int, coords_reorder_set, @@ -43,7 +42,6 @@ from linopy.constants import ( FACTOR_DIM, GREATER_EQUAL, - HELPER_DIMS, LESS_EQUAL, SOS_BIG_M_ATTR, SOS_DIM_ATTR, @@ -1157,18 +1155,19 @@ def _constraint_from_lhs( ) def _allocate_constraint_labels( - self, data: Dataset, name: str, mask: DataArray | None = None - ) -> Dataset: - """Assign label ranges from the constraint counter and apply an optional mask.""" + self, like: DataArray, mask: DataArray | None = None + ) -> tuple[DataArray, tuple[int, int]]: + """ + Labels shaped like ``like`` from the constraint counter, ``-1`` where + the optional ``mask`` is False, and their label range. + """ start = self._cCounter - end = start + data.labels.size - data.labels.values = self._allocate_labels(start, end).reshape( - data.labels.shape - ) - self._cCounter += data.labels.size + end = start + like.size + labels = like.copy(data=self._allocate_labels(start, end).reshape(like.shape)) + self._cCounter = end if mask is not None: - data.labels.values = np.where(mask.values, data.labels.values, -1) - return data.assign_attrs(label_range=(start, end), name=name) + labels = labels.where(mask, -1) + return labels, (start, end) @overload def add_constraints( @@ -1335,18 +1334,18 @@ def add_constraints( _densify_notice( "constraint added unfrozen, `freeze=False`", self, explicit=True ) - data = con.data - - _check_infinities(data.sign, data.rhs, name) + dense = ( + con + if isinstance(con, Constraint) + else Constraint.from_dataset(con.data, self) + ) - # ensure helper dimensions are not set as coordinates - if drop_dims := set(HELPER_DIMS).intersection(data.coords): - # TODO: add a warning here, routines should be safe against this - data = data.drop_vars(drop_dims) + _check_infinities(dense.sign, dense.rhs, name) - rhs_nan = data.rhs.isnull() + rhs = dense.rhs + rhs_nan = rhs.isnull() if rhs_nan.any(): - data = assign_fields(data, rhs=data.rhs.fillna(0)) + rhs = rhs.fillna(0) rhs_mask = ~rhs_nan mask = ( rhs_mask @@ -1354,30 +1353,34 @@ def add_constraints( else (as_dataarray(mask).astype(bool) & rhs_mask) ) - data["labels"] = -1 - (data,) = xr.broadcast(data, exclude=[TERM_DIM]) - row_coords = data.labels.coords - data = assign_fields( - data, - scaling=validate_scaling( - broadcast_to_coords(scaling, row_coords, label="constraint scaling"), - "constraint scaling", - ), + row_scaling = validate_scaling( + broadcast_to_coords(scaling, rhs.coords, label="constraint scaling"), + "constraint scaling", ) if mask is not None: - mask = broadcast_to_coords(mask, data.coords, label="mask").astype(bool) - - self.check_force_dim_names(data) - self._check_valid_dim_names(data) - - data = self._allocate_constraint_labels(data, name, mask) + mask = broadcast_to_coords(mask, dense.coords, label="mask").astype(bool) + + lhs = dense.lhs + self.check_force_dim_names(lhs) + self._check_valid_dim_names(lhs) + + labels, label_range = self._allocate_constraint_labels(rhs, mask) + constraint = Constraint( + lhs, + dense.sign, + rhs, + self, + name, + labels=labels, + scaling=row_scaling, + attrs={**dense.attrs, "label_range": label_range}, + ) + enforce_no_multiindex(constraint, context=f"constraint {name!r}") if self.chunk: - data = data.chunk(self.chunk) + constraint = constraint.chunk(self.chunk) - enforce_no_multiindex(data, context=f"constraint {name!r}") - constraint = Constraint(data, name=name, model=self, skip_broadcast=True) return self._soften_added(self.constraints.add(constraint, freeze), penalty) @staticmethod @@ -1447,28 +1450,28 @@ def add_indicator_constraints( if sign is not None: sign = maybe_replace_signs(as_dataarray(sign)) - data = self._constraint_from_lhs(lhs, sign, rhs).data - - data["binary_var"] = binary_var.labels - data["binary_val"] = binary_val - - data["labels"] = -1 - (data,) = xr.broadcast(data, exclude=[TERM_DIM]) - row_coords = data.labels.coords - data = assign_fields( - data, - scaling=validate_scaling( - broadcast_to_coords( - scaling, row_coords, label="indicator constraint scaling" - ), - "indicator constraint scaling", + base = self._constraint_from_lhs(lhs, sign, rhs) + rhs_da = base.rhs + row_scaling = validate_scaling( + broadcast_to_coords( + scaling, rhs_da.coords, label="indicator constraint scaling" ), + "indicator constraint scaling", ) - - data = self._allocate_constraint_labels(data, name) - - enforce_no_multiindex(data, context=f"constraint {name!r}") - con = Constraint(data, name=name, model=self, skip_broadcast=True) + labels, label_range = self._allocate_constraint_labels(rhs_da) + con = Constraint( + base.lhs, + base.sign, + rhs_da, + self, + name, + labels=labels, + scaling=row_scaling, + binary_var=binary_var.labels, + binary_val=binary_val, + attrs={**base.attrs, "label_range": label_range}, + ) + enforce_no_multiindex(con, context=f"constraint {name!r}") return self.constraints.add(con, freeze=self._sparse) def remove_indicator_constraints(self, name: str) -> None: diff --git a/linopy/store.py b/linopy/store.py index d70323bef..9f0cb7406 100644 --- a/linopy/store.py +++ b/linopy/store.py @@ -42,6 +42,7 @@ if TYPE_CHECKING: from xarray.core.types import JoinOptions + from linopy.common import VariableLabelIndex from linopy.csr import AuxCoords, CSRStore, Grid from linopy.model import Model @@ -206,6 +207,16 @@ def prune(self) -> Self: def canonical(self) -> Self: """The canonical form a constraint is built from.""" + @abstractmethod + def positional_csr( + self, label_index: VariableLabelIndex, active_rows: np.ndarray, name: str + ) -> scipy.sparse.csr_array: + """ + The cells selected by the boolean ``active_rows`` as CSR rows over the + dense variable positions of ``label_index``, raising for labels of + removed variables with the constraint ``name``. + """ + @abstractmethod def with_grid(self, grid: Grid) -> Self: """The same cells on ``grid``, a grid of the same shape.""" @@ -550,6 +561,44 @@ def prune(self) -> DenseStore: def canonical(self) -> DenseStore: return self + def active_terms( + self, active_rows: np.ndarray + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Coefficients, raw variable labels and row pointers of the live terms + of the cells selected by the boolean ``active_rows``, in CSR layout. + """ + from linopy.csr import index_dtype + + dims = (*self.dims, TERM_DIM) + vars_ = self.ds.vars.transpose(*dims).values + coeffs = self.ds.coeffs.transpose(*dims).values + valid = (vars_ != -1) & (coeffs != 0) + valid &= active_rows.reshape(valid.shape[:-1])[..., None] + flat = np.flatnonzero(valid) + labels = _take_flat(vars_, flat) + data = _take_flat(coeffs, flat) + counts = np.bincount(flat // max(self.nterm, 1), minlength=self.n_cells) + counts = counts[active_rows] + indptr = np.empty( + len(counts) + 1, index_dtype(len(data), counts.shape, self.model) + ) + indptr[0] = 0 + np.cumsum(counts, out=indptr[1:]) + return data, labels, indptr + + def positional_csr( + self, label_index: VariableLabelIndex, active_rows: np.ndarray, name: str + ) -> scipy.sparse.csr_array: + from linopy.csr import _csr_from_label_columns + + data, labels, indptr = self.active_terms(active_rows) + csr = _csr_from_label_columns( + data, labels, indptr, len(indptr) - 1, label_index, name + ) + csr.sum_duplicates() + return csr + def with_grid(self, grid: Grid) -> DenseStore: coeffs, vars_ = self._term_arrays() return self.from_arrays(grid, coeffs, vars_, self.const, self.model) @@ -584,6 +633,13 @@ def concat(self, others: Iterable[DenseStore], dim: str, grid: Grid) -> DenseSto return self.merged((self, *others), dim, "override", np.nan, {}) +def _take_flat(values: np.ndarray, flat: np.ndarray) -> np.ndarray: + """Gather C-order flat positions, copying a non-contiguous array only if dense.""" + if values.flags.c_contiguous or 4 * flat.size > values.size: + return values.ravel()[flat] + return values.flat[flat] + + def as_sparse(store: Store) -> CSRStore: """``store`` as a sparse store, raising where it cannot convert.""" res = store.to_sparse() diff --git a/test/test_constraint.py b/test/test_constraint.py index d11fb8a26..bbd921871 100644 --- a/test/test_constraint.py +++ b/test/test_constraint.py @@ -33,6 +33,7 @@ ConstraintBase, Constraints, ) +from linopy.csr import canonical from linopy.testing import assert_linequal, assert_varequal @@ -298,7 +299,7 @@ def test_constraint_wrapped_methods(x: linopy.Variable, y: linopy.Variable) -> N assert isinstance(con, Constraint) # Test wrapped methods - con.assign(new_var=xr.DataArray(np.zeros((2, 2)), coords=[range(2), range(2)])) + con.assign(rhs=xr.zeros_like(con.rhs)) con.assign_attrs({"new_attr": "value"}) con.assign_coords( {"new_coord": xr.DataArray(np.zeros((2, 2)), coords=[range(2), range(2)])} @@ -400,7 +401,7 @@ def test_constraint_update_variables( mc: linopy.constraints.Constraint, x: linopy.Variable ) -> None: mc.update(variables=x) - assert_equal(mc.vars, x.labels) + assert_equal(mc.vars.isel({TERM_DIM: 0}), x.labels) @pytest.mark.parametrize("as_type", [lambda v: v, pd.DataFrame]) @@ -657,7 +658,7 @@ def test_constraint_to_polars_mixed_signs(m: Model, x: linopy.Variable) -> None: # Replace sign data with mixed signs across the first dimension n = con.sizes["first"] signs = np.array(["<=" if i % 2 == 0 else ">=" for i in range(n)]) - con.data["sign"] = xr.DataArray(signs, dims=con.data["sign"].dims) + con.update(sign=xr.DataArray(signs, dims=con.sign.dims)) df = con.to_polars() assert isinstance(df, pl.DataFrame) assert set(df["sign"].to_list()) == {"<=", ">="} @@ -797,7 +798,7 @@ def test_constraint_with_helper_dims_as_coords(m: Model) -> None: data = xr.Dataset({"coeffs": coeffs, "vars": vars, "sign": sign, "rhs": rhs}) assert set(HELPER_DIMS).intersection(set(data.coords)) - con = Constraint(data, m, "c") + con = Constraint.from_dataset(data, m, "c") expr = m.add_constraints(con) assert not set(HELPER_DIMS).intersection(set(expr.coords)) @@ -977,7 +978,7 @@ def test_from_dense_mixed_signs() -> None: m.add_constraints(x >= 0, name="mixed", freeze=False) mc = m.constraints["mixed"] assert isinstance(mc, Constraint) - mc._data["sign"] = xr.DataArray(["<=", ">=", "<="], dims=["i"]) + mc.update(sign=xr.DataArray(["<=", ">=", "<="], dims=["i"])) frozen = linopy.constraints.CSRConstraint.from_dense(mc) assert isinstance(frozen._sign, np.ndarray) assert list(frozen._sign) == ["<=", ">=", "<="] @@ -1092,8 +1093,10 @@ def test_mixed_sign_sanitize_infinities() -> None: m.add_constraints(x >= 0, name="c", freeze=False) mc = m.constraints["c"] assert isinstance(mc, Constraint) - mc._data["sign"] = xr.DataArray(["<=", ">=", "<=", ">="], dims=["i"]) - mc._data["rhs"] = xr.DataArray([np.inf, -np.inf, 1.0, 2.0], dims=["i"]) + mc.update( + sign=xr.DataArray(["<=", ">=", "<=", ">="], dims=["i"]), + rhs=xr.DataArray([np.inf, -np.inf, 1.0, 2.0], dims=["i"]), + ) frozen = mc.freeze() frozen.sanitize_infinities() assert frozen.ncons == 2 @@ -1371,3 +1374,86 @@ def test_constraint_slack_matches_returned_slack_for_eq( assert slack.negative is not None assert resolved.negative is not None assert_varequal(resolved.negative, slack.negative) + + +def test_constraint_lhs_is_new_wrapper_on_held_store( + m: Model, x: linopy.Variable +) -> None: + con = m.add_constraints(x >= 0, name="d", freeze=False) + assert con.lhs is not con.lhs + assert con.lhs._store is con.lhs._store + + +def test_constraint_canonicalizes_sparse_lhs() -> None: + m = Model(sparse=True) + x = m.add_variables(coords=[pd.RangeIndex(3, name="i")], name="x") + lhs = x + 2 * x + assert lhs.is_sparse and lhs.nterm == 2 + con = Constraint(lhs, "<=", 1, m) + store = con.lhs._store + assert canonical(store.csr) is store.csr + assert store.nterm == 1 + con.update(lhs=x - 3 * x) + updated = con.lhs._store + assert updated.nterm == 1 + assert (updated.csr.data == -2).all() + + +def test_constraint_data_roundtrips_through_from_dataset( + m: Model, x: linopy.Variable +) -> None: + con = m.add_constraints(2 * x >= 1, name="d", freeze=False) + con.dual = xr.zeros_like(con.rhs) + restored = Constraint.from_dataset(con.data, m, con.name) + xr.testing.assert_identical(restored.data, con.data) + assert restored.attrs == con.attrs + + +def test_from_dataset_rejects_unknown_field(m: Model, x: linopy.Variable) -> None: + con = m.add_constraints(x >= 1, name="d", freeze=False) + with pytest.raises(ValueError, match="unknown constraint fields"): + Constraint.from_dataset(con.data.assign(foo=con.rhs), m) + + +def test_constraint_rejects_lhs_with_constant(m: Model, x: linopy.Variable) -> None: + with pytest.raises(ValueError, match="zero constant"): + Constraint(x + 1, "<=", 0, m) + + +def test_set_blocks_writes_blocks_field() -> None: + m = Model() + x = m.add_variables(coords=[pd.RangeIndex(3, name="i")], name="x") + y = m.add_variables(coords=[pd.RangeIndex(3, name="i")], name="y") + m.add_constraints(x >= 0, name="cx") + m.add_constraints(x + y >= 0, name="cxy") + m.constraints.set_blocks(np.array([1, 1, 1, 2, 2, 2])) + assert m.constraints["cx"].data.blocks.values.tolist() == [1, 1, 1] + assert m.constraints["cxy"].data.blocks.values.tolist() == [3, 3, 3] + + +@pytest.mark.parametrize( + "rhs", + [ + xr.DataArray(np.ones((10, 2)), dims=["first", "extra"]), + xr.DataArray([5.0, 6.0], coords={"first": [0, 1]}), + xr.DataArray(np.ones(10), coords={"first": np.arange(1, 11)}), + ], + ids=["extra-dim", "subset-coords", "shifted-coords"], +) +def test_constraint_rejects_field_off_lhs_grid( + m: Model, x: linopy.Variable, rhs: xr.DataArray +) -> None: + with pytest.raises(ValueError, match="Constraint rhs"): + Constraint(1 * x, "<=", rhs, m) + + +@pytest.mark.parametrize("freeze", [True, False]) +def test_to_matrix_names_constraint_with_removed_variable(freeze: bool) -> None: + m = Model() + i = pd.RangeIndex(3, name="i") + a = m.add_variables(coords=[i], name="a") + b = m.add_variables(coords=[i], name="b") + con = m.add_constraints(a + b >= 1, name="c", freeze=freeze) + m.variables.remove("a") + with pytest.raises(ValueError, match="Constraint 'c' references variable"): + con.to_matrix(m.variables.label_index) diff --git a/test/test_constraint_label_index.py b/test/test_constraint_label_index.py index 3652d5637..34546a3c8 100644 --- a/test/test_constraint_label_index.py +++ b/test/test_constraint_label_index.py @@ -4,7 +4,7 @@ import linopy import linopy.constants -from linopy.constraints import Constraint +from linopy.store import DenseStore @pytest.fixture @@ -30,13 +30,13 @@ def test_assign_result_does_not_build_matrix( monkeypatch: pytest.MonkeyPatch, model_with_mask: linopy.Model ) -> None: calls = {"n": 0} - original = Constraint._matrix_export_data + original = DenseStore.positional_csr - def counting(self, label_index): # type: ignore[no-untyped-def] + def counting(self, label_index, active_rows, name): # type: ignore[no-untyped-def] calls["n"] += 1 - return original(self, label_index) + return original(self, label_index, active_rows, name) - monkeypatch.setattr(Constraint, "_matrix_export_data", counting) + monkeypatch.setattr(DenseStore, "positional_csr", counting) model_with_mask.solve("highs") diff --git a/test/test_io.py b/test/test_io.py index 34a0ff78c..74fa029bf 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -1151,7 +1151,9 @@ def test_to_file_lp_unsorted_constraint_labels(tmp_path: Path) -> None: expected = _lp_constraint_section(m, tmp_path / "sorted.lp") con = m.constraints["a"] - m.constraints.data["a"] = Constraint(con.data.isel(i=slice(None, None, -1)), m, "a") + m.constraints.data["a"] = Constraint.from_dataset( + con.data.isel(i=slice(None, None, -1)), m, "a" + ) assert not m.constraints["a"].to_polars()["labels"].is_sorted() assert _lp_constraint_section(m, tmp_path / "unsorted.lp") == expected diff --git a/test/test_store_conformance.py b/test/test_store_conformance.py index 92a2527a1..3a3472056 100644 --- a/test/test_store_conformance.py +++ b/test/test_store_conformance.py @@ -612,3 +612,48 @@ def test_named_expression_keeps_name_through_assignment( assert named.data.attrs["name"] == "named" assign(named) assert named.name == backend.expressions["named"].name == "named" + + +def positional_model() -> tuple[Model, LinearExpression]: + m = Model() + x = m.add_variables(coords=[A], name="x") + y = m.add_variables(coords=[A], name="y") + return m, 2 * dense_linexpr(x) + dense_linexpr(y) - dense_linexpr(x) + + +ROWS = {"all": [True, True, True], "masked": [True, False, True]} + + +@pytest.mark.parametrize("rows", list(ROWS)) +def test_positional_csr_agrees(rows: str) -> None: + m, expr = positional_model() + active = np.array(ROWS[rows]) + label_index = m.variables.label_index + dense = expr._store.positional_csr(label_index, active, "c") + sparse = ( + expr._store.to_sparse().canonical().positional_csr(label_index, active, "c") + ) + assert dense.shape == sparse.shape == (active.sum(), 6) + assert (dense != sparse).nnz == 0 + + +@pytest.mark.parametrize("rows", list(ROWS)) +def test_sparse_positional_csr_cached_while_held(rows: str) -> None: + m, expr = positional_model() + store = expr._store.to_sparse().canonical() + active = np.array(ROWS[rows]) + first = store.positional_csr(m.variables.label_index, active, "c") + again = store.positional_csr(m.variables.label_index, active, "c") + assert all( + a is b + for a, b in zip( + (first.data, first.indices, first.indptr), + (again.data, again.indices, again.indptr), + ) + ) + other_rows = store.positional_csr(m.variables.label_index, ~active, "c") + assert other_rows.indices is not first.indices + m.add_variables(name="z") + rebuilt = store.positional_csr(m.variables.label_index, active, "c") + assert rebuilt.indices is not first.indices + np.testing.assert_array_equal(rebuilt.indices, first.indices) From b7f7824e19d419ae6e36fed5cecf89cd1075ad76 Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 23 Oct 2026 15:39:19 +0200 Subject: [PATCH 2/8] refactor(constraints): one Constraint class, CSRConstraint/ConstraintBase/freeze API removed (#1047) A constraint in a sparse model is a Constraint whose lhs holds a CSRStore. update(), soften and set_blocks keep a sparse lhs (closes #1033). MatrixAccessor, dual write-back and Model.copy lose their isinstance branches. Twin-build test helpers deleted; tests parametrized over backends. Step 4b of #1019. --- doc/api.rst | 61 +- doc/release_notes.rst | 5 + doc/user-guide.rst | 7 +- examples/creating-constraints.ipynb | 52 +- examples/sparse-models.ipynb | 31 +- linopy/__init__.py | 9 +- linopy/common.py | 12 +- linopy/constraints.py | 2336 ++++++---------------- linopy/csr.py | 35 +- linopy/expressions.py | 24 +- linopy/io.py | 44 +- linopy/matrices.py | 96 +- linopy/model.py | 152 +- linopy/persistent/diff.py | 7 +- linopy/persistent/snapshot.py | 27 +- linopy/scaling.py | 5 - linopy/store.py | 27 +- linopy/testing.py | 10 +- linopy/types.py | 7 +- linopy/variables.py | 14 +- test/test_assign_coords.py | 14 +- test/test_common.py | 2 +- test/test_constraint.py | 310 ++- test/test_csr.py | 579 +++--- test/test_indicator_constraints.py | 56 +- test/test_io.py | 151 +- test/test_matrices.py | 17 +- test/test_model.py | 11 +- test/test_optimization.py | 2 +- test/test_persistent_snapshot_buffers.py | 52 +- test/test_scalar_constraint.py | 18 +- test/test_scaling.py | 27 +- test/test_solvers.py | 16 +- test/test_store_conformance.py | 40 +- 34 files changed, 1457 insertions(+), 2799 deletions(-) diff --git a/doc/api.rst b/doc/api.rst index db873545d..bd207cef2 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -358,8 +358,10 @@ Post-solve access Constraint ========== -Subclass of ``xarray.DataArray`` carrying labels for a multi-dimensional -constraint. +A multi-dimensional constraint ``lhs sign rhs`` with labels. Its lhs holds a +dense store, or in a sparse model (``Model(sparse=True)``) a memory-efficient +CSR store; ``to_sparse`` and ``to_dense`` convert it. See the +:doc:`creating-constraints` and :doc:`sparse-models` guides for usage. .. autosummary:: :toctree: generated/ @@ -378,6 +380,8 @@ Structure constraints.Constraint.scaling constraints.Constraint.coeffs constraints.Constraint.vars + constraints.Constraint.ncons + constraints.Constraint.nterm Modification ------------ @@ -407,52 +411,8 @@ Conversion :toctree: generated/ constraints.Constraint.to_polars - - -CSRConstraint -============= - -Memory-efficient, immutable constraint representation backed by a scipy -CSR sparse matrix. Opt in via ``Model(sparse=True)`` or -``Model.add_constraints(..., freeze=True)``. See the -:doc:`creating-constraints` and :doc:`sparse-models` guides for usage. - -.. autosummary:: - :toctree: generated/ - - constraints.CSRConstraint - -Structure ---------- - -.. autosummary:: - :toctree: generated/ - - constraints.CSRConstraint.coeffs - constraints.CSRConstraint.vars - constraints.CSRConstraint.sign - constraints.CSRConstraint.rhs - constraints.CSRConstraint.scaling - constraints.CSRConstraint.ncons - constraints.CSRConstraint.nterm - -Post-solve access ------------------ - -.. autosummary:: - :toctree: generated/ - - constraints.CSRConstraint.dual - -Conversion ----------- - -.. autosummary:: - :toctree: generated/ - - constraints.CSRConstraint.to_polars - constraints.CSRConstraint.to_dense - constraints.CSRConstraint.mutable + constraints.Constraint.to_sparse + constraints.Constraint.to_dense Constraints @@ -694,8 +654,9 @@ Utilities - ``warn_on_densify``: the conversion notice, a ``PerformanceWarning`` when a sparse (CSR) backing is converted to dense or a dense operand is converted to sparse for a merge with a sparse one. ``None`` (default) - warns in a sparse model on implicit conversions only, ``True`` always - warns, ``False`` never warns. See :doc:`sparse-models`. + warns in a sparse model, ``True`` always warns, ``False`` never warns. + Only implicit conversions warn; ``Constraint.to_dense()`` and + ``Constraint.to_sparse()`` never do. See :doc:`sparse-models`. Warnings diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 169b5d00b..8a9e3faad 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -51,6 +51,11 @@ Version 1.0.0 * The ``Constraint`` constructor takes the fields instead of a dataset: ``Constraint(lhs, sign, rhs, model, name="", labels=None, scaling=1.0, dual=None, binary_var=None, binary_val=None, blocks=None, attrs=None)``; an lhs with a nonzero constant raises. Build a constraint from its dataset layout with ``Constraint.from_dataset(ds, model, name)``; the ``skip_broadcast`` argument is removed. A field with dimensions outside the lhs or with coordinates differing from those of the lhs raises, instead of broadcasting the lhs or filling the field with NaN. (`#1046 `__) * ``Constraint.assign`` and the other wrapped xarray methods reject data variables other than the constraint fields, and ``Constraint.broadcast_like`` no longer broadcasts the fields over ``_term``. After ``Constraint.update(coeffs=...)`` or ``update(variables=...)``, ``coeffs`` and ``vars`` keep the ``_term`` dimension, a value without it is spread over the terms. (`#1046 `__) * ``Constraints.set_blocks`` stores the block numbers as the ``blocks`` field of each constraint, read as ``con.data.blocks``. (`#1046 `__) +* ``CSRConstraint``, ``ConstraintBase`` and the freeze API are removed without aliases: ``Model.add_constraints(freeze=...)``, ``Constraints.add(freeze=...)``, ``Constraint.freeze()``, ``CSRConstraint.mutable()``, ``CSRConstraint.from_dense``/``from_csr`` and the "frozen constraint" ``AttributeError``. A constraint in a sparse model is a ``Constraint`` whose lhs holds a ``CSRStore``; replace ``isinstance(c, CSRConstraint)`` with ``c.lhs.is_sparse``. It supports the whole ``Constraint`` API, e.g. ``update``, the ``scaling`` setter, ``sanitize_*``, ``set_blocks`` and the wrapped xarray methods; ``update`` and ``soften`` keep a sparse lhs sparse (`#1033 `__). Reading ``coeffs`` or ``vars`` of a sparse constraint, ``update(coeffs=..., variables=...)`` and ``set_blocks`` work on a dense copy with a densify notice. (`#1047 `__) +* In a sparse model, ``add_constraints`` and ``add_indicator_constraints`` hold the lhs on a ``CSRStore`` whenever the grid can hold its coordinates, also for a dense lhs, and ``Constraint.from_dataset`` reads into the model's store class. ``Constraint.to_sparse()`` and ``Constraint.to_dense()`` replace the lhs store in place. Explicit conversions emit no notice at all, so ``warn_on_densify=True`` no longer flags them; the ``explicit`` argument of the conversion notice is removed. (`#1047 `__) +* A sparse constraint keeps ``rhs``, ``sign``, ``scaling`` and ``dual`` on its full grid like a dense one, instead of on its active rows; masked rows have label ``-1`` and rows without terms stay labelled but inactive. A masked row with an invalid infinite rhs now raises in a sparse model too. ``Constraint.active_labels``, ``to_matrix`` and ``to_polars`` raise ``ValueError`` on a constraint without labels. A sparse lhs reads its coefficients back as ``float64``. (`#1047 `__) +* ``Constraint.sanitize_zeros`` drops near-zero terms from the lhs store (``Store.prune`` takes a tolerance), so the ``_term`` dimension of a dense constraint can shrink, instead of marking the terms ``-1``/NaN. ``Constraint.iterate_slices`` gathers row blocks with ``Store.take`` and keeps a sparse lhs sparse. ``MatrixAccessor`` and the dual write-back of ``Model.assign_result`` treat both stores alike. (`#1047 `__) +* Netcdf: a sparse constraint is still written in the CSR layout (``_linopy_format="csr"``), now over every cell of its grid with ``labels``, ``rhs``, ``scaling``, the sign and ``dual`` per cell; files written with the active-row layout of ``CSRConstraint`` still load, into a ``Constraint`` with a sparse lhs. (`#1047 `__) **Deprecations** diff --git a/doc/user-guide.rst b/doc/user-guide.rst index 6720507fa..dcb09dd4a 100644 --- a/doc/user-guide.rst +++ b/doc/user-guide.rst @@ -25,8 +25,7 @@ you're unsure what a particular operator or argument does. quadratic) expressions; arithmetic, broadcasting, ``sum``, ``groupby``, ``rolling``, ``where``. - :doc:`creating-constraints` — turning expressions into ``≤`` / ``≥`` - / ``==`` constraints, and the ``CSRConstraint`` memory-efficient - alternative. + / ``==`` constraints, and their memory-efficient sparse (CSR) storage. - :doc:`coordinate-alignment` — how linopy lines up operands that live on different coordinates, and how to control it with ``join``. @@ -78,8 +77,8 @@ Sparse models Models that aggregate many entities into few rows, such as generators into buses, can store their expressions as sparse matrices instead of padded dense arrays. Opt in with ``Model(sparse=True)``. ``groupby(...).sum()`` and ``@`` then return sparse-backed -expressions, and by default every constraint is stored as a frozen -``CSRConstraint``. Operations without a sparse path convert an +expressions, and every constraint holds its left-hand side in CSR form. +Operations without a sparse path convert an expression to dense and emit a ``PerformanceWarning``. .. code-block:: python diff --git a/examples/creating-constraints.ipynb b/examples/creating-constraints.ipynb index 52f894e0b..73700a110 100644 --- a/examples/creating-constraints.ipynb +++ b/examples/creating-constraints.ipynb @@ -249,11 +249,9 @@ "source": [ "## CSR Backend (Advanced)\n", "\n", - "By default, linopy stores each constraint as an `xarray.Dataset` (`Constraint`). This is flexible and allows full label-based indexing, but can use significant memory when constraints have many terms.\n", + "By default, linopy stores the left-hand side of each constraint as a dense `xarray.Dataset` padded along the `_term` dimension. This is flexible, but can use significant memory when constraints have many terms of uneven length.\n", "\n", - "For large models, linopy provides an alternative **CSR backend** via the `CSRConstraint` class. It stores the constraint coefficients as a [scipy CSR sparse matrix](https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.csr_array.html) with flat numpy arrays for the right-hand side and signs. This can reduce memory usage by up to **90%** and speeds up matrix generation for direct solver APIs by **30–120x**.\n", - "\n", - "`CSRConstraint` is **immutable** — once frozen, the constraint data cannot be modified in place. You can always convert back to the mutable xarray-backed `Constraint` if needed." + "For large models, the left-hand side can instead be held as a [scipy CSR sparse matrix](https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.csr_array.html). The constraint stays a `Constraint` with the same API; only the storage of its lhs differs. This can reduce memory usage by up to **90%** and speeds up matrix generation for direct solver APIs by **30–120x**." ] }, { @@ -261,9 +259,9 @@ "id": "25", "metadata": {}, "source": [ - "### Freezing individual constraints\n", + "### Converting individual constraints\n", "\n", - "Pass `freeze=True` to `add_constraints` to store a single constraint in CSR format:" + "`Constraint.to_sparse()` holds the lhs of a constraint in CSR form, in place; `lhs.is_sparse` reports the storage:" ] }, { @@ -280,10 +278,10 @@ "m2 = Model()\n", "y = m2.add_variables(coords=[np.arange(100)], name=\"y\")\n", "\n", - "m2.add_constraints(y <= 10, name=\"upper\", freeze=True)\n", + "con = m2.add_constraints(y <= 10, name=\"upper\").to_sparse()\n", "\n", - "print(type(m2.constraints[\"upper\"]))\n", - "m2.constraints[\"upper\"]" + "print(con.lhs.is_sparse)\n", + "con" ] }, { @@ -293,7 +291,7 @@ "source": [ "### Building a sparse model\n", "\n", - "Set `sparse=True` on the `Model` to build it on the sparse path: every constraint added via `add_constraints` is frozen, and `groupby(...).sum()` and `@` return CSR-backed expressions. Pass `freeze=False` to keep a single constraint mutable. The setting is fixed at creation:" + "Set `sparse=True` on the `Model` to build it on the sparse path: every constraint added via `add_constraints` holds its lhs in CSR form, and `groupby(...).sum()` and `@` return CSR-backed expressions. The setting is fixed at creation:" ] }, { @@ -308,8 +306,8 @@ "m3.add_constraints(z >= 0, name=\"lower\")\n", "m3.add_constraints(z <= 100, name=\"upper\")\n", "\n", - "print(type(m3.constraints[\"lower\"]))\n", - "print(type(m3.constraints[\"upper\"]))" + "print(m3.constraints[\"lower\"].lhs.is_sparse)\n", + "print(m3.constraints[\"upper\"].lhs.is_sparse)" ] }, { @@ -319,7 +317,7 @@ "source": [ "### Converting between representations\n", "\n", - "Use `.freeze()` and `.mutable()` to convert between the two representations. The conversion is lossless:" + "Use `.to_dense()` and `.to_sparse()` to convert the lhs between the two representations, in place. The conversion is lossless:" ] }, { @@ -329,14 +327,14 @@ "metadata": {}, "outputs": [], "source": [ - "frozen = m3.constraints[\"lower\"]\n", - "print(f\"Frozen type: {type(frozen).__name__}\")\n", + "con = m3.constraints[\"lower\"]\n", + "print(f\"Sparse: {con.lhs.is_sparse}\")\n", "\n", - "thawed = frozen.mutable()\n", - "print(f\"Mutable type: {type(thawed).__name__}\")\n", + "con.to_dense()\n", + "print(f\"Sparse after to_dense(): {con.lhs.is_sparse}\")\n", "\n", - "refrozen = thawed.freeze()\n", - "print(f\"Re-frozen type: {type(refrozen).__name__}\")" + "con.to_sparse()\n", + "print(f\"Sparse after to_sparse(): {con.lhs.is_sparse}\")" ] }, { @@ -344,20 +342,14 @@ "id": "31", "metadata": {}, "source": [ - "### API differences from `Constraint`\n", - "\n", - "`CSRConstraint` deliberately exposes a narrower API than the xarray-backed `Constraint`:\n", - "\n", - "- **No in-place mutation.** `Constraint.update(...)` is only available on `Constraint`. (The legacy setters — `con.coeffs = ...`, `con.vars = ...`, `con.sign = ...`, `con.rhs = ...`, `con.lhs = ...` — still forward to `update` on `Constraint` but emit a `DeprecationWarning` and will be removed in a future release.)\n", - "- **No label-based indexing.** `con.loc[...]` is only available on `Constraint`.\n", - "- **Accessing `.coeffs` / `.vars` triggers reconstruction.** On a `CSRConstraint` these properties rebuild the full xarray `Dataset` on demand and emit a `PerformanceWarning`. For solver-oriented workflows prefer `con.to_matrix()` or work with the CSR data directly.\n", + "### Working with a sparse constraint\n", "\n", - "If you need any of the above, call `.mutable()` first to get a `Constraint`:\n", + "A constraint with a sparse lhs supports the full `Constraint` API, e.g. `update(...)`, `loc[...]` and the xarray methods. `update(lhs=...)` and `soften(...)` keep the lhs sparse. Reading `.coeffs` / `.vars` or calling `update(coeffs=..., variables=...)` works on a dense copy and, in a sparse model, emits a `PerformanceWarning`; for solver-oriented workflows prefer `con.to_matrix()`:\n", "\n", "```python\n", - "con = m.constraints[\"my_constraint\"].mutable()\n", - "con.loc[{\"time\": 0}] # label-based indexing now available\n", - "con.update(rhs=5) # mutation now available\n", + "con = m.constraints[\"my_constraint\"]\n", + "con.loc[{\"time\": 0}] # label-based indexing\n", + "con.update(rhs=5) # mutation, the lhs stays sparse\n", "```" ] }, diff --git a/examples/sparse-models.ipynb b/examples/sparse-models.ipynb index 9ef2e2548..c0a136d5b 100644 --- a/examples/sparse-models.ipynb +++ b/examples/sparse-models.ipynb @@ -16,8 +16,8 @@ "A **sparse model** stores these expressions as a CSR matrix (compressed sparse\n", "rows, stores only nonzeros) instead: `A @ x + c` over the coordinate grid. The\n", "expression keeps its public type, `LinearExpression`, and behaves the same. Only\n", - "the backing differs, similar to a dask-backed xarray object. Constraints are\n", - "stored frozen, as `CSRConstraint`, and export reads the CSR data directly." + "the backing differs, similar to a dask-backed xarray object. Constraints hold\n", + "their left-hand side the same way, and export reads the CSR data directly." ] }, { @@ -31,7 +31,7 @@ "turns on the whole sparse path. In a sparse model\n", "`groupby(...).sum()`, `@`/`dot` against a constant and linear arithmetic on\n", "variables such as `2 * x` or `x + y` return CSR-backed expressions, and\n", - "`add_constraints` freezes every constraint unless you pass `freeze=False`." + "`add_constraints` holds the left-hand side of every constraint in CSR form." ] }, { @@ -101,9 +101,10 @@ "id": "6", "metadata": {}, "source": [ - "Constraints in a sparse model are frozen `CSRConstraint` objects, also when\n", - "the expression is dense. A frozen constraint is read-only: call `.mutable()` to\n", - "get a modifiable `Constraint`." + "Constraints in a sparse model are ordinary `Constraint` objects whose\n", + "left-hand side is CSR-backed, also when the expression is dense.\n", + "`con.lhs.is_sparse` reports it. They support the whole `Constraint` API, and\n", + "`update(...)` and `soften(...)` keep the left-hand side sparse." ] }, { @@ -115,7 +116,7 @@ "source": [ "con = m.add_constraints(balance == demand, name=\"balance\")\n", "m.add_objective((p * xr.DataArray(np.arange(6.0), coords=[gens])).sum())\n", - "type(con).__name__" + "con.lhs.is_sparse" ] }, { @@ -145,7 +146,7 @@ "| aggregation | `sum`, `groupby(...).sum()`, `@` / `dot` against a constant |\n", "| arithmetic | `+`, `-` and `linopy.merge` with expressions, variables or constants, also over different dimensions, which broadcast as on the dense path; negation; multiplication and division by a scalar or by a `DataArray` over existing dimensions; adding a constant; `fillna` with a scalar |\n", "| concatenation | `linopy.merge(..., dim=...)` along an existing dimension |\n", - "| constraints | comparison to a scalar or `DataArray` gives a `CSRConstraint` |" + "| constraints | comparison to a scalar or `DataArray` gives a `Constraint` with a CSR-backed lhs |" ] }, { @@ -235,12 +236,13 @@ "\n", "| Value | Behaviour |\n", "|---|---|\n", - "| `None` (default) | warn in a sparse model, on implicit conversions only |\n", - "| `True` | warn on every conversion, also in a dense model and on explicit ones |\n", + "| `None` (default) | warn in a sparse model |\n", + "| `True` | warn on every conversion, also in a dense model |\n", "| `False` | never warn |\n", "\n", - "Explicit conversions are `CSRConstraint.mutable()`, `CSRConstraint.to_dense()`\n", - "and `add_constraints(..., freeze=False)`. They stay quiet under `None`.\n", + "Only implicit conversions warn. The explicit ones, `Constraint.to_dense()` and\n", + "`Constraint.to_sparse()`, replace the store of a constraint's left-hand side\n", + "in place and never warn.\n", "\n", "The option also governs the opposite conversion: a dense operand merged with\n", "a sparse one, e.g. `sparse_expr + dense_expr`, is converted to sparse with one\n", @@ -252,7 +254,7 @@ "`warnings.simplefilter(\"error\", linopy.PerformanceWarning)`; the traceback\n", "then points to it.\n", "\n", - "An explicit conversion warns only under `True`:" + "An explicit conversion stays quiet, also under `True`:" ] }, { @@ -262,10 +264,9 @@ "metadata": {}, "outputs": [], "source": [ - "m.constraints[\"balance\"].mutable()\n", "with linopy.options as opts:\n", " opts[\"warn_on_densify\"] = True\n", - " m.constraints[\"balance\"].mutable()" + " m.constraints[\"balance\"].to_dense().to_sparse().lhs.is_sparse" ] }, { diff --git a/linopy/__init__.py b/linopy/__init__.py index 9ed6112ef..2122b52d2 100644 --- a/linopy/__init__.py +++ b/linopy/__init__.py @@ -21,12 +21,7 @@ EvolvingAPIWarning, PerformanceWarning, ) -from linopy.constraints import ( - Constraint, - ConstraintBase, - Constraints, - CSRConstraint, -) +from linopy.constraints import Constraint, Constraints from linopy.expressions import LinearExpression, QuadraticExpression, merge from linopy.io import read_netcdf from linopy.model import Model, Variable, Variables @@ -49,8 +44,6 @@ __all__ = ( "ABSENT", - "CSRConstraint", - "ConstraintBase", "Constraints", "Constraint", "EQUAL", diff --git a/linopy/common.py b/linopy/common.py index 844361ffd..bc4176c61 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -39,7 +39,7 @@ ) if TYPE_CHECKING: - from linopy.constraints import ConstraintBase + from linopy.constraints import Constraint from linopy.expressions import LinearExpression, QuadraticExpression from linopy.variables import Variable @@ -312,7 +312,7 @@ def validate_coords_reassignment( Single-sourced per-item invariants for coordinate reassignment, used by the Dataset-backed (:func:`reassign_coords`) and the - CSR-backed (``CSRConstraint``) paths alike, so every entry point is safe + CSR-backed (``CSRStore``) paths alike, so every entry point is safe on its own. Model-level aggregate checks (dimension exists somewhere in the model, consistent lengths across containers) live on top of this. @@ -390,7 +390,7 @@ def reassign_coords(ds: Dataset, **coords: Any) -> Dataset: return Dataset(data_vars, coords=ordered, attrs=new.attrs) -T = TypeVar("T", Dataset, "Variable", "LinearExpression", "ConstraintBase") +T = TypeVar("T", Dataset, DataArray, "Variable", "LinearExpression", "Constraint") def iterate_slices( @@ -459,7 +459,7 @@ def iterate_slices( start = i * chunk_size end = min(start + chunk_size, size_of_leading_dim) slice_dict = {leading_dim: slice(start, end)} - yield ds.isel(slice_dict) # type: ignore[attr-defined] + yield ds.isel(slice_dict) def _remap(array: np.ndarray, mapping: np.ndarray) -> np.ndarray: @@ -1079,7 +1079,7 @@ def check_common_keys_values(list_of_dicts: list[dict[str, Any]]) -> bool: "Variable", "LinearExpression", "QuadraticExpression", - "ConstraintBase", + "Constraint", ) @@ -1097,7 +1097,7 @@ def __getitem__( # expand the indexer so we can handle Ellipsis labels = indexing.expanded_indexer(key, self.object.ndim) key = dict(zip(self.object.dims, labels)) - return self.object.sel(key) # type: ignore[attr-defined] + return self.object.sel(key) def coords_to_dataset_vars(coords: list[pd.Index]) -> dict[str, DataArray]: diff --git a/linopy/constraints.py b/linopy/constraints.py index 5ee492494..991abe405 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -7,9 +7,6 @@ from __future__ import annotations import functools -import warnings -import weakref -from abc import ABC, abstractmethod from collections.abc import ( Callable, Generator, @@ -23,7 +20,6 @@ TYPE_CHECKING, Any, NamedTuple, - NoReturn, overload, ) @@ -60,14 +56,11 @@ has_optimized_model, iterate_slices, maybe_group_terms_polars, - maybe_replace_sign, maybe_replace_signs, reassign_coords, replace_by_map, save_join, to_dataframe, - to_polars, - validate_coords_reassignment, ) from linopy.config import options from linopy.constants import ( @@ -76,22 +69,13 @@ HELPER_DIMS, LESS_EQUAL, TERM_DIM, - PerformanceWarning, SIGNS_pretty, ) from linopy.containers import LabelContainer -from linopy.csr import ( - Grid, - _csr_from_label_columns, - _densify_notice, - _positional_csr, - canonical, - csr_nterm, - csr_to_term_arrays, -) +from linopy.csr import CSRStore, Grid, _densify_notice from linopy.scaling import validate_scaling from linopy.semantics import check_user_nan -from linopy.store import DenseStore, Store, _take_flat +from linopy.store import DenseStore, Reason, Store, as_sparse from linopy.types import ( ConstantLike, ConstraintLike, @@ -101,1556 +85,107 @@ VariableLike, ) -if TYPE_CHECKING: - from linopy.csr import CSRStore - from linopy.model import Model - - -FILL_VALUE = { - "labels": -1, - "rhs": np.nan, - "coeffs": np.nan, - "vars": -1, - "sign": "=", - "scaling": 1.0, -} - - -def conwrap( - method: Callable, *default_args: Any, **new_default_kwargs: Any -) -> Callable: - @functools.wraps(method) - def _conwrap(con: Constraint, *args: Any, **kwargs: Any) -> Constraint: - for k, v in new_default_kwargs.items(): - kwargs.setdefault(k, v) - return con.from_dataset( - method(con.data, *default_args, *args, **kwargs), con.model, con.name - ) - - _conwrap.__doc__ = ( - f"Wrapper for the xarray {method.__qualname__} function for linopy.Constraint" - ) - if new_default_kwargs: - _conwrap.__doc__ += f" with default arguments: {new_default_kwargs}" - - return _conwrap - - -def _con_unwrap(con: ConstraintBase | Dataset) -> Dataset: - return con.data if isinstance(con, ConstraintBase) else con - - -class ConstraintBase(ABC): - """ - Abstract base class for Constraint and CSRConstraint. - - Provides all read-only properties and methods shared by both the frozen - Constraint (CSR-backed) and the standard Constraint (Dataset-backed). - """ - - _fill_value = FILL_VALUE - - @property - @abstractmethod - def data(self) -> Dataset: - """Get the underlying xarray Dataset representation.""" - - @property - @abstractmethod - def model(self) -> Model: - """Get the model reference.""" - - @property - @abstractmethod - def name(self) -> str: - """Get the constraint name.""" - - @property - @abstractmethod - def is_assigned(self) -> bool: - """Whether the constraint has been assigned labels by the model.""" - - @property - @abstractmethod - def labels(self) -> DataArray: - """Get the labels DataArray.""" - - @property - @abstractmethod - def range(self) -> tuple[int, int]: - """Return the label range of the constraint.""" - - @property - @abstractmethod - def coeffs(self) -> DataArray: - """Get the LHS coefficients DataArray.""" - - @property - @abstractmethod - def vars(self) -> DataArray: - """Get the LHS variable labels DataArray.""" - - @property - @abstractmethod - def sign(self) -> DataArray: - """Get the constraint sign DataArray.""" - - @property - @abstractmethod - def rhs(self) -> DataArray: - """Get the RHS DataArray.""" - - @property - @abstractmethod - def scaling(self) -> DataArray: - """Get the row scaling DataArray.""" - - @property - @abstractmethod - def dual(self) -> DataArray: - """Get the dual values DataArray.""" - - @dual.setter - @abstractmethod - def dual(self, value: DataArray) -> None: - """Set the dual values DataArray.""" - - @abstractmethod - def _assign_coords(self, **coords: Any) -> ConstraintBase: - """ - Reassign coordinate values on the constraint, keeping the shape. - - Internal: values-only replacement of existing dimension coordinates, - used by :meth:`linopy.Model.assign_coords`. No relabeling, no - reindexing, no shape change, and the order of the underlying data is - preserved. - """ - - @property - @abstractmethod - def is_indicator(self) -> bool: - """Whether the constraint is an indicator constraint.""" - - @property - @abstractmethod - def binary_var(self) -> DataArray | None: - """Get the indicator binary variable labels, or None.""" - - @property - @abstractmethod - def binary_val(self) -> int | np.ndarray | None: - """Get the indicator triggering value(s), or None.""" - - @property - def data_attrs(self) -> list[str]: - """Data variables that define this constraint's persistent state.""" - base = list(Constraints.dataset_attrs) - return base + ["binary_var", "binary_val"] if self.is_indicator else base - - def has_variable(self, variable: variables.Variable) -> bool: - """ - Check if the constraint references any of the given variable labels. - - Masked variable entries (label -1) are ignored: they are not part of - the model and would otherwise match every empty term slot. - """ - return self.has_labels(assigned_labels(variable.labels)) - - @abstractmethod - def has_labels(self, labels: np.ndarray) -> bool: - """ - Check if the constraint references any of the given variable labels. - - ``labels`` must not contain the ``-1`` sentinel, see - :func:`linopy.common.assigned_labels`. - """ - - @abstractmethod - def sanitize_zeros(self) -> ConstraintBase: - """Remove terms with zero or near-zero coefficients.""" - - @abstractmethod - def sanitize_missings(self) -> ConstraintBase: - """Mask out rows where all variables are missing (-1).""" - - @abstractmethod - def sanitize_infinities(self) -> ConstraintBase: - """Mask out rows with invalid infinite RHS values.""" - - @abstractmethod - def to_polars(self) -> pl.DataFrame: - """Convert constraint to a polars DataFrame.""" - - @abstractmethod - def freeze(self) -> CSRConstraint: - """Return an immutable Constraint (CSR-backed).""" - - @abstractmethod - def mutable(self) -> Constraint: - """Return a mutable Constraint.""" - - @abstractmethod - def to_matrix_with_rhs( - self, label_index: VariableLabelIndex - ) -> tuple[scipy.sparse.csr_array, np.ndarray, np.ndarray, np.ndarray]: - """ - Return (csr, con_labels, b, sense) in one pass. - - Avoids computing the CSR matrix twice when both the matrix and - the RHS/sense vectors are needed. - """ - - @abstractmethod - def active_labels(self) -> np.ndarray: - """Active constraint labels in build order, without building the CSR.""" - - @abstractmethod - def active_row_mask(self) -> np.ndarray: - """Boolean mask over raveled rows selecting active constraint rows.""" - - @property - def slack(self) -> Slack | None: - """ - Slack variable(s) added via :meth:`soften`, or ``None`` if the - constraint has never been softened. - """ - names = _slack_names(self.attrs) - if names is None: - return None - positive, negative = names - return Slack( - positive=self.model.variables[positive], - negative=self.model.variables[negative] if negative else None, - ) - - def soften( - self, - penalty: ConstantLike, - *, - max_violation: ConstantLike | None = None, - name: str | None = None, - ) -> Slack: - """ - Soften a constraint, adding a slack variable and a penalty to the objective function. - - Parameters - ---------- - penalty : constant-like - The penalty that will match the slack variable inside the objective function. Must be bigger than 0. - max_violation: constant-like - The max violation possible that caps the slack (upper bound). If None, the slack will be unbounded. - name: string - The name for the slack variable. If None, it will reuse the constraint name and add a '_slack'. - - Returns - ------- - Slack - Named tuple with the `positive` slack variable, and the `negative` one for equality constraints (`None` - for inequality constraints). - - Notes - ----- - On a frozen CSRConstraint the slack terms are appended to the sparse rows in place, without densifying. - A detached copy (from `.mutable()`, `.sel()`, `.isel()` or `.freeze()`) is not registered in - model.constraints, so soften raises ValueError on it. - - Softening an already-softened constraint raises ValueError instead of stacking a second, redundant slack term - onto the same lhs. - - Examples - -------- - >>> from linopy import Model - >>> import pandas as pd - - >>> m = Model() - >>> investments = pd.Index(["A", "B", "C"], name="investments") - >>> expected_return = pd.Series( - ... [0.08, 0.03, 0.1], index=investments, name="expected_return" - ... ) - >>> w = m.add_variables(lower=0, upper=1, coords=[investments], name="weights") - >>> m.add_objective((expected_return * w).sum(), sense="max") - >>> budget_penalty = 2 - - >>> budget_constraint = m.add_constraints(w.sum() == 1, name="budget") - >>> slack = budget_constraint.soften(penalty=budget_penalty) - """ - if not bool(np.all(np.asarray(penalty) > 0)): - raise ValueError("Penalty is not positive.") - - model = self.model - if model.objective.expression.empty: - raise ValueError( - "Objective must be defined via `model.add_objective` before calling `soften` on constraints." - ) - - if model.constraints.data.get(self.name) is not self: - raise ValueError( - f"Constraint {self.name!r} is not the constraint registered in the model, so " - "`soften` would not affect it (it may be a detached copy from `.mutable()`, " - "`.sel()`, `.isel()`, or `.freeze()`). Call `soften` on `model.constraints[name]` " - "directly." - ) - - existing_slack = self.slack - if existing_slack is not None: - raise ValueError( - f"Constraint {self.name!r} was already softened (existing slack " - f"variable {existing_slack.positive.name!r})" - ) - - sign_values = pd.unique(self.sign.values.ravel()) - if len(sign_values) > 1: - raise NotImplementedError( - "Constraint.soften does not support constraints with mixed signs." - ) - sign = sign_values.item() - - name = name or f"{self.name}_slack" - upper = np.inf if max_violation is None else max_violation - - def add_slack(suffix: str) -> variables.Variable: - return model.add_variables( - lower=0, - upper=upper, - coords=self.coords, - mask=self.mask, - name=f"{name}_{suffix}", - ) - - positive = add_slack("pos") - slack = Slack(positive, add_slack("neg") if sign == EQUAL else None) - self._attach_slack(slack, sign) - - violation = positive if slack.negative is None else positive + slack.negative - direction = 1 if model.sense == "min" else -1 - model.objective += direction * (penalty * violation).sum() - return slack - - @abstractmethod - def _attach_slack(self, slack: Slack, sign: str) -> None: - """Add the slack terms to the lhs and record the slack variable names.""" - - def __getitem__( - self, selector: str | int | slice | list | tuple | dict - ) -> Constraint: - """ - Get selection from the constraint. - Returns a Constraint with the selected data. - """ - ds = self.data - data = Dataset({k: ds[k][selector] for k in ds}, attrs=ds.attrs) - return Constraint.from_dataset(data, self.model, self.name) - - @property - def attrs(self) -> dict[str, Any]: - """Get the attributes of the constraint.""" - return self.data.attrs - - @property - def coords(self) -> DatasetCoordinates: - """Get the coordinates of the constraint.""" - return self.data.coords - - @property - def indexes(self) -> Indexes: - """Get the indexes of the constraint.""" - return self.data.indexes - - @property - def dims(self) -> Frozen[Hashable, int]: - """Get the dimensions of the constraint.""" - return self.data.dims - - @property - def sizes(self) -> Frozen[Hashable, int]: - """Get the sizes of the constraint.""" - return self.data.sizes - - @property - def nterm(self) -> int: - """Get the number of terms in the constraint.""" - return self.data.sizes.get(TERM_DIM, 1) - - @property - def ndim(self) -> int: - """Get the number of dimensions of the constraint.""" - return self.rhs.ndim - - @property - def shape(self) -> tuple[int, ...]: - """Get the shape of the constraint.""" - return self.rhs.shape - - @property - def size(self) -> int: - """Get the size of the constraint.""" - return self.rhs.size - - @property - @abstractmethod - def ncons(self) -> int: - """Get the number of active constraints (non-masked, with at least one valid variable).""" - - @property - def coord_dims(self) -> tuple[Hashable, ...]: - return tuple(k for k in self.dims if k not in HELPER_DIMS) - - @property - def coord_sizes(self) -> dict[Hashable, int]: - return {k: self.sizes[k] for k in self.coord_dims} - - @property - def coord_names(self) -> list[str]: - """Get the names of the coordinates.""" - return get_dims_with_index_levels(self.data, self.coord_dims) - - @property - def type(self) -> str: - """Get the type string of the constraint.""" - return "Constraint" if self.is_assigned else "Constraint (unassigned)" - - @property - def term_dim(self) -> str: - """Return the term dimension of the constraint.""" - return TERM_DIM - - @property - def mask(self) -> DataArray | None: - """ - Get the mask of the constraint. - - The mask indicates on which coordinates the constraint is enabled - (True) and disabled (False). - """ - if self.is_assigned: - result: DataArray = self.labels != FILL_VALUE["labels"] # type: ignore[assignment] - return result.astype(bool) - return None - - @property - @abstractmethod - def lhs(self) -> expressions.LinearExpression: - """Get the left-hand-side linear expression of the constraint.""" - - def __contains__(self, value: Any) -> bool: - return self.data.__contains__(value) - - def __repr__(self) -> str: - """Print the constraint arrays.""" - max_lines = options["display_max_rows"] - dims = list(self.coord_dims) - ndim = len(dims) - dim_names = self.coord_names - dim_sizes = list(self.coord_sizes.values()) - size = np.prod(dim_sizes) - masked_entries = (~self.mask).sum().values if self.mask is not None else 0 - lines = [] - - header_string = f"{self.type} `{self.name}`" if self.name else f"{self.type}" - - if size > 1 or ndim > 0: - row_labels = get_printout_labels(self.data, dims) - for indices in generate_indices_for_printout(dim_sizes, max_lines): - if indices is None: - lines.append("\t\t...") - else: - coord = [row_labels[i][int(ind)] for i, ind in enumerate(indices)] - if self.mask is None or self.mask.values[indices]: - expr = format_single_expression( - self.coeffs.values[indices], - self.vars.values[indices], - 0, - self.model, - ) - sign = SIGNS_pretty[self.sign.values[indices]] - rhs = self.rhs.values[indices] - line = format_coord(coord) + f": {expr} {sign} {rhs}" - else: - line = format_coord(coord) + ": None" - lines.append(line) - lines = align_lines_by_delimiter(lines, list(SIGNS_pretty.values())) - - shape_str = ", ".join(f"{d}: {s}" for d, s in zip(dim_names, dim_sizes)) - mask_str = f" - {masked_entries} masked entries" if masked_entries else "" - underscore = "-" * (len(shape_str) + len(mask_str) + len(header_string) + 4) - lines.insert(0, f"{header_string} [{shape_str}]{mask_str}:\n{underscore}") - elif size == 1: - expr = format_single_expression( - self.coeffs.values, self.vars.values, 0, self.model - ) - lines.append( - f"{header_string}\n{'-' * len(header_string)}\n{expr} {SIGNS_pretty[self.sign.item()]} {self.rhs.item()}" - ) - else: - lines.append(f"{header_string}\n{'-' * len(header_string)}\n") - - return "\n".join(lines) - - def print(self, display_max_rows: int = 20, display_max_terms: int = 20) -> None: - """ - Print the linear expression. - - Parameters - ---------- - display_max_rows : int - Maximum number of rows to be displayed. - display_max_terms : int - Maximum number of terms to be displayed. - """ - with options as opts: - opts.set_value( - display_max_rows=display_max_rows, display_max_terms=display_max_terms - ) - print(self) - - def _flat(self) -> pd.DataFrame: - ds = self.data - - def mask_func(data: dict) -> pd.Series: - mask = (data["vars"] != -1) & (data["coeffs"] != 0) - if "labels" in data: - mask &= data["labels"] != -1 - return mask - - df = to_dataframe(ds, mask_func=mask_func) - - # Group repeated variables in the same constraint - agg_custom = {k: "first" for k in list(df.columns)} - agg_standards = dict(coeffs="sum", rhs="first", sign="first") - agg = {**agg_custom, **agg_standards} - df = df.groupby(["labels", "vars"], as_index=False).aggregate(agg) - check_has_nulls(df, name=f"{self.type} {self.name}") - return df - - @abstractmethod - def to_matrix( - self, label_index: VariableLabelIndex - ) -> tuple[scipy.sparse.csr_array, np.ndarray]: - """ - Construct a CSR matrix for this constraint. - - Only active (non-masked) rows are included. Column indices are dense - positions in the active variable array, as given by ``label_index``. - - Returns - ------- - csr : scipy.sparse.csr_array - Shape (n_active_cons, n_active_vars). - con_labels : np.ndarray - Active constraint labels in row order. - """ - - def to_netcdf_ds(self) -> Dataset: - """Return a Dataset representation suitable for netcdf serialization.""" - return self.data - - iterate_slices = iterate_slices - - -def _equal_nnz_slices( - indptr: np.ndarray, slice_size: int -) -> Generator[slice, None, None]: - """Yield row slices such that each slice contains at most slice_size non-zeros.""" - n_rows = len(indptr) - 1 - start = 0 - while start < n_rows: - offset = np.searchsorted( - indptr[start + 1 :], indptr[start] + slice_size, side="left" - ) - end = min(start + 1 + offset, n_rows) - yield slice(start, end) - start = end - - -_UNSUPPORTED = "is not supported on a frozen constraint" - - -def _is_uniform(values: np.ndarray) -> bool: - return values.size > 0 and bool((values == values[0]).all()) - - -def _slack_names(attrs: Mapping[str, Any]) -> tuple[str, str] | None: - """Names of the positive and negative slack variables recorded by ``soften``.""" - positive = attrs.get("slack_positive") - if positive is None: - return None - return str(positive), str(attrs.get("slack_negative", "")) - - -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]"] - - -class CSRConstraint(ConstraintBase): - """ - Frozen constraint backed by a CSR sparse matrix. - - The structure is frozen: rows, coeffs, sign and rhs cannot be mutated - (see the raising setters below). The one exception is ``soften``, which - appends slack terms to the CSR matrix in place. - - Parameters - ---------- - csr : scipy.sparse.csr_array - Shape (n_active_cons, model._xCounter). Rows are the active rows of - the constraint grid, column indices are raw variable labels. They are - mapped to dense positions in ``to_matrix``/``to_matrix_with_rhs``, so - adding or removing variables after freezing stays consistent. - active_positions : np.ndarray - Shape (n_active_cons,). Flat positions of the active rows in the - constraint grid; labels are ``cindex + active_positions``. - rhs : np.ndarray - Shape (n_active_cons,). Right-hand-side values. - 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. - grid : Grid - The coordinate grid the constraint rows live on. - model : Model - The linopy model this constraint belongs to. - name : str - Name of the constraint. - cindex : int or None - Starting label assigned by the model. None if not yet assigned. - dual : np.ndarray or None - Shape (n_flat,). Dual values after solving, or None. - """ - - __slots__ = ( - "_csr", - "_active_positions", - "_rhs", - "_sign", - "_scaling", - "_grid", - "_model", - "_name", - "_cindex", - "_dual", - "_binvar_labels", - "_binval", - "_slack", - "_positional_cache", - ) - - def __init__( - self, - csr: scipy.sparse.csr_array, - active_positions: np.ndarray, - rhs: np.ndarray, - sign: str | np.ndarray, - grid: Grid, - model: Model, - name: str = "", - cindex: int | None = None, - dual: np.ndarray | None = None, - binvar_labels: np.ndarray | None = None, - binval: int | np.ndarray | None = None, - scaling: np.ndarray | None = None, - slack: tuple[str, str] | None = None, - ) -> None: - self._csr = csr - self._active_positions = active_positions - self._rhs = rhs - self._sign = sign - self._scaling = ( - np.asarray(scaling, dtype=float) - if scaling is not None - else np.ones_like(rhs, dtype=float) - ) - self._grid = grid - self._model = model - self._name = name - self._cindex = cindex - self._dual = dual - self._binvar_labels = binvar_labels - self._binval = binval - self._slack = slack - self._positional_cache: _PositionalCache | None = None - - @property - def model(self) -> Model: - return self._model - - @property - def name(self) -> str: - return self._name - - @property - def is_assigned(self) -> bool: - return self._cindex is not None - - @property - def shape(self) -> tuple[int, ...]: - return self._grid.shape - - @property - def full_size(self) -> int: - return self._grid.size - - @property - def range(self) -> tuple[int, int]: - """Return the (start, end) label range of the constraint.""" - if self._cindex is None: - raise AttributeError("Constraint has not been assigned labels yet.") - return (self._cindex, self._cindex + self.full_size) - - @property - def ncons(self) -> int: - return self._csr.shape[0] - - @property - def attrs(self) -> dict[str, Any]: - d: dict[str, Any] = {"name": self._name} - if self._cindex is not None: - d["label_range"] = (self._cindex, self._cindex + self.full_size) - if self._slack is not None: - d["slack_positive"], d["slack_negative"] = self._slack - return d - - @property - def coords(self) -> DatasetCoordinates: - return self._grid.to_dataset().coords - - @property - def dims(self) -> Frozen[Hashable, int]: - d: dict[Hashable, int] = dict(zip(self._grid.dims, self._grid.shape)) - d[TERM_DIM] = self.nterm - return Frozen(d) - - @property - def active_positions(self) -> np.ndarray: - """Flat positions of active (non-masked) rows in the original coord shape.""" - return self._active_positions - - @property - def sizes(self) -> Frozen[Hashable, int]: - return self.dims - - @property - def indexes(self) -> Indexes: - return Dataset(coords=self._grid.indexes).indexes - - @property - def nterm(self) -> int: - return csr_nterm(self._csr) - - @property - def coord_names(self) -> list[str]: - return list(self._grid.dims) - - def __getstate__(self) -> dict[str, Any]: - return self._init_kwargs() - - def __setstate__(self, state: dict[str, Any]) -> None: - self.__init__(**state) # type: ignore[misc] - - def _init_kwargs(self) -> dict[str, Any]: - return dict( - csr=self._csr, - active_positions=self._active_positions, - rhs=self._rhs, - sign=self._sign, - grid=self._grid, - model=self._model, - name=self._name, - cindex=self._cindex, - dual=self._dual, - binvar_labels=self._binvar_labels, - binval=self._binval, - scaling=self._scaling, - slack=self._slack, - ) - - def _replace(self, **changes: Any) -> CSRConstraint: - """Copy with the given constructor arguments replaced.""" - kwargs = self._init_kwargs() - kwargs.update(changes) - new = CSRConstraint(**kwargs) - if kwargs["csr"] is self._csr: - new._positional_cache = self._positional_cache - return new - - def assign_labels( - self, - cindex: int, - name: str, - scaling: float | DataArray = 1.0, - mask: np.ndarray | None = None, - ) -> CSRConstraint: - """ - Return a copy labelled from ``cindex`` and named ``name``. - - Rows without terms are dropped, as when freezing a dense constraint; - a zero coefficient counts as a term. Active rows where the boolean - ``mask`` is False are dropped in the same gather. ``scaling`` is a - scalar or a row scaling broadcast on the grid; its distinct values are - validated without expanding a broadcast view. - """ - values = np.asarray(scaling) - distinct = values[tuple(slice(None) if s else 0 for s in values.strides)] - validate_scaling(distinct, "constraint scaling") - keep = np.diff(self._csr.indptr) > 0 - if mask is not None: - keep &= mask - kept = self._kept(keep) - csr = kept._csr - if not csr.data.all(): - csr = csr.copy() if csr is self._csr else csr - csr.eliminate_zeros() - if isinstance(scaling, DataArray): - row_scaling = kept.active_values(scaling) - else: - row_scaling = np.full(csr.shape[0], float(scaling)) - return kept._replace(csr=csr, cindex=cindex, name=name, scaling=row_scaling) - - def active_values(self, values: DataArray) -> np.ndarray: - """Values of ``values``, broadcast on the grid, at the active rows.""" - return _take_flat( - values.transpose(*self._grid.dims).values, self._active_positions - ) - - def _kept(self, keep: np.ndarray) -> CSRConstraint: - """ - Copy holding only the active rows where ``keep`` is True, sharing the - row arrays when every row is kept. - """ - if keep.all(): - return self._replace() - - 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=rows(self._sign), - scaling=self._scaling[keep], - dual=rows(self._dual), - binvar_labels=rows(self._binvar_labels), - binval=rows(self._binval), - ) - - def _assign_coords(self, **coords: Any) -> CSRConstraint: - """ - Reassign coordinate values on the constraint, keeping the shape. - - Internal: values-only replacement of existing dimension coordinates, - used by :meth:`linopy.Model.assign_coords`. No relabeling, no - reindexing, no shape change, and the order of the underlying data is - preserved. - """ - new_indexes = validate_coords_reassignment( - {dim: len(index) for dim, index in self._grid.indexes.items()}, - coords, - f"constraint '{self.name}'", - ) - self._grid = self._grid.with_indexes(new_indexes) - return self - - def _active_to_dataarray( - self, active_values: np.ndarray, fill: float | int | str = -1 - ) -> DataArray: - full = np.full(self.full_size, fill, dtype=active_values.dtype) - full[self.active_positions] = active_values - return self._grid.dataarray(full) - - @property - def labels(self) -> DataArray: - """Get labels DataArray, shape (*coord_dims).""" - if self._cindex is None: - return DataArray([]) - return self._active_to_dataarray(self.active_labels(), fill=-1) - - @property - def coeffs(self) -> DataArray: - """Get coefficients DataArray, shape (*coord_dims, _term).""" - warnings.warn( - "Accessing .coeffs on a Constraint triggers full Dataset reconstruction. " - "Use .to_matrix() for efficient access.", - PerformanceWarning, - stacklevel=2, - ) - 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).""" - warnings.warn( - "Accessing .vars on a Constraint triggers full Dataset reconstruction. " - "Use .to_matrix() for efficient access.", - PerformanceWarning, - stacklevel=2, - ) - 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 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).""" - return self._active_to_dataarray(self._rhs, fill=np.nan) - - @rhs.setter - def rhs(self, value: ConstantLike) -> None: - raise _frozen_error("rhs") - - @property - def scaling(self) -> DataArray: - """Get row scaling DataArray, shape (*coord_dims).""" - return self._active_to_dataarray(self._scaling, fill=1.0) - - @scaling.setter - def scaling(self, value: ConstantLike) -> None: - raise _frozen_error("scaling") - - @property - def lhs(self) -> expressions.LinearExpression: - """Get LHS as LinearExpression (triggers Dataset reconstruction).""" - ds = self._to_dataset(self.nterm) - return expressions.LinearExpression(ds[["coeffs", "vars"]], self._model) - - @lhs.setter - def lhs(self, value: ExpressionLike | VariableLike | ConstantLike) -> None: - 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 _attach_slack(self, slack: Slack, sign: str) -> None: - slacks = [v for v in slack if v is not None] - sign_coeff = 1.0 if sign == GREATER_EQUAL else -1.0 - coeffs = np.array([sign_coeff, 1.0][: len(slacks)]) - labels = [v.labels.transpose(*self._grid.dims).values.ravel() for v in slacks] - cols = np.stack(labels, axis=1)[self._active_positions].ravel() - n = self.ncons - data = np.tile(coeffs.astype(self._csr.dtype), n) - shape = (n, self._model._xCounter) - indptr = np.arange(n + 1) * len(slacks) - extra = scipy.sparse.csr_array((data, cols, indptr), shape=shape) - widened = scipy.sparse.csr_array( - (self._csr.data, self._csr.indices, self._csr.indptr), shape=shape - ) - self._csr = widened + extra - self._positional_cache = None - self._slack = ( - slack.positive.name, - slack.negative.name if slack.negative is not None else "", - ) - - @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 - @has_optimized_model - def dual(self) -> DataArray: - """Get dual values DataArray, shape (*coord_dims).""" - if self._dual is None: - raise AttributeError( - "Underlying is optimized but does not have dual values stored." - ) - return self._active_to_dataarray(self._dual, fill=np.nan) - - @dual.setter - def dual(self, value: DataArray) -> None: - """Set dual values from a DataArray aligned with the full coord shape.""" - vals = np.asarray(value).ravel() - self._dual = vals[self.active_positions] - - @property - def is_indicator(self) -> bool: - return self._binvar_labels is not None - - @property - def binary_var(self) -> DataArray | None: - if self._binvar_labels is None: - return None - return self._active_to_dataarray(self._binvar_labels, fill=-1) - - @property - def binary_val(self) -> int | np.ndarray | None: - return self._binval - - def _to_dataset(self, nterm: int) -> Dataset: - """ - Reconstruct labels/coeffs/vars Dataset from the CSR matrix. - - Parameters - ---------- - nterm : int - Number of terms per row (width of the dense term block). - - Returns - ------- - Dataset with variables ``labels``, ``coeffs``, ``vars``. - """ - shape = self.shape - full_size = self.full_size - active_positions = self.active_positions - 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 = self._grid.indexes | self._grid.aux - dims_with_term = dim_names + [TERM_DIM] - coeffs_da = DataArray( - coeffs_2d.reshape(shape + (nterm,)), - coords=xr_coords, - dims=dims_with_term, - ) - vars_da = DataArray( - vars_2d.reshape(shape + (nterm,)), - coords=xr_coords, - dims=dims_with_term, - ) - ds = Dataset({"coeffs": coeffs_da, "vars": vars_da}) - if self._cindex is not None: - labels_flat = np.full(full_size, -1, dtype=self._model._dtypes["labels"]) - labels_flat[active_positions] = self.active_labels() - ds = assign_fields( - ds, - labels=self._grid.dataarray(labels_flat), - ) - return ds - - @property - def data(self) -> Dataset: - """Reconstruct the xarray Dataset from the CSR representation.""" - ds = self._to_dataset(self.nterm) - extra: dict[str, Any] = { - "sign": self.sign, - "rhs": self.rhs, - "scaling": self.scaling, - } - if self._dual is not None: - extra["dual"] = self._active_to_dataarray(self._dual, fill=np.nan) - if self._binvar_labels is not None: - extra["binary_var"] = self._active_to_dataarray( - self._binvar_labels, fill=-1 - ) - binval = self._binval - if isinstance(binval, np.ndarray) and binval.ndim > 0: - extra["binary_val"] = self._active_to_dataarray(binval, fill=-1) - else: - extra["binary_val"] = binval - return assign_fields(ds, **extra).assign_attrs(self.attrs) - - def __repr__(self) -> str: - """Print the constraint without reconstructing the full Dataset.""" - max_lines = options["display_max_rows"] - coords = self._grid.coords - shape = self.shape - dim_names = self.coord_names - size = self.full_size - # Map active rows (CSR is active-only) back to full flat positions - csr = self._csr - nterm = self.nterm - active_positions = self.active_positions - masked_entries = size - len(active_positions) - - # pos_to_row[flat_idx] -> CSR row index, or -1 if masked - # active_positions is sorted (labels assigned in order) - pos_to_row = np.full(size, -1, dtype=np.intp) - pos_to_row[active_positions] = np.arange(len(active_positions), dtype=np.intp) - - header_string = f"{self.type} `{self._name}`" if self._name else f"{self.type}" - lines = [] - - def row_expr(row: int) -> str: - start, end = int(csr.indptr[row]), int(csr.indptr[row + 1]) - vars_row = np.full(nterm, -1, dtype=np.int64) - coeffs_row = np.zeros(nterm, dtype=csr.dtype) - vars_row[: end - start] = csr.indices[start:end] - coeffs_row[: end - start] = csr.data[start:end] - sign = self._sign if isinstance(self._sign, str) else self._sign[row] - return f"{format_single_expression(coeffs_row, vars_row, 0, self._model)} {SIGNS_pretty[sign]} {self._rhs[row]}" - - if size > 1: - for indices in generate_indices_for_printout(shape, max_lines): - if indices is None: - lines.append("\t\t...") - else: - coord = [coords[i][int(ind)] for i, ind in enumerate(indices)] - flat_idx = int(np.ravel_multi_index(indices, shape)) - row = pos_to_row[flat_idx] - body = row_expr(row) if row >= 0 else "None" - lines.append(format_coord(coord) + f": {body}") - lines = align_lines_by_delimiter(lines, list(SIGNS_pretty.values())) - - shape_str = ", ".join(f"{d}: {s}" for d, s in zip(dim_names, shape)) - mask_str = f" - {masked_entries} masked entries" if masked_entries else "" - underscore = "-" * (len(shape_str) + len(mask_str) + len(header_string) + 4) - lines.insert(0, f"{header_string} [{shape_str}]{mask_str}:\n{underscore}") - elif size == 1: - body = row_expr(0) if len(active_positions) > 0 else "None" - lines.append(f"{header_string}\n{'-' * len(header_string)}\n{body}") - else: - lines.append(f"{header_string}\n{'-' * len(header_string)}\n") - - return "\n".join(lines) - - def to_matrix( - self, label_index: VariableLabelIndex - ) -> tuple[scipy.sparse.csr_array, np.ndarray]: - """Return the CSR matrix with dense variable positions and con_labels.""" - return self._to_positional_csr(label_index), self.active_labels() - - def _to_positional_csr( - self, label_index: VariableLabelIndex - ) -> scipy.sparse.csr_array: - """ - Return the stored CSR with label columns replaced by dense positions. - - ``indptr`` and ``data`` stay the stored arrays. The positional - ``indices`` are cached by weak reference, keyed on the identity of the - stored CSR and of ``label_index.label_to_pos``, which is rebuilt - whenever variables are added or removed. While a caller such as a - ``ModelSnapshot`` holds the array, repeated calls on an unchanged model - return the same object; otherwise it is freed and rebuilt on demand. - """ - csr = self._csr - label_to_pos = label_index.label_to_pos - if self._positional_cache is not None: - cached_csr, cached_label_to_pos, ref = self._positional_cache - cols = ref() - if ( - cached_csr is csr - and cached_label_to_pos is label_to_pos - and cols is not None - ): - return _positional_csr( - csr.data, cols, csr.indptr, csr.shape[0], label_index.n_active_vars - ) - positional = _csr_from_label_columns( - csr.data, csr.indices, csr.indptr, csr.shape[0], label_index, self.name - ) - self._positional_cache = (csr, label_to_pos, weakref.ref(positional.indices)) - return positional - - def to_netcdf_ds(self) -> Dataset: - """Return a Dataset with raw CSR components for netcdf serialization.""" - csr = self._csr - data_vars: dict[str, DataArray] = { - "indptr": DataArray(csr.indptr, dims=["_indptr"]), - "indices": DataArray(csr.indices, dims=["_nnz"]), - "data": DataArray(csr.data, dims=["_nnz"]), - "rhs": DataArray(self._rhs, dims=["_flat"]), - "scaling": DataArray(self._scaling, dims=["_flat"]), - "_active_positions": DataArray(self._active_positions, dims=["_flat"]), - } - if isinstance(self._sign, np.ndarray): - data_vars["_sign"] = DataArray(self._sign, dims=["_flat"]) - 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) - attrs: dict[str, Any] = { - "_linopy_format": "csr", - "_csr_columns": "labels", - "cindex": self._cindex if self._cindex is not None else -1, - "shape": list(csr.shape), - "coord_dims": dim_names, - "name": self._name, - } - if isinstance(self._sign, str): - attrs["sign"] = self._sign - if self._slack is not None: - attrs["slack_positive"], attrs["slack_negative"] = self._slack - if self._binvar_labels is not None: - attrs["is_indicator"] = True - data_vars["_binvar_labels"] = DataArray(self._binvar_labels, dims=["_flat"]) - binval = self._binval - if isinstance(binval, np.ndarray) and binval.ndim > 0: - data_vars["_binval"] = DataArray(binval, dims=["_flat"]) - else: - attrs["binval"] = int(binval) # type: ignore[arg-type] - return Dataset(data_vars, attrs=attrs) - - @classmethod - 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 hold dense variable - positions instead of labels. - """ - attrs = ds.attrs - shape = tuple(attrs["shape"]) - indices = ds["indices"].values - columns = attrs.get("_csr_columns", "positions") - if columns == "positions": - vlabels = model.variables.label_index.vlabels - indices = vlabels[indices] - shape = (shape[0], int(vlabels[-1]) + 1 if len(vlabels) else 1) - elif columns != "labels": - raise ValueError( - f"Unknown CSR column format '{columns}' for constraint '{name}'." - ) - csr = scipy.sparse.csr_array( - (ds["data"].values, indices, ds["indptr"].values), - shape=shape, - ) - rhs = ds["rhs"].values - scaling = ds["scaling"].values if "scaling" in ds else None - sign: str | np.ndarray = ds["_sign"].values if "_sign" in ds else attrs["sign"] - _cindex_raw = int(attrs["cindex"]) - cindex: int | None = _cindex_raw if _cindex_raw >= 0 else None - coord_dims = attrs["coord_dims"] - if isinstance(coord_dims, str): - coord_dims = [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 - elif "_con_labels" in ds: - active_positions = ds["_con_labels"].values - (cindex or 0) - else: - active_positions = np.arange(len(rhs), dtype=np.intp) - binvar_labels: np.ndarray | None = None - binval: int | np.ndarray | None = None - if "_binvar_labels" in ds: - binvar_labels = ds["_binvar_labels"].values - binval = ds["_binval"].values if "_binval" in ds else attrs["binval"] - slack = _slack_names(attrs) - return cls( - csr, - active_positions, - rhs, - sign, - grid, - model, - name, - cindex=cindex, - dual=dual, - binvar_labels=binvar_labels, - binval=binval, - scaling=scaling, - slack=slack, - ) - - def has_labels(self, labels: np.ndarray) -> bool: - return contains_labels(self._csr.indices, labels) - - def to_matrix_with_rhs( - self, label_index: VariableLabelIndex - ) -> tuple[scipy.sparse.csr_array, np.ndarray, np.ndarray, np.ndarray]: - """Return (csr, con_labels, b, sense); only the columns are recomputed.""" - if isinstance(self._sign, str): - sense = np.full(len(self._rhs), self._sign[0]) - else: - sense = self._sign.astype("U1") - return ( - self._to_positional_csr(label_index), - self.active_labels(), - self._rhs, - sense, - ) - - def active_labels(self) -> np.ndarray: - if self._cindex is None: - raise ValueError(f"Constraint {self._name!r} has not been assigned labels.") - return self._active_positions + self._cindex - - def active_row_mask(self) -> np.ndarray: - return np.ones(self._csr.shape[0], dtype=bool) - - def sanitize_zeros(self) -> CSRConstraint: - """ - Remove terms with zero or near-zero coefficients. - - Copy-on-write: rebinds ``_csr`` instead of mutating its arrays, so - external holders of the previous arrays (e.g. a ModelSnapshot - sharing them) keep a valid baseline. - """ - data = self._csr.data - zeros = data <= 1e-10 - zeros &= data >= -1e-10 - if zeros.any(): - csr = self._csr.copy() - csr.data[zeros] = 0 - csr.eliminate_zeros() - self._csr = csr - self._positional_cache = None - return self - - def sanitize_missings(self) -> CSRConstraint: - """No-op: missing rows are already excluded during freezing.""" - return self - - def sanitize_infinities(self) -> CSRConstraint: - """Mask out rows with invalid infinite RHS values (mutates in-place).""" - if isinstance(self._sign, str): - if self._sign == LESS_EQUAL: - invalid = self._rhs == np.inf - elif self._sign == GREATER_EQUAL: - invalid = self._rhs == -np.inf - else: - return self - else: - invalid = ((self._sign == LESS_EQUAL) & (self._rhs == np.inf)) | ( - (self._sign == GREATER_EQUAL) & (self._rhs == -np.inf) - ) - if not invalid.any(): - return self - keep = ~invalid - self._csr = self._csr[keep] - self._positional_cache = None - self._active_positions = self._active_positions[keep] - self._rhs = self._rhs[keep] - self._scaling = self._scaling[keep] - if not isinstance(self._sign, str): - self._sign = self._sign[keep] - return self - - def freeze(self) -> CSRConstraint: - """Return self (structure is already frozen).""" - return self - - def to_dense(self) -> Constraint: - """Convert to a Constraint.""" - _densify_notice( - "frozen constraint converted by `to_dense()`/`mutable()`", - self._model, - explicit=True, - ) - return Constraint.from_dataset(self.data, self._model, self._name) +if TYPE_CHECKING: + from linopy.model import Model - def mutable(self) -> Constraint: - """Return a modifiable Constraint (same as to_dense).""" - return self.to_dense() - def to_polars(self) -> pl.DataFrame: - """Convert frozen constraint to polars DataFrame directly from CSR.""" - csr = self._csr - sign_dtype = pl.Enum(["=", "<=", ">="]) - if csr.nnz == 0: - return pl.DataFrame( - schema={ - "labels": pl.Int64, - "coeffs": pl.Float64, - "vars": pl.Int64, - "sign": sign_dtype, - "rhs": pl.Float64, - } - ) +FILL_VALUE = { + "labels": -1, + "rhs": np.nan, + "coeffs": np.nan, + "vars": -1, + "sign": "=", + "scaling": 1.0, +} - rows = np.repeat(np.arange(csr.shape[0]), np.diff(csr.indptr)) - data: dict[str, Any] = { - "labels": self.active_labels()[rows], - "coeffs": csr.data, - "vars": csr.indices, - "rhs": self._rhs[rows], - } - sign_expr: pl.Expr | pl.Series = ( - pl.lit(self._sign, dtype=sign_dtype) - if isinstance(self._sign, str) - else pl.Series("sign", self._sign[rows], dtype=sign_dtype) +def conwrap( + method: Callable, *default_args: Any, **new_default_kwargs: Any +) -> Callable: + @functools.wraps(method) + def _conwrap(con: Constraint, *args: Any, **kwargs: Any) -> Constraint: + for k, v in new_default_kwargs.items(): + kwargs.setdefault(k, v) + return con.from_dataset( + method(con.data, *default_args, *args, **kwargs), con.model, con.name ) - df = pl.DataFrame(data).with_columns(sign=sign_expr) - return df[["labels", "coeffs", "vars", "sign", "rhs"]] - def iterate_slices( - self, - slice_size: int | None = 2_000_000, - slice_dims: list | None = None, - ) -> Generator[CSRConstraint, None, None]: - """ - Yield row-batched sub-Constraints without Dataset reconstruction. - - Batches are raw CSR slices suitable only for ``to_polars()``. They are - 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 - if slice_size is None or nnz <= slice_size: - yield self - return - - for rows in _equal_nnz_slices(self._csr.indptr, slice_size): - sign = self._sign if isinstance(self._sign, str) else self._sign[rows] - yield CSRConstraint( - csr=self._csr[rows], - active_positions=self._active_positions[rows], - rhs=self._rhs[rows], - sign=sign, - scaling=self._scaling[rows], - grid=Grid({}), - model=self._model, - name=self._name, - cindex=self._cindex, - ) + _conwrap.__doc__ = ( + f"Wrapper for the xarray {method.__qualname__} function for linopy.Constraint" + ) + if new_default_kwargs: + _conwrap.__doc__ += f" with default arguments: {new_default_kwargs}" - @classmethod - def from_dense( - cls, - con: Constraint, - cindex: int | None = None, - ) -> CSRConstraint: - """ - Create a CSRConstraint from a Constraint. + return _conwrap - Parameters - ---------- - con : Constraint - cindex : int or None - Starting label index; defaults to the one recorded on ``con``. - """ - if cindex is None and "label_range" in con.attrs: - cindex = int(con.attrs["label_range"][0]) - active_mask = con.active_row_mask() - store = con.lhs._store.to_dense() - data, vlabel_cols, indptr = store.active_terms(active_mask) - active_positions = np.flatnonzero(active_mask) - csr = scipy.sparse.csr_array( - (data, vlabel_cols, indptr), - shape=(len(active_positions), con.model._xCounter), - ) - if con.nterm > 1: - csr.sum_duplicates() - csr.eliminate_zeros() - 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() - active_signs = sign_vals[active_mask] - sign: str | np.ndarray - if not active_signs.size: - sign = str(sign_vals[0]) if _is_uniform(sign_vals) else "=" - elif _is_uniform(active_signs): - sign = str(active_signs[0]) - else: - sign = active_signs - dual = ( - con.data["dual"].values.ravel()[active_mask] if "dual" in con.data else None - ) - binvar_labels: np.ndarray | None = None - binval: int | np.ndarray | None = None - binary_var = con.binary_var - if binary_var is not None: - binvar_labels = binary_var.values.ravel()[active_mask] - bv = con.binary_val - if isinstance(bv, np.ndarray) and bv.ndim > 0: - binval = bv.ravel()[active_mask] - else: - binval = bv - return cls( - csr, - active_positions, - rhs, - sign, - grid, - con.model, - con.name, - cindex=cindex, - dual=dual, - binvar_labels=binvar_labels, - binval=binval, - scaling=scaling, - slack=_slack_names(con.attrs), - ) - @classmethod - def from_csr(cls, expr: CSRStore, sign: str, rhs: DataArray) -> CSRConstraint: - """ - Staple sign and rhs onto a CSR-backed lhs to form an unassigned CSRConstraint. - - Builds directly from a :class:`~linopy.csr.CSRStore`. The - expression's label columns are kept as they are, its rows brought to - canonical form (:func:`~linopy.csr.canonical`), 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 - active = np.flatnonzero(~np.isnan(rhs_flat)) - all_active = active.size == rhs_flat.size - lhs = canonical(expr.csr) if all_active else canonical(expr.csr[active], True) - return cls( - lhs, - active, - rhs_flat if all_active else rhs_flat[active], - sign, - grid=expr.grid, - model=expr.model, - ) +def _con_unwrap(con: Constraint | Dataset) -> Dataset: + return con._dataset() if isinstance(con, Constraint) else con + +def _slack_names(attrs: Mapping[str, Any]) -> tuple[str, str] | None: + """Names of the positive and negative slack variables recorded by ``soften``.""" + positive = attrs.get("slack_positive") + if positive is None: + return None + return str(positive), str(attrs.get("slack_negative", "")) -def csr_rhs(expr: CSRStore, rhs: Any) -> DataArray | str: + +def _sparse_if(store: Store, sparse: bool, reason: Reason | None = None) -> Store: """ - Return the constant ``rhs`` as a DataArray on the expression grid, or the - reason the sparse path cannot take it: an rhs 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`. + ``store`` converted to sparse if ``sparse`` and it converts, else as is; + with a ``reason``, a store that stays dense emits the densify notice. """ - da = as_dataarray(rhs) - if set(da.dims) & set(HELPER_DIMS) or not set(da.dims) <= set(expr.grid.dims): - return "to_constraint(rhs over dimensions outside the grid)" - return da + if not sparse: + return store + converted = store.to_sparse() + if not isinstance(converted, str): + return converted + if reason is not None: + _densify_notice(f"{reason} ({converted})", store.model) + return store + +def model_store(store: Store, model: Model) -> Store: + """``store`` on the store class of ``model`` where it converts, without a notice.""" + return _sparse_if(store, model.store_cls is CSRStore) -def _rhs_grid_values(expr: CSRStore, rhs: DataArray) -> np.ndarray: + +def rhs_on_grid(store: CSRStore, rhs: Any) -> DataArray | Reason: """ - Broadcast the rhs onto the expression grid and flatten it, with dense parity: - NaN in the rhs raises (§5) and a reordered or differing index on a - shared dim raises (§8), as on the dense path. + The constant ``rhs`` as a DataArray on the grid of ``store``, with dense + parity: NaN in the rhs raises (§5) and a reordered or differing index on + a shared dim raises (§8). An rhs with helper dims or dims outside the + grid returns the reason to fall back to the dense path. An expression + rhs is moved to the lhs before, see :meth:`LinearExpression.to_constraint`. """ - if bool(rhs.isnull().any()): + da = as_dataarray(rhs) + grid = store.grid + if set(da.dims) & set(HELPER_DIMS) or not set(da.dims) <= set(grid.dims): + return "to_constraint(rhs over dimensions outside the grid)" + if bool(da.isnull().any()): check_user_nan() - for d in rhs.dims: - if not rhs.get_index(d).equals(expr.grid.indexes[str(d)]): + for d in da.dims: + if not da.get_index(d).equals(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: i for d, i in expr.grid.indexes.items() if d not in rhs.dims} + missing = {d: i for d, i in grid.indexes.items() if d not in da.dims} if missing: - rhs = rhs.expand_dims(missing) - return rhs.transpose(*expr.grid.dims).to_numpy().reshape(-1) + da = da.expand_dims(missing) + values = da.transpose(*grid.dims).to_numpy().reshape(-1) + return grid.dataarray(values.astype(float), name="rhs") + + +def _rows_on_grid( + grid: Grid, positions: np.ndarray, values: np.ndarray, fill: Any +) -> DataArray: + """``values`` at the flat ``positions`` of ``grid``, ``fill`` elsewhere.""" + full = np.full(grid.size, fill, dtype=np.result_type(values, np.asarray(fill))) + full[positions] = values + return grid.dataarray(full) class Slack(NamedTuple): @@ -1698,12 +233,14 @@ def _on_grid(value: Any, template: DataArray, label: str) -> DataArray: return da.assign_coords(missing) if missing else da -class Constraint(ConstraintBase): +class Constraint: """ Constraint ``lhs sign rhs`` on the grid of its left-hand side. The lhs is a :class:`~linopy.expressions.LinearExpression` with a zero - constant, its store in canonical form. ``sign``, ``rhs``, ``labels``, + constant, its store in canonical form: a + :class:`~linopy.store.DenseStore`, or a :class:`~linopy.csr.CSRStore` in a + sparse model (``Model(sparse=True)``). ``sign``, ``rhs``, ``labels``, ``scaling``, ``dual``, ``binary_var``, ``binary_val`` and ``blocks`` are DataArrays on its grid; :attr:`data` is the Dataset derived from them. Supports setters, xarray operations via conwrap, and from_rule @@ -1726,6 +263,8 @@ class Constraint(ConstraintBase): "_coef_dirty", ) + _fill_value = FILL_VALUE + _lhs: expressions.LinearExpression _sign: DataArray _rhs: DataArray @@ -1783,6 +322,7 @@ def from_dataset(cls, ds: Dataset, model: Model, name: str = "") -> Constraint: The Dataset holds ``coeffs`` and ``vars`` on the grid dims followed by ``_term``, and ``sign``, ``rhs`` and optionally ``labels``, ``scaling``, ``dual``, ``binary_var``, ``binary_val`` and ``blocks``. + The lhs is held on the store class of ``model`` where it converts. """ if not isinstance(ds, Dataset): raise ValueError(f"data must be a Dataset, got {type(ds)}") @@ -1793,10 +333,133 @@ def from_dataset(cls, ds: Dataset, model: Model, name: str = "") -> Constraint: raise ValueError(f"unknown constraint fields {sorted(unknown)} in data") (ds,) = xr.broadcast(ds, exclude=[TERM_DIM]) - lhs = expressions.LinearExpression(ds[["coeffs", "vars"]], model) + dense = expressions.LinearExpression(ds[["coeffs", "vars"]], model) + lhs = expressions.LinearExpression._from_store( + model_store(dense._store, model), model + ) fields = {k: ds[k] for k in _FIELDS if k in ds} return cls(lhs, model=model, name=name, attrs=ds.attrs, **fields) + @classmethod + def from_netcdf_ds(cls, ds: Dataset, model: Model, name: str) -> Constraint: + """ + Read a constraint written by :meth:`to_netcdf_ds`: the dense layout, + or the CSR layout (``_linopy_format="csr"``) into a sparse lhs. + + CSR files without the ``_csr_columns`` attribute hold dense variable + positions instead of labels. + """ + if ds.attrs.get("_linopy_format") != "csr": + return cls.from_dataset(ds, model, name) + attrs = ds.attrs + shape = tuple(attrs["shape"]) + indices = ds["indices"].values + columns = attrs.get("_csr_columns", "positions") + if columns == "positions": + vlabels = model.variables.label_index.vlabels + indices = vlabels[indices] + shape = (shape[0], int(vlabels[-1]) + 1 if len(vlabels) else 1) + elif columns != "labels": + raise ValueError( + f"Unknown CSR column format '{columns}' for constraint '{name}'." + ) + coord_dims = attrs["coord_dims"] + if isinstance(coord_dims, str): + coord_dims = [coord_dims] + grid = Grid.from_netcdf_vars(ds, coord_dims) + cindex_raw = int(attrs["cindex"]) + cindex = cindex_raw if cindex_raw >= 0 else None + rhs = ds["rhs"].values + if "_active_positions" in ds: + positions = ds["_active_positions"].values + elif "_con_labels" in ds: + positions = ds["_con_labels"].values - (cindex or 0) + else: + positions = np.arange(len(rhs), dtype=np.intp) + + counts = np.zeros(grid.size, dtype=np.int64) + counts[positions] = np.diff(ds["indptr"].values) + indptr = np.concatenate([[0], np.cumsum(counts)]) + csr = scipy.sparse.csr_array( + (ds["data"].values, indices, indptr), shape=(grid.size, shape[1]) + ) + store = CSRStore(csr, np.zeros(grid.size), grid, model) + + def on_grid(values: np.ndarray, fill: Any) -> DataArray: + return _rows_on_grid(grid, positions, values, fill) + + sign = ds["_sign"].values if "_sign" in ds else np.full(len(rhs), attrs["sign"]) + fields: dict[str, Any] = { + "sign": on_grid(sign, EQUAL), + "rhs": on_grid(rhs, np.nan), + } + if "scaling" in ds: + fields["scaling"] = on_grid(ds["scaling"].values, 1.0) + if cindex is not None: + labels = ds["labels"].values if "labels" in ds else positions + cindex + fields["labels"] = on_grid(labels, -1).astype(model._dtypes["labels"]) + if "dual" in ds: + fields["dual"] = on_grid(ds["dual"].values, np.nan) + if "_binvar_labels" in ds: + fields["binary_var"] = on_grid(ds["_binvar_labels"].values, -1) + binval = ds["_binval"].values if "_binval" in ds else attrs["binval"] + fields["binary_val"] = on_grid(np.broadcast_to(binval, len(rhs)), -1) + if "blocks" in ds: + fields["blocks"] = on_grid(ds["blocks"].values, -1) + con_attrs: dict[str, Any] = {} + if cindex is not None: + con_attrs["label_range"] = (cindex, cindex + grid.size) + if (slack := _slack_names(attrs)) is not None: + con_attrs["slack_positive"], con_attrs["slack_negative"] = slack + lhs = expressions.LinearExpression._from_store(store, model) + return cls(lhs, model=model, name=name, attrs=con_attrs, **fields) + + def to_netcdf_ds(self) -> Dataset: + """ + A Dataset for netcdf serialization: :attr:`data`, or for a sparse lhs + its raw CSR rows next to the flat fields, read by :meth:`from_netcdf_ds`. + """ + store = self._lhs._store + if not isinstance(store, CSRStore): + return self.data + csr = store.csr + + def flat(da: DataArray) -> DataArray: + return DataArray(da.values.ravel(), dims=["_flat"]) + + sign = self._sign.values.ravel() + data_vars: dict[str, DataArray] = { + "indptr": DataArray(csr.indptr, dims=["_indptr"]), + "indices": DataArray(csr.indices, dims=["_nnz"]), + "data": DataArray(csr.data, dims=["_nnz"]), + "rhs": flat(self._rhs), + "scaling": flat(self._scaling), + **store.grid.to_netcdf_vars(), + } + attrs: dict[str, Any] = { + "_linopy_format": "csr", + "_csr_columns": "labels", + "cindex": self.range[0] if self.is_assigned else -1, + "shape": list(csr.shape), + "coord_dims": list(store.grid.dims), + "name": self.name, + } + if sign.size and (sign == sign[0]).all(): + attrs["sign"] = str(sign[0]) + else: + data_vars["_sign"] = flat(self._sign) + optional = { + "labels": self._labels, + "dual": self._dual, + "_binvar_labels": self._binary_var, + "_binval": self._binary_val, + "blocks": self._blocks, + } + data_vars |= {k: flat(v) for k, v in optional.items() if v is not None} + if (slack := _slack_names(self._attrs)) is not None: + attrs["slack_positive"], attrs["slack_negative"] = slack + return Dataset(data_vars, attrs=attrs) + def _bind( self, lhs: expressions.LinearExpression, fields: Mapping[Any, Any] ) -> None: @@ -1822,6 +485,14 @@ def _bind( self._binary_val = self._optional_on_grid(fields, "binary_val", template) self._blocks = self._optional_on_grid(fields, "blocks", template) + def _replace_store(self, store: Store) -> None: + """Internal: hold ``store`` as the lhs, the fields kept. Sets ``_coef_dirty``.""" + self._bind( + expressions.LinearExpression._from_store(store, self._model), + self._fields(), + ) + self._coef_dirty = True + @staticmethod def _optional_on_grid( fields: Mapping[Any, Any], key: str, template: DataArray @@ -1860,21 +531,51 @@ def _copy(self, model: Model, deep: bool, include_solution: bool) -> Constraint: fields = {k: v.copy(deep=deep) for k, v in self._fields().items() if k in keep} return self._on_store(self._lhs._store.copy(model, deep), model, fields) + def _take(self, positions: DataArray) -> Constraint: + """The cells at the flat ``positions``, laid out on their grid.""" + rows = positions.values.ravel() + grid = Grid.from_dataset(positions, map(str, positions.dims)) + fields = { + k: positions.copy(data=v.values.ravel()[rows].reshape(positions.shape)) + for k, v in self._fields().items() + } + return self._on_store(self._lhs._store.take(rows, grid), self._model, fields) + def __getitem__( self, selector: str | int | slice | list | tuple | dict ) -> Constraint: - terms = self._lhs._store.to_dataset() - selected = Dataset({k: terms[k][selector] for k in terms}) + """ + Get selection from the constraint. + Returns a Constraint with the selected data. + """ + terms = self._lhs._dense_store("__getitem__").ds + selected = DenseStore( + Dataset({k: terms[k][selector] for k in terms}), self._model + ) fields = {k: v[selector] for k, v in self._fields().items()} - return self._on_store(DenseStore(selected, self._model), self._model, fields) + store = model_store(selected, self._model) + return self._on_store(store, self._model, fields) @property def data(self) -> Dataset: - """The dense Dataset derived from the lhs and the fields.""" - terms = self._lhs._store.to_dataset() + """ + The dense Dataset derived from the lhs and the fields, for a sparse + lhs with a densify notice. + """ + return self._dataset("data(read)") + + def _dataset(self, reason: Reason | None = None) -> Dataset: + """:attr:`data`, with a densify notice naming ``reason`` if given.""" + terms = self._lhs._dense_store(reason).ds fields = {"coeffs": terms.coeffs, "vars": terms.vars, **self._fields()} return Dataset(fields, attrs=self._attrs)[list(fields)] + @property + def data_attrs(self) -> list[str]: + """Data variables that define this constraint's persistent state.""" + base = list(Constraints.dataset_attrs) + return base + ["binary_var", "binary_val"] if self.is_indicator else base + @property def model(self) -> Model: return self._model @@ -1903,10 +604,43 @@ def dims(self) -> Frozen[Hashable, int]: def sizes(self) -> Frozen[Hashable, int]: return Frozen({**self._grid_template().sizes, TERM_DIM: self.nterm}) + @property + def ndim(self) -> int: + """Get the number of dimensions of the constraint.""" + return self.rhs.ndim + + @property + def shape(self) -> tuple[int, ...]: + """Get the shape of the constraint.""" + return self.rhs.shape + + @property + def size(self) -> int: + """Get the size of the constraint.""" + return self.rhs.size + + @property + def coord_dims(self) -> tuple[Hashable, ...]: + return tuple(k for k in self.dims if k not in HELPER_DIMS) + + @property + def coord_sizes(self) -> dict[Hashable, int]: + return {k: self.sizes[k] for k in self.coord_dims} + @property def coord_names(self) -> list[str]: return get_dims_with_index_levels(self._grid, self.coord_dims) + @property + def type(self) -> str: + """Get the type string of the constraint.""" + return "Constraint" if self.is_assigned else "Constraint (unassigned)" + + @property + def term_dim(self) -> str: + """Return the term dimension of the constraint.""" + return TERM_DIM + @property def is_assigned(self) -> bool: return self._labels is not None @@ -1925,6 +659,19 @@ def range(self) -> tuple[int, int]: """Return the range of the constraint.""" return self._attrs["label_range"] + @property + def mask(self) -> DataArray | None: + """ + Get the mask of the constraint. + + The mask indicates on which coordinates the constraint is enabled + (True) and disabled (False). + """ + if self.is_assigned: + result: DataArray = self.labels != FILL_VALUE["labels"] # type: ignore[assignment] + return result.astype(bool) + return None + @property def loc(self) -> LocIndexer: return LocIndexer(self) @@ -1949,58 +696,168 @@ def sign(self) -> DataArray: def rhs(self) -> DataArray: return self._rhs - @property - def scaling(self) -> DataArray: - return self._scaling + @property + def scaling(self) -> DataArray: + return self._scaling + + @scaling.setter + def scaling(self, value: ConstantLike) -> None: + scaling = broadcast_to_coords( + value, + coords=self.labels.coords, + dims=self.coord_dims, + label=f"scaling for constraint '{self.name}'", + ) + scaling = validate_scaling(scaling, f"scaling for constraint '{self.name}'") + self._scaling = _on_grid(scaling, self._grid_template(), "scaling") + + @property + def is_indicator(self) -> bool: + return self._binary_var is not None + + @property + def binary_var(self) -> DataArray | None: + return self._binary_var + + @property + def binary_val(self) -> int | np.ndarray | None: + return None if self._binary_val is None else self._binary_val.values + + @property + def lhs(self) -> expressions.LinearExpression: + """The lhs, a new expression on the store the constraint holds.""" + return expressions.LinearExpression._from_store(self._lhs._store, self._model) + + def __contains__(self, value: Any) -> bool: + return value in ("coeffs", "vars", *self._fields()) or value in self.coords + + def __repr__(self) -> str: + """Print the constraint arrays.""" + max_lines = options["display_max_rows"] + dims = list(self.coord_dims) + ndim = len(dims) + dim_names = self.coord_names + dim_sizes = list(self.coord_sizes.values()) + size = np.prod(dim_sizes) + mask = self.mask + masked_entries = (~mask).sum().values if mask is not None else 0 + store = self._lhs._store + lines = [] + + header_string = f"{self.type} `{self.name}`" if self.name else f"{self.type}" + + def row(indices: tuple[Any, ...]) -> str: + coeffs, vars_, _ = store.cell(indices) + expr = format_single_expression(coeffs, vars_, 0, self.model) + sign = SIGNS_pretty[self.sign.values[indices]] + return f"{expr} {sign} {self.rhs.values[indices]}" + + if size > 1 or ndim > 0: + row_labels = get_printout_labels(self._grid, dims) + for indices in generate_indices_for_printout(dim_sizes, max_lines): + if indices is None: + lines.append("\t\t...") + else: + coord = [row_labels[i][int(ind)] for i, ind in enumerate(indices)] + if mask is None or mask.values[indices]: + line = format_coord(coord) + f": {row(indices)}" + else: + line = format_coord(coord) + ": None" + lines.append(line) + lines = align_lines_by_delimiter(lines, list(SIGNS_pretty.values())) + + shape_str = ", ".join(f"{d}: {s}" for d, s in zip(dim_names, dim_sizes)) + mask_str = f" - {masked_entries} masked entries" if masked_entries else "" + underscore = "-" * (len(shape_str) + len(mask_str) + len(header_string) + 4) + lines.insert(0, f"{header_string} [{shape_str}]{mask_str}:\n{underscore}") + elif size == 1: + lines.append(f"{header_string}\n{'-' * len(header_string)}\n{row(())}") + else: + lines.append(f"{header_string}\n{'-' * len(header_string)}\n") + + return "\n".join(lines) + + def print(self, display_max_rows: int = 20, display_max_terms: int = 20) -> None: + """ + Print the linear expression. - @scaling.setter - def scaling(self, value: ConstantLike) -> None: - scaling = broadcast_to_coords( - value, - coords=self.labels.coords, - dims=self.coord_dims, - label=f"scaling for constraint '{self.name}'", - ) - scaling = validate_scaling(scaling, f"scaling for constraint '{self.name}'") - self._scaling = _on_grid(scaling, self._grid_template(), "scaling") + Parameters + ---------- + display_max_rows : int + Maximum number of rows to be displayed. + display_max_terms : int + Maximum number of terms to be displayed. + """ + with options as opts: + opts.set_value( + display_max_rows=display_max_rows, display_max_terms=display_max_terms + ) + print(self) - @property - def is_indicator(self) -> bool: - return self._binary_var is not None + def _flat(self) -> pd.DataFrame: + ds = self._dataset("flat") - @property - def binary_var(self) -> DataArray | None: - return self._binary_var + def mask_func(data: dict) -> pd.Series: + mask = (data["vars"] != -1) & (data["coeffs"] != 0) + if "labels" in data: + mask &= data["labels"] != -1 + return mask - @property - def binary_val(self) -> int | np.ndarray | None: - return None if self._binary_val is None else self._binary_val.values + df = to_dataframe(ds, mask_func=mask_func) - @property - def lhs(self) -> expressions.LinearExpression: - """The lhs, a new expression on the store the constraint holds.""" - return expressions.LinearExpression._from_store(self._lhs._store, self._model) + agg_custom = {k: "first" for k in list(df.columns)} + agg_standards = dict(coeffs="sum", rhs="first", sign="first") + agg = {**agg_custom, **agg_standards} + df = df.groupby(["labels", "vars"], as_index=False).aggregate(agg) + check_has_nulls(df, name=f"{self.type} {self.name}") + return df def _assign_lhs( - self, expr: expressions.LinearExpression, rhs: DataArray | None = None + self, + expr: expressions.LinearExpression, + rhs: DataArray | None = None, + reason: Reason | None = None, ) -> None: """ Internal: replace the lhs by ``expr``, moving its constant part to - the rhs. Sets ``_coef_dirty``. + the rhs and keeping a sparse lhs sparse where ``expr`` converts, else + with a densify notice naming ``reason`` if given. Sets ``_coef_dirty``. """ base_rhs = self.rhs if rhs is None else rhs fields = self._fields() | {"rhs": base_rhs - expr.const} - self._bind(expr.reset_const(), fields) + store = _sparse_if(expr.reset_const()._store, self._lhs.is_sparse, reason) + self._bind(expressions.LinearExpression._from_store(store, self._model), fields) self._coef_dirty = True - def _assign_terms(self, **terms: DataArray) -> None: + def _assign_terms( + self, + coeffs: ConstantLike | None, + variables: variables.Variable | None, + ) -> None: """ Internal: replace ``coeffs`` and/or ``vars`` of the dense lhs - layout. Sets ``_coef_dirty``. + layout, a sparse lhs converted with a notice and back. Sets + ``_coef_dirty``. """ - ds = assign_fields(self.lhs.data[["coeffs", "vars"]], **terms) - self._bind(expressions.LinearExpression(ds, self._model), self._fields()) - self._coef_dirty = True + from linopy.variables import Variable + + store = self._lhs._store + ds = store.to_dense("update(coeffs=, variables=)").ds[["coeffs", "vars"]] + if coeffs is not None: + new_coeffs = DataArray(coeffs).broadcast_like(ds.vars, exclude=[TERM_DIM]) + ds = assign_fields(ds, coeffs=new_coeffs) + if variables is not None: + if not isinstance(variables, Variable): + raise TypeError( + "Constraint.update(variables=...) expects a Variable; " + f"got {type(variables).__name__}." + ) + new_vars = variables.labels.broadcast_like(ds.coeffs, exclude=[TERM_DIM]) + ds = assign_fields(ds, vars=new_vars) + dense = expressions.LinearExpression(ds, self._model)._store + self._replace_store( + _sparse_if(dense, self._lhs.is_sparse, "update(coeffs=, variables=)") + ) def update( self, @@ -2090,10 +947,10 @@ def update( "Constraint.update: positional `constraint` argument " "cannot be combined with keyword arguments." ) - con: ConstraintBase + con: Constraint if isinstance(constraint, AnonymousScalarConstraint): con = constraint.to_constraint() - elif isinstance(constraint, ConstraintBase): + elif isinstance(constraint, Constraint): con = constraint else: raise TypeError( @@ -2121,7 +978,7 @@ def update( raise TypeError( "Constraint.update: lhs must be linear; got a quadratic expression." ) - self._assign_lhs(expr) + self._assign_lhs(expr, reason="update(lhs=)") if rhs is not None: expr = expressions.as_expression( @@ -2133,23 +990,8 @@ def update( else: self._rhs = _on_grid(expr.const, self._grid_template(), "rhs") - if coeffs is not None: - new_coeffs = DataArray(coeffs).broadcast_like( - self.vars, exclude=[self.term_dim] - ) - self._assign_terms(coeffs=new_coeffs) - if variables is not None: - from linopy.variables import Variable as _Variable - - if not isinstance(variables, _Variable): - raise TypeError( - "Constraint.update(variables=...) expects a Variable; " - f"got {type(variables).__name__}." - ) - new_vars = variables.labels.broadcast_like( - self.coeffs, exclude=[self.term_dim] - ) - self._assign_terms(vars=new_vars) + if coeffs is not None or variables is not None: + self._assign_terms(coeffs, variables) if sign is not None: new_sign = maybe_replace_signs(DataArray(sign)).broadcast_like(self.sign) @@ -2183,8 +1025,29 @@ def dual(self) -> DataArray: def dual(self, value: ConstantLike) -> None: self._dual = _on_grid(value, self._grid_template(), "dual") + def has_variable(self, variable: variables.Variable) -> bool: + """ + Check if the constraint references any of the given variable labels. + + Masked variable entries (label -1) are ignored: they are not part of + the model and would otherwise match every empty term slot. + """ + return self.has_labels(assigned_labels(variable.labels)) + def has_labels(self, labels: np.ndarray) -> bool: - return contains_labels(self.vars.values.ravel(), labels) + """ + Check if the constraint references any of the given variable labels. + + ``labels`` must not contain the ``-1`` sentinel, see + :func:`linopy.common.assigned_labels`. + """ + return contains_labels(self._lhs._store.variable_labels(), labels) + + def _label_values(self) -> np.ndarray: + """The flat labels; raises for a constraint without labels.""" + if self._labels is None: + raise ValueError(f"Constraint {self.name!r} has not been assigned labels.") + return self._labels.values.ravel() def to_matrix( self, label_index: VariableLabelIndex @@ -2204,39 +1067,41 @@ def to_matrix( """ active = self.active_row_mask() csr = self._lhs._store.positional_csr(label_index, active, self.name) - return csr, self.labels.values.ravel()[active] + return csr, self._label_values()[active] def active_row_mask(self) -> np.ndarray: """Boolean mask over raveled rows: label set and at least one variable present.""" - return (self.labels.values.ravel() != -1) & self._lhs._store.has_terms() + return (self._label_values() != -1) & self._lhs._store.has_terms() def active_labels(self) -> np.ndarray: - return self.labels.values.ravel()[self.active_row_mask()] + """Active constraint labels in build order, without building the CSR.""" + return self._label_values()[self.active_row_mask()] def to_matrix_with_rhs( self, label_index: VariableLabelIndex ) -> tuple[scipy.sparse.csr_array, np.ndarray, np.ndarray, np.ndarray]: - """Return (csr, con_labels, b, sense) in one pass.""" + """ + Return (csr, con_labels, b, sense) in one pass. + + Avoids computing the CSR matrix twice when both the matrix and + the RHS/sense vectors are needed. + """ active = self.active_row_mask() csr = self._lhs._store.positional_csr(label_index, active, self.name) - con_labels = self.labels.values.ravel()[active] + con_labels = self._label_values()[active] b = self.rhs.values.ravel()[active] sense = self.sign.values.ravel()[active].astype("U1") return csr, con_labels, b, sense def sanitize_zeros(self) -> Constraint: """Remove terms with zero or near-zero coefficients.""" - not_zero = abs(self.coeffs) > 1e-10 - self._assign_terms( - vars=self.vars.where(not_zero, -1), - coeffs=self.coeffs.where(not_zero), - ) + self._replace_store(self._lhs._store.prune(1e-10)) return self def sanitize_missings(self) -> Constraint: """Mask out rows where all variables are missing (-1).""" - contains_non_missing = (self.vars != -1).any(self.term_dim) - self._labels = self.labels.where(contains_non_missing, -1) + has_terms = self._lhs._store.dataarray(self._lhs._store.has_terms(), "labels") + self._labels = self.labels.where(has_terms, -1) return self def sanitize_infinities(self) -> Constraint: @@ -2247,14 +1112,146 @@ def sanitize_infinities(self) -> Constraint: self._labels = self.labels.where(~valid_infinity_values, -1) return self - def freeze(self) -> CSRConstraint: - """Convert to an immutable Constraint.""" - return CSRConstraint.from_dense(self) + def to_dense(self) -> Constraint: + """Hold the lhs on a dense store, in place.""" + self._replace_store(self._lhs._store.to_dense()) + return self - def mutable(self) -> Constraint: - """Return self (already mutable).""" + def to_sparse(self) -> Constraint: + """Hold the lhs on a sparse (CSR) store, in place; raises where it cannot convert.""" + self._replace_store(as_sparse(self._lhs._store)) return self + @property + def slack(self) -> Slack | None: + """ + Slack variable(s) added via :meth:`soften`, or ``None`` if the + constraint has never been softened. + """ + names = _slack_names(self.attrs) + if names is None: + return None + positive, negative = names + return Slack( + positive=self.model.variables[positive], + negative=self.model.variables[negative] if negative else None, + ) + + def soften( + self, + penalty: ConstantLike, + *, + max_violation: ConstantLike | None = None, + name: str | None = None, + ) -> Slack: + """ + Soften a constraint, adding a slack variable and a penalty to the objective function. + + Parameters + ---------- + penalty : constant-like + The penalty that will match the slack variable inside the objective function. Must be bigger than 0. + max_violation: constant-like + The max violation possible that caps the slack (upper bound). If None, the slack will be unbounded. + name: string + The name for the slack variable. If None, it will reuse the constraint name and add a '_slack'. + + Returns + ------- + Slack + Named tuple with the `positive` slack variable, and the `negative` one for equality constraints (`None` + for inequality constraints). + + Notes + ----- + On a sparse lhs the slack terms are merged into the sparse rows, without densifying. + A detached copy (from `.sel()`, `.isel()` or indexing) is not registered in + model.constraints, so soften raises ValueError on it. + + Softening an already-softened constraint raises ValueError instead of stacking a second, redundant slack term + onto the same lhs. + + Examples + -------- + >>> from linopy import Model + >>> import pandas as pd + + >>> m = Model() + >>> investments = pd.Index(["A", "B", "C"], name="investments") + >>> expected_return = pd.Series( + ... [0.08, 0.03, 0.1], index=investments, name="expected_return" + ... ) + >>> w = m.add_variables(lower=0, upper=1, coords=[investments], name="weights") + >>> m.add_objective((expected_return * w).sum(), sense="max") + >>> budget_penalty = 2 + + >>> budget_constraint = m.add_constraints(w.sum() == 1, name="budget") + >>> slack = budget_constraint.soften(penalty=budget_penalty) + """ + if not bool(np.all(np.asarray(penalty) > 0)): + raise ValueError("Penalty is not positive.") + + model = self.model + if model.objective.expression.empty: + raise ValueError( + "Objective must be defined via `model.add_objective` before calling `soften` on constraints." + ) + + if model.constraints.data.get(self.name) is not self: + raise ValueError( + f"Constraint {self.name!r} is not the constraint registered in the model, so " + "`soften` would not affect it (it may be a detached copy from `.sel()`, " + "`.isel()` or indexing). Call `soften` on `model.constraints[name]` " + "directly." + ) + + existing_slack = self.slack + if existing_slack is not None: + raise ValueError( + f"Constraint {self.name!r} was already softened (existing slack " + f"variable {existing_slack.positive.name!r})" + ) + + sign_values = pd.unique(self.sign.values.ravel()) + if len(sign_values) > 1: + raise NotImplementedError( + "Constraint.soften does not support constraints with mixed signs." + ) + sign = sign_values.item() + + name = name or f"{self.name}_slack" + upper = np.inf if max_violation is None else max_violation + + def add_slack(suffix: str) -> variables.Variable: + return model.add_variables( + lower=0, + upper=upper, + coords=self.coords, + mask=self.mask, + name=f"{name}_{suffix}", + ) + + positive = add_slack("pos") + slack = Slack(positive, add_slack("neg") if sign == EQUAL else None) + self._attach_slack(slack, sign) + + violation = positive if slack.negative is None else positive + slack.negative + direction = 1 if model.sense == "min" else -1 + model.objective += direction * (penalty * violation).sum() + return slack + + def _attach_slack(self, slack: Slack, sign: str) -> None: + """Add the slack terms to the lhs and record the slack variable names.""" + positive = slack.positive + lhs = self.lhs + positive if sign == GREATER_EQUAL else self.lhs - positive + if slack.negative is not None: + lhs = lhs + slack.negative + self.update(lhs=lhs) + self._attrs.update( + slack_positive=slack.positive.name, + slack_negative=slack.negative.name if slack.negative is not None else "", + ) + @classmethod def from_rule(cls, model: Model, rule: Callable, coords: CoordsLike) -> Constraint: """ @@ -2321,17 +1318,6 @@ def from_rule(cls, model: Model, rule: Callable, coords: CoordsLike) -> Constrai rhs = DataArray(array([c.rhs for c in cons]).reshape(shape), coords) return cls(lhs, sign, rhs, model) - def _attach_slack(self, slack: Slack, sign: str) -> None: - positive = slack.positive - lhs = self.lhs + positive if sign == GREATER_EQUAL else self.lhs - positive - if slack.negative is not None: - lhs = lhs + slack.negative - self.update(lhs=lhs) - self._attrs.update( - slack_positive=slack.positive.name, - slack_negative=slack.negative.name if slack.negative is not None else "", - ) - def to_polars(self) -> pl.DataFrame: """ Convert the constraint to a polars DataFrame. @@ -2344,31 +1330,25 @@ def to_polars(self) -> pl.DataFrame: ------- df : polars.DataFrame """ - ds = self.data - - keys = [k for k in ds if ("_term" in ds[k].dims) or (k == "labels")] - long = to_polars(ds[keys]) - + cells, vars_, coeffs = self._lhs._store.terms() + labels_flat = self._label_values() + long = pl.DataFrame( + {"labels": labels_flat[cells], "coeffs": coeffs, "vars": vars_} + ) long = filter_nulls_polars(long) - if ds.sizes.get("_term", 1) > 1: + if self.nterm > 1: long = maybe_group_terms_polars(long) check_has_nulls_polars(long, name=f"{self.type} {self.name}") - labels_flat = ds["labels"].values.reshape(-1) mask = labels_flat != -1 - labels_masked = labels_flat[mask] - rhs_flat = np.broadcast_to(ds["rhs"].values, ds["labels"].shape).reshape(-1) - - sign_values = ds["sign"].values - sign_flat = np.broadcast_to(sign_values, ds["labels"].shape).reshape(-1) + short_data: dict = { + "labels": labels_flat[mask], + "rhs": self.rhs.values.ravel()[mask], + } + sign_flat = self.sign.values.ravel() all_same_sign = len(sign_flat) > 0 and ( sign_flat[0] == sign_flat[-1] and (sign_flat[0] == sign_flat).all() ) - - short_data: dict = { - "labels": labels_masked, - "rhs": rhs_flat[mask], - } if all_same_sign: short = pl.DataFrame(short_data).with_columns( pl.lit(sign_flat[0]).cast(pl.Enum(["=", "<=", ">="])).alias("sign") @@ -2382,7 +1362,27 @@ def to_polars(self) -> pl.DataFrame: df = long.join(short, on="labels", how="inner") return df[["labels", "coeffs", "vars", "sign", "rhs"]] - # Wrapped xarray methods — only available on Constraint + def iterate_slices( + self, + slice_size: int | None = 10_000, + slice_dims: list | None = None, + ) -> Generator[Constraint, None, None]: + """ + Yield sub-constraints along the largest of ``slice_dims``, each with at + most about ``slice_size`` term slots, see + :func:`linopy.common.iterate_slices`. The lhs store is gathered by + :meth:`~linopy.store.Store.take` and keeps its backing. + """ + template = self._grid_template() + positions = template.copy(data=np.arange(template.size).reshape(template.shape)) + slots = positions.expand_dims({TERM_DIM: self.nterm}, axis=-1) + dims = list(self.coord_dims) if slice_dims is None else slice_dims + for block in iterate_slices(slots, slice_size, dims): + if block.sizes == slots.sizes: + yield self + else: + yield self._take(block.isel({TERM_DIM: 0})) + assign = conwrap(assign_fields) assign_attrs = conwrap(Dataset.assign_attrs) assign_coords = conwrap(Dataset.assign_coords) @@ -2406,7 +1406,7 @@ def to_polars(self) -> pl.DataFrame: @dataclass(repr=False) -class Constraints(LabelContainer[ConstraintBase]): +class Constraints(LabelContainer[Constraint]): """ A constraint container used for storing multiple constraint arrays. """ @@ -2422,9 +1422,9 @@ def __getitem__(self, names: str) -> Constraint: ... def __getitem__(self, names: list[str]) -> Constraints: ... def __getitem__(self, names: str | list[str]) -> Constraint | Constraints: - return super().__getitem__(names) # type: ignore[return-value] + return super().__getitem__(names) - def add(self, constraint: ConstraintBase, freeze: bool = False) -> ConstraintBase: + def add(self, constraint: Constraint) -> Constraint: """ Add a constraint to the constraints container. """ @@ -2433,8 +1433,6 @@ def add(self, constraint: ConstraintBase, freeze: bool = False) -> ConstraintBas f"Constraint {constraint.name!r} has not been assigned labels; " "use `Model.add_constraints`." ) - if freeze and isinstance(constraint, Constraint): - constraint = constraint.freeze() self.data[constraint.name] = constraint self._invalidate_label_position_index() return constraint @@ -2628,15 +1626,10 @@ def set_blocks(self, block_map: np.ndarray) -> None: """ N = block_map.max() - for name, constraint in self.items(): - if not isinstance(constraint, Constraint): - _densify_notice( - "frozen constraint converted by `set_blocks`", self.model - ) - constraint = Constraint.from_dataset(constraint.data, self.model, name) - self.data[name] = constraint + for constraint in self.data.values(): res = xr.full_like(constraint.labels, N + 1, dtype=block_map.dtype) - entries = replace_by_map(constraint.vars, block_map) + terms = constraint._lhs._store.to_dense("`set_blocks`").ds + entries = replace_by_map(terms.vars, block_map) not_zero = entries != 0 not_missing = entries != -1 @@ -2702,15 +1695,8 @@ def reset_dual(self) -> None: """ Reset the stored solution of variables. """ - for k, c in self.items(): - if isinstance(c, CSRConstraint): - if c._dual is not None: - self.data[k] = c._replace(dual=None) - elif isinstance(c, Constraint): - c._dual = None - else: - msg = f"reset_dual encountered an unknown constraint type: {type(c)}" - raise NotImplementedError(msg) + for c in self.data.values(): + c._dual = None class AnonymousScalarConstraint: diff --git a/linopy/csr.py b/linopy/csr.py index 7302a6094..ec9768eb8 100644 --- a/linopy/csr.py +++ b/linopy/csr.py @@ -11,7 +11,7 @@ coefficients of a repeated variable adding up. Operations stack, interleave and gather the stored terms as they are; :func:`canonical` (duplicate variables summed, terms label-ordered) runs once where a consumer needs it, -when a constraint is added (:meth:`linopy.constraints.CSRConstraint.from_csr`), +when a constraint is built on it (:class:`linopy.constraints.Constraint`), which also covers matrix assembly and export. ``nterm`` is the widest stored row, as for the dense term layout of the same operation. The stored arrays (``data``, ``indices``, ``indptr`` and ``const``) are read-only views, so @@ -32,8 +32,7 @@ This module documents the CSR structure only, working on plain datasets. The one bridge back to the dense store is :meth:`CSRStore.to_dense`, which wraps the expanded dataset in a :class:`~linopy.store.DenseStore`. The reverse -bridges are :meth:`CSRStore.from_dense` and -:meth:`linopy.constraints.CSRConstraint.from_csr`. +bridge is :meth:`CSRStore.from_dense`. """ from __future__ import annotations @@ -287,6 +286,9 @@ def __post_init__(self) -> None: object.__setattr__(self, "csr", csr) object.__setattr__(self, "const", _readonly(self.const, copy=False)) + def __getstate__(self) -> dict[str, Any]: + return {**self.__dict__, "_positional_cache": None} + @property def shape(self) -> tuple[int, ...]: return self.grid.shape @@ -547,6 +549,9 @@ def terms(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: keep = self.live_terms() return cells[keep], self.csr.indices[keep], self.csr.data[keep] + def variable_labels(self) -> np.ndarray: + return self.csr.indices + def solution( self, label_solution: Callable[[np.ndarray], np.ndarray] ) -> np.ndarray: @@ -607,9 +612,13 @@ def live_terms(self) -> np.ndarray: present = np.repeat(~np.isnan(self.const), np.diff(self.csr.indptr)) return present & (self.csr.data != 0) - def prune(self) -> CSRStore: - """Drop explicit zero coefficients; cell activeness stays with ``const``.""" + def prune(self, tol: float = 0.0) -> CSRStore: + """ + Drop the coefficients of at most ``tol`` in magnitude; cell activeness + stays with ``const``. + """ csr = self.csr.copy() + csr.data[np.abs(csr.data) <= tol] = 0 csr.eliminate_zeros() return replace(self, csr=csr) @@ -1096,26 +1105,26 @@ def csr_to_term_arrays( return vars_, coeffs -def _conversion_notice(message: str, model: Model, explicit: bool) -> None: +def _conversion_notice(message: str, model: Model) -> None: """ Emit ``message`` as a :class:`~linopy.constants.PerformanceWarning` if - ``options["warn_on_densify"]`` is set, or left at its default ``None``, - ``model`` is sparse and the conversion is an implicit fallback rather - than an ``explicit`` user request. + ``options["warn_on_densify"]`` is set, or left at its default ``None`` + and ``model`` is sparse. Only implicit fallbacks emit it; an explicit + conversion, e.g. :meth:`linopy.constraints.Constraint.to_dense`, does not. """ enabled = options["warn_on_densify"] - if model.sparse and not explicit if enabled is None else enabled: + if model.sparse if enabled is None else enabled: warn_outside_linopy(message, PerformanceWarning) -def _densify_notice(reason: str, model: Model, explicit: bool = False) -> None: +def _densify_notice(reason: str, model: Model) -> None: """The conversion notice naming why a sparse (CSR) backing is dropped.""" - _conversion_notice(f"Sparse (CSR) backing densified: {reason}.", model, explicit) + _conversion_notice(f"Sparse (CSR) backing densified: {reason}.", model) def _sparsify_notice(reason: str, model: Model) -> None: """The conversion notice naming why a dense backing is converted to sparse.""" - _conversion_notice(f"Dense backing sparsified: {reason}.", model, False) + _conversion_notice(f"Dense backing sparsified: {reason}.", model) def _readonly(values: np.ndarray, copy: bool = True) -> np.ndarray: diff --git a/linopy/expressions.py b/linopy/expressions.py index cfaa3cb1d..9b307850c 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -147,7 +147,7 @@ if TYPE_CHECKING: from linopy.constraints import ( AnonymousScalarConstraint, - ConstraintBase, + Constraint, ) from linopy.model import Model from linopy.variables import ScalarVariable, Variable @@ -1151,13 +1151,13 @@ def __truediv__(self, other: SideLike) -> Self: except TypeError: return NotImplemented - def __le__(self, rhs: SideLike) -> ConstraintBase: + def __le__(self, rhs: SideLike) -> Constraint: return self.to_constraint(LESS_EQUAL, rhs) - def __ge__(self, rhs: SideLike) -> ConstraintBase: + def __ge__(self, rhs: SideLike) -> Constraint: return self.to_constraint(GREATER_EQUAL, rhs) - def __eq__(self, rhs: SideLike) -> ConstraintBase: # type: ignore[override] + def __eq__(self, rhs: SideLike) -> Constraint: # type: ignore[override] return self.to_constraint(EQUAL, rhs) def __gt__(self, other: Any) -> NotImplementedType: @@ -1319,7 +1319,7 @@ def le( self, rhs: SideLike, join: JoinOptions | None = None, - ) -> ConstraintBase: + ) -> Constraint: """ Less than or equal constraint. @@ -1339,7 +1339,7 @@ def ge( self, rhs: SideLike, join: JoinOptions | None = None, - ) -> ConstraintBase: + ) -> Constraint: """ Greater than or equal constraint. @@ -1359,7 +1359,7 @@ def eq( self, rhs: SideLike, join: JoinOptions | None = None, - ) -> ConstraintBase: + ) -> Constraint: """ Equality constraint. @@ -1630,7 +1630,7 @@ def cumsum( def to_constraint( self, sign: SignLike, rhs: SideLike, join: JoinOptions | None = None - ) -> ConstraintBase: + ) -> Constraint: """ Convert a linear expression to a constraint. @@ -2467,7 +2467,7 @@ def fillna( def to_constraint( self, sign: SignLike, rhs: SideLike, join: JoinOptions | None = None - ) -> ConstraintBase: + ) -> Constraint: store = self._store if not isinstance(store, CSRStore): return super().to_constraint(sign, rhs, join) @@ -2477,10 +2477,12 @@ def to_constraint( 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(store, rhs) + rhs_da = constraints.rhs_on_grid(store, rhs) if isinstance(rhs_da, str): return self._fallback(rhs_da).to_constraint(sign, rhs, join) - return constraints.CSRConstraint.from_csr(store, sign, rhs_da) + return constraints.Constraint( + self.reset_const(), sign, rhs_da - self.const, self.model + ) @overload def __add__( diff --git a/linopy/io.py b/linopy/io.py index fcb87151a..9b4473c81 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -21,7 +21,6 @@ import numpy as np import pandas as pd import polars as pl -import scipy.sparse import xarray as xr from tqdm import tqdm @@ -584,7 +583,7 @@ def indicator_constraints_to_file( ) for con in m.constraints.indicator.data.values(): - ic_data = con.data + ic_data = con._dataset() labels_flat = ic_data.labels.values.flatten() binary_var_flat = ic_data.binary_var.values.flatten() binary_val_flat = np.broadcast_to( @@ -1187,12 +1186,7 @@ def read_netcdf(path: Path | str, **kwargs: Any) -> Model: the deserialized Model is too, but :meth:`Model.undo_sos_reformulation` is a no-op on it. """ - from linopy.constraints import ( - Constraint, - ConstraintBase, - Constraints, - CSRConstraint, - ) + from linopy.constraints import Constraint, Constraints from linopy.expressions import Expressions, LinearExpression, QuadraticExpression from linopy.model import Model from linopy.variables import Variable, Variables @@ -1282,14 +1276,10 @@ def owner(k: str) -> str | None: m._expressions = Expressions(expressions, m) - constraints: dict[str, ConstraintBase] = {} + constraints: dict[str, Constraint] = {} for k in containers["constraints"]: name = remove_prefix(k, "constraints") - con_ds = get_prefix(ds, k) - if con_ds.attrs.get("_linopy_format") == "csr": - constraints[name] = CSRConstraint.from_netcdf_ds(con_ds, m, name) - else: - constraints[name] = Constraint.from_dataset(con_ds, m, name) + constraints[name] = Constraint.from_netcdf_ds(get_prefix(ds, k), m, name) m._constraints = Constraints(constraints, m) objective = get_prefix(ds, "objective") @@ -1372,12 +1362,7 @@ def copy(m: Model, include_solution: bool = False, deep: bool = True) -> Model: Model A deep or shallow copy of the model. """ - from linopy.constraints import ( - Constraint, - ConstraintBase, - Constraints, - CSRConstraint, - ) + from linopy.constraints import Constraints from linopy.expressions import Expressions, LinearExpression, QuadraticExpression from linopy.model import Model, Objective from linopy.variables import Variable, Variables @@ -1427,22 +1412,11 @@ def _copy_expr( new_model, ) - def _copy_con(con: ConstraintBase) -> ConstraintBase: - if isinstance(con, Constraint): - return con._copy(new_model, deep, include_solution) - if not isinstance(con, CSRConstraint): - raise TypeError(f"Cannot copy constraint of type {type(con)}") - buffer_types = (np.ndarray, scipy.sparse.sparray) - changes: dict[str, Any] = {"model": new_model} - if deep: - kwargs = con._init_kwargs().items() - changes |= {k: v.copy() for k, v in kwargs if isinstance(v, buffer_types)} - if not include_solution: - changes["dual"] = None - return con._replace(**changes) - new_model._constraints = Constraints( - {name: _copy_con(con) for name, con in m.constraints.items()}, + { + name: con._copy(new_model, deep, include_solution) + for name, con in m.constraints.items() + }, new_model, ) diff --git a/linopy/matrices.py b/linopy/matrices.py index f7beaaf07..e4d22bbfa 100644 --- a/linopy/matrices.py +++ b/linopy/matrices.py @@ -9,67 +9,46 @@ from collections.abc import Callable from functools import cached_property, partial -from typing import TYPE_CHECKING, NamedTuple, cast +from typing import TYPE_CHECKING, NamedTuple import numpy as np import scipy.sparse from numpy import ndarray from linopy import expressions -from linopy.constraints import CSRConstraint from linopy.csr import index_dtype if TYPE_CHECKING: from linopy.common import VariableLabelIndex - from linopy.constraints import ConstraintBase + from linopy.constraints import Constraint from linopy.model import Model class _RowBlock(NamedTuple): """ - Rows of one constraint. ``positional`` returns its positional CSR; for a - frozen constraint it is built only while the rows are gathered. + Active rows of one constraint. ``positional`` returns their positional + CSR, built only while the rows are gathered; ``nnz`` bounds its nonzeros. """ positional: Callable[[], scipy.sparse.csr_array] n_rows: int nnz: int rhs: ndarray - sign: str | ndarray + sign: ndarray row_scaling: ndarray | None -def _unit_or(scaling: ndarray) -> ndarray | None: - return None if (scaling == 1).all() else scaling - - -def _row_scaling(c: ConstraintBase) -> ndarray: - """Scaling of the active rows of a constraint, in row order.""" - if isinstance(c, CSRConstraint): - return c._scaling - return c.scaling.values.ravel()[c.active_row_mask()] - - -def _row_block(con: ConstraintBase, label_index: VariableLabelIndex) -> _RowBlock: - row_scaling = _unit_or(_row_scaling(con)) - if isinstance(con, CSRConstraint): - stored = con._csr - return _RowBlock( - partial(con._to_positional_csr, label_index), - stored.shape[0], - int(np.count_nonzero(stored.data)), - con._rhs, - con._sign, - row_scaling, - ) - csr, _, rhs, sense = con.to_matrix_with_rhs(label_index) +def _row_block(con: Constraint, label_index: VariableLabelIndex) -> _RowBlock: + active = con.active_row_mask() + store = con.lhs._store + scaling = con.scaling.values.ravel()[active] return _RowBlock( - lambda: csr, - csr.shape[0], - int(np.count_nonzero(csr.data)), - rhs, - sense, - row_scaling, + partial(store.positional_csr, label_index, active, con.name), + int(active.sum()), + store.variable_labels().size, + con.rhs.values.ravel()[active], + con.sign.values.ravel()[active].astype("U1"), + None if (scaling == 1).all() else scaling, ) @@ -106,7 +85,7 @@ def _stack( pos = row = 0 for block in blocks: source = block.positional() - count = block.nnz + count = int(np.count_nonzero(source.data)) block_data, block_indices, block_indptr = ( source.data, source.indices, @@ -135,15 +114,18 @@ def _stack( block_offsets += pos sense[rows] = block.sign pos, row = entries.stop, rows.stop - A = scipy.sparse.csr_array((data, indices, indptr), shape=(n_rows, n_cols)) + A = scipy.sparse.csr_array( + (data[:pos], indices[:pos], indptr), shape=(n_rows, n_cols) + ) return A, b, sense -def _binval_per_row(binval: int | np.ndarray, n: int) -> ndarray: - """Broadcast an indicator triggering value to one entry per active row.""" - if np.ndim(binval) == 0: - return np.full(n, int(binval), dtype=np.intp) - return np.asarray(binval, dtype=np.intp).ravel() +def _indicator_rows(con: Constraint, active: ndarray) -> tuple[ndarray, ndarray]: + """Binary variable labels and triggering values of the active indicator rows.""" + if con.binary_var is None or con.binary_val is None: + raise ValueError(f"Constraint {con.name!r} is not an indicator constraint.") + var_labels = con.binary_var.values.ravel()[active] + return var_labels, np.ravel(con.binary_val)[active].astype(np.intp) class MatrixAccessor: @@ -208,11 +190,11 @@ def _build_cons(self) -> None: if not c.is_indicator: regular.append(_row_block(c, label_index)) continue - cc = c if isinstance(c, CSRConstraint) else c.freeze() - indicator.append(_row_block(cc, label_index)) - binvar.append(label_to_pos[cc._binvar_labels]) - cc_binval = cast("int | np.ndarray", cc._binval) - binval.append(_binval_per_row(cc_binval, len(cc._rhs))) + active = c.active_row_mask() + indicator.append(_row_block(c, label_index)) + var_labels, val = _indicator_rows(c, active) + binvar.append(label_to_pos[var_labels]) + binval.append(val) self.clabels: ndarray = m.constraints.label_index.clabels self.A: scipy.sparse.csr_array | None @@ -286,20 +268,12 @@ def dual(self) -> ndarray: for c in m.constraints.data.values(): if c.is_indicator: continue - if isinstance(c, CSRConstraint): - # _dual is active-only - if c._dual is not None: - dual_list.append(c._dual) - has_dual = True - else: - dual_list.append(np.full(c.ncons, np.nan)) + active = c.active_row_mask() + if "dual" in c: + dual_list.append(c.dual.values.ravel()[active]) + has_dual = True else: - active_rows = np.flatnonzero(c.active_row_mask()) - if "dual" in c.data: - dual_list.append(c.dual.values.ravel()[active_rows]) - has_dual = True - else: - dual_list.append(np.full(len(active_rows), np.nan)) + dual_list.append(np.full(int(active.sum()), np.nan)) if not has_dual: raise AttributeError( "Underlying is optimized but does not have dual values stored." diff --git a/linopy/model.py b/linopy/model.py index 2add8e245..730e8e994 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -54,11 +54,10 @@ from linopy.constraints import ( AnonymousScalarConstraint, Constraint, - ConstraintBase, Constraints, - CSRConstraint, + model_store, ) -from linopy.csr import CSRStore, _densify_notice +from linopy.csr import CSRStore from linopy.dualization import dualize from linopy.expressions import ( Expressions, @@ -268,8 +267,7 @@ def __init__( Build the model on the sparse (CSR) path, exposed read-only as ``Model.sparse``. ``groupby(...).sum()`` and ``@``/``dot`` against a constant return CSR-backed expressions, and ``add_constraints`` - freezes constraints to a CSRConstraint unless ``freeze=False`` is - passed. Requires no ``chunk``. See + holds the constraint lhs on a CSR store. Requires no ``chunk``. See :mod:`linopy.csr`. The default is False. Returns @@ -1115,7 +1113,7 @@ def _constraint_from_lhs( sign: SignLike | None, rhs: ConstantLike | VariableLike | ExpressionLike | None, coords: Sequence[Sequence | pd.Index] | Mapping | None = None, - ) -> ConstraintBase: + ) -> Constraint: """Build the anonymous constraint from an ``lhs`` and optional ``sign``/``rhs``.""" msg_required = ( f"`sign` and `rhs` are required when `lhs` is a {type(lhs).__name__}." @@ -1140,7 +1138,7 @@ def _constraint_from_lhs( if sign is not None or rhs is not None: raise ValueError(msg_must_be_none) return lhs.to_constraint() - elif isinstance(lhs, ConstraintBase): + elif isinstance(lhs, Constraint): if sign is not None or rhs is not None: raise ValueError(msg_must_be_none) return lhs @@ -1169,60 +1167,6 @@ def _allocate_constraint_labels( labels = labels.where(mask, -1) return labels, (start, end) - @overload - def add_constraints( - self, - lhs: VariableLike - | ExpressionLike - | ConstraintLike - | Sequence[tuple[ConstantLike, VariableLike | str]] - | Callable, - sign: SignLike | None = ..., - rhs: ConstantLike | VariableLike | ExpressionLike | None = ..., - name: str | None = ..., - coords: Sequence[Sequence | pd.Index] | Mapping | None = ..., - mask: MaskLike | None = ..., - freeze: Literal[False] = ..., - scaling: ConstantLike = ..., - penalty: ConstantLike | None = ..., - ) -> Constraint: ... - - @overload - def add_constraints( - self, - lhs: VariableLike - | ExpressionLike - | ConstraintLike - | Sequence[tuple[ConstantLike, VariableLike | str]] - | Callable, - sign: SignLike | None = ..., - rhs: ConstantLike | VariableLike | ExpressionLike | None = ..., - name: str | None = ..., - coords: Sequence[Sequence | pd.Index] | Mapping | None = ..., - mask: MaskLike | None = ..., - freeze: Literal[True] = ..., - scaling: ConstantLike = ..., - penalty: None = ..., - ) -> CSRConstraint: ... - - @overload - def add_constraints( - self, - lhs: VariableLike - | ExpressionLike - | ConstraintLike - | Sequence[tuple[ConstantLike, VariableLike | str]] - | Callable, - sign: SignLike | None = ..., - rhs: ConstantLike | VariableLike | ExpressionLike | None = ..., - name: str | None = ..., - coords: Sequence[Sequence | pd.Index] | Mapping | None = ..., - mask: MaskLike | None = ..., - freeze: bool | None = ..., - scaling: ConstantLike = ..., - penalty: ConstantLike | None = ..., - ) -> ConstraintBase: ... - def add_constraints( self, lhs: VariableLike @@ -1235,10 +1179,9 @@ def add_constraints( name: str | None = None, coords: Sequence[Sequence | pd.Index] | Mapping | None = None, mask: MaskLike | None = None, - freeze: bool | None = None, scaling: ConstantLike = 1, penalty: ConstantLike | None = None, - ) -> ConstraintBase: + ) -> Constraint: """ Assign a new, possibly multi-dimensional array of constraints to the model. @@ -1250,7 +1193,7 @@ def add_constraints( Parameters ---------- - lhs : linopy.LinearExpression/linopy.ConstraintBase/callable + lhs : linopy.LinearExpression/linopy.Constraint/callable Left hand side of the constraint(s) or optionally full constraint. In case a linear expression is passed, `sign` and `rhs` must not be None. @@ -1272,10 +1215,6 @@ def add_constraints( Boolean mask with False values for constraints which are skipped. The shape of the mask has to match the shape the added constraints. Default is None. - freeze : bool, optional - If True, convert the constraint to an immutable CSR-backed CSRConstraint - for better memory efficiency. If None, freezes if the model is - sparse (``Model(sparse=True)``). scaling : float/array_like, optional Positive finite scaling factor(s) for constraint rows. Solver-side left-hand-side coefficients and right-hand-side values are multiplied @@ -1283,8 +1222,8 @@ def add_constraints( penalty : constant-like, optional If given, soften the constraint right away by calling :meth:`Constraint.soften` with this penalty, adding a slack variable - and a penalty term to the objective. A frozen constraint is - softened natively on its sparse rows. The resulting Slack is not + and a penalty term to the objective. A sparse lhs is softened + on its sparse rows. The resulting Slack is not returned by this shortcut; retrieve the slack variable(s) from model.variables using the derived name f"{name}_slack_pos" (and f"{name}_slack_neg" for @@ -1292,14 +1231,12 @@ def add_constraints( Returns ------- - constraint : linopy.ConstraintBase - The added constraint (Constraint by default, or CSRConstraint if freeze=True). + constraint : linopy.Constraint + The added constraint, its lhs on the model's store class + (:attr:`store_cls`) where it converts. """ name = self._resolve_constraint_name(name) - if freeze is None: - freeze = self._sparse - freeze = freeze and not self.chunk if isinstance(sign, str): sign = maybe_replace_sign(sign) @@ -1307,42 +1244,9 @@ def add_constraints( sign = maybe_replace_signs(as_dataarray(sign)) con = self._constraint_from_lhs(lhs, sign, rhs, coords) - if isinstance(con, CSRConstraint) and freeze: - row_mask = ( - None - if mask is None - else con.active_values( - broadcast_to_coords(mask, con.coords, label="mask") - ).astype(bool) - ) - _check_infinities(con._sign, con._rhs, name, row_mask) - self.check_force_dim_names(con.coords.to_dataset()) - self._check_valid_dim_names(con.coords.to_dataset()) - enforce_no_multiindex(con, context=f"constraint {name!r}") - row_scaling = ( - float(scaling) - if isinstance(scaling, int | float | np.number) - else broadcast_to_coords( - scaling, con.coords, label="constraint scaling" - ) - ) - cindex = self._cCounter - self._cCounter += con.full_size - con = con.assign_labels(cindex, name, row_scaling, row_mask) - return self._soften_added(self.constraints.add(con), penalty) - if isinstance(con, CSRConstraint): - _densify_notice( - "constraint added unfrozen, `freeze=False`", self, explicit=True - ) - dense = ( - con - if isinstance(con, Constraint) - else Constraint.from_dataset(con.data, self) - ) + _check_infinities(con.sign, con.rhs, name) - _check_infinities(dense.sign, dense.rhs, name) - - rhs = dense.rhs + rhs = con.rhs rhs_nan = rhs.isnull() if rhs_nan.any(): rhs = rhs.fillna(0) @@ -1359,34 +1263,34 @@ def add_constraints( ) if mask is not None: - mask = broadcast_to_coords(mask, dense.coords, label="mask").astype(bool) + mask = broadcast_to_coords(mask, con.coords, label="mask").astype(bool) - lhs = dense.lhs + lhs = con.lhs self.check_force_dim_names(lhs) self._check_valid_dim_names(lhs) labels, label_range = self._allocate_constraint_labels(rhs, mask) constraint = Constraint( - lhs, - dense.sign, + LinearExpression._from_store(model_store(lhs._store, self), self), + con.sign, rhs, self, name, labels=labels, scaling=row_scaling, - attrs={**dense.attrs, "label_range": label_range}, + attrs={**con.attrs, "label_range": label_range}, ) enforce_no_multiindex(constraint, context=f"constraint {name!r}") if self.chunk: constraint = constraint.chunk(self.chunk) - return self._soften_added(self.constraints.add(constraint, freeze), penalty) + return self._soften_added(self.constraints.add(constraint), penalty) @staticmethod def _soften_added( - constraint: ConstraintBase, penalty: ConstantLike | None - ) -> ConstraintBase: + constraint: Constraint, penalty: ConstantLike | None + ) -> Constraint: if penalty is not None: constraint.soften(penalty=penalty) return constraint @@ -1400,7 +1304,7 @@ def add_indicator_constraints( rhs: ConstantLike | None = None, name: str | None = None, scaling: ConstantLike = 1, - ) -> ConstraintBase: + ) -> Constraint: """ Add indicator constraints to the model. @@ -1434,7 +1338,7 @@ def add_indicator_constraints( Returns ------- - linopy.constraints.ConstraintBase + linopy.constraints.Constraint The added indicator constraint. """ if not binary_var.attrs.get("binary", False): @@ -1460,7 +1364,7 @@ def add_indicator_constraints( ) labels, label_range = self._allocate_constraint_labels(rhs_da) con = Constraint( - base.lhs, + LinearExpression._from_store(model_store(base.lhs._store, self), self), base.sign, rhs_da, self, @@ -1472,7 +1376,7 @@ def add_indicator_constraints( attrs={**base.attrs, "label_range": label_range}, ) enforce_no_multiindex(con, context=f"constraint {name!r}") - return self.constraints.add(con, freeze=self._sparse) + return self.constraints.add(con) def remove_indicator_constraints(self, name: str) -> None: """ @@ -2445,10 +2349,6 @@ def assign_result( if con.is_indicator: continue start, end = con.range - if isinstance(con, CSRConstraint): - active = dual[start:end][con.active_positions] - con._dual = active * con._scaling / self.objective.scaling - continue coords = {dim: con.coords[dim] for dim in con.coord_dims} values = ( dual[start:end].reshape(con.shape) diff --git a/linopy/persistent/diff.py b/linopy/persistent/diff.py index 52058d433..14d7c5ed3 100644 --- a/linopy/persistent/diff.py +++ b/linopy/persistent/diff.py @@ -9,7 +9,6 @@ import numpy as np from linopy.constants import short_GREATER_EQUAL, short_LESS_EQUAL -from linopy.constraints import Constraint from linopy.persistent.snapshot import ( ContainerConBuffers, ContainerVarBuffers, @@ -25,7 +24,7 @@ from numpy.typing import DTypeLike from linopy.common import ConstraintLabelIndex, VariableLabelIndex - from linopy.constraints import ConstraintBase + from linopy.constraints import Constraint from linopy.model import Model from linopy.variables import Variable @@ -303,7 +302,7 @@ def from_snapshot( return reason for name, con in model.constraints.items(): - skip = same_model and isinstance(con, Constraint) and not con._coef_dirty + skip = same_model and not con._coef_dirty reason = builder.diff_con( name, con, @@ -449,7 +448,7 @@ def diff_var( def diff_con( self, name: str, - con: ConstraintBase, + con: Constraint, base_buf: ContainerConBuffers, base_coords: dict[str, tuple[str | None, np.ndarray]], skip_coef_compare: bool, diff --git a/linopy/persistent/snapshot.py b/linopy/persistent/snapshot.py index ee09390cf..91b576135 100644 --- a/linopy/persistent/snapshot.py +++ b/linopy/persistent/snapshot.py @@ -7,10 +7,8 @@ import numpy as np import pandas as pd -from linopy.constraints import Constraint - if TYPE_CHECKING: - from linopy.constraints import ConstraintBase + from linopy.constraints import Constraint from linopy.model import Model from linopy.variables import Variable, VariableLabelIndex @@ -57,19 +55,19 @@ def _extract_var_buffers(var: Variable) -> ContainerVarBuffers: def _extract_con_buffers( - con: ConstraintBase, var_label_index: VariableLabelIndex + con: Constraint, var_label_index: VariableLabelIndex ) -> ContainerConBuffers: """ Extract flat constraint buffers without copying. - Mutable ``Constraint`` objects build fresh arrays in - ``to_matrix_with_rhs``, so the buffers are exclusively owned. - ``CSRConstraint`` returns its stored ``indptr``/``data`` and a weakly - cached positional ``indices`` array: as long as a snapshot holds it and - neither the constraint nor the variable label index changed, the same - array is returned. The buffers share memory with the constraint, every - mutation path rebinds whole arrays (copy-on-write), and the diff uses - object identity to skip comparisons on untouched containers. + A dense lhs builds fresh arrays in ``to_matrix_with_rhs``, so the + buffers are exclusively owned. A sparse (CSR) lhs returns its stored + ``indptr``/``data`` and a positional ``indices`` array its store caches + by weak reference: as long as a snapshot holds it and neither the lhs + store nor the variable label index changed, the same array is returned. + The buffers share memory with the store, which is immutable, so every + mutation replaces it, and the diff uses object identity to skip + comparisons on untouched containers. """ csr, con_labels, b, sense = con.to_matrix_with_rhs(var_label_index) return ContainerConBuffers( @@ -120,7 +118,7 @@ class ContainerConBuffers: def _coord_snapshot( - obj: Variable | ConstraintBase, + obj: Variable | Constraint, ) -> dict[str, tuple[str | None, np.ndarray]]: """ Snapshot a container's coordinate indexes. @@ -152,8 +150,7 @@ def clear_coef_dirty(model: Model) -> None: a later ``same_model=True`` diff silently skip changed coefficients. """ for con in model.constraints.data.values(): - if isinstance(con, Constraint): - con._coef_dirty = False + con._coef_dirty = False @dataclass diff --git a/linopy/scaling.py b/linopy/scaling.py index c81a83520..c3d741148 100644 --- a/linopy/scaling.py +++ b/linopy/scaling.py @@ -67,13 +67,8 @@ def variable_scaling_lookup(model: Model) -> np.ndarray: def constraint_scaling_lookup(model: Model) -> np.ndarray: """Return constraint scaling indexed by raw constraint label.""" - from linopy.constraints import CSRConstraint - scaling = np.ones(model._cCounter, dtype=float) for con in model.constraints.data.values(): - if isinstance(con, CSRConstraint): - scaling[con.active_labels()] = con._scaling - continue labels = con.labels.values.ravel() _scatter_active(scaling, labels, con.scaling.values.ravel()) return scaling diff --git a/linopy/store.py b/linopy/store.py index 9f0cb7406..fe28a36fd 100644 --- a/linopy/store.py +++ b/linopy/store.py @@ -154,6 +154,13 @@ def cell(self, indices: tuple[Any, ...]) -> tuple[np.ndarray, np.ndarray, float] def terms(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """Cell, variable label and coefficient of the live terms.""" + @abstractmethod + def variable_labels(self) -> np.ndarray: + """ + Variable labels of all stored terms, zero coefficients included, as + :meth:`positional_csr` reads them; ``-1`` may mark an empty slot. + """ + @abstractmethod def solution( self, label_solution: Callable[[np.ndarray], np.ndarray] @@ -200,8 +207,8 @@ def absorbed(self) -> Self: """Drop the terms of absent cells.""" @abstractmethod - def prune(self) -> Self: - """Drop the terms with a zero coefficient.""" + def prune(self, tol: float = 0.0) -> Self: + """Drop the terms with a coefficient of at most ``tol`` in magnitude.""" @abstractmethod def canonical(self) -> Self: @@ -509,6 +516,9 @@ def terms(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: cells = np.broadcast_to(np.arange(self.n_cells)[:, None], keep.shape) return cells[keep], vars_[keep], coeffs[keep] + def variable_labels(self) -> np.ndarray: + return self.ds.vars.values.ravel() + def solution( self, label_solution: Callable[[np.ndarray], np.ndarray] ) -> np.ndarray: @@ -554,8 +564,8 @@ def combine( def absorbed(self) -> DenseStore: return self._derived(absorb_absence(self.ds)) - def prune(self) -> DenseStore: - pruned = pruned_terms(self.ds) + def prune(self, tol: float = 0.0) -> DenseStore: + pruned = pruned_terms(self.ds, tol) return self if pruned is self.ds else self._derived(pruned) def canonical(self) -> DenseStore: @@ -704,15 +714,16 @@ def summed(data: Dataset, dim: DimsLike | None = None) -> Dataset: return assign_fields(ds, const=data.const.sum(dim)) -def pruned_terms(ds: Dataset) -> Dataset: +def pruned_terms(ds: Dataset, tol: float = 0.0) -> Dataset: """ - ``ds`` with the nonzero terms moved to the front of the term dimension - and the all-zero tail cut off; ``ds`` itself if no coefficient is zero. + ``ds`` with the terms of a coefficient above ``tol`` in magnitude moved + to the front of the term dimension and the tail cut off; ``ds`` itself + if there is none to drop. """ coeffs = ds.coeffs.transpose(..., TERM_DIM) vars = ds.vars.transpose(*coeffs.dims, ...) cdata = coeffs.data - mask = cdata != 0 + mask = ~(np.abs(cdata) <= tol) if mask.all(): return ds diff --git a/linopy/testing.py b/linopy/testing.py index fba122bdb..900dfcf1a 100644 --- a/linopy/testing.py +++ b/linopy/testing.py @@ -6,7 +6,7 @@ from xarray.testing import assert_equal from linopy.constants import TERM_DIM -from linopy.constraints import ConstraintBase, _con_unwrap +from linopy.constraints import Constraint, _con_unwrap from linopy.expressions import LinearExpression, QuadraticExpression, _expr_unwrap from linopy.model import Model from linopy.variables import Variable, _var_unwrap @@ -125,7 +125,7 @@ def assert_exprequal( assert_linequal(a, b) -def assert_conequal(a: ConstraintBase, b: ConstraintBase, strict: bool = True) -> None: +def assert_conequal(a: Constraint, b: Constraint, strict: bool = True) -> None: """ Assert that two constraints are equal. @@ -136,12 +136,14 @@ def assert_conequal(a: ConstraintBase, b: ConstraintBase, strict: bool = True) - b: Constraint The second constraint. strict: bool - Whether to compare the constraints strictly. If not, only compare mathematically relevant parts. + Whether to compare the constraints strictly. If not, only compare + mathematically relevant parts: the lhs by :func:`assert_equivalent`, + independently of its store, and the sign and rhs. """ if strict: assert_equal(_con_unwrap(a), _con_unwrap(b)) else: - assert_linequal(a.lhs, b.lhs) + assert_equivalent(a.lhs, b.lhs) assert_equal(a.sign, b.sign) assert_equal(a.rhs, b.rhs) diff --git a/linopy/types.py b/linopy/types.py index ffd9cf686..6ca69efba 100644 --- a/linopy/types.py +++ b/linopy/types.py @@ -10,10 +10,7 @@ from xarray.core.coordinates import DataArrayCoordinates, DatasetCoordinates if TYPE_CHECKING: - from linopy.constraints import ( - AnonymousScalarConstraint, - ConstraintBase, - ) + from linopy.constraints import AnonymousScalarConstraint, Constraint from linopy.expressions import ( LinearExpression, QuadraticExpression, @@ -50,6 +47,6 @@ ExpressionLike: TypeAlias = Union[ "ScalarLinearExpression", "LinearExpression", "QuadraticExpression" ] -ConstraintLike = Union["ConstraintBase", "AnonymousScalarConstraint"] +ConstraintLike = Union["Constraint", "AnonymousScalarConstraint"] LinExprLike = Union["Variable", "LinearExpression"] SideLike = Union[ConstantLike, VariableLike, ExpressionLike] # noqa: UP007 diff --git a/linopy/variables.py b/linopy/variables.py index 08c23e32a..794de240b 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -81,7 +81,7 @@ ) if TYPE_CHECKING: - from linopy.constraints import AnonymousScalarConstraint, ConstraintBase + from linopy.constraints import AnonymousScalarConstraint, Constraint from linopy.expressions import ( GenericExpression, LinearExpression, @@ -579,13 +579,13 @@ def __rsub__(self, other: ConstantLike) -> LinearExpression: except TypeError: return NotImplemented - def __le__(self, other: SideLike) -> ConstraintBase: + def __le__(self, other: SideLike) -> Constraint: return self.to_linexpr().__le__(other) - def __ge__(self, other: SideLike) -> ConstraintBase: + def __ge__(self, other: SideLike) -> Constraint: return self.to_linexpr().__ge__(other) - def __eq__(self, other: SideLike) -> ConstraintBase: # type: ignore[override] + def __eq__(self, other: SideLike) -> Constraint: # type: ignore[override] return self.to_linexpr().__eq__(other) def __gt__(self, other: Any) -> NotImplementedType: @@ -701,7 +701,7 @@ def div( """ return self.to_linexpr().div(other, join=join, fill_value=fill_value) - def le(self, rhs: SideLike, join: JoinOptions | None = None) -> ConstraintBase: + def le(self, rhs: SideLike, join: JoinOptions | None = None) -> Constraint: """ Less than or equal constraint. @@ -717,7 +717,7 @@ def le(self, rhs: SideLike, join: JoinOptions | None = None) -> ConstraintBase: """ return self.to_linexpr().le(rhs, join=join) - def ge(self, rhs: SideLike, join: JoinOptions | None = None) -> ConstraintBase: + def ge(self, rhs: SideLike, join: JoinOptions | None = None) -> Constraint: """ Greater than or equal constraint. @@ -733,7 +733,7 @@ def ge(self, rhs: SideLike, join: JoinOptions | None = None) -> ConstraintBase: """ return self.to_linexpr().ge(rhs, join=join) - def eq(self, rhs: SideLike, join: JoinOptions | None = None) -> ConstraintBase: + def eq(self, rhs: SideLike, join: JoinOptions | None = None) -> Constraint: """ Equality constraint. diff --git a/test/test_assign_coords.py b/test/test_assign_coords.py index 0f71908cc..2ab5013fd 100644 --- a/test/test_assign_coords.py +++ b/test/test_assign_coords.py @@ -15,7 +15,6 @@ from linopy import Model from linopy.constants import Result, Solution, Status -from linopy.constraints import CSRConstraint sns0 = pd.date_range("2026-01-01", periods=3, freq="h", name="snapshot") sns1 = sns0 + pd.Timedelta("1h") @@ -30,7 +29,7 @@ def m() -> Model: coords=[sns0, ["a", "b"]], dims=["snapshot", "spatial"], name="y" ) model.add_constraints(x >= 0, name="dense_c") - model.add_constraints(x >= 1, name="frozen_c", freeze=True) + model.add_constraints(x >= 1, name="sparse_c").to_sparse() model.add_expressions(x * 2 + y.sum(), name="e") model.add_objective(1.0 * x) model.parameters = xr.Dataset( @@ -56,14 +55,13 @@ def test_assign_coords_csr_constraint(m: Model) -> None: CSR-backed constraints get the new labels on their grid, not only in their reconstructed Dataset view. """ - con = m.constraints["frozen_c"] - assert isinstance(con, CSRConstraint) + con = m.constraints["sparse_c"] + assert con.lhs.is_sparse m.assign_coords(snapshot=sns1) - con = m.constraints["frozen_c"] + con = m.constraints["sparse_c"] assert (con.coords["snapshot"].values == sns1.values).all() - # CSR-backed container still reconstructs its Dataset with the new coords assert (con.data.coords["snapshot"].values == sns1.values).all() assert con.labels.sizes["snapshot"] == 3 @@ -235,8 +233,8 @@ def test_assign_coords_rejects_scalar_values(m: Model) -> None: def test_csr_assign_coords_rejects_wrong_length(m: Model) -> None: """The CSR path validates lengths and rejects unknown coordinates too.""" short = pd.date_range("2026-01-01", periods=2, freq="h", name="snapshot") - con = m.constraints["frozen_c"] - assert isinstance(con, CSRConstraint) + con = m.constraints["sparse_c"] + assert con.lhs.is_sparse with pytest.raises(ValueError, match="length"): con._assign_coords(snapshot=short) diff --git a/test/test_common.py b/test/test_common.py index 011d49411..9588e8fb4 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -103,7 +103,7 @@ def test_coords_from_dataset_reads_name_keyed_format() -> None: ], ) def test_coords_from_dataset_loads_multiindex_flat(ds: xr.Dataset) -> None: - """A frozen constraint over a MultiIndex from linopy < 1.0 loads flat.""" + """A sparse (CSR) constraint over a MultiIndex from linopy < 1.0 loads flat.""" (index,), levels = coords_from_dataset(ds, ["d"]) assert index.equals(pd.RangeIndex(3)) and index.name == "d" assert {n: (d, v.tolist()) for n, (d, v) in levels.items()} == { diff --git a/test/test_constraint.py b/test/test_constraint.py index bbd921871..eac8ec5b8 100644 --- a/test/test_constraint.py +++ b/test/test_constraint.py @@ -30,7 +30,6 @@ from linopy.constraints import ( AnonymousScalarConstraint, Constraint, - ConstraintBase, Constraints, ) from linopy.csr import canonical @@ -43,7 +42,7 @@ def m() -> Model: x = m.add_variables(coords=[pd.RangeIndex(10, name="first")], name="x") m.add_variables(coords=[pd.Index([1, 2, 3], name="second")], name="y") m.add_variables(0, 10, name="z") - m.add_constraints(x >= 0, name="c", freeze=True) + m.add_constraints(x >= 0, name="c") return m @@ -58,35 +57,37 @@ def y(m: Model) -> linopy.Variable: @pytest.fixture -def c(m: Model) -> linopy.constraints.ConstraintBase: +def c(m: Model) -> linopy.constraints.Constraint: return m.constraints["c"] @pytest.fixture def mc(m: Model) -> linopy.constraints.Constraint: - return m.constraints["c"].mutable() + return m.constraints["c"] -def test_constraint_repr(c: linopy.constraints.CSRConstraint) -> None: +def test_constraint_repr(c: linopy.constraints.Constraint) -> None: c.__repr__() -def test_constraint_repr_equivalent_to_mutable( - c: linopy.constraints.CSRConstraint, +def test_constraint_repr_equivalent_across_stores( + c: linopy.constraints.Constraint, ) -> None: - """Constraint (CSR-backed) and Constraint repr must be identical.""" - frozen = c.freeze() - assert repr(frozen) == repr(c) + """The repr is the same for a dense and a sparse lhs.""" + dense = repr(c) + assert repr(c.to_sparse()) == dense def test_constraints_repr(m: Model) -> None: m.constraints.__repr__() -def test_add_constraints_freeze(m: Model, x: linopy.Variable) -> None: - c = m.add_constraints(x >= 1, name="frozen_c", freeze=True) - assert isinstance(c, linopy.constraints.CSRConstraint) - assert isinstance(m.constraints["frozen_c"], linopy.constraints.CSRConstraint) +def test_add_constraints_holds_lhs_on_model_store(backend: Model) -> None: + x = backend.add_variables(coords=[pd.RangeIndex(10, name="first")], name="x") + c = backend.add_constraints(1 * x >= 1, name="c") + assert type(c) is Constraint + assert c.lhs.is_sparse == backend.sparse + assert backend.constraints["c"] is c assert c.ncons == 10 @@ -107,7 +108,7 @@ def test_add_constraints_penalty_softens_constraint( assert_linequal(m.objective.expression, expected_objective) -def test_constraint_name(c: linopy.constraints.CSRConstraint) -> None: +def test_constraint_name(c: linopy.constraints.Constraint) -> None: assert c.name == "c" @@ -116,11 +117,10 @@ def test_empty_constraints_repr() -> None: Model().constraints.__repr__() -@pytest.mark.parametrize("freeze", [True, False]) -def test_constraint_handles_empty_rows(freeze: bool) -> None: +def test_constraint_handles_empty_rows(backend: Model) -> None: """An empty constraint group must be accepted and solve cleanly.""" - m = Model() + m = backend x = m.add_variables( lower=0.0, coords=[range(3), range(2)], @@ -128,8 +128,8 @@ def test_constraint_handles_empty_rows(freeze: bool) -> None: name="x", ) empty = x.isel(time=range(1, 1)) - c = m.add_constraints(empty == 0, name="empty", freeze=freeze) - assert isinstance(c, linopy.constraints.ConstraintBase) + c = m.add_constraints(empty == 0, name="empty") + assert isinstance(c, linopy.constraints.Constraint) assert c.size == 0 # Solving a model with only an empty constraint group is also fine. m.add_objective(x.sum()) @@ -143,7 +143,7 @@ def test_cannot_create_constraint_without_variable() -> None: _ = linopy.LinearExpression(12, model) == linopy.LinearExpression(13, model) -def test_constraints_getter(m: Model, c: linopy.constraints.CSRConstraint) -> None: +def test_constraints_getter(m: Model, c: linopy.constraints.Constraint) -> None: assert c.shape == (10,) assert isinstance(m.constraints[["c"]], Constraints) @@ -324,7 +324,7 @@ def test_anonymous_constraint_sel(x: linopy.Variable, y: linopy.Variable) -> Non expr = 10 * x + y con = expr <= 10 assert isinstance(con, Constraint) - assert isinstance(con.sel(first=[1, 2]), ConstraintBase) + assert isinstance(con.sel(first=[1, 2]), Constraint) def test_anonymous_constraint_swap_dims(x: linopy.Variable, y: linopy.Variable) -> None: @@ -333,7 +333,7 @@ def test_anonymous_constraint_swap_dims(x: linopy.Variable, y: linopy.Variable) assert isinstance(con, Constraint) con = con.assign_coords({"third": ("second", con.indexes["second"] + 100)}) con = con.swap_dims({"second": "third"}) - assert isinstance(con, ConstraintBase) + assert isinstance(con, Constraint) assert con.coord_dims == ("first", "third") @@ -341,13 +341,13 @@ def test_anonymous_constraint_loc(x: linopy.Variable, y: linopy.Variable) -> Non expr = 10 * x + y con = expr <= 10 assert isinstance(con, Constraint) - assert isinstance(con.loc[[1, 2]], ConstraintBase) + assert isinstance(con.loc[[1, 2]], Constraint) def test_anonymous_constraint_getitem(x: linopy.Variable, y: linopy.Variable) -> None: expr = 10 * x + y con = expr <= 10 - assert isinstance(con[1], ConstraintBase) + assert isinstance(con[1], Constraint) def test_constraint_from_rule(m: Model, x: linopy.Variable, y: linopy.Variable) -> None: @@ -356,7 +356,7 @@ def bound(m: Model, i: int, j: int) -> AnonymousScalarConstraint: coords = [x.coords["first"], y.coords["second"]] con = Constraint.from_rule(m, bound, coords) - assert isinstance(con, ConstraintBase) + assert isinstance(con, Constraint) assert con.lhs.nterm == 2 repr(con) # test repr @@ -371,7 +371,7 @@ def bound(m: Model, i: int, j: int) -> AnonymousScalarConstraint | None: coords = [x.coords["first"], y.coords["second"]] con = Constraint.from_rule(m, bound, coords) - assert isinstance(con, ConstraintBase) + assert isinstance(con, Constraint) assert isinstance(con.lhs.vars, xr.DataArray) assert con.lhs.nterm == 2 assert (con.lhs.vars.loc[0, :] == -1).all() @@ -389,11 +389,11 @@ def test_constraint_coeffs_getter(mc: linopy.constraints.Constraint) -> None: assert (mc.coeffs == 1).all() -def test_constraint_sign_getter(c: linopy.constraints.CSRConstraint) -> None: +def test_constraint_sign_getter(c: linopy.constraints.Constraint) -> None: assert (c.sign == GREATER_EQUAL).all() -def test_constraint_rhs_getter(c: linopy.constraints.CSRConstraint) -> None: +def test_constraint_rhs_getter(c: linopy.constraints.Constraint) -> None: assert (c.rhs == 0).all() @@ -625,36 +625,42 @@ def test_constraint_update_rhs_broadcasts_missing_dim() -> None: assert (con.rhs.sel(i=1) == 2.0).all() -def test_constraint_labels_setter_invalid(c: linopy.constraints.CSRConstraint) -> None: - # Test that assigning labels raises AttributeError (Constraint is frozen) +def test_constraint_labels_setter_invalid(c: linopy.constraints.Constraint) -> None: with pytest.raises(AttributeError): c.labels = c.labels # type: ignore -def test_constraint_sel(c: linopy.constraints.CSRConstraint) -> None: - assert isinstance(c.mutable().sel(first=[1, 2]), ConstraintBase) - assert isinstance(c.mutable().isel(first=[1, 2]), ConstraintBase) +def test_constraint_sel(c: linopy.constraints.Constraint) -> None: + assert isinstance(c.sel(first=[1, 2]), Constraint) + assert isinstance(c.isel(first=[1, 2]), Constraint) -def test_constraint_flat_removed(c: linopy.constraints.CSRConstraint) -> None: +def test_constraint_flat_removed(c: linopy.constraints.Constraint) -> None: with pytest.raises(AttributeError): c.flat # type: ignore[attr-defined] -def test_iterate_slices(mc: linopy.constraints.Constraint) -> None: - for i in mc.iterate_slices(slice_size=2): - assert isinstance(i, ConstraintBase) - assert mc.coord_dims == i.coord_dims +@pytest.mark.parametrize("sparse", [False, True], ids=["dense", "sparse"]) +def test_iterate_slices(mc: linopy.constraints.Constraint, sparse: bool) -> None: + con = mc.to_sparse() if sparse else mc + slices = list(con.iterate_slices(slice_size=2)) + assert len(slices) > 1 + for i in slices: + assert isinstance(i, Constraint) + assert con.coord_dims == i.coord_dims + assert i.lhs.is_sparse == sparse + sliced = pl.concat([i.to_polars() for i in slices]) + assert sliced.equals(con.to_polars()) -def test_constraint_to_polars(c: linopy.constraints.CSRConstraint) -> None: +def test_constraint_to_polars(c: linopy.constraints.Constraint) -> None: assert isinstance(c.to_polars(), pl.DataFrame) def test_constraint_to_polars_mixed_signs(m: Model, x: linopy.Variable) -> None: """Test to_polars when a constraint has mixed sign values across dims.""" # Use Constraint so sign data can be patched - con = m.add_constraints(x >= 0, name="mixed", freeze=False) + con = m.add_constraints(x >= 0, name="mixed") # Replace sign data with mixed signs across the first dimension n = con.sizes["first"] signs = np.array(["<=" if i % 2 == 0 else ">=" for i in range(n)]) @@ -667,15 +673,18 @@ def test_constraint_to_polars_mixed_signs(m: Model, x: linopy.Variable) -> None: def test_constraint_assignment_with_anonymous_constraints( m: Model, x: linopy.Variable, y: linopy.Variable ) -> None: - m.add_constraints(x + y == 0, name="c2", freeze=False) + m.add_constraints(x + y == 0, name="c2") assert m.constraints["c2"].vars.notnull().all() assert m.constraints["c2"].coeffs.notnull().all() +@pytest.mark.parametrize("sparse", [False, True], ids=["dense", "sparse"]) def test_constraint_assignment_sanitize_zeros( - m: Model, x: linopy.Variable, y: linopy.Variable + m: Model, x: linopy.Variable, y: linopy.Variable, sparse: bool ) -> None: - m.add_constraints(0 * x + y == 0, name="c2", freeze=True) + con = m.add_constraints(0 * x + y == 0, name="c2") + if sparse: + con.to_sparse() m.constraints.sanitize_zeros() c2 = m.constraints["c2"] assert c2.nterm == 1 @@ -890,99 +899,90 @@ def test_constraints_equalities(m: Model) -> None: assert isinstance(m.constraints.equalities, Constraints) -def test_freeze_mutable_roundtrip(m: Model) -> None: - frozen = m.constraints["c"] - assert isinstance(frozen, linopy.constraints.CSRConstraint) - mc = frozen.mutable() - assert isinstance(mc, Constraint) - 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) - np.testing.assert_array_equal(frozen._csr.toarray(), refrozen._csr.toarray()) - np.testing.assert_array_equal(frozen.active_labels(), refrozen.active_labels()) - - -def test_frozen_coeff_dtype_preserved() -> None: +@pytest.mark.parametrize("masked", [False, True], ids=["all", "masked"]) +def test_to_sparse_to_dense_roundtrip(masked: bool) -> None: + m = Model() + x = m.add_variables(coords=[pd.RangeIndex(5, name="i")], name="x") + mask = xr.DataArray([True, False, True, False, True], dims=["i"]) + con = m.add_constraints((x.where(mask) if masked else 1 * x) >= 0, name="c") + before = con.data + assert con.to_sparse() is con + assert con.lhs.is_sparse + assert con.ncons == (3 if masked else 5) + assert_equal(con.labels, before.labels) + assert_equal(con.rhs, before.rhs) + assert_equal(con.sign, before.sign) + assert not con.to_dense().lhs.is_sparse + assert_equal(con.data, before) + + +def test_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 + con = m.add_constraints(coeff * x >= 1, name="c") + assert con.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) - csr, _ = frozen.to_matrix(m.variables.label_index) +def test_sparse_to_matrix_has_active_variable_columns( + m: Model, c: linopy.constraints.Constraint +) -> None: + csr, _ = c.to_sparse().to_matrix(m.variables.label_index) assert csr.shape[1] == m.variables.label_index.n_active_vars -def _model_with_frozen_and_mutation(freeze: bool, mutation: str) -> Model: - m = Model() +def _model_with_mutation(sparse: bool, mutation: str) -> Model: + m = Model(sparse=sparse) i = pd.RangeIndex(3, name="i") m.add_variables(coords=[i], name="a") b = m.add_variables(coords=[i], name="b") - m.add_constraints(2 * b >= 1, name="c1", freeze=freeze) + m.add_constraints(2 * b >= 1, name="c1") d = m.add_variables(coords=[pd.RangeIndex(2, name="i")], name="d") - m.add_constraints(d <= 5, name="c2", freeze=freeze) + m.add_constraints(d <= 5, name="c2") if mutation == "remove": m.remove_variables("a") return m @pytest.mark.parametrize("mutation", ["add", "remove"]) -def test_frozen_matrices_after_variable_mutation(mutation: str) -> None: - frozen = _model_with_frozen_and_mutation(True, mutation) - mutable = _model_with_frozen_and_mutation(False, mutation) - a_frozen, a_mutable = frozen.matrices.A, mutable.matrices.A - assert a_frozen is not None and a_mutable is not None - np.testing.assert_array_equal(a_frozen.toarray(), a_mutable.toarray()) - np.testing.assert_array_equal(frozen.matrices.b, mutable.matrices.b) - np.testing.assert_array_equal(frozen.matrices.clabels, mutable.matrices.clabels) - np.testing.assert_array_equal(frozen.matrices.vlabels, mutable.matrices.vlabels) - - -@pytest.mark.parametrize("freeze", [True, False]) -def test_constraint_removed_with_referenced_variable(freeze: bool) -> None: - m = Model() +def test_sparse_matrices_after_variable_mutation(mutation: str) -> None: + sparse = _model_with_mutation(True, mutation) + dense = _model_with_mutation(False, mutation) + a_sparse, a_dense = sparse.matrices.A, dense.matrices.A + assert a_sparse is not None and a_dense is not None + np.testing.assert_array_equal(a_sparse.toarray(), a_dense.toarray()) + np.testing.assert_array_equal(sparse.matrices.b, dense.matrices.b) + np.testing.assert_array_equal(sparse.matrices.clabels, dense.matrices.clabels) + np.testing.assert_array_equal(sparse.matrices.vlabels, dense.matrices.vlabels) + + +@pytest.mark.parametrize("coeff", [1.0, 0.0], ids=["term", "zero-term"]) +def test_constraint_removed_with_referenced_variable( + backend: Model, coeff: float +) -> None: + m = backend i = pd.RangeIndex(3, name="i") a = m.add_variables(coords=[i], name="a") b = m.add_variables(coords=[i], name="b") - m.add_constraints(a + b >= 1, name="c", freeze=freeze) + m.add_constraints(coeff * a + b >= 1, name="c") + m.add_constraints(1 * b >= 0, name="d") with pytest.warns(UserWarning, match="also removes constraints"): m.remove_variables("a") assert "c" not in m.constraints + assert m.matrices.A is not None -def test_freeze_mutable_roundtrip_with_masking() -> None: - m = Model() - x = m.add_variables(coords=[pd.RangeIndex(5, name="i")], name="x") - mask = xr.DataArray([True, False, True, False, True], dims=["i"]) - m.add_constraints(x.where(mask) >= 0, name="c", freeze=True) - frozen = m.constraints["c"] - assert isinstance(frozen, linopy.constraints.CSRConstraint) - mc = frozen.mutable() - 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_dense_mixed_signs() -> None: +def test_to_sparse_keeps_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.update(sign=xr.DataArray(["<=", ">=", "<="], dims=["i"])) - frozen = linopy.constraints.CSRConstraint.from_dense(mc) - assert isinstance(frozen._sign, np.ndarray) - assert list(frozen._sign) == ["<=", ">=", "<="] - assert_equal(frozen.sign, mc.sign) + con = m.add_constraints(x >= 0, name="mixed") + con.update(sign=xr.DataArray(["<=", ">=", "<="], dims=["i"])) + sign = con.sign + con.to_sparse() + assert_equal(con.sign, sign) + _, _, _, sense = con.to_matrix_with_rhs(m.variables.label_index) + assert list(sense) == ["<", ">", "<"] def test_variable_label_index(m: Model) -> None: @@ -1003,25 +1003,18 @@ def test_variable_label_index_invalidation(m: Model) -> None: assert len(li.vlabels) > len(old_vlabels) -def test_to_matrix_with_rhs(m: Model) -> None: +@pytest.mark.parametrize("sparse", [False, True], ids=["dense", "sparse"]) +def test_to_matrix_with_rhs(m: Model, sparse: bool) -> None: c = m.constraints["c"] - assert isinstance(c, linopy.constraints.CSRConstraint) + if sparse: + c.to_sparse() li = m.variables.label_index csr, con_labels, b, sense = c.to_matrix_with_rhs(li) assert csr.shape[0] == len(con_labels) assert csr.shape[0] == len(b) assert csr.shape[0] == len(sense) assert all(s in ("<", ">", "=") for s in sense) - np.testing.assert_array_equal(b, c._rhs) - - -def test_to_matrix_with_rhs_mutable(m: Model) -> None: - mc = m.constraints["c"].mutable() - li = m.variables.label_index - csr, con_labels, b, sense = mc.to_matrix_with_rhs(li) - assert csr.shape[0] == len(con_labels) - assert csr.shape[0] == len(b) - assert csr.shape[0] == len(sense) + np.testing.assert_array_equal(b, c.rhs.values.ravel()[c.active_row_mask()]) def test_constraint_repr_shows_variable_names(m: Model) -> None: @@ -1030,8 +1023,8 @@ def test_constraint_repr_shows_variable_names(m: Model) -> None: assert "x" in r -def test_freeze_mixed_signs_from_rule() -> None: - m = Model() +def test_sparse_mixed_signs_from_rule() -> None: + m = Model(sparse=True) x = m.add_variables(coords=[pd.RangeIndex(4, name="i")], name="x") coords = [pd.RangeIndex(4, name="i")] @@ -1040,34 +1033,10 @@ def bound(m: Model, i: int) -> AnonymousScalarConstraint: return x.at[i] >= i return x.at[i] == 0.0 - con = m.add_constraints(bound, coords=coords, name="mixed_rule", freeze=True) - assert isinstance(con, linopy.constraints.CSRConstraint) - assert isinstance(con._sign, np.ndarray) + con = m.add_constraints(bound, coords=coords, name="mixed_rule") + assert con.lhs.is_sparse assert con.ncons == 4 - expected_signs = ["=", ">=", "=", ">="] - assert list(con._sign) == expected_signs - np.testing.assert_array_equal(con.sign.values, expected_signs) - - -def test_frozen_lhs_setter_raises() -> None: - m = Model() - time = pd.RangeIndex(5, name="t") - x = m.add_variables(lower=0, coords=[time], name="x") - y = m.add_variables(lower=0, coords=[time], name="y") - con = m.add_constraints(x >= 0, name="c", freeze=True) - assert isinstance(con, linopy.constraints.CSRConstraint) - with pytest.raises(AttributeError, match="read-only"): - con.lhs = 3 * x + 2 * y - - -def test_frozen_rhs_setter_raises() -> None: - m = Model() - time = pd.RangeIndex(5, name="t") - x = m.add_variables(lower=0, coords=[time], name="x") - con = m.add_constraints(x >= 0, name="c", freeze=True) - assert isinstance(con, linopy.constraints.CSRConstraint) - with pytest.raises(AttributeError, match="read-only"): - con.rhs = 10 + np.testing.assert_array_equal(con.sign.values, ["=", ">=", "=", ">="]) def test_mixed_sign_to_matrix_with_rhs() -> None: @@ -1087,20 +1056,19 @@ def bound(m: Model, i: int) -> AnonymousScalarConstraint: assert list(sense) == ["=", ">", "=", ">"] -def test_mixed_sign_sanitize_infinities() -> None: - m = Model() +@pytest.mark.parametrize("sparse", [False, True], ids=["dense", "sparse"]) +def test_mixed_sign_sanitize_infinities(sparse: bool) -> None: + m = Model(sparse=sparse) x = m.add_variables(coords=[pd.RangeIndex(4, name="i")], name="x") - m.add_constraints(x >= 0, name="c", freeze=False) - mc = m.constraints["c"] - assert isinstance(mc, Constraint) + mc = m.add_constraints(x >= 0, name="c") mc.update( sign=xr.DataArray(["<=", ">=", "<=", ">="], dims=["i"]), rhs=xr.DataArray([np.inf, -np.inf, 1.0, 2.0], dims=["i"]), ) - frozen = mc.freeze() - frozen.sanitize_infinities() - assert frozen.ncons == 2 - np.testing.assert_array_equal(frozen._rhs, [1.0, 2.0]) + mc.sanitize_infinities() + assert mc.lhs.is_sparse == sparse + assert mc.ncons == 2 + np.testing.assert_array_equal(mc.rhs.values[mc.active_row_mask()], [1.0, 2.0]) def test_mixed_sign_repr() -> None: @@ -1155,19 +1123,18 @@ def test_constraint_soften_returns_slack_for_eq(m: Model, y: linopy.Variable) -> assert isinstance(y_slack.negative, linopy.Variable) -def test_constraint_soften_raises_on_detached_mutable_constraint( - m: Model, x: linopy.Variable, mc: linopy.Constraint +def test_constraint_soften_raises_on_detached_constraint( + m: Model, x: linopy.Variable, c: linopy.Constraint ) -> None: """ - `.mutable()` on a frozen constraint returns a detached copy that is not - registered in `model.constraints`; softening it would silently fail to - affect the actual model, so `soften` must raise instead. + A selection returns a detached copy that is not registered in + `model.constraints`; softening it would silently fail to affect the + actual model, so `soften` must raise instead. """ m.add_objective(x.sum(), sense="min") with pytest.raises(ValueError, match="not the constraint registered"): - # mc := m.constraints["c"].mutable(), detached from the model - mc.soften(penalty=10) + c.isel(first=slice(None)).soften(penalty=10) def test_constraint_soften_twice_raises(m: Model, y: linopy.Variable) -> None: @@ -1379,7 +1346,7 @@ def test_constraint_slack_matches_returned_slack_for_eq( def test_constraint_lhs_is_new_wrapper_on_held_store( m: Model, x: linopy.Variable ) -> None: - con = m.add_constraints(x >= 0, name="d", freeze=False) + con = m.add_constraints(x >= 0, name="d") assert con.lhs is not con.lhs assert con.lhs._store is con.lhs._store @@ -1402,7 +1369,7 @@ def test_constraint_canonicalizes_sparse_lhs() -> None: def test_constraint_data_roundtrips_through_from_dataset( m: Model, x: linopy.Variable ) -> None: - con = m.add_constraints(2 * x >= 1, name="d", freeze=False) + con = m.add_constraints(2 * x >= 1, name="d") con.dual = xr.zeros_like(con.rhs) restored = Constraint.from_dataset(con.data, m, con.name) xr.testing.assert_identical(restored.data, con.data) @@ -1410,7 +1377,7 @@ def test_constraint_data_roundtrips_through_from_dataset( def test_from_dataset_rejects_unknown_field(m: Model, x: linopy.Variable) -> None: - con = m.add_constraints(x >= 1, name="d", freeze=False) + con = m.add_constraints(x >= 1, name="d") with pytest.raises(ValueError, match="unknown constraint fields"): Constraint.from_dataset(con.data.assign(foo=con.rhs), m) @@ -1447,13 +1414,12 @@ def test_constraint_rejects_field_off_lhs_grid( Constraint(1 * x, "<=", rhs, m) -@pytest.mark.parametrize("freeze", [True, False]) -def test_to_matrix_names_constraint_with_removed_variable(freeze: bool) -> None: - m = Model() +def test_to_matrix_names_constraint_with_removed_variable(backend: Model) -> None: + m = backend i = pd.RangeIndex(3, name="i") a = m.add_variables(coords=[i], name="a") b = m.add_variables(coords=[i], name="b") - con = m.add_constraints(a + b >= 1, name="c", freeze=freeze) + con = m.add_constraints(a + b >= 1, name="c") m.variables.remove("a") with pytest.raises(ValueError, match="Constraint 'c' references variable"): con.to_matrix(m.variables.label_index) diff --git a/test/test_csr.py b/test/test_csr.py index 8d5cc4a56..17b63360f 100644 --- a/test/test_csr.py +++ b/test/test_csr.py @@ -1,6 +1,6 @@ """ Tests for sparse groupby-sum (linopy.csr): type stability, transparent -materialization, and direct CSR realization under freeze. +materialization, and constraints on a sparse lhs. """ from __future__ import annotations @@ -16,7 +16,6 @@ import numpy as np import pandas as pd -import polars as pl import pytest import scipy.sparse import xarray as xr @@ -34,7 +33,7 @@ import linopy from linopy import LinearExpression, Model, QuadraticExpression, Variable from linopy.constants import TERM_DIM -from linopy.constraints import Constraint, ConstraintBase, CSRConstraint +from linopy.constraints import Constraint from linopy.csr import CSRStore, Grid, canonical, column_compacted_matmul from linopy.testing import ( assert_conequal, @@ -103,31 +102,6 @@ def base_model( return Case(m, gen_p, flow, flow_t, eff, gbus, bus0, bus1, load) -def twin_models(**kwargs: Any) -> tuple[Case, Case]: - """A dense model and its sparse twin, built identically.""" - return base_model(**kwargs), base_model(sparse=True, **kwargs) - - -def canon(df: pl.DataFrame) -> pl.DataFrame: - return ( - df.group_by(["labels", "vars"]) - .agg(pl.col("coeffs").sum(), pl.col("sign").first(), pl.col("rhs").first()) - .filter(pl.col("coeffs") != 0) - .sort(["labels", "vars"]) - ) - - -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"]) - assert np.allclose(d1["coeffs"], d2["coeffs"]) - assert (d1["sign"] == d2["sign"]).all() - assert np.allclose(d1["rhs"], d2["rhs"]) - labels = con1.labels.values.ravel() - assert np.array_equal(np.sort(labels[labels != -1]), np.sort(con2.active_labels())) - - def sparse_store(e: LinearExpression) -> CSRStore: assert isinstance(e._store, CSRStore) return e._store @@ -143,7 +117,7 @@ def reindexed_balance(c: Case) -> LinearExpression: def test_csr_is_plain_linear_expression_and_materializes_equivalently() -> None: - c1, c2 = twin_models() + c1, c2 = base_model(), base_model(sparse=True) sparse = c2.gen_sum() assert type(sparse) is LinearExpression assert sparse.is_sparse @@ -151,7 +125,7 @@ def test_csr_is_plain_linear_expression_and_materializes_equivalently() -> None: def test_csr_composition_materializes_equivalently() -> None: - c1, c2 = twin_models() + c1, c2 = base_model(), base_model(sparse=True) sparse = c2.balance_lhs() assert type(sparse) is LinearExpression assert sparse.is_sparse @@ -159,7 +133,7 @@ def test_csr_composition_materializes_equivalently() -> None: def test_scalar_ops_stay_csr() -> None: - c1, c2 = twin_models() + c1, c2 = base_model(), base_model(sparse=True) sparse = -2.0 * c2.gen_sum() assert sparse.is_sparse assert_linequal(sparse, -2.0 * c1.gen_sum()) @@ -171,13 +145,13 @@ def test_zero_coefficient_rows_stay_active(sparse: bool) -> None: lhs = (0.0 * c.gen_p).groupby(c.gbus).sum() lhs = lhs + (0.0 * c.flow).groupby(c.bus0).sum() nterm = lhs.nterm - con = c.m.add_constraints(lhs == c.load, name="bal", freeze=True) + con = c.m.add_constraints(lhs == c.load, name="bal") assert len(con.active_labels()) == c.load.size assert lhs.nterm == nterm def test_merge_keeps_absent_cell_absent() -> None: - c1, c2 = twin_models() + c1, c2 = base_model(), base_model(sparse=True) mask = xr.DataArray(np.arange(len(c1.load.bus)) % 2 == 0, coords=[c1.load.bus]) flow1, flow2 = ( (1.0 * c.flow).groupby(c.bus0).sum(use_fallback=True).where(mask) @@ -191,7 +165,7 @@ def test_merge_keeps_absent_cell_absent() -> None: ) con = c2.m.add_constraints(tot >= c2.load, name="bal") - assert isinstance(con, CSRConstraint) + assert con.lhs.is_sparse assert con.ncons == int(mask.sum()) * c2.load.sizes["snapshot"] @@ -238,7 +212,7 @@ def test_namelist_sparse_matches_dense(observed: bool, member_first: bool) -> No @pytest.mark.parametrize("as_namelist", [False, True]) def test_single_key_sparse_ignores_observed(as_namelist: bool) -> None: - c1, c2 = twin_models() + c1, c2 = base_model(), base_model(sparse=True) dense, sparse = ( (c.eff * c.gen_p).assign_coords(bus=("gen", c.gbus.to_numpy())) for c in (c1, c2) @@ -271,18 +245,16 @@ def test_dataframe_grouper_sparse_stays_compact() -> None: assert_linequal(sparse, dense.groupby(df).sum()) -def test_namelist_sparse_observed_freezes_compact() -> None: +def test_namelist_sparse_observed_constraint_compact() -> None: m1, e1, rhs = keyed_model() m2, e2, _ = keyed_model(sparse=True) dense = m1.add_constraints( - e1.groupby(["period", "season"]).sum(observed=True) == rhs, - name="c", - freeze=True, + e1.groupby(["period", "season"]).sum(observed=True) == rhs, name="c" ) sparse = m2.add_constraints( e2.groupby(["period", "season"]).sum(observed=True) == rhs, name="c" ) - assert isinstance(sparse, CSRConstraint) + assert sparse.lhs.is_sparse assert sparse.ncons == rhs.size assert_conequal(dense, sparse, strict=False) @@ -291,14 +263,12 @@ def test_namelist_sparse_grid_absent_cells_inactive() -> None: m1, e1, _ = keyed_model() m2, e2, _ = keyed_model(sparse=True) keys = ["period", "season"] - dense = m1.add_constraints(e1.groupby(keys).sum() == 0, name="c", freeze=True) + dense = m1.add_constraints(e1.groupby(keys).sum() == 0, name="c") sparse = m2.add_constraints(e2.groupby(keys).sum() == 0, name="c") - assert isinstance(sparse, CSRConstraint) + assert sparse.lhs.is_sparse assert sparse.ncons == 4 * 2 assert_conequal(dense, sparse, strict=False) - assert np.array_equal( - np.sort(dense.active_labels()), np.sort(sparse.active_labels()) - ) + np.testing.assert_array_equal(dense.active_labels(), sparse.active_labels()) def test_namelist_sparse_observed_keeps_aux_coords_through_merge() -> None: @@ -354,23 +324,16 @@ def test_nan_multikey_grouper_raises_eagerly() -> None: expr.groupby(["period", "season"]).sum() -def test_freeze_realizes_csr_without_dense_rectangle() -> None: - c1, c2 = twin_models() +@pytest.mark.parametrize("to_dense", [False, True], ids=["sparse", "to-dense"]) +def test_sparse_constraint_matches_dense(to_dense: bool) -> None: + c1, c2 = base_model(), base_model(sparse=True) con1 = c1.m.add_constraints(c1.balance_lhs() == c1.load, name="bal") con2 = c2.m.add_constraints(c2.balance_lhs() == c2.load, name="bal") - - assert isinstance(con2, CSRConstraint) - assert_frozen_equal(con1, con2) - - -def test_freeze_false_falls_back_to_identical_dense_constraint() -> None: - """The fallback is canonical-form, so compare mathematically (strict=False).""" - c1, c2 = twin_models() - con1 = c1.m.add_constraints(c1.balance_lhs() == c1.load, name="bal") - con2 = c2.m.add_constraints(c2.balance_lhs() == c2.load, name="bal", freeze=False) - assert isinstance(con2, Constraint) + assert con2.lhs.is_sparse + if to_dense: + assert not con2.to_dense().lhs.is_sparse assert_conequal(con1, con2, strict=False) - assert np.array_equal(con1.labels.values, con2.labels.values) + xr.testing.assert_equal(con1.labels, con2.labels) def gen_mask(c: Case) -> xr.DataArray: @@ -391,74 +354,66 @@ def mixed_signs(c: Case) -> xr.DataArray: return xr.DataArray(signs, coords=[c.gbus.index]) -DENSE_FREEZES: dict[str, Callable[[Case], Constraint]] = { - "one-term": lambda c: c.m.add_constraints(c.gen_p <= 1, freeze=False), - "one-term-masked": lambda c: c.m.add_constraints( - c.gen_p <= 1, mask=gen_mask(c), freeze=False - ), +CONSTRAINT_BUILDS: dict[str, Callable[[Case], Constraint]] = { + "one-term": lambda c: c.m.add_constraints(c.gen_p <= 1), + "one-term-masked": lambda c: c.m.add_constraints(c.gen_p <= 1, mask=gen_mask(c)), "one-term-all-masked": lambda c: c.m.add_constraints( - c.gen_p <= 1, mask=xr.zeros_like(gen_mask(c)), freeze=False - ), - "one-term-absent": lambda c: c.m.add_constraints( - c.gen_p.shift(snapshot=1) <= 1, freeze=False + c.gen_p <= 1, mask=xr.zeros_like(gen_mask(c)) ), + "one-term-absent": lambda c: c.m.add_constraints(c.gen_p.shift(snapshot=1) <= 1), "one-term-zeros": lambda c: c.m.add_constraints( - c.eff.where(gen_mask(c), 0) * c.gen_p <= 1, freeze=False + c.eff.where(gen_mask(c), 0) * c.gen_p <= 1 ), "terms-absent": lambda c: c.m.add_constraints( - c.gen_p - c.gen_p.shift(snapshot=1) <= 1, freeze=False + c.gen_p - c.gen_p.shift(snapshot=1) <= 1 ), "terms-duplicates": lambda c: c.m.add_constraints( c.gen_p.shift(snapshot=1) + 2 * c.gen_p + c.gen_p.shift(snapshot=1) >= 0, - freeze=False, ), "terms-zeros-masked": lambda c: c.m.add_constraints( - 0 * c.gen_p + c.gen_p.shift(snapshot=1) >= 0, mask=gen_mask(c), freeze=False + 0 * c.gen_p + c.gen_p.shift(snapshot=1) >= 0, mask=gen_mask(c) ), "terms-cancelling": lambda c: c.m.add_constraints( - c.gen_p - c.gen_p + c.gen_p.shift(snapshot=1) >= 0, freeze=False + c.gen_p - c.gen_p + c.gen_p.shift(snapshot=1) >= 0 ), "terms-empty-trailing-rows": lambda c: c.m.add_constraints( all_but_last_gen(c) * (c.gen_p + c.gen_p.shift(snapshot=1)) >= 0, - freeze=False, ), "terms-sparse-strided": lambda c: c.m.add_constraints( - (c.gen_p * gen_line_eye(c)).sum("gen") >= 0, freeze=False + (c.gen_p * gen_line_eye(c)).sum("gen") >= 0 ), "terms-dense-strided": lambda c: c.m.add_constraints( - (c.gen_p * xr.ones_like(gen_line_eye(c))).sum("gen") >= 0, freeze=False + (c.gen_p * xr.ones_like(gen_line_eye(c))).sum("gen") >= 0 ), "terms-mixed-signs": lambda c: c.m.add_constraints( - c.gen_p - c.gen_p.shift(snapshot=1), mixed_signs(c), 1, freeze=False + c.gen_p - c.gen_p.shift(snapshot=1), mixed_signs(c), 1 ), - "grouped": lambda c: c.m.add_constraints(c.balance_lhs() == c.load, freeze=False), + "grouped": lambda c: c.m.add_constraints(c.balance_lhs() == c.load), } -@pytest.mark.parametrize("build", list(DENSE_FREEZES)) -def test_freeze_of_dense_constraint_is_canonical_and_equal(build: str) -> None: - con = DENSE_FREEZES[build](base_model()) - frozen = con.freeze() - assert_frozen_equal(con, frozen) - np.testing.assert_array_equal(frozen.active_labels(), con.active_labels()) - csr = frozen._csr - ref = scipy.sparse.csr_array( - (csr.data.copy(), csr.indices.copy(), csr.indptr.copy()), shape=csr.shape - ) - ref.sum_duplicates() - ref.eliminate_zeros() - assert csr.has_canonical_format - for got, want in zip( - (csr.indptr, csr.indices, csr.data), (ref.indptr, ref.indices, ref.data) - ): - np.testing.assert_array_equal(got, want) +@pytest.mark.parametrize("convert", [False, True], ids=["sparse-model", "to-sparse"]) +@pytest.mark.parametrize("build", list(CONSTRAINT_BUILDS)) +def test_sparse_constraint_is_canonical_and_matches_dense( + build: str, convert: bool +) -> None: + dense = CONSTRAINT_BUILDS[build](base_model()) + if convert: + sparse = CONSTRAINT_BUILDS[build](base_model()).to_sparse() + else: + sparse = CONSTRAINT_BUILDS[build](base_model(sparse=True)) + assert sparse.lhs.is_sparse + assert_conequal(dense, sparse, strict=False) + np.testing.assert_array_equal(sparse.active_labels(), dense.active_labels()) + csr = sparse_store(sparse.lhs).csr + assert canonical(csr) is csr -def test_to_constraint_on_csr_lhs_is_unassigned_csr_constraint() -> None: - c1, c2 = twin_models() +def test_to_constraint_on_csr_lhs_is_unassigned_sparse_constraint() -> None: + c1, c2 = base_model(), base_model(sparse=True) dense = c1.balance_lhs() == c1.load con = c2.balance_lhs() == c2.load - assert isinstance(con, CSRConstraint) + assert con.lhs.is_sparse assert not con.is_assigned assert con.type == "Constraint (unassigned)" assert "None" not in repr(con) @@ -472,42 +427,35 @@ def test_to_constraint_on_csr_lhs_is_unassigned_csr_constraint() -> None: con1 = c1.m.add_constraints(dense, name="bal") con2 = c2.m.add_constraints(con, name="bal") - assert isinstance(con2, CSRConstraint) + assert con2.lhs.is_sparse assert con2.is_assigned - assert np.array_equal( - np.sort(con1.labels.values.ravel()), np.sort(con2.active_labels()) - ) assert_conequal(con1, con2, strict=False) + xr.testing.assert_equal(con1.labels, con2.labels) -@pytest.mark.parametrize("freeze", [False, True]) -def test_group_without_terms_matches_dense_labels(freeze: bool) -> None: +def test_group_without_terms_matches_dense_labels() -> None: - def build(sparse: bool) -> ConstraintBase: + def build(sparse: bool) -> Constraint: c = base_model(sparse=sparse) gens = c.gen_p.indexes["gen"] gen_p = c.gen_p.where(xr.DataArray(gens != "gen7", coords=[gens])) lhs = (c.eff * gen_p).groupby(c.gbus).sum() - return c.m.add_constraints(lhs == c.load, name="bal", freeze=freeze) + return c.m.add_constraints(lhs == c.load, name="bal") dense, sparse = build(False), build(True) - assert isinstance(sparse, CSRConstraint if freeze else Constraint) + assert sparse.lhs.is_sparse assert_conequal(dense, sparse, strict=False) - np.testing.assert_array_equal(dense.labels.values, sparse.labels.values) + xr.testing.assert_equal(dense.labels, sparse.labels) @pytest.mark.parametrize("masked", [False, True], ids=["unmasked", "masked"]) @pytest.mark.parametrize("sparse", [True, False], ids=["sparse", "dense"]) -def test_frozen_invalid_infinite_rhs_raises(sparse: bool, masked: bool) -> None: +def test_invalid_infinite_rhs_raises(sparse: bool, masked: bool) -> None: c = base_model(sparse=sparse) valid = c.load.bus != "bus0" rhs = c.load.where(valid, -np.inf) mask = valid if masked else None - add = partial(c.m.add_constraints, name="bal", freeze=True, mask=mask) - if masked and sparse: - con = add(c.balance_lhs() <= rhs) - assert con.ncons == int(valid.sum()) * c.load.sizes["snapshot"] - return + add = partial(c.m.add_constraints, name="bal", mask=mask) with pytest.raises(ValueError, match="incorrect infinite values"): add(c.balance_lhs() <= rhs) @@ -524,7 +472,7 @@ def test_frozen_invalid_infinite_rhs_raises(sparse: bool, masked: bool) -> None: @pytest.mark.parametrize("masked", [False, True], ids=["all", "masked"]) @pytest.mark.parametrize("scaling", list(ROW_SCALINGS)) @pytest.mark.parametrize("sparse", [True, False], ids=["sparse", "dense"]) -def test_frozen_constraint_applies_row_scaling( +def test_constraint_applies_row_scaling( sparse: bool, scaling: str, masked: bool ) -> None: c = base_model(n_snap=6, sparse=sparse) @@ -533,11 +481,10 @@ def test_frozen_constraint_applies_row_scaling( con = c.m.add_constraints( c.balance_lhs() == c.load, name="bal", - freeze=True, scaling=row_scaling, mask=mask, ) - assert isinstance(con, CSRConstraint) + assert con.lhs.is_sparse == sparse active = con.labels != -1 expected = xr.DataArray(row_scaling).broadcast_like(c.load).astype(float) xr.testing.assert_equal( @@ -548,16 +495,14 @@ def test_frozen_constraint_applies_row_scaling( @pytest.mark.parametrize("grid", [False, True], ids=["scalar", "grid"]) @pytest.mark.parametrize("invalid", [-1.0, np.nan]) @pytest.mark.parametrize("sparse", [True, False], ids=["sparse", "dense"]) -def test_frozen_constraint_rejects_invalid_row_scaling_on_masked_rows( +def test_constraint_rejects_invalid_row_scaling_on_masked_rows( sparse: bool, invalid: float, grid: bool ) -> None: c = base_model(sparse=sparse) mask = c.load.bus == "bus1" scaling = c.load.where(mask, invalid) if grid else invalid with pytest.raises(ValueError, match="finite positive"): - c.m.add_constraints( - c.balance_lhs() == c.load, freeze=True, scaling=scaling, mask=mask - ) + c.m.add_constraints(c.balance_lhs() == c.load, scaling=scaling, mask=mask) @pytest.mark.parametrize("masked", [False, True], ids=["all", "masked"]) @@ -566,7 +511,7 @@ def test_frozen_constraint_rejects_invalid_row_scaling_on_masked_rows( [(Case.gen_sum, 0.25), (Case.balance_lhs, 1.5)], ids=["canonical", "unsorted"], ) -def test_freeze_copies_lhs_matrix_only_to_canonicalize( +def test_sparse_constraint_copies_lhs_matrix_only_to_canonicalize( masked: bool, build: Callable[[Case], LinearExpression], copies: float ) -> None: c = base_model(gens_per_bus=(400,) * 50, n_snap=20, sparse=True) @@ -578,7 +523,7 @@ def test_freeze_copies_lhs_matrix_only_to_canonicalize( mask = c.load.bus == "bus1" if masked else None tracemalloc.start() try: - c.m.add_constraints(lhs == c.load, name="bal", freeze=True, mask=mask) + c.m.add_constraints(lhs == c.load, name="bal", mask=mask) _, peak = tracemalloc.get_traced_memory() finally: tracemalloc.stop() @@ -597,7 +542,6 @@ def sparse_model_results(c: Case) -> dict[str, Any]: "expression @": (1.0 * c.gen_p) @ weights, "variable @": c.gen_p @ weights, "constraint": c.m.add_constraints(lhs == c.load, name="bal"), - "unfrozen": c.m.add_constraints(lhs >= 0, name="free", freeze=False), "copy": c.m.copy().constraints["bal"], } @@ -605,13 +549,9 @@ def sparse_model_results(c: Case) -> dict[str, Any]: @pytest.mark.parametrize("sparse", [False, True]) def test_model_sparse_switches_the_whole_build(sparse: bool) -> None: res = sparse_model_results(base_model(sparse=sparse)) - unfrozen = res.pop("unfrozen") - assert isinstance(unfrozen, Constraint) for key, obj in res.items(): - if isinstance(obj, LinearExpression): - assert obj.is_sparse is sparse, key - else: - assert isinstance(obj, CSRConstraint) is sparse, key + expr = obj if isinstance(obj, LinearExpression) else obj.lhs + assert expr.is_sparse is sparse, key def test_model_sparse_persists_through_netcdf(tmp_path: Path) -> None: @@ -622,24 +562,6 @@ def test_model_sparse_persists_through_netcdf(tmp_path: Path) -> None: assert read.copy().sparse -@pytest.mark.parametrize("sparse", [False, True]) -@pytest.mark.parametrize("freeze", [0, 1]) -def test_file_with_freeze_constraints_attr_loads( - tmp_path: Path, sparse: bool, freeze: int -) -> None: - c = base_model(sparse=sparse) - c.m.add_constraints(c.balance_lhs() == c.load, name="bal") - c.m.to_netcdf(tmp_path / "m.nc") - ds = xr.load_dataset(tmp_path / "m.nc") - ds.attrs["freeze_constraints"] = freeze - ds.to_netcdf(tmp_path / "old.nc") - with warnings.catch_warnings(): - warnings.simplefilter("error") - read = linopy.read_netcdf(tmp_path / "old.nc") - assert read.sparse is sparse - assert_conequal(read.constraints["bal"], c.m.constraints["bal"]) - - @pytest.mark.parametrize( "build", [ @@ -659,12 +581,12 @@ def test_sparse_model_notices_groupby_without_sparse_path() -> None: (1.0 * c.gen_p).groupby(c.gbus).sum(use_fallback=True) -def test_materialized_csr_still_freezes_via_dense_path() -> None: +def test_materialized_csr_constraint_stays_sparse() -> None: c = base_model(sparse=True) lhs = c.balance_lhs() _ = lhs.nterm con = c.m.add_constraints(lhs == c.load, name="bal") - assert isinstance(con, CSRConstraint) + assert con.lhs.is_sparse @pytest.mark.parametrize("sparse", [True, False], ids=["sparse", "dense"]) @@ -673,7 +595,7 @@ def test_nan_rhs_raises(sparse: bool) -> None: load = c.load.copy() load[0, 0] = np.nan with pytest.raises(ValueError, match="NaN"): - c.m.add_constraints(c.balance_lhs() == load, name="bal", freeze=True) + c.m.add_constraints(c.balance_lhs() == load, name="bal") @pytest.mark.parametrize("sparse", [True, False], ids=["sparse", "dense"]) @@ -681,7 +603,7 @@ def test_reordered_rhs_raises(sparse: bool) -> None: c = base_model(sparse=sparse) load = c.load.isel(bus=slice(None, None, -1)) with pytest.raises(ValueError, match="[Cc]oordinate"): - c.m.add_constraints(c.balance_lhs() == load, name="bal", freeze=True) + c.m.add_constraints(c.balance_lhs() == load, name="bal") def test_nan_grouper_raises_eagerly() -> None: @@ -710,7 +632,7 @@ def canon_lp(text: str) -> list[str]: def test_lp_files_identical(tmp_path: Path) -> None: sizes = (7, 1, 3, 1, 2, 1, 1, 4, 1, 2, 1, 1) - c1, c2 = twin_models(gens_per_bus=sizes) + c1, c2 = base_model(gens_per_bus=sizes), base_model(gens_per_bus=sizes, sparse=True) c1.m.add_constraints(c1.balance_lhs() == c1.load, name="bal") c1.m.add_objective((1.0 * c1.gen_p).sum()) c2.m.add_constraints(c2.balance_lhs() == c2.load, name="bal") @@ -726,7 +648,7 @@ def case_twins( build: Callable[[Case], LinearExpression], ) -> Callable[[], tuple[LinearExpression, LinearExpression]]: def twins() -> tuple[LinearExpression, LinearExpression]: - c1, c2 = twin_models() + c1, c2 = base_model(), base_model(sparse=True) return build(c1), build(c2) return twins @@ -788,20 +710,20 @@ def test_csr_solution_raises_on_removed_variable() -> None: def test_reindex_falls_back_to_dense_for_unsupported_kwargs() -> None: - c1, c2 = twin_models() + c1, c2 = base_model(), base_model(sparse=True) res = c2.gen_sum().reindex(bus=["bus3", "bus0"], copy=False) assert not res.is_sparse assert_linequal(res, c1.gen_sum().reindex(bus=["bus3", "bus0"])) -def test_reindex_merge_chain_freezes_csr() -> None: - c1, c2 = twin_models() +def test_reindex_merge_chain_constraint_stays_sparse() -> None: + c1, c2 = base_model(), base_model(sparse=True) con1 = c1.m.add_constraints(reindexed_balance(c1) == c1.load, name="bal") tot = reindexed_balance(c2) assert tot.is_sparse con2 = c2.m.add_constraints(tot == c2.load, name="bal") - assert isinstance(con2, CSRConstraint) - assert_frozen_equal(con1, con2) + assert con2.lhs.is_sparse + assert_conequal(con1, con2, strict=False) def test_reindex_merge_chain_peak_memory() -> None: @@ -820,7 +742,7 @@ def test_reindex_merge_chain_peak_memory() -> None: @pytest.mark.parametrize("value", [0.0, 3.5]) def test_fillna_stays_csr_and_matches_dense(value: float) -> None: - c1, c2 = twin_models() + c1, c2 = base_model(), base_model(sparse=True) wide = {"bus": [f"bus{i}" for i in range(7)]} filled = c2.gen_sum().reindex(wide).fillna(value) assert filled.is_sparse @@ -829,7 +751,7 @@ def test_fillna_stays_csr_and_matches_dense(value: float) -> None: def test_fillna_with_array_falls_back_to_dense() -> None: - c1, c2 = twin_models() + c1, c2 = base_model(), base_model(sparse=True) fill = xr.zeros_like(c1.load) res = c2.gen_sum().fillna(fill) assert not res.is_sparse @@ -837,7 +759,7 @@ def test_fillna_with_array_falls_back_to_dense() -> None: def test_rename_stays_csr_and_matches_dense() -> None: - c1, c2 = twin_models() + c1, c2 = base_model(), base_model(sparse=True) sparse = c2.gen_sum().rename(bus="node") assert sparse.is_sparse dense = c1.gen_sum().rename(bus="node") @@ -874,7 +796,7 @@ def cross_grid_parts( def twin_cross_grid_parts() -> tuple[list[LinearExpression], list[LinearExpression]]: """``cross_grid_parts`` of a sparse model and of its dense twin.""" - c1, c2 = twin_models() + c1, c2 = base_model(), base_model(sparse=True) return cross_grid_parts(c2), cross_grid_parts(c1) @@ -911,7 +833,7 @@ def test_cross_grid_merge_stays_csr_and_matches_dense( 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.""" - c1, c2 = twin_models() + c1, c2 = base_model(), base_model(sparse=True) a = (1.0 * c2.flow).groupby(c2.bus0).sum() b = (1.0 * c2.flow_t).groupby(c2.bus0).sum() assert a.coord_dims == b.coord_dims[::-1] != b.coord_dims @@ -929,7 +851,7 @@ def test_transposed_grid_exact_merge_stays_csr_and_matches_dense() -> None: def test_three_operand_cross_grid_merge_matches_dense( join: JoinOptions, fill_value: Any ) -> None: - c1, c2 = twin_models() + c1, c2 = base_model(), base_model(sparse=True) third_lines = ("line3", "line4") sparse = cross_grid_parts(c2) + cross_grid_parts(c2, third_lines)[1:] dense = cross_grid_parts(c1) + cross_grid_parts(c1, third_lines)[1:] @@ -1007,7 +929,8 @@ def test_cross_grid_merge_with_aux_coord_operand_stays_csr(join: JoinOptions) -> want_expr = linopy.merge( [dense[0], dense[1].assign_coords(tag=tag)], join=join, cls=LinearExpression ) - assert_sparse_matches(res, want_expr) + assert res.is_sparse + assert_equivalent(res, want_expr) def test_cross_grid_merge_aux_coord_conflict_raises_like_dense() -> None: @@ -1022,7 +945,7 @@ def test_cross_grid_merge_aux_coord_conflict_raises_like_dense() -> None: def test_cross_grid_merge_with_duplicate_labels_raises_like_dense() -> None: dup = pd.Index(["bus1", "bus1", "bus2"], name="bus") - for c in twin_models(): + for c in (base_model(), base_model(sparse=True)): shed = 1.0 * c.m.add_variables( coords=[dup, c.load.indexes["snapshot"]], name="shed" ) @@ -1033,22 +956,22 @@ def test_cross_grid_merge_with_duplicate_labels_raises_like_dense() -> None: def test_override_merge_same_shape_stays_csr() -> None: dense, sparse = ( [c.gen_sum(), (1.0 * c.flow).groupby(c.bus1.str.upper()).sum()] - for c in twin_models() + for c in (base_model(), base_model(sparse=True)) ) res = linopy.merge(sparse, join="override", cls=LinearExpression) assert res.is_sparse assert_terms_equal(res, linopy.merge(dense, join="override", cls=LinearExpression)) -def test_cross_grid_balance_freezes_csr() -> None: - c1, c2 = twin_models() +def test_cross_grid_balance_constraint_stays_sparse() -> None: + c1, c2 = base_model(), base_model(sparse=True) lhs1 = linopy.merge(cross_grid_parts(c1), join="outer") con1 = c1.m.add_constraints(lhs1 == c1.load, name="bal") lhs2 = linopy.merge(cross_grid_parts(c2), join="outer", cls=LinearExpression) assert lhs2.is_sparse con2 = c2.m.add_constraints(lhs2 == c2.load, name="bal") - assert isinstance(con2, CSRConstraint) - assert_frozen_equal(con1, con2) + assert con2.lhs.is_sparse + assert_conequal(con1, con2, strict=False) def test_cross_grid_merge_peak_memory() -> None: @@ -1458,7 +1381,7 @@ def test_matmul_output_dim_order_matches_dense() -> None: def test_matmul_keeps_csr_backing_through_chain() -> None: - c1, c2 = twin_models() + c1, c2 = base_model(), base_model(sparse=True) operand = loc_operand(c2.gen_p.indexes["snapshot"]) gen = c2.gen_sum() flow = (1.0 * c2.flow).groupby(c2.bus0).sum() @@ -1471,7 +1394,7 @@ def test_matmul_keeps_csr_backing_through_chain() -> None: total = res + (flow @ operand) assert total.is_sparse con = c2.m.add_constraints(total <= 0, name="matmul") - assert isinstance(con, CSRConstraint) + assert con.lhs.is_sparse def test_quadratic_matmul_stays_on_dense_path() -> None: @@ -1632,12 +1555,11 @@ def set_blocks(e: LinearExpression, c: Case) -> None: lambda e, c: e <= xr.DataArray([1.0, 2.0], coords=[LOC]), r"to_constraint\(rhs over dimensions outside the grid\)", ), - "mutable": (lambda e, c: (e == c.load).mutable(), r"`mutable\(\)`"), - "add-unfrozen": ( - lambda e, c: c.m.add_constraints(e >= 1, freeze=False), - "constraint added unfrozen, `freeze=False`", + "set-blocks": (set_blocks, "`set_blocks`"), + "update-coeffs": ( + lambda e, c: c.m.add_constraints(e >= 1).update(coeffs=2.0), + r"update\(coeffs=, variables=\)", ), - "set-blocks": (set_blocks, "frozen constraint converted by `set_blocks`"), "new-dim": ( lambda e, c: e * xr.DataArray([1.0, 2.0], coords=[LOC]), r"mul\(constant over new dimensions\)", @@ -1697,8 +1619,7 @@ def test_warn_on_densify_names_the_reason(op: str, enabled: bool | None) -> None opts.set_value(warn_on_densify=enabled) func(sparse, c) notices = [w for w in caught if issubclass(w.category, linopy.PerformanceWarning)] - implicit_in_sparse_model = op not in {"mutable", "add-unfrozen"} - if enabled is False or (enabled is None and not implicit_in_sparse_model): + if enabled is False: assert notices == [] return assert len(notices) == 1 @@ -1726,9 +1647,12 @@ def test_sparse_store_indices_follow_model_label_dtype(build: str) -> None: m = sparse.model con = m.add_constraints(sparse >= 1, name="con") A = m.matrices.A - assert isinstance(con, CSRConstraint) and sparse.is_sparse + assert con.lhs.is_sparse and sparse.is_sparse assert A is not None - dtypes = {sparse_store(sparse).csr.indices.dtype, con._csr.indices.dtype} + dtypes = { + sparse_store(sparse).csr.indices.dtype, + sparse_store(con.lhs).csr.indices.dtype, + } assert dtypes == {np.dtype(m.dtypes["labels"]), A.indices.dtype} @@ -1917,7 +1841,8 @@ def test_merge_with_dense_expression_stays_csr_and_matches_dense( other = DENSE_ADDENDS[addend](c, dense) with sparsified(): res = other + sparse if swap else sparse + other - assert_sparse_matches(res, other + dense if swap else dense + other) + assert res.is_sparse + assert_equivalent(res, other + dense if swap else dense + other) def unaligned(e: LinearExpression) -> LinearExpression: @@ -1943,7 +1868,8 @@ def test_merge_over_different_dims_joins_like_dense( want = want_first.add(want_second, join=join, fill_value=fill_value) assert isinstance(res, LinearExpression) assert isinstance(want, LinearExpression) - assert_sparse_matches(res, want) + assert res.is_sparse + assert_equivalent(res, want) @pytest.mark.parametrize( @@ -1995,10 +1921,11 @@ def sum_then_reuse(e: LinearExpression) -> LinearExpression: def test_variable_arithmetic_is_csr_in_sparse_model_and_matches_dense( op: str, ) -> None: - dense_case, sparse_case = twin_models() + dense_case, sparse_case = base_model(), base_model(sparse=True) with no_densify(): res = VARIABLE_OPS[op](sparse_case) - assert_sparse_matches(res, VARIABLE_OPS[op](dense_case)) + assert res.is_sparse + assert_equivalent(res, VARIABLE_OPS[op](dense_case)) DENSE_VARIABLE_OPS: dict[str, Callable[[Case], Any]] = { @@ -2012,7 +1939,7 @@ def test_variable_arithmetic_is_csr_in_sparse_model_and_matches_dense( @pytest.mark.parametrize("op", list(DENSE_VARIABLE_OPS)) def test_variable_ops_without_sparse_path_build_dense_silently(op: str) -> None: - dense_case, sparse_case = twin_models() + dense_case, sparse_case = base_model(), base_model(sparse=True) with no_densify(): res = DENSE_VARIABLE_OPS[op](sparse_case) want = DENSE_VARIABLE_OPS[op](dense_case) @@ -2100,34 +2027,24 @@ def test_join_with_constant_keeps_aux_coords_like_dense( ) with no_densify(): res = getattr(sparse, op)(x, join=join) - assert_sparse_matches(res, getattr(dense, op)(x, join=join)) + assert res.is_sparse + assert_equivalent(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: cons = [] - for c in twin_models(): + for c in (base_model(), base_model(sparse=True)): lhs = c.gen_sum() with no_densify(): con = lhs.to_constraint("<=", c.load.isel(bus=[0, 2]), join=join) cons.append(c.m.add_constraints(con, name="c")) - dense, frozen = cons - assert isinstance(frozen, CSRConstraint) + dense, sparse = cons + assert sparse.lhs.is_sparse xr.testing.assert_identical( - xr.Dataset(coords=frozen.coords), xr.Dataset(coords=dense.coords) + xr.Dataset(coords=sparse.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 - assert res.coord_dims == want.coord_dims - xr.testing.assert_identical( - xr.Dataset(coords=res.coords), xr.Dataset(coords=want.coords) - ) - xr.testing.assert_identical(res.isnull(), want.isnull()) - assert_cells_equal(res, want, tuple(map(str, want.coord_dims))) + assert_conequal(dense, sparse, strict=False) SUM_DIMS: dict[str, Callable[[tuple[str, ...]], Any]] = { @@ -2149,7 +2066,8 @@ def test_sum_stays_csr_and_matches_dense(build: str, dims: str) -> None: dim = SUM_DIMS[dims](tuple(map(str, sparse.coord_dims))) with no_densify(): res = sparse.sum(dim) - assert_sparse_matches(res, dense.sum(dim)) + assert res.is_sparse + assert_equivalent(res, dense.sum(dim)) @pytest.mark.parametrize("drop_zeros", [False, True]) @@ -2161,7 +2079,8 @@ def test_sum_keeps_explicit_zeros_unless_dropped(drop_zeros: bool) -> None: res = sparse.sum("snapshot", drop_zeros=drop_zeros) assert res.is_sparse assert (sparse_store(res).csr.nnz == 0) == drop_zeros - assert_sparse_matches(res, dense.sum("snapshot", drop_zeros=drop_zeros)) + assert res.is_sparse + assert_equivalent(res, dense.sum("snapshot", drop_zeros=drop_zeros)) def test_sum_unknown_dim_raises_like_dense() -> None: @@ -2191,7 +2110,8 @@ def test_chained_groupby_stays_csr_and_matches_dense(build: str, chain: str) -> func = CHAINS[chain] with no_densify(): res = func(sparse) - assert_sparse_matches(res, func(dense_build(build))) + assert res.is_sparse + assert_equivalent(res, func(dense_build(build))) @pytest.mark.parametrize("observed", [False, True]) @@ -2199,7 +2119,8 @@ def test_chained_namelist_groupby_on_aux_coords_matches_dense(observed: bool) -> sparse, _ = sparse_and_dense("aux") dense = dense_build("aux") res = sparse.groupby(["bus", "tag"]).sum(observed=observed) - assert_sparse_matches(res, dense.groupby(["bus", "tag"]).sum(observed=observed)) + assert res.is_sparse + assert_equivalent(res, dense.groupby(["bus", "tag"]).sum(observed=observed)) def dim0(e: LinearExpression) -> str: @@ -2247,7 +2168,8 @@ def test_selection_stays_csr_and_matches_dense(build: str, select: str) -> None: func = SELECTIONS[select] with no_densify(): res = func(sparse) - assert_sparse_matches(res, func(dense)) + assert res.is_sparse + assert_equivalent(res, func(dense)) def dim_last(e: LinearExpression) -> str: @@ -2283,7 +2205,8 @@ def test_reindex_matches_dense(build: str, reindexer: str, scale: float) -> None res = (scale * sparse).reindex(indexers) want = (scale * dense).reindex(indexers) xr.testing.assert_identical(res.has_terms, want.has_terms) - assert_sparse_matches(res, want) + assert res.is_sparse + assert_equivalent(res, want) @pytest.mark.parametrize("build", ["grouped", "aux", "absent"]) @@ -2312,7 +2235,8 @@ def test_n_ary_merge_keeps_absent_cells_and_explicit_zeros(build: str, n: int) - assert (np.diff(csr.csr.indptr)[np.isnan(csr.const)] == 0).all() want = linopy.merge(merge_operands(dense, n), cls=LinearExpression) xr.testing.assert_identical(res.has_terms, want.has_terms) - assert_sparse_matches(res, want) + assert res.is_sparse + assert_equivalent(res, want) def tagged_by(c: Case, **aux: pd.Series) -> LinearExpression: @@ -2327,7 +2251,8 @@ def test_n_ary_merge_unites_aux_coords_and_raises_on_conflict() -> None: with no_densify(): res = linopy.merge(parts, cls=LinearExpression) want = linopy.merge([densified(p) for p in parts], cls=LinearExpression) - assert_sparse_matches(res, want) + assert res.is_sparse + assert_equivalent(res, want) with pytest.raises(ValueError, match="conflicting values"): linopy.merge([*parts, tagged_by(c, tag=c.gbus + "z")], cls=LinearExpression) @@ -2354,7 +2279,8 @@ def test_scalar_selection_coords_merge_like_dense() -> None: flow = densified((1.0 * c.flow).groupby(c.bus0).sum()) with sparsified(): res = sparse.sel(snapshot=1) + flow.sel(snapshot=1) - assert_sparse_matches(res, dense.sel(snapshot=1) + flow.sel(snapshot=1)) + assert res.is_sparse + assert_equivalent(res, dense.sel(snapshot=1) + flow.sel(snapshot=1)) with pytest.raises(ValueError) as want: dense.sel(snapshot=1) + flow.sel(snapshot=2) with pytest.raises(ValueError) as got: @@ -2395,10 +2321,10 @@ 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( +def test_add_constraints_mask_keeps_sparse_and_matches_dense( mask: str, prebuilt: bool ) -> None: - c1, c2 = twin_models() + c1, c2 = base_model(), base_model(sparse=True) lhs = c2.balance_lhs() m = MASKS[mask](lhs) con1 = c1.m.add_constraints(c1.balance_lhs(), ">=", c1.load, "bal", mask=m) @@ -2407,8 +2333,8 @@ def test_add_constraints_mask_freezes_sparse_and_matches_dense( con2 = c2.m.add_constraints(lhs >= c2.load, name="bal", mask=m) else: con2 = c2.m.add_constraints(lhs, ">=", c2.load, "bal", mask=m) - assert isinstance(con2, CSRConstraint) - assert_frozen_equal(con1, con2) + assert con2.lhs.is_sparse + assert_conequal(con1, con2, strict=False) def observed_keys_group() -> LinearExpression: @@ -2442,65 +2368,68 @@ def dashed_group() -> LinearExpression: @pytest.mark.parametrize("build", list(AUX_BUILDS)) -def test_frozen_constraint_keeps_aux_coords_like_dense( +def test_sparse_constraint_keeps_aux_coords_like_dense( build: str, tmp_path: Path ) -> None: sparse = AUX_BUILDS[build]() m = sparse.model assert sparse.is_sparse - ref = m.add_constraints(densified(sparse) >= 1, name="dense", freeze=False) + ref = densified(sparse) >= 1 with no_densify(): con = m.add_constraints(sparse >= 1, name="sparse") - assert isinstance(con, CSRConstraint) + assert con.lhs.is_sparse and not ref.lhs.is_sparse 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) + assert set(want.coords) > set(map(str, con.coord_dims)) + xr.testing.assert_identical(xr.Dataset(coords=con.coords), want) m.to_netcdf(tmp_path / "m.nc") read = linopy.read_netcdf(tmp_path / "m.nc").constraints["sparse"] - assert isinstance(read, CSRConstraint) + assert read.lhs.is_sparse 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: - c = base_model(sparse=True) - con = c.m.add_constraints(c.balance_lhs() >= c.load) - 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]] = { - "from_rule": ( - lambda con: CSRConstraint.from_rule(con.model, lambda m, i: None, [[0]]), - r"Constraint\.from_rule .*\.freeze\(\)", +SPARSE_UPDATES: dict[str, Callable[[Constraint, Case], Any]] = { + "lhs": lambda con, c: con.update(lhs=2 * c.balance_lhs()), + "lhs-dense": lambda con, c: con.update( + lhs=(1.0 * c.flow).groupby(c.bus0).sum(use_fallback=True) ), + "rhs": lambda con, c: con.update(rhs=2 * c.load), + "rhs-expr": lambda con, c: con.update(rhs=(1.0 * c.flow).groupby(c.bus1).sum()), + "sign": lambda con, c: con.update(sign="<="), + "constraint": lambda con, c: con.update(c.gen_sum() <= c.load), + "coeffs": lambda con, c: con.update(coeffs=2.0), + "scaling": lambda con, c: setattr(con, "scaling", 2.0), } -@pytest.mark.parametrize("op", list(FROZEN_UNSUPPORTED)) -def test_frozen_unsupported_names_the_working_route(op: str) -> None: - c = base_model(sparse=True) - con = c.m.add_constraints(c.balance_lhs() >= c.load) - assert isinstance(con, CSRConstraint) - call, remedy = FROZEN_UNSUPPORTED[op] - with pytest.raises(AttributeError, match=rf"CSRConstraint\.{op} .*{remedy}"): - call(con) +@pytest.mark.parametrize("op", list(SPARSE_UPDATES)) +def test_update_keeps_sparse_lhs_and_matches_dense(op: str) -> None: + cons = [] + for sparse in (False, True): + c = base_model(sparse=sparse) + con = c.m.add_constraints(c.balance_lhs() >= c.load, name="c") + SPARSE_UPDATES[op](con, c) + assert con.lhs.is_sparse == sparse + assert c.m.constraints["c"] is con + cons.append(con) + dense, sparse_con = cons + assert_conequal(dense, sparse_con, strict=False) + xr.testing.assert_equal(dense.scaling, sparse_con.scaling) + + +def test_set_blocks_on_sparse_model_matches_dense(tmp_path: Path) -> None: + blocks = [] + for sparse in (False, True): + c = base_model(sparse=sparse) + con = c.m.add_constraints(c.balance_lhs() >= c.load, name="c") + c.m.constraints.set_blocks(np.arange(c.m._xCounter) % 3) + assert con.lhs.is_sparse == sparse + path = tmp_path / f"{sparse}.nc" + c.m.to_netcdf(path) + read = linopy.read_netcdf(path).constraints["c"] + assert read.lhs.is_sparse == sparse + xr.testing.assert_equal(read.data["blocks"], con.data["blocks"]) + blocks.append(con.data["blocks"]) + xr.testing.assert_equal(*blocks) EXPR_RHS: dict[str, Callable[[Case], LinearExpression]] = { @@ -2515,11 +2444,9 @@ def test_frozen_unsupported_names_the_working_route(op: str) -> None: @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: - c1, c2 = twin_models() - con1 = c1.m.add_constraints( - c1.gen_sum() <= EXPR_RHS[rhs](c1), name="c", freeze=True - ) +def test_expression_rhs_keeps_sparse_and_matches_dense(rhs: str, form: str) -> None: + c1, c2 = base_model(), base_model(sparse=True) + con1 = c1.m.add_constraints(c1.gen_sum() <= EXPR_RHS[rhs](c1), name="c") lhs2 = c2.gen_sum() rhs2 = EXPR_RHS[rhs](c2) with sparsified() if rhs == "dense-expr-const" else no_densify(): @@ -2527,8 +2454,8 @@ def test_expression_rhs_freezes_sparse_and_matches_dense(rhs: str, form: str) -> con2 = c2.m.add_constraints(lhs2 <= rhs2, name="c") else: con2 = c2.m.add_constraints(lhs2, "<=", rhs2, name="c") - assert isinstance(con2, CSRConstraint) - assert_frozen_equal(con1, con2) + assert con2.lhs.is_sparse + assert_conequal(con1, con2, strict=False) MASK_KINDS: dict[str, Callable[[Case], xr.DataArray | None]] = { @@ -2544,7 +2471,7 @@ def softened( route: str, max_violation: float | None, sparse: bool, -) -> tuple[Case, ConstraintBase]: +) -> tuple[Case, Constraint]: c = base_model(sparse=sparse) c.m.add_objective(1.0 * c.gen_p.sum()) mask = MASK_KINDS[mask_kind](c) @@ -2554,7 +2481,6 @@ def softened( con = c.m.add_constraints(*args, **kwargs) con.soften(penalty=2.0, max_violation=max_violation) else: - kwargs["freeze"] = None if route == "penalty-default" else sparse con = c.m.add_constraints(*args, **kwargs, penalty=2.0) return c, con @@ -2563,32 +2489,31 @@ def softened( ("soften", None), ("soften", 5.0), ("penalty", None), - ("penalty-default", None), ] @pytest.mark.parametrize(("route", "max_violation"), SOFTEN_ROUTES) @pytest.mark.parametrize("mask_kind", list(MASK_KINDS)) @pytest.mark.parametrize("sign", ["<=", ">=", "=="]) -def test_frozen_soften_matches_dense( +def test_sparse_soften_matches_dense( sign: str, mask_kind: str, route: str, max_violation: float | None, tmp_path: Path ) -> None: dense_case, dense = softened(sign, mask_kind, route, max_violation, sparse=False) with no_densify(): case, con = softened(sign, mask_kind, route, max_violation, sparse=True) case.m.to_netcdf(tmp_path / "m.nc") - assert isinstance(con, CSRConstraint) + assert con.lhs.is_sparse assert case.m.constraints["c"] is con - assert_frozen_equal(dense, con) + assert_conequal(dense, con, strict=False) obj, want_obj = case.m.objective.expression, dense_case.m.objective.expression assert isinstance(obj, LinearExpression) and isinstance(want_obj, LinearExpression) assert_cells_equal(obj, want_obj, ()) assert dense.slack is not None read = linopy.read_netcdf(tmp_path / "m.nc").constraints["c"] - assert isinstance(read, CSRConstraint) - assert_frozen_equal(dense, read) - for frozen_con in (con, read, con.mutable(), dense.freeze()): - slack = frozen_con.slack + assert read.lhs.is_sparse + assert_conequal(dense, read, strict=False) + for sparse_con in (con, read): + slack = sparse_con.slack assert slack is not None assert_varequal(slack.positive, dense.slack.positive) assert (slack.negative is None) == (dense.slack.negative is None) @@ -2596,7 +2521,7 @@ def test_frozen_soften_matches_dense( assert_varequal(slack.negative, dense.slack.negative) -def test_frozen_soften_keeps_sparse_objective() -> None: +def test_sparse_soften_keeps_sparse_objective() -> None: c = base_model(sparse=True) with no_densify(): c.m.add_objective((1.0 * c.gen_p).groupby(c.gbus).sum().sum()) @@ -2604,12 +2529,12 @@ def test_frozen_soften_keeps_sparse_objective() -> None: assert c.m.objective.expression.is_sparse -def test_frozen_soften_max_sense_with_array_penalty() -> None: +def test_sparse_soften_max_sense_with_array_penalty() -> None: c = base_model(sparse=True) c.m.add_objective(1.0 * c.gen_p.sum(), sense="max") with no_densify(): con = c.m.add_constraints(c.balance_lhs(), ">=", c.load, name="c") - assert isinstance(con, CSRConstraint) + assert con.lhs.is_sparse penalty = xr.full_like(c.load, 2.0) slack = con.soften(penalty=penalty) expected_objective = 1.0 * c.gen_p.sum() - (penalty * slack.positive).sum() @@ -2620,7 +2545,7 @@ def test_frozen_soften_max_sense_with_array_penalty() -> None: @pytest.mark.skipif("highs" not in linopy.licensed_solvers, reason="needs highs") @pytest.mark.parametrize("sign", ["<=", ">=", "=="]) -def test_frozen_soften_solves_like_dense(sign: str) -> None: +def test_sparse_soften_solves_like_dense(sign: str) -> None: values = [] for sparse in (False, True): c, _ = softened(sign, "masked", "soften", 20.0, sparse) @@ -2631,12 +2556,12 @@ def test_frozen_soften_solves_like_dense(sign: str) -> None: assert values[0] == pytest.approx(values[1]) -def soften_twice(con: ConstraintBase) -> Any: +def soften_twice(con: Constraint) -> Any: con.soften(penalty=1.0) return con.soften(penalty=1.0) -SOFTEN_ERRORS: dict[str, tuple[Callable[[ConstraintBase], Any], type, str]] = { +SOFTEN_ERRORS: dict[str, tuple[Callable[[Constraint], Any], type, str]] = { "twice": (soften_twice, ValueError, "already softened"), "zero-penalty": (lambda con: con.soften(penalty=0.0), ValueError, "not positive"), "no-objective": ( @@ -2652,32 +2577,32 @@ def soften_twice(con: ConstraintBase) -> Any: } # "twice", "zero-penalty" and "no-objective" are already covered for dense -# constraints in test/test_constraint.py; only the frozen path needs them here. -FROZEN_ONLY_ERRORS = {"twice", "zero-penalty", "no-objective"} +# constraints in test/test_constraint.py; only the sparse path needs them here. +SPARSE_ONLY_ERRORS = {"twice", "zero-penalty", "no-objective"} SOFTEN_REJECT_CASES = [ - pytest.param(error, freeze, id=f"{error}-{'frozen' if freeze else 'dense'}") + pytest.param(error, sparse, id=f"{error}-{'sparse' if sparse else 'dense'}") for error in SOFTEN_ERRORS - for freeze in ((True,) if error in FROZEN_ONLY_ERRORS else (False, True)) + for sparse in ((True,) if error in SPARSE_ONLY_ERRORS else (False, True)) ] -@pytest.mark.parametrize(("error", "freeze"), SOFTEN_REJECT_CASES) -def test_soften_rejects(error: str, freeze: bool) -> None: - c = base_model() +@pytest.mark.parametrize(("error", "sparse"), SOFTEN_REJECT_CASES) +def test_soften_rejects(error: str, sparse: bool) -> None: + c = base_model(sparse=sparse) if error != "no-objective": c.m.add_objective(1.0 * c.gen_p.sum()) sign: Any = ">=" if error == "mixed-signs": sign = xr.DataArray(np.where(c.load > 4, ">=", "<="), coords=c.load.coords) - con = c.m.add_constraints(c.balance_lhs(), sign, c.load, name="c", freeze=freeze) - assert isinstance(con, CSRConstraint) == freeze + con = c.m.add_constraints(c.balance_lhs(), sign, c.load, name="c") + assert con.lhs.is_sparse == sparse call, exc, match = SOFTEN_ERRORS[error] with pytest.raises(exc, match=match): call(con) -def frozen_model(soften: bool) -> tuple[Model, CSRConstraint]: +def sparse_case_model(soften: bool) -> tuple[Model, Constraint]: c = base_model(sparse=True) for var in (c.gen_p, c.flow): var.update(lower=0, upper=1) @@ -2686,30 +2611,31 @@ def frozen_model(soften: bool) -> tuple[Model, CSRConstraint]: con = c.m.add_constraints(c.balance_lhs(), ">=", c.load, name="c") if soften: con.soften(penalty=2.0, max_violation=20.0) - assert isinstance(con, CSRConstraint) + assert con.lhs.is_sparse return c.m, con @pytest.mark.skipif("highs" not in linopy.licensed_solvers, reason="needs highs") @pytest.mark.parametrize("include_solution", [True, False]) @pytest.mark.parametrize("deep", [True, False]) -def test_copy_keeps_frozen_constraints(deep: bool, include_solution: bool) -> None: - m, con = frozen_model(soften=True) +def test_copy_keeps_sparse_constraints(deep: bool, include_solution: bool) -> None: + m, con = sparse_case_model(soften=True) m.solve("highs") with no_densify(): c = m.copy(deep=deep, include_solution=include_solution) copied = c.constraints["c"] assert copied.model is c and copied.name == "c" + assert copied.lhs.is_sparse assert c.objective.expression.is_sparse assert con.slack is not None and copied.slack is not None assert copied.slack.positive.labels.equals(con.slack.positive.labels) assert copied.slack.positive.model is c - assert ("dual" in copied.mutable().data) == include_solution + assert ("dual" in copied) == include_solution @pytest.mark.parametrize("deep", [True, False]) -def test_softening_frozen_copy_leaves_original(deep: bool) -> None: - m, con = frozen_model(soften=False) +def test_softening_sparse_copy_leaves_original(deep: bool) -> None: + m, con = sparse_case_model(soften=False) with no_densify(): copied = m.copy(deep=deep) copied.constraints["c"].soften(penalty=2.0) @@ -2719,7 +2645,7 @@ def test_softening_frozen_copy_leaves_original(deep: bool) -> None: @pytest.mark.parametrize("deep", [True, False]) -def test_copy_frozen_matrices(deep: bool) -> None: +def test_copy_sparse_matrices(deep: bool) -> None: m = Model(sparse=True) i = pd.RangeIndex(4, name="i") x = m.add_variables(lower=0, coords=[i], name="x") @@ -2732,8 +2658,9 @@ def test_copy_frozen_matrices(deep: bool) -> None: c = m.copy(deep=deep) for name, con in m.constraints.items(): copied = c.constraints[name] - assert isinstance(con, CSRConstraint) and isinstance(copied, CSRConstraint) - assert np.shares_memory(con._csr.data, copied._csr.data) != deep + assert con.lhs.is_sparse and copied.lhs.is_sparse + stored, copied_csr = sparse_store(con.lhs).csr, sparse_store(copied.lhs).csr + assert np.shares_memory(stored.data, copied_csr.data) != deep got, want = c.matrices, m.matrices assert got.A is not None and want.A is not None assert got.indicator_A is not None and want.indicator_A is not None @@ -2748,8 +2675,11 @@ def test_copy_frozen_matrices(deep: bool) -> None: ] for g, w in pairs: assert np.array_equal(g, w) - with pytest.raises(ValueError, match="read-only"): - c.constraints["c"].coords["aux"].values[0] = -1 + copied_aux = c.constraints["c"].coords["aux"].values + original_aux = m.constraints["c"].coords["aux"].values + assert not ( + copied_aux.flags.writeable and np.shares_memory(copied_aux, original_aux) + ) assert c.constraints["c"].coords["aux"].equals(m.constraints["c"].coords["aux"]) @@ -2774,7 +2704,8 @@ def test_shift_roll_diff_stay_csr_and_match_dense( sparse, dense = scale * sparse, scale * dense with no_densify(): res = SHIFT_OPS[op](sparse) - assert_sparse_matches(res, SHIFT_OPS[op](dense)) + assert res.is_sparse + assert_equivalent(res, SHIFT_OPS[op](dense)) @pytest.mark.parametrize( @@ -2827,7 +2758,8 @@ def test_concat_stays_csr_and_matches_dense( with no_densify(): res = concat(operands, dim=dim) assert all(e.is_sparse for e in operands) - assert_sparse_matches(res, concat(split(dense, dim, *positions), dim=dim)) + assert res.is_sparse + assert_equivalent(res, concat(split(dense, dim, *positions), dim=dim)) @pytest.mark.parametrize("transposed", [False, True], ids=["same-order", "transposed"]) @@ -2842,7 +2774,8 @@ def test_concat_with_dense_operand_stays_csr_and_matches_dense( ) with sparsified(): res = concat([sparse.isel(snapshot=[0]), tail]) - assert_sparse_matches(res, concat([dense.isel(snapshot=[0]), tail])) + assert res.is_sparse + assert_equivalent(res, concat([dense.isel(snapshot=[0]), tail])) def joined_parts(e: LinearExpression) -> list[LinearExpression]: @@ -2855,7 +2788,8 @@ def test_concat_join_stays_csr_and_matches_dense(join: JoinOptions) -> None: sparse, dense = sparse_and_dense("grouped") with no_densify(): res = concat(joined_parts(sparse), join) - assert_sparse_matches(res, concat(joined_parts(dense), join)) + assert res.is_sparse + assert_equivalent(res, concat(joined_parts(dense), join)) def override_parts(e: LinearExpression) -> list[LinearExpression]: @@ -2883,20 +2817,22 @@ def test_concat_coords_follow_dense( sparse, dense = sparse_and_dense(build) with no_densify(): res = concat(parts(sparse), join) - assert_sparse_matches(res, concat(parts(dense), join)) + assert res.is_sparse + assert_equivalent(res, concat(parts(dense), join)) -def test_concat_freezes_like_dense() -> None: - dense_c, sparse_c = twin_models() +def test_concat_constraint_matches_dense() -> None: + dense_c, sparse_c = base_model(), base_model(sparse=True) def lhs(c: Case) -> LinearExpression: parts = split(c.gen_sum(), "bus", [0, 1], [2, 3, 4]) return linopy.merge(parts, dim="bus", cls=LinearExpression) with no_densify(): - frozen = sparse_c.m.add_constraints(lhs(sparse_c) >= 1, name="c") - assert isinstance(frozen, CSRConstraint) - assert_frozen_equal(dense_c.m.add_constraints(lhs(dense_c) >= 1, name="c"), frozen) + sparse = sparse_c.m.add_constraints(lhs(sparse_c) >= 1, name="c") + assert sparse.lhs.is_sparse + dense = dense_c.m.add_constraints(lhs(dense_c) >= 1, name="c") + assert_conequal(dense, sparse, strict=False) @pytest.mark.parametrize( @@ -2940,7 +2876,7 @@ def stored_arrays(e: LinearExpression) -> list[np.ndarray]: "groupby": lambda e, c: e.groupby(halves(e)).sum(), "reindex": lambda e, c: e.reindex(bus=["bus3", "bus0", "bus9"]), "repr": lambda e, c: repr(e), - "freeze": lambda e, c: c.m.add_constraints(e >= 1, name="c", freeze=True), + "constraint": lambda e, c: c.m.add_constraints(e >= 1, name="c"), "matrices": lambda e, c: (c.m.add_constraints(e >= 1, name="c"), c.m.matrices.A), } @@ -2980,15 +2916,16 @@ def build(sparse: bool) -> tuple[LinearExpression, LinearExpression]: assert_linequal(sparse, dense_repeated + dense_y) -def test_duplicate_terms_kept_like_dense_and_summed_at_freeze() -> None: - dense_c, sparse_c = twin_models() +def test_duplicate_terms_kept_like_dense_and_summed_in_constraint() -> None: + dense_c, sparse_c = base_model(), base_model(sparse=True) dense, sparse = duplicated_lhs(dense_c), duplicated_lhs(sparse_c) assert sparse.nterm == dense.nterm - assert_sparse_matches(sparse, dense) - frozen = sparse_c.m.add_constraints(sparse == sparse_c.load, name="bal") + assert sparse.is_sparse + assert_equivalent(sparse, dense) + con = sparse_c.m.add_constraints(sparse == sparse_c.load, name="bal") dense_c.m.add_constraints(dense == dense_c.load, name="bal") - assert isinstance(frozen, CSRConstraint) - assert frozen.nterm == max(dense_c.gbus.value_counts()) + 2 + assert con.lhs.is_sparse + assert con.nterm == max(dense_c.gbus.value_counts()) + 2 got, want = sparse_c.m.matrices.A, dense_c.m.matrices.A assert got is not None and want is not None assert got.nnz == want.nnz diff --git a/test/test_indicator_constraints.py b/test/test_indicator_constraints.py index ece30bd1c..7ab33d891 100644 --- a/test/test_indicator_constraints.py +++ b/test/test_indicator_constraints.py @@ -7,8 +7,7 @@ import pytest import linopy -from linopy import Model, available_solvers, licensed_solvers -from linopy.constraints import Constraint, CSRConstraint +from linopy import Constraint, Model, available_solvers, licensed_solvers from linopy.variables import Variable requires_gurobi = pytest.mark.skipif( @@ -184,18 +183,13 @@ def test_remove(self, mbx: tuple[Model, Variable, Variable]) -> None: def test_regular_constraint_has_no_indicator_fields( self, mbx: tuple[Model, Variable, Variable] ) -> None: - """Regular constraints report no indicator metadata, frozen or not.""" + """Regular constraints report no indicator metadata, sparse or not.""" m, _, x = mbx - m.add_constraints(x <= 5, name="c") - con = m.constraints["c"] - assert not con.is_indicator - assert con.binary_var is None - assert con.binary_val is None - - frozen = con.freeze() - assert not frozen.is_indicator - assert frozen.binary_var is None - assert frozen.binary_val is None + con = m.add_constraints(x <= 5, name="c") + for c in (con, con.to_sparse()): + assert not c.is_indicator + assert c.binary_var is None + assert c.binary_val is None def test_with_coords( self, coords_mbx: tuple[Model, Variable, Variable, pd.RangeIndex] @@ -207,24 +201,20 @@ def test_with_coords( class TestPersistence: - """Indicator metadata survives freeze, copy, and netCDF round-trips.""" + """Indicator metadata survives store conversion, copy, and netCDF round-trips.""" - def test_freeze_roundtrip(self, mbx: tuple[Model, Variable, Variable]) -> None: - """is_indicator and binary fields survive freeze and freeze->mutable.""" + @pytest.mark.parametrize("to", ["to_sparse", "to_dense"]) + def test_store_conversion_roundtrip( + self, mbx: tuple[Model, Variable, Variable], to: str + ) -> None: + """is_indicator and binary fields survive converting the lhs store.""" m, b, x = mbx - m.add_indicator_constraints(b, 1, x, "<=", 5, name="ic0") - - frozen = m.constraints["ic0"].freeze() - assert isinstance(frozen, CSRConstraint) - assert frozen.is_indicator - assert frozen.binary_var is not None - assert frozen.binary_val == 1 - - mutable = frozen.mutable() - assert isinstance(mutable, Constraint) - assert mutable.is_indicator - assert mutable.binary_var is not None - assert np.all(mutable.binary_val == 1) + con = m.add_indicator_constraints(b, 1, x, "<=", 5, name="ic0") + converted = getattr(con.to_sparse(), to)() + assert converted.lhs.is_sparse == (to == "to_sparse") + assert converted.is_indicator + assert converted.binary_var is not None + assert np.all(converted.binary_val == 1) def test_copy_preserves(self, mbx: tuple[Model, Variable, Variable]) -> None: """Model.copy preserves is_indicator and the binary fields.""" @@ -261,10 +251,10 @@ def test_array_binval_roundtrip(self, tmp_path: Path) -> None: x = m.add_variables(coords=[idx], lower=0, upper=10, name="x") m.add_indicator_constraints(b, 1, x, "<=", 5, name="ic0") - frozen = m.constraints["ic0"] - assert isinstance(frozen, CSRConstraint) - assert isinstance(frozen.binary_val, np.ndarray) and frozen.binary_val.ndim == 1 - assert "binary_val" in frozen.data + con = m.constraints["ic0"] + assert con.lhs.is_sparse + assert isinstance(con.binary_val, np.ndarray) and con.binary_val.ndim == 1 + assert "binary_val" in con.data fn = tmp_path / "model.nc" m.to_netcdf(fn) diff --git a/test/test_io.py b/test/test_io.py index 74fa029bf..db9817bb9 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -171,20 +171,18 @@ def test_read_netcdf_inconsistent_order_attr_raises( read_netcdf(fn) -def test_model_to_netcdf_frozen_constraint(tmp_path: Path) -> None: - from linopy.constraints import CSRConstraint - - m = Model() +def test_model_to_netcdf_keeps_constraint_store(tmp_path: Path, backend: Model) -> None: + m = backend x = m.add_variables(coords=[pd.RangeIndex(5, name="i")], name="x") - m.add_constraints(x >= 1, name="c", freeze=True) + m.add_constraints(x >= 1, name="c") - assert isinstance(m.constraints["c"], CSRConstraint) + assert m.constraints["c"].lhs.is_sparse == m.sparse - fn = tmp_path / "test_frozen.nc" + fn = tmp_path / "test_store.nc" m.to_netcdf(fn) p = read_netcdf(fn) - assert isinstance(p.constraints["c"], CSRConstraint) + assert p.constraints["c"].lhs.is_sparse == m.sparse assert_model_equal(m, p) @@ -192,7 +190,7 @@ def test_model_to_netcdf_frozen_constraint(tmp_path: Path) -> None: def test_model_to_netcdf_compression(tmp_path: Path) -> None: m = Model() x = m.add_variables(lower=0, coords=[pd.RangeIndex(2000, name="i")], name="x") - m.add_constraints(x + x.shift(i=1) >= 1, name="c", freeze=True) + m.add_constraints(x + x.shift(i=1) >= 1, name="c").to_sparse() m.add_constraints(x <= 10, name="d") m.add_objective(x.sum()) @@ -219,7 +217,7 @@ def test_model_to_netcdf_compression_non_numeric_coords(tmp_path: Path) -> None: d = pd.timedelta_range("1h", periods=3, freq="h", name="d") n = pd.Index(["north", "south"], name="n") x = m.add_variables(lower=0, coords=[t, d, n], name="x") - m.add_constraints(x >= 1, name="c", freeze=True) + m.add_constraints(x >= 1, name="c").to_sparse() m.add_objective(x.sum()) fn = tmp_path / "packed.nc" @@ -272,18 +270,16 @@ def test_model_to_netcdf_compression_scipy( assert_model_equal(model, read_netcdf(fn)) -def test_model_from_netcdf_frozen_constraint_legacy_positions(tmp_path: Path) -> None: +def test_model_from_netcdf_sparse_constraint_legacy_positions(tmp_path: Path) -> None: """Files written before #926 stored dense positions as CSR columns.""" - from linopy.constraints import CSRConstraint - - m = Model() + m = Model(sparse=True) i = pd.RangeIndex(3, name="i") mask = xr.DataArray([True, False, False], dims=["i"]) m.add_variables(coords=[i], name="a", mask=mask) b = m.add_variables(coords=[i], name="b") - m.add_constraints(2 * b >= 1, name="c", freeze=True) + m.add_constraints(2 * b >= 1, name="c") - fn = tmp_path / "test_frozen_legacy.nc" + fn = tmp_path / "test_sparse_legacy.nc" m.to_netcdf(fn) ds = xr.load_dataset(fn) @@ -293,14 +289,12 @@ def test_model_from_netcdf_frozen_constraint_legacy_positions(tmp_path: Path) -> assert not np.array_equal(labels, positions) ds["constraints-c-indices"] = xr.DataArray(positions, dims=["constraints-c-_nnz"]) del ds.attrs["constraints-c-_csr_columns"] - cindex = ds.attrs["constraints-c-cindex"] - ds = ds.rename({"constraints-c-_active_positions": "constraints-c-_con_labels"}) - ds["constraints-c-_con_labels"] = ds["constraints-c-_con_labels"] + cindex + ds = ds.rename({"constraints-c-labels": "constraints-c-_con_labels"}) legacy_fn = tmp_path / "legacy.nc" ds.to_netcdf(legacy_fn) p = read_netcdf(legacy_fn) - assert isinstance(p.constraints["c"], CSRConstraint) + assert p.constraints["c"].lhs.is_sparse A, A_read = m.matrices.A, p.matrices.A assert A is not None and A_read is not None np.testing.assert_array_equal(A.toarray(), A_read.toarray()) @@ -323,10 +317,8 @@ def test_model_to_netcdf_masked_variable_nan_bounds(tmp_path: Path) -> None: assert_model_equal(m, p.copy()) -def test_model_to_netcdf_mixed_sign_constraint(tmp_path: Path) -> None: - from linopy.constraints import CSRConstraint - - m = Model() +def test_model_to_netcdf_mixed_sign_constraint(tmp_path: Path, backend: Model) -> None: + m = backend x = m.add_variables(coords=[pd.RangeIndex(4, name="i")], name="x") def bound(m: Model, i: int) -> object: @@ -334,17 +326,15 @@ def bound(m: Model, i: int) -> object: return x.at[i] >= i return x.at[i] == 0.0 - m.add_constraints(bound, coords=[pd.RangeIndex(4, name="i")], name="c", freeze=True) - assert isinstance(m.constraints["c"], CSRConstraint) + m.add_constraints(bound, coords=[pd.RangeIndex(4, name="i")], name="c") + assert m.constraints["c"].lhs.is_sparse == m.sparse fn = tmp_path / "test_mixed_sign.nc" m.to_netcdf(fn) p = read_netcdf(fn) - assert isinstance(p.constraints["c"], CSRConstraint) - import numpy as np - - np.testing.assert_array_equal(m.constraints["c"]._sign, p.constraints["c"]._sign) + assert p.constraints["c"].lhs.is_sparse == m.sparse + xr.testing.assert_equal(m.constraints["c"].sign, p.constraints["c"].sign) assert_model_equal(m, p) @@ -1029,64 +1019,44 @@ def test_lp_and_direct_agree( assert m_direct.objective.value == m_lp.objective.value == -2 -def test_to_file_lp_frozen_vs_mutable(tmp_path: Path) -> None: - """Test that frozen and mutable constraints produce identical LP output.""" - m_frozen = Model() +def test_to_file_lp_sparse_vs_dense(tmp_path: Path) -> None: + """Sparse and dense models produce identical LP output.""" N = np.arange(5) - x = m_frozen.add_variables(coords=[N], name="x") - y = m_frozen.add_variables(coords=[N], name="y") - m_frozen.add_constraints(x + y <= 10, name="upper") - m_frozen.add_constraints(x >= 1, name="lower") - m_frozen.add_constraints(2 * x + y == 8, name="eq") - m_frozen.add_objective(x.sum() + 2 * y.sum()) - - m_mutable = Model() - x2 = m_mutable.add_variables(coords=[N], name="x") - y2 = m_mutable.add_variables(coords=[N], name="y") - m_mutable.add_constraints(x2 + y2 <= 10, name="upper", freeze=False) - m_mutable.add_constraints(x2 >= 1, name="lower", freeze=False) - m_mutable.add_constraints(2 * x2 + y2 == 8, name="eq", freeze=False) - m_mutable.add_objective(x2.sum() + 2 * y2.sum()) - - fn_frozen = tmp_path / "frozen.lp" - fn_mutable = tmp_path / "mutable.lp" - m_frozen.to_file(fn_frozen) - m_mutable.to_file(fn_mutable) - - assert fn_frozen.read_text() == fn_mutable.read_text() - - -def test_to_file_lp_frozen_mixed_sign(tmp_path: Path) -> None: - """Test LP writing for frozen constraint with per-row signs.""" - m_frozen = Model() - N = pd.RangeIndex(4, name="i") - x = m_frozen.add_variables(coords=[N], name="x") - - def bound(m: Model, i: int) -> object: - if i % 2: - return x.at[i] >= i - return x.at[i] <= 10 - - m_frozen.add_constraints(bound, coords=[N], name="mixed", freeze=True) - m_frozen.add_objective(x.sum()) - - m_mutable = Model() - x2 = m_mutable.add_variables(coords=[N], name="x") - - def bound2(m: Model, i: int) -> object: - if i % 2: - return x2.at[i] >= i - return x2.at[i] <= 10 - - m_mutable.add_constraints(bound2, coords=[N], name="mixed", freeze=False) - m_mutable.add_objective(x2.sum()) + texts = [] + for sparse in (True, False): + m = Model(sparse=sparse) + x = m.add_variables(coords=[N], name="x") + y = m.add_variables(coords=[N], name="y") + m.add_constraints(x + y <= 10, name="upper") + m.add_constraints(x >= 1, name="lower") + m.add_constraints(2 * x + y == 8, name="eq") + m.add_objective(x.sum() + 2 * y.sum()) + fn = tmp_path / f"{sparse}.lp" + m.to_file(fn) + texts.append(fn.read_text()) + assert texts[0] == texts[1] - fn_frozen = tmp_path / "frozen_mixed.lp" - fn_mutable = tmp_path / "mutable_mixed.lp" - m_frozen.to_file(fn_frozen) - m_mutable.to_file(fn_mutable) - assert fn_frozen.read_text() == fn_mutable.read_text() +def test_to_file_lp_sparse_mixed_sign(tmp_path: Path) -> None: + """LP writing for a mixed-sign constraint matches between backends.""" + N = pd.RangeIndex(4, name="i") + texts = [] + for sparse in (True, False): + m = Model(sparse=sparse) + x = m.add_variables(coords=[N], name="x") + + def bound(m: Model, i: int) -> object: + if i % 2: + return x.at[i] >= i + return x.at[i] <= 10 + + m.add_constraints(bound, coords=[N], name="mixed") + assert m.constraints["mixed"].lhs.is_sparse == sparse + m.add_objective(x.sum()) + fn = tmp_path / f"mixed_{sparse}.lp" + m.to_file(fn) + texts.append(fn.read_text()) + assert texts[0] == texts[1] def _lp_constraint_section(m: Model, path: Path, slice_size: int = 2_000_000) -> str: @@ -1094,7 +1064,6 @@ def _lp_constraint_section(m: Model, path: Path, slice_size: int = 2_000_000) -> return path.read_text().split("s.t.\n\n")[1].split("\n\nbounds")[0] -@pytest.mark.parametrize("freeze", [True, False]) @pytest.mark.parametrize( "slice_size", [ @@ -1123,18 +1092,18 @@ def _lp_constraint_section(m: Model, path: Path, slice_size: int = 2_000_000) -> ) def test_to_file_lp_constraint_section( tmp_path: Path, - freeze: bool, + backend: Model, slice_size: int, scaled: bool, expected: str, ) -> None: - m = Model() + m = backend i = pd.RangeIndex(2, name="i") row_scaling = 2.0 if scaled else 1.0 x = m.add_variables(coords=[i], name="x", scaling=[1.0, 2.0] if scaled else 1.0) y = m.add_variables(coords=[i], name="y") - m.add_constraints(x - 2 * y <= 1.5, name="a", freeze=freeze, scaling=row_scaling) - m.add_constraints(x.sum() >= -0.0, name="b", freeze=freeze, scaling=row_scaling) + m.add_constraints(x - 2 * y <= 1.5, name="a", scaling=row_scaling) + m.add_constraints(x.sum() >= -0.0, name="b", scaling=row_scaling) m.add_objective(x.sum()) fn = tmp_path / "constraints.lp" @@ -1146,7 +1115,7 @@ def test_to_file_lp_unsorted_constraint_labels(tmp_path: Path) -> None: i = pd.RangeIndex(3, name="i") x = m.add_variables(coords=[i], name="x") y = m.add_variables(coords=[i], name="y") - m.add_constraints(x + 2 * y <= 1, name="a", freeze=False) + m.add_constraints(x + 2 * y <= 1, name="a") m.add_objective(x.sum()) expected = _lp_constraint_section(m, tmp_path / "sorted.lp") diff --git a/test/test_matrices.py b/test/test_matrices.py index eeb86f226..83bc4e966 100644 --- a/test/test_matrices.py +++ b/test/test_matrices.py @@ -69,21 +69,18 @@ def test_matrices_duplicated_variables() -> None: assert np.isin(np.unique(np.array(A)), [0.0, 2.0]).all() -@pytest.mark.parametrize("freeze", [False, True]) -def test_matrices_drops_explicit_zeros(freeze: bool) -> None: +def test_matrices_drops_explicit_zeros(backend: Model) -> None: # https://github.com/PyPSA/linopy/issues/814 # Expressions that broadcast against a dense coordinate store one coefficient # per pair, most of them structurally zero. Those must not reach A, whose # stored nnz drives the direct-solver handoff (e.g. highspy.addRows). - m = Model() + m = backend x = m.add_variables(coords=[range(4)], name="x") coeff = xr.DataArray( np.eye(4), dims=["j", "i"], coords={"j": range(4), "i": range(4)} ) - m.add_constraints( - (coeff * x.rename(dim_0="i")).sum("i") <= 1, name="c", freeze=freeze - ) + m.add_constraints((coeff * x.rename(dim_0="i")).sum("i") <= 1, name="c") A = m.matrices.A assert A is not None @@ -120,8 +117,8 @@ def test_matrices_sol_aligned_with_vlabels() -> None: np.testing.assert_array_equal(M.sol, [1.0, 3.0, 4.0, 5.0, 6.0]) -@pytest.mark.parametrize("freeze", [False, True]) -def test_matrices_scaled_masked_mixed_signs(freeze: bool) -> None: +@pytest.mark.parametrize("sparse", [False, True], ids=["dense", "sparse"]) +def test_matrices_scaled_masked_mixed_signs(sparse: bool) -> None: m = Model() i = pd.RangeIndex(3, name="i") x = m.add_variables(0, 4, coords=[i], name="x", scaling=[1.0, 2.0, 4.0]) @@ -135,9 +132,9 @@ def test_matrices_scaled_masked_mixed_signs(freeze: bool) -> None: name="c", scaling=xr.DataArray([10.0, 1.0, 5.0], coords=[i]), mask=xr.DataArray([True, False, True], coords=[i]), - freeze=freeze, ) - m.add_constraints(x + 3 * y >= 1, name="d", freeze=not freeze) + d = m.add_constraints(x + 3 * y >= 1, name="d") + (m.constraints["c"] if sparse else d).to_sparse() m.add_indicator_constraints(z, 1, x <= 2, name="ind") M = m.matrices diff --git a/test/test_model.py b/test/test_model.py index f718fdaf2..3baf7f3a5 100644 --- a/test/test_model.py +++ b/test/test_model.py @@ -178,13 +178,12 @@ def test_remove_variable() -> None: assert not m.objective.vars.isin(x.labels).any() -@pytest.mark.parametrize("freeze", [False, True]) @pytest.mark.parametrize("quadratic", [False, True]) def test_remove_masked_variable_keeps_unrelated_constraints( - freeze: bool, quadratic: bool + backend: Model, quadratic: bool ) -> None: # https://github.com/PyPSA/linopy/issues/883 - m: Model = Model() + m = backend i = pd.Index(range(3), name="i") mask = [True, False, True] @@ -194,12 +193,10 @@ def test_remove_masked_variable_keeps_unrelated_constraints( # `b` is masked, so the constraint carries empty term slots (-1), but it # never references `a` - without_a = m.add_constraints( - b.sum() + c, EQUAL, 0, name="without_a", freeze=freeze - ) + without_a = m.add_constraints(b.sum() + c, EQUAL, 0, name="without_a") assert not without_a.has_variable(a) - with_a = m.add_constraints(a.sum() + c, EQUAL, 0, name="with_a", freeze=freeze) + with_a = m.add_constraints(a.sum() + c, EQUAL, 0, name="with_a") assert with_a.has_variable(a) if quadratic: diff --git a/test/test_optimization.py b/test/test_optimization.py index 91b44bf12..dfcfedf8a 100644 --- a/test/test_optimization.py +++ b/test/test_optimization.py @@ -292,7 +292,7 @@ def modified_model() -> Model: x = m.add_variables(coords=[lower.index], name="x", binary=True) y = m.add_variables(lower, name="y") - c = m.add_constraints(x + y, GREATER_EQUAL, 10, freeze=False) + c = m.add_constraints(x + y, GREATER_EQUAL, 10) y.update(lower=9) c.update(lhs=2 * x + y) diff --git a/test/test_persistent_snapshot_buffers.py b/test/test_persistent_snapshot_buffers.py index 949919fdf..25a612c9b 100644 --- a/test/test_persistent_snapshot_buffers.py +++ b/test/test_persistent_snapshot_buffers.py @@ -126,9 +126,9 @@ def test_csr_capture_deterministic(baseline_model: Model) -> None: np.testing.assert_array_equal(b1.data, b2.data) -def test_frozen_positional_indices_freed_without_holder(frozen_model: Model) -> None: - con = frozen_model.constraints["c2"] - label_index = frozen_model.variables.label_index +def test_sparse_positional_indices_freed_without_holder(sparse_model: Model) -> None: + con = sparse_model.constraints["c2"] + label_index = sparse_model.variables.label_index csr, _ = con.to_matrix(label_index) ref = weakref.ref(csr.indices) del csr @@ -137,17 +137,17 @@ def test_frozen_positional_indices_freed_without_holder(frozen_model: Model) -> np.testing.assert_array_equal(csr.indices, [0, 1, 2, 3, 4]) -def test_frozen_model_pickles_after_matrix_assembly(frozen_model: Model) -> None: - frozen_model.constraints.to_matrix() - restored = pickle.loads(pickle.dumps(frozen_model)) +def test_sparse_model_pickles_after_matrix_assembly(sparse_model: Model) -> None: + sparse_model.constraints.to_matrix() + restored = pickle.loads(pickle.dumps(sparse_model)) csr, _ = restored.constraints["c2"].to_matrix(restored.variables.label_index) np.testing.assert_array_equal(csr.indices, [0, 1, 2, 3, 4]) -def test_frozen_capture_keeps_buffer_identity(frozen_model: Model) -> None: - s1 = ModelSnapshot.capture(frozen_model) - frozen_model.reset_solution() - s2 = ModelSnapshot.capture(frozen_model) +def test_sparse_capture_keeps_buffer_identity(sparse_model: Model) -> None: + s1 = ModelSnapshot.capture(sparse_model) + sparse_model.reset_solution() + s2 = ModelSnapshot.capture(sparse_model) for name in s1.con_buffers: b1, b2 = s1.con_buffers[name], s2.con_buffers[name] assert b1.indptr is b2.indptr, name @@ -156,19 +156,19 @@ def test_frozen_capture_keeps_buffer_identity(frozen_model: Model) -> None: @pytest.fixture -def frozen_model() -> Model: - m = Model() +def sparse_model() -> Model: + m = Model(sparse=True) x = m.add_variables(0, 10, coords=[range(3)], name="x") y = m.add_variables(0, 5, coords=[range(2)], name="y") - m.add_constraints(2 * x >= 4, name="c1", freeze=True) - m.add_constraints(x.sum() + y.sum() <= 20, name="c2", freeze=True) + m.add_constraints(2 * x >= 4, name="c1") + m.add_constraints(x.sum() + y.sum() <= 20, name="c2") m.add_objective(x.sum()) return m -def test_frozen_con_buffers_keep_data_identity(frozen_model: Model) -> None: - label_index = frozen_model.variables.label_index - for name, con in frozen_model.constraints.items(): +def test_sparse_con_buffers_keep_data_identity(sparse_model: Model) -> None: + label_index = sparse_model.variables.label_index + for name, con in sparse_model.constraints.items(): b1 = _extract_con_buffers(con, label_index) b2 = _extract_con_buffers(con, label_index) assert b1.data is b2.data, name @@ -176,11 +176,11 @@ def test_frozen_con_buffers_keep_data_identity(frozen_model: Model) -> None: assert b1.indices is b2.indices, name -def test_frozen_con_positional_csr_follows_variable_changes() -> None: - m = Model() +def test_sparse_con_positional_csr_follows_variable_changes() -> None: + m = Model(sparse=True) m.add_variables(0, 10, coords=[range(3)], name="x") y = m.add_variables(0, 5, coords=[range(2)], name="y") - m.add_constraints(y >= 1, name="c", freeze=True) + m.add_constraints(y >= 1, name="c") con = m.constraints["c"] label_index = m.variables.label_index before = _extract_con_buffers(con, label_index).indices @@ -191,11 +191,11 @@ def test_frozen_con_positional_csr_follows_variable_changes() -> None: assert _extract_con_buffers(con, label_index).indices is after -def test_frozen_con_positional_csr_follows_csr_rebind() -> None: - m = Model() +def test_sparse_con_positional_csr_follows_csr_rebind() -> None: + m = Model(sparse=True) x = m.add_variables(0, 10, coords=[range(3)], name="x") z = m.add_variables(0, 1, name="z") - m.add_constraints(x + 1e-12 * z >= 1, name="c", freeze=True) + m.add_constraints(x + 1e-12 * z >= 1, name="c") con = m.constraints["c"] label_index = m.variables.label_index before = _extract_con_buffers(con, label_index) @@ -205,9 +205,9 @@ def test_frozen_con_positional_csr_follows_csr_rebind() -> None: np.testing.assert_array_equal(after.indices, [0, 1, 2]) -def test_untouched_frozen_constraint_needs_no_rebuild(frozen_model: Model) -> None: - snap = ModelSnapshot.capture(frozen_model) - diff = ModelDiff.from_snapshot(snap, frozen_model) +def test_untouched_sparse_constraint_needs_no_rebuild(sparse_model: Model) -> None: + snap = ModelSnapshot.capture(sparse_model) + diff = ModelDiff.from_snapshot(snap, sparse_model) assert isinstance(diff, ModelDiff) assert diff.is_empty diff --git a/test/test_scalar_constraint.py b/test/test_scalar_constraint.py index 4872fada2..cf5b37241 100644 --- a/test/test_scalar_constraint.py +++ b/test/test_scalar_constraint.py @@ -6,7 +6,7 @@ import linopy from linopy import GREATER_EQUAL, Model, Variable -from linopy.constraints import AnonymousScalarConstraint, ConstraintBase +from linopy.constraints import AnonymousScalarConstraint, Constraint @pytest.fixture @@ -32,20 +32,20 @@ def test_anonymous_scalar_constraint_type(x: Variable) -> None: def test_simple_constraint_type(m: Model, x: Variable) -> None: - c: ConstraintBase = m.add_constraints(x.at[0] >= 0) - assert isinstance(c, linopy.constraints.ConstraintBase) + c: Constraint = m.add_constraints(x.at[0] >= 0) + assert isinstance(c, linopy.constraints.Constraint) def test_compound_constraint_type(m: Model, x: Variable) -> None: - c: ConstraintBase = m.add_constraints(x.at[0] + x.at[1] >= 0) - assert isinstance(c, linopy.constraints.ConstraintBase) + c: Constraint = m.add_constraints(x.at[0] + x.at[1] >= 0) + assert isinstance(c, linopy.constraints.Constraint) def test_explicit_simple_constraint_type(m: Model, x: Variable) -> None: - c: ConstraintBase = m.add_constraints(x.at[0], GREATER_EQUAL, 0) - assert isinstance(c, linopy.constraints.ConstraintBase) + c: Constraint = m.add_constraints(x.at[0], GREATER_EQUAL, 0) + assert isinstance(c, linopy.constraints.Constraint) def test_explicit_compound_constraint_type(m: Model, x: Variable) -> None: - c: ConstraintBase = m.add_constraints(x.at[0] + x.at[1], GREATER_EQUAL, 0) - assert isinstance(c, linopy.constraints.ConstraintBase) + c: Constraint = m.add_constraints(x.at[0] + x.at[1], GREATER_EQUAL, 0) + assert isinstance(c, linopy.constraints.Constraint) diff --git a/test/test_scaling.py b/test/test_scaling.py index e605e6db4..40e7b565d 100644 --- a/test/test_scaling.py +++ b/test/test_scaling.py @@ -10,7 +10,6 @@ from linopy import Model, read_netcdf from linopy.constants import Result, Solution, Status -from linopy.constraints import CSRConstraint from linopy.solvers import licensed_solvers @@ -19,9 +18,8 @@ def _dense(matrix: Any) -> np.ndarray: return matrix.toarray() -@pytest.mark.parametrize("freeze", [False, True]) -def test_masked_scaling_in_matrices(freeze: bool) -> None: - m = Model() +def test_masked_scaling_in_matrices(backend: Model) -> None: + m = backend i = pd.RangeIndex(3, name="i") mask = xr.DataArray([True, False, True], coords=[i]) x = m.add_variables( @@ -38,7 +36,6 @@ def test_masked_scaling_in_matrices(freeze: bool) -> None: name="c", mask=mask, scaling=xr.DataArray([5.0, 6.0, 7.0], coords=[i]), - freeze=freeze, ) matrices = m.matrices @@ -137,8 +134,8 @@ def test_indicator_constraint_scaling_in_matrix() -> None: ) -def test_frozen_constraint_scaling_in_matrix() -> None: - m = Model() +def test_constraint_scaling_in_matrix(backend: Model) -> None: + m = backend i = pd.Index(["a", "b"], name="i") x = m.add_variables(coords=[i], name="x", scaling=[10.0, 100.0]) @@ -147,10 +144,9 @@ def test_frozen_constraint_scaling_in_matrix() -> None: 2 * x >= xr.DataArray([20.0, 40.0], coords=[i]), name="c", scaling=row_scaling, - freeze=True, ) - assert isinstance(con, CSRConstraint) + assert con.lhs.is_sparse == m.sparse np.testing.assert_allclose(m.matrices.b, [40.0, 160.0]) np.testing.assert_allclose( _dense(m.matrices.A), @@ -212,14 +208,13 @@ def test_indicator_constraint_lp_export_uses_scaled_values(tmp_path: Path) -> No assert f"<= {40.0 * 4}" in text -@pytest.mark.parametrize("freeze", [False, True]) -def test_assign_result_unscales_solution_objective_and_dual(freeze: bool) -> None: - m = Model() +def test_assign_result_unscales_solution_objective_and_dual(backend: Model) -> None: + m = backend i = pd.Index(["a", "b"], name="i") x = m.add_variables(coords=[i], name="x", scaling=[10.0, 100.0]) b = m.add_variables(binary=True, name="b", scaling=50.0) - m.add_constraints(x + b >= 1, name="c", scaling=[2.0, 4.0], freeze=freeze) + m.add_constraints(x + b >= 1, name="c", scaling=[2.0, 4.0]) m.add_objective(x.sum() + b, scaling=10.0) primal = np.full(m._xCounter, np.nan) @@ -241,11 +236,11 @@ def test_assign_result_unscales_solution_objective_and_dual(freeze: bool) -> Non np.testing.assert_allclose(m.constraints["c"].dual.values, [0.8, 3.2]) -def test_scaling_is_preserved_in_netcdf(tmp_path: Path) -> None: - m = Model() +def test_scaling_is_preserved_in_netcdf(tmp_path: Path, backend: Model) -> None: + m = backend i = pd.Index(["a", "b"], name="i") x = m.add_variables(coords=[i], name="x", scaling=[10.0, 100.0]) - m.add_constraints(x >= 1, name="c", scaling=[2.0, 4.0], freeze=True) + m.add_constraints(x >= 1, name="c", scaling=[2.0, 4.0]) m.add_objective(x.sum(), scaling=5.0) path = tmp_path / "scaled.nc" diff --git a/test/test_solvers.py b/test/test_solvers.py index a83d257a4..1ca85ed06 100644 --- a/test/test_solvers.py +++ b/test/test_solvers.py @@ -14,7 +14,7 @@ from linopy import GREATER_EQUAL, Model, solvers from linopy.constants import Result, Solution, Status -from linopy.constraints import CSRConstraint +from linopy.constraints import Constraint from linopy.solver_capabilities import ( SOLVER_REGISTRY, SolverFeature, @@ -167,14 +167,14 @@ def test_assign_result_explicit(simple_model: Model) -> None: assert float(simple_model.variables["y"].solution) == 2.0 -def test_assign_result_with_csr_constraints_avoids_data_reconstruction( +def test_assign_result_with_sparse_constraints_avoids_data_reconstruction( monkeypatch: pytest.MonkeyPatch, ) -> None: - m = Model() + m = Model(sparse=True) x = m.add_variables(coords=[range(3)], name="x") - m.add_constraints(x >= 0, name="c", freeze=True) + m.add_constraints(x >= 0, name="c") con = m.constraints["c"] - assert isinstance(con, CSRConstraint) + assert con.lhs.is_sparse primal = np.arange(m._xCounter, dtype=float) dual = np.arange(m._cCounter, dtype=float) + 10 @@ -184,10 +184,10 @@ def test_assign_result_with_csr_constraints_avoids_data_reconstruction( solver_name="mock", ) - def fail_data(self: CSRConstraint) -> None: - raise AssertionError("CSRConstraint.data was accessed") + def fail_data(self: Constraint) -> None: + raise AssertionError("Constraint.data was accessed") - monkeypatch.setattr(CSRConstraint, "data", property(fail_data)) + monkeypatch.setattr(Constraint, "data", property(fail_data)) m.assign_result(result) np.testing.assert_array_equal(m.variables["x"].solution.values, primal) diff --git a/test/test_store_conformance.py b/test/test_store_conformance.py index 3a3472056..893bf8fab 100644 --- a/test/test_store_conformance.py +++ b/test/test_store_conformance.py @@ -21,7 +21,7 @@ import linopy from linopy import LinearExpression, Model, QuadraticExpression, Variable from linopy.constants import TERM_DIM -from linopy.constraints import ConstraintBase, CSRConstraint +from linopy.constraints import Constraint from linopy.store import DenseStore from linopy.testing import assert_equivalent @@ -135,11 +135,8 @@ class Case(NamedTuple): build: Callable[[Model], LinearExpression] op: Callable[[LinearExpression], Any] expected: str - sparse_xfail: str | None = None -SIGN_EQ = "CSRConstraint normalizes the sign '==' to '=' while Constraint keeps it" - CASES: dict[str, Case] = { "sel": Case(grid, lambda e: e.sel(a=[0, 2]), KEPT), "isel": Case(grid, lambda e: e.isel(a=slice(0, 2)), KEPT), @@ -354,7 +351,26 @@ class Case(NamedTuple): grid, lambda e: e.to_constraint("<=", 1 * var(e, "y")), KEPT ), "to-constraint-expression-eq": Case( - grid, lambda e: e.to_constraint("==", 1 * var(e, "y")), KEPT, SIGN_EQ + grid, lambda e: e.to_constraint("==", 1 * var(e, "y")), KEPT + ), + "to-constraint-update-coeffs": Case( + grid, + lambda e: e.to_constraint("<=", 1).update(coeffs=2.0), + "update(coeffs=, variables=)", + ), + "to-constraint-update-lhs-off-grid": Case( + grid, + lambda e: e.to_constraint("<=", 1).update( + lhs=with_2d_coord(dense_linexpr(var(e, "x"))) + ), + "update(lhs=)", + ), + "constraint-data": Case(grid, lambda e: e.to_constraint("<=", 1).data, DATA_READ), + "constraint-sel": Case( + grid, lambda e: e.to_constraint("<=", 1).sel(a=[0, 1]), DATA_READ + ), + "constraint-getitem": Case( + grid, lambda e: e.to_constraint("<=", 1)[{"a": [0, 1]}], "__getitem__" ), "to-constraint-array-sign": Case( grid, @@ -446,7 +462,7 @@ def assert_same(res: Any, want: Any) -> None: pd.testing.assert_frame_equal(res.flat, want.flat) elif isinstance(res, LinearExpression): assert_equivalent(res, want) - elif isinstance(res, ConstraintBase): + elif isinstance(res, Constraint): assert_equivalent(res.lhs, want.lhs) xr.testing.assert_equal(res.sign, want.sign) xr.testing.assert_allclose(res.rhs, want.rhs) @@ -459,12 +475,8 @@ def assert_same(res: Any, want: Any) -> None: @pytest.mark.parametrize("case", list(CASES)) -def test_operation(backend: Model, case: str, request: pytest.FixtureRequest) -> None: - build, op, expected, sparse_xfail = CASES[case] - if backend.sparse and sparse_xfail is not None: - request.applymarker( - pytest.mark.xfail(strict=True, raises=AssertionError, reason=sparse_xfail) - ) +def test_operation(backend: Model, case: str) -> None: + build, op, expected = CASES[case] want = op(build(Model())) expr = build(backend) assert expr.is_sparse == backend.sparse @@ -482,8 +494,8 @@ def test_operation(backend: Model, case: str, request: pytest.FixtureRequest) -> if backend.sparse and expected in (KEPT, SPARSIFIED): if isinstance(res, LinearExpression): assert res.is_sparse - elif isinstance(res, ConstraintBase): - assert isinstance(res, CSRConstraint) + elif isinstance(res, Constraint): + assert res.lhs.is_sparse CONST_ASSIGNERS: dict[str, Callable[[LinearExpression, Any], Any]] = { From d75d04a6470bfe1fd8050c19a4084c6786355ae4 Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 23 Oct 2026 16:18:30 +0200 Subject: [PATCH 3/8] refactor(io): one netcdf layout written from the store, pre-1.0 files refused (#1048) Store.to_netcdf_vars/from_netcdf_vars dispatch on a _store_kind attribute for constraint lhs, model expressions and the objective; the csr/dense legacy reader branches are deleted. read_netcdf refuses files without the 1.0 layout attribute. Round trips parametrized over both backends. Step 4c of #1019. --- doc/api.rst | 5 +- doc/migrating-to-v1.rst | 8 ++ doc/release_notes.rst | 4 +- doc/user-guide.rst | 3 +- examples/creating-constraints.ipynb | 2 +- examples/sparse-models.ipynb | 6 +- linopy/common.py | 57 ---------- linopy/constraints.py | 134 ++-------------------- linopy/csr.py | 62 +++++------ linopy/expressions.py | 7 ++ linopy/io.py | 61 +++++----- linopy/objective.py | 10 +- linopy/store.py | 36 ++++++ test/data/pre-1.0-model.nc | Bin 0 -> 27215 bytes test/test_common.py | 60 ---------- test/test_io.py | 165 ++++++++++------------------ 16 files changed, 186 insertions(+), 434 deletions(-) create mode 100644 test/data/pre-1.0-model.nc diff --git a/doc/api.rst b/doc/api.rst index bd207cef2..0a56c6d3b 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -386,10 +386,7 @@ Structure Modification ------------ -``Constraint.update`` is the canonical mutation API. The legacy ``lhs`` / -``sign`` / ``rhs`` / ``coeffs`` / ``vars`` setters still forward to -``update`` but emit a ``DeprecationWarning`` and will be removed in a -future release. +``Constraint.update`` is the canonical mutation API. .. autosummary:: :toctree: generated/ diff --git a/doc/migrating-to-v1.rst b/doc/migrating-to-v1.rst index 45e8f6f46..2b64fb07b 100644 --- a/doc/migrating-to-v1.rst +++ b/doc/migrating-to-v1.rst @@ -124,3 +124,11 @@ Every row is a legacy guess that is now an explicit rule. The full, normative rule list lives in the :doc:`arithmetic convention `. + +Saved models (netcdf) +--------------------- + +The netcdf file layout changed: dense and sparse stores now share one layout. +:func:`~linopy.io.read_netcdf` refuses files written by linopy < 1.0 and raises a +``ValueError``. Rebuild the model with linopy 1.0 and save it again with +:meth:`~linopy.model.Model.to_netcdf`. diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 8a9e3faad..0f71794fb 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -55,7 +55,9 @@ Version 1.0.0 * In a sparse model, ``add_constraints`` and ``add_indicator_constraints`` hold the lhs on a ``CSRStore`` whenever the grid can hold its coordinates, also for a dense lhs, and ``Constraint.from_dataset`` reads into the model's store class. ``Constraint.to_sparse()`` and ``Constraint.to_dense()`` replace the lhs store in place. Explicit conversions emit no notice at all, so ``warn_on_densify=True`` no longer flags them; the ``explicit`` argument of the conversion notice is removed. (`#1047 `__) * A sparse constraint keeps ``rhs``, ``sign``, ``scaling`` and ``dual`` on its full grid like a dense one, instead of on its active rows; masked rows have label ``-1`` and rows without terms stay labelled but inactive. A masked row with an invalid infinite rhs now raises in a sparse model too. ``Constraint.active_labels``, ``to_matrix`` and ``to_polars`` raise ``ValueError`` on a constraint without labels. A sparse lhs reads its coefficients back as ``float64``. (`#1047 `__) * ``Constraint.sanitize_zeros`` drops near-zero terms from the lhs store (``Store.prune`` takes a tolerance), so the ``_term`` dimension of a dense constraint can shrink, instead of marking the terms ``-1``/NaN. ``Constraint.iterate_slices`` gathers row blocks with ``Store.take`` and keeps a sparse lhs sparse. ``MatrixAccessor`` and the dual write-back of ``Model.assign_result`` treat both stores alike. (`#1047 `__) -* Netcdf: a sparse constraint is still written in the CSR layout (``_linopy_format="csr"``), now over every cell of its grid with ``labels``, ``rhs``, ``scaling``, the sign and ``dual`` per cell; files written with the active-row layout of ``CSRConstraint`` still load, into a ``Constraint`` with a sparse lhs. (`#1047 `__) +* Netcdf: one layout for every linear expression, constraint lhs and the objective, written from its store: a ``_store_kind`` attribute (``"dense"`` or ``"csr"``), ``const`` on the grid dimensions with their coordinates, and either ``coeffs``/``vars`` on the grid dimensions and ``_term`` or the CSR arrays ``data``/``indices``/``indptr``. Constraint fields (``labels``, ``sign``, ``rhs``, ``scaling``, ``dual``, ``binary_var``, ``binary_val``, ``blocks``) lie on the grid dimensions. Each lhs, expression and objective reads back on the store class it was written with, so a dense constraint in a sparse model stays dense and vice versa. The ``_linopy_format``/``_csr_columns`` CSR layout and the active-row and multiindex legacy readers are removed. (`#1048 `__) +* ``read_netcdf`` refuses files written by linopy <1.0 with a ``ValueError``; files carry the layout version in the ``_linopy_netcdf_layout`` attribute next to ``_linopy_version``. Rebuild the model and write it again with linopy >=1.0. (`#1048 `__) +* ``linopy.testing.assert_conequal(strict=False)`` compares the lhs with ``assert_equivalent``, independently of its store and term layout, instead of comparing the dense datasets. (`#1047 `__) **Deprecations** diff --git a/doc/user-guide.rst b/doc/user-guide.rst index dcb09dd4a..f2e706373 100644 --- a/doc/user-guide.rst +++ b/doc/user-guide.rst @@ -77,7 +77,8 @@ Sparse models Models that aggregate many entities into few rows, such as generators into buses, can store their expressions as sparse matrices instead of padded dense arrays. Opt in with ``Model(sparse=True)``. ``groupby(...).sum()`` and ``@`` then return sparse-backed -expressions, and every constraint holds its left-hand side in CSR form. +expressions, and every constraint holds its left-hand side in CSR form +whenever the grid can hold its coordinates. Operations without a sparse path convert an expression to dense and emit a ``PerformanceWarning``. diff --git a/examples/creating-constraints.ipynb b/examples/creating-constraints.ipynb index 73700a110..a6c88c837 100644 --- a/examples/creating-constraints.ipynb +++ b/examples/creating-constraints.ipynb @@ -291,7 +291,7 @@ "source": [ "### Building a sparse model\n", "\n", - "Set `sparse=True` on the `Model` to build it on the sparse path: every constraint added via `add_constraints` holds its lhs in CSR form, and `groupby(...).sum()` and `@` return CSR-backed expressions. The setting is fixed at creation:" + "Set `sparse=True` on the `Model` to build it on the sparse path: every constraint added via `add_constraints` holds its lhs in CSR form whenever the grid can hold its coordinates, and `groupby(...).sum()` and `@` return CSR-backed expressions. The setting is fixed at creation:" ] }, { diff --git a/examples/sparse-models.ipynb b/examples/sparse-models.ipynb index c0a136d5b..0ccdcc094 100644 --- a/examples/sparse-models.ipynb +++ b/examples/sparse-models.ipynb @@ -31,7 +31,8 @@ "turns on the whole sparse path. In a sparse model\n", "`groupby(...).sum()`, `@`/`dot` against a constant and linear arithmetic on\n", "variables such as `2 * x` or `x + y` return CSR-backed expressions, and\n", - "`add_constraints` holds the left-hand side of every constraint in CSR form." + "`add_constraints` holds the left-hand side of every constraint in CSR form\n", + "whenever the grid can hold its coordinates." ] }, { @@ -102,7 +103,8 @@ "metadata": {}, "source": [ "Constraints in a sparse model are ordinary `Constraint` objects whose\n", - "left-hand side is CSR-backed, also when the expression is dense.\n", + "left-hand side is CSR-backed whenever the grid can hold its coordinates, also\n", + "when the expression is dense.\n", "`con.lhs.is_sparse` reports it. They support the whole `Constraint` API, and\n", "`update(...)` and `soften(...)` keep the left-hand side sparse." ] diff --git a/linopy/common.py b/linopy/common.py index bc4176c61..417c07f65 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -7,7 +7,6 @@ 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 @@ -1100,62 +1099,6 @@ def __getitem__( return self.object.sel(key) -def coords_to_dataset_vars(coords: list[pd.Index]) -> dict[str, DataArray]: - """ - Serialize a list of pd.Index to a DataArray dict. - - 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 names, since the netcdf reader splits variable names on ``-``. - Reconstruct with :func:`coords_from_dataset`. - """ - return { - f"_index{i}": DataArray(np.array(c), dims=[f"_indexdim{i}"]) - for i, c in enumerate(coords) - } - - -def coords_from_dataset( - ds: Dataset, coord_dims: list[str] -) -> tuple[list[pd.Index], dict[str, tuple[str, np.ndarray]]]: - """ - Deserialize a list of pd.Index from a Dataset. - - Reconstructs coordinates previously serialized by :func:`coords_to_dataset_vars`, - or by its earlier name-keyed format (``_coord_``). A MultiIndex - written by linopy < 1.0 in either format loads flat: a positional index, - with its levels returned as auxiliary coords on the dim. - """ - coords: list[pd.Index] = [] - levels: dict[str, tuple[str, np.ndarray]] = {} - for i, d in enumerate(coord_dims): - mi_levels = _multiindex_levels(ds, i, d) - if mi_levels is None: - key = f"_index{i}" if f"_index{i}" in ds else f"_coord_{d}" - coords.append(pd.Index(ds[key].values, name=d)) - continue - size = len(next(iter(mi_levels.values()))) - coords.append(pd.RangeIndex(size, name=d)) - levels |= {name: (d, values) for name, values in mi_levels.items()} - return coords, levels - - -def _multiindex_levels(ds: Dataset, i: int, d: str) -> dict[str, np.ndarray] | None: - if f"_index{i}_codes" in ds: - codes = ds[f"_index{i}_codes"] - names = [str(n) for n in json.loads(codes.attrs["level_names"])] - keys = [f"_index{i}_level{j}" for j in range(len(names))] - elif f"_coord_{d}_codes" in ds: - codes = ds[f"_coord_{d}_codes"] - prefix = f"_coord_{d}_level_" - names = [str(k)[len(prefix) :] for k in ds if str(k).startswith(prefix)] - keys = [prefix + n for n in names] - else: - return None - return {n: ds[k].values[c] for n, k, c in zip(names, keys, codes.values.T)} - - 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 991abe405..8b4810257 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -75,7 +75,7 @@ from linopy.csr import CSRStore, Grid, _densify_notice from linopy.scaling import validate_scaling from linopy.semantics import check_user_nan -from linopy.store import DenseStore, Reason, Store, as_sparse +from linopy.store import NETCDF_STORE_ATTRS, DenseStore, Reason, Store, as_sparse from linopy.types import ( ConstantLike, ConstraintLike, @@ -179,15 +179,6 @@ def rhs_on_grid(store: CSRStore, rhs: Any) -> DataArray | Reason: return grid.dataarray(values.astype(float), name="rhs") -def _rows_on_grid( - grid: Grid, positions: np.ndarray, values: np.ndarray, fill: Any -) -> DataArray: - """``values`` at the flat ``positions`` of ``grid``, ``fill`` elsewhere.""" - full = np.full(grid.size, fill, dtype=np.result_type(values, np.asarray(fill))) - full[positions] = values - return grid.dataarray(full) - - class Slack(NamedTuple): """ Slack variable(s) added by :meth:`Constraint.soften`. @@ -342,123 +333,20 @@ def from_dataset(cls, ds: Dataset, model: Model, name: str = "") -> Constraint: @classmethod def from_netcdf_ds(cls, ds: Dataset, model: Model, name: str) -> Constraint: - """ - Read a constraint written by :meth:`to_netcdf_ds`: the dense layout, - or the CSR layout (``_linopy_format="csr"``) into a sparse lhs. - - CSR files without the ``_csr_columns`` attribute hold dense variable - positions instead of labels. - """ - if ds.attrs.get("_linopy_format") != "csr": - return cls.from_dataset(ds, model, name) - attrs = ds.attrs - shape = tuple(attrs["shape"]) - indices = ds["indices"].values - columns = attrs.get("_csr_columns", "positions") - if columns == "positions": - vlabels = model.variables.label_index.vlabels - indices = vlabels[indices] - shape = (shape[0], int(vlabels[-1]) + 1 if len(vlabels) else 1) - elif columns != "labels": - raise ValueError( - f"Unknown CSR column format '{columns}' for constraint '{name}'." - ) - coord_dims = attrs["coord_dims"] - if isinstance(coord_dims, str): - coord_dims = [coord_dims] - grid = Grid.from_netcdf_vars(ds, coord_dims) - cindex_raw = int(attrs["cindex"]) - cindex = cindex_raw if cindex_raw >= 0 else None - rhs = ds["rhs"].values - if "_active_positions" in ds: - positions = ds["_active_positions"].values - elif "_con_labels" in ds: - positions = ds["_con_labels"].values - (cindex or 0) - else: - positions = np.arange(len(rhs), dtype=np.intp) - - counts = np.zeros(grid.size, dtype=np.int64) - counts[positions] = np.diff(ds["indptr"].values) - indptr = np.concatenate([[0], np.cumsum(counts)]) - csr = scipy.sparse.csr_array( - (ds["data"].values, indices, indptr), shape=(grid.size, shape[1]) - ) - store = CSRStore(csr, np.zeros(grid.size), grid, model) - - def on_grid(values: np.ndarray, fill: Any) -> DataArray: - return _rows_on_grid(grid, positions, values, fill) - - sign = ds["_sign"].values if "_sign" in ds else np.full(len(rhs), attrs["sign"]) - fields: dict[str, Any] = { - "sign": on_grid(sign, EQUAL), - "rhs": on_grid(rhs, np.nan), - } - if "scaling" in ds: - fields["scaling"] = on_grid(ds["scaling"].values, 1.0) - if cindex is not None: - labels = ds["labels"].values if "labels" in ds else positions + cindex - fields["labels"] = on_grid(labels, -1).astype(model._dtypes["labels"]) - if "dual" in ds: - fields["dual"] = on_grid(ds["dual"].values, np.nan) - if "_binvar_labels" in ds: - fields["binary_var"] = on_grid(ds["_binvar_labels"].values, -1) - binval = ds["_binval"].values if "_binval" in ds else attrs["binval"] - fields["binary_val"] = on_grid(np.broadcast_to(binval, len(rhs)), -1) - if "blocks" in ds: - fields["blocks"] = on_grid(ds["blocks"].values, -1) - con_attrs: dict[str, Any] = {} - if cindex is not None: - con_attrs["label_range"] = (cindex, cindex + grid.size) - if (slack := _slack_names(attrs)) is not None: - con_attrs["slack_positive"], con_attrs["slack_negative"] = slack + """Read a constraint written by :meth:`to_netcdf_ds`.""" + store = Store.from_netcdf_vars(ds, model) lhs = expressions.LinearExpression._from_store(store, model) - return cls(lhs, model=model, name=name, attrs=con_attrs, **fields) + fields = {k: ds[k] for k in _FIELDS if k in ds} + attrs = {k: v for k, v in ds.attrs.items() if k not in NETCDF_STORE_ATTRS} + return cls(lhs, model=model, name=name, attrs=attrs, **fields) def to_netcdf_ds(self) -> Dataset: """ - A Dataset for netcdf serialization: :attr:`data`, or for a sparse lhs - its raw CSR rows next to the flat fields, read by :meth:`from_netcdf_ds`. + A Dataset for netcdf serialization: the lhs store arrays, see + :meth:`Store.to_netcdf_vars`, next to the fields on the grid dims. """ - store = self._lhs._store - if not isinstance(store, CSRStore): - return self.data - csr = store.csr - - def flat(da: DataArray) -> DataArray: - return DataArray(da.values.ravel(), dims=["_flat"]) - - sign = self._sign.values.ravel() - data_vars: dict[str, DataArray] = { - "indptr": DataArray(csr.indptr, dims=["_indptr"]), - "indices": DataArray(csr.indices, dims=["_nnz"]), - "data": DataArray(csr.data, dims=["_nnz"]), - "rhs": flat(self._rhs), - "scaling": flat(self._scaling), - **store.grid.to_netcdf_vars(), - } - attrs: dict[str, Any] = { - "_linopy_format": "csr", - "_csr_columns": "labels", - "cindex": self.range[0] if self.is_assigned else -1, - "shape": list(csr.shape), - "coord_dims": list(store.grid.dims), - "name": self.name, - } - if sign.size and (sign == sign[0]).all(): - attrs["sign"] = str(sign[0]) - else: - data_vars["_sign"] = flat(self._sign) - optional = { - "labels": self._labels, - "dual": self._dual, - "_binvar_labels": self._binary_var, - "_binval": self._binary_val, - "blocks": self._blocks, - } - data_vars |= {k: flat(v) for k, v in optional.items() if v is not None} - if (slack := _slack_names(self._attrs)) is not None: - attrs["slack_positive"], attrs["slack_negative"] = slack - return Dataset(data_vars, attrs=attrs) + ds = self._lhs._store.to_netcdf_vars().assign(self._fields()) + return ds.assign_attrs(self._attrs) def _bind( self, lhs: expressions.LinearExpression, fields: Mapping[Any, Any] @@ -720,7 +608,7 @@ def binary_var(self) -> DataArray | None: return self._binary_var @property - def binary_val(self) -> int | np.ndarray | None: + def binary_val(self) -> np.ndarray | None: return None if self._binary_val is None else self._binary_val.values @property diff --git a/linopy/csr.py b/linopy/csr.py index ec9768eb8..f7b852154 100644 --- a/linopy/csr.py +++ b/linopy/csr.py @@ -50,11 +50,17 @@ import scipy.sparse from xarray import DataArray, Dataset, broadcast -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 check_arithmetic_nan, warn_outside_linopy -from linopy.store import DenseStore, Reason, Store, merged_grid +from linopy.store import ( + CSR_COLUMNS_ATTR, + STORE_KIND_ATTR, + DenseStore, + Reason, + Store, + merged_grid, +) if TYPE_CHECKING: from linopy.common import VariableLabelIndex @@ -66,7 +72,6 @@ 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" @@ -91,11 +96,6 @@ def __post_init__(self) -> None: self, "aux", {n: (d, _readonly(v)) for n, (d, v) in self.aux.items()} ) - @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 | DataArray, dims: Iterable[str]) -> Grid: """Build from the indexes and auxiliary coordinates ``ds`` carries on ``dims``.""" @@ -103,32 +103,6 @@ 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) - } - indexes, levels = coords_from_dataset(ds, list(dims)) - return cls(cls.from_coords(indexes).indexes, aux | levels) - - 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) @@ -538,6 +512,26 @@ def copy(self, model: Model, deep: bool = True) -> CSRStore: return replace(self, model=model) return replace(self, csr=self.csr.copy(), const=self.const.copy(), model=model) + def to_netcdf_vars(self) -> Dataset: + csr = self.csr + data_vars = { + "data": DataArray(csr.data, dims=["_nnz"]), + "indices": DataArray(csr.indices, dims=["_nnz"]), + "indptr": DataArray(csr.indptr, dims=["_indptr"]), + "const": self.grid.dataarray(self.const), + } + attrs = {STORE_KIND_ATTR: "csr", CSR_COLUMNS_ATTR: csr.shape[1]} + return Dataset(data_vars, attrs=attrs) + + @classmethod + def read_netcdf_vars(cls, ds: Dataset, model: Model) -> CSRStore: + const = ds["const"] + grid = Grid.from_dataset(const, map(str, const.dims)) + arrays = (ds["data"].values, ds["indices"].values, ds["indptr"].values) + shape = (grid.size, int(ds.attrs[CSR_COLUMNS_ATTR])) + csr = scipy.sparse.csr_array(arrays, shape=shape) + return cls(csr, const.values.reshape(-1), grid, model) + def dataarray(self, values: np.ndarray, name: str) -> DataArray: return self.grid.dataarray(values, name=name) diff --git a/linopy/expressions.py b/linopy/expressions.py index 9b307850c..47a33927a 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -1459,6 +1459,10 @@ def _to_dataset(self) -> Dataset: """The dense dataset, without converting a sparse expression.""" return self.data + def _to_netcdf_ds(self) -> Dataset: + """The dataset written to netcdf by :func:`linopy.io.to_netcdf`.""" + return self._to_dataset() + @property def model(self) -> Model: return self._model @@ -2140,6 +2144,9 @@ def data(self) -> Dataset: def _to_dataset(self) -> Dataset: return self._store.to_dataset() + def _to_netcdf_ds(self) -> Dataset: + return self._store.to_netcdf_vars() + def _fallback(self, reason: Reason | None) -> Self: """ A dense copy for an operation without a sparse path, with a densify diff --git a/linopy/io.py b/linopy/io.py index 9b4473c81..f6a344368 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -29,11 +29,10 @@ sos_weights, to_polars, ) -from linopy.constants import FACTOR_DIM, SOS_DIM_ATTR, SOS_TYPE_ATTR -from linopy.csr import CSRStore +from linopy.constants import SOS_DIM_ATTR, SOS_TYPE_ATTR from linopy.objective import Objective, linear_part from linopy.scaling import constraint_scaling_lookup, variable_scaling_lookup -from linopy.store import DenseStore, expression_dataset +from linopy.store import Store if TYPE_CHECKING: from cuopt.linear_programming import DataModel as cuoptDataModel @@ -47,6 +46,8 @@ logger = logging.getLogger(__name__) NETCDF_VERSION_ATTR = "_linopy_version" +NETCDF_LAYOUT_ATTR = "_linopy_netcdf_layout" +NETCDF_LAYOUT = 1 EXPR_TYPE_ATTR = "_linopy_expr_type" CONTAINER_ORDER_ATTR = "_linopy_{}_order" @@ -1106,7 +1107,7 @@ def with_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: ] exprs = [ with_prefix( - expr._to_dataset().assign_attrs(name=name, **{EXPR_TYPE_ATTR: expr.type}), + expr._to_netcdf_ds().assign_attrs(name=name, **{EXPR_TYPE_ATTR: expr.type}), f"expressions-{name}", ) for name, expr in m.expressions.items() @@ -1126,6 +1127,7 @@ def with_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: ds = xr.merge(vars + cons + exprs + obj + params, combine_attrs="drop_conflicts") ds = ds.assign_attrs(scalars) ds.attrs[NETCDF_VERSION_ATTR] = version("linopy") + ds.attrs[NETCDF_LAYOUT_ATTR] = NETCDF_LAYOUT for kind, container in ( ("variables", m.variables), ("expressions", m.expressions), @@ -1178,8 +1180,8 @@ def read_netcdf(path: Path | str, **kwargs: Any) -> Model: Notes ----- Variables, expressions and constraints are restored in the insertion - order stored by :func:`to_netcdf`. Files written by earlier versions - carry no order and are loaded in sorted name order. + order stored by :func:`to_netcdf`. Files written by linopy <1.0 raise + ``ValueError``. The SOS reformulation lifecycle token is not persisted by :func:`to_netcdf`. If the saved model was in reformulated form, @@ -1195,7 +1197,13 @@ def read_netcdf(path: Path | str, **kwargs: Any) -> Model: path = Path(path) ds = xr.load_dataset(path, **kwargs) - m = Model(sparse=bool(ds.attrs.get("sparse", False))) + if ds.attrs.get(NETCDF_LAYOUT_ATTR) != NETCDF_LAYOUT: + written_by = ds.attrs.get(NETCDF_VERSION_ATTR, "unknown") + raise ValueError( + f"netcdf written by linopy <1.0 (version {written_by}) cannot be read; " + "rebuild the model and write it with linopy >=1.0." + ) + m = Model(sparse=bool(ds.attrs["sparse"])) def remove_prefix(k: str, prefix: str) -> str: return k[len(prefix) + 1 :] @@ -1212,12 +1220,13 @@ def get_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: for k, v in ds.attrs.items() if owners[k] == prefix } - - for dim in list(ds.dims): - if ds.attrs.pop(f"{dim}_multiindex", None) is not None: - ds = ds.assign_coords({dim: pd.RangeIndex(ds.sizes[dim], name=dim)}) return ds + def read_expression(ds: xr.Dataset) -> LinearExpression | QuadraticExpression: + if ds.attrs[EXPR_TYPE_ATTR] == QuadraticExpression.__name__: + return QuadraticExpression(xr.Dataset(ds.data_vars, ds.coords), m) + return LinearExpression._from_store(Store.from_netcdf_vars(ds, m), m) + def container_names(kind: str) -> list[str]: found = {str(k).rsplit("-", 1)[0] for k in ds if str(k).startswith(kind)} order_attr = ds.attrs.get(CONTAINER_ORDER_ATTR.format(kind)) @@ -1257,20 +1266,7 @@ def owner(k: str) -> str | None: expressions: dict[str, LinearExpression | QuadraticExpression] = {} for k in containers["expressions"]: name = remove_prefix(k, "expressions") - expr_ds = get_prefix(ds, k) - expr_type = expr_ds.attrs.pop(EXPR_TYPE_ATTR, None) - expr_ds.attrs.pop("name", None) # re-attached below, after construction - expr: LinearExpression | QuadraticExpression - if expr_type == "QuadraticExpression" or ( - expr_type is None and FACTOR_DIM in expr_ds.dims - ): - expr = QuadraticExpression(expr_ds, m) - else: - dense = DenseStore(expression_dataset(expr_ds, m), m) - store = dense.to_sparse() if m.store_cls is CSRStore else dense - expr = LinearExpression._from_store( - dense if isinstance(store, str) else store, m - ) + expr = read_expression(get_prefix(ds, k)) expr.attrs["name"] = name expressions[name] = expr @@ -1283,20 +1279,13 @@ def owner(k: str) -> str | None: m._constraints = Constraints(constraints, m) objective = get_prefix(ds, "objective") - obj_type = objective.attrs.pop(EXPR_TYPE_ATTR, None) - obj_cls: type[LinearExpression] | type[QuadraticExpression] = ( - QuadraticExpression - if obj_type == "QuadraticExpression" - or (obj_type is None and FACTOR_DIM in objective.dims) - else LinearExpression - ) m.objective = Objective( - obj_cls(objective, m), + read_expression(objective), m, - objective.attrs.pop("sense"), - objective.attrs.pop("scaling", 1), + objective.attrs["sense"], + objective.attrs["scaling"], ) - m.objective._value = objective.attrs.pop("value", None) + m.objective._value = objective.attrs.get("value") m.parameters = get_prefix(ds, "parameters") diff --git a/linopy/objective.py b/linopy/objective.py index 32b1b64c6..46b698b22 100644 --- a/linopy/objective.py +++ b/linopy/objective.py @@ -159,14 +159,8 @@ def linear_terms(self) -> tuple[np.ndarray, np.ndarray]: 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. - """ - return self.expression._to_dataset().assign_attrs(name=self.name) + """Returns the objective expression in its netcdf layout, see :func:`linopy.io.to_netcdf`.""" + return self.expression._to_netcdf_ds().assign_attrs(name=self.name) @property def coeffs(self) -> DataArray: diff --git a/linopy/store.py b/linopy/store.py index fe28a36fd..da9e12328 100644 --- a/linopy/store.py +++ b/linopy/store.py @@ -52,6 +52,11 @@ FILL_VALUE = {"vars": -1, "coeffs": np.nan, "const": np.nan} """What a dense expression dataset holds at a cell or term it lacks.""" +STORE_KIND_ATTR = "_store_kind" +CSR_COLUMNS_ATTR = "_n_columns" +NETCDF_STORE_ATTRS = frozenset((STORE_KIND_ATTR, CSR_COLUMNS_ATTR)) +"""The netcdf attributes a store writes next to its arrays.""" + class Store(ABC): """ @@ -128,6 +133,29 @@ def to_dataset(self) -> Dataset: def copy(self, model: Model, deep: bool = True) -> Self: """The same store on ``model``, its arrays copied if ``deep``.""" + @abstractmethod + def to_netcdf_vars(self) -> Dataset: + """ + The arrays for netcdf with ``const`` on the grid dims and the store + kind in the ``_store_kind`` attribute, read by :meth:`from_netcdf_vars`. + """ + + @classmethod + @abstractmethod + def read_netcdf_vars(cls, ds: Dataset, model: Model) -> Self: + """Rebuild this store kind from the arrays of :meth:`to_netcdf_vars`.""" + + @staticmethod + def from_netcdf_vars(ds: Dataset, model: Model) -> Store: + """The store written by :meth:`to_netcdf_vars`, by its ``_store_kind``.""" + from linopy.csr import CSRStore + + kinds: dict[str, type[Store]] = { + "dense": DenseStore, + "csr": CSRStore, + } + return kinds[ds.attrs[STORE_KIND_ATTR]].read_netcdf_vars(ds, model) + @property @abstractmethod def nterm(self) -> int: @@ -497,6 +525,14 @@ def to_dataset(self) -> Dataset: def copy(self, model: Model, deep: bool = True) -> DenseStore: return DenseStore(self.ds.copy(deep=deep), model) + def to_netcdf_vars(self) -> Dataset: + ds = Dataset({k: self.ds[k] for k in ("coeffs", "vars", "const")}) + return ds.assign_attrs({STORE_KIND_ATTR: "dense"}) + + @classmethod + def read_netcdf_vars(cls, ds: Dataset, model: Model) -> DenseStore: + return cls(Dataset({k: ds[k] for k in ("coeffs", "vars", "const")}), model) + def dataarray(self, values: np.ndarray, name: str) -> DataArray: const = self.ds.const.transpose(*self.dims) return const.copy(data=values.reshape(const.shape)).rename(name) diff --git a/test/data/pre-1.0-model.nc b/test/data/pre-1.0-model.nc new file mode 100644 index 0000000000000000000000000000000000000000..c8406efbda56de9ae4239b9be7e63a1a7e2e6556 GIT binary patch literal 27215 zcmeHP3v`s#wcdYb5<-ZC2qd5phae$@Bqk(UXyp+y35jCjB`K*=7=~mL1}8IFGXn&Q zN?kr^uh$k(*IkxJq4>O3>!X&XSZlSuSlHIey69Dxr4K8&Shd$%m)DDT@3YUI`TxmG z$iU=M>KrBK|IfeA&9ob$&$WIT>Kklkr7h$IMm?uVgUg3LED)=^VbSA8>K{ZyJH5 z5E%9OxlcAOZn^?UkTYORw6e(a))8FxfJmIsg@~Ko18cz6zzO`)>V33kHG-xrS0vk^A*)dqB;vJBUL81Q5>NvS+)I=#x`sV z?$x80myn!m(2%ZxejEv~;=y+-_{oKuK=S1~M~4@*FKWcmxv^uo3}7T?n?8KSfC9pB z0e}Oud4F!rfF8VnY24)Aw`Yxp2WW&yL27<@n{xsjr#5N83$JDY6k!%D>Z^jQiJ}gm0JEY`kIe!%;fn?Nm_52RXD-}s z;8h0z)tqo2!tq18Y63VHGuLN@&9IvQt_Ct>S`IK4 zvyAzhuZ9b$jX3}lF+U(U{0De>#wW%9 zkqU44C;c{NR$GvOSBU~SswA$EIB~hwfG)a0(20K*9>JgK$(I3$2~a|Sq6Cm+ScDE9 zpeKs~a_EU!lq}GLhq^`qCM|*6=mv#*G~GNyH>Ci1bW=baMtN5HZI&U}LrW$sE3nH3)-7okQ?UFW|u6#{3q_-!=l%_eBa(~(YamMAKY!x)eZH{4_Y*=W6M831ZER*U|;!T z7WD@&SRd)S=jE(+6H&1G%(`sXgD}hMNP0^s#q7@ zo+zO`C0{}%bbK&5&;fcw57B_a-+8Z5QNv84qo7+?=~WueqDU{3wl+SlsIrKPVk1bU zViFW}(y8uCnfNm z#`jjhQwVPbJiG7VRq%wWcJZ)8fvqSg6 zHLBg)UR(N8_ycxa4v!oh`3toSoenp>H1{CP$I|KWyyMb;ht=3c862DII1FA@n)yz_ zZ-D|dCD3xw3%>)J$jadNzq0=Y#^cB;p#4DjLwFP~C&RrTz5f}!i3Bs?1;@(g@Ztzio0RtO4( ztHh%BT-$Ep&$q%d5~mK%{`kvSl}5{jesY5H;soVW5|m#M zQ~t7M+|-!z(b~F{Bq*O2Q@&U;vNS>Y^qBHHG^5L6%B3AtI+X;UH?$zgJRe)G$>B#4 z11lxzJdHpafiwbX1kwnk5lADDMj(wq8i6zdrxyW=QBd$V3;dVzNPeydlzhAazEbtv zu@h{08|4PPM!|s{Z@#z@Cdo2c_xb!>&q(Gu* za3GZ{>X+zI8b^gS^FAxS@q_dG4uu2 zd4j@G!!F@S3(cCoK2I>u$?P`l5~CbzZYR{~_Idq1#f?HGq06()R6h4gk1y2F#6KqJ70r7Zgu{zOFT<$1lHWV+T z5FEGAQh!A-fG-sG;xQenK~WH@wbm1)(++cHmg4~kJkK+ywz9gavbua#SNBD8W>(Iu ztU{u6N+Xa)VE7QQ>OWc5f~Lb>8=ibtq4=nA(*`HWi2MihY`A}@3_-LhF)E(^zMw0- zzRzPhgty4+_qc-%_{Mn;$p z1!7e)TU|&I@T(pugip0_e;`MnB3_v0G9`b~_D*6VlP{4;%iOU-}Arlpf3dj+<9pzypQE!Z%)bzBdr!Xa5w5dSS zg;Q8nWmRQgAQY~y?DDLwuAb2|v#MrRt;ch5jSctbcv_Seu8ut(grvauAtlB;~9V$bxlcfz9W(4M53$S z+uR^A-=diPHoP6s?@H9>;Q9ewstq}8UGiv``k1h}D?E1U=^28WM;j>`IFOL1*GncP z?Dek+pd~UNcaCQ>qKy^u%bUC?eFioHf$&^8Qa1d!L%}exg7WB2nhC`N$02^RmPb_0 zUUyh=g?LB&>J#D^6{3j=n+mOR z<2(|}5fMkpqb=e1FcHv7EX_d6;OKNY87Y1%nLKcA)G8dOPK2F;0k<_+m5{fGXD1pX z{ksHr`SW;G(J3?1FfAB}D)NYC74-^{l8O)=uePudCq%3hFEutbFNl0=s|a5I64Sk2 zqw@yS=np zHs90cv-w^&silqBT-qG0UwTG#l+5Flx*}(C&y+Afem+i@$L8_hnc9f$oBAW)P1d_a z!zq3nG2;^Abe={ajX)ZKGy-V^21S4me0x%U;CtF7+^rn+l218?{lNEWJb}ok%2zb= zo5wj-zOUbXK%J6YM3X!@hhylJWb|XjMa^yP>Q@Qq93P#XqqTIr1qn_2;wF)+Q1GiC z)F}thyaDjKKCInwIZ|Lcx{`D8|J=4^$8>C(9{}v%G1(X{k(l>|6^SPji$n={e=VC* z+3-mRGvbKmM^c+o>`La1{NY}aF*-k+8X0rofZChi@nQBi=cC_a5y>0OYD7LGCU2I` zTxaHxvui7(d%V~rpN4d3v5<`_u@r-Rd{{>pau=xdlKK$Bc>`z zj(Vt8y(yeDzfLVCV|HXLlGB)EYAeEP^F2joleUR_9)7yWEQi}P(qeF^!6DlL8ck^2 zW&eel8X4v5KRbHJy8Djr^2>(&g@#=pQwPhoHGeeU8Sr#>hpH~47O*Y5e1uY9jz%#l z#y>+HtwJbg{^Jl*2x^x8g#p7(>OpDzz!rn6LOqN5sbQBdTH*|nBY$_c3#oiJFmYg$ zdRyt{B|?ijt&)XShvfJmz#Cr^TKTRZ-l+kk^>%tfH*14y*hQ=3of;r>QHwH2$*_|P zMacq{9s)wJa5+CV?Bb@z>tDavJ{J?m|I22gCK)qv!Wok1oBu{lAG0+zRED)pRZnQ>Q6)K%=#(rsaW&ejp8k*YMzj>Ngh!`H6HxE zG0`=B|IohfVrKA^!CW3?_a_>Lnzu;#=v~Q&&n%!zTbU@Nm^ad7nH_g&t5l`xpwjB! zPaIY^hz6Ln`D3yC{X!mS95<--;49x?xKh8a{H2CT6jopw)^(+V$;q1V|NLbeb9j`%l_Ltgtho%HQIfU zlHG51w+xqk`0E$eJ{)dieB+4FuZ@Qc^9(y`4wH%z&D~Ew1)9p0n4FZ6{8a5jtO$vh zDI#^ttux;oXg|b*kUaX^W?YguC^1DNBr~?ykBr1lR~?Pzl!6kpkjzDIFFA*cW(Z2= zTpWxGu>8}fvi`@+d~vVj)pC6_>{K$}JQ5n2xcPKVY02WP*A8TsETbd&3H@VnooY0* zVvYT+aub>d9BxaKPkAe ziA7@UU?GXf94d!$%^_$V{)p_v+?+rZhIJk#4>7g#=yYuH>+>3&){rck{otPL8w|Vr z#7MkQAAO+Y>93MG(6FDQQI8kHmT) zg&q2e literal 0 HcmV?d00001 diff --git a/test/test_common.py b/test/test_common.py index 9588e8fb4..58195b502 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -15,8 +15,6 @@ from linopy.common import ( assign_fields, best_int, - coords_from_dataset, - coords_to_dataset_vars, get_dims_with_index_levels, get_printout_labels, is_constant, @@ -54,64 +52,6 @@ def test_assign_fields_replaces_without_aligning() -> None: assert result.attrs == {"a": 1} -def test_coords_dataset_vars_roundtrip() -> None: - """Dash-named coords survive serialization to Dataset vars and back.""" - coords = [pd.Index([1, 2, 3], name="my-simple"), pd.Index(["a"], name="b-c")] - ds = xr.Dataset(coords_to_dataset_vars(coords)) - assert not any("-" in str(k) for k in ds) - restored, levels = coords_from_dataset(ds, ["my-simple", "b-c"]) - assert levels == {} - for got, want in zip(restored, coords): - assert got.equals(want) - assert got.name == want.name - - -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_simple": ("e", [1, 2, 3])}) - (plain,), _ = coords_from_dataset(ds, ["simple"]) - assert plain.equals(pd.Index([1, 2, 3])) and plain.name == "simple" - - -@pytest.mark.parametrize( - "ds", - [ - pytest.param( - xr.Dataset( - { - "_coord_d_level_p": ("l0", [1, 2]), - "_coord_d_level_q": ("l1", ["a", "b"]), - "_coord_d_codes": (("n", "k"), [[0, 0], [0, 1], [1, 0]]), - } - ), - id="name-keyed-0.9", - ), - pytest.param( - xr.Dataset( - { - "_index0_level0": ("l0", [1, 2]), - "_index0_level1": ("l1", ["a", "b"]), - "_index0_codes": ( - ("n", "k"), - [[0, 0], [0, 1], [1, 0]], - {"level_names": '["p", "q"]'}, - ), - } - ), - id="positional-0.10", - ), - ], -) -def test_coords_from_dataset_loads_multiindex_flat(ds: xr.Dataset) -> None: - """A sparse (CSR) constraint over a MultiIndex from linopy < 1.0 loads flat.""" - (index,), levels = coords_from_dataset(ds, ["d"]) - assert index.equals(pd.RangeIndex(3)) and index.name == "d" - assert {n: (d, v.tolist()) for n, (d, v) in levels.items()} == { - "p": ("d", [1, 1, 2]), - "q": ("d", ["a", "b", "a"]), - } - - def test_iterate_slices_basic() -> None: ds = xr.Dataset( {"var": (("x", "y"), np.random.rand(10, 10))}, # noqa: NPY002 diff --git a/test/test_io.py b/test/test_io.py index db9817bb9..f12011f1c 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -24,7 +24,12 @@ from linopy.constraints import Constraint from linopy.expressions import LinearExpression, QuadraticExpression from linopy.io import CONTAINER_ORDER_ATTR, signed_number -from linopy.testing import assert_equivalent, assert_exprequal, assert_model_equal +from linopy.testing import ( + assert_conequal, + assert_equivalent, + assert_exprequal, + assert_model_equal, +) HAS_NETCDF4 = importlib.util.find_spec("netCDF4") is not None @@ -171,21 +176,64 @@ def test_read_netcdf_inconsistent_order_attr_raises( read_netcdf(fn) -def test_model_to_netcdf_keeps_constraint_store(tmp_path: Path, backend: Model) -> None: +@pytest.fixture +def mixed_store_model(backend: Model) -> Model: m = backend - x = m.add_variables(coords=[pd.RangeIndex(5, name="i")], name="x") - m.add_constraints(x >= 1, name="c") + i = pd.RangeIndex(4, name="i") + x = m.add_variables(0, 10, coords=[i], name="x") + b = m.add_variables(binary=True, coords=[i], name="b") + m.add_expressions(3 * x + 1, name="e") + flipped = m.add_constraints(x - b >= 0, name="flipped") + (flipped.to_dense if m.sparse else flipped.to_sparse)() + m.add_expressions(flipped.lhs, name="flipped") + m.add_objective(2 * x.sum() + b.sum()) + mask = xr.DataArray([True, False, True, True], coords=[i]) + m.add_constraints(2 * x >= 1, name="masked", mask=mask) + sign = xr.DataArray(["<=", ">=", "=", "<="], coords=[i]) + m.add_constraints(x + b, sign, 3, name="mixed") + m.add_indicator_constraints(b, 1, x <= 5, name="indicator") + m.add_constraints(x.sum() <= 20, name="soft").soften(10) + m.constraints["masked"].dual = xr.DataArray(np.arange(4.0), coords=[i]) + m.constraints.set_blocks(np.arange(m._xCounter) % 2) + return m - assert m.constraints["c"].lhs.is_sparse == m.sparse - fn = tmp_path / "test_store.nc" +def test_model_to_netcdf_round_trip_keeps_stores( + mixed_store_model: Model, tmp_path: Path +) -> None: + m = mixed_store_model + fn = tmp_path / "stores.nc" m.to_netcdf(fn) p = read_netcdf(fn) - assert p.constraints["c"].lhs.is_sparse == m.sparse + assert p.constraints["flipped"].lhs.is_sparse != m.sparse + assert p.expressions["flipped"].is_sparse != m.sparse + for name, con in m.constraints.items(): + got = p.constraints[name] + assert got.lhs.is_sparse == con.lhs.is_sparse + assert_conequal(con, got, strict=True) + xr.testing.assert_equal(got.labels, con.labels) + xr.testing.assert_equal(got.data["blocks"], con.data["blocks"]) + got, want = p.constraints["masked"].data, m.constraints["masked"].data + xr.testing.assert_equal(got["dual"], want["dual"]) + pairs = [ + (p.expressions["e"], m.expressions["e"]), + (p.expressions["flipped"], m.expressions["flipped"]), + (p.objective.expression, m.objective.expression), + ] + for got_expr, want in pairs: + assert isinstance(got_expr, LinearExpression) + assert isinstance(want, LinearExpression) + assert got_expr.is_sparse == want.is_sparse + assert_equivalent(got_expr, want) assert_model_equal(m, p) +def test_read_netcdf_pre_1_0_file_raises() -> None: + with pytest.raises(ValueError, match="linopy <1.0"): + read_netcdf(Path(__file__).parent / "data" / "pre-1.0-model.nc") + + @pytest.mark.skipif(not HAS_NETCDF4, reason="netCDF4 not installed") def test_model_to_netcdf_compression(tmp_path: Path) -> None: m = Model() @@ -224,7 +272,7 @@ def test_model_to_netcdf_compression_non_numeric_coords(tmp_path: Path) -> None: m.to_netcdf(fn, engine="netcdf4", compression={"zlib": True}) with xr.open_dataset(fn, engine="netcdf4") as ds: - for k in ["variables-x-t", "variables-x-d", "constraints-c-_index0"]: + for k in ["variables-x-t", "variables-x-d", "constraints-c-t"]: assert ds[k].encoding["zlib"], k assert_model_equal(m, read_netcdf(fn)) @@ -270,37 +318,6 @@ def test_model_to_netcdf_compression_scipy( assert_model_equal(model, read_netcdf(fn)) -def test_model_from_netcdf_sparse_constraint_legacy_positions(tmp_path: Path) -> None: - """Files written before #926 stored dense positions as CSR columns.""" - m = Model(sparse=True) - i = pd.RangeIndex(3, name="i") - mask = xr.DataArray([True, False, False], dims=["i"]) - m.add_variables(coords=[i], name="a", mask=mask) - b = m.add_variables(coords=[i], name="b") - m.add_constraints(2 * b >= 1, name="c") - - fn = tmp_path / "test_sparse_legacy.nc" - m.to_netcdf(fn) - - ds = xr.load_dataset(fn) - label_to_pos = m.variables.label_index.label_to_pos - labels = ds["constraints-c-indices"].values - positions = label_to_pos[labels] - assert not np.array_equal(labels, positions) - ds["constraints-c-indices"] = xr.DataArray(positions, dims=["constraints-c-_nnz"]) - del ds.attrs["constraints-c-_csr_columns"] - ds = ds.rename({"constraints-c-labels": "constraints-c-_con_labels"}) - legacy_fn = tmp_path / "legacy.nc" - ds.to_netcdf(legacy_fn) - - p = read_netcdf(legacy_fn) - assert p.constraints["c"].lhs.is_sparse - A, A_read = m.matrices.A, p.matrices.A - assert A is not None and A_read is not None - np.testing.assert_array_equal(A.toarray(), A_read.toarray()) - assert_model_equal(m, p) - - def test_model_to_netcdf_masked_variable_nan_bounds(tmp_path: Path) -> None: m = Model() i = pd.RangeIndex(3, name="i") @@ -317,27 +334,6 @@ def test_model_to_netcdf_masked_variable_nan_bounds(tmp_path: Path) -> None: assert_model_equal(m, p.copy()) -def test_model_to_netcdf_mixed_sign_constraint(tmp_path: Path, backend: Model) -> None: - m = backend - x = m.add_variables(coords=[pd.RangeIndex(4, name="i")], name="x") - - def bound(m: Model, i: int) -> object: - if i % 2: - return x.at[i] >= i - return x.at[i] == 0.0 - - m.add_constraints(bound, coords=[pd.RangeIndex(4, name="i")], name="c") - assert m.constraints["c"].lhs.is_sparse == m.sparse - - fn = tmp_path / "test_mixed_sign.nc" - m.to_netcdf(fn) - p = read_netcdf(fn) - - assert p.constraints["c"].lhs.is_sparse == m.sparse - xr.testing.assert_equal(m.constraints["c"].sign, p.constraints["c"].sign) - assert_model_equal(m, p) - - def test_model_to_netcdf_with_sense(model: Model, tmp_path: Path) -> None: m = model m.objective.sense = "max" @@ -403,8 +399,8 @@ def test_pickle_model(model_with_dash_names: Model, tmp_path: Path) -> None: assert_model_equal(m, p) -def test_model_to_netcdf_aux_coords(tmp_path: Path) -> None: - m = Model() +def test_model_to_netcdf_aux_coords(tmp_path: Path, backend: Model) -> None: + m = backend coords = xr.Coordinates( { "snapshot": np.arange(4), @@ -425,37 +421,6 @@ def test_model_to_netcdf_aux_coords(tmp_path: Path) -> None: assert_model_equal(m, read_netcdf(fn)) -def test_read_netcdf_old_multiindex_format_loads_flat(tmp_path: Path) -> None: - m = Model() - coords = xr.Coordinates( - { - "d": pd.RangeIndex(4), - "first": ("d", [1, 1, 2, 2]), - "second": ("d", ["a", "b", "a", "b"]), - } - ) - x = m.add_variables(lower=0, coords=coords, name="x") - m.add_constraints(x >= 1, name="c") - m.add_expressions(2 * x, name="e") - m.add_objective(x.sum()) - - fn = tmp_path / "flat.nc" - m.to_netcdf(fn) - ds = xr.load_dataset(fn) - mi_dims = [d for d in ds.dims if str(d).endswith("-d")] - ds = ds.drop_vars(mi_dims) - ds.attrs.update( - {f"{d}_multiindex": json.dumps(["first", "second"]) for d in mi_dims} - ) - old = tmp_path / "old_multiindex.nc" - ds.to_netcdf(old) - - p = read_netcdf(old) - for obj in (p.variables["x"], p.constraints["c"], p.expressions["e"]): - assert isinstance(obj.indexes["d"], pd.RangeIndex) - assert_model_equal(m, p) - - def test_model_to_netcdf_with_expressions( model_with_expressions: Model, tmp_path: Path ) -> None: @@ -601,20 +566,6 @@ def test_pickle_model_with_expressions( assert p.expressions["lin"].model is p -def test_read_netcdf_without_version_stamp(model: Model, tmp_path: Path) -> None: - from linopy.io import NETCDF_VERSION_ATTR - - fn = tmp_path / "test.nc" - model.to_netcdf(fn) - - ds = xr.load_dataset(fn).load() - del ds.attrs[NETCDF_VERSION_ATTR] - fn_legacy = tmp_path / "legacy.nc" - ds.to_netcdf(fn_legacy) - - assert_model_equal(model, read_netcdf(fn_legacy)) - - @pytest.mark.skipif("gurobi" not in licensed_solvers, reason="Gurobipy not installed") def test_to_file_lp(model: Model, tmp_path: Path) -> None: import gurobipy From bd8be88a8b30acaeb404ac976fb46f26482fedac Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 23 Oct 2026 19:23:12 +0200 Subject: [PATCH 4/8] docs/io: review fixes for the 1.0 netcdf layout Docs no longer promise that pre-1.0 netcdf files load. read_netcdf tells a foreign file from a pre-1.0 one; missing store kind raises ValueError; layout attrs carry the _linopy_ prefix; dead order fallback removed. --- doc/migrating-to-v1.rst | 5 +++-- doc/release_notes.rst | 8 +++----- linopy/io.py | 13 ++++++++----- linopy/store.py | 16 +++++++++++----- test/test_io.py | 34 ++++++++++++++++++---------------- 5 files changed, 43 insertions(+), 33 deletions(-) diff --git a/doc/migrating-to-v1.rst b/doc/migrating-to-v1.rst index 2b64fb07b..ec85ce5a0 100644 --- a/doc/migrating-to-v1.rst +++ b/doc/migrating-to-v1.rst @@ -107,8 +107,9 @@ Every row is a legacy guess that is now an explicit rule. - Rejected; linopy uses a flat dimension with the levels as aux coords, and a per-level input must be mapped onto the flat dimension explicitly (no implicit projection). Affects PyPSA multi-investment models. Netcdf - files written with such a dimension load flat: a positional index with - the levels as aux coords. + files written with such a dimension by linopy < 1.0 are refused by + :func:`~linopy.io.read_netcdf`; rebuild the model and save it again + (see `Saved models (netcdf)`_). - ``.reset_index(dim)`` to flatten, then project the per-level input by its level aux coord. * - A multi-key ``groupby(...).sum()`` result diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 0f71794fb..99ad2a4e4 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -12,11 +12,10 @@ Version 1.0.0 * ``broadcast_to_coords`` no longer accepts ``warn_reorder``. * The legacy helpers ``is_v1``, ``warn_legacy``, ``legacy_warnings_muted`` and ``warn_legacy_absence`` are removed from ``linopy.semantics``, and ``check_user_nan`` no longer accepts ``op_kind=``, takes an optional ``context`` message argument and always raises. * ``LinearExpressionGroupby.map`` no longer accepts the deprecated ``shortcut`` argument. -* The deprecated sparse opt-ins are removed: ``Model(freeze_constraints=...)``, the ``Model.freeze_constraints`` property and setter, ``groupby(...).sum(sparse=...)`` and ``linopy.options["sparse_groupby"]``. Use ``Model(sparse=True)`` instead and read ``Model.sparse``. Netcdf files written with a ``freeze_constraints`` attribute still load; the attribute is ignored. -* ``Model(auto_mask=...)`` and the ``Model.auto_mask`` property are removed: user-supplied NaN no longer masks variables or constraints. NaN in coefficients or RHS values raises; pass ``mask=`` instead. Netcdf files written with an ``auto_mask`` attribute still load; the attribute is ignored. +* The deprecated sparse opt-ins are removed: ``Model(freeze_constraints=...)``, the ``Model.freeze_constraints`` property and setter, ``groupby(...).sum(sparse=...)`` and ``linopy.options["sparse_groupby"]``. Use ``Model(sparse=True)`` instead and read ``Model.sparse``. +* ``Model(auto_mask=...)`` and the ``Model.auto_mask`` property are removed: user-supplied NaN no longer masks variables or constraints. NaN in coefficients or RHS values raises; pass ``mask=`` instead. * NaN in a user-supplied variable bound (``lower``/``upper`` in ``add_variables``, ``Variable.update``, ``fix``) now raises a ``ValueError``; exclude positions explicitly with ``mask=``. * ``MultiIndex`` coordinates are unsupported everywhere: besides ``add_variables``, ``add_constraints`` and ``add_objective``, ``Model.parameters`` and constant operands reject them with the ``reset_index`` hint. ``broadcast_to_coords`` no longer projects inputs indexed by the levels of a ``MultiIndex`` dimension onto it. -* Netcdf files written with ``MultiIndex`` dimensions load in flat form: the dimension gets a positional index and the former levels become auxiliary coordinates. This includes frozen (CSR) constraints. * The ``assign_multiindex_safe`` methods of ``Variable``, ``LinearExpression``, ``QuadraticExpression`` and ``Constraint`` are removed; use ``assign``. The helpers ``linopy.common.assign_multiindex_safe`` and ``linopy.common.assign_coords_multiindex_safe`` are renamed to ``assign_fields`` and ``reassign_coords``. * Long-deprecated shims are removed: @@ -36,7 +35,6 @@ Version 1.0.0 * An operation without a sparse path on a CSR-backed ``LinearExpression``, including reading ``.data``, now works on a dense copy and no longer converts the expression itself, which stays sparse; only assigning ``vars``, ``coeffs`` or an off-grid ``const`` replaces its store. The densify notice names the operation and the case as ``op(case)``, e.g. ``Sparse (CSR) backing densified: reindex(method=...)``, so filters matching the former reason texts need updating. A named sparse expression stays sparse through ``Model.add_expressions`` and ``Model.copy`` and keeps its name through assignments. ``densify_terms``, ``reset_const`` (so also ``Constraint.update(rhs=...)``) and ``Model.assign_coords`` keep a sparse expression sparse. ``LinearExpression.where`` and ``QuadraticExpression.where`` take ``drop`` as their only keyword argument instead of forwarding ``**kwargs`` to ``xarray.Dataset.where``, which accepts no other. (`#1020 `__) * ``Model.store_cls`` names the store class a model builds expressions on, ``CSRStore`` in a sparse model and ``DenseStore`` otherwise. ``Variable.to_linexpr`` always returns an expression on it; the private ``_dense`` keyword is removed. (`#1045 `__) * A dense operand in a merge with a CSR-backed one is converted up through the new ``Store.to_sparse(reason)``; a merge that stays sparse emits one ``Dense backing sparsified: merge with a sparse operand`` notice, under the same ``warn_on_densify`` option as the densify notice. A dense operand with coordinates the sparse grid cannot hold, e.g. two-dimensional auxiliary ones, now falls the whole merge back to dense with the reason ``merge(coordinates the Grid cannot hold: )``. A merge of dense operands stays dense, also in a sparse model. (`#1045 `__) -* A named linear expression of a sparse model reads back from netcdf as a CSR-backed expression, unless it carries coordinates the sparse grid cannot hold, e.g. two-dimensional auxiliary ones; then it reads back dense with those coordinates. (`#1045 `__) * A constant operation that produces NaN in a live cell, e.g. ``0 * inf``, ``0 / 0`` or ``inf - inf`` from multiplying, dividing, adding or subtracting a constant, now raises ``ValueError`` on both backends. Before, the cell kept its terms under a NaN constant. linopy trusts NaN only as absence, so mark cells to drop with ``.where(...)`` or ``mask=`` instead. Cells that were absent before stay absent. A sparse division by zero no longer emits numpy ``RuntimeWarning``\ s. (`#1045 `__) * Cells with a NaN constant now drop their terms when a ``LinearExpression`` or ``QuadraticExpression`` is built from a dataset or its ``vars`` or ``coeffs`` are assigned, so neither a dataset with terms under a NaN constant nor a ``vars`` assignment can revive an absent cell. (`#1045 `__) @@ -55,7 +53,7 @@ Version 1.0.0 * In a sparse model, ``add_constraints`` and ``add_indicator_constraints`` hold the lhs on a ``CSRStore`` whenever the grid can hold its coordinates, also for a dense lhs, and ``Constraint.from_dataset`` reads into the model's store class. ``Constraint.to_sparse()`` and ``Constraint.to_dense()`` replace the lhs store in place. Explicit conversions emit no notice at all, so ``warn_on_densify=True`` no longer flags them; the ``explicit`` argument of the conversion notice is removed. (`#1047 `__) * A sparse constraint keeps ``rhs``, ``sign``, ``scaling`` and ``dual`` on its full grid like a dense one, instead of on its active rows; masked rows have label ``-1`` and rows without terms stay labelled but inactive. A masked row with an invalid infinite rhs now raises in a sparse model too. ``Constraint.active_labels``, ``to_matrix`` and ``to_polars`` raise ``ValueError`` on a constraint without labels. A sparse lhs reads its coefficients back as ``float64``. (`#1047 `__) * ``Constraint.sanitize_zeros`` drops near-zero terms from the lhs store (``Store.prune`` takes a tolerance), so the ``_term`` dimension of a dense constraint can shrink, instead of marking the terms ``-1``/NaN. ``Constraint.iterate_slices`` gathers row blocks with ``Store.take`` and keeps a sparse lhs sparse. ``MatrixAccessor`` and the dual write-back of ``Model.assign_result`` treat both stores alike. (`#1047 `__) -* Netcdf: one layout for every linear expression, constraint lhs and the objective, written from its store: a ``_store_kind`` attribute (``"dense"`` or ``"csr"``), ``const`` on the grid dimensions with their coordinates, and either ``coeffs``/``vars`` on the grid dimensions and ``_term`` or the CSR arrays ``data``/``indices``/``indptr``. Constraint fields (``labels``, ``sign``, ``rhs``, ``scaling``, ``dual``, ``binary_var``, ``binary_val``, ``blocks``) lie on the grid dimensions. Each lhs, expression and objective reads back on the store class it was written with, so a dense constraint in a sparse model stays dense and vice versa. The ``_linopy_format``/``_csr_columns`` CSR layout and the active-row and multiindex legacy readers are removed. (`#1048 `__) +* Netcdf: one layout for every linear expression, constraint lhs and the objective, written from its store: a ``_linopy_store_kind`` attribute (``"dense"`` or ``"csr"``), ``const`` on the grid dimensions with their coordinates, and either ``coeffs``/``vars`` on the grid dimensions and ``_term`` or the CSR arrays ``data``/``indices``/``indptr``. Constraint fields (``labels``, ``sign``, ``rhs``, ``scaling``, ``dual``, ``binary_var``, ``binary_val``, ``blocks``) lie on the grid dimensions. Each lhs, expression and objective reads back on the store class it was written with, so a dense constraint in a sparse model stays dense and vice versa. The ``_linopy_format``/``_csr_columns`` CSR layout and the active-row and multiindex legacy readers are removed. (`#1048 `__) * ``read_netcdf`` refuses files written by linopy <1.0 with a ``ValueError``; files carry the layout version in the ``_linopy_netcdf_layout`` attribute next to ``_linopy_version``. Rebuild the model and write it again with linopy >=1.0. (`#1048 `__) * ``linopy.testing.assert_conequal(strict=False)`` compares the lhs with ``assert_equivalent``, independently of its store and term layout, instead of comparing the dense datasets. (`#1047 `__) diff --git a/linopy/io.py b/linopy/io.py index f6a344368..ef95daf1d 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -1197,8 +1197,13 @@ def read_netcdf(path: Path | str, **kwargs: Any) -> Model: path = Path(path) ds = xr.load_dataset(path, **kwargs) + if NETCDF_VERSION_ATTR not in ds.attrs: + raise ValueError( + f"{path} is not a linopy netcdf: it has no {NETCDF_VERSION_ATTR!r} " + "attribute." + ) if ds.attrs.get(NETCDF_LAYOUT_ATTR) != NETCDF_LAYOUT: - written_by = ds.attrs.get(NETCDF_VERSION_ATTR, "unknown") + written_by = ds.attrs[NETCDF_VERSION_ATTR] raise ValueError( f"netcdf written by linopy <1.0 (version {written_by}) cannot be read; " "rebuild the model and write it with linopy >=1.0." @@ -1229,10 +1234,8 @@ def read_expression(ds: xr.Dataset) -> LinearExpression | QuadraticExpression: def container_names(kind: str) -> list[str]: found = {str(k).rsplit("-", 1)[0] for k in ds if str(k).startswith(kind)} - order_attr = ds.attrs.get(CONTAINER_ORDER_ATTR.format(kind)) - if order_attr is None: - return sorted(found) - ordered = [f"{kind}-{name}" for name in json.loads(order_attr)] + order = json.loads(ds.attrs[CONTAINER_ORDER_ATTR.format(kind)]) + ordered = [f"{kind}-{name}" for name in order] if len(ordered) != len(found) or set(ordered) != found: mismatch = sorted(set(ordered) ^ found) raise ValueError( diff --git a/linopy/store.py b/linopy/store.py index da9e12328..67f145906 100644 --- a/linopy/store.py +++ b/linopy/store.py @@ -52,8 +52,8 @@ FILL_VALUE = {"vars": -1, "coeffs": np.nan, "const": np.nan} """What a dense expression dataset holds at a cell or term it lacks.""" -STORE_KIND_ATTR = "_store_kind" -CSR_COLUMNS_ATTR = "_n_columns" +STORE_KIND_ATTR = "_linopy_store_kind" +CSR_COLUMNS_ATTR = "_linopy_n_columns" NETCDF_STORE_ATTRS = frozenset((STORE_KIND_ATTR, CSR_COLUMNS_ATTR)) """The netcdf attributes a store writes next to its arrays.""" @@ -137,7 +137,7 @@ def copy(self, model: Model, deep: bool = True) -> Self: def to_netcdf_vars(self) -> Dataset: """ The arrays for netcdf with ``const`` on the grid dims and the store - kind in the ``_store_kind`` attribute, read by :meth:`from_netcdf_vars`. + kind in the ``_linopy_store_kind`` attribute, read by :meth:`from_netcdf_vars`. """ @classmethod @@ -147,14 +147,20 @@ def read_netcdf_vars(cls, ds: Dataset, model: Model) -> Self: @staticmethod def from_netcdf_vars(ds: Dataset, model: Model) -> Store: - """The store written by :meth:`to_netcdf_vars`, by its ``_store_kind``.""" + """The store written by :meth:`to_netcdf_vars`, by its ``_linopy_store_kind``.""" from linopy.csr import CSRStore kinds: dict[str, type[Store]] = { "dense": DenseStore, "csr": CSRStore, } - return kinds[ds.attrs[STORE_KIND_ATTR]].read_netcdf_vars(ds, model) + kind = ds.attrs.get(STORE_KIND_ATTR) + if kind not in kinds: + raise ValueError( + f"Unknown or missing {STORE_KIND_ATTR!r} attribute {kind!r} in " + f"dataset with variables {list(ds.data_vars)}; expected one of {sorted(kinds)}." + ) + return kinds[kind].read_netcdf_vars(ds, model) @property @abstractmethod diff --git a/test/test_io.py b/test/test_io.py index f12011f1c..b646bf155 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -24,6 +24,7 @@ from linopy.constraints import Constraint from linopy.expressions import LinearExpression, QuadraticExpression from linopy.io import CONTAINER_ORDER_ATTR, signed_number +from linopy.store import STORE_KIND_ATTR from linopy.testing import ( assert_conequal, assert_equivalent, @@ -147,22 +148,6 @@ def test_model_to_netcdf_keeps_insertion_order( np.testing.assert_array_equal(A.toarray(), A_read.toarray()) -def test_read_netcdf_without_order_attrs_falls_back_to_sorted( - unsorted_model: Model, tmp_path: Path -) -> None: - fn = tmp_path / "test.nc" - unsorted_model.to_netcdf(fn) - ds = xr.load_dataset(fn) - for kind in ("variables", "expressions", "constraints"): - del ds.attrs[CONTAINER_ORDER_ATTR.format(kind)] - ds.to_netcdf(fn) - p = read_netcdf(fn) - - assert list(p.variables) == ["a", "z"] - assert list(p.expressions) == ["b", "y"] - assert list(p.constraints) == ["con1", "con2"] - - def test_read_netcdf_inconsistent_order_attr_raises( unsorted_model: Model, tmp_path: Path ) -> None: @@ -234,6 +219,23 @@ def test_read_netcdf_pre_1_0_file_raises() -> None: read_netcdf(Path(__file__).parent / "data" / "pre-1.0-model.nc") +def test_read_netcdf_foreign_file_raises(tmp_path: Path) -> None: + fn = tmp_path / "foreign.nc" + xr.Dataset({"x": ("i", [1.0, 2.0])}).to_netcdf(fn) + with pytest.raises(ValueError, match="not a linopy netcdf"): + read_netcdf(fn) + + +def test_read_netcdf_missing_store_kind_raises(model: Model, tmp_path: Path) -> None: + fn = tmp_path / "test.nc" + model.to_netcdf(fn) + ds = xr.load_dataset(fn) + ds.attrs = {k: v for k, v in ds.attrs.items() if not k.endswith(STORE_KIND_ATTR)} + ds.to_netcdf(fn) + with pytest.raises(ValueError, match=STORE_KIND_ATTR): + read_netcdf(fn) + + @pytest.mark.skipif(not HAS_NETCDF4, reason="netCDF4 not installed") def test_model_to_netcdf_compression(tmp_path: Path) -> None: m = Model() From 5037b812f4f30160b4371f9154918ee4bb53af66 Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 23 Oct 2026 19:26:54 +0200 Subject: [PATCH 5/8] test: narrow store types for mypy over the test suite --- test/test_constraint.py | 4 +++- test/test_io.py | 12 ++++++------ test/test_store_conformance.py | 11 +++++++---- 3 files changed, 16 insertions(+), 11 deletions(-) diff --git a/test/test_constraint.py b/test/test_constraint.py index eac8ec5b8..991c404cc 100644 --- a/test/test_constraint.py +++ b/test/test_constraint.py @@ -32,7 +32,7 @@ Constraint, Constraints, ) -from linopy.csr import canonical +from linopy.csr import CSRStore, canonical from linopy.testing import assert_linequal, assert_varequal @@ -1358,10 +1358,12 @@ def test_constraint_canonicalizes_sparse_lhs() -> None: assert lhs.is_sparse and lhs.nterm == 2 con = Constraint(lhs, "<=", 1, m) store = con.lhs._store + assert isinstance(store, CSRStore) assert canonical(store.csr) is store.csr assert store.nterm == 1 con.update(lhs=x - 3 * x) updated = con.lhs._store + assert isinstance(updated, CSRStore) assert updated.nterm == 1 assert (updated.csr.data == -2).all() diff --git a/test/test_io.py b/test/test_io.py index b646bf155..d6524ebcd 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -199,18 +199,18 @@ def test_model_to_netcdf_round_trip_keeps_stores( assert_conequal(con, got, strict=True) xr.testing.assert_equal(got.labels, con.labels) xr.testing.assert_equal(got.data["blocks"], con.data["blocks"]) - got, want = p.constraints["masked"].data, m.constraints["masked"].data - xr.testing.assert_equal(got["dual"], want["dual"]) + got_data, want_data = p.constraints["masked"].data, m.constraints["masked"].data + xr.testing.assert_equal(got_data["dual"], want_data["dual"]) pairs = [ (p.expressions["e"], m.expressions["e"]), (p.expressions["flipped"], m.expressions["flipped"]), (p.objective.expression, m.objective.expression), ] - for got_expr, want in pairs: + for got_expr, want_expr in pairs: assert isinstance(got_expr, LinearExpression) - assert isinstance(want, LinearExpression) - assert got_expr.is_sparse == want.is_sparse - assert_equivalent(got_expr, want) + assert isinstance(want_expr, LinearExpression) + assert got_expr.is_sparse == want_expr.is_sparse + assert_equivalent(got_expr, want_expr) assert_model_equal(m, p) diff --git a/test/test_store_conformance.py b/test/test_store_conformance.py index 893bf8fab..0210120bc 100644 --- a/test/test_store_conformance.py +++ b/test/test_store_conformance.py @@ -22,6 +22,7 @@ from linopy import LinearExpression, Model, QuadraticExpression, Variable from linopy.constants import TERM_DIM from linopy.constraints import Constraint +from linopy.csr import CSRStore from linopy.store import DenseStore from linopy.testing import assert_equivalent @@ -642,9 +643,9 @@ def test_positional_csr_agrees(rows: str) -> None: active = np.array(ROWS[rows]) label_index = m.variables.label_index dense = expr._store.positional_csr(label_index, active, "c") - sparse = ( - expr._store.to_sparse().canonical().positional_csr(label_index, active, "c") - ) + sparse_store = expr._store.to_sparse() + assert isinstance(sparse_store, CSRStore) + sparse = sparse_store.canonical().positional_csr(label_index, active, "c") assert dense.shape == sparse.shape == (active.sum(), 6) assert (dense != sparse).nnz == 0 @@ -652,7 +653,9 @@ def test_positional_csr_agrees(rows: str) -> None: @pytest.mark.parametrize("rows", list(ROWS)) def test_sparse_positional_csr_cached_while_held(rows: str) -> None: m, expr = positional_model() - store = expr._store.to_sparse().canonical() + sparse_store = expr._store.to_sparse() + assert isinstance(sparse_store, CSRStore) + store = sparse_store.canonical() active = np.array(ROWS[rows]) first = store.positional_csr(m.variables.label_index, active, "c") again = store.positional_csr(m.variables.label_index, active, "c") From d9a0351276d5a3717566c1cd5e2eaabefcbd8529 Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 23 Oct 2026 19:58:42 +0200 Subject: [PATCH 6/8] perf: size export buffers by store.nnz, avoid full-cube copies in DenseStore _row_block sized buffers by term slots instead of nonzeros; DenseStore has_terms/take/terms copied broadcast term arrays; constraints_to_file held two slices; constraint construction copied const and coords. Peak memory of the export benchmarks is back to master level. --- linopy/constraints.py | 12 +++-- linopy/csr.py | 3 ++ linopy/expressions.py | 2 +- linopy/io.py | 6 +-- linopy/matrices.py | 2 +- linopy/model.py | 4 +- linopy/store.py | 80 +++++++++++++++++++++------------- test/test_store_conformance.py | 5 ++- 8 files changed, 72 insertions(+), 42 deletions(-) diff --git a/linopy/constraints.py b/linopy/constraints.py index 8b4810257..3302ca602 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -206,7 +206,7 @@ class Slack(NamedTuple): def _on_grid(value: Any, template: DataArray, label: str) -> DataArray: """``value`` broadcast onto the dims of ``template`` and in their order.""" - da = DataArray(value) + da = value if isinstance(value, DataArray) else DataArray(value) if extra := set(da.dims) - set(template.dims): raise ValueError( f"Constraint {label} has dimensions {sorted(map(str, extra))} outside " @@ -357,7 +357,7 @@ def _bind( """ store = lhs._store.canonical() const = store.const - if np.any(const[~np.isnan(const)]): + if np.any((const != 0) & ~np.isnan(const)): raise ValueError( "The lhs of a constraint must have a zero constant; move it to the rhs." ) @@ -424,7 +424,9 @@ def _take(self, positions: DataArray) -> Constraint: rows = positions.values.ravel() grid = Grid.from_dataset(positions, map(str, positions.dims)) fields = { - k: positions.copy(data=v.values.ravel()[rows].reshape(positions.shape)) + k: positions.copy( + deep=False, data=v.values.ravel()[rows].reshape(positions.shape) + ) for k, v in self._fields().items() } return self._on_store(self._lhs._store.take(rows, grid), self._model, fields) @@ -1262,7 +1264,9 @@ def iterate_slices( :meth:`~linopy.store.Store.take` and keeps its backing. """ template = self._grid_template() - positions = template.copy(data=np.arange(template.size).reshape(template.shape)) + positions = template.copy( + deep=False, data=np.arange(template.size).reshape(template.shape) + ) slots = positions.expand_dims({TERM_DIM: self.nterm}, axis=-1) dims = list(self.coord_dims) if slice_dims is None else slice_dims for block in iterate_slices(slots, slice_size, dims): diff --git a/linopy/csr.py b/linopy/csr.py index f7b852154..f2571ece5 100644 --- a/linopy/csr.py +++ b/linopy/csr.py @@ -555,6 +555,9 @@ def solution( def canonical(self) -> CSRStore: return replace(self, csr=canonical(self.csr)) + def nnz(self, active_rows: np.ndarray) -> int: + return int(np.diff(self.csr.indptr)[active_rows].sum()) + def positional_csr( self, label_index: VariableLabelIndex, active_rows: np.ndarray, name: str ) -> scipy.sparse.csr_array: diff --git a/linopy/expressions.py b/linopy/expressions.py index 47a33927a..483c6c598 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -2438,7 +2438,7 @@ def densify_terms(self) -> Self: def reset_const(self) -> Self: store = self._store return type(self)._from_store( - store.assign_const(np.zeros(store.n_cells)), self._model + store.assign_const(np.broadcast_to(0.0, store.n_cells)), self._model ) @property diff --git a/linopy/io.py b/linopy/io.py index ef95daf1d..9ee5ecd14 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -30,6 +30,7 @@ to_polars, ) from linopy.constants import SOS_DIM_ATTR, SOS_TYPE_ATTR +from linopy.constraints import Constraint from linopy.objective import Objective, linear_part from linopy.scaling import constraint_scaling_lookup, variable_scaling_lookup from linopy.store import Store @@ -652,8 +653,7 @@ def constraints_to_file( for name in names: con = regular[name] - for con_slice in con.iterate_slices(slice_size): - df = con_slice.to_polars() + for df in map(Constraint.to_polars, con.iterate_slices(slice_size)): if df.height == 0: continue if scaled: @@ -1188,7 +1188,7 @@ def read_netcdf(path: Path | str, **kwargs: Any) -> Model: the deserialized Model is too, but :meth:`Model.undo_sos_reformulation` is a no-op on it. """ - from linopy.constraints import Constraint, Constraints + from linopy.constraints import Constraints from linopy.expressions import Expressions, LinearExpression, QuadraticExpression from linopy.model import Model from linopy.variables import Variable, Variables diff --git a/linopy/matrices.py b/linopy/matrices.py index e4d22bbfa..7ecbb1d36 100644 --- a/linopy/matrices.py +++ b/linopy/matrices.py @@ -45,7 +45,7 @@ def _row_block(con: Constraint, label_index: VariableLabelIndex) -> _RowBlock: return _RowBlock( partial(store.positional_csr, label_index, active, con.name), int(active.sum()), - store.variable_labels().size, + store.nnz(active), con.rhs.values.ravel()[active], con.sign.values.ravel()[active].astype("U1"), None if (scaling == 1).all() else scaling, diff --git a/linopy/model.py b/linopy/model.py index 730e8e994..9c1b59a6d 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -1161,7 +1161,9 @@ def _allocate_constraint_labels( """ start = self._cCounter end = start + like.size - labels = like.copy(data=self._allocate_labels(start, end).reshape(like.shape)) + labels = like.copy( + deep=False, data=self._allocate_labels(start, end).reshape(like.shape) + ) self._cCounter = end if mask is not None: labels = labels.where(mask, -1) diff --git a/linopy/store.py b/linopy/store.py index 67f145906..eb18bbeca 100644 --- a/linopy/store.py +++ b/linopy/store.py @@ -248,6 +248,13 @@ def prune(self, tol: float = 0.0) -> Self: def canonical(self) -> Self: """The canonical form a constraint is built from.""" + @abstractmethod + def nnz(self, active_rows: np.ndarray) -> int: + """ + Upper bound on the nonzeros :meth:`positional_csr` returns for the + cells selected by the boolean ``active_rows``. + """ + @abstractmethod def positional_csr( self, label_index: VariableLabelIndex, active_rows: np.ndarray, name: str @@ -465,12 +472,25 @@ def _on_cells(self, values: float | np.ndarray) -> float | DataArray: return DataArray(values.reshape(self.shape), dims=self.dims) return float(values) + def _term_views(self) -> tuple[np.ndarray, np.ndarray]: + """Coefficients and variable labels on ``(*dims, _term)``, uncopied.""" + dims = (*self.dims, TERM_DIM) + return ( + self.ds.coeffs.transpose(*dims).values, + self.ds.vars.transpose(*dims).values, + ) + + def _term_rows(self, rows: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Coefficients and variable labels of the cells ``rows``, ``(len(rows), nterm)``.""" + index = np.unravel_index(rows, (1, *self.shape)) + coeffs, vars_ = self._term_views() + return coeffs[None][index], vars_[None][index] + def _term_arrays(self) -> tuple[np.ndarray, np.ndarray]: """Coefficients and variable labels as ``(n_cells, nterm)`` arrays.""" - dims = (*self.dims, TERM_DIM) shape = (self.n_cells, self.nterm) - coeffs = self.ds.coeffs.transpose(*dims).to_numpy().reshape(shape) - return coeffs, self.ds.vars.transpose(*dims).to_numpy().reshape(shape) + coeffs, vars_ = self._term_views() + return coeffs.reshape(shape), vars_.reshape(shape) @classmethod def from_dense(cls, ds: Dataset, model: Model) -> DenseStore: @@ -541,10 +561,10 @@ def read_netcdf_vars(cls, ds: Dataset, model: Model) -> DenseStore: def dataarray(self, values: np.ndarray, name: str) -> DataArray: const = self.ds.const.transpose(*self.dims) - return const.copy(data=values.reshape(const.shape)).rename(name) + return const.copy(deep=False, data=values.reshape(const.shape)).rename(name) def has_terms(self) -> np.ndarray: - return (self._term_arrays()[1] != -1).any(axis=1) + return (self._term_views()[1] != -1).any(axis=-1).reshape(-1) def cell(self, indices: tuple[Any, ...]) -> tuple[np.ndarray, np.ndarray, float]: ds = self.ds @@ -552,11 +572,9 @@ def cell(self, indices: tuple[Any, ...]) -> tuple[np.ndarray, np.ndarray, float] return coeffs, vars_, float(ds.const.values[indices]) def terms(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - coeffs, vars_ = self._term_arrays() - present = ~np.isnan(self.const)[:, None] - keep = present & (vars_ != -1) & (coeffs != 0) - cells = np.broadcast_to(np.arange(self.n_cells)[:, None], keep.shape) - return cells[keep], vars_[keep], coeffs[keep] + coeffs, vars_, keep = self._active_mask(~np.isnan(self.const)) + counts = keep.sum(axis=-1).reshape(-1) + return np.repeat(np.arange(self.n_cells), counts), vars_[keep], coeffs[keep] def variable_labels(self) -> np.ndarray: return self.ds.vars.values.ravel() @@ -565,7 +583,9 @@ def solution( self, label_solution: Callable[[np.ndarray], np.ndarray] ) -> np.ndarray: labels = self.ds.vars - values = labels.copy(data=label_solution(labels.values)[labels.values]) + values = labels.copy( + deep=False, data=label_solution(labels.values)[labels.values] + ) sol = (self.ds.coeffs * values).sum(TERM_DIM) + self.ds.const return sol.transpose(*self.dims).to_numpy().reshape(-1) @@ -613,6 +633,18 @@ def prune(self, tol: float = 0.0) -> DenseStore: def canonical(self) -> DenseStore: return self + def _active_mask( + self, active_rows: np.ndarray + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + coeffs, vars_ = self._term_views() + valid = vars_ != -1 + valid &= coeffs != 0 + valid &= active_rows.reshape(valid.shape[:-1])[..., None] + return coeffs, vars_, valid + + def nnz(self, active_rows: np.ndarray) -> int: + return int(np.count_nonzero(self._active_mask(active_rows)[2])) + def active_terms( self, active_rows: np.ndarray ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: @@ -622,16 +654,9 @@ def active_terms( """ from linopy.csr import index_dtype - dims = (*self.dims, TERM_DIM) - vars_ = self.ds.vars.transpose(*dims).values - coeffs = self.ds.coeffs.transpose(*dims).values - valid = (vars_ != -1) & (coeffs != 0) - valid &= active_rows.reshape(valid.shape[:-1])[..., None] - flat = np.flatnonzero(valid) - labels = _take_flat(vars_, flat) - data = _take_flat(coeffs, flat) - counts = np.bincount(flat // max(self.nterm, 1), minlength=self.n_cells) - counts = counts[active_rows] + coeffs, vars_, valid = self._active_mask(active_rows) + labels, data = vars_[valid], coeffs[valid] + counts = valid.sum(axis=-1).reshape(-1)[active_rows] indptr = np.empty( len(counts) + 1, index_dtype(len(data), counts.shape, self.model) ) @@ -658,10 +683,10 @@ def with_grid(self, grid: Grid) -> DenseStore: def take(self, rows: np.ndarray, grid: Grid) -> DenseStore: present = rows >= 0 safe = np.where(present, rows, 0) - coeffs, vars_ = self._term_arrays() + coeffs, vars_ = self._term_rows(safe) const = np.where(present, self.const[safe], np.nan) - coeffs = np.where(present[:, None], coeffs[safe], np.nan) - vars_ = np.where(present[:, None], vars_[safe], -1) + coeffs[~present] = np.nan + vars_[~present] = -1 return self.from_arrays(grid, coeffs, vars_, const, self.model) def reindex(self, grid: Grid, fill: float = np.nan) -> DenseStore: @@ -685,13 +710,6 @@ def concat(self, others: Iterable[DenseStore], dim: str, grid: Grid) -> DenseSto return self.merged((self, *others), dim, "override", np.nan, {}) -def _take_flat(values: np.ndarray, flat: np.ndarray) -> np.ndarray: - """Gather C-order flat positions, copying a non-contiguous array only if dense.""" - if values.flags.c_contiguous or 4 * flat.size > values.size: - return values.ravel()[flat] - return values.flat[flat] - - def as_sparse(store: Store) -> CSRStore: """``store`` as a sparse store, raising where it cannot convert.""" res = store.to_sparse() diff --git a/test/test_store_conformance.py b/test/test_store_conformance.py index 0210120bc..306aeaeb2 100644 --- a/test/test_store_conformance.py +++ b/test/test_store_conformance.py @@ -645,9 +645,12 @@ def test_positional_csr_agrees(rows: str) -> None: dense = expr._store.positional_csr(label_index, active, "c") sparse_store = expr._store.to_sparse() assert isinstance(sparse_store, CSRStore) - sparse = sparse_store.canonical().positional_csr(label_index, active, "c") + canonical = sparse_store.canonical() + sparse = canonical.positional_csr(label_index, active, "c") assert dense.shape == sparse.shape == (active.sum(), 6) assert (dense != sparse).nnz == 0 + assert expr._store.nnz(active) == 3 * active.sum() + assert canonical.nnz(active) >= sparse.nnz @pytest.mark.parametrize("rows", list(ROWS)) From f535e29a61a4bc93918eb810c056d3a721f8ee6f Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 23 Oct 2026 20:41:16 +0200 Subject: [PATCH 7/8] perf: slice dense stores by view in iterate_slices, free cell buffer in to_polars Store.isel returns views on DenseStore; Constraint.iterate_slices uses it instead of take, which copied coeffs and vars per LP slice. --- linopy/common.py | 49 ++++++++++++++++++----------------------- linopy/constraints.py | 37 ++++++++++--------------------- linopy/store.py | 18 +++++++++++++-- test/test_constraint.py | 16 +++++++++----- 4 files changed, 60 insertions(+), 60 deletions(-) diff --git a/linopy/common.py b/linopy/common.py index 417c07f65..81d1b7a22 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -424,41 +424,34 @@ def iterate_slices( if slice_dims is None: slice_dims = list(getattr(ds, "coord_dims", ds.dims)) - if not set(slice_dims).issubset(ds.dims): + for indexers in slice_indexers(dict(ds.sizes), slice_size, slice_dims): + yield ds.isel(indexers) if indexers else ds + + +def slice_indexers( + sizes: Mapping[Hashable, int], slice_size: int | None, slice_dims: list +) -> Generator[dict[Hashable, slice], None, None]: + """ + Positional slices along the largest of ``slice_dims`` cutting an object + of ``sizes`` into parts of about ``slice_size`` elements; a single empty + indexer where no cut is needed. + """ + if not set(slice_dims).issubset(sizes): raise ValueError( "Invalid slice dimensions. Must be a subset of the dataset dimensions." ) - - # Calculate the total number of elements in the dataset - size = np.prod([ds.sizes[dim] for dim in ds.dims], dtype=int) - - if slice_size is None or size <= slice_size: - yield ds + size = np.prod(list(sizes.values()), dtype=int) + dims = {dim: sizes[dim] for dim in slice_dims} + if slice_size is None or size <= slice_size or not dims: + yield {} return - # number of slices - n_slices = max((size + slice_size - 1) // slice_size, 1) - - # leading dimension (the dimension with the largest size) - sizes = {dim: ds.sizes[dim] for dim in slice_dims} - if not sizes: - yield ds - return - - leading_dim = max(sizes, key=sizes.get) # type: ignore - size_of_leading_dim = ds.sizes[leading_dim] - - if size_of_leading_dim < n_slices: - n_slices = size_of_leading_dim - - chunk_size = (ds.sizes[leading_dim] + n_slices - 1) // n_slices - - # Iterate over the Cartesian product of slice indices + leading_dim = max(dims, key=dims.get) # type: ignore + n_slices = min(max((size + slice_size - 1) // slice_size, 1), dims[leading_dim]) + chunk_size = (dims[leading_dim] + n_slices - 1) // n_slices for i in range(n_slices): start = i * chunk_size - end = min(start + chunk_size, size_of_leading_dim) - slice_dict = {leading_dim: slice(start, end)} - yield ds.isel(slice_dict) + yield {leading_dim: slice(start, min(start + chunk_size, dims[leading_dim]))} def _remap(array: np.ndarray, mapping: np.ndarray) -> np.ndarray: diff --git a/linopy/constraints.py b/linopy/constraints.py index 3302ca602..7d28f8a4b 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -54,12 +54,12 @@ get_dims_with_index_levels, get_printout_labels, has_optimized_model, - iterate_slices, maybe_group_terms_polars, maybe_replace_signs, reassign_coords, replace_by_map, save_join, + slice_indexers, to_dataframe, ) from linopy.config import options @@ -72,7 +72,7 @@ SIGNS_pretty, ) from linopy.containers import LabelContainer -from linopy.csr import CSRStore, Grid, _densify_notice +from linopy.csr import CSRStore, _densify_notice from linopy.scaling import validate_scaling from linopy.semantics import check_user_nan from linopy.store import NETCDF_STORE_ATTRS, DenseStore, Reason, Store, as_sparse @@ -419,17 +419,10 @@ def _copy(self, model: Model, deep: bool, include_solution: bool) -> Constraint: fields = {k: v.copy(deep=deep) for k, v in self._fields().items() if k in keep} return self._on_store(self._lhs._store.copy(model, deep), model, fields) - def _take(self, positions: DataArray) -> Constraint: - """The cells at the flat ``positions``, laid out on their grid.""" - rows = positions.values.ravel() - grid = Grid.from_dataset(positions, map(str, positions.dims)) - fields = { - k: positions.copy( - deep=False, data=v.values.ravel()[rows].reshape(positions.shape) - ) - for k, v in self._fields().items() - } - return self._on_store(self._lhs._store.take(rows, grid), self._model, fields) + def _isel(self, indexers: Mapping[Hashable, slice]) -> Constraint: + """The cells at the positional slices ``indexers`` of grid dims.""" + fields = {k: v.isel(indexers) for k, v in self._fields().items()} + return self._on_store(self._lhs._store.isel(indexers), self._model, fields) def __getitem__( self, selector: str | int | slice | list | tuple | dict @@ -1225,6 +1218,7 @@ def to_polars(self) -> pl.DataFrame: long = pl.DataFrame( {"labels": labels_flat[cells], "coeffs": coeffs, "vars": vars_} ) + del cells, vars_, coeffs long = filter_nulls_polars(long) if self.nterm > 1: long = maybe_group_terms_polars(long) @@ -1260,20 +1254,13 @@ def iterate_slices( """ Yield sub-constraints along the largest of ``slice_dims``, each with at most about ``slice_size`` term slots, see - :func:`linopy.common.iterate_slices`. The lhs store is gathered by - :meth:`~linopy.store.Store.take` and keeps its backing. + :func:`linopy.common.slice_indexers`. The lhs store is sliced by + :meth:`~linopy.store.Store.isel` and keeps its backing. """ - template = self._grid_template() - positions = template.copy( - deep=False, data=np.arange(template.size).reshape(template.shape) - ) - slots = positions.expand_dims({TERM_DIM: self.nterm}, axis=-1) + sizes = {**self._grid_template().sizes, TERM_DIM: self.nterm} dims = list(self.coord_dims) if slice_dims is None else slice_dims - for block in iterate_slices(slots, slice_size, dims): - if block.sizes == slots.sizes: - yield self - else: - yield self._take(block.isel({TERM_DIM: 0})) + for indexers in slice_indexers(sizes, slice_size, dims): + yield self._isel(indexers) if indexers else self assign = conwrap(assign_fields) assign_attrs = conwrap(Dataset.assign_attrs) diff --git a/linopy/store.py b/linopy/store.py index eb18bbeca..b86f4fd8f 100644 --- a/linopy/store.py +++ b/linopy/store.py @@ -276,6 +276,16 @@ def take(self, rows: np.ndarray, grid: Grid) -> Self: absent without terms. """ + def isel(self, indexers: Mapping[Hashable, slice]) -> Self: + """The cells at the positional slices ``indexers`` of grid dims.""" + from linopy.csr import Grid + + grid = self.grid + positions = grid.dataarray(np.arange(grid.size)).isel(indexers) + return self.take( + positions.values.ravel(), Grid.from_dataset(positions, grid.dims) + ) + def broadcast(self, grid: Grid) -> Self: """Spread the cells over ``grid``, whose dims include this grid's.""" rows, valid = self.grid.indexer(grid) @@ -573,8 +583,9 @@ def cell(self, indices: tuple[Any, ...]) -> tuple[np.ndarray, np.ndarray, float] def terms(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: coeffs, vars_, keep = self._active_mask(~np.isnan(self.const)) - counts = keep.sum(axis=-1).reshape(-1) - return np.repeat(np.arange(self.n_cells), counts), vars_[keep], coeffs[keep] + cells = np.flatnonzero(keep) + cells //= self.nterm + return cells, vars_[keep], coeffs[keep] def variable_labels(self) -> np.ndarray: return self.ds.vars.values.ravel() @@ -689,6 +700,9 @@ def take(self, rows: np.ndarray, grid: Grid) -> DenseStore: vars_[~present] = -1 return self.from_arrays(grid, coeffs, vars_, const, self.model) + def isel(self, indexers: Mapping[Hashable, slice]) -> DenseStore: + return DenseStore(self.ds.isel(indexers), self.model) + def reindex(self, grid: Grid, fill: float = np.nan) -> DenseStore: fill_value = FILL_VALUE | {"const": fill} ds = self.ds.reindex(grid.indexes, fill_value=fill_value) diff --git a/test/test_constraint.py b/test/test_constraint.py index 991c404cc..5d1848b56 100644 --- a/test/test_constraint.py +++ b/test/test_constraint.py @@ -640,15 +640,21 @@ def test_constraint_flat_removed(c: linopy.constraints.Constraint) -> None: c.flat # type: ignore[attr-defined] -@pytest.mark.parametrize("sparse", [False, True], ids=["dense", "sparse"]) -def test_iterate_slices(mc: linopy.constraints.Constraint, sparse: bool) -> None: - con = mc.to_sparse() if sparse else mc - slices = list(con.iterate_slices(slice_size=2)) +@pytest.mark.parametrize("slice_size", [7, 100, 500]) +def test_iterate_slices(backend: Model, slice_size: int) -> None: + s, t = pd.RangeIndex(60, name="s"), pd.RangeIndex(40, name="t") + x = backend.add_variables(coords=[s, t], name="x") + x = x.assign_coords(lvl=("s", np.arange(60) % 3)) + y = backend.add_variables(coords=[t], name="y") + rng = np.random.default_rng(0) + mask = x.labels.copy(data=rng.random(x.labels.shape) > 0.3) + con = backend.add_constraints(2 * x + x.shift(t=1) - y <= 3, name="c", mask=mask) + slices = list(con.iterate_slices(slice_size=slice_size)) assert len(slices) > 1 for i in slices: assert isinstance(i, Constraint) assert con.coord_dims == i.coord_dims - assert i.lhs.is_sparse == sparse + assert i.lhs.is_sparse == con.lhs.is_sparse sliced = pl.concat([i.to_polars() for i in slices]) assert sliced.equals(con.to_polars()) From 455f7a1eaa770ed4f00be6a753502934e3fa67f5 Mon Sep 17 00:00:00 2001 From: Fabian Date: Sat, 24 Oct 2026 13:05:54 +0200 Subject: [PATCH 8/8] fix(constraints): keep grid dim order on update, raise when an update adds dims Reorder an incoming lhs to the constraint grid in Constraint._rebind so labels stay contiguous for range-based dual write-back; Store.transpose. Review nits: sizes reuse, dead _fill_value, Sequence[Hashable] annotation, release-note wording; tests for store swap dirtiness, masked duals on both backends, dense 1e-10 prune tolerance. --- doc/release_notes.rst | 3 +- linopy/common.py | 4 +- linopy/constraints.py | 28 +++++++++----- linopy/store.py | 21 +++++++++-- test/test_constraint.py | 26 ++++++++----- test/test_constraint_coef_dirty.py | 24 +++++++++++- test/test_constraint_label_index.py | 57 ++++++++++++++++++++++++++--- 7 files changed, 132 insertions(+), 31 deletions(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 99ad2a4e4..55f780578 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -52,10 +52,11 @@ Version 1.0.0 * ``CSRConstraint``, ``ConstraintBase`` and the freeze API are removed without aliases: ``Model.add_constraints(freeze=...)``, ``Constraints.add(freeze=...)``, ``Constraint.freeze()``, ``CSRConstraint.mutable()``, ``CSRConstraint.from_dense``/``from_csr`` and the "frozen constraint" ``AttributeError``. A constraint in a sparse model is a ``Constraint`` whose lhs holds a ``CSRStore``; replace ``isinstance(c, CSRConstraint)`` with ``c.lhs.is_sparse``. It supports the whole ``Constraint`` API, e.g. ``update``, the ``scaling`` setter, ``sanitize_*``, ``set_blocks`` and the wrapped xarray methods; ``update`` and ``soften`` keep a sparse lhs sparse (`#1033 `__). Reading ``coeffs`` or ``vars`` of a sparse constraint, ``update(coeffs=..., variables=...)`` and ``set_blocks`` work on a dense copy with a densify notice. (`#1047 `__) * In a sparse model, ``add_constraints`` and ``add_indicator_constraints`` hold the lhs on a ``CSRStore`` whenever the grid can hold its coordinates, also for a dense lhs, and ``Constraint.from_dataset`` reads into the model's store class. ``Constraint.to_sparse()`` and ``Constraint.to_dense()`` replace the lhs store in place. Explicit conversions emit no notice at all, so ``warn_on_densify=True`` no longer flags them; the ``explicit`` argument of the conversion notice is removed. (`#1047 `__) * A sparse constraint keeps ``rhs``, ``sign``, ``scaling`` and ``dual`` on its full grid like a dense one, instead of on its active rows; masked rows have label ``-1`` and rows without terms stay labelled but inactive. A masked row with an invalid infinite rhs now raises in a sparse model too. ``Constraint.active_labels``, ``to_matrix`` and ``to_polars`` raise ``ValueError`` on a constraint without labels. A sparse lhs reads its coefficients back as ``float64``. (`#1047 `__) -* ``Constraint.sanitize_zeros`` drops near-zero terms from the lhs store (``Store.prune`` takes a tolerance), so the ``_term`` dimension of a dense constraint can shrink, instead of marking the terms ``-1``/NaN. ``Constraint.iterate_slices`` gathers row blocks with ``Store.take`` and keeps a sparse lhs sparse. ``MatrixAccessor`` and the dual write-back of ``Model.assign_result`` treat both stores alike. (`#1047 `__) +* ``Constraint.sanitize_zeros`` drops near-zero terms from the lhs store (``Store.prune`` takes a tolerance), so the ``_term`` dimension of a dense constraint can shrink, instead of marking the terms ``-1``/NaN. ``Constraint.iterate_slices`` slices the lhs store with ``Store.isel`` and keeps a sparse lhs sparse. ``MatrixAccessor`` and the dual write-back of ``Model.assign_result`` treat both stores alike. (`#1047 `__) * Netcdf: one layout for every linear expression, constraint lhs and the objective, written from its store: a ``_linopy_store_kind`` attribute (``"dense"`` or ``"csr"``), ``const`` on the grid dimensions with their coordinates, and either ``coeffs``/``vars`` on the grid dimensions and ``_term`` or the CSR arrays ``data``/``indices``/``indptr``. Constraint fields (``labels``, ``sign``, ``rhs``, ``scaling``, ``dual``, ``binary_var``, ``binary_val``, ``blocks``) lie on the grid dimensions. Each lhs, expression and objective reads back on the store class it was written with, so a dense constraint in a sparse model stays dense and vice versa. The ``_linopy_format``/``_csr_columns`` CSR layout and the active-row and multiindex legacy readers are removed. (`#1048 `__) * ``read_netcdf`` refuses files written by linopy <1.0 with a ``ValueError``; files carry the layout version in the ``_linopy_netcdf_layout`` attribute next to ``_linopy_version``. Rebuild the model and write it again with linopy >=1.0. (`#1048 `__) * ``linopy.testing.assert_conequal(strict=False)`` compares the lhs with ``assert_equivalent``, independently of its store and term layout, instead of comparing the dense datasets. (`#1047 `__) +* ``Constraint.update`` raises a ``ValueError`` when the new ``lhs`` or ``rhs`` expression adds dimensions to the constraint; broadcasting a constraint to a larger grid is not supported. Before, the update was accepted and the labels stayed on the old dimensions, which left an inconsistent constraint. **Deprecations** diff --git a/linopy/common.py b/linopy/common.py index 81d1b7a22..3fbe8247f 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -429,7 +429,9 @@ def iterate_slices( def slice_indexers( - sizes: Mapping[Hashable, int], slice_size: int | None, slice_dims: list + sizes: Mapping[Hashable, int], + slice_size: int | None, + slice_dims: Sequence[Hashable], ) -> Generator[dict[Hashable, slice], None, None]: """ Positional slices along the largest of ``slice_dims`` cutting an object diff --git a/linopy/constraints.py b/linopy/constraints.py index 7d28f8a4b..9762d468f 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -254,8 +254,6 @@ class Constraint: "_coef_dirty", ) - _fill_value = FILL_VALUE - _lhs: expressions.LinearExpression _sign: DataArray _rhs: DataArray @@ -373,12 +371,24 @@ def _bind( self._binary_val = self._optional_on_grid(fields, "binary_val", template) self._blocks = self._optional_on_grid(fields, "blocks", template) + def _rebind(self, store: Store, fields: Mapping[Any, Any]) -> None: + """ + Internal: hold ``store`` as the lhs of this constraint, its grid dims + kept in the current order so the labels stay contiguous. Raises if + ``store`` adds dims. + """ + current = self._grid_template().dims + extra = [d for d in store.dims if d not in current] + if extra: + raise ValueError( + f"Constraint.update with an expression that adds dimensions {extra} is unsupported." + ) + store = store.transpose([d for d in current if d in store.dims]) + self._bind(expressions.LinearExpression._from_store(store, self._model), fields) + def _replace_store(self, store: Store) -> None: """Internal: hold ``store`` as the lhs, the fields kept. Sets ``_coef_dirty``.""" - self._bind( - expressions.LinearExpression._from_store(store, self._model), - self._fields(), - ) + self._rebind(store, self._fields()) self._coef_dirty = True @staticmethod @@ -709,7 +719,7 @@ def _assign_lhs( base_rhs = self.rhs if rhs is None else rhs fields = self._fields() | {"rhs": base_rhs - expr.const} store = _sparse_if(expr.reset_const()._store, self._lhs.is_sparse, reason) - self._bind(expressions.LinearExpression._from_store(store, self._model), fields) + self._rebind(store, fields) self._coef_dirty = True def _assign_terms( @@ -892,7 +902,7 @@ def _assign_coords(self, **coords: Any) -> Constraint: preserved. """ fields = reassign_coords(Dataset(self._fields()), **coords) - self._bind(self.lhs._assign_coords(**coords), fields) + self._rebind(self.lhs._assign_coords(**coords)._store, fields) return self @property @@ -1257,7 +1267,7 @@ def iterate_slices( :func:`linopy.common.slice_indexers`. The lhs store is sliced by :meth:`~linopy.store.Store.isel` and keeps its backing. """ - sizes = {**self._grid_template().sizes, TERM_DIM: self.nterm} + sizes = dict(self.sizes) dims = list(self.coord_dims) if slice_dims is None else slice_dims for indexers in slice_indexers(sizes, slice_size, dims): yield self._isel(indexers) if indexers else self diff --git a/linopy/store.py b/linopy/store.py index b86f4fd8f..f91d977a2 100644 --- a/linopy/store.py +++ b/linopy/store.py @@ -276,14 +276,24 @@ def take(self, rows: np.ndarray, grid: Grid) -> Self: absent without terms. """ + def _take_positions(self, positions: DataArray) -> Self: + from linopy.csr import Grid + + dims = tuple(map(str, positions.dims)) + return self.take(positions.values.ravel(), Grid.from_dataset(positions, dims)) + def isel(self, indexers: Mapping[Hashable, slice]) -> Self: """The cells at the positional slices ``indexers`` of grid dims.""" - from linopy.csr import Grid + grid = self.grid + return self._take_positions(grid.dataarray(np.arange(grid.size)).isel(indexers)) + def transpose(self, dims: Sequence[str]) -> Self: + """The cells on the grid dims reordered to ``dims``.""" + if tuple(dims) == self.dims: + return self grid = self.grid - positions = grid.dataarray(np.arange(grid.size)).isel(indexers) - return self.take( - positions.values.ravel(), Grid.from_dataset(positions, grid.dims) + return self._take_positions( + grid.dataarray(np.arange(grid.size)).transpose(*dims) ) def broadcast(self, grid: Grid) -> Self: @@ -703,6 +713,9 @@ def take(self, rows: np.ndarray, grid: Grid) -> DenseStore: def isel(self, indexers: Mapping[Hashable, slice]) -> DenseStore: return DenseStore(self.ds.isel(indexers), self.model) + def transpose(self, dims: Sequence[str]) -> DenseStore: + return DenseStore(self.ds.transpose(*dims, TERM_DIM), self.model) + def reindex(self, grid: Grid, fill: float = np.nan) -> DenseStore: fill_value = FILL_VALUE | {"const": fill} ds = self.ds.reindex(grid.indexes, fill_value=fill_value) diff --git a/test/test_constraint.py b/test/test_constraint.py index 5d1848b56..cddb64802 100644 --- a/test/test_constraint.py +++ b/test/test_constraint.py @@ -56,6 +56,11 @@ def y(m: Model) -> linopy.Variable: return m.variables["y"] +@pytest.fixture +def z(m: Model) -> linopy.Variable: + return m.variables["z"] + + @pytest.fixture def c(m: Model) -> linopy.constraints.Constraint: return m.constraints["c"] @@ -430,9 +435,9 @@ def test_constraint_update_coeffs(mc: linopy.constraints.Constraint) -> None: def test_constraint_update_lhs( - mc: linopy.constraints.Constraint, x: linopy.Variable, y: linopy.Variable + mc: linopy.constraints.Constraint, x: linopy.Variable, z: linopy.Variable ) -> None: - mc.update(lhs=x + y) + mc.update(lhs=x + z) assert mc.lhs.nterm == 2 assert mc.vars.notnull().all().item() assert mc.coeffs.notnull().all().item() @@ -511,10 +516,10 @@ def test_constraint_update_returns_self( def test_constraint_update_positional_constraint_expression( - mc: linopy.constraints.Constraint, x: linopy.Variable, y: linopy.Variable + mc: linopy.constraints.Constraint, x: linopy.Variable, z: linopy.Variable ) -> None: """``c.update(x + 5 <= 3)`` replaces lhs / sign / rhs in one call.""" - mc.update(x + y <= 7) + mc.update(x + z <= 7) assert (mc.rhs == 7).all() assert (mc.sign == LESS_EQUAL).all() assert mc.lhs.nterm == 2 @@ -537,12 +542,12 @@ def test_constraint_update_positional_rejects_non_constraint( def test_constraint_update_lhs_only( - mc: linopy.constraints.Constraint, x: linopy.Variable, y: linopy.Variable + mc: linopy.constraints.Constraint, x: linopy.Variable, z: linopy.Variable ) -> None: """lhs= alone replaces the expression; rhs and sign untouched.""" old_rhs = mc.rhs.copy() old_sign = mc.sign.copy() - mc.update(lhs=5 * x + 7 * y) + mc.update(lhs=5 * x + 7 * z) assert (mc.rhs == old_rhs).all() assert (mc.sign == old_sign).all() assert mc.lhs.nterm == 2 @@ -594,9 +599,9 @@ def test_constraint_update_rhs_with_variable( def test_constraint_update_rhs_with_expression( - mc: linopy.constraints.Constraint, x: linopy.Variable, y: linopy.Variable + mc: linopy.constraints.Constraint, x: linopy.Variable, z: linopy.Variable ) -> None: - mc.update(rhs=x + y) + mc.update(rhs=x + z) assert (mc.rhs == 0).all() assert (mc.coeffs.isel({mc.term_dim: -1}) == -1).all() assert mc.lhs.nterm == 3 @@ -684,11 +689,12 @@ def test_constraint_assignment_with_anonymous_constraints( assert m.constraints["c2"].coeffs.notnull().all() +@pytest.mark.parametrize("coeff", [0.0, 1e-12]) @pytest.mark.parametrize("sparse", [False, True], ids=["dense", "sparse"]) def test_constraint_assignment_sanitize_zeros( - m: Model, x: linopy.Variable, y: linopy.Variable, sparse: bool + m: Model, x: linopy.Variable, y: linopy.Variable, sparse: bool, coeff: float ) -> None: - con = m.add_constraints(0 * x + y == 0, name="c2") + con = m.add_constraints(coeff * x + y == 0, name="c2") if sparse: con.to_sparse() m.constraints.sanitize_zeros() diff --git a/test/test_constraint_coef_dirty.py b/test/test_constraint_coef_dirty.py index 6e32217b6..8a0fae192 100644 --- a/test/test_constraint_coef_dirty.py +++ b/test/test_constraint_coef_dirty.py @@ -1,8 +1,10 @@ from __future__ import annotations +from collections.abc import Callable + import pytest -from linopy import Model +from linopy import Constraint, Model @pytest.fixture @@ -89,3 +91,23 @@ def test_update_noop_does_not_set_dirty(m_with_c: tuple[Model, str]) -> None: c = m.constraints[name] c.update() assert c._coef_dirty is False + + +@pytest.mark.parametrize( + ("source", "target"), + [ + (Constraint.to_dense, Constraint.to_sparse), + (Constraint.to_sparse, Constraint.to_dense), + ], + ids=["dense-to-sparse", "sparse-to-dense"], +) +def test_store_conversion_sets_dirty( + m_with_c: tuple[Model, str], + source: Callable[[Constraint], Constraint], + target: Callable[[Constraint], Constraint], +) -> None: + m, name = m_with_c + c = source(m.constraints[name]) + c._coef_dirty = False + target(c) + assert c._coef_dirty is True diff --git a/test/test_constraint_label_index.py b/test/test_constraint_label_index.py index 34546a3c8..b74b4480b 100644 --- a/test/test_constraint_label_index.py +++ b/test/test_constraint_label_index.py @@ -1,15 +1,17 @@ +from collections.abc import Callable + import numpy as np import pandas as pd import pytest +import xarray as xr import linopy import linopy.constants -from linopy.store import DenseStore @pytest.fixture -def model_with_mask() -> linopy.Model: - m = linopy.Model() +def model_with_mask(backend: linopy.Model) -> linopy.Model: + m = backend coords = pd.Index(range(5), name="i") mask = np.array([True, False, True, True, False]) x = m.add_variables(lower=0, coords=[coords], name="x") @@ -30,13 +32,14 @@ def test_assign_result_does_not_build_matrix( monkeypatch: pytest.MonkeyPatch, model_with_mask: linopy.Model ) -> None: calls = {"n": 0} - original = DenseStore.positional_csr + store_cls = type(model_with_mask.constraints["c_xy"].lhs._store) + original = store_cls.positional_csr def counting(self, label_index, active_rows, name): # type: ignore[no-untyped-def] calls["n"] += 1 return original(self, label_index, active_rows, name) - monkeypatch.setattr(DenseStore, "positional_csr", counting) + monkeypatch.setattr(store_cls, "positional_csr", counting) model_with_mask.solve("highs") @@ -84,3 +87,47 @@ def test_assign_result_correctness_with_mask(model_with_mask: linopy.Model) -> N mask = np.array([True, False, True, True, False]) assert np.isfinite(dual[mask]).all() assert np.isnan(dual[~mask]).all() + + +def test_update_with_reordered_lhs_keeps_label_order(backend: linopy.Model) -> None: + m = backend + i = pd.RangeIndex(2, name="i") + j = pd.Index(list("abc"), name="j") + x = m.add_variables(0, 10, coords=[i, j], name="x") + y = m.add_variables(0, 10, coords=[j], name="y") + z = m.add_variables(0, 10, coords=[i], name="z") + rhs = xr.DataArray(np.arange(6.0).reshape(2, 3), coords=[i, j]) + c = m.add_constraints(x >= rhs, name="c") + c.update(lhs=y + 0 * z + x) + assert c.labels.dims == ("i", "j") + start, end = c.range + np.testing.assert_array_equal(c.labels.values.ravel(), np.arange(start, end)) + m.add_objective(x.sum() + y.sum() + z.sum()) + m.solve("highs") + expected = [[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]] + np.testing.assert_allclose(c.dual.transpose("i", "j").values, expected, atol=1e-9) + + +@pytest.mark.parametrize( + "update", + [ + lambda c, x: c.update(lhs=x), + lambda c, x: c.update(rhs=x), + lambda c, x: c.update(x >= 1), + ], + ids=["lhs", "rhs", "constraint"], +) +def test_update_adding_dims_raises( + backend: linopy.Model, + update: Callable[[linopy.constraints.Constraint, linopy.Variable], object], +) -> None: + m = backend + i = pd.RangeIndex(2, name="i") + j = pd.Index(list("abc"), name="j") + x = m.add_variables(0, 10, coords=[j, i], name="x") + z = m.add_variables(0, 10, coords=[i], name="z") + c = m.add_constraints(z >= 1, name="c") + with pytest.raises( + ValueError, match=r"expression that adds dimensions \['j'\] is unsupported" + ): + update(c, x)