From d5ae01d02996d11a5abc2e85e4bed78efe4b64a1 Mon Sep 17 00:00:00 2001
From: odow <o.dowson@gmail.com>
Date: Mon, 20 Jan 2025 14:02:58 +1300
Subject: [PATCH 01/10] Add variable bounds to incoming states in
 LagrangianDuality

---
 src/plugins/duality_handlers.jl |  7 +++
 src/user_interface.jl           | 77 +++++++++++++++++++++++++++++++++
 test/user_interface.jl          | 41 ++++++++++++++++++
 3 files changed, 125 insertions(+)

diff --git a/src/plugins/duality_handlers.jl b/src/plugins/duality_handlers.jl
index dc3391d3c..bc9e1b3ce 100644
--- a/src/plugins/duality_handlers.jl
+++ b/src/plugins/duality_handlers.jl
@@ -179,6 +179,13 @@ function get_dual_solution(node::Node, lagrange::LagrangianDuality)
         x_in_value[i] = JuMP.fix_value(state.in)
         h_expr[i] = @expression(node.subproblem, state.in - x_in_value[i])
         JuMP.unfix(state.in)
+        l, u = node.incoming_state_bounds[key]
+        if l !== nothing
+            JuMP.set_lower_bound(state.in, l)
+        end
+        if u !== nothing
+            JuMP.set_upper_bound(state.in, u)
+        end
         λ_star[i] = conic_dual[key]
     end
     # Check that the conic dual is feasible for the subproblem. Sometimes it
diff --git a/src/user_interface.jl b/src/user_interface.jl
index 645c7e0e6..5ee925300 100644
--- a/src/user_interface.jl
+++ b/src/user_interface.jl
@@ -655,6 +655,7 @@ mutable struct Node{T}
     ext::Dict{Symbol,Any}
     # Lock for threading
     lock::ReentrantLock
+    incoming_state_bounds::Dict{Symbol,NTuple{2,Union{Nothing,Float64}}}
 end
 
 function Base.show(io::IO, node::Node)
@@ -987,6 +988,7 @@ function PolicyGraph(
             # The extension dictionary.
             Dict{Symbol,Any}(),
             ReentrantLock(),
+            Dict{Symbol,NTuple{2,Union{Nothing,Float64}}}(),
         )
         subproblem.ext[:sddp_policy_graph] = policy_graph
         policy_graph.nodes[node_index] = subproblem.ext[:sddp_node] = node
