From f457b3480d53920a0d8b0e3b8cdb2b18601088ee Mon Sep 17 00:00:00 2001
From: Thibaut Horel <thibaut.horel@gmail.com>
Date: Mon, 30 Nov 2015 20:01:39 -0500
Subject: Add constant source

---
 simulation/utils.py | 9 +++++++++
 1 file changed, 9 insertions(+)

(limited to 'simulation')

diff --git a/simulation/utils.py b/simulation/utils.py
index aad7771..37cabf0 100644
--- a/simulation/utils.py
+++ b/simulation/utils.py
@@ -49,6 +49,15 @@ def random_source(graph, node_p=None):
     return x0
 
 
+def constant_source(graph, source):
+    if type(source) == int:
+        x = np.zeros(graph.shape[0], dtype=bool)
+        x[source] = True
+        return x
+    else:
+        return source
+
+
 def simulate_cascades(n_obs, graph, source=random_source):
     n_nodes = graph.shape[0]
     x_obs = np.zeros((n_obs, n_nodes), dtype=bool)
-- 
cgit v1.2.3-70-g09d2