From 38279950e316c87968104e7dc6c5439b46dae356 Mon Sep 17 00:00:00 2001
From: GabrielKS <23368820+GabrielKS@users.noreply.github.com>
Date: Tue, 21 Jan 2025 14:22:21 -0700
Subject: [PATCH] Add fields to `PowerFlowEvaluationModel`s

---
 src/PowerFlowData.jl   |  2 +-
 src/PowerFlows.jl      |  1 +
 src/powerflow_types.jl | 41 +++++++++++++++++++++++++++++++++++------
 src/psse_export.jl     |  1 +
 4 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/src/PowerFlowData.jl b/src/PowerFlowData.jl
index 021b99d8..749188b7 100644
--- a/src/PowerFlowData.jl
+++ b/src/PowerFlowData.jl
@@ -1,4 +1,4 @@
-Base.@kwdef mutable struct SolverData
+@kwdef mutable struct SolverData
     J::Union{SparseMatrixCSC{Float64, Int}, Nothing} = nothing
     dSbus_dV_ref::Union{Vector{Float64}, Nothing} = nothing
 end
diff --git a/src/PowerFlows.jl b/src/PowerFlows.jl
index 13672ffe..c0179221 100644
--- a/src/PowerFlows.jl
+++ b/src/PowerFlows.jl
@@ -18,6 +18,7 @@ export write_export
 export get_psse_export_paths
 export penalty_factors
 
+import Base: @kwdef
 import Logging
 import DataFrames
 import PowerSystems
diff --git a/src/powerflow_types.jl b/src/powerflow_types.jl
index 96685f68..4926ef8d 100644
--- a/src/powerflow_types.jl
+++ b/src/powerflow_types.jl
@@ -6,24 +6,53 @@ struct NLSolveACPowerFlow <: ACPowerFlowSolverType end
 
 struct ACPowerFlow{ACSolver <: ACPowerFlowSolverType} <: PowerFlowEvaluationModel
     check_reactive_power_limits::Bool
+    exporter::Union{Nothing, PowerFlowEvaluationModel}
 end
 
 ACPowerFlow{ACSolver}(;
     check_reactive_power_limits::Bool = false,
+    exporter::Union{Nothing, PowerFlowEvaluationModel} = nothing,
 ) where {ACSolver <: ACPowerFlowSolverType} =
-    ACPowerFlow{ACSolver}(check_reactive_power_limits)
+    ACPowerFlow{ACSolver}(check_reactive_power_limits, exporter)
 
 ACPowerFlow(
     ACSolver::Type{<:ACPowerFlowSolverType} = KLUACPowerFlow;
     check_reactive_power_limits::Bool = false,
-) = ACPowerFlow{ACSolver}(check_reactive_power_limits)
+    exporter::Union{Nothing, PowerFlowEvaluationModel} = nothing,
+) = ACPowerFlow{ACSolver}(check_reactive_power_limits, exporter)
 
-struct DCPowerFlow <: PowerFlowEvaluationModel end
-struct PTDFDCPowerFlow <: PowerFlowEvaluationModel end
-struct vPTDFDCPowerFlow <: PowerFlowEvaluationModel end
+@kwdef struct DCPowerFlow <: PowerFlowEvaluationModel
+    exporter::Union{Nothing, PowerFlowEvaluationModel} = nothing
+end
+
+@kwdef struct PTDFDCPowerFlow <: PowerFlowEvaluationModel
+    exporter::Union{Nothing, PowerFlowEvaluationModel} = nothing
+end
+
+@kwdef struct vPTDFDCPowerFlow <: PowerFlowEvaluationModel
+    exporter::Union{Nothing, PowerFlowEvaluationModel} = nothing
+end
 
-Base.@kwdef struct PSSEExportPowerFlow <: PowerFlowEvaluationModel
+@kwdef struct PSSEExportPowerFlow <: PowerFlowEvaluationModel
     psse_version::Symbol
     export_dir::AbstractString
     write_comments::Bool = false
+    overwrite::Bool = false
+end
+
+get_exporter(pfem::PowerFlowEvaluationModel) = pfem.exporter
+get_exporter(::PSSEExportPowerFlow) = nothing
+
+"""
+Expand a single `PowerFlowEvaluationModel` into its possibly multiple parts for separate
+evaluation. Namely, if `pfem` contains a non-nothing `exporter`, return `[pfem, exporter]`,
+else return `[pfem]`.
+"""
+function flatten_power_flow_evaluation_model(pfem::PowerFlowEvaluationModel)
+    exporter = get_exporter(pfem)
+    return if isnothing(exporter)
+        PowerFlowEvaluationModel[pfem]
+    else
+        PowerFlowEvaluationModel[pfem, exporter]
+    end
 end
diff --git a/src/psse_export.jl b/src/psse_export.jl
index b22ccaff..2168ad9f 100644
--- a/src/psse_export.jl
+++ b/src/psse_export.jl
@@ -1337,6 +1337,7 @@ make_power_flow_container(pfem::PSSEExportPowerFlow, sys::PSY.System; kwargs...)
         pfem.export_dir;
         write_comments = pfem.write_comments,
         step = (0, 0),
+        overwrite = pfem.overwrite,
     )
 
 solve_powerflow!(exporter::PSSEExporter) = write_export(exporter)