continuous_distribution_measures.py

"""

"""

import sympy as sp

from lbmpy.moments import polynomial_to_exponent_representation
from pystencils.cache import disk_cache, memorycache
from pystencils.sympyextensions import complete_the_squares_in_exp


@memorycache()
def moment_generating_function(generating_function, symbols, symbols_in_result):
    """
    Computes the moment generating function of a probability distribution. It is defined as:

    .. math ::
        F[f(\mathbf{x})](\mathbf{t}) = \int e^{<\mathbf{x}, \mathbf{t}>} f(x)\; dx

    :param generating_function: sympy expression
    :param symbols: a sequence of symbols forming the vector x
    :param symbols_in_result: a sequence forming the vector t
    :return: transformation result F: an expression that depends now on symbols_in_result
             (symbols have been integrated out)

    .. note::
         This function uses sympys symbolic integration mechanism, which may not work or take a large
         amount of time for some functions.
         Therefore this routine does some transformations/simplifications on the function first, which are
         taylored to expressions of the form exp(polynomial) i.e. Maxwellian distributions, so that these kinds
         of functions can be integrated quickly.

    """
    assert len(symbols) == len(symbols_in_result)
    for t_i, v_i in zip(symbols_in_result, symbols):
        generating_function *= sp.exp(t_i * v_i)

    # This is a custom transformation that speeds up the integrating process
    # of a MaxwellBoltzmann distribution
    # without this transformation the symbolic integration is sometimes not possible (e.g. in 2D without assumptions)
    # or is really slow
    # other functions should not be affected by this transformation
    # Without this transformation the following assumptions are required for the u and v variables of Maxwell Boltzmann
    #  2D: real=True ( without assumption it will not work)
    #  3D: no assumption ( with assumptions it will not work )
    generating_function = complete_the_squares_in_exp(generating_function.simplify(), symbols)
    generating_function = generating_function.collect(symbols)

    bounds = [(s_i, -sp.oo, sp.oo) for s_i in symbols]
    result = sp.integrate(generating_function, *bounds)

    return sp.simplify(result)


def cumulant_generating_function(func, symbols, symbols_in_result):
    """
    Computes cumulant generating func, which is the logarithm of the moment generating func.
    For parameter description see :func:`moment_generating_function`.
    """
    return sp.ln(moment_generating_function(func, symbols, symbols_in_result))


@disk_cache
def multi_differentiation(generating_function, index, symbols):
    """
    Computes moment from moment-generating function or cumulant from cumulant-generating function,
    by differentiating the generating function, as specified by index and evaluating the derivative at symbols=0

    :param generating_function: function with is differentiated
    :param index: the i'th index specifies how often to differentiate w.r.t. to symbols[i]
    :param symbols: symbol to differentiate
    """
    assert len(index) == len(symbols), "Length of index and length of symbols has to match"

    diff_args = []
    for order, t_i in zip(index, symbols):
        for i in range(order):
            diff_args.append(t_i)

    if len(diff_args) > 0:
        r = sp.diff(generating_function, *diff_args)
    else:
        r = generating_function

    for t_i in symbols:
        r = r.subs(t_i, 0)

    return r


@memorycache(maxsize=512)
def __continuous_moment_or_cumulant(func, moment, symbols, generating_function):
    if type(moment) is tuple and not symbols:
        symbols = sp.symbols("xvar yvar zvar")

    dim = len(moment) if type(moment) is tuple else len(symbols)

    # not using sp.Dummy here - since it prohibits caching
    t = tuple([sp.Symbol("tmpvar_%d" % i, ) for i in range(dim)])
    symbols = symbols[:dim]
    generating_function = generating_function(func, symbols, t)

    if type(moment) is tuple:
        return multi_differentiation(generating_function, moment, t)
    else:
        assert symbols is not None, "When passing a polynomial as moment, also the moment symbols have to be passed"
        moment = sp.sympify(moment)

        result = 0
        for coefficient, exponents in polynomial_to_exponent_representation(moment, dim=dim):
            result += coefficient * multi_differentiation(generating_function, exponents, t)

        return result


def continuous_moment(func, moment, symbols=None):
    """Computes moment of given function.

    :param func: function to compute moments of
    :param moment: tuple or polynomial describing the moment
    :param symbols: if moment is given as polynomial, pass the moment symbols, i.e. the dof of the polynomial
    """
    return __continuous_moment_or_cumulant(func, moment, symbols, moment_generating_function)


def continuous_cumulant(func, moment, symbols=None):
    """Computes cumulant of continuous function.

    for parameter description see :func:`continuous_moment`
    """
    return __continuous_moment_or_cumulant(func, moment, symbols, cumulant_generating_function)