nested_blockmodel.py

#! /usr/bin/env python
# -*- coding: utf-8 -*-
#
# graph_tool -- a general graph manipulation python module
#
# Copyright (C) 2006-2015 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/>.

from __future__ import division, absolute_import, print_function
import sys
if sys.version_info < (3,):
    range = xrange

from .. import _degree, _prop, Graph, GraphView, libcore, _get_rng, PropertyMap, \
    infect_vertex_property
from .. stats import label_self_loops
from .. generation import graph_union
from .. topology import shortest_path
import random
from numpy import *
import numpy
from scipy.optimize import fsolve, fminbound
import scipy.special
from collections import defaultdict
import copy

from . blockmodel import *
from . blockmodel import _bm_test

class NestedBlockState(object):
    r"""This class encapsulates the nested block state of a given graph.

    This must be instantiated and used by functions such as :func:`nested_mcmc_sweep`.

    The instances of this class contain a data member called ``levels``, which
    is a list of :class:`~graph_tool.community.BlockState` (or
    :class:`~graph_tool.community.OverlapBlockState`) instances, containing the
    entire nested hierarchy.

    Parameters
    ----------
    g : :class:`~graph_tool.Graph`
        Graph to be used.
    eweight : :class:`~graph_tool.PropertyMap` (optional, default: ``None``)
        Edge weights (i.e. multiplicity).
    vweight : :class:`~graph_tool.PropertyMap` (optional, default: ``None``)
        Vertex weights (i.e. multiplicity).
    bs : list of :class:`~graph_tool.PropertyMap` or :class:`~numpy.ndarray` instances (optional, default: ``None``)
        Initial block labels on the vertices, for each hierarchy level.
    Bs : list of ``int`` (optional, default: ``None``)
        Number of blocks for each hierarchy level.
    deg_corr : ``bool`` (optional, default: ``True``)
        If ``True``, the degree-corrected version of the blockmodel ensemble will
        be used in the bottom level, otherwise the traditional variant will be used.
    overlap : ``bool`` (optional, default: ``False``)
        If ``True``, the mixed-membership version of the blockmodel will be used
        at the lowest level.
    clabel : :class:`~graph_tool.PropertyMap` (optional, default: ``None``)
        Constraint labels on the vertices. If supplied, vertices with different
        label values will not be clustered in the same group.
    max_BE : ``int`` (optional, default: ``1000``)
        If the number of blocks exceeds this number, a sparse representation of
        the block graph is used, which is slightly less efficient, but uses less
        memory,
    """

    def __init__(self, g, eweight=None, vweight=None, ec=None, bs=None, Bs=None,
                 deg_corr=True, overlap=False, layers=False, clabel=None,
                 max_BE=1000, **kwargs):
        L = len(Bs) if Bs is not None else len(bs)
        self.g = cg = g
        self.vweight = vcount = vweight
        self.eweight = ecount = eweight
        self.ec = ec
        self.layers = layers

        self.levels = []
        self.overlap = overlap
        self.deg_corr = deg_corr
        self.clabel = clabel if clabel is not None else g.new_vertex_property("int")
        self.ignore_degrees = kwargs.get("ignore_degrees", None)

        for l in range(L):
            Bl = Bs[l] if Bs is not None else None
            bl = None
            if bs is not None:
                if isinstance(bs[l], PropertyMap):
                    bl = cg.own_property(bs[l])
                else:
                    bl = bs[l]

            if l == 0:
                if ec is None:
                    if overlap:
                        state = OverlapBlockState(g, B=Bl, b=bl,
                                                  eweight=ecount,
                                                  vweight=vcount,
                                                  deg_corr=deg_corr != False,
                                                  clabel=self.clabel,
                                                  max_BE=max_BE)
                        self.clabel = state.clabel.copy()
                        state.clabel.a = 0
                    else:
                        state = BlockState(g, B=Bl, b=bl,
                                           eweight=ecount,
                                           vweight=vcount,
                                           deg_corr=deg_corr != False,
                                           #clabel=self.clabel,
                                           max_BE=max_BE,
                                           ignore_degrees=self.ignore_degrees)
                else:
                    state = CovariateBlockState(g, B=Bl, b=bl,
                                                ec=ec,
                                                layers=layers,
                                                eweight=ecount,
                                                vweight=vcount,
                                                deg_corr=deg_corr != False,
                                                clabel=self.clabel,
                                                overlap=overlap,
                                                max_BE=max_BE)
                    if overlap:
                        self.clabel = state.clabel.copy()
                        state.clabel.a = 0


            else:
                state = self.levels[-1].get_block_state(b=bl,
                                                        overlap=self.overlap == "full",
                                                        deg_corr=self.deg_corr == "full")[0]
                if _bm_test():
                    assert not state.deg_corr, "upper levels must be non-deg-corr"

            self.levels.append(state)

        # for l in range(len(self.levels) - 1):
        #     clabel = self.__project_partition(l, l+1)
        #     self.levels[l].clabel = clabel

        if ec is not None:
            self.ec = self.levels[0].ec

    def __repr__(self):
        return "<NestedBlockState object with %d %sblocks,%s%s for graph %s, with %d levels of sizes %s at 0x%x>" % \
            (self.levels[0].B, "overlapping " if self.overlap else "",
             " with %d %s," % (self.levels[0].C, "layers" if self.layers else "covariates") if self.ec is not None else "",
             " degree corrected," if self.deg_corr else "",
             str(self.g), len(self.levels), str([(s.N, s.B) for s in self.levels]), id(self))


    def __copy__(self):
        return self.copy()

    def __deepcopy__(self, memo):
        g = self.g.copy()
        eweight = g.own_property(self.eweight.copy()) if self.eweight is not None else None
        vweight = g.own_property(self.vweight.copy()) if self.vweight is not None else None
        clabel = g.own_property(self.clabel.copy())  if self.clabel is not None else None
        ec = g.own_property(self.ec.copy()) if self.ec is not None else None
        bstack = self.get_bstack()
        return self.copy(g=g, eweight=eweight, vweight=vweight, clabel=clabel,
                         ec=ec, bs=[s.vp.b.a for s in bstack])

    def copy(self, g=None, eweight=None, vweight=None, bs=None, ec=None,
             layers=None, deg_corr=None, overlap=None, clabel=None, **kwargs):
        r"""Copies the block state. The parameters override the state properties, and
         have the same meaning as in the constructor.."""

        if bs is None:
            bs = [s.b.a for s in self.levels]
        if overlap is None:
            overlap = self.overlap
        elif self.overlap and not overlap:
            raise ValueError("Cannot automatically convert overlapping nested state to nonoverlapping")
        elif not self.overlap and overlap:
            s = self.levels[0].copy(overlap=True)
            bs[0] = s.b.a
        if deg_corr is None:
            deg_corr = self.deg_corr
        if layers is None:
            layers = self.layers
        return NestedBlockState(self.g if g is None else g,
                                self.eweight if eweight is None else eweight,
                                self.vweight if vweight is None else vweight,
                                self.ec if ec is None else ec, bs,
                                layers=layers, deg_corr=deg_corr,
                                overlap=overlap, clabel=clabel,
                                max_BE=self.levels[0].max_BE,
                                ignore_degrees=kwargs.pop("ignore_degrees", self.ignore_degrees),
                                **kwargs)

    def __getstate__(self):
        state = dict(g=self.g,
                     ec=self.ec,
                     eweight=self.eweight,
                     vweight=self.vweight,
                     overlap=self.overlap,
                     bs=[array(s.b.a) for s in self.levels],
                     clabel=self.clabel,
                     deg_corr=self.deg_corr,
                     max_BE=self.levels[0].max_BE,
                     ignore_degrees=self.ignore_degrees)
        return state

    def __setstate__(self, state):
        self.__init__(**state)
        return state

    def __project_partition(self, l, j):
        """Project partition of level 'j' onto level 'l'"""
        if self.overlap != "full":
            b = self.levels[l].b.copy()
            for i in range(l + 1, j + 1):
                clabel = self.levels[i].b.copy()
                pmap(b, clabel)
        else:
            b = self.levels[j].b.copy()
        return b

    def __propagate_clabel(self, l):
        clabel = self.clabel.copy()
        for j in range(l):
            bg = self.levels[j].bg
            bclabel = bg.new_vertex_property("int")
            reverse_map(self.levels[j].b, bclabel)
            pmap(bclabel, clabel)
            clabel = bclabel
        return clabel

    def __consistency_check(self, op, l):

        print("consistency check after", op, "at level", l)

        for j in range(len(self.levels)):

            c_state = self.levels[j].copy()
            S1 = self.levels[j].entropy()
            S2 = c_state.entropy()

            assert abs(S1 - S2) < 1e-8 and not isnan(S2) and not isnan(S1), "inconsistency at level %d after %s of level %d, different entropies of copies! (%g, %g)" % (j, op, l, S1, S2)


            if self.levels[j].wr.a.min() == 0:
                print("WARNING: empty blocks at level", j)
            if self.levels[j].b.a.max() + 1 != self.levels[j].B:
                print("WARNING: b.max() + 1 != B at level", j, self.levels[j].b.max() + 1, self.levels[j].B)

        for j in range(len(self.levels) - 1):

            B = self.levels[j].b.a.max() + 1
            bg_state = self.levels[j].get_block_state(b=self.levels[j+1].b.copy(),
                                                      overlap=self.overlap == "full",
                                                      deg_corr=self.deg_corr == "full")[0]

            S1 = bg_state.entropy(dense=True and self.deg_corr != "full", multigraph=True)
            S2 = self.levels[j+1].entropy(dense=True and self.deg_corr != "full", multigraph=True)

            if self.levels[j].B != self.levels[j+1].N or S1 != S2:
                self.print_summary()

                from graph_tool.topology import similarity
                print(bg_state)
                print(self.levels[j+1])
                print("N, B:", bg_state.N, bg_state.B)
                print("N, B:", self.levels[j + 1].N, self.levels[j + 1].B)
                print("similarity:", similarity(bg_state.g, self.levels[j+1].g))
                print("b:", bg_state.b.a)
                print("b:", self.levels[j+1].b.a)

                print("wr:", bg_state.wr.a)
                print("wr:", self.levels[j+1].wr.a)

                print("mrs:", bg_state.mrs.a)
                print("mrs:", self.levels[j+1].mrs.a)

                print("eweight:", bg_state.eweight.a)
                print("eweight:", self.levels[j+1].eweight.a)

                print("vweight:", bg_state.vweight.a)
                print("vweight:", self.levels[j+1].vweight.a)


            assert abs(S1 - S2) < 1e-6 and not isnan(S2) and not isnan(S1), "inconsistency at level %d after %s of level %d, different entropies (%g, %g)" % (j, op, l, S1, S2)
            assert self.levels[j].B == self.levels[j+1].N, "inconsistency  at level %d after %s of level %d, different sizes" % (j + 1, op, l)


            # verify hierarchy / clabel consistency
            clabel = self.__project_partition(0, j)
            self.levels[0].clabel.a = self.clabel.a
            assert self.levels[0]._BlockState__check_clabel(),  "inconsistency at level %d after %s of level %d, clabel invalidated" % (j + 1, op, l)
            self.levels[0].clabel.a = 0

            # verify hierarchy consistency
            clabel = self.__project_partition(j, j + 1)
            self.levels[0].clabel.a = self.clabel.a
            assert self.levels[0]._BlockState__check_clabel(),  "inconsistency at level %d after %s of level %d, partition not compatible with upper level" % (j + 1, op, l)
            self.levels[0].clabel.a = 0


    def __rebuild_level(self, l, b, clabel=None):
        r"""Replace level ``l`` given the new partition ``b``, and the
        projected upper level partition clabel."""

        if _bm_test():
            assert clabel is not None or l == len(self.levels) - 1, "clabel not given for intermediary level"

        if clabel is None:
            clabel = self.levels[l].g.new_vertex_property("int")

        old_b = b.copy()

        state = self.levels[l].copy(b=b, clabel=clabel.a)
        self.levels[l] = state
        if l == 0:
            self.clabel = state.g.own_property(self.clabel)  # for CovariateBlockState

        # upper level
        bclabel = state.get_bclabel()
        bstate = self.levels[l].get_block_state(b=bclabel,
                                                overlap=self.overlap == "full",
                                                deg_corr=self.deg_corr == "full")[0]
        if l == len(self.levels) - 1:
            self.levels.append(None)
        self.levels[l + 1] = bstate

        self.levels[l].clabel.a = 0
        self.levels[l + 1].clabel.a = 0

        # if l + 1 < len(self.levels) - 1:
        #     self.levels[l + 1].clabel = self.__project_partition(l + 1, l + 2)


        if l + 1 < len(self.levels) - 1:
            bstate = self.levels[l + 1].get_block_state(b=self.levels[l + 2].b,
                                                        overlap=self.overlap == "full",
                                                        deg_corr=self.deg_corr == "full")[0]

            if _bm_test():
                from graph_tool.topology import similarity
                print("- similarity:", similarity(bstate.g, self.levels[l + 2].g))


                if abs(bstate.entropy() - self.levels[l + 2].entropy()) > 1e-6:
                    print("********************** inconsistent rebuild! **************************")

                    print(bstate.b.a)
                    print(self.levels[l + 2].b.a)

                    print(bstate.wr.a)
                    print(self.levels[l + 2].wr.a)

                    print(bstate.eweight.a)
                    print(self.levels[l + 2].eweight.a)

                    nclabel = self.__project_partition(l, l + 1)
                    print(nclabel.a)
                    print(clabel.a)
                    print(self.levels[l].b.a)
                    print(self.levels[l+1].b.a)
                    print(self.levels[l+2].b.a)
                    print(bstate.b.a)
                print ("DS", l, l + 1, bstate.entropy(), self.levels[l + 2].entropy())

        B = self.levels[l].B

        if _bm_test():
            self.__consistency_check("rebuild", l)

    def __delete_level(self, l):
        if l == 0:
            raise ValueError("cannot delete level l=0")
        b = self.__project_partition(l - 1, l)
        if l < len(self.levels) - 1:
            clabel = self.__project_partition(l - 1, l + 1)
        else:
            clabel = None

        del self.levels[l]

        self.__rebuild_level(l - 1, b=b, clabel=clabel)

        if _bm_test():
            self.__consistency_check("delete", l)

    def __duplicate_level(self, l):
        assert l > 0, "attempted to duplicate level 0"
        if not self.levels[l].overlap:
            bstate = self.levels[l].copy(b=self.levels[l].g.vertex_index.copy("int"))
        else:
            bstate = self.levels[l].copy(b=arange(self.levels[l].g.num_vertices()))
        self.levels.insert(l, bstate)

        if _bm_test():
            self.__consistency_check("duplicate", l)


    def level_entropy(self, l, complete=True, dense=False, multigraph=True,
                      norm=True, dl_ent=False):
        r"""Compute the description length of hierarchy level l.

        Parameters
        ----------
        l : ``int``
            Hierarchy level.
        complete : ``bool`` (optional, default: ``False``)
            If ``True``, the complete entropy will be returned, including constant
            terms not relevant to the block partition.
        dense : ``bool`` (optional, default: ``False``)
            If ``True``, the "dense" variant of the entropy will be computed.
        multigraph : ``bool`` (optional, default: ``True``)
            If ``True``, the multigraph entropy will be used.
        norm : ``bool`` (optional, default: ``True``)
            If ``True``, the entropy will be "normalized" by dividing by the
            number of edges.
        dl_ent : ``bool`` (optional, default: ``False``)
            If ``True``, the description length of the degree sequence will be
            approximated by its entropy.
        """

        bstate = self.levels[l]

        S = bstate.entropy(dl=True, edges_dl=False,
                           dense=dense or (l > 0 and self.deg_corr != "full"),
                           multigraph=multigraph or l > 0,
                           complete=complete or (l > 0 and self.deg_corr == "full"),
                           norm=norm, dl_ent=dl_ent)
        return S

    def entropy(self, complete=True, dense=False, multigraph=True, norm=False,
                dl_ent=False):
        r"""Compute the description length of the entire hierarchy.

        Parameters
        ----------
        complete : ``bool`` (optional, default: ``False``)
            If ``True``, the complete entropy will be returned, including constant
            terms not relevant to the block partition.
        dense : ``bool`` (optional, default: ``False``)
            If ``True``, the "dense" variant of the entropy will be computed.
        multigraph : ``bool`` (optional, default: ``True``)
            If ``True``, the multigraph entropy will be used.
        norm : ``bool`` (optional, default: ``True``)
            If ``True``, the entropy will be "normalized" by dividing by the
            number of edges.
        dl_ent : ``bool`` (optional, default: ``False``)
            If ``True``, the description length of the degree sequence will be
            approximated by its entropy.
        """

        S = 0
        for l in range(len(self.levels)):
            S += self.level_entropy(l, complete=complete, dense=dense,
                                    multigraph=multigraph, norm=norm,
                                    dl_ent=dl_ent)
        return S

    def get_bstack(self):
        r"""Return the nested levels as individual graphs.

        This returns a list of :class:`~graph_tool.Graph` instances
        representing the inferred hierarchy at each level. Each graph has two
        internal vertex and edge property maps named "count" which correspond to
        the vertex and edge counts at the lower level, respectively. Additionally,
        an internal vertex property map named "b" specifies the block partition.
        """

        bstack = []
        for l, bstate in enumerate(self.levels):
            cg = bstate.g
            if l == 0:
                cg = GraphView(cg, skip_properties=True)
            cg.vp["b"] = bstate.b.copy()
            cg.ep["count"] = bstate.eweight
            if bstate.overlap:
                if self.ec is None:
                    cg.vp["node_index"] = bstate.node_index.copy()
                else:
                    cg.vp["node_index"] = bstate.total_state.node_index.copy()

            bstack.append(cg)
            if bstate.N == 1:
                break
        if bstack[-1].num_vertices() > 1:
            cg = Graph(directed=bstack[-1].is_directed())
            cg.add_vertex()
            cg.vp["b"] = cg.new_vertex_property("int")
            e = cg.add_edge(0, 0)
            ew = cg.new_edge_property("int")
            ew[e] = self.levels[-1].E
            cg.ep["count"] = ew
            bstack.append(cg)
        return bstack


    def project_level(self, l):
        r"""Project the partition at level ``l`` onto the lowest level, and return the
        corresponding :class:`~graph_tool.community.BlockState` (or
        :class:`~graph_tool.community.OverlapBlockState`).  """
        if self.overlap != "full":
            clabel = b = self.levels[l].b.copy()
            while l - 1 >= 0:
                clabel = b
                b = self.levels[l - 1].b.copy()
                pmap(b, clabel)
                l -= 1
        else:
            b = self.levels[l].b.copy()
        state = self.levels[0].copy(b=b.a)
        return state

    def merge_layers(self, l_src, l_tgt, revert=False):
        ctxs = []
        for state in self.levels:
            ctxs.append(state.merge_layers(l_src, l_tgt, revert))
        if revert:
            if hasattr(contextlib, "nested"):
                return contextlib.nested(*ctxs)
            else:
                with contextlib.ExitStack() as stack:
                    for ctx in ctxs:
                        stack.enter_context(ctx)
                    return stack.pop_all()

    def print_summary(self):
        for l, state in enumerate(self.levels):
            print("l: %d, N: %d, B: %d" % (l, state.N, state.B))


