Jd/parameter update fixes

src/core/dataset.jl

@@ -20,9 +20,9 @@ end
 
 # Values field is accessed with dot syntax to avoid type instability
 
-mutable struct InMemoryDataset <: AbstractDataset
-    "Data with dimensions (column names, row indexes)"
-    values::DenseAxisArray{Float64, 2}
+mutable struct InMemoryDataset{N} <: AbstractDataset
+    "Data with dimensions (N column names, row indexes)"
+    values::DenseAxisArray{Float64, N}
     # We use Array here to allow for overwrites when updating the state
     timestamps::Vector{Dates.DateTime}
     # Resolution is needed because AbstractDataset might have just one row
@@ -33,12 +33,12 @@ mutable struct InMemoryDataset <: AbstractDataset
 end
 
 function InMemoryDataset(
-    values::DenseAxisArray{Float64, 2},
+    values::DenseAxisArray{Float64, N},
     timestamps::Vector{Dates.DateTime},
     resolution::Dates.Millisecond,
     end_of_step_index::Int,
-)
-    return InMemoryDataset(
+) where {N}
+    return InMemoryDataset{N}(
         values,
         timestamps,
         resolution,
@@ -48,8 +48,8 @@ function InMemoryDataset(
     )
 end
 
-function InMemoryDataset(values::DenseAxisArray{Float64, 2})
-    return InMemoryDataset(
+function InMemoryDataset(values::DenseAxisArray{Float64, N}) where {N}
+    return InMemoryDataset{N}(
         values,
         Vector{Dates.DateTime}(),
         Dates.Second(0.0),
@@ -59,34 +59,95 @@ function InMemoryDataset(values::DenseAxisArray{Float64, 2})
     )
 end
 
-get_num_rows(s::InMemoryDataset) = size(s.values)[2]
+# Helper method for one dimensional cases
+function InMemoryDataset(
+    fill_val::Float64,
+    initial_time::Dates.DateTime,
+    resolution::Dates.Millisecond,
+    end_of_step_index::Int,
+    row_count::Int,
+    column_names::Vector{String})
+    return InMemoryDataset(
+        fill_val,
+        initial_time,
+        resolution,
+        end_of_step_index,
+        row_count,
+        (column_names,),
+    )
+end
+
+function InMemoryDataset(
+    fill_val::Float64,
+    initial_time::Dates.DateTime,
+    resolution::Dates.Millisecond,
+    end_of_step_index::Int,
+    row_count::Int,
+    column_names::NTuple{N, <:Any}) where {N}
+    return InMemoryDataset(
+        fill!(
+            DenseAxisArray{Float64}(undef, column_names..., 1:row_count),
+            fill_val,
+        ),
+        collect(
+            range(
+                initial_time;
+                step = resolution,
+                length = row_count,
+            ),
+        ),
+        resolution,
+        end_of_step_index,
+    )
+end
+
+get_num_rows(s::InMemoryDataset{N}) where {N} = size(s.values)[N]
 
 function make_system_state(
-    values::DenseAxisArray{Float64, 2},
     timestamp::Dates.DateTime,
     resolution::Dates.Millisecond,
-)
-    return InMemoryDataset(values, [timestamp], resolution, 0, 1, UNSET_INI_TIME)
+    columns::NTuple{N, <:Any},
+) where {N}
+    return InMemoryDataset(NaN, timestamp, resolution, 0, 1, columns)
 end
 
-function get_dataset_value(s::InMemoryDataset, date::Dates.DateTime)
+function get_dataset_value(
+    s::T,
+    date::Dates.DateTime,
+) where {T <: Union{InMemoryDataset{1}, InMemoryDataset{2}}}
     s_index = find_timestamp_index(s.timestamps, date)
     if isnothing(s_index)
         error("Request time stamp $date not in the state")
     end
     return s.values[:, s_index]
 end
 
-get_column_names(s::InMemoryDataset) = axes(s.values)[1]
-get_column_names(::OptimizationContainerKey, s::InMemoryDataset) = get_column_names(s)
+function get_dataset_value(s::InMemoryDataset{3}, date::Dates.DateTime)
+    s_index = find_timestamp_index(s.timestamps, date)
+    if isnothing(s_index)
+        error("Request time stamp $date not in the state")
+    end
+    return s.values[:, :, s_index]
+end
+
+function get_column_names(k::OptimizationContainerKey, s::InMemoryDataset)
+    return get_column_names(k, s.values)
+end
 
-function get_last_recorded_value(s::InMemoryDataset)
+function get_last_recorded_value(s::InMemoryDataset{2})
     if get_last_recorded_row(s) == 0
         error("The Dataset hasn't been written yet")
     end
     return s.values[:, get_last_recorded_row(s)]
 end
 
+function get_last_recorded_value(s::InMemoryDataset{3})
+    if get_last_recorded_row(s) == 0
+        error("The Dataset hasn't been written yet")
+    end
+    return s.values[:, :, get_last_recorded_row(s)]
+end
+
 function get_end_of_step_timestamp(s::InMemoryDataset)
     return s.timestamps[s.end_of_step_index]
 end
@@ -110,50 +171,124 @@ function get_value_timestamp(s::InMemoryDataset, date::Dates.DateTime)
     return s.timestamps[s_index]
 end
 
-function set_value!(s::InMemoryDataset, vals::DenseAxisArray{Float64, 2}, index::Int)
+# These set_value! methods expect a single time_step value because they are used to update
+#the state so the incoming vals will have one dimension less than the DataSet. The exception
+# is for vals of Dimension 1 which are still stored in DataSets of dimension 2.
+function set_value!(s::InMemoryDataset{2}, vals::DenseAxisArray{Float64, 2}, index::Int)
     s.values[:, index] = vals[:, index]
     return
 end
 
-function set_value!(s::InMemoryDataset, vals::DenseAxisArray{Float64, 1}, index::Int)
+function set_value!(s::InMemoryDataset{2}, vals::DenseAxisArray{Float64, 1}, index::Int)
     s.values[:, index] = vals
     return
 end
 
+function set_value!(s::InMemoryDataset{3}, vals::DenseAxisArray{Float64, 2}, index::Int)
+    s.values[:, :, index] = vals
+    return
+end
+
 # HDF5Dataset does not account of overwrites in the data. Values are written sequentially.
-mutable struct HDF5Dataset <: AbstractDataset
+mutable struct HDF5Dataset{N} <: AbstractDataset
     values::HDF5.Dataset
     column_dataset::HDF5.Dataset
     write_index::Int
     last_recorded_row::Int
     resolution::Dates.Millisecond
     initial_timestamp::Dates.DateTime
     update_timestamp::Dates.DateTime
-    column_names::Vector{String}
+    column_names::NTuple{N, Vector{String}}
 
-    function HDF5Dataset(values, column_dataset, write_index, last_recorded_row, resolution,
+    function HDF5Dataset{N}(values,
+        column_dataset,
+        write_index,
+        last_recorded_row,
+        resolution,
         initial_timestamp,
-        update_timestamp, column_names,
-    )
-        new(values, column_dataset, write_index, last_recorded_row, resolution,
+        update_timestamp,
+        column_names::NTuple{N, Vector{String}},
+    ) where {N}
+        new{N}(values, column_dataset, write_index, last_recorded_row, resolution,
             initial_timestamp,
             update_timestamp, column_names)
     end
 end
 
-HDF5Dataset(values, column_dataset, resolution, initial_time) =
-    HDF5Dataset(
+function HDF5Dataset{1}(
+    values::HDF5.Dataset,
+    column_dataset::HDF5.Dataset,
+    ::Tuple,
+    resolution::Dates.Millisecond,
+    initial_time::Dates.DateTime,
+)
+    HDF5Dataset{1}(
         values,
         column_dataset,
         1,
         0,
         resolution,
         initial_time,
         UNSET_INI_TIME,
-        column_dataset[:],
+        (column_dataset[:],),
     )
+end
 
-get_column_names(::OptimizationContainerKey, s::HDF5Dataset) = s.column_names
+function HDF5Dataset{2}(
+    values::HDF5.Dataset,
+    column_dataset::HDF5.Dataset,
+    dims::NTuple{4, Int},
+    resolution::Dates.Period,
+    initial_time::Dates.DateTime,
+)
+    # The indexing is done in this way because we save all the names in an
+    # adjacent column entry in the HDF5 Datatset. The indexes for each column
+    # are known because we know how many elements are in each dimension.
+    # the names for the first column are store in the 1:first_column_number_of_elements.
+    col1 = column_dataset[1:dims[2]]
+    # the names for the second column are store in the first_column_number_of elements + 1:end of the column with the names.
+    col2 = column_dataset[(dims[2] + 1):end]
+    HDF5Dataset{2}(
+        values,
+        column_dataset,
+        1,
+        0,
+        resolution,
+        initial_time,
+        UNSET_INI_TIME,
+        (col1, col2),
+    )
+end
+
+function HDF5Dataset{2}(
+    values::HDF5.Dataset,
+    column_dataset::HDF5.Dataset,
+    dims::NTuple{5, Int},
+    resolution::Dates.Period,
+    initial_time::Dates.DateTime,
+)
+    # The indexing is done in this way because we save all the names in an
+    # adjacent column entry in the HDF5 Datatset. The indexes for each column
+    # are known because we know how many elements are in each dimension.
+    # the names for the first column are store in the 1:first_column_number_of_elements.
+    col1 = column_dataset[1:dims[2]]
+    # the names for the second column are store in the first_column_number_of elements + 1:end of the column with the names.
+    col2 = column_dataset[(dims[2] + 1):end]
+    HDF5Dataset{2}(
+        values,
+        column_dataset,
+        1,
+        0,
+        resolution,
+        initial_time,
+        UNSET_INI_TIME,
+        (col1, col2),
+    )
+end
+
+function get_column_names(::OptimizationContainerKey, s::HDF5Dataset)
+    return s.column_names
+end
 
 """
 Return the timestamp from most recent data row updated in the dataset. This value may not be the same as the result from `get_update_timestamp`

src/core/definitions.jl

@@ -24,6 +24,7 @@ const JuMPVariableMatrix = DenseAxisArray{
         JuMP.Containers._AxisLookup{Tuple{Int64, Int64}},
     },
 }
+const JuMPFloatMatrix = DenseAxisArray{Float64, 2}
 const JuMPFloatArray = DenseAxisArray{Float64}
 const JuMPVariableArray = DenseAxisArray{JuMP.VariableRef}
 

src/core/optimizer_stats.jl

@@ -100,5 +100,5 @@ function to_dict(stats::OptimizerStats)
 end
 
 function get_column_names(::Type{OptimizerStats})
-    return collect(string.(fieldnames(OptimizerStats)))
+    return (collect(string.(fieldnames(OptimizerStats))),)
 end

src/core/results_by_time.jl

@@ -1,37 +1,53 @@
-mutable struct ResultsByTime{T}
+mutable struct ResultsByTime{T, N}
     key::OptimizationContainerKey
     data::SortedDict{Dates.DateTime, T}
     resolution::Dates.Period
-    column_names::Vector{String}
+    column_names::NTuple{N, Vector{String}}
 end
 
-function ResultsByTime(key, data, resolution, column_names)
+function ResultsByTime(
+    key::OptimizationContainerKey,
+    data::SortedDict{Dates.DateTime, T},
+    resolution::Dates.Period,
+    column_names,
+) where {T}
     _check_column_consistency(data, column_names)
     ResultsByTime(key, data, resolution, column_names)
 end
 
 function _check_column_consistency(
     data::SortedDict{Dates.DateTime, DenseAxisArray{Float64, 2}},
-    cols::Vector{String},
+    cols::Tuple{Vector{String}},
 )
     for val in values(data)
-        if axes(val)[1] != cols
+        if axes(val)[1] != cols[1]
             error("Mismatch in DenseAxisArray column names: $(axes(val)[1]) $cols")
         end
     end
 end
 
 function _check_column_consistency(
     data::SortedDict{Dates.DateTime, Matrix{Float64}},
-    cols::Vector{String},
+    cols::Tuple{Vector{String}},
 )
     for val in values(data)
-        if size(val)[2] != length(cols)
-            error("Mismatch in length of Matrix columns: $(size(val)[2]) $(length(cols))")
+        if size(val)[2] != length(cols[1])
+            error(
+                "Mismatch in length of Matrix columns: $(size(val)[2]) $(length(cols[1]))",
+            )
         end
     end
 end
 
+function _check_column_consistency(
+    data::SortedDict{Dates.DateTime, DenseAxisArray{Float64, 2}},
+    cols::NTuple{N, Vector{String}},
+) where {N}
+    # TODO:
+end
+
+# TODO: Implement consistency check for other sizes
+
 # This struct behaves like a dict, delegating to its 'data' field.
 Base.length(res::ResultsByTime) = length(res.data)
 Base.iterate(res::ResultsByTime) = iterate(res.data)
@@ -73,6 +89,19 @@ function make_dataframe(
     return df
 end
 
+function make_dataframe(
+    results::ResultsByTime{DenseAxisArray{Float64, 3}},
+    timestamp::Dates.DateTime,
+)
+    df = DataFrames.DataFrame()
+    array = results.data[timestamp]
+    for idx in Iterators.product(array.axes[1:2]...)
+        df[!, "$(idx)"] = array[idx..., :].data
+    end
+    # _add_timestamps!(df, results, timestamp, array)
+    return df
+end
+
 function make_dataframe(results::ResultsByTime{Matrix{Float64}}, timestamp::Dates.DateTime)
     array = results.data[timestamp]
     df = DataFrames.DataFrame(array, results.column_names)

src/operation/decision_model_store.jl

@@ -49,7 +49,7 @@ function initialize_storage!(
             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),
+                    DenseAxisArray{Float64}(undef, column_names..., 1:time_steps_count),
                     NaN,
                 )
             end
@@ -133,5 +133,5 @@ end
 
 function get_column_names(store::DecisionModelStore, key::OptimizationContainerKey)
     container = getfield(store, get_store_container_type(key))
-    return axes(first(values(container[key])))[1]
+    return get_column_names(key, first(values(container[key])))
 end