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
from .. stats import label_self_loops
from .. spectral import adjacency
import random
from numpy import *
import numpy
from scipy.optimize import fsolve, fminbound
import scipy.special
from collections import defaultdict
import copy
import heapq

from .. dl_import import dl_import
dl_import("from . import libgraph_tool_community as libcommunity")

__test__ = False

def set_test(test):
    global __test__
    __test__ = test

def _bm_test():
    global __test__
    return __test__

def get_block_graph(g, B, b, vcount, ecount):
    cg, br, vcount, ecount = condensation_graph(g, b,
                                                vweight=vcount,
                                                eweight=ecount,
                                                self_loops=True)[:4]
    cg.vp["count"] = vcount
    cg.ep["count"] = ecount
    cg = Graph(cg, vorder=br)

    cg.add_vertex(B - cg.num_vertices())
    return cg

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

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

    Parameters
    ----------
    g : :class:`~graph_tool.Graph`
        Graph to be modelled.
    eweight : :class:`~graph_tool.PropertyMap` (optional, default: ``None``)
        Edge multiplicities (for multigraphs or block graphs).
    vweight : :class:`~graph_tool.PropertyMap` (optional, default: ``None``)
        Vertex multiplicities (for block graphs).
    b : :class:`~graph_tool.PropertyMap` (optional, default: ``None``)
        Initial block labels on the vertices. If not supplied, it will be
        randomly sampled.
    B : ``int`` (optional, default: ``None``)
        Number of blocks. If not supplied it will be either obtained from the
        parameter ``b``, or set to the maximum possible value according to the
        minimum description length.
    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.
    deg_corr : ``bool`` (optional, default: ``True``)
        If ``True``, the degree-corrected version of the blockmodel ensemble will
        be assumed, otherwise the traditional variant will be used.
    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,
    """

    _state_ref_count = 0

    def __init__(self, g, eweight=None, vweight=None, b=None,
                 B=None, clabel=None, deg_corr=True,
                 max_BE=1000, **kwargs):
        BlockState._state_ref_count += 1

        # initialize weights to unity, if necessary
        if eweight is None:
            eweight = g.new_edge_property("int")
            eweight.fa = 1
        elif eweight.value_type() != "int32_t":
            eweight = eweight.copy(value_type="int32_t")
        if vweight is None:
            vweight = g.new_vertex_property("int")
            vweight.fa = 1
        elif vweight.value_type() != "int32_t":
            vweight = vweight.copy(value_type="int32_t")
        self.eweight = g.own_property(eweight)
        self.vweight = g.own_property(vweight)

        self.is_weighted = False
        if ((g.num_edges() > 0 and self.eweight.fa.max() > 1) or
            kwargs.get("force_weighted", False)):
            self.is_weighted = True

        # configure the main graph and block model parameters
        self.g = g

        self.E = int(self.eweight.fa.sum())
        self.N = int(self.vweight.fa.sum())

        self.deg_corr = deg_corr

        # ensure we have at most as many blocks as nodes
        if B is not None and b is None:
            B = min(B, self.g.num_vertices())

        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())
            B = min(B, self.g.num_vertices())
            ba = random.randint(0, B, self.g.num_vertices())
            ba[:B] = arange(B)        # avoid empty blocks
            if B < self.g.num_vertices():
                random.shuffle(ba)
            b = g.new_vertex_property("int")
            b.fa = ba
            self.b = b
        else:
            # if a partition is available, we will incorporate it.
            if isinstance(b, numpy.ndarray):
                self.b = g.new_vertex_property("int")
                self.b.fa = b
            else:
                self.b = b = g.own_property(b.copy(value_type="int"))
            if B is None:
                B = int(self.b.fa.max()) + 1

        # if B > self.N:
        #     raise ValueError("B > N!")

        if self.b.fa.max() >= B:
            raise ValueError("Maximum value of b is larger or equal to B! (%d vs %d)" % (self.b.fa.max(), B))

        # Construct block-graph
        self.bg = get_block_graph(g, B, self.b, self.vweight, self.eweight)
        self.bg.set_fast_edge_removal()

        self.mrs = self.bg.ep["count"]
        self.wr = self.bg.vp["count"]

        del self.bg.ep["count"]
        del self.bg.vp["count"]

        self.mrp = self.bg.degree_property_map("out", weight=self.mrs)

        if g.is_directed():
            self.mrm = self.bg.degree_property_map("in", weight=self.mrs)
        else:
            self.mrm = self.mrp

        self.vertices = libcommunity.get_vector(B)
        self.vertices.a = arange(B)
        self.B = B

        if clabel is not None:
            if isinstance(clabel, PropertyMap):
                self.clabel = self.g.own_property(clabel.copy("int"))
            else:
                self.clabel = self.g.new_vertex_property("int")
                self.clabel.a = clabel
        else:
            self.clabel = self.g.new_vertex_property("int")

        self.emat = None
        if max_BE is None:
            max_BE = 1000
        self.max_BE = max_BE

        self.overlap = False
        self.ignore_degrees = kwargs.get("ignore_degrees", None)
        if self.ignore_degrees is None:
            self.ignore_degrees = g.new_vertex_property("bool", False)

        # used by mcmc_sweep()
        self.egroups = None
        self.nsampler = None
        self.sweep_vertices = None
        self.overlap_stats = libcommunity.overlap_stats()
        self.partition_stats = libcommunity.partition_stats()
        self.edges_dl = False

        # computation cache
        E = g.num_edges()
        N = g.num_vertices()
        libcommunity.init_safelog(int(5 * max(E, N)))
        libcommunity.init_xlogx(int(5 * max(E, N)))
        libcommunity.init_lgamma(int(3 * max(E, N)))

    def __del__(self):
        try:
            BlockState._state_ref_count -= 1
            if BlockState._state_ref_count == 0:
                libcommunity.clear_safelog()
                libcommunity.clear_xlogx()
                libcommunity.clear_lgamma()
        except (ValueError, AttributeError, TypeError):
            pass

    def __repr__(self):
        return "<BlockState object with %d blocks,%s for graph %s, at 0x%x>" % \
            (self.B, " degree corrected," if self.deg_corr else "", str(self.g),
             id(self))


    def __init_partition_stats(self, empty=True, edges_dl=False):
        self.edges_dl = edges_dl
        if not empty:
            self.partition_stats = libcommunity.init_partition_stats(self.g._Graph__graph,
                                                                     _prop("v", self.g, self.b),
                                                                     _prop("e", self.g, self.eweight),
                                                                     self.N, self.B,
                                                                     edges_dl,
                                                                     _prop("v", self.g, self.ignore_degrees))
        else:
            self.partition_stats = libcommunity.partition_stats()



    def copy(self, g=None, eweight=None, vweight=None, b=None, B=None,
             deg_corr=None, clabel=None, overlap=False, **kwargs):
        r"""Copies the block state. The parameters override the state properties, and
         have the same meaning as in the constructor. If ``overlap=True`` an
         instance of :class:`~graph_tool.community.OverlapBlockState` is
         returned."""

        if not overlap:
            state = BlockState(self.g if g is None else g,
                               eweight=self.eweight if eweight is None else eweight,
                               vweight=self.vweight if vweight is None else vweight,
                               b=self.b.copy() if b is None else b,
                               B=(self.B if b is None else None) if B is None else B,
                               clabel=self.clabel if clabel is None else clabel,
                               deg_corr=self.deg_corr if deg_corr is None else deg_corr,
                               max_BE=self.max_BE,
                               ignore_degrees=kwargs.pop("ignore_degrees", self.ignore_degrees),
                               **kwargs)
        else:
            state = OverlapBlockState(self.g if g is None else g,
                                      b=b if b is not None else self.b,
                                      B=(self.B if b is None else None) if B is None else B,
                                      clabel=self.clabel if clabel is None else clabel,
                                      deg_corr=self.deg_corr if deg_corr is None else deg_corr,
                                      max_BE=self.max_BE, **kwargs)

        if not state.__check_clabel():
            b = state.b.a + state.clabel.a * state.B
            continuous_map(b)
            state = state.copy(b=b)

            if _bm_test():
                assert state.__check_clabel()

        return state


    def __getstate__(self):
        state = dict(g=self.g,
                     eweight=self.eweight,
                     vweight=self.vweight,
                     b=self.b,
                     B=self.B,
                     clabel=self.clabel,
                     deg_corr=self.deg_corr,
                     max_BE=self.max_BE,
                     ignore_degrees=self.ignore_degrees)
        return state

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

    def get_block_state(self, b=None, vweight=False, deg_corr=False,
                        overlap=False, **kwargs):
        r"""Returns a :class:`~graph_tool.community.BlockState`` corresponding to the
        block graph. The parameters have the same meaning as the in the constructor."""


        state = BlockState(self.bg, eweight=self.mrs,
                           vweight=self.wr if vweight else None,
                           b=self.bg.vertex_index.copy("int") if b is None else b,
                           clabel=self.get_bclabel(),
                           deg_corr=deg_corr,
                           max_BE=self.max_BE)
        if overlap:
            state = state.copy(overlap=True)
        n_map = self.b.copy()
        return state, n_map

    def get_bclabel(self):
        r"""Returns a :class:`~graph_tool.PropertyMap`` corresponding to constraint
        labels for the block graph."""

        bclabel = self.bg.new_vertex_property("int")
        reverse_map(self.b, bclabel)
        pmap(bclabel, self.clabel)
        return bclabel

    def __check_clabel(self):
        b = self.b.a + self.clabel.a * self.B
        continuous_map(b)
        b2 = self.b.copy()
        continuous_map(b2.a)
        return (b == b2.a).all()

    def __get_emat(self):
        if self.emat is None:
            self.__regen_emat()
        return self.emat

    def __regen_emat(self):
        if self.B <= self.max_BE:
            self.emat = libcommunity.create_emat(self.bg._Graph__graph)
        else:
            self.emat = libcommunity.create_ehash(self.bg._Graph__graph)

    def __build_egroups(self, empty=False):
        self.esrcpos = self.g.new_edge_property("int")
        self.etgtpos = self.g.new_edge_property("int")

        self.egroups = libcommunity.build_egroups(self.g._Graph__graph,
                                                  self.bg._Graph__graph,
                                                  _prop("v", self.g, self.b),
                                                  _prop("e", self.g, self.eweight),
                                                  _prop("e", self.g, self.esrcpos),
                                                  _prop("e", self.g, self.etgtpos),
                                                  self.is_weighted, empty)

    def __build_nsampler(self, empty=False):
        self.nsampler = libcommunity.init_neighbour_sampler(self.g._Graph__graph,
                                                            _prop("e", self.g, self.eweight),
                                                            True, empty)

    def __cleanup_bg(self):
        emask = self.bg.new_edge_property("bool")
        emask.a = self.mrs.a[:len(emask.a)] > 0
        self.bg.set_edge_filter(emask)
        self.bg.purge_edges()
        self.emat = None

    def get_blocks(self):
        r"""Returns the property map which contains the block labels for each vertex."""
        return self.b

    def get_bg(self):
        r"""Returns the block graph."""
        return self.bg

    def get_ers(self):
        r"""Returns the edge property map of the block graph which contains the :math:`e_{rs}` matrix entries.
        For undirected graphs, the diagonal values (self-loops) contain :math:`e_{rr}/2`."""
        return self.mrs

    def get_er(self):
        r"""Returns the vertex property map of the block graph which contains the number
        :math:`e_r` of half-edges incident on block :math:`r`. If the graph is
        directed, a pair of property maps is returned, with the number of
        out-edges :math:`e^+_r` and in-edges :math:`e^-_r`, respectively."""
        if self.bg.is_directed():
            return self.mrp, self.mrm
        else:
            return self.mrp

    def get_nr(self):
        r"""Returns the vertex property map of the block graph which contains the block sizes :math:`n_r`."""
        return self.wr

    def entropy(self, complete=True, dl=False, partition_dl=True,
                degree_dl=True, edges_dl=True, dense=False, multigraph=True,
                norm=False, dl_ent=False, **kwargs):
        r"""Calculate the entropy associated with the current block partition.

        Parameters
        ----------
        complete : ``bool`` (optional, default: ``False``)
            If ``True``, the complete entropy will be returned, including constant
            terms not relevant to the block partition.
        dl : ``bool`` (optional, default: ``False``)
            If ``True``, the full description length will be returned.
        partition_dl : ``bool`` (optional, default: ``True``)
            If ``True``, and ``dl == True`` the partition description length
            will be considered.
        edges_dl : ``bool`` (optional, default: ``True``)
            If ``True``, and ``dl == True`` the edge matrix description length
            will be considered.
        degree_dl : ``bool`` (optional, default: ``True``)
            If ``True``, and ``dl == True`` the degree sequence description
            length will be considered.
        dense : ``bool`` (optional, default: ``False``)
            If ``True``, the "dense" variant of the entropy will be computed.
        multigraph : ``bool`` (optional, default: ``False``)
            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.

        Notes
        -----
        For the traditional blockmodel (``deg_corr == False``), the entropy is
        given by

        .. math::

          \mathcal{S}_t &\cong E - \frac{1}{2} \sum_{rs}e_{rs}\ln\left(\frac{e_{rs}}{n_rn_s}\right), \\
          \mathcal{S}^d_t &\cong E - \sum_{rs}e_{rs}\ln\left(\frac{e_{rs}}{n_rn_s}\right),

        for undirected and directed graphs, respectively, where :math:`e_{rs}`
        is the number of edges from block :math:`r` to :math:`s` (or the number
        of half-edges for the undirected case when :math:`r=s`), and :math:`n_r`
        is the number of vertices in block :math:`r` .

        For the degree-corrected variant with "hard" degree constraints the
        equivalent expressions are

        .. math::

            \mathcal{S}_c &\cong -E -\sum_kN_k\ln k! - \frac{1}{2} \sum_{rs}e_{rs}\ln\left(\frac{e_{rs}}{e_re_s}\right), \\
            \mathcal{S}^d_c &\cong -E -\sum_{k^+}N_{k^+}\ln k^+!  -\sum_{k^-}N_{k^-}\ln k^-! - \sum_{rs}e_{rs}\ln\left(\frac{e_{rs}}{e^+_re^-_s}\right),

        where :math:`e_r = \sum_se_{rs}` is the number of half-edges incident on
        block :math:`r`, and :math:`e^+_r = \sum_se_{rs}` and :math:`e^-_r =
        \sum_se_{sr}` are the numbers of out- and in-edges adjacent to block
        :math:`r`, respectively.

        If ``dense == False`` and ``multigraph == True``, the entropy used will
        be of the "Poisson" model, with the additional term:

        .. math::

            {\mathcal{S}_{cm}^{(d)}} = \mathcal{S}_c^{(d)} + \sum_{i>j} \ln A_{ij}! + \sum_i \ln A_{ii}!!


        If ``dl == True``, the description length :math:`\mathcal{L}_t` of the
        model will be returned as well, as described in
        :func:`model_entropy`. Note that for the degree-corrected version the
        description length is

        .. math::

            \mathcal{L}_c = \mathcal{L}_t + \sum_r\min\left(\mathcal{L}^{(1)}_r, \mathcal{L}^{(2)}_r\right),

        with

        .. math::

              \mathcal{L}^{(1)}_r &= \ln{\left(\!\!{n_r \choose e_r}\!\!\right)}, \\
              \mathcal{L}^{(2)}_r &= \ln\Xi_r + \ln n_r! - \sum_k \ln n^r_k!,

        and :math:`\ln\Xi_r \simeq 2\sqrt{\zeta(2)e_r}`, where :math:`\zeta(x)`
        is the `Riemann zeta function
        <https://en.wikipedia.org/wiki/Riemann_zeta_function>`_, and
        :math:`n^r_k` is the number of nodes in block :math:`r` with degree
        :math:`k`. For directed graphs we have instead :math:`k \to (k^-, k^+)`,
        and :math:`\ln\Xi_r \to \ln\Xi^+_r + \ln\Xi^-_r` with :math:`\ln\Xi_r^+
        \simeq 2\sqrt{\zeta(2)e^+_r}` and :math:`\ln\Xi_r^- \simeq
        2\sqrt{\zeta(2)e^-_r}`.

        If ``dl_ent=True`` is passed, this will be approximated instead by

        .. math::

            \mathcal{L}_c \simeq \mathcal{L}_t - \sum_rn_r\sum_kp^r_k\ln p^r_k,

        where :math:`p^r_k = n^r_k / n_r`.

        If the "dense" entropies are requested (``dense=True``), they will be
        computed as

        .. math::

            \mathcal{S}_t  &= \sum_{r>s} \ln{\textstyle {n_rn_s \choose e_{rs}}} + \sum_r \ln{\textstyle {{n_r\choose 2}\choose e_{rr}/2}}\\
            \mathcal{S}^d_t  &= \sum_{rs} \ln{\textstyle {n_rn_s \choose e_{rs}}},

        for simple graphs, and

        .. math::

            \mathcal{S}_m  &= \sum_{r>s} \ln{\textstyle \left(\!\!{n_rn_s \choose e_{rs}}\!\!\right)} + \sum_r \ln{\textstyle \left(\!\!{\left(\!{n_r\choose 2}\!\right)\choose e_{rr}/2}\!\!\right)}\\
            \mathcal{S}^d_m  &= \sum_{rs} \ln{\textstyle \left(\!\!{n_rn_s \choose e_{rs}}\!\!\right)},

        for multigraphs (i.e. ``multigraph == True``). A dense entropy for the
        degree-corrected model is not available, and if requested will return a
        :exc:`NotImplementedError`.

        If ``complete == False`` constants in the above equations which do not
        depend on the partition of the nodes will be omitted.

        Note that in all cases if ``norm==True`` the value returned corresponds
        to the entropy `per edge`, i.e. :math:`(\mathcal{S}_{t/c}\; [\,+\,\mathcal{L}_{t/c}])/ E`.
        """

        xi_fast = kwargs.get("xi_fast", False)
        dl_deg_alt = kwargs.get("dl_deg_alt", True)

        E = self.E
        N = self.N

        if dense:
            if self.deg_corr:
                raise NotImplementedError('A degree-corrected "dense" entropy is not yet implemented')

            S = libcommunity.entropy_dense(self.bg._Graph__graph,
                                            _prop("e", self.bg, self.mrs),
                                            _prop("v", self.bg, self.wr),
                                            multigraph)
        else:
            S = libcommunity.entropy(self.bg._Graph__graph,
                                      _prop("e", self.bg, self.mrs),
                                      _prop("v", self.bg, self.mrp),
                                      _prop("v", self.bg, self.mrm),
                                      _prop("v", self.bg, self.wr),
                                      self.deg_corr)

            if _bm_test():
                assert not isnan(S) and not isinf(S), "invalid entropy %g (%s) " % (S, str(dict(complete=complete,
                                                                                                random=random, dl=dl,
                                                                                                partition_dl=partition_dl,
                                                                                                edges_dl=edges_dl,
                                                                                                dense=dense, multigraph=multigraph,
                                                                                                norm=norm)))
            if self.deg_corr:
                S -= E
            else:
                S += E

            if complete:
                if self.deg_corr:
                    S += libcommunity.deg_entropy_term(self.g._Graph__graph,
                                                       libcore.any(),
                                                       self.overlap_stats,
                                                       self.N,
                                                       _prop("e", self.g, self.eweight),
                                                       _prop("v", self.g, self.ignore_degrees))

                if multigraph:
                    S += libcommunity.entropy_parallel(self.g._Graph__graph,
                                                       _prop("e", self.g, self.eweight))

                if _bm_test():
                    assert not isnan(S) and not isinf(S), "invalid entropy %g (%s) " % (S, str(dict(complete=complete,
                                                                                                    random=random, dl=dl,
                                                                                                    partition_dl=partition_dl,
                                                                                                    edges_dl=edges_dl,
                                                                                                    dense=dense, multigraph=multigraph,
                                                                                                    norm=norm)))
        if dl:
            if partition_dl:
                if self.partition_stats.is_enabled():
                    S += self.partition_stats.get_partition_dl()
                else:
                    self.__init_partition_stats(empty=False)
                    S += self.partition_stats.get_partition_dl()
                    self.__init_partition_stats(empty=True)

                if _bm_test():
                    assert not isnan(S) and not isinf(S), "invalid entropy %g (%s) " % (S, str(dict(complete=complete,
                                                                                                    random=random, dl=dl,
                                                                                                    partition_dl=partition_dl,
                                                                                                    edges_dl=edges_dl,
                                                                                                    dense=dense, multigraph=multigraph,
                                                                                                    norm=norm)))
            if edges_dl:
                actual_B = (self.wr.a > 0).sum()
                S += model_entropy(actual_B, N, E, directed=self.g.is_directed(), nr=False)

            if self.deg_corr and degree_dl:
                if self.partition_stats.is_enabled():
                    S_seq = self.partition_stats.get_deg_dl(dl_ent, dl_deg_alt, xi_fast)
                else:
                    self.__init_partition_stats(empty=False)
                    S_seq = self.partition_stats.get_deg_dl(dl_ent, dl_deg_alt, xi_fast)
                    self.__init_partition_stats(empty=True)

                S += S_seq

                if _bm_test():
                    assert not isnan(S_seq) and not isinf(S_seq), "invalid entropy %g (%s) " % (S_seq, str(dict(complete=complete,
                                                                                                                random=random, dl=dl,
                                                                                                                partition_dl=partition_dl,
                                                                                                                edges_dl=edges_dl,
                                                                                                                dense=dense, multigraph=multigraph,
                                                                                                                norm=norm)))

        if _bm_test():
            assert not isnan(S) and not isinf(S), "invalid entropy %g (%s) " % (S, str(dict(complete=complete,
                                                                                            random=random, dl=dl,
                                                                                            partition_dl=partition_dl,
                                                                                            edges_dl=edges_dl,
                                                                                            dense=dense, multigraph=multigraph,
                                                                                            norm=norm)))

        if norm:
            return S / E
        else:
            return S

    def get_matrix(self):
        r"""Returns the block matrix (as a sparse :class:`~scipy.sparse.csr_matrix`),
        which contains the number of edges between each block pair.

        .. warning::

           This corresponds to the adjacency matrix of the block graph, which by
           convention includes twice the amount of edges in the diagonal entries
           if the graph is undirected.

        Examples
        --------

        .. testsetup:: get_matrix

           gt.seed_rng(42)
           np.random.seed(42)
           from pylab import *

        .. doctest:: get_matrix

           >>> g = gt.collection.data["polbooks"]
           >>> state = gt.BlockState(g, B=5, deg_corr=True)
           >>> for i in range(1000):
           ...     ds, nmoves = gt.mcmc_sweep(state)
           >>> m = state.get_matrix()
           >>> figure()
           <...>
           >>> matshow(m.todense())
           <...>
           >>> savefig("bloc_mat.pdf")

        .. testcleanup:: get_matrix

           savefig("bloc_mat.png")

        .. figure:: bloc_mat.*
           :align: center

           A  5x5 block matrix.

       """

        return adjacency(self.bg, weight=self.mrs)


def model_entropy(B, N, E, directed=False, nr=None):
    r"""Computes the amount of information necessary for the parameters of the traditional blockmodel ensemble, for ``B`` blocks, ``N`` vertices, ``E`` edges, and either a directed or undirected graph.

    A traditional blockmodel is defined as a set of :math:`N` vertices which can
    belong to one of :math:`B` blocks, and the matrix :math:`e_{rs}` describes
    the number of edges from block :math:`r` to :math:`s` (or twice that number
    if :math:`r=s` and the graph is undirected).

    For an undirected graph, the number of distinct :math:`e_{rs}` matrices is given by,

    .. math::

       \Omega_m = \left(\!\!{\left(\!{B \choose 2}\!\right) \choose E}\!\!\right)

    and for a directed graph,

    .. math::
       \Omega_m = \left(\!\!{B^2 \choose E}\!\!\right)


    where :math:`\left(\!{n \choose k}\!\right) = {n+k-1\choose k}` is the
    number of :math:`k` combinations with repetitions from a set of size :math:`n`.

    The total information necessary to describe the model is then,

    .. math::

       \mathcal{L}_t = \ln\Omega_m + \ln\left(\!\!{B \choose N}\!\!\right) + \ln N! - \sum_r \ln n_r!,


    where the remaining term is the information necessary to describe the
    block partition, where :math:`n_r` is the number of nodes in block :math:`r`.

    If ``nr`` is ``None``, it is assumed :math:`n_r=N/B`.

    References
    ----------

    .. [peixoto-parsimonious-2013] Tiago P. Peixoto, "Parsimonious module inference in large networks",
       Phys. Rev. Lett. 110, 148701 (2013), :doi:`10.1103/PhysRevLett.110.148701`, :arxiv:`1212.4794`.
    .. [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`.

    """

    if directed:
        x = (B * B);
    else:
        x = (B * (B + 1)) / 2;
    if nr is False:
        L = lbinom(x + E - 1, E)
    else:
        L = lbinom(x + E - 1, E) + partition_entropy(B, N, nr)
    return L

def lbinom(n, k):
    return scipy.special.gammaln(float(n + 1)) - scipy.special.gammaln(float(n - k + 1)) - scipy.special.gammaln(float(k + 1))

def lbinom_careful(n, k):
    return libcommunity.lbinom_careful(n, k)

def lbinom_fast(n, k):
    return libcommunity.lbinom_fast(n, k)

def partition_entropy(B, N, nr=None):
    if nr is None:
        S = N * log(B) + log1p(-(1 - 1./B) ** N)
    else:
        S = lbinom(B + N - 1, N) + scipy.special.gammaln(N + 1) - scipy.special.gammaln(nr + 1).sum()
    return S

def get_max_B(N, E, directed=False):
    r"""Return the maximum detectable number of blocks, obtained by minimizing:

    .. math::

        \mathcal{L}_t(B, N, E) - E\ln B

    where :math:`\mathcal{L}_t(B, N, E)` is the information necessary to
    describe a traditional blockmodel with `B` blocks, `N` nodes and `E`
    edges (see :func:`model_entropy`).

    Examples
    --------

    >>> gt.get_max_B(N=1e6, E=5e6)
    1572

    References
    ----------
    .. [peixoto-parsimonious-2013] Tiago P. Peixoto, "Parsimonious module inference in large networks",
       Phys. Rev. Lett. 110, 148701 (2013), :doi:`10.1103/PhysRevLett.110.148701`, :arxiv:`1212.4794`.


    """

    def Sdl(B, S, N, E, directed=False):
        return S + model_entropy(B, N, E, directed) / E

    B = fminbound(lambda B: Sdl(B, -log(B), N, E, directed), 1, E,
                  xtol=1e-6, maxfun=1500, disp=0)
    if isnan(B):
        B = 1
    return min(N, max(int(ceil(B)), 2))

def get_akc(B, I, N=numpy.inf, directed=False):
    r"""Return the minimum value of the average degree of the network, so that some block structure with :math:`B` blocks can be detected, according to the minimum description length criterion.

    This is obtained by solving

    .. math::

       \Sigma_b = \mathcal{L}_t(B, N, E) - E\mathcal{I}_{t/c} = 0,

    where :math:`\mathcal{L}_{t}` is the necessary information to describe the
    blockmodel parameters (see :func:`model_entropy`), and
    :math:`\mathcal{I}_{t/c}` characterizes the planted block structure, and is
    given by

    .. math::

        \mathcal{I}_t &= \sum_{rs}m_{rs}\ln\left(\frac{m_{rs}}{w_rw_s}\right),\\
        \mathcal{I}_c &= \sum_{rs}m_{rs}\ln\left(\frac{m_{rs}}{m_rm_s}\right),

    where :math:`m_{rs} = e_{rs}/2E` (or :math:`m_{rs} = e_{rs}/E` for directed
    graphs) and :math:`w_r=n_r/N`. We note that :math:`\mathcal{I}_{t/c}\in[0,
    \ln B]`. In the case where :math:`E \gg B^2`, this simplifies to

    .. math::

       \left<k\right>_c &= \frac{2\ln B}{\mathcal{I}_{t/c}},\\
       \left<k^{-/+}\right>_c &= \frac{\ln B}{\mathcal{I}_{t/c}},

    for undirected and directed graphs, respectively. This limit is assumed if
    ``N == inf``.

    Examples
    --------

    >>> gt.get_akc(10, log(10) / 100, N=100)
    2.414413200430159

    References
    ----------
    .. [peixoto-parsimonious-2013] Tiago P. Peixoto, "Parsimonious module inference in large networks",
       Phys. Rev. Lett. 110, 148701 (2013), :doi:`10.1103/PhysRevLett.110.148701`, :arxiv:`1212.4794`.

    """
    if N != numpy.inf:
        if directed:
            get_dl = lambda ak: model_entropy(B, N, N * ak, directed) / N * ak - N * ak * I
        else:
            get_dl = lambda ak: model_entropy(B, N, N * ak / 2., directed) * 2 / (N * ak)  - N * ak * I / 2.
        ak = fsolve(lambda ak: get_dl(ak), 10)
        ak = float(ak)
    else:
        ak = 2 * log(B) / S
        if directed:
            ak /= 2
    return ak

def mcmc_sweep(state, beta=1., c=1., niter=1, dl=False, dense=False,
               multigraph=False, node_coherent=False, confine_layers=False,
               sequential=True, parallel=False, vertices=None,
               target_blocks=None, verbose=False, **kwargs):
    r"""Performs a Markov chain Monte Carlo sweep on the network, to sample the block partition according to a probability :math:`\propto e^{-\beta \mathcal{S}_{t/c}}`, where :math:`\mathcal{S}_{t/c}` is the blockmodel entropy.

    Parameters
    ----------
    state : :class:`~graph_tool.community.BlockState`, :class:`~graph_tool.community.OverlapBlockState` or :class:`~graph_tool.community.CovariateBlockState`
        The 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.
    niter : ``int`` (optional, default: ``1``)
        Number of sweeps to perform.
    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.
    dense : ``bool`` (optional, default: ``False``)
        If ``True``, the "dense" variant of the entropy will be computed.
    multigraph : ``bool`` (optional, default: ``False``)
        If ``True``, the multigraph entropy will be used. Only has an effect
        if ``dense == True``.
    node_coherent : ``bool`` (optional, default: ``False``)
        If ``True``, and if the ``state`` is an instance of
        :class:`~graph_tool.community.OverlapBlockState`, then all half-edges
        incident on the same node are moved simultaneously.
    confine_layers : ``bool`` (optional, default: ``False``)
        If ``True``, and if the ``state`` is an instance of
        :class:`~graph_tool.community.CovariateBlockState`, with an
        *overlapping* partition, the half edges will only be moved in such a way
         that inside each layer the group membership remains non-overlapping.
    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.
    parallel : ``bool`` (optional, default: ``False``)
        If ``True``, the updates are performed in parallel (multiple
        threads).

        .. warning::

            If this is used, the Markov Chain is not guaranteed to be sampled with
            the correct probabilities. This is better used in conjunction with
            ``beta=float('inf')``, where this is not an issue.

    vertices : ``list of ints`` (optional, default: ``None``)
        A list of vertices which will be attempted to be moved. If ``None``, all
        vertices will be attempted.
    target_blocks : ``list of ints`` (optional, default: ``None``)
        A list of groups to which the corresponding vertices will to be forcibly
        moved. If ``None``, the standard MCMC rules will be applied.
    verbose : ``bool`` (optional, default: ``False``)
        If ``True``, verbose information is displayed.

    Returns
    -------

    dS : ``float``
       The entropy difference (in nats) after the sweeps.
    nmoves : ``int``
       The number of accepted block membership moves.


    Notes
    -----

    This algorithm performs a Markov chain Monte Carlo sweep on the network,
    where the block memberships are randomly moved, and either accepted or
    rejected, so that after sufficiently many sweeps the partitions are sampled
    with probability proportional to :math:`e^{-\beta\mathcal{S}_{t/c}}`, where
    :math:`\mathcal{S}_{t/c}` is the blockmodel entropy, given by

    .. math::

      \mathcal{S}_t &\cong - \frac{1}{2} \sum_{rs}e_{rs}\ln\left(\frac{e_{rs}}{n_rn_s}\right), \\
      \mathcal{S}^d_t &\cong - \sum_{rs}e_{rs}\ln\left(\frac{e_{rs}}{n_rn_s}\right),

    for undirected and directed traditional blockmodels (``deg_corr == False``),
    respectively, where :math:`e_{rs}` is the number of edges from block
    :math:`r` to :math:`s` (or the number of half-edges for the undirected case
    when :math:`r=s`), and :math:`n_r` is the number of vertices in block
    :math:`r`, and constant terms which are independent of the block partition
    were dropped (see :meth:`BlockState.entropy` for the complete entropy). For
    the degree-corrected variant with "hard" degree constraints the equivalent
    expressions are

    .. math::

       \mathcal{S}_c &\cong  - \frac{1}{2} \sum_{rs}e_{rs}\ln\left(\frac{e_{rs}}{e_re_s}\right), \\
       \mathcal{S}^d_c &\cong - \sum_{rs}e_{rs}\ln\left(\frac{e_{rs}}{e^+_re^-_s}\right),

    where :math:`e_r = \sum_se_{rs}` is the number of half-edges incident on
    block :math:`r`, and :math:`e^+_r = \sum_se_{rs}` and :math:`e^-_r =
    \sum_se_{sr}` are the number of out- and in-edges adjacent to block
    :math:`r`, respectively.

    The Monte Carlo algorithm employed attempts to improve the mixing time of
    the Markov chain by proposing membership moves :math:`r\to s` with
    probability :math:`p(r\to s|t) \propto e_{ts} + c`, where :math:`t` is the
    block label of a random neighbour of the vertex being moved. See
    [peixoto-efficient-2014]_ for more details.

    This algorithm has a complexity of :math:`O(E)`, where :math:`E` is the
    number of edges in the network.

    Examples
    --------
    .. testsetup:: mcmc

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

    .. doctest:: mcmc

       >>> g = gt.collection.data["polbooks"]
       >>> state = gt.BlockState(g, B=3, deg_corr=True)
       >>> pv = None
       >>> for i in range(1000):        # remove part of the transient
       ...     ds, nmoves = gt.mcmc_sweep(state)
       >>> for i in range(1000):
       ...     ds, nmoves = gt.mcmc_sweep(state)
       ...     pv = gt.collect_vertex_marginals(state, pv)
       >>> gt.graph_draw(g, pos=g.vp["pos"], vertex_shape="pie", vertex_pie_fractions=pv, output="polbooks_blocks_soft.pdf")
       <...>

    .. testcleanup:: mcmc

       gt.graph_draw(g, pos=g.vp["pos"], vertex_shape="pie", vertex_pie_fractions=pv, output="polbooks_blocks_soft.png")

    .. figure:: polbooks_blocks_soft.*
       :align: center

       "Soft" block partition of a political books network with :math:`B=3`.

     References
    ----------

    .. [holland-stochastic-1983] Paul W. Holland, Kathryn Blackmond Laskey,
       Samuel Leinhardt, "Stochastic blockmodels: First steps",
       Carnegie-Mellon University, Pittsburgh, PA 15213, U.S.A.,
       :doi:`10.1016/0378-8733(83)90021-7`
    .. [faust-blockmodels-1992] Katherine Faust, and Stanley
       Wasserman. "Blockmodels: Interpretation and Evaluation." Social Networks
       14, no. 1-2 (1992): 5-61. :doi:`10.1016/0378-8733(92)90013-W`
    .. [karrer-stochastic-2011] Brian Karrer, and M. E. J. Newman. "Stochastic
       Blockmodels and Community Structure in Networks." Physical Review E 83,
       no. 1 (2011): 016107. :doi:`10.1103/PhysRevE.83.016107`.
    .. [peixoto-entropy-2012] Tiago P. Peixoto "Entropy of Stochastic Blockmodel
       Ensembles." Physical Review E 85, no. 5 (2012): 056122.
       :doi:`10.1103/PhysRevE.85.056122`, :arxiv:`1112.6028`.
    .. [peixoto-parsimonious-2013] Tiago P. Peixoto, "Parsimonious module
       inference in large networks", Phys. Rev. Lett. 110, 148701 (2013),
       :doi:`10.1103/PhysRevLett.110.148701`, :arxiv:`1212.4794`.
    .. [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-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`.
    """