fix(variables): Variable.where works with solution attached (#790)

* fix(variables): Variable.where works with solution attached

Variable.where passed a fill dict covering only labels/lower/upper to
Dataset.where, which requires an exact match against the dataset's data
variables. After solve() or fix() the variable carries extra columns
(solution, stashed bounds), so where() raised a ValueError. Augment the
fill dict to cover all present data variables, defaulting extras to NaN.

* fix(variables): annotate where fill dict for mypy

Author: FabianHofmann

doc/release_notes.rst

@@ -30,6 +30,7 @@ Upcoming Version
 
 * LP file export now honors bounds tightened below ``[0, 1]`` on a binary variable via the ``.lower``/``.upper`` setters after creation (e.g. ``upper = 0``). Previously such bounds were written only by ``io_api="direct"`` and dropped by ``io_api="lp"``. (https://github.com/PyPSA/linopy/issues/776)
 * Freezing an empty constraint group (e.g. an empty ``isel`` slice) no longer raises ``ValueError: cannot reshape array of size 0``. ``Model(freeze_constraints=True)`` and ``Constraint.freeze()`` now round-trip zero-row constraints losslessly.
+* ``Variable.where`` no longer raises ``ValueError: exact match required for all data variable names`` once a solution is attached (after ``Model.solve``) or the variable is fixed. The fill value now covers auxiliary data variables (``solution``, stashed bounds) instead of only ``labels``/``lower``/``upper``.
 
 Version 0.8.0
 -------------

linopy/variables.py

@@ -1223,6 +1223,9 @@ def where(
             raise ValueError(
                 f"other must be a Variable, ScalarVariable, dict or Dataset, got {type(other)}"
             )
+        if isinstance(_other, dict):
+            fill: dict[str, float] = {str(k): np.nan for k in self.data}
+            _other = {**fill, **_other}
         return self.__class__(
             self.data.where(cond, _other, **kwargs), self.model, self.name
         )

test/test_variable.py

@@ -276,6 +276,15 @@ def test_variable_where(x: linopy.Variable) -> None:
         x.where([True] * 4 + [False] * 6, 0)  # type: ignore
 
 
+def test_variable_where_with_solution(x: linopy.Variable) -> None:
+    x.solution = xr.DataArray(np.arange(10.0), coords=x.labels.coords)
+    cond = [True] * 4 + [False] * 6
+    filtered = x.where(cond)
+    assert filtered.labels[9] == x._fill_value["labels"]
+    assert filtered.data["solution"][0] == 0.0
+    assert np.isnan(filtered.data["solution"][9])
+
+
 def test_variable_shift(x: linopy.Variable) -> None:
     x = x.shift(first=3)
     assert isinstance(x, linopy.variables.Variable)