perf(model): write solve results instead of rebuilding DataArrays

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.

linopy/constraints.py

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

linopy/model.py

@@ -2343,21 +2343,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
 

linopy/variables.py

@@ -48,6 +48,7 @@
     iterate_slices,
     reassign_coords,
     save_join,
+    set_data_var,
     to_dataframe,
     to_polars,
 )
@@ -1097,6 +1098,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:
         """

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