Big code cleanup

1 files changed, 64 insertions, 0 deletions

diff --git a/simulation/utils.py b/simulation/utils.py
new file mode 100644
index 0000000..aad7771
--- /dev/null
+++ b/simulation/utils.py
@@ -0,0 +1,64 @@
+import numpy as np
+import numpy.random as nr
+from six.moves import range
+
+
+def create_random_graph(n_nodes, p=.5):
+    graph = .5 * np.random.binomial(2, p=.5, size=(n_nodes, n_nodes))
+    for k in range(len(graph)):
+        graph[k, k] = 0
+    return np.log(1. / (1 - p * graph))
+
+
+def create_wheel(n_nodes, p=.5):
+    graph = np.zeros((n_nodes, n_nodes))
+    graph[0] = np.ones(n_nodes)
+    graph[0, 0] = 0
+    for i in range(1, n_nodes-1):
+        graph[i, i + 1] = 1
+    graph[n_nodes-1, 1] = 1
+    return np.log(1. / (1 - p * graph))
+
+
+def simulate_cascade(x, graph):
+    """
+    Simulate an IC cascade given a graph and initial state.
+
+    For each time step we yield:
+        - susc: the nodes susceptible at the beginning of this time step
+        - x: the subset of susc who became infected
+    """
+    yield x, np.zeros(graph.shape[0], dtype=bool)
+    susc = np.ones(graph.shape[0], dtype=bool)
+    while np.any(x):
+        susc = susc ^ x  # nodes infected at previous step are now inactive
+        if not np.any(susc):
+            break
+        x = 1 - np.exp(-np.dot(graph.T, x))
+        y = nr.random(x.shape[0])
+        x = (x >= y) & susc
+        yield x, susc
+
+
+def random_source(graph, node_p=None):
+    n_nodes = graph.shape[0]
+    if node_p is None:
+        node_p = np.ones(n_nodes) / n_nodes
+    x0 = np.zeros(graph.shape[0], dtype=bool)
+    x0[nr.choice(n_nodes, p=node_p)] = True
+    return x0
+
+
+def simulate_cascades(n_obs, graph, source=random_source):
+    n_nodes = graph.shape[0]
+    x_obs = np.zeros((n_obs, n_nodes), dtype=bool)
+    s_obs = np.zeros((n_obs, n_nodes), dtype=bool)
+    i = 0
+    while i < n_obs:
+        for x, s in simulate_cascade(source(graph), graph):
+            x_obs[i] = x
+            s_obs[i] = s
+            i += 1
+            if i >= n_obs:
+                break
+    return x_obs, s_obs