From e5d1072b7fb067c17abe09d21f899fb193960c0d Mon Sep 17 00:00:00 2001
From: Lane Smith <lane.smith@breakthroughenergy.org>
Date: Mon, 5 Apr 2021 19:48:55 -0700
Subject: [PATCH] fix: update load_shed and demand flexibility constraints each
 iteration

---
 src/loop.jl  | 19 +++++++++++++++++--
 src/model.jl | 21 ++++++++++++---------
 src/types.jl |  1 +
 3 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/src/loop.jl b/src/loop.jl
index a76f2a7b..7e25dae7 100644
--- a/src/loop.jl
+++ b/src/loop.jl
@@ -53,6 +53,10 @@ function interval_loop(factory_like, model_kwargs::Dict,
         interval_start = start_index + (i - 1) * interval
         interval_end = interval_start + interval - 1
         model_kwargs["start_index"] = interval_start
+        bus_flex_amt = _make_bus_flexibility_amount(
+            case, flexibility, interval_start, interval_end
+        )
+        flexibility_enabled = (bus_flex_amt != zeros(sets.num_bus, interval))
         if i == 1
             # Build a model with no initial ramp constraint
             if storage_enabled
@@ -87,8 +91,9 @@ function interval_loop(factory_like, model_kwargs::Dict,
             if (("load_shed_enabled" in keys(model_kwargs))
                 && (model_kwargs["load_shed_enabled"] == true))
                 for t in 1:interval, i in 1:length(sets.load_bus_idx)
-                    JuMP.set_upper_bound(voi.load_shed[i, t],
-                                         bus_demand[sets.load_bus_idx[i], t])
+                    JuMP.set_normalized_rhs(
+                        voi.load_shed_ub[i, t], bus_demand[sets.load_bus_idx[i], t]
+                    )
                 end
             end
             for t in 1:interval, g in 1:sets.num_hydro
@@ -116,6 +121,16 @@ function interval_loop(factory_like, model_kwargs::Dict,
                     JuMP.set_normalized_rhs(voi.initial_soc[s], storage_e0[s])
                 end
             end
+            if flexibility_enabled
+                for t in 1:interval, i in 1:length(sets.load_bus_idx)
+                    JuMP.set_upper_bound(
+                        voi.load_shift_up[i, t], bus_flex_amt[sets.load_bus_idx[i], t]
+                    )
+                    JuMP.set_upper_bound(
+                        voi.load_shift_dn[i, t], bus_flex_amt[sets.load_bus_idx[i], t]
+                    )
+                end
+            end
         end
 
         while true
diff --git a/src/model.jl b/src/model.jl
index 9d4d0779..ee97f82c 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -277,9 +277,11 @@ function _build_model(
         theta[sets.bus_idx, hour_idx], (container=Array)
     end)
     if load_shed_enabled
-        JuMP.@variable(m,
-            load_shed[i in 1:sets.num_load_bus, j in 1:interval_length],
-            container=Array)
+        JuMP.@variable(
+            m,
+            load_shed[i in 1:sets.num_load_bus, j in 1:interval_length] >= 0,
+            container=Array
+        )
     end
     if trans_viol_enabled
         JuMP.@variable(m,
@@ -344,15 +346,15 @@ function _build_model(
             JuMP.@constraint(
                 m, 
                 load_shed_ub[i in 1:sets.num_load_bus, j in 1:interval_length], 
-                0 <= load_shed[i, j] 
-                    <= bus_demand[sets.load_bus_idx[i], j] + load_shift_up[i, j] 
-                        - load_shift_dn[i, j]
+                load_shed[i, j] <= bus_demand[sets.load_bus_idx[i], j] 
+                    + load_shift_up[i, j]
+                    - load_shift_dn[i, j]
             )
         else
             JuMP.@constraint(
                 m, 
                 load_shed_ub[i in 1:sets.num_load_bus, j in 1:interval_length], 
-                0 <= load_shed[i, j] <= bus_demand[sets.load_bus_idx[i], j]
+                load_shed[i, j] <= bus_demand[sets.load_bus_idx[i], j]
             )
         end
     end
@@ -523,14 +525,15 @@ function _build_model(
     initial_soc = storage_enabled ? initial_soc : nothing
     initial_rampup = initial_ramp_enabled ? initial_rampup : nothing
     initial_rampdown = initial_ramp_enabled ? initial_rampdown : nothing
+    load_shed_ub = load_shed_enabled ? load_shed_ub : nothing
     voi = VariablesOfInterest(;
         # Variables
         pg=pg, pf=pf, 
         load_shed=load_shed, load_shift_up=load_shift_up, load_shift_dn=load_shift_dn, 
         storage_soc=storage_soc, storage_dis=storage_dis, storage_chg=storage_chg,
         # Constraints
-        branch_min=branch_min, branch_max=branch_max,
-        powerbalance=powerbalance, initial_soc=initial_soc,
+        branch_min=branch_min, branch_max=branch_max, powerbalance=powerbalance,
+        initial_soc=initial_soc, load_shed_ub=load_shed_ub, 
         initial_rampup=initial_rampup, initial_rampdown=initial_rampdown,
         hydro_fixed=hydro_fixed, solar_max=solar_max, wind_max=wind_max)
     return (m, voi)
diff --git a/src/types.jl b/src/types.jl
index 7d8e1f33..e29453eb 100644
--- a/src/types.jl
+++ b/src/types.jl
@@ -70,6 +70,7 @@ Base.@kwdef struct VariablesOfInterest
     pg::Array{JuMP.VariableRef,2}
     pf::Array{JuMP.VariableRef,2}
     load_shed::Union{Array{JuMP.VariableRef,2},Nothing}
+    load_shed_ub::Union{JuMP.Containers.DenseAxisArray,Nothing}
     load_shift_up::Union{Array{JuMP.VariableRef,2},Nothing}
     load_shift_dn::Union{Array{JuMP.VariableRef,2},Nothing}
     powerbalance::Array{JuMP.ConstraintRef,2}