def nested_mcmc_sweep(state, beta=1., c=1., dl=False, sequential=True,
                      parallel=False, verbose=False):
    r"""Performs a Markov chain Monte Carlo sweep on all levels of the hierarchy.

    Parameters
    ----------
    state : :class:`~graph_tool.community.NestedBlockState`
        The nested block state.
    beta : `float` (optional, default: `1.0`)
        The inverse temperature parameter :math:`\beta`.
    c : ``float`` (optional, default: ``1.0``)
        This parameter specifies how often fully random moves are attempted,
        instead of more likely moves based on the inferred block partition.
        For ``c == 0``, no fully random moves are attempted, and for ``c == inf``
        they are always attempted.
    dl : ``bool`` (optional, default: ``False``)
        If ``True``, the change in the whole description length will be
        considered after each vertex move, not only the entropy.
    sequential : ``bool`` (optional, default: ``True``)
        If ``True``, the move attempts on the vertices are done in sequential
        random order. Otherwise a total of `N` moves attempts are made, where
        `N` is the number of vertices, where each vertex can be selected with
        equal probability.
    verbose : ``bool`` (optional, default: ``False``)
        If ``True``, verbose information is displayed.

    Returns
    -------

    dS_moves : list of (``float``, ``int``) tuples
       The entropy difference (per edge) and number of accepted block membership
       moves after a full sweep for each level.


    Notes
    -----

    This algorithm performs a Markov chain Monte Carlo sweep on each level of the
    network, via the function :func:`~graph_tool.community.mcmc_sweep`.

    This algorithm has a worse-case complexity of :math:`O(E \times L)`, where
    :math:`E` is the number of edges in the network, and :math:`L` is the depth
    of the hierarchy.

    Examples
    --------
    .. testsetup:: nested_mcmc

       gt.seed_rng(42)
       np.random.seed(42)

    .. doctest:: nested_mcmc

       >>> g = gt.collection.data["polbooks"]
       >>> state = gt.NestedBlockState(g, Bs=[10, 5, 3, 2, 1], deg_corr=True)
       >>> ret = gt.nested_mcmc_sweep(state)
       >>> print(ret)
       [(0.0, 0), (0.0, 0), (0.0, 0), (0.0, 0), (0.0, 0)]

    References
    ----------

    .. [peixoto-efficient-2014] Tiago P. Peixoto, "Efficient Monte Carlo and greedy
       heuristic for the inference of stochastic block models", Phys. Rev. E 89, 012804 (2014),
       :doi:`10.1103/PhysRevE.89.012804`, :arxiv:`1310.4378`.
    .. [peixoto-hierarchical-2014] Tiago P. Peixoto, "Hierarchical block structures and high-resolution
       model selection in large networks ", Phys. Rev. X 4, 011047 (2014), :doi:`10.1103/PhysRevX.4.011047`,
       :arxiv:`1310.4377`.
    .. [peixoto-model-2015] Tiago P. Peixoto, "Model selection and hypothesis
       testing for large-scale network models with overlapping groups",
       Phys. Rev. X 5, 011033 (2015), :doi:`10.1103/PhysRevX.5.011033`,
       :arxiv:`1409.3059`.
    """

    rets = []
    for l, bstate in enumerate(state.levels):
        if verbose:
            print("Level:", l, "N:", bstate.N, "B:", bstate.B)

        # constraint partitions not to invalidate upper layers
        if l < len(state.levels) - 1:
            clabel = state._NestedBlockState__project_partition(l, l + 1)
        else:
            clabel = bstate.g.new_vertex_property("int")

        # propagate externally imposed clabel at the bottom
        cclabel = state._NestedBlockState__propagate_clabel(l)
        cclabel.a += clabel.a * (cclabel.a.max() + 1)
        continuous_map(cclabel)

        bstate.clabel = cclabel

        ret = mcmc_sweep(bstate, beta=beta, c=c, dl=dl,
                         dense = l > 0 and state.deg_corr != "full",
                         multigraph = l > 0,
                         sequential=sequential, parallel=parallel,
                         verbose=verbose)

        bstate.clabel.a = 0

        rets.append(ret)
    return rets

