fix: use multiindex safe assignment (#456)

Author: FabianHofmann

.gitignore

@@ -16,6 +16,8 @@ doc/api
 Highs.log
 paper/
 monkeytype.sqlite3
+.github/copilot-instructions.md
+uv.lock
 
 # Environments
 .env

doc/release_notes.rst

@@ -12,6 +12,9 @@ Upcoming Version
   gap tolerance.
 * Improve the mapping of termination conditions for the SCIP solver
 * Treat GLPK's `integer undefined` status as not-OK
+* Internally assign new data fields to `Variable` and `Constraint` with a multiindexed-safe routine. Before the
+  assignment when using multi-indexed coordinates, an deprecation warning was raised. This is fixed now.
+
 
 Version 0.5.3
 --------------

linopy/constraints.py

@@ -7,13 +7,12 @@
 from __future__ import annotations
 
 import functools
-from collections.abc import Hashable, ItemsView, Iterator, Sequence
+from collections.abc import Callable, Hashable, ItemsView, Iterator, Sequence
 from dataclasses import dataclass
 from itertools import product
 from typing import (
     TYPE_CHECKING,
     Any,
-    Callable,
     overload,
 )
 
@@ -460,7 +459,7 @@ def sign(self) -> DataArray:
     @is_constant
     def sign(self, value: SignLike) -> None:
         value = maybe_replace_signs(DataArray(value)).broadcast_like(self.sign)
-        self.data["sign"] = value
+        self._data = assign_multiindex_safe(self.data, sign=value)
 
     @property
     def rhs(self) -> DataArray:
@@ -478,7 +477,7 @@ def rhs(self, value: ExpressionLike) -> None:
             value, self.model, coords=self.coords, dims=self.coord_dims
         )
         self.lhs = self.lhs - value.reset_const()
-        self.data["rhs"] = value.const
+        self._data = assign_multiindex_safe(self.data, rhs=value.const)
 
     @property
     @has_optimized_model
@@ -1006,7 +1005,7 @@ def set_blocks(self, block_map: np.ndarray) -> None:
 
             res = res.where(not_missing.any(constraint.term_dim), -1)
             res = res.where(not_zero.any(constraint.term_dim), 0)
-            constraint.data["blocks"] = res
+            constraint._data = assign_multiindex_safe(constraint.data, blocks=res)
 
     @property
     def flat(self) -> pd.DataFrame:

linopy/expressions.py

@@ -313,7 +313,7 @@ def sum(self, **kwargs: Any) -> LinearExpression:
             .rename({TERM_DIM: STACKED_TERM_DIM})
             .stack({TERM_DIM: [STACKED_TERM_DIM, "_rolling_term"]}, create_index=False)
         )
-        ds["const"] = data.const.sum("_rolling_term")
+        ds = assign_multiindex_safe(ds, const=data.const.sum("_rolling_term"))
         return LinearExpression(ds, self.model)
 
 
@@ -783,23 +783,23 @@ def vars(self) -> DataArray:
 
     @vars.setter
     def vars(self, value: DataArray) -> None:
-        self.data["vars"] = value
+        self._data = assign_multiindex_safe(self.data, vars=value)
 
     @property
     def coeffs(self) -> DataArray:
         return self.data.coeffs
 
     @coeffs.setter
     def coeffs(self, value: DataArray) -> None:
-        self.data["coeffs"] = value
+        self._data = assign_multiindex_safe(self.data, coeffs=value)
 
     @property
     def const(self) -> DataArray:
         return self.data.const
 
     @const.setter
     def const(self, value: DataArray) -> None:
-        self.data["const"] = value
+        self._data = assign_multiindex_safe(self.data, const=value)
 
     @property
     def has_constant(self) -> DataArray:
@@ -1458,7 +1458,7 @@ def to_polars(self) -> pl.DataFrame:
         return df
 
     # Wrapped function which would convert variable to dataarray
