add stores

src/PowerSystemsInvestments.jl

@@ -24,6 +24,7 @@ include("base/capital_model.jl")
 include("base/operation_model.jl")
 include("base/feasibility_model.jl")
 include("base/investment_problem.jl")
+include("base/investment_model_store.jl")
 include("base/investment_model.jl")
 
 @template (FUNCTIONS, METHODS) = """

src/base/investment_model.jl

@@ -1,4 +1,6 @@
 mutable struct InvestmentModel{<: InvestmentProblem}
     portfolio::PSIP.Portfolio
-    container::OptimizationContainer
+    internal::Union{Nothing, IS.ModelInternal}
+    store::InvestmentModelStore
+    ext::Dict{String, Any}
 end

src/base/investment_model_store.jl

@@ -0,0 +1,140 @@
+"""
+Stores results data for one InvestmentModel
+"""
+mutable struct InvestmentModelStore <: IS.AbstractModelStore
+    # All DenseAxisArrays have axes (column names, row indexes)
+    duals::Dict{ConstraintKey, OrderedDict{Dates.DateTime, DenseAxisArray{Float64, 2}}}
+    parameters::Dict{
+        ParameterKey,
+        OrderedDict{Dates.DateTime, DenseAxisArray{Float64, 2}},
+    }
+    variables::Dict{VariableKey, OrderedDict{Dates.DateTime, DenseAxisArray{Float64, 2}}}
+    aux_variables::Dict{AuxVarKey, OrderedDict{Dates.DateTime, DenseAxisArray{Float64, 2}}}
+    expressions::Dict{
+        ExpressionKey,
+        OrderedDict{Dates.DateTime, DenseAxisArray{Float64, 2}},
+    }
+    optimizer_stats::OrderedDict{Dates.DateTime, OptimizerStats}
+end
+
+function InvestmentModelStore()
+    return InvestmentModelStore(
+        Dict{ConstraintKey, OrderedDict{Dates.DateTime, DenseAxisArray{Float64, 2}}}(),
+        Dict{ParameterKey, OrderedDict{Dates.DateTime, DenseAxisArray{Float64, 2}}}(),
+        Dict{VariableKey, OrderedDict{Dates.DateTime, DenseAxisArray{Float64, 2}}}(),
+        Dict{AuxVarKey, OrderedDict{Dates.DateTime, DenseAxisArray{Float64, 2}}}(),
+        Dict{ExpressionKey, OrderedDict{Dates.DateTime, DenseAxisArray{Float64, 2}}}(),
+        OrderedDict{Dates.DateTime, OptimizerStats}(),
+    )
+end
+
+function initialize_storage!(
+    store::InvestmentModelStore,
+    container::AbstractOptimizationContainer,
+    params::ModelStoreParams,
+)
+    num_of_executions = get_num_executions(params)
+    if length(get_time_steps(container)) < 1
+        error("The time step count in the optimization container is not defined")
+    end
+    time_steps_count = get_time_steps(container)[end]
+    initial_time = get_initial_time(container)
+    model_interval = get_interval(params)
+    for type in STORE_CONTAINERS
+        field_containers = getfield(container, type)
+        results_container = getfield(store, type)
+        for (key, field_container) in field_containers
+            !should_write_resulting_value(key) && continue
+            @debug "Adding $(encode_key_as_string(key)) to InvestmentModelStore" _group =
+                LOG_GROUP_MODEL_STORE
+            column_names = get_column_names(key, field_container)
+            data = OrderedDict{Dates.DateTime, DenseAxisArray{Float64, 2}}()
+            for timestamp in
+                range(initial_time; step = model_interval, length = num_of_executions)
+                data[timestamp] = fill!(
+                    DenseAxisArray{Float64}(undef, column_names..., 1:time_steps_count),
+                    NaN,
+                )
+            end
+            results_container[key] = data
+        end
+    end
+    return
+end
+
+function write_result!(
+    store::InvestmentModelStore,
+    name::Symbol,
+    key::OptimizationContainerKey,
+    index::InvestmentModelIndexType,
+    update_timestamp::Dates.DateTime,
+    array::DenseAxisArray{<:Any, 2},
+)
+    columns = axes(array)[1]
+    if eltype(columns) !== String
+        # TODO: This happens because buses are stored by indexes instead of name.
+        columns = string.(columns)
+    end
+    container = getfield(store, get_store_container_type(key))
+    container[key][index] = DenseAxisArray(array.data, columns, 1:size(array)[2])
+    return
+end
+
+function write_result!(
+    store::InvestmentModelStore,
+    name::Symbol,
+    key::OptimizationContainerKey,
+    index::InvestmentModelIndexType,
+    update_timestamp::Dates.DateTime,
+    array::DenseAxisArray{<:Any, 1},
+)
+    columns = axes(array)[1]
+    if eltype(columns) !== String
+        # TODO: This happens because buses are stored by indexes instead of name.
+        columns = string.(columns)
+    end
+    container = getfield(store, get_store_container_type(key))
+    container[key][index] =
+        DenseAxisArray(reshape(array.data, 1, length(columns)), ["1"], columns)
+    return
+end
+
+function read_results(
+    store::InvestmentModelStore,
+    key::OptimizationContainerKey;
+    index::Union{InvestmentModelIndexType, Nothing} = nothing,
+)
+    container = getfield(store, get_store_container_type(key))
+    data = container[key]
+    if isnothing(index)
+        @assert length(data) == 1
+        index = first(keys(data))
+    end
+
+    # Return a copy because callers may mutate it.
+    return deepcopy(data[index])
+end
+
+function write_optimizer_stats!(
+    store::InvestmentModelStore,
+    stats::OptimizerStats,
+    index::InvestmentModelIndexType,
+)
+    if index in keys(store.optimizer_stats)
+        @warn "Overwriting optimizer stats"
+    end
+    store.optimizer_stats[index] = stats
+    return
+end
+
+function read_optimizer_stats(store::InvestmentModelStore)
+    stats = [to_namedtuple(x) for x in values(store.optimizer_stats)]
+    df = DataFrames.DataFrame(stats)
+    DataFrames.insertcols!(df, 1, :DateTime => keys(store.optimizer_stats))
+    return df
+end
+
+function get_column_names(store::InvestmentModelStore, key::OptimizationContainerKey)
+    container = getfield(store, get_store_container_type(key))
+    return get_column_names(key, first(values(container[key])))
+end

src/base/variables.jl

@@ -1 +1,4 @@
+"""
+Total installed capacity for a technology
+"""
 struct InstalledCapacity <: IS.VariableType end