Slight improvements to Wang-Landau sampling

src/graph/inference/graph_blockmodel_multicanonical.hh

@@ -49,6 +49,7 @@ using namespace std;
     ((entropy_args,, entropy_args_t, 0))                                       \
     ((allow_vacate,, bool, 0))                                                 \
     ((verbose,, bool, 0))                                                      \
+    ((target_bin,, int, 0))                                                 \
     ((niter,, size_t, 0))
 
 
@@ -82,6 +83,12 @@ struct Multicanonical
 
         typename state_t::g_t& _g;
 
+        int get_bin(double S)
+        {
+            return round((_hist.size() - 1) *
+                         ((S - _S_min) / (_S_max - _S_min)));
+        };
+
         size_t node_state(size_t v)
         {
             return _state._b[v];

src/graph/inference/multicanonical_loop.hh

@@ -47,15 +47,8 @@ auto multicanonical_sweep(MulticanonicalState& state, RNG& rng)
     auto& hist = state._hist;
     auto& dens = state._dens;
     int M = hist.size();
-    double S_min = state._S_min;
-    double S_max = state._S_max;
 
-    auto get_bin = [&](double x) -> int
-        {
-            return round((M - 1) * (x - S_min) / (S_max - S_min));
-        };
-
-    int i = get_bin(S);
+    int i = state.get_bin(S);
 
     if (i < 0 || i >= M)
         throw ValueException("current state lies outside the allowed entropy range");
@@ -71,7 +64,7 @@ auto multicanonical_sweep(MulticanonicalState& state, RNG& rng)
 
         std::pair<double, double> dS = state.virtual_move_dS(v, s);
 
-        int j = get_bin(S + dS.first);
+        int j = state.get_bin(S + dS.first);
 
         bool accept;
         if (j < 0 || j >= M)
@@ -107,6 +100,9 @@ auto multicanonical_sweep(MulticanonicalState& state, RNG& rng)
         state._time += 1./M;
         if (state._refine)
             state._f *= (state._time - 1. / M) / state._time;
+
+        if (i == state._target_bin)
+            break;
     }
     return make_pair(S, nmoves);
 }

src/graph_tool/inference/blockmodel.py

@@ -1499,7 +1499,7 @@ class BlockState(object):
 
     def multicanonical_sweep(self, m_state, c=numpy.inf, niter=1,
                              entropy_args={}, allow_vacate=True, vertices=None,
-                             verbose=False, **kwargs):
+                             target_bin=-1, verbose=False, **kwargs):
         r"""Perform ``niter`` sweeps of a non-Markovian multicanonical sampling using the
         Wang-Landau algorithm.
 
@@ -1525,6 +1525,8 @@ class BlockState(object):
         vertices : ``list`` of ints (optional, default: ``None``)
             If provided, this should be a list of vertices which will be
             moved. Otherwise, all vertices will.
+        target_bin : ``int`` (optional, default: ``-1``)
+            If this energy bin is reached, stop sweep prematurely.
         verbose : ``bool`` (optional, default: ``False``)
             If ``verbose == True``, detailed information will be displayed.
 
@@ -1579,6 +1581,7 @@ class BlockState(object):
         multi_state.S_max = m_state._S_max
         multi_state.hist = m_state._hist
         multi_state.dens = m_state._density
+        multi_state.target_bin = target_bin
 
         S, nmoves, f, time = \
                 self._multicanonical_sweep_dispatch(multi_state)

src/graph_tool/inference/mcmc.py

@@ -456,10 +456,13 @@ class MulticanonicalState(object):
 
     def get_entropy(self, S, B=None):
         r = self.get_density(B)
-        dS = (self._S_max - self._S_min) / len(r)
-        j = int(round((S - self._S_min) / dS))
+        j = self.get_bin()
         return r[j]
 
+    def get_bin(self, S):
+        return int(round((len(self._hist) - 1) * ((S - self._S_min) /
+                                                  (self._S_max - self._S_min))))
+
     def get_hist(self):
         "Get energy histogram."
         return numpy.array(self._hist.a)
@@ -585,7 +588,7 @@ def multicanonical_equilibrate(state, m_state, f_range=(1., 1e-6),
             print(verbose_pad(verbose) +
                   "count: %d  time: %#8.8g  f: %#8.8g  flatness: %#8.8g  S: %#8.8g" % \
                   (count, m_state._time, m_state._f, hf,
-                   state.entropy(multicanonical_args.get("entropy_args", {}))))
+                   state.entropy(**multicanonical_args.get("entropy_args", {}))))
 
         if not m_state._refine:
             if hf > flatness: