Commit bd375efd authored by Tiago Peixoto's avatar Tiago Peixoto

Improve checkpointing in blockmodel.py

parent bc4b3707
...@@ -943,34 +943,63 @@ def mc_get_dl(state, nsweep, greedy, rng, checkpoint, checkpoint_state, ...@@ -943,34 +943,63 @@ def mc_get_dl(state, nsweep, greedy, rng, checkpoint, checkpoint_state,
if verbose: if verbose:
print("beta = %g" % beta) print("beta = %g" % beta)
min_dl = S delta, nmoves = mcmc_sweep(state, beta=beta)
count = 0 S += delta
while True:
delta, nmoves = mcmc_sweep(state, beta=float("inf")) if S < min_dl:
min_dl = S
count = 0
elif S > max_dl:
max_dl = S
count = 0
else:
count += 1
checkpoint_state[B]["S"] = S
checkpoint_state[B]["min_dl"] = min_dl
checkpoint_state[B]["max_dl"] = max_dl
checkpoint_state[B]["count"] = count
if checkpoint is not None: if checkpoint is not None:
checkpoint(state, S, delta, nmoves) checkpoint(state, S, delta, nmoves, checkpoint_state)
if verbose:
print("beta = inf")
if not greedy:
checkpoint_state[B]["greedy"] = True
min_dl = S
count = 0
while count <= abs(nsweep):
delta, nmoves = mcmc_sweep(state, beta=float("inf"))
S += delta S += delta
if S < min_dl: if S < min_dl:
min_dl = S min_dl = S
count = 0 count = 0
else: else:
count += 1 count += 1
if count > abs(nsweep): checkpoint_state[B]["S"] = S
break checkpoint_state[B]["min_dl"] = min_dl
checkpoint_state[B]["count"] = count
if checkpoint is not None:
checkpoint(state, S, delta, nmoves, checkpoint_state)
return state._BlockState__min_dl() return state._BlockState__min_dl()
def get_b_dl(g, bs, bs_start, B, nsweep, anneal, greedy, clabel, deg_corr, rng, def get_b_dl(g, bs, B, nsweep, anneal, greedy, clabel, deg_corr, rng,
checkpoint=None, verbose=False): checkpoint=None, checkpoint_state=None, verbose=False):
prev_dl = float("inf") if B not in checkpoint_state:
if B in bs: checkpoint_state[B] = {}
if B in bs and checkpoint_state[B].get("done", False):
return bs[B][0] return bs[B][0]
elif B in bs_start: elif B in bs:
if verbose: if verbose:
print("starting from previous result for B=%d" % B) print("starting from previous result for B=%d" % B)
prev_dl, b = bs_start[B] b = bs[B][1]
state = BlockState(g, b=b.copy(), clabel=clabel, deg_corr=deg_corr) state = BlockState(g, b=b.copy(), clabel=clabel, deg_corr=deg_corr)
else: else:
checkpoint_state[B] = {}
n_iter = 0 n_iter = 0
bs_keys = [k for k in bs.keys() if type(k) != str] bs_keys = [k for k in bs.keys() if type(k) != str]
B_sup = max(bs_keys) if len(bs_keys) > 0 else B B_sup = max(bs_keys) if len(bs_keys) > 0 else B
...@@ -994,8 +1023,9 @@ def get_b_dl(g, bs, bs_start, B, nsweep, anneal, greedy, clabel, deg_corr, rng, ...@@ -994,8 +1023,9 @@ def get_b_dl(g, bs, bs_start, B, nsweep, anneal, greedy, clabel, deg_corr, rng,
bg_state = BlockState(cg, B=B, clabel=blabel, vweight=vcount, eweight=ecount, bg_state = BlockState(cg, B=B, clabel=blabel, vweight=vcount, eweight=ecount,
deg_corr=deg_corr) deg_corr=deg_corr)
mc_get_dl(bg_state, nsweep=nsweep, greedy=greedy, rng=rng, dl = mc_get_dl(bg_state, nsweep=nsweep, greedy=greedy, rng=rng,
checkpoint=checkpoint, anneal=anneal, verbose=verbose) checkpoint=None, checkpoint_state=None,
anneal=anneal, verbose=verbose)
### FIXME: the following could be improved by moving it to the C++ ### FIXME: the following could be improved by moving it to the C++
### side ### side
...@@ -1005,16 +1035,16 @@ def get_b_dl(g, bs, bs_start, B, nsweep, anneal, greedy, clabel, deg_corr, rng, ...@@ -1005,16 +1035,16 @@ def get_b_dl(g, bs, bs_start, B, nsweep, anneal, greedy, clabel, deg_corr, rng,
for v in g.vertices(): for v in g.vertices():
b[v] = bg_state.b[bmap[b[v]]] b[v] = bg_state.b[bmap[b[v]]]
checkpoint_state[B] = {}
bs[B] = [dl, b.copy()]
state = BlockState(g, b=b, B=B, clabel=clabel, deg_corr=deg_corr) state = BlockState(g, b=b, B=B, clabel=clabel, deg_corr=deg_corr)
dl = mc_get_dl(state, nsweep=nsweep, greedy=greedy, rng=rng, dl = mc_get_dl(state, nsweep=nsweep, greedy=greedy, rng=rng,
checkpoint=checkpoint, anneal=anneal, verbose=verbose) checkpoint=checkpoint, checkpoint_state=checkpoint_state,
anneal=anneal, verbose=verbose)
if dl < prev_dl: bs[B] = [dl, state.b.copy()]
bs[B] = [dl, state.b.copy()] checkpoint_state[B]["done"] = True
else:
bs[B] = bs_start[B]
dl = prev_dl
return dl return dl
def fibo(n): def fibo(n):
...@@ -1035,8 +1065,8 @@ def is_fibo(x): ...@@ -1035,8 +1065,8 @@ def is_fibo(x):
def minimize_blockmodel_dl(g, deg_corr=True, nsweeps=100, adaptive_convergence=True, def minimize_blockmodel_dl(g, deg_corr=True, nsweeps=100, adaptive_convergence=True,
anneal=1., greedy_cooling=True, max_B=None, min_B=1, anneal=1., greedy_cooling=True, max_B=None, min_B=1,
mid_B=None, b_cache=None, b_start=None, clabel=None, clabel=None, mid_B=None, b_cache=None, checkpoint=None,
checkpoint=None, verbose=False): checkpoint_state=None, verbose=False):
r"""Find the block partition of an unspecified size which minimizes the description r"""Find the block partition of an unspecified size which minimizes the description
length of the network, according to the stochastic blockmodel ensemble which length of the network, according to the stochastic blockmodel ensemble which
best describes it. best describes it.
...@@ -1072,6 +1102,9 @@ def minimize_blockmodel_dl(g, deg_corr=True, nsweeps=100, adaptive_convergence=T ...@@ -1072,6 +1102,9 @@ def minimize_blockmodel_dl(g, deg_corr=True, nsweeps=100, adaptive_convergence=T
mid_B : ``int`` (optional, default: ``None``) mid_B : ``int`` (optional, default: ``None``)
Middle of the range which brackets the minimum. If not supplied, will be Middle of the range which brackets the minimum. If not supplied, will be
automatically determined. automatically determined.
clabel : :class:`~graph_tool.PropertyMap` (optional, default: ``None``)
Constraint labels on the vertices, such that vertices with different
labels cannot belong to the same block.
b_cache : :class:`dict` with ``int`` keys and (``float``, :class:`~graph_tool.PropertyMap`) values (optional, default: ``None``) b_cache : :class:`dict` with ``int`` keys and (``float``, :class:`~graph_tool.PropertyMap`) values (optional, default: ``None``)
If provided, this corresponds to a dictionary where the keys are the If provided, this corresponds to a dictionary where the keys are the
number of blocks, and the values are tuples containing two values: the number of blocks, and the values are tuples containing two values: the
...@@ -1079,12 +1112,6 @@ def minimize_blockmodel_dl(g, deg_corr=True, nsweeps=100, adaptive_convergence=T ...@@ -1079,12 +1112,6 @@ def minimize_blockmodel_dl(g, deg_corr=True, nsweeps=100, adaptive_convergence=T
in this dictionary will not be computed, and will be used unmodified as in this dictionary will not be computed, and will be used unmodified as
the solution for the corresponding number of blocks. This can be used to the solution for the corresponding number of blocks. This can be used to
continue from a previously unfinished run. continue from a previously unfinished run.
b_start : :class:`dict` with ``int`` keys and (``float``, :class:`~graph_tool.PropertyMap`) values (optional, default: ``None``)
Like `b_cache`, but the partitions present in the dictionary will be
used as the starting point of the minimization.
clabel : :class:`~graph_tool.PropertyMap` (optional, default: ``None``)
Constraint labels on the vertices, such that vertices with different
labels cannot belong to the same block.
checkpoint : function (optional, default: ``None``) checkpoint : function (optional, default: ``None``)
If provided, this function will be called after each call to If provided, this function will be called after each call to
:func:`mcmc_sweep`. This can be used to store the current state, so it :func:`mcmc_sweep`. This can be used to store the current state, so it
...@@ -1092,17 +1119,26 @@ def minimize_blockmodel_dl(g, deg_corr=True, nsweeps=100, adaptive_convergence=T ...@@ -1092,17 +1119,26 @@ def minimize_blockmodel_dl(g, deg_corr=True, nsweeps=100, adaptive_convergence=T
.. code-block:: python .. code-block:: python
def checkpoint(state, L, delta, nmoves): def checkpoint(state, L, delta, nmoves, checkpoint_state):
... ...
where `state` is either a :class:`~graph_tool.community.BlockState` where `state` is either a :class:`~graph_tool.community.BlockState`
instance or ``None``, `L` is the current description length, `delta` is instance or ``None``, `L` is the current description length, `delta` is
the entropy difference in the last MCMC sweep, and `nmoves` is the the entropy difference in the last MCMC sweep, and `nmoves` is the
number of accepted block membership moves. number of accepted block membership moves. The ``checkpoint_state``
argument is an opaque object which specifies the current state of the
algorithm, which can be stored via :mod:`pickle`, and supplied via the
``checkpoint_state`` option below to continue from an interrupted run.
This function will also be called when the MCMC has finished for the This function will also be called when the MCMC has finished for the
current value of :math:`B`, in which case ``state == None``, and the current value of :math:`B`, in which case ``state == None``, and the
remaining parameters will be zero. remaining parameters will be zero, except the last.
checkpoint_state : object (optional, default: ``None``)
If provided, this will specify an exact point of execution from which
the algorithm will continue. The expected object is an opaque type which
wiil be passed to the callback of the ``checkpoint`` option above, and
can be stored by :mod:`pickle`. This must be used in conjunction with
the option ``b_cache`` to continue from an interrupted run.
verbose : ``bool`` (optional, default: ``False``) verbose : ``bool`` (optional, default: ``False``)
If ``True``, verbose information is displayed. If ``True``, verbose information is displayed.
...@@ -1208,25 +1244,25 @@ def minimize_blockmodel_dl(g, deg_corr=True, nsweeps=100, adaptive_convergence=T ...@@ -1208,25 +1244,25 @@ def minimize_blockmodel_dl(g, deg_corr=True, nsweeps=100, adaptive_convergence=T
greedy = greedy_cooling greedy = greedy_cooling
if b_start is None:
b_start = {}
bs = b_cache bs = b_cache
if bs is None: if bs is None:
bs = {} bs = {}
if checkpoint_state is None:
checkpoint_state = {}
while True: while True:
f_max = get_b_dl(g, bs, b_start, max_B, nsweeps, anneal, greedy, clabel, f_max = get_b_dl(g, bs, max_B, nsweeps, anneal, greedy, clabel,
deg_corr, rng, checkpoint, verbose) deg_corr, rng, checkpoint, checkpoint_state, verbose)
f_mid = get_b_dl(g, bs, b_start, mid_B, nsweeps, anneal, greedy, clabel, f_mid = get_b_dl(g, bs, mid_B, nsweeps, anneal, greedy, clabel,
deg_corr, rng, checkpoint, verbose) deg_corr, rng, checkpoint, checkpoint_state, verbose)
f_min = get_b_dl(g, bs, b_start, min_B, nsweeps, anneal, greedy, clabel, f_min = get_b_dl(g, bs, min_B, nsweeps, anneal, greedy, clabel,
deg_corr, rng, checkpoint, verbose) deg_corr, rng, checkpoint, checkpoint_state, verbose)
if verbose: if verbose:
print("bracket:", min_B, mid_B, max_B, f_min, f_mid, f_max) print("bracket:", min_B, mid_B, max_B, f_min, f_mid, f_max)
if checkpoint is not None: if checkpoint is not None:
checkpoint(None, 0, 0, 0) checkpoint(None, 0, 0, 0, checkpoint_state)
if f_max > f_mid > f_min: if f_max > f_mid > f_min:
max_B = mid_B max_B = mid_B
...@@ -1248,10 +1284,10 @@ def minimize_blockmodel_dl(g, deg_corr=True, nsweeps=100, adaptive_convergence=T ...@@ -1248,10 +1284,10 @@ def minimize_blockmodel_dl(g, deg_corr=True, nsweeps=100, adaptive_convergence=T
else: else:
x = get_mid(min_B, mid_B) x = get_mid(min_B, mid_B)
f_x = get_b_dl(g, bs, b_start, x, nsweeps, anneal, greedy, clabel, f_x = get_b_dl(g, bs, x, nsweeps, anneal, greedy, clabel,
deg_corr, rng, checkpoint, verbose) deg_corr, rng, checkpoint, checkpoint_state, verbose)
f_mid = get_b_dl(g, bs, b_start, mid_B, nsweeps, anneal, greedy, clabel, f_mid = get_b_dl(g, bs, mid_B, nsweeps, anneal, greedy, clabel,
deg_corr, rng, checkpoint, verbose) deg_corr, rng, checkpoint, checkpoint_state, verbose)
if verbose: if verbose:
print("bisect: (", min_B, mid_B, max_B, ") ->", x, f_x) #, is_fibo((mid_B - min_B)), is_fibo((max_B - mid_B))) print("bisect: (", min_B, mid_B, max_B, ") ->", x, f_x) #, is_fibo((mid_B - min_B)), is_fibo((max_B - mid_B)))
...@@ -1268,7 +1304,7 @@ def minimize_blockmodel_dl(g, deg_corr=True, nsweeps=100, adaptive_convergence=T ...@@ -1268,7 +1304,7 @@ def minimize_blockmodel_dl(g, deg_corr=True, nsweeps=100, adaptive_convergence=T
return bs[best_B][1], bs[best_B][0], bs return bs[best_B][1], bs[best_B][0], bs
if checkpoint is not None: if checkpoint is not None:
checkpoint(None, 0, 0, 0) checkpoint(None, 0, 0, 0, checkpoint_state)
if f_x < f_mid: if f_x < f_mid:
if max_B - mid_B > mid_B - min_B: if max_B - mid_B > mid_B - min_B:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment