more changes

Author: RobbieKiwi

linopy/model.py

@@ -268,11 +268,7 @@ def objective(self) -> Objective:
     def objective(
         self, obj: Objective | LinearExpression | QuadraticExpression
     ) -> Objective:
-        if not isinstance(obj, Objective):
-            expr = obj
-        else:
-            expr = obj.expression
-        self.add_objective(expr=expr, overwrite=True, allow_constant=False)
+        self.add_objective(expr=obj, overwrite=True, allow_constant=False)
         return self._objective
 
     @property
@@ -782,23 +778,46 @@ def add_constraints(
         self.constraints.add(constraint)
         return constraint
 
+    @overload
+    def add_objective(
+        self,
+        expr: Objective,
+        sense: None = None,
+        overwrite: bool = False,
+        allow_constant: bool = False,
+    ) -> None: ...
+
+    @overload
     def add_objective(
         self,
         expr: Variable
         | LinearExpression
         | QuadraticExpression
         | Sequence[tuple[ConstantLike, VariableLike]],
+        sense: Literal["min", "max"] | None = None,
+        overwrite: bool = False,
+        allow_constant: bool = False,
+    ) -> None: ...
+
+    def add_objective(
+        self,
+        expr: Variable
+        | LinearExpression
+        | QuadraticExpression
+        | Sequence[tuple[ConstantLike, VariableLike]]
+        | Objective,
+        sense: Literal["min", "max"] | None = None,
         overwrite: bool = False,
-        sense: str = "min",
         allow_constant: bool = False,
     ) -> None:
         """
         Add an objective function to the model.
 
         Parameters
         ----------
-        expr : linopy.LinearExpression, linopy.QuadraticExpression
+        expr : linopy.Variable, linopy.LinearExpression, linopy.QuadraticExpression, Objective
             Expression describing the objective function.
+        sense: "min" or "max", the sense to optimize for. Defaults to min. Cannot be set if passing Objective directly
         overwrite : False, optional
             Whether to overwrite the existing objective. The default is False.
         allow_constant: bool, optional
@@ -814,15 +833,20 @@ def add_objective(
                 "Objective already defined."
                 " Set `overwrite` to True to force overwriting."
             )
-        if isinstance(expr, Variable):
-            expr = 1 * expr
 
-        if not allow_constant and expr.has_constant:
+        if isinstance(expr, Objective):
+            assert sense is None, "Cannot set sense if objective object is passed"
+            objective = expr
+            assert objective.model == self
+        else:
+            sense = sense or "min"
+            objective = Objective(expression=expr, model=self, sense=sense)
+
+        if not allow_constant and objective.expression.has_constant:
             raise ConstantObjectiveError(
                 "Objective contains constant term. Either remove constants from the expression with expr.drop_constants() or use model.add_objective(..., allow_constant=True).",
             )
 
-        objective = Objective(expression=expr, model=self, sense=sense)
         self._objective = objective
 
     def remove_variables(self, name: str) -> None:

linopy/objective.py

@@ -9,7 +9,7 @@
 
 import functools
 from collections.abc import Callable, Sequence
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Literal
 
 import numpy as np
 import polars as pl
@@ -24,6 +24,7 @@
 
 from linopy import expressions
 from linopy.types import ConstantLike
+from linopy.variables import Variable
 
 if TYPE_CHECKING:
     from linopy.expressions import LinearExpression, QuadraticExpression
@@ -64,13 +65,19 @@ class Objective:
 
     def __init__(
         self,
-        expression: expressions.LinearExpression | expressions.QuadraticExpression,
+        expression: Variable
+        | expressions.LinearExpression
+        | expressions.QuadraticExpression,
         model: Model,
-        sense: str = "min",
+        sense: Literal["min", "max"] = "min",
     ) -> None:
         self._model: Model = model
         self._value: float | None = None
 
+        if isinstance(expression, Variable):
+            expression = 1 * expression
+
+        assert sense in ["min", "max"]
         self.sense: str = sense
         self.expression: (
             expressions.LinearExpression | expressions.QuadraticExpression

test/test_model.py

@@ -9,10 +9,12 @@
 from tempfile import gettempdir
 
 import numpy as np
+import pandas as pd
 import pytest
 import xarray as xr
 
 from linopy import EQUAL, Model
+from linopy.model import ConstantObjectiveError
 from linopy.testing import assert_model_equal
 
 target_shape: tuple[int, int] = (10, 10)
@@ -67,7 +69,7 @@ def test_objective() -> None:
     y = m.add_variables(lower, upper, name="y")
 
     obj1 = (10 * x + 5 * y).sum()
-    m.add_objective(obj1)
+    m.add_objective(obj1, allow_constant=True)
     assert m.objective.vars.size == 200
 
     # test overwriting
@@ -82,8 +84,9 @@ def test_objective() -> None:
     assert m.objectiverange.min() == 2
     assert m.objectiverange.max() == 2
 
-    # test objective with constant which is supported
-    m.objective = m.objective + 3
+    # test setting constant term in objective with explicitly allowing it
+    with pytest.raises(ConstantObjectiveError):
+        m.objective = m.objective + 3
 
 
 def test_remove_variable() -> None:
@@ -162,3 +165,20 @@ def test_assert_model_equal() -> None:
     m.add_objective(obj)
 
     assert_model_equal(m, m)
+
+
+def test_constant_not_allowed_in_objective_unless_specified_explicitly() -> None:
+    model = Model()
+    days = pd.Index(["Mon", "Tue", "Wed", "Thu", "Fri"], name="day")
+    x = model.add_variables(name="x", coords=[days])
+    non_linear = x + 1
+
+    with pytest.raises(ConstantObjectiveError):
+        model.add_objective(expr=non_linear, overwrite=True, allow_constant=False)
+    with pytest.raises(ConstantObjectiveError):
+        model.add_objective(expr=non_linear, overwrite=True)
+
+    with pytest.raises(ConstantObjectiveError):
+        model.objective = non_linear
+
+    model.add_objective(expr=non_linear, overwrite=True, allow_constant=True)

test/test_objective.py

@@ -187,9 +187,3 @@ def test_repr(linear_objective: Objective, quadratic_objective: Objective) -> No
 
     assert "Linear" in linear_objective.__repr__()
     assert "Quadratic" in quadratic_objective.__repr__()
-
-
-def test_objective_constant() -> None:
-    m = Model()
-    linear_expr = LinearExpression(None, m) + 1
-    m.objective = Objective(linear_expr, m)

test/test_optimization.py

@@ -20,7 +20,6 @@
 from linopy import GREATER_EQUAL, LESS_EQUAL, Model, solvers
 from linopy.common import to_path
 from linopy.expressions import LinearExpression
-from linopy.model import ConstantObjectiveError
 from linopy.solver_capabilities import (
     SolverFeature,
     get_available_solvers_with_feature,
@@ -956,20 +955,6 @@ def test_model_resolve(
     assert np.isclose(model.objective.value or 0, 5.25)
 
 
-def test_constant_not_allowed_unless_specified_explicitly(model: Model) -> None:
-    objective = model.objective.expression + 1
-
-    with pytest.raises(ConstantObjectiveError):
-        model.add_objective(expr=objective, overwrite=True, allow_constant=False)
-    with pytest.raises(ConstantObjectiveError):
-        model.add_objective(expr=objective, overwrite=True)
-
-    with pytest.raises(ConstantObjectiveError):
-        model.objective = objective
-
-    model.add_objective(expr=objective, overwrite=True, allow_constant=True)
-
-
 def test_constant_feasible(model: Model) -> None:
     objective = model.objective.expression + 1
     model.add_objective(expr=objective, overwrite=True, allow_constant=True)