From d6c7d8a0b2cdd79efcccfb3ea5615d5532aa1407 Mon Sep 17 00:00:00 2001 From: James Tomlinson Date: Fri, 12 Apr 2024 16:15:15 +0100 Subject: [PATCH] feat: Add additional metadata to HDF5 outputs. (#157) --- pywr-core/src/recorders/csv.rs | 4 +- pywr-core/src/recorders/hdf.rs | 97 ++++++++++++++++++++++++++++++++-- pywr-core/src/scenario.rs | 15 +++--- 3 files changed, 101 insertions(+), 15 deletions(-) diff --git a/pywr-core/src/recorders/csv.rs b/pywr-core/src/recorders/csv.rs index c39500b6..6858441f 100644 --- a/pywr-core/src/recorders/csv.rs +++ b/pywr-core/src/recorders/csv.rs @@ -106,8 +106,8 @@ impl Recorder for CsvWideFmtOutput { // This is a vec of vec for each scenario group let mut header_scenario_groups = Vec::new(); - for group_name in domain.scenarios().group_names() { - header_scenario_groups.push(vec![format!("scenario-group: {}", group_name)]); + for group in domain.scenarios().groups() { + header_scenario_groups.push(vec![format!("scenario-group: {}", group.name())]); } for scenario_index in domain.scenarios().indices().iter() { diff --git a/pywr-core/src/recorders/hdf.rs b/pywr-core/src/recorders/hdf.rs index 79ae9535..2be1d02a 100644 --- a/pywr-core/src/recorders/hdf.rs +++ b/pywr-core/src/recorders/hdf.rs @@ -2,7 +2,7 @@ use super::{MetricSetState, PywrError, Recorder, RecorderMeta, Timestep}; use crate::models::ModelDomain; use crate::network::Network; use crate::recorders::MetricSetIndex; -use crate::scenario::ScenarioIndex; +use crate::scenario::{ScenarioDomain, ScenarioIndex}; use crate::state::State; use chrono::{Datelike, Timelike}; use hdf5::{Extents, Group}; @@ -10,7 +10,14 @@ use ndarray::{s, Array1}; use std::any::Any; use std::ops::Deref; use std::path::PathBuf; +use std::str::FromStr; +/// A recorder that saves model outputs to an HDF5 file. +/// +/// This recorder saves the model outputs to an HDF5 file. The file will contain a number of groups +/// and datasets that correspond to the metrics in the metric set. Additionally, the file will +/// contain metadata about the time steps and scenarios that were used in the model simulation. +/// #[derive(Clone, Debug)] pub struct HDF5Recorder { meta: RecorderMeta, @@ -69,7 +76,9 @@ impl Recorder for HDF5Recorder { Ok(f) => f, Err(e) => return Err(PywrError::HDF5Error(e.to_string())), }; - let mut datasets = Vec::new(); + + write_pywr_metadata(&file)?; + write_scenarios_metadata(&file, domain.scenarios())?; // Create the time table let dates: Array1<_> = domain.time().timesteps().iter().map(DateTime::from_timestamp).collect(); @@ -83,12 +92,14 @@ impl Recorder for HDF5Recorder { let metric_set = network.get_metric_set(self.metric_set_idx)?; + let mut datasets = Vec::new(); + for metric in metric_set.iter_metrics() { let name = metric.name(network)?; let sub_name = metric.sub_name(network)?; let attribute = metric.attribute(); - let ds = require_node_dataset(root_grp, shape, name, sub_name, attribute)?; + let ds = require_metric_dataset(root_grp, shape, name, sub_name, attribute)?; datasets.push(ds); } @@ -162,7 +173,7 @@ fn require_dataset>(parent: &Group, shape: S, name: &str) -> Re } /// Create a node dataset in /parent/name/sub_name/attribute -fn require_node_dataset>( +fn require_metric_dataset>( parent: &Group, shape: S, name: &str, @@ -193,3 +204,81 @@ fn require_group(parent: &Group, name: &str) -> Result { } } } + +fn write_pywr_metadata(file: &hdf5::File) -> Result<(), PywrError> { + let root = file.deref(); + + let grp = require_group(root, "pywr")?; + + // Write the Pywr version as an attribute + const VERSION: &str = env!("CARGO_PKG_VERSION"); + let version = hdf5::types::VarLenUnicode::from_str(VERSION).map_err(|e| PywrError::HDF5Error(e.to_string()))?; + + let attr = grp + .new_attr::() + .shape(()) + .create("pywr-version") + .map_err(|e| PywrError::HDF5Error(e.to_string()))?; + attr.as_writer() + .write_scalar(&version) + .map_err(|e| PywrError::HDF5Error(e.to_string()))?; + + Ok(()) +} + +#[derive(hdf5::H5Type, Clone, PartialEq, Debug)] +#[repr(C)] +pub struct ScenarioGroupEntry { + pub name: hdf5::types::VarLenUnicode, + pub size: usize, +} + +#[derive(hdf5::H5Type, Clone, PartialEq, Debug)] +#[repr(C)] +pub struct H5ScenarioIndex { + index: usize, + indices: hdf5::types::VarLenArray, +} + +/// Write scenario metadata to the HDF5 file. +/// +/// This function will create the `/scenarios` group in the HDF5 file and write the scenario +/// groups and indices into `/scenarios/groups` and `/scenarios/indices` respectively. +fn write_scenarios_metadata(file: &hdf5::File, domain: &ScenarioDomain) -> Result<(), PywrError> { + // Create the scenario group and associated datasets + let grp = require_group(file.deref(), "scenarios")?; + + let scenario_groups: Array1 = domain + .groups() + .iter() + .map(|s| { + let name = + hdf5::types::VarLenUnicode::from_str(s.name()).map_err(|e| PywrError::HDF5Error(e.to_string()))?; + + Ok(ScenarioGroupEntry { name, size: s.size() }) + }) + .collect::>()?; + + if let Err(e) = grp.new_dataset_builder().with_data(&scenario_groups).create("groups") { + return Err(PywrError::HDF5Error(e.to_string())); + } + + let scenarios: Array1 = domain + .indices() + .iter() + .map(|s| { + let indices = hdf5::types::VarLenArray::from_slice(&s.indices); + + Ok(H5ScenarioIndex { + index: s.index, + indices, + }) + }) + .collect::>()?; + + if let Err(e) = grp.new_dataset_builder().with_data(&scenarios).create("indices") { + return Err(PywrError::HDF5Error(e.to_string())); + } + + Ok(()) +} diff --git a/pywr-core/src/scenario.rs b/pywr-core/src/scenario.rs index 698cac12..76a96b43 100644 --- a/pywr-core/src/scenario.rs +++ b/pywr-core/src/scenario.rs @@ -112,7 +112,7 @@ impl ScenarioIndex { #[derive(Debug)] pub struct ScenarioDomain { scenario_indices: Vec, - scenario_group_names: Vec, + scenario_groups: Vec, } impl ScenarioDomain { @@ -131,12 +131,11 @@ impl ScenarioDomain { /// Return the index of a scenario group by name pub fn group_index(&self, name: &str) -> Option { - self.scenario_group_names.iter().position(|n| n == name) + self.scenario_groups.iter().position(|g| g.name == name) } - /// Return the name of each scenario group - pub fn group_names(&self) -> &[String] { - &self.scenario_group_names + pub fn groups(&self) -> &[ScenarioGroup] { + &self.scenario_groups } } @@ -144,16 +143,14 @@ impl From for ScenarioDomain { fn from(value: ScenarioGroupCollection) -> Self { // Handle creating at-least one scenario if the collection is empty. if !value.is_empty() { - let scenario_group_names = value.groups.iter().map(|g| g.name.clone()).collect(); - Self { scenario_indices: value.scenario_indices(), - scenario_group_names, + scenario_groups: value.groups, } } else { Self { scenario_indices: vec![ScenarioIndex::new(0, vec![0])], - scenario_group_names: vec!["default".to_string()], + scenario_groups: vec![ScenarioGroup::new("default", 1)], } } }