calculate sampling time for averaging and improve logging (#524)

* calculate sampling time for averaging and improve logging

* adjust tests to changes from sum to mean

---------

Co-authored-by: Eric Hartmann <hartmaec@rh05659.villa-bosch.de>

Author: ehhartmann

src/kimmdy/kmc.py

@@ -133,23 +133,36 @@ def rf_kmc(
     reaction_probability = []
 
     # 1. Calculate the probability for each reaction
+    # get overall sampling time
+    # TODO: Can be replaced if we parse mdps for the simulation time
+    t_min = recipe_collection.recipes[0].timespans[0][0]
+    t_max = recipe_collection.recipes[0].timespans[0][1]
+    for recipe in recipe_collection.recipes:
+        for ts in recipe.timespans:
+            t1 = ts[0]
+            t2 = ts[1]
+            t_min = t1 if t1 < t_min else t_min
+            t_max = t2 if t2 > t_max else t_max
+    sampling_time = t_max - t_min
+    logger.debug(f"Sampling time: {sampling_time}ps")
+
     for recipe in recipe_collection.recipes:
         dt = [x[1] - x[0] for x in recipe.timespans]
-        reaction_probability.append(sum(np.multiply(dt, recipe.rates)))
+        reaction_probability.append(sum(np.multiply(dt, recipe.rates)) / sampling_time)
 
     # 2. Set the total rate to the sum of individual rates
-    probability_cumulative = np.cumsum(reaction_probability)
-    probability_sum = probability_cumulative[-1]
+    rate_cumulative = np.cumsum(reaction_probability)
+    total_rate = rate_cumulative[-1]
 
     # 3. Generate two independent uniform (0,1) random numbers u1,u2
     u = rng.random(2)
     logger.debug(
-        f"\tRandom values u: {u}, number cumulative probabilities "
-        f"{len(probability_cumulative)}, probability sum {probability_sum}"
+        f"\tRandom values u: {u}, number cumulative rates "
+        f"{len(rate_cumulative)}, total rate {total_rate}"
     )
 
     # 4. Find the even to carry out, mu, using binary search (np.searchsorted)
-    pos = np.searchsorted(probability_cumulative, u[0] * probability_sum)
+    pos = np.searchsorted(rate_cumulative, u[0] * total_rate)
     recipe: Recipe = recipe_collection.recipes[pos]
     time_index = np.argmax(recipe.rates)
     reaction_time = recipe.timespans[time_index][1]
@@ -158,10 +171,8 @@ def rf_kmc(
     )
 
     # 5. Calculate the time step associated with uu
-    time_delta = np.log(1 / u[1]) / probability_sum
-    logger.info(
-        f"Time delta: {time_delta} ps with cumulative probability: {probability_sum}"
-    )
+    time_delta = np.log(1 / u[1]) / total_rate
+    logger.info(f"Time delta: {time_delta} ps with total rate: {total_rate}")
 
     return KMCAccept(
         recipe=recipe,
@@ -213,11 +224,27 @@ def frm(
     # 1. Generate M independent uniform (0,1) random numbers
     u = rng.random(len(recipe_collection.recipes))
     # 2. Calculate the cumulative probabilities for each event
+    # get overall sampling time
+    # TODO: Can be replaced if we parse mdps for the simulation time
+    t_min = recipe_collection.recipes[0].timespans[0][0]
+    t_max = recipe_collection.recipes[0].timespans[0][1]
+    for recipe in recipe_collection.recipes:
+        for ts in recipe.timespans:
+            t1 = ts[0]
+            t2 = ts[1]
+            t_min = t1 if t1 < t_min else t_min
+            t_max = t2 if t2 > t_max else t_max
+    sampling_time = t_max - t_min
+    logger.debug(f"Sampling time: {sampling_time}ps")
+
     for i, recipe in enumerate(recipe_collection.recipes):
         dt = [x[1] - x[0] for x in recipe.timespans]
         cumulative_probability = np.cumsum(
             list(map(lambda x, y: x * y, dt, recipe.rates))
         )
+        cumulative_probability = (
+            cumulative_probability / sampling_time
+        )  # calculate mean instead of sum
         reaction_probability.append(cumulative_probability[-1])
         # 3. For each event k, find the time until the this reaction takes place tau_k,
         #    if it is during the time when the propensities are defined (simulation time)
@@ -534,9 +561,24 @@ def multi_rfkmc(
         reaction_probability = []
 
         # 1. Calculate the probability for each reaction
-        for recipe in recipes:
+        # get overall sampling time
+        # TODO: Can be replaced if we parse mdps for the simulation time
+        t_min = recipe_collection.recipes[0].timespans[0][0]
+        t_max = recipe_collection.recipes[0].timespans[0][1]
+        for recipe in recipe_collection.recipes:
+            for ts in recipe.timespans:
+                t1 = ts[0]
+                t2 = ts[1]
+                t_min = t1 if t1 < t_min else t_min
+                t_max = t2 if t2 > t_max else t_max
+        sampling_time = t_max - t_min
+        logger.debug(f"Sampling time: {sampling_time}ps")
+
+        for recipe in recipe_collection.recipes:
             dt = [x[1] - x[0] for x in recipe.timespans]
-            reaction_probability.append(sum(np.multiply(dt, recipe.rates)))
+            reaction_probability.append(
+                sum(np.multiply(dt, recipe.rates)) / sampling_time
+            )
 
         # 2. Set the total rate to the sum of individual rates
         probability_cumulative = np.cumsum(reaction_probability)

tests/test_kmc.py

@@ -36,8 +36,8 @@ def reference_KMC() -> KMCAccept:
         recipe=Recipe(
             [Bind(2, 3)], rates=[0.12, 0.0], timespans=[(0.0, 6.0), (6.0, 10.0)]
         ),
-        time_delta=0.04032167624965666,
-        reaction_probability=[0.0, 0.72, 0.54, 0.0],
+        time_delta=0.4032167624965666,
+        reaction_probability=[0.0, 0.072, 0.054, 0.0],
         time_start=0,
         time_start_index=0,
     )
@@ -51,8 +51,8 @@ def reference_multi2_KMC() -> KMCAccept:
             rates=[0.12, 0.0, 0.15, 0.06],
             timespans=[(0.0, 6.0), (6.0, 10.0), (2.0, 4.0), (4.0, 8.0)],
         ),
-        time_delta=0.09762211196506274,
-        reaction_probability=[0.0, 0.72, 0.54, 0.0],
+        time_delta=0.7809768957205019,
+        reaction_probability=[0.0, 0.072, 0.054, 0.0],
         time_start=6.0,
         time_start_index=0,
     )