Skip to content

Commit

Permalink
refactor datasets for emulator
Browse files Browse the repository at this point in the history
  • Loading branch information
jd-lara committed Mar 26, 2024
1 parent 19f614d commit b8af24a
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 61 deletions.
2 changes: 1 addition & 1 deletion src/core/dataset.jl
Original file line number Diff line number Diff line change
Expand Up @@ -313,6 +313,6 @@ end

function set_value!(s::HDF5Dataset, vals, index::Int)
# Temporary while there is no implementation of caching of em_data
_write_dataset!(s.values, vals, index:index)
_write_dataset!(s.values, vals, index)

Check warning on line 316 in src/core/dataset.jl

View check run for this annotation

Codecov / codecov/patch

src/core/dataset.jl#L316

Added line #L316 was not covered by tests
return
end
78 changes: 18 additions & 60 deletions src/simulation/hdf_simulation_store.jl
Original file line number Diff line number Diff line change
Expand Up @@ -507,7 +507,7 @@ function _read_result(
#end
columns = get_column_names(key, dataset)
data = permutedims(data)
@assert_op size(data)[2] == length(columns)
@assert_op size(data)[2] == length(columns[1])

Check warning on line 510 in src/simulation/hdf_simulation_store.jl

View check run for this annotation

Codecov / codecov/patch

src/simulation/hdf_simulation_store.jl#L510

Added line #L510 was not covered by tests
@assert_op size(data)[1] == 1
return data, columns
end
Expand Down Expand Up @@ -644,41 +644,15 @@ function write_result!(
key::OptimizationContainerKey,
index::EmulationModelIndexType,
simulation_time::Dates.DateTime,
data::Array{Float64},
data::DenseAxisArray,
)
dataset = _get_em_dataset(store, key)
_write_dataset!(dataset.values, data, index)
_write_dataset!(dataset.values, to_matrix(data), index)
set_last_recorded_row!(dataset, index)
set_update_timestamp!(dataset, simulation_time)
return
end

function write_result!(
store::HdfSimulationStore,
model_name::Symbol,
key::OptimizationContainerKey,
index::EmulationModelIndexType,
simulation_time::Dates.DateTime,
data::DenseAxisArray{Float64, 2},
)
data_array = Array{Float64, 3}(undef, size(data)[1], size(data)[2], 1)
data_array[:, :, 1] = data
write_result!(store, model_name, key, index, simulation_time, data_array)
return
end

function write_result!(
store::HdfSimulationStore,
model_name::Symbol,
key::OptimizationContainerKey,
index::EmulationModelIndexType,
simulation_time::Dates.DateTime,
data::DenseAxisArray{Float64, 1},
)
write_result!(store, model_name, key, index, simulation_time, to_matrix(data))
return
end

function _check_state(store::HdfSimulationStore)
if has_dirty(store.cache)
error("BUG!!! dirty cache is present at shutdown: $(store.file)")
Expand Down Expand Up @@ -1014,60 +988,44 @@ function _read_length(::Type{OptimizerStats}, store::HdfSimulationStore)
return HDF5.read(HDF5.attributes(dataset), "columns")
end

# Specific data set writing function that writes decision model data. It dispatches on the index type of the dataset as a range
function _write_dataset!(
dataset::HDF5.Dataset,
array::Matrix{Float64},
row_range::UnitRange{Int64},
::Val{3},
)
dataset[:, 1, row_range] = array
@debug "wrote dataset" dataset row_range
return
end

function _write_dataset!(
dataset::HDF5.Dataset,
array::Matrix{Float64},
array::Array{Float64, 3},
row_range::UnitRange{Int64},
::Val{2},
)
dataset[row_range, :] = array
@debug "wrote dataset" dataset row_range
dataset[:, :, row_range] = array
@debug "wrote dm dataset" dataset row_range
return
end

function _write_dataset!(
dataset::HDF5.Dataset,
array::Array{Float64, 3},
array::Array{Float64, 4},
row_range::UnitRange{Int64},
::Val{3},
)
dataset[row_range, :, :] = array
@debug "wrote dataset" dataset row_range
return
end

function _write_dataset!(dataset::HDF5.Dataset, array::Array{Float64}, index::Int)
_write_dataset!(dataset, array, index:index, Val{ndims(dataset)}())
dataset[:, :, :, row_range] = array
@debug "wrote dm dataset" dataset row_range

Check warning on line 1008 in src/simulation/hdf_simulation_store.jl

View check run for this annotation

Codecov / codecov/patch

src/simulation/hdf_simulation_store.jl#L1007-L1008

Added lines #L1007 - L1008 were not covered by tests
return
end

# Specific data set writing function that writes emulation model data. It dispatches on the index type of the dataset
function _write_dataset!(
dataset::HDF5.Dataset,
array::Array{Float64, 3},
row_range::UnitRange{Int64},
array::Array{Float64, 2},
index::EmulationModelIndexType,
)
dataset[:, :, row_range] = array
@debug "wrote dataset" dataset row_range
dataset[index, :] = array
@debug "wrote em dataset" dataset index
return
end

function _write_dataset!(
dataset::HDF5.Dataset,
array::Array{Float64, 4},
row_range::UnitRange{Int64},
index::EmulationModelIndexType,
)
dataset[:, :, :, row_range] = array
@debug "wrote dataset" dataset row_range
dataset[index, :, :] = array
@debug "wrote em dataset" dataset index

Check warning on line 1029 in src/simulation/hdf_simulation_store.jl

View check run for this annotation

Codecov / codecov/patch

src/simulation/hdf_simulation_store.jl#L1028-L1029

Added lines #L1028 - L1029 were not covered by tests
return
end

0 comments on commit b8af24a

Please sign in to comment.