Add plotting function (#54)

* Add plotting function

* Plot notebooks

* add empty line at the end of environment

Author: jan-janssen

environment.yml

@@ -7,3 +7,4 @@ dependencies:
 - pygraphviz=1.14
 - aiida-workgraph=0.5.1
 - conda_subprocess =0.0.6
+- networkx=3.4.2

python_workflow_definition/src/python_workflow_definition/plot.py

@@ -0,0 +1,43 @@
+import json
+
+from IPython.display import SVG, display
+import networkx as nx
+
+
+from python_workflow_definition.purepython import group_edges
+from python_workflow_definition.shared import (
+    get_kwargs,
+    convert_nodes_list_to_dict,
+    NODES_LABEL,
+    EDGES_LABEL,
+    SOURCE_LABEL,
+    SOURCE_PORT_LABEL,
+)
+
+
+def plot(file_name):
+    with open(file_name, "r") as f:
+        content = json.load(f)
+
+    graph = nx.DiGraph()
+    node_dict = convert_nodes_list_to_dict(nodes_list=content[NODES_LABEL])
+    total_lst = group_edges(edges_lst=content[EDGES_LABEL])
+
+    for node_id, node_name in node_dict.items():
+        graph.add_node(node_id, name=str(node_name), label=str(node_name))
+
+    for edge_tuple in total_lst:
+        target_node, edge_dict = edge_tuple
+        edge_label_dict = {}
+        for k, v in edge_dict.items():
+            if v[SOURCE_LABEL] not in edge_label_dict:
+                edge_label_dict[v[SOURCE_LABEL]] = []
+            if v[SOURCE_PORT_LABEL] is None:
+                edge_label_dict[v[SOURCE_LABEL]].append(k)
+            else:
+                edge_label_dict[v[SOURCE_LABEL]].append(v[SOURCE_PORT_LABEL] + "=" + k)
+        for k, v in edge_label_dict.items():
+            graph.add_edge(str(k), str(target_node), label=", ".join(v))
+
+    svg = nx.nx_agraph.to_agraph(graph).draw(prog="dot", format="svg")
+    display(SVG(svg))

universal_workflow_nfdi.ipynb

@@ -7,6 +7,25 @@
     "# Load Simple Workflow"
    ]
   },
+  {
+   "metadata": {},
+   "cell_type": "markdown",
+   "source": "## Plot"
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "outputs": [],
+   "execution_count": null,
+   "source": "from python_workflow_definition.plot import plot"
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "outputs": [],
+   "execution_count": null,
+   "source": "plot(file_name='workflow_nfdi.json')"
+  },
   {
    "cell_type": "markdown",
    "metadata": {},

universal_workflow_qe.ipynb

@@ -7,6 +7,25 @@
     "# Load Quantum Espresso Energy Volume Curve Workflow"
    ]
   },
+  {
+   "metadata": {},
+   "cell_type": "markdown",
+   "source": "## Plot"
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "outputs": [],
+   "execution_count": null,
+   "source": "from python_workflow_definition.plot import plot"
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "outputs": [],
+   "execution_count": null,
+   "source": "plot(file_name='workflow_qe.json')"
+  },
   {
    "cell_type": "markdown",
    "metadata": {},

universal_workflow_simple.ipynb

@@ -7,6 +7,25 @@
     "# Load Simple Workflow"
    ]
   },
+  {
+   "metadata": {},
+   "cell_type": "markdown",
+   "source": "## Plot"
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "outputs": [],
+   "execution_count": null,
+   "source": "from python_workflow_definition.plot import plot"
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "outputs": [],
+   "execution_count": null,
+   "source": "plot(file_name='workflow_simple.json')"
+  },
   {
    "cell_type": "markdown",
    "metadata": {},