feat: Add Edge metric

A new schema and core level metric for returning edge flows. The
schema level metric returns core a metric that covers the case
where multiple core edges are created from a single schema
edge.

Includes some additional tests for picewise-link1.json. This
covers the case described above.

pywr-core/src/metric.rs

@@ -85,6 +85,7 @@ pub enum MetricF64 {
     AggregatedNodeOutFlow(AggregatedNodeIndex),
     AggregatedNodeVolume(AggregatedStorageNodeIndex),
     EdgeFlow(EdgeIndex),
+    MultiEdgeFlow { indices: Vec<EdgeIndex>, name: String },
     ParameterValue(GeneralParameterIndex<f64>),
     IndexParameterValue(GeneralParameterIndex<usize>),
     MultiParameterValue((GeneralParameterIndex<MultiValue>, String)),
@@ -119,6 +120,13 @@ impl MetricF64 {
             }
 
             MetricF64::EdgeFlow(idx) => Ok(state.get_network_state().get_edge_flow(idx)?),
+            MetricF64::MultiEdgeFlow { indices, .. } => {
+                let flow = indices
+                    .iter()
+                    .map(|idx| state.get_network_state().get_edge_flow(idx))
+                    .sum::<Result<_, _>>()?;
+                Ok(flow)
+            }
             MetricF64::ParameterValue(idx) => Ok(state.get_parameter_value(*idx)?),
             MetricF64::IndexParameterValue(idx) => Ok(state.get_parameter_index(*idx)? as f64),
             MetricF64::MultiParameterValue((idx, key)) => Ok(state.get_multi_parameter_value(*idx, key)?),

pywr-core/src/network.rs

@@ -802,6 +802,17 @@ impl Network {
         self.edges.get(index)
     }
 
+    /// Get an [`EdgeIndex`] from connecting node indices.
+    pub fn get_edge_index(&self, from_node_index: NodeIndex, to_node_index: NodeIndex) -> Result<EdgeIndex, PywrError> {
+        match self
+            .edges
+            .iter()
+            .find(|edge| edge.from_node_index == from_node_index && edge.to_node_index == to_node_index)
+        {
+            Some(edge) => Ok(edge.index),
+            None => Err(PywrError::EdgeIndexNotFound),
+        }
+    }
     /// Get a Node from a node's name
     pub fn get_node_index_by_name(&self, name: &str, sub_name: Option<&str>) -> Result<NodeIndex, PywrError> {
         Ok(self.get_node_by_name(name, sub_name)?.index())

pywr-schema/src/edge.rs

@@ -1,6 +1,13 @@
+#[cfg(feature = "core")]
+use crate::model::LoadArgs;
+#[cfg(feature = "core")]
+use crate::SchemaError;
+#[cfg(feature = "core")]
+use pywr_core::{edge::EdgeIndex, metric::MetricF64, node::NodeIndex};
 use schemars::JsonSchema;
+use std::fmt::{Display, Formatter};
 
-#[derive(serde::Deserialize, serde::Serialize, Clone, JsonSchema)]
+#[derive(serde::Deserialize, serde::Serialize, Clone, JsonSchema, Debug)]
 pub struct Edge {
     pub from_node: String,
     pub to_node: String,
@@ -20,3 +27,97 @@ impl From<pywr_v1_schema::edge::Edge> for Edge {
         }
     }
 }
+
+const EDGE_SYMBOL: &str = "->";
+
+impl Display for Edge {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match (&self.from_slot, &self.to_slot) {
+            (Some(from_slot), Some(to_slot)) => {
+                write!(
+                    f,
+                    "{}[{}]{}{}[{}]",
+                    self.from_node, from_slot, EDGE_SYMBOL, self.to_node, to_slot
+                )
+            }
+            (Some(from_slot), None) => write!(f, "{}[{}]{}{}", self.from_node, from_slot, EDGE_SYMBOL, self.to_node),
+            (None, Some(to_slot)) => {
+                write!(f, "{}{}{}[{}]", self.from_node, EDGE_SYMBOL, self.to_node, to_slot)
+            }
+            (None, None) => write!(f, "{}{}{}", self.from_node, EDGE_SYMBOL, self.to_node),
+        }
+    }
+}
+
+#[cfg(feature = "core")]
+impl Edge {
+    /// Returns an iterator of the pairs (from, to) of `NodeIndex` that represent this
+    /// edge when added to a model. In general this can be several nodes because some nodes
+    /// have multiple internal nodes when connected from or to.
+    fn iter_node_index_pairs(
+        &self,
+        network: &pywr_core::network::Network,
+        args: &LoadArgs,
+    ) -> Result<impl Iterator<Item = (NodeIndex, NodeIndex)>, SchemaError> {
+        let from_node = args
+            .schema
+            .get_node_by_name(self.from_node.as_str())
+            .ok_or_else(|| SchemaError::NodeNotFound(self.from_node.clone()))?;
+
+        let to_node = args
+            .schema
+            .get_node_by_name(self.to_node.as_str())
+            .ok_or_else(|| SchemaError::NodeNotFound(self.to_node.clone()))?;
+
+        let from_slot = self.from_slot.as_deref();
+
+        // Collect the node indices at each end of the edge
+        let from_node_indices: Vec<NodeIndex> = from_node
+            .output_connectors(from_slot)
+            .into_iter()
+            .map(|(name, sub_name)| network.get_node_index_by_name(name, sub_name.as_deref()))
+            .collect::<Result<_, _>>()?;
+
+        let to_node_indices: Vec<NodeIndex> = to_node
+            .input_connectors()
+            .into_iter()
+            .map(|(name, sub_name)| network.get_node_index_by_name(name, sub_name.as_deref()))
+            .collect::<Result<_, _>>()?;
+
+        let pairs: Vec<_> = from_node_indices
+            .into_iter()
+            .flat_map(|from_node_index| std::iter::repeat(from_node_index).zip(to_node_indices.iter().copied()))
+            .collect();
+
+        Ok(pairs.into_iter())
+    }
+
+    /// Add the edge to the network
+    pub fn add_to_model(&self, network: &mut pywr_core::network::Network, args: &LoadArgs) -> Result<(), SchemaError> {
+        // Connect each "from" connector to each "to" connector
+        for (from_node_index, to_node_index) in self.iter_node_index_pairs(network, args)? {
+            network.connect_nodes(from_node_index, to_node_index)?;
+        }
+
+        Ok(())
+    }
+
+    /// Create a metric that will return this edge's total flow in the model.
+    pub fn create_metric(
+        &self,
+        network: &pywr_core::network::Network,
+        args: &LoadArgs,
+    ) -> Result<MetricF64, SchemaError> {
+        let indices: Vec<EdgeIndex> = self
+            .iter_node_index_pairs(network, args)?
+            .map(|(from_node_index, to_node_index)| network.get_edge_index(from_node_index, to_node_index))
+            .collect::<Result<_, _>>()?;
+
+        let metric = MetricF64::MultiEdgeFlow {
+            indices,
+            name: self.to_string(),
+        };
+
+        Ok(metric)
+    }
+}

