From 4f0c48e29c9f3817fb6e382210e4beb49a9e0f71 Mon Sep 17 00:00:00 2001
From: Jose Daniel Lara <jdlara@berkeley.edu>
Date: Mon, 26 Feb 2024 22:05:04 -0700
Subject: [PATCH 01/17] add new methods and abstraction for MPtemplate

---
 src/PowerSimulationsDecomposition.jl |  1 +
 src/multiproblem_template.jl         | 17 ++++++++++++++++-
 src/print.jl                         |  7 +++++++
 3 files changed, 24 insertions(+), 1 deletion(-)
 create mode 100644 src/print.jl

diff --git a/src/PowerSimulationsDecomposition.jl b/src/PowerSimulationsDecomposition.jl
index 6d24947..176ba4e 100644
--- a/src/PowerSimulationsDecomposition.jl
+++ b/src/PowerSimulationsDecomposition.jl
@@ -29,5 +29,6 @@ include("multi_optimization_container.jl")
 include("algorithms/sequential_algorithm.jl")
 include("algorithms/mpi_parallel_algorithm.jl")
 include("problems/multi_region_problem.jl")
+include("print.jl")
 
 end
diff --git a/src/multiproblem_template.jl b/src/multiproblem_template.jl
index 222cf36..0e2f19a 100644
--- a/src/multiproblem_template.jl
+++ b/src/multiproblem_template.jl
@@ -1,8 +1,17 @@
-struct MultiProblemTemplate
+struct MultiProblemTemplate <: PSI.AbstractProblemTemplate
     base_template::PSI.ProblemTemplate
     sub_templates::Dict{String, PSI.ProblemTemplate}
 end
 
+function Base.isempty(template::MultiProblemTemplate)
+    for template in values(template.sub_templates)
+        if !isempty(template.sub_templates)
+            return false
+        end
+    end
+    return isempty(template.base_template)
+end
+
 function MultiProblemTemplate(
     base_template::PSI.ProblemTemplate,
     problem_keys::Vector{String},
@@ -142,3 +151,9 @@ function PSI.set_service_model!(
     end
     return
 end
+
+function PSI.finalize_template!(template::MultiProblemTemplate, sys::PSY.System)
+    for sub_template in get_sub_templates(template)
+        PSI.finalize_template!(sub_template, sys)
+    end
+end
diff --git a/src/print.jl b/src/print.jl
new file mode 100644
index 0000000..e1de52b
--- /dev/null
+++ b/src/print.jl
@@ -0,0 +1,7 @@
+function Base.show(io::IO, ::MIME"text/plain", input::MultiProblemTemplate)
+    println(io, "Print somenthing clever here. Template")
+end
+
+function Base.show(io::IO, ::MIME"text/plain", input::PSI.DecisionModel{MultiRegionProblem})
+    println(io, "Print somenthing clever here. Problem")
+end

From 9d9840ac5553bbd6fe44fee3c652fab429304ee5 Mon Sep 17 00:00:00 2001
From: Jose Daniel Lara <jdlara@berkeley.edu>
Date: Mon, 26 Feb 2024 22:36:07 -0700
Subject: [PATCH 02/17] code clean up

---
 .github/workflows/format-check.yml     |  2 +-
 docs/src/index.md                      |  3 +-
 src/algorithms/sequential_algorithm.jl |  3 +-
 src/multi_decision_model.jl            |  2 +-
 src/multiproblem_template.jl           | 44 ++++++++++------
 src/problems/multi_region_problem.jl   | 69 +++++++++++++++-----------
 test/Project.toml                      |  3 +-
 test/first_test.jl                     | 45 ++++++++---------
 test/runtests.jl                       |  1 -
 test/test_core.jl                      |  1 -
 10 files changed, 97 insertions(+), 76 deletions(-)

diff --git a/.github/workflows/format-check.yml b/.github/workflows/format-check.yml
index 75786f5..cf2bed0 100644
--- a/.github/workflows/format-check.yml
+++ b/.github/workflows/format-check.yml
@@ -3,7 +3,7 @@ name: Format Check
 on:
   push:
     branches:
-      - 'master'
+      - 'main'
       - 'release-'
     tags: '*'
   pull_request:
diff --git a/docs/src/index.md b/docs/src/index.md
index 441f586..ec468db 100644
--- a/docs/src/index.md
+++ b/docs/src/index.md
@@ -8,7 +8,8 @@ CurrentModule = PowerSimulationsDecomposition
 
 `PowerSimulationsDecomposition .jl` is a [`Julia`](http://www.julialang.org) package that provides blah blah
 
-------------
+* * *
+
 PowerSimulationsDecomposition  has been developed as part of the Scalable Integrated Infrastructure Planning
 (SIIP) initiative at the U.S. Department of Energy's National Renewable Energy
 Laboratory ([NREL](https://www.nrel.gov/)).
diff --git a/src/algorithms/sequential_algorithm.jl b/src/algorithms/sequential_algorithm.jl
index 3907c49..cdd9feb 100644
--- a/src/algorithms/sequential_algorithm.jl
+++ b/src/algorithms/sequential_algorithm.jl
@@ -23,8 +23,7 @@ function build_main_problem!(
     sys::PSY.System,
 ) end
 
-function write_results_to_main_container(container::MultiOptimizationContainer)
-end
+function write_results_to_main_container(container::MultiOptimizationContainer) end
 
 function solve_impl!(
     container::MultiOptimizationContainer{SequentialAlgorithm},
diff --git a/src/multi_decision_model.jl b/src/multi_decision_model.jl
index 248a500..a1d6d6d 100644
--- a/src/multi_decision_model.jl
+++ b/src/multi_decision_model.jl
@@ -1 +1 @@
-# ! needs the definition of decision problem that can get type MultiProblemTemplate
\ No newline at end of file
+# ! needs the definition of decision problem that can get type MultiProblemTemplate
diff --git a/src/multiproblem_template.jl b/src/multiproblem_template.jl
index 0e2f19a..d6cca2c 100644
--- a/src/multiproblem_template.jl
+++ b/src/multiproblem_template.jl
@@ -3,29 +3,37 @@ struct MultiProblemTemplate <: PSI.AbstractProblemTemplate
     sub_templates::Dict{String, PSI.ProblemTemplate}
 end
 
-function Base.isempty(template::MultiProblemTemplate)
-    for template in values(template.sub_templates)
-        if !isempty(template.sub_templates)
-            return false
-        end
-    end
-    return isempty(template.base_template)
-end
-
 function MultiProblemTemplate(
     base_template::PSI.ProblemTemplate,
     problem_keys::Vector{String},
 )
-    sub_templates = Dict{String, PSI.ProblemTemplate}(k => deepcopy(base_template) for k in problem_keys)
+    sub_templates = Dict{String, PSI.ProblemTemplate}(
+        k => deepcopy(base_template) for k in problem_keys
+    )
     return MultiProblemTemplate(base_template, sub_templates)
 end
 
 function MultiProblemTemplate(
     network::PSI.NetworkModel{T},
-    problem_keys::Vector{String},) where {T <: PM.AbstractPowerModel}
+    problem_keys::Vector{String},
+) where {T <: PM.AbstractPowerModel}
     return MultiProblemTemplate(PSI.ProblemTemplate(network), problem_keys)
 end
 
+function Base.isempty(template::MultiProblemTemplate)
+    for template in values(template.sub_templates)
+        if !isempty(template.sub_templates)
+            return false
+        end
+    end
+    return isempty(template.base_template)
+end
+
+function PSI.get_network_formulation(template::MultiProblemTemplate)
+    bt = template.base_template
+    return PSI.get_network_formulation(PSI.get_network_model(bt))
+end
+
 function get_sub_templates(template::MultiProblemTemplate)
     return values(template.sub_templates)
 end
@@ -53,7 +61,10 @@ function PSI.set_device_model!(
     component_type::Type{<:PSY.Device},
     formulation::Type{<:PSI.AbstractDeviceFormulation},
 )
-    PSI.set_device_model!(template.base_template, PSI.DeviceModel(component_type, formulation))
+    PSI.set_device_model!(
+        template.base_template,
+        PSI.DeviceModel(component_type, formulation),
+    )
     for sub_template in get_sub_templates(template)
         PSI.set_device_model!(sub_template, PSI.DeviceModel(component_type, formulation))
     end
@@ -98,13 +109,13 @@ function PSI.set_service_model!(
     PSI.set_service_model!(
         template.base_template,
         service_name,
-        ServiceModel(service_type, formulation; use_service_name = true),
+        ServiceModel(service_type, formulation; use_service_name=true),
     )
     for sub_template in get_sub_templates(template)
         PSI.set_service_model!(
             sub_template,
             service_name,
-            ServiceModel(service_type, formulation; use_service_name = true),
+            ServiceModel(service_type, formulation; use_service_name=true),
         )
     end
     return
@@ -118,7 +129,10 @@ function PSI.set_service_model!(
     service_type::Type{<:PSY.Service},
     formulation::Type{<:PSI.AbstractServiceFormulation},
 )
-    PSI.set_service_model!(template.base_template, PSI.ServiceModel(service_type, formulation))
+    PSI.set_service_model!(
+        template.base_template,
+        PSI.ServiceModel(service_type, formulation),
+    )
     for sub_template in get_sub_templates(template)
         PSI.set_service_model!(
             sub_template,
diff --git a/src/problems/multi_region_problem.jl b/src/problems/multi_region_problem.jl
index 77947a7..1e794fe 100644
--- a/src/problems/multi_region_problem.jl
+++ b/src/problems/multi_region_problem.jl
@@ -20,25 +20,26 @@ function get_n_subregions(sys::PSY.System)
 end
 
 function PSI.DecisionModel{MultiRegionProblem}(
-    template::PSI.ProblemTemplate,
+    template::MultiProblemTemplate,
     sys::PSY.System,
-    settings::PSI.Settings,
-    ::Union{Nothing,JuMP.Model}=nothing;
-    name=nothing,
+    ::Union{Nothing, JuMP.Model}=nothing;
+    kwargs...,
 )
-    if name === nothing
-        name = nameof(MultiRegionProblem)
-    elseif name isa String
-        name = Symbol(name)
-    end
-
+    name = Symbol(get(kwargs, :name, nameof(MultiRegionProblem)))
+    settings = PSI.Settings(sys; [k for k in kwargs if first(k) ∉ [:name]]...)
     # get number of system subregions
     # NOTE: `ext` field for each component must be filled with a "subregion" key in the dictionary
     region_keys = get_n_subregions(sys)
 
     # define the optimization container with master and subproblems
     internal = PSI.ModelInternal(
-        MultiOptimizationContainer(SequentialAlgorithm, sys, settings, PSY.Deterministic, region_keys),
+        MultiOptimizationContainer(
+            SequentialAlgorithm,
+            sys,
+            settings,
+            PSY.Deterministic,
+            region_keys,
+        ),
     )
     template_ = deepcopy(template)
 
@@ -51,7 +52,7 @@ function PSI.DecisionModel{MultiRegionProblem}(
         sys,
         internal,
         PSI.DecisionModelStore(),
-        Dict{String,Any}(),
+        Dict{String, Any}(),
     )
 end
 
@@ -80,8 +81,7 @@ function build_pre_step!(model::PSI.DecisionModel{MultiRegionProblem})
     return
 end
 
-function handle_initial_conditions!(model::PSI.DecisionModel{MultiRegionProblem})
-end
+function handle_initial_conditions!(model::PSI.DecisionModel{MultiRegionProblem}) end
 
 function instantiate_network_model(model::PSI.DecisionModel{MultiRegionProblem})
     PSI.instantiate_network_model(model)
@@ -89,7 +89,11 @@ function instantiate_network_model(model::PSI.DecisionModel{MultiRegionProblem})
 end
 
 function PSI.build_model!(model::PSI.DecisionModel{MultiRegionProblem})
-    build_impl!(PSI.get_optimization_container(model), PSI.get_template(model), PSI.get_system(model))
+    build_impl!(
+        PSI.get_optimization_container(model),
+        PSI.get_template(model),
+        PSI.get_system(model),
+    )
     return
 end
 
@@ -99,33 +103,38 @@ function PSI.solve_impl!(model::PSI.DecisionModel{MultiRegionProblem})
     return
 end
 
-function PSI.write_model_dual_results!(store,
+function PSI.write_model_dual_results!(
+    store,
     model::PSI.DecisionModel{MultiRegionProblem},
     index::PSI.DecisionModelIndexType,
     update_timestamp::Dates.DateTime,
-    export_params::Union{Dict{Symbol,Any},Nothing},)
-end
-function PSI.write_model_parameter_results!(store,
+    export_params::Union{Dict{Symbol, Any}, Nothing},
+) end
+function PSI.write_model_parameter_results!(
+    store,
     model::PSI.DecisionModel{MultiRegionProblem},
     index::PSI.DecisionModelIndexType,
     update_timestamp::Dates.DateTime,
-    export_params::Union{Dict{Symbol,Any},Nothing},)
-end
-function PSI.write_model_variable_results!(store,
+    export_params::Union{Dict{Symbol, Any}, Nothing},
+) end
+function PSI.write_model_variable_results!(
+    store,
     model::PSI.DecisionModel{MultiRegionProblem},
     index::PSI.DecisionModelIndexType,
     update_timestamp::Dates.DateTime,
-    export_params::Union{Dict{Symbol,Any},Nothing},)
-end
-function PSI.write_model_aux_variable_results!(store,
+    export_params::Union{Dict{Symbol, Any}, Nothing},
+) end
+function PSI.write_model_aux_variable_results!(
+    store,
     model::PSI.DecisionModel{MultiRegionProblem},
     index::PSI.DecisionModelIndexType,
     update_timestamp::Dates.DateTime,
-    export_params::Union{Dict{Symbol,Any},Nothing},)
-end
-function PSI.write_model_expression_results!(store,
+    export_params::Union{Dict{Symbol, Any}, Nothing},
+) end
+function PSI.write_model_expression_results!(
+    store,
     model::PSI.DecisionModel{MultiRegionProblem},
     index::PSI.DecisionModelIndexType,
     update_timestamp::Dates.DateTime,
-    export_params::Union{Dict{Symbol,Any},Nothing},)
-end
+    export_params::Union{Dict{Symbol, Any}, Nothing},
+) end
diff --git a/test/Project.toml b/test/Project.toml
index a79c3c2..8cbe0c2 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -1,5 +1,6 @@
 [deps]
-HydroPowerSimulations = "fc1677e0-6ad7-4515-bf3a-bd6bf20a0b1b"
+InfrastructureSystems = "2cd47ed4-ca9b-11e9-27f2-ab636a7671f1"
+PowerSimulations = "e690365d-45e2-57bb-ac84-44ba829e73c4"
 PowerSimulationsDecomposition = "bed98974-b02a-5e2f-9ee0-a103f5c450dd"
 PowerSystemCaseBuilder = "f00506e0-b84f-492a-93c2-c0a9afc4364e"
 PowerSystems = "bcd98974-b02a-5e2f-9ee0-a103f5c450dd"
diff --git a/test/first_test.jl b/test/first_test.jl
index 401bb77..dec1828 100644
--- a/test/first_test.jl
+++ b/test/first_test.jl
@@ -29,7 +29,8 @@ const PSB = PowerSystemCaseBuilder
 
 # consider the use of custom system used for GDO case
 name_ = "AC_inter"
-sys_twin_rts_DA = PSY.System("GDO systems/saved_main_RTS_GMLC_DA_final_sys_" * name_ * ".json")   # day ahead 
+sys_twin_rts_DA =
+    PSY.System("GDO systems/saved_main_RTS_GMLC_DA_final_sys_" * name_ * ".json")   # day ahead 
 
 # ! check reserves
 
@@ -57,7 +58,6 @@ for b in [get_from(arc_), get_to(arc_)]
     PSY.set_ext!(b, Dict("subregion" => Set(["1", "2"])))
 end
 
-
 if name_ == "HVDC_inter"
     HVDC_inter = true
 else
@@ -75,17 +75,16 @@ storage_model = DeviceModel(
         "cycling_limits" => false,
         "energy_target" => false,
         "complete_coverage" => false,
-        "regularization" => true
+        "regularization" => true,
     ),
 )
 
 # UC model
-template_uc =
-    MultiProblemTemplate(
-        # NetworkModel(StandardPTDFModel; PTDF_matrix = PTDF(sys_twin_rts)),
-        NetworkModel(DCPPowerModel; use_slacks=true),
-        ["1", "2"]
-    )
+template_uc = MultiProblemTemplate(
+    # NetworkModel(StandardPTDFModel; PTDF_matrix = PTDF(sys_twin_rts)),
+    NetworkModel(DCPPowerModel; use_slacks=true),
+    ["1", "2"],
+)
 set_device_model!(template_uc, ThermalStandard, ThermalStandardUnitCommitment)
 set_device_model!(template_uc, RenewableDispatch, RenewableFullDispatch)
 set_device_model!(template_uc, RenewableFix, FixedOutput)
@@ -98,11 +97,11 @@ set_device_model!(template_uc, HydroEnergyReservoir, FixedOutput)
 set_device_model!(template_uc, storage_model)
 set_service_model!(
     template_uc,
-    ServiceModel(VariableReserve{ReserveUp}, RangeReserve; use_slacks = true),
+    ServiceModel(VariableReserve{ReserveUp}, RangeReserve; use_slacks=true),
 )
 set_service_model!(
     template_uc,
-    ServiceModel(VariableReserve{ReserveDown}, RangeReserve; use_slacks = true),
+    ServiceModel(VariableReserve{ReserveDown}, RangeReserve; use_slacks=true),
 )
 
 # add the HVDC line in case is present
@@ -116,19 +115,19 @@ model = DecisionModel(
     MultiRegionProblem,
     template_uc,
     sys_twin_rts_DA;
-    name = "UC",
-    optimizer = optimizer_with_attributes(
-        Xpress.Optimizer, 
+    name="UC",
+    optimizer=optimizer_with_attributes(
+        Xpress.Optimizer,
         "MIPRELSTOP" => 0.01,       # Set the relative mip gap tolerance
         "MAXMEMORYSOFT" => 600000,   # Set the maximum amount of memory the solver can use (in MB)
     ),
-    system_to_file = false,
-    initialize_model = true,
-    optimizer_solve_log_print = true,
-    direct_mode_optimizer = true,
-    rebuild_model = false,
-    store_variable_names = true,
-    calculate_conflict = true,
+    system_to_file=false,
+    initialize_model=true,
+    optimizer_solve_log_print=true,
+    direct_mode_optimizer=true,
+    rebuild_model=false,
+    store_variable_names=true,
+    calculate_conflict=true,
 )
 
 # for b in get_components(ACBus, model.sys)
@@ -145,5 +144,5 @@ model = DecisionModel(
 
 # PowerSimulationsDecomposition.instantiate_network_model(model)
 
-build!(model; console_level = Logging.Info, output_dir = mktempdir())
-solve!(model; console_level = Logging.Info)
\ No newline at end of file
+build!(model; console_level=Logging.Info, output_dir=mktempdir())
+solve!(model; console_level=Logging.Info)
diff --git a/test/runtests.jl b/test/runtests.jl
index b8c5bb4..dc2ef72 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -14,7 +14,6 @@ using StorageSystemsSimulations
 using Test
 using Logging
 
-
 import Aqua
 Aqua.test_unbound_args(PowerSimulationsDecomposition)
 Aqua.test_undefined_exports(PowerSimulationsDecomposition)
diff --git a/test/test_core.jl b/test/test_core.jl
index 8ab2c21..9aec8f9 100644
--- a/test/test_core.jl
+++ b/test/test_core.jl
@@ -1,4 +1,3 @@
 @test 1 == 1
 
-
 #sys = build_system(PSISystems, "sys10_pjm_ac_dc")

From caa85fde72094ce73ad0b2105b911954c325863e Mon Sep 17 00:00:00 2001
From: Jose Daniel Lara <jdlara@berkeley.edu>
Date: Mon, 26 Feb 2024 22:43:00 -0700
Subject: [PATCH 03/17] remove hacks

---
 src/algorithms/sequential_algorithm.jl |  4 +-
 src/multi_optimization_container.jl    | 70 --------------------------
 src/problems/multi_region_problem.jl   |  6 ++-
 3 files changed, 6 insertions(+), 74 deletions(-)

diff --git a/src/algorithms/sequential_algorithm.jl b/src/algorithms/sequential_algorithm.jl
index cdd9feb..4c1536e 100644
--- a/src/algorithms/sequential_algorithm.jl
+++ b/src/algorithms/sequential_algorithm.jl
@@ -3,10 +3,8 @@ function build_impl!(
     template::MultiProblemTemplate,
     sys::PSY.System,
 )
-    sub_templates = get_subtemplates(template)
-    for (index, sub_problem) in container.subproblems
+    for (index, sub_problem) in get_sub_templates(template)
         @debug "Building Subproblem $index" _group = PSI.LOG_GROUP_OPTIMIZATION_CONTAINER
-        _system_modification!(sys, index)
         PSI.build_impl!(sub_problem, sub_templates[index], sys)
     end
 
diff --git a/src/multi_optimization_container.jl b/src/multi_optimization_container.jl
index aade861..c0e57b8 100644
--- a/src/multi_optimization_container.jl
+++ b/src/multi_optimization_container.jl
@@ -154,66 +154,6 @@ function _finalize_jump_model!(
     return
 end
 
-function _system_modification!(sys::PSY.System, index)
-    for component in PSY.get_components(PSY.Component, sys)
-        if typeof(component) <: PSY.Bus || :available in fieldnames(typeof(component))
-            sb_ = PSY.get_ext(component)["subregion"]
-            if typeof(component) <: PSY.Bus
-                if "original_type" ∉ keys(PSY.get_ext(component))
-                    component.ext["original_type"] = PSY.get_bustype(component)
-                end
-                if PSY.get_bustype(component) == PSY.ACBusTypes.ISOLATED
-                    if component.ext["original_type"] != PSY.ACBusTypes.ISOLATED
-                        PSY.set_bustype!(component, component.ext["original_type"])
-                    end
-                end
-                if index ∉ sb_
-                    #@error "Changed $(summary(component)) to ISOLATED"
-                    PSY.set_bustype!(component, PSY.ACBusTypes.ISOLATED)
-                end
-            else
-                if "original_type" ∉ keys(PSY.get_ext(component))
-                    component.ext["original_available"] = PSY.get_available(component)
-                end
-                if index ∈ sb_
-                    #@error "Changed $(summary(component)) to True"
-                    PSY.set_available!(component, true)
-                else
-                    #@error "Changed $(summary(component)) to False"
-                    PSY.set_available!(component, false)
-                end
-                if typeof(component) <: PSY.Reserve
-                    @show (index, PSY.get_name(component), PSY.get_available(component))
-                end
-            end
-        end
-    end
-    total_number_of_gens =
-        length(PSY.get_components(x -> PSY.get_available(x), PSY.ThermalStandard, sys))
-    total_number_of_ac_buses = length(
-        PSY.get_components(
-            x -> PSY.get_bustype(x) != PSY.ACBusTypes.ISOLATED,
-            PSY.ACBus,
-            sys,
-        ),
-    )
-    @show "Inside Mod Function Components using PSY functions"
-    @show total_number_of_gens
-    @show total_number_of_ac_buses
-    return
-end
-
-function _restore_system!(sys::PSY.System)
-    for component in PSY.get_components(PSY.Component, sys)
-        if typeof(component) <: PSY.Bus
-            PSY.set_bustype!(component, PSY.get_ext(component)["original_type"])
-        elseif :available in fieldnames(typeof(component))
-            PSY.set_available!(component, PSY.get_ext(component)["original_available"])
-        end
-    end
-    return
-end
-
 function init_optimization_container!(
     container::MultiOptimizationContainer,
     ::Type{T},
@@ -244,27 +184,17 @@ function init_optimization_container!(
 
     total_number_of_gens = length(PSI.get_available_components(PSY.ThermalStandard, sys))
     total_number_of_ac_buses = length(PSI.get_available_components(PSY.ACBus, sys))
-    @show "Total Components"
-    @show total_number_of_gens
-    @show total_number_of_ac_buses
 
     for (index, sub_problem) in container.subproblems
         @debug "Initializing Container Subproblem $index" _group =
             PSI.LOG_GROUP_OPTIMIZATION_CONTAINER
         sub_problem.settings = deepcopy(settings)
-        _system_modification!(sys, index)
         total_number_of_gens =
             length(PSI.get_available_components(PSY.ThermalStandard, sys))
         total_number_of_ac_buses = length(PSI.get_available_components(PSY.ACBus, sys))
-        @show "Before Init number of components using PSI functions"
-        @show total_number_of_gens
-        @show total_number_of_ac_buses
         PSI.init_optimization_container!(sub_problem, T, sys)
     end
 
-    # restore original system
-    _restore_system!(sys)
-
     return
 end
 
diff --git a/src/problems/multi_region_problem.jl b/src/problems/multi_region_problem.jl
index 1e794fe..7b575d3 100644
--- a/src/problems/multi_region_problem.jl
+++ b/src/problems/multi_region_problem.jl
@@ -84,7 +84,11 @@ end
 function handle_initial_conditions!(model::PSI.DecisionModel{MultiRegionProblem}) end
 
 function instantiate_network_model(model::PSI.DecisionModel{MultiRegionProblem})
-    PSI.instantiate_network_model(model)
+    template = PSI.get_template(model)
+    for sub_template in get_sub_templates(template)
+        network_model = PSI.get_network_model(sub_template)
+        PSI.instantiate_network_model(network_model, PSI.get_system(model))
+    end
     return
 end
 

From f6191322ebdbfb8d32cfcbb912faf2277eac56ef Mon Sep 17 00:00:00 2001
From: Jose Daniel Lara <jdlara@berkeley.edu>
Date: Tue, 27 Feb 2024 00:10:25 -0700
Subject: [PATCH 04/17] multi container build updates

---
 src/algorithms/sequential_algorithm.jl |  6 +++---
 src/multi_optimization_container.jl    |  4 ++++
 src/multiproblem_template.jl           | 21 ++++++++++++------
 src/problems/multi_region_problem.jl   | 30 +++-----------------------
 4 files changed, 24 insertions(+), 37 deletions(-)

diff --git a/src/algorithms/sequential_algorithm.jl b/src/algorithms/sequential_algorithm.jl
index 4c1536e..a596e49 100644
--- a/src/algorithms/sequential_algorithm.jl
+++ b/src/algorithms/sequential_algorithm.jl
@@ -3,9 +3,9 @@ function build_impl!(
     template::MultiProblemTemplate,
     sys::PSY.System,
 )
-    for (index, sub_problem) in get_sub_templates(template)
-        @debug "Building Subproblem $index" _group = PSI.LOG_GROUP_OPTIMIZATION_CONTAINER
-        PSI.build_impl!(sub_problem, sub_templates[index], sys)
+    for (index, sub_template) in get_sub_templates(template)
+        @info "Building Subproblem $index" _group = PSI.LOG_GROUP_OPTIMIZATION_CONTAINER
+        PSI.build_impl!(get_subproblem(container, index), sub_template, sys)
     end
 
     build_main_problem!(container, template, sys)
diff --git a/src/multi_optimization_container.jl b/src/multi_optimization_container.jl
index c0e57b8..a58ea81 100644
--- a/src/multi_optimization_container.jl
+++ b/src/multi_optimization_container.jl
@@ -110,6 +110,10 @@ PSI.get_aux_variables(container::MultiOptimizationContainer) = container.aux_var
 PSI.get_base_power(container::MultiOptimizationContainer) = container.base_power
 PSI.get_constraints(container::MultiOptimizationContainer) = container.constraints
 
+function get_subproblem(container::MultiOptimizationContainer, id::String)
+    return container.subproblems[id]
+end
+
 function check_optimization_container(container::MultiOptimizationContainer)
     # Solve main problem
     for (index, sub_problem) in container.subproblems
diff --git a/src/multiproblem_template.jl b/src/multiproblem_template.jl
index d6cca2c..334b766 100644
--- a/src/multiproblem_template.jl
+++ b/src/multiproblem_template.jl
@@ -35,7 +35,11 @@ function PSI.get_network_formulation(template::MultiProblemTemplate)
 end
 
 function get_sub_templates(template::MultiProblemTemplate)
-    return values(template.sub_templates)
+    return template.sub_templates
+end
+
+function get_sub_problem_keys(template::MultiProblemTemplate)
+    return sort!(collect(keys(get_sub_templates(template))))
 end
 
 """
@@ -65,7 +69,7 @@ function PSI.set_device_model!(
         template.base_template,
         PSI.DeviceModel(component_type, formulation),
     )
-    for sub_template in get_sub_templates(template)
+    for (_, sub_template) in get_sub_templates(template)
         PSI.set_device_model!(sub_template, PSI.DeviceModel(component_type, formulation))
     end
     return
@@ -150,7 +154,7 @@ function PSI.set_service_model!(
 )
     PSI.set_service_model!(template.base_template, service_name, model)
     for sub_template in get_sub_templates(template)
-        PSI.set_service_model!(sub_template, service_name, model)
+        PSI.set_service_model!(sub_template, service_name, deepcopy(model))
     end
     return
 end
@@ -160,14 +164,17 @@ function PSI.set_service_model!(
     model::PSI.ServiceModel{<:PSY.Service, <:PSI.AbstractServiceFormulation},
 )
     PSI.set_service_model!(template.base_template, model)
-    for sub_template in get_sub_templates(template)
-        PSI.set_service_model!(sub_template, model)
+    for (_, sub_template) in get_sub_templates(template)
+        PSI.set_service_model!(sub_template, deepcopy(model))
     end
     return
 end
 
-function PSI.finalize_template!(template::MultiProblemTemplate, sys::PSY.System)
-    for sub_template in get_sub_templates(template)
+function finalize_template!(template::MultiProblemTemplate, sys::PSY.System)
+    PSI.finalize_template!(template.base_template, sys)
+    for (ix, sub_template) in get_sub_templates(template)
+        @debug "Finalizing template for sub probem $ix"
         PSI.finalize_template!(sub_template, sys)
     end
+    return
 end
diff --git a/src/problems/multi_region_problem.jl b/src/problems/multi_region_problem.jl
index 7b575d3..011b66f 100644
--- a/src/problems/multi_region_problem.jl
+++ b/src/problems/multi_region_problem.jl
@@ -1,24 +1,5 @@
 struct MultiRegionProblem <: PSI.DecisionProblem end
 
-function get_n_subregions(sys::PSY.System)
-    # TODO: this function can be improved (less for loops)
-    subregions_ = String[]
-    for d in PSY.get_components(PSY.Component, sys)
-        if typeof(d) ∉ [PSY.Area, PSY.Arc, PSY.LoadZone]
-            if isempty(PSY.get_ext(d))
-                error("Field `ext` must be non empty if spatial decomposition is used.")
-            else
-                for sb in PSY.get_ext(d)["subregion"]
-                    if sb ∉ subregions_
-                        push!(subregions_, sb)
-                    end
-                end
-            end
-        end
-    end
-    return subregions_
-end
-
 function PSI.DecisionModel{MultiRegionProblem}(
     template::MultiProblemTemplate,
     sys::PSY.System,
@@ -27,23 +8,18 @@ function PSI.DecisionModel{MultiRegionProblem}(
 )
     name = Symbol(get(kwargs, :name, nameof(MultiRegionProblem)))
     settings = PSI.Settings(sys; [k for k in kwargs if first(k) ∉ [:name]]...)
-    # get number of system subregions
-    # NOTE: `ext` field for each component must be filled with a "subregion" key in the dictionary
-    region_keys = get_n_subregions(sys)
-
-    # define the optimization container with master and subproblems
     internal = PSI.ModelInternal(
         MultiOptimizationContainer(
             SequentialAlgorithm,
             sys,
             settings,
             PSY.Deterministic,
-            region_keys,
+            get_sub_problem_keys(template),
         ),
     )
     template_ = deepcopy(template)
 
-    PSI.finalize_template!(template_, sys)
+    finalize_template!(template_, sys)
 
     # return multi-region decision model container
     return PSI.DecisionModel{MultiRegionProblem}(
@@ -85,7 +61,7 @@ function handle_initial_conditions!(model::PSI.DecisionModel{MultiRegionProblem}
 
 function instantiate_network_model(model::PSI.DecisionModel{MultiRegionProblem})
     template = PSI.get_template(model)
-    for sub_template in get_sub_templates(template)
+    for (_, sub_template) in get_sub_templates(template)
         network_model = PSI.get_network_model(sub_template)
         PSI.instantiate_network_model(network_model, PSI.get_system(model))
     end

From cdb1f6eaf59b69041b4dd1729d008db4351c071e Mon Sep 17 00:00:00 2001
From: Jose Daniel Lara <jdlara@berkeley.edu>
Date: Tue, 27 Feb 2024 17:58:48 -0700
Subject: [PATCH 05/17] remove stale code

---
 src/multi_optimization_container.jl | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/src/multi_optimization_container.jl b/src/multi_optimization_container.jl
index a58ea81..7864680 100644
--- a/src/multi_optimization_container.jl
+++ b/src/multi_optimization_container.jl
@@ -186,16 +186,10 @@ function init_optimization_container!(
 
     _finalize_jump_model!(container, settings)
 
-    total_number_of_gens = length(PSI.get_available_components(PSY.ThermalStandard, sys))
-    total_number_of_ac_buses = length(PSI.get_available_components(PSY.ACBus, sys))
-
     for (index, sub_problem) in container.subproblems
         @debug "Initializing Container Subproblem $index" _group =
             PSI.LOG_GROUP_OPTIMIZATION_CONTAINER
         sub_problem.settings = deepcopy(settings)
-        total_number_of_gens =
-            length(PSI.get_available_components(PSY.ThermalStandard, sys))
-        total_number_of_ac_buses = length(PSI.get_available_components(PSY.ACBus, sys))
         PSI.init_optimization_container!(sub_problem, T, sys)
     end
 

From f1c8bd4fc81e931b6f1a7819961fef7e62edfdf5 Mon Sep 17 00:00:00 2001
From: Jose Daniel Lara <jdlara@berkeley.edu>
Date: Wed, 28 Feb 2024 15:53:40 -0700
Subject: [PATCH 06/17] set subs systems ids

---
 src/multiproblem_template.jl         | 33 ++++++++++++++++++----------
 src/problems/multi_region_problem.jl |  5 +++--
 2 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/src/multiproblem_template.jl b/src/multiproblem_template.jl
index 334b766..1f80e40 100644
--- a/src/multiproblem_template.jl
+++ b/src/multiproblem_template.jl
@@ -50,7 +50,8 @@ function PSI.set_network_model!(
     model::PSI.NetworkModel{<:PM.AbstractPowerModel},
 )
     PSI.set_network_model!(template.base_template, model)
-    for sub_template in get_sub_templates(template)
+    for (id, sub_template) in get_sub_templates(template)
+        PSI.set_subsystem!(model, id)
         PSI.set_network_model!(sub_template, model)
     end
     return
@@ -69,8 +70,10 @@ function PSI.set_device_model!(
         template.base_template,
         PSI.DeviceModel(component_type, formulation),
     )
-    for (_, sub_template) in get_sub_templates(template)
-        PSI.set_device_model!(sub_template, PSI.DeviceModel(component_type, formulation))
+    for (id, sub_template) in get_sub_templates(template)
+        network_model = PSI.DeviceModel(component_type, formulation)
+        PSI.set_subsystem!(network_model, id)
+        PSI.set_device_model!(sub_template, network_model)
     end
     return
 end
@@ -83,7 +86,8 @@ function PSI.set_device_model!(
     model::PSI.DeviceModel{<:PSY.Device, <:PSI.AbstractDeviceFormulation},
 )
     PSI.set_device_model!(template.base_template, model)
-    for sub_template in get_sub_templates(template)
+    for (id, sub_template) in get_sub_templates(template)
+        PSI.set_subsystem!(model, id)
         PSI.set_device_model!(sub_template, model)
     end
     return
@@ -94,7 +98,8 @@ function PSI.set_device_model!(
     model::PSI.DeviceModel{<:PSY.Branch, <:PSI.AbstractDeviceFormulation},
 )
     PSI.set_device_model!(template.base_template, model)
-    for sub_template in get_sub_templates(template)
+    for (id, sub_template) in get_sub_templates(template)
+        PSI.set_subsystem!(model, id)
         PSI.set_device_model!(sub_template, PSI.DeviceModel(component_type, formulation))
     end
     return
@@ -115,11 +120,13 @@ function PSI.set_service_model!(
         service_name,
         ServiceModel(service_type, formulation; use_service_name=true),
     )
-    for sub_template in get_sub_templates(template)
+    for (id, sub_template) in get_sub_templates(template)
+        service_model = ServiceModel(service_type, formulation; use_service_name=true)
+        PSI.set_subsystem!(service_model, id)
         PSI.set_service_model!(
             sub_template,
             service_name,
-            ServiceModel(service_type, formulation; use_service_name=true),
+            service_model,
         )
     end
     return
@@ -137,11 +144,13 @@ function PSI.set_service_model!(
         template.base_template,
         PSI.ServiceModel(service_type, formulation),
     )
-    for sub_template in get_sub_templates(template)
+    for (id, sub_template) in get_sub_templates(template)
+        service_model = ServiceModel(service_type, formulation)
+        PSI.set_subsystem!(service_model, id)
         PSI.set_service_model!(
             sub_template,
             service_name,
-            PSI.ServiceModel(service_type, formulation),
+            service_model,
         )
     end
     return
@@ -153,7 +162,8 @@ function PSI.set_service_model!(
     model::PSI.ServiceModel{<:PSY.Service, <:PSI.AbstractServiceFormulation},
 )
     PSI.set_service_model!(template.base_template, service_name, model)
-    for sub_template in get_sub_templates(template)
+    for (id, sub_template) in get_sub_templates(template)
+        PSI.set_subsystem!(model, id)
         PSI.set_service_model!(sub_template, service_name, deepcopy(model))
     end
     return
@@ -164,7 +174,8 @@ function PSI.set_service_model!(
     model::PSI.ServiceModel{<:PSY.Service, <:PSI.AbstractServiceFormulation},
 )
     PSI.set_service_model!(template.base_template, model)
-    for (_, sub_template) in get_sub_templates(template)
+    for (id, sub_template) in get_sub_templates(template)
+        PSI.set_subsystem!(model, id)
         PSI.set_service_model!(sub_template, deepcopy(model))
     end
     return
diff --git a/src/problems/multi_region_problem.jl b/src/problems/multi_region_problem.jl
index 011b66f..a78a04d 100644
--- a/src/problems/multi_region_problem.jl
+++ b/src/problems/multi_region_problem.jl
@@ -39,7 +39,7 @@ function PSI.build_impl!(model::PSI.DecisionModel{MultiRegionProblem})
     handle_initial_conditions!(model)
     PSI.build_model!(model)
     # Might need custom implementation for this container type
-    #serialize_metadata!(get_optimization_container(model), get_output_dir(model))
+    # serialize_metadata!(get_optimization_container(model), get_output_dir(model))
     PSI.log_values(PSI.get_settings(model))
     return
 end
@@ -61,8 +61,9 @@ function handle_initial_conditions!(model::PSI.DecisionModel{MultiRegionProblem}
 
 function instantiate_network_model(model::PSI.DecisionModel{MultiRegionProblem})
     template = PSI.get_template(model)
-    for (_, sub_template) in get_sub_templates(template)
+    for (id, sub_template) in get_sub_templates(template)
         network_model = PSI.get_network_model(sub_template)
+        PSI.set_subsystem!(network_model, id)
         PSI.instantiate_network_model(network_model, PSI.get_system(model))
     end
     return

From 1677242ee084034231cc158484fde725771f9767 Mon Sep 17 00:00:00 2001
From: Jose Daniel Lara <jdlara@berkeley.edu>
Date: Wed, 28 Feb 2024 16:51:46 -0700
Subject: [PATCH 07/17] use model in build functions

---
 src/multi_optimization_container.jl  | 6 +++---
 src/multiproblem_template.jl         | 5 +++++
 src/problems/multi_region_problem.jl | 2 +-
 3 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/src/multi_optimization_container.jl b/src/multi_optimization_container.jl
index 7864680..75bd083 100644
--- a/src/multi_optimization_container.jl
+++ b/src/multi_optimization_container.jl
@@ -160,9 +160,9 @@ end
 
 function init_optimization_container!(
     container::MultiOptimizationContainer,
-    ::Type{T},
+    network_model::PSI.NetworkModel{<:PM.AbstractPowerModel},
     sys::PSY.System,
-) where {T <: PM.AbstractPowerModel}
+)
     PSY.set_units_base_system!(sys, "SYSTEM_BASE")
     # The order of operations matter
     settings = PSI.get_settings(container)
@@ -190,7 +190,7 @@ function init_optimization_container!(
         @debug "Initializing Container Subproblem $index" _group =
             PSI.LOG_GROUP_OPTIMIZATION_CONTAINER
         sub_problem.settings = deepcopy(settings)
-        PSI.init_optimization_container!(sub_problem, T, sys)
+        PSI.init_optimization_container!(sub_problem, network_model, sys)
     end
 
     return
diff --git a/src/multiproblem_template.jl b/src/multiproblem_template.jl
index 1f80e40..68a22a6 100644
--- a/src/multiproblem_template.jl
+++ b/src/multiproblem_template.jl
@@ -3,6 +3,11 @@ struct MultiProblemTemplate <: PSI.AbstractProblemTemplate
     sub_templates::Dict{String, PSI.ProblemTemplate}
 end
 
+PSI.get_device_models(template::MultiProblemTemplate) = template.base_template.devices
+PSI.get_branch_models(template::MultiProblemTemplate) = template.base_template.branches
+PSI.get_service_models(template::MultiProblemTemplate) = template.base_template.services
+PSI.get_network_model(template::MultiProblemTemplate) = template.base_template.network_model
+
 function MultiProblemTemplate(
     base_template::PSI.ProblemTemplate,
     problem_keys::Vector{String},
diff --git a/src/problems/multi_region_problem.jl b/src/problems/multi_region_problem.jl
index a78a04d..6bcbfbb 100644
--- a/src/problems/multi_region_problem.jl
+++ b/src/problems/multi_region_problem.jl
@@ -48,7 +48,7 @@ function build_pre_step!(model::PSI.DecisionModel{MultiRegionProblem})
     @info "Initializing Optimization Container For a DecisionModel"
     init_optimization_container!(
         PSI.get_optimization_container(model),
-        PSI.get_network_formulation(PSI.get_template(model)),
+        PSI.get_network_model(PSI.get_template(model)),
         PSI.get_system(model),
     )
     @info "Initializing ModelStoreParams"

From 302b5350ab80f4300d02be74dd0bd136611c71c9 Mon Sep 17 00:00:00 2001
From: Jose Daniel Lara <jdlara@berkeley.edu>
Date: Thu, 29 Feb 2024 11:02:09 -0700
Subject: [PATCH 08/17] formatter

---
 src/multiproblem_template.jl | 12 ++----------
 1 file changed, 2 insertions(+), 10 deletions(-)

diff --git a/src/multiproblem_template.jl b/src/multiproblem_template.jl
index 68a22a6..57d48de 100644
--- a/src/multiproblem_template.jl
+++ b/src/multiproblem_template.jl
@@ -128,11 +128,7 @@ function PSI.set_service_model!(
     for (id, sub_template) in get_sub_templates(template)
         service_model = ServiceModel(service_type, formulation; use_service_name=true)
         PSI.set_subsystem!(service_model, id)
-        PSI.set_service_model!(
-            sub_template,
-            service_name,
-            service_model,
-        )
+        PSI.set_service_model!(sub_template, service_name, service_model)
     end
     return
 end
@@ -152,11 +148,7 @@ function PSI.set_service_model!(
     for (id, sub_template) in get_sub_templates(template)
         service_model = ServiceModel(service_type, formulation)
         PSI.set_subsystem!(service_model, id)
-        PSI.set_service_model!(
-            sub_template,
-            service_name,
-            service_model,
-        )
+        PSI.set_service_model!(sub_template, service_name, service_model)
     end
     return
 end

From dd9b873f70cd77ae639e3dc154bf4566b8ebca8d Mon Sep 17 00:00:00 2001
From: Jose Daniel Lara <jdlara@berkeley.edu>
Date: Wed, 6 Mar 2024 17:08:38 -0700
Subject: [PATCH 09/17] add place holder for serialize

---
 src/problems/multi_region_problem.jl | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/problems/multi_region_problem.jl b/src/problems/multi_region_problem.jl
index 6bcbfbb..c251532 100644
--- a/src/problems/multi_region_problem.jl
+++ b/src/problems/multi_region_problem.jl
@@ -69,6 +69,10 @@ function instantiate_network_model(model::PSI.DecisionModel{MultiRegionProblem})
     return
 end
 
+function PSI.serialize_problem(model::PSI.DecisionModel{MultiRegionProblem};
+    optimizer::Nothing)
+end
+
 function PSI.build_model!(model::PSI.DecisionModel{MultiRegionProblem})
     build_impl!(
         PSI.get_optimization_container(model),

From 698c2ad9f5f5f2ccc467d469e7f8ad86eea6a3db Mon Sep 17 00:00:00 2001
From: Jose Daniel Lara <jdlara@berkeley.edu>
Date: Wed, 6 Mar 2024 17:21:36 -0700
Subject: [PATCH 10/17] formatter

---
 src/problems/multi_region_problem.jl | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/src/problems/multi_region_problem.jl b/src/problems/multi_region_problem.jl
index c251532..d229423 100644
--- a/src/problems/multi_region_problem.jl
+++ b/src/problems/multi_region_problem.jl
@@ -69,9 +69,10 @@ function instantiate_network_model(model::PSI.DecisionModel{MultiRegionProblem})
     return
 end
 
-function PSI.serialize_problem(model::PSI.DecisionModel{MultiRegionProblem};
-    optimizer::Nothing)
-end
+function PSI.serialize_problem(
+    model::PSI.DecisionModel{MultiRegionProblem};
+    optimizer::Nothing,
+) end
 
 function PSI.build_model!(model::PSI.DecisionModel{MultiRegionProblem})
     build_impl!(

From 361598827cabf9913c02c11b7b2cf3fbcfb00870 Mon Sep 17 00:00:00 2001
From: Jose Daniel Lara <jdlara@berkeley.edu>
Date: Wed, 6 Mar 2024 17:46:43 -0700
Subject: [PATCH 11/17] use a main container instead of a jump model

---
 src/multi_optimization_container.jl | 42 ++++++-----------------------
 1 file changed, 8 insertions(+), 34 deletions(-)

diff --git a/src/multi_optimization_container.jl b/src/multi_optimization_container.jl
index 75bd083..c83c9c3 100644
--- a/src/multi_optimization_container.jl
+++ b/src/multi_optimization_container.jl
@@ -1,6 +1,6 @@
 mutable struct MultiOptimizationContainer{T <: DecompositionAlgorithm} <:
                PSI.AbstractModelContainer
-    main_JuMPmodel::JuMP.Model
+    main_problem::PSI.OptimizationContainer
     subproblems::Dict{String, PSI.OptimizationContainer}
     time_steps::UnitRange{Int}
     resolution::Dates.TimePeriod
@@ -45,7 +45,7 @@ function MultiOptimizationContainer(
     end
 
     return MultiOptimizationContainer{T}(
-        JuMP.Model(),
+        PSI.OptimizationContainer(sys, settings, nothing, U),
         subproblems,
         1:1,
         IS.time_period_conversion(resolution),
@@ -88,7 +88,7 @@ PSI.get_initial_conditions_data(container::MultiOptimizationContainer) =
     container.initial_conditions_data
 PSI.get_initial_time(container::MultiOptimizationContainer) =
     PSI.get_initial_time(container.settings)
-PSI.get_jump_model(container::MultiOptimizationContainer) = container.main_JuMPmodel
+PSI.get_jump_model(container::MultiOptimizationContainer) = PSI.get_jump_model(container.main_problem)
 PSI.get_metadata(container::MultiOptimizationContainer) = container.metadata
 PSI.get_optimizer_stats(container::MultiOptimizationContainer) = container.optimizer_stats
 PSI.get_parameters(container::MultiOptimizationContainer) = container.parameters
@@ -126,35 +126,7 @@ function _finalize_jump_model!(
     settings::PSI.Settings,
 )
     @debug "Instantiating the JuMP model" _group = PSI.LOG_GROUP_OPTIMIZATION_CONTAINER
-    #=
-    if PSI.built_for_recurrent_solves(container) && PSI.get_optimizer(settings) === nothing
-        throw(
-            IS.ConflictingInputsError(
-                "Optimizer can not be nothing when building for recurrent solves",
-            ),
-        )
-    end
-    =#
-
-    if PSI.get_direct_mode_optimizer(settings)
-        optimizer = () -> MOI.instantiate(PSI.get_optimizer(settings))
-        container.main_JuMPmodel = JuMP.direct_model(optimizer())
-    elseif PSI.get_optimizer(settings) === nothing
-        @debug "The optimization model has no optimizer attached" _group =
-            LOG_GROUP_OPTIMIZATION_CONTAINER
-    else
-        JuMP.set_optimizer(PSI.get_jump_model(container), PSI.get_optimizer(settings))
-    end
-
-    JuMPmodel = PSI.get_jump_model(container)
-
-    if PSI.get_optimizer_solve_log_print(settings)
-        JuMP.unset_silent(JuMPmodel)
-        @debug "optimizer unset to silent" _group = PSI.LOG_GROUP_OPTIMIZATION_CONTAINER
-    else
-        JuMP.set_silent(JuMPmodel)
-        @debug "optimizer set to silent" _group = PSI.LOG_GROUP_OPTIMIZATION_CONTAINER
-    end
+    PSI._finalize_jump_model!(container.main_problem, settings)
     return
 end
 
@@ -184,7 +156,9 @@ function init_optimization_container!(
     stats = PSI.get_optimizer_stats(container)
     stats.detailed_stats = PSI.get_detailed_optimizer_stats(settings)
 
-    _finalize_jump_model!(container, settings)
+    # need a special method for the main problem to initialize the optimization container
+    # without actually caring about the subnetworks
+    # PSI.init_optimization_container!(sub_problem, network_model, sys)
 
     for (index, sub_problem) in container.subproblems
         @debug "Initializing Container Subproblem $index" _group =
@@ -192,7 +166,7 @@ function init_optimization_container!(
         sub_problem.settings = deepcopy(settings)
         PSI.init_optimization_container!(sub_problem, network_model, sys)
     end
-
+    _finalize_jump_model!(container, settings)
     return
 end
 

From ebcb423c43b4692aa52f2a007fb985661a515632 Mon Sep 17 00:00:00 2001
From: Jose Daniel Lara <jdlara@berkeley.edu>
Date: Thu, 7 Mar 2024 02:11:03 -0700
Subject: [PATCH 12/17] update container mappings

---
 Project.toml                         |  4 +-
 src/PowerSimulationsDecomposition.jl |  2 +
 src/definitions.jl                   |  1 +
 src/multi_optimization_container.jl  |  6 ++-
 src/problems/multi_region_problem.jl | 76 +++++++++++++++-------------
 5 files changed, 51 insertions(+), 38 deletions(-)
 create mode 100644 src/definitions.jl

diff --git a/Project.toml b/Project.toml
index ae41079..b20cca5 100644
--- a/Project.toml
+++ b/Project.toml
@@ -4,6 +4,7 @@ authors = ["Jose Daniel Lara"]
 version = "0.1.0"
 
 [deps]
+DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 InfrastructureSystems = "2cd47ed4-ca9b-11e9-27f2-ab636a7671f1"
@@ -15,10 +16,11 @@ PowerSystems = "bcd98974-b02a-5e2f-9ee0-a103f5c450dd"
 
 [compat]
 Dates = "1"
+DataStructures = "^0.18"
 DocStringExtensions = "~0.8, ~0.9"
 InfrastructureSystems = "^1.21"
 JuMP = "1"
-MathOptInterface = "1"
 MPI = "^0.20"
+MathOptInterface = "1"
 PowerSystems = "^3"
 julia = "^1.6"
diff --git a/src/PowerSimulationsDecomposition.jl b/src/PowerSimulationsDecomposition.jl
index 176ba4e..f05b112 100644
--- a/src/PowerSimulationsDecomposition.jl
+++ b/src/PowerSimulationsDecomposition.jl
@@ -10,6 +10,7 @@ import JuMP
 import Dates
 import MPI
 import MathOptInterface
+import DataStructures: SortedDict
 
 const PSI = PowerSimulations
 const PSY = PowerSystems
@@ -23,6 +24,7 @@ using DocStringExtensions
                                  $(DOCSTRING)
                                  """
 
+include("definitions.jl")
 include("core.jl")
 include("multiproblem_template.jl")
 include("multi_optimization_container.jl")
diff --git a/src/definitions.jl b/src/definitions.jl
new file mode 100644
index 0000000..19aa200
--- /dev/null
+++ b/src/definitions.jl
@@ -0,0 +1 @@
+const CONTAINER_FIELDS = [:variables, :aux_variables, :constraints, :expressions, :duals]
diff --git a/src/multi_optimization_container.jl b/src/multi_optimization_container.jl
index c83c9c3..ca175b9 100644
--- a/src/multi_optimization_container.jl
+++ b/src/multi_optimization_container.jl
@@ -88,7 +88,8 @@ PSI.get_initial_conditions_data(container::MultiOptimizationContainer) =
     container.initial_conditions_data
 PSI.get_initial_time(container::MultiOptimizationContainer) =
     PSI.get_initial_time(container.settings)
-PSI.get_jump_model(container::MultiOptimizationContainer) = PSI.get_jump_model(container.main_problem)
+PSI.get_jump_model(container::MultiOptimizationContainer) =
+    PSI.get_jump_model(container.main_problem)
 PSI.get_metadata(container::MultiOptimizationContainer) = container.metadata
 PSI.get_optimizer_stats(container::MultiOptimizationContainer) = container.optimizer_stats
 PSI.get_parameters(container::MultiOptimizationContainer) = container.parameters
@@ -115,10 +116,11 @@ function get_subproblem(container::MultiOptimizationContainer, id::String)
 end
 
 function check_optimization_container(container::MultiOptimizationContainer)
-    # Solve main problem
     for (index, sub_problem) in container.subproblems
         PSI.check_optimization_container(sub_problem)
     end
+    PSI.check_optimization_container(container.main_problem)
+    return
 end
 
 function _finalize_jump_model!(
diff --git a/src/problems/multi_region_problem.jl b/src/problems/multi_region_problem.jl
index d229423..8cff855 100644
--- a/src/problems/multi_region_problem.jl
+++ b/src/problems/multi_region_problem.jl
@@ -32,12 +32,51 @@ function PSI.DecisionModel{MultiRegionProblem}(
     )
 end
 
+function _get_axes!(common_axes::Dict, container::PSI.OptimizationContainer)
+    for field in CONTAINER_FIELDS
+        field_data = getfield(container, field)
+        for (key, value_container) in field_data
+            if isa(value_container, JuMP.Containers.SparseAxisArray)
+                continue
+            end
+            axes_data = get!(common_axes[field], key, SortedDict{Int, Set}())
+            for (ix, vals) in enumerate(axes(value_container))
+                union!(get!(axes_data, ix, Set(vals)))
+            end
+        end
+    end
+    return
+end
+
+function _map_containers(model::PSI.DecisionModel{MultiRegionProblem})
+    common_axes =
+        Dict(key => Dict{PSI.OptimizationContainerKey, Any}() for key in CONTAINER_FIELDS)
+    container = PSI.get_optimization_container(model)
+    for subproblem in values(container.subproblems)
+        _get_axes!(common_axes, subproblem)
+    end
+
+    for (field, vals) in common_axes
+        field_data = getfield(container, field)
+        for (key, axes_data) in vals
+            ax = [sort!(collect(v)) for v in values(axes_data)]
+            field_data[key] =
+                PSI.remove_undef!(JuMP.Containers.DenseAxisArray{Float64}(undef, ax...))
+        end
+    end
+
+    #TODO: Parameters Requires a different approach
+
+    return
+end
+
 function PSI.build_impl!(model::PSI.DecisionModel{MultiRegionProblem})
     build_pre_step!(model)
     @info "Instantiating Network Model"
     instantiate_network_model(model)
     handle_initial_conditions!(model)
     PSI.build_model!(model)
+    _map_containers(model)
     # Might need custom implementation for this container type
     # serialize_metadata!(get_optimization_container(model), get_output_dir(model))
     PSI.log_values(PSI.get_settings(model))
@@ -89,38 +128,5 @@ function PSI.solve_impl!(model::PSI.DecisionModel{MultiRegionProblem})
     return
 end
 
-function PSI.write_model_dual_results!(
-    store,
-    model::PSI.DecisionModel{MultiRegionProblem},
-    index::PSI.DecisionModelIndexType,
-    update_timestamp::Dates.DateTime,
-    export_params::Union{Dict{Symbol, Any}, Nothing},
-) end
-function PSI.write_model_parameter_results!(
-    store,
-    model::PSI.DecisionModel{MultiRegionProblem},
-    index::PSI.DecisionModelIndexType,
-    update_timestamp::Dates.DateTime,
-    export_params::Union{Dict{Symbol, Any}, Nothing},
-) end
-function PSI.write_model_variable_results!(
-    store,
-    model::PSI.DecisionModel{MultiRegionProblem},
-    index::PSI.DecisionModelIndexType,
-    update_timestamp::Dates.DateTime,
-    export_params::Union{Dict{Symbol, Any}, Nothing},
-) end
-function PSI.write_model_aux_variable_results!(
-    store,
-    model::PSI.DecisionModel{MultiRegionProblem},
-    index::PSI.DecisionModelIndexType,
-    update_timestamp::Dates.DateTime,
-    export_params::Union{Dict{Symbol, Any}, Nothing},
-) end
-function PSI.write_model_expression_results!(
-    store,
-    model::PSI.DecisionModel{MultiRegionProblem},
-    index::PSI.DecisionModelIndexType,
-    update_timestamp::Dates.DateTime,
-    export_params::Union{Dict{Symbol, Any}, Nothing},
-) end
+function PSI._check_numerical_bounds(model::PSI.DecisionModel{MultiRegionProblem})
+end

From 4893734fbbbacd219452e7e92d010648d1b2f527 Mon Sep 17 00:00:00 2001
From: Jose Daniel Lara <jdlara@berkeley.edu>
Date: Thu, 7 Mar 2024 11:41:29 -0700
Subject: [PATCH 13/17] shared axes code

---
 src/problems/multi_region_problem.jl | 37 ++++++++++++++++++++++------
 1 file changed, 30 insertions(+), 7 deletions(-)

diff --git a/src/problems/multi_region_problem.jl b/src/problems/multi_region_problem.jl
index 8cff855..bb9dfaa 100644
--- a/src/problems/multi_region_problem.jl
+++ b/src/problems/multi_region_problem.jl
@@ -32,7 +32,22 @@ function PSI.DecisionModel{MultiRegionProblem}(
     )
 end
 
-function _get_axes!(common_axes::Dict, container::PSI.OptimizationContainer)
+function _join_axes!(axes_data::SortedDict{Int, Set}, ix::Int, axes_value::UnitRange{Int})
+    _axes_data = get!(axes_data, ix, Set{UnitRange{Int}}())
+    if _axes_data == axes_value
+        return
+    end
+    union!(_axes_data, [axes_value])
+    return
+end
+
+function _join_axes!(axes_data::SortedDict{Int, Set}, ix::Int, axes_value::Vector)
+    _axes_data = get!(axes_data, ix, Set{eltype(axes_value)}())
+    union!(_axes_data, axes_value)
+    return
+end
+
+function _get_axes!(common_axes::Dict{Symbol, Dict{PSI.OptimizationContainerKey, SortedDict{Int, Set}}}, container::PSI.OptimizationContainer)
     for field in CONTAINER_FIELDS
         field_data = getfield(container, field)
         for (key, value_container) in field_data
@@ -41,30 +56,38 @@ function _get_axes!(common_axes::Dict, container::PSI.OptimizationContainer)
             end
             axes_data = get!(common_axes[field], key, SortedDict{Int, Set}())
             for (ix, vals) in enumerate(axes(value_container))
-                union!(get!(axes_data, ix, Set(vals)))
+                _join_axes!(axes_data, ix, vals)
             end
         end
     end
     return
 end
 
+function _make_joint_axes!(dim1::Set{T}, dim2::Set{UnitRange{Int}}) where T <: Union{Int, String}
+    return (collect(dim1), first(dim2))
+end
+
+function _make_joint_axes!(dim1::Set{UnitRange{Int}})
+    return (first(dim1),)
+end
+
+
 function _map_containers(model::PSI.DecisionModel{MultiRegionProblem})
     common_axes =
-        Dict(key => Dict{PSI.OptimizationContainerKey, Any}() for key in CONTAINER_FIELDS)
+        Dict{Symbol, Dict{PSI.OptimizationContainerKey, SortedDict{Int, Set}}}(key => Dict{PSI.OptimizationContainerKey, SortedDict{Int, Set}}() for key in CONTAINER_FIELDS)
     container = PSI.get_optimization_container(model)
-    for subproblem in values(container.subproblems)
+    for (_, subproblem) in container.subproblems
         _get_axes!(common_axes, subproblem)
     end
 
     for (field, vals) in common_axes
-        field_data = getfield(container, field)
+        field_data = getproperty(container, field)
         for (key, axes_data) in vals
-            ax = [sort!(collect(v)) for v in values(axes_data)]
+            ax = _make_joint_axes!(collect(values(axes_data))...)
             field_data[key] =
                 PSI.remove_undef!(JuMP.Containers.DenseAxisArray{Float64}(undef, ax...))
         end
     end
-
     #TODO: Parameters Requires a different approach
 
     return

From f15143c9c7662bbba296476d11c4f34fcf76acd5 Mon Sep 17 00:00:00 2001
From: Jose Daniel Lara <jdlara@berkeley.edu>
Date: Thu, 7 Mar 2024 11:47:56 -0700
Subject: [PATCH 14/17] formatter

---
 src/problems/multi_region_problem.jl | 20 +++++++++++++-------
 1 file changed, 13 insertions(+), 7 deletions(-)

diff --git a/src/problems/multi_region_problem.jl b/src/problems/multi_region_problem.jl
index bb9dfaa..dda4aeb 100644
--- a/src/problems/multi_region_problem.jl
+++ b/src/problems/multi_region_problem.jl
@@ -47,7 +47,10 @@ function _join_axes!(axes_data::SortedDict{Int, Set}, ix::Int, axes_value::Vecto
     return
 end
 
-function _get_axes!(common_axes::Dict{Symbol, Dict{PSI.OptimizationContainerKey, SortedDict{Int, Set}}}, container::PSI.OptimizationContainer)
+function _get_axes!(
+    common_axes::Dict{Symbol, Dict{PSI.OptimizationContainerKey, SortedDict{Int, Set}}},
+    container::PSI.OptimizationContainer,
+)
     for field in CONTAINER_FIELDS
         field_data = getfield(container, field)
         for (key, value_container) in field_data
@@ -63,7 +66,10 @@ function _get_axes!(common_axes::Dict{Symbol, Dict{PSI.OptimizationContainerKey,
     return
 end
 
-function _make_joint_axes!(dim1::Set{T}, dim2::Set{UnitRange{Int}}) where T <: Union{Int, String}
+function _make_joint_axes!(
+    dim1::Set{T},
+    dim2::Set{UnitRange{Int}},
+) where {T <: Union{Int, String}}
     return (collect(dim1), first(dim2))
 end
 
@@ -71,10 +77,11 @@ function _make_joint_axes!(dim1::Set{UnitRange{Int}})
     return (first(dim1),)
 end
 
-
 function _map_containers(model::PSI.DecisionModel{MultiRegionProblem})
-    common_axes =
-        Dict{Symbol, Dict{PSI.OptimizationContainerKey, SortedDict{Int, Set}}}(key => Dict{PSI.OptimizationContainerKey, SortedDict{Int, Set}}() for key in CONTAINER_FIELDS)
+    common_axes = Dict{Symbol, Dict{PSI.OptimizationContainerKey, SortedDict{Int, Set}}}(
+        key => Dict{PSI.OptimizationContainerKey, SortedDict{Int, Set}}() for
+        key in CONTAINER_FIELDS
+    )
     container = PSI.get_optimization_container(model)
     for (_, subproblem) in container.subproblems
         _get_axes!(common_axes, subproblem)
@@ -151,5 +158,4 @@ function PSI.solve_impl!(model::PSI.DecisionModel{MultiRegionProblem})
     return
 end
 
-function PSI._check_numerical_bounds(model::PSI.DecisionModel{MultiRegionProblem})
-end
+function PSI._check_numerical_bounds(model::PSI.DecisionModel{MultiRegionProblem}) end

From a4bf11e16f3d97817de53c0a36f07e64e28feb08 Mon Sep 17 00:00:00 2001
From: Jose Daniel Lara <jdlara@berkeley.edu>
Date: Thu, 7 Mar 2024 12:16:11 -0700
Subject: [PATCH 15/17] add write to container loop

---
 src/algorithms/sequential_algorithm.jl | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/src/algorithms/sequential_algorithm.jl b/src/algorithms/sequential_algorithm.jl
index a596e49..c89a7b7 100644
--- a/src/algorithms/sequential_algorithm.jl
+++ b/src/algorithms/sequential_algorithm.jl
@@ -21,7 +21,20 @@ function build_main_problem!(
     sys::PSY.System,
 ) end
 
-function write_results_to_main_container(container::MultiOptimizationContainer) end
+function write_results_to_main_container(container::MultiOptimizationContainer)
+    # TODO: This process needs to work in parallel almost right away
+    for (index, sub_problem) in container.subproblems
+        for field in CONTAINER_FIELDS
+            sub_problem_data_field = getproperty(sub_problem, field)
+            main_container_data_field = getproperty(container, field)
+            for (key, value_container) in sub_problem_data_field
+                # write PSI._jump_value() from the value container to the main_container_data_field
+            end
+        end
+    end
+    # Parameters need a separate approach due to the way the containers work
+    return
+end
 
 function solve_impl!(
     container::MultiOptimizationContainer{SequentialAlgorithm},
@@ -33,6 +46,6 @@ function solve_impl!(
         @info "Solving problem $index"
         status = PSI.solve_impl!(sub_problem, sys)
     end
-    #write_results_to_main_container()
+    write_results_to_main_container(container)
     return status
 end

From 3957b012efd8896f77f4bbb0fa982520267f9572 Mon Sep 17 00:00:00 2001
From: Jose Daniel Lara <jdlara@berkeley.edu>
Date: Thu, 7 Mar 2024 12:18:27 -0700
Subject: [PATCH 16/17] add comment

---
 src/algorithms/sequential_algorithm.jl | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/src/algorithms/sequential_algorithm.jl b/src/algorithms/sequential_algorithm.jl
index c89a7b7..ed82d14 100644
--- a/src/algorithms/sequential_algorithm.jl
+++ b/src/algorithms/sequential_algorithm.jl
@@ -21,6 +21,11 @@ function build_main_problem!(
     sys::PSY.System,
 ) end
 
+
+# The drawback of this approach is that it will loop over the results twice
+# once to write into the main container and a second time when writing into the
+# store. The upside of this approach is that doesn't require overloading write_model_XXX_results!
+# methods from PowerSimulations.
 function write_results_to_main_container(container::MultiOptimizationContainer)
     # TODO: This process needs to work in parallel almost right away
     for (index, sub_problem) in container.subproblems

From 6840d30f508beffff71cf9e94ff7d23a3df9507c Mon Sep 17 00:00:00 2001
From: Daniel Thom <daniel.thom@nrel.gov>
Date: Thu, 14 Mar 2024 09:39:13 -0600
Subject: [PATCH 17/17] Write results from subproblem containers into main
 container

---
 src/algorithms/sequential_algorithm.jl | 44 +++++++++++---
 src/multi_optimization_container.jl    | 80 ++++++++++++--------------
 src/problems/multi_region_problem.jl   |  4 +-
 3 files changed, 77 insertions(+), 51 deletions(-)

diff --git a/src/algorithms/sequential_algorithm.jl b/src/algorithms/sequential_algorithm.jl
index ed82d14..fb3042c 100644
--- a/src/algorithms/sequential_algorithm.jl
+++ b/src/algorithms/sequential_algorithm.jl
@@ -21,19 +21,49 @@ function build_main_problem!(
     sys::PSY.System,
 ) end
 
-
 # The drawback of this approach is that it will loop over the results twice
 # once to write into the main container and a second time when writing into the
 # store. The upside of this approach is that doesn't require overloading write_model_XXX_results!
 # methods from PowerSimulations.
 function write_results_to_main_container(container::MultiOptimizationContainer)
     # TODO: This process needs to work in parallel almost right away
-    for (index, sub_problem) in container.subproblems
+    # TODO: This doesn't handle the case where subproblems have an overlap in axis names.
+
+    for subproblem in values(container.subproblems)
         for field in CONTAINER_FIELDS
-            sub_problem_data_field = getproperty(sub_problem, field)
+            subproblem_data_field = getproperty(subproblem, field)
             main_container_data_field = getproperty(container, field)
-            for (key, value_container) in sub_problem_data_field
-                # write PSI._jump_value() from the value container to the main_container_data_field
+            for (key, src) in subproblem_data_field
+                if src isa JuMP.Containers.SparseAxisArray
+                    @warn "Skip SparseAxisArray" field key
+                    continue
+                end
+                num_dims = ndims(src)
+                num_dims > 2 && error("ndims = $(num_dims) is not supported yet")
+                data = nothing
+                try
+                    data = PSI.jump_value.(src)
+                catch e
+                    if e isa UndefRefError
+                        @warn "Skip UndefRefError for" field key
+                        continue
+                    end
+                    rethrow()
+                end
+                dst = main_container_data_field[key]
+                if num_dims == 1
+                    dst[1:length(axes(src)[1])] = data
+                elseif num_dims == 2
+                    columns = axes(src)[1]
+                    len = length(axes(src)[2])
+                    dst[columns, 1:len] = PSI.jump_value.(src[:, :])
+                elseif num_dims == 3
+                    # TODO: untested
+                    axis1 = axes(src)[1]
+                    axis2 = axes(src)[2]
+                    len = length(axes(src)[3])
+                    dst[axis1, axis2, 1:len] = PSI.jump_value.(src[:, :, :])
+                end
             end
         end
     end
@@ -47,9 +77,9 @@ function solve_impl!(
 )
     # Solve main problem
     status = PSI.RunStatus.SUCCESSFUL
-    for (index, sub_problem) in container.subproblems
+    for (index, subproblem) in container.subproblems
         @info "Solving problem $index"
-        status = PSI.solve_impl!(sub_problem, sys)
+        status = PSI.solve_impl!(subproblem, sys)
     end
     write_results_to_main_container(container)
     return status
diff --git a/src/multi_optimization_container.jl b/src/multi_optimization_container.jl
index ca175b9..683dc46 100644
--- a/src/multi_optimization_container.jl
+++ b/src/multi_optimization_container.jl
@@ -1,5 +1,5 @@
-mutable struct MultiOptimizationContainer{T <: DecompositionAlgorithm} <:
-               PSI.AbstractModelContainer
+Base.@kwdef mutable struct MultiOptimizationContainer{T <: DecompositionAlgorithm} <:
+                           PSI.AbstractModelContainer
     main_problem::PSI.OptimizationContainer
     subproblems::Dict{String, PSI.OptimizationContainer}
     time_steps::UnitRange{Int}
@@ -16,10 +16,8 @@ mutable struct MultiOptimizationContainer{T <: DecompositionAlgorithm} <:
     primal_values_cache::PSI.PrimalValuesCache
     initial_conditions::Dict{PSI.ICKey, Vector{<:PSI.InitialCondition}}
     initial_conditions_data::PSI.InitialConditionsData
-    infeasibility_conflict::Dict{Symbol, Array}
-    pm::Union{Nothing, PM.AbstractPowerModel}
     base_power::Float64
-    optimizer_stats::PSI.OptimizerStats
+    optimizer_stats::PSI.OptimizerStats  # TODO: needs custom struct
     built_for_recurrent_solves::Bool
     metadata::PSI.OptimizationContainerMetadata
     default_time_series_type::Type{<:PSY.TimeSeriesData}  # Maybe isn't needed here
@@ -31,7 +29,7 @@ function MultiOptimizationContainer(
     sys::PSY.System,
     settings::PSI.Settings,
     ::Type{U},
-    sub_problem_keys::Vector{String},
+    subproblem_keys::Vector{String},
 ) where {T <: DecompositionAlgorithm, U <: PSY.TimeSeriesData}
     resolution = PSY.get_time_series_resolution(sys)
     if isabstracttype(U)
@@ -39,36 +37,33 @@ function MultiOptimizationContainer(
     end
 
     # define dictionary containing the optimization container for the subregion
-    subproblems = Dict{String, PSI.OptimizationContainer}()
-    for k in sub_problem_keys
-        subproblems[k] = PSI.OptimizationContainer(sys, settings, nothing, U)
-    end
+    subproblems = Dict(
+        k => PSI.OptimizationContainer(sys, settings, nothing, U) for k in subproblem_keys
+    )
 
     return MultiOptimizationContainer{T}(
-        PSI.OptimizationContainer(sys, settings, nothing, U),
-        subproblems,
-        1:1,
-        IS.time_period_conversion(resolution),
-        settings,
-        PSI.copy_for_serialization(settings),
-        Dict{PSI.VariableKey, AbstractArray}(),
-        Dict{PSI.AuxVarKey, AbstractArray}(),
-        Dict{PSI.ConstraintKey, AbstractArray}(),
-        Dict{PSI.ConstraintKey, AbstractArray}(),
-        PSI.ObjectiveFunction(),
-        Dict{PSI.ExpressionKey, AbstractArray}(),
-        Dict{PSI.ParameterKey, PSI.ParameterContainer}(),
-        PSI.PrimalValuesCache(),
-        Dict{PSI.ICKey, Vector{PSI.InitialCondition}}(),
-        PSI.InitialConditionsData(),
-        Dict{Symbol, Array}(),
-        nothing,
-        PSY.get_base_power(sys),
-        PSI.OptimizerStats(),
-        false,
-        PSI.OptimizationContainerMetadata(),
-        U,
-        nothing,
+        main_problem=PSI.OptimizationContainer(sys, settings, nothing, U),
+        subproblems=subproblems,
+        time_steps=1:1,
+        resolution=IS.time_period_conversion(resolution),
+        settings=settings,
+        settings_copy=PSI.copy_for_serialization(settings),
+        variables=Dict{PSI.VariableKey, AbstractArray}(),
+        aux_variables=Dict{PSI.AuxVarKey, AbstractArray}(),
+        duals=Dict{PSI.ConstraintKey, AbstractArray}(),
+        constraints=Dict{PSI.ConstraintKey, AbstractArray}(),
+        objective_function=PSI.ObjectiveFunction(),
+        expressions=Dict{PSI.ExpressionKey, AbstractArray}(),
+        parameters=Dict{PSI.ParameterKey, PSI.ParameterContainer}(),
+        primal_values_cache=PSI.PrimalValuesCache(),
+        initial_conditions=Dict{PSI.ICKey, Vector{PSI.InitialCondition}}(),
+        initial_conditions_data=PSI.InitialConditionsData(),
+        base_power=PSY.get_base_power(sys),
+        optimizer_stats=PSI.OptimizerStats(),
+        built_for_recurrent_solves=false,
+        metadata=PSI.OptimizationContainerMetadata(),
+        default_time_series_type=U,
+        mpi_info=nothing,
     )
 end
 
@@ -80,8 +75,6 @@ PSI.get_default_time_series_type(container::MultiOptimizationContainer) =
     container.default_time_series_type
 PSI.get_duals(container::MultiOptimizationContainer) = container.duals
 PSI.get_expressions(container::MultiOptimizationContainer) = container.expressions
-PSI.get_infeasibility_conflict(container::MultiOptimizationContainer) =
-    container.infeasibility_conflict
 PSI.get_initial_conditions(container::MultiOptimizationContainer) =
     container.initial_conditions
 PSI.get_initial_conditions_data(container::MultiOptimizationContainer) =
@@ -116,8 +109,8 @@ function get_subproblem(container::MultiOptimizationContainer, id::String)
 end
 
 function check_optimization_container(container::MultiOptimizationContainer)
-    for (index, sub_problem) in container.subproblems
-        PSI.check_optimization_container(sub_problem)
+    for subproblem in values(container.subproblems)
+        PSI.check_optimization_container(subproblem)
     end
     PSI.check_optimization_container(container.main_problem)
     return
@@ -147,9 +140,12 @@ function init_optimization_container!(
         elseif PSI.get_default_time_series_type(container) <: PSY.SingleTimeSeries
             ini_time, _ = PSY.check_time_series_consistency(sys, PSY.SingleTimeSeries)
             PSI.set_initial_time!(settings, ini_time)
+        else
+            error("Bug: unhandled $(PSI.get_default_time_series_type(container))")
         end
     end
 
+    # TODO DT: what if the time series type is SingleTimeSeries?
     if PSI.get_horizon(settings) == PSI.UNSET_HORIZON
         PSI.set_horizon!(settings, PSY.get_forecast_horizon(sys))
     end
@@ -160,13 +156,13 @@ function init_optimization_container!(
 
     # need a special method for the main problem to initialize the optimization container
     # without actually caring about the subnetworks
-    # PSI.init_optimization_container!(sub_problem, network_model, sys)
+    # PSI.init_optimization_container!(subproblem, network_model, sys)
 
-    for (index, sub_problem) in container.subproblems
+    for (index, subproblem) in container.subproblems
         @debug "Initializing Container Subproblem $index" _group =
             PSI.LOG_GROUP_OPTIMIZATION_CONTAINER
-        sub_problem.settings = deepcopy(settings)
-        PSI.init_optimization_container!(sub_problem, network_model, sys)
+        subproblem.settings = deepcopy(settings)
+        PSI.init_optimization_container!(subproblem, network_model, sys)
     end
     _finalize_jump_model!(container, settings)
     return
diff --git a/src/problems/multi_region_problem.jl b/src/problems/multi_region_problem.jl
index dda4aeb..f691b12 100644
--- a/src/problems/multi_region_problem.jl
+++ b/src/problems/multi_region_problem.jl
@@ -83,8 +83,8 @@ function _map_containers(model::PSI.DecisionModel{MultiRegionProblem})
         key in CONTAINER_FIELDS
     )
     container = PSI.get_optimization_container(model)
-    for (_, subproblem) in container.subproblems
-        _get_axes!(common_axes, subproblem)
+    for subproblem_container in values(container.subproblems)
+        _get_axes!(common_axes, subproblem_container)
     end
 
     for (field, vals) in common_axes