diff --git a/doc/api.rst b/doc/api.rst index 0656a99a..5a0aafa2 100644 --- a/doc/api.rst +++ b/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 diff --git a/doc/release_notes.rst b/doc/release_notes.rst index dae3a6dd..be4920f5 100644 --- a/doc/release_notes.rst +++ b/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* diff --git a/examples/manipulating-models.ipynb b/examples/manipulating-models.ipynb index 2a1eb460..b7f067ff 100644 --- a/examples/manipulating-models.ipynb +++ b/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": { diff --git a/linopy/common.py b/linopy/common.py index 5cee22db..ba055859 100644 --- a/linopy/common.py +++ b/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") diff --git a/linopy/constants.py b/linopy/constants.py index 7936ef1c..77ef5ff4 100644 --- a/linopy/constants.py +++ b/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" diff --git a/linopy/constraints.py b/linopy/constraints.py index f3b301cc..8aa23f6b 100644 --- a/linopy/constraints.py +++ b/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: diff --git a/linopy/csr.py b/linopy/csr.py index 424ecf71..df9d7d17 100644 --- a/linopy/csr.py +++ b/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( { diff --git a/linopy/expressions.py b/linopy/expressions.py index 65472232..e1c459eb 100644 --- a/linopy/expressions.py +++ b/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() diff --git a/linopy/model.py b/linopy/model.py index 614adf38..6b40c6f8 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -10,7 +10,7 @@ import os import re import warnings -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Hashable, Mapping, Sequence from pathlib import Path from tempfile import NamedTemporaryFile, gettempdir from types import MappingProxyType @@ -30,6 +30,8 @@ from linopy import solvers from linopy.alignment import as_dataarray, broadcast_to_coords from linopy.common import ( + _as_renamed_index, + assign_coords_multiindex_safe, assign_multiindex_safe, assigned_labels, best_int, @@ -87,7 +89,7 @@ ) from linopy.remote import RemoteHandler from linopy.scaling import validate_scaling -from linopy.semantics import enforce_no_multiindex +from linopy.semantics import enforce_no_multiindex, is_v1 try: from linopy.remote import OetcHandler @@ -2114,6 +2116,8 @@ def solve( sanitize_zeros=sanitize_zeros, sanitize_infinities=sanitize_infinities ) + self._check_coord_consistency() + # check io_api if io_api is not None and io_api not in IO_APIS: raise ValueError( @@ -2248,6 +2252,179 @@ def solve( return self.assign_result(result) + def assign_coords(self, **coords_kwargs: Any) -> Model: + """ + Reassign coordinate values across the whole model, keeping the shape. + + Mirrors :meth:`xarray.Dataset.assign_coords` semantics for an existing + model: for every named dimension, the coordinate values of every + variable, constraint (dense and CSR-backed), expression and the + parameters carrying that dimension are replaced in place. + Values-only — no relabeling, no reindexing, no shape change. The order + of each dataset's variables and coordinates is preserved. + + Containers may hold *subsets* of a dimension (e.g. a piecewise + commitment variable on a subset of generators): the new values must + match the length of the full-index carrier (the master), and every + container's labels are mapped through the master's ``old -> new`` + correspondence, preserving subset relations. + + Typical use is rolling-horizon optimization with the persistent solver + interface, where the model structure stays identical between + iterations while the window's coordinate labels advance. + + Parameters + ---------- + **coords_kwargs : Any + New coordinate values, keyed by an existing dimension name, e.g. + ``m.assign_coords(snapshot=new_snapshots)``. Accepted are + index-likes: numpy arrays, pandas Index objects, DataArrays or + lists. + + Returns + ------- + Model + ``self`` for chaining. + + Raises + ------ + ValueError + If a named dimension does not exist anywhere in the model, the + new values do not match any container's dimension length, or a + container carries labels outside the master's index. + + Examples + -------- + >>> import pandas as pd + >>> import linopy + >>> + >>> sns = pd.date_range("2026-01-01", periods=3, freq="h", name="snapshot") + >>> m = linopy.Model() + >>> x = m.add_variables(coords=[sns], name="x") + >>> _ = m.add_constraints(x >= 0, name="c") + >>> + >>> _ = m.assign_coords(snapshot=sns + pd.Timedelta("1h")) + """ + # validate everything up front, then mutate + mapped: dict[Any, dict[str, pd.Index]] = {} + for dim, values in coords_kwargs.items(): + carriers = [ + (item.name, item.indexes[dim]) + for item in self._coordinate_carriers() + if dim in item.sizes and dim in item.indexes + ] + if dim in self.parameters.sizes and dim in self.parameters.indexes: + carriers.append(("parameters", self.parameters.indexes[dim])) + if not carriers: + raise ValueError( + f"Cannot assign coordinates to dimension '{dim}': " + "not found in the model." + ) + new = _as_renamed_index(values, dim, "model") + + masters = [index for _, index in carriers if len(index) == len(new)] + if not masters: + lengths = sorted({len(index) for _, index in carriers}) + raise ValueError( + f"Cannot assign coordinates to dimension '{dim}' with " + f"length {len(new)}: no container carries it with a " + f"matching length ({lengths})." + ) + master = masters[0] + if any(not master.equals(other) for other in masters[1:]): + raise ValueError( + f"Cannot assign coordinates to dimension '{dim}': " + "containers of matching length carry different values." + ) + + for name, index in carriers: + if not index.isin(master).all(): + raise ValueError( + f"Cannot assign coordinates to dimension '{dim}': " + f"container '{name}' carries labels outside the " + "dimension's index. Relabel model-wide first, or " + "align the container with `.sel`." + ) + mapped.setdefault(dim, {})[name] = index.map( + dict(zip(master, new)) + ).rename(dim) + + for container in (self.variables, self.constraints, self.expressions): + for name, item in container.items(): + applicable = { + dim: mapped[dim][name] + for dim in mapped + if name in container.data and dim in item.sizes + } + if applicable: + item._assign_coords(**applicable) + + # note: the objective is stored as a reduced (coords-less) expression + # by design, so it never carries dimensions to reassign + + parameters_applicable = { + dim: indexes["parameters"] + for dim, indexes in mapped.items() + if "parameters" in indexes + } + if parameters_applicable: + self._parameters = assign_coords_multiindex_safe( + self.parameters, **parameters_applicable + ) + + return self + + def _coordinate_carriers(self) -> list[Any]: + """All items carrying coordinates: variables, constraints, expressions.""" + return [ + *self.variables.data.values(), + *self.constraints.data.values(), + *self.expressions.data.values(), + ] + + def _check_coord_consistency(self) -> None: + """ + Raise if containers carry incompatible labels on a shared dimension. + + v1-only guard: under v1 semantics (convention §8) shared dimensions + must carry identical labels, and models can only diverge through + internal-state corruption. Under legacy, non-aligned containers are + documented positional behavior and the check is a no-op. + + Allowed are label sets that nest by inclusion (a container may hold a + subset of the dimension, like a piecewise commitment gate); what + cannot be aligned by subset — same-length relabelings, disjoint or + partially overlapping labels — raises. Internal dims + (underscore-prefixed, like ``_term`` or ``_breakpoint_piece``) are + exempt: they are per-container bookkeeping which the piecewise + machinery labels differently on purpose. + """ + if not is_v1(): + return + indexes_by_dim: dict[Hashable, list[tuple[str, pd.Index]]] = {} + for item in self._coordinate_carriers(): + for dim, index in item.indexes.items(): + if str(dim).startswith("_"): + continue + indexes_by_dim.setdefault(dim, []).append((str(item.name), index)) + + for dim, entries in indexes_by_dim.items(): + (first_name, first_index), *rest = entries + for other_name, other_index in rest: + nested = ( + first_index.isin(other_index).all() + or other_index.isin(first_index).all() + ) + if nested: + continue + raise ValueError( + f"Coordinates for dimension '{dim}' are incompatible " + f"across the model: '{first_name}' and '{other_name}' " + "carry labels that are neither equal nor subsets of one " + "another. Use Model.assign_coords to relabel the model, " + "or align the containers with `.sel`." + ) + def assign_result( self, result: Result, diff --git a/linopy/variables.py b/linopy/variables.py index e0ad70bc..4a97057f 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -36,6 +36,7 @@ LabelPositionIndex, LocIndexer, VariableLabelIndex, + assign_coords_multiindex_safe, assign_multiindex_safe, check_has_nulls, check_has_nulls_polars, @@ -1133,6 +1134,18 @@ def _validate_update( ) return updates + def _assign_coords(self, **coords: Any) -> Variable: + """ + Reassign coordinate values on the variable, 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 solution(self) -> DataArray: diff --git a/test/test_assign_coords.py b/test/test_assign_coords.py new file mode 100644 index 00000000..1e31e21b --- /dev/null +++ b/test/test_assign_coords.py @@ -0,0 +1,396 @@ +#!/usr/bin/env python3 +""" +Test Model.assign_coords and the per-container assign_coords machinery. + +Coordinate reassignment on an existing model: values-only replacement with +unchanged shape, dataset variable order preserved, model-wide propagation. +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest +import xarray as xr + +from linopy import Model +from linopy.constants import Result, Solution, Status +from linopy.constraints import CSRConstraint + +sns0 = pd.date_range("2026-01-01", periods=3, freq="h", name="snapshot") +sns1 = sns0 + pd.Timedelta("1h") + + +@pytest.fixture +def m() -> Model: + """Model with a ``snapshot`` dim in every coordinate-carrying container.""" + model = Model() + x = model.add_variables(coords=[sns0], name="x") + y = model.add_variables( + coords=[sns0, ["a", "b"]], dims=["snapshot", "spatial"], name="y" + ) + model.add_constraints(x >= 0, name="dense_c") + model.add_constraints(x >= 1, name="frozen_c", freeze=True) + model.add_expressions(x * 2 + y.sum(), name="e") + model.add_objective(1.0 * x) + model.parameters = xr.Dataset( + {"w": (("snapshot",), [1.0, 2.0, 3.0])}, coords={"snapshot": sns0} + ) + return model + + +def test_assign_coords_reaches_all_containers(m: Model) -> None: + """Reassignment lands on every container type and returns the model.""" + result = m.assign_coords(snapshot=sns1) + + assert result is m + assert (m.variables["x"].coords["snapshot"].values == sns1.values).all() + assert (m.variables["y"].coords["snapshot"].values == sns1.values).all() + assert (m.constraints["dense_c"].coords["snapshot"].values == sns1.values).all() + assert (m.expressions["e"].coords["snapshot"].values == sns1.values).all() + assert (m.parameters.coords["snapshot"].values == sns1.values).all() + + +def test_assign_coords_csr_constraint(m: Model) -> None: + """ + CSR-backed constraints get the new labels on their grid, not only in + their reconstructed Dataset view. + """ + con = m.constraints["frozen_c"] + assert isinstance(con, CSRConstraint) + + m.assign_coords(snapshot=sns1) + + con = m.constraints["frozen_c"] + assert (con.coords["snapshot"].values == sns1.values).all() + # CSR-backed container still reconstructs its Dataset with the new coords + assert (con.data.coords["snapshot"].values == sns1.values).all() + assert con.labels.sizes["snapshot"] == 3 + + +def test_assign_coords_preserves_order(m: Model) -> None: + """ + Coordinate and data-variable order survive the reassignment. + + Plain ``Dataset.assign_coords`` moves the updated coord to the end, which + breaks dim inference downstream (e.g. in ``Model.assign_result``) — the + order-safe rebuild must not. + """ + var_before = m.variables["y"] + con_before = m.constraints["dense_c"] + coords_before = list(var_before.coords) + var_order_before = list(var_before.data.data_vars) + con_coords_before = list(con_before.coords) + con_order_before = list(con_before.data.data_vars) + + m.assign_coords(snapshot=sns1) + + assert list(m.variables["y"].coords) == coords_before + assert list(m.variables["y"].data.data_vars) == var_order_before + assert list(m.constraints["dense_c"].coords) == con_coords_before + assert list(m.constraints["dense_c"].data.data_vars) == con_order_before + + +def test_assign_coords_values_only(m: Model) -> None: + """Only index labels move: shape, dims and data variables are untouched.""" + y = m.variables["y"] + dims_before = y.dims + data_vars_before = list(y.data.data_vars) + + m.assign_coords(snapshot=sns1) + + assert y.shape == (3, 2) + assert y.dims == dims_before + # labels and bounds are untouched, only the index labels moved + assert y.labels.dims == y.dims + assert list(y.data.data_vars) == data_vars_before + + +def test_assign_result_after_reassign(m: Model) -> None: + """ + Solution and dual values land on the *new* labels after reassignment. + + This is the rolling-horizon flow from the issue: reassign the window, + then write a solver result back — solution/dual arrays must carry the + updated coordinates, not the stale ones. + """ + n_labels = 1 + max(int(v.labels.max()) for _, v in m.variables.items()) + n_cons = 1 + max(int(c.labels.max()) for _, c in m.constraints.items()) + result = Result( + status=Status.process("ok", "optimal"), + solution=Solution( + primal=np.arange(n_labels, dtype=float), + dual=np.arange(n_cons, dtype=float), + objective=42.0, + ), + ) + + m.assign_coords(snapshot=sns1) + m.assign_result(result) + + solution = m.variables["x"].solution + assert (solution.coords["snapshot"].values == sns1.values).all() + assert solution.dims == ("snapshot",) + + dual = m.constraints["dense_c"].dual + assert (dual.coords["snapshot"].values == sns1.values).all() + assert dual.dims == ("snapshot",) + + +def test_mock_solve_after_reassign(m: Model) -> None: + """Mock solving still infers dims correctly from the reordered-safe coords.""" + m.assign_coords(snapshot=sns1) + m._mock_solve() + assert (m.variables["x"].solution.coords["snapshot"].values == sns1.values).all() + + +def test_assign_coords_container_methods_return_self(m: Model) -> None: + """The private per-container methods mutate in place and chain.""" + assert m.variables["x"]._assign_coords(snapshot=sns1) is m.variables["x"] + assert ( + m.constraints["dense_c"]._assign_coords(snapshot=sns1) + is (m.constraints["dense_c"]) + ) + assert m.expressions["e"]._assign_coords(snapshot=sns1) is m.expressions["e"] + + +def test_assign_coords_rejects_wrong_length(m: Model) -> None: + """ + New values shorter than the dimension raise, at both API levels. + + Validates the issue's "same length" contract and that a failed call + leaves the model untouched. + """ + short = pd.date_range("2026-01-01", periods=2, freq="h", name="snapshot") + with pytest.raises(ValueError, match="length"): + m.assign_coords(snapshot=short) + + with pytest.raises(ValueError, match="length"): + m.variables["x"]._assign_coords(snapshot=short) + + # nothing was mutated + assert (m.variables["x"].coords["snapshot"].values == sns0.values).all() + + +def test_assign_coords_rejects_unknown_dim(m: Model) -> None: + """ + Unknown dimensions raise instead of being added as new coordinates. + + The model-level message reports the dimension as not found in the model; + the per-container message reports the missing coordinate. + """ + with pytest.raises(ValueError, match="not found"): + m.assign_coords(nonexistent=[1, 2, 3]) + + with pytest.raises(ValueError, match="missing"): + m.variables["x"]._assign_coords(nonexistent=[1, 2, 3]) + + +def test_assign_coords_accepts_index_like_values(m: Model) -> None: + """ + Lists, pandas Index objects and DataArrays are accepted as values. + + A DataArray contributes its values, like ``xarray.assign_coords``; a + passed Index that is named differently is renamed to the target + dimension (the keyword key is authoritative). + """ + m.assign_coords(snapshot=list(range(3))) + assert (m.variables["x"].coords["snapshot"].values == np.arange(3)).all() + + m.assign_coords(snapshot=pd.Index([9, 8, 7], name="snapshot")) + assert (m.variables["x"].coords["snapshot"].values == [9, 8, 7]).all() + + # a DataArray contributes its values, like xarray.assign_coords + da = xr.DataArray([1, 2, 3], coords={"snapshot": sns1}) + m.assign_coords(snapshot=da) + assert (m.variables["x"].coords["snapshot"].values == [1, 2, 3]).all() + + +def test_assign_coords_multiple_dims(m: Model) -> None: + """Several dimensions can be reassigned in one call.""" + m.assign_coords(snapshot=sns1, spatial=["c", "d"]) + assert (m.variables["y"].coords["snapshot"].values == sns1.values).all() + assert (m.variables["y"].coords["spatial"].values == ["c", "d"]).all() + + +def test_model_assign_coords_skips_carriers_without_dim(m: Model) -> None: + """Containers not carrying the dimension are left untouched.""" + m.add_variables(coords=[pd.RangeIndex(2, name="other")], name="scalar_var") + + m.assign_coords(snapshot=sns1) + + assert (m.variables["scalar_var"].coords["other"].values == [0, 1]).all() + assert (m.variables["x"].coords["snapshot"].values == sns1.values).all() + + +def test_assign_coords_rejects_scalar_values(m: Model) -> None: + """Scalars are not index-like and raise instead of broadcasting.""" + with pytest.raises(ValueError, match="index-like"): + m.assign_coords(snapshot=5) + + with pytest.raises(ValueError, match="index-like"): + m.variables["x"]._assign_coords(snapshot=5) + + +def test_csr_assign_coords_rejects_wrong_length(m: Model) -> None: + """The CSR path validates lengths and rejects unknown coordinates too.""" + short = pd.date_range("2026-01-01", periods=2, freq="h", name="snapshot") + con = m.constraints["frozen_c"] + assert isinstance(con, CSRConstraint) + + with pytest.raises(ValueError, match="length"): + con._assign_coords(snapshot=short) + + with pytest.raises(ValueError, match="missing"): + con._assign_coords(nonexistent=[1, 2, 3]) + + +def test_assign_coords_reaches_quadratic_expression(m: Model) -> None: + """Stored quadratic expressions get the new labels like linear ones.""" + x = m.variables["x"] + m.add_expressions(x * x, name="q") + + m.assign_coords(snapshot=sns1) + + q = m.expressions["q"] + assert (q.coords["snapshot"].values == sns1.values).all() + + +def test_assign_coords_skips_uncoordinated_parameters(m: Model) -> None: + """A parameters dim without an index coord has no labels to move: skip it.""" + m.parameters = xr.Dataset({"w": ("snapshot", [1.0, 2.0, 3.0])}) + + m.assign_coords(snapshot=sns1) + + assert (m.variables["x"].coords["snapshot"].values == sns1.values).all() + + +def test_assign_coords_renames_mismatched_index(m: Model) -> None: + """A passed Index named differently is renamed to the target dimension.""" + m.assign_coords(snapshot=pd.Index([9, 8, 7], name="timestep")) + coord = m.variables["x"].coords["snapshot"] + assert coord.name == "snapshot" + assert (coord.values == [9, 8, 7]).all() + + +def test_assign_coords_maps_subset_carriers(m: Model) -> None: + """ + A container holding a subset of the dimension is mapped by label. + + Case 1 from the design: a piecewise commitment gate ``u`` on a subset of + ``gen`` must follow the master's relabeling (``a -> g1``, ``c -> g3``), + preserving the subset relation — not be relabeled positionally. + """ + from linopy.variables import Variable + + full = pd.Index(["a", "b", "c"], name="gen") + subset = pd.Index(["a", "c"], name="gen") + m2 = Model() + x2 = m2.add_variables(coords=[full], name="x") + m2.add_variables(binary=True, coords=[subset], name="u") + m2.add_constraints(x2 >= 0, name="c") + m2.add_constraints(x2.sel(gen=["a", "c"]) >= 1, name="c_subset") + + m2.assign_coords(gen=["g1", "g2", "g3"]) + + assert (m2.variables["x"].coords["gen"].values == ["g1", "g2", "g3"]).all() + assert (m2.variables["u"].coords["gen"].values == ["g1", "g3"]).all() + assert (m2.constraints["c_subset"].coords["gen"].values == ["g1", "g3"]).all() + + # a container with labels outside the master's index is corruption: raise + odd_data = x2.data.isel(gen=slice(1, 2)).assign_coords(gen=["z"]) + m2.variables.data["odd"] = Variable(odd_data, m2, "odd") + with pytest.raises(ValueError, match="outside the"): + m2.assign_coords(gen=["h1", "h2", "h3"]) + + +def test_assign_coords_rejects_unmatched_master_length(m: Model) -> None: + """New values matching no carrier's length raise, listing the lengths.""" + with pytest.raises(ValueError, match="no container carries it"): + m.assign_coords(snapshot=pd.date_range("2026-01-01", periods=7, freq="h")) + + +def test_assign_coords_rejects_ambiguous_masters(m: Model) -> None: + """Same-length containers with different values are ambiguous: raise.""" + m.variables["x"]._assign_coords(snapshot=sns0 + pd.Timedelta("5h")) + + with pytest.raises(ValueError, match="matching length carry different"): + m.assign_coords(snapshot=sns1) + + +@pytest.mark.v1 +def test_solve_rejects_diverged_coords(m: Model) -> None: + """ + v1 solving raises when containers carry incompatible labels on a dim. + + Same length, different values — the mislabeling trap no construction-time + §8 check catches (and which the public API cannot build); the pre-solve + guard is the backstop. A re-aligned model passes again. + """ + shifted = sns0 + pd.Timedelta("5h") + m.variables["x"]._assign_coords(snapshot=shifted) + + with pytest.raises(ValueError, match="incompatible"): + m._check_coord_consistency() + + # a consistent model passes + m.variables["x"]._assign_coords(snapshot=sns0) + m._check_coord_consistency() + + +@pytest.mark.legacy +def test_check_coord_consistency_legacy_noop(m: Model) -> None: + """ + Under legacy, non-aligned containers are documented positional + behavior (convention §8), so the guard is a no-op. + """ + shifted = sns0 + pd.Timedelta("5h") + m.variables["x"]._assign_coords(snapshot=shifted) + + m._check_coord_consistency() + + +@pytest.mark.v1 +def test_check_coord_consistency_allows_helper_dim_divergence(m: Model) -> None: + """ + Internal (underscore-prefixed) dims are exempt from the guard. + + The piecewise machinery labels ``_breakpoint_piece`` differently on each + container on purpose — that must not trip the guard. + """ + delta = m.add_variables( + binary=True, + coords=[sns0, [0, 1, 2]], + dims=["snapshot", "_breakpoint_piece"], + name="delta", + ) + delta_hi = delta.isel(_breakpoint_piece=slice(1, None), drop=True) + delta_hi._assign_coords(_breakpoint_piece=[0, 1]) + + m._check_coord_consistency() + + +@pytest.mark.v1 +def test_check_coord_consistency_allows_subset_labels(m: Model) -> None: + """Containers nesting by inclusion (a gate on a gen subset) pass the guard.""" + m.add_variables( + binary=True, + coords=[pd.Index(["a", "c"], name="gen")], + name="u", + ) + m.add_variables( + coords=[pd.Index(["a", "b", "c"], name="gen")], + name="x_gen", + ) + + m._check_coord_consistency() + + +@pytest.mark.v1 +def test_check_coord_consistency_rejects_partial_overlap(m: Model) -> None: + """Labels neither equal nor nested (partial overlap) raise under v1.""" + m.add_variables(coords=[pd.Index(["a", "c"], name="gen")], name="u") + m.add_variables(coords=[pd.Index(["b", "c"], name="gen")], name="v") + + with pytest.raises(ValueError, match="incompatible"): + m._check_coord_consistency()