From 3f67db897f1f876975aafc386c370ad53bd6e6ed Mon Sep 17 00:00:00 2001
From: Jose Daniel Lara <jdlara@berkeley.edu>
Date: Wed, 12 Jun 2024 19:06:20 -0600
Subject: [PATCH] fix branch map for interchange

---
 .../device_constructors/branch_constructor.jl      | 14 +++++++++-----
 src/devices_models/devices/area_interchange.jl     |  6 +++++-
 src/operation/problem_template.jl                  |  4 ++--
 3 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/src/devices_models/device_constructors/branch_constructor.jl b/src/devices_models/device_constructors/branch_constructor.jl
index 288aa2cb8f..4813083f12 100644
--- a/src/devices_models/device_constructors/branch_constructor.jl
+++ b/src/devices_models/device_constructors/branch_constructor.jl
@@ -931,13 +931,17 @@ function construct_device!(
 end
 
 function _get_branch_map(
-    modeled_branch_types::Vector{DataType},
+    network_model::NetworkModel{AreaPTDFPowerModel},
     sys::PSY.System,
 )
+    @assert !isempty(network_model.modeled_branch_types)
     inter_area_branch_map =
         Dict{Tuple{PSY.Area, PSY.Area}, Dict{DataType, Vector{<:PSY.ACBranch}}}()
-    for branch_type in modeled_branch_types
-        for d in get_available_components(branch_type, sys)
+    for branch_type in network_model.modeled_branch_types
+        if branch_type == PSY.AreaInterchange
+            continue
+        end
+        for d in get_available_components(network_model, branch_type, sys)
             area_from = PSY.get_area(PSY.get_arc(d).from)
             area_to = PSY.get_area(PSY.get_arc(d).to)
             if area_from != area_to
@@ -969,7 +973,7 @@ function construct_device!(
     # Not ideal to do this here, but it is a not terrible workaround
     # The area interchanges are like a services/device mix.
     # Doesn't include the possibility of Multi-terminal HVDC
-    inter_area_branch_map = _get_branch_map(network_model.modeled_branch_types, sys)
+    inter_area_branch_map = _get_branch_map(network_model, sys)
 
     add_constraints!(
         container,
@@ -999,7 +1003,7 @@ function construct_device!(
     model::DeviceModel{PSY.AreaInterchange, StaticBranchUnbounded},
     network_model::NetworkModel{AreaPTDFPowerModel},
 )
-    inter_area_branch_map = _get_branch_map(network_model.modeled_branch_types, sys)
+    inter_area_branch_map = _get_branch_map(network_model, sys)
     # Not ideal to do this here, but it is a not terrible workaround
     # The area interchanges are like a services/device mix.
     # Doesn't include the possibility of Multi-terminal HVDC
diff --git a/src/devices_models/devices/area_interchange.jl b/src/devices_models/devices/area_interchange.jl
index f433bc439f..59d1686b0b 100644
--- a/src/devices_models/devices/area_interchange.jl
+++ b/src/devices_models/devices/area_interchange.jl
@@ -122,8 +122,12 @@ function add_constraints!(
     devices::IS.FlattenIteratorWrapper{PSY.AreaInterchange},
     model::DeviceModel{PSY.AreaInterchange, <:AbstractBranchFormulation},
     network_model::NetworkModel{AreaPTDFPowerModel},
-    inter_area_branch_map,
+    inter_area_branch_map::Dict{
+        Tuple{PSY.Area, PSY.Area},
+        Dict{DataType, Vector{<:PSY.ACBranch}},
+    },
 )
+    @assert !isempty(inter_area_branch_map)
     time_steps = get_time_steps(container)
     device_names = [PSY.get_name(d) for d in devices]
 
diff --git a/src/operation/problem_template.jl b/src/operation/problem_template.jl
index 32b0ab7f44..6902f2be39 100644
--- a/src/operation/problem_template.jl
+++ b/src/operation/problem_template.jl
@@ -311,8 +311,8 @@ end
 function _add_modeled_lines!(template::ProblemTemplate, sys::PSY.System)
     network_model = get_network_model(template)
     branch_models = get_branch_models(template)
-    for k in keys(branch_models)
-        push!(network_model.modeled_branch_types)
+    for v in values(branch_models)
+        push!(network_model.modeled_branch_types, get_component_type(v))
     end
     return
 end