Fix error on inconsistent input to add_objective_state (odow#540)

src/user_interface.jl

@@ -923,38 +923,33 @@ If the objective state is `N`-dimensional, each keyword argument must be an
 function add_objective_state(
     update::Function,
     subproblem::JuMP.Model;
-    initial_value,
-    lipschitz,
-    lower_bound = -Inf,
-    upper_bound = Inf,
+    initial_value::Union{Real,Tuple},
+    lipschitz::Union{Real,Tuple},
+    lower_bound::Union{Real,Tuple} = -Inf,
+    upper_bound::Union{Real,Tuple} = Inf,
 )
+    tup_initial_value = _to_tuple(initial_value)
+    N = length(tup_initial_value)
     return add_objective_state(
         update,
         subproblem,
-        initial_value,
-        lower_bound,
-        upper_bound,
-        lipschitz,
+        tup_initial_value,
+        _to_tuple(lower_bound, N),
+        _to_tuple(upper_bound, N),
+        _to_tuple(lipschitz, N),
     )
 end
 
-# Internal function: add_objective_state with positional Float64 arguments.
-function add_objective_state(
-    update::Function,
-    subproblem::JuMP.Model,
-    initial_value::Float64,
-    lower_bound::Float64,
-    upper_bound::Float64,
-    lipschitz::Float64,
-)
-    return add_objective_state(
-        update,
-        subproblem,
-        (initial_value,),
-        (lower_bound,),
-        (upper_bound,),
-        (lipschitz,),
-    )
+_to_tuple(x::Real, N::Int = 1) = ntuple(i -> Float64(x), N)
+
+function _to_tuple(x::Tuple, N::Int = length(x))
+    if length(x) != N
+        error(
+            "Invalid dimension in the input to `add_objective_state`. Got: ",
+            "`$x`, but expected it to have length `$N`.",
+        )
+    end
+    return Float64.(x)
 end
 
 # Internal function: add_objective_state with positional NTuple arguments.

test/user_interface.jl

@@ -671,6 +671,96 @@ function test_initial_feasibility()
     return
 end
 
+function test_objective_state_error_missing_upper_bound()
+    model = SDDP.LinearPolicyGraph(
+        stages = 3,
+        sense = :Max,
+        upper_bound = 50.0,
+        optimizer = HiGHS.Optimizer,
+    ) do sp, _
+        @variable(sp, x >= 0, SDDP.State, initial_value = 2)
+        @constraint(sp, x.out <= x.in)
+        SDDP.add_objective_state(
+            sp,
+            initial_value = (1.5,),
+            lower_bound = (0.75,),
+            lipschitz = (100.0,),
+        ) do y, ω
+            return y + ω
+        end
+        SDDP.parameterize(sp, [-0.25, -0.125, 0.125, 0.25]) do ω
+            price = SDDP.objective_state(sp)
+            @stageobjective(sp, price * (x.in - x.out))
+        end
+    end
+    state = model[1].objective_state
+    @test state.lower_bound == (0.75,)
+    @test state.upper_bound == (Inf,)
+    @test state.initial_value == (1.5,)
+    @test JuMP.lower_bound(state.μ[1]) == -100.0
+    @test JuMP.upper_bound(state.μ[1]) == 100.0
+    return
+end
+
+function test_objective_state_error_missing_lower_bound()
+    model = SDDP.LinearPolicyGraph(
+        stages = 3,
+        sense = :Max,
+        upper_bound = 50.0,
+        optimizer = HiGHS.Optimizer,
+    ) do sp, _
+        @variable(sp, x >= 0, SDDP.State, initial_value = 2)
+        @constraint(sp, x.out <= x.in)
+        SDDP.add_objective_state(
+            sp,
+            initial_value = (1,),
+            lipschitz = (100,),
+        ) do y, ω
+            return y + ω
+        end
+        SDDP.parameterize(sp, [-0.25, -0.125, 0.125, 0.25]) do ω
+            price = SDDP.objective_state(sp)
+            @stageobjective(sp, price * (x.in - x.out))
+        end
+    end
+    state = model[1].objective_state
+    @test state.lower_bound == (-Inf,)
+    @test state.upper_bound == (Inf,)
+    @test state.initial_value == (1.0,)
+    @test JuMP.lower_bound.(state.μ) == (-100.0,)
+    @test JuMP.upper_bound.(state.μ) == (100.0,)
+    return
+end
+
+function test_objective_state_error_dimension_two_missing_lower_bound()
+    err = ErrorException(
+        "Invalid dimension in the input to `add_objective_state`. Got: " *
+        "`(100,)`, but expected it to have length `2`.",
+    )
+    @test_throws err SDDP.LinearPolicyGraph(
+        stages = 3,
+        sense = :Max,
+        upper_bound = 50.0,
+        optimizer = HiGHS.Optimizer,
+    ) do sp, _
+        @variable(sp, x >= 0, SDDP.State, initial_value = 2)
+        @constraint(sp, x.out <= x.in)
+        SDDP.add_objective_state(
+            sp,
+            initial_value = (1, 0),
+            lipschitz = (100,),
+            upper_bound = (1, Inf),
+        ) do y, ω
+            return y + ω
+        end
+        SDDP.parameterize(sp, [-0.25, -0.125, 0.125, 0.25]) do ω
+            price = SDDP.objective_state(sp)
+            @stageobjective(sp, price * (x.in - x.out))
+        end
+    end
+    return
+end
+
 function test_no_stage_objective()
     model = SDDP.LinearPolicyGraph(
         stages = 2,