simple heuristic + star graph + rmse computation

1 files changed, 10 insertions, 0 deletions

diff --git a/simulation/main.py b/simulation/main.py
index a916034..81133c7 100644
--- a/simulation/main.py
+++ b/simulation/main.py
@@ -19,6 +19,16 @@ def create_random_graph(n_nodes, p=.5):
     return np.log(1. / (1 - p * graph))
 
 
+def create_star(n_nodes, p=.5):
+    graph = np.zeros((n_nodes, n_nodes))
+    graph[0] = np.ones((n_nodes,))
+    graph[0, 0] = 0
+    for index, row in enumerate(graph[1:-1]):
+        row[index + 1] = 1
+    graph[-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.