Add link prediction to layered and nested SBMs

src/graph/inference/graph_blockmodel_layers.cc

@@ -463,12 +463,18 @@ void export_layered_blockmodel_state()
                           = &state_t::set_partition;
                       void (state_t::*move_vertices)(python::object, python::object) =
                           &state_t::move_vertices;
+                      void (state_t::*remove_vertices)(python::object) =
+                          &state_t::remove_vertices;
+                      void (state_t::*add_vertices)(python::object, python::object) =
+                          &state_t::add_vertices;
 
                       class_<state_t> c(name_demangle(typeid(state_t).name()).c_str(),
                                         no_init);
                       c.def("remove_vertex", &state_t::remove_vertex)
                           .def("add_vertex", &state_t::add_vertex)
                           .def("move_vertex", &state_t::move_vertex)
+                          .def("add_vertices", add_vertices)
+                          .def("remove_vertices", remove_vertices)
                           .def("move_vertices", move_vertices)
                           .def("set_partition", set_partition)
                           .def("virtual_move", virtual_move)

src/graph/inference/graph_blockmodel_layers.hh

@@ -250,6 +250,44 @@ struct Layers
             BaseState::remove_vertex(v);
         }
 
+        template <class Vec>
+        void remove_vertices(Vec& vs)
+        {
+            gt_hash_map<size_t, vector<size_t>> lvs;
+            for (auto v : vs)
+                for (auto l : _vc[v])
+                    lvs[l].push_back(v);
+            for (auto& lv : lvs)
+            {
+                auto l = lv.first;
+                auto& state = _layers[l];
+                vector<size_t> us;
+                gt_hash_map<size_t, size_t> rus;
+                for (auto v : lv.second)
+                {
+                    auto u = _vmap[v][l];
+                    us.push_back(u);
+                    size_t r = _b[v];
+                    size_t r_u = state._b[u];
+                    rus[r] = r_u;
+                }
+                state.remove_vertices(us);
+
+                for (auto rr_u : rus)
+                {
+                    if (state._wr[rr_u.second] == 0)
+                        state.remove_block_map(rr_u.first);
+                }
+            }
+            BaseState::remove_vertices(vs);
+        }
+
+        void remove_vertices(python::object ovs)
+        {
+            multi_array_ref<uint64_t, 1> vs = get_array<uint64_t, 1>(ovs);
+            remove_vertices(vs);
+        }
+
         void add_vertex(size_t v, size_t r)
         {
             auto& ls = _vc[v];
@@ -265,6 +303,45 @@ struct Layers
             BaseState::add_vertex(v, r);
         }
 
+        template <class Vs, class Rs>
+        void add_vertices(Vs& vs, Rs& rs)
+        {
+            if (vs.size() != rs.size())
+                throw ValueException("vertex and group lists do not have the same size");
+
+            gt_hash_map<size_t, vector<size_t>> lvs;
+            gt_hash_map<size_t, size_t> vrs;
+            for (size_t i = 0; i < vs.size(); ++i)
+            {
+                auto v = vs[i];
+                vrs[v] = rs[i];
+                for (auto l : _vc[v])
+                    lvs[l].push_back(v);
+            }
+
+            for (auto& lv : lvs)
+            {
+                auto l = lv.first;
+                auto& state = _layers[l];
+                vector<size_t> us;
+                vector<size_t> rus;
+                for (auto v : lv.second)
+                {
+                    us.emplace_back(_vmap[v][l]);
+                    rus.emplace_back(state.get_block_map(vrs[v]));
+                }
+                state.add_vertices(us, rus);
+            }
+            BaseState::add_vertices(vs, rs);
+        }
+
+        void add_vertices(python::object ovs, python::object ors)
+        {
+            multi_array_ref<uint64_t, 1> vs = get_array<uint64_t, 1>(ovs);
+            multi_array_ref<uint64_t, 1> rs = get_array<uint64_t, 1>(ors);
+            add_vertices(vs, rs);
+        }
+
         template <class VMap>
         void set_partition(VMap&& b)
         {

src/graph_tool/inference/layered_blockmodel.py

@@ -637,6 +637,123 @@ class LayeredBlockState(OverlapBlockState, BlockState):
 
         return S
 
+    def _get_lvertex(self, v, l):
+        i = numpy.searchsorted(self.vc[v].a, l)
+        if i >= len(self.vc[v]) or l != self.vc[v][i]:
+            raise ValueError("vertex %d not present in layer %d" % (v, l))
+        u = self.vmap[v][i]
+        return u
+
+    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, ec)`` tuples, or
+        :meth:`~graph_tool.Edge` instances). The values in ``entropy_args`` are
+        passed to :meth:`graph_tool.LayeredBlockState.entropy()` to calculate the
+        log-probability.
+        """
+        pos = {}
+        nes = []
+        for e in edge_list:
+            try:
+                u, v = e
+                l = self.ec[e]
+            except ValueError:
+                u, v, l = e
+
+            pos[u] = self.b[u]
+            pos[v] = self.b[v]
+
+            nes.append((u, v, (l, False)))
+            nes.append((self._get_lvertex(u, l),
+                        self._get_lvertex(v, l), (l, True)))
+
+        edge_list = nes
+
+        Si = self.entropy(**entropy_args)
+
+        self.remove_vertex(pos.keys())
+
+        agg_state = self.agg_state
+
+        try:
+            if missing:
+                new_es = []
+                for u, v, l in edge_list:
+                    if not l[1]:
+                        state = self.agg_state
+                    else:
+                        state = self.layer_states[l[0]]
+                    e = state.g.add_edge(u, v)
+                    if not l[1]:
+                        self.ec[e] = l[0]
+                    if state.is_weighted:
+                        state.eweight[e] = 1
+                    new_es.append((e, l))
+            else:
+                old_es = []
+                for u, v, l in edge_list:
+                    if not l[1]:
+                        state = self.agg_state
+                        es = state.g.edge(u, v, all_edges=True)
+                        es = [e for e in es if self.ec[e] == l[0]]
+                        if len(es) > 0:
+                            e = es[0]
+                        else:
+                            e = None
+                    else:
+                        state = self.layer_states[l[0]]
+                        e = state.g.edge(u, v)
+                    if e is None:
+                        raise ValueError("edge not found: (%d, %d, %d)" % \
+                                         (int(u), int(v), l[0]))
+
+                    if state.is_weighted:
+                        staete.eweight[e] -= 1
+                        if state.eweight[e] == 0:
+                            state.g.remove_edge(e)
+                    else:
+                        state.g.remove_edge(e)
+                    old_es.append((u, v, l))
+
+            self.add_vertex(pos.keys(), pos.values())
+
+            Sf = self.entropy(**entropy_args)
+
+            self.remove_vertex(pos.keys())
+
+        finally:
+            if missing:
+                for e, l in new_es:
+                    if not l[1]:
+                        state = self.agg_state
+                    else:
+                        state = self.layer_states[l[0]]
+                    state.g.remove_edge(e)
+            else:
+                for u, v, l in old_es:
+                    if not l[1]:
+                        state = self.agg_state
+                    else:
+                        state = self.layer_states[l[0]]
+                    if state.is_weighted:
+                        e = state.g.edge(u, v)
+                        if e is None:
+                            e = state.g.add_edge(u, v)
+                            state.eweight[e] = 0
+                            if not l[1]:
+                                self.ec[e] = l[0]
+                        state.eweight[e] += 1
+                    else:
+                        e = state.g.add_edge(u, v)
+                        if not l[1]:
+                            self.ec[e] = l[0]
+            self.add_vertex(pos.keys(), pos.values())
+
+        if missing:
+            return Si - Sf
+        else:
+            return Sf - Si
+
     def _mcmc_sweep_dispatch(self, mcmc_state):
         if not self.overlap:
             return libinference.mcmc_layered_sweep(mcmc_state, self._state,

src/graph_tool/inference/nested_blockmodel.py

@@ -296,7 +296,7 @@ class NestedBlockState(object):
         """
         L = 0
         for l, state in enumerate(self.levels):
-            eargs = overlay(entropy_args, dl=True,
+            eargs = overlay(entropy_args,
                             edges_dl=(l == (len(self.levels) - 1)))
             if l > 0:
                 eargs = overlay(eargs, **self.hentropy_args)
@@ -501,6 +501,27 @@ class NestedBlockState(object):
         """
         return self._h_sweep(lambda s, **a: s.multicanonical_sweep(**a))
 
+    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.
+        """
+        S = 0
+        for l, lstate in enumerate(self.levels):
+            if l > 0:
+                eargs = overlay(self.hentropy_args,
+                                edges_dl=(l == len(self.levels) - 1))
+            else:
+                eargs = entropy_args
+            S += lstate.get_edges_prob(edge_list, missing, entropy_args=eargs)
+            if isinstance(self.levels[0], LayeredBlockState):
+                edge_list = [(lstate.b[u], lstate.b[v], l) for u, v, l in edge_list]
+            else:
+                edge_list = [(lstate.b[u], lstate.b[v]) for u, v in edge_list]
+        return S
+
     def draw(self, **kwargs):
         r"""Convenience wrapper to :func:`~graph_tool.draw.draw_hierarchy` that
         draws the hierarchical state."""