Merge pull request #88 from Heuro-labs/develop

Fix tau adaptive

Author: dileep-kishore

HISTORY.rst

@@ -2,6 +2,19 @@
 History
 =======
 
+0.7.1 (2019-03-23)
+------------------
+
+Changed
++++++++
+- Refactor ``tau_adaptive``
+- Rename ``direct_naive`` to ``direct``
+
+Fixed
++++++
+- SSA part of ``tau_adaptive``
+- Bug in linux compatibility of ``tau_adaptive``
+
 0.7.0 (2019-02-02)
 ------------------
 

codecov.yml

@@ -1,6 +1,6 @@
 coverage:
     ignore:
-        - "pyssa/direct_naive.py"
+        - "pyssa/direct.py"
         - "pyssa/utils.py"
         - "pyssa/tau_leaping.py"
         - "pyssa/tau_adaptive.py"

docs/tutorial.rst

@@ -74,7 +74,7 @@ The results of the simulation can be retrieved by accessing the ``Results`` obje
 
 .. code-block:: python
 
-    <Results n_rep=10 algorithm=direct_naive seed=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]>
+    <Results n_rep=10 algorithm=direct seed=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]>
 
 The ``Results`` object provides abstractions for easy retrieval and iteration over the simulation results. For example you can iterate over every run of the simulation using ::
 
@@ -112,7 +112,7 @@ Algorithms
 
 The ``Simulation`` class currently supports the following algorithms:
 
-1. Direct naive
+1. Direct
 2. Tau leaping
 
 You can change the algorithm used to perform a simulation using the ``simulation`` flag

pyssa/direct_naive.py renamed to pyssa/direct.py

@@ -1,5 +1,5 @@
 """
-    Naive implementation of the Gillespie algorithm (direct method) in Numba
+    Implementation of the direct method.
 """
 
 from typing import Tuple
@@ -11,7 +11,7 @@
 
 
 @njit(nogil=True, cache=False)
-def direct_naive(
+def direct(
     react_stoic: np.ndarray,
     prod_stoic: np.ndarray,
     init_state: np.ndarray,

pyssa/simulation.py

@@ -11,7 +11,7 @@
 import matplotlib.pyplot as plt
 import matplotlib.lines as mlines
 
-from .direct_naive import direct_naive
+from .direct import direct
 from .tau_leaping import tau_leaping
 from .tau_adaptive import tau_adaptive
 from .results import Results
@@ -135,7 +135,7 @@ def simulate(
         seed: Optional[List[int]] = None,
         n_rep: int = 1,
         n_procs: int = 1,
-        algorithm: str = "direct_naive",
+        algorithm: str = "direct",
         **kwargs,
     ):
         """
@@ -164,7 +164,7 @@ def simulate(
             The default value is 1
         algorithm : str, optional
             The algorithm to be used to run the simulation
-            The default value is "direct_naive"
+            The default value is "direct"
 
         Returns
         -------
@@ -194,7 +194,7 @@ def simulate(
             raise TypeError("max_iter should be of type int")
 
         algo_args = []
-        if algorithm == "direct_naive":
+        if algorithm == "direct":
             for index in range(n_rep):
                 algo_args.append(
                     (
@@ -209,7 +209,7 @@ def simulate(
                         self._chem_flag,
                     )
                 )
-                algo = direct_naive
+                algo = direct
         elif algorithm == "tau_leaping":
             if "tau" in kwargs.keys():
                 tau = kwargs["tau"]

pyssa/tau_adaptive.py

@@ -6,7 +6,7 @@
 from numba import njit, jit
 import numpy as np
 from .utils import get_kstoc, roulette_selection, HIGH, TINY
-from .direct_naive import direct_naive
+from .direct import direct
 
 
 @njit(nogil=True, cache=False)
@@ -68,6 +68,102 @@ def get_HOR(react_stoic: np.ndarray):
     return HOR
 
 
+@njit(nogil=True, cache=False)
+def step1(kstoc, xt, react_stoic, v, nc):
+    """ Determine critical reactions """
+    nr = react_stoic.shape[1]
+    ns = react_stoic.shape[0]
+    prop = np.copy(kstoc)
+    L = np.zeros(nr, dtype=np.int64)
+    # Calculate the propensities
+    for ind1 in range(nr):
+        for ind2 in range(ns):
+            # prop = kstoc * product of (number raised to order)
+            prop[ind1] *= np.power(xt[ind2], react_stoic[ind2, ind1])
+    for ind in range(nr):
+        vis = v[:, ind]
+        L[ind] = np.nanmin(xt[vis < 0] / np.abs(vis[vis < 0]))
+    # A reaction j is critical if Lj <nc. However criticality is
+    # considered only for reactions with propensity greater than
+    # 0 (`prop > 0`).
+    crit = (L < nc) * (prop > 0)
+    # To get the non-critical reactions, we use the bitwise not operator.
+    not_crit = ~crit
+    return prop, crit, not_crit
+
+
+@njit(nogil=True, cache=False)
+def step2(not_crit, react_species, v, xt, HOR, prop, epsilon):
+    """ 2. Generate candidate taup """
+    n_react_species = react_species.shape[0]
+    mup = np.zeros(n_react_species, dtype=np.float64)
+    sigp = np.zeros(n_react_species, dtype=np.float64)
+    tau_num = np.zeros(n_react_species, dtype=np.float64)
+    if np.sum(not_crit) == 0:
+        taup = HIGH
+    else:
+        # Compute mu from eqn 32a and sig from eqn 32b
+        for ind, species_index in enumerate(react_species):
+            this_v = v[species_index, :]
+            for i in range(len(not_crit)):
+                if not_crit[i]:
+                    mup[ind] += this_v[i] * prop[i]
+                    sigp[ind] += this_v[i] * prop[i] * this_v[i]
+            if mup[ind] == 0:
+                mup[ind] = TINY
+            if sigp[ind] == 0:
+                sigp[ind] = TINY
+            if HOR[species_index] > 0:
+                g = HOR[species_index]
+            elif HOR[species_index] == -2:
+                if xt[species_index] is not 1:
+                    g = 1 + 2 / (xt[species_index] - 1)
+                else:
+                    g = 2
+            elif HOR[species_index] == -3:
+                if xt[species_index] not in [1, 2]:
+                    g = 3 + 1 / (xt[species_index] - 1) + 2 / (xt[species_index] - 2)
+                else:
+                    g = 3
+            elif HOR[species_index] == -32:
+                if xt[species_index] is not 1:
+                    g = 3 / 2 * (2 + 1 / (xt[species_index] - 1))
+                else:
+                    g = 3
+            tau_num[ind] = max(epsilon * xt[species_index] / g, 1)
+        taup = np.nanmin(
+            np.concatenate((tau_num / np.abs(mup), np.power(tau_num, 2) / np.abs(sigp)))
+        )
+    return taup
+
+
+@njit(nogil=True, cache=False)
+def step5(taup, taupp, nr, not_crit, prop, xt):
+    K = np.zeros(nr, dtype=np.int64)
+    if taup < taupp:
+        tau = taup
+        for ind in range(nr):
+            if not_crit[ind]:
+                K[ind] = np.random.poisson(prop[ind] * tau)
+            else:
+                K[ind] = 0
+    else:
+        tau = taupp
+        # Identify the only critical reaction to fire
+        # Send in xt to match signature of roulette_selection
+        temp = prop.copy()
+        temp[not_crit] = 0
+        j_crit, _ = roulette_selection(temp, xt)
+        for ind in range(nr):
+            if not_crit[ind]:
+                K[ind] = np.random.poisson(prop[ind] * tau)
+            elif ind == j_crit:
+                K[ind] = 1
+            else:
+                K[ind] = 0
+    return tau, K
+
+
 @njit(nogil=True, cache=False)
 def tau_adaptive(
     react_stoic: np.ndarray,
@@ -159,107 +255,47 @@ def tau_adaptive(
 
     M = nr
     N = ns
-    L = np.zeros(M, dtype=np.int64)
-    K = np.zeros(M, dtype=np.int64)
     vis = np.zeros(M, dtype=np.int64)
     react_species = np.where(np.sum(react_stoic, axis=1) > 0)[0]
-    n_react_species = react_species.shape[0]
-    mup = np.zeros(n_react_species, dtype=np.float64)
-    sigp = np.zeros(n_react_species, dtype=np.float64)
-    tau_num = np.zeros(n_react_species, dtype=np.float64)
     skip_flag = False
 
     while ite < max_iter:
 
         # If negatives are not detected in step 6, perform steps 1 and 2.
         # Else skip to step 3.
         if not skip_flag:
-
             xt = x[ite - 1, :]
-            # 1. Determine critical reactions
-            # -------------------------------
-            # Calculate the propensities
-            prop = np.copy(kstoc)
-            for ind1 in range(nr):
-                for ind2 in range(ns):
-                    # prop = kstoc * product of (number raised to order)
-                    prop[ind1] *= np.power(x[ite - 1, ind2], react_stoic[ind2, ind1])
+
+            # Step 1:
+            prop, crit, not_crit = step1(kstoc, xt, react_stoic, v, nc)
             prop_sum = np.sum(prop)
-            if prop_sum < 1e-30:
+            if prop_sum < TINY:
                 status = 3
                 return t[:ite], x[:ite, :], status
-            for ind in range(M):
-                vis = v[:, ind]
-                L[ind] = np.nanmin(xt[vis < 0] / np.abs(vis[vis < 0]))
-            # A reaction j is critical if Lj <nc. However criticality is
-            # considered only for reactions with propensity greater than
-            # 0 (`prop > 0`).
-            crit = (L < nc) * (prop > 0)
-            # To get the non-critical reactions, we use the bitwise not operator.
-            not_crit = ~crit
 
-            # 2. Generate candidate taup
-            # --------------------------
-            if np.sum(not_crit) == 0:
-                taup = HIGH
-            else:
-                # Compute mu from eqn 32a and sig from eqn 32b
-                for ind, species_index in enumerate(react_species):
-                    this_v = v[species_index, :]
-                    for i in range(len(not_crit)):
-                        if not_crit[i]:
-                            mup[ind] += this_v[i] * prop[i]
-                            sigp[ind] += this_v[i] * prop[i] * this_v[i]
-                    if mup[ind] == 0:
-                        mup[ind] = TINY
-                    if sigp[ind] == 0:
-                        sigp[ind] = TINY
-                    if HOR[species_index] > 0:
-                        g = HOR[species_index]
-                    elif HOR[species_index] == -2:
-                        if xt[species_index] is not 1:
-                            g = 1 + 2 / (xt[species_index] - 1)
-                        else:
-                            g = 2
-                    elif HOR[species_index] == -3:
-                        if xt[species_index] not in [1, 2]:
-                            g = (
-                                3
-                                + 1 / (xt[species_index] - 1)
-                                + 2 / (xt[species_index] - 2)
-                            )
-                        else:
-                            g = 3
-                    elif HOR[species_index] == -32:
-                        if xt[species_index] is not 1:
-                            g = 3 / 2 * (2 + 1 / (xt[species_index] - 1))
-                        else:
-                            g = 3
-                    tau_num[ind] = max(epsilon * xt[species_index] / g, 1)
-                taup = np.nanmin(
-                    np.concatenate(
-                        (tau_num / np.abs(mup), np.power(tau_num, 2) / np.abs(sigp))
-                    )
-                )
+            # Step 2:
+            taup = step2(not_crit, react_species, v, xt, HOR, prop, epsilon)
 
         # 3. For small taup, do SSA
         # -------------------------
         skip_flag = False
         if taup < 10 / prop_sum:
-            t_ssa, x_ssa, status = direct_naive(
+            t_ssa, x_ssa, status = direct(
                 react_stoic,
                 prod_stoic,
                 x[ite - 1, :],
                 k_det,
                 max_t=max_t - t[ite - 1],
-                max_iter=min(100, max_iter - ite),
+                max_iter=min(101, max_iter - ite),
                 volume=volume,
                 seed=seed,
                 chem_flag=chem_flag,
             )
-            len_simulation = len(t_ssa)
-            t[ite : ite + len_simulation] = t_ssa + t[ite-1]
-            x[ite : ite + len_simulation, :] = x_ssa
+            # t_ssa first element is 0. x_ssa first element is x[ite - 1, :].
+            # Both should be dropped while logging the results.
+            len_simulation = len(t_ssa) - 1  # Since t_ssa[0] is 0
+            t[ite : ite + len_simulation] = t_ssa[1:] + t[ite-1]
+            x[ite : ite + len_simulation, :] = x_ssa[1:]
             ite += len_simulation
             if status == 3 or status == 2:
                 return t, x, status
@@ -274,27 +310,7 @@ def tau_adaptive(
 
         # 5. Leap
         # -------
-        if taup < taupp:
-            tau = taup
-            for ind in range(M):
-                if not_crit[ind]:
-                    K[ind] = np.random.poisson(prop[ind] * tau)
-                else:
-                    K[ind] = 0
-        else:
-            tau = taupp
-            # Identify the only critical reaction to fire
-            # Send in xt to match signature of roulette_selection
-            temp = prop.copy()
-            temp[not_crit] = 0
-            j_crit, _ = roulette_selection(temp, x[ite - 1, :])
-            for ind in range(M):
-                if not_crit[ind]:
-                    K[ind] = np.random.poisson(prop[ind] * tau)
-                elif ind == j_crit:
-                    K[ind] = 1
-                else:
-                    K[ind] = 0
+        tau, K = step5(taup, taupp, nr, not_crit, prop, xt)
 
         # 6. Handle negatives, update and exit conditions
         # -----------------------------------------------
@@ -308,11 +324,11 @@ def tau_adaptive(
         if np.any(x_new < 0):
             taup = taup / 2
             skip_flag = True
-
-        # Update states if nothing is negative
-        x[ite, :] = x[ite - 1, :] + vdotK
-        t[ite] = t[ite - 1] + tau
-        ite += 1
+        else:
+            # Update states if nothing is negative
+            x[ite, :] = x[ite - 1, :] + vdotK
+            t[ite] = t[ite - 1] + tau
+            ite += 1
 
         # Exit conditions
         if t[ite - 1] > max_t:

pyssa/utils.py

@@ -83,7 +83,7 @@ def roulette_selection(prop_list: np.ndarray, Xt: np.ndarray) -> Tuple[int, int]
     choice : int
         Index of the chosen reaction.
     status : int
-        Status of the simulation as described in `direct_naive`.
+        Status of the simulation as described in `direct`.
     """
     prop0 = np.sum(prop_list)  # Sum of propensities
     # choice = 0