diff --git a/linopy/common.py b/linopy/common.py index 3fbe8247..e93ed04d 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -20,6 +20,7 @@ from numpy import nan, signedinteger from numpy.typing import DTypeLike from xarray import DataArray, Dataset, apply_ufunc, broadcast +from xarray import Variable as XrVariable from xarray import align as xr_align from xarray.core import indexing from xarray.namedarray.utils import is_dict_like @@ -271,6 +272,32 @@ def save_join(*dataarrays: DataArray, fill_dtype: DTypeLike | None = None) -> Da return Dataset({ds.name: ds for ds in arrs}) +def set_data_var( + ds: Dataset, name: str, dims: tuple[Hashable, ...], values: Any +) -> Dataset: + """ + Return a copy of ``ds`` with one data variable set, values-only. + + Equivalent to ``ds[name] = (dims, values)``, but writes the variable straight + into the copied dataset. The public form runs xarray's merge machinery, which + dominates repeated writes of a same-shaped variable, such as a solution or a + dual after a solve. + """ + out = ds.copy(deep=False) + data = np.asarray(values) + existing = out._variables.get(name) + if existing is None: + out._variables[name] = XrVariable(dims, data) + return out + if data.shape != existing.shape: + raise ValueError( + f"Cannot write {name!r}: shape {data.shape} does not match " + f"existing shape {existing.shape}." + ) + out._variables[name] = existing._replace(data=data) + return out + + def assign_fields(ds: Dataset, **fields: Any) -> Dataset: """ Assign data variables to a Dataset, replacing existing ones wholesale. diff --git a/linopy/constraints.py b/linopy/constraints.py index 8ba02819..c8bf1d50 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -960,6 +960,17 @@ def has_variable(self, variable: variables.Variable) -> bool: """ return self.has_labels(assigned_labels(variable.labels)) + def _assign_dual(self, values: np.ndarray) -> None: + """ + Write dual values without rebuilding the DataArray. + + Internal: values-only overwrite used by :meth:`Model.assign_result`. The + values already have the row shape and order, so the alignment and + broadcasting checks in the public setter are pure overhead. + """ + template = self._grid_template() + self._dual = DataArray(values, coords=template.coords, dims=template.dims) + def has_labels(self, labels: np.ndarray) -> bool: """ Check if the constraint references any of the given variable labels. diff --git a/linopy/model.py b/linopy/model.py index 2e390232..ba830df9 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -2475,21 +2475,22 @@ def assign_result( for _, var in self.variables.items(): start, end = var.range values = primal[start:end].reshape(var.shape) / var.solver_scaling.values - var.solution = xr.DataArray(values, var.coords, dims=var.dims) + var._assign_solution(values) if len(result.solution.dual): dual = result.solution.dual + objective_scaling = self.objective.scaling for _, con in self.constraints.items(): if con.is_indicator: continue start, end = con.range - coords = {dim: con.coords[dim] for dim in con.coord_dims} values = ( dual[start:end].reshape(con.shape) * con.scaling.values - / self.objective.scaling + / objective_scaling ) - con.dual = xr.DataArray(values, coords, dims=con.coord_dims) + assert isinstance(con, Constraint) + con._assign_dual(values) return status_value, termination_condition diff --git a/linopy/variables.py b/linopy/variables.py index a6e8b13a..e33714bd 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -49,6 +49,7 @@ iterate_slices, reassign_coords, save_join, + set_data_var, to_dataframe, to_polars, ) @@ -1150,6 +1151,16 @@ def solution(self, value: ConstantLike) -> None: value = DataArray(value).broadcast_like(self.labels) self._data = assign_fields(self.data, solution=value) + def _assign_solution(self, values: np.ndarray) -> None: + """ + Write solution values without rebuilding the dataset. + + Internal: values-only overwrite used by :meth:`Model.assign_result`. The + variable's shape and coordinates are already correct, so constructing a + ``DataArray`` and merging it costs more than writing the array. + """ + self._data = set_data_var(self._data, "solution", self.dims, values) + @has_optimized_model def get_solver_attribute(self, attr: str) -> DataArray: """ diff --git a/test/test_common.py b/test/test_common.py index 58195b50..70d3b51e 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -21,6 +21,7 @@ iterate_slices, label_coord, maybe_group_terms_polars, + set_data_var, ) @@ -52,6 +53,21 @@ def test_assign_fields_replaces_without_aligning() -> None: assert result.attrs == {"a": 1} +@pytest.mark.parametrize("name", ["solution", "dual"]) +def test_set_data_var(name: str) -> None: + ds = xr.Dataset({"solution": (("a", "b"), np.zeros((2, 3)))}) + out = set_data_var(ds, name, ("a", "b"), np.ones((2, 3))) + assert out[name].dims == ("a", "b") + assert (out[name] == 1).all() + assert (ds["solution"] == 0).all() + + +def test_set_data_var_shape_mismatch_raises() -> None: + ds = xr.Dataset({"solution": (("a", "b"), np.zeros((2, 3)))}) + with pytest.raises(ValueError, match="does not match"): + set_data_var(ds, "solution", ("a", "b"), np.ones((3, 2))) + + def test_iterate_slices_basic() -> None: ds = xr.Dataset( {"var": (("x", "y"), np.random.rand(10, 10))}, # noqa: NPY002