def stationary_distribution(edges=None, exact=False, logspace=False,
                            initial_state=None, iterations=None, lim=1e-8,
                            states=None):
    """
    Convenience function to route to different stationary distribution
    computations.

    Parameters
    ----------
    edges: list of tuples or function
        The transitions of the process, either a list of (source, target,
        transition_probability), or an edge_function that takes two parameters,
        the source and target states, to the transition transition probability.
        If using an edge_function you must supply the states of the process.
    exact: bool, False
        Try to use the exact formula. Only works for reversible processes (no
        check) on a discretization of simplex
    logspace: bool False
        Carry out the calculation in logspace
    iterations: int, None
        Maximum number of iterations for stationary approximation
    lim: float, 1e-8
        Approximate algorithm breaks when successive iterations have a
        kl_divergence less than lim
    initial_state, None
        A distribution over the states of the process. If None, the uniform
        distiribution is used.
    states: list, None
        States for use with the edge_function.
    """

    if isinstance(edges, list):
        if not exact:
            return approx_stationary(
                edges, logspace=logspace, iterations=iterations, lim=lim,
                initial_state=initial_state)
        else:
            return exact_stationary(
                edges, initial_state=initial_state, logspace=logspace)
    elif isinstance(edges, Callable):
        if not states:
            raise ValueError(
                "Keyword argument `states` required with edge_func")
        return approx_stationary_func(
            edges, states, iterations=iterations, lim=lim, logspace=logspace)
    # Still here?
    raise Exception("Parameter combination not implemented") def is_valid_state(state, lower, upper):
def multivariate_transitions(N, incentive, num_types=3, mu=0.001,
                             no_boundary=False):
def multivariate_transitions_gen(N, incentive, num_types=3, mu=0.001,
                                 no_boundary=False):
def log_multivariate_transitions(N, logincentive, num_types=3, mu=0.001,
                                 no_boundary=False):
def compute_edges(N=30, num_types=None, m=None, incentive_func=logit, beta=1.,
                  q=1., mu=None):
def k_fold_incentive_transitions(N, incentive, num_types, mu=None, k=None):
def neutral_stationary(N, alpha, n=3, logspace=False):
def log_neutral_stationary(N, alpha, n=3): def cache_multinomial_coefficients(N, num_types=3):
def multivariate_transitions_sub(N, incentive, mu=0.001, low_memory=False):
    def multinomial_probability(xs, ps):
    def g(current_state, next_state):
    def h(current_state, next_state):
