From aebfabed9aa4f652b075415f56272a453b94e1c0 Mon Sep 17 00:00:00 2001
From: Tiago de Paula Peixoto <tiago@skewed.de>
Date: Sat, 30 Mar 2019 14:22:30 +0000
Subject: [PATCH] Add 'edges_only' options to UncertainBlockModel

---
 .../uncertain/graph_blockmodel_sample_edge.hh | 51 ++++++++++++++-----
 .../graph_blockmodel_uncertain_mcmc.hh        |  3 +-
 .../inference/uncertain_blockmodel.py         |  1 +
 3 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/src/graph/inference/uncertain/graph_blockmodel_sample_edge.hh b/src/graph/inference/uncertain/graph_blockmodel_sample_edge.hh
index 3a73da1a..3ddb2985 100644
--- a/src/graph/inference/uncertain/graph_blockmodel_sample_edge.hh
+++ b/src/graph/inference/uncertain/graph_blockmodel_sample_edge.hh
@@ -38,11 +38,25 @@ template <class State>
 class SBMEdgeSampler
 {
 public:
-    SBMEdgeSampler(State& state)
+    SBMEdgeSampler(State& state, bool edges_only=false)
         : _state(state),
           _v_in_sampler(graph_tool::is_directed(state._g) ?
-                        __v_in_sampler : _v_out_sampler)
+                        __v_in_sampler : _v_out_sampler),
+          _edges_only(edges_only)
     {
+        _N = num_vertices(state._g);
+
+        for (auto e : edges_range(state._g))
+        {
+            size_t u = source(e, state._g);
+            size_t v = target(e, state._g);
+            _edges.push_back(get_edge(u, v));
+            _edge_pos[_edges.back()] = _edges.size() - 1;
+        }
+
+        if (_edges_only)
+            return;
+
         for (auto me : edges_range(state._bg))
         {
             auto r = source(me, state._bg);
@@ -87,15 +101,6 @@ public:
 
         auto B = _groups.size();
         _NB = B * B;
-        _N = num_vertices(state._g);
-
-        for (auto e : edges_range(state._g))
-        {
-            size_t u = source(e, state._g);
-            size_t v = target(e, state._g);
-            _edges.push_back(get_edge(u, v));
-            _edge_pos[_edges.back()] = _edges.size() - 1;
-        }
     }
 
     std::tuple<size_t, size_t> get_edge(size_t u, size_t v)
@@ -109,6 +114,9 @@ public:
     template <bool add>
     void update_edge(size_t u, size_t v, size_t m)
     {
+        if (_edges_only)
+            return;
+
         if (add)
         {
             if (m == 0)
@@ -203,6 +211,22 @@ public:
         // std::uniform_int_distribution<size_t> sample(0, _N-1);
         // return {sample(rng), sample(rng)};
 
+        if (_edges_only)
+        {
+            std::bernoulli_distribution coin(_edges.size() /
+                                             double(_edges.size() + _N));
+            if (coin(rng))
+            {
+                return uniform_sample(_edges, rng);
+            }
+            else
+            {
+                std::uniform_int_distribution<size_t> vsample(0, _N-1);
+                auto v = vsample(rng);
+                return {v, v};
+            }
+        }
+
         if (!_edges.empty())
         {
             std::bernoulli_distribution coin(.5);
@@ -231,7 +255,8 @@ public:
 
     double log_prob(size_t u, size_t v, size_t m, int delta)
     {
-        // return 0;
+        if (_edges_only)
+            return 0;
 
         auto& g = _state._g;
         size_t r = _state._b[u];
@@ -344,6 +369,8 @@ private:
     size_t _NB = 0;
     size_t _E = 0;
     size_t _N = 0;
+
+    bool _edges_only;
 };
 
 
diff --git a/src/graph/inference/uncertain/graph_blockmodel_uncertain_mcmc.hh b/src/graph/inference/uncertain/graph_blockmodel_uncertain_mcmc.hh
index 88d8a424..c7250d42 100644
--- a/src/graph/inference/uncertain/graph_blockmodel_uncertain_mcmc.hh
+++ b/src/graph/inference/uncertain/graph_blockmodel_uncertain_mcmc.hh
@@ -39,6 +39,7 @@ typedef std::vector<size_t> vlist_t;
     ((state, &, State&, 0))                                                    \
     ((beta,, double, 0))                                                       \
     ((entropy_args,, uentropy_args_t, 0))                                      \
+    ((edges_only,, bool, 0))                                                   \
     ((verbose,, bool, 0))                                                      \
     ((niter,, size_t, 0))
 
@@ -62,7 +63,7 @@ struct MCMC
                                             sizeof...(Ts)>* = nullptr>
         MCMCUncertainState(ATs&&... as)
             : MCMCUncertainStateBase<Ts...>(as...),
-              _edge_sampler(_state._block_state),
+              _edge_sampler(_state._block_state, _edges_only),
               _vlist(num_vertices(_state._u))
         {
         }
diff --git a/src/graph_tool/inference/uncertain_blockmodel.py b/src/graph_tool/inference/uncertain_blockmodel.py
index f2446ec5..d907f7cf 100644
--- a/src/graph_tool/inference/uncertain_blockmodel.py
+++ b/src/graph_tool/inference/uncertain_blockmodel.py
@@ -159,6 +159,7 @@ class UncertainBaseState(object):
         kwargs = kwargs.copy()
         beta = kwargs.get("beta", 1.)
         niter = kwargs.get("niter", 1)
+        edges_only = kwargs.pop("edges_only", False)
         verbose = kwargs.get("verbose", False)
         slist = self.slist
         tlist = self.tlist
-- 
2.26.2