pylint work

examples/example2.py

@@ -7,29 +7,31 @@ from collections import namedtuple
 from os import path
 from multiprocessing import Pool
 from typing import ChainMap
-from numpy import mean
 
-import yaml
 import sys
+import copy
 import csv
-import matplotlib.pyplot as plt
 import time
 import pandas as pd
-import copy
+import matplotlib.pyplot as plt
+import yaml
 
-from pydtnsim import Network, RandomTraffic, Node, EpidemicNode, CSVTrace
-from pydtnsim.community import BubbleKCliqueNode, BubbleLouvainNode
-from pydtnsim.community import HCBFKCliqueNode, HCBFLouvainNode
+from numpy import mean
+from pydtnsim import Network, RandomTraffic, CSVTrace
 from gui2 import Gui
 
-Simulation = namedtuple("Simulation", ["tag", "trace", "node_type", "seed", "node_options", "traffic_options"])
+Simulation = namedtuple(
+    "Simulation",
+    ["tag", "trace", "node_type", "seed", "node_options", "traffic_options"],
+)
+
 
 def run_simulation(simulation):
     """Run a simulation"""
-    
-    csv = path.join(simulation.trace, "contact.csv")
+
+    trace_csv = path.join(simulation.trace, "contact.csv")
     metadata = path.join(simulation.trace, "metadata.json")
-    trace = CSVTrace(csv, metadata=metadata)
+    trace = CSVTrace(trace_csv, metadata=metadata)
 
     nodes = {
         node_id: simulation.node_type(**simulation.node_options)
@@ -45,43 +47,49 @@ def run_simulation(simulation):
         "trace": simulation.trace,
         "node_type": simulation.node_type.__name__,
         "seed": simulation.seed,
-        "tag" : simulation.tag,
-        "ttl" : simulation.traffic_options["time_to_live"]
+        "tag": simulation.tag,
+        "ttl": simulation.traffic_options["time_to_live"],
     }
     stats.update(network.stats_summary)
 
     # return stats because we can't pickle the network as it is a generator.
     return stats
 
+
 def bar_plot(graphing_options, results):
+    """
+    bar plot sim results
+    todo: 2+ plots
+    """
     choice = graphing_options["choice"]
     dependant = graphing_options["dependant"]
-    data = [ [inp] for inp in graphing_options["input"] ]
+    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))]
+        batch_y = [[] for _ in range(len(data))]
         # split the stats into batches
         for stat in sorted_results:
-            yyys[stat['tag']].append([stat[dependant[0]]])
+            batch_y[stat["tag"]].append([stat[dependant[0]]])
 
-        for dat, y in zip(data, yyys):
-            dat.append(mean(y))
+        for dat, batch in zip(data, batch_y):
+            dat.append(mean(batch))
 
-    df = pd.DataFrame(data, columns=columns)
-    print(df)
-    df.plot(x=choice, kind='bar', stacked=False, rot=0, ylabel=dependant[0])
+    pd.DataFrame(data, columns=columns).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.legend(bbox_to_anchor=(1.05, 1), loc="upper left", borderaxespad=0)
     plt.show()
 
 
 def line_plot(graphing_options, results):
+    """line plot sim results."""
     dependant = graphing_options["dependant"]
     (_, axes) = plt.subplots(len(dependant))
     x = graphing_options["input"]
@@ -89,58 +97,44 @@ def line_plot(graphing_options, results):
     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
+        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))]
+                batch_y = [[] for _ in range(len(x))]
                 # split the stats into batches
                 for stat in sorted_results:
-                    yyys[stat['tag']].append([stat[dep]])
+                    batch_y[stat["tag"]].append([stat[dep]])
 
-                y = [mean(stat) for stat in yyys]
+                y = [mean(stat) for stat in batch_y]
                 ax.plot(x, y, "o-", label=node)
                 ax.set_ylabel(dep)
         else:
-            yyys = [ [] for _ in range(len(x))]
+            batch_y = [[] for _ in range(len(x))]
             # split the stats into batches
             for stat in sorted_results:
-                yyys[stat['tag']].append([stat[dependant[0]]])
+                batch_y[stat["tag"]].append([stat[dependant[0]]])
 
-            y = [mean(stat) for stat in yyys]
+            y = [mean(stat) for stat in batch_y]
             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.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"]:
-        with open(args["config"], "r", newline="") as f:
-            config = yaml.load(f, Loader=yaml.FullLoader)
-    else:
-        # put up GUI for user input
-        gui = Gui(None)
-        gui.title("Pydtnsim")
-
-        if args["config"]:
-            gui.read_config(args["config"])
-
-        gui.mainloop()
-        config = gui.param
-
-    pool = Pool()
-    simulations = []
-
+def create_sim_list(config):
+    """create a list of Simulations."""
     # create base list of sims
     base_sims = []
     trace = config["SimOptions"]["contact_dir"]
     seeds = config["SimOptions"]["seeds"]
     node_options = ChainMap(config["NodeOptions"], {"context": {}})
