skip virtual nodes when creating d3 graph data from model json (#885)

* skip virtual nodes when creating d3 graph data from model json

* added test to check that the from_model and from_json funcs return similar data

* added docstring to graph data test

* Tidy up test docstring.

Co-authored-by: James Batchelor <james.batchelor@atkinsglobal.com>
Co-authored-by: James Tomlinson <tomo.bbe@gmail.com>

pywr/notebook/__init__.py

@@ -285,18 +285,14 @@ def pywr_json_to_d3_json(model, attributes=False):
         with open(model) as d:
             model = json.load(d)
 
-    nodes = [node["name"] for node in model["nodes"]]
-
-    edges = []
-    for edge in model["edges"]:
-        sourceindex = nodes.index(edge[0])
-        targetindex = nodes.index(edge[1])
-        edges.append({'source': sourceindex, 'target': targetindex})
-
     nodes = []
     node_classes = create_node_class_trees()
 
     for node in model["nodes"]:
+
+        if node["type"].lower() in ["annualvirtualstorage", "virtualstorage"]:
+            continue
+
         json_node = {'name': node["name"], 'clss': node_classes[node["type"].lower()]}
         try:
             json_node['position'] = node['position']['schematic']
@@ -327,6 +323,14 @@ def pywr_json_to_d3_json(model, attributes=False):
 
         nodes.append(json_node)
 
+    nodes_names = [node["name"] for node in nodes]
+
+    edges = []
+    for edge in model["edges"]:
+        sourceindex = nodes_names.index(edge[0])
+        targetindex = nodes_names.index(edge[1])
+        edges.append({'source': sourceindex, 'target': targetindex})
+
     graph = {
         "nodes": nodes,
         "links": edges}

tests/test_notebook.py

@@ -43,3 +43,22 @@ def test_from_json(from_json):
     demand_max_flow = get_node_attribute(demand, "max_flow")
 
     assert demand_max_flow["value"] == "demand_max_flow - AggregatedParameter"
+
+
+def test_d3_data():
+    """Test returned by `pywr_json_to_d3_json` and `pywr_model_to_d3_json` is similar.
+
+    These return graph data from a JSON file and Model instance respectively. Here we test that each returns the
+    same node names and number on links. The data won't match exactly due to differences in node ordering.
+    """
+    json_path = os.path.join(os.path.dirname(__file__), "models", "demand_saving2_with_variables.json")
+    model = load_model("demand_saving2_with_variables.json")
+
+    d3_data_from_json = pywr_json_to_d3_json(json_path)
+    d3_data_from_model = pywr_model_to_d3_json(model)
+
+    json_nodes = {n["name"]: n for n in d3_data_from_json["nodes"]}
+    model_nodes = {n["name"]: n for n in d3_data_from_model["nodes"]}
+
+    assert json_nodes == model_nodes
+    assert len(d3_data_from_json["links"]) == len(d3_data_from_model["links"])