diff --git a/doc/api.rst b/doc/api.rst index 028fdc0a..2c7a5c97 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -78,6 +78,7 @@ Modifying a model model.Model.remove_objective model.Model.remove_sos_constraints model.Model.assign_coords + model.Model.restrict model.Model.copy model.Model.apply_sos_reformulation model.Model.undo_sos_reformulation diff --git a/doc/index.rst b/doc/index.rst index 17426c59..e73a1c0d 100644 --- a/doc/index.rst +++ b/doc/index.rst @@ -138,6 +138,7 @@ This package is published under MIT license. sos-constraints piecewise-linear-constraints + restrict numerical-scaling testing-framework diff --git a/doc/release_notes.rst b/doc/release_notes.rst index d9510584..18de2209 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -38,6 +38,7 @@ Version 1.0.0 * 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 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 `__) +* ``Solver.apply_update`` takes ``(diff, n_vars, n_cons)``, the number of active variables and constraints, instead of the variable and constraint label indexes. * ``expand_dims`` on an expression now puts the new dim first in ``coord_dims`` and in the dims of its dataset, matching the layout of its arrays. Before, it came after the existing dims. Selecting by dim name is unaffected. (`#1078 `__) **Breaking Changes: constraints on stores** @@ -110,11 +111,17 @@ Version 1.0.0 * ``Model.add_constraints`` rejects dimensions named after reserved fields (``sign``, ``rhs``, ``scaling``, ...) with a ``ValueError``, consistent with ``add_variables`` and ``add_expressions``. * Multiplying a constant or one-term ``LinearExpression`` by a multi-term ``LinearExpression`` no longer raises ``TypeError`` and gives the same result as the mirrored product, on both the dense and sparse backends. (`#1055 `__) +**Features** + +* :meth:`linopy.Model.restrict` cuts a subproblem out of a model, e.g. a time slice for rolling-horizon optimization, a region for spatial decomposition or a part of the model for fix-and-optimize heuristics. ``m.restrict(snapshot=window, fix=...)`` selects labels along one or several dimensions with ``Dataset.sel`` semantics and returns a :class:`~linopy.restrict.WindowModel` built from a slice of the parent's matrices: constraint terms referencing variables outside the window are folded into the right-hand side at their actual matrix coefficient, ``rhs -= coeff * fixed_value``, and rows that lose every term are checked for feasibility and dropped. Because the window is a row and column selection of the parent's ``A``, ``b``, ``c``, ``lb`` and ``ub``, creating or moving it costs a slice rather than a rebuild of every variable and constraint. ``window.solve()`` hands those arrays to the solver and updates the existing solver model in place when consecutive windows share the sparsity pattern of ``A``, through the same :class:`~linopy.persistent.diff.ModelDiff` delta and per-backend apply hooks as the persistent interface. ``window.advance(snapshot=next, fix=...)`` moves the window to a new selection. Solutions and duals are built per container only when read. See the new :doc:`restrict` guide. (https://github.com/PyPSA/linopy/issues/768) + **Internal** * The core modules are compacted after the legacy removal: ``Variables`` and ``Constraints`` share a ``LabelContainer`` base, the piecewise formulations share one link builder, and uncalled helpers are removed. ``merge`` no longer accepts bare ``Variable`` operands (pass ``var.to_linexpr()``), ``SolverReport.barrier_iterations`` (never populated) is removed, and ``linopy.common.set_int_index``, ``infer_schema_polars`` and ``require_constant`` are removed. * The methods of ``CSRStore`` and ``Grid`` in ``linopy.csr`` use pandas/xarray-style verbs instead of past participles, e.g. ``summed`` is ``sum``, ``reindexed`` is ``reindex``, ``filled`` is ``fillna``, ``contracted`` is ``dot``, ``with_const`` is ``assign_const`` and ``Grid.reordered`` is ``Grid.select``. (`#1018 `__) * Store fallbacks convert a copy, never the receiver, and merges go through ``linopy.store.merge_stores``; ``_densify_all`` is removed. (`#1052 `__) +* ``linopy.matrices.Matrices`` is a frozen dataclass of a model's flat arrays (``A``, ``b``, ``c``, ``Q``, bounds, labels and scalings) with ``Matrices.take(rows, cols)`` to select rows and columns. ``model.matrices`` (``MatrixAccessor``) subclasses it and builds ``c`` and ``Q`` eagerly. Matrix windows reuse it: ``MatrixWindow.select`` returns the window's ``Matrices`` and the selected labels. +* The direct solver builders consume ``Matrices`` instead of a ``Model``, so windows reach every backend without a model stand-in. The Gurobi and Xpress builders no longer call ``Constraints.sanitize_missings``, so the builders do not mutate the model; empty constraint rows are excluded from the matrices instead. Version 0.10.1 -------------- diff --git a/doc/restrict.rst b/doc/restrict.rst new file mode 100644 index 00000000..8d384c73 --- /dev/null +++ b/doc/restrict.rst @@ -0,0 +1,197 @@ +.. _restrict: + +Restricting a Model +=================== + +:meth:`linopy.Model.restrict` cuts a subproblem out of a model. You select +labels along one or several dimensions, and pass fixed values for the +variables the selection leaves out. Linopy returns a window of the model: +the selected variables and constraints, with every term on a left-out +variable moved to the right-hand side. + +This one operation serves several patterns: + +* **Rolling-horizon optimization**: solve a long horizon one time slice at a + time, carrying the end state into the next slice. +* **Spatial or sectoral decomposition**: solve one region or one sector with + the exchange to its neighbours held fixed. +* **Fix-and-optimize heuristics**: reoptimize a part of the model while the + rest stays at a known solution. +* **Decomposition subproblems**: condition a subproblem on master-problem + values, as in Benders-type or ADMM-type schemes. + +Three methods cover the workflow: + +* :meth:`linopy.Model.restrict` returns a window of the model, +* :meth:`linopy.restrict.WindowModel.solve` solves it, +* :meth:`linopy.restrict.WindowModel.advance` moves it to a new selection. + +A window is a row and column slice of the model's matrices. Creating or +moving it does not rebuild variables or constraints. + +.. contents:: + :local: + :depth: 2 + +Selecting a window +------------------ + +:meth:`~linopy.Model.restrict` follows :meth:`xarray.Dataset.sel` semantics: +every container's labels are selected along the dimensions it holds, containers +without them stay whole. Labels, lists, slices and scalars work, and several +dimensions can be selected at once, e.g. +``m.restrict(snapshot=slice(t0, t1), bus=["north"])``. The parent model is +unchanged. + +Boundary values +--------------- + +A constraint in the window can reference a variable outside it. +:meth:`~linopy.Model.restrict` moves those terms to the right-hand side, +multiplied by their coefficient: + +.. math:: + + A_{WW}\, x_W \sim b_W - A_{WB}\, x_B + +where :math:`x_W` are the window's variables, :math:`x_B` the variables +outside it, and :math:`x_B` on the right-hand side the values you pass in +``fix``. In words, each outside term adds ``coefficient * fixed value`` to the +right-hand side. Rows that lose every term are checked and dropped when they +are satisfied. + +Without ``fix``, a window that cuts a coupling raises +:class:`~linopy.restrict.RestrictionError` naming the constraints. A window +with no outside terms needs no ``fix``. + +``fix`` values are keyed by variable name, aligned with the variable by +coordinate label and broadcast over the dimensions they lack, so a +DataArray over the full model, a ``pandas.Series`` or a scalar works directly. +Every coordinate of the variable must be present, otherwise ``restrict`` +raises. Values for in-window variables are ignored. A ``NaN`` at an outside +label counts as missing. + +Example: one region of a network +-------------------------------- + +Two regions balance their demand with local generation and imports. The +exchange constraint couples the regions: what one imports, the other exports. + +.. code-block:: python + + import pandas as pd + import xarray as xr + import linopy + + buses = pd.Index(["north", "south"], name="bus") + sns = pd.RangeIndex(24, name="snapshot") + demand = xr.DataArray([[30.0] * 24, [50.0] * 24], coords=[buses, sns]) + cost = xr.DataArray([1.0, 2.0], coords=[buses]) + + m = linopy.Model() + gen = m.add_variables(0, 100, coords=[buses, sns], name="gen") + imports = m.add_variables(-20, 20, coords=[buses, sns], name="imports") + m.add_constraints(gen + imports == demand, name="balance") + m.add_constraints(imports.sum("bus") == 0, name="exchange") + m.add_objective((cost * gen).sum()) + + # solve the north with the south importing 15 in every snapshot + south_imports = xr.DataArray([0.0, 15.0], coords=[buses]) + north = m.restrict(bus=["north"], fix={"imports": south_imports}) + north.solve(solver_name="highs") + +The ``exchange`` constraint has no ``bus`` dimension, so it stays whole. Its +term on the south's imports moves to the right-hand side, which forces the +north to export 15. + +Example: rolling horizon +------------------------ + +A long-horizon problem is often solved one time slice at a time: solve a slice +(say, one day), carry the end state into the next slice, then move forward. + +.. code-block:: python + + import pandas as pd + import xarray as xr + import linopy + + sns = pd.date_range("2026-01-01", periods=8760, freq="h", name="snapshot") + m = linopy.Model() + soc = m.add_variables(0, 100, coords=[sns], name="soc") + flow = m.add_variables(0, 1000, coords=[sns], name="flow") + # the balance references the previous snapshot + m.add_constraints(soc - 0.5 * soc.shift(snapshot=1) - flow <= 10, name="balance") + m.add_objective(flow.sum(), sense="min") + + # initial state of charge over the full horizon + initial = xr.DataArray([50.0] * 8760, dims=["snapshot"]) + + window = m.restrict(snapshot=sns[:24], fix={"soc": initial}) + window.solve(solver_name="highs") + + for start in range(24, 8760 - 24, 24): + window.advance(snapshot=sns[start : start + 24], fix={"soc": initial}) + window.solve(solver_name="highs") + + # containers are built when you read them + solution = window.solution + +In the first row of each window, the balance references the previous +snapshot's state of charge, which lies outside the window. That term moves to +the right-hand side with the value from ``fix``. + +.. note:: + + Folding uses the constraint's *actual* coefficient on the outside term. + Hand-written rolling-horizon code often uses a coefficient of 1 for the + initial state of charge, instead of the loss coefficient used on the other + rows. ``restrict`` reproduces a slightly different model. + +Working with a window +--------------------- + +:class:`~linopy.restrict.WindowModel` hands the solver the window's matrices +directly: + +* :attr:`~linopy.restrict.WindowModel.matrices` holds ``A``, ``b``, ``sense``, + ``c``, ``lb``, ``ub`` and the labels of its rows and columns. +* :attr:`~linopy.restrict.WindowModel.solution` and + :attr:`~linopy.restrict.WindowModel.dual` are ``Dataset`` objects with one + array per variable or constraint. They are built when you read them. +* :attr:`~linopy.restrict.WindowModel.objective` is the last solve's value. + +Terms outside the window are dropped from the objective, so the window's value +differs from the full model's by a constant. + +Moving the window +----------------- + +:meth:`~linopy.restrict.WindowModel.advance` slices the matrices for a new +selection along the same dimensions and moves the new boundary values to the +right-hand side. When the new window's matrix has the same sparsity pattern, +the solver updates its existing model instead of building a new one, so a +chain of equal-sized windows keeps that fast path. A changed pattern, e.g. a +coefficient that is zero in one window only, rebuilds the solver model, or +raises :class:`~linopy.persistent.errors.RebuildRequiredError` with +``disallow_rebuild=True``. + +Limits +------ + +Every container holding a selected dimension must contain the selected labels; +otherwise ``restrict`` raises :class:`~linopy.restrict.RestrictionError`. A +slice selects whatever labels each container holds within its bounds. + +Indicators and quadratic objectives raise ``NotImplementedError``. + +Native SOS constraints also raise: they are not part of the matrices. +Reformulate them first with :meth:`Model.apply_sos_reformulation`. A +cardinality row from the reformulation that spans the selected dimension +still cannot be windowed. + +Piecewise formulations work; they expand into ordinary variables and constraints +when added. + +A window is not a :class:`~linopy.Model`: it cannot take new variables or +constraints. diff --git a/linopy/common.py b/linopy/common.py index 3fbe8247..78330b80 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -499,6 +499,24 @@ def assigned_labels(labels: np.ndarray | DataArray) -> np.ndarray: return flat[flat != -1] +def lookup_table( + keys: np.ndarray, + size: int, + values: np.ndarray | None = None, + dtype: DTypeLike = np.intp, +) -> np.ndarray: + """ + Table of length ``size`` mapping ``keys`` to ``values``, -1 elsewhere. + + Without ``values``, each key maps to its position in ``keys``, which inverts + a label-to-position array. Index the table only with keys in ``[0, size)``: + a -1 sentinel would read the last entry. + """ + table = np.full(size, -1, dtype=dtype) + table[keys] = np.arange(keys.size) if values is None else values + return table + + def contains_labels(values: np.ndarray, labels: np.ndarray) -> bool: """ Whether any entry of ``values`` is one of ``labels``. @@ -765,9 +783,7 @@ def label_to_pos(self) -> np.ndarray: vlabels = self.vlabels model = self._variables.model dtype = model._dtypes["labels"] - label_to_pos = np.full(model._xCounter, -1, dtype=dtype) - label_to_pos[vlabels] = np.arange(len(vlabels), dtype=dtype) - return label_to_pos + return lookup_table(vlabels, model._xCounter, dtype=dtype) @property def n_active_vars(self) -> int: @@ -808,9 +824,7 @@ def label_to_pos(self) -> np.ndarray: clabels = self.clabels model = self._constraints.model dtype = model._dtypes["labels"] - label_to_pos = np.full(model._cCounter, -1, dtype=dtype) - label_to_pos[clabels] = np.arange(len(clabels), dtype=dtype) - return label_to_pos + return lookup_table(clabels, model._cCounter, dtype=dtype) @property def n_active_cons(self) -> int: diff --git a/linopy/csr.py b/linopy/csr.py index e8e5a218..3098510f 100644 --- a/linopy/csr.py +++ b/linopy/csr.py @@ -48,8 +48,10 @@ import numpy as np import pandas as pd import scipy.sparse +from numpy.typing import ArrayLike, DTypeLike from xarray import DataArray, Dataset, broadcast +from linopy.common import lookup_table from linopy.config import options from linopy.constants import HELPER_DIMS, TERM_DIM, PerformanceWarning from linopy.semantics import check_arithmetic_nan, warn_outside_linopy @@ -393,8 +395,7 @@ def from_labels(cls, ds: Dataset, model: Model) -> CSRStore: shape = (vars_.size, model._xCounter) dtype = index_dtype(vars_.size, shape, model) present = vars_ != -1 - indptr = np.zeros(vars_.size + 1, dtype=dtype) - np.cumsum(present, out=indptr[1:]) + indptr = indptr_from_counts(present, dtype) if not present.all(): vars_, coeffs = vars_[present], coeffs[present] indices = vars_.astype(dtype, copy=False) @@ -539,7 +540,7 @@ def has_terms(self) -> np.ndarray: return np.diff(self.csr.indptr) > 0 def terms(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - cells = np.repeat(np.arange(self.n_cells), np.diff(self.csr.indptr)) + cells = repeat_per_row(np.arange(self.n_cells), self.csr.indptr) keep = self.live_terms() return cells[keep], self.csr.indices[keep], self.csr.data[keep] @@ -606,7 +607,7 @@ def live_terms(self) -> np.ndarray: Mask over the stored terms: the cell is present and the coefficient is nonzero. """ - present = np.repeat(~np.isnan(self.const), np.diff(self.csr.indptr)) + present = repeat_per_row(~np.isnan(self.const), self.csr.indptr) return present & (self.csr.data != 0) def prune(self, tol: float = 0.0) -> CSRStore: @@ -640,7 +641,7 @@ def _scaled_terms( """Apply ``op`` with a scalar or per-cell ``factor`` to every coefficient.""" per_term = factor if np.ndim(factor): - per_term = np.repeat(factor, np.diff(self.csr.indptr)) + per_term = repeat_per_row(factor, self.csr.indptr) with np.errstate(divide="ignore", invalid="ignore"): data = op(self.csr.data, per_term) check_arithmetic_nan(data, op, per_term) @@ -687,7 +688,8 @@ def _without_terms(self, cells: np.ndarray) -> CSRStore: if not counts[cells].any(): return self keep = np.repeat(~cells, counts) - indptr = np.concatenate([[0], np.cumsum(np.where(cells, 0, counts))]) + kept = np.where(cells, 0, counts) + indptr = indptr_from_counts(kept, self.csr.indptr.dtype) csr = scipy.sparse.csr_array( (self.csr.data[keep], self.csr.indices[keep], indptr), shape=self.csr.shape, @@ -752,8 +754,7 @@ def reindex(self, grid: Grid, fill: float = np.nan) -> CSRStore: """ row_map, valid = grid.indexer(self.grid) if valid.all() and grid.size == self.n_cells: - source = np.full(grid.size, -1, dtype=row_map.dtype) - source[row_map] = np.arange(grid.size) + source = lookup_table(row_map, grid.size, dtype=row_map.dtype) if (source >= 0).all(): return self.take(source, self.grid.conform(grid)) @@ -942,6 +943,26 @@ def index_dtype(nnz: int, shape: tuple[int, ...], model: Model) -> np.dtype: return dtype if max(nnz, *shape) <= np.iinfo(dtype).max else np.dtype(np.int64) +def indptr_from_counts(counts: np.ndarray, dtype: DTypeLike) -> np.ndarray: + """Row pointers of rows holding ``counts`` entries, in ``dtype``.""" + indptr = np.zeros(counts.size + 1, dtype=dtype) + np.cumsum(counts, out=indptr[1:]) + return indptr + + +def repeat_per_row(values: ArrayLike, indptr: np.ndarray) -> np.ndarray: + """``values``, one per row, repeated for every entry of its row.""" + return np.repeat(values, np.diff(indptr)) + + +def keep_entries( + data: np.ndarray, indices: np.ndarray, indptr: np.ndarray, keep: np.ndarray +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """CSR arrays holding only the entries in the mask ``keep``, rows unchanged.""" + kept = indptr_from_counts(keep, indptr.dtype) + return data[keep], indices[keep], kept[indptr] + + def column_compacted_matmul( left: scipy.sparse.csr_array, right: scipy.sparse.csr_array ) -> scipy.sparse.csr_array: @@ -1055,7 +1076,7 @@ def _label_ordered(csr: scipy.sparse.csr_array) -> scipy.sparse.csr_array: """ if _is_sorted(csr, strict=False): return csr - rows = np.repeat(np.arange(csr.shape[0], dtype=np.int64), np.diff(csr.indptr)) + rows = repeat_per_row(np.arange(csr.shape[0], dtype=np.int64), csr.indptr) order = np.argsort(rows * csr.shape[1] + csr.indices, kind="stable") arrays = (csr.data[order], csr.indices[order], csr.indptr) return scipy.sparse.csr_array(arrays, shape=csr.shape) @@ -1093,14 +1114,13 @@ def csr_to_term_arrays( output row ``row_positions[i]`` of an ``n_rows``-row output, for a CSR that holds only the active rows of a larger grid. """ - counts = np.diff(csr.indptr) n_rows = csr.shape[0] if n_rows is None else n_rows vars_ = np.full((n_rows, nterm), -1, dtype=label_dtype) coeffs = np.full((n_rows, nterm), np.nan, dtype=coeff_dtype) if csr.nnz: positions = np.arange(csr.shape[0]) if row_positions is None else row_positions - rows = np.repeat(positions, counts) - cols = np.arange(csr.nnz) - np.repeat(csr.indptr[:-1], counts) + rows = repeat_per_row(positions, csr.indptr) + cols = np.arange(csr.nnz) - repeat_per_row(csr.indptr[:-1], csr.indptr) vars_[rows, cols] = csr.indices coeffs[rows, cols] = csr.data return vars_, coeffs diff --git a/linopy/dualization.py b/linopy/dualization.py index 43189c37..2f339dea 100644 --- a/linopy/dualization.py +++ b/linopy/dualization.py @@ -12,6 +12,7 @@ import numpy as np import xarray as xr +from linopy.common import lookup_table from linopy.expressions import LinearExpression if TYPE_CHECKING: @@ -265,9 +266,8 @@ def _build_label_to_flat_index_lookup(labels: np.ndarray) -> np.ndarray: if not valid.any(): return np.array([], dtype=np.int64) - lookup = np.full(int(labels[valid].max()) + 1, -1, dtype=np.int64) - lookup[labels[valid].astype(np.int64)] = np.where(valid)[0].astype(np.int64) - return lookup + keys = labels[valid].astype(np.int64) + return lookup_table(keys, int(keys.max()) + 1, np.flatnonzero(valid), np.int64) def _gather_with_default( diff --git a/linopy/examples.py b/linopy/examples.py index 549574bf..2a0d8acf 100644 --- a/linopy/examples.py +++ b/linopy/examples.py @@ -2,6 +2,8 @@ This module contains examples of linear programming models using the linopy library. """ +from typing import Any + import xarray as xr from numpy import arange @@ -68,6 +70,8 @@ def benchmark_model(n: int = 10, integerlabels: bool = False) -> Model: Model: The created linear programming model. """ m = Model() + naxis: Any + maxis: Any if integerlabels: naxis, maxis = [arange(n, dtype=int), arange(n, dtype=int)] else: diff --git a/linopy/io.py b/linopy/io.py index 6039c26a..4ba0cde5 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -821,7 +821,11 @@ def to_file( # Use very fast highspy implementation # Might be replaced by custom writer, however needs C/Rust bindings for performance h = solvers.Highs._build_solver_model( - m, explicit_coordinate_names=explicit_coordinate_names, set_names=True + m.matrices, + names=get_printers_scalar( + m, explicit_coordinate_names=explicit_coordinate_names + ), + sos=list(solvers._iter_sos_sets(m)), ) h.writeModel(str(fn)) else: diff --git a/linopy/matrices.py b/linopy/matrices.py index 7ecbb1d3..1c647494 100644 --- a/linopy/matrices.py +++ b/linopy/matrices.py @@ -8,15 +8,16 @@ from __future__ import annotations from collections.abc import Callable +from dataclasses import dataclass from functools import cached_property, partial -from typing import TYPE_CHECKING, NamedTuple +from typing import TYPE_CHECKING, NamedTuple, TypedDict import numpy as np import scipy.sparse from numpy import ndarray from linopy import expressions -from linopy.csr import index_dtype +from linopy.csr import index_dtype, keep_entries, repeat_per_row if TYPE_CHECKING: from linopy.common import VariableLabelIndex @@ -35,20 +36,29 @@ class _RowBlock(NamedTuple): nnz: int rhs: ndarray sign: ndarray - row_scaling: ndarray | None + scaling: ndarray -def _row_block(con: Constraint, label_index: VariableLabelIndex) -> _RowBlock: - active = con.active_row_mask() +def _non_unit(scaling: ndarray) -> ndarray | None: + """``scaling`` where any factor differs from 1, else None.""" + return None if (scaling == 1).all() else scaling + + +def _concat(arrays: list[ndarray]) -> ndarray: + return np.concatenate(arrays) if arrays else np.empty(0, dtype=float) + + +def _row_block( + con: Constraint, active: ndarray, label_index: VariableLabelIndex +) -> _RowBlock: store = con.lhs._store - scaling = con.scaling.values.ravel()[active] return _RowBlock( partial(store.positional_csr, label_index, active, con.name), int(active.sum()), store.nnz(active), con.rhs.values.ravel()[active], con.sign.values.ravel()[active].astype("U1"), - None if (scaling == 1).all() else scaling, + con.scaling.values.ravel()[active], ) @@ -92,20 +102,19 @@ def _stack( source.indptr, ) if count < source.nnz: - keep = block_data != 0 - block_data, block_indices = block_data[keep], block_indices[keep] - block_indptr = np.concatenate([[0], np.cumsum(keep)])[block_indptr] + block_data, block_indices, block_indptr = keep_entries( + block_data, block_indices, block_indptr, block_data != 0 + ) rows = slice(row, row + source.shape[0]) entries = slice(pos, pos + count) - if block.row_scaling is None: + row_scaling = _non_unit(block.scaling) + if row_scaling is None: data[entries] = block_data b[rows] = block.rhs else: - row_counts = np.diff(block_indptr) - np.multiply( - block_data, np.repeat(block.row_scaling, row_counts), out=data[entries] - ) - np.multiply(block.rhs, block.row_scaling, out=b[rows]) + per_row = repeat_per_row(row_scaling, block_indptr) + np.multiply(block_data, per_row, out=data[entries]) + np.multiply(block.rhs, row_scaling, out=b[rows]) if col_scaling is not None: data[entries] /= col_scaling[block_indices] indices[entries] = block_indices @@ -128,133 +137,235 @@ def _indicator_rows(con: Constraint, active: ndarray) -> tuple[ndarray, ndarray] return var_labels, np.ravel(con.binary_val)[active].astype(np.intp) -class MatrixAccessor: +@dataclass(frozen=True, eq=False) +class Matrices: """ - Helper class to quickly access model related vectors and matrices. + A model as flat arrays. All arrays are compact — only active (non-masked) entries are included. Position i in variable-side arrays corresponds to vlabels[i]. Position i in constraint-side arrays corresponds to clabels[i]. The matrix members (``A``, ``b``, ``c``, ``Q``, ``lb``, ``ub``) are in - solver-scaled units while the read-back members (``sol``, ``dual``) are in - user units, so ``A @ sol`` and ``c @ sol`` do not reconstruct ``b`` or the - objective when scaling factors differ from 1. + solver-scaled units. ``con_scaling`` is the row scaling applied to ``A``, + aligned with ``clabels``; consumers that report duals need it to invert the + scaling on the way back. """ - def __init__(self, model: Model) -> None: - self._parent = model - self._build_vars() - self._build_cons() + vlabels: ndarray + clabels: ndarray + A: scipy.sparse.csr_array | None + b: ndarray + sense: ndarray + c: ndarray + Q: scipy.sparse.csc_matrix | None + lb: ndarray + ub: ndarray + vtypes: ndarray + var_scaling: ndarray + con_scaling: ndarray + objective_scaling: float + objective_sense: str + indicator_A: scipy.sparse.csr_array | None + indicator_b: ndarray + indicator_sense: ndarray + indicator_binvar: ndarray + indicator_binval: ndarray + + @property + def type(self) -> str: + """Short string identifying the problem type, as :attr:`Model.type`.""" + discrete = self.vtypes != "C" + variable_type = ("I" if discrete.all() else "MI") if discrete.any() else "" + return f"{variable_type}{'L' if self.Q is None else 'Q'}P" + + def take(self, rows: ndarray, cols: ndarray) -> Matrices: + """ + The matrices restricted to the given row and column positions. + + Raises when indicator constraints are present, since their binary + variables would have to be remapped onto the new columns. + """ + if self.indicator_b.size: + raise NotImplementedError( + "Taking rows and columns of matrices with indicator constraints " + "is not supported." + ) + return Matrices( + vlabels=self.vlabels[cols], + clabels=self.clabels[rows], + A=None if self.A is None else self.A[rows][:, cols], + b=self.b[rows], + sense=self.sense[rows], + c=self.c[cols], + Q=None if self.Q is None else self.Q[cols][:, cols], + lb=self.lb[cols], + ub=self.ub[cols], + vtypes=self.vtypes[cols], + var_scaling=self.var_scaling[cols], + con_scaling=self.con_scaling[rows], + objective_scaling=self.objective_scaling, + objective_sense=self.objective_sense, + indicator_A=None, + indicator_b=self.indicator_b, + indicator_sense=self.indicator_sense, + indicator_binvar=self.indicator_binvar, + indicator_binval=self.indicator_binval, + ) - def _build_vars(self) -> None: - m = self._parent - self.vlabels: ndarray = m.variables.label_index.vlabels - n = len(self.vlabels) - self.var_scaling: ndarray = np.empty(n, dtype=float) - self.lb: ndarray = np.empty(n, dtype=float) - self.ub: ndarray = np.empty(n, dtype=float) - self.vtypes: ndarray = np.empty(n, dtype="U1") - - pos = 0 - for name, var in m.variables.items(): - mask = var.labels.values.ravel() != -1 - cols = slice(pos, pos + int(np.count_nonzero(mask))) - if name in m.binaries: - vtype = "B" - elif name in m.integers: - vtype = "I" - elif name in m.semi_continuous: - vtype = "S" - else: - vtype = "C" - self.vtypes[cols] = vtype - self.var_scaling[cols] = var.solver_scaling.values.ravel()[mask] - self.lb[cols] = var.lower.values.ravel()[mask] - self.ub[cols] = var.upper.values.ravel()[mask] - pos = cols.stop - - if not (self.var_scaling == 1).all(): - self.lb *= self.var_scaling - self.ub *= self.var_scaling - - def _build_cons(self) -> None: - m = self._parent - label_index = m.variables.label_index - label_to_pos = label_index.label_to_pos - col_scaling = None if (self.var_scaling == 1).all() else self.var_scaling - regular, indicator = [], [] - binvar, binval = [], [] - for c in m.constraints.data.values(): - if not c.is_indicator: - regular.append(_row_block(c, label_index)) - continue - 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 - self.A, self.b, self.sense = _stack(regular, label_index, col_scaling, m) - self.indicator_A: scipy.sparse.csr_array | None - self.indicator_A, self.indicator_b, self.indicator_sense = _stack( - indicator, label_index, col_scaling, m - ) - label_dtype = m._dtypes["labels"] - self.indicator_binvar: ndarray = np.concatenate( - [np.array([], dtype=label_dtype), *binvar] - ) - self.indicator_binval: ndarray = ( +class _VarArrays(TypedDict): + vlabels: ndarray + var_scaling: ndarray + lb: ndarray + ub: ndarray + vtypes: ndarray + + +class _ConArrays(TypedDict): + clabels: ndarray + con_scaling: ndarray + A: scipy.sparse.csr_array | None + b: ndarray + sense: ndarray + indicator_A: scipy.sparse.csr_array | None + indicator_b: ndarray + indicator_sense: ndarray + indicator_binvar: ndarray + indicator_binval: ndarray + + +def _build_vars(m: Model) -> _VarArrays: + vlabels = m.variables.label_index.vlabels + n = len(vlabels) + var_scaling = np.empty(n, dtype=float) + lb = np.empty(n, dtype=float) + ub = np.empty(n, dtype=float) + vtypes = np.empty(n, dtype="U1") + + pos = 0 + for name, var in m.variables.items(): + mask = var.labels.values.ravel() != -1 + cols = slice(pos, pos + int(np.count_nonzero(mask))) + if name in m.binaries: + vtype = "B" + elif name in m.integers: + vtype = "I" + elif name in m.semi_continuous: + vtype = "S" + else: + vtype = "C" + vtypes[cols] = vtype + var_scaling[cols] = var.solver_scaling.values.ravel()[mask] + lb[cols] = var.lower.values.ravel()[mask] + ub[cols] = var.upper.values.ravel()[mask] + pos = cols.stop + + if _non_unit(var_scaling) is not None: + lb *= var_scaling + ub *= var_scaling + return _VarArrays( + vlabels=vlabels, var_scaling=var_scaling, lb=lb, ub=ub, vtypes=vtypes + ) + + +def _build_cons(m: Model, col_scaling: ndarray | None) -> _ConArrays: + label_index = m.variables.label_index + label_to_pos = label_index.label_to_pos + + regular, indicator = [], [] + binvar, binval = [], [] + for c in m.constraints.data.values(): + active = c.active_row_mask() + block = _row_block(c, active, label_index) + if not c.is_indicator: + regular.append(block) + continue + indicator.append(block) + var_labels, val = _indicator_rows(c, active) + binvar.append(label_to_pos[var_labels]) + binval.append(val) + + A, b, sense = _stack(regular, label_index, col_scaling, m) + indicator_A, indicator_b, indicator_sense = _stack( + indicator, label_index, col_scaling, m + ) + label_dtype = m._dtypes["labels"] + return _ConArrays( + clabels=m.constraints.label_index.clabels, + con_scaling=_concat([block.scaling for block in regular]), + A=A, + b=b, + sense=sense, + indicator_A=indicator_A, + indicator_b=indicator_b, + indicator_sense=indicator_sense, + indicator_binvar=np.concatenate([np.array([], dtype=label_dtype), *binvar]), + indicator_binval=( np.concatenate(binval) if binval else np.array([], dtype=np.intp) - ) + ), + ) - @cached_property - def c(self) -> ndarray: - """Objective coefficients aligned with vlabels.""" - m = self._parent - result = np.zeros(len(self.vlabels)) - - label_index = m.variables.label_index - label_to_pos = label_index.label_to_pos - var_labels, coeffs = m.objective.linear_terms() - positions = label_to_pos[var_labels] - scaled_coeffs = coeffs / self.var_scaling[positions] - scaled_coeffs = scaled_coeffs * m.objective.scaling - np.add.at(result, positions, scaled_coeffs) - return result - @cached_property - def Q(self) -> scipy.sparse.csc_matrix | None: - """Quadratic objective matrix, shape (n_active_vars, n_active_vars).""" - m = self._parent - expr = m.objective.expression - if not isinstance(expr, expressions.QuadraticExpression): - return None - q = expr.to_matrix()[self.vlabels][:, self.vlabels] - if q.nnz: - # Quadratic coefficients get one inverse column scaling per factor. - q = q.multiply(1 / self.var_scaling[:, np.newaxis]) - q = q.multiply(1 / self.var_scaling[np.newaxis, :]) - q = q * self._parent.objective.scaling - return q.tocsc() +def _objective( + m: Model, vlabels: ndarray, var_scaling: ndarray +) -> tuple[ndarray, scipy.sparse.csc_matrix | None]: + """Linear coefficients and quadratic matrix of the objective, aligned with vlabels.""" + c = np.zeros(len(vlabels)) + var_labels, coeffs = m.objective.linear_terms() + positions = m.variables.label_index.label_to_pos[var_labels] + np.add.at(c, positions, coeffs / var_scaling[positions] * m.objective.scaling) + + expr = m.objective.expression + if not isinstance(expr, expressions.QuadraticExpression): + return c, None + q = expr.to_matrix()[vlabels][:, vlabels] + if q.nnz: + q = q.multiply(1 / var_scaling[:, np.newaxis]) + q = q.multiply(1 / var_scaling[np.newaxis, :]) + q = q * m.objective.scaling + return c, q.tocsc() + + +class MatrixAccessor(Matrices): + """ + The flat matrices of a model, with read-back of its solution and duals. + + The read-back members (``sol``, ``dual``) are in user units, so ``A @ sol`` + and ``c @ sol`` do not reconstruct ``b`` or the objective when scaling + factors differ from 1. + """ + + _parent: Model + + def __init__(self, model: Model) -> None: + object.__setattr__(self, "_parent", model) + variables = _build_vars(model) + c, Q = _objective(model, variables["vlabels"], variables["var_scaling"]) + super().__init__( + **variables, + **_build_cons(model, _non_unit(variables["var_scaling"])), + c=c, + Q=Q, + objective_scaling=model.objective.scaling, + objective_sense=model.objective.sense, + ) @cached_property def sol(self) -> ndarray: """Solution values aligned with vlabels.""" if not self._parent.status == "ok": raise ValueError("Model is not optimized.") - m = self._parent - result = np.full(len(self.vlabels), np.nan) - label_index = m.variables.label_index - label_to_pos = label_index.label_to_pos - for _, var in m.variables.items(): - labels = var.labels.values.ravel() - mask = labels != -1 - positions = label_to_pos[labels[mask]] - result[positions] = var.solution.values.ravel()[mask] + values = [ + var.solution.values.ravel()[var.labels.values.ravel() != -1] + for _, var in self._parent.variables.items() + ] + result = _concat(values) + if result.size != len(self.vlabels): + raise ValueError( + f"Solution holds {result.size} active values, expected " + f"{len(self.vlabels)} aligned with vlabels." + ) return result @cached_property diff --git a/linopy/model.py b/linopy/model.py index 2e390232..1482fc54 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -123,6 +123,7 @@ if TYPE_CHECKING: from linopy.piecewise import PiecewiseFormulation + from linopy.restrict.window import WindowModel from linopy.spec.api import SpecLike logger = logging.getLogger(__name__) @@ -2427,6 +2428,76 @@ def _coordinate_carriers(self) -> list[Any]: *self.expressions.data.values(), ] + def restrict( + self, + fix: Mapping[str, Any] | None = None, + **selection: Any, + ) -> WindowModel: + """ + Build a matrix-backed window of the model by label selection. + + A window is a slice of the model's flattened matrices: the parent is + built once, and each selection picks the rows and columns of the + selected labels. The selection follows :meth:`xarray.Dataset.sel` + semantics per container: labels, lists, slices and scalars work, along + one or several dimensions, and containers without a selected dimension + stay whole. Terms whose variables fall outside the window are folded + into the right-hand side at their actual coefficient when ``fix`` + supplies their values. Rows that lose every term are dropped when their + folded right-hand side satisfies their sign, and raise otherwise. + + The returned :class:`~linopy.restrict.WindowModel` exposes ``matrices`` + directly, so solving never rebuilds a container. Variables and + constraints are materialised only when their values are read. + + Parameters + ---------- + fix : mapping, optional + Fixed values for out-of-window variables, keyed by variable name. + Values are aligned by coordinate label and broadcast over the + variable's dimensions they lack; every coordinate of the variable + must be present. ``NaN`` counts as unfixed. Values for in-window + variables are ignored. + **selection : Any + Dimensions and the labels of the window, e.g. + ``snapshot=sns[:24]`` or ``snapshot=slice(t0, t1), bus=["a"]``. + + Returns + ------- + WindowModel + The window, with its matrices computed. + + Raises + ------ + RestrictionError + If a dimension does not exist, a carrier does not contain the + selected labels, the fixed values do not match a variable's + coordinates, a cut coupling has no fixed values, or a fully folded + row is unsatisfiable. + NotImplementedError + If the model has indicators, a quadratic objective or native SOS + constraints. + + Examples + -------- + >>> import pandas as pd + >>> import linopy + >>> + >>> sns = pd.date_range("2026-01-01", periods=4, freq="h", name="snapshot") + >>> m = linopy.Model() + >>> storage = m.add_variables(0, 10, coords=[sns], name="storage") + >>> flow = m.add_variables(0, 100, coords=[sns], name="flow") + >>> _ = m.add_constraints(storage - flow <= 0, name="cap") + >>> _ = m.add_objective(flow.sum(), sense="min") + >>> + >>> window = m.restrict(snapshot=sns[:2]) + >>> window.matrices.A.shape + (2, 4) + """ + from linopy.restrict import restrict_model + + return restrict_model(self, fix, selection) + def assign_result( self, result: Result, diff --git a/linopy/persistent/diff.py b/linopy/persistent/diff.py index 14d7c5ed..000c4dbd 100644 --- a/linopy/persistent/diff.py +++ b/linopy/persistent/diff.py @@ -9,11 +9,13 @@ import numpy as np from linopy.constants import short_GREATER_EQUAL, short_LESS_EQUAL +from linopy.csr import repeat_per_row from linopy.persistent.snapshot import ( ContainerConBuffers, ContainerVarBuffers, ModelSnapshot, StructuralKey, + VarKind, _coord_snapshot, _extract_con_buffers, _extract_var_buffers, @@ -21,14 +23,25 @@ ) if TYPE_CHECKING: + import scipy.sparse from numpy.typing import DTypeLike from linopy.common import ConstraintLabelIndex, VariableLabelIndex from linopy.constraints import Constraint + from linopy.matrices import Matrices from linopy.model import Model from linopy.variables import Variable +#: Matrix column kinds (``vtypes``) mapped onto the persistent diff's enum. +_VAR_KIND = { + "C": VarKind.CONTINUOUS, + "B": VarKind.BINARY, + "I": VarKind.INTEGER, + "S": VarKind.SEMI_CONTINUOUS, +} + + class RebuildReason(enum.Enum): STRUCTURAL_LABELS = "vlabels/clabels mismatch" STRUCTURAL_CONTAINERS = "container set changed" @@ -57,6 +70,12 @@ def _cat(parts: list[np.ndarray], dtype: DTypeLike) -> np.ndarray: return np.concatenate(parts).astype(dtype, copy=False) +def _window_A(window: Matrices) -> scipy.sparse.csr_array: + if window.A is None: + raise ValueError("Matrix windows require a constraint matrix A.") + return window.A + + def _same(a: np.ndarray, b: np.ndarray) -> bool: return a is b or np.array_equal(a, b) @@ -143,7 +162,7 @@ class ModelDiff: #: buffers the diff walk already extracted — adopting it after a #: successful apply replaces a full re-capture. Note: holding a diff #: therefore pins all container buffers for its lifetime. - snapshot: ModelSnapshot + snapshot: ModelSnapshot | None @property def is_empty(self) -> bool: @@ -347,6 +366,84 @@ def from_models( ignore_dims=ignore_dims, ) + @staticmethod + def same_window_pattern(previous: Matrices, current: Matrices) -> bool: + """Whether two matrix windows share the shape and sparsity of ``A``.""" + before, after = _window_A(previous), _window_A(current) + return ( + before.shape == after.shape + and np.array_equal(before.indptr, after.indptr) + and np.array_equal(before.indices, after.indices) + ) + + @classmethod + def from_windows(cls, previous: Matrices, current: Matrices) -> ModelDiff: + """ + Delta between two matrix windows sharing the sparsity pattern of ``A``. + + A rolling-horizon window is already flat, so the difference is a set of + changed bound, rhs, cost and coefficient entries. Building it here lets + the window path reuse :meth:`Solver.apply_update` and every backend's + in-place hooks, instead of a parallel update implementation. + """ + if not cls.same_window_pattern(previous, current): + raise ValueError("Matrix windows differ in the sparsity pattern of A.") + columns = np.flatnonzero( + (current.lb != previous.lb) | (current.ub != previous.ub) + ) + rows = np.flatnonzero(current.b != previous.b) + cost = np.flatnonzero(current.c != previous.c) + + A = _window_A(current) + changed = np.flatnonzero(A.data != _window_A(previous).data) + changed_rows = np.unique(np.searchsorted(A.indptr, changed, side="right") - 1) + n_coef_updates = int(np.diff(A.indptr)[changed_rows].sum()) + coef_deltas: list[_CoefDelta] = [] + if changed_rows.size: + buffer = ContainerConBuffers( + indptr=A.indptr, + indices=A.indices, + data=A.data, + rhs=current.b, + sign=current.sense, + active_labels=np.empty(0, dtype=np.int64), + ) + coef_deltas = [ + _CoefDelta(buffer, changed_rows, changed_rows, n_coef_updates) + ] + + sign_changed = np.flatnonzero(current.sense != previous.sense) + + types = np.flatnonzero(current.vtypes != previous.vtypes) + kinds = np.array( + [_VAR_KIND[kind] for kind in current.vtypes[types]], dtype=object + ) + + return cls( + var_bounds_indices=columns, + var_bounds_lower=current.lb[columns], + var_bounds_upper=current.ub[columns], + var_type_positions=types, + var_type_kinds=kinds, + coef_deltas=coef_deltas, + n_coef_updates=n_coef_updates, + con_rhs_indices=rows, + con_rhs_values=current.b[rows], + con_rhs_signs=current.sense[rows], + con_sign_indices=sign_changed, + con_sign_values=current.sense[sign_changed], + obj_c_indices=cost if cost.size else None, + obj_c_values=current.c[cost] if cost.size else None, + obj_sense=( + current.objective_sense + if current.objective_sense != previous.objective_sense + else None + ), + var_slices={}, + con_slices={}, + snapshot=None, + ) + class _DiffBuilder: """Accumulates per-container deltas and finalizes them into a ModelDiff.""" @@ -472,8 +569,7 @@ def diff_con( if not (skip_coef_compare or new_buf.data is base_buf.data): data_diff = new_buf.data != base_buf.data if data_diff.any(): - nnz_per_row = np.diff(new_buf.indptr) - row_of_nnz = np.repeat(np.arange(n_rows), nnz_per_row) + row_of_nnz = repeat_per_row(np.arange(n_rows), new_buf.indptr) changed_rows = np.unique(row_of_nnz[data_diff]) rhs_idx = None diff --git a/linopy/restrict/__init__.py b/linopy/restrict/__init__.py new file mode 100644 index 00000000..091616a8 --- /dev/null +++ b/linopy/restrict/__init__.py @@ -0,0 +1,30 @@ +""" +Model windows: subproblems cut from a model's flattened matrices. + +:meth:`linopy.Model.restrict` returns a :class:`~linopy.restrict.WindowModel`, +the rows and columns of the parent's flattened ``A``, ``b``, ``c``, ``lb`` and +``ub`` that belong to the selected labels. Terms whose variables fall outside +the window are folded into the right-hand side at their actual coefficient: + + ``A_WW @ x_W ~ b_W - A_WB @ x_bar_B`` + +with ``x_bar_B`` the fixed boundary values, i.e. the subproblem-conditioning +primitive behind rolling-horizon, spatial decomposition and +fix-and-optimize schemes. + +* :mod:`linopy.restrict.matrix` - :class:`MatrixWindow`, the flattened view. +* :mod:`linopy.restrict.window` - :class:`WindowModel`, the user-facing window. +""" + +from __future__ import annotations + +from linopy.restrict.matrix import MatrixWindow, RestrictionError +from linopy.restrict.window import WindowModel, WindowObjective, restrict_model + +__all__ = [ + "MatrixWindow", + "RestrictionError", + "WindowModel", + "WindowObjective", + "restrict_model", +] diff --git a/linopy/restrict/matrix.py b/linopy/restrict/matrix.py new file mode 100644 index 00000000..8978579f --- /dev/null +++ b/linopy/restrict/matrix.py @@ -0,0 +1,281 @@ +""" +Window a model's flattened matrices. + +:class:`MatrixWindow` is built once from a parent model. :meth:`MatrixWindow.select` +resolves the selection to positions once per distinct coordinate index, as +:meth:`xarray.DataArray.sel` would, selects every container's labels with +:meth:`xarray.DataArray.isel`, maps the selected labels onto rows and columns of +the parent's flattened matrices and folds the out-of-window columns into the right-hand side at their fixed values. +The window is a :class:`~linopy.matrices.Matrices` taken from the parent's with +:meth:`~linopy.matrices.Matrices.take`. No container is rebuilt. + +Indicator constraints, quadratic objectives and native SOS constraints raise. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import fields, replace +from typing import TYPE_CHECKING, Any + +import numpy as np +import pandas as pd +import scipy.sparse +from xarray import DataArray +from xarray.indexes import PandasIndex + +from linopy.alignment import as_dataarray +from linopy.common import lookup_table +from linopy.csr import indptr_from_counts, repeat_per_row +from linopy.matrices import Matrices + +if TYPE_CHECKING: + from linopy.model import Model + +__all__ = ["MatrixWindow", "RestrictionError"] + + +class RestrictionError(ValueError): + """ + Raised when a model cannot be windowed. + + A window requires dimensions that exist in the model, carriers that contain + every selected label, fixed values aligned with their variables, and folded + boundary terms whose constant rows are satisfiable. + """ + + +def _indexer(index: pd.Index, dim: str, value: Any) -> Any: + """Positional indexer of ``value`` on ``index``, a slice when contiguous.""" + indexer = PandasIndex(index, dim).sel({dim: value}).dim_indexers[dim] + if not isinstance(indexer, np.ndarray) or indexer.ndim != 1 or not indexer.size: + return indexer + if (np.diff(indexer) == 1).all(): + return slice(int(indexer[0]), int(indexer[-1]) + 1) + return indexer + + +class _Indexers: + """Positional indexers of a selection, resolved once per distinct index.""" + + def __init__(self, selection: Mapping[str, Any]) -> None: + self._selection = selection + self._cache: dict[str, list[tuple[pd.Index, Any]]] = {} + self.distinct = True + + def __call__(self, labels: DataArray) -> dict[str, Any]: + return { + d: self._resolve(labels.indexes[d], d) + for d in self._selection + if d in labels.dims + } + + def _resolve(self, index: pd.Index, dim: str) -> Any: + cached = self._cache.setdefault(dim, []) + for known, indexer in cached: + if known is index or known.equals(index): + return indexer + indexer = _indexer(index, dim, self._selection[dim]) + self.distinct &= isinstance(indexer, int | slice) + cached.append((index, indexer)) + return indexer + + +def _select(containers: Any, indexers: _Indexers) -> dict[str, DataArray]: + """Each container's labels, selected along the dimensions it carries.""" + selected = {} + for name, item in containers.items(): + labels = item.labels + try: + selected[name] = labels.isel(indexers(labels)) + except KeyError as error: + if (labels.values != -1).any(): + raise RestrictionError( + f"Cannot window {name!r}: it does not carry the selected " + f"labels ({error})." + ) from error + selected[name] = labels + return selected + + +def _positions( + selected: dict[str, DataArray], label_to_pos: np.ndarray, distinct: bool +) -> np.ndarray: + """Parent matrix positions of the selected active labels, in container order.""" + empty = np.empty(0, dtype=label_to_pos.dtype) + flat = np.concatenate([empty, *(lab.values.ravel() for lab in selected.values())]) + positions = label_to_pos[flat[flat != -1]] + positions = positions[positions >= 0] + if not distinct and np.unique(positions).size != positions.size: + raise RestrictionError("Cannot restrict model: the selection repeats labels.") + return positions + + +def _align(given: Any, labels: DataArray, name: str) -> DataArray: + """Fixed values of one variable, aligned and broadcast onto its labels.""" + try: + value = ( + given + if isinstance(given, DataArray) + else as_dataarray(given, coords=labels.coords, dims=labels.dims) + ) + shared = {d: labels.indexes[d] for d in value.indexes if d in labels.indexes} + same_shape = value.dims == labels.dims and value.shape == labels.shape + if same_shape and _equal(value, shared): + return value + return value.sel(shared).broadcast_like(labels).transpose(*labels.dims) + except (KeyError, ValueError) as error: + raise RestrictionError( + f"Fixed values for variable {name!r} do not match its coordinates " + f"({error})." + ) from error + + +def _equal(value: DataArray, indexes: dict[str, pd.Index]) -> bool: + """Whether ``value`` carries exactly ``indexes`` on every dimension.""" + return value.indexes.keys() == indexes.keys() and all( + value.indexes[d].equals(index) for d, index in indexes.items() + ) + + +def _check_dropped( + rhs: np.ndarray, sense: np.ndarray, magnitude: np.ndarray, clabels: np.ndarray +) -> None: + """Raise when a row that lost every term violates its sign.""" + tolerance = 1e-9 * np.maximum(1.0, magnitude) + violated = ( + ((sense == "=") & (np.abs(rhs) > tolerance)) + | ((sense == "<") & (rhs < -tolerance)) + | ((sense == ">") & (rhs > tolerance)) + ) + if violated.any(): + raise RestrictionError( + "Cannot window the matrices: folding the dangling terms makes " + f"constraint labels {clabels[violated].tolist()} unsatisfiable " + f"(folded rhs {rhs[violated].tolist()} against sign " + f"{sense[violated].tolist()})." + ) + + +class MatrixWindow: + """ + A matrix-level view of a model, windowed by label selection. + + Build once, then call :meth:`select` per window. The parent model is read + only; every selection allocates its own arrays. + """ + + def __init__(self, model: Model) -> None: + if model.variables.sos: + raise NotImplementedError( + "Matrix windows do not support unreformulated SOS constraints: " + "call `model.apply_sos_reformulation()` first." + ) + matrices = model.matrices + if matrices.Q is not None: + raise NotImplementedError( + "Matrix windows do not support quadratic objectives." + ) + if matrices.indicator_b.size: + raise NotImplementedError( + "Matrix windows do not support indicator constraints." + ) + if matrices.A is None: + raise NotImplementedError("Matrix windows require at least one constraint.") + self._model = model + self._matrices = matrices + self._A = matrices.A + frame = {f.name: vars(matrices)[f.name] for f in fields(Matrices)} + self._frame = Matrices(**{**frame, "A": None}) + self._dims = {d for item in model._coordinate_carriers() for d in item.sizes} + self._dims |= set(model.parameters.sizes) + + def select( + self, selection: Mapping[str, Any], fix: Mapping[str, Any] | None = None + ) -> tuple[Matrices, dict[str, DataArray], dict[str, DataArray]]: + """ + The window's matrices and each container's selected labels. + + ``selection`` maps dimensions to labels, as for :meth:`Model.restrict`. + ``fix`` maps variable names to the values of out-of-window terms, which + are folded into the right-hand side at their actual coefficient; a row + that loses every term is dropped when its folded right-hand side + satisfies its sign. + + Columns follow the parent variable labels in ``vlabels``; rows follow the + parent constraint labels in ``clabels``. The selected labels are where + flat results are scattered back. Out-of-window objective terms are + dropped, so the objective value differs from the parent's by a constant. + """ + m, mat = self._model, self._matrices + unknown = [d for d in selection if d not in self._dims] + if unknown: + raise RestrictionError( + f"Cannot restrict the model along {unknown}: the dimension does " + "not exist in the model." + ) + indexers = _Indexers(selection) + variables = _select(m.variables, indexers) + constraints = _select(m.constraints, indexers) + cols = _positions( + variables, m.variables.label_index.label_to_pos, indexers.distinct + ) + rows = _positions( + constraints, m.constraints.label_index.label_to_pos, indexers.distinct + ) + if not cols.size: + raise RestrictionError( + f"Cannot restrict model: the selection {dict(selection)} is empty." + ) + + A = self._A[rows] + position = lookup_table(cols, mat.vlabels.size, dtype=A.indices.dtype) + mapped = position[A.indices] + inside = mapped >= 0 + entry_rows = repeat_per_row(np.arange(rows.size), A.indptr) + outer_rows = entry_rows[~inside] + outer_data = A.data[~inside] + fixed = self._fixed(fix, A.indices[~inside]) + missing = np.isnan(fixed) & (outer_data != 0) + if missing.any(): + bad = mat.clabels[np.unique(rows[outer_rows[missing]])] + owners = [ + n for n, lab in constraints.items() if np.isin(lab.values, bad).any() + ] + raise RestrictionError( + f"Cannot restrict model: constraint(s) {owners} have dangling " + "terms on variables outside the selection without fixed values. " + "Pass their values via `fix`." + ) + folded = outer_data * np.nan_to_num(fixed) + rhs = mat.b[rows] - np.bincount(outer_rows, folded, minlength=rows.size) + counts = np.bincount(entry_rows[inside], minlength=rows.size) + keep = counts > 0 + magnitude = np.bincount(outer_rows, np.abs(folded), minlength=rows.size) + dropped = rows[~keep] + _check_dropped( + rhs[~keep], mat.sense[dropped], magnitude[~keep], mat.clabels[dropped] + ) + indptr = indptr_from_counts(counts[keep], A.indptr.dtype) + inner = scipy.sparse.csr_array( + (A.data[inside], mapped[inside], indptr), shape=(indptr.size - 1, cols.size) + ) + window = replace(self._frame.take(rows[keep], cols), A=inner, b=rhs[keep]) + return window, variables, constraints + + def _fixed(self, fix: Mapping[str, Any] | None, cols: np.ndarray) -> np.ndarray: + """Fixed values of the given parent columns, scaled, NaN where unfixed.""" + mat = self._matrices + labels = mat.vlabels[cols] + values = np.full(cols.size, np.nan) + for name, given in (fix or {}).items(): + if name not in self._model.variables: + raise RestrictionError( + f"Cannot fix {name!r}: it does not name a variable of the model." + ) + variable = self._model.variables[name] + start, stop = variable.range + own = (labels >= start) & (labels < stop) + aligned = _align(given, variable.labels, name).values.ravel() + values[own] = aligned[labels[own] - start] + return values * mat.var_scaling[cols] diff --git a/linopy/restrict/window.py b/linopy/restrict/window.py new file mode 100644 index 00000000..501af65a --- /dev/null +++ b/linopy/restrict/window.py @@ -0,0 +1,167 @@ +""" +A matrix-backed window of a parent model. + +:class:`WindowModel` is what :meth:`Model.restrict` returns. It exposes the +window's flattened matrices directly, so solving never rebuilds a container, and +materialises variables and constraints only when their values are read. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any + +import numpy as np +from xarray import DataArray, Dataset + +from linopy.restrict.matrix import MatrixWindow, RestrictionError + +if TYPE_CHECKING: + from linopy.constants import Solution + from linopy.matrices import Matrices + from linopy.model import Model + from linopy.solvers import Solver + +__all__ = ["WindowModel", "WindowObjective", "restrict_model"] + + +class WindowObjective: + """The window's objective value, as read off the last solve.""" + + def __init__(self, window: WindowModel) -> None: + self._window = window + + @property + def value(self) -> float | None: + result = self._window._result + if result is None: + return None + return result.objective / self._window.matrices.objective_scaling + + @property + def sense(self) -> str: + return self._window.matrices.objective_sense + + +class WindowModel: + """ + A window of a parent model, backed by its flattened matrices. + + ``matrices`` is available immediately. ``solution``, ``dual`` and access to a + named variable materialise that container from the flat result vectors, so a + caller pays only for what it reads. + """ + + def __init__( + self, + parent: Model, + selection: Mapping[str, Any], + fix: Mapping[str, Any] | None = None, + ) -> None: + self._parent = parent + self._matrix_window = MatrixWindow(parent) + self._selection = dict(selection) + self._fix = fix + self._matrices, self._variables, self._constraints = self._matrix_window.select( + self._selection, fix + ) + self.objective = WindowObjective(self) + self._result: Solution | None = None + self._solver: Solver | None = None + + @property + def matrices(self) -> Matrices: + return self._matrices + + def advance( + self, fix: Mapping[str, Any] | None = None, **selection: Any + ) -> WindowModel: + """Move the window, recomputing the matrices for the new selection.""" + if selection: + if selection.keys() != self._selection.keys(): + raise RestrictionError( + "Cannot advance model: the selection must name exactly the " + f"dimensions {list(self._selection)} of the previous selection, " + f"got {list(selection)}." + ) + self._selection = selection + self._fix = fix if fix is not None else self._fix + self._matrices, self._variables, self._constraints = self._matrix_window.select( + self._selection, self._fix + ) + self._result = None + return self + + def solve( + self, + solver_name: str | None = None, + disallow_rebuild: bool = False, + **options: Any, + ) -> WindowModel: + """Solve the window, updating an existing solver model in place when possible.""" + from linopy.solvers import Solver + + if self._solver is None: + self._solver = Solver.from_name( + solver_name or "highs", io_api="direct", options=options + ) + elif options: + self._solver.solver_options.update(options) + self._result = self._solver.solve_window( + self._matrices, disallow_rebuild=disallow_rebuild + ) + return self + + @property + def solution(self) -> Dataset: + """The window's variable values, one array per variable.""" + return Dataset(self._primal(self._variables)) + + @property + def dual(self) -> Dataset: + """The window's constraint duals, one array per constraint.""" + dual = self._solved().dual + mat = self._matrices + if dual.size != mat.clabels.size: + raise AttributeError("The last solve returned no dual values.") + values = dual * mat.con_scaling / mat.objective_scaling + size = self._parent._cCounter + return Dataset(_scatter(self._constraints, mat.clabels, values, size)) + + def __getitem__(self, name: str) -> DataArray: + """A variable of the window, materialised on demand.""" + return self._primal({name: self._variables[name]})[name] + + def _primal(self, variables: dict[str, DataArray]) -> dict[str, DataArray]: + """Primal values in model units, laid out like the given labels.""" + mat = self._matrices + values = self._solved().primal / mat.var_scaling + return _scatter(variables, mat.vlabels, values, self._parent._xCounter) + + def _solved(self) -> Solution: + if self._result is None: + raise AttributeError( + "The window has not been solved. Call `solve()` before reading values." + ) + return self._result + + +def _scatter( + containers: dict[str, DataArray], labels: np.ndarray, values: np.ndarray, size: int +) -> dict[str, DataArray]: + """Flat values by label, laid out like each container's windowed labels.""" + by_label = np.full(size, np.nan) + by_label[labels] = values + return { + name: raw.copy(data=np.where(raw.values != -1, by_label[raw.values], np.nan)) + for name, raw in containers.items() + } + + +def restrict_model( + model: Model, + fix: Mapping[str, Any] | None, + selection: Mapping[str, Any], +) -> WindowModel: + """Build a matrix-backed window of ``model``, see :meth:`linopy.Model.restrict`.""" + return WindowModel(model, selection, fix) diff --git a/linopy/solvers.py b/linopy/solvers.py index 1361b49e..2d14a4c4 100644 --- a/linopy/solvers.py +++ b/linopy/solvers.py @@ -181,9 +181,12 @@ def _installed_version_in(pkg: str, spec: str) -> bool: import highspy import mosek + from linopy.matrices import Matrices from linopy.model import Model EnvType = TypeVar("EnvType") +Printers = tuple[Callable, Callable] +SosSet = tuple[int, np.ndarray, np.ndarray] FILE_IO_APIS = ["lp", "lp-polars", "mps"] IO_APIS = FILE_IO_APIS + ["direct"] @@ -458,6 +461,7 @@ class Solver(ABC, Generic[EnvType]): _n_vars: int = field(init=False, default=0, repr=False) _n_cons: int = field(init=False, default=0, repr=False) _problem_fn: Path | None = field(init=False, default=None, repr=False) + _window: Matrices | None = field(init=False, default=None, repr=False) snapshot: ModelSnapshot | None = field(init=False, default=None, repr=False) _rebuilds: int = field(init=False, default=0, repr=False) @@ -485,12 +489,105 @@ def __post_init__(self) -> None: raise ImportError(msg) self._lock: threading.Lock = threading.Lock() - def apply_update( + def solve_window( + self, + M: Matrices, + disallow_rebuild: bool = False, + **run_kwargs: Any, + ) -> Solution: + """ + Solve one matrix window, reusing this solver's model when possible. + + When the solver's native model was built from the window solved last + and both share the sparsity pattern of ``A``, the difference is applied + in place. Otherwise, or when the solver rejects the update, the native + model is rebuilt from ``M`` unless ``disallow_rebuild`` is set, with + positional labels so results stay aligned with the window's arrays. + """ + if self.io_api != "direct": + raise NotImplementedError( + f"{self.display_name} window solving requires io_api='direct', " + f"got {self.io_api!r}." + ) + previous = self._window + self.model = None + self.snapshot = None + rebuild = functools.partial(self._build_window, M) + if previous is None: + rebuild() + elif not ModelDiff.same_window_pattern(previous, M): + self._fall_back_to_rebuild( + RebuildReason.STRUCTURAL_LABELS, + rebuild, + disallow_rebuild, + "the window changes the sparsity pattern", + ) + else: + self._apply_or_rebuild( + ModelDiff.from_windows(previous, M), + len(M.vlabels), + len(M.clabels), + rebuild, + disallow_rebuild, + ) + self._window = M + solution = self._run_direct(**run_kwargs).solution + if solution is None: + raise RuntimeError("the solver returned no solution for the window") + return solution + + def _build_window(self, M: Matrices) -> None: + self.close() + self._window = None + self._validate_features( + quadratic=M.Q is not None, + semi_continuous=bool((M.vtypes == "S").any()), + indicator=M.indicator_A is not None, + ) + self._load_matrices(M, names=None, sos=[]) + n_vars, n_cons = len(M.vlabels), len(M.clabels) + self._cache_labels(np.arange(n_vars), np.arange(n_cons), n_vars, n_cons) + + def _apply_or_rebuild( self, diff: ModelDiff, - var_label_index: Any, - con_label_index: Any, + n_vars: int, + n_cons: int, + rebuild: Callable[[], None], + disallow_rebuild: bool, + ) -> bool: + try: + self.apply_update(diff, n_vars, n_cons) + except Exception as exc: + self._fall_back_to_rebuild( + RebuildReason.BACKEND_REJECTED, + rebuild, + disallow_rebuild, + str(exc), + cause=exc, + ) + return False + self._in_place_updates += 1 + self._last_rebuild_reason = None + return True + + def _fall_back_to_rebuild( + self, + reason: RebuildReason, + rebuild: Callable[[], None], + disallow_rebuild: bool, + detail: str | None = None, + cause: Exception | None = None, ) -> None: + if disallow_rebuild: + raise RebuildRequiredError(reason, detail) from cause + if cause is not None and not isinstance(cause, UnsupportedUpdate): + logger.warning("Solver rebuild required (%s): %s", reason.name, detail) + rebuild() + self._rebuilds += 1 + self._last_rebuild_reason = reason + + def apply_update(self, diff: ModelDiff, n_vars: int, n_cons: int) -> None: """ Apply an in-place :class:`ModelDiff` to the built native model. @@ -501,7 +598,7 @@ def apply_update( if not self.supports_persistent_update: raise UnsupportedUpdate(type(self).__name__) self._validate_update(diff) - ctx = self._apply_begin(var_label_index, con_label_index) + ctx = self._apply_begin(n_vars, n_cons) if diff.var_bounds_indices.size: self._apply_var_bounds( ctx, @@ -537,7 +634,7 @@ def _validate_update(self, diff: ModelDiff) -> None: f"{self.display_name} does not support in-place constraint sign change" ) - def _apply_begin(self, var_label_index: Any, con_label_index: Any) -> Any: + def _apply_begin(self, n_vars: int, n_cons: int) -> Any: """Backend prep + validation; the return value is passed to every hook.""" return self.solver_model @@ -705,13 +802,14 @@ def _build(self, **build_kwargs: Any) -> None: The Solver never mutates ``self.model``. Constraint sanitization (``model.constraints.sanitize_zeros()`` / - ``.sanitize_infinities()``) and SOS reformulation + ``.sanitize_infinities()`` / ``.sanitize_missings()``) and SOS reformulation (``model.apply_sos_reformulation()``) are Model-level operations the caller applies first; this builder consumes whatever shape it is handed. """ if self.model is None: raise RuntimeError("Solver has no model attached; cannot build.") + self._window = None self._validate_model() if self.io_api == "direct": if build_kwargs.get("set_names") is None: @@ -727,15 +825,31 @@ def _validate_model(self) -> None: """Pre-build checks on whether this solver can handle ``self.model``.""" model = self.model assert model is not None + self._validate_features( + quadratic=model.is_quadratic, + semi_continuous=bool(model.variables.semi_continuous), + indicator=bool(model.indicator_constraints), + ) + if model.variables.sos and not self.supports(SolverFeature.SOS_CONSTRAINTS): + raise ValueError( + f"Solver {self.solver_name.value} does not support SOS constraints. " + "Reformulate first via `Model.solve(reformulate_sos=True)` or " + "`model.apply_sos_reformulation()`, or use a solver that supports SOS." + ) + + def _validate_features( + self, *, quadratic: bool, semi_continuous: bool, indicator: bool + ) -> None: + """Check the problem features this solver must support, model or window.""" solver_name = self.solver_name.value cls = type(self) - if model.is_quadratic and not cls.supports(SolverFeature.QUADRATIC_OBJECTIVE): + if quadratic and not cls.supports(SolverFeature.QUADRATIC_OBJECTIVE): raise ValueError( f"Solver {solver_name} does not support quadratic problems." ) - if model.variables.semi_continuous and not cls.supports( + if semi_continuous and not cls.supports( SolverFeature.SEMI_CONTINUOUS_VARIABLES ): raise ValueError( @@ -743,23 +857,57 @@ def _validate_model(self) -> None: "Use a solver that supports them (gurobi, cplex, highs)." ) - if model.variables.sos and not cls.supports(SolverFeature.SOS_CONSTRAINTS): - raise ValueError( - f"Solver {solver_name} does not support SOS constraints. " - "Reformulate first via `Model.solve(reformulate_sos=True)` or " - "`model.apply_sos_reformulation()`, or use a solver that supports SOS." - ) - - if model.indicator_constraints and not cls.supports( - SolverFeature.INDICATOR_CONSTRAINTS - ): + if indicator and not cls.supports(SolverFeature.INDICATOR_CONSTRAINTS): raise ValueError( f"Solver {solver_name} does not support indicator constraints. " "Use a solver that supports them." ) - def _build_direct(self, *, set_names: bool, **build_kwargs: Any) -> None: - """Build the native solver model from ``self.model``. Override per-solver.""" + def _build_direct( + self, + *, + set_names: bool = False, + explicit_coordinate_names: bool = False, + **build_kwargs: Any, + ) -> None: + """Build the native solver model from the matrices of ``self.model``.""" + model = self.model + assert model is not None + M = model.matrices + names = ( + linopy.io.get_printers_scalar( + model, explicit_coordinate_names=explicit_coordinate_names + ) + if set_names + else None + ) + sos = list(_iter_sos_sets(model)) + self._load_matrices(M, names=names, sos=sos, **build_kwargs) + self._cache_labels(M.vlabels, M.clabels, model._xCounter, model._cCounter) + + def _load_matrices( + self, + M: Matrices, + *, + names: Printers | None, + sos: list[SosSet], + **build_kwargs: Any, + ) -> None: + self.solver_model = self._new_solver_model( + M, names=names, sos=sos, **build_kwargs + ) + self.io_api = "direct" + self.sense = M.objective_sense + + def _new_solver_model( + self, + M: Matrices, + *, + names: Printers | None, + sos: list[SosSet], + **build_kwargs: Any, + ) -> Any: + """Create the native solver model from ``M``. Override per-solver.""" raise NotImplementedError( f"Solver {self.solver_name.value} does not support direct API model export." ) @@ -928,7 +1076,11 @@ def _compute_diff( snapshot, model, same_model=same_model, ignore_dims=ignore_dims ) baseline = self.model - assert baseline is not None + if baseline is None: + raise RuntimeError( + "The solver model was built from a matrix window and cannot be " + "updated from a Model. Build a new solver for the Model." + ) return ModelDiff.from_models(baseline, model, ignore_dims=ignore_dims) def _apply_locked( @@ -937,45 +1089,28 @@ def _apply_locked( ignore_dims: Iterable[str] = (), disallow_rebuild: bool = False, ) -> ModelDiff | RebuildReason: + rebuild = functools.partial(self._rebuild, model) if not self.supports_persistent_update: - if disallow_rebuild: - raise RebuildRequiredError(RebuildReason.BACKEND_REJECTED) - self._rebuild(model, RebuildReason.BACKEND_REJECTED) - return RebuildReason.BACKEND_REJECTED + reason = RebuildReason.BACKEND_REJECTED + self._fall_back_to_rebuild(reason, rebuild, disallow_rebuild) + return reason diff = self._compute_diff(model, ignore_dims, same_model=model is self.model) if isinstance(diff, RebuildReason): - if disallow_rebuild: - raise RebuildRequiredError(diff) - self._rebuild(model, diff) - return diff - try: - self.apply_update( - diff, - model.variables.label_index, - model.constraints.label_index, - ) - except Exception as exc: - if disallow_rebuild: - raise RebuildRequiredError( - RebuildReason.BACKEND_REJECTED, str(exc) - ) from exc - self._last_rebuild_reason = RebuildReason.BACKEND_REJECTED - self._rebuild(model, RebuildReason.BACKEND_REJECTED) + self._fall_back_to_rebuild(diff, rebuild, disallow_rebuild) return diff - self.model = model - if self.track_updates: - self.snapshot = diff.snapshot - clear_coef_dirty(model) - self._in_place_updates += 1 - self._last_rebuild_reason = None + n_vars = model.variables.label_index.n_active_vars + n_cons = model.constraints.label_index.n_active_cons + if self._apply_or_rebuild(diff, n_vars, n_cons, rebuild, disallow_rebuild): + self.model = model + if self.track_updates: + self.snapshot = diff.snapshot + clear_coef_dirty(model) return diff - def _rebuild(self, model: Model, reason: RebuildReason) -> None: + def _rebuild(self, model: Model) -> None: self.close() self.model = model self._build() - self._rebuilds += 1 - self._last_rebuild_reason = reason def _run_direct(self, **run_kwargs: Any) -> Result: """Run the pre-built native solver model. Override per-solver.""" @@ -1103,12 +1238,14 @@ def solve_problem_from_file( env=env, ) - def _cache_model_labels(self, model: Model) -> None: + def _cache_labels( + self, vlabels: np.ndarray, clabels: np.ndarray, n_vars: int, n_cons: int + ) -> None: """Cache vlabels/clabels and total label counts for label-indexed solutions.""" - self._vlabels = model.variables.label_index.vlabels - self._clabels = model.constraints.label_index.clabels - self._n_vars = model._xCounter - self._n_cons = model._cCounter + self._vlabels = vlabels + self._clabels = clabels + self._n_vars = n_vars + self._n_cons = n_cons def _cache_model_sizes(self, model: Model) -> None: """Cache total label counts only (file-based solvers parse names).""" @@ -1151,7 +1288,7 @@ def close(self) -> None: self._env_stack = None def __getstate__(self) -> dict[str, Any]: - drop = {"solver_model", "env", "_env_stack", "snapshot", "_lock"} + drop = {"solver_model", "env", "_env_stack", "snapshot", "_lock", "_window"} return {k: v for k, v in self.__dict__.items() if k not in drop} def __setstate__(self, state: dict[str, Any]) -> None: @@ -1160,6 +1297,7 @@ def __setstate__(self, state: dict[str, Any]) -> None: self.env = None self._env_stack = None self.snapshot = None + self._window = None self._lock = threading.Lock() def __repr__(self) -> str: @@ -1613,6 +1751,11 @@ def _apply_var_bounds( ) -> None: ctx.changeColsBounds(indices.size, indices, lower, upper) + def _apply_con_rhs(self, ctx: Any, diff: ModelDiff) -> None: + lower, upper = diff.con_rhs_as_bounds() + for pos, lo, up in zip(diff.con_rhs_indices, lower, upper): + ctx.changeRowBounds(int(pos), float(lo), float(up)) + def _apply_var_types( self, ctx: Any, positions: np.ndarray, kinds: np.ndarray ) -> None: @@ -1624,11 +1767,6 @@ def _apply_var_types( ) ctx.changeColsIntegrality(positions.size, positions, integrality) - def _apply_con_rhs(self, ctx: Any, diff: ModelDiff) -> None: - lower, upper = diff.con_rhs_as_bounds() - for pos, lo, up in zip(diff.con_rhs_indices, lower, upper): - ctx.changeRowBounds(int(pos), float(lo), float(up)) - def _apply_con_coefs( self, ctx: Any, rows: np.ndarray, cols: np.ndarray, vals: np.ndarray ) -> None: @@ -1646,21 +1784,20 @@ def _apply_obj_sense(self, ctx: Any, sense: str) -> None: ) ctx.changeObjectiveSense(native) - def _build_direct( + def _new_solver_model( self, - explicit_coordinate_names: bool = False, + M: Matrices, *, - set_names: bool, + names: Printers | None, + sos: list[SosSet], log_fn: Path | None = None, **kwargs: Any, - ) -> None: - model = self.model - assert model is not None + ) -> highspy.Highs: if self.solver_options.get("solver") in [ "simplex", "ipm", "pdlp", - ] and model.type in [ + ] and M.type in [ "QP", "MILP", ]: @@ -1668,44 +1805,31 @@ def _build_direct( "The HiGHS solver ignores quadratic terms / integrality if the solver is set to 'simplex', 'ipm' or 'pdlp'. " "Drop the solver option or use 'choose' to enable quadratic terms / integrality." ) - - h = self._build_solver_model( - model, - explicit_coordinate_names=explicit_coordinate_names, - set_names=set_names, - ) + h = self._build_solver_model(M, names=names, sos=sos) self._set_solver_params(h, log_fn) - self.solver_model = h - self.io_api = "direct" - self.sense = model.sense - self._cache_model_labels(model) + return h @staticmethod def _build_solver_model( - model: Model, - explicit_coordinate_names: bool = False, + M: Matrices, *, - set_names: bool, + names: Printers | None, + sos: list[SosSet], ) -> highspy.Highs: - """Build a highspy.Highs instance that mirrors the linopy `model`.""" - if model.variables.sos: + """Build a highspy.Highs instance from the matrices ``M``.""" + if sos: raise NotImplementedError( "SOS constraints are not supported by the HiGHS direct API. " "Use io_api='lp' instead." ) - M = model.matrices h = highspy.Highs() h.addVars(len(M.vlabels), M.lb, M.ub) - if ( - len(model.binaries) - + len(model.integers) - + len(list(model.variables.semi_continuous)) - ): - vtypes = M.vtypes - integrality_map = {"C": 0, "B": 1, "I": 1, "S": 2} - int_mask = (vtypes == "B") | (vtypes == "I") | (vtypes == "S") - labels = np.arange(len(vtypes))[int_mask] + vtypes = M.vtypes + int_mask = vtypes != "C" + if int_mask.any(): + integrality_map = {"B": 1, "I": 1, "S": 2} + labels = np.flatnonzero(int_mask) integrality = np.array( [integrality_map[v] for v in vtypes[int_mask]], dtype=np.uint8 ) @@ -1722,10 +1846,8 @@ def _build_solver_model( upper = np.where(M.sense != ">", M.b, np.inf) h.addRows(num_cons, lower, upper, A.nnz, A.indptr, A.indices, A.data) - if set_names: - print_variables, print_constraints = linopy.io.get_printers_scalar( - model, explicit_coordinate_names=explicit_coordinate_names - ) + if names is not None: + print_variables, print_constraints = names mdl = h.getModel() mdl.lp_.col_names_ = print_variables(M.vlabels) if len(M.clabels): @@ -1738,7 +1860,7 @@ def _build_solver_model( num_vars = Q.shape[0] h.passHessian(num_vars, Q.nnz, 1, Q.indptr, Q.indices, Q.data) - if model.objective.sense == "max": + if M.objective_sense == "max": h.changeObjectiveSense(highspy.ObjSense.kMaximize) return h @@ -1988,72 +2110,53 @@ def _register_solver_model(self, m: gurobipy.Model) -> gurobipy.Model: assert self._env_stack is not None return self._env_stack.enter_context(m) - def _build_direct( + def _new_solver_model( self, - explicit_coordinate_names: bool = False, - env: gurobipy.Env | dict[str, Any] | None = None, + M: Matrices, *, - set_names: bool, + names: Printers | None, + sos: list[SosSet], + env: gurobipy.Env | dict[str, Any] | None = None, **kwargs: Any, - ) -> None: - model = self.model - assert model is not None - env_ = self._resolve_env(env) + ) -> gurobipy.Model: m = self._build_solver_model( - model, - env=env_, - explicit_coordinate_names=explicit_coordinate_names, - set_names=set_names, + M, env=self._resolve_env(env), names=names, sos=sos ) - self.solver_model = self._register_solver_model(m) - self.io_api = "direct" - self.sense = model.sense - self._cache_model_labels(model) + return self._register_solver_model(m) @staticmethod def _build_solver_model( - model: Model, + M: Matrices, env: gurobipy.Env | None = None, - explicit_coordinate_names: bool = False, *, - set_names: bool, + names: Printers | None, + sos: list[SosSet], ) -> gurobipy.Model: - """Build a gurobipy.Model that mirrors the linopy `model`.""" - model.constraints.sanitize_missings() + """Build a gurobipy.Model from the matrices ``M``.""" gm = gurobipy.Model(env=env) - M = model.matrices - kwargs: dict[str, Any] = {} - if set_names: - print_variables, print_constraints = linopy.io.get_printers_scalar( - model, explicit_coordinate_names=explicit_coordinate_names - ) + if names is not None: + print_variables, print_constraints = names kwargs["name"] = print_variables(M.vlabels) - if ( - len(model.binaries.labels) - + len(model.integers.labels) - + len(list(model.variables.semi_continuous)) - ): + if (M.vtypes != "C").any(): kwargs["vtype"] = M.vtypes x = gm.addMVar(M.vlabels.shape, M.lb, M.ub, **kwargs) - if model.is_quadratic: - assert M.Q is not None + if M.Q is not None: gm.setObjective(0.5 * x.T @ M.Q @ x + M.c @ x) else: gm.setObjective(M.c @ x) - if model.objective.sense == "max": + if M.objective_sense == "max": gm.ModelSense = -1 if M.A is not None: c = gm.addMConstr(M.A, x, M.sense, M.b) - if set_names: - names = print_constraints(M.clabels) - c.setAttr("ConstrName", names) + if names is not None: + c.setAttr("ConstrName", print_constraints(M.clabels)) - for sos_type, positions, weights in _iter_sos_sets(model): + for sos_type, positions, weights in sos: gm.addSOS(sos_type, x[positions.tolist()].tolist(), weights.tolist()) if M.indicator_A is not None: @@ -2093,13 +2196,13 @@ def _build_solver_model( } _GUROBI_SENSE_MAP: ClassVar[dict[str, int]] = {"min": 1, "max": -1} - def _apply_begin(self, var_label_index: Any, con_label_index: Any) -> Any: + def _apply_begin(self, n_vars: int, n_cons: int) -> Any: gm = self.solver_model gurobi_vars = gm.getVars() gurobi_cons = gm.getConstrs() - if len(gurobi_vars) != var_label_index.n_active_vars: + if len(gurobi_vars) != n_vars: raise UnsupportedUpdate("gurobi var count mismatch") - if len(gurobi_cons) != con_label_index.n_active_cons: + if len(gurobi_cons) != n_cons: raise UnsupportedUpdate("gurobi con count mismatch") return _GurobiApplyCtx(gm, gurobi_vars, gurobi_cons) @@ -2747,37 +2850,29 @@ def _apply_obj_sense(self, ctx: Any, sense: str) -> None: except AttributeError: # Fallback to old API ctx.chgobjsense(direction) - def _build_direct( + def _new_solver_model( self, - explicit_coordinate_names: bool = False, + M: Matrices, *, - set_names: bool, + names: Printers | None, + sos: list[SosSet], **kwargs: Any, - ) -> None: - model = self.model - assert model is not None + ) -> xpress.problem: self.close() self._env_stack = contextlib.ExitStack() - problem = self._build_solver_model( - model, - explicit_coordinate_names=explicit_coordinate_names, - set_names=set_names, - ) + problem = self._build_solver_model(M, names=names, sos=sos) self._env_stack.enter_context(problem) - self.solver_model = problem - self.io_api = "direct" - self.sense = model.sense - self._cache_model_labels(model) + return problem @staticmethod def _build_solver_model( - model: Model, - explicit_coordinate_names: bool = False, + M: Matrices, *, - set_names: bool, + names: Printers | None, + sos: list[SosSet], ) -> xpress.problem: """ - Build an ``xpress.problem`` that mirrors the linopy ``model`` via ``loadproblem``. + Build an ``xpress.problem`` from the matrices ``M`` via ``loadproblem``. ``loadproblem`` is Xpress' universal native-array entry point loading LP/QP/MIQP in a single call; see the parameter reference at @@ -2785,10 +2880,8 @@ def _build_solver_model( SOS arguments are left ``None`` and sets are added afterwards via ``addSOS`` so multi-dim ``add_sos_constraints`` can be grouped natively. """ - model.constraints.sanitize_missings() problem = xpress.problem() - M = model.matrices A = M.A Q = M.Q @@ -2912,16 +3005,14 @@ def _build_solver_model( refval=None, ) - if model.objective.sense == "max": + if M.objective_sense == "max": try: # Try new API first (Xpress 9.8+) problem.chgObjSense(xpress.maximize) except AttributeError: # Fallback to old API problem.chgobjsense(xpress.maximize) - if set_names: - print_variable, print_constraint = linopy.io.get_printers_scalar( - model, explicit_coordinate_names=explicit_coordinate_names - ) + if names is not None: + print_variable, print_constraint = names vnames = print_variable(M.vlabels) if vnames: try: # Try new API first (Xpress 9.8+) @@ -2939,7 +3030,7 @@ def _build_solver_model( except AttributeError: # Fallback to old API problem.addnames(xpress_Namespaces.ROW, cnames, 0, len(cnames) - 1) - for sos_type, positions, weights in _iter_sos_sets(model): + for sos_type, positions, weights in sos: problem.addSOS(positions.tolist(), weights.tolist(), type=sos_type) return problem @@ -3471,58 +3562,45 @@ def _run_direct( sense=self.sense, ) - def _build_direct( + def _new_solver_model( self, - explicit_coordinate_names: bool = False, + M: Matrices, *, - set_names: bool, + names: Printers | None, + sos: list[SosSet], task: mosek.Task | None = None, **kwargs: Any, - ) -> None: - model = self.model - assert model is not None + ) -> mosek.Task: self.close() self._env_stack = contextlib.ExitStack() if task is None: env = self._env_stack.enter_context(mosek.Env()) task = self._env_stack.enter_context(env.Task(0, 0)) - m = self._build_solver_model( - model, - task, - explicit_coordinate_names=explicit_coordinate_names, - set_names=set_names, - ) - self.solver_model = m - self.io_api = "direct" - self.sense = model.sense - self._cache_model_labels(model) + return self._build_solver_model(M, task, names=names, sos=sos) @staticmethod def _build_solver_model( - model: Model, + M: Matrices, task: mosek.Task, - explicit_coordinate_names: bool = False, *, - set_names: bool, + names: Printers | None, + sos: list[SosSet], ) -> mosek.Task: - """Populate an empty MOSEK task with the contents of `model`.""" - if model.variables.sos: + """Populate an empty MOSEK task with the matrices ``M``.""" + if sos: raise NotImplementedError("SOS constraints are not supported by MOSEK.") - if model.variables.semi_continuous: + if (M.vtypes == "S").any(): raise NotImplementedError( "Semi-continuous variables are not supported by MOSEK. " "Use a solver that supports them (gurobi, cplex, highs)." ) - task.appendvars(model.nvars) - task.appendcons(model.ncons) + n_vars, n_cons = len(M.vlabels), len(M.clabels) + task.appendvars(n_vars) + task.appendcons(n_cons) - M = model.matrices - - if set_names: - print_variables, print_constraints = linopy.io.get_printers_scalar( - model, explicit_coordinate_names=explicit_coordinate_names - ) + if names is not None: + print_variables, print_constraints = names labels = print_variables(M.vlabels) task.generatevarnames( np.arange(0, len(labels)), "%0", [len(labels)], None, [0], labels @@ -3539,17 +3617,17 @@ def _build_solver_model( ) blx = np.where(lb_fin, M.lb, 0.0) bux = np.where(ub_fin, M.ub, 0.0) - task.putvarboundslice(0, model.nvars, keys[bkx_code].tolist(), blx, bux) + task.putvarboundslice(0, n_vars, keys[bkx_code].tolist(), blx, bux) - if len(model.binaries.labels) + len(model.integers.labels) > 0: - idx = np.flatnonzero(np.isin(M.vtypes, ["B", "I"])).astype(np.int32) + idx = np.flatnonzero(np.isin(M.vtypes, ["B", "I"])).astype(np.int32) + if idx.size: task.putvartypelist(idx, [mosek.variabletype.type_int] * len(idx)) - if len(model.constraints) > 0: - if set_names: - names = print_constraints(M.clabels) + if n_cons: + if names is not None: + cnames = print_constraints(M.clabels) task.generateconnames( - np.arange(0, len(names)), "%0", [len(names)], None, [0], names + np.arange(0, len(cnames)), "%0", [len(cnames)], None, [0], cnames ) b_lo = M.b > -np.inf b_up = M.b < np.inf @@ -3563,16 +3641,16 @@ def _build_solver_model( if M.A is not None: A = M.A.tocsr() task.putarowslice( - 0, model.ncons, A.indptr[:-1], A.indptr[1:], A.indices, A.data + 0, n_cons, A.indptr[:-1], A.indptr[1:], A.indices, A.data ) - task.putconboundslice(0, model.ncons, keys[bkc_code].tolist(), blc, buc) + task.putconboundslice(0, n_cons, keys[bkc_code].tolist(), blc, buc) if M.Q is not None: Q = (0.5 * tril(M.Q + M.Q.transpose())).tocoo() task.putqobj(Q.row, Q.col, Q.data) - task.putclist(np.arange(model.nvars, dtype=np.int32), M.c) + task.putclist(np.arange(n_vars, dtype=np.int32), M.c) - if model.objective.sense == "max": + if M.objective_sense == "max": task.putobjsense(mosek.objsense.maximize) else: task.putobjsense(mosek.objsense.minimize) @@ -4234,35 +4312,40 @@ def _run_file( env=env, ) - def _build_direct(self, **kwargs: Any) -> None: - model = self.model - assert model is not None - if model.type in ["QP", "MILP"]: - msg = "cuPDLPx does not currently support QP or MILP problems." - raise NotImplementedError(msg) - if kwargs.get("explicit_coordinate_names"): + def _build_direct( + self, *, explicit_coordinate_names: bool = False, **kwargs: Any + ) -> None: + if explicit_coordinate_names: warnings.warn( "cuPDLPx does not support named variables/constraints. " "The explicit_coordinate_names parameter is ignored.", UserWarning, stacklevel=2, ) - cu_model = self._build_solver_model(model) - self.solver_model = cu_model - self.io_api = "direct" - self.sense = model.sense - self._cache_model_labels(model) + super()._build_direct(**kwargs) + + def _new_solver_model( + self, + M: Matrices, + *, + names: Printers | None, + sos: list[SosSet], + **kwargs: Any, + ) -> cupdlpx.Model: + if M.type in ["QP", "MILP"]: + msg = "cuPDLPx does not currently support QP or MILP problems." + raise NotImplementedError(msg) + return self._build_solver_model(M) @staticmethod - def _build_solver_model(model: Model) -> cupdlpx.Model: - """Build a cupdlpx.Model that mirrors the linopy `model`.""" - if model.variables.semi_continuous: + def _build_solver_model(M: Matrices) -> cupdlpx.Model: + """Build a cupdlpx.Model from the matrices ``M``.""" + if (M.vtypes == "S").any(): raise NotImplementedError( "Semi-continuous variables are not supported by cuPDLPx. " "Use a solver that supports them (gurobi, cplex, highs)." ) - M = model.matrices if M.A is None: raise ValueError("Model has no constraints, cannot export to cuPDLPx.") A = M.A.tocsr() @@ -4286,7 +4369,7 @@ def _build_solver_model(model: Model) -> cupdlpx.Model: variable_upper_bound=M.ub, ) - if model.objective.sense == "max": + if M.objective_sense == "max": cu_model.ModelSense = cupdlpx.PDLP.MAXIMIZE return cu_model @@ -4619,37 +4702,43 @@ def _build_file(self, **build_kwargs: Any) -> None: build_kwargs.pop("progress", None) self._build_direct(**build_kwargs) - def _build_direct(self, **kwargs: Any) -> None: - model = self.model - assert model is not None - if model.type in ("MIQP", "IQP"): - msg = ( - "cuOpt does not support quadratic objectives together with " - "integer or semi-continuous variables. Use a solver that " - "supports MIQP (gurobi, xpress, mosek)." - ) - raise NotImplementedError(msg) - if kwargs.get("explicit_coordinate_names"): + def _build_direct( + self, *, explicit_coordinate_names: bool = False, **kwargs: Any + ) -> None: + if explicit_coordinate_names: warnings.warn( "cuOpt does not support named variables/constraints. " "The explicit_coordinate_names parameter is ignored.", UserWarning, stacklevel=2, ) - self.solver_model = self._build_solver_model(model) - self.io_api = "direct" - self.sense = model.sense - self._cache_model_labels(model) + super()._build_direct(**kwargs) + + def _new_solver_model( + self, + M: Matrices, + *, + names: Printers | None, + sos: list[SosSet], + **kwargs: Any, + ) -> cuopt.linear_programming.DataModel: + if M.type in ("MIQP", "IQP"): + msg = ( + "cuOpt does not support quadratic objectives together with " + "integer or semi-continuous variables. Use a solver that " + "supports MIQP (gurobi, xpress, mosek)." + ) + raise NotImplementedError(msg) + return self._build_solver_model(M) @staticmethod - def _build_solver_model(model: Model) -> cuopt.linear_programming.DataModel: - """Build a cuopt DataModel that mirrors the linopy `model`.""" - M = model.matrices + def _build_solver_model(M: Matrices) -> cuopt.linear_programming.DataModel: + """Build a cuopt DataModel from the matrices ``M``.""" # Always transform to minimize problem: `min -obj`. Don't use cuOpt's # `set_maximize(True)`, whose presolve path returns negated duals # (see class Notes). `_solve` flips back the sign when needed # (negates objective, duals and MIP bound - primal untouched). - sign = -1.0 if model.objective.sense == "max" else 1.0 + sign = -1.0 if M.objective_sense == "max" else 1.0 lower = np.where( np.logical_or(np.equal(M.sense, ">"), np.equal(M.sense, "=")), diff --git a/linopy/store.py b/linopy/store.py index 9facf686..44a85547 100644 --- a/linopy/store.py +++ b/linopy/store.py @@ -864,17 +864,13 @@ def active_terms( 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 + from linopy.csr import index_dtype, indptr_from_counts 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) - ) - indptr[0] = 0 - np.cumsum(counts, out=indptr[1:]) - return data, labels, indptr + dtype = index_dtype(len(data), counts.shape, self.model) + return data, labels, indptr_from_counts(counts, dtype) def positional_csr( self, label_index: VariableLabelIndex, active_rows: np.ndarray, name: str diff --git a/test/test_cuopt.py b/test/test_cuopt.py index b1faa003..6b4be534 100644 --- a/test/test_cuopt.py +++ b/test/test_cuopt.py @@ -499,7 +499,7 @@ def test_status_map_empty_constraint_matrix() -> None: def test_status_map_mismatched_variable_types() -> None: """A malformed input array must not be reported as a solved model.""" model = milp_model() - data_model = solvers.cuOpt._build_solver_model(model) + data_model = solvers.cuOpt._build_solver_model(model.matrices) data_model.set_variable_types(np.array(["I"], dtype=" None: def test_status_map_numerical_error() -> None: """A concave quadratic objective makes cuOpt fail numerically.""" model = sign_matrix_model("min", "<=") - data_model = solvers.cuOpt._build_solver_model(model) + data_model = solvers.cuOpt._build_solver_model(model.matrices) Q = csr_array(np.diag([-1.0, 0.0, 0.0])) data_model.set_quadratic_objective_matrix(Q.data, Q.indices, Q.indptr) @@ -696,7 +696,7 @@ def test_naive_hessian_changes_the_answer() -> None: reachable and that the tolerances above would reject it. """ model = three_variable_qp("min") - data_model = solvers.cuOpt._build_solver_model(model) + data_model = solvers.cuOpt._build_solver_model(model.matrices) assert model.matrices.Q is not None naive = csr_array(model.matrices.Q.tocsr()) data_model.set_quadratic_objective_matrix(naive.data, naive.indices, naive.indptr) diff --git a/test/test_matrices.py b/test/test_matrices.py index 83bc4e96..9531c1ac 100644 --- a/test/test_matrices.py +++ b/test/test_matrices.py @@ -11,6 +11,7 @@ import xarray as xr from linopy import EQUAL, GREATER_EQUAL, Model +from linopy.matrices import Matrices def test_basic_matrices() -> None: @@ -101,6 +102,27 @@ def test_matrices_float_c() -> None: assert np.all(c == np.array([1.5, 1.5])) +@pytest.mark.parametrize( + ("integer", "binary", "quadratic", "expected"), + [ + (False, False, False, "LP"), + (False, True, False, "MILP"), + (True, True, False, "ILP"), + (False, False, True, "QP"), + (False, True, True, "MIQP"), + ], +) +def test_matrices_type( + integer: bool, binary: bool, quadratic: bool, expected: str +) -> None: + m = Model() + x = m.add_variables(0, 10, coords=[range(2)], name="x", integer=integer) + y = m.add_variables(name="y", binary=binary) + m.add_constraints(x.sum() + y >= 1) + m.add_objective(x * x + y if quadratic else x.sum() + y) + assert m.matrices.type == m.type == expected + + def test_matrices_sol_aligned_with_vlabels() -> None: m = Model() i = pd.RangeIndex(3, name="i") @@ -116,6 +138,12 @@ def test_matrices_sol_aligned_with_vlabels() -> None: np.testing.assert_array_equal(M.vlabels, [0, 2, 3, 4, 5]) np.testing.assert_array_equal(M.sol, [1.0, 3.0, 4.0, 5.0, 6.0]) + stale = m.matrices + z = m.add_variables(coords=[i], name="z") + z.solution = xr.DataArray([7.0, 8.0, 9.0], coords=[i]) + with pytest.raises(ValueError, match="expected 5"): + stale.sol + @pytest.mark.parametrize("sparse", [False, True], ids=["dense", "sparse"]) def test_matrices_scaled_masked_mixed_signs(sparse: bool) -> None: @@ -156,3 +184,57 @@ def test_matrices_scaled_masked_mixed_signs(sparse: bool) -> None: np.testing.assert_array_equal(M.indicator_binvar, [6, 7, 8]) np.testing.assert_array_equal(M.lb[:3], [0.0, 0.0, 0.0]) np.testing.assert_array_equal(M.ub[:3], [4.0, 8.0, 16.0]) + + +@pytest.mark.parametrize( + ("rows", "cols"), + [([0, 2], [1, 5, 6]), ([4, 1, 3], [7, 0, 2]), ([], [3])], + ids=["subset", "reordered", "no-rows"], +) +def test_matrices_take(rows: list[int], cols: list[int]) -> None: + m = Model() + i = pd.RangeIndex(4, name="i") + x = m.add_variables(0, 4, coords=[i], name="x", scaling=[1.0, 2.0, 4.0, 1.0]) + y = m.add_variables(coords=[i], name="y", integer=True) + m.add_constraints(x + 2 * y >= 1, name="c", scaling=[1.0, 3.0, 1.0, 2.0]) + m.add_constraints(x.sum() <= 3, name="d") + m.add_objective((x * x).sum() + 3 * y.sum(), sense="max") + M = m.matrices + r, c = np.array(rows, dtype=np.intp), np.array(cols, dtype=np.intp) + + taken = M.take(r, c) + + assert type(taken) is Matrices + assert taken.A is not None and M.A is not None + assert taken.Q is not None and M.Q is not None + np.testing.assert_array_equal(taken.A.toarray(), M.A.toarray()[np.ix_(r, c)]) + np.testing.assert_array_equal(taken.Q.toarray(), M.Q.toarray()[np.ix_(c, c)]) + for column_side, full in [ + (taken.vlabels, M.vlabels), + (taken.c, M.c), + (taken.lb, M.lb), + (taken.ub, M.ub), + (taken.vtypes, M.vtypes), + (taken.var_scaling, M.var_scaling), + ]: + np.testing.assert_array_equal(column_side, full[c]) + for row_side, full in [ + (taken.clabels, M.clabels), + (taken.b, M.b), + (taken.sense, M.sense), + (taken.con_scaling, M.con_scaling), + ]: + np.testing.assert_array_equal(row_side, full[r]) + assert (taken.objective_scaling, taken.objective_sense) == (1.0, "max") + + +def test_matrices_take_rejects_indicator_constraints() -> None: + m = Model() + x = m.add_variables(0, 4, coords=[range(2)], name="x") + z = m.add_variables(coords=[range(2)], name="z", binary=True) + m.add_constraints(x >= 0, name="c") + m.add_indicator_constraints(z, 1, x <= 2, name="ind") + rows, cols = np.arange(1), np.arange(2) + + with pytest.raises(NotImplementedError, match="indicator"): + m.matrices.take(rows, cols) diff --git a/test/test_persistent_snapshot_diff.py b/test/test_persistent_snapshot_diff.py index 2d84b3d7..e76b3e4e 100644 --- a/test/test_persistent_snapshot_diff.py +++ b/test/test_persistent_snapshot_diff.py @@ -334,6 +334,7 @@ def test_diff_snapshot_matches_capture(baseline: Model, mutate: str) -> None: baseline.add_objective(3 * x.sum() + 2 * y.sum(), overwrite=True) diff = ModelDiff.from_snapshot(snap, baseline) assert isinstance(diff, ModelDiff) + assert diff.snapshot is not None _assert_snapshot_equal(diff.snapshot, ModelSnapshot.capture(baseline)) @@ -350,6 +351,7 @@ def build(t0: int) -> Model: snap = ModelSnapshot.capture(m1) diff = ModelDiff.from_snapshot(snap, m2, ignore_dims={"t"}) assert isinstance(diff, ModelDiff) + assert diff.snapshot is not None _assert_snapshot_equal(diff.snapshot, ModelSnapshot.capture(m2)) @@ -364,6 +366,7 @@ def build(rhs: float) -> Model: m1, m2 = build(4.0), build(7.0) diff = ModelDiff.from_models(m1, m2) assert isinstance(diff, ModelDiff) + assert diff.snapshot is not None _assert_snapshot_equal(diff.snapshot, ModelSnapshot.capture(m2)) diff --git a/test/test_persistent_solver_orchestrator.py b/test/test_persistent_solver_orchestrator.py index 35f3638b..8d8e07bb 100644 --- a/test/test_persistent_solver_orchestrator.py +++ b/test/test_persistent_solver_orchestrator.py @@ -121,9 +121,7 @@ def test_unmutated_resolve_diff_is_empty(model: Model) -> None: class FakePersistentSolver(FakeSolver): supports_persistent_update = True - def apply_update( - self, diff: ModelDiff, var_label_index: Any, con_label_index: Any - ) -> None: + def apply_update(self, diff: ModelDiff, n_vars: int, n_cons: int) -> None: return None diff --git a/test/test_restrict_matrix.py b/test/test_restrict_matrix.py new file mode 100644 index 00000000..a04bd404 --- /dev/null +++ b/test/test_restrict_matrix.py @@ -0,0 +1,687 @@ +"""Matrix-level windows agree with the container-level restriction.""" + +from __future__ import annotations + +import logging +from dataclasses import replace +from typing import Any + +import numpy as np +import pandas as pd +import pytest +import xarray as xr + +from linopy import LinearExpression, Model, merge +from linopy.matrices import Matrices +from linopy.persistent.diff import ModelDiff, RebuildReason +from linopy.persistent.errors import RebuildRequiredError +from linopy.restrict import MatrixWindow, RestrictionError +from linopy.solvers import Solver + +TIME = pd.RangeIndex(8, name="time") +BUS = pd.Index(["a", "b"], name="bus") +BOUNDARY = xr.DataArray(np.where(TIME == 0, 7.0, 0.0), coords=[TIME]) + + +def storage_model() -> Model: + """Two coupled constraints: a shifted balance and an elementwise cap.""" + m = Model() + soc = m.add_variables(lower=0, upper=10, coords=[TIME], name="soc") + charge = m.add_variables(lower=0, upper=5, coords=[TIME], name="charge") + m.add_constraints(soc - 0.9 * soc.shift(time=1) - charge == 0, name="balance") + m.add_constraints(charge <= 4, name="cap") + m.add_objective((charge + soc).sum()) + return m + + +def initial(soc_value: float = 3.0) -> pd.Series: + return pd.Series(np.full(len(TIME), soc_value), index=TIME) + + +def assert_same_window(window: Matrices, reference: Matrices) -> None: + assert window.A is not None and reference.A is not None + assert window.A.shape == reference.A.shape + assert np.array_equal(window.A.indptr, reference.A.indptr) + assert np.array_equal(window.A.indices, reference.A.indices) + assert np.allclose(window.A.data, reference.A.data) + assert np.allclose(window.b, reference.b) + assert np.allclose(window.c, reference.c) + assert np.allclose(window.lb, reference.lb) + assert np.allclose(window.ub, reference.ub) + assert np.array_equal(window.sense, reference.sense) + + +@pytest.mark.parametrize( + "window", + [TIME[0:3], TIME[1:4], TIME[3:6], [3, 1], [1, 3, 6], [6, 2, 4]], + ids=["head", "inner", "tail", "reversed", "gapped", "unordered"], +) +def test_select_matches_dense_reference(window: Any) -> None: + m = storage_model() + mat = m.matrices + selection = {"time": window} + + def positions(containers: Any) -> np.ndarray: + flat = [ + c.labels.sel(selection).values.ravel() for c in containers.data.values() + ] + labels = np.concatenate(flat) + return containers.label_index.label_to_pos[labels[labels != -1]] + + rows, cols = positions(m.constraints), positions(m.variables) + fixed = np.zeros(mat.vlabels.size) + fixed[m.variables.label_index.label_to_pos[m.variables["soc"].labels.values]] = 3.0 + fixed[cols] = 0.0 + assert mat.A is not None + A = mat.A.toarray()[rows] + + window_mat = MatrixWindow(m).select(selection, fix={"soc": initial()})[0] + + assert window_mat.A is not None + assert np.allclose(window_mat.A.toarray(), A[:, cols]) + assert np.allclose(window_mat.b, mat.b[rows] - A @ fixed) + assert np.allclose(window_mat.c, mat.c[cols]) + + +def test_select_without_fix_matches_restrict() -> None: + """A model with no cross-window coupling needs no boundary values.""" + m = Model() + x = m.add_variables(lower=0, coords=[TIME], name="x") + m.add_constraints(x <= 4, name="cap") + m.add_objective(x.sum()) + + window = TIME[2:5] + matrix = MatrixWindow(m).select({"time": window})[0] + reference = m.restrict(time=window).matrices + + assert_same_window(matrix, reference) + + +def test_window_keeps_parent_labels() -> None: + m = storage_model() + window = TIME[1:4] + selected, _, _ = MatrixWindow(m).select({"time": window}, fix={"soc": initial()}) + + assert selected.A is not None + assert selected.A.shape == (selected.clabels.size, selected.vlabels.size) + assert np.isin(selected.vlabels, m.variables.label_index.vlabels).all() + + +@pytest.mark.parametrize( + ("sign", "rhs", "satisfied"), + [("==", 2, False), (">=", 2, False), ("<=", -2, False), ("<=", 2, True)], +) +def test_folded_row_is_checked_against_its_sign( + sign: str, rhs: float, satisfied: bool +) -> None: + """A row that loses every term is dropped only when its sign holds.""" + m = Model() + x = m.add_variables(lower=0, coords=[TIME], name="x") + m.add_constraints(x == 1, name="eq") + m.add_constraints(x.isel(time=0), sign, rhs, name="pin") + m.add_objective(x.sum()) + + window = MatrixWindow(m) + fix = {"x": pd.Series(np.zeros(len(TIME)), index=TIME)} + if satisfied: + assert window.select({"time": TIME[1:3]}, fix=fix)[0].clabels.size == 2 + return + with pytest.raises(RestrictionError, match="unsatisfiable"): + window.select({"time": TIME[1:3]}, fix=fix) + + +def test_binary_integrality_is_preserved() -> None: + """A discrete column stays discrete through the window.""" + m = Model() + x = m.add_variables(binary=True, coords=[TIME], name="x") + m.add_constraints(x >= 0.5, name="half") + m.add_objective(x.sum(), sense="min") + + window = m.restrict(time=TIME[:3]) + window.solve(solver_name="highs", output_flag=False) + + assert set(np.unique(window.solution["x"].values)) <= {0.0, 1.0} + assert window.solution["x"].values.min() == 1.0 + + +def test_variable_scaling_is_undone() -> None: + """A scaled column reports the model's units, not the solver's.""" + m = Model() + x = m.add_variables(lower=0, upper=100, coords=[TIME], scaling=10, name="x") + m.add_constraints(x >= 50, name="lb") + m.add_objective(x.sum(), sense="min") + + window = m.restrict(time=TIME[:3]) + window.solve(solver_name="highs", output_flag=False) + + assert np.allclose(window.solution["x"].values, 50.0) + + +def test_constraint_scaling_is_undone() -> None: + """A scaled row reports the model's dual units.""" + m = Model() + x = m.add_variables(lower=0, coords=[TIME], name="x") + m.add_constraints(x >= 5, name="lb", scaling=2) + m.add_objective(x.sum(), sense="min") + + window = m.restrict(time=TIME[:3]) + window.solve(solver_name="highs", output_flag=False) + + assert np.allclose(window.dual["lb"].values, 1.0) + + +def test_subset_carrier_raises() -> None: + """Every carrier of the selected dimension must share its index.""" + m = Model() + x = m.add_variables(coords=[TIME], name="x") + subset = m.add_variables(coords=[TIME[::2]], name="subset_x") + m.add_constraints(x <= 1, name="cap") + m.add_constraints(subset <= 1, name="subset_cap") + m.add_objective(m.variables["x"].sum() + subset.sum()) + + with pytest.raises(RestrictionError, match="does not carry the selected labels"): + MatrixWindow(m).select({"time": TIME[1:3]}) + + +def linear_model() -> Model: + m = Model() + x = m.add_variables(lower=0, coords=[TIME], name="x") + m.add_constraints(x <= 5, name="cap") + m.add_objective(x.sum(), sense="min") + return m + + +def test_restrict_rejects_unknown_dimension() -> None: + with pytest.raises(RestrictionError, match="does not exist"): + linear_model().restrict(period=[0]) + + +def test_window_rejects_quadratic_objective() -> None: + m = Model() + x = m.add_variables(lower=0, coords=[TIME], name="x") + m.add_constraints(x <= 5, name="cap") + m.add_objective((x * x).sum()) + + with pytest.raises(NotImplementedError, match="quadratic"): + MatrixWindow(m) + + +def test_window_rejects_indicator_constraints() -> None: + m = Model() + x = m.add_variables(lower=0, coords=[TIME], name="x") + flag = m.add_variables(binary=True, coords=[TIME], name="flag") + m.add_indicator_constraints(flag, 1, x, "<=", 1, name="ind") + m.add_objective(x.sum()) + + with pytest.raises(NotImplementedError, match="indicator"): + MatrixWindow(m) + + +def test_window_requires_a_constraint() -> None: + m = Model() + m.add_variables(lower=0, coords=[TIME], name="x") + m.add_objective(m.variables["x"].sum()) + + with pytest.raises(NotImplementedError, match="at least one constraint"): + MatrixWindow(m) + + +def test_fix_must_name_a_variable() -> None: + with pytest.raises(RestrictionError, match="does not name a variable"): + MatrixWindow(linear_model()).select({"time": TIME[:2]}, fix={"nope": 1.0}) + + +@pytest.mark.parametrize( + "given", + [ + [1.0, 2.0], + BOUNDARY.isel(time=slice(0, 4)), + BOUNDARY.expand_dims(other=[0]), + xr.DataArray(np.ones(len(TIME) - 1), dims="time"), + ], + ids=["wrong-length", "missing-coords", "extra-dim", "coordless-wrong-length"], +) +def test_fix_must_match_the_variable_coordinates(given: object) -> None: + with pytest.raises(RestrictionError, match="do not match its coordinates"): + MatrixWindow(linear_model()).select({"time": TIME[:2]}, fix={"x": given}) + + +def ramp_model(scaling: float = 1.0, mask: xr.DataArray | None = None) -> Model: + """A two-dimensional ramp: every step must not fall below the previous one.""" + m = Model() + x = m.add_variables( + 0, 100, coords=[TIME, BUS], scaling=scaling, mask=mask, name="x" + ) + m.add_constraints(x - x.shift(time=1) >= 0, name="ramp") + m.add_objective(x.sum()) + return m + + +def solve_window(m: Model, fix: object, **selection: object) -> xr.DataArray: + window = m.restrict(fix={"x": fix}, **selection) + window.solve(solver_name="highs", output_flag=False) + return window.solution["x"] + + +MASK = xr.DataArray([[True, True]] + [[True, False]] * 7, coords=[TIME, BUS]) +ACTIVE = xr.DataArray([[True, True]] * 8, coords=[TIME, BUS]) + + +@pytest.mark.parametrize("mask", [ACTIVE, MASK], ids=["active", "masked"]) +@pytest.mark.parametrize("scaling", [1.0, 10.0]) +@pytest.mark.parametrize( + "fix", + [ + BOUNDARY, + BOUNDARY.isel(time=slice(None, None, -1)), + BOUNDARY.broadcast_like(xr.DataArray(BUS, coords=[BUS])).transpose(), + BOUNDARY.to_series(), + xr.DataArray(BOUNDARY.values, dims=["time"]), + 7.0, + ], + ids=["broadcast", "reversed", "transposed", "series", "positional", "scalar"], +) +def test_fix_is_aligned_by_coordinates( + scaling: float, fix: object, mask: xr.DataArray +) -> None: + """Fixed values align by label, broadcast over missing dims, in model units.""" + solution = solve_window(ramp_model(scaling, mask), fix, time=TIME[1:3]) + + assert solution.dims == ("time", "bus") + assert solution.where(mask).equals(xr.full_like(solution, 7.0).where(mask)) + + +@pytest.mark.parametrize( + ("selection", "sizes"), + [ + ({"time": slice(1, 2)}, {"time": 2, "bus": 2}), + ({"time": [1, 2]}, {"time": 2, "bus": 2}), + ({"time": slice(1, 2), "bus": ["a"]}, {"time": 2, "bus": 1}), + ({"time": 1, "bus": "b"}, {}), + ], + ids=["slice", "list", "multi-dim", "scalars"], +) +def test_selection_follows_sel_semantics( + selection: dict[str, object], sizes: dict[str, int] +) -> None: + solution = solve_window(ramp_model(), 7.0, **selection) + + assert dict(solution.sizes) == sizes + assert np.allclose(solution.values, 7.0) + assert solution.coords.to_dataset().equals( + ramp_model().variables["x"].labels.sel(selection).coords.to_dataset() + ) + + +def test_slice_selection_matches_label_selection() -> None: + window = MatrixWindow(storage_model()) + fix = {"soc": initial()} + + assert_same_window( + window.select({"time": slice(1, 3)}, fix=fix)[0], + window.select({"time": TIME[1:4]}, fix=fix)[0], + ) + + +def sparse_coupling_model() -> Model: + """A coefficient that is zero at one step, so windows differ in sparsity.""" + m = Model() + x = m.add_variables(0, 100, coords=[TIME[:4]], name="x") + y = m.add_variables(0, 100, coords=[TIME[:4]], name="y") + k = xr.DataArray([1.0, 1.0, 0.0, 1.0], coords=[TIME[:4]]) + m.add_constraints(x + k * y >= 3, name="c") + m.add_objective((x + 2 * y).sum()) + return m + + +def alternating_sign_model() -> Model: + """Constraint signs alternate in time, so shifted windows differ in sign.""" + m = Model() + x = m.add_variables(0, 10, coords=[TIME], name="x") + sign = xr.DataArray(np.where(TIME % 2 == 0, ">=", "<="), coords=[TIME]) + m.add_constraints(x, sign, 5, name="c") + m.add_objective(x.sum()) + return m + + +@pytest.mark.parametrize( + ("m", "first", "second", "reason", "expected"), + [ + ( + sparse_coupling_model(), + TIME[:2], + TIME[2:4], + RebuildReason.STRUCTURAL_LABELS, + {"x": [3, 3], "y": [0, 0]}, + ), + ( + alternating_sign_model(), + TIME[:2], + TIME[1:3], + RebuildReason.BACKEND_REJECTED, + {"x": [0, 5]}, + ), + ], + ids=["sparsity", "sign"], +) +def test_advance_rebuilds_when_an_update_is_impossible( + m: Model, + first: pd.Index, + second: pd.Index, + reason: RebuildReason, + expected: dict[str, list[float]], + caplog: pytest.LogCaptureFixture, +) -> None: + window = m.restrict(time=first) + window.solve(solver_name="highs", output_flag=False) + window.advance(time=second) + + with pytest.raises(RebuildRequiredError) as excinfo: + window.solve(disallow_rebuild=True) + assert excinfo.value.reason is reason + rejected = reason is RebuildReason.BACKEND_REJECTED + assert (excinfo.value.__cause__ is not None) is rejected + with caplog.at_level(logging.WARNING, logger="linopy.solvers"): + window.solve() + assert "rebuild required" not in caplog.text + + assert window._solver is not None + assert window._solver._rebuilds == 1 + assert window._solver._last_rebuild_reason is reason + for name, values in expected.items(): + assert np.allclose(window.solution[name].values, values) + + +def test_from_windows_rejects_a_changed_sparsity_pattern() -> None: + m = sparse_coupling_model() + window = MatrixWindow(m) + first = window.select({"time": TIME[:2]})[0] + second = window.select({"time": TIME[2:4]})[0] + + assert first.A is not None and second.A is not None + assert first.A.shape == second.A.shape + with pytest.raises(ValueError, match="sparsity pattern"): + ModelDiff.from_windows(first, second) + + +@pytest.mark.parametrize( + "selection", + [{"other": [0]}, {"time": TIME[2:4], "bus": ["a"]}], + ids=["missing", "extra"], +) +def test_advance_must_name_exactly_the_selected_dimensions( + selection: dict[str, Any], +) -> None: + window = ramp_model().restrict(fix={"x": 0.0}, time=TIME[:2]) + with pytest.raises(RestrictionError, match="name exactly the dimensions"): + window.advance(**selection) + + +@pytest.mark.parametrize( + ("selection", "match"), + [ + ({"time": slice(20, 30)}, "is empty"), + ({"time": [1, 1]}, "repeats labels"), + ({"time": [1, 2], "bus": ["a", "a"]}, "repeats labels"), + ], + ids=["empty", "repeated", "repeated-multi-dim"], +) +def test_select_rejects_degenerate_selections( + selection: dict[str, object], match: str +) -> None: + with pytest.raises(RestrictionError, match=match): + ramp_model().restrict(fix={"x": 0.0}, **selection) + + +def test_advance_without_selection_keeps_the_window() -> None: + window = linear_model().restrict(time=TIME[:2]) + before = window.matrices.b.copy() + + window.advance() + + assert np.array_equal(window.matrices.b, before) + + +def test_solve_updates_options_on_later_calls() -> None: + window = linear_model().restrict(time=TIME[:2]) + window.solve(solver_name="highs", output_flag=False) + window.solve(output_flag=False) + window.solve() + + assert window.objective.value == pytest.approx(0.0) + + +@pytest.mark.parametrize( + "read", + [lambda w: w.solution, lambda w: w.dual, lambda w: w["x"]], + ids=["solution", "dual", "getitem"], +) +def test_reading_values_before_solve_raises(read: Any) -> None: + window = linear_model().restrict(time=TIME[:2]) + with pytest.raises(AttributeError, match="has not been solved"): + read(window) + + +def test_dual_size_mismatch_raises() -> None: + window = linear_model().restrict(time=TIME[:2]) + window.solve(solver_name="highs", output_flag=False) + assert window._result is not None + window._result = replace(window._result, dual=np.empty(0)) + + with pytest.raises(AttributeError, match="no dual values"): + window.dual + + +def test_getitem_after_solve_returns_values() -> None: + window = linear_model().restrict(time=TIME[:2]) + window.solve(solver_name="highs", output_flag=False) + + assert np.allclose(window["x"].values, 0.0) + + +def test_getitem_unknown_variable_raises() -> None: + window = linear_model().restrict(time=TIME[:2]) + with pytest.raises(KeyError): + window["nope"] + + +def kinds_model(kinds: str, time: pd.Index, scaling: float = 1.0) -> Model: + """Unit demand split across continuous ``C`` and cheaper integer ``I`` supply.""" + m = Model() + supply = { + kind: m.add_variables(0, 5, coords=[time], integer=kind == "I", name=kind) + for kind in kinds + } + cost = {"C": 1.0, "I": 0.6} + lhs = merge([var.to_linexpr() for var in supply.values()], cls=LinearExpression) + m.add_constraints(lhs >= 1.5, name="demand") + m.add_objective( + merge([cost[k] * var.sum() for k, var in supply.items()], cls=LinearExpression), + scaling=scaling, + ) + return m + + +@pytest.mark.parametrize("scaling", [1.0, 10.0]) +@pytest.mark.parametrize( + "kinds", ["C", "CI", "I"], ids=["continuous", "mixed", "discrete"] +) +def test_window_solve_matches_the_restricted_model(kinds: str, scaling: float) -> None: + window = kinds_model(kinds, TIME, scaling).restrict(time=TIME[:2]) + window.solve(solver_name="highs", output_flag=False) + reference = kinds_model(kinds, TIME[:2], scaling) + reference.solve(solver_name="highs", output_flag=False) + + assert window.objective.value == pytest.approx(reference.objective.value) + for kind in kinds: + assert np.allclose(window[kind].values, reference.solution[kind].values) + + +def _window_matrices() -> Matrices: + return linear_model().restrict(time=TIME[:2]).matrices + + +def test_from_windows_reports_coefficient_changes() -> None: + previous = _window_matrices() + assert previous.A is not None + matrix = previous.A.copy() + matrix.data = matrix.data + 1.0 + + diff = ModelDiff.from_windows(previous, replace(previous, A=matrix)) + + assert diff.n_coef_updates > 0 + + +def test_from_windows_reports_variable_type_changes() -> None: + previous = _window_matrices() + vtypes = previous.vtypes.copy() + vtypes[0] = "B" + + diff = ModelDiff.from_windows(previous, replace(previous, vtypes=vtypes)) + + assert diff.var_type_positions.tolist() == [0] + + +def test_from_windows_reports_constraint_sign_changes() -> None: + previous = _window_matrices() + sense = previous.sense.copy() + sense[0] = ">=" + + diff = ModelDiff.from_windows(previous, replace(previous, sense=sense)) + + assert diff.con_sign_indices.tolist() == [0] + + +def test_from_windows_reports_objective_sense_changes() -> None: + previous = _window_matrices() + + diff = ModelDiff.from_windows(previous, replace(previous, objective_sense="max")) + + assert diff.obj_sense == "max" + + +def test_from_windows_reports_bounds_and_cost_changes() -> None: + previous = _window_matrices() + + diff = ModelDiff.from_windows( + previous, replace(previous, lb=previous.lb + 1.0, c=previous.c + 1.0) + ) + + assert diff.var_bounds_indices.size == previous.lb.size + assert diff.obj_c_indices is not None + + +def test_select_without_fix_rejects_dangling_terms() -> None: + """A window that cuts a coupling must be given the boundary values.""" + with pytest.raises(RestrictionError, match="dangling terms"): + MatrixWindow(storage_model()).select({"time": TIME[1:3]}) + + +def test_nan_fix_counts_as_missing() -> None: + fix = {"soc": pd.Series(np.full(len(TIME), np.nan), index=TIME)} + with pytest.raises(RestrictionError, match="dangling terms"): + MatrixWindow(storage_model()).select({"time": TIME[1:3]}, fix=fix) + + +WINDOW_SOLVERS = ("highs", "gurobi", "cplex", "scip", "xpress", "mosek") + + +@pytest.mark.parametrize("solver_name", WINDOW_SOLVERS) +def test_window_solve_across_solvers(solver_name: str) -> None: + """Every available direct solver updates a same-pattern window in place.""" + from linopy import licensed_solvers + + if solver_name not in licensed_solvers: + pytest.skip(f"{solver_name} is not available with a license") + + m = Model() + x = m.add_variables(lower=0, upper=10, coords=[TIME], name="x") + y = m.add_variables(lower=0, coords=[TIME], name="y") + m.add_constraints(x - 0.9 * x.shift(time=1) - y == 0, name="balance") + m.add_objective((x + y).sum()) + fix = {"x": pd.Series(TIME + 1.0, index=TIME)} + decay = 0.9 ** np.arange(1, 4) + + window = m.restrict(time=TIME[1:4], fix=fix) + try: + window.solve(solver_name=solver_name) + except NotImplementedError as exc: # a file-only backend + pytest.skip(str(exc)) + assert np.allclose(window.solution["x"].values, decay) + + window.advance(time=TIME[4:7], fix=fix) + window.solve(disallow_rebuild=True) + assert window._solver is not None + assert window._solver._in_place_updates == 1 + assert np.allclose(window.solution["x"].values, 4 * decay) + assert window.objective.value == pytest.approx(4 * decay.sum()) + + +def test_window_solve_detaches_the_solver_from_its_model() -> None: + solver = Solver.from_name("highs", io_api="direct", track_updates=True) + solver.solve(model=linear_model()) + solver.solve_window(linear_model().restrict(time=TIME[:2]).matrices) + + with pytest.raises(RuntimeError, match="matrix window"): + solver.update(linear_model()) + + +def test_solve_refuses_a_rebuild_when_asked() -> None: + m = linear_model() + fix = {"x": pd.Series(np.zeros(len(TIME)), index=TIME)} + window = m.restrict(time=TIME[:3], fix=fix) + window.solve(solver_name="highs", disallow_rebuild=True) + + window.advance(time=TIME[3:5], fix=fix) # a different width changes the shape + with pytest.raises(RebuildRequiredError): + window.solve(solver_name="highs", disallow_rebuild=True) + + +def test_solve_rebuilds_a_changed_shape() -> None: + m = linear_model() + fix = {"x": pd.Series(np.zeros(len(TIME)), index=TIME)} + window = m.restrict(time=TIME[:3], fix=fix) + window.solve(solver_name="highs") + + window.advance(time=TIME[3:5], fix=fix) + window.solve(solver_name="highs") + + assert window.solution["x"].sizes["time"] == 2 + + +def test_window_rejects_sos_constraints() -> None: + m = linear_model() + m.add_sos_constraints(m.variables["x"], 1, "time") + + with pytest.raises(NotImplementedError, match="SOS"): + MatrixWindow(m) + + +def test_window_solves_piecewise_formulations() -> None: + """Piecewise expands into ordinary variables and constraints.""" + from linopy.piecewise import add_piecewise_formulation + + m = Model() + x = m.add_variables(lower=0, upper=10, coords=[TIME], name="x") + y = m.add_variables(lower=0, coords=[TIME], name="y") + m.add_constraints(x <= 5, name="cap") + m.add_objective(y.sum()) + add_piecewise_formulation(m, (x, [0, 5, 10]), (y, [0, 4, 9])) + + window = m.restrict(time=TIME[:2]) + window.solve(solver_name="highs", output_flag=False) + + assert np.allclose(window.solution["y"].values, 0.0) + + +def test_window_rejects_horizon_wide_sos_rows() -> None: + """Reformulating SOS adds a horizon-wide cardinality row, which cannot fold.""" + m = Model() + x = m.add_variables(lower=0, upper=10, coords=[TIME], name="x") + m.add_constraints(x <= 5, name="cap") + m.add_objective(x.sum()) + m.add_sos_constraints(x, 1, "time") + m.apply_sos_reformulation() + + with pytest.raises(RestrictionError, match="dangling terms"): + m.restrict(time=TIME[:2]) diff --git a/test/test_solvers.py b/test/test_solvers.py index 1ca85ed0..57f77509 100644 --- a/test/test_solvers.py +++ b/test/test_solvers.py @@ -6,6 +6,7 @@ """ from pathlib import Path +from typing import Any from unittest.mock import MagicMock import numpy as np @@ -220,7 +221,7 @@ def test_copt_closes_env_per_solve( closes: list[object] = [] original = coptpy.Envr.close - def close(self: object) -> None: + def close(self: Any) -> None: closes.append(self) original(self) diff --git a/test/test_windowing_integration.py b/test/test_windowing_integration.py new file mode 100644 index 00000000..49392b1b --- /dev/null +++ b/test/test_windowing_integration.py @@ -0,0 +1,275 @@ +#!/usr/bin/env python3 +""" +Auto-executed windowing-equivalence integration tests. + +Each case registers a pre-existing linopy model — from the test suite, the +examples module or the benchmark suite — via the :func:`windowed` decorator. +A single parametrized driver then runs every case under both pattern modes +(strict and union), and assesses equivalence against the monolithic solve: + +* **decomposable** cases (no cross-tile terms): the chunked objectives sum to + the monolithic objective, and primal and dual values match per position; +* **replay** cases (coupled models): the :meth:`Model.advance` chain must be + indistinguishable from re-solving every window as a fresh + :meth:`Model.restrict` with the same carried fix sequence. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any, Literal + +import numpy as np +import pandas as pd +import pytest +import xarray as xr +from numpy.testing import assert_allclose + +from linopy import Model, available_solvers +from linopy.examples import benchmark_model + +needs_highs = pytest.mark.skipif( + "highs" not in available_solvers, reason="highs solver not available" +) + + +@dataclass(frozen=True) +class WindowCase: + """A model builder registered for windowing-equivalence testing.""" + + name: str + build: Callable[[], Model] + dim: str + tile: int + coupling: Literal["decomposable", "replay"] + fix_vars: tuple[str, ...] = () + + +CASES: list[WindowCase] = [] + + +def windowed( + dim: str, + tile: int, + coupling: Literal["decomposable", "replay"], + fix_vars: tuple[str, ...] = (), +) -> Callable[[Callable[[], Model]], Callable[[], Model]]: + """Register a model builder as a windowing-equivalence case.""" + + def deco(build: Callable[[], Model]) -> Callable[[], Model]: + CASES.append(WindowCase(build.__name__, build, dim, tile, coupling, fix_vars)) + return build + + return deco + + +# --------------------------------------------------------------------------- +# cases: pre-existing models, sources cited +# --------------------------------------------------------------------------- + + +@windowed( + dim="snapshot", + tile=3, + coupling="decomposable", + fix_vars=("x",), +) +def mi_period_demand() -> Model: + """ + From ``test/test_mi_feasibility.py``: per-period demand via groupby-sum. + Windows align with the period boundaries; a period's demand row leaks + into later windows only as a satisfied folded constant. + """ + m = Model() + x = m.add_variables( + lower=0, upper=4.0, coords=[pd.RangeIndex(6, name="snapshot")], name="x" + ) + period = xr.DataArray( + np.repeat([2020, 2030], 3), + dims="snapshot", + coords={"snapshot": x.indexes["snapshot"]}, + ) + e = (1.0 * x).assign_coords(period=period) + rhs = xr.DataArray([5.0, 7.0], dims="period", coords={"period": [2020, 2030]}) + m.add_constraints(e.groupby("period").sum() >= rhs) + m.add_objective((np.arange(1.0, 7.0) * x).sum()) + return m + + +@windowed(dim="time", tile=1, coupling="decomposable") +def constraint_empty_rows() -> Model: + """From ``test/test_constraint.py``: an empty constraint group solves.""" + m = Model() + x = m.add_variables( + lower=0.0, coords=[range(3), range(2)], dims=["time", "product"], name="x" + ) + m.add_constraints(x.isel(time=range(1, 1)) == 0, name="empty") + m.add_objective(x.sum()) + return m + + +@windowed(dim="dim_0", tile=5, coupling="decomposable") +def examples_benchmark() -> Model: + """From ``linopy/examples.py``: the JOSS benchmark model.""" + return benchmark_model() + + +@windowed( + dim="time", + tile=24, + coupling="replay", + fix_vars=("soc",), +) +def benchmarks_storage() -> Model: + """ + Replicated from ``benchmarks/models/storage.py`` (not importable from + the test tree): the 168h shift-coupled SoC fleet — a real, non-test toy + model with intertemporal coupling. + """ + m = Model() + storages = pd.RangeIndex(10, name="storage") + time = pd.RangeIndex(168, name="time") + soc = m.add_variables(lower=0, upper=100, coords=[storages, time], name="soc") + charge = m.add_variables(lower=0, upper=50, coords=[storages, time], name="charge") + discharge = m.add_variables( + lower=0, upper=50, coords=[storages, time], name="discharge" + ) + prev = soc.shift(time=1) + m.add_constraints( + soc - 0.99 * prev - 0.95 * charge + discharge / 0.95 == 0, name="soc_balance" + ) + m.add_objective((charge + discharge).sum()) + return m + + +# --------------------------------------------------------------------------- +# the harness +# --------------------------------------------------------------------------- + + +def _dim_index(model: Model, case: WindowCase) -> pd.Index: + """The case's windowing dimension index, from any carrier.""" + for _, var in model.variables.items(): + if case.dim in var.indexes: + return var.indexes[case.dim] + raise ValueError(f"No variable carries the dimension '{case.dim}'.") + + +def _initial_carried(model: Model, case: WindowCase) -> dict[str, xr.DataArray]: + """ + Full-horizon carry arrays for the case's fix variables, initialized to a + feasible guess (the variable's upper bound, 0 where unbounded): folded + constant rows from not-yet-solved windows must come out satisfied. + """ + carried = {} + for name in case.fix_vars: + var = model.variables[name] + labels = var.labels + upper = np.asarray(var.upper.values, dtype=float) + fill = np.where(np.isfinite(upper), upper, 0.0) + carried[name] = xr.DataArray( + fill, coords=labels.coords, dims=labels.dims, name=name + ) + return carried + + +def _solve_windows( + model: Model, + case: WindowCase, + windows: list[pd.Index], + carried: dict[str, xr.DataArray], +) -> list[dict[str, Any]]: + """Solve the windows sequentially via restrict + advance.""" + fix = {name: carried[name] for name in case.fix_vars} + w = model.restrict(**{case.dim: windows[0]}, fix=fix) + results = [] + for i, window in enumerate(windows): + if i > 0: + w.advance(**{case.dim: window}, fix=fix) + w.solve(solver_name="highs", output_flag=False) + results.append( + { + "objective": w.objective.value, + "solution": w.solution, + "dual": w.dual, + } + ) + for name in case.fix_vars: + carried[name].loc[{case.dim: window}] = w.solution[name].values + return results + + +def _assert_chunked_equals_full( + results: list[dict[str, Any]], full: Model, case: WindowCase +) -> None: + """ + Objectives sum to the monolithic one; primal/dual values match per + position for every array carrying the windowing dim. + + Arrays WITHOUT the dim (shared variables re-optimized per window, + aggregate rows enforced per window) are per-window quantities: their + equality with the monolithic solve is not implied by decomposability and + is not asserted. + """ + total = sum(r["objective"] for r in results) + assert total == pytest.approx(full.objective.value) + + for name, da in full.solution.data_vars.items(): + if case.dim not in da.dims: + continue + parts = [r["solution"][name] for r in results] + merged = xr.concat(parts, dim=da.indexes[case.dim]) + assert_allclose(merged.values, da.values, rtol=1e-7) + for name, da in full.dual.data_vars.items(): + if case.dim not in da.dims: + continue + parts = [r["dual"][name] for r in results] + merged = xr.concat(parts, dim=da.indexes[case.dim]) + assert_allclose(merged.values, da.values, rtol=1e-6) + + +def _assert_replay_equals_fresh( + case: WindowCase, + windows: list[pd.Index], + carried_init: dict[str, xr.DataArray], + results: list[dict[str, Any]], +) -> None: + """ + Replay the chain's fix sequence with fresh restricts: every window must + re-solve to the chain's own result. + """ + carried = {name: da.copy() for name, da in carried_init.items()} + for i, window in enumerate(windows): + fix = {name: carried[name] for name in case.fix_vars} + fresh = case.build().restrict(**{case.dim: window}, fix=fix) + fresh.solve(solver_name="highs", output_flag=False) + + assert fresh.objective.value == pytest.approx(results[i]["objective"]) + for name, da in fresh.solution.data_vars.items(): + assert_allclose(da.values, results[i]["solution"][name].values, rtol=1e-7) + for name, da in fresh.dual.data_vars.items(): + assert_allclose(da.values, results[i]["dual"][name].values, rtol=1e-6) + for name in case.fix_vars: + carried[name].loc[{case.dim: window}] = fresh.solution[name].values + + +@needs_highs +@pytest.mark.parametrize("case", CASES, ids=lambda c: c.name) +def test_windowed_equivalence(case: WindowCase) -> None: + """Every registered model is solved equivalently through the windows.""" + full = case.build() + full.solve(solver_name="highs", output_flag=False) + + index = _dim_index(full, case) + n_tiles = int(np.ceil(len(index) / case.tile)) + windows = [index[i * case.tile : (i + 1) * case.tile] for i in range(n_tiles)] + + model = case.build() + carried = _initial_carried(model, case) + results = _solve_windows(model, case, windows, carried) + + if case.coupling == "decomposable": + _assert_chunked_equals_full(results, full, case) + else: + _assert_replay_equals_fresh(case, windows, carried, results)