Source code for dit.algorithms.maxentropy

"""
Maximum entropy with marginal distribution constraints.

Note: We are actually doing the maximum entropy optimization. So we have not
built in the fact that the solution is an exponential family.

Also, this doesn't seem to work that well in practice. The optimization
simply fails to converge for many distributions. Xor() works great, but And()
fails to converge for 2-way marginals. Random distributions seem to work.
Jittering the distributions sometimes helps.

We might need to assume the exponential form and then fit the params to match
the marginals. Perhaps exact gradient and Hessians might help, or maybe even
some rescaling of the linear constraints.


TODO:

This code for moment-based maximum entropy needs to be updated so that it can
handle any Cartesian product sample space, rather than just homogeneous ones.

"""

import itertools

import numpy as np
from debtcollector import removals
from loguru import logger

import dit
from dit.abstractdist import AbstractDenseDistribution, get_abstract_dist

from ..helpers import parse_rvs
from ..utils import flatten
from .optutil import Bunch, CVXOPT_Template, _as_full_rank, as_full_rank, prepare_dist

# from ..utils import powerset

__all__ = (
    # 'MarginalMaximumEntropy',
    # 'MomentMaximumEntropy',
    # Use version provided by maxentropyfw.py
    # 'marginal_maxent_dists',
    # 'moment_maxent_dists',
    "degrees_of_freedom",
)


@removals.remove(message="Please see methods in dit.algorithms.distribution_optimizers.py.", version="1.0.1")
def isolate_zeros_generic(dist, rvs):
    """
    Determines if there are any elements of the optimization vector that must
    be zero.

    If p(marginal) = 0, then every component of the joint that contributes to
    that marginal probability must be exactly zero for all feasible solutions.
    """
    assert dist.is_dense()
    assert dist.get_base() == "linear"

    rvs_, indexes = parse_rvs(dist, set(flatten(rvs)), unique=True, sort=True)
    rvs = [[indexes[rvs_.index(rv)] for rv in subrv] for subrv in rvs]

    d = get_abstract_dist(dist)
    n_elements = d.n_elements

    zero_elements = np.zeros(n_elements, dtype=int)
    cache = {}
    pmf = dist.pmf

    for subrvs in rvs:
        marray = d.parameter_array(subrvs, cache=cache)
        for idx in marray:
            # Convert the sparse nonzero elements to a dense boolean array
            bvec = np.zeros(n_elements, dtype=int)
            bvec[idx] = 1
            p = pmf[idx].sum()
            if np.isclose(p, 0):
                zero_elements += bvec

    zero = np.nonzero(zero_elements)[0]
    zeroset = set(zero)
    nonzero = [i for i in range(n_elements) if i not in zeroset]
    variables = Bunch(nonzero=nonzero, zero=zero)

    return variables


@removals.remove(message="Please see methods in dit.algorithms.distribution_optimizers.py.", version="1.0.1")
def isolate_zeros(dist, k):
    """
    Determines if there are any elements of the optimization vector that must
    be zero.

    If :math:`p(marginal) = 0`, then every component of the joint that
    contributes to that marginal probability must be exactly zero for all
    feasible solutions.
    """
    assert dist.is_dense()
    assert dist.get_base() == "linear"

    d = get_abstract_dist(dist)
    n_variables = d.n_variables
    n_elements = d.n_elements

    rvs = range(n_variables)
    zero_elements = np.zeros(n_elements, dtype=int)
    cache = {}
    pmf = dist.pmf
    if k > 0:
        for subrvs in itertools.combinations(rvs, k):
            marray = d.parameter_array(subrvs, cache=cache)
            for idx in marray:
                # Convert the sparse nonzero elements to a dense boolean array
                bvec = np.zeros(n_elements, dtype=int)
                bvec[idx] = 1
                p = pmf[idx].sum()
                if np.isclose(p, 0):
                    zero_elements += bvec

    zero = np.nonzero(zero_elements)[0]
    zeroset = set(zero)
    nonzero = [i for i in range(n_elements) if i not in zeroset]
    variables = Bunch(nonzero=nonzero, zero=zero)

    return variables


