Commit ef39caab authored by Tiago Peixoto's avatar Tiago Peixoto
Browse files

dynamics: add vertex-specific parameters for epidemic models

parent a1b97fbe
...@@ -64,6 +64,15 @@ public: ...@@ -64,6 +64,15 @@ public:
std::shared_ptr<std::vector<size_t>> _active; 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> template <bool exposed, bool weighted, bool constant_beta>
class SI_state: public discrete_state_base<> class SI_state: public discrete_state_base<>
{ {
...@@ -73,12 +82,13 @@ public: ...@@ -73,12 +82,13 @@ public:
typedef typename eprop_map_t<double>::type::unchecked_t bmap_t; typedef typename eprop_map_t<double>::type::unchecked_t bmap_t;
typedef std::conditional_t<weighted, bmap_t, double> beta_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> template <class Graph, class RNG>
SI_state(Graph& g, smap_t s, smap_t s_temp, python::dict params, RNG&) SI_state(Graph& g, smap_t s, smap_t s_temp, python::dict params, RNG&)
: discrete_state_base(s, s_temp), : discrete_state_base(s, s_temp),
_epsilon(python::extract<double>(params["epsilon"])), _epsilon(get_pmap<rmap_t>(params["epsilon"])),
_r(python::extract<double>(params["r"])), _r(get_pmap<rmap_t>(params["r"])),
_m(num_vertices(g)), _m(num_vertices(g)),
_m_temp(num_vertices(g)) _m_temp(num_vertices(g))
{ {
...@@ -202,8 +212,9 @@ public: ...@@ -202,8 +212,9 @@ public:
if (exposed && _s[v] == State::E) if (exposed && _s[v] == State::E)
{ {
std::bernoulli_distribution einfect(_epsilon); auto epsilon = _epsilon[v];
if (_epsilon > 0 && einfect(rng)) std::bernoulli_distribution einfect(epsilon);
if (epsilon > 0 && einfect(rng))
{ {
infect<sync>(g, v, s_out); infect<sync>(g, v, s_out);
return 1; return 1;
...@@ -211,8 +222,9 @@ public: ...@@ -211,8 +222,9 @@ public:
return 0; return 0;
} }
std::bernoulli_distribution spontaneous(_r); auto r = _r[v];
if (_r > 0 && spontaneous(rng)) std::bernoulli_distribution spontaneous(r);
if (r > 0 && spontaneous(rng))
{ {
if constexpr (exposed) if constexpr (exposed)
expose(g, v, s_out); expose(g, v, s_out);
...@@ -255,8 +267,8 @@ public: ...@@ -255,8 +267,8 @@ public:
protected: protected:
beta_t _beta; beta_t _beta;
double _epsilon; rmap_t _epsilon;
double _r; rmap_t _r;
typedef std::conditional_t<weighted, typedef std::conditional_t<weighted,
typename vprop_map_t<double>::type::unchecked_t, typename vprop_map_t<double>::type::unchecked_t,
...@@ -272,6 +284,7 @@ public: ...@@ -272,6 +284,7 @@ public:
typedef SI_state<exposed, weighted, constant_beta> base_t; typedef SI_state<exposed, weighted, constant_beta> base_t;
typedef typename base_t::smap_t smap_t; typedef typename base_t::smap_t smap_t;
typedef typename base_t::rmap_t rmap_t;
typedef typename base_t::State State; typedef typename base_t::State State;
using base_t::_s; using base_t::_s;
using base_t::_m; using base_t::_m;
...@@ -280,7 +293,7 @@ public: ...@@ -280,7 +293,7 @@ public:
template <class Graph, class RNG> template <class Graph, class RNG>
SIS_state(Graph& g, smap_t s, smap_t s_temp, python::dict params, RNG& 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), : 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> template <bool sync, class Graph>
...@@ -333,8 +346,9 @@ public: ...@@ -333,8 +346,9 @@ public:
{ {
if (_s[v] == State::I) if (_s[v] == State::I)
{ {
std::bernoulli_distribution srecover(_gamma); auto gamma = _gamma[v];
if (_gamma > 0 && srecover(rng)) std::bernoulli_distribution srecover(gamma);
if (gamma > 0 && srecover(rng))
{ {
recover<sync>(g, v, s_out); recover<sync>(g, v, s_out);
return 1; return 1;
...@@ -351,7 +365,7 @@ public: ...@@ -351,7 +365,7 @@ public:
protected: protected:
double _gamma; rmap_t _gamma;
}; };
template <bool exposed, bool weighted, bool constant_beta> template <bool exposed, bool weighted, bool constant_beta>
...@@ -360,13 +374,14 @@ class SIRS_state: public SIS_state<exposed, true, weighted, constant_beta> ...@@ -360,13 +374,14 @@ class SIRS_state: public SIS_state<exposed, true, weighted, constant_beta>
public: public:
typedef SIS_state<exposed, true, weighted, constant_beta> base_t; typedef SIS_state<exposed, true, weighted, constant_beta> base_t;
typedef typename base_t::smap_t smap_t; typedef typename base_t::smap_t smap_t;
typedef typename base_t::rmap_t rmap_t;
typedef typename base_t::State State; typedef typename base_t::State State;
using base_t::_s; using base_t::_s;
template <class Graph, class RNG> template <class Graph, class RNG>
SIRS_state(Graph& g, smap_t s, smap_t s_temp, python::dict params, RNG& 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), : 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: ...@@ -375,8 +390,9 @@ public:
{ {
if (_s[v] == State::R) if (_s[v] == State::R)
{ {
std::bernoulli_distribution srecover(_mu); auto mu = _mu[v];
if (_mu > 0 && srecover(rng)) std::bernoulli_distribution srecover(mu);
if (mu > 0 && srecover(rng))
{ {
s_out[v] = State::S; s_out[v] = State::S;
return 1; return 1;
...@@ -392,7 +408,7 @@ public: ...@@ -392,7 +408,7 @@ public:
constexpr bool has_absorbing() { return false; } constexpr bool has_absorbing() { return false; }
private: private:
double _mu; rmap_t _mu;
}; };
......
...@@ -70,7 +70,7 @@ Contents ...@@ -70,7 +70,7 @@ Contents
from __future__ import division, absolute_import, print_function from __future__ import division, absolute_import, print_function
from .. import _degree, _prop, Graph, GraphView, _get_rng, PropertyMap, \ from .. import _degree, _prop, Graph, GraphView, _get_rng, PropertyMap, \
EdgePropertyMap, _check_prop_scalar EdgePropertyMap, VertexPropertyMap, _check_prop_scalar
from .. stats import label_self_loops from .. stats import label_self_loops
import numpy import numpy
import numpy.random import numpy.random
...@@ -146,7 +146,7 @@ class DiscreteStateBase(object): ...@@ -146,7 +146,7 @@ class DiscreteStateBase(object):
return self._state.iterate_async(niter, _get_rng()) return self._state.iterate_async(niter, _get_rng())
class EpidemicStateBase(DiscreteStateBase): 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 r"""Base state for epidemic dynamics. This class it not meant to be
instantiated directly.""" instantiated directly."""
...@@ -159,6 +159,7 @@ class EpidemicStateBase(DiscreteStateBase): ...@@ -159,6 +159,7 @@ class EpidemicStateBase(DiscreteStateBase):
v0 = g.vertex(v0, use_index=False) v0 = g.vertex(v0, use_index=False)
self.s[v0] = 1 self.s[v0] = 1
beta = params["beta"]
weighted = isinstance(beta, EdgePropertyMap) weighted = isinstance(beta, EdgePropertyMap)
if weighted: if weighted:
_check_prop_scalar(beta, "beta") _check_prop_scalar(beta, "beta")
...@@ -166,7 +167,18 @@ class EpidemicStateBase(DiscreteStateBase): ...@@ -166,7 +167,18 @@ class EpidemicStateBase(DiscreteStateBase):
if constant_beta: if constant_beta:
raise ValueError("if constant_beta == True, the type of beta must be double") raise ValueError("if constant_beta == True, the type of beta must be double")
beta = beta.copy("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) self.make_state = lambda *args: make_state(*args, weighted, constant_beta)
...@@ -183,12 +195,12 @@ class SIState(EpidemicStateBase): ...@@ -183,12 +195,12 @@ class SIState(EpidemicStateBase):
Transmission probability. If an :class:`~graph_tool.EdgePropertyMap` Transmission probability. If an :class:`~graph_tool.EdgePropertyMap`
object is passed, it must contain the transmission probability for object is passed, it must contain the transmission probability for
every edge. every edge.
r : ``float`` (optional, default: ``0.``) r : ``float`` or :class:`~graph_tool.VertexPropertyMap` (optional, default: ``0.``)
Spontaneous infection probability. Spontaneous infection probability.
exposed : ``boolean`` (optional, default: ``False``) exposed : ``boolean`` (optional, default: ``False``)
If ``True``, an SEI model is simulated, with an additional "exposed" If ``True``, an SEI model is simulated, with an additional "exposed"
state. 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 Susceptible to exposed transition probability. This only has an
effect if ``exposed=True``. effect if ``exposed=True``.
v0 : ``int`` or :class:`~graph_tool.Vertex` (optional, default: ``None``) v0 : ``int`` or :class:`~graph_tool.Vertex` (optional, default: ``None``)
...@@ -217,7 +229,7 @@ class SIState(EpidemicStateBase): ...@@ -217,7 +229,7 @@ class SIState(EpidemicStateBase):
1. If :math:`s_i(t) = 0`, we have :math:`s_i(t+1) = 1` with probability 1. If :math:`s_i(t) = 0`, we have :math:`s_i(t+1) = 1` with probability
.. math:: .. 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`. otherwise :math:`s_i(t+1) = 0`.
...@@ -226,7 +238,7 @@ class SIState(EpidemicStateBase): ...@@ -226,7 +238,7 @@ class SIState(EpidemicStateBase):
If the option ``exposed == True`` is given, then the states transit If the option ``exposed == True`` is given, then the states transit
first from 0 to -1 (exposed) with probability given by 1. above, and 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 Examples
-------- --------
...@@ -270,11 +282,11 @@ class SIState(EpidemicStateBase): ...@@ -270,11 +282,11 @@ class SIState(EpidemicStateBase):
:doi:`10.1103/RevModPhys.87.925`, :arxiv:`1408.2701` :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, lib_dynamics.make_SEI_state if exposed else lib_dynamics.make_SI_state,
dict(beta=beta, r=r, epsilon=epsilon),
v0, s) v0, s)
DiscreteStateBase.__init__(self, g, self.make_state, DiscreteStateBase.__init__(self, g, self.make_state, self.params,
dict(beta=self.beta, r=r, epsilon=epsilon),
self.s) self.s)
class SISState(DiscreteStateBase): class SISState(DiscreteStateBase):
...@@ -325,17 +337,17 @@ 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 1. If :math:`s_i(t) = 0`, we have :math:`s_i(t+1) = 1` with probability
.. math:: .. 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`. otherwise :math:`s_i(t+1) = 0`.
2. If :math:`s_i(t) = 1`, we have :math:`s_i(t+1) = 0` with probability 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:`\gamma_i`, or :math:`s_i(t+1) = 1` with probability
:math:`1-\gamma`. :math:`1-\gamma_i`.
If the option ``exposed == True`` is given, then the states transit If the option ``exposed == True`` is given, then the states transit
first from 0 to -1 (exposed) with probability given by 1. above, and 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 Examples
-------- --------
...@@ -379,12 +391,12 @@ class SISState(DiscreteStateBase): ...@@ -379,12 +391,12 @@ class SISState(DiscreteStateBase):
:doi:`10.1103/RevModPhys.87.925`, :arxiv:`1408.2701` :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, lib_dynamics.make_SEIS_state if exposed else lib_dynamics.make_SIS_state,
dict(beta=beta, gamma=gamma, r=r,
epsilon=epsilon),
v0, s) v0, s)
DiscreteStateBase.__init__(self, g, self.make_state, DiscreteStateBase.__init__(self, g, self.make_state, self.params, self.s)
dict(beta=self.beta, gamma=gamma, r=r,
epsilon=epsilon), self.s)
class SIRState(DiscreteStateBase): class SIRState(DiscreteStateBase):
def __init__(self, g, beta=1., gamma=.1, r=0, exposed=False, epsilon=.1, def __init__(self, g, beta=1., gamma=.1, r=0, exposed=False, epsilon=.1,
...@@ -434,17 +446,17 @@ class SIRState(DiscreteStateBase): ...@@ -434,17 +446,17 @@ class SIRState(DiscreteStateBase):
1. If :math:`s_i(t) = 0`, we have :math:`s_i(t+1) = 1` with probability 1. If :math:`s_i(t) = 0`, we have :math:`s_i(t+1) = 1` with probability
.. math:: .. 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`. otherwise :math:`s_i(t+1) = 0`.
2. If :math:`s_i(t) = 1`, we have :math:`s_i(t+1) = 2` with probability 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:`\gamma_i`, or :math:`s_i(t+1) = 1` with probability
:math:`1-\gamma`. :math:`1-\gamma_i`.
If the option ``exposed == True`` is given, then the states transit If the option ``exposed == True`` is given, then the states transit
first from 0 to -1 (exposed) with probability given by 1. above, and 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 Examples
-------- --------
...@@ -498,12 +510,13 @@ class SIRState(DiscreteStateBase): ...@@ -498,12 +510,13 @@ class SIRState(DiscreteStateBase):
:doi:`10.1103/RevModPhys.87.925`, :arxiv:`1408.2701` :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, lib_dynamics.make_SEIR_state if exposed else lib_dynamics.make_SIR_state,
dict(beta=beta, gamma=gamma, r=r,
epsilon=epsilon),
v0, s) v0, s)
DiscreteStateBase.__init__(self, g, self.make_state, DiscreteStateBase.__init__(self, g, self.make_state,
dict(beta=self.beta, gamma=gamma, r=r, self.params, self.s)
epsilon=epsilon), self.s)
class SIRSState(DiscreteStateBase): class SIRSState(DiscreteStateBase):
def __init__(self, g, beta=1, gamma=.1, mu=.1, r=0, exposed=False, def __init__(self, g, beta=1, gamma=.1, mu=.1, r=0, exposed=False,
...@@ -556,21 +569,21 @@ class SIRSState(DiscreteStateBase): ...@@ -556,21 +569,21 @@ class SIRSState(DiscreteStateBase):
1. If :math:`s_i(t) = 0`, we have :math:`s_i(t+1) = 1` with probability 1. If :math:`s_i(t) = 0`, we have :math:`s_i(t+1) = 1` with probability
.. math:: .. 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`. otherwise :math:`s_i(t+1) = 0`.
2. If :math:`s_i(t) = 1`, we have :math:`s_i(t+1) = 2` with probability 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:`\gamma_i`, or :math:`s_i(t+1) = 1` with probability
:math:`1-\gamma`. :math:`1-\gamma_i`.
3. If :math:`s_i(t) = 2`, we have :math:`s_i(t+1) = 1` with probability 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:`\mu_i`, or :math:`s_i(t+1) = 2` with probability
:math:`1-\mu`. :math:`1-\mu_i`.
If the option ``exposed == True`` is given, then the states transit If the option ``exposed == True`` is given, then the states transit
first from 0 to -1 (exposed) with probability given by 1. above, and 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 Examples
-------- --------
...@@ -624,12 +637,13 @@ class SIRSState(DiscreteStateBase): ...@@ -624,12 +637,13 @@ class SIRSState(DiscreteStateBase):
:doi:`10.1103/RevModPhys.87.925`, :arxiv:`1408.2701` :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, 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) v0, s)
DiscreteStateBase.__init__(self, g, self.make_state, DiscreteStateBase.__init__(self, g, self.make_state,
dict(beta=self.beta, gamma=gamma, mu=mu, r=r, self.params, self.s)
epsilon=epsilon), self.s)
class VoterState(DiscreteStateBase): class VoterState(DiscreteStateBase):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment