Merge branch 'master' into oetc-support

Author: FabianHofmann

doc/release_notes.rst

@@ -3,6 +3,7 @@ Release Notes
 
 .. Upcoming Version
 .. ----------------
+.. * Improved constraint equality check in `linopy.testing.assert_conequal` to less strict optionally
 
 Version 0.5.6
 --------------

linopy/expressions.py

@@ -1941,49 +1941,40 @@ def merge(
     dim : str
         Dimension along which the expressions should be concatenated.
     cls : type
-        Type of the resulting expression.
+        Explicitly set the type of the resulting expression (So that the type checker will know the return type)
     **kwargs
         Additional keyword arguments passed to xarray.concat. Defaults to
         {coords: "minimal", compat: "override"} or, in the special case described
         above, to {coords: "minimal", compat: "override", "join": "override"}.
 
     Returns
     -------
-    res : linopy.LinearExpression
+    res : linopy.LinearExpression or linopy.QuadraticExpression
     """
-    if cls is None:
-        warn(
-            "Using merge without specifying the class is deprecated",
-            DeprecationWarning,
-        )
-        cls = LinearExpression
-
-    linopy_types = (variables.Variable, LinearExpression, QuadraticExpression)
-
     if not isinstance(exprs, list) and len(add_exprs):
         warn(
             "Passing a tuple to the merge function is deprecated. Please pass a list of objects to be merged",
             DeprecationWarning,
         )
         exprs = [exprs] + list(add_exprs)  # type: ignore
-    model = exprs[0].model
 
-    if (
-        cls is QuadraticExpression
-        and dim == TERM_DIM
-        and any(type(e) is LinearExpression for e in exprs)
-    ):
+    has_quad_expression = any(type(e) is QuadraticExpression for e in exprs)
+    has_linear_expression = any(type(e) is LinearExpression for e in exprs)
+    if cls is None:
+        cls = QuadraticExpression if has_quad_expression else LinearExpression
+
+    if cls is QuadraticExpression and dim == TERM_DIM and has_linear_expression:
         raise ValueError(
             "Cannot merge linear and quadratic expressions along term dimension."
             "Convert to QuadraticExpression first."
         )
 
-    if cls is not QuadraticExpression and any(
-        type(e) is QuadraticExpression for e in exprs
-    ):
-        raise ValueError(
-            "Cannot merge linear and quadratic expressions to QuadraticExpression"
-        )
+    if has_quad_expression and cls is not QuadraticExpression:
+        raise ValueError("Cannot merge linear expressions to QuadraticExpression")
+
+    linopy_types = (variables.Variable, LinearExpression, QuadraticExpression)
+
+    model = exprs[0].model
 
     if cls in linopy_types and dim in HELPER_DIMS:
         coord_dims = [

linopy/testing.py

@@ -29,9 +29,25 @@ def assert_quadequal(
     return assert_equal(_expr_unwrap(a), _expr_unwrap(b))
 
 
-def assert_conequal(a: Constraint, b: Constraint) -> None:
-    """Assert that two constraints are equal."""
-    return assert_equal(_con_unwrap(a), _con_unwrap(b))
+def assert_conequal(a: Constraint, b: Constraint, strict: bool = True) -> None:
+    """
+    Assert that two constraints are equal.
+
+    Parameters
+    ----------
+        a: Constraint
+            The first constraint.
+        b: Constraint
+            The second constraint.
+        strict: bool
+            Whether to compare the constraints strictly. If not, only compare mathematically relevant parts.
+    """
+    if strict:
+        assert_equal(_con_unwrap(a), _con_unwrap(b))
+    else:
+        assert_linequal(a.lhs, b.lhs)
+        assert_equal(a.sign, b.sign)
+        assert_equal(a.rhs, b.rhs)
 
 
 def assert_model_equal(a: Model, b: Model) -> None:

test/test_constraints.py

@@ -42,6 +42,31 @@ def test_constraint_assignment() -> None:
     assert_conequal(m.constraints.con0, con0)
 
 
+def test_constraint_equality() -> None:
+    m: Model = Model()
+
+    lower: xr.DataArray = xr.DataArray(
+        np.zeros((10, 10)), coords=[range(10), range(10)]
+    )
+    upper: xr.DataArray = xr.DataArray(np.ones((10, 10)), coords=[range(10), range(10)])
+    x = m.add_variables(lower, upper, name="x")
+    y = m.add_variables(name="y")
+
+    con0 = m.add_constraints(1 * x + 10 * y, EQUAL, 0)
+
+    assert_conequal(con0, 1 * x + 10 * y == 0, strict=False)
+    assert_conequal(1 * x + 10 * y == 0, 1 * x + 10 * y == 0, strict=False)
+
+    with pytest.raises(AssertionError):
+        assert_conequal(con0, 1 * x + 10 * y <= 0, strict=False)
+
+    with pytest.raises(AssertionError):
+        assert_conequal(con0, 1 * x + 10 * y >= 0, strict=False)
+
+    with pytest.raises(AssertionError):
+        assert_conequal(10 * y + 2 * x == 0, 1 * x + 10 * y == 0, strict=False)
+
+
 def test_constraints_getattr_formatted() -> None:
     m: Model = Model()
     x = m.add_variables(0, 10, name="x")

test/test_linear_expression.py

@@ -1128,9 +1128,8 @@ def test_merge(x: Variable, y: Variable, z: Variable) -> None:
     res = merge([expr1, expr2], cls=LinearExpression)
     assert res.nterm == 6
 
-    with pytest.warns(DeprecationWarning):
-        res: LinearExpression = merge([expr1, expr2])  # type: ignore
-        assert res.nterm == 6
+    res: LinearExpression = merge([expr1, expr2])  # type: ignore
+    assert isinstance(res, LinearExpression)
 
     # now concat with same length of terms
     expr1 = z.sel(dim_0=0).sum("dim_1")

test/test_quadratic_expression.py

@@ -238,6 +238,9 @@ def test_merge_linear_expression_and_quadratic_expression(
     with pytest.raises(ValueError):
         merge([linexpr, quadexpr], cls=QuadraticExpression)
 
+    new_quad_ex = merge([linexpr.to_quadexpr(), quadexpr])  # type: ignore
+    assert isinstance(new_quad_ex, QuadraticExpression)
+
     with pytest.warns(DeprecationWarning):
         merge(quadexpr, quadexpr, cls=QuadraticExpression)  # type: ignore