def marginal_constraints_generic(dist, rvs, with_normalization=True):
    """
    Returns :math:`A` and :math:`b` in :math:`A x = b`, for a system of marginal
    constraints.

    In general, the resulting matrix :math:`A` will not have full rank.

    Parameters
    ----------
    dist : distribution
        The distribution used to calculate the marginal constraints.

    rvs : sequence
        A sequence whose elements are also sequences.  Each inner sequence
        specifies a marginal distribution as a set of random variable from
        `dist`. The inner sequences need not be pairwise mutually exclusive
        with one another. A random variable can only appear once within
        each inner sequence, but it can occur in multiple inner sequences.
    """
    assert dist.is_dense()
    assert dist.get_base() == "linear"

    parse = lambda rv: parse_rvs(dist, rv, unique=True, sort=True)[1]

    # potential inclusion: include implied constraints
    # rvs = set().union(*[set(r for r in powerset(rv) if r) for rv in rvs])
    indexes = [parse(rv) for rv in rvs]

    pmf = dist.pmf

    d = get_abstract_dist(dist)

    A = []
    b = []

    # Begin with the normalization constraint.
    if with_normalization:
        A.append(np.ones(d.n_elements))
        b.append(1)

    # Now add all the marginal constraints.
    cache = {}
    for rvec in indexes:
        for idx in d.parameter_array(rvec, cache=cache):
            bvec = np.zeros(d.n_elements)
            bvec[idx] = 1
            A.append(bvec)
            b.append(pmf[idx].sum())

    A = np.asarray(A, dtype=float)
    b = np.asarray(b, dtype=float)

    return A, b


def marginal_constraints(dist, m, with_normalization=True):
    """
    Returns :math:`A` and :math:`b` in :math:`A x = b`, for a system of marginal
    constraints.

    The resulting matrix :math:`A` is not guaranteed to have full rank.

    Parameters
    ----------
    dist : distribution
        The distribution from which the marginal constraints are constructed.

    m : int
        The size of the marginals to constrain. When `m=2`, pairwise marginals
        are constrained to equal the pairwise marginals in `pmf`. When `m=3`,
        three-way marginals are constrained to equal those in `pmf.

    with_normalization : bool
        If true, include a constraint for normalization.

    Returns
    -------
    A : array-like, shape (p, q)
        The matrix defining the marginal equality constraints and also the
        normalization constraint. The number of rows is:
            :math:`p = C(n_variables, m) * n_symbols ** m + 1`
        where C() is the choose formula. The number of columns is:
            :math:`q = n_symbols ** n_variables`

    b : array-like, (p,)
        The RHS of the linear equality constraints.
    """
    n_variables = dist.outcome_length()

    if m > n_variables:
        msg = "Cannot constrain {0}-way marginals"
        msg += " with only {1} random variables."
        msg = msg.format(m, n_variables)
        raise ValueError(msg)

    variables = dist.get_rv_names()
    if variables is not None:
        rvs = list(itertools.combinations(variables, m))
    else:
        rvs = list(itertools.combinations(range(n_variables), m))

    A, b = marginal_constraints_generic(dist, rvs, with_normalization=with_normalization)
    return A, b


def marginal_constraint_rank(dist, m):
    """
    Returns the rank of the marginal constraint matrix.
    """
    dist = prepare_dist(dist)

    A, b = marginal_constraints(dist, m)
    _, _, rank = _as_full_rank(A, b)
    return rank


def marginal_constraint_rank_generic(dist, rvs):
    """
    Returns the rank of the marginal constraint matrix for an arbitrary set of
    marginals.

    Parameters
    ----------
    dist : distribution
        The distribution used to calculate the marginal constraints.
    rvs : sequence
        A sequence whose elements are also sequences. Each inner sequence
        specifies a marginal distribution as a set of random variables from
        `dist`.

    Returns
    -------
    rank : int
        The rank of the marginal constraint matrix, including the normalization
        constraint.
    """
    dist = prepare_dist(dist)

    A, _ = marginal_constraints_generic(dist, rvs)
    return int(np.linalg.matrix_rank(A))


