dist.py

import numpy as np
import scipy.stats as stats

__all__ = [
    'gaussian_1d_mixture',
    'gaussian_2d_mixture']


def gaussian_1d_mixture(arr, cov_factor=0.2):
    """A Gaussian kernel density estimation with a tuneable covariance factor.

    Parameters
    ----------
    arr : array-like
        Observations of variable X.
    cov_factor : float
        Covariance factor of the densities in the Gaussian mixture.

    Returns
    -------
    kernal : scipy.stats.gaussian_kde
        Representation of a kernel-density estimate using Gaussian kernels.

    """

    kernel = stats.gaussian_kde(arr)
    kernel.covariance_factor = lambda: cov_factor
    kernel._compute_covariance()
    return kernel


def gaussian_2d_mixture(x, y):
    """A Gaussian kernel density estimation for two-dimensional observations.
    Returns density values over two-dimensions.

    Parameters
    ----------
    x : array-like
        Observations of variable X.
    y : array-like
        Observations of variable Y.

    Returns
    -------
    xx : array-like
        Linearly spaced X coordinates.
    yy : array-like
        Linearly spaced Y coordinates.
    zz : array-like
        Density values in xx, yy space.
    kernel : scipy.stats.gaussian_kde
        Representation ofa kernel-density estimate using Gaussian kernels.

    """
    xx, yy = np.mgrid[min(x): max(x): 100j, min(y): max(y): 100j]
    positions = np.vstack([xx.ravel(), yy.ravel()])
    values = np.vstack([x, y])
    kernel = stats.gaussian_kde(values)

    zz = np.reshape(kernel(positions).T, xx.shape)

    return xx, yy, zz, kernel