def replace_level(l, state, min_B=None, max_B=None, max_b=None, nsweeps=10,
                  nmerge_sweeps=10, adaptive_sweeps=True, r=2, c=0, epsilon=0.,
                  sequential=True, parallel=False, dl=False, dense=False,
                  multigraph=True, sparse_thresh=100, verbose=False,
                  checkpoint=None, minimize_state=None, dl_ent=False,
                  confine_layers=False):
    r"""Replaces level l with another state with a possibly different number of
    groups. This may change not only the state at level l, but also the one at
    level l + 1, which needs to be 'rebuilt' because of the label changes at
    level l."""

    if _bm_test():
        state._NestedBlockState__consistency_check("(before) replace level", l)

    bstate = state.levels[l]
    g = bstate.g
    base_g = g if not bstate.overlap else bstate.base_g
    eweight = bstate.eweight if not bstate.overlap else None
    ec = None
    if state.ec is not None:
        if bstate.overlap:
            ec = bstate.base_ec
        else:
            ec = bstate.ec

    if l > 0 or min_B is None:
        if l + 1 < len(state.levels):
            min_B = state.levels[l + 1].B
        else:
            min_B = 1
    if l > 0 or max_B is None:
        max_B = bstate.N

    min_B = max(min_B, state.clabel.a.max() + 1)

    # constraint partitions not to invalidate upper layers
    if l < len(state.levels) - 1:
        clabel = state._NestedBlockState__project_partition(l, l + 1)
    else:
        clabel = bstate.g.new_vertex_property("int")

    assert min_B <= max_B, (min_B, max_B, bstate.B, bstate.N, g.num_vertices(),
                            state.clabel.a.max() + 1, clabel.a.max() + 1, l)

    # propagate externally imposed clabel at the bottom
    cclabel = state._NestedBlockState__propagate_clabel(l)
    assert cclabel.a.max() <= state.clabel.a.max(),  (cclabel.a.max(), state.clabel.a.max())
    cclabel.a += clabel.a * (cclabel.a.max() + 1)
    continuous_map(cclabel)

    min_B = max(min_B, cclabel.a.max() + 1)

    assert min_B <= max_B, (min_B, max_B, bstate.B, bstate.N, g.num_vertices(),
                            cclabel.a.max() + 1, state.clabel.a.max() + 1, clabel.a.max() + 1, l)

    if _bm_test():
        assert bstate._BlockState__check_clabel(), "invalid clabel before minimize!"

    nested_dl = l < len(state.levels) - 1

    state.levels[l].clabel.a = cclabel.a

    Sb = get_b_dl(state.levels[l],
                  dense=(l > 0 and state.deg_corr != "full") or dense,
                  multigraph=l > 0 or multigraph,
                  nested_dl=nested_dl,
                  nested_overlap=state.overlap == "full",
                  dl_ent=dl_ent)

    state.levels[l].clabel.a = 0

    if _bm_test():
        assert clabel.a.max() + 1 <= min_B

    if l == 0 and max_b is not None:
        istate = state.levels[0].copy(b=max_b.copy(),
                                      clabel=cclabel.a if state.overlap else cclabel)
        if _bm_test():
            assert istate._BlockState__check_clabel(), "invalid clabel!"
        if istate.B > max_B:
            istate = multilevel_minimize(istate, B=max_B, nsweeps=nsweeps,
                                         nmerge_sweeps=nmerge_sweeps,
                                         adaptive_sweeps=adaptive_sweeps, c=c,
                                         r=r, dl=dl, sequential=sequential,
                                         parallel=parallel,
                                         greedy_cooling=True, epsilon=epsilon,
                                         confine_layers=confine_layers,
                                         verbose=verbose=="full")
            if _bm_test():
                assert istate._BlockState__check_clabel(), "invalid clabel!"
        init_states = [istate]
    else:
        init_states = None

    res = minimize_blockmodel_dl(base_g,
                                 ec=ec,
                                 layers=state.layers,
                                 confine_layers=confine_layers,
                                 eweight=eweight,
                                 deg_corr=bstate.deg_corr,
                                 nsweeps=nsweeps,
                                 nmerge_sweeps=nmerge_sweeps,
                                 adaptive_sweeps=adaptive_sweeps,
                                 c=c, r=r, sequential=sequential,
                                 parallel=parallel,
                                 greedy_cooling=True,
                                 epsilon=epsilon,
                                 max_B=max_B,
                                 min_B=min_B,
                                 clabel=cclabel if not bstate.overlap else cclabel.a,
                                 max_BE=bstate.max_BE,
                                 dl=dl,
                                 dense=(l > 0 and state.deg_corr != "full") or dense,
                                 multigraph=l > 0 or multigraph,
                                 sparse_heuristic=base_g.num_vertices() > sparse_thresh,
                                 nested_dl=nested_dl,
                                 overlap=bstate.overlap,
                                 nested_overlap=state.overlap == "full",
                                 nonoverlap_compare=False,
                                 nonoverlap_init=False,
                                 init_states=init_states,
                                 verbose=verbose=="full",
                                 ##exaustive=g.num_vertices() <= 100,
                                 #minimize_state=minimize_state.minimize_state, >>>>>> HERE <<<<<
                                 checkpoint=checkpoint,
                                 dl_ent=dl_ent,
                                 ignore_degrees=state.ignore_degrees if l == 0 else None)

    if _bm_test():
        assert (res.clabel.a == cclabel.a).all(), (res.clabel.a, cclabel.a)
        assert res._BlockState__check_clabel(), "invalid clabel after minimize!"

    Sf = get_b_dl(res,
                  dense=(l > 0 and state.deg_corr != "full") or dense,
                  multigraph=l > 0 or multigraph,
                  nested_dl=nested_dl,
                  nested_overlap=state.overlap == "full",
                  dl_ent=dl_ent)

    res.clabel.a = 0
    if state.ec is not None:
        for s in res.states:
            s.clabel.a = 0
    b = res.b

    kept = False
    if Sf - Sb >= -1e-10:
        kept = True
    if Sf - Sb == 0 and bstate.B != state.levels[l].B:
        kept = False
    if res.B == res.N:
        kept = True

    if kept:
        Sf_rej = Sf
        Sf = Sb
    else:
        # rebuild current level
        state._NestedBlockState__rebuild_level(l, b=b, clabel=clabel)

    if verbose:
        print("level", l, ": resizing", bstate.B, "->", state.levels[l].B, ", dS:", Sf - Sb, end="")
        if kept:
            print(" [kept, rejected (%d, %g) vs (%d, %g)]" % (res.B, Sf_rej, bstate.B, Sb))
        else:
            print()

    if _bm_test():
        state._NestedBlockState__consistency_check("replace level", l)

    dS = Sf - Sb
    return dS, kept


