moved graphing code to function

examples/example2.py

@@ -15,7 +15,6 @@ import csv
 import matplotlib.pyplot as plt
 import time
 import pandas as pd
-import pprint
 import copy
 
 from pydtnsim import Network, RandomTraffic, Node, EpidemicNode, CSVTrace
@@ -26,7 +25,6 @@ from gui2 import Gui
 Simulation = namedtuple("Simulation", ["tag", "trace", "node_type", "seed", "node_options", "traffic_options"])
 
 def run_simulation(simulation):
-    #print(simulation.node_options)
     """Run a simulation"""
     
     csv = path.join(simulation.trace, "contact.csv")
@@ -55,6 +53,69 @@ def run_simulation(simulation):
     # return stats because we can't pickle the network as it is a generator.
     return stats
 
+def bar_plot(graphing_options, results):
+    choice = graphing_options["choice"]
+    dependant = graphing_options["dependant"]
+    data = [ [inp] for inp in graphing_options["input"] ]
+    columns = [choice]
+
+    for node, stats in results.items():
+        columns.append(node)
+        sorted_results = sorted(stats, key=lambda x: x["tag"])
+
+        yyys = [ [] for _ in range(len(data))]
+        # split the stats into batches
+        for stat in sorted_results:
+            yyys[stat['tag']].append([stat[dependant[0]]])
+
+        for dat, y in zip(data, yyys):
+            dat.append(mean(y))
+
+    df = pd.DataFrame(data, columns=columns)
+    print(df)
+    df.plot(x=choice, kind='bar', stacked=False, rot=0, ylabel=dependant[0])
+    plt.subplots_adjust(right=0.75)
+            
+    plt.xlabel(choice)
+    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left', borderaxespad=0)
+    plt.show()
+
+
+def line_plot(graphing_options, results):
+    dependant = graphing_options["dependant"]
+    (_, axes) = plt.subplots(len(dependant))
+    x = graphing_options["input"]
+
+    for node, stats in results.items():
+        sorted_results = sorted(stats, key=lambda x: x["tag"])
+
+        if len(dependant) > 1:      # avoids axes is not subscriptable error when plotting one graph
+            for ax, dep in zip(axes, dependant):
+                yyys = [ [] for _ in range(len(x))]
+                # split the stats into batches
+                for stat in sorted_results:
+                    yyys[stat['tag']].append([stat[dep]])
+
+                y = [mean(stat) for stat in yyys]
+                ax.plot(x, y, "o-", label=node)
+                ax.set_ylabel(dep)
+        else:
+            yyys = [ [] for _ in range(len(x))]
+            # split the stats into batches
+            for stat in sorted_results:
+                yyys[stat['tag']].append([stat[dependant[0]]])
+
+            y = [mean(stat) for stat in yyys]
+            axes.plot(x, y, "o-", label=node)
+            axes.set_ylabel(dependant[0])
+
+    plt.subplots_adjust(right=0.75)
+    plt.xlabel(graphing_options["choice"])
+    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left', borderaxespad=0)
+    plt.show()
+
+
+
 def main(args):
     """Run a simulation for each seed, for each independant variable, graph the results."""
     if args["no_gui"]:
@@ -136,79 +197,7 @@ def main(args):
                 writer.writerow(result)
 
     # graph results
-
-    """     line plot
-    dependant = config["GraphingOptions"]["dependant"]
-    (_, axes) = plt.subplots(len(dependant))
-    x = input
-
-    for node, stats in results.items():
-        sorted_results = sorted(stats, key=lambda x: x["tag"])
-
-        if len(dependant) > 1:      # avoids axes is not subscriptable error when plotting one graph
-            for ax, dep in zip(axes, dependant):
-                yyys = [ [] for _ in range(len(input))]
-                # split the stats into batches
-                for stat in sorted_results:
-                    yyys[stat['tag']].append([stat[dep]])
-
-                y = [mean(stat) for stat in yyys]
-                ax.plot(x, y, "o-", label=node)
-                ax.set_ylabel(dep)
-        else:
-            yyys = [ [] for _ in range(len(input))]
-            # split the stats into batches
-            for stat in sorted_results:
-                yyys[stat['tag']].append([stat[dependant[0]]])
-
-            y = [mean(stat) for stat in yyys]
-            axes.plot(x, y, "o-", label=node)
-            axes.set_ylabel(dependant[0])
-    """
-
-    # bar plot
-
-    # just one plot for now
-    choice = config["GraphingOptions"]["choice"]
-    dependant = config["GraphingOptions"]["dependant"]
-    data = [ [inp] for inp in input ]
-    columns = [choice]
-
-    for node, stats in results.items():
-        columns.append(node)
-        sorted_results = sorted(stats, key=lambda x: x["tag"])
-
-        yyys = [ [] for _ in range(len(input))]
-        # split the stats into batches
-        for stat in sorted_results:
-            yyys[stat['tag']].append([stat[dependant[0]]])
-
-        for dat, y in zip(data, yyys):
-            dat.append(mean(y))
-
-    df = pd.DataFrame(data, columns=columns)
-    print(df)
-    df.plot(x=choice, kind='bar', stacked=False, title="test", rot=0)
-    plt.subplots_adjust(right=0.75)
-            
-        
-
-
-
-
-
-
-
-
-
-
-    plt.xlabel(choice)
-    plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left', borderaxespad=0)
-    plt.show()
-
-
-
-
+    line_plot(config["GraphingOptions"], results)
 
     
 def parse_args(args):