Improve memory performance of collect_edge_marginals() and bethe_entropy()

src/graph/inference/Makefile.am

@@ -34,6 +34,7 @@ libgraph_tool_inference_la_SOURCES = \
     graph_blockmodel_layers_overlap_gibbs.cc \
     graph_blockmodel_layers_overlap_multicanonical.cc \
     graph_blockmodel_layers_overlap_vacate.cc \
+    graph_blockmodel_marginals.cc \
     graph_blockmodel_mcmc.cc \
     graph_blockmodel_multicanonical.cc \
     graph_blockmodel_merge.cc \

src/graph/inference/graph_blockmodel.cc

@@ -141,103 +141,6 @@ simple_degs_t copy_simple_degs(simple_degs_t& degs)
     return degs;
 }
 
-double mf_entropy(GraphInterface& gi, boost::any opv)
-{
-    double H=0;
-    run_action<>()
-        (gi,
-         [&](auto& g, auto pv)
-         {
-             for (auto v : vertices_range(g))
-             {
-                 double sum = 0;
-                 for (auto p : pv[v])
-                     sum += p;
-                 for (double p : pv[v])
-                 {
-                     if (p == 0)
-                         continue;
-                     p /= sum;
-                     H -= p * log(p);
-                 }
-             }
-         },
-         vertex_scalar_vector_properties())(opv);
-
-    return H;
-}
-
-boost::python::tuple bethe_entropy(GraphInterface& gi, size_t B, boost::any op,
-                                   boost::any opv)
-{
-    typedef vprop_map_t<vector<double>>::type vmap_t;
-    vmap_t pv = any_cast<vmap_t>(opv);
-
-    double H=0, sH=0, Hmf=0, sHmf=0;
-    run_action<>()
-        (gi,
-         [&](auto& g, auto p)
-         {
-             for (auto v : vertices_range(g))
-             {
-                 pv[v].resize(B);
-                 for (size_t i = 0; i < B; ++i)
-                     pv[v][i] = 0;
-             }
-
-             for (auto e : edges_range(g))
-             {
-                 auto u = min(source(e, g), target(e, g));
-                 auto v = max(source(e, g), target(e, g));
-
-                 double sum = 0;
-                 for (size_t r = 0; r < B; ++r)
-                     for (size_t s = 0; s < B; ++s)
-                     {
-                         size_t i = r + B * s;
-                         pv[u][r] += p[e][i];
-                         pv[v][s] += p[e][i];
-                         sum += p[e][i];
-                     }
-
-                 for (size_t i = 0; i < B * B; ++i)
-                 {
-                     if (p[e][i] == 0)
-                         continue;
-                     double pi = double(p[e][i]) / sum;
-                     H -= pi * log(pi);
-                     sH += pow((log(pi) + 1) * sqrt(pi / sum), 2);
-                 }
-             }
-
-             for (auto v : vertices_range(g))
-             {
-                 double sum = 0;
-                 for (size_t i = 0; i < B; ++i)
-                     sum += pv[v][i];
-                 for (size_t i = 0; i < B; ++i)
-                 {
-                     if (pv[v][i] == 0)
-                         continue;
-                     pv[v][i] /= sum;
-                     double pi = pv[v][i];
-                     double kt = (1 - double(in_degreeS()(v, g)) - double(out_degree(v, g)));
-                     if (kt != 0)
-                     {
-                         H -= kt * (pi * log(pi));
-                         sH += pow(kt * (log(pi) + 1) * sqrt(pi / sum), 2);
-                     }
-
-                     Hmf -= pi * log(pi);
-                     sHmf += pow((log(pi) + 1) * sqrt(pi / sum), 2);
-                 }
-             }
-         },
-         edge_scalar_vector_properties())(op);
-
-    return boost::python::make_tuple(H, sH, Hmf, sHmf);
-}
-
 void export_blockmodel_state()
 {
     using namespace boost::python;
@@ -335,9 +238,6 @@ void export_blockmodel_state()
     class_<simple_degs_t>("simple_degs_t")
         .def("copy", &copy_simple_degs);
 
-    def("mf_entropy", &mf_entropy);
-    def("bethe_entropy", &bethe_entropy);
-
     def("init_q_cache", init_q_cache);
     def("log_q", log_q<size_t>);
     def("q_rec", q_rec);

src/graph/inference/graph_blockmodel_marginals.cc

+// graph-tool -- a general graph modification and manipulation thingy
+//
+// Copyright (C) 2006-2016 Tiago de Paula Peixoto <tiago@skewed.de>
+//
+// This program is free software; you can redistribute it and/or
+// modify it under the terms of the GNU General Public License
+// as published by the Free Software Foundation; either version 3
+// of the License, or (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+//
+// You should have received a copy of the GNU General Public License
+// along with this program. If not, see <http://www.gnu.org/licenses/>.
+
+#include "graph_tool.hh"
+
+#include <boost/python.hpp>
+#include "numpy_bind.hh"
+#include "hash_map_wrap.hh"
+
+using namespace std;
+using namespace boost;
+using namespace graph_tool;
+
+void collect_vertex_marginals(GraphInterface& gi, boost::any ob,
+                              boost::any op, double update)
+{
+    typedef vprop_map_t<int32_t>::type vmap_t;
+    auto b = any_cast<vmap_t>(ob).get_unchecked();
+
+    run_action<>()
+        (gi, [&](auto& g, auto p)
+         {
+             typename property_traits<decltype(p)>::value_type::value_type
+                 up = update;
+             parallel_vertex_loop
+                 (g,
+                  [&](auto v)
+                  {
+                      auto r = b[v];
+                      auto& pv = p[v];
+                      if (pv.size() <= size_t(r))
+                          pv.resize(r + 1);
+                      pv[r] += up;
+                  });
+         },
+         vertex_scalar_vector_properties())(op);
+}
+
+class BlockPairHist
+    : public gt_hash_map<std::pair<int32_t,int32_t>,double>
+{
+public:
+
+    boost::python::dict get_state()
+    {
+        boost::python::dict state;
+        for (auto& kv : *this)
+        {
+            auto k = python::make_tuple(kv.first.first,
+                                        kv.first.second);
+            state[k] = kv.second;
+        }
+        return state;
+    }
+
+    void set_state(boost::python::dict state)
+    {
+        auto keys = state.keys();
+        for (int i = 0; i < python::len(keys); ++i)
+        {
+            auto k = keys[i];
+            int32_t r = python::extract<int32_t>(k[0]);
+            int32_t s = python::extract<int32_t>(k[1]);
+            double v = python::extract<double>(state[k]);
+            (*this)[make_pair(r, s)] = v;
+        }
+    }
+
+    double get_item(boost::python::object k)
+    {
+        int32_t r = python::extract<int32_t>(k[0]);
+        int32_t s = python::extract<int32_t>(k[1]);
+        auto iter = this->find(make_pair(r, s));
+        if (iter == this->end())
+            return 0;
+        return iter->second;
+    }
+
+    void set_item(boost::python::object k, double v)
+    {
+        int32_t r = python::extract<int32_t>(k[0]);
+        int32_t s = python::extract<int32_t>(k[1]);
+        (*this)[make_pair(r, s)] = v;
+    }
+
+};
+
+void collect_edge_marginals(GraphInterface& gi, boost::any ob,
+                            boost::any op, double update)
+{
+    typedef vprop_map_t<int32_t>::type vmap_t;
+    auto b = any_cast<vmap_t>(ob).get_unchecked();
+
+    typedef eprop_map_t<python::object>::type
+        emap_t;
+    auto pe = any_cast<emap_t>(op).get_unchecked(gi.get_edge_index_range());
+
+    run_action<>()
+        (gi,
+         [&](auto& g)
+         {
+            parallel_edge_loop
+                 (g,
+                  [&](const auto& e)
+                  {
+                      auto u = min(source(e, g), target(e, g));
+                      auto v = max(source(e, g), target(e, g));
+
+                      auto r = b[u];
+                      auto s = b[v];
+
+                      BlockPairHist& h =
+                          boost::python::extract<BlockPairHist&>(pe[e]);
+
+                      h[make_pair(r, s)] += update;
+                  });
+         })();
+}
+
+double mf_entropy(GraphInterface& gi, boost::any opv)
+{
+    double H=0;
+    run_action<>()
+        (gi,
+         [&](auto& g, auto pv)
+         {
+             for (auto v : vertices_range(g))
+             {
+                 double sum = 0;
+                 for (auto p : pv[v])
+                     sum += p;
+                 for (double p : pv[v])
+                 {
+                     if (p == 0)
+                         continue;
+                     p /= sum;
+                     H -= p * log(p);
+                 }
+             }
+         },
+         vertex_scalar_vector_properties())(opv);
+
+    return H;
+}
+
+boost::python::tuple bethe_entropy(GraphInterface& gi, boost::any op,
+                                   boost::any opv)
+{
+    typedef vprop_map_t<vector<double>>::type vmap_t;
+    vmap_t pv = any_cast<vmap_t>(opv);
+
+    typedef eprop_map_t<boost::python::object>::type emap_t;
+    auto pe = any_cast<emap_t>(op).get_unchecked();
+
+    double H=0, Hmf=0;
+    run_action<>()
+        (gi,
+         [&](auto& g)
+         {
+             for (auto e : edges_range(g))
+             {
+                 auto u = min(source(e, g), target(e, g));
+                 auto v = max(source(e, g), target(e, g));
+
+                 double sum = 0;
+
+                 BlockPairHist& h =
+                     boost::python::extract<BlockPairHist&>(pe[e]);
+
+                 for (auto& prs : h)
+                 {
+                     sum += prs.second;
+                     size_t r = prs.first.first;
+                     size_t s = prs.first.second;
+                     if (r >= pv[u].size())
+                         pv[u].resize(r + 1);
+                     pv[u][r] += prs.second;
+                     if (s >= pv[v].size())
+                         pv[v].resize(s + 1);
+                     pv[v][s] += prs.second;
+                 }
+
+                 for (auto& prs : h)
+                 {
+                     if (prs.second == 0)
+                         continue;
+                     double pi = prs.second / sum;
+                     H -= pi * log(pi);
+                 }
+             }
+
+             for (auto v : vertices_range(g))
+             {
+                 double sum = std::accumulate(pv[v].begin(), pv[v].end(), 0);
+                 for (double pi : pv[v])
+                 {
+                     if (pi == 0)
+                         continue;
+                     pi /= sum;
+                     int kt = 1 - total_degreeS()(v, g);
+                     if (kt != 0)
+                         H -= kt * (pi * log(pi));
+                     Hmf -= pi * log(pi);
+                 }
+             }
+         })();
+
+    return boost::python::make_tuple(H, Hmf);
+}
+
+void export_marginals()
+{
+    using namespace boost::python;
+
+    class_<BlockPairHist>("BlockPairHist")
+        .def("__setitem__", &BlockPairHist::set_item)
+        .def("__getitem__", &BlockPairHist::get_item)
+        .def("__setstate__", &BlockPairHist::set_state)
+        .def("__getstate__", &BlockPairHist::get_state).enable_pickling();
+
+    def("vertex_marginals", &collect_vertex_marginals);
+    def("edge_marginals", &collect_edge_marginals);
+    def("mf_entropy", &mf_entropy);
+    def("bethe_entropy", &bethe_entropy);
+}

src/graph/inference/graph_inference.cc

@@ -25,64 +25,6 @@ using namespace std;
 using namespace boost;
 using namespace graph_tool;
 
-void collect_vertex_marginals(GraphInterface& gi, boost::any ob,
-                              boost::any op, double update)
-{
-    typedef vprop_map_t<int32_t>::type vmap_t;
-    auto b = any_cast<vmap_t>(ob).get_unchecked();
-
-    run_action<>()
-        (gi, [&](auto& g, auto p)
-         {
-             typename property_traits<decltype(p)>::value_type::value_type
-                 up = update;
-             parallel_vertex_loop
-                 (g,
-                  [&](auto v)
-                  {
-                      auto r = b[v];
-                      auto& pv = p[v];
-                      if (pv.size() <= size_t(r))
-                          pv.resize(r + 1);
-                      pv[r] += up;
-                  });
-         },
-         vertex_scalar_vector_properties())(op);
-}
-
-void collect_edge_marginals(GraphInterface& gi, size_t B, boost::any ob,
-                            boost::any op, double update)
-{
-    typedef vprop_map_t<int32_t>::type vmap_t;
-    auto b = any_cast<vmap_t>(ob).get_unchecked();
-
-    run_action<>()
-        (gi,
-         [&](auto& g, auto p)
-         {
-             typename property_traits<decltype(p)>::value_type::value_type
-                 up = update;
-             parallel_edge_loop
-                 (g,
-                  [&](const auto& e)
-                  {
-                      auto u = min(source(e, g), target(e, g));
-                      auto v = max(source(e, g), target(e, g));
-
-                      auto r = b[u];
-                      auto s = b[v];
-
-                      auto& pv = p[e];
-                      if (pv.size() < B * B)
-                          pv.resize(B * B);
-                      size_t j = r + B * s;
-                      pv[j] += up;
-                  });
-         },
-         edge_scalar_vector_properties())(op);
-}
-
-
 template <class Value>
 void vector_map(boost::python::object ovals, boost::python::object omap)
 {
@@ -157,6 +99,7 @@ extern void export_blockmodel_exhaustive();
 extern void export_overlap_blockmodel_exhaustive();
 extern void export_layered_blockmodel_exhaustive();
 extern void export_layered_overlap_blockmodel_exhaustive();
+extern void export_marginals();
 
 BOOST_PYTHON_MODULE(libgraph_tool_inference)
 {
@@ -189,9 +132,7 @@ BOOST_PYTHON_MODULE(libgraph_tool_inference)
     export_overlap_blockmodel_exhaustive();
     export_layered_blockmodel_exhaustive();
     export_layered_overlap_blockmodel_exhaustive();
-
-    def("vertex_marginals", collect_vertex_marginals);
-    def("edge_marginals", collect_edge_marginals);
+    export_marginals();
 
     def("vector_map", vector_map<int32_t>);
     def("vector_map64", vector_map<int64_t>);

src/graph_tool/inference/blockmodel.py

@@ -1572,21 +1572,16 @@ class BlockState(object):
         Parameters
         ----------
         p : :class:`~graph_tool.PropertyMap` (optional, default: ``None``)
-            Edge property map with vector-type values, storing the previous block
-            membership counts.  Each vector entry corresponds to ``b[i] + B *
-            b[j]``, where ``b`` is the block membership and ``i = min(source(e),
-            target(e))`` and ``j = max(source(e), target(e))``. If not provided, an
-            empty histogram will be created
+            Edge property map with edge marginals to be updated.  If not
+            provided, an empty histogram will be created.
         update : float (optional, default: ``1.``)
             Each call increases the current count by the amount given by this
             parameter.
 
         Returns
         -------
-        p : :class:`~graph_tool.PropertyMap` (optional, default: ``None``)
-            Vertex property map with vector-type values, storing the accumulated
-            block membership counts.
-
+        p : :class:`~graph_tool.libcore.any`
+            Edge property map with updated edge marginals.
 
         Examples
         --------
@@ -1605,15 +1600,16 @@ class BlockState(object):
            >>> for i in range(1000):
            ...     ds, nmoves = state.mcmc_sweep(niter=10)
            ...     pe = state.collect_edge_marginals(pe)
-           >>> gt.bethe_entropy(g, state.B, pe)[0]
+           >>> gt.bethe_entropy(g, pe)[0]
            6.65429696440...
         """
 
         if p is None:
-            p = self.g.new_ep("vector<double>")
+            p = self.g.new_ep("python::object",
+                              vals=[libinference.BlockPairHist()
+                                    for i in range(self.g.num_edges())])
 
         libinference.edge_marginals(self.g._Graph__graph,
-                                    self.B,
                                     _prop("v", self.g, self.b),
                                     _prop("e", self.g, p),
                                     update)
@@ -1787,20 +1783,15 @@ def model_entropy(B, N, E, directed=False, nr=None, allow_empty=True):
         L = lbinom(x + E - 1, E) + partition_entropy(B, N, nr, allow_empty)
     return L
 
-def bethe_entropy(g, B, p):
+def bethe_entropy(g, p):
     r"""Compute the Bethe entropy given the edge block membership marginals.
 
     Parameters
     ----------
     g : :class:`~graph_tool.Graph`
         The graph.
-    B : int
-        The number of blocks.
     p : :class:`~graph_tool.PropertyMap`
-        Edge property map with vector-type values, storing the previous block
-        membership counts.  Each vector entry corresponds to ``b[i] + B *
-        b[j]``, where ``b`` is the block membership and ``i = min(source(e),
-        target(e))`` and ``j = max(source(e), target(e))``.
+       Edge property map with edge marginals.
 
     Returns
     -------
@@ -1838,9 +1829,9 @@ def bethe_entropy(g, B, p):
     H = 0
     pv =  g.new_vertex_property("vector<double>")
 
-    H, sH, Hmf, sHmf  = libinference.bethe_entropy(g._Graph__graph, B,
-                                                   _prop("e", g, p),
-                                                   _prop("v", g, pv))
+    H, Hmf  = libinference.bethe_entropy(g._Graph__graph,
+                                         _prop("e", g, p),
+                                         _prop("v", g, pv))
     return H, Hmf, pv