Merge pull request #5 from voduchuy/update_api

Update api

.gitignore

@@ -0,0 +1,8 @@
+.idea
+build
+dist
+*.pdf
+*.png
+*.npz
+*.cpp
+*.h

examples/gene_expression_pmpdmsr.py

@@ -1,5 +1,5 @@
 from pypacmensl.fsp_solver import FspSolverMultiSinks
-from pypacmensl.ssa.ssa import PmPdmsrSampler
+from pypacmensl.ssa.ssa import PmPdmsrSampler, SSASolver
 import mpi4py.MPI as mpi
 import numpy as np
 import matplotlib.pyplot as plt
@@ -51,11 +51,11 @@ def gene_transition_propensity(reaction, x, out):
 solver.SetVerbosity(2)
 solver.SetOdeSolver("KRYLOV")
 solver.SetUp()
-solutions = solver.SolveTspan(tspan, 1.0e-4)
+cme_solutions = solver.SolveTspan(tspan, 1.0e-4)
 
 marginals_fsp = []
 for j in range(0, n_t):
-    marginals_fsp.append(solutions[j].Marginal(2))
+    marginals_fsp.append(cme_solutions[j].Marginal(2))
 #%% Monte Carlo approximation with PMPDMSR
 def rna_dist_from_samples(comm: mpi.Comm, poisson_samples: np.ndarray, nmax: int) -> np.ndarray:
 
@@ -88,3 +88,4 @@ def rna_dist_from_samples(comm: mpi.Comm, poisson_samples: np.ndarray, nmax: int
         axs[j].plot(marginals_fsp[j], color="b", label="FSP")
         axs[j].legend()
     plt.show()
+

examples/gene_expression_pmpdmsr2.py

@@ -0,0 +1,162 @@
+from pypacmensl.fsp_solver import FspSolverMultiSinks
+from pypacmensl.ssa.ssa import PmPdmsrSampler, SSASolver
+from pypacmensl.smfish.snapshot import SmFishSnapshot
+import mpi4py.MPI as mpi
+import numpy as np
+import matplotlib.pyplot as plt
+from scipy.stats import poisson
+
+#%% Rate constants
+K01 = 0.01
+K10 = 0.001
+BETA1 = 1.0
+BETA2 = 0.5
+DEGRADATION = 0.1
+#%% This block of code defines the telegraph model to be solved with FSP
+stoich_mat = np.array(
+    [
+        [-1, 1, 0, 0],
+        [1, -1, 0, 0],
+        [0, 0, 1, 0],
+        [0, 0, 0, 1],
+        [0, 0, -1, 0],
+        [0, 0, 0, -1],
+    ]
+)
+x0 = np.array([[1, 0, 0, 0]])
+p0 = np.array([1.0])
+constr_init = np.array([1, 1, 100, 100])
+
+def propensity_t(t, out):
+    out[:] = 1.0
+    # out[1] = 1.0+np.cos(2*t)
+
+def propensity(reaction, x, out):
+    if reaction == 0:
+        out[:] = K01 * x[:, 0]
+    if reaction == 1:
+        out[:] = K10 * x[:, 1]
+    if reaction == 2:
+        out[:] = BETA1 * x[:, 1]
+    if reaction == 3:
+        out[:] = BETA2 * x[:, 1]
+    if reaction == 4:
+        out[:] = DEGRADATION * x[:, 2]
+    if reaction == 5:
+        out[:] = DEGRADATION * x[:, 3]
+
+
+n_t = 5
+tspan = np.linspace(0, 100, n_t)
+#%%
+def transcription_rate(gene_state):
+    return np.array([(gene_state[1] == 1) * BETA1, (gene_state[1] == 1) * BETA2])
+
+
+gene_transition_stoich = np.array([[-1, 1], [1, -1]])
+
+
+def gene_transition_propensity(reaction, x, out):
+    if reaction == 0:
+        out[:] = K01 * x[:, 0]
+    if reaction == 1:
+        out[:] = K10 * x[:, 1]
+
+
+#%% Find the marginal mRNA distributions with FSP
+solver = FspSolverMultiSinks(mpi.COMM_WORLD)
+solver.SetModel(stoich_mat, propensity_t, propensity)
+solver.SetFspShape(None, constr_init)
+solver.SetInitialDist(x0, p0)
+solver.SetVerbosity(2)
+solver.SetOdeSolver("PETSC")
+solver.SetUp()
+solutions = solver.SolveTspan(tspan, 1.0e-4)
+
+marginals_fsp = []
+for j in range(0, n_t):
+    marginals_fsp.append([solutions[j].Marginal(2), solutions[j].Marginal(3)])
+#%% Monte Carlo approximation with PMPDMSR
+
+def joint_dist_from_poisson_parameters(
+    comm: mpi.Comm, poisson_parameters: np.ndarray, nmax: int
+):
+    nsamples_local = poisson_parameters.shape[0]
+    nsamples = comm.allreduce(nsamples_local, mpi.SUM)
+    nmax = comm.allreduce(nmax, mpi.MAX)
+
+    p_local = np.zeros((nmax + 1, nmax + 1), dtype=float)
+
+    x_eval = np.arange(0, nmax + 1, dtype=int)
+    p_marginals = np.zeros((2, nmax + 1), dtype=float)
+
+    for j in range(0, nsamples_local):
+        for ispecies in range(0, 2):
+            p_marginals[ispecies, :] = poisson.pmf(
+                x_eval, mu=poisson_parameters[j, ispecies]
+            )
+        p_local += np.kron(p_marginals[0, :], p_marginals[1, :]).reshape(
+            (nmax + 1, nmax + 1)
+        )
+    p_global = comm.allreduce(p_local, mpi.SUM)
+    p_global = (1.0 / nsamples) * p_global
+    return p_global
+
+
+sampler = PmPdmsrSampler(mpi.COMM_WORLD)
+sampler.SetModel(
+    gene_transition_stoich,
+    propensity_t,
+    gene_transition_propensity,
+    transcription_rate,
+    np.array([DEGRADATION]*2),
+)
+poisson_samples = []
+pjoint_pmpdmsr = []
+marginals_pmpdmsr = []
+for j in range(0, n_t):
+    _, p_samples = sampler.SolveTv(
+        tspan[j], np.array([[1, 0]], dtype=int), 10000, update_rates=1.0E-7, send_to_root=False
+    )
+    poisson_samples.append(p_samples)
+    pjoint = joint_dist_from_poisson_parameters(mpi.COMM_WORLD, p_samples, 100)
+    marginals_pmpdmsr.append(
+        [
+            np.sum(pjoint, axis=1), np.sum(pjoint, axis=0)
+        ]
+    )
+    pjoint_pmpdmsr.append(pjoint)
+#
+# if mpi.COMM_WORLD.Get_rank() == 0:
+#     import matplotlib.colors as colors
+#     fig, axs = plt.subplots(2, n_t)
+#     for j in range(0, n_t):
+#         for i in range(0, 2):
+#             axs[i, j].plot(np.arange(0, 101), marginals_pmpdmsr[j][i], color="r", label="PMPDMSR")
+#             axs[i, j].plot(marginals_fsp[j][i], color="b", label="FSP", linestyle=":")
+#             axs[i, j].legend()
+#     plt.show()
+#
+#     fig, axs = plt.subplots(1, n_t)
+#     for j in range(0, n_t):
+#         axs[j].pcolorfast(pjoint_pmpdmsr[j], norm=colors.LogNorm(vmin=1.0E-5, vmax=1.0))
+#     plt.show()
+
+#%% Likelihood computation
+
+ssa = SSASolver(mpi.COMM_WORLD)
+ssa.SetModel(stoich_mat, propensity_t, propensity)
+data = []
+raw_obs = []
+for i in range(0, n_t):
+    samples = ssa.SolveTv(tspan[j], np.array([[1, 0, 0, 0]], dtype=int), 1000, update_rates=1.0E-7, send_to_root=True)
+    samples = mpi.COMM_WORLD.bcast(samples)
+    data.append(SmFishSnapshot(samples[:, 2:]))
+    raw_obs.append(samples[:,2:])
+
+for i in range(0, n_t):
+    ll_fsp = data[i].LogLikelihood(solutions[i], np.array([2,3]))
+    ll_pm = sampler.compute_loglike(poisson_samples[i], raw_obs[i])
+    if mpi.COMM_WORLD.Get_rank() == 0:
+        print([ll_fsp, ll_pm, ll_fsp - ll_pm])
+

examples/toggle_fim.py

@@ -1,66 +1,53 @@
-
-
 import pypacmensl.sensitivity.multi_sinks as sensfsp
 import mpi4py.MPI as mpi
 import numpy as np
 import matplotlib.pyplot as plt
 
 # %%
-stoich_matrix = np.array(
-    [
-        [1, 0],
-        [-1, 0],
-        [0, 1],
-        [0, -1]
-    ]
-)
-
-def tcoeff(t, out):
-    out[0] = 35.0
-    out[1] = 0.20
-    out[2] = 50.0
-    out[3] = 1.0
+stoich_matrix = np.array([[1, 0],
+                          [-1, 0],
+                          [0, 1],
+                          [0, -1]])
 
 
 def propensity(reaction, states, outs):
-    if reaction is 0:
-        outs[:] = np.reciprocal(1 + states[:, 1] ** 1.5)
+    if reaction == 0:
+        outs[:] = 35.0*np.reciprocal(1 + states[:, 1] ** 1.5)
         return
-    if reaction is 1:
-        outs[:] = states[:, 0]
+    if reaction == 1:
+        outs[:] = 0.20*states[:, 0]
         return
-    if reaction is 2:
-        outs[:] = np.reciprocal(1 + states[:, 0] ** 2.5)
+    if reaction == 2:
+        outs[:] = 50.0*np.reciprocal(1 + states[:, 0] ** 2.5)
         return
-    if reaction is 3:
-        outs[:] = states[:, 1]
+    if reaction == 3:
+        outs[:] = 1.0*states[:, 1]
+
+def dpropensity(parameter, reaction, states, outs):
+    if parameter == 0:
+        if reaction == 0:
+            outs[:] = np.reciprocal(1 + states[:, 1] ** 1.5)
+            return
+    if parameter == 1:
+        if reaction == 1:
+            outs[:] = states[:, 0]
+            return
+    if parameter == 2:
+        if reaction == 2:
+            outs[:] = np.reciprocal(1 + states[:, 0] ** 2.5)
+            return
+    if parameter == 3:
+        if reaction == 3:
+            outs[:] = states[:, 1]
 
 
 def simple_constr(X, out):
     out[:, 0] = X[:, 0]
     out[:, 1] = X[:, 1]
     out[:, 2] = X[:, 0] + X[:, 1]
 
-
 init_bounds = np.array([10, 10, 10])
 
-
-def d_tcoeff_factory(i):
-    def d_tcoeff(t, out):
-        out[i] = 1.0
-
-    return d_tcoeff
-
-
-dtcoeff = []
-for i in range(0, 4):
-    dtcoeff.append(d_tcoeff_factory(i))
-dpropensity = [propensity] * 4
-
-dprop_sparsity = np.zeros((4, 4), dtype=np.intc)
-for i in range(0, 4):
-    dprop_sparsity[i][i] = 1
-
 X0 = np.array([[0, 0]])
 p0 = np.array([1.0])
 s0 = np.array([0.0])
@@ -74,8 +61,17 @@ def d_tcoeff(t, out):
 my_rank = comm.rank
 # %%
 my_solver = sensfsp.SensFspSolverMultiSinks(comm)
-my_solver.SetModel(stoich_matrix, tcoeff, propensity,
-                dtcoeff, dpropensity, dprop_sparsity)
+my_solver.SetModel(
+    4,
+    stoich_matrix=stoich_matrix,
+    propensity_t=None,
+    propensity_x=propensity,
+    tv_reactions=[],
+    d_propensity_t=None,
+    d_propensity_t_sp=None,
+    d_propensity_x=dpropensity,
+    d_propensity_x_sp=[[i] for i in range(4)]
+)
 my_solver.SetFspShape(constr_fun=simple_constr, constr_bound=init_bounds)
 my_solver.SetInitialDist(X0, p0, [s0] * 4)
 my_solver.SetVerbosity(2)
@@ -126,12 +122,13 @@ def d_tcoeff(t, out):
 if comm.Get_rank() == 0:
     fig = plt.figure()
     ax = fig.add_subplot(1, 1, 1)
-    plt.rc('text', usetex=True)
+    plt.rc("text", usetex=True)
     plt.plot(t_meas, np.log10(DetFIMs), label="Observing both species")
     plt.plot(t_meas, np.log10(DetFIMs0), label="Observing species 0")
     plt.plot(t_meas, np.log10(DetFIMs1), label="Observing species 1")
-    plt.set_title("Log10-determinant of the Fisher Information Matrix for different combinations of observables")
-    plt.xlabel('Time (sec)')
-    plt.ylabel(r'$\log_{10}(|FIM|)$')
+    plt.title(
+        "Log10-determinant of the Fisher Information Matrix for different combinations of observables"
+    )
+    plt.xlabel("Time (sec)")
+    plt.ylabel(r"$\log_{10}(|FIM|)$")
     plt.show()
-

examples/toggle_state_space.py

@@ -81,9 +81,10 @@ def simple_constr(X, out):
 f2, ax2 = plot_state_set(comm, my_set_hypergraph)
 ax2.set_title('Hypergraph-partitioned FSP')
 
+if (my_rank == 0):
+    plt.show()
 #%% Uncomment if you want to save the plots to files
-# if (my_rank == 0):
-#     f0.savefig('naive_fsp.eps', format='eps')
-#     f1.savefig('graph_fsp.eps', format='eps')
-#     f2.savefig('hypergraph_fsp.eps', format='eps')
-#     plt.show()
+if (my_rank == 0):
+    f0.savefig('naive_fsp.eps', format='eps')
+    f1.savefig('graph_fsp.eps', format='eps')
+    f2.savefig('hypergraph_fsp.eps', format='eps')

setup.py

@@ -1,8 +1,6 @@
 import os
-from sys import platform
-from distutils.core import setup, Extension
+from setuptools import setup
 from Cython.Build import cythonize
-from glob import glob
 import numpy as np
 from os.path import join
 
@@ -55,7 +53,6 @@
                      'pypacmensl.fsp_solver':  'src/pypacmensl/fsp_solver',
                      'pypacmensl.state_set':   'src/pypacmensl/state_set',
                      'pypacmensl.sensitivity': 'src/pypacmensl/sensitivity',
-                     # 'pypacmensl.stationary':  'src/pypacmensl/stationary'
                      },
         ext_modules= pypacmensl_ext,
         **metadata

src/pypacmensl/callbacks/pacmensl_callbacks.pxd

@@ -1,9 +1,13 @@
 # distutils: language = c++
 from cpython.ref cimport PyObject
 
-cdef public int call_py_prop_obj(const int reaction, const int num_species, const int num_states, const int* states, double* outputs, void* args)
+cdef public int call_py_propx_obj(const int reaction_idx, const int num_species, const int num_states, const int *states, double *outputs, void * args)
 
-cdef public int call_py_t_fun_obj (double t, int num_coefs, double* outputs, void* args)
+cdef public int call_py_propt_obj(double t, int num_coefs, double* outputs, void* args)
+
+cdef public int call_py_dpropx_obj(const int parameter_idx, const int reaction_idx, const int num_species, const int num_states, const int* states, double* outputs, void* args)
+
+cdef public int call_py_dpropt_obj(const int parameter_idx, double t, int num_coefs, double* outputs, void* args)
 
 cdef public int call_py_constr_obj(int num_species, int num_constr, int n_states, int *states, int *outputs, void *args)