Update storage variable creation using indices table (#904)

src/create-model.jl

@@ -71,13 +71,6 @@ function create_model(
             DataFrames.groupby(dataframes[:storage_level_inter_rp], [:asset, :year])
     end
 
-    # Unpacking dataframes with variable indices
-    @timeit to "unpacking variable indices" begin
-        storage_level_intra_rp_indices = variables[:storage_level_intra_rp].indices
-        storage_level_inter_rp_indices = variables[:storage_level_inter_rp].indices
-        is_charging_indices = variables[:is_charging].indices
-    end
-
     ## Model
     model = JuMP.Model()
 
@@ -89,14 +82,7 @@ function create_model(
         sets,
         variables,
     )
-    @timeit to "add_storage_variables!" add_storage_variables!(
-        model,
-        graph,
-        sets,
-        storage_level_intra_rp_indices,
-        storage_level_inter_rp_indices,
-        is_charging_indices,
-    )
+    @timeit to "add_storage_variables!" add_storage_variables!(model, graph, sets, variables)
 
     # TODO: This should change heavily, so I just moved things to the function and unpack them here from model
     assets_decommission_compact_method = model[:assets_decommission_compact_method]
@@ -106,12 +92,14 @@ function create_model(
     assets_investment_energy = model[:assets_investment_energy]
     flows_decommission_using_simple_method = model[:flows_decommission_using_simple_method]
     flows_investment = model[:flows_investment]
-    storage_level_inter_rp = model[:storage_level_inter_rp]
-    storage_level_intra_rp = model[:storage_level_intra_rp]
 
     # TODO: This should disapear after the changes on add_expressions_to_dataframe! and storing the solution
     model[:flow] = df_flows.flow = variables[:flow].container
     model[:units_on] = df_units_on.units_on = variables[:units_on].container
+    storage_level_intra_rp =
+        model[:storage_level_intra_rp] = variables[:storage_level_intra_rp].container
+    storage_level_inter_rp =
+        model[:storage_level_inter_rp] = variables[:storage_level_inter_rp].container
 
     ## Add expressions to dataframes
     # TODO: What will improve this? Variables (#884)?, Constraints?
@@ -128,12 +116,12 @@ function create_model(
         outgoing_flow_storage_inter_rp_balance,
     ) = add_expressions_to_dataframe!(
         dataframes,
+        variables,
         model,
         expression_workspace,
         representative_periods,
         timeframe,
         graph,
-        is_charging_indices,
     )
 
     ## Expressions for multi-year investment

src/model-preparation.jl

@@ -248,6 +248,7 @@ end
 """
     add_expression_is_charging_terms_intra_rp_constraints!(df_cons,
                                                        is_charging_indices,
+                                                       is_charging_variables,
                                                        workspace
                                                        )
 
@@ -263,6 +264,7 @@ This strategy is based on the replies in this discourse thread:
 function add_expression_is_charging_terms_intra_rp_constraints!(
     df_cons,
     is_charging_indices,
+    is_charging_variables,
     workspace,
 )
     # Aggregating function: We have to compute the proportion of each variable is_charging in the constraint timesteps_block.
@@ -284,7 +286,7 @@ function add_expression_is_charging_terms_intra_rp_constraints!(
         for row in eachrow(grouped_is_charging[(year, rep_period, asset)])
             asset = row[:asset]
             for t in row.timesteps_block
-                JuMP.add_to_expression!(workspace[t], row.is_charging)
+                JuMP.add_to_expression!(workspace[t], is_charging_variables[row.index])
             end
         end
         # Apply the agg funtion to the corresponding variables from the workspace
@@ -433,14 +435,17 @@ end
 
 function add_expressions_to_dataframe!(
     dataframes,
+    variables,
     model,
     expression_workspace,
     representative_periods,
     timeframe,
     graph,
-    is_charging_indices,
 )
     @timeit to "add_expression_terms_to_df" begin
+        # Unpack variables
+        is_charging_indices = variables[:is_charging].indices
+        is_charging_variables = variables[:is_charging].container
 
         # Creating the incoming and outgoing flow expressions
         add_expression_terms_intra_rp_constraints!(
@@ -526,11 +531,13 @@ function add_expressions_to_dataframe!(
         add_expression_is_charging_terms_intra_rp_constraints!(
             dataframes[:highest_in],
             is_charging_indices,
+            is_charging_variables,
             expression_workspace,
         )
         add_expression_is_charging_terms_intra_rp_constraints!(
             dataframes[:highest_out],
             is_charging_indices,
+            is_charging_variables,
             expression_workspace,
         )
         if !isempty(dataframes[:units_on_and_outflows])

src/variables/storage.jl

@@ -7,39 +7,35 @@ Adds storage-related variables to the optimization `model`, including storage le
 The function also optionally sets binary constraints for certain charging variables based on storage methods.
 
 """
-function add_storage_variables!(
-    model,
-    graph,
-    sets,
-    storage_level_intra_rp_indices,
-    storage_level_inter_rp_indices,
-    is_charging_indices,
-)
-    model[:storage_level_intra_rp] = [
+function add_storage_variables!(model, graph, sets, variables)
+    storage_level_intra_rp_indices = variables[:storage_level_intra_rp].indices
+    storage_level_inter_rp_indices = variables[:storage_level_inter_rp].indices
+    is_charging_indices = variables[:is_charging].indices
+
+    variables[:storage_level_intra_rp].container = [
         @variable(
             model,
             lower_bound = 0.0,
             base_name = "storage_level_intra_rp[$(row.asset),$(row.year),$(row.rep_period),$(row.timesteps_block)]"
         ) for row in eachrow(storage_level_intra_rp_indices)
     ]
 
-    model[:storage_level_inter_rp] = [
+    variables[:storage_level_inter_rp].container = [
         @variable(
             model,
             lower_bound = 0.0,
             base_name = "storage_level_inter_rp[$(row.asset),$(row.year),$(row.periods_block)]"
         ) for row in eachrow(storage_level_inter_rp_indices)
     ]
 
-    model[:is_charging] =
-        is_charging_indices.is_charging = [
-            @variable(
-                model,
-                lower_bound = 0.0,
-                upper_bound = 1.0,
-                base_name = "is_charging[$(row.asset),$(row.year),$(row.rep_period),$(row.timesteps_block)]"
-            ) for row in eachrow(is_charging_indices)
-        ]
+    variables[:is_charging].container = [
+        @variable(
+            model,
+            lower_bound = 0.0,
+            upper_bound = 1.0,
+            base_name = "is_charging[$(row.asset),$(row.year),$(row.rep_period),$(row.timesteps_block)]"
+        ) for row in eachrow(is_charging_indices)
+    ]
 
     ### Binary Charging Variables
     is_charging_indices.use_binary_storage_method = [
@@ -54,7 +50,7 @@ function add_storage_variables!(
     )
 
     for row in eachrow(sub_df_is_charging_indices)
-        JuMP.set_binary(model[:is_charging][row.index])
+        JuMP.set_binary(variables[:is_charging].container[row.index])
     end
 
     return