performance improvement for flattening lists|

Author: geraintpalmer

ciw/auxiliary.py

@@ -44,4 +44,10 @@ def truncated_normal(mean, sd):
 	sample = random.normalvariate(mean, sd)
 	while sample <= 0.0:
 		sample = random.normalvariate(mean, sd)
-	return sample
+	return sample
+
+def flatten_list(list_of_lists):
+	flat = []
+	for a_list in list_of_lists:
+		flat += a_list
+	return flat

ciw/node.py

@@ -6,7 +6,7 @@
 
 import networkx as nx
 
-from .auxiliary import random_choice
+from .auxiliary import random_choice, flatten_list
 from .data_record import DataRecord
 from .server import Server
 
@@ -67,8 +67,7 @@ def __init__(self, id_, simulation):
 
     @property
     def all_individuals(self):
-        return [i for priority_class in self.individuals
-                for i in priority_class]
+        return flatten_list(self.individuals)
 
     def __repr__(self):
         """

ciw/tests/test_auxiliary.py

@@ -77,4 +77,10 @@ def test_randomchoice(self):
         self.assertEqual(choice_counts, {'Exit Node': 100})
         self.assertEqual(r1, r2)
 
-
+    def test_flatten_list(self):
+        for seed in range(20):
+            random.seed(seed)
+            all_classes = [[random.random() for _ in range(random.randrange(3, 30, 1))] for _ in range(random.randrange(5, 20, 1))]
+            A = [i for priority in all_classes for i in priority]
+            B = ciw.flatten_list(all_classes)
+            self.assertEqual(A, B)