feat(restrict): window a model's flattened matrices

doc/api.rst

@@ -74,6 +74,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

doc/index.rst

@@ -134,6 +134,7 @@ This package is published under MIT license.
 
    sos-constraints
    piecewise-linear-constraints
+   rolling-horizon
    numerical-scaling
    sparse-models
    testing-framework

doc/release_notes.rst

@@ -90,6 +90,7 @@ Version 0.10.1
 * ``read_netcdf`` no longer drops parameters, parameter attributes or ``MultiIndex`` dimensions whose name contains ``-``. (`#1025 <https://github.com/PyPSA/linopy/pull/1025>`__)
 * ``to_netcdf``/``read_netcdf`` keep the auxiliary (non-index) coordinates of dense variables, constraints, expressions and ``Model.parameters``, such as the levels on a flat dimension that v1 recommends instead of a ``MultiIndex``, and scalar coordinates left by a scalar selection. (`#1024 <https://github.com/PyPSA/linopy/pull/1024>`__)
 * ``licensed_solvers`` and ``check_solver_licenses()`` no longer report Xpress as licensed when its licence is rejected, e.g. a licence that only covers an older Xpress release. Xpress now probes its licence with ``xpress.init()``. (`#994 <https://github.com/PyPSA/linopy/issues/994>`__)
+* :meth:`linopy.Model.restrict` turns rolling-horizon optimization into a pure linopy pattern. ``m.restrict(snapshot=window, fix=...)`` 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 slice and a column range of the parent's ``A``, ``b``, ``c``, ``lb`` and ``ub``, a step 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 a shape, 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. Solutions and duals are built per container only when read. See the new :doc:`rolling-horizon` guide. (https://github.com/PyPSA/linopy/issues/768)
 
 Version 0.10.0
 --------------

doc/rolling-horizon.rst

@@ -0,0 +1,127 @@
+.. _rolling-horizon:
+
+Rolling-Horizon Optimization
+============================
+
+A long-horizon problem is often solved one window at a time: solve a slice of
+the horizon (say, one day), carry the end state into the next window, then move
+the window forward. Linopy has three methods for this:
+
+* :meth:`linopy.Model.restrict` returns a window of the model,
+* :meth:`linopy.restrict.WindowModel.solve` solves it,
+* :meth:`linopy.restrict.WindowModel.advance` moves it forward.
+
+A window is a row and column slice of the model's matrices. Moving it does not
+rebuild variables or constraints.
+
+.. contents::
+   :local:
+   :depth: 2
+
+The rolling-horizon recipe
+--------------------------
+
+.. 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
+
+:meth:`~linopy.Model.restrict` follows :meth:`xarray.Dataset.sel` semantics:
+every container holding the selected dimension is sliced, containers without it
+stay whole. The parent model is unchanged.
+
+Boundary values
+---------------
+
+The balance constraint references the previous snapshot's state of charge. In
+the first row of a window, that reference points outside the window.
+: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``.
+
+.. 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.
+
+``fix`` values are keyed by variable name and broadcast over the variable's
+coordinates, so a full-horizon DataArray works directly. Values for in-window
+variables are ignored. A ``NaN`` at an outside label counts as missing.
+
+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 the new
+selection and moves the new boundary values to the right-hand side. When the
+new window has the same shape, the solver updates its existing model instead of
+building a new one. Nothing is rebuilt, so a chain of equal-width windows
+keeps that fast path.
+
+Limits
+------
+
+A window slices one dimension. Every container holding that dimension must
+contain the window's labels; otherwise ``restrict`` raises.
+
+Several selected dimensions, indicators and quadratic objectives raise
+:class:`~linopy.restrict.RestrictionError` or ``NotImplementedError``.
+
+Native SOS constraints also raise: they are not part of the matrices. Reformulate
+them first with :meth:`Model.apply_sos_reformulation`, though the cardinality row
+that adds spans the horizon and still cannot be windowed.
+
+Piecewise formulations work; they expand into ordinary variables and constraints
+when added.

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:

linopy/matrices.py

@@ -38,10 +38,16 @@ class _RowBlock(NamedTuple):
     row_scaling: ndarray | None
 
 
-def _row_block(con: Constraint, label_index: VariableLabelIndex) -> _RowBlock:
-    active = con.active_row_mask()
+def _row_scaling(con: Constraint, active: ndarray) -> ndarray:
+    """Scaling of the active rows of a constraint, in row order."""
+    return con.scaling.values.ravel()[active]
+
+
+def _row_block(
+    con: Constraint, label_index: VariableLabelIndex, active: ndarray
+) -> _RowBlock:
     store = con.lhs._store
-    scaling = con.scaling.values.ravel()[active]
+    scaling = _row_scaling(con, active)
     return _RowBlock(
         partial(store.positional_csr, label_index, active, con.name),
         int(active.sum()),
@@ -186,17 +192,24 @@ def _build_cons(self) -> None:
 
         regular, indicator = [], []
         binvar, binval = [], []
+        scalings = []
         for c in m.constraints.data.values():
+            active = c.active_row_mask()
             if not c.is_indicator:
-                regular.append(_row_block(c, label_index))
+                regular.append(_row_block(c, label_index, active))
+                scalings.append(_row_scaling(c, active))
                 continue
-            active = c.active_row_mask()
-            indicator.append(_row_block(c, label_index))
+            indicator.append(_row_block(c, label_index, active))
             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
+        #: Row scaling actually applied to ``A``, aligned with ``clabels``. Consumers
+        #: that report duals need it to invert the scaling on the way back.
+        self.con_scaling: ndarray = (
+            np.concatenate(scalings) if scalings else np.empty(0, dtype=float)
+        )
         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

linopy/model.py

@@ -117,6 +117,7 @@
 
 if TYPE_CHECKING:
     from linopy.piecewise import PiecewiseFormulation
+    from linopy.restrict.window import WindowModel
 
 logger = logging.getLogger(__name__)
 
@@ -2295,6 +2296,69 @@ 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 along one coordinate dimension.
+
+        A window is a slice of the model's flattened matrices: the parent is
+        built once, and each selection slices its rows and columns. 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 and
+            broadcast against the variable's coordinates. Values for in-window
+            variables are ignored.
+        **selection : Any
+            Exactly one dimension and the labels of the window, e.g.
+            ``snapshot=sns[:24]``.
+
+        Returns
+        -------
+        WindowModel
+            The window, with its matrices computed.
+
+        Raises
+        ------
+        RestrictionError
+            If the dimension does not exist, a carrier does not contain the
+            window's labels, or a fully folded row is unsatisfiable.
+        NotImplementedError
+            If the selection names several dimensions, or the model has
+            indicators, a quadratic objective, SOS or piecewise formulations.
+
+        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,