fix: Clamp storage volumes to be within min & max volume ranges.

pywr-core/src/derived_metric.rs

@@ -78,7 +78,7 @@ impl DerivedMetric {
     pub fn compute(&self, network: &Network, state: &State) -> Result<f64, PywrError> {
         match self {
             Self::NodeProportionalVolume(idx) => {
-                let max_volume = network.get_node(idx)?.get_current_max_volume(state)?;
+                let max_volume = network.get_node(idx)?.get_max_volume(state)?;
                 Ok(state
                     .get_network_state()
                     .get_node_proportional_volume(idx, max_volume)?)
@@ -100,7 +100,7 @@ impl DerivedMetric {
                 let max_volume: f64 = node
                     .nodes
                     .iter()
-                    .map(|idx| network.get_node(idx)?.get_current_max_volume(state))
+                    .map(|idx| network.get_node(idx)?.get_max_volume(state))
                     .sum::<Result<_, _>>()?;
                 // TODO handle divide by zero
                 Ok(volume / max_volume)

pywr-core/src/node.rs

@@ -418,7 +418,7 @@ impl Node {
         }
     }
 
-    pub fn get_current_min_volume(&self, state: &State) -> Result<f64, PywrError> {
+    pub fn get_min_volume(&self, state: &State) -> Result<f64, PywrError> {
         match self {
             Self::Input(_) => Err(PywrError::StorageConstraintsUndefined),
             Self::Link(_) => Err(PywrError::StorageConstraintsUndefined),
@@ -439,7 +439,7 @@ impl Node {
         }
     }
 
-    pub fn get_current_max_volume(&self, state: &State) -> Result<f64, PywrError> {
+    pub fn get_max_volume(&self, state: &State) -> Result<f64, PywrError> {
         match self {
             Self::Input(_) => Err(PywrError::StorageConstraintsUndefined),
             Self::Link(_) => Err(PywrError::StorageConstraintsUndefined),
@@ -448,10 +448,11 @@ impl Node {
         }
     }
 
-    pub fn get_current_volume_bounds(&self, state: &State) -> Result<(f64, f64), PywrError> {
-        match (self.get_current_min_volume(state), self.get_current_max_volume(state)) {
+    /// Return the current min and max volumes as a tuple.
+    pub fn get_volume_bounds(&self, state: &State) -> Result<(f64, f64), PywrError> {
+        match (self.get_min_volume(state), self.get_max_volume(state)) {
             (Ok(min_vol), Ok(max_vol)) => Ok((min_vol, max_vol)),
-            _ => Err(PywrError::FlowConstraintsUndefined),
+            _ => Err(PywrError::StorageConstraintsUndefined),
         }
     }
 

pywr-core/src/solvers/builder.rs

@@ -531,7 +531,7 @@ where
             .iter()
             .zip(network.virtual_storage_nodes().deref())
         {
-            let (avail, missing) = match node.get_current_available_volume_bounds(state) {
+            let (avail, missing) = match node.get_available_volume_bounds(state) {
                 Ok(bnds) => bnds,
                 Err(e) => return Err(e),
             };

pywr-core/src/solvers/cbc/mod.rs

@@ -288,7 +288,7 @@ impl Solver for CbcSolver {
             let flow = if flow.abs() < 1e-10 { 0.0 } else { flow };
             network_state.add_flow(edge, timestep, flow)?;
         }
-        network_state.complete(model, timestep)?;
+        state.complete(model, timestep)?;
         timings.save_solution += start_save_solution.elapsed();
 
         Ok(timings)

pywr-core/src/solvers/clp/mod.rs

@@ -287,7 +287,7 @@ impl Solver for ClpSolver {
             let flow = solution[col];
             network_state.add_flow(edge, timestep, flow)?;
         }
-        network_state.complete(model, timestep)?;
+        state.complete(model, timestep)?;
         timings.save_solution += start_save_solution.elapsed();
 
         Ok(timings)

pywr-core/src/solvers/highs/mod.rs

@@ -288,7 +288,7 @@ impl Solver for HighsSolver {
             let flow = solution[col];
             network_state.add_flow(edge, timestep, flow)?;
         }
-        network_state.complete(network, timestep)?;
+        state.complete(network, timestep)?;
         timings.save_solution += start_save_solution.elapsed();
 
         Ok(timings)

pywr-core/src/state.rs

@@ -63,6 +63,12 @@ impl NodeState {
             Self::Storage(s) => s.add_out_flow(flow, timestep),
         };
     }
+
+    fn clamp_volume(&mut self, min_volume: f64, max_volume: f64) {
+        if let Self::Storage(s) = self {
+            s.clamp(min_volume, max_volume);
+        }
+    }
 }
 
 #[derive(Clone, Copy, Debug, Default)]
@@ -124,6 +130,25 @@ impl StorageState {
         // TODO handle divide by zero (is it full or empty?)
         self.volume / max_volume
     }
+
+    /// Ensure the volume is within the min and max volume range (inclusive). If the volume
+    /// is more than 1E6 outside the min or max volume then this function will panic,
+    /// reporting a mass-balance message.
+    fn clamp(&mut self, min_volume: f64, max_volume: f64) {
+        if (self.volume - min_volume) < -1e-6 {
+            panic!(
+                "Mass-balance error detected. Volume ({}) is smaller than minimum volume ({}).",
+                self.volume, min_volume
+            );
+        }
+        if (self.volume - max_volume) > 1e-6 {
+            panic!(
+                "Mass-balance error detected. Volume ({}) is greater than maximum volume ({}).",
+                self.volume, max_volume,
+            );
+        }
+        self.volume = self.volume.clamp(min_volume, max_volume);
+    }
 }
 
 /// Stores the history of virtual storage flows.
@@ -215,6 +240,10 @@ impl VirtualStorageState {
     fn proportional_volume(&self, max_volume: f64) -> f64 {
         self.storage.proportional_volume(max_volume)
     }
+
+    fn clamp_volume(&mut self, min_volume: f64, max_volume: f64) {
+        self.storage.clamp(min_volume, max_volume);
+    }
 }
 
 #[derive(Clone, Copy, Debug, Default)]
@@ -530,8 +559,8 @@ impl NetworkState {
     /// Complete a timestep after all the flow has been added.
     ///
     /// This final step ensures that derived states (e.g. virtual storage volume) are updated
-    /// once all of the flows have been updated.
-    pub fn complete(&mut self, model: &Network, timestep: &Timestep) -> Result<(), PywrError> {
+    /// once all the flows have been updated.
+    fn update_derived_states(&mut self, model: &Network, timestep: &Timestep) -> Result<(), PywrError> {
         // Update virtual storage node states
         for (state, node) in self
             .virtual_storage_states
@@ -562,6 +591,32 @@ impl NetworkState {
         Ok(())
     }
 
+    /// Clamp the volume of `node_index` to be within the bounds provided.
+    fn clamp_node_volume(&mut self, node_index: &NodeIndex, min_volume: f64, max_volume: f64) -> Result<(), PywrError> {
+        match self.node_states.get_mut(*node_index.deref()) {
+            Some(s) => {
+                s.clamp_volume(min_volume, max_volume);
+                Ok(())
+            }
+            None => Err(PywrError::NodeIndexNotFound),
+        }
+    }
+
+    /// Clamp the volume of `node_index` to be within the bounds provided.
+    fn clamp_virtual_storage_node_volume(
+        &mut self,
+        node_index: &VirtualStorageIndex,
+        min_volume: f64,
+        max_volume: f64,
+    ) -> Result<(), PywrError> {
+        match self.virtual_storage_states.get_mut(*node_index.deref()) {
+            Some(s) => {
+                s.clamp_volume(min_volume, max_volume);
+                Ok(())
+            }
+            None => Err(PywrError::VirtualStorageIndexNotFound(*node_index)),
+        }
+    }
     pub fn get_node_in_flow(&self, node_index: &NodeIndex) -> Result<f64, PywrError> {
         match self.node_states.get(*node_index.deref()) {
             Some(s) => Ok(s.get_in_flow()),
@@ -913,6 +968,32 @@ impl State {
             None => Err(PywrError::MultiNetworkTransferIndexNotFound(idx)),
         }
     }
+
+    /// Complete a timestep after all the flow has been added.
+    ///
+    /// This final step ensures, once all the flows have been updated, that:
+    ///   - Derived states (e.g. virtual storage volume) are updated
+    ///   - Volumes are within bounds
+    pub fn complete(&mut self, model: &Network, timestep: &Timestep) -> Result<(), PywrError> {
+        for node in model.nodes().iter() {
+            if let Node::Storage(_) = node {
+                let node_index = node.index();
+                let min_volume = node.get_min_volume(self)?;
+                let max_volume = node.get_max_volume(self)?;
+                self.network.clamp_node_volume(&node_index, min_volume, max_volume)?;
+            }
+        }
+
+        for node in model.virtual_storage_nodes().iter() {
+            let node_index = node.index();
+            let min_volume = node.get_min_volume(self)?;
+            let max_volume = node.get_max_volume(self)?;
+            self.network
+                .clamp_virtual_storage_node_volume(&node_index, min_volume, max_volume)?;
+        }
+
+        self.network.update_derived_states(model, timestep)
+    }
 }
 
 /// Builder for the [`State`] struct.

pywr-core/src/virtual_storage.rs

@@ -272,7 +272,7 @@ impl VirtualStorage {
             .get_max_volume(&state.get_simple_parameter_values())
     }
 
-    pub fn get_current_available_volume_bounds(&self, state: &State) -> Result<(f64, f64), PywrError> {
+    pub fn get_available_volume_bounds(&self, state: &State) -> Result<(f64, f64), PywrError> {
         let min_vol = self.get_min_volume(state)?;
         let max_vol = self.get_max_volume(state)?;