class NestedMinimizeState(object):
    r"""This object stores information regarding the current entropy minimization
    state, so that the algorithms can resume previously started runs.
    This object can be saved to disk via the :mod:`pickle` interface."""

    def __init__(self):
        self.minimize_state = MinimizeState()
        self.l = 0
        self.bs = []
        self.done = []
        self.init = True

    def clear(self):
        self.minimize_state.clear()
        self.l = 0
        del self.bs[:]
        del self.done[:]

    def sync(self, state):
        if len(self.bs) == 0:
            for s in state.levels:
                self.bs.append(array(s.b.fa))
        while len(self.done) < len(state.levels):
            self.done.append(False)

    def delete(self, l):
        del self.done[l]
        del self.bs[l]

    def insert(self, l, state):
        self.done.insert(l + 1, False)
        ba = array(state.levels[l].b.fa)
        self.bs.insert(l + 1, ba)

    def mark_level(self, l, done, state):
        while len(state.levels) > len(self.bs):
            ba = array(state.levels[len(self.bs)].b.fa)
            self.bs.append(ba)
            self.done.append(False)

        self.done[l] = done
        if done:
            self.bs[l] = array(state.levels[l].b.fa)

    def clear_mstate(self):
        self.minimize_state.clear()

def get_checkpoint_wrap(checkpoint, state, minimize_state, dl_ent):
    S_total = state.entropy(complete=True, dl_ent=dl_ent)

    if checkpoint is not None:
        def check_wrap(bstate, Sb, delta, nmoves, ms):
            l = minimize_state.l
            bstate = state.levels[l]
            S_l = bstate.entropy()
            S = S_total - S_l + Sb
            if bstate is None:
                checkpoint(None, S, delta, nmoves, minimize_state)
            else:
                checkpoint(state, S, delta, nmoves, minimize_state)
        chkp = check_wrap
    else:
        chkp = None

    return chkp


