small fixes

1 files changed, 5 insertions, 1 deletions

diff --git a/simulation/main.py b/simulation/main.py
index 3e458a7..402aa8d 100644
--- a/simulation/main.py
+++ b/simulation/main.py
@@ -57,8 +57,12 @@ def build_cascade_list(cascades, collapse=False):
 
 def cascadeLkl(graph, infect, sus):
     # There is a problem with the current implementation
+    # Note that you need to take into account the time diff between the label
+    # and the values being conditioned. Note also that the matrix if stacked as
+    # such will require to keep track of the state 0 of each cascade.
     a = np.dot(infect, graph)
-    return np.log(1. - np.exp(-a[(infect)*sus])).sum() - a[(~infect)*sus].sum()
+    return np.log(1. - np.exp(-a[(infect[1:])*sus[1:]])).sum() \
+            - a[(~infect[1:])*sus].sum()
 
 
 if __name__ == "__main__":