fix: Fix calculation of proportional volume on VS nodes after reset.

Derived metrics, such as proportional volume, must be recalculated
at the start of the time-step if the volume has been reset. This
ensures that parameters that depend on the proportional volume
(e.g., control curve calculations) use the correct value after the
volume is reset.

pywr-core/src/derived_metric.rs

@@ -67,8 +67,19 @@ pub enum DerivedMetric {
 
 impl DerivedMetric {
     pub fn before(&self, timestep: &Timestep, network: &Network, state: &State) -> Result<Option<f64>, PywrError> {
+        // Virtual storage nodes can reset their volume. If this has happened then the
+        // proportional volume should also be recalculated.
+        let has_reset = if let Self::VirtualStorageProportionalVolume(idx) = self {
+            match state.get_network_state().get_virtual_storage_last_reset(*idx)? {
+                Some(last_reset) => last_reset == timestep,
+                None => false,
+            }
+        } else {
+            false
+        };
+
         // On the first time-step set the initial value
-        if timestep.is_first() {
+        if timestep.is_first() || has_reset {
             self.compute(network, state).map(Some)
         } else {
             Ok(None)

pywr-core/src/metric.rs

@@ -229,6 +229,12 @@ where
     }
 }
 
+impl From<DerivedMetricIndex> for MetricF64 {
+    fn from(idx: DerivedMetricIndex) -> Self {
+        Self::DerivedMetric(idx)
+    }
+}
+
 impl From<ParameterIndex<f64>> for MetricF64 {
     fn from(idx: ParameterIndex<f64>) -> Self {
         match idx {

pywr-core/src/network.rs

@@ -1050,6 +1050,22 @@ impl Network {
         }
     }
 
+    /// Get a `VirtualStorageNode` from a node's name
+    pub fn get_mut_virtual_storage_node_by_name(
+        &mut self,
+        name: &str,
+        sub_name: Option<&str>,
+    ) -> Result<&mut VirtualStorage, PywrError> {
+        match self
+            .virtual_storage_nodes
+            .iter_mut()
+            .find(|n| n.full_name() == (name, sub_name))
+        {
+            Some(node) => Ok(node),
+            None => Err(PywrError::NodeNotFound(name.to_string())),
+        }
+    }
+
     /// Get a `VirtualStorageNode` from a node's name
     pub fn get_virtual_storage_node_index_by_name(
         &self,
@@ -1339,6 +1355,17 @@ impl Network {
         Ok(vs_node_index)
     }
 
+    pub fn set_virtual_storage_node_cost(
+        &mut self,
+        name: &str,
+        sub_name: Option<&str>,
+        value: Option<MetricF64>,
+    ) -> Result<(), PywrError> {
+        let node = self.get_mut_virtual_storage_node_by_name(name, sub_name)?;
+        node.set_cost(value);
+        Ok(())
+    }
+
     /// Add a [`parameters::GeneralParameter`] to the network
     pub fn add_parameter(
         &mut self,
@@ -1771,7 +1798,7 @@ mod tests {
     #[test]
     fn test_step() {
         const NUM_SCENARIOS: usize = 2;
-        let model = simple_model(NUM_SCENARIOS);
+        let model = simple_model(NUM_SCENARIOS, None);
 
         let mut timings = RunTimings::default();
 
@@ -1796,7 +1823,7 @@ mod tests {
     #[test]
     /// Test running a simple model
     fn test_run() {
-        let mut model = simple_model(10);
+        let mut model = simple_model(10, None);
 
         // Set-up assertion for "input" node
         let idx = model.network().get_node_by_name("input", None).unwrap().index();
@@ -1992,7 +2019,7 @@ mod tests {
     #[test]
     /// Test the variable API
     fn test_variable_api() {
-        let mut model = simple_model(1);
+        let mut model = simple_model(1, None);
 
         let variable = ActivationFunction::Unit { min: 0.0, max: 10.0 };
         let input_max_flow = parameters::ConstantParameter::new("my-constant".into(), 10.0);

pywr-core/src/parameters/control_curves/piecewise.rs

@@ -82,7 +82,7 @@ mod test {
     /// Basic functional test of the piecewise interpolation.
     #[test]
     fn test_basic() {
-        let mut model = simple_model(1);
+        let mut model = simple_model(1, None);
 
         // Create an artificial volume series to use for the interpolation test
         let volume = Array1Parameter::new("test-x".into(), Array1::linspace(1.0, 0.0, 21), None);

pywr-core/src/parameters/delay.rs

@@ -170,7 +170,7 @@ mod test {
     /// Basic functional test of the delay parameter.
     #[test]
     fn test_basic() {
-        let mut model = simple_model(1);
+        let mut model = simple_model(1, None);
 
         // Create an artificial volume series to use for the delay test
         let volumes = Array1::linspace(1.0, 0.0, 21);

pywr-core/src/parameters/discount_factor.rs

@@ -61,7 +61,7 @@ mod test {
     /// Basic functional test of the delay parameter.
     #[test]
     fn test_basic() {
-        let mut model = simple_model(1);
+        let mut model = simple_model(1, None);
         let network = model.network_mut();
 
         // Create an artificial volume series to use for the delay test

pywr-core/src/recorders/mod.rs

@@ -354,7 +354,7 @@ mod tests {
 
     #[test]
     fn test_array2_recorder() {
-        let mut model = simple_model(2);
+        let mut model = simple_model(2, None);
 
         let node_idx = model.network().get_node_index_by_name("input", None).unwrap();
 

pywr-core/src/test_utils.rs

@@ -91,11 +91,11 @@ pub fn simple_network(network: &mut Network, inflow_scenario_index: usize, num_i
     output_node.set_cost(Some(demand_cost.into()));
 }
 /// Create a simple test model with three nodes.
-pub fn simple_model(num_scenarios: usize) -> Model {
+pub fn simple_model(num_scenarios: usize, timestepper: Option<Timestepper>) -> Model {
     let mut scenario_collection = ScenarioGroupCollection::default();
     scenario_collection.add_group("test-scenario", num_scenarios);
 
-    let domain = ModelDomain::from(default_timestepper(), scenario_collection).unwrap();
+    let domain = ModelDomain::from(timestepper.unwrap_or_else(default_timestepper), scenario_collection).unwrap();
     let mut network = Network::default();
 
     let idx = domain

pywr-core/src/timestep.rs

@@ -105,7 +105,7 @@ impl PywrDuration {
 pub type TimestepIndex = usize;
 
 #[cfg_attr(feature = "pyo3", pyclass)]
-#[derive(Debug, Copy, Clone)]
+#[derive(Debug, Copy, Clone, PartialEq)]
 pub struct Timestep {
     pub date: NaiveDateTime,
     pub index: TimestepIndex,

pywr-core/src/virtual_storage.rs

@@ -212,6 +212,10 @@ impl VirtualStorage {
         }
     }
 
+    pub fn set_cost(&mut self, cost: Option<MetricF64>) {
+        self.cost = cost;
+    }
+
     pub fn before(&self, timestep: &Timestep, state: &mut State) -> Result<(), PywrError> {
         let do_reset = if timestep.is_first() {
             // Set the initial volume if it is the first timestep.
@@ -296,16 +300,18 @@ fn months_since_last_reset(current: &NaiveDateTime, last_reset: &NaiveDateTime)
 
 #[cfg(test)]
 mod tests {
+    use crate::derived_metric::DerivedMetric;
     use crate::metric::MetricF64;
     use crate::models::Model;
     use crate::network::Network;
     use crate::node::StorageInitialVolume;
+    use crate::parameters::ControlCurveInterpolatedParameter;
     use crate::recorders::{AssertionFnRecorder, AssertionRecorder};
     use crate::scenario::ScenarioIndex;
     use crate::test_utils::{default_timestepper, run_all_solvers, simple_model};
-    use crate::timestep::Timestep;
+    use crate::timestep::{Timestep, TimestepDuration, Timestepper};
     use crate::virtual_storage::{months_since_last_reset, VirtualStorageBuilder, VirtualStorageReset};
-    use chrono::NaiveDate;
+    use chrono::{Datelike, NaiveDate};
     use ndarray::Array;
     use std::num::NonZeroUsize;
 
@@ -433,7 +439,7 @@ mod tests {
     #[test]
     /// Test virtual storage node costs
     fn test_virtual_storage_node_costs() {
-        let mut model = simple_model(1);
+        let mut model = simple_model(1, None);
         let network = model.network_mut();
 
         let nodes = vec![network.get_node_index_by_name("input", None).unwrap()];
@@ -449,6 +455,7 @@ mod tests {
         network.add_virtual_storage_node(vs_builder).unwrap();
 
         let expected = Array::zeros((366, 1));
+
         let idx = network.get_node_by_name("output", None).unwrap().index();
         let recorder = AssertionRecorder::new("output-flow", MetricF64::NodeInFlow(idx), expected, None, None);
         network.add_recorder(Box::new(recorder)).unwrap();
@@ -457,10 +464,77 @@ mod tests {
         run_all_solvers(&model, &["ipm-ocl", "ipm-simd"], &[], &[]);
     }
 
+    #[test]
+    /// Virtual storage node resets every month. This test will check that a parameter which
+    /// uses the derived proportional volume receives the correct value after each reset.
+    fn test_virtual_storage_node_cost_dynamic() {
+        // This test needs to be run over a period of time to see the cost change
+        let start = NaiveDate::from_ymd_opt(2020, 1, 1)
+            .unwrap()
+            .and_hms_opt(0, 0, 0)
+            .unwrap();
+        let end = NaiveDate::from_ymd_opt(2020, 12, 31)
+            .unwrap()
+            .and_hms_opt(0, 0, 0)
+            .unwrap();
+        let duration = TimestepDuration::Days(1);
+
+        let mut model = simple_model(1, Some(Timestepper::new(start, end, duration)));
+        let network = model.network_mut();
+
+        let nodes = vec![network.get_node_index_by_name("input", None).unwrap()];
+
+        let vs_builder = VirtualStorageBuilder::new("vs", &nodes)
+            .initial_volume(StorageInitialVolume::Proportional(1.0))
+            .min_volume(Some(0.0.into()))
+            .max_volume(Some(100.0.into()))
+            .reset(VirtualStorageReset::NumberOfMonths { months: 1 });
+
+        let vs_idx = network.add_virtual_storage_node(vs_builder).unwrap();
+        let vs_vol_metric = network.add_derived_metric(DerivedMetric::VirtualStorageProportionalVolume(vs_idx));
+
+        // Virtual storage node cost increases with decreasing volume
+        let cost_param = ControlCurveInterpolatedParameter::new(
+            "cost".into(),
+            vs_vol_metric.into(),
+            vec![],
+            vec![0.0.into(), 20.0.into()],
+        );
+
+        let cost_param = network.add_parameter(Box::new(cost_param)).unwrap();
+
+        network
+            .set_virtual_storage_node_cost("vs", None, Some(cost_param.into()))
+            .unwrap();
+
+        let expected = |ts: &Timestep, _si: &ScenarioIndex| {
+            // Calculate the current volume within each month
+            // This should be the absolute volume at the start of each, which is then used
+            // to calculate the cost.
+            let mut volume = 100.0;
+            for dom in 1..ts.date.day() {
+                let demand_met = (1.0 + ts.index as f64 - (ts.date.day() - dom) as f64).min(12.0);
+                volume -= demand_met;
+            }
+
+            if volume >= 50.0 {
+                (1.0 + ts.index as f64).min(12.0)
+            } else {
+                0.0
+            }
+        };
+        let idx = network.get_node_by_name("output", None).unwrap().index();
+        let recorder = AssertionFnRecorder::new("output-flow", MetricF64::NodeInFlow(idx), expected, None, None);
+        network.add_recorder(Box::new(recorder)).unwrap();
+
+        // Test all solvers
+        run_all_solvers(&model, &["ipm-ocl", "ipm-simd"], &[], &[]);
+    }
+
     #[test]
     /// Test virtual storage rolling window constraint
     fn test_virtual_storage_node_rolling_constraint() {
-        let mut model = simple_model(1);
+        let mut model = simple_model(1, None);
         let network = model.network_mut();
 
         let nodes = vec![network.get_node_index_by_name("input", None).unwrap()];