inference.blockmodel: implement "waiting times"

src/graph/inference/graph_blockmodel.hh

@@ -56,6 +56,9 @@ typedef mpl::vector2<simple_degs_t, degs_map_t> degs_tr;
     ((pclabel,, vmap_t, 0))                                                    \
     ((merge_map,, vmap_t, 0))                                                  \
     ((deg_corr,, bool, 0))                                                     \
+    ((use_waiting,, bool, 0))                                                  \
+    ((vtfield,, vprop_map_t<double>::type, 0))                                 \
+    ((btfield,, vprop_map_t<double>::type, 0))                                 \
     ((ignore_degrees,, typename vprop_map_t<uint8_t>::type, 0))
 
 GEN_STATE_BASE(BlockStateBase, BLOCK_STATE_params)
@@ -795,6 +798,22 @@ public:
                 dS += ps.get_delta_edges_dl(v, r, nr, _vweight, _g);
         }
 
+        if (_use_waiting && (r != nr) && _ignore_degrees[v] > 0)
+        {
+            double dt = _vtfield[v];
+            int k = out_degreeS()(v, _g, _eweight);
+            if (_mrp[r] > 0)
+                dS -= _mrp[r] * log(_btfield[r]) - lgamma_fast(_mrp[r]);
+            if (_mrp[nr] > 0)
+                dS -= _mrp[nr] * log(_btfield[nr]) - lgamma_fast(_mrp[nr]);
+            if (_mrp[r] > k)
+                dS += (_mrp[r] - k) * log(_btfield[r] - dt)
+                    - lgamma_fast(_mrp[r] - k);
+            if (_mrp[nr] + k > 0)
+                dS += (_mrp[nr] + k) * log(_btfield[nr] + dt)
+                    - lgamma_fast(_mrp[nr] + k);
+        }
+
         return dS;
     }
 
@@ -1053,10 +1072,19 @@ public:
 
     double entropy(bool dense, bool multigraph, bool deg_entropy)
     {
+        double S = 0;
         if (dense)
-            return dense_entropy(multigraph);
+            S = dense_entropy(multigraph);
         else
-            return sparse_entropy(multigraph, deg_entropy);
+            S = sparse_entropy(multigraph, deg_entropy);
+
+        if (_use_waiting)
+        {
+            for (auto r : vertices_range(_bg))
+                if (_btfield[r] > 0)
+                    S += _mrp[r] * log(_btfield[r]) - lgamma_fast(_mrp[r]);
+        }
+        return S;
     }
 
     double get_partition_dl()

src/graph_tool/inference/blockmodel.py

@@ -49,17 +49,25 @@ def _bm_test():
     global __test__
     return __test__
 
-def get_block_graph(g, B, b, vcount=None, ecount=None):
+def get_block_graph(g, B, b, vcount=None, ecount=None, edt=None):
     if isinstance(ecount, libinference.unity_eprop_t):
         ecount = None
     if isinstance(vcount, libinference.unity_vprop_t):
         vcount = None
-    cg, br, vcount, ecount = condensation_graph(g, b,
-                                                vweight=vcount,
-                                                eweight=ecount,
-                                                self_loops=True)[:4]
-    cg.vp["count"] = vcount
-    cg.ep["count"] = ecount
+    aeprops = []
+    if edt is not None:
+        aeprops.append(edt)
+    cg, br, vcount, ecount, av, ae = condensation_graph(g, b,
+                                                        vweight=vcount,
+                                                        eweight=ecount,
+                                                        aeprops=aeprops,
+                                                        self_loops=True)
+    cg.vp.count = vcount
+    cg.ep.count = ecount
+    if edt is not None:
+        cg.ep.dt = ae[0]
+        del ae[0]
+
     cg = Graph(cg, vorder=br)
 
     cg.add_vertex(B - cg.num_vertices())
@@ -137,7 +145,7 @@ class BlockState(object):
         if b is None:
             # create a random partition into B blocks.
             if B is None:
-                B = get_max_B(self.N, self.E, directed=g.is_directed())
+                raise ValueError("either 'b' or 'B' must be specified")
             B = min(B, self.g.num_vertices())
             ba = random.randint(0, B, self.g.num_vertices())
             ba[:B] = arange(B)        # avoid empty blocks
@@ -161,7 +169,8 @@ class BlockState(object):
                              (self.b.fa.max(), B))
 
         # Construct block-graph
-        self.bg = get_block_graph(g, B, self.b, self.vweight, self.eweight)
+        self.bg = get_block_graph(g, B, self.b, self.vweight, self.eweight,
+                                  edt=kwargs.get("edt", None))
         self.bg.set_fast_edge_removal()
 
         self.mrs = self.bg.ep["count"]
@@ -223,6 +232,19 @@ class BlockState(object):
         self.block_list = Vector_size_t()
         self.block_list.extend(arange(self.B, dtype="int"))
 
+        self.edt = extract_arg(kwargs, "edt", None)
+        self.use_waiting = self.edt is not None
+        if self.use_waiting:
+            self.vtfield = self.g.degree_property_map("out", self.edt)
+            self.btfield = self.bg.degree_property_map("out", self.bg.ep.dt)
+            tokens = self.ignore_degrees.fa == 0
+            token_groups = bincount(self.b.fa[tokens]) > 0
+            token_groups = token_groups.resize(self.bg.num_vertices())
+            self.btfield.a[token_groups] = 0
+        else:
+            self.vtfield = self.g.new_vp("double")
+            self.btfield = self.bg.new_vp("double")
+
         self._abg = self.bg._get_any()
         self._state = libinference.make_block_state(self, _get_rng())
 
@@ -262,6 +284,7 @@ class BlockState(object):
                                ignore_degrees=kwargs.pop("ignore_degrees", self.ignore_degrees),
                                degs=self.degs.copy(),
                                merge_map=kwargs.get("merge_map", self.merge_map.copy()),
+                               edt=self.edt,
                                **dmask(kwargs, ["ignore_degrees", "merge_map"]))
         else:
             state = OverlapBlockState(self.g if g is None else g,