def nested_tree_sweep(state, min_B=None, max_B=None, max_b=None, nsweeps=10,
                      epsilon=0., r=2., nmerge_sweeps=10, adaptive_sweeps=True,
                      c=0, dl=False, dense=False, multigraph=True,
                      sequential=True, parallel=False, sparse_thresh=100,
                      checkpoint=None, minimize_state=None, frozen_levels=None,
                      confine_layers=False, verbose=False, **kwargs):
    r"""Performs one greedy sweep in the entire hierarchy tree, attempting to
    decrease its description length.

    Parameters
    ----------
    state : :class:`~graph_tool.community.NestedBlockState`
        The nested block state.
    min_B : ``int`` (optional, default: ``None``)
        Minimum number of blocks at the lowest level.
    max_B : ``int`` (optional, default: ``None``)
        Maximum number of blocks at the lowest level.
    max_b : ``int`` (optional, default: ``None``)
        Block partition used for the maximum number of blocks at the lowest
        level.
    nsweeps : ``int`` (optional, default: ``10``)
        The number of sweeps done after each merge step to reach the local
        minimum.
    epsilon : ``float`` (optional, default: ``0``)
        Converge criterion for ``adaptive_sweeps``.
    r : ``float`` (optional, default: ``2.``)
        Agglomeration ratio for the merging steps. Each merge step will attempt
        to find the best partition into :math:`B_{i-1} / r` blocks, where
        :math:`B_{i-1}` is the number of blocks in the last step.
    nmerge_sweeps : `int` (optional, default: `10`)
        The number of merge sweeps done, where in each sweep a better merge
        candidate is searched for every block.
    c : ``float`` (optional, default: ``1.0``)
        This parameter specifies how often fully random moves are attempted,
        instead of more likely moves based on the inferred block partition.
        For ``c == 0``, no fully random moves are attempted, and for ``c == inf``
        they are always attempted.
    dense : ``bool`` (optional, default: ``False``)
        If ``True``, the "dense" variant of the entropy will be computed.
    sequential : ``bool`` (optional, default: ``True``)
        If ``True``, the move attempts on the vertices are done in sequential
        random order. Otherwise a total of `N` moves attempts are made, where
        `N` is the number of vertices, where each vertex can be selected with
        equal probability.
    sparse_thresh : ``int`` (optional, default: ``100``)
        If the number of nodes in the higher level multigraphs exceeds this
        number, the sparse entropy will be used to find the best partition,
        but the dense entropy will be used to compare different partitions.
    checkpoint : function (optional, default: ``None``)
        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
        can be continued later. The function must have the following signature:

        .. code-block:: python

            def checkpoint(state, S, delta, nmoves, minimize_state):
                ...

        where `state` is either a :class:`~graph_tool.community.NestedBlockState`
        instance or ``None``, `S` is the current entropy value, `delta` is
        the entropy difference in the last MCMC sweep, and `nmoves` is the
        number of accepted block membership moves. The ``minimize_state``
        argument is a :class:`NestedMinimizeState` instance which specifies
        the current state of the algorithm, which can be stored via :mod:`pickle`,
        and supplied via the ``minimize_state`` option below to continue from an
        interrupted run.

        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
        remaining parameters will be zero, except the last.
    minimize_state : :class:`NestedMinimizeState` (optional, default: ``None``)
        If provided, this will specify an exact point of execution from which
        the algorithm will continue. The expected object is a
        :class:`NestedMinimizeState` instance which will be passed to the
        callback of the ``checkpoint`` option above, and  can be stored by
        :mod:`pickle`.
    frozen_levels : :class:`list` (optional, default: ``None``)
        List of levels (``int``s) which will remain unmodified during the
        algorithm.
    verbose : ``bool`` (optional, default: ``False``)
        If ``True``, verbose information is displayed.

    Returns
    -------

    dS : ``float``
       The description length difference (per edge) after the move.


    Notes
    -----

    This algorithm performs a constrained agglomerative heuristic on each level
    of the network, via the function :func:`~graph_tool.community.multilevel_minimize`.

    This algorithm has worst-case complexity of :math:`O(N\ln^2 N \times L)`,
    where  :math:`N` is the number of nodes in the network, and :math:`L` is
    the depth of the hierarchy.

    References
    ----------

    .. [peixoto-efficient-2014] Tiago P. Peixoto, "Efficient Monte Carlo and greedy
       heuristic for the inference of stochastic block models", Phys. Rev. E 89, 012804 (2014),
       :doi:`10.1103/PhysRevE.89.012804`, :arxiv:`1310.4378`.
    .. [peixoto-hierarchical-2014] Tiago P. Peixoto, "Hierarchical block structures and high-resolution
       model selection in large networks ", Phys. Rev. X 4, 011047 (2014), :doi:`10.1103/PhysRevX.4.011047`,
       :arxiv:`1310.4377`.
    .. [peixoto-model-2015] Tiago P. Peixoto, "Model selection and hypothesis
       testing for large-scale network models with overlapping groups",
       Phys. Rev. X 5, 011033 (2015), :doi:`10.1103/PhysRevX.5.011033`,
       :arxiv:`1409.3059`.
    """

    dl_ent = kwargs.get("dl_ent", False)

    if minimize_state is None: