dynamics: add vertex-specific parameters for epidemic models

src/graph/dynamics/graph_discrete.hh

@@ -64,6 +64,15 @@ public:
     std::shared_ptr<std::vector<size_t>> _active;
 };
 
+template <class Map>
+Map get_pmap(python::object o)
+{
+    o = o.attr("_get_any")();
+    boost::any& a = python::extract<boost::any&>(o);
+    Map m = boost::any_cast<typename Map::checked_t>(a).get_unchecked();
+    return m;
+}
+
 template <bool exposed, bool weighted, bool constant_beta>
 class SI_state: public discrete_state_base<>
 {
@@ -73,12 +82,13 @@ public:
 
     typedef typename eprop_map_t<double>::type::unchecked_t bmap_t;
     typedef std::conditional_t<weighted, bmap_t, double> beta_t;
+    typedef typename vprop_map_t<double>::type::unchecked_t rmap_t;
 
     template <class Graph, class RNG>
     SI_state(Graph& g, smap_t s, smap_t s_temp, python::dict params, RNG&)
         : discrete_state_base(s, s_temp),
-          _epsilon(python::extract<double>(params["epsilon"])),
-          _r(python::extract<double>(params["r"])),
+          _epsilon(get_pmap<rmap_t>(params["epsilon"])),
+          _r(get_pmap<rmap_t>(params["r"])),
           _m(num_vertices(g)),
           _m_temp(num_vertices(g))
     {
@@ -202,8 +212,9 @@ public:
 
         if (exposed && _s[v] == State::E)
         {
-            std::bernoulli_distribution einfect(_epsilon);
-            if (_epsilon > 0 && einfect(rng))
+            auto epsilon = _epsilon[v];
+            std::bernoulli_distribution einfect(epsilon);
+            if (epsilon > 0 && einfect(rng))
             {
                 infect<sync>(g, v, s_out);
                 return 1;
@@ -211,8 +222,9 @@ public:
             return 0;
         }
 
-        std::bernoulli_distribution spontaneous(_r);
-        if (_r > 0 && spontaneous(rng))
+        auto r = _r[v];
+        std::bernoulli_distribution spontaneous(r);
+        if (r > 0 && spontaneous(rng))
         {
             if constexpr (exposed)
                 expose(g, v, s_out);
@@ -255,8 +267,8 @@ public:
 
 protected:
     beta_t _beta;
-    double _epsilon;
-    double _r;
+    rmap_t _epsilon;
+    rmap_t _r;
 
     typedef std::conditional_t<weighted,
                                typename vprop_map_t<double>::type::unchecked_t,
@@ -272,6 +284,7 @@ public:
 
     typedef SI_state<exposed, weighted, constant_beta> base_t;
     typedef typename base_t::smap_t smap_t;
+    typedef typename base_t::rmap_t rmap_t;
     typedef typename base_t::State State;
     using base_t::_s;
     using base_t::_m;
@@ -280,7 +293,7 @@ public:
     template <class Graph, class RNG>
     SIS_state(Graph& g, smap_t s, smap_t s_temp, python::dict params, RNG& rng)
         : base_t(g, s, s_temp, params, rng),
-          _gamma(python::extract<double>(params["gamma"]))
+          _gamma(get_pmap<rmap_t>(params["gamma"]))
     {};
 
     template <bool sync, class Graph>
@@ -333,8 +346,9 @@ public:
     {
         if (_s[v] == State::I)
         {
-            std::bernoulli_distribution srecover(_gamma);
-            if (_gamma > 0 && srecover(rng))
+            auto gamma = _gamma[v];
+            std::bernoulli_distribution srecover(gamma);
+            if (gamma > 0 && srecover(rng))
             {
                 recover<sync>(g, v, s_out);
                 return 1;
@@ -351,7 +365,7 @@ public:
 
 
 protected:
-    double _gamma;
+    rmap_t _gamma;
 };
 
 template <bool exposed, bool weighted, bool constant_beta>
@@ -360,13 +374,14 @@ class SIRS_state: public SIS_state<exposed, true, weighted, constant_beta>
 public:
     typedef SIS_state<exposed, true, weighted, constant_beta> base_t;
     typedef typename base_t::smap_t smap_t;
+    typedef typename base_t::rmap_t rmap_t;
     typedef typename base_t::State State;
     using base_t::_s;
 
     template <class Graph, class RNG>
     SIRS_state(Graph& g, smap_t s, smap_t s_temp, python::dict params, RNG& rng)
         : base_t(g, s, s_temp, params, rng),
-          _mu(python::extract<double>(params["mu"]))
+          _mu(get_pmap<rmap_t>(params["mu"]))
     {};
 
 
@@ -375,8 +390,9 @@ public:
     {
         if (_s[v] == State::R)
         {
-            std::bernoulli_distribution srecover(_mu);
-            if (_mu > 0 && srecover(rng))
+            auto mu = _mu[v];
+            std::bernoulli_distribution srecover(mu);
+            if (mu > 0 && srecover(rng))
             {
                 s_out[v] = State::S;
                 return 1;
@@ -392,7 +408,7 @@ public:
     constexpr bool has_absorbing() { return false; }
 
 private:
-    double _mu;
+    rmap_t _mu;
 };
 
 

src/graph_tool/dynamics/__init__.py

@@ -70,7 +70,7 @@ Contents
 from __future__ import division, absolute_import, print_function
 
 from .. import _degree, _prop, Graph, GraphView, _get_rng, PropertyMap, \
-    EdgePropertyMap, _check_prop_scalar
+    EdgePropertyMap, VertexPropertyMap, _check_prop_scalar
 from .. stats import label_self_loops
 import numpy
 import numpy.random
@@ -146,7 +146,7 @@ class DiscreteStateBase(object):
         return self._state.iterate_async(niter, _get_rng())
 
 class EpidemicStateBase(DiscreteStateBase):
-    def __init__(self, g, beta, constant_beta, make_state, v0=None, s=None):
+    def __init__(self, g, constant_beta, make_state, params, v0=None, s=None):
         r"""Base state for epidemic dynamics. This class it not meant to be
         instantiated directly."""
 
@@ -159,6 +159,7 @@ class EpidemicStateBase(DiscreteStateBase):
                 v0 = g.vertex(v0, use_index=False)
             self.s[v0] = 1
 
+        beta = params["beta"]
         weighted = isinstance(beta, EdgePropertyMap)
         if weighted:
             _check_prop_scalar(beta, "beta")
@@ -166,7 +167,18 @@ class EpidemicStateBase(DiscreteStateBase):
                 if constant_beta:
                     raise ValueError("if constant_beta == True, the type of beta must be double")
                 beta = beta.copy("double")
-        self.beta = beta
+        params["beta"] = beta
+
+        for p in ["r", "epsilon", "gamma", "mu"]:
+            if p not in params:
+                continue
+            if not isinstance(params[p], VertexPropertyMap):
+                params[p] = g.new_vp("double", val=params[p])
+            _check_prop_scalar(params[p], p)
+            if params[p].value_type() != "double":
+                params[p] = params[p].copy("double")
+        self.params = params
+
         self.make_state = lambda *args: make_state(*args, weighted, constant_beta)
 
 
@@ -183,12 +195,12 @@ class SIState(EpidemicStateBase):
            Transmission probability. If an :class:`~graph_tool.EdgePropertyMap`
            object is passed, it must contain the transmission probability for
            every edge.
-        r : ``float`` (optional, default: ``0.``)
+        r : ``float`` or :class:`~graph_tool.VertexPropertyMap` (optional, default: ``0.``)
            Spontaneous infection probability.
         exposed : ``boolean`` (optional, default: ``False``)
            If ``True``, an SEI model is simulated, with an additional "exposed"
            state.
-        epsilon : ``float`` (optional, default: ``.1``)
+        epsilon : ``float`` or :class:`~graph_tool.VertexPropertyMap` (optional, default: ``.1``)
            Susceptible to exposed transition probability. This only has an
            effect if ``exposed=True``.
         v0 : ``int`` or :class:`~graph_tool.Vertex` (optional, default: ``None``)
@@ -217,7 +229,7 @@ class SIState(EpidemicStateBase):
         1. If :math:`s_i(t) = 0`, we have :math:`s_i(t+1) = 1` with probability
             .. math::
 
-               (1-r)\left(1-\prod_j(1-\beta_{ij})^{A_{ij}\delta_{s_j(t),1}}\right) + r,
+               (1-r_i)\left(1-\prod_j(1-\beta_{ij})^{A_{ij}\delta_{s_j(t),1}}\right) + r_i,
 
            otherwise :math:`s_i(t+1) = 0`.
 
@@ -226,7 +238,7 @@ class SIState(EpidemicStateBase):
 
         If the option ``exposed == True`` is given, then the states transit
         first from 0 to -1 (exposed) with probability given by 1. above, and
-        then finally from -1 to 1 with probability :math:`\epsilon`.
+        then finally from -1 to 1 with probability :math:`\epsilon_i`.
 
         Examples
         --------
@@ -270,11 +282,11 @@ class SIState(EpidemicStateBase):
            :doi:`10.1103/RevModPhys.87.925`, :arxiv:`1408.2701`
 
         """
-        EpidemicStateBase.__init__(self, g, beta, constant_beta,
+        EpidemicStateBase.__init__(self, g, constant_beta,
                                    lib_dynamics.make_SEI_state if exposed else lib_dynamics.make_SI_state,
+                                   dict(beta=beta, r=r, epsilon=epsilon),
                                    v0, s)
-        DiscreteStateBase.__init__(self, g, self.make_state,
-                                   dict(beta=self.beta, r=r, epsilon=epsilon),
+        DiscreteStateBase.__init__(self, g, self.make_state, self.params,
                                    self.s)
 
 class SISState(DiscreteStateBase):
@@ -325,17 +337,17 @@ class SISState(DiscreteStateBase):
         1. If :math:`s_i(t) = 0`, we have :math:`s_i(t+1) = 1` with probability
             .. math::
 
-               (1-r)\left(1-\prod_j(1-\beta_{ij})^{A_{ij}\delta_{s_j(t),1}}\right) + r,
+               (1-r_i)\left(1-\prod_j(1-\beta_{ij})^{A_{ij}\delta_{s_j(t),1}}\right) + r_i,
 
            otherwise :math:`s_i(t+1) = 0`.
 
         2. If :math:`s_i(t) = 1`, we have :math:`s_i(t+1) = 0` with probability
-           :math:`\gamma`, or :math:`s_i(t+1) = 1` with probability
-           :math:`1-\gamma`.
+           :math:`\gamma_i`, or :math:`s_i(t+1) = 1` with probability
+           :math:`1-\gamma_i`.
 
         If the option ``exposed == True`` is given, then the states transit
         first from 0 to -1 (exposed) with probability given by 1. above, and
-        then finally from -1 to 1 with probability :math:`\epsilon`.
+        then finally from -1 to 1 with probability :math:`\epsilon_i`.
 
         Examples
         --------
@@ -379,12 +391,12 @@ class SISState(DiscreteStateBase):
            :doi:`10.1103/RevModPhys.87.925`, :arxiv:`1408.2701`
 
         """
-        EpidemicStateBase.__init__(self, g, beta, constant_beta,
+        EpidemicStateBase.__init__(self, g, constant_beta,
                                    lib_dynamics.make_SEIS_state if exposed else lib_dynamics.make_SIS_state,
+                                   dict(beta=beta, gamma=gamma, r=r,
+                                        epsilon=epsilon),
                                    v0, s)
-        DiscreteStateBase.__init__(self, g, self.make_state,
-                                   dict(beta=self.beta, gamma=gamma, r=r,
-                                        epsilon=epsilon), self.s)
+        DiscreteStateBase.__init__(self, g, self.make_state, self.params, self.s)
 
 class SIRState(DiscreteStateBase):
     def __init__(self, g, beta=1., gamma=.1, r=0, exposed=False, epsilon=.1,
@@ -434,17 +446,17 @@ class SIRState(DiscreteStateBase):
         1. If :math:`s_i(t) = 0`, we have :math:`s_i(t+1) = 1` with probability
             .. math::
 
-               (1-r)\left(1-\prod_j(1-\beta_{ij})^{A_{ij}\delta_{s_j(t),1}}\right) + r,
+               (1-r_i)\left(1-\prod_j(1-\beta_{ij})^{A_{ij}\delta_{s_j(t),1}}\right) + r_i,
 
            otherwise :math:`s_i(t+1) = 0`.
 
         2. If :math:`s_i(t) = 1`, we have :math:`s_i(t+1) = 2` with probability
-           :math:`\gamma`, or :math:`s_i(t+1) = 1` with probability
-           :math:`1-\gamma`.
+           :math:`\gamma_i`, or :math:`s_i(t+1) = 1` with probability
+           :math:`1-\gamma_i`.
 
         If the option ``exposed == True`` is given, then the states transit
         first from 0 to -1 (exposed) with probability given by 1. above, and
-        then finally from -1 to 1 with probability :math:`\epsilon`.
+        then finally from -1 to 1 with probability :math:`\epsilon_i`.
 
         Examples
         --------
@@ -498,12 +510,13 @@ class SIRState(DiscreteStateBase):
            :doi:`10.1103/RevModPhys.87.925`, :arxiv:`1408.2701`
 
         """
-        EpidemicStateBase.__init__(self, g, beta, constant_beta,
+        EpidemicStateBase.__init__(self, g, constant_beta,
                                    lib_dynamics.make_SEIR_state if exposed else lib_dynamics.make_SIR_state,
+                                   dict(beta=beta, gamma=gamma, r=r,
+                                        epsilon=epsilon),
                                    v0, s)
         DiscreteStateBase.__init__(self, g, self.make_state,
-                                   dict(beta=self.beta, gamma=gamma, r=r,
-                                        epsilon=epsilon), self.s)
+                                   self.params, self.s)
 
 class SIRSState(DiscreteStateBase):
     def __init__(self, g, beta=1, gamma=.1, mu=.1, r=0, exposed=False,
@@ -556,21 +569,21 @@ class SIRSState(DiscreteStateBase):
         1. If :math:`s_i(t) = 0`, we have :math:`s_i(t+1) = 1` with probability
             .. math::
 
-               (1-r)\left(1-\prod_j(1-\beta_{ij})^{A_{ij}\delta_{s_j(t),1}}\right) + r,
+               (1-r_i)\left(1-\prod_j(1-\beta_{ij})^{A_{ij}\delta_{s_j(t),1}}\right) + r_i,
 
            otherwise :math:`s_i(t+1) = 0`.
 
         2. If :math:`s_i(t) = 1`, we have :math:`s_i(t+1) = 2` with probability
-           :math:`\gamma`, or :math:`s_i(t+1) = 1` with probability
-           :math:`1-\gamma`.
+           :math:`\gamma_i`, or :math:`s_i(t+1) = 1` with probability
+           :math:`1-\gamma_i`.
 
         3. If :math:`s_i(t) = 2`, we have :math:`s_i(t+1) = 1` with probability
-           :math:`\mu`, or :math:`s_i(t+1) = 2` with probability
-           :math:`1-\mu`.
+           :math:`\mu_i`, or :math:`s_i(t+1) = 2` with probability
+           :math:`1-\mu_i`.
 
         If the option ``exposed == True`` is given, then the states transit
         first from 0 to -1 (exposed) with probability given by 1. above, and
-        then finally from -1 to 1 with probability :math:`\epsilon`.
+        then finally from -1 to 1 with probability :math:`\epsilon_i`.
 
         Examples
         --------
@@ -624,12 +637,13 @@ class SIRSState(DiscreteStateBase):
            :doi:`10.1103/RevModPhys.87.925`, :arxiv:`1408.2701`
 
         """
-        EpidemicStateBase.__init__(self, g, beta, constant_beta,
+        EpidemicStateBase.__init__(self, g, constant_beta,
                                    lib_dynamics.make_SEIRS_state if exposed else lib_dynamics.make_SIRS_state,
+                                   dict(beta=beta, gamma=gamma, mu=mu, r=r,
+                                        epsilon=epsilon),
                                    v0, s)
         DiscreteStateBase.__init__(self, g, self.make_state,
-                                   dict(beta=self.beta, gamma=gamma, mu=mu, r=r,
-                                        epsilon=epsilon), self.s)
+                                   self.params, self.s)
 
 
 class VoterState(DiscreteStateBase):