pywr-schema/src/metric.rs

@@ -1,4 +1,5 @@
 use crate::data_tables::TableDataRef;
+use crate::edge::Edge;
 #[cfg(feature = "core")]
 use crate::error::SchemaError;
 #[cfg(feature = "core")]
@@ -28,6 +29,7 @@ pub enum Metric {
     Table(TableDataRef),
     /// An attribute of a node.
     Node(NodeReference),
+    Edge(EdgeReference),
     Timeseries(TimeseriesReference),
     Parameter(ParameterReference),
     InlineParameter {
@@ -116,18 +118,20 @@ impl Metric {
                     None => Err(SchemaError::InterNetworkTransferNotFound(name.to_string())),
                 }
             }
+            Self::Edge(edge_ref) => edge_ref.load(network, args),
         }
     }
 
-    fn name(&self) -> Result<&str, SchemaError> {
+    fn name(&self) -> Result<String, SchemaError> {
         match self {
-            Self::Node(node_ref) => Ok(&node_ref.name),
-            Self::Parameter(parameter_ref) => Ok(&parameter_ref.name),
+            Self::Node(node_ref) => Ok(node_ref.name.to_string()),
+            Self::Parameter(parameter_ref) => Ok(parameter_ref.name.clone()),
             Self::Constant { .. } => Err(SchemaError::LiteralConstantOutputNotSupported),
-            Self::Table(table_ref) => Ok(&table_ref.table),
-            Self::Timeseries(ts_ref) => Ok(&ts_ref.name),
-            Self::InlineParameter { definition } => Ok(definition.name()),
-            Self::InterNetworkTransfer { name } => Ok(name),
+            Self::Table(table_ref) => Ok(table_ref.table.clone()),
+            Self::Timeseries(ts_ref) => Ok(ts_ref.name.clone()),
+            Self::InlineParameter { definition } => Ok(definition.name().to_string()),
+            Self::InterNetworkTransfer { name } => Ok(name.clone()),
+            Self::Edge(edge_ref) => Ok(edge_ref.edge.to_string()),
         }
     }
 
@@ -140,6 +144,7 @@ impl Metric {
             Self::Timeseries(_) => "value".to_string(),
             Self::InlineParameter { .. } => "value".to_string(),
             Self::InterNetworkTransfer { .. } => "value".to_string(),
+            Self::Edge { .. } => "Flow".to_string(),
         };
 
         Ok(attribute)
@@ -158,6 +163,7 @@ impl Metric {
             Self::Timeseries(_) => None,
             Self::InlineParameter { definition } => Some(definition.parameter_type().to_string()),
             Self::InterNetworkTransfer { .. } => None,
+            Self::Edge { .. } => None,
         };
 
         Ok(sub_type)
@@ -174,7 +180,7 @@ impl Metric {
         let sub_type = self.sub_type(args)?;
 
         Ok(OutputMetric::new(
-            self.name()?,
+            self.name()?.as_str(),
             &self.attribute(args)?,
             &ty,
             sub_type.as_deref(),
@@ -409,3 +415,18 @@ impl ParameterReference {
         Ok(parameter.parameter_type())
     }
 }
+
+#[derive(serde::Deserialize, serde::Serialize, Debug, Clone, JsonSchema)]
+#[serde(deny_unknown_fields)]
+pub struct EdgeReference {
+    /// The edge referred to by this reference.
+    pub edge: Edge,
+}
+
+#[cfg(feature = "core")]
+impl EdgeReference {
+    pub fn load(&self, network: &mut pywr_core::network::Network, args: &LoadArgs) -> Result<MetricF64, SchemaError> {
+        // This is the associated node in the schema
+        self.edge.create_metric(network, args)
+    }
+}

pywr-schema/src/model.rs

@@ -451,24 +451,7 @@ impl PywrNetwork {
 
         // Create the edges
         for edge in &self.edges {
-            let from_node = self
-                .get_node_by_name(edge.from_node.as_str())
-                .ok_or_else(|| SchemaError::NodeNotFound(edge.from_node.clone()))?;
-            let to_node = self
-                .get_node_by_name(edge.to_node.as_str())
-                .ok_or_else(|| SchemaError::NodeNotFound(edge.to_node.clone()))?;
-
-            let from_slot = edge.from_slot.as_deref();
-
-            // Connect each "from" connector to each "to" connector
-            for from_connector in from_node.output_connectors(from_slot) {
-                for to_connector in to_node.input_connectors() {
-                    let from_node_index =
-                        network.get_node_index_by_name(from_connector.0, from_connector.1.as_deref())?;
-                    let to_node_index = network.get_node_index_by_name(to_connector.0, to_connector.1.as_deref())?;
-                    network.connect_nodes(from_node_index, to_node_index)?;
-                }
-            }
+            edge.add_to_model(&mut network, &args)?;
         }
 
         // Create all the parameters

pywr-schema/src/nodes/piecewise_link.rs

@@ -203,23 +203,34 @@ impl TryFrom<PiecewiseLinkNodeV1> for PiecewiseLinkNode {
 mod tests {
     use crate::model::PywrModel;
     use ndarray::Array2;
+    use pywr_core::test_utils::ExpectedOutputs;
     use pywr_core::{metric::MetricF64, recorders::AssertionRecorder, test_utils::run_all_solvers};
+    use tempfile::TempDir;
 
     fn model_str() -> &'static str {
         include_str!("../test_models/piecewise_link1.json")
     }
 
+    fn pwl_node_outputs_str() -> &'static str {
+        include_str!("../test_models/piecewise-link1-nodes.csv")
+    }
+
+    fn pwl_edges_outputs_str() -> &'static str {
+        include_str!("../test_models/piecewise-link1-edges.csv")
+    }
+
     #[test]
     fn test_model_run() {
         let data = model_str();
         let schema: PywrModel = serde_json::from_str(data).unwrap();
-        let mut model = schema.build_model(None, None).unwrap();
+        let temp_dir = TempDir::new().unwrap();
+
+        let mut model = schema.build_model(None, Some(temp_dir.path())).unwrap();
 
         let network = model.network_mut();
         assert_eq!(network.nodes().len(), 5);
         assert_eq!(network.edges().len(), 6);
 
-        // TODO put this assertion data in the test model file.
         let idx = network.get_node_by_name("link1", Some("step-00")).unwrap().index();
         let expected = Array2::from_elem((366, 1), 1.0);
         let recorder = AssertionRecorder::new("link1-s0-flow", MetricF64::NodeOutFlow(idx), expected, None, None);
@@ -235,7 +246,18 @@ mod tests {
         let recorder = AssertionRecorder::new("link1-s0-flow", MetricF64::NodeOutFlow(idx), expected, None, None);
         network.add_recorder(Box::new(recorder)).unwrap();
 
+        let expected_outputs = [
+            ExpectedOutputs::new(
+                temp_dir.path().join("piecewise-link1-nodes.csv"),
+                pwl_node_outputs_str(),
+            ),
+            ExpectedOutputs::new(
+                temp_dir.path().join("piecewise-link1-edges.csv"),
+                pwl_edges_outputs_str(),
+            ),
+        ];
+
         // Test all solvers
-        run_all_solvers(&model, &[], &[]);
+        run_all_solvers(&model, &[], &expected_outputs);
     }
 }

pywr-schema/src/test_models/piecewise-link1-edges.csv

@@ -0,0 +1,4 @@
+time_start,time_end,scenario_index,metric_set,name,attribute,value
+2015-01-01T00:00:00,2015-01-02T00:00:00,0,edges,input1->link1,Flow,4.0
+2015-01-02T00:00:00,2015-01-03T00:00:00,0,edges,input1->link1,Flow,4.0
+2015-01-03T00:00:00,2015-01-04T00:00:00,0,edges,input1->link1,Flow,4.0

pywr-schema/src/test_models/piecewise-link1-nodes.csv

@@ -0,0 +1,4 @@
+time_start,time_end,scenario_index,metric_set,name,attribute,value
+2015-01-01T00:00:00,2015-01-02T00:00:00,0,nodes,link1,Inflow,4.0
+2015-01-02T00:00:00,2015-01-03T00:00:00,0,nodes,link1,Inflow,4.0
+2015-01-03T00:00:00,2015-01-04T00:00:00,0,nodes,link1,Inflow,4.0

pywr-schema/src/test_models/piecewise_link1.json

@@ -6,7 +6,7 @@
   },
   "timestepper": {
     "start": "2015-01-01",
-    "end": "2015-12-31",
+    "end": "2015-01-03",
     "timestep": 1
   },
   "network": {
@@ -79,6 +79,52 @@
         "from_node": "link1",
         "to_node": "demand1"
       }
+    ],
+    "metric_sets": [
+      {
+        "name": "nodes",
+        "metrics": [
+          {
+            "type": "Node",
+            "name": "link1",
+            "attribute": "Inflow"
+          }
+        ]
+      },
+      {
+        "name": "edges",
+        "metrics": [
+          {
+            "type": "Edge",
+            "edge": {
+              "from_node": "input1",
+              "to_node": "link1"
+            }
+          }
+        ]
+      }
+    ],
+    "outputs": [
+      {
+        "name": "node-outputs",
+        "type": "CSV",
+        "format": "long",
+        "filename": "piecewise-link1-nodes.csv",
+        "metric_set": [
+          "nodes"
+        ],
+        "decimal_places": 1
+      },
+      {
+        "name": "edge-outputs",
+        "type": "CSV",
+        "format": "long",
+        "filename": "piecewise-link1-edges.csv",
+        "metric_set": [
+          "edges"
+        ],
+        "decimal_places": 1
+      }
     ]
   }
 }