def multivariate_transitions(N, incentive, mu=0.001, num_types=3,
                             low_memory=False):
    def multinomial_probability(xs, ps):
    def g(current_state, next_state):
    def h(current_state, next_state):
        M = numpy.zeros(shape=N+1)
        M = numpy.zeros(shape=(N+1, N+1))
    M = cache_multinomial_coefficients(N, num_types=num_types)
    M = cache_multinomial_coefficients(N, num_types=num_types) def simplex_generator(N, d=2):
    """
    Generates a discretation of the simplex.

    Parameters
    ----------
    N: int
        The number of subdivsions in each dimension
    d: int, 2
        The dimension of the simplex (the number of population types is d+1

    Yields
    ------
    (d+1)-tuples of numbers summing to N. The total number of yielded tuples is
    equal to the simplicial polytopic number corresponding to N and d,
    binom{N + d - 1}{d} (see https://en.wikipedia.org/wiki/Figurate_number )
    """

    if d == 1:
        for i in range(N+1):
            yield (i, N - i)
    if d > 1:
        for j in range(N+1):
            for s in simplex_generator(N - j, d - 1):
                t = [j]
                t.extend(s)
                yield tuple(t) def slice_dictionary(d, N, slice_index=0, slice_value=0):
    """
    Take a three dimensional slice from a four dimensional
    dictionary.
    """

    slice_dict = dict()
    for state in simplex_generator(N, 2):
        new_state = list(state)
        new_state.insert(slice_index, slice_value)
        slice_dict[state] = d[tuple(new_state)]
    return slice_dict def edges_to_edge_dict(edges):
    """
    Converts a list of edges to a transition dictionary taking (source, target)
    to the transition between the states.

    Parameters
    ----------
    edges: list of tuples
        Transition probabilities of the form [(source, target, transition
        probability

    Returns
    -------
    edges, maps 2-tuples of nodes to floats
    """

    edge_dict = dict()
    for e1, e2, v in edges:
        edge_dict[(e1, e2)] = v
    return edge_dict def plot_dictionary(s, ax=None):
    """
    Plot two or three dimensional dictionary on a simplex partition, such as a
    stationary distribution.
    """

    num_types = len(s.keys()[0])
    N = sum(s.keys()[0])

    if not ax:
        fig, ax = pyplot.subplots()

    if num_types == 2:
        domain = list(range(0, N+1))
        values = []
        for i in domain:
            state = (i, N-i)
            values.append(s[state])
        if not ax:
            fig, ax = pyplot.subplots()
        pyplot.plot(domain, values)

    if num_types == 3:
        fig, tax = ternary.figure(scale=N, ax=ax)
        tax.heatmap(s, style='d') def expected_divergence(edges, states=None, q_d=1, boundary=True):
    """
    Computes the KL-div of the expected state with the state, for all states.

    Parameters
    ----------
    edges: list of tuples or function
        The transitions of the process, either a list of (source, target,
        transition_probability), or an edge_function that takes two parameters,
        the source and target states, to the transition transition probability.
        If using an edge_function you must supply the states of the process.
    q_d: float, 1
        parameter that specifies which divergence function to use
    states: list, None
        States for use with the edge_func
    boundary: bool, False
        Exclude the boundary states

    Returns
    -------
    Dictionary mapping states to D(E(state), state)
    """

    dist = q_divergence(q_d)
    e = defaultdict(float)

    if isinstance(edges, list):
        for x, y, w in edges:
            e[x] += numpy.array(y) * w
    elif isinstance(edges, Callable):
        if not states:
            raise ValueError(
                "Keyword argument `states` required with edge_func")
        for x in states:
            e[x] = 0.
            for y in states:
                w = edges(x, y)
                e[x] += numpy.array(y) * w
    d = dict()
    for state, v in e.items():
        # Some divergences do not play well on the boundary
        if not boundary:
            p = 1.
            for s in state:
                p *= s
            if p == 0:
                continue
        d[state] = dist(normalize(v), normalize(list(state)))
    return d def bomze_matrices(filename="bomze.txt"):
    """
    Yields the 48 matrices from I.M. Bomze's classification of three player phase
    portraits.
    """

    this_dir, this_filename = os.path.split(__file__)

    handle = open(os.path.join(this_dir, filename))
    for line in handle:
        a, b, c, d, e, f, g, h, i = map(float, line.split())
        yield [[a, b, c], [d, e, f], [g, h, i]] def ensure_directory(directory):
    if not os.path.isdir(directory):
        os.mkdir(directory)
import math
import os
import matplotlib
import matplotlib.gridspec as gridspec
import ternary
from matplotlib import pyplot
from stationary import stationary_distribution
from stationary.processes.incentives import *
from stationary.processes import incentive_process, wright_fisher
from stationary.utils.math_helpers import simplex_generator, slice_dictionary
from stationary.utils.edges import edges_to_edge_dict
from stationary.utils.plotting import plot_dictionary
from stationary.utils import expected_divergence
from stationary.utils.bomze import bomze_matrices
from stationary.utils.files import ensure_directory
    edges = incentive_process.multivariate_transitions(N, incentive, num_types=num_types, mu=mu)

    d1 = expected_divergence(edges, q_d=0, boundary=True)
    d2 = stationary_distribution(edges)

    # We need to slice the 4dim dictionary into three-dim slices for plotting.
    for slice_index in range(4):
        for d in [d1, d2]:
            slice_dict = slice_dictionary(d, N, slice_index=3)
            figure, tax = ternary.figure(scale=N)
            tax.heatmap(slice_dict, style="d")
    pyplot.show()


