From 1373c612ad01292cfd1b7e7c2bdcfe16b84145f2 Mon Sep 17 00:00:00 2001
From: barneydobson <barnaby.dobson1@gmail.com>
Date: Sat, 16 Mar 2024 15:22:26 +0000
Subject: [PATCH] Add new metrics

---
 swmmanywhere/metric_utilities.py |  5 +++-
 tests/test_metric_utilities.py   | 48 +++++++++++++++++++++++++-------
 2 files changed, 42 insertions(+), 11 deletions(-)

diff --git a/swmmanywhere/metric_utilities.py b/swmmanywhere/metric_utilities.py
index a00d9928..245ffc5f 100644
--- a/swmmanywhere/metric_utilities.py
+++ b/swmmanywhere/metric_utilities.py
@@ -441,6 +441,7 @@ def outlet_nse_flooding(synthetic_G: nx.Graph,
                          list(sg_syn.nodes),
                          list(sg_real.nodes))
 
+@metrics.register
 def outlet_kstest_diameters(real_G: nx.Graph,
                             synthetic_G: nx.Graph,
                             real_results: pd.DataFrame,
@@ -453,7 +454,6 @@ def outlet_kstest_diameters(real_G: nx.Graph,
     'real' network is calculated by dominant_outlet, while the dominant outlet
     node of the 'synthetic' network is calculated by best_outlet_match.
     """
-
     # Identify synthetic and real outlet arcs
     sg_syn, _ = best_outlet_match(synthetic_G, real_subs)
     sg_real, _ = dominant_outlet(real_G, real_results)
@@ -464,6 +464,7 @@ def outlet_kstest_diameters(real_G: nx.Graph,
     return stats.ks_2samp(list(syn_diameters.values()),
                          list(real_diameters.values())).statistic
 
+@metrics.register
 def outlet_pbias_length(real_G: nx.Graph,
                         synthetic_G: nx.Graph,
                         real_results: pd.DataFrame,
@@ -485,6 +486,7 @@ def outlet_pbias_length(real_G: nx.Graph,
     real_length = sum([d['length'] for u,v,d in sg_real.edges(data=True)])
     return (syn_length - real_length) / real_length
 
+@metrics.register
 def outlet_pbias_nmanholes(real_G: nx.Graph,
                            synthetic_G: nx.Graph,
                            real_results: pd.DataFrame,
@@ -504,6 +506,7 @@ def outlet_pbias_nmanholes(real_G: nx.Graph,
     return (sg_syn.number_of_nodes() - sg_real.number_of_nodes()) \
         / sg_real.number_of_nodes()
 
+@metrics.register
 def outlet_pbias_npipes(real_G: nx.Graph,
                         synthetic_G: nx.Graph,
                         real_results: pd.DataFrame,
diff --git a/tests/test_metric_utilities.py b/tests/test_metric_utilities.py
index e5d7c6f1..da1542b4 100644
--- a/tests/test_metric_utilities.py
+++ b/tests/test_metric_utilities.py
@@ -282,27 +282,55 @@ def test_outlet_nse_flooding():
 def test_design_params():
     """Test the design param related metrics."""
     G = load_graph(Path(__file__).parent / 'test_data' / 'graph_topo_derived.json')
+    nx.set_edge_attributes(G, 0.15, 'diameter')
     subs = get_subs()
 
-    # Mock results
-    results = pd.DataFrame([{'object' : 4253560,
+    # Mock results (only needed for dominant outlet)
+    results = pd.DataFrame([{'id' : 4253560,
                              'variable' : 'flow',
                              'value' : 10,
                              'date' : pd.to_datetime('2021-01-01 00:00:00')},
-                             {'object' : 4253560,
+                             {'id' : 4253560,
                              'variable' : 'flow',
                              'value' : 5,
                              'date' : pd.to_datetime('2021-01-01 00:00:05')},
                              ])
     
-    # Calculate NSE (perfect results)
-    val = mu.metrics.outlet_kstest_diameters(synthetic_G = G,
-                                    synthetic_results = results,
-                                    real_G = G,
-                                    real_results = results,
-                                    real_subs = subs)
-    assert val == 0.0
+    # Target results
+    design_results = {'outlet_kstest_diameters' : 0.0625,
+               'outlet_pbias_length' : -0.15088965,
+               'outlet_pbias_nmanholes' : -0.05,
+               'outlet_pbias_npipes' : -0.15789473}
+    
+    # Iterate for G = G, i.e., perfect results
+    metrics = mu.iterate_metrics(synthetic_G = G,
+                                 synthetic_subs = None,
+                                 synthetic_results = None,
+                                 real_G = G,
+                                 real_subs = subs,
+                                 real_results = results,
+                                 metric_list = design_results.keys())
+    for metric, val in metrics.items():
+        assert metric in design_results
+        assert np.isclose(val, 0)
+
+    # edit the graph for target results
+    G_ = G.copy()
+    G_.remove_node(list(G.nodes)[0])
+    G_.edges[list(G_.edges)[0]]['diameter'] = 0.3
 
+    metrics = mu.iterate_metrics(synthetic_G = G_,
+                                 synthetic_subs = None,
+                                 synthetic_results = None,
+                                 real_G = G,
+                                 real_subs = subs,
+                                 real_results = results,
+                                 metric_list = design_results.keys())
+
+    for metric, val in metrics.items():
+        assert metric in design_results
+        assert np.isclose(val, design_results[metric]), metric
+        
 def test_netcomp_iterate():
     """Test the netcomp metrics and iterate_metrics."""
     netcomp_results = {'nc_deltacon0' : 0.00129408,