-    traffic_options = ChainMap(config["TrafficOptions"], {"start": config["NodeOptions"]["epoch"]})
+    traffic_options = ChainMap(
+        config["TrafficOptions"], {"start": config["NodeOptions"]["epoch"]}
+    )
 
     for seed in seeds:
         for node_type in config["NodeChoices"]["Nodes"]:
@@ -148,58 +142,106 @@ def main(args):
             base_sims.append(sim)
 
     # expand base list by the independant variable
+    simulations = []
     choice = config["GraphingOptions"]["choice"]
-    input = config["GraphingOptions"]["input"]
+    idp_var = config["GraphingOptions"]["input"]
     if choice in config["NodeOptions"].keys():
-        input.sort()
-        for tag, option in enumerate(input):
+        idp_var.sort()
+        for tag, option in enumerate(idp_var):
             new_node_options = copy.deepcopy(node_options)
             new_node_options[choice] = option
             for base_sim in base_sims:
-                sim = Simulation(tag, base_sim.trace, base_sim.node_type, base_sim.seed, new_node_options, base_sim.traffic_options)
+                sim = Simulation(
+                    tag,
+                    base_sim.trace,
+                    base_sim.node_type,
+                    base_sim.seed,
+                    new_node_options,
+                    base_sim.traffic_options,
+                )
                 simulations.append(sim)
     elif choice in config["TrafficOptions"].keys():
-        input.sort()
-        for tag, option in enumerate(input):
+        idp_var.sort()
+        for tag, option in enumerate(idp_var):
             new_traffic_options = copy.deepcopy(traffic_options)
             new_traffic_options[choice] = option
             for base_sim in base_sims:
-                sim = Simulation(tag, base_sim.trace, base_sim.node_type, base_sim.seed, base_sim.node_options, new_traffic_options)
+                sim = Simulation(
+                    tag,
+                    base_sim.trace,
+                    base_sim.node_type,
+                    base_sim.seed,
+                    base_sim.node_options,
+                    new_traffic_options,
+                )
                 simulations.append(sim)
     elif choice == "contact_dir":
-        for tag, trace in enumerate(input):     # just trust that they list the traces in order
+        for tag, trace in enumerate(
+            idp_var
+        ):  # just trust that they list the traces in order
             for base_sim in base_sims:
-                sim = Simulation(tag, trace, base_sim.node_type, base_sim.seed, base_sim.node_options, base_sim.traffic_options)
+                sim = Simulation(
+                    tag,
+                    trace,
+                    base_sim.node_type,
+                    base_sim.seed,
+                    base_sim.node_options,
+                    base_sim.traffic_options,
+                )
                 simulations.append(sim)
     else:
         simulations = base_sims
 
+    return simulations
+
+
+def main(args):
+    """Run a simulation for each seed, for each independant variable, graph the results."""
+    if args["no_gui"]:
+        with open(args["config"], "r", newline="", encoding="utf8") as f:
+            config = yaml.load(f, Loader=yaml.FullLoader)
+    else:
+        # put up GUI for user input
+        gui = Gui(None)
+        gui.title("Pydtnsim")
+
+        if args["config"]:
+            gui.read_config(args["config"])
+
+        gui.mainloop()
+        config = gui.param
+
+    pool = Pool()
+    simulations = create_sim_list(config)
+
     results = {}
 
     print("sim running, please wait :)")
     start = time.time()
     for _ in range(config["SimOptions"]["batch"]):
         for stats in pool.imap_unordered(run_simulation, simulations):
-            type = stats["node_type"]
-            if type not in results:
-                results[type] = []
-            results[type].append(stats)
+            node_type = stats["node_type"]
+            if node_type not in results:
+                results[node_type] = []
+            results[node_type].append(stats)
     end = time.time()
-    print("sim runtime:", end-start)
+    print("sim runtime:", end - start)
 
     # dump stats in csv
     with open("testing.csv", "w", newline="") as results_file:
         for node_type in results:
             fieldnames = results[node_type][0].keys()
-            writer = csv.DictWriter(results_file, fieldnames=fieldnames, extrasaction='ignore')  # extrasaction avoids intermittent dictwriter valueerror
+            writer = csv.DictWriter(
+                results_file, fieldnames=fieldnames, extrasaction="ignore"
+            )  # extrasaction avoids intermittent dictwriter valueerror
             writer.writeheader()
             for result in results[node_type]:
                 writer.writerow(result)
 
     # graph results
-    line_plot(config["GraphingOptions"], results)
+    bar_plot(config["GraphingOptions"], results)
+
 
-    
 def parse_args(args):
     """Parse arguments."""
     parser = ArgumentParser()
@@ -217,5 +259,6 @@ def parse_args(args):
     args = parser.parse_args(args)
     return vars(args)
 
+
 if __name__ == "__main__":
     sys.exit(main(parse_args(sys.argv[1:])))