def graphical_abstract_figures(N=60, q=1, beta=0.1):
    """
    Three dimensional process examples.
    """

    a = 0
    b = 1
    m = [[a, b, b], [b, a, b], [b, b, a]]
    mu = (3. / 2 ) * 1. / N
    fitness_landscape = linear_fitness_landscape(m)
    incentive = fermi(fitness_landscape, beta=beta, q=q)
    edges = incentive_process.multivariate_transitions(N, incentive, num_types=3, mu=mu)
    d = stationary_distribution(edges, iterations=None)

    figure, tax = ternary.figure(scale=N)
    tax.heatmap(d, scale=N)
    tax.savefig(filename="ga_stationary.eps", dpi=600)

    d = expected_divergence(edges, q_d=0)
    figure, tax = ternary.figure(scale=N)
    tax.heatmap(d, scale=N)
    tax.savefig(filename="ga_d_0.eps", dpi=600)

    d = expected_divergence(edges, q_d=1)
    figure, tax = ternary.figure(scale=N)
    tax.heatmap(d, scale=N)
    tax.savefig(filename="ga_d_1.eps", dpi=600)


def rps_figures(N=60, q=1, beta=1.):
    """
    Three rock-paper-scissors examples.
    """

    m = [[0, -1, 1], [1, 0, -1], [-1, 1, 0]]
    num_types = len(m[0])
    fitness_landscape = linear_fitness_landscape(m)
    for i, mu in enumerate([1./math.sqrt(N), 1./N, 1./N**(3./2)]):
        # Approximate calculation
        mu = 3/2. * mu
        incentive = fermi(fitness_landscape, beta=beta, q=q)
        edges = incentive_process.multivariate_transitions(N, incentive, num_types=num_types, mu=mu)
        d = stationary_distribution(edges, lim=1e-10)

        figure, tax = ternary.figure()
        tax.heatmap(d, scale=N)
        tax.savefig(filename="rsp_mu_" + str(i) + ".eps", dpi=600)


def tournament_stationary_3(N, mu=None):
    """
    Example for a tournament selection matrix.
    """

    if not mu:
        mu = 3./2 * 1./N
    m = [[1,1,1], [0,1,1], [0,0,1]]
    num_types = len(m[0])
    fitness_landscape = linear_fitness_landscape(m)
    incentive = replicator(fitness_landscape)
    edges = incentive_process.multivariate_transitions(N, incentive, num_types=num_types, mu=mu)
    s = stationary_distribution(edges)
    ternary.heatmap(s, scale=N, scientific=True)
    d = expected_divergence(edges, q_d=0)
    ternary.heatmap(d, scale=N, scientific=True)
    pyplot.show()


def two_dim_transitions(edges):
    """
    Compute the + and - transitions for plotting in two_dim_transitions_figure
    """

    d = edges_to_edge_dict(edges)
    N = sum(edges[0][0])
    ups = []
    downs = []
    stays = []
    for i in range(0, N+1):
        try:
            up = d[((i, N-i), (i+1, N-i-1))]
        except KeyError:
            up = 0
        try:
            down = d[((i, N-i), (i-1, N-i+1))]
        except KeyError:
            down = 0
        ups.append(up)
        downs.append(down)
        stays.append(1 - up - down)
    return ups, downs, stays


def two_dim_transitions_figure(N, m, mu=0.01, incentive_func=replicator):
    """
    Plot transition entropies and stationary distributions.
    """

    n = len(m[0])
    fitness_landscape = linear_fitness_landscape(m)
    incentive = incentive_func(fitness_landscape)
    if not mu:
        mu = 1./ N
    edges = incentive_process.multivariate_transitions(N, incentive, num_types=n, mu=mu)

    s = stationary_distribution(edges, exact=True)

    d = edges_to_edge_dict(edges)