[docs] def degrees_of_freedom(dist, structure=None): """ Returns the degrees of freedom of a marginal model: the number of free parameters needed to specify the maximum entropy distribution consistent with the marginals in `structure`. This is the "complexity" of a reconstructability-analysis model. It is the rank of the marginal constraint matrix (which includes the normalization constraint) minus one for that normalization constraint. Parameters ---------- dist : distribution The distribution defining the sample space and the marginal values. structure : sequence, None A sequence whose elements are also sequences, each specifying a marginal (a "projection") to constrain, e.g. ``[[0, 1], [1, 2]]`` for ``AB:BC``. If None, the independence model (each variable on its own) is used. Returns ------- df : int The degrees of freedom of the model. Examples -------- >>> d = dit.uniform_distribution(3, 2) >>> dit.algorithms.degrees_of_freedom(d, [[0, 1], [2]]) 4 >>> dit.algorithms.degrees_of_freedom(d, [[0, 1], [1, 2]]) 5 """ if structure is None: structure = [[v] for v in range(dist.outcome_length())] return marginal_constraint_rank_generic(dist, structure) - 1
def moment(f, pmf, center=0, n=1): """ Return the nth moment of `f` about `center`, distributed by `pmf`. Explicitly: :math:`\\sum_i (f(i) - center)**n p(i)` Note, `pmf` is the joint distribution. So n=1 can be used even when calculating covariances such as <xx> and <xy>. The first would actually be a 2nd moment, while the second would be a mixed 1st moment. Parameters ---------- f : array-like The numerical values assigned to each outcome of `p`. pmf : array-like The pmf for a distribution, linear-distributed values. center : float Calculate a centered moment. n : int The moment to calculate. """ return ((f - center) ** n * pmf).sum() def moment_constraints(pmf, n_variables, m, symbol_map, with_replacement=True): """ Returns :math:`A` and :math:`b` in :math:`A x = b`, for an Ising-like system. If without replacement, we include only m-way first-moment constraints where each element is distinct. So <xx> and <yy> would not be included if `n_variables = 2` and `m = 2`. The function we take means of is: :math:`f(x) = \\prod_i x_i` The resulting matrix :math:`A` is not guaranteed to have full rank. Parameters ---------- pmf : array-like, shape ( n_symbols ** n_variables, ) The probability mass function of the distribution. The pmf must have a Cartesian product sample space with the same sample space used for each random variable. n_variables : int The number of random variables. m : int | list The size of the moments to constrain. When `m = 2`, pairwise means are constrained to equal the pairwise means in `pmf`. When `m = 3`, three-way means are constrained to equal those in `pmf. If m is a list, then include all m-way moments in the list. symbol_map : array-like A mapping from the ith symbol to a real number that is to be used in the calculation of moments. For example, `symbol_map=[-1, 1]` corresponds to the typical Ising model. with_replacement : bool If `True`, variables are selected with replacement. The standard Ising does not select with replacement, and so terms like <xx>, <yy> do not appear for m=2. When `True`, we are constraining the entire moment matrix. Returns ------- A : array-like, shape (p, q) The matrix defining the marginal equality constraints and also the normalization constraint. The number of rows is: :math:`p = C(n_variables, m) * n_symbols ** m + 1` where C() is the choose formula. The number of columns is: :math:`q = n_symbols ** n_variables` b : array-like, (p,) The RHS of the linear equality constraints. """ n_symbols = len(symbol_map) d = AbstractDenseDistribution(n_variables, n_symbols) if len(pmf) != d.n_elements: msg = "Length of `pmf` != n_symbols ** n_variables. Symbol map: {0!r}" raise ValueError(msg.format(symbol_map)) # Begin with the normalization constraint. A = [np.ones(d.n_elements)] b = [1] try: m[0] except TypeError: mvals = [m] except IndexError: # m is empty list pass else: mvals = m combinations = itertools.combinations_with_replacement if with_replacement else itertools.combinations # Now add all the moment constraints. for m in mvals: if m < 1: continue outcomes = list(itertools.product(symbol_map, repeat=n_variables)) outcomes = np.asarray(outcomes) for rvs in combinations(range(n_variables), m): # Make it a list for NumPy indexing rvs = list(rvs) f = np.array([outcome[rvs].prod() for outcome in outcomes]) mean = moment(f, pmf, n=1) A.append(f) b.append(mean) A = np.asarray(A, dtype=float) b = np.asarray(b, dtype=float) return A, b def moment_constraint_rank(dist, m, symbol_map=None, cumulative=True, with_replacement=True): """ Returns the rank of the moment constraint matrix. """ mvals = range(m + 1) if cumulative else [m] dist = prepare_dist(dist) n_variables = dist.outcome_length() n_symbols = len(dist.alphabet[0]) pmf = dist.pmf # Symbol map if symbol_map is None: symbol_map = range(n_symbols) A, b = moment_constraints(pmf, n_variables, mvals, symbol_map, with_replacement=with_replacement) _, _, rank = _as_full_rank(A, b) return rank def ising_constraint_rank(dist, m, symbol_map=None, cumulative=True): """ Returns the rank of the Ising constraint matrix. """ return moment_constraint_rank(dist, m, symbol_map, cumulative, with_replacement=False) def negentropy(p): """ Entropy which operates on vectors of length `N`. """ negH = np.nansum(p * np.log2(p)) return negH @removals.removed_class( "MaximumEntropy", replacement="dit.algorithms.scipy_optimizers.MaxEntOptimizer", message="Please see methods in dit.algorithms.distribution_optimizers.py.", version="1.0.1", ) class MaximumEntropy(CVXOPT_Template): """ Find maximum entropy distribution. """ def build_function(self): self.func = negentropy @removals.removed_class( "MarginalMaximumEntropy", replacement="dit.algorithms.scipy_optimizers.MaxEntOptimizer", message="Please see methods in dit.algorithms.distribution_optimizers.py.", version="1.0.1", ) class MarginalMaximumEntropy(MaximumEntropy): """ Find maximum entropy distribution subject to `k`-way marginal constraints. `k = 0` should reproduce the behavior of MaximumEntropy. """ def __init__(self, dist, k, tol=None, prng=None): """ Initialize optimizer. Parameters ---------- dist : distribution The distribution used to specify the marginal constraints. k : int The number of variables in the constrained marginals. """ self.k = k super().__init__(dist, tol=tol, prng=prng) def prep(self): # We are only removing elements which should be fixed at zero. # This means they don't contribute to the entropy, so there is no # need to adjust the function. Also, we are using numdifftools. self.variables = isolate_zeros(self.dist, self.k) # Make self.n reflect only the size of the nonzero elements. This # automatically adjusts the size of G for the inequality constraint. self.n = len(self.variables.nonzero) # pylint: disable=no-member def build_linear_equality_constraints(self): from cvxopt import matrix A, b = marginal_constraints(self.dist, self.k) # Reduce the size of the constraint matrix # Since we are only removing elements which are exactly zero, then # the constraint equations are unchanged. E.g. the normalization is # still that the nonzero values should add to 1. Asmall = A[:, self.variables.nonzero] # pylint: disable=no-member Asmall, b, rank = as_full_rank(Asmall, b) if rank > Asmall.shape[1]: raise ValueError("More independent constraints than free parameters.") Asmall = matrix(Asmall) b = matrix(b) # now a column vector self.A = Asmall self.b = b def initial_dist(self): from .maxentropyfw import initial_point initial_x, _ = initial_point( self.dist, self.k, A=self.A, b=self.b, isolated=self.variables, show_progress=False ) return initial_x def build_gradient_hessian(self): ln2 = np.log(2) def gradient(xarr): # This operates only on nonzero elements. # All of the optimization elements should be greater than zero # But occasional they might go slightly negative or zero. # In those cases, we will just set the gradient to zero and keep the # value fixed from that point forward. bad_x = xarr <= 0 grad = np.log2(xarr) + 1 / ln2 grad[bad_x] = 0 return grad def hessian(xarr): bad_x = xarr <= 0 diag = 1 / xarr / ln2 diag[bad_x] = 0 return np.diag(diag) self.gradient = gradient self.hessian = hessian @removals.removed_class( "MomentMaximumEntropy", replacement="dit.algorithms.scipy_optimizers.MaxEntOptimizer", message="Please see methods in dit.algorithms.distribution_optimizers.py.", version="1.0.1", ) class MomentMaximumEntropy(MaximumEntropy): """ Find maximum entropy distribution subject to `k`-way marginal constraints. `k = 0` should reproduce the behavior of MaximumEntropy. """ def __init__(self, dist, k, symbol_map, cumulative=True, with_replacement=True, tol=None, prng=None): """ Initialize optimizer. Parameters ---------- dist : distribution The distribution used to specify the marginal constraints. k : int The number of variables in the constrained marginals. symbol_map : list The mapping from states to real numbers. This is used while taking moments. cumulative : bool If `True`, include all moments less than or equal to `k`. with_replacement : bool If `True`, then variables are selected for moments with replacement. The standard Ising model selects without replacement. tol : float | None The desired convergence tolerance. prng : RandomState A pseudorandom number generator. """ self.k = k self.symbol_map = symbol_map self.cumulative = cumulative self.with_replacement = with_replacement super().__init__(dist, tol=tol, prng=prng) def build_linear_equality_constraints(self): from cvxopt import matrix # Dimension of optimization variable n = self.n k = range(self.k + 1) if self.cumulative else [self.k] args = (self.pmf, self.n_variables, k, self.symbol_map) kwargs = {"with_replacement": self.with_replacement} A, b = moment_constraints(*args, **kwargs) AA, bb, rank = _as_full_rank(A, b) if rank > n: raise ValueError("More independent constraints than parameters.") AA = matrix(AA) bb = matrix(bb) # now a column vector self.A = AA self.b = bb @removals.remove(message="Please see methods in dit.algorithms.distribution_optimizers.py.", version="1.0.1") def marginal_maxent_dists(dist, k_max=None, jitter=True, show_progress=True): """ Return the marginal-constrained maximum entropy distributions. Parameters ---------- dist : distribution The distribution used to constrain the maxent distributions. k_max : int The maximum order to calculate. jitter : bool | float When `True` or a float, we perturb the distribution slightly before proceeding. This can sometimes help with convergence. show-progress : bool If `True`, show convergence progress to stdout. """ dist = prepare_dist(dist) if jitter: # This is sometimes necessary. If your distribution does not have # full support than convergence can be difficult to come by. dist.pmf = dit.math.pmfops.jittered(dist.pmf) n_variables = dist.outcome_length() if k_max is None: k_max = n_variables outcomes = list(dist.sample_space()) dists = [] for k in range(k_max + 1): logger.info("Constraining maxent dist to match {k}-way marginals.", k=k) opt = MarginalMaximumEntropy(dist, k) pmf_opt = opt.optimize(show_progress=show_progress) pmf_opt = pmf_opt.reshape(pmf_opt.shape[0]) pmf = np.zeros(len(dist.pmf)) pmf[opt.variables.nonzero] = pmf_opt # pylint: disable=no-member d = dit.Distribution(outcomes, pmf) dists.append(d) return dists @removals.remove(message="Please see methods in dit.algorithms.distribution_optimizers.py.", version="1.0.1") def moment_maxent_dists(dist, symbol_map, k_max=None, jitter=True, with_replacement=True, show_progress=True): """ Return the marginal-constrained maximum entropy distributions. Parameters ---------- dist : distribution The distribution used to constrain the maxent distributions. symbol_map : iterable A list whose elements are the real values that each state is assigned while calculating moments. Typical values are [-1, 1] or [0, 1]. k_max : int The maximum order to calculate. jitter : bool | float When `True` or a float, we perturb the distribution slightly before proceeding. This can sometimes help with convergence. with_replacement : bool If `True`, then variables are selected for moments with replacement. The standard Ising model selects without replacement. show-progress : bool If `True`, show convergence progress to stdout. """ dist = prepare_dist(dist) if jitter: # This is sometimes necessary. If your distribution does not have # full support than convergence can be difficult to come by. dist.pmf = dit.math.pmfops.jittered(dist.pmf) n_variables = dist.outcome_length() symbols = dist.alphabet[0] if k_max is None: k_max = n_variables outcomes = list(dist._product(symbols, repeat=n_variables)) text = "with replacement" if with_replacement else "without replacement" dists = [] for k in range(k_max + 1): logger.info("Constraining maxent dist to match {k}-way moments, {text}.", k=k, text=text) opt = MomentMaximumEntropy(dist, k, symbol_map, with_replacement=with_replacement) pmf_opt = opt.optimize(show_progress=show_progress) pmf_opt = pmf_opt.reshape(pmf_opt.shape[0]) d = dit.Distribution(outcomes, pmf_opt) dists.append(d) return dists