dynamics: implement weighted versions of epidemic models

This fixes issue #624

src/graph/dynamics/graph_discrete.cc

@@ -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>);

src/graph/dynamics/graph_discrete.hh

@@ -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>

src/graph/graph_python_interface.cc

@@ -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

src/graph/graph_python_interface.hh

@@ -738,6 +738,7 @@ private:
     std::istream& _s;
 };
 
+bool hasattr(boost::python::object obj, std::string const& attrName);
 
 } //graph_tool namespace
 

src/graph/inference/uncertain/graph_blockmodel_dynamics_continuous.hh

@@ -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:

src/graph/inference/uncertain/graph_blockmodel_dynamics_discrete.hh

@@ -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:

src/graph/inference/uncertain/graph_blockmodel_dynamics_pseudo_ising.cc

@@ -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;

src/graph/inference/uncertain/graph_blockmodel_dynamics_pseudo_ising_mcmc_h.hh

@@ -25,6 +25,7 @@
 #include "graph_tool.hh"
 #include "../support/graph_state.hh"