-    assign = exprwrap(Dataset.assign)
+    assign = exprwrap(assign_multiindex_safe)
 
     assign_multiindex_safe = exprwrap(assign_multiindex_safe)
 

linopy/model.py

@@ -27,6 +27,7 @@
 from linopy import solvers
 from linopy.common import (
     as_dataarray,
+    assign_multiindex_safe,
     best_int,
     maybe_replace_signs,
     replace_by_map,
@@ -734,7 +735,9 @@ def remove_variables(self, name: str) -> None:
         for k in list(self.constraints):
             vars = self.constraints[k].data["vars"]
             vars = vars.where(~vars.isin(labels), -1)
-            self.constraints[k].data["vars"] = vars
+            self.constraints[k]._data = assign_multiindex_safe(
+                self.constraints[k].data, vars=vars
+            )
 
         self.objective = self.objective.sel(
             {TERM_DIM: ~self.objective.vars.isin(labels)}

linopy/variables.py

@@ -8,12 +8,11 @@
 
 import functools
 import logging
-from collections.abc import Hashable, ItemsView, Iterator, Mapping
+from collections.abc import Callable, Hashable, ItemsView, Iterator, Mapping
 from dataclasses import dataclass
 from typing import (
     TYPE_CHECKING,
     Any,
-    Callable,
     overload,
 )
 from warnings import warn
@@ -755,7 +754,7 @@ def upper(self, value: ConstantLike) -> None:
         value = DataArray(value).broadcast_like(self.upper)
         if not set(value.dims).issubset(self.model.variables[self.name].dims):
             raise ValueError("Cannot assign new dimensions to existing variable.")
-        self.data["upper"] = value
+        self._data = assign_multiindex_safe(self.data, upper=value)
 
     @property
     def lower(self) -> DataArray:
@@ -779,7 +778,7 @@ def lower(self, value: ConstantLike) -> None:
         value = DataArray(value).broadcast_like(self.lower)
         if not set(value.dims).issubset(self.model.variables[self.name].dims):
             raise ValueError("Cannot assign new dimensions to existing variable.")
-        self.data["lower"] = value
+        self._data = assign_multiindex_safe(self.data, lower=value)
 
     @property
     @has_optimized_model
@@ -1059,8 +1058,7 @@ def bfill(self, dim: str, limit: None = None) -> Variable:
             .map(DataArray.bfill, dim=dim, limit=limit)
             .fillna(self._fill_value)
         )
-        data = data.assign(labels=data.labels.astype(int))
-        return self.__class__(data, self.model, self.name)
+        return self.assign(labels=data.labels.astype(int))
 
     def sanitize(self) -> Variable:
         """
@@ -1071,8 +1069,7 @@ def sanitize(self) -> Variable:
         linopy.Variable
         """
         if issubdtype(self.labels.dtype, floating):
-            data = self.data.assign(labels=self.labels.fillna(-1).astype(int))
-            return self.__class__(data, self.model, self.name)
+            return self.assign(labels=self.labels.fillna(-1).astype(int))
         return self
 
     def equals(self, other: Variable) -> bool:
@@ -1083,6 +1080,8 @@ def equals(self, other: Variable) -> bool:
 
     assign_coords = varwrap(Dataset.assign_coords)
 
+    assign = varwrap(assign_multiindex_safe)
+
     assign_multiindex_safe = varwrap(assign_multiindex_safe)
 
     broadcast_like = varwrap(Dataset.broadcast_like)
@@ -1450,7 +1449,9 @@ def set_blocks(self, blocks: DataArray) -> None:
 
         for name, variable in self.items():
             if dim in variable.dims:
-                variable.data["blocks"] = blocks.broadcast_like(variable.labels)
+                variable._data = assign_multiindex_safe(
+                    variable.data, blocks=blocks.broadcast_like(variable.labels)
+                )
 
     def get_blockmap(self, dtype: type = np.int8) -> ndarray:
         """