feat(model): Model.assign_coords — reassign coordinates on an existing model

doc/api.rst

@@ -72,6 +72,7 @@ Modifying a model
    model.Model.remove_expressions
    model.Model.remove_objective
    model.Model.remove_sos_constraints
+   model.Model.assign_coords
    model.Model.copy
    model.Model.apply_sos_reformulation
    model.Model.undo_sos_reformulation

doc/release_notes.rst

@@ -31,6 +31,8 @@ Upcoming Version
 
 *Other*
 
+* New method :meth:`linopy.Model.assign_coords` reassigns coordinate values across an existing model — variables, constraints (dense and CSR-backed), expressions and parameters — without changing the model's shape: ``m.assign_coords(snapshot=new_snapshots)``. Values-only: the new values must match the length of the dimension's full-index container, and containers holding subsets of the dimension are mapped by label, preserving the subset relation. Dataset variable order is preserved. Under v1 semantics, ``Model.solve()`` raises when containers carry labels on a shared dimension that are neither equal nor subsets of one another. Typical use is advancing the window in rolling-horizon optimization with the persistent solver interface. (https://github.com/PyPSA/linopy/issues/767)
+
 * ``add_piecewise_formulation`` gained a ``mask`` parameter declaring which breakpoint slots hold a real breakpoint. It is needed for **ragged** curves — entities with different numbers of breakpoints — which are stored densely with the surplus slots left absent. Under v1 that absence must be declared (``mask=x_pts.notnull()``) rather than read off the NaN padding. (https://github.com/PyPSA/linopy/issues/884)
 
 *Internal*

examples/manipulating-models.ipynb

@@ -390,6 +390,51 @@
     "# z is binary again\n",
     "m.variables[\"z\"].attrs[\"binary\"]"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "32",
+   "metadata": {},
+   "source": [
+    "## Reassigning coordinates\n",
+    "\n",
+    "Some workflows keep a model's *structure* fixed while the coordinate labels move —\n",
+    "the typical case is rolling-horizon optimization, where each iteration solves the\n",
+    "same model over the next time window. For that, `Model.assign_coords` replaces the\n",
+    "*values* of an existing dimension's coordinates on every variable, constraint\n",
+    "(dense and CSR-backed), expression and parameter carrying it — without relabeling,\n",
+    "reindexing or changing the shape:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "33",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "shifted_time = time + 1\n",
+    "\n",
+    "m.assign_coords(time=shifted_time)\n",
+    "m.variables[\"x\"].coords[\"time\"]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "34",
+   "metadata": {},
+   "source": [
+    "Unlike the xarray method it resembles, this operates on the *whole* model:\n",
+    "the container-level ``assign_coords`` methods (e.g. ``x.assign_coords(time=...)``)\n",
+    "are plain xarray passthroughs that return a detached copy and leave the model\n",
+    "untouched — always use ``m.assign_coords(...)`` to reassign coordinates on the\n",
+    "model itself.\n",
+    "\n",
+    ".. note::\n",
+    "   Coordinate reassignment is values-only: the new values must match the length\n",
+    "   of the existing dimension. Typical use is rolling-horizon optimization, where\n",
+    "   each iteration re-solves the same model over the next time window."
+   ]
   }
  ],
  "metadata": {

linopy/common.py

@@ -8,7 +8,7 @@
 from __future__ import annotations
 
 import operator
-from collections.abc import Callable, Generator, Hashable, Iterable, Sequence
+from collections.abc import Callable, Generator, Hashable, Iterable, Mapping, Sequence
 from functools import cached_property, reduce, wraps
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Generic, TypeVar, overload
@@ -25,6 +25,7 @@
 from xarray.core import indexing
 from xarray.namedarray.utils import is_dict_like
 
+from linopy.alignment import _as_index
 from linopy.config import options
 from linopy.constants import (
     SIGNS,
@@ -337,6 +338,112 @@ def assign_multiindex_safe(ds: Dataset, **fields: Any) -> Dataset:
     return Dataset({**ds[remainders], **fields}, attrs=ds.attrs)
 
 
+def _as_renamed_index(values: Any, name: Hashable, context: str) -> pd.Index:
+    """
+    Convert coordinate-like values to a pandas Index named ``name``.
+
+    The dimension name is authoritative: a passed Index or DataArray with a
+    different (or missing) name is renamed to ``name``, like
+    :meth:`xarray.Dataset.assign_coords` does for keyword-assigned coords.
+    """
+    try:
+        index = _as_index(values)
+        if index.name != name:
+            index = index.rename(name)
+    except (TypeError, ValueError) as e:
+        raise ValueError(
+            f"New coordinates for dimension '{name}' on {context} must be "
+            f"index-like, got {type(values).__name__}: {values!r}."
+        ) from e
+    return index
+
+
+def validate_coords_reassignment(
+    sizes: Mapping[str, int], coords: Mapping[str, Any], context: str
+) -> dict[str, pd.Index]:
+    """
+    Validate values-only coordinate reassignment and convert the new values.
+
+    Single-sourced per-item invariants for coordinate reassignment, used by
+    the Dataset-backed (:func:`assign_coords_multiindex_safe`) and the
+    CSR-backed (``CSRConstraint``) paths alike, so every entry point is safe
+    on its own. Model-level aggregate checks (dimension exists somewhere in
+    the model, consistent lengths across containers) live on top of this.
+
+    The keyword key is authoritative: values converted to an Index with a
+    different (or missing) name are renamed to the target dimension, like
+    :meth:`xarray.Dataset.assign_coords` does.
+
+    Parameters
+    ----------
+    sizes : Mapping
+        Existing dimension name to length, defining which coordinates may be
+        reassigned.
+    coords : Mapping
+        New coordinate values, keyed by existing dimension name.
+    context : str
+        Name of the object being reassigned, used in error messages.
+
+    Returns
+    -------
+    dict
+        New coordinate values converted to named pandas Index objects, ready
+        for assignment.
+
+    Raises
+    ------
+    ValueError
+        If a named coordinate does not exist in ``sizes``, the new values are
+        not index-like, or their length differs from the existing dimension.
+    """
+    missing = [name for name in coords if name not in sizes]
+    if missing:
+        raise ValueError(f"Cannot assign missing coordinates {missing} to {context}.")
+    new_indexes: dict[str, pd.Index] = {}
+    for name, values in coords.items():
+        index = _as_renamed_index(values, name, context)
+        if len(index) != sizes[name]:
+            raise ValueError(
+                f"Cannot assign coordinates to dimension '{name}' on {context} "
+                f"with a different length: expected {sizes[name]}, got "
+                f"{len(index)}."
+            )
+        new_indexes[name] = index
+    return new_indexes
+
+
+def assign_coords_multiindex_safe(ds: Dataset, **coords: Any) -> Dataset:
+    """
+    Reassign coordinate values on an existing Dataset, keeping the shape.
+
+    Values-only replacement of existing dimension coordinates: each new value
+    must match the length of the dimension it replaces. Neither the order of
+    the dataset's variables nor the order of its coordinates is altered —
+    plain :meth:`xarray.Dataset.assign_coords` moves every reassigned
+    coordinate to the end, which breaks downstream dimension inference.
+
+    Parameters
+    ----------
+    ds : Dataset
+        Dataset to reassign the coordinates on.
+    **coords : Any
+        New coordinate values, keyed by existing dimension name. Accepted
+        like in :meth:`xarray.Dataset.assign_coords`: index-likes such as
+        numpy arrays, pandas Index objects, DataArrays or lists.
+
+    Returns
+    -------
+    Dataset
+        Dataset with reassigned coordinate values.
+    """
+    sizes = {str(name): coord.size for name, coord in ds.coords.items()}
+    new_indexes = validate_coords_reassignment(sizes, coords, "dataset")
+    new = ds.assign_coords(new_indexes)
+    ordered = {name: new[name] for name in ds.coords}
+    data_vars = {name: new[name].variable for name in new.data_vars}
+    return Dataset(data_vars, coords=ordered, attrs=new.attrs)
+
+
 T = TypeVar("T", Dataset, "Variable", "LinearExpression", "ConstraintBase")
 
 

linopy/constants.py

@@ -71,6 +71,8 @@ class PerformanceWarning(UserWarning):
 
 PWL_CONVEXITIES: frozenset[str] = frozenset(get_args(PWL_CONVEXITY))
 """Set of valid :data:`~linopy.constants.PWL_CONVEXITY` values."""
+# All internal dims are underscore-prefixed — user-facing dims never start
+# with "_". ``Model._check_coord_consistency`` relies on this to exempt them.
 BREAKPOINT_DIM = "_breakpoint"
 SEGMENT_DIM = "_segment"
 LP_PIECE_DIM = f"{BREAKPOINT_DIM}_piece"

linopy/constraints.py

@@ -39,6 +39,7 @@
     LocIndexer,
     VariableLabelIndex,
     align_lines_by_delimiter,
+    assign_coords_multiindex_safe,
     assign_multiindex_safe,
     assigned_labels,
     check_has_nulls,
@@ -65,6 +66,7 @@
     save_join,
     to_dataframe,
     to_polars,
+    validate_coords_reassignment,
 )
 from linopy.config import options
 from linopy.constants import (
@@ -202,6 +204,17 @@ def dual(self) -> DataArray:
     def dual(self, value: DataArray) -> None:
         """Set the dual values DataArray."""
 
+    @abstractmethod
+    def _assign_coords(self, **coords: Any) -> ConstraintBase:
+        """
+        Reassign coordinate values on the constraint, keeping the shape.
+
+        Internal: values-only replacement of existing dimension coordinates,
+        used by :meth:`linopy.Model.assign_coords`. No relabeling, no
+        reindexing, no shape change, and the order of the underlying data is
+        preserved.
+        """
+
     @property
     @abstractmethod
     def is_indicator(self) -> bool:
@@ -787,6 +800,23 @@ def assign_labels(
             changes["scaling"] = scaling[positions]
         return self._replace(**changes)
 
+    def _assign_coords(self, **coords: Any) -> CSRConstraint:
+        """
+        Reassign coordinate values on the constraint, keeping the shape.
+
+        Internal: values-only replacement of existing dimension coordinates,
+        used by :meth:`linopy.Model.assign_coords`. No relabeling, no
+        reindexing, no shape change, and the order of the underlying data is
+        preserved.
+        """
+        new_indexes = validate_coords_reassignment(
+            {dim: len(index) for dim, index in self._grid.indexes.items()},
+            coords,
+            f"constraint '{self.name}'",
+        )
+        self._grid = self._grid.with_indexes(new_indexes)
+        return self
+
     def _active_to_dataarray(
         self, active_values: np.ndarray, fill: float | int | str = -1
     ) -> DataArray:
@@ -1814,6 +1844,18 @@ def update(
 
         return self
 
+    def _assign_coords(self, **coords: Any) -> Constraint:
+        """
+        Reassign coordinate values on the constraint, keeping the shape.
+
+        Internal: values-only replacement of existing dimension coordinates,
+        used by :meth:`linopy.Model.assign_coords`. No relabeling, no
+        reindexing, no shape change, and the order of the underlying data is
+        preserved.
+        """
+        self._data = assign_coords_multiindex_safe(self.data, **coords)
+        return self
+
     @property
     @has_optimized_model
     def dual(self) -> DataArray:

linopy/csr.py

@@ -22,7 +22,7 @@
 
 from __future__ import annotations
 
-from collections.abc import Hashable, Iterable, Mapping
+from collections.abc import Iterable, Mapping
 from dataclasses import dataclass, field, replace
 from typing import TYPE_CHECKING, Any
 
@@ -119,7 +119,7 @@ def reordered(self, dims: Iterable[str]) -> Grid:
         """Select and order the given dimensions; labels unchanged."""
         return Grid({d: self.indexes[d] for d in dims})
 
-    def with_indexes(self, indexers: Mapping[Hashable, Any]) -> Grid:
+    def with_indexes(self, indexers: Mapping[Any, Any]) -> Grid:
         """Replace the labels of the named dimensions; the rest unchanged."""
         return Grid(
             {

linopy/expressions.py

@@ -70,6 +70,7 @@
 from linopy.common import (
     EmptyDeprecationWrapper,
     LocIndexer,
+    assign_coords_multiindex_safe,
     assign_multiindex_safe,
     check_common_keys_values,
     check_has_nulls,
@@ -1639,6 +1640,18 @@ def const(self) -> DataArray:
     def const(self, value: DataArray) -> None:
         self._data = assign_multiindex_safe(self.data, const=value)
 
+    def _assign_coords(self, **coords: Any) -> Self:
+        """
+        Reassign coordinate values on the expression, keeping the shape.
+
+        Internal: values-only replacement of existing dimension coordinates,
+        used by :meth:`linopy.Model.assign_coords`. No relabeling, no
+        reindexing, no shape change, and the order of the underlying data is
+        preserved.
+        """
+        self._data = assign_coords_multiindex_safe(self.data, **coords)
+        return self
+
     @property
     def has_constant(self) -> DataArray:
         return self.const.any()