diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 2d7a9fcb..2bf6965d 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -4,6 +4,11 @@ Release Notes Upcoming Version ---------------- +**Fix Regression** + +* Reinsert broadcasting logic of mask object to be fully compatible with performance improvements in version 0.6.2 using `np.where` instead of `xr.where`. + + Version 0.6.2 -------------- diff --git a/linopy/common.py b/linopy/common.py index e6eef583..0823deac 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -286,6 +286,32 @@ def as_dataarray( return arr +def broadcast_mask(mask: DataArray, labels: DataArray) -> DataArray: + """ + Broadcast a boolean mask to match the shape of labels. + + Ensures that mask dimensions are a subset of labels dimensions, broadcasts + the mask accordingly, and fills any NaN values (from missing coordinates) + with False while emitting a FutureWarning. + """ + assert set(mask.dims).issubset(labels.dims), ( + "Dimensions of mask not a subset of resulting labels dimensions." + ) + mask = mask.broadcast_like(labels) + if mask.isnull().any(): + warn( + "Mask contains coordinates not covered by the data dimensions. " + "Missing values will be filled with False (masked out). " + "In a future version, this will raise an error. " + "Use mask.reindex() or `linopy.align()` to explicitly handle missing " + "coordinates.", + FutureWarning, + stacklevel=3, + ) + mask = mask.fillna(False).astype(bool) + return mask + + # TODO: rename to to_pandas_dataframe def to_dataframe( ds: Dataset, diff --git a/linopy/model.py b/linopy/model.py index af171ae4..d5d4830a 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -9,7 +9,6 @@ import logging import os import re -import warnings from collections.abc import Callable, Mapping, Sequence from pathlib import Path from tempfile import NamedTemporaryFile, gettempdir @@ -30,6 +29,7 @@ as_dataarray, assign_multiindex_safe, best_int, + broadcast_mask, maybe_replace_signs, replace_by_map, set_int_index, @@ -552,16 +552,7 @@ def add_variables( if mask is not None: mask = as_dataarray(mask, coords=data.coords, dims=data.dims).astype(bool) - if set(mask.dims) != set(data["labels"].dims): - warnings.warn( - f"Mask dimensions {set(mask.dims)} do not match the data " - f"dimensions {set(data['labels'].dims)}. The mask will be " - f"broadcast across the missing dimensions " - f"{set(data['labels'].dims) - set(mask.dims)}. In a future " - "version, this will raise an error.", - FutureWarning, - stacklevel=2, - ) + mask = broadcast_mask(mask, data.labels) # Auto-mask based on NaN in bounds (use numpy for speed) if self.auto_mask: @@ -582,7 +573,7 @@ def add_variables( self._xCounter += data.labels.size if mask is not None: - data.labels.values = data.labels.where(mask, -1).values + data.labels.values = np.where(mask.values, data.labels.values, -1) data = data.assign_attrs( label_range=(start, end), name=name, binary=binary, integer=integer @@ -756,20 +747,7 @@ def add_constraints( if mask is not None: mask = as_dataarray(mask).astype(bool) - # TODO: simplify - assert set(mask.dims).issubset(data.dims), ( - "Dimensions of mask not a subset of resulting labels dimensions." - ) - if set(mask.dims) != set(data["labels"].dims): - warnings.warn( - f"Mask dimensions {set(mask.dims)} do not match the data " - f"dimensions {set(data['labels'].dims)}. The mask will be " - f"broadcast across the missing dimensions " - f"{set(data['labels'].dims) - set(mask.dims)}. In a future " - "version, this will raise an error.", - FutureWarning, - stacklevel=2, - ) + mask = broadcast_mask(mask, data.labels) # Auto-mask based on null expressions or NaN RHS (use numpy for speed) if self.auto_mask: @@ -780,11 +758,9 @@ def add_constraints( auto_mask_values = ~vars_all_invalid if original_rhs_mask is not None: coords, dims, rhs_notnull = original_rhs_mask - # Broadcast RHS mask to match data shape if needed if rhs_notnull.shape != auto_mask_values.shape: rhs_da = DataArray(rhs_notnull, coords=coords, dims=dims) - rhs_da, _ = xr.broadcast(rhs_da, data.labels) - rhs_notnull = rhs_da.values + rhs_notnull = rhs_da.broadcast_like(data.labels).values auto_mask_values = auto_mask_values & rhs_notnull auto_mask_arr = DataArray( auto_mask_values, coords=data.labels.coords, dims=data.labels.dims @@ -802,7 +778,7 @@ def add_constraints( self._cCounter += data.labels.size if mask is not None: - data.labels.values = data.labels.where(mask, -1).values + data.labels.values = np.where(mask.values, data.labels.values, -1) data = data.assign_attrs(label_range=(start, end), name=name) diff --git a/test/test_constraints.py b/test/test_constraints.py index afd2d77d..01aebb69 100644 --- a/test/test_constraints.py +++ b/test/test_constraints.py @@ -157,14 +157,12 @@ def test_masked_constraints() -> None: y = m.add_variables() mask = pd.Series([True] * 5 + [False] * 5) - with pytest.warns(FutureWarning, match="Mask dimensions"): - m.add_constraints(1 * x + 10 * y, EQUAL, 0, mask=mask) + m.add_constraints(1 * x + 10 * y, EQUAL, 0, mask=mask) assert (m.constraints.labels.con0[0:5, :] != -1).all() assert (m.constraints.labels.con0[5:10, :] == -1).all() def test_masked_constraints_broadcast() -> None: - """Test that a constraint mask with fewer dimensions broadcasts correctly.""" m: Model = Model() lower = xr.DataArray(np.zeros((10, 10)), coords=[range(10), range(10)]) @@ -172,39 +170,31 @@ def test_masked_constraints_broadcast() -> None: x = m.add_variables(lower, upper) y = m.add_variables() - # 1D mask applied to 2D constraint — must broadcast over second dim mask = pd.Series([True] * 5 + [False] * 5) - with pytest.warns(FutureWarning, match="Mask dimensions"): - m.add_constraints(1 * x + 10 * y, EQUAL, 0, name="bc1", mask=mask) + m.add_constraints(1 * x + 10 * y, EQUAL, 0, name="bc1", mask=mask) assert (m.constraints.labels.bc1[0:5, :] != -1).all() assert (m.constraints.labels.bc1[5:10, :] == -1).all() - # Mask along second dimension only mask2 = xr.DataArray([True] * 5 + [False] * 5, dims=["dim_1"]) - with pytest.warns(FutureWarning, match="Mask dimensions"): - m.add_constraints(1 * x + 10 * y, EQUAL, 0, name="bc2", mask=mask2) + m.add_constraints(1 * x + 10 * y, EQUAL, 0, name="bc2", mask=mask2) assert (m.constraints.labels.bc2[:, 0:5] != -1).all() assert (m.constraints.labels.bc2[:, 5:10] == -1).all() - -def test_constraints_mask_no_warning_when_aligned() -> None: - """Test that no FutureWarning is emitted when mask has same dims as data.""" - m: Model = Model() - - lower = xr.DataArray(np.zeros((10, 10)), coords=[range(10), range(10)]) - upper = xr.DataArray(np.ones((10, 10)), coords=[range(10), range(10)]) - x = m.add_variables(lower, upper) - y = m.add_variables() - - mask = xr.DataArray( - np.array([[True] * 10] * 5 + [[False] * 10] * 5), - coords=[range(10), range(10)], + mask3 = xr.DataArray( + [True, True, False, False, False], + dims=["dim_0"], + coords={"dim_0": range(5)}, ) - import warnings - - with warnings.catch_warnings(): - warnings.simplefilter("error", FutureWarning) - m.add_constraints(1 * x + 10 * y, EQUAL, 0, mask=mask) + with pytest.warns(FutureWarning, match="Missing values will be filled"): + m.add_constraints(1 * x + 10 * y, EQUAL, 0, name="bc3", mask=mask3) + assert (m.constraints.labels.bc3[0:2, :] != -1).all() + assert (m.constraints.labels.bc3[2:5, :] == -1).all() + assert (m.constraints.labels.bc3[5:10, :] == -1).all() + + # Mask with extra dimension not in data should raise + mask4 = xr.DataArray([True, False], dims=["extra_dim"]) + with pytest.raises(AssertionError, match="not a subset"): + m.add_constraints(1 * x + 10 * y, EQUAL, 0, name="bc4", mask=mask4) def test_non_aligned_constraints() -> None: diff --git a/test/test_variable_assignment.py b/test/test_variable_assignment.py index ec68b1e0..02da32df 100644 --- a/test/test_variable_assignment.py +++ b/test/test_variable_assignment.py @@ -227,8 +227,7 @@ def test_variable_assigment_masked() -> None: lower = pd.DataFrame(np.zeros((10, 10))) upper = pd.Series(np.ones(10)) mask = pd.Series([True] * 5 + [False] * 5) - with pytest.warns(FutureWarning, match="Mask dimensions"): - m.add_variables(lower, upper, mask=mask) + m.add_variables(lower, upper, mask=mask) assert m.variables.labels.var0[-1, -1].item() == -1 diff --git a/test/test_variables.py b/test/test_variables.py index 8b6c71ed..37de6aff 100644 --- a/test/test_variables.py +++ b/test/test_variables.py @@ -108,43 +108,36 @@ def test_variables_nvars(m: Model) -> None: def test_variables_mask_broadcast() -> None: - """Test that a mask with fewer dimensions broadcasts correctly.""" m = Model() lower = xr.DataArray(np.zeros((10, 10)), coords=[range(10), range(10)]) upper = xr.DataArray(np.ones((10, 10)), coords=[range(10), range(10)]) - # 1D mask applied to 2D variable — must broadcast over second dim mask = pd.Series([True] * 5 + [False] * 5) - with pytest.warns(FutureWarning, match="Mask dimensions"): - x = m.add_variables(lower, upper, name="x", mask=mask) + x = m.add_variables(lower, upper, name="x", mask=mask) assert (x.labels[0:5, :] != -1).all() assert (x.labels[5:10, :] == -1).all() - # Mask along second dimension only mask2 = xr.DataArray([True] * 5 + [False] * 5, dims=["dim_1"]) - with pytest.warns(FutureWarning, match="Mask dimensions"): - y = m.add_variables(lower, upper, name="y", mask=mask2) + y = m.add_variables(lower, upper, name="y", mask=mask2) assert (y.labels[:, 0:5] != -1).all() assert (y.labels[:, 5:10] == -1).all() - -def test_variables_mask_no_warning_when_aligned() -> None: - """Test that no FutureWarning is emitted when mask has same dims as data.""" - m = Model() - - lower = xr.DataArray(np.zeros((10, 10)), coords=[range(10), range(10)]) - upper = xr.DataArray(np.ones((10, 10)), coords=[range(10), range(10)]) - - mask = xr.DataArray( - np.array([[True] * 10] * 5 + [[False] * 10] * 5), - coords=[range(10), range(10)], + mask3 = xr.DataArray( + [True, True, False, False, False], + dims=["dim_0"], + coords={"dim_0": range(5)}, ) - import warnings - - with warnings.catch_warnings(): - warnings.simplefilter("error", FutureWarning) - m.add_variables(lower, upper, name="x", mask=mask) + with pytest.warns(FutureWarning, match="Missing values will be filled"): + z = m.add_variables(lower, upper, name="z", mask=mask3) + assert (z.labels[0:2, :] != -1).all() + assert (z.labels[2:5, :] == -1).all() + assert (z.labels[5:10, :] == -1).all() + + # Mask with extra dimension not in data should raise + mask4 = xr.DataArray([True, False], dims=["extra_dim"]) + with pytest.raises(AssertionError, match="not a subset"): + m.add_variables(lower, upper, name="w", mask=mask4) def test_variables_get_name_by_label(m: Model) -> None: