From cc4d6ab7036fc4283bae69b1e29cc4e2d1aa2275 Mon Sep 17 00:00:00 2001 From: Fabian Date: Mon, 9 Feb 2026 14:28:33 +0100 Subject: [PATCH 1/2] refac: introduce consistent convention for linopy operations with subsets and supersets --- CLAUDE.md | 22 +- doc/release_notes.rst | 6 + linopy/expressions.py | 62 ++++-- linopy/model.py | 10 + linopy/monkey_patch_xarray.py | 64 +++--- linopy/variables.py | 3 +- test/test_linear_expression.py | 365 +++++++++++++++++++++++++++++++++ 7 files changed, 464 insertions(+), 68 deletions(-) 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 edf67076..311b93d0 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 Upcoming Version ---------------- diff --git a/linopy/expressions.py b/linopy/expressions.py index 848067cf..372a5c9f 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -532,15 +532,31 @@ def _multiply_by_linear_expression( res = res + self.reset_const() * other.const return res + def _add_constant( + self: GenericExpression, other: ConstantLike + ) -> GenericExpression: + 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( @@ -556,7 +572,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 @@ -862,7 +878,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 ------- @@ -875,6 +894,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 @@ -1313,9 +1341,11 @@ def __add__( 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__) + elif 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 @@ -1853,13 +1883,15 @@ def __add__(self, other: SideLike) -> QuadraticExpression: try: if np.isscalar(other): return self.assign(const=self.const + other) + elif isinstance(other, SUPPORTED_CONSTANT_TYPES): + return self._add_constant(other) + else: + other = as_expression(other, model=self.model, dims=self.coord_dims) - 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 @@ -1877,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 657b2d45..fc5472ae 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -699,6 +699,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) From e408b8e856c714e972e2d565133181d2d99aa120 Mon Sep 17 00:00:00 2001 From: Fabian Date: Wed, 11 Feb 2026 10:08:48 +0100 Subject: [PATCH 2/2] move scalar addition to add_constant --- linopy/expressions.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/linopy/expressions.py b/linopy/expressions.py index b9eb5579..11730e5d 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -95,6 +95,7 @@ SUPPORTED_CONSTANT_TYPES = ( np.number, + np.bool_, int, float, DataArray, @@ -536,6 +537,8 @@ def _multiply_by_linear_expression( 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) @@ -1340,9 +1343,7 @@ def __add__( return other.__add__(self) try: - if np.isscalar(other): - return self.assign(const=self.const + other) - elif isinstance(other, SUPPORTED_CONSTANT_TYPES): + if isinstance(other, SUPPORTED_CONSTANT_TYPES): return self._add_constant(other) else: other = as_expression(other, model=self.model, dims=self.coord_dims) @@ -1882,9 +1883,7 @@ 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) - elif isinstance(other, SUPPORTED_CONSTANT_TYPES): + if isinstance(other, SUPPORTED_CONSTANT_TYPES): return self._add_constant(other) else: other = as_expression(other, model=self.model, dims=self.coord_dims)