Refactor variables to use TulipaVariable (#925)

docs/src/20-tutorials.md

@@ -121,7 +121,7 @@ We need the sets and the variables indices.
 
 ```@example manual
 sets = create_sets(graph, years)
-variables = compute_variables_indices(dataframes)
+variables = compute_variables_indices(connection, dataframes)
 ```
 
 Now we can compute the model.
@@ -133,13 +133,13 @@ model = create_model(graph, sets, variables, representative_periods, dataframes,
 Finally, we can compute the solution.
 
 ```@example manual
-solution = solve_model(model)
+solution = solve_model(model, variables)
 ```
 
 or, if we want to store the `flow`, `storage_level_intra_rp`, and `storage_level_inter_rp` optimal value in the dataframes:
 
 ```@example manual
-solution = solve_model!(dataframes, model)
+solution = solve_model!(dataframes, model, variables)
 ```
 
 This `solution` structure is the same as the one returned when using an `EnergyProblem`.
@@ -174,7 +174,7 @@ or
 ```@example manual
 using GLPK
 
-solution = solve_model(model, GLPK.Optimizer)
+solution = solve_model(model, variables, GLPK.Optimizer)
 ```
 
 Notice that, in any of these cases, we need to explicitly add the GLPK package

src/constraints/group.jl

@@ -5,11 +5,13 @@ export add_group_constraints!
 
 Adds group constraints for assets that share a common limits or bounds
 """
-function add_group_constraints!(model, graph, sets, assets_investment, groups)
+function add_group_constraints!(model, variables, graph, sets, groups)
     # unpack from sets
     Ai = sets[:Ai]
     Y = sets[:Y]
 
+    assets_investment = variables[:assets_investment].lookup
+
     # - Group constraints for investments at each year
     assets_at_year_in_group = Dict(
         group => (

src/constraints/investment.jl

@@ -1,24 +1,25 @@
 export add_investment_constraints!
 
 """
-add_investment_constraints!(graph, Ai, Ase, Fi, assets_investment, assets_investment_energy, flows_investment)
+    add_investment_constraints!(graph, Ai, Ase, Fi, assets_investment, assets_investment_energy, flows_investment)
 
 Adds the investment constraints for all asset types and transport flows to the model
 """
-
-function add_investment_constraints!(
-    graph,
-    sets,
-    assets_investment,
-    assets_investment_energy,
-    flows_investment,
-)
+function add_investment_constraints!(graph, sets, variables)
+    # TODO: Since this function is defining bound constraints, it doesn't need the `model`
+    # When we refactor the signatures to look the same, we should consider naming it differently
+    # TODO: Verify if it's possible and reasonable to move the bound definition to when the
+    # indices are created
     # unpack from sets
     Ai = sets[:Ai]
     Ase = sets[:Ase]
     Fi = sets[:Fi]
     Y = sets[:Y]
 
+    assets_investment = variables[:assets_investment].lookup
+    assets_investment_energy = variables[:assets_investment_energy].lookup
+    flows_investment = variables[:flows_investment].lookup
+
     # - Maximum (i.e., potential) investment limit for assets
     for y in Y, a in Ai[y]
         if graph[a].capacity > 0 && !ismissing(graph[a].investment_limit[y])

src/create-model.jl

@@ -76,23 +76,14 @@ function create_model(
 
     ## Variables
     @timeit to "add_flow_variables!" add_flow_variables!(model, variables)
-    @timeit to "add_investment_variables!" add_investment_variables!(model, graph, sets)
+    @timeit to "add_investment_variables!" add_investment_variables!(model, graph, sets, variables)
     @timeit to "add_unit_commitment_variables!" add_unit_commitment_variables!(
         model,
         sets,
         variables,
     )
     @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]
-    assets_decommission_simple_method = model[:assets_decommission_simple_method]
-    assets_decommission_energy_simple_method = model[:assets_decommission_energy_simple_method]
-    assets_investment = model[:assets_investment]
-    assets_investment_energy = model[:assets_investment_energy]
-    flows_decommission_using_simple_method = model[:flows_decommission_using_simple_method]
-    flows_investment = model[:flows_investment]
-
     # 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
@@ -125,7 +116,7 @@ function create_model(
     )
 
     ## Expressions for multi-year investment
-    create_multi_year_expressions!(model, graph, sets)
+    create_multi_year_expressions!(model, graph, sets, variables)
     accumulated_flows_export_units = model[:accumulated_flows_export_units]
     accumulated_flows_import_units = model[:accumulated_flows_import_units]
     accumulated_initial_units = model[:accumulated_initial_units]
@@ -136,12 +127,20 @@ function create_model(
     accumulated_units_simple_method = model[:accumulated_units_simple_method]
 
     ## Expressions for storage assets
-    add_storage_expressions!(model, graph, sets)
+    add_storage_expressions!(model, graph, sets, variables)
     accumulated_energy_units_simple_method = model[:accumulated_energy_units_simple_method]
     accumulated_energy_capacity = model[:accumulated_energy_capacity]
 
     ## Expressions for the objective function
-    add_objective!(model, graph, dataframes, representative_periods, sets, model_parameters)
+    add_objective!(
+        model,
+        variables,
+        graph,
+        dataframes,
+        representative_periods,
+        sets,
+        model_parameters,
+    )
 
     # TODO: Pass sets instead of the explicit values
     ## Constraints
@@ -210,20 +209,14 @@ function create_model(
         accumulated_flows_import_units,
     )
 
-    @timeit to "add_investment_constraints!" add_investment_constraints!(
-        graph,
-        sets,
-        assets_investment,
-        assets_investment_energy,
-        flows_investment,
-    )
+    @timeit to "add_investment_constraints!" add_investment_constraints!(graph, sets, variables)
 
     if !isempty(groups)
         @timeit to "add_group_constraints!" add_group_constraints!(
             model,
+            variables,
             graph,
             sets,
-            assets_investment,
             groups,
         )
     end

src/expressions/multi-year.jl

@@ -1,11 +1,12 @@
-function create_multi_year_expressions!(model, graph, sets)
+function create_multi_year_expressions!(model, graph, sets, variables)
     @timeit to "multi-year investment expressions" begin
         # Unpacking
-        assets_investment = model[:assets_investment]
-        assets_decommission_simple_method = model[:assets_decommission_simple_method]
-        assets_decommission_compact_method = model[:assets_decommission_compact_method]
-        flows_investment = model[:flows_investment]
-        flows_decommission_using_simple_method = model[:flows_decommission_using_simple_method]
+        assets_investment = variables[:assets_investment].lookup
+        assets_decommission_simple_method = variables[:assets_decommission_simple_method].lookup
+        assets_decommission_compact_method = variables[:assets_decommission_compact_method].lookup
+        flows_investment = variables[:flows_investment].lookup
+        flows_decommission_using_simple_method =
+            variables[:flows_decommission_using_simple_method].lookup
 
         accumulated_initial_units = @expression(
             model,

src/expressions/storage.jl

@@ -1,6 +1,7 @@
-function add_storage_expressions!(model, graph, sets)
-    assets_investment_energy = model[:assets_investment_energy]
-    assets_decommission_energy_simple_method = model[:assets_decommission_energy_simple_method]
+function add_storage_expressions!(model, graph, sets, variables)
+    assets_investment_energy = variables[:assets_investment_energy].lookup
+    assets_decommission_energy_simple_method =
+        variables[:assets_decommission_energy_simple_method].lookup
     accumulated_investment_units_using_simple_method =
         model[:accumulated_investment_units_using_simple_method]
     accumulated_decommission_units_using_simple_method =

src/objective.jl

@@ -1,10 +1,18 @@
-function add_objective!(model, graph, dataframes, representative_periods, sets, model_parameters)
-    assets_investment = model[:assets_investment]
+function add_objective!(
+    model,
+    variables,
+    graph,
+    dataframes,
+    representative_periods,
+    sets,
+    model_parameters,
+)
+    assets_investment = variables[:assets_investment].lookup
     accumulated_units_simple_method = model[:accumulated_units_simple_method]
     accumulated_units_compact_method = model[:accumulated_units_compact_method]
-    assets_investment_energy = model[:assets_investment_energy]
+    assets_investment_energy = variables[:assets_investment_energy].lookup
     accumulated_energy_units_simple_method = model[:accumulated_energy_units_simple_method]
-    flows_investment = model[:flows_investment]
+    flows_investment = variables[:flows_investment].lookup
     accumulated_flows_export_units = model[:accumulated_flows_export_units]
     accumulated_flows_import_units = model[:accumulated_flows_import_units]
 

src/solve-model.jl

@@ -17,8 +17,13 @@ function solve_model!(
             error("Model is not created, run create_model(energy_problem) first.")
         end
 
-        energy_problem.solution =
-            solve_model!(energy_problem.dataframes, model, optimizer; parameters = parameters)
+        energy_problem.solution = solve_model!(
+            energy_problem.dataframes,
+            model,
+            energy_problem.variables,
+            optimizer;
+            parameters = parameters,
+        )
         energy_problem.termination_status = JuMP.termination_status(model)
         if energy_problem.solution === nothing
             # Warning has been given at internal function
@@ -34,7 +39,7 @@ function solve_model!(
 
         for ((y, a), value) in energy_problem.solution.assets_investment_energy
             graph[a].investment_energy[y] =
-                graph[a].investment_integer_storage_energy[y] ? round(Int, value) : value
+                graph[a].investment_integer_storage_energy ? round(Int, value) : value
         end
 
         for row in eachrow(energy_problem.dataframes[:storage_level_intra_rp])
@@ -158,11 +163,12 @@ The `solution` object is a mutable struct with the following fields:
 
 ```julia
 parameters = Dict{String,Any}("presolve" => "on", "time_limit" => 60.0, "output_flag" => true)
-solution = solve_model(model, HiGHS.Optimizer; parameters = parameters)
+solution = solve_model(model, variables, HiGHS.Optimizer; parameters = parameters)
 ```
 """
 function solve_model(
     model::JuMP.Model,
+    variables,
     optimizer = HiGHS.Optimizer;
     parameters = default_parameters(optimizer),
 )
@@ -181,9 +187,9 @@ function solve_model(
     end
 
     return Solution(
-        JuMP.value.(model[:assets_investment]).data,   # .data returns a OrderedDict since variable is SparseAxisArray
-        JuMP.value.(model[:assets_investment_energy]).data,
-        JuMP.value.(model[:flows_investment]).data,
+        Dict(k => JuMP.value(v) for (k, v) in variables[:assets_investment].lookup),
+        Dict(k => JuMP.value(v) for (k, v) in variables[:assets_investment_energy].lookup),
+        Dict(k => JuMP.value(v) for (k, v) in variables[:flows_investment].lookup),
         JuMP.value.(model[:storage_level_intra_rp]),
         JuMP.value.(model[:storage_level_inter_rp]),
         JuMP.value.(model[:max_energy_inter_rp]),

src/structures.jl

@@ -38,9 +38,10 @@ Structure to hold the JuMP variables for the TulipaEnergyModel
 mutable struct TulipaVariable
     indices::DataFrame
     container::Vector{JuMP.VariableRef}
+    lookup::OrderedDict # TODO: This is probably not type stable so it's only used for strangling
 
-    function TulipaVariable(indices, container)
-        return new(indices, container)
+    function TulipaVariable(indices, container = JuMP.VariableRef[])
+        return new(indices, container, Dict())
     end
 end
 
@@ -218,7 +219,7 @@ end
 mutable struct Solution
     assets_investment::Dict{Tuple{Int,String},Float64}
     assets_investment_energy::Dict{Tuple{Int,String},Float64} # for storage assets with energy method
-    flows_investment::Dict{Tuple{Int,Tuple{String,String}},Float64}
+    flows_investment::Any # TODO: Fix this type
     storage_level_intra_rp::Vector{Float64}
     storage_level_inter_rp::Vector{Float64}
     max_energy_inter_rp::Vector{Float64}
@@ -307,7 +308,7 @@ mutable struct EnergyProblem
         end
 
         elapsed_time_vars = @elapsed begin
-            variables = compute_variables_indices(dataframes)
+            variables = compute_variables_indices(connection, dataframes)
         end
 
         energy_problem = new(

src/time-resolution.jl

@@ -1,4 +1,4 @@
-export compute_rp_partition, compute_constraints_partitions, compute_variables_indices
+export compute_rp_partition, compute_constraints_partitions
 
 using SparseArrays
 
@@ -259,17 +259,3 @@ function compute_rp_partition(
     end
     return rp_partition
 end
-
-function compute_variables_indices(dataframes)
-    variables = Dict(
-        :flow => TulipaVariable(dataframes[:flows], Vector()),
-        :units_on => TulipaVariable(dataframes[:units_on], Vector()),
-        :storage_level_intra_rp =>
-            TulipaVariable(dataframes[:storage_level_intra_rp], Vector()),
-        :storage_level_inter_rp =>
-            TulipaVariable(dataframes[:storage_level_inter_rp], Vector()),
-        :is_charging => TulipaVariable(dataframes[:lowest_in_out], Vector()),
-    )
-
-    return variables
-end