changed createStanfordGraph -> import_from_file & created Cascade class

1 files changed, 36 insertions, 15 deletions

diff --git a/jpa_test/cascade_creation.py b/jpa_test/cascade_creation.py
index 19e799f..01332f8 100644
--- a/jpa_test/cascade_creation.py
+++ b/jpa_test/cascade_creation.py
@@ -1,5 +1,6 @@
 import networkx as nx
 import numpy as np
+import collections
 
 class InfluenceGraph(nx.Graph):
     """
@@ -14,23 +15,15 @@ class InfluenceGraph(nx.Graph):
         self.add_nodes_from(G.nodes())
         self.add_edges_from(G.edges())
 
-    def createStanfordGraph(self, file):
+    def import_from_file(self, file_name):
         """
         Takes a file from the Stanford collection of networks
-        Need to remove comments on top of the file
-        Graph still needs to be weighted on the edges
         """
-        f = open(file, 'r')
-        data = f.readlines()
-        G = nx.DiGraph()
-        for edge in data:
-            split1 = edge.split('\t')
-            split2 = split1[1].split('\n')
-            u = int(split1[0])
-            v = int(split2[0])
-            G.add_edge(u,v)
-        self.add_nodes_from(G.nodes())
-        self.add_edges_from(G.edges())
+        with open(file_name, 'r') as f:
+            for edge in f:
+                if "#" not in edge:
+                    u, v = [int(node) for node in edge.split()]
+                    self.add_edge(u, v)
 
     @property
     def mat(self):
@@ -46,6 +39,34 @@ class InfluenceGraph(nx.Graph):
         return self._logmat
 
 
+class Cascade(list):
+    """
+    Cascade object: list with attributes
+    """
+    def __init__(self, *args, **kwargs):
+        super(Cascade, self).__init__(*args, **kwargs)
+
+    def infection_time(self, node):
+        """
+        Returns lists of infections times for node i in cascade
+        """
+        infected_times = []
+        for t, infected_set in zip(xrange(len(self)), self):
+            if infected_set[node]:
+                infected_times.append(t)
+        return infected_times
+
+    def candidate_infectors(self, node):
+        """
+        Returns Counter of nodes infected just before node i was
+        """
+        candidate_infectors = collections.Counter()
+        for t in self.infection_time(node):
+            if t > 0:
+                candidate_infectors.update(np.where(self[t-1])[0])
+        return candidate_infectors
+
+
 def icc_cascade(G, p_init):
     """
     Returns boolean vectors for one cascade
@@ -55,7 +76,7 @@ def icc_cascade(G, p_init):
     susceptible = np.ones(G.number_of_nodes(), dtype=bool)
     active = np.random.rand(G.number_of_nodes()) < p_init
     susceptible = susceptible - active
-    cascade = []
+    cascade = Cascade()
     while active.any() and susceptible.any():
         cascade.append(active)
         active = np.exp(np.dot(G.logmat, active)) \