Implement BlockState.get_edges_prob()

src/graph/inference/graph_blockmodel_layers.hh

@@ -167,8 +167,6 @@ struct Layers
         size_t _total_B;
         bool _is_partition_stats_enabled;
 
-        //TODO: remove_vertex and add_vertex, etc.
-
         void move_vertex(size_t v, size_t s)
         {
             if (BaseState::_vweight[v] == 0)
@@ -217,6 +215,39 @@ struct Layers
             }
         }
 
+        void remove_vertex(size_t v)
+        {
+            size_t r = _b[v];
+            auto& ls = _vc[v];
+            auto& vs = _vmap[v];
+            for (size_t j = 0; j < ls.size(); ++j)
+            {
+                int l = ls[j];
+                size_t u = vs[j];
+                auto& state = _layers[l];
+                size_t r_u = state._b[u];
+                state.remove_vertex(u);
+                if (state._wr[r_u] == 0)
+                    state.remove_block_map(r);
+            }
+            BaseState::remove_vertex(v);
+        }
+
+        void add_vertex(size_t v, size_t r)
+        {
+            auto& ls = _vc[v];
+            auto& vs = _vmap[v];
+            for (size_t j = 0; j < ls.size(); ++j)
+            {
+                int l = ls[j];
+                size_t u = vs[j];
+                auto& state = _layers[l];
+                size_t r_u = state.get_block_map(r);
+                state.add_vertex(u, r_u);
+            }
+            BaseState::add_vertex(v, r);
+        }
+
         template <class VMap>
         void set_partition(VMap&& b)
         {

src/graph_tool/inference/blockmodel.py

@@ -654,6 +654,70 @@ class BlockState(object):
         """
         return self._state.get_move_prob(int(v), self.b[v], s, c, reverse)
 
+    def get_edges_prob(self, edge_list, missing=True, entropy_args={}):
+        """Compute the log-probability of the missing (or spurious if ``missing=False``)
+        edges given by ``edge_list`` (a list of ``(source, target)`` tuples, or
+        :meth:`~graph_tool.Edge` instances). The values in ``entropy_args`` are
+        passed to :meth:`graph_tool.BlockState.entropy()` to calculate the
+        log-probability.
+        """
+        pos = {}
+        for u, v in edge_list:
+            pos[u] = self.b[u]
+            pos[v] = self.b[v]
+
+        Si = self.entropy(**entropy_args)
+
+        for v in pos.keys():
+            self.remove_vertex(v)
+
+        try:
+            if missing:
+                new_es = []
+                for u, v in edge_list:
+                    e = self.g.add_edge(u, v)
+                    new_es.append(e)
+                    self.E += 1
+            else:
+                old_es = []
+                for e in edge_list:
+                    if isinstance(e, tuple):
+                        u, v = e
+                        tmp = self.g.edge(u, v)
+                        if tmp is None:
+                            raise ValueError("edge not found: (%d, %d)" % (int(u),
+                                                                           int(v)))
+                        self.g.remove_edge(tmp)
+                    else:
+                        u, v = e
+                        self.g.remove_edge(e)
+                    old_es.append((u, v))
+                    self.E -= 1
+
+            for v in pos.keys():
+                self.add_vertex(v, pos[v])
+
+            Sf = self.entropy(**entropy_args)
+
+            for v in pos.keys():
+                self.remove_vertex(v)
+
+        finally:
+            if missing:
+                for e in new_es:
+                    self.g.remove_edge(e)
+                    self.E -= 1
+            else:
+                for u, v in old_es:
+                    self.g.add_edge(u, v)
+                    self.E += 1
+            for v in pos.keys():
+                self.add_vertex(v, pos[v])
+
+        if missing:
+            return Si - Sf
+        else:
+            return Sf - Si
 
     def _mcmc_sweep_dispatch(self, mcmc_state):
         if (mcmc_state.multigraph and not mcmc_state.dense and

src/graph_tool/inference/nested_blockmodel.py

@@ -177,7 +177,8 @@ class NestedBlockState(object):
         if _bm_test():
             self._consistency_check()
 
-    def level_entropy(self, l, dense=False, multigraph=True, bstate=None, **kwargs):
+    def level_entropy(self, l, dense=False, multigraph=True, bstate=None,
+                      **kwargs):
         """Compute the entropy of level ``l``."""
 
         if bstate is None:
@@ -198,6 +199,24 @@ class NestedBlockState(object):
                                     **kwargs)
         return S
 
+    def get_edges_prob(self, edge_list, missing=True, entropy_args={}):
+        """Compute the log-probability of the missing (or spurious if ``missing=False``)
+        edges given by ``edge_list`` (a list of ``(source, target)`` tuples, or
+        :meth:`~graph_tool.Edge` instances). The values in ``entropy_args`` are
+        passed to :meth:`graph_tool.NestedBlockState.entropy()` to calculate the
+        log-probability.
+        """
+        L = 0
+        for l, state in enumerate(self.levels):
+            eargs = overlay(entropy_args, dl=True,
+                            edges_dl=(l == (len(self.levels) - 1)))
+            if l > 0:
+                eargs = overlay(eargs, dense=True, multigraph=True)
+            L += state.get_edges_prob(edge_list, missing=missing,
+                                      entropy_args=eargs)
+            edge_list = [(state.b[u], state.b[v]) for u, v in edge_list]
+        return L
+
     def get_bstack(self):
         """Return the nested levels as individual graphs.