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

dynamics: implement weighted versions of epidemic models

This fixes issue #624
parent 3fb37596
......@@ -120,7 +120,6 @@ struct add_ptr
};
};
template <class State>
void export_state()
{
......@@ -132,31 +131,57 @@ void export_state()
});
}
template <template <bool...> class Base, bool... As>
void export_SI_state()
{
export_state<Base<As..., false, false>>();
export_state<Base<As..., true, false>>();
export_state<Base<As..., true, true>>();
}
template <template <bool...> class Base, bool... As>
python::object make_SI_state(GraphInterface& gi, boost::any as,
boost::any as_temp, python::dict params,
rng_t& rng, bool weighted, bool constant_beta)
{
if (weighted)
{
if (constant_beta)
return make_state<Base<As..., true, true>>(gi, as, as_temp, params, rng);
else
return make_state<Base<As..., true, false>>(gi, as, as_temp, params, rng);
}
else
{
return make_state<Base<As..., false, false>>(gi, as, as_temp, params, rng);
}
}
void export_discrete()
{
export_state<SI_state<false>>();
def("make_SI_state", &make_state<SI_state<false>>);
export_SI_state<SI_state,false>();
def("make_SI_state", &make_SI_state<SI_state,false>);
export_state<SIS_state<false,false>>();
def("make_SIS_state", &make_state<SIS_state<false,false>>);
export_SI_state<SIS_state,false,false>();
def("make_SIS_state", &make_SI_state<SIS_state,false,false>);
export_state<SIS_state<false,true>>();
def("make_SIR_state", &make_state<SIS_state<false,true>>);
export_SI_state<SIS_state,false,true>();
def("make_SIR_state", &make_SI_state<SIS_state,false,true>);
export_state<SIRS_state<false>>();
def("make_SIRS_state", &make_state<SIRS_state<false>>);
export_SI_state<SIRS_state,false>();
def("make_SIRS_state", &make_SI_state<SIRS_state,false>);
export_state<SI_state<true>>();
def("make_SEI_state", &make_state<SI_state<true>>);
export_SI_state<SI_state,true>();
def("make_SEI_state", &make_SI_state<SI_state,true>);
export_state<SIS_state<true,false>>();
def("make_SEIS_state", &make_state<SIS_state<true,false>>);
export_SI_state<SIS_state,true,false>();
def("make_SEIS_state", &make_SI_state<SIS_state,true,false>);
export_state<SIS_state<true,true>>();
def("make_SEIR_state", &make_state<SIS_state<true,true>>);
export_SI_state<SIS_state,true,true>();
def("make_SEIR_state", &make_SI_state<SIS_state,true,true>);
export_state<SIRS_state<true>>();
def("make_SEIRS_state", &make_state<SIRS_state<true>>);
export_SI_state<SIRS_state,true>();
def("make_SEIRS_state", &make_SI_state<SIRS_state,true>);
export_state<voter_state>();
def("make_voter_state", &make_state<voter_state>);
......
......@@ -21,6 +21,7 @@
#include "graph.hh"
#include "graph_filtering.hh"
#include "graph_util.hh"
#include "graph_python_interface.hh"
#ifdef _OPENMP
#include <omp.h>
#endif
......@@ -63,36 +64,74 @@ public:
std::shared_ptr<std::vector<size_t>> _active;
};
template <bool exposed>
template <bool exposed, bool weighted, bool constant_beta>
class SI_state: public discrete_state_base<>
{
public:
enum State { S, I, R, E};
enum State { S, I, R, E };
typedef typename eprop_map_t<double>::type::unchecked_t bmap_t;
typedef std::conditional_t<weighted, bmap_t, double> beta_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),
_beta(python::extract<double>(params["beta"])),
_epsilon(python::extract<double>(params["epsilon"])),
_r(python::extract<double>(params["r"])),
_m(num_vertices(g)),
_m_temp(num_vertices(g))
{
size_t M = 0;
for (auto v : vertices_range(g))
python::object obeta = params["beta"];
if constexpr (weighted)
{
obeta = obeta.attr("_get_any")();
boost::any& abeta = python::extract<boost::any&>(obeta);
_beta = boost::any_cast<typename beta_t::checked_t>(abeta).get_unchecked();
}
else
{
size_t k = 0;
for (auto w : in_or_out_neighbors_range(v, g))
_beta = python::extract<beta_t>(obeta);
}
if constexpr (!weighted)
{
size_t M = 0;
for (auto v : vertices_range(g))
{
_m[v] += (_s[w] == State::I);
++k;
size_t k = 0;
for (auto w : in_or_out_neighbors_range(v, g))
{
_m[v] += (_s[w] == State::I);
++k;
}
_m_temp[v] = _m[v];
M = std::max(M, k);
}
for (size_t m = 0; m < M + 1; ++m)
_prob.push_back(1-std::pow(1-_beta, m));
}
else
{
if constexpr (constant_beta)
{
eprop_map_t<double>::type beta(get(edge_index_t(), g));
for (auto e : edges_range(g))
beta[e] = std::log1p(-_beta[e]);
_beta = beta;
}
for (auto v : vertices_range(g))
{
for (auto e : in_or_out_edges_range(v, g))
{
auto w = (source(e, g) != v) ? source(e, g) : target(e, g);
if (_s[w] == State::I)
_m[v] += get_p(e);
}
_m_temp[v] = _m[v];
}
_m_temp[v] = _m[v];
M = std::max(M, k);
}
for (size_t m = 0; m < M + 1; ++m)
_prob.push_back(1-std::pow(1-_beta, m));
};
template <class Graph>
......@@ -101,23 +140,57 @@ public:
s_out[v] = State::E;
}
template <class Edge>
constexpr double get_p(Edge& e)
{
if constexpr (constant_beta)
return _beta[e];
else
return std::log1p(-_beta[e]);
}
template <bool sync, class Graph>
void infect(Graph& g, size_t v, smap_t& s_out)
{
s_out[v] = State::I;
if (sync)
if constexpr (!weighted)
{
for (auto w : out_neighbors_range(v, g))
if constexpr (sync)
{
auto& m = _m_temp[w];
#pragma omp atomic
m++;
for (auto w : out_neighbors_range(v, g))
{
auto& m = _m_temp[w];
#pragma omp atomic
m++;
}
}
else
{
for (auto w : out_neighbors_range(v, g))
_m[w]++;
}
}
else
{
for (auto w : out_neighbors_range(v, g))
_m[w]++;
if constexpr (sync)
{
for (auto e : out_edges_range(v, g))
{
auto w = target(e, g);
auto& m = _m_temp[w];
auto p = get_p(e);
#pragma omp atomic
m += p;
}
}
else
{
for (auto e : out_edges_range(v, g))
{
auto w = target(e, g);
_m[w] += get_p(e);
}
}
}
}
......@@ -150,8 +223,14 @@ public:
auto m = _m[v];
std::bernoulli_distribution minfect(_prob[m]);
if (m > 0 && minfect(rng))
double prob = 0;
if constexpr (!weighted)
prob = _prob[m];
else
prob = 1 - std::exp(m);
std::bernoulli_distribution minfect(prob);
if (prob > 0 && minfect(rng))
{
if constexpr (exposed)
expose(g, v, s_out);
......@@ -175,27 +254,32 @@ public:
constexpr bool has_absorbing() { return true; }
protected:
double _beta;
beta_t _beta;
double _epsilon;
double _r;
discrete_state_base<>::smap_t _m, _m_temp;
typedef std::conditional_t<weighted,
typename vprop_map_t<double>::type::unchecked_t,
discrete_state_base<>::smap_t> m_t;
m_t _m, _m_temp;
std::vector<double> _prob;
};
template <bool exposed, bool recovered>
class SIS_state: public SI_state<exposed>
template <bool exposed, bool recovered, bool weighted, bool constant_beta>
class SIS_state: public SI_state<exposed, weighted, constant_beta>
{
public:
typedef typename SI_state<exposed>::smap_t smap_t;
typedef typename SI_state<exposed>::State State;
using SI_state<exposed>::_s;
using SI_state<exposed>::_m;
using SI_state<exposed>::_m_temp;
typedef SI_state<exposed, weighted, constant_beta> base_t;
typedef typename base_t::smap_t smap_t;
typedef typename base_t::State State;
using base_t::_s;
using base_t::_m;
using base_t::_m_temp;
template <class Graph, class RNG>
SIS_state(Graph& g, smap_t s, smap_t s_temp, python::dict params, RNG& rng)
: SI_state<exposed>(g, s, s_temp, params, rng),
: base_t(g, s, s_temp, params, rng),
_gamma(python::extract<double>(params["gamma"]))
{};
......@@ -203,19 +287,44 @@ public:
void recover(Graph& g, size_t v, smap_t& s_out)
{
s_out[v] = recovered ? State::R : State::S;
if constexpr (sync)
if constexpr (!weighted)
{
for (auto w : out_neighbors_range(v, g))
if constexpr (sync)
{
auto& m = _m_temp[w];
#pragma omp atomic
m--;
for (auto w : out_neighbors_range(v, g))
{
auto& m = _m_temp[w];
#pragma omp atomic
m--;
}
}
else
{
for (auto w : out_neighbors_range(v, g))
_m[w]--;
}
}
else
{
for (auto w : out_neighbors_range(v, g))
_m[w]--;
if constexpr (sync)
{
for (auto e : out_edges_range(v, g))
{
auto w = target(e, g);
auto& m = _m_temp[w];
auto p = base_t::get_p(e);
#pragma omp atomic
m -= p;
}
}
else
{
for (auto e : out_edges_range(v, g))
{
auto w = target(e, g);
_m[w] -= base_t::get_p(e);
}
}
}
}
......@@ -232,7 +341,7 @@ public:
}
return 0;
}
return SI_state<exposed>::template update_node<sync>(g, v, s_out, rng);
return base_t::template update_node<sync>(g, v, s_out, rng);
}
template <class Graph>
......@@ -245,18 +354,18 @@ protected:
double _gamma;
};
template <bool exposed>
class SIRS_state: public SIS_state<exposed, true>
template <bool exposed, bool weighted, bool constant_beta>
class SIRS_state: public SIS_state<exposed, true, weighted, constant_beta>
{
public:
typedef typename SI_state<exposed>::smap_t smap_t;
typedef typename SI_state<exposed>::State State;
using SI_state<exposed>::_s;
typedef SIS_state<exposed, true, weighted, constant_beta> base_t;
typedef typename base_t::smap_t smap_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)
: SIS_state<exposed,true>(g, s, s_temp, params, rng),
: base_t(g, s, s_temp, params, rng),
_mu(python::extract<double>(params["mu"]))
{};
......@@ -274,7 +383,7 @@ public:
}
return 0;
}
return SIS_state<exposed,true>::template update_node<sync>(g, v, s_out, rng);
return base_t::template update_node<sync>(g, v, s_out, rng);
}
template <class Graph>
......
......@@ -805,6 +805,11 @@ void do_add_edge_list_hashed(GraphInterface& gi, python::object aedge_list,
void do_add_edge_list_iter(GraphInterface& gi, python::object edge_list,
python::object eprops);
bool hasattr(boost::python::object obj, std::string const& attrName)
{
return PyObject_HasAttrString(obj.ptr(), attrName.c_str());
}
} // namespace graph_tool
// register everything
......
......@@ -738,6 +738,7 @@ private:
std::istream& _s;
};
bool hasattr(boost::python::object obj, std::string const& attrName);
} //graph_tool namespace
......
......@@ -23,6 +23,7 @@
#include <vector>
#include "graph_blockmodel_dynamics.hh"
#include "graph_python_interface.hh"
namespace graph_tool
{
......@@ -234,8 +235,6 @@ double l2sinha(T x) // log((exp(x) - exp(-x))/x)
return x + log1p(-exp(-2*x)) - log(x);
}
bool hasattr(boost::python::object obj, std::string const& attrName);
class CIsingBaseState
{
public:
......
......@@ -23,6 +23,7 @@
#include <vector>
#include "graph_blockmodel_dynamics.hh"
#include "graph_python_interface.hh"
namespace graph_tool
{
......@@ -697,8 +698,6 @@ double l1p2cosh(T x) // log(1 + exp(x) + exp(-x))
}
bool hasattr(boost::python::object obj, std::string const& attrName);
class IsingBaseState
{
public:
......
......@@ -60,14 +60,6 @@ python::object make_pseudo_ising_state(boost::python::object oblock_state,
return state;
}
namespace graph_tool
{
bool hasattr(boost::python::object obj, std::string const& attrName)
{
return PyObject_HasAttrString(obj.ptr(), attrName.c_str());
}
}
void export_pseudo_ising_state()
{
using namespace boost::python;
......
......@@ -25,6 +25,7 @@
#include "graph_tool.hh"
#include "../support/graph_state.hh"
#include "graph_blockmodel_dynamics.hh"
#include "graph_python_interface.hh"
namespace graph_tool
{
......
......@@ -69,7 +69,8 @@ Contents
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
from .. stats import label_self_loops
import numpy
import numpy.random
......@@ -145,7 +146,7 @@ class DiscreteStateBase(object):
return self._state.iterate_async(niter, _get_rng())
class EpidemicStateBase(DiscreteStateBase):
def __init__(self, g, v0=None, s=None):
def __init__(self, g, beta, constant_beta, make_state, v0=None, s=None):
r"""Base state for epidemic dynamics. This class it not meant to be
instantiated directly."""
......@@ -158,16 +159,30 @@ class EpidemicStateBase(DiscreteStateBase):
v0 = g.vertex(v0, use_index=False)
self.s[v0] = 1
weighted = isinstance(beta, EdgePropertyMap)
if weighted:
_check_prop_scalar(beta, "beta")
if beta.value_type() != "double":
if constant_beta:
raise ValueError("if constant_beta == True, the type of beta must be double")
beta = beta.copy("double")
self.beta = beta
self.make_state = lambda *args: make_state(*args, weighted, constant_beta)
class SIState(EpidemicStateBase):
def __init__(self, g, beta=1., r=0, exposed=False, epsilon=.1, v0=None, s=None):
def __init__(self, g, beta=1., r=0, exposed=False, epsilon=.1, v0=None,
s=None, constant_beta=True):
r"""SI compartmental epidemic model.
Parameters
----------
g : :class:`~graph_tool.Graph`
Graph to be used for the dynamics
beta : ``float`` (optional, default: ``1.``)
Transmission probability.
beta : ``float`` or :class:`~graph_tool.EdgePropertyMap` (optional, default: ``1.``)
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.``)
Spontaneous infection probability.
exposed : ``boolean`` (optional, default: ``False``)
......@@ -182,6 +197,11 @@ class SIState(EpidemicStateBase):
s : :class:`~graph_tool.VertexPropertyMap` (optional, default: ``None``)
Initial global state. If not provided, all vertices will be
initialized to the susceptible state.
constant_beta : ``boolean`` (optional, default: ``True``)
If ``True``, and ``beta`` is an edge property map, it will be assumed
that the ``beta`` values do not change, such that the probability
values can be pre-computed for efficiency. If ``beta`` is a
``float``, this option has no effect.
Notes
-----
......@@ -197,7 +217,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)^{A_{ij}\delta_{s_j(t),1}}\right) + r,
(1-r)\left(1-\prod_j(1-\beta_{ij})^{A_{ij}\delta_{s_j(t),1}}\right) + r,
otherwise :math:`s_i(t+1) = 0`.
......@@ -250,14 +270,16 @@ class SIState(EpidemicStateBase):
:doi:`10.1103/RevModPhys.87.925`, :arxiv:`1408.2701`
"""
EpidemicStateBase.__init__(self, g, v0, s)
DiscreteStateBase.__init__(self, g,
EpidemicStateBase.__init__(self, g, beta, constant_beta,
lib_dynamics.make_SEI_state if exposed else lib_dynamics.make_SI_state,
dict(beta=beta, r=r, epsilon=epsilon), self.s)
v0, s)
DiscreteStateBase.__init__(self, g, self.make_state,
dict(beta=self.beta, r=r, epsilon=epsilon),
self.s)
class SISState(DiscreteStateBase):
def __init__(self, g, beta=1., gamma=.1, r=0, exposed=False, epsilon=.1,
v0=None, s=None):
v0=None, s=None, constant_beta=True):
r"""SIS compartmental epidemic model.
Parameters
......@@ -282,6 +304,11 @@ class SISState(DiscreteStateBase):
s : :class:`~graph_tool.VertexPropertyMap` (optional, default: ``None``)
Initial global state. If not provided, all vertices will be
initialized to the susceptible state.
constant_beta : ``boolean`` (optional, default: ``True``)
If ``True``, and ``beta`` is an edge property map, it will be assumed
that the ``beta`` values do not change, such that the probability
values can be pre-computed for efficiency. If ``beta`` is a
``float``, this option has no effect.
Notes
-----
......@@ -298,7 +325,7 @@ 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)^{A_{ij}\delta_{s_j(t),1}}\right) + r,
(1-r)\left(1-\prod_j(1-\beta_{ij})^{A_{ij}\delta_{s_j(t),1}}\right) + r,
otherwise :math:`s_i(t+1) = 0`.
......@@ -352,14 +379,16 @@ class SISState(DiscreteStateBase):
:doi:`10.1103/RevModPhys.87.925`, :arxiv:`1408.2701`
"""
EpidemicStateBase.__init__(self, g, v0, s)
DiscreteStateBase.__init__(self, g,
EpidemicStateBase.__init__(self, g, beta, constant_beta,
lib_dynamics.make_SEIS_state if exposed else lib_dynamics.make_SIS_state,
dict(beta=beta, gamma=gamma, r=r, epsilon=epsilon), self.s)