add basic code

src/PowerSimulationsDecomposition.jl

@@ -1,10 +1,26 @@
 module PowerSimulationsDecomposition
 
-using DocStringExtensions
+import PowerSimulations
+import JuMP
+import Dates
+import MPI
+
+const PSI = PowerSimulations
+const PM = PSI.PM
+const PSY = PSI.PSY
+const IS = PSI.IS
 
+using DocStringExtensions
 @template (FUNCTIONS, METHODS) = """
                                  $(TYPEDSIGNATURES)
                                  $(DOCSTRING)
                                  """
 
+
+include("core.jl")
+include("multi_optimization_container.jl")
+include("algorithms/sequential_algorithm.jl")
+include("algorithms/mpi_parallel_algorithm.jl")
+include("problems/multi_region_problem.jl")
+
 end

src/algorithms/mpi_parallel_algorithm.jl

@@ -0,0 +1,22 @@
+"""
+Default solve method for MultiOptimizationContainer
+"""
+function solve_impl!(
+    container::MultiOptimizationContainer{MPIParallelAlgorithm},
+    system::PSY.System
+    )
+    # 1. Initialize MPI
+    MPI.Init()
+    try
+        mpi = MpiInfo(MPI.COMM_WORLD)
+        # TODO: Look for for loop MPI over dicts to solve subproblems
+        solution = solve_subproblem(sp, params, method.inner_method)
+        MPI.Barrier(mpi.comm)
+        compute_main_problem!(container, mpi, system, solution)
+        # Finish loop?
+    finally
+        update_results!(container, system)
+        MPI.Finalize()
+    end
+    return
+end

src/algorithms/sequential_algorithm.jl

@@ -0,0 +1,31 @@
+function build_impl!(
+    container::MultiOptimizationContainer{SequentialAlgorithm},
+    template::PSI.ProblemTemplate,
+    sys::PSY.System,
+)
+
+    for (index, sub_problem) in container.subproblems
+        @debug "Building Subproblem $index" _group = PSI.LOG_GROUP_OPTIMIZATION_CONTAINER
+        PSI.build_impl!(sub_problem, template, sys)
+    end
+
+    build_main_problem!(container, template, sys)
+
+    check_optimization_container(container)
+
+    return
+end
+
+function build_main_problem!(
+    container::MultiOptimizationContainer{SequentialAlgorithm},
+    template::PSI.ProblemTemplate,
+    sys::PSY.System,)
+end
+
+function solve_impl!(container::MultiOptimizationContainer{SequentialAlgorithm}, sys::PSY.System)
+    # Solve main problem
+    for (index, sub_problem) in container.subproblems
+        status = PSI.solve_impl!(sub_problem, sys)
+    end
+    #write_results_to_main_container()
+end

src/core.jl

@@ -0,0 +1,21 @@
+struct MultiRegionProblem <: PSI.DecisionProblem end
+
+abstract type DecompositionAlgorithm end
+
+struct SequentialAlgorithm <: DecompositionAlgorithm end
+struct MPIParallelAlgorithm <: DecompositionAlgorithm end
+
+# Taken from https://github.com/ANL-CEEESA/UnitCommitment.jl/blob/dev/src/solution/methods/ProgressiveHedging/structs.jl
+struct MpiInfo
+    comm::Any
+    rank::Int
+    root::Bool
+    nprocs::Int
+
+    function MpiInfo(comm)
+        rank = MPI.Comm_rank(comm) + 1
+        is_root = (rank == 1)
+        nprocs = MPI.Comm_size(comm)
+        return new(comm, rank, is_root, nprocs)
+    end
+end

src/multi_optimization_container.jl

@@ -0,0 +1,185 @@
+mutable struct MultiOptimizationContainer{T<:DecompositionAlgorithm} <: PSI.AbstractModelContainer
+    main_JuMPmodel::JuMP.Model
+    subproblems::Dict{String, PSI.OptimizationContainer}
+    time_steps::UnitRange{Int}
+    resolution::Dates.TimePeriod
+    settings::PSI.Settings
+    settings_copy::PSI.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
+    infeasibility_conflict::Dict{Symbol, Array}
+    pm::Union{Nothing, PM.AbstractPowerModel}
+    base_power::Float64
+    optimizer_stats::PSI.OptimizerStats
+    built_for_recurrent_solves::Bool
+    metadata::PSI.OptimizationContainerMetadata # Unclear about how to extend this correctly
+    default_time_series_type::Type{<:PSY.TimeSeriesData}  # Maybe isn't needed here
+    mpi_info::Union{Nothing, MpiInfo}
+end
+
+function MultiOptimizationContainer(
+    ::Type{T},
+    sys::PSY.System,
+    settings::PSI.Settings,
+    ::Type{U},
+    sub_problem_keys::Vector{String}
+) where {T <: DecompositionAlgorithm, U <: PSY.TimeSeriesData}
+    resolution = PSY.get_time_series_resolution(sys)
+    if isabstracttype(U)
+        error("Default Time Series Type $U can't be abstract")
+    end
+
+    subproblems = Dict{String, PSI.OptimizationContainer}()
+    for k in sub_problem_keys
+        subproblems[k] = PSI.OptimizationContainer(sys, settings, nothing, U)
+    end
+
+    return MultiOptimizationContainer{T}(
+        JuMP.Model(),
+        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,
+    )
+end
+
+function get_container_keys(container::MultiOptimizationContainer)
+    return Iterators.flatten(keys(getfield(container, f)) for f in STORE_CONTAINERS)
+end
+
+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) =
+    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_metadata(container::MultiOptimizationContainer) = container.metadata
+PSI.get_optimizer_stats(container::MultiOptimizationContainer) = container.optimizer_stats
+PSI.get_parameters(container::MultiOptimizationContainer) = container.parameters
+PSI.get_resolution(container::MultiOptimizationContainer) = container.resolution
+PSI.get_settings(container::MultiOptimizationContainer) = container.settings
+PSI.get_time_steps(container::MultiOptimizationContainer) = container.time_steps
+PSI.get_variables(container::MultiOptimizationContainer) = container.variables
+
+PSI.set_initial_conditions_data!(container::MultiOptimizationContainer, data) =
+    container.initial_conditions_data = data
+    PSI.get_objective_expression(container::MultiOptimizationContainer) = container.objective_function
+    PSI.is_synchronized(container::MultiOptimizationContainer) =
+    container.objective_function.synchronized
+PSI.set_time_steps!(container::MultiOptimizationContainer, time_steps::UnitRange{Int64}) =
+    container.time_steps = time_steps
+
+PSI.get_aux_variables(container::MultiOptimizationContainer) = container.aux_variables
+PSI.get_base_power(container::MultiOptimizationContainer) = container.base_power
+PSI.get_constraints(container::MultiOptimizationContainer) = container.constraints
+
+
+function check_optimization_container(container::MultiOptimizationContainer)
+    # Solve main problem
+    for (index, sub_problem) in container.subproblems
+        PSI.check_optimization_container(sub_problem)
+    end
+end
+
+function _finalize_jump_model!(container::MultiOptimizationContainer, 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
+    return
+end
+
+function init_optimization_container!(
+    container::MultiOptimizationContainer,
+    ::Type{T},
+    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)
+
+    if PSI.get_initial_time(settings) == PSI.UNSET_INI_TIME
+        if PSI.get_default_time_series_type(container) <: PSY.AbstractDeterministic
+            PSI.set_initial_time!(settings, PSY.get_forecast_initial_timestamp(sys))
+        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)
+        end
+    end
+
+    if PSI.get_horizon(settings) == PSI.UNSET_HORIZON
+        PSI.set_horizon!(settings, PSY.get_forecast_horizon(sys))
+    end
+    container.time_steps = 1:PSI.get_horizon(settings)
+
+    stats = PSI.get_optimizer_stats(container)
+    stats.detailed_stats = PSI.get_detailed_optimizer_stats(settings)
+
+    _finalize_jump_model!(container, settings)
+
+    for (index, sub_problem) 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, T, sys)
+    end
+
+    return
+end

src/problems/multi_region_problem.jl

@@ -0,0 +1,99 @@
+function PSI.DecisionModel{MultiRegionProblem}(
+    template::PSI.ProblemTemplate,
+    sys::PSY.System,
+    settings::PSI.Settings,
+    ::Union{Nothing,JuMP.Model}=nothing;
+    name=nothing,
+)
+    if name === nothing
+        name = nameof(MultiRegionProblem)
+    elseif name isa String
+        name = Symbol(name)
+    end
+    # Get these from the system later with subsystems
+    region_keys = ["1", "2,", "3"]
+    internal = PSI.ModelInternal(
+        MultiOptimizationContainer(SequentialAlgorithm, sys, settings, PSY.Deterministic, region_keys),
+    )
+    template_ = deepcopy(template)
+    PSI.finalize_template!(template_, sys)
+    return PSI.DecisionModel{MultiRegionProblem}(
+        name,
+        template_,
+        sys,
+        internal,
+        PSI.DecisionModelStore(),
+        Dict{String,Any}(),
+    )
+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)
+    # 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))
+    return
+end
+
+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_system(model),
+    )
+    @info "Initializing ModelStoreParams"
+    PSI.init_model_store_params!(model)
+    PSI.set_status!(model, PSI.BuildStatus.IN_PROGRESS)
+    return
+end
+
+function handle_initial_conditions!(model::PSI.DecisionModel{MultiRegionProblem})
+end
+
+function instantiate_network_model(model::PSI.DecisionModel{MultiRegionProblem})
+    PSI.instantiate_network_model(model)
+end
+
+function PSI.build_model!(model::PSI.DecisionModel{MultiRegionProblem})
+    build_impl!(PSI.get_optimization_container(model), PSI.get_template(model), PSI.get_system(model))
+end
+
+function PSI.solve_impl!(model::PSI.DecisionModel{MultiRegionProblem})
+    solve_impl!(PSI.get_optimization_container(model), PSI.get_system(model))
+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