diff --git a/CLAUDE.md b/CLAUDE.md index 67155ae3..1f696a0b 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -110,27 +110,6 @@ When modifying the codebase, maintain consistency with these patterns and ensure * Always create a feature branch for new features or bug fixes. * Use the github cli (gh) to interact with the Github repository. -### GitHub Claude Code Integration - -This repository includes Claude Code GitHub Actions for automated assistance: - -1. **Automated PR Reviews** (`claude-code-review.yml`): - - Automatically reviews PRs only when first created (opened) - - Subsequent reviews require manual `@claude` mention - - Focuses on Python best practices, xarray patterns, and optimization correctness - - Can run tests and linting as part of the review - - **Skip initial review by**: Adding `[skip-review]` or `[WIP]` to PR title, or using draft PRs - -2. **Manual Claude Assistance** (`claude.yml`): - - Trigger by mentioning `@claude` in any: - - Issue comments - - Pull request comments - - Pull request reviews - - New issue body or title - - Claude can help with bug fixes, feature implementation, code explanations, etc. - -**Note**: Both workflows require the `ANTHROPIC_API_KEY` secret to be configured in the repository settings. - ## Development Guidelines @@ -140,3 +119,4 @@ This repository includes Claude Code GitHub Actions for automated assistance: 4. Use type hints and mypy for type checking. 5. Always write tests into the `test` directory, following the naming convention `test_*.py`. 6. Always write temporary and non git-tracked code in the `dev-scripts` directory. +7. In test scripts use linopy assertions from the testing.py module where useful (assert_linequal, assert_varequal, etc.) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 13df0267..b75c8728 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -6,6 +6,12 @@ Upcoming Version * Fix docs (pick highs solver) * Add the `sphinx-copybutton` to the documentation +* Harmonize coordinate alignment for operations with subset/superset objects: + - Multiplication and division fill missing coords with 0 (variable doesn't participate) + - Addition and subtraction of constants fill missing coords with 0 (identity element) and pin result to LHS coords + - Comparison operators (``==``, ``<=``, ``>=``) fill missing RHS coords with NaN (no constraint created) + - Fixes crash on ``subset + var`` / ``subset + expr`` reverse addition + - Fixes superset DataArrays expanding result coords beyond the variable's coordinate space * Add ``auto_mask`` parameter to ``Model`` class that automatically masks variables and constraints where bounds, coefficients, or RHS values contain NaN. This eliminates the need to manually create mask arrays when working with sparse or incomplete data. * Speed up LP file writing by 2-2.7x on large models through Polars streaming engine, join-based constraint assembly, and reduced per-constraint overhead * Fix multiplication of constant-only ``LinearExpression`` with other expressions diff --git a/linopy/expressions.py b/linopy/expressions.py index 649989f7..11730e5d 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -95,6 +95,7 @@ SUPPORTED_CONSTANT_TYPES = ( np.number, + np.bool_, int, float, DataArray, @@ -533,15 +534,33 @@ def _multiply_by_linear_expression( res = res + self.reset_const() * other.const return res + def _add_constant( + self: GenericExpression, other: ConstantLike + ) -> GenericExpression: + if np.isscalar(other): + return self.assign(const=self.const + other) + da = as_dataarray(other, coords=self.coords, dims=self.coord_dims) + da = da.reindex_like(self.const, fill_value=0) + return self.assign(const=self.const + da) + def _multiply_by_constant( self: GenericExpression, other: ConstantLike ) -> GenericExpression: multiplier = as_dataarray(other, coords=self.coords, dims=self.coord_dims) + multiplier = multiplier.reindex_like(self.const, fill_value=0) coeffs = self.coeffs * multiplier - assert all(coeffs.sizes[d] == s for d, s in self.coeffs.sizes.items()) const = self.const * multiplier return self.assign(coeffs=coeffs, const=const) + def _divide_by_constant( + self: GenericExpression, other: ConstantLike + ) -> GenericExpression: + divisor = as_dataarray(other, coords=self.coords, dims=self.coord_dims) + divisor = divisor.reindex_like(self.const, fill_value=1) + coeffs = self.coeffs / divisor + const = self.const / divisor + return self.assign(coeffs=coeffs, const=const) + def __div__(self: GenericExpression, other: SideLike) -> GenericExpression: try: if isinstance( @@ -557,7 +576,7 @@ def __div__(self: GenericExpression, other: SideLike) -> GenericExpression: f"{type(self)} and {type(other)}" "Non-linear expressions are not yet supported." ) - return self._multiply_by_constant(other=1 / other) + return self._divide_by_constant(other) except TypeError: return NotImplemented @@ -863,7 +882,10 @@ def to_constraint(self, sign: SignLike, rhs: SideLike) -> Constraint: sign : str, array-like Sign(s) of the constraints. rhs : constant, Variable, LinearExpression - Right-hand side of the constraint. + Right-hand side of the constraint. If a DataArray, it is + reindexed to match expression coordinates (fill_value=np.nan). + Extra dimensions in the RHS not present in the expression + raise a ValueError. NaN entries in the RHS mean "no constraint". Returns ------- @@ -876,6 +898,15 @@ def to_constraint(self, sign: SignLike, rhs: SideLike) -> Constraint: f"Both sides of the constraint are constant. At least one side must contain variables. {self} {rhs}" ) + if isinstance(rhs, DataArray): + extra_dims = set(rhs.dims) - set(self.coord_dims) + if extra_dims: + raise ValueError( + f"RHS DataArray has dimensions {extra_dims} not present " + f"in the expression. Cannot create constraint." + ) + rhs = rhs.reindex_like(self.const, fill_value=np.nan) + all_to_lhs = (self - rhs).data data = assign_multiindex_safe( all_to_lhs[["coeffs", "vars"]], sign=sign, rhs=-all_to_lhs.const @@ -1312,11 +1343,11 @@ def __add__( return other.__add__(self) try: - if np.isscalar(other): - return self.assign(const=self.const + other) - - other = as_expression(other, model=self.model, dims=self.coord_dims) - return merge([self, other], cls=self.__class__) + if isinstance(other, SUPPORTED_CONSTANT_TYPES): + return self._add_constant(other) + else: + other = as_expression(other, model=self.model, dims=self.coord_dims) + return merge([self, other], cls=self.__class__) except TypeError: return NotImplemented @@ -1852,15 +1883,15 @@ def __add__(self, other: SideLike) -> QuadraticExpression: dimension names of self will be filled in other """ try: - if np.isscalar(other): - return self.assign(const=self.const + other) - - other = as_expression(other, model=self.model, dims=self.coord_dims) + if isinstance(other, SUPPORTED_CONSTANT_TYPES): + return self._add_constant(other) + else: + other = as_expression(other, model=self.model, dims=self.coord_dims) - if isinstance(other, LinearExpression): - other = other.to_quadexpr() + if isinstance(other, LinearExpression): + other = other.to_quadexpr() - return merge([self, other], cls=self.__class__) + return merge([self, other], cls=self.__class__) except TypeError: return NotImplemented @@ -1878,13 +1909,7 @@ def __sub__(self, other: SideLike) -> QuadraticExpression: dimension names of self will be filled in other """ try: - if np.isscalar(other): - return self.assign(const=self.const - other) - - other = as_expression(other, model=self.model, dims=self.coord_dims) - if type(other) is LinearExpression: - other = other.to_quadexpr() - return merge([self, -other], cls=self.__class__) + return self.__add__(-other) except TypeError: return NotImplemented diff --git a/linopy/model.py b/linopy/model.py index a2fa8e4e..a8ef8782 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -741,6 +741,16 @@ def add_constraints( # TODO: add a warning here, routines should be safe against this data = data.drop_vars(drop_dims) + rhs_nan = data.rhs.isnull() + if rhs_nan.any(): + data["rhs"] = data.rhs.fillna(0) + rhs_mask = ~rhs_nan + mask = ( + rhs_mask + if mask is None + else (as_dataarray(mask).astype(bool) & rhs_mask) + ) + data["labels"] = -1 (data,) = xr.broadcast(data, exclude=[TERM_DIM]) diff --git a/linopy/monkey_patch_xarray.py b/linopy/monkey_patch_xarray.py index dc60608c..1e526c92 100644 --- a/linopy/monkey_patch_xarray.py +++ b/linopy/monkey_patch_xarray.py @@ -1,37 +1,45 @@ from __future__ import annotations from collections.abc import Callable -from functools import partialmethod, update_wrapper -from types import NotImplementedType +from functools import update_wrapper from typing import Any from xarray import DataArray from linopy import expressions, variables - -def monkey_patch(cls: type[DataArray], pass_unpatched_method: bool = False) -> Callable: - def deco(func: Callable) -> Callable: - func_name = func.__name__ - wrapped = getattr(cls, func_name) - update_wrapper(func, wrapped) - if pass_unpatched_method: - func = partialmethod(func, unpatched_method=wrapped) # type: ignore - setattr(cls, func_name, func) - return func - - return deco - - -@monkey_patch(DataArray, pass_unpatched_method=True) -def __mul__( - da: DataArray, other: Any, unpatched_method: Callable -) -> DataArray | NotImplementedType: - if isinstance( - other, - variables.Variable - | expressions.LinearExpression - | expressions.QuadraticExpression, - ): - return NotImplemented - return unpatched_method(da, other) +_LINOPY_TYPES = ( + variables.Variable, + variables.ScalarVariable, + expressions.LinearExpression, + expressions.ScalarLinearExpression, + expressions.QuadraticExpression, +) + + +def _make_patched_op(op_name: str) -> None: + """Patch a DataArray operator to return NotImplemented for linopy types, enabling reflected operators.""" + original = getattr(DataArray, op_name) + + def patched( + da: DataArray, other: Any, unpatched_method: Callable = original + ) -> Any: + if isinstance(other, _LINOPY_TYPES): + return NotImplemented + return unpatched_method(da, other) + + update_wrapper(patched, original) + setattr(DataArray, op_name, patched) + + +for _op in ( + "__mul__", + "__add__", + "__sub__", + "__truediv__", + "__le__", + "__ge__", + "__eq__", +): + _make_patched_op(_op) +del _op diff --git a/linopy/variables.py b/linopy/variables.py index d90a4775..83b4246e 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -315,6 +315,7 @@ def to_linexpr( Linear expression with the variables and coefficients. """ coefficient = as_dataarray(coefficient, coords=self.coords, dims=self.dims) + coefficient = coefficient.reindex_like(self.labels, fill_value=0) ds = Dataset({"coeffs": coefficient, "vars": self.labels}).expand_dims( TERM_DIM, -1 ) @@ -454,7 +455,7 @@ def __div__( f"{type(self)} and {type(other)}. " "Non-linear expressions are not yet supported." ) - return self.to_linexpr(1 / other) + return self.to_linexpr()._divide_by_constant(other) def __truediv__( self, coefficient: float | int | LinearExpression | Variable diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index 0da9ec7f..93a02f45 100644 --- a/test/test_linear_expression.py +++ b/test/test_linear_expression.py @@ -575,6 +575,371 @@ def test_linear_expression_multiplication_invalid( expr / x +class TestSubsetCoordinateAlignment: + @pytest.fixture + def subset(self) -> xr.DataArray: + return xr.DataArray([10.0, 30.0], dims=["dim_2"], coords={"dim_2": [1, 3]}) + + @pytest.fixture + def superset(self) -> xr.DataArray: + return xr.DataArray( + np.arange(25, dtype=float), dims=["dim_2"], coords={"dim_2": range(25)} + ) + + @pytest.fixture + def expected_fill(self) -> np.ndarray: + arr = np.zeros(20) + arr[1] = 10.0 + arr[3] = 30.0 + return arr + + def test_var_mul_subset( + self, v: Variable, subset: xr.DataArray, expected_fill: np.ndarray + ) -> None: + result = v * subset + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.coeffs.values).any() + np.testing.assert_array_equal(result.coeffs.squeeze().values, expected_fill) + + def test_expr_mul_subset( + self, v: Variable, subset: xr.DataArray, expected_fill: np.ndarray + ) -> None: + expr = 1 * v + result = expr * subset + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.coeffs.values).any() + np.testing.assert_array_equal(result.coeffs.squeeze().values, expected_fill) + + @pytest.mark.parametrize( + "make_lhs,make_rhs", + [ + (lambda v, s: s * v, lambda v, s: v * s), + (lambda v, s: s * (1 * v), lambda v, s: (1 * v) * s), + (lambda v, s: s + v, lambda v, s: v + s), + (lambda v, s: s + (v + 5), lambda v, s: (v + 5) + s), + ], + ids=["subset*var", "subset*expr", "subset+var", "subset+expr"], + ) + def test_commutativity( + self, v: Variable, subset: xr.DataArray, make_lhs: object, make_rhs: object + ) -> None: + assert_linequal(make_lhs(v, subset), make_rhs(v, subset)) + + def test_var_add_subset( + self, v: Variable, subset: xr.DataArray, expected_fill: np.ndarray + ) -> None: + result = v + subset + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.const.values).any() + np.testing.assert_array_equal(result.const.values, expected_fill) + + def test_var_sub_subset( + self, v: Variable, subset: xr.DataArray, expected_fill: np.ndarray + ) -> None: + result = v - subset + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.const.values).any() + np.testing.assert_array_equal(result.const.values, -expected_fill) + + def test_subset_sub_var(self, v: Variable, subset: xr.DataArray) -> None: + assert_linequal(subset - v, -v + subset) + + def test_expr_add_subset( + self, v: Variable, subset: xr.DataArray, expected_fill: np.ndarray + ) -> None: + expr = v + 5 + result = expr + subset + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.const.values).any() + np.testing.assert_array_equal(result.const.values, expected_fill + 5) + + def test_expr_sub_subset( + self, v: Variable, subset: xr.DataArray, expected_fill: np.ndarray + ) -> None: + expr = v + 5 + result = expr - subset + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.const.values).any() + np.testing.assert_array_equal(result.const.values, 5 - expected_fill) + + def test_subset_sub_expr(self, v: Variable, subset: xr.DataArray) -> None: + expr = v + 5 + assert_linequal(subset - expr, -(expr - subset)) + + def test_var_div_subset(self, v: Variable, subset: xr.DataArray) -> None: + result = v / subset + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.coeffs.values).any() + assert result.coeffs.squeeze().sel(dim_2=1).item() == pytest.approx(0.1) + assert result.coeffs.squeeze().sel(dim_2=0).item() == pytest.approx(1.0) + + def test_var_le_subset(self, v: Variable, subset: xr.DataArray) -> None: + con = v <= subset + assert con.sizes["dim_2"] == v.sizes["dim_2"] + assert con.rhs.sel(dim_2=1).item() == 10.0 + assert con.rhs.sel(dim_2=3).item() == 30.0 + assert np.isnan(con.rhs.sel(dim_2=0).item()) + + @pytest.mark.parametrize("sign", ["<=", ">=", "=="]) + def test_var_comparison_subset( + self, v: Variable, subset: xr.DataArray, sign: str + ) -> None: + if sign == "<=": + con = v <= subset + elif sign == ">=": + con = v >= subset + else: + con = v == subset + assert con.sizes["dim_2"] == v.sizes["dim_2"] + assert con.rhs.sel(dim_2=1).item() == 10.0 + assert np.isnan(con.rhs.sel(dim_2=0).item()) + + def test_expr_le_subset(self, v: Variable, subset: xr.DataArray) -> None: + expr = v + 5 + con = expr <= subset + assert con.sizes["dim_2"] == v.sizes["dim_2"] + assert con.rhs.sel(dim_2=1).item() == pytest.approx(5.0) + assert con.rhs.sel(dim_2=3).item() == pytest.approx(25.0) + assert np.isnan(con.rhs.sel(dim_2=0).item()) + + def test_add_commutativity_full_coords(self, v: Variable) -> None: + full = xr.DataArray( + np.arange(20, dtype=float), + dims=["dim_2"], + coords={"dim_2": range(20)}, + ) + assert_linequal(v + full, full + v) + + def test_superset_addition_pins_to_lhs( + self, v: Variable, superset: xr.DataArray + ) -> None: + result = v + superset + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.const.values).any() + + def test_superset_add_var(self, v: Variable, superset: xr.DataArray) -> None: + assert_linequal(superset + v, v + superset) + + def test_superset_sub_var(self, v: Variable, superset: xr.DataArray) -> None: + assert_linequal(superset - v, -v + superset) + + def test_superset_mul_var(self, v: Variable, superset: xr.DataArray) -> None: + assert_linequal(superset * v, v * superset) + + @pytest.mark.parametrize("sign", ["<=", ">="]) + def test_superset_comparison_var( + self, v: Variable, superset: xr.DataArray, sign: str + ) -> None: + if sign == "<=": + con = superset <= v + else: + con = superset >= v + assert con.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(con.lhs.coeffs.values).any() + assert not np.isnan(con.rhs.values).any() + + def test_disjoint_addition_pins_to_lhs(self, v: Variable) -> None: + disjoint = xr.DataArray( + [100.0, 200.0], dims=["dim_2"], coords={"dim_2": [50, 60]} + ) + result = v + disjoint + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.const.values).any() + np.testing.assert_array_equal(result.const.values, np.zeros(20)) + + def test_expr_div_subset(self, v: Variable, subset: xr.DataArray) -> None: + expr = 1 * v + result = expr / subset + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.coeffs.values).any() + assert result.coeffs.squeeze().sel(dim_2=1).item() == pytest.approx(0.1) + assert result.coeffs.squeeze().sel(dim_2=0).item() == pytest.approx(1.0) + + def test_subset_add_var_coefficients( + self, v: Variable, subset: xr.DataArray + ) -> None: + result = subset + v + np.testing.assert_array_equal(result.coeffs.squeeze().values, np.ones(20)) + + def test_subset_sub_var_coefficients( + self, v: Variable, subset: xr.DataArray + ) -> None: + result = subset - v + np.testing.assert_array_equal(result.coeffs.squeeze().values, -np.ones(20)) + + @pytest.mark.parametrize("sign", ["<=", ">=", "=="]) + def test_subset_comparison_var( + self, v: Variable, subset: xr.DataArray, sign: str + ) -> None: + if sign == "<=": + con = subset <= v + elif sign == ">=": + con = subset >= v + else: + con = subset == v + assert con.sizes["dim_2"] == v.sizes["dim_2"] + assert np.isnan(con.rhs.sel(dim_2=0).item()) + assert con.rhs.sel(dim_2=1).item() == pytest.approx(10.0) + + def test_superset_mul_pins_to_lhs( + self, v: Variable, superset: xr.DataArray + ) -> None: + result = v * superset + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.coeffs.values).any() + + def test_superset_div_pins_to_lhs(self, v: Variable) -> None: + superset_nonzero = xr.DataArray( + np.arange(1, 26, dtype=float), + dims=["dim_2"], + coords={"dim_2": range(25)}, + ) + result = v / superset_nonzero + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.coeffs.values).any() + + def test_quadexpr_add_subset( + self, v: Variable, subset: xr.DataArray, expected_fill: np.ndarray + ) -> None: + qexpr = v * v + result = qexpr + subset + assert isinstance(result, QuadraticExpression) + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.const.values).any() + np.testing.assert_array_equal(result.const.values, expected_fill) + + def test_quadexpr_sub_subset( + self, v: Variable, subset: xr.DataArray, expected_fill: np.ndarray + ) -> None: + qexpr = v * v + result = qexpr - subset + assert isinstance(result, QuadraticExpression) + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.const.values).any() + np.testing.assert_array_equal(result.const.values, -expected_fill) + + def test_quadexpr_mul_subset( + self, v: Variable, subset: xr.DataArray, expected_fill: np.ndarray + ) -> None: + qexpr = v * v + result = qexpr * subset + assert isinstance(result, QuadraticExpression) + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.coeffs.values).any() + np.testing.assert_array_equal(result.coeffs.squeeze().values, expected_fill) + + def test_subset_mul_quadexpr( + self, v: Variable, subset: xr.DataArray, expected_fill: np.ndarray + ) -> None: + qexpr = v * v + result = subset * qexpr + assert isinstance(result, QuadraticExpression) + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.coeffs.values).any() + np.testing.assert_array_equal(result.coeffs.squeeze().values, expected_fill) + + def test_subset_add_quadexpr(self, v: Variable, subset: xr.DataArray) -> None: + qexpr = v * v + assert_quadequal(subset + qexpr, qexpr + subset) + + def test_multidim_subset_mul(self, m: Model) -> None: + coords_a = pd.RangeIndex(4, name="a") + coords_b = pd.RangeIndex(5, name="b") + w = m.add_variables(coords=[coords_a, coords_b], name="w") + + subset_2d = xr.DataArray( + [[2.0, 3.0], [4.0, 5.0]], + dims=["a", "b"], + coords={"a": [1, 3], "b": [0, 4]}, + ) + result = w * subset_2d + assert result.sizes["a"] == 4 + assert result.sizes["b"] == 5 + assert not np.isnan(result.coeffs.values).any() + assert result.coeffs.squeeze().sel(a=1, b=0).item() == pytest.approx(2.0) + assert result.coeffs.squeeze().sel(a=3, b=4).item() == pytest.approx(5.0) + assert result.coeffs.squeeze().sel(a=0, b=0).item() == pytest.approx(0.0) + assert result.coeffs.squeeze().sel(a=1, b=2).item() == pytest.approx(0.0) + + def test_multidim_subset_add(self, m: Model) -> None: + coords_a = pd.RangeIndex(4, name="a") + coords_b = pd.RangeIndex(5, name="b") + w = m.add_variables(coords=[coords_a, coords_b], name="w") + + subset_2d = xr.DataArray( + [[2.0, 3.0], [4.0, 5.0]], + dims=["a", "b"], + coords={"a": [1, 3], "b": [0, 4]}, + ) + result = w + subset_2d + assert result.sizes["a"] == 4 + assert result.sizes["b"] == 5 + assert not np.isnan(result.const.values).any() + assert result.const.sel(a=1, b=0).item() == pytest.approx(2.0) + assert result.const.sel(a=3, b=4).item() == pytest.approx(5.0) + assert result.const.sel(a=0, b=0).item() == pytest.approx(0.0) + + def test_constraint_rhs_extra_dims_raises(self, v: Variable) -> None: + rhs = xr.DataArray( + [[1.0, 2.0]], dims=["extra", "dim_2"], coords={"dim_2": [0, 1]} + ) + with pytest.raises(ValueError, match="not present in the expression"): + v <= rhs + + def test_da_truediv_var_raises(self, v: Variable) -> None: + da = xr.DataArray(np.ones(20), dims=["dim_2"], coords={"dim_2": range(20)}) + with pytest.raises(TypeError): + da / v + + def test_disjoint_mul_produces_zeros(self, v: Variable) -> None: + disjoint = xr.DataArray( + [10.0, 20.0], dims=["dim_2"], coords={"dim_2": [50, 60]} + ) + result = v * disjoint + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.coeffs.values).any() + np.testing.assert_array_equal(result.coeffs.squeeze().values, np.zeros(20)) + + def test_disjoint_div_preserves_coeffs(self, v: Variable) -> None: + disjoint = xr.DataArray( + [10.0, 20.0], dims=["dim_2"], coords={"dim_2": [50, 60]} + ) + result = v / disjoint + assert result.sizes["dim_2"] == v.sizes["dim_2"] + assert not np.isnan(result.coeffs.values).any() + np.testing.assert_array_equal(result.coeffs.squeeze().values, np.ones(20)) + + def test_da_eq_da_still_works(self) -> None: + da1 = xr.DataArray([1, 2, 3]) + da2 = xr.DataArray([1, 2, 3]) + result = da1 == da2 + assert result.values.all() + + def test_da_eq_scalar_still_works(self) -> None: + da = xr.DataArray([1, 2, 3]) + result = da == 2 + np.testing.assert_array_equal(result.values, [False, True, False]) + + def test_subset_constraint_solve_integration(self) -> None: + from linopy import available_solvers + + if not available_solvers: + pytest.skip("No solver available") + m = Model() + coords = pd.RangeIndex(5, name="i") + x = m.add_variables(lower=0, upper=100, coords=[coords], name="x") + subset_ub = xr.DataArray([10.0, 20.0], dims=["i"], coords={"i": [1, 3]}) + m.add_constraints(x <= subset_ub, name="subset_ub") + m.add_objective(x.sum(), sense="max") + m.solve(solver_name=available_solvers[0]) + sol = m.solution["x"] + assert sol.sel(i=1).item() == pytest.approx(10.0) + assert sol.sel(i=3).item() == pytest.approx(20.0) + assert sol.sel(i=0).item() == pytest.approx(100.0) + assert sol.sel(i=2).item() == pytest.approx(100.0) + assert sol.sel(i=4).item() == pytest.approx(100.0) + + def test_expression_inherited_properties(x: Variable, y: Variable) -> None: expr = 10 * x + y assert isinstance(expr.attrs, dict)