@@ -1032,9 +1034,84 @@ function PolicyGraph(
     if length(graph.belief_partition) > 0
         initialize_belief_states(policy_graph, graph)
     end
+    domain = _get_incoming_domain(policy_graph)
+    for (node_name, node) in policy_graph.nodes
+        for (k, v) in domain[node_name]
+            node.incoming_state_bounds[k] = v
+        end
+    end
     return policy_graph
 end
 
+function _get_incoming_domain(model::PolicyGraph{T}) where {T}
+    function _bounds(x)
+        l, u = nothing, nothing
+        if has_lower_bound(x)
+            l = lower_bound(x)
+        end
+        if has_upper_bound(x)
+            u = upper_bound(x)
+        end
+        if is_fixed(x)
+            l = u = fix_value(x)
+        end
+        if is_binary(x)
+            l, u = max(something(l, 0.0), 0.0), min(something(u, 1.0), 1.0)
+        end
+        return l, u
+    end
+    compare(f, a, b) = f(a, b)
+    compare(f, ::Nothing, b) = b
+    compare(f, a, ::Nothing) = a
+    compare(f, ::Nothing, ::Nothing) = nothing
+    outgoing_bounds = Dict{Tuple{T,Symbol},NTuple{2,Union{Nothing,Float64}}}(
+        (k, state_name) => (nothing, nothing)
+        for (k, node) in model.nodes
+        for (state_name, state) in node.states
+    )
+    for (k, node) in model.nodes
+        for noise in node.noise_terms
+            SDDP.parameterize(node, noise.term)
+            for (state_name, state) in node.states
+                l, u = outgoing_bounds[(k, state_name)]
+                l_new, u_new = _bounds(state.out)
+                outgoing_bounds[(k, state_name)] =
+                    compare(min, l, l_new), compare(max, u, u_new)
+            end
+        end
+    end
+    incoming_bounds = Dict{T,Dict{Symbol,NTuple{2,Union{Nothing,Float64}}}}()
+    for (k, node) in model.nodes
+        incoming_bounds[k] = Dict{Symbol,NTuple{2,Union{Nothing,Float64}}}(
+            state_name => (nothing, nothing)
+            for (state_name, state) in node.states
+        )
+    end
+    for (parent_name, parent) in model.nodes
+        for (state_name, state) in parent.states
+            l_new, u_new = outgoing_bounds[(parent_name, state_name)]
+            for child in parent.children
+                l, u = incoming_bounds[child.term][state_name]
+                incoming_bounds[child.term][state_name] =
+                    compare(min, l, l_new), compare(max, u, u_new)
+            end
+        end
+    end
+    for (state_name, value) in model.initial_root_state
+        for child in model.root_children
+            l, u = incoming_bounds[child.term][state_name]
+            if l !== nothing
+                l = compare(min, l, value)
+            end
+            if u !== nothing
+                u = compare(max, u, value)
+            end
+            incoming_bounds[child.term][state_name] = (l, u)
+        end
+    end
+    return incoming_bounds
+end
+
 # Internal function: set up ::BeliefState for each node.
 function initialize_belief_states(
     policy_graph::PolicyGraph{T},
diff --git a/test/user_interface.jl b/test/user_interface.jl
index e83c3e33a..6a93323ec 100644
--- a/test/user_interface.jl
+++ b/test/user_interface.jl
@@ -911,6 +911,47 @@ function test_print_problem_statistics()
     return
 end
 
+function test_incoming_state_bounds()
+    graph = SDDP.Graph(0)
+    SDDP.add_node.((graph,), 1:3)
+    SDDP.add_edge(graph, 0 => 1, 0.5)
+    SDDP.add_edge(graph, 0 => 2, 0.5)
+    SDDP.add_edge(graph, 1 => 2, 1.0)
+    SDDP.add_edge(graph, 2 => 2, 0.45)
+    SDDP.add_edge(graph, 2 => 3, 0.5)
+    SDDP.add_edge(graph, 3 => 1, 1.0)
+
+    model = SDDP.PolicyGraph(graph; lower_bound = 0.0) do sp, node
+        @variable(sp, x, SDDP.State, initial_value = 0)
+        @variable(sp, y, SDDP.State, initial_value = -1)
+        @variable(sp, z, SDDP.State, initial_value = 1)
+        if node == 1
+            set_binary(z.out)
+            SDDP.parameterize(sp, 1:2) do w
+                JuMP.set_lower_bound(x.out, -w)
+                return JuMP.set_upper_bound(y.out, w)
+            end
+        elseif node == 2
+        elseif node == 3
+            set_lower_bound(x.out, 1)
+            set_upper_bound(x.out, 2)
+            fix(y.out, 3)
+        end
+        @stageobjective(sp, 0)
+        return
+    end
+    @test model[1].incoming_state_bounds ==
+          Dict(:x => (0.0, 2.0), :y => (-1.0, 3.0), :z => (nothing, nothing))
+    @test model[2].incoming_state_bounds ==
+          Dict(:x => (-2.0, nothing), :y => (nothing, 2.0), :z => (0.0, 1.0))
+    @test model[3].incoming_state_bounds == Dict(
+        :x => (nothing, nothing),
+        :y => (nothing, nothing),
+        :z => (nothing, nothing),
+    )
+    return
+end
+
 end  # module
 
 TestUserInterface.runtests()

From 1201be141051749688c4e3840c2b992f7bad11f5 Mon Sep 17 00:00:00 2001
From: odow <o.dowson@gmail.com>
Date: Mon, 20 Jan 2025 14:28:42 +1300
Subject: [PATCH 02/10] Update

---
 src/user_interface.jl            | 9 ++++-----
 test/plugins/sampling_schemes.jl | 4 ++--
 test/user_interface.jl           | 3 ++-
 3 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/src/user_interface.jl b/src/user_interface.jl
index 5ee925300..92f75aa3b 100644
--- a/src/user_interface.jl
+++ b/src/user_interface.jl
@@ -1065,13 +1065,12 @@ function _get_incoming_domain(model::PolicyGraph{T}) where {T}
     compare(f, a, ::Nothing) = a
     compare(f, ::Nothing, ::Nothing) = nothing
     outgoing_bounds = Dict{Tuple{T,Symbol},NTuple{2,Union{Nothing,Float64}}}(
-        (k, state_name) => (nothing, nothing)
-        for (k, node) in model.nodes
+        (k, state_name) => (nothing, nothing) for (k, node) in model.nodes
         for (state_name, state) in node.states
     )
     for (k, node) in model.nodes
         for noise in node.noise_terms
-            SDDP.parameterize(node, noise.term)
+            parameterize(node, noise.term)
             for (state_name, state) in node.states
                 l, u = outgoing_bounds[(k, state_name)]
                 l_new, u_new = _bounds(state.out)
@@ -1083,8 +1082,8 @@ function _get_incoming_domain(model::PolicyGraph{T}) where {T}
     incoming_bounds = Dict{T,Dict{Symbol,NTuple{2,Union{Nothing,Float64}}}}()
     for (k, node) in model.nodes
         incoming_bounds[k] = Dict{Symbol,NTuple{2,Union{Nothing,Float64}}}(
-            state_name => (nothing, nothing)
-            for (state_name, state) in node.states
+            state_name => (nothing, nothing) for
+            (state_name, state) in node.states
         )
     end
     for (parent_name, parent) in model.nodes
diff --git a/test/plugins/sampling_schemes.jl b/test/plugins/sampling_schemes.jl
index 0958aa522..b75f3c31b 100644
--- a/test/plugins/sampling_schemes.jl
+++ b/test/plugins/sampling_schemes.jl
@@ -380,8 +380,8 @@ function test_SimulatorSamplingScheme_with_noise()
     ) do sp, node
         t, price = node
         @variable(sp, 0 <= x <= 1, SDDP.State, initial_value = 0)
-        SDDP.parameterize(sp, [(price, i) for i in 1:2]) do ω
-            return SDDP.@stageobjective(sp, price * x.out + i)
+        SDDP.parameterize(sp, [(price, i) for i in 1:2]) do (p, i)
+            return SDDP.@stageobjective(sp, p * x.out + i)
         end
     end
     sampler = SDDP.SimulatorSamplingScheme(simulator)
diff --git a/test/user_interface.jl b/test/user_interface.jl
index 6a93323ec..4d9a02fbb 100644
--- a/test/user_interface.jl
+++ b/test/user_interface.jl
@@ -929,7 +929,8 @@ function test_incoming_state_bounds()
             set_binary(z.out)
             SDDP.parameterize(sp, 1:2) do w
                 JuMP.set_lower_bound(x.out, -w)
-                return JuMP.set_upper_bound(y.out, w)
+                JuMP.set_upper_bound(y.out, w)
+                return
             end
         elseif node == 2
         elseif node == 3

From 24d34e29d9825ad0ead49b156ffa6c9473b9dd75 Mon Sep 17 00:00:00 2001
From: odow <o.dowson@gmail.com>
Date: Mon, 20 Jan 2025 14:46:32 +1300
Subject: [PATCH 03/10] Update

---
 test/user_interface.jl | 26 +++++++++++++-------------
 1 file changed, 13 insertions(+), 13 deletions(-)

diff --git a/test/user_interface.jl b/test/user_interface.jl
index 4d9a02fbb..a38cabbae 100644
--- a/test/user_interface.jl
+++ b/test/user_interface.jl
@@ -305,7 +305,8 @@ function test_parameterize()
     end
     node = model[2]
     @test length(node.noise_terms) == 3
-    @test JuMP.upper_bound(node.subproblem[:x]) == 1
+    # SDDP is allowed to call `parameterize` during construction
+    # @test JuMP.upper_bound(node.subproblem[:x]) == 1
     node.parameterize(node.noise_terms[2].term)
     @test JuMP.upper_bound(node.subproblem[:x]) == 2
     node.parameterize(3)
@@ -486,20 +487,19 @@ function test_numerical_stability_report()
 end
 
 function test_objective_state()
-    model = SDDP.LinearPolicyGraph(;
-        stages = 2,
-        lower_bound = 0,
-        optimizer = HiGHS.Optimizer,
-    ) do subproblem, t
-        @variable(subproblem, x, SDDP.State, initial_value = 0)
-        SDDP.parameterize(subproblem, [1, 2]) do ω
-            price = SDDP.objective_state(subproblem)
-            @stageobjective(subproblem, price * x.out)
-        end
-    end
     @test_throws(
         ErrorException("No objective state defined."),
-        SDDP.simulate(model, 1; parallel_scheme = SDDP.Serial()),
+        SDDP.LinearPolicyGraph(;
+            stages = 2,
+            lower_bound = 0,
+            optimizer = HiGHS.Optimizer,
+        ) do subproblem, t
+            @variable(subproblem, x, SDDP.State, initial_value = 0)
+            SDDP.parameterize(subproblem, [1, 2]) do ω
+                price = SDDP.objective_state(subproblem)
+                @stageobjective(subproblem, price * x.out)
+            end
+        end,
     )
     @test_throws(
         ErrorException("add_objective_state can only be called once."),

From 229273b8c436c7d16ee3b58c3d60054fb746884b Mon Sep 17 00:00:00 2001
From: odow <o.dowson@gmail.com>
Date: Mon, 20 Jan 2025 15:19:41 +1300
Subject: [PATCH 04/10] Update

---
 docs/src/guides/add_multidimensional_noise.md | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/docs/src/guides/add_multidimensional_noise.md b/docs/src/guides/add_multidimensional_noise.md
index 57d21165b..15da27a5e 100644
--- a/docs/src/guides/add_multidimensional_noise.md
+++ b/docs/src/guides/add_multidimensional_noise.md
@@ -26,11 +26,12 @@ julia> model = SDDP.LinearPolicyGraph(
            SDDP.parameterize(subproblem, Ω, P) do ω
                JuMP.fix(x.out, ω.value)
                @stageobjective(subproblem, ω.coefficient * x.out)
-               println("ω is: ", ω)
            end
        end;
 
-julia> SDDP.simulate(model, 1);
+julia> for s in SDDP.simulate(model, 1)[1]
+           println("ω is: ", s[:noise_term_term])
+       end
 ω is: (value = 1, coefficient = 4)
 ω is: (value = 1, coefficient = 3)
 ω is: (value = 2, coefficient = 4)
@@ -46,7 +47,7 @@ For sampling multidimensional random variates, it can be useful to use the
 There are several ways to go about this. If the sample space is finite, and
 small enough that it makes sense to enumerate each element, we can use
 `Base.product` and `Distributions.support` to generate the entire sample space
-`Ω` from each of the marginal distributions. 
+`Ω` from each of the marginal distributions.
 
 We can then evaluate the density function of the product distribution on each
 element of this space to get the vector of corresponding probabilities, `P`.
@@ -75,11 +76,12 @@ julia> model = SDDP.LinearPolicyGraph(
            SDDP.parameterize(subproblem, Ω, P) do ω
                JuMP.fix(x.out, ω[1])
                @stageobjective(subproblem, ω[2] * x.out + ω[3])
-               println("ω is: ", ω)
            end
        end;
 
-julia> SDDP.simulate(model, 1);
+julia> for s in SDDP.simulate(model, 1)[1]
+           println("ω is: ", s[:noise_term])
+       end
 ω is: [10, 0, 3]
 ω is: [0, 1, 6]
 ω is: [6, 0, 5]
@@ -117,11 +119,12 @@ julia> model = SDDP.LinearPolicyGraph(
            SDDP.parameterize(subproblem, Ω, P) do ω
                JuMP.fix(x.out, ω[1])
                @stageobjective(subproblem, ω[2] * x.out + ω[3])
-               println("ω is: ", ω)
            end
        end;
 
-julia> SDDP.simulate(model, 1);
+julia> for s in SDDP.simulate(model, 1)[1]
+           println("ω is: ", s[:noise_term])
+       end
 ω is: [54, 38, 19]
 ω is: [43, 3, 13]
 ω is: [43, 4, 17]

From a6f0bbceb8a6583a39ea09f4e22c0b63ae15e8ba Mon Sep 17 00:00:00 2001
From: odow <o.dowson@gmail.com>
Date: Mon, 20 Jan 2025 15:43:23 +1300
Subject: [PATCH 05/10] Update

---
 docs/src/guides/add_multidimensional_noise.md         | 2 +-
 docs/src/guides/add_noise_in_the_constraint_matrix.md | 9 ++-------
 2 files changed, 3 insertions(+), 8 deletions(-)

diff --git a/docs/src/guides/add_multidimensional_noise.md b/docs/src/guides/add_multidimensional_noise.md
index 15da27a5e..52f4d8403 100644
--- a/docs/src/guides/add_multidimensional_noise.md
+++ b/docs/src/guides/add_multidimensional_noise.md
@@ -30,7 +30,7 @@ julia> model = SDDP.LinearPolicyGraph(
        end;
 
 julia> for s in SDDP.simulate(model, 1)[1]
-           println("ω is: ", s[:noise_term_term])
+           println("ω is: ", s[:noise_term])
        end
 ω is: (value = 1, coefficient = 4)
 ω is: (value = 1, coefficient = 3)
diff --git a/docs/src/guides/add_noise_in_the_constraint_matrix.md b/docs/src/guides/add_noise_in_the_constraint_matrix.md
index 2bec219fe..6e3dd7646 100644
--- a/docs/src/guides/add_noise_in_the_constraint_matrix.md
+++ b/docs/src/guides/add_noise_in_the_constraint_matrix.md
@@ -15,20 +15,15 @@ julia> model = SDDP.LinearPolicyGraph(
                stages=3, lower_bound = 0, optimizer = HiGHS.Optimizer
                ) do subproblem, t
            @variable(subproblem, x, SDDP.State, initial_value = 0.0)
-           @constraint(subproblem, emissions, 1x.out <= 1)
+           @constraint(subproblem, emissions, 1 * x.out <= 1)
            SDDP.parameterize(subproblem, [0.2, 0.5, 1.0]) do ω
+               ## rewrite 1 * x.out <= 1 to ω * x.out <= 1
                JuMP.set_normalized_coefficient(emissions, x.out, ω)
-               println(emissions)
            end
            @stageobjective(subproblem, -x.out)
        end
 A policy graph with 3 nodes.
  Node indices: 1, 2, 3
-
-julia> SDDP.simulate(model, 1);
-emissions : x_out <= 1
-emissions : 0.2 x_out <= 1
-emissions : 0.5 x_out <= 1
 ```
 
 !!! note

From 1e0c2252df17b8a8665d80a4655ce846dd122932 Mon Sep 17 00:00:00 2001
From: odow <o.dowson@gmail.com>
Date: Wed, 22 Jan 2025 10:54:47 +1300
Subject: [PATCH 06/10] Update

---
 src/plugins/duality_handlers.jl |  4 ++--
 src/user_interface.jl           | 25 +++++++++----------------
 2 files changed, 11 insertions(+), 18 deletions(-)

diff --git a/src/plugins/duality_handlers.jl b/src/plugins/duality_handlers.jl
index bc9e1b3ce..93fc94df1 100644
--- a/src/plugins/duality_handlers.jl
+++ b/src/plugins/duality_handlers.jl
@@ -180,10 +180,10 @@ function get_dual_solution(node::Node, lagrange::LagrangianDuality)
         h_expr[i] = @expression(node.subproblem, state.in - x_in_value[i])
         JuMP.unfix(state.in)
         l, u = node.incoming_state_bounds[key]
-        if l !== nothing
+        if l > -Inf
             JuMP.set_lower_bound(state.in, l)
         end
-        if u !== nothing
+        if u < Inf
             JuMP.set_upper_bound(state.in, u)
         end
         λ_star[i] = conic_dual[key]
diff --git a/src/user_interface.jl b/src/user_interface.jl
index 92f75aa3b..56f107c97 100644
--- a/src/user_interface.jl
+++ b/src/user_interface.jl
@@ -655,7 +655,7 @@ mutable struct Node{T}
     ext::Dict{Symbol,Any}
     # Lock for threading
     lock::ReentrantLock
-    incoming_state_bounds::Dict{Symbol,NTuple{2,Union{Nothing,Float64}}}
+    incoming_state_bounds::Dict{Symbol,NTuple{2,Float64}}
 end
 
 function Base.show(io::IO, node::Node)
@@ -988,7 +988,7 @@ function PolicyGraph(
             # The extension dictionary.
             Dict{Symbol,Any}(),
             ReentrantLock(),
-            Dict{Symbol,NTuple{2,Union{Nothing,Float64}}}(),
+            Dict{Symbol,NTuple{2,Float64}}(),
         )
         subproblem.ext[:sddp_policy_graph] = policy_graph
         policy_graph.nodes[node_index] = subproblem.ext[:sddp_node] = node
@@ -1045,7 +1045,7 @@ end
 
 function _get_incoming_domain(model::PolicyGraph{T}) where {T}
     function _bounds(x)
-        l, u = nothing, nothing
+        l, u = -Inf, Inf
         if has_lower_bound(x)
             l = lower_bound(x)
         end
@@ -1061,10 +1061,8 @@ function _get_incoming_domain(model::PolicyGraph{T}) where {T}
         return l, u
     end
     compare(f, a, b) = f(a, b)
-    compare(f, ::Nothing, b) = b
-    compare(f, a, ::Nothing) = a
-    compare(f, ::Nothing, ::Nothing) = nothing
-    outgoing_bounds = Dict{Tuple{T,Symbol},NTuple{2,Union{Nothing,Float64}}}(
+    compare(::Function, ::Nothing, b) = b
+    outgoing_bounds = Dict{Tuple{T,Symbol},Any}(
         (k, state_name) => (nothing, nothing) for (k, node) in model.nodes
         for (state_name, state) in node.states
     )
@@ -1079,9 +1077,9 @@ function _get_incoming_domain(model::PolicyGraph{T}) where {T}
             end
         end
     end
-    incoming_bounds = Dict{T,Dict{Symbol,NTuple{2,Union{Nothing,Float64}}}}()
+    incoming_bounds = Dict{T,Dict{Symbol,Any}}()
     for (k, node) in model.nodes
-        incoming_bounds[k] = Dict{Symbol,NTuple{2,Union{Nothing,Float64}}}(
+        incoming_bounds[k] = Dict{Symbol,Any}(
             state_name => (nothing, nothing) for
             (state_name, state) in node.states
         )
@@ -1099,13 +1097,8 @@ function _get_incoming_domain(model::PolicyGraph{T}) where {T}
     for (state_name, value) in model.initial_root_state
         for child in model.root_children
             l, u = incoming_bounds[child.term][state_name]
-            if l !== nothing
-                l = compare(min, l, value)
-            end
-            if u !== nothing
-                u = compare(max, u, value)
-            end
-            incoming_bounds[child.term][state_name] = (l, u)
+            incoming_bounds[child.term][state_name] =
+                (compare(min, l, value), compare(max, u, value))
         end
     end
     return incoming_bounds

From 7a2671399ce5f1d8e8d249b29441d5cca6f8346e Mon Sep 17 00:00:00 2001
From: odow <o.dowson@gmail.com>
Date: Wed, 22 Jan 2025 11:19:45 +1300
Subject: [PATCH 07/10] Update

---
 src/user_interface.jl  |  5 ++---
 test/user_interface.jl | 37 ++++++++++++++++++++++++-------------
 2 files changed, 26 insertions(+), 16 deletions(-)

diff --git a/src/user_interface.jl b/src/user_interface.jl
index 56f107c97..0ae80f655 100644
--- a/src/user_interface.jl
+++ b/src/user_interface.jl
@@ -1094,11 +1094,10 @@ function _get_incoming_domain(model::PolicyGraph{T}) where {T}
             end
         end
     end
+    # The incoming state from the root node can be anything
     for (state_name, value) in model.initial_root_state
         for child in model.root_children
-            l, u = incoming_bounds[child.term][state_name]
-            incoming_bounds[child.term][state_name] =
-                (compare(min, l, value), compare(max, u, value))
+            incoming_bounds[child.term][state_name] = (-Inf, Inf)
         end
     end
     return incoming_bounds
diff --git a/test/user_interface.jl b/test/user_interface.jl
index a38cabbae..d2189f3a0 100644
--- a/test/user_interface.jl
+++ b/test/user_interface.jl
@@ -913,14 +913,14 @@ end
 
 function test_incoming_state_bounds()
     graph = SDDP.Graph(0)
-    SDDP.add_node.((graph,), 1:3)
-    SDDP.add_edge(graph, 0 => 1, 0.5)
-    SDDP.add_edge(graph, 0 => 2, 0.5)
+    SDDP.add_node.((graph,), 1:4)
+    SDDP.add_edge(graph, 0 => 4, 1.0)
+    SDDP.add_edge(graph, 4 => 1, 0.5)
+    SDDP.add_edge(graph, 4 => 2, 0.5)
     SDDP.add_edge(graph, 1 => 2, 1.0)
-    SDDP.add_edge(graph, 2 => 2, 0.45)
     SDDP.add_edge(graph, 2 => 3, 0.5)
-    SDDP.add_edge(graph, 3 => 1, 1.0)
-
+    SDDP.add_edge(graph, 3 => 1, 0.5)
+    SDDP.add_edge(graph, 3 => 4, 0.5)
     model = SDDP.PolicyGraph(graph; lower_bound = 0.0) do sp, node
         @variable(sp, x, SDDP.State, initial_value = 0)
         @variable(sp, y, SDDP.State, initial_value = -1)
@@ -933,23 +933,34 @@ function test_incoming_state_bounds()
                 return
             end
         elseif node == 2
+            SDDP.parameterize(sp, 1:2) do w
+                if w == 1
+                    set_upper_bound(y.out, 1.0)
+                elseif w == 2 && has_upper_bound(y.out)
+                    delete_upper_bound(y.out)
+                end
+                return
+            end
         elseif node == 3
             set_lower_bound(x.out, 1)
             set_upper_bound(x.out, 2)
             fix(y.out, 3)
+        elseif node == 4
+            fix(x.out, 0)
+            fix(y.out, -1)
+            fix(z.out, 1)
         end
         @stageobjective(sp, 0)
         return
     end
     @test model[1].incoming_state_bounds ==
-          Dict(:x => (0.0, 2.0), :y => (-1.0, 3.0), :z => (nothing, nothing))
+          Dict(:x => (0.0, 2.0), :y => (-1.0, 3.0), :z => (-Inf, Inf))
     @test model[2].incoming_state_bounds ==
-          Dict(:x => (-2.0, nothing), :y => (nothing, 2.0), :z => (0.0, 1.0))
-    @test model[3].incoming_state_bounds == Dict(
-        :x => (nothing, nothing),
-        :y => (nothing, nothing),
-        :z => (nothing, nothing),
-    )
+          Dict(:x => (-2.0, Inf), :y => (-Inf, 2.0), :z => (0.0, 1.0))
+    @test model[3].incoming_state_bounds ==
+          Dict(:x => (-Inf, Inf), :y => (-Inf, Inf), :z => (-Inf, Inf))
+    @test model[4].incoming_state_bounds ==
+          Dict(:x => (-Inf, Inf), :y => (-Inf, Inf), :z => (-Inf, Inf))
     return
 end
 

From bede536235804e6ca88520c0f0c9b98f3374fc11 Mon Sep 17 00:00:00 2001
From: odow <o.dowson@gmail.com>
Date: Wed, 22 Jan 2025 11:43:34 +1300
Subject: [PATCH 08/10] Update

---
 src/user_interface.jl | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/user_interface.jl b/src/user_interface.jl
index 0ae80f655..354541a2d 100644
--- a/src/user_interface.jl
+++ b/src/user_interface.jl
@@ -1037,6 +1037,9 @@ function PolicyGraph(
     domain = _get_incoming_domain(policy_graph)
     for (node_name, node) in policy_graph.nodes
         for (k, v) in domain[node_name]
+            if v === (nothing, nothing)
+                v = (-Inf, Inf)
+            end
             node.incoming_state_bounds[k] = v
         end
     end

From 375e72ab298abe576e6b000e4e2570d94f3f0566 Mon Sep 17 00:00:00 2001
From: odow <o.dowson@gmail.com>
Date: Wed, 22 Jan 2025 13:37:24 +1300
Subject: [PATCH 09/10] Update

---
 src/plugins/duality_handlers.jl |  8 ++++-
 src/user_interface.jl           | 53 +++++++++++++++++++--------------
 test/user_interface.jl          | 30 +++++++++++++------
 3 files changed, 58 insertions(+), 33 deletions(-)

diff --git a/src/plugins/duality_handlers.jl b/src/plugins/duality_handlers.jl
index 93fc94df1..de5a77c97 100644
--- a/src/plugins/duality_handlers.jl
+++ b/src/plugins/duality_handlers.jl
@@ -179,13 +179,16 @@ function get_dual_solution(node::Node, lagrange::LagrangianDuality)
         x_in_value[i] = JuMP.fix_value(state.in)
         h_expr[i] = @expression(node.subproblem, state.in - x_in_value[i])
         JuMP.unfix(state.in)
-        l, u = node.incoming_state_bounds[key]
+        l, u, is_integer = node.incoming_state_bounds[key]
         if l > -Inf
             JuMP.set_lower_bound(state.in, l)
         end
         if u < Inf
             JuMP.set_upper_bound(state.in, u)
         end
+        if is_integer
+            JuMP.set_integer(state.in)
+        end
         λ_star[i] = conic_dual[key]
     end
     # Check that the conic dual is feasible for the subproblem. Sometimes it
@@ -201,6 +204,9 @@ function get_dual_solution(node::Node, lagrange::LagrangianDuality)
             return L_k === nothing ? nothing : (s * L_k, s * h_k)
         end
     for (i, (_, state)) in enumerate(node.states)
+        if JuMP.is_integer(state.in)
+            JuMP.unset_integer(state.in)
+        end
         JuMP.fix(state.in, x_in_value[i]; force = true)
     end
     λ_solution = Dict{Symbol,Float64}(
diff --git a/src/user_interface.jl b/src/user_interface.jl
index 354541a2d..2993ff706 100644
--- a/src/user_interface.jl
+++ b/src/user_interface.jl
@@ -655,7 +655,8 @@ mutable struct Node{T}
     ext::Dict{Symbol,Any}
     # Lock for threading
     lock::ReentrantLock
-    incoming_state_bounds::Dict{Symbol,NTuple{2,Float64}}
+    # (lower, upper, is_integer)
+    incoming_state_bounds::Dict{Symbol,Tuple{Float64,Float64,Bool}}
 end
 
 function Base.show(io::IO, node::Node)
@@ -988,7 +989,7 @@ function PolicyGraph(
             # The extension dictionary.
             Dict{Symbol,Any}(),
             ReentrantLock(),
-            Dict{Symbol,NTuple{2,Float64}}(),
+            Dict{Symbol,Tuple{Float64,Float64,Bool}}(),
         )
         subproblem.ext[:sddp_policy_graph] = policy_graph
         policy_graph.nodes[node_index] = subproblem.ext[:sddp_node] = node
@@ -1036,11 +1037,9 @@ function PolicyGraph(
     end
     domain = _get_incoming_domain(policy_graph)
     for (node_name, node) in policy_graph.nodes
-        for (k, v) in domain[node_name]
-            if v === (nothing, nothing)
-                v = (-Inf, Inf)
-            end
-            node.incoming_state_bounds[k] = v
+        for (k, (lower, upper, is_integer)) in domain[node_name]
+            node.incoming_state_bounds[k] =
+                (something(lower, -Inf), something(upper, Inf), is_integer)
         end
     end
     return policy_graph
@@ -1058,49 +1057,57 @@ function _get_incoming_domain(model::PolicyGraph{T}) where {T}
         if is_fixed(x)
             l = u = fix_value(x)
         end
+        is_int = is_integer(x)
         if is_binary(x)
             l, u = max(something(l, 0.0), 0.0), min(something(u, 1.0), 1.0)
+            is_int = true
         end
-        return l, u
+        return l, u, is_int
     end
-    compare(f, a, b) = f(a, b)
-    compare(::Function, ::Nothing, b) = b
     outgoing_bounds = Dict{Tuple{T,Symbol},Any}(
-        (k, state_name) => (nothing, nothing) for (k, node) in model.nodes
-        for (state_name, state) in node.states
+        (k, state_name) => nothing for (k, node) in model.nodes
+        for (state_name, _) in node.states
     )
     for (k, node) in model.nodes
         for noise in node.noise_terms
             parameterize(node, noise.term)
             for (state_name, state) in node.states
-                l, u = outgoing_bounds[(k, state_name)]
-                l_new, u_new = _bounds(state.out)
-                outgoing_bounds[(k, state_name)] =
-                    compare(min, l, l_new), compare(max, u, u_new)
+                domain = outgoing_bounds[(k, state_name)]
+                l_new, u_new, is_int_new = _bounds(state.out)
+                outgoing_bounds[(k, state_name)] = if domain === nothing
+                    (l_new, u_new, is_int_new)
+                else
+                    l, u, is_int = domain::Tuple{Float64,Float64,Bool}
+                    (min(l, l_new), max(u, u_new), is_int & is_int_new)
+                end
             end
         end
     end
     incoming_bounds = Dict{T,Dict{Symbol,Any}}()
     for (k, node) in model.nodes
         incoming_bounds[k] = Dict{Symbol,Any}(
-            state_name => (nothing, nothing) for
-            (state_name, state) in node.states
+            state_name => nothing for (state_name, _) in node.states
         )
     end
     for (parent_name, parent) in model.nodes
         for (state_name, state) in parent.states
-            l_new, u_new = outgoing_bounds[(parent_name, state_name)]
+            domain_new = outgoing_bounds[(parent_name, state_name)]
+            l_new, u_new, is_int_new = domain_new::Tuple{Float64,Float64,Bool}
             for child in parent.children
-                l, u = incoming_bounds[child.term][state_name]
-                incoming_bounds[child.term][state_name] =
-                    compare(min, l, l_new), compare(max, u, u_new)
+                domain = incoming_bounds[child.term][state_name]
+                incoming_bounds[child.term][state_name] = if domain === nothing
+                    (l_new, u_new, is_int_new)
+                else
+                    l, u, is_int = domain::Tuple{Float64,Float64,Bool}
+                    (min(l, l_new), max(u, u_new), is_int & is_int_new)
+                end
             end
         end
     end
     # The incoming state from the root node can be anything
     for (state_name, value) in model.initial_root_state
         for child in model.root_children
-            incoming_bounds[child.term][state_name] = (-Inf, Inf)
+            incoming_bounds[child.term][state_name] = (-Inf, Inf, false)
         end
     end
     return incoming_bounds
diff --git a/test/user_interface.jl b/test/user_interface.jl
index d2189f3a0..6e31f565e 100644
--- a/test/user_interface.jl
+++ b/test/user_interface.jl
@@ -922,7 +922,7 @@ function test_incoming_state_bounds()
     SDDP.add_edge(graph, 3 => 1, 0.5)
     SDDP.add_edge(graph, 3 => 4, 0.5)
     model = SDDP.PolicyGraph(graph; lower_bound = 0.0) do sp, node
-        @variable(sp, x, SDDP.State, initial_value = 0)
+        @variable(sp, x, Int, SDDP.State, initial_value = 0)
         @variable(sp, y, SDDP.State, initial_value = -1)
         @variable(sp, z, SDDP.State, initial_value = 1)
         if node == 1
@@ -953,14 +953,26 @@ function test_incoming_state_bounds()
         @stageobjective(sp, 0)
         return
     end
-    @test model[1].incoming_state_bounds ==
-          Dict(:x => (0.0, 2.0), :y => (-1.0, 3.0), :z => (-Inf, Inf))
-    @test model[2].incoming_state_bounds ==
-          Dict(:x => (-2.0, Inf), :y => (-Inf, 2.0), :z => (0.0, 1.0))
-    @test model[3].incoming_state_bounds ==
-          Dict(:x => (-Inf, Inf), :y => (-Inf, Inf), :z => (-Inf, Inf))
-    @test model[4].incoming_state_bounds ==
-          Dict(:x => (-Inf, Inf), :y => (-Inf, Inf), :z => (-Inf, Inf))
+    @test model[1].incoming_state_bounds == Dict(
+        :x => (0.0, 2.0, true),
+        :y => (-1.0, 3.0, false),
+        :z => (-Inf, Inf, false),
+    )
+    @test model[2].incoming_state_bounds == Dict(
+        :x => (-2.0, Inf, true),
+        :y => (-Inf, 2.0, false),
+        :z => (0.0, 1.0, false),
+    )
+    @test model[3].incoming_state_bounds == Dict(
+        :x => (-Inf, Inf, true),
+        :y => (-Inf, Inf, false),
+        :z => (-Inf, Inf, false),
+    )
+    @test model[4].incoming_state_bounds == Dict(
+        :x => (-Inf, Inf, false),
+        :y => (-Inf, Inf, false),
+        :z => (-Inf, Inf, false),
+    )
     return
 end
 

From 2bb0a1a25bcfda11b034d6b56b4d0d4afc370f15 Mon Sep 17 00:00:00 2001
From: odow <o.dowson@gmail.com>
Date: Wed, 22 Jan 2025 13:51:11 +1300
Subject: [PATCH 10/10] Update

---
 src/user_interface.jl | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/src/user_interface.jl b/src/user_interface.jl
index 2993ff706..ac81073a1 100644
--- a/src/user_interface.jl
+++ b/src/user_interface.jl
@@ -1037,9 +1037,8 @@ function PolicyGraph(
     end
     domain = _get_incoming_domain(policy_graph)
     for (node_name, node) in policy_graph.nodes
-        for (k, (lower, upper, is_integer)) in domain[node_name]
-            node.incoming_state_bounds[k] =
-                (something(lower, -Inf), something(upper, Inf), is_integer)
+        for (k, v) in domain[node_name]
+            node.incoming_state_bounds[k] = something(v, (-Inf, Inf, false))
         end
     end
     return policy_graph
@@ -1065,8 +1064,8 @@ function _get_incoming_domain(model::PolicyGraph{T}) where {T}
         return l, u, is_int
     end
     outgoing_bounds = Dict{Tuple{T,Symbol},Any}(
-        (k, state_name) => nothing for (k, node) in model.nodes
-        for (state_name, _) in node.states
+        (k, state_name) => nothing for (k, node) in model.nodes for
+        (state_name, _) in node.states
     )
     for (k, node) in model.nodes
         for noise in node.noise_terms