feat: Several output additions and improvements. (#103)

* feat: Several output additions and improvements.

- Fixed an issue where metric sets were not finalised at the end of a model run.
- Added a new Memory output and associated recorder. This recorder stores data in memory.
- Add new functions to Recorder API to allow accessing an aggregated metric.
  - Additional work is needed to give access to different levels of aggregation.
  - Currently only the new Memory recorder supports this API.
  - The order of the aggregation operations probably needs looking at.
- Changed serialised format of OutputMetric. This will allow it it eventually contain additional information.
- Changed serialised format of MetricAggFunc. Again, this will allow definition of functions with additional arguments.
- Added CountNonZero aggregation function.

Cargo.toml

@@ -36,6 +36,7 @@ serde = { version = "1", features = ["derive"] }
 serde_json = "1.0"
 thiserror = "1.0.25"
 num = "0.4.0"
+float-cmp = "0.9.0"
 ndarray = "0.15.3"
 polars = { version = "0.37.0", features = ["lazy", "rows", "ndarray"] }
 pyo3-polars = "0.11.1"

pywr-core/Cargo.toml

@@ -18,7 +18,7 @@ libc = "0.2.97"
 thiserror = { workspace = true }
 ndarray = { workspace = true }
 num = { workspace = true }
-float-cmp = "0.9.0"
+float-cmp = {  workspace = true }
 hdf5 = { workspace = true }
 csv = { workspace = true }
 clp-sys = { path = "../clp-sys", version = "0.1.0" }

pywr-core/src/lib.rs

@@ -6,7 +6,7 @@ use crate::derived_metric::DerivedMetricIndex;
 use crate::models::MultiNetworkTransferIndex;
 use crate::node::NodeIndex;
 use crate::parameters::{IndexParameterIndex, InterpolationError, MultiValueParameterIndex, ParameterIndex};
-use crate::recorders::{MetricSetIndex, RecorderIndex};
+use crate::recorders::{AggregationError, MetricSetIndex, RecorderIndex};
 use crate::virtual_storage::VirtualStorageIndex;
 use pyo3::exceptions::{PyException, PyRuntimeError};
 use pyo3::{create_exception, PyErr};
@@ -123,6 +123,8 @@ pub enum PywrError {
     InvalidMetricValue(String),
     #[error("recorder not initialised")]
     RecorderNotInitialised,
+    #[error("recorder does not supported aggregation")]
+    RecorderDoesNotSupportAggregation,
     #[error("hdf5 error: {0}")]
     HDF5Error(String),
     #[error("csv error: {0}")]
@@ -161,6 +163,8 @@ pub enum PywrError {
     TimestepRangeGenerationError(String),
     #[error("Could not create timesteps for frequency '{0}'")]
     TimestepGenerationError(String),
+    #[error("aggregation error: {0}")]
+    Aggregation(#[from] AggregationError),
 }
 
 // Python errors

pywr-core/src/models/multi.rs

@@ -298,8 +298,9 @@ impl MultiNetworkModel {
         }
 
         for (idx, entry) in self.networks.iter().enumerate() {
+            let sub_model_ms_states = state.states.get_mut(idx).unwrap().all_metric_set_internal_states_mut();
             let sub_model_recorder_states = state.recorder_states.get_mut(idx).unwrap();
-            entry.network.finalise(sub_model_recorder_states)?;
+            entry.network.finalise(sub_model_ms_states, sub_model_recorder_states)?;
         }
         // End the global timer and print the run statistics
         timings.finish(count);

pywr-core/src/models/simple.rs

@@ -22,6 +22,10 @@ impl<S> ModelState<S> {
     pub fn network_state_mut(&mut self) -> &mut NetworkState {
         &mut self.state
     }
+
+    pub fn recorder_state(&self) -> &Vec<Option<Box<dyn Any>>> {
+        &self.recorder_state
+    }
 }
 
 /// A standard Pywr model containing a single network.
@@ -196,7 +200,7 @@ impl Model {
     /// Run a model through the given time-steps.
     ///
     /// This method will setup state and solvers, and then run the model through the time-steps.
-    pub fn run<S>(&self, settings: &S::Settings) -> Result<(), PywrError>
+    pub fn run<S>(&self, settings: &S::Settings) -> Result<Vec<Option<Box<dyn Any>>>, PywrError>
     where
         S: Solver,
         <S as Solver>::Settings: SolverSettings,
@@ -205,7 +209,7 @@ impl Model {
 
         self.run_with_state::<S>(&mut state, settings)?;
 
-        Ok(())
+        Ok(state.recorder_state)
     }
 
     /// Run the model with the provided states and solvers.
@@ -243,7 +247,10 @@ impl Model {
             count += self.domain.scenarios.indices().len();
         }
 
-        self.network.finalise(&mut state.recorder_state)?;
+        self.network.finalise(
+            state.state.all_metric_set_internal_states_mut(),
+            &mut state.recorder_state,
+        )?;
         // End the global timer and print the run statistics
         timings.finish(count);
         timings.print_table();
@@ -296,7 +303,10 @@ impl Model {
             count += self.domain.scenarios.indices().len();
         }
 
-        self.network.finalise(&mut state.recorder_state)?;
+        self.network.finalise(
+            state.state.all_metric_set_internal_states_mut(),
+            &mut state.recorder_state,
+        )?;
 
         // End the global timer and print the run statistics
         timings.finish(count);

pywr-core/src/network.rs

@@ -182,6 +182,10 @@ impl NetworkState {
     pub fn iter_parameter_states_mut(&mut self) -> IterMut<'_, ParameterStates> {
         self.parameter_internal_states.iter_mut()
     }
+
+    pub fn all_metric_set_internal_states_mut(&mut self) -> &mut [Vec<MetricSetState>] {
+        &mut self.metric_set_internal_states
+    }
 }
 
 /// A Pywr network containing nodes, edges, parameters, metric sets, etc.
@@ -358,10 +362,22 @@ impl Network {
         S::setup(self, scenario_indices.len(), settings)
     }
 
-    pub fn finalise(&self, recorder_internal_states: &mut [Option<Box<dyn Any>>]) -> Result<(), PywrError> {
+    pub fn finalise(
+        &self,
+        metric_set_states: &mut [Vec<MetricSetState>],
+        recorder_internal_states: &mut [Option<Box<dyn Any>>],
+    ) -> Result<(), PywrError> {
+        // Finally, save new data to the metric set
+
+        for ms_states in metric_set_states.iter_mut() {
+            for (metric_set, ms_state) in self.metric_sets.iter().zip(ms_states.iter_mut()) {
+                metric_set.finalise(ms_state);
+            }
+        }
+
         // Setup recorders
         for (recorder, internal_state) in self.recorders.iter().zip(recorder_internal_states) {
-            recorder.finalise(internal_state)?;
+            recorder.finalise(metric_set_states, internal_state)?;
         }
 
         Ok(())
@@ -1136,6 +1152,20 @@ impl Network {
         }
     }
 
+    pub fn get_recorder_index_by_name(&self, name: &str) -> Result<RecorderIndex, PywrError> {
+        match self.recorders.iter().position(|r| r.name() == name) {
+            Some(idx) => Ok(RecorderIndex::new(idx)),
+            None => Err(PywrError::RecorderNotFound),
+        }
+    }
+
+    pub fn get_aggregated_value(&self, name: &str, recorder_states: &[Option<Box<dyn Any>>]) -> Result<f64, PywrError> {
+        match self.recorders.iter().enumerate().find(|(_, r)| r.name() == name) {
+            Some((idx, recorder)) => recorder.aggregated_value(&recorder_states[idx]),
+            None => Err(PywrError::RecorderNotFound),
+        }
+    }
+
     /// Add a new Node::Input to the network.
     pub fn add_input_node(&mut self, name: &str, sub_name: Option<&str>) -> Result<NodeIndex, PywrError> {
         // Check for name.