"""
Gaussian Mixture Models.

This implementation corresponds to frequentist (non-Bayesian) formulation
of Gaussian Mixture Models.
"""

# Author: Ron Weiss <ronweiss@gmail.com>
#         Fabian Pedregosa <fabian.pedregosa@inria.fr>
#         Bertrand Thirion <bertrand.thirion@inria.fr>

import warnings
import numpy as np
from scipy import linalg
from time import time

from ..base import BaseEstimator
from ..utils import check_random_state, check_array
from ..utils.extmath import logsumexp
from ..utils.validation import check_is_fitted
from .. import cluster

from sklearn.externals.six.moves import zip

EPS = np.finfo(float).eps

def log_multivariate_normal_density(X, means, covars, covariance_type='diag'):
    """Compute the log probability under a multivariate Gaussian distribution.

    Parameters
    ----------
    X : array_like, shape (n_samples, n_features)
        List of n_features-dimensional data points.  Each row corresponds to a
        single data point.

    means : array_like, shape (n_components, n_features)
        List of n_features-dimensional mean vectors for n_components Gaussians.
        Each row corresponds to a single mean vector.

    covars : array_like
        List of n_components covariance parameters for each Gaussian. The shape
        depends on `covariance_type`:
            (n_components, n_features)      if 'spherical',
            (n_features, n_features)    if 'tied',
            (n_components, n_features)    if 'diag',
            (n_components, n_features, n_features) if 'full'

    covariance_type : string
        Type of the covariance parameters.  Must be one of
        'spherical', 'tied', 'diag', 'full'.  Defaults to 'diag'.

    Returns
    -------
    lpr : array_like, shape (n_samples, n_components)
        Array containing the log probabilities of each data point in
        X under each of the n_components multivariate Gaussian distributions.
    """
    log_multivariate_normal_density_dict = {
        'spherical': _log_multivariate_normal_density_spherical,
        'tied': _log_multivariate_normal_density_tied,
        'diag': _log_multivariate_normal_density_diag,
        'full': _log_multivariate_normal_density_full}
    return log_multivariate_normal_density_dict[covariance_type](
        X, means, covars)

def sample_gaussian(mean, covar, covariance_type='diag', n_samples=1,
                    random_state=None):
    """Generate random samples from a Gaussian distribution.

    Parameters
    ----------
    mean : array_like, shape (n_features,)
        Mean of the distribution.

    covar : array_like, optional
        Covariance of the distribution. The shape depends on `covariance_type`:
            scalar if 'spherical',
            (n_features) if 'diag',
            (n_features, n_features)  if 'tied', or 'full'

    covariance_type : string, optional
        Type of the covariance parameters.  Must be one of
        'spherical', 'tied', 'diag', 'full'.  Defaults to 'diag'.

    n_samples : int, optional
        Number of samples to generate. Defaults to 1.

    Returns
    -------
    X : array, shape (n_features, n_samples)
        Randomly generated sample
    """
    rng = check_random_state(random_state)
    n_dim = len(mean)
    rand = rng.randn(n_dim, n_samples)
    if n_samples == 1:
        rand.shape = (n_dim,)

    if covariance_type == 'spherical':
        rand *= np.sqrt(covar)
    elif covariance_type == 'diag':
        rand = np.dot(np.diag(np.sqrt(covar)), rand)
    else:
        s, U = linalg.eigh(covar)
        s.clip(0, out=s)        # get rid of tiny negatives
        np.sqrt(s, out=s)
        U *= s
        rand = np.dot(U, rand)

    return (rand.T + mean).T

class GMM(BaseEstimator):
    """Gaussian Mixture Model

    Representation of a Gaussian mixture model probability distribution.
    This class allows for easy evaluation of, sampling from, and
    maximum-likelihood estimation of the parameters of a GMM distribution.

    Initializes parameters such that every mixture component has zero
    mean and identity covariance.

    Read more in the :ref:`User Guide <gmm>`.

    Parameters
    ----------
    n_components : int, optional
        Number of mixture components. Defaults to 1.

    covariance_type : string, optional
        String describing the type of covariance parameters to
        use.  Must be one of 'spherical', 'tied', 'diag', 'full'.
        Defaults to 'diag'.

    random_state: RandomState or an int seed (None by default)
        A random number generator instance

    min_covar : float, optional
        Floor on the diagonal of the covariance matrix to prevent
        overfitting.  Defaults to 1e-3.

    tol : float, optional
        Convergence threshold. EM iterations will stop when average
        gain in log-likelihood is below this threshold.  Defaults to 1e-3.

    n_iter : int, optional
        Number of EM iterations to perform.

    n_init : int, optional
        Number of initializations to perform. the best results is kept

    params : string, optional
        Controls which parameters are updated in the training
        process.  Can contain any combination of 'w' for weights,
        'm' for means, and 'c' for covars.  Defaults to 'wmc'.

    init_params : string, optional
        Controls which parameters are updated in the initialization
        process.  Can contain any combination of 'w' for weights,
        'm' for means, and 'c' for covars.  Defaults to 'wmc'.

    verbose : int, default: 0
        Enable verbose output. If 1 then it always prints the current
        initialization and iteration step. If greater than 1 then
        it prints additionally the change and time needed for each step.

    Attributes
    ----------
    weights_ : array, shape (`n_components`,)
        This attribute stores the mixing weights for each mixture component.

    means_ : array, shape (`n_components`, `n_features`)
        Mean parameters for each mixture component.

    covars_ : array
        Covariance parameters for each mixture component.  The shape
        depends on `covariance_type`::

            (n_components, n_features)             if 'spherical',
            (n_features, n_features)               if 'tied',
            (n_components, n_features)             if 'diag',
            (n_components, n_features, n_features) if 'full'

    converged_ : bool
        True when convergence was reached in fit(), False otherwise.

    See Also
    --------

    DPGMM : Infinite gaussian mixture model, using the dirichlet
        process, fit with a variational algorithm

    VBGMM : Finite gaussian mixture model fit with a variational
        algorithm, better for situations where there might be too little
        data to get a good estimate of the covariance matrix.

    Examples
    --------

    >>> import numpy as np
    >>> from sklearn import mixture
    >>> np.random.seed(1)
    >>> g = mixture.GMM(n_components=2)
    >>> # Generate random observations with two modes centered on 0
    >>> # and 10 to use for training.
    >>> obs = np.concatenate((np.random.randn(100, 1),
    ...                       10 + np.random.randn(300, 1)))
    >>> g.fit(obs) # doctest: +NORMALIZE_WHITESPACE
    GMM(covariance_type='diag', init_params='wmc', min_covar=0.001,
            n_components=2, n_init=1, n_iter=100, params='wmc',
            random_state=None, thresh=None, tol=0.001, verbose=0)
    >>> np.round(g.weights_, 2)
    array([ 0.75,  0.25])
    >>> np.round(g.means_, 2)
    array([[ 10.05],
           [  0.06]])
    >>> np.round(g.covars_, 2) #doctest: +SKIP
    array([[[ 1.02]],
           [[ 0.96]]])
    >>> g.predict([[0], [2], [9], [10]]) #doctest: +ELLIPSIS
    array([1, 1, 0, 0]...)
    >>> np.round(g.score([[0], [2], [9], [10]]), 2)
    array([-2.19, -4.58, -1.75, -1.21])
    >>> # Refit the model on new data (initial parameters remain the
    >>> # same), this time with an even split between the two modes.
    >>> g.fit(20 * [[0]] +  20 * [[10]]) # doctest: +NORMALIZE_WHITESPACE
    GMM(covariance_type='diag', init_params='wmc', min_covar=0.001,
            n_components=2, n_init=1, n_iter=100, params='wmc',
            random_state=None, thresh=None, tol=0.001, verbose=0)
    >>> np.round(g.weights_, 2)
    array([ 0.5,  0.5])

    """

    def __init__(self, n_components=1, covariance_type='diag',
                 random_state=None, thresh=None, tol=1e-3, min_covar=1e-3,
                 n_iter=100, n_init=1, params='wmc', init_params='wmc',
                 verbose=0):
        if thresh is not None:
            warnings.warn("'thresh' has been replaced by 'tol' in 0.16 "
                          " and will be removed in 0.18.",
                          DeprecationWarning)
        self.n_components = n_components
        self.covariance_type = covariance_type
        self.thresh = thresh
        self.tol = tol
        self.min_covar = min_covar
        self.random_state = random_state
        self.n_iter = n_iter
        self.n_init = n_init
        self.params = params
        self.init_params = init_params
        self.verbose = verbose

        if covariance_type not in ['spherical', 'tied', 'diag', 'full']:
            raise ValueError('Invalid value for covariance_type: %s' %
                             covariance_type)

        if n_init < 1:
            raise ValueError('GMM estimation requires at least one run')

        self.weights_ = np.ones(self.n_components) / self.n_components

        # flag to indicate exit status of fit() method: converged (True) or
        # n_iter reached (False)
        self.converged_ = False

    def _get_covars(self):
        """Covariance parameters for each mixture component.

        The shape depends on ``cvtype``::

            (n_states, n_features)                if 'spherical',
            (n_features, n_features)              if 'tied',
            (n_states, n_features)                if 'diag',
            (n_states, n_features, n_features)    if 'full'

        """
        if self.covariance_type == 'full':
            return self.covars_
        elif self.covariance_type == 'diag':
            return [np.diag(cov) for cov in self.covars_]
        elif self.covariance_type == 'tied':
            return [self.covars_] * self.n_components
        elif self.covariance_type == 'spherical':
            return [np.diag(cov) for cov in self.covars_]

    def _set_covars(self, covars):
        """Provide values for covariance"""
        covars = np.asarray(covars)
        _validate_covars(covars, self.covariance_type, self.n_components)
        self.covars_ = covars

    def score_samples(self, X):
        """Return the per-sample likelihood of the data under the model.

        Compute the log probability of X under the model and
        return the posterior distribution (responsibilities) of each
        mixture component for each element of X.

        Parameters
        ----------
        X: array_like, shape (n_samples, n_features)
            List of n_features-dimensional data points. Each row
            corresponds to a single data point.

        Returns
        -------
        logprob : array_like, shape (n_samples,)
            Log probabilities of each data point in X.

        responsibilities : array_like, shape (n_samples, n_components)
            Posterior probabilities of each mixture component for each
            observation
        """
        check_is_fitted(self, 'means_')

        X = check_array(X)
        if X.ndim == 1:
            X = X[:, np.newaxis]
        if X.size == 0:
            return np.array([]), np.empty((0, self.n_components))
        if X.shape[1] != self.means_.shape[1]:
            raise ValueError('The shape of X  is not compatible with self')

        lpr = (log_multivariate_normal_density(X, self.means_, self.covars_,
                                               self.covariance_type) +
               np.log(self.weights_))
        logprob = logsumexp(lpr, axis=1)
        responsibilities = np.exp(lpr - logprob[:, np.newaxis])
        return logprob, responsibilities

    def score(self, X, y=None):
        """Compute the log probability under the model.

        Parameters
        ----------
        X : array_like, shape (n_samples, n_features)
            List of n_features-dimensional data points.  Each row
            corresponds to a single data point.

        Returns
        -------
        logprob : array_like, shape (n_samples,)
            Log probabilities of each data point in X
        """
        logprob, _ = self.score_samples(X)
        return logprob

    def predict(self, X):
        """Predict label for data.

        Parameters
        ----------
        X : array-like, shape = [n_samples, n_features]

        Returns
        -------
        C : array, shape = (n_samples,) component memberships
        """
        logprob, responsibilities = self.score_samples(X)
        return responsibilities.argmax(axis=1)

    def predict_proba(self, X):
        """Predict posterior probability of data under each Gaussian
        in the model.

        Parameters
        ----------
        X : array-like, shape = [n_samples, n_features]

        Returns
        -------
        responsibilities : array-like, shape = (n_samples, n_components)
            Returns the probability of the sample for each Gaussian
            (state) in the model.
        """
        logprob, responsibilities = self.score_samples(X)
        return responsibilities

    def sample(self, n_samples=1, random_state=None):
        """Generate random samples from the model.

        Parameters
        ----------
        n_samples : int, optional
            Number of samples to generate. Defaults to 1.

        Returns
        -------
        X : array_like, shape (n_samples, n_features)
            List of samples
        """
        check_is_fitted(self, 'means_')

        if random_state is None:
            random_state = self.random_state
        random_state = check_random_state(random_state)
        weight_cdf = np.cumsum(self.weights_)

        X = np.empty((n_samples, self.means_.shape[1]))
        rand = random_state.rand(n_samples)
        # decide which component to use for each sample
        comps = weight_cdf.searchsorted(rand)
        # for each component, generate all needed samples
        for comp in range(self.n_components):
            # occurrences of current component in X
            comp_in_X = (comp == comps)
            # number of those occurrences
            num_comp_in_X = comp_in_X.sum()
            if num_comp_in_X > 0:
                if self.covariance_type == 'tied':
                    cv = self.covars_
                elif self.covariance_type == 'spherical':
                    cv = self.covars_[comp][0]
                else:
                    cv = self.covars_[comp]
                X[comp_in_X] = sample_gaussian(
                    self.means_[comp], cv, self.covariance_type,
                    num_comp_in_X, random_state=random_state).T
        return X

    def fit_predict(self, X, y=None):
        """Fit and then predict labels for data.

        Warning: due to the final maximization step in the EM algorithm,
        with low iterations the prediction may not be 100% accurate

        Parameters
        ----------
        X : array-like, shape = [n_samples, n_features]

        Returns
        -------
        C : array, shape = (n_samples,) component memberships
        """
        return self._fit(X, y).argmax(axis=1)

    def _fit(self, X, y=None, do_prediction=False):
        """Estimate model parameters with the EM algorithm.

        A initialization step is performed before entering the
        expectation-maximization (EM) algorithm. If you want to avoid
        this step, set the keyword argument init_params to the empty
        string '' when creating the GMM object. Likewise, if you would
        like just to do an initialization, set n_iter=0.

        Parameters
        ----------
        X : array_like, shape (n, n_features)
            List of n_features-dimensional data points.  Each row
            corresponds to a single data point.

        Returns
        -------
        responsibilities : array, shape (n_samples, n_components)
            Posterior probabilities of each mixture component for each
            observation.
        """

        # initialization step
        X = check_array(X, dtype=np.float64, ensure_min_samples=2)
        if X.shape[0] < self.n_components:
            raise ValueError(
                'GMM estimation with %s components, but got only %s samples' %
                (self.n_components, X.shape[0]))

        max_log_prob = -np.infty

        if self.verbose > 0:
            print('Expectation-maximization algorithm started.')

        for init in range(self.n_init):
            if self.verbose > 0:
                print('Initialization ' + str(init + 1))
                start_init_time = time()

            if 'm' in self.init_params or not hasattr(self, 'means_'):
                self.means_ = cluster.KMeans(
                    n_clusters=self.n_components,
                    random_state=self.random_state).fit(X).cluster_centers_
                if self.verbose > 1:
                    print('\tMeans have been initialized.')

            if 'w' in self.init_params or not hasattr(self, 'weights_'):
                self.weights_ = np.tile(1.0 / self.n_components,
                                        self.n_components)
                if self.verbose > 1:
                    print('\tWeights have been initialized.')

            if 'c' in self.init_params or not hasattr(self, 'covars_'):
                cv = np.cov(X.T) + self.min_covar * np.eye(X.shape[1])
                if not cv.shape:
                    cv.shape = (1, 1)
                self.covars_ = \
                    distribute_covar_matrix_to_match_covariance_type(
                        cv, self.covariance_type, self.n_components)
                if self.verbose > 1:
                    print('\tCovariance matrices have been initialized.')

            # EM algorithms
            current_log_likelihood = None
            # reset self.converged_ to False
            self.converged_ = False

            # this line should be removed when 'thresh' is removed in v0.18
            tol = (self.tol if self.thresh is None
                   else self.thresh / float(X.shape[0]))

            for i in range(self.n_iter):
                if self.verbose > 0:
                    print('\tEM iteration ' + str(i + 1))
                    start_iter_time = time()
                prev_log_likelihood = current_log_likelihood
                # Expectation step
                log_likelihoods, responsibilities = self.score_samples(X)
                current_log_likelihood = log_likelihoods.mean()

                # Check for convergence.
                # (should compare to self.tol when deprecated 'thresh' is
                # removed in v0.18)
                if prev_log_likelihood is not None:
                    change = abs(current_log_likelihood - prev_log_likelihood)
                    if self.verbose > 1:
                        print('\t\tChange: ' + str(change))
                    if change < tol:
                        self.converged_ = True
                        if self.verbose > 0:
                            print('\t\tEM algorithm converged.')
                        break

                # Maximization step
                self._do_mstep(X, responsibilities, self.params,
                               self.min_covar)
                if self.verbose > 1:
                    print('\t\tEM iteration ' + str(i + 1) + ' took {0:.5f}s'.format(
                        time() - start_iter_time))

            # if the results are better, keep it
            if self.n_iter:
                if current_log_likelihood > max_log_prob:
                    max_log_prob = current_log_likelihood
                    best_params = {'weights': self.weights_,
                                   'means': self.means_,
                                   'covars': self.covars_}
                    if self.verbose > 1:
                        print('\tBetter parameters were found.')

            if self.verbose > 1:
                print('\tInitialization ' + str(init + 1) + ' took {0:.5f}s'.format(
                    time() - start_init_time))

        # check the existence of an init param that was not subject to
        # likelihood computation issue.
        if np.isneginf(max_log_prob) and self.n_iter:
            raise RuntimeError(
                "EM algorithm was never able to compute a valid likelihood " +
                "given initial parameters. Try different init parameters " +
                "(or increasing n_init) or check for degenerate data.")

        if self.n_iter:
            self.covars_ = best_params['covars']
            self.means_ = best_params['means']
            self.weights_ = best_params['weights']
        else:  # self.n_iter == 0 occurs when using GMM within HMM
            # Need to make sure that there are responsibilities to output
            # Output zeros because it was just a quick initialization
            responsibilities = np.zeros((X.shape[0], self.n_components))

        return responsibilities

    def fit(self, X, y=None):
        """Estimate model parameters with the EM algorithm.

        A initialization step is performed before entering the
        expectation-maximization (EM) algorithm. If you want to avoid
        this step, set the keyword argument init_params to the empty
        string '' when creating the GMM object. Likewise, if you would
        like just to do an initialization, set n_iter=0.

        Parameters
        ----------
        X : array_like, shape (n, n_features)
            List of n_features-dimensional data points.  Each row
            corresponds to a single data point.

        Returns
        -------
        self
        """
        self._fit(X, y)
        return self

    def _do_mstep(self, X, responsibilities, params, min_covar=0):
        """ Perform the Mstep of the EM algorithm and return the class weights
        """
        weights = responsibilities.sum(axis=0)
        weighted_X_sum = np.dot(responsibilities.T, X)
        inverse_weights = 1.0 / (weights[:, np.newaxis] + 10 * EPS)

        if 'w' in params:
            self.weights_ = (weights / (weights.sum() + 10 * EPS) + EPS)
        if 'm' in params:
            self.means_ = weighted_X_sum * inverse_weights
        if 'c' in params:
            covar_mstep_func = _covar_mstep_funcs[self.covariance_type]
            self.covars_ = covar_mstep_func(
                self, X, responsibilities, weighted_X_sum, inverse_weights,
                min_covar)
        return weights

    def _n_parameters(self):
        """Return the number of free parameters in the model."""
        ndim = self.means_.shape[1]
        if self.covariance_type == 'full':
            cov_params = self.n_components * ndim * (ndim + 1) / 2.
        elif self.covariance_type == 'diag':
            cov_params = self.n_components * ndim
        elif self.covariance_type == 'tied':
            cov_params = ndim * (ndim + 1) / 2.
        elif self.covariance_type == 'spherical':
            cov_params = self.n_components
        mean_params = ndim * self.n_components
        return int(cov_params + mean_params + self.n_components - 1)

    def bic(self, X):
        """Bayesian information criterion for the current model fit
        and the proposed data

        Parameters
        ----------
        X : array of shape(n_samples, n_dimensions)

        Returns
        -------
        bic: float (the lower the better)
        """
        return (-2 * self.score(X).sum() +
                self._n_parameters() * np.log(X.shape[0]))

    def aic(self, X):
        """Akaike information criterion for the current model fit
        and the proposed data

        Parameters
        ----------
        X : array of shape(n_samples, n_dimensions)

        Returns
        -------
        aic: float (the lower the better)
        """
        return - 2 * self.score(X).sum() + 2 * self._n_parameters()

#########################################################################
# some helper routines
#########################################################################

def _log_multivariate_normal_density_diag(X, means, covars):
    """Compute Gaussian log-density at X for a diagonal model"""
    n_samples, n_dim = X.shape
    lpr = -0.5 * (n_dim * np.log(2 * np.pi) + np.sum(np.log(covars), 1)
                  + np.sum((means ** 2) / covars, 1)
                  - 2 * np.dot(X, (means / covars).T)
                  + np.dot(X ** 2, (1.0 / covars).T))
    return lpr

def _log_multivariate_normal_density_spherical(X, means, covars):
    """Compute Gaussian log-density at X for a spherical model"""
    cv = covars.copy()
    if covars.ndim == 1:
        cv = cv[:, np.newaxis]
    if covars.shape[1] == 1:
        cv = np.tile(cv, (1, X.shape[-1]))
    return _log_multivariate_normal_density_diag(X, means, cv)

def _log_multivariate_normal_density_tied(X, means, covars):
    """Compute Gaussian log-density at X for a tied model"""
    cv = np.tile(covars, (means.shape[0], 1, 1))
    return _log_multivariate_normal_density_full(X, means, cv)

def _log_multivariate_normal_density_full(X, means, covars, min_covar=1.e-7):
    """Log probability for full covariance matrices."""
    n_samples, n_dim = X.shape
    nmix = len(means)
    log_prob = np.empty((n_samples, nmix))
    for c, (mu, cv) in enumerate(zip(means, covars)):
        try:
            cv_chol = linalg.cholesky(cv, lower=True)
        except linalg.LinAlgError:
            # The model is most probably stuck in a component with too
            # few observations, we need to reinitialize this components
            try:
                cv_chol = linalg.cholesky(cv + min_covar * np.eye(n_dim),
                                          lower=True)
            except linalg.LinAlgError:
                raise ValueError("'covars' must be symmetric, "
                                 "positive-definite")

        cv_log_det = 2 * np.sum(np.log(np.diagonal(cv_chol)))
        cv_sol = linalg.solve_triangular(cv_chol, (X - mu).T, lower=True).T
        log_prob[:, c] = - .5 * (np.sum(cv_sol ** 2, axis=1) +
                                 n_dim * np.log(2 * np.pi) + cv_log_det)

    return log_prob

def _validate_covars(covars, covariance_type, n_components):
    """Do basic checks on matrix covariance sizes and values
    """
    from scipy import linalg
    if covariance_type == 'spherical':
        if len(covars) != n_components:
            raise ValueError("'spherical' covars have length n_components")
        elif np.any(covars <= 0):
            raise ValueError("'spherical' covars must be non-negative")
    elif covariance_type == 'tied':
        if covars.shape[0] != covars.shape[1]:
            raise ValueError("'tied' covars must have shape (n_dim, n_dim)")
        elif (not np.allclose(covars, covars.T)
              or np.any(linalg.eigvalsh(covars) <= 0)):
            raise ValueError("'tied' covars must be symmetric, "
                             "positive-definite")
    elif covariance_type == 'diag':
        if len(covars.shape) != 2:
            raise ValueError("'diag' covars must have shape "
                             "(n_components, n_dim)")
        elif np.any(covars <= 0):
            raise ValueError("'diag' covars must be non-negative")
    elif covariance_type == 'full':
        if len(covars.shape) != 3:
            raise ValueError("'full' covars must have shape "
                             "(n_components, n_dim, n_dim)")
        elif covars.shape[1] != covars.shape[2]:
            raise ValueError("'full' covars must have shape "
                             "(n_components, n_dim, n_dim)")
        for n, cv in enumerate(covars):
            if (not np.allclose(cv, cv.T)
                    or np.any(linalg.eigvalsh(cv) <= 0)):
                raise ValueError("component %d of 'full' covars must be "
                                 "symmetric, positive-definite" % n)
    else:
        raise ValueError("covariance_type must be one of " +
                         "'spherical', 'tied', 'diag', 'full'")

def distribute_covar_matrix_to_match_covariance_type(
        tied_cv, covariance_type, n_components):
    """Create all the covariance matrices from a given template"""
    if covariance_type == 'spherical':
        cv = np.tile(tied_cv.mean() * np.ones(tied_cv.shape[1]),
                     (n_components, 1))
    elif covariance_type == 'tied':
        cv = tied_cv
    elif covariance_type == 'diag':
        cv = np.tile(np.diag(tied_cv), (n_components, 1))
    elif covariance_type == 'full':
        cv = np.tile(tied_cv, (n_components, 1, 1))
    else:
        raise ValueError("covariance_type must be one of " +
                         "'spherical', 'tied', 'diag', 'full'")
    return cv

def _covar_mstep_diag(gmm, X, responsibilities, weighted_X_sum, norm,
                      min_covar):
    """Performing the covariance M step for diagonal cases"""
    avg_X2 = np.dot(responsibilities.T, X * X) * norm
    avg_means2 = gmm.means_ ** 2
    avg_X_means = gmm.means_ * weighted_X_sum * norm
    return avg_X2 - 2 * avg_X_means + avg_means2 + min_covar

def _covar_mstep_spherical(*args):
    """Performing the covariance M step for spherical cases"""
    cv = _covar_mstep_diag(*args)
    return np.tile(cv.mean(axis=1)[:, np.newaxis], (1, cv.shape[1]))

def _covar_mstep_full(gmm, X, responsibilities, weighted_X_sum, norm,
                      min_covar):
    """Performing the covariance M step for full cases"""
    # Eq. 12 from K. Murphy, "Fitting a Conditional Linear Gaussian
    # Distribution"
    n_features = X.shape[1]
    cv = np.empty((gmm.n_components, n_features, n_features))
    for c in range(gmm.n_components):
        post = responsibilities[:, c]
        mu = gmm.means_[c]
        diff = X - mu
        with np.errstate(under='ignore'):
            # Underflow Errors in doing post * X.T are  not important
            avg_cv = np.dot(post * diff.T, diff) / (post.sum() + 10 * EPS)
        cv[c] = avg_cv + min_covar * np.eye(n_features)
    return cv

def _covar_mstep_tied(gmm, X, responsibilities, weighted_X_sum, norm,
                      min_covar):
    # Eq. 15 from K. Murphy, "Fitting a Conditional Linear Gaussian
    # Distribution"
    avg_X2 = np.dot(X.T, X)
    avg_means2 = np.dot(gmm.means_.T, weighted_X_sum)
    out = avg_X2 - avg_means2
    out *= 1. / X.shape[0]
    out.flat[::len(out) + 1] += min_covar
    return out

_covar_mstep_funcs = {'spherical': _covar_mstep_spherical,
                      'diag': _covar_mstep_diag,
                      'tied': _covar_mstep_tied,
                      'full': _covar_mstep_full,
                      }


"""
Soft Voting/Majority Rule classifier.

This module contains a Soft Voting/Majority Rule classifier for
classification estimators.

"""

# Authors: Sebastian Raschka <se.raschka@gmail.com>,
#          Gilles Louppe <g.louppe@gmail.com>
#
# Licence: BSD 3 clause

import numpy as np

from ..base import BaseEstimator
from ..base import ClassifierMixin
from ..base import TransformerMixin
from ..base import clone
from ..preprocessing import LabelEncoder
from ..externals import six

class VotingClassifier(BaseEstimator, ClassifierMixin, TransformerMixin):
    """Soft Voting/Majority Rule classifier for unfitted estimators.

    Read more in the :ref:`User Guide <voting_classifier>`.

    Parameters
    ----------
    estimators : list of (string, estimator) tuples
        Invoking the `fit` method on the `VotingClassifier` will fit clones
        of those original estimators that will be stored in the class attribute
        `self.estimators_`.

    voting : str, {'hard', 'soft'} (default='hard')
        If 'hard', uses predicted class labels for majority rule voting.
        Else if 'soft', predicts the class label based on the argmax of
        the sums of the predicted probalities, which is recommended for
        an ensemble of well-calibrated classifiers.

    weights : array-like, shape = [n_classifiers], optional (default=`None`)
        Sequence of weights (`float` or `int`) to weight the occurances of
        predicted class labels (`hard` voting) or class probabilities
        before averaging (`soft` voting). Uses uniform weights if `None`.

    Attributes
    ----------
    classes_ : array-like, shape = [n_predictions]

    Examples
    --------
    >>> import numpy as np
    >>> from sklearn.linear_model import LogisticRegression
    >>> from sklearn.naive_bayes import GaussianNB
    >>> from sklearn.ensemble import RandomForestClassifier
    >>> clf1 = LogisticRegression(random_state=1)
    >>> clf2 = RandomForestClassifier(random_state=1)
    >>> clf3 = GaussianNB()
    >>> X = np.array([[-1, -1], [-2, -1], [-3, -2], [1, 1], [2, 1], [3, 2]])
    >>> y = np.array([1, 1, 1, 2, 2, 2])
    >>> eclf1 = VotingClassifier(estimators=[
    ...         ('lr', clf1), ('rf', clf2), ('gnb', clf3)], voting='hard')
    >>> eclf1 = eclf1.fit(X, y)
    >>> print(eclf1.predict(X))
    [1 1 1 2 2 2]
    >>> eclf2 = VotingClassifier(estimators=[
    ...         ('lr', clf1), ('rf', clf2), ('gnb', clf3)],
    ...         voting='soft')
    >>> eclf2 = eclf2.fit(X, y)
    >>> print(eclf2.predict(X))
    [1 1 1 2 2 2]
    >>> eclf3 = VotingClassifier(estimators=[
    ...        ('lr', clf1), ('rf', clf2), ('gnb', clf3)],
    ...        voting='soft', weights=[2,1,1])
    >>> eclf3 = eclf3.fit(X, y)
    >>> print(eclf3.predict(X))
    [1 1 1 2 2 2]
    >>>
    """

    def __init__(self, estimators, voting='hard', weights=None):

        self.estimators = estimators
        self.named_estimators = dict(estimators)
        self.voting = voting
        self.weights = weights

    def fit(self, X, y):
        """ Fit the estimators.

        Parameters
        ----------
        X : {array-like, sparse matrix}, shape = [n_samples, n_features]
            Training vectors, where n_samples is the number of samples and
            n_features is the number of features.

        y : array-like, shape = [n_samples]
            Target values.

        Returns
        -------
        self : object
        """
        if isinstance(y, np.ndarray) and len(y.shape) > 1 and y.shape[1] > 1:
            raise NotImplementedError('Multilabel and multi-output'
                                      ' classification is not supported.')

        if self.voting not in ('soft', 'hard'):
            raise ValueError("Voting must be 'soft' or 'hard'; got (voting=%r)"
                             % self.voting)

        if self.weights and len(self.weights) != len(self.estimators):
            raise ValueError('Number of classifiers and weights must be equal'
                             '; got %d weights, %d estimators'
                             % (len(self.weights), len(self.estimators)))

        self.le_ = LabelEncoder()
        self.le_.fit(y)
        self.classes_ = self.le_.classes_
        self.estimators_ = []

        for name, clf in self.estimators:
            fitted_clf = clone(clf).fit(X, self.le_.transform(y))
            self.estimators_.append(fitted_clf)

        return self

    def predict(self, X):
        """ Predict class labels for X.

        Parameters
        ----------
        X : {array-like, sparse matrix}, shape = [n_samples, n_features]
            Training vectors, where n_samples is the number of samples and
            n_features is the number of features.

        Returns
        ----------
        maj : array-like, shape = [n_samples]
            Predicted class labels.
        """
        if self.voting == 'soft':
            maj = np.argmax(self.predict_proba(X), axis=1)

        else:  # 'hard' voting
            predictions = self._predict(X)
            maj = np.apply_along_axis(lambda x:
                                      np.argmax(np.bincount(x,
                                                weights=self.weights)),
                                      axis=1,
                                      arr=predictions)

        maj = self.le_.inverse_transform(maj)

        return maj

    def _collect_probas(self, X):
        """Collect results from clf.predict calls. """
        return np.asarray([clf.predict_proba(X) for clf in self.estimators_])

    def _predict_proba(self, X):
        """Predict class probabilities for X in 'soft' voting """
        avg = np.average(self._collect_probas(X), axis=0, weights=self.weights)
        return avg

    @property
    def predict_proba(self):
        """Compute probabilities of possible outcomes for samples in X.

        Parameters
        ----------
        X : {array-like, sparse matrix}, shape = [n_samples, n_features]
            Training vectors, where n_samples is the number of samples and
            n_features is the number of features.

        Returns
        ----------
        avg : array-like, shape = [n_samples, n_classes]
            Weighted average probability for each class per sample.
        """
        if self.voting == 'hard':
            raise AttributeError("predict_proba is not available when"
                                 " voting=%r" % self.voting)
        return self._predict_proba

    def transform(self, X):
        """Return class labels or probabilities for X for each estimator.

        Parameters
        ----------
        X : {array-like, sparse matrix}, shape = [n_samples, n_features]
            Training vectors, where n_samples is the number of samples and
            n_features is the number of features.

        Returns
        -------
        If `voting='soft'`:
          array-like = [n_classifiers, n_samples, n_classes]
            Class probabilties calculated by each classifier.
        If `voting='hard'`:
          array-like = [n_classifiers, n_samples]
            Class labels predicted by each classifier.
        """
        if self.voting == 'soft':
            return self._collect_probas(X)
        else:
            return self._predict(X)

    def get_params(self, deep=True):
        """Return estimator parameter names for GridSearch support"""
        if not deep:
            return super(VotingClassifier, self).get_params(deep=False)
        else:
            out = super(VotingClassifier, self).get_params(deep=False)
            out.update(self.named_estimators.copy())
            for name, step in six.iteritems(self.named_estimators):
                for key, value in six.iteritems(step.get_params(deep=True)):
                    out['%s__%s' % (name, key)] = value
            return out

    def _predict(self, X):
        """Collect results from clf.predict calls. """
        return np.asarray([clf.predict(X) for clf in self.estimators_]).T


#!/usr/bin/env python3

"""Generate a binary file from a sample image of the MNIST dataset.

Pixel of the sample are stored as float32, images have size 28x28.
"""

import os
import argparse

import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf

SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))

def main(args):
    _, (mnist_test, _) = tf.keras.datasets.mnist.load_data()
    data = mnist_test[args.index]

    output_path = os.path.join(SCRIPT_DIR, args.output)
    np.ndarray.tofile(data.astype('float32'), output_path)

    if args.no_plot is False:
        plt.gray()
        plt.imshow(data.reshape(28, 28))
        plt.show()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-i", "--index", type=int, default=0,
                        help="Image index in MNIST test dataset")
    parser.add_argument("-o", "--output", type=str, default='digit',
                        help="Output filename")
    parser.add_argument("--no-plot", default=False, action='store_true',
                        help="Disable image display in matplotlib")
    main(parser.parse_args())


#!/usr/bin/env python
#
# Copyright 2007,2008,2011 Free Software Foundation, Inc.
#
# This file is part of GNU Radio
#
# GNU Radio 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, or (at your option)
# any later version.
#
# GNU Radio 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 GNU Radio; see the file COPYING.  If not, write to
# the Free Software Foundation, Inc., 51 Franklin Street,
# Boston, MA 02110-1301, USA.
#

try:
    import scipy
    from scipy import fftpack
except ImportError:
    print "Please install SciPy to run this script (http://www.scipy.org/)"
    raise SystemExit, 1

try:
    from pylab import *
except ImportError:
    print "Please install Matplotlib to run this script (http://matplotlib.sourceforge.net/)"
    raise SystemExit, 1

from optparse import OptionParser

class plot_fft_base:
    def __init__(self, datatype, filename, options):
        self.hfile = open(filename, "r")
        self.block_length = options.block
        self.start = options.start
        self.sample_rate = options.sample_rate

        self.datatype = getattr(scipy, datatype)
        self.sizeof_data = self.datatype().nbytes    # number of bytes per sample in file

        self.axis_font_size = 16
        self.label_font_size = 18
        self.title_font_size = 20
        self.text_size = 22

        # Setup PLOT
        self.fig = figure(1, figsize=(16, 12), facecolor='w')
        rcParams['xtick.labelsize'] = self.axis_font_size
        rcParams['ytick.labelsize'] = self.axis_font_size

        self.text_file     = figtext(0.10, 0.94, ("File: %s" % filename), weight="heavy", size=self.text_size)
        self.text_file_pos = figtext(0.10, 0.88, "File Position: ", weight="heavy", size=self.text_size)
        self.text_block    = figtext(0.35, 0.88, ("Block Size: %d" % self.block_length),
                                     weight="heavy", size=self.text_size)
        self.text_sr       = figtext(0.60, 0.88, ("Sample Rate: %.2f" % self.sample_rate),
                                     weight="heavy", size=self.text_size)
        self.make_plots()

        self.button_left_axes = self.fig.add_axes([0.45, 0.01, 0.05, 0.05], frameon=True)
        self.button_left = Button(self.button_left_axes, "<")
        self.button_left_callback = self.button_left.on_clicked(self.button_left_click)

        self.button_right_axes = self.fig.add_axes([0.50, 0.01, 0.05, 0.05], frameon=True)
        self.button_right = Button(self.button_right_axes, ">")
        self.button_right_callback = self.button_right.on_clicked(self.button_right_click)

        self.xlim = self.sp_iq.get_xlim()

        self.manager = get_current_fig_manager()
        connect('draw_event', self.zoom)
        connect('key_press_event', self.click)
        show()

    def get_data(self):
        self.position = self.hfile.tell()/self.sizeof_data
        self.text_file_pos.set_text("File Position: %d" % (self.position))
        try:
            self.iq = scipy.fromfile(self.hfile, dtype=self.datatype, count=self.block_length)
        except MemoryError:
            print "End of File"
        else:
            self.iq_fft = self.dofft(self.iq)

            tstep = 1.0 / self.sample_rate
            #self.time = scipy.array([tstep*(self.position + i) for i in xrange(len(self.iq))])
            self.time = scipy.array([tstep*(i) for i in xrange(len(self.iq))])

            self.freq = self.calc_freq(self.time, self.sample_rate)

    def dofft(self, iq):
        N = len(iq)
        iq_fft = scipy.fftpack.fftshift(scipy.fft(iq))       # fft and shift axis
        iq_fft = 20*scipy.log10(abs((iq_fft+1e-15)/N)) # convert to decibels, adjust power
        # adding 1e-15 (-300 dB) to protect against value errors if an item in iq_fft is 0
        return iq_fft

    def calc_freq(self, time, sample_rate):
        N = len(time)
        Fs = 1.0 / (time.max() - time.min())
        Fn = 0.5 * sample_rate
        freq = scipy.array([-Fn + i*Fs for i in xrange(N)])
        return freq

    def make_plots(self):
        # if specified on the command-line, set file pointer
        self.hfile.seek(self.sizeof_data*self.start, 1)

        # Subplot for real and imaginary parts of signal
        self.sp_iq = self.fig.add_subplot(2,2,1, position=[0.075, 0.2, 0.4, 0.6])
        self.sp_iq.set_title(("I&Q"), fontsize=self.title_font_size, fontweight="bold")
        self.sp_iq.set_xlabel("Time (s)", fontsize=self.label_font_size, fontweight="bold")
        self.sp_iq.set_ylabel("Amplitude (V)", fontsize=self.label_font_size, fontweight="bold")

        # Subplot for FFT plot
        self.sp_fft = self.fig.add_subplot(2,2,2, position=[0.575, 0.2, 0.4, 0.6])
        self.sp_fft.set_title(("FFT"), fontsize=self.title_font_size, fontweight="bold")
        self.sp_fft.set_xlabel("Frequency (Hz)", fontsize=self.label_font_size, fontweight="bold")
        self.sp_fft.set_ylabel("Power Spectrum (dBm)", fontsize=self.label_font_size, fontweight="bold")

        self.get_data()

        self.plot_iq  = self.sp_iq.plot([], 'bo-') # make plot for reals
        self.plot_iq += self.sp_iq.plot([], 'ro-') # make plot for imags
        self.draw_time()                           # draw the plot

        self.plot_fft = self.sp_fft.plot([], 'bo-')  # make plot for FFT
        self.draw_fft()                              # draw the plot

        draw()

    def draw_time(self):
        reals = self.iq.real
        imags = self.iq.imag
        self.plot_iq[0].set_data([self.time, reals])
        self.plot_iq[1].set_data([self.time, imags])
        self.sp_iq.set_xlim(self.time.min(), self.time.max())
        self.sp_iq.set_ylim([1.5*min([reals.min(), imags.min()]),
                             1.5*max([reals.max(), imags.max()])])

    def draw_fft(self):
        self.plot_fft[0].set_data([self.freq, self.iq_fft])
        self.sp_fft.set_xlim(self.freq.min(), self.freq.max())
        self.sp_fft.set_ylim([self.iq_fft.min()-10, self.iq_fft.max()+10])

    def update_plots(self):
        self.draw_time()
        self.draw_fft()

        self.xlim = self.sp_iq.get_xlim()
        draw()

    def zoom(self, event):
        newxlim = scipy.array(self.sp_iq.get_xlim())
        curxlim = scipy.array(self.xlim)
        if(newxlim[0] != curxlim[0] or newxlim[1] != curxlim[1]):
            self.xlim = newxlim
            #xmin = max(0, int(ceil(self.sample_rate*(self.xlim[0] - self.position))))
            #xmax = min(int(ceil(self.sample_rate*(self.xlim[1] - self.position))), len(self.iq))
            xmin = max(0, int(ceil(self.sample_rate*(self.xlim[0]))))
            xmax = min(int(ceil(self.sample_rate*(self.xlim[1]))), len(self.iq))

            iq = self.iq[xmin : xmax]
            time = self.time[xmin : xmax]

            iq_fft = self.dofft(iq)
            freq = self.calc_freq(time, self.sample_rate)

            self.plot_fft[0].set_data(freq, iq_fft)
            self.sp_fft.axis([freq.min(), freq.max(),
                              iq_fft.min()-10, iq_fft.max()+10])

            draw()

    def click(self, event):
        forward_valid_keys = [" ", "down", "right"]
        backward_valid_keys = ["up", "left"]

        if(find(event.key, forward_valid_keys)):
            self.step_forward()

        elif(find(event.key, backward_valid_keys)):
            self.step_backward()

    def button_left_click(self, event):
        self.step_backward()

    def button_right_click(self, event):
        self.step_forward()

    def step_forward(self):
        self.get_data()
        self.update_plots()

    def step_backward(self):
        # Step back in file position
        if(self.hfile.tell() >= 2*self.sizeof_data*self.block_length ):
            self.hfile.seek(-2*self.sizeof_data*self.block_length, 1)
        else:
            self.hfile.seek(-self.hfile.tell(),1)
        self.get_data()
        self.update_plots()

    @staticmethod
    def setup_options():
        usage="%prog: [options] input_filename"
        description = "Takes a GNU Radio complex binary file and displays the I&Q data versus time as well as the frequency domain (FFT) plot. The y-axis values are plotted assuming volts as the amplitude of the I&Q streams and converted into dBm in the frequency domain (the 1/N power adjustment out of the FFT is performed internally). The script plots a certain block of data at a time, specified on the command line as -B or --block. This value defaults to 1000. The start position in the file can be set by specifying -s or --start and defaults to 0 (the start of the file). By default, the system assumes a sample rate of 1, so in time, each sample is plotted versus the sample number. To set a true time and frequency axis, set the sample rate (-R or --sample-rate) to the sample rate used when capturing the samples."

        parser = OptionParser(conflict_handler="resolve", usage=usage, description=description)
        parser.add_option("-d", "--data-type", type="string", default="complex64",
                          help="Specify the data type (complex64, float32, (u)int32, (u)int16, (u)int8) [default=%default]")
        parser.add_option("-B", "--block", type="int", default=1000,
                          help="Specify the block size [default=%default]")
        parser.add_option("-s", "--start", type="int", default=0,
                          help="Specify where to start in the file [default=%default]")
        parser.add_option("-R", "--sample-rate", type="float", default=1.0,
                          help="Set the sampler rate of the data [default=%default]")
        return parser

def find(item_in, list_search):
    try:
	return list_search.index(item_in) != None
    except ValueError:
	return False

def main():
    parser = plot_fft_base.setup_options()
    (options, args) = parser.parse_args ()
    if len(args) != 1:
        parser.print_help()
        raise SystemExit, 1
    filename = args[0]

    dc = plot_fft_base(options.data_type, filename, options)

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        pass


from astropy.io import fits
import aplpy
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import astropy.units as u
import astropy.constants as c
import warnings
import numpy as np
from astropy.visualization import hist
from config import plottingDictionary
"""
Make histogram plots of NH3-derived properties for DR1 regions
"""
def mask_hist(par_data,epar_data,epar_lim,par_max,par_min=0):
    mask1 = np.isfinite(epar_data)
    mask2 = epar_data < epar_lim
    mask3 = epar_data > 0
    mask4 = par_data < par_max
    mask5 = par_data > par_min
    return par_data * mask1 * mask2 * mask3 * mask4 * mask5

region_list = ['B18','NGC1333','L1688','OrionA']
par_list = ['Vlsr','Sigma','Tkin','Tex','N_NH3']
epar_ext = [10,9,6,7,8]
epar_limits = [0.05,0.1,1,2,0.5]
extension = 'DR1_rebase3'
label_list=['$v_{LSR}$ (km s$^{-1}$)','$\sigma_v$ (km s$^{-1}$)','$T_K$ (K)','$T_{ex}$ (K)','log N(para-NH$_3$) (cm$^{-2}$)']
label_short = ['$v_{LSR}$','$\sigma_v$','$T_K$','$T_{ex}$','log N(para-NH$_3$)']
plot_colours = ['black','blue','green','orange']
#plot_colours = ['black','darkblue','blue','cornflowerblue']
#plot_colours = ['#a6cee3', '#fdbf6f', '#33a02c', '#fb9a99']
hist_minx_list = [2,0,5,2.7,13]
hist_maxx_list = [13,1.5,37,12,15.7]
hist_maxy_list = [2.1,8,0.65,1.6,2]
ytick_int_maj = [0.4,2,0.2,0.4,0.5]
ytick_int_min = [0.1,0.5,0.025,0.1,0.25]
dataDir = ''

hist_kwds1 = dict(histtype='stepfilled',alpha=0.2,normed=True)

# Separate regions in plots to better show distributions
for par_i in range(len(par_list)):
    fig,axes = plt.subplots(len(region_list),1,figsize=(4,5))
    par = par_list[par_i]
    label = label_list[par_i]
    ylabel = label_short[par_i]
    for i, ax in enumerate(fig.axes):
        region_i = i
        region = region_list[region_i]
        plot_param=plottingDictionary[region]
        par_file = dataDir + '{0}/parameterMaps/{0}_{1}_{2}_flag.fits'.format(region,par,extension)
        epar_file = dataDir + '{0}/{0}_parameter_maps_{1}_trim.fits'.format(region,extension)
        epar_hdu = fits.open(epar_file)
        epar_data = epar_hdu[0].data[epar_ext[par_i],:,:]
        epar_hdu.close()
        par_hdu = fits.open(par_file)
        par_data = par_hdu[0].data
        par_hdu.close()
        pmin_list = plot_param['pmin_list']
        pmax_list = plot_param['pmax_list']
        pmin = np.max([pmin_list[par_i],np.nanmin(par_data)])
        pmax = np.min([pmax_list[par_i],np.nanmax(par_data)])
        par_masked = mask_hist(par_data,epar_data,epar_limits[par_i],pmax,par_min=pmin)
        par_masked = par_masked[np.isfinite(par_masked)]
        if par == 'Vlsr':
            bin_width = 0.3
            nbins = np.int((np.max(par_masked) - np.min(par_masked !=0))/bin_width)
            hist(par_masked[par_masked !=0],bins=nbins,ax=ax,histtype='stepfilled',alpha=0.3,
                 color=plot_colours[region_i],label=region,normed=True)
        else:
            hist(par_masked[par_masked !=0],bins='knuth',ax=ax,histtype='stepfilled',alpha=0.3,
                 color=plot_colours[region_i],label=region,normed=True)
        if (i+1) != len(region_list):
            ax.set_xticklabels([])
        ax.set_xlim(hist_minx_list[par_i],hist_maxx_list[par_i])
        #ax.set_ylim(0,hist_maxy_list[par_i])
        if par == 'Tkin':
            if region == 'OrionA' or region == 'L1688' or region == 'NGC1333':
                ax.set_ylim(0,0.25)
        ax.yaxis.set_major_locator(ticker.MultipleLocator(ytick_int_maj[par_i]))
        ax.yaxis.set_minor_locator(ticker.MultipleLocator(ytick_int_min[par_i]))
        ax.annotate('{0}'.format(region),xy=(0.97,0.7),xycoords='axes fraction',horizontalalignment='right')
    #ax.legend(frameon=False)
    ax.set_xlabel(label)
    fig.text(0.001,0.5,'P ({0})'.format(ylabel),va='center',rotation='vertical')
    #fig.tight_layout()
    fig.savefig('figures/{0}_histogram_separated.pdf'.format(par))
    plt.close('all')

# Same plot for X(NH3)
# Need to apply Tex-based mask to get rid of some noisy data
fig,axes = plt.subplots(len(region_list),1,figsize=(4,5))
for i, ax in enumerate(fig.axes):
    region_i = i
    region = region_list[region_i]
    plot_param=plottingDictionary[region]
    par_file = dataDir + '{0}/parameterMaps/{0}_XNH3_{1}.fits'.format(region,extension)
    epar_file = dataDir + '{0}/parameterMaps/{0}_eTex_{1}_flag.fits'.format(region,extension)
    par_hdu = fits.open(par_file)
    par_data = par_hdu[0].data
    par_hdu.close()
    epar_hdu = fits.open(epar_file)
    epar_data = epar_hdu[0].data
    epar_hdu.close()
    pmin = -9.5
    pmax = -6.5
    #par_data[par_data == 0] = np.nan
    par_masked = mask_hist(par_data,epar_data,epar_limits[3],pmax,par_min=pmin)
    par_masked = par_masked[np.isfinite(par_masked)]
    hist(par_masked[par_masked !=0],bins='knuth',ax=ax,histtype='stepfilled',alpha=0.3,
         color=plot_colours[region_i],label=region,normed=True)
    if (i+1) != len(region_list):
        ax.set_xticklabels([])
    ax.set_xlim(pmin,pmax)
    #ax.set_ylim(0,hist_maxy_list[par_i])
    ax.yaxis.set_major_locator(ticker.MultipleLocator(1))
    ax.yaxis.set_minor_locator(ticker.MultipleLocator(0.5))
    ax.annotate('{0}'.format(region),xy=(0.97,0.7),xycoords='axes fraction',horizontalalignment='right')
#ax.legend(frameon=False)
ax.set_xlabel('log $X$(NH$_3$)')
fig.text(0.01,0.5,'P($X$(NH$_3$))',va='center',rotation='vertical')
#fig.tight_layout()
fig.savefig('figures/XNH3_histogram_separated.pdf',bbox_inches='tight')
plt.close('all')

'''
# All together

fig = plt.figure(figsize=(6.5,8))
for par_i in range(len(par_list)):
    par = par_list[par_i]
    label = label_list[par_i]
    #fig = plt.figure()
    #ax = plt.gca()
    ax = plt.subplot(3,2,par_i+1)
    for region_i in range(len(region_list)):
        region = region_list[region_i]
        plot_param=plottingDictionary[region]
        par_file = dataDir + '{0}/parameterMaps/{0}_{1}_{2}_flag.fits'.format(region,par,extension)
        epar_file = dataDir + '{0}/{0}_parameter_maps_{1}_trim.fits'.format(region,extension)
        epar_hdu = fits.open(epar_file)
        epar_data = epar_hdu[0].data[epar_ext[par_i],:,:]
        epar_hdu.close()
        par_hdu = fits.open(par_file)
        par_data = par_hdu[0].data
        par_hdu.close()
        pmin_list = plot_param['pmin_list']
        pmax_list = plot_param['pmax_list']
        pmin = np.max([pmin_list[par_i],np.nanmin(par_data)])
        pmax = np.min([pmax_list[par_i],np.nanmax(par_data)])
        par_masked = mask_hist(par_data,epar_data,epar_limits[par_i],pmax,par_min=pmin)
        par_masked = par_masked[np.isfinite(par_masked)]
        if par == 'Vlsr':
            bin_width = 0.3
            nbins = np.int((np.max(par_masked) - np.min(par_masked !=0))/bin_width)
            hist(par_masked[par_masked !=0],bins=nbins,ax=ax,histtype='stepfilled',alpha=0.3,
                 normed=True,color=plot_colours[region_i],label=region)
        else:
            hist(par_masked[par_masked !=0],bins='knuth',ax=ax,histtype='stepfilled',alpha=0.3,
                 normed=True,color=plot_colours[region_i],label=region)

    ax.set_xlabel(label)
    ax.set_ylabel('P(t)')
    #ax.set_ylabel('N')
    ax.set_xlim(hist_minx_list[par_i],hist_maxx_list[par_i])
    ax.set_ylim(0,hist_maxy_list[par_i])
    #ax.legend(frameon=False)
    #fig.savefig('figures/{0}_histogram_number.pdf'.format(par))
ax.legend(frameon=False,bbox_to_anchor=(2.0,1.0))
fig.tight_layout()
fig.savefig('figures/all_histograms.pdf')
plt.close('all')
'''


'''create scatterplot with confidence ellipsis

Author: Josef Perktold
License: BSD-3

TODO: update script to use sharex, sharey, and visible=False
    see http://www.scipy.org/Cookbook/Matplotlib/Multiple_Subplots_with_One_Axis_Label
    for sharex I need to have the ax of the last_row when editing the earlier
    rows. Or you axes_grid1, imagegrid
    http://matplotlib.sourceforge.net/mpl_toolkits/axes_grid/users/overview.html
'''

from statsmodels.compat.python import range
import numpy as np
from scipy import stats

from . import utils

__all__ = ['scatter_ellipse']

def _make_ellipse(mean, cov, ax, level=0.95, color=None):
    """Support function for scatter_ellipse."""
    from matplotlib.patches import Ellipse

    v, w = np.linalg.eigh(cov)
    u = w[0] / np.linalg.norm(w[0])
    angle = np.arctan(u[1]/u[0])
    angle = 180 * angle / np.pi # convert to degrees
    v = 2 * np.sqrt(v * stats.chi2.ppf(level, 2)) #get size corresponding to level
    ell = Ellipse(mean[:2], v[0], v[1], 180 + angle, facecolor='none',
                  edgecolor=color,
                  #ls='dashed',  #for debugging
                  lw=1.5)
    ell.set_clip_box(ax.bbox)
    ell.set_alpha(0.5)
    ax.add_artist(ell)

def scatter_ellipse(data, level=0.9, varnames=None, ell_kwds=None,
                    plot_kwds=None, add_titles=False, keep_ticks=False,
                    fig=None):
    """Create a grid of scatter plots with confidence ellipses.

    ell_kwds, plot_kdes not used yet

    looks ok with 5 or 6 variables, too crowded with 8, too empty with 1

    Parameters
    ----------
    data : array_like
        Input data.
    level : scalar, optional
        Default is 0.9.
    varnames : list of str, optional
        Variable names.  Used for y-axis labels, and if `add_titles` is True
        also for titles.  If not given, integers 1..data.shape[1] are used.
    ell_kwds : dict, optional
        UNUSED
    plot_kwds : dict, optional
        UNUSED
    add_titles : bool, optional
        Whether or not to add titles to each subplot.  Default is False.
        Titles are constructed from `varnames`.
    keep_ticks : bool, optional
        If False (default), remove all axis ticks.
    fig : Matplotlib figure instance, optional
        If given, this figure is simply returned.  Otherwise a new figure is
        created.

    Returns
    -------
    fig : Matplotlib figure instance
        If `fig` is None, the created figure.  Otherwise `fig` itself.

    """
    fig = utils.create_mpl_fig(fig)
    import matplotlib.ticker as mticker

    data = np.asanyarray(data)  #needs mean and cov
    nvars = data.shape[1]
    if varnames is None:
        #assuming single digit, nvars<=10  else use 'var%2d'
        varnames = ['var%d' % i for i in range(nvars)]

    plot_kwds_ = dict(ls='none', marker='.', color='k', alpha=0.5)
    if plot_kwds:
        plot_kwds_.update(plot_kwds)

    ell_kwds_= dict(color='k')
    if ell_kwds:
        ell_kwds_.update(ell_kwds)

    dmean = data.mean(0)
    dcov = np.cov(data, rowvar=0)

    for i in range(1, nvars):
        #print '---'
        ax_last=None
        for j in range(i):
            #print i,j, i*(nvars-1)+j+1
            ax = fig.add_subplot(nvars-1, nvars-1, (i-1)*(nvars-1)+j+1)
##                                 #sharey=ax_last) #sharey doesn't allow empty ticks?
##            if j == 0:
##                print 'new ax_last', j
##                ax_last = ax
##                ax.set_ylabel(varnames[i])
            #TODO: make sure we have same xlim and ylim

            formatter = mticker.FormatStrFormatter('% 3.1f')
            ax.yaxis.set_major_formatter(formatter)
            ax.xaxis.set_major_formatter(formatter)

            idx = np.array([j,i])
            ax.plot(*data[:,idx].T, **plot_kwds_)

            if np.isscalar(level):
                level = [level]
            for alpha in level:
                _make_ellipse(dmean[idx], dcov[idx[:,None], idx], ax, level=alpha,
                         **ell_kwds_)

            if add_titles:
                ax.set_title('%s-%s' % (varnames[i], varnames[j]))
            if not ax.is_first_col():
                if not keep_ticks:
                    ax.set_yticks([])
                else:
                    ax.yaxis.set_major_locator(mticker.MaxNLocator(3))
            else:
                ax.set_ylabel(varnames[i])
            if ax.is_last_row():
                ax.set_xlabel(varnames[j])
            else:
                if not keep_ticks:
                    ax.set_xticks([])
                else:
                    ax.xaxis.set_major_locator(mticker.MaxNLocator(3))

            dcorr = np.corrcoef(data, rowvar=0)
            dc = dcorr[idx[:,None], idx]
            xlim = ax.get_xlim()
            ylim = ax.get_ylim()
##            xt = xlim[0] + 0.1 * (xlim[1] - xlim[0])
##            yt = ylim[0] + 0.1 * (ylim[1] - ylim[0])
##            if dc[1,0] < 0 :
##                yt = ylim[0] + 0.1 * (ylim[1] - ylim[0])
##            else:
##                yt = ylim[1] - 0.2 * (ylim[1] - ylim[0])
            yrangeq = ylim[0] + 0.4 * (ylim[1] - ylim[0])
            if dc[1,0] < -0.25 or (dc[1,0] < 0.25 and dmean[idx][1] > yrangeq):
                yt = ylim[0] + 0.1 * (ylim[1] - ylim[0])
            else:
                yt = ylim[1] - 0.2 * (ylim[1] - ylim[0])
            xt = xlim[0] + 0.1 * (xlim[1] - xlim[0])
            ax.text(xt, yt, '$\\rho=%0.2f$'% dc[1,0])

    for ax in fig.axes:
        if ax.is_last_row(): # or ax.is_first_col():
            ax.xaxis.set_major_locator(mticker.MaxNLocator(3))
        if ax.is_first_col():
           ax.yaxis.set_major_locator(mticker.MaxNLocator(3))

    return fig


# Authors: Tom Dupre la Tour <tom.dupre-la-tour@m4x.org>
#          Olivier Grisel <olivier.grisel@ensta.org>
#
# License: BSD 3 clause

import matplotlib.pyplot as plt
import numpy as np
import gc
import time

from sklearn.externals.joblib import Memory
from sklearn.linear_model import (LogisticRegression, SGDClassifier)
from sklearn.datasets import fetch_rcv1
from sklearn.linear_model.sag import get_auto_step_size
from sklearn.linear_model.sag_fast import get_max_squared_sum

try:
    import lightning.classification as lightning_clf
except ImportError:
    lightning_clf = None

m = Memory(cachedir='.', verbose=0)

# compute logistic loss
def get_loss(w, intercept, myX, myy, C):
    n_samples = myX.shape[0]
    w = w.ravel()
    p = np.mean(np.log(1. + np.exp(-myy * (myX.dot(w) + intercept))))
    print("%f + %f" % (p, w.dot(w) / 2. / C / n_samples))
    p += w.dot(w) / 2. / C / n_samples
    return p

# We use joblib to cache individual fits. Note that we do not pass the dataset
# as argument as the hashing would be too slow, so we assume that the dataset
# never changes.
@m.cache()
def bench_one(name, clf_type, clf_params, n_iter):
    clf = clf_type(**clf_params)
    try:
        clf.set_params(max_iter=n_iter, random_state=42)
    except:
        clf.set_params(n_iter=n_iter, random_state=42)

    st = time.time()
    clf.fit(X, y)
    end = time.time()

    try:
        C = 1.0 / clf.alpha / n_samples
    except:
        C = clf.C

    try:
        intercept = clf.intercept_
    except:
        intercept = 0.

    train_loss = get_loss(clf.coef_, intercept, X, y, C)
    train_score = clf.score(X, y)
    test_score = clf.score(X_test, y_test)
    duration = end - st

    return train_loss, train_score, test_score, duration

def bench(clfs):
    for (name, clf, iter_range, train_losses, train_scores,
         test_scores, durations) in clfs:
        print("training %s" % name)
        clf_type = type(clf)
        clf_params = clf.get_params()

        for n_iter in iter_range:
            gc.collect()

            train_loss, train_score, test_score, duration = bench_one(
                name, clf_type, clf_params, n_iter)

            train_losses.append(train_loss)
            train_scores.append(train_score)
            test_scores.append(test_score)
            durations.append(duration)
            print("classifier: %s" % name)
            print("train_loss: %.8f" % train_loss)
            print("train_score: %.8f" % train_score)
            print("test_score: %.8f" % test_score)
            print("time for fit: %.8f seconds" % duration)
            print("")

        print("")
    return clfs

def plot_train_losses(clfs):
    plt.figure()
    for (name, _, _, train_losses, _, _, durations) in clfs:
        plt.plot(durations, train_losses, '-o', label=name)
        plt.legend(loc=0)
        plt.xlabel("seconds")
        plt.ylabel("train loss")

def plot_train_scores(clfs):
    plt.figure()
    for (name, _, _, _, train_scores, _, durations) in clfs:
        plt.plot(durations, train_scores, '-o', label=name)
        plt.legend(loc=0)
        plt.xlabel("seconds")
        plt.ylabel("train score")
        plt.ylim((0.92, 0.96))

def plot_test_scores(clfs):
    plt.figure()
    for (name, _, _, _, _, test_scores, durations) in clfs:
        plt.plot(durations, test_scores, '-o', label=name)
        plt.legend(loc=0)
        plt.xlabel("seconds")
        plt.ylabel("test score")
        plt.ylim((0.92, 0.96))

def plot_dloss(clfs):
    plt.figure()
    pobj_final = []
    for (name, _, _, train_losses, _, _, durations) in clfs:
        pobj_final.append(train_losses[-1])

    indices = np.argsort(pobj_final)
    pobj_best = pobj_final[indices[0]]

    for (name, _, _, train_losses, _, _, durations) in clfs:
        log_pobj = np.log(abs(np.array(train_losses) - pobj_best)) / np.log(10)

        plt.plot(durations, log_pobj, '-o', label=name)
        plt.legend(loc=0)
        plt.xlabel("seconds")
        plt.ylabel("log(best - train_loss)")

rcv1 = fetch_rcv1()
X = rcv1.data
n_samples, n_features = X.shape

# consider the binary classification problem 'CCAT' vs the rest
ccat_idx = rcv1.target_names.tolist().index('CCAT')
y = rcv1.target.tocsc()[:, ccat_idx].toarray().ravel().astype(np.float64)
y[y == 0] = -1

# parameters
C = 1.
fit_intercept = True
tol = 1.0e-14

# max_iter range
sgd_iter_range = list(range(1, 121, 10))
newton_iter_range = list(range(1, 25, 3))
lbfgs_iter_range = list(range(1, 242, 12))
liblinear_iter_range = list(range(1, 37, 3))
liblinear_dual_iter_range = list(range(1, 85, 6))
sag_iter_range = list(range(1, 37, 3))

clfs = [
    ("LR-liblinear",
     LogisticRegression(C=C, tol=tol,
                        solver="liblinear", fit_intercept=fit_intercept,
                        intercept_scaling=1),
     liblinear_iter_range, [], [], [], []),
    ("LR-liblinear-dual",
     LogisticRegression(C=C, tol=tol, dual=True,
                        solver="liblinear", fit_intercept=fit_intercept,
                        intercept_scaling=1),
     liblinear_dual_iter_range, [], [], [], []),
    ("LR-SAG",
     LogisticRegression(C=C, tol=tol,
                        solver="sag", fit_intercept=fit_intercept),
     sag_iter_range, [], [], [], []),
    ("LR-newton-cg",
     LogisticRegression(C=C, tol=tol, solver="newton-cg",
                        fit_intercept=fit_intercept),
     newton_iter_range, [], [], [], []),
    ("LR-lbfgs",
     LogisticRegression(C=C, tol=tol,
                        solver="lbfgs", fit_intercept=fit_intercept),
     lbfgs_iter_range, [], [], [], []),
    ("SGD",
     SGDClassifier(alpha=1.0 / C / n_samples, penalty='l2', loss='log',
                   fit_intercept=fit_intercept, verbose=0),
     sgd_iter_range, [], [], [], [])]

if lightning_clf is not None and not fit_intercept:
    alpha = 1. / C / n_samples
    # compute the same step_size than in LR-sag
    max_squared_sum = get_max_squared_sum(X)
    step_size = get_auto_step_size(max_squared_sum, alpha, "log",
                                   fit_intercept)

    clfs.append(
        ("Lightning-SVRG",
         lightning_clf.SVRGClassifier(alpha=alpha, eta=step_size,
                                      tol=tol, loss="log"),
         sag_iter_range, [], [], [], []))
    clfs.append(
        ("Lightning-SAG",
         lightning_clf.SAGClassifier(alpha=alpha, eta=step_size,
                                     tol=tol, loss="log"),
         sag_iter_range, [], [], [], []))

    # We keep only 200 features, to have a dense dataset,
    # and compare to lightning SAG, which seems incorrect in the sparse case.
    X_csc = X.tocsc()
    nnz_in_each_features = X_csc.indptr[1:] - X_csc.indptr[:-1]
    X = X_csc[:, np.argsort(nnz_in_each_features)[-200:]]
    X = X.toarray()
    print("dataset: %.3f MB" % (X.nbytes / 1e6))

# Split training and testing. Switch train and test subset compared to
# LYRL2004 split, to have a larger training dataset.
n = 23149
X_test = X[:n, :]
y_test = y[:n]
X = X[n:, :]
y = y[n:]

clfs = bench(clfs)

plot_train_scores(clfs)
plot_test_scores(clfs)
plot_train_losses(clfs)
plot_dloss(clfs)
plt.show()


from __future__ import division

from math import log, sqrt, pi
from barak.utilities import adict
from barak.absorb import split_trans_name
from barak.io import parse_config, loadobj
from barak.interp import AkimaSpline, MapCoord_Interpolator
from cloudy.utils import read_observed
import numpy as np
import os
from glob import glob
from barak.plot import get_nrows_ncols, puttext

from matplotlib.ticker import AutoMinorLocator

import astropy.constants as c
import astropy.units as u
from astropy.table import Table

import pylab as plt

import sys

# dex 1 sigma error in UVB (and so nH)
Unorm_sig = 0.3

USE_HEXBIN = True

def make_cmap_red():
    from matplotlib.colors import LinearSegmentedColormap
    x = np.linspace(0,1,9)
    cm = plt.cm.Reds(x)
    r,g,b = cm[:,0], cm[:,1], cm[:,2]

    g[0] = 1
    b[0] = 1

    cdict = dict(red=zip(x, r, r), green=zip(x, g, g), blue=zip(x, b, b))

    return LinearSegmentedColormap('red_nhmc', cdict)

def make_cmap_blue():
    from matplotlib.colors import LinearSegmentedColormap
    x = np.linspace(0,1,15)
    cm = plt.cm.Blues(x)
    r,g,b = cm[:,0], cm[:,1], cm[:,2]

    g[0] = 1
    b[0] = 1
    r[1:10] = r[4:13]
    g[1:10] = g[4:13]
    b[1:10] = b[4:13]

    cdict = dict(red=zip(x, r, r), green=zip(x, g, g), blue=zip(x, b, b))

    return LinearSegmentedColormap('blue_nhmc', cdict)

def find_min_interval(x, alpha):
    """ Determine the minimum interval containing a given probability.

    x is an array of parameter values (such as from an MCMC trace).

    alpha (0 -> 1) is the desired probability encompassed by the
    interval.

    Inspired by the pymc function of the same name.
    """

    assert len(x) > 1
    x = np.sort(x)

    # Initialize interval
    min_int = None, None

    # Number of elements in trace
    n = len(x)

    # Start at far left
    end0 = int(n*alpha)
    start, end = 0, end0

    # Initialize minimum width to large value
    min_width = np.inf

    for i in xrange(n - end0):
        hi, lo = x[end+i], x[start+i]
        width = hi - lo
        if width < min_width:
            min_width = width
            min_int = lo, hi

    return min_int

def make_interpolators_uvbtilt(trans, simnames):
    """ Make interpolators including different UV slopes, given by the
    simulation names.

    simname naming scheme should be (uvb_k00, uvb_k01, uvb_k02, ...),

    uvb k values must be sorted in ascending order!
    """

    Models = []
    aUV = []
    for simname in simnames:
        # need to define prefix, SIMNAME
        gridname = os.path.join(simname, 'grid.cfg')

        #print 'Reading', gridname
        cfg = parse_config(gridname)
        aUV.append(cfg.uvb_tilt)

        name = os.path.join(simname, cfg.prefix + '_grid.sav.gz')
        #print 'Reading', name
        M = loadobj(name)
        M = adict(M)

        Uconst = (M.U + M.nH)[0]
        #print 'Uconst', Uconst, cfg.uvb_tilt
        assert np.allclose(Uconst, M.U + M.nH)
        Models.append(M)

    ##########################################################################
    # Interpolate cloudy grids onto a finer scale for plotting and
    # likelihood calculation
    ##########################################################################

    roman_map = {'I':0, 'II':1, 'III':2, 'IV':3, 'V':4, 'VI':5,
                 'VII':6, 'VIII':7, 'IX':8, 'X':9, '2':2}
    Ncloudy = {}
    Ncloudy_raw = {}
    #print 'Interpolating...'
    for tr in trans:
        shape = len(M.NHI), len(M.nH), len(M.Z), len(aUV)
        Nvals = np.zeros(shape)
        if tr == 'Tgas':
            for i,M in enumerate(Models):
                Nvals[:,:,:,i] = M['Tgas'][:,:,:,0]
        elif tr == 'NH':
            for i,M in enumerate(Models):
                logNHI = M.N['H'][:,:,:,0]
                logNHII = M.N['H'][:,:,:,1]
                logNHtot = np.log10(10**logNHI + 10**logNHII)
                Nvals[:,:,:,i] = logNHtot
        elif tr in ['CII*']:
            for i,M in enumerate(Models):
                Nvals[:,:,:,i] = M.Nex[tr][:,:,:]
        else:
            atom, stage = split_trans_name(tr)
            ind = roman_map[stage]
            for i,M in enumerate(Models):
                Nvals[:,:,:,i] = M.N[atom][:,:,:,ind]

        # use ndimage.map_coordinates (which is spline interpolation)
        coord = M.NHI, M.nH, M.Z, aUV
        try:
            Ncloudy[tr] = MapCoord_Interpolator(Nvals, coord)
        except:
            import pdb; pdb.set_trace()

        Ncloudy_raw[tr] = Nvals

    #print 'done'
    return Ncloudy, Ncloudy_raw, Models, np.array(aUV, np.float)

def triplot(names, vals, sigvals, fig, indirect={}, labels=None, fontsize=14):

    from barak.plot import hist_yedge, hist_xedge, puttext
    npar = len(names)
    bins = {}
    for n in names:
        x0, x1 = vals[n].min(), vals[n].max()
        dx = x1 - x0
        lo = x0 - 0.1*dx
        hi = x1 + 0.1*dx
        bins[n] = np.linspace(lo, hi, 20)

    axes = {}
    for i0,n0 in enumerate(names):
        for i1,n1 in enumerate(names):
            if i0 == i1:# or i1 < i0: # uncomment to keep just one triangle.
                continue
            ax = fig.add_subplot(npar,npar, i0 * npar + i1 + 1)
            ax.locator_params(tight=True, nbins=8)
            ax.xaxis.set_minor_locator(AutoMinorLocator())
            ax.yaxis.set_minor_locator(AutoMinorLocator())
            axes[(n0 + ' ' + n1)] = ax

            y,x = vals[n0], vals[n1]
            if USE_HEXBIN:
                ax.hexbin(x,y,cmap=CM, gridsize=40,linewidths=0.1)
            else:
                ax.plot(x,y,'r.', ms=0.5, mew=0)#, alpha=0.5)

            color = 'k' if n0 not in indirect else 'g'
            text = labels[n0] if labels is not None else n0
            puttext(0.05, 0.95, text, ax, color=color ,fontsize=fontsize, va='top')
            color = 'k' if n1 not in indirect else 'g'
            text = labels[n1] if labels is not None else n1
            puttext(0.95, 0.08, text, ax, color=color ,fontsize=fontsize, ha='right')
            # set limits
            y0, y1 = np.percentile(vals[n0], [5, 95])
            dy = y1 - y0
            ax.set_ylim(y0 - dy, y1 + dy)
            x0, x1 = np.percentile(vals[n1], [5, 95])
            dx = x1 - x0
            ax.set_xlim(x0 - dx, x1 + dx)

            c = 'k'
            if i0 == 0:
                ax.xaxis.set_tick_params(labeltop='on')
                ax.xaxis.set_tick_params(labelbottom='off')
                for t in ax.get_xticklabels():
                    t.set_rotation(60)
            elif i0 == npar-1 or (i0 == npar-2 and i1 == npar-1):
                hist_xedge(vals[n1], ax, color='forestgreen',
                           fill=dict(color='forestgreen',alpha=0.3),
                           bins=bins[n1], loc='bottom')
                ax.axvline(sigvals[n1][0], ymax=0.2, color=c, lw=0.5)
                ax.axvline(sigvals[n1][1], ymax=0.2, color=c, lw=0.5)
                cen = sum(sigvals[n1]) / 2.
                ax.axvline(cen, ymax=0.2, color=c, lw=1.5)
                for t in ax.get_xticklabels():
                    t.set_rotation(60)
            else:
                ax.set_xticklabels('')

            if not (i1 == 0 or (i0 == 0 and i1 == 1) or i1 == npar-1):
                ax.set_yticklabels('')

            if (i0 == 0 and i1 == 1) or i1 == 0:
                hist_yedge(vals[n0], ax, color='forestgreen',
                           fill=dict(color='forestgreen',alpha=0.3),
                           bins=bins[n0], loc='left')
                ax.axhline(sigvals[n0][0], xmax=0.2, color=c, lw=0.5)
                ax.axhline(sigvals[n0][1], xmax=0.2, color=c, lw=0.5)
                cen = sum(sigvals[n0]) / 2.
                ax.axhline(cen, xmax=0.2, color=c, lw=1.5)

            if i1 == npar - 1:
                ax.yaxis.set_tick_params(labelright='on')
                ax.yaxis.set_tick_params(labelleft='off')

            #ax.minorticks_on()

    return axes

if 1:
    print_header = False
    if len(sys.argv[1:]) > 0 and sys.argv[1] ==  '--header':
        print_header = True
if 1:
    ##################################################
    # Read configuration file, set global variables
    ##################################################
    testing = 0

    cfgname = 'model.cfg'
    # we only need the cfg file for the prefix of the cloudy runs and
    # the name of the file with the observed column densities.
    opt = parse_config(cfgname)

    simnames = sorted(glob(opt['simname']))
    #print opt['simname']
    #print simnames

    #CM = make_cmap_blue()  # plt.cm.binary
    #CM = make_cmap_red()  # plt.cm.binary
    CM = plt.cm.gist_heat_r # plt.cm.binary
    #CM = plt.cm.afmhot_r # plt.cm.binary

    #CM = plt.cm.bone_r # plt.cm.binary
    #CM = plt.cm.terrain_r # plt.cm.binary
    #CM = plt.cm.ocean_r # plt.cm.binary

    trans = 'Tgas', 'NH'

if 1:
    ################################################################
    # Read the cloudy grids and make the interpolators
    ################################################################
    Ncloudy, Ncloudy_raw, Models, aUV = make_interpolators_uvbtilt(
        trans, simnames)
    M = Models[0]
    #import pdb; pdb.set_trace()
    Uconst_vals = []
    for model in Models:
        Uconst_vals.append((model['U'] + model['nH'])[0])

    # note it's a function of aUV!
    Uconst = AkimaSpline(aUV, Uconst_vals)

    # Now find the parameter chains
    samples = loadobj('samples_mcmc.sav.gz')
    nwalkers, nsamples, npar = samples['chain'].shape
    parvals = samples['chain'].reshape(-1, npar)

    PAR = samples['par']
    assert PAR['names'][-1] == 'aUV'
    assert PAR['names'][-2] == 'Z'
    assert PAR['names'][-3] == 'nH'
    assert PAR['names'][-4] == 'NHI'

    aUV = parvals[:,-1]
    logZ = parvals[:,-2]
    lognH = parvals[:,-3]
    logNHI = parvals[:,-4]

    logU = Uconst(aUV) - lognH

    #import pdb; pdb.set_trace()

    # call the interpolators with these parameter values.
    logT = Ncloudy['Tgas'](parvals[:,-4:].T)
    logNtot = Ncloudy['NH'](parvals[:,-4:].T)
    # note this is log of D in kpc
    logD = logNtot - lognH - np.log10(c.kpc.to(u.cm).value)

    logP = logT + lognH

    #import pdb; pdb.set_trace()

    H_massfrac = 0.76  # (1 / mu)

    # Joe's mass calculation
    mass = 4./3. * pi * (3./4. * 10**logD * u.kpc)**3 * 10**lognH * \
           u.cm**-3 * u.M_p / H_massfrac

    # D = NH / nH

    logM = np.log10(mass.to(u.M_sun).value)

if 1:
    # print out the results and uncertainties
    vals = dict(U=logU, T=logT, N=logNtot, D=logD, P=logP, M=logM,
                nH=lognH, aUV=aUV, NHI=logNHI, Z=logZ)

    levels = 0.6827, 0.9545

    sigvals = {}
    for key in vals:
        sigvals[key] = find_min_interval(vals[key], levels[0])

    if print_header:
        print r'$\log(Z/Z_\odot)$&$\alpha_{UV}$ & $\log \nH$ & $\log U$& $\log \NHI$ & $\log \NH$& $\log T$ & $\log (P/k)$& $\log D$ & $\log M$ \\'
        print r'                 &              & (\cmmm)    &         & (\cmm)      &  (\cmm)   & (K)      &  (\cmmm K)  &  (kpc)   & (\msun)  \\'

    s = ''
    ALLPAR = 'Z aUV nH U NHI N T P D M'.split()
    for key in ALLPAR:
        sig = 0.5 * (sigvals[key][1] - sigvals[key][0])
        val = 0.5 * (sigvals[key][1] + sigvals[key][0])
        if key in {'nH', 'D', 'P'}:
            sig1 = np.hypot(sig, Unorm_sig)
            s += '$%.2f\\pm%.2f(%.2f)$ &' % (val, sig1, sig)
        elif key == 'M':
            sig1 = np.hypot(sig, 2*Unorm_sig)
            s += '$%.2f\\pm%.2f(%.2f)$ &' % (val, sig1, sig)
        else:
            s += '$%.2f\\pm%.2f$ &' % (val, sig)

    print s[:-1] +  r'\\'

if 1:
    labels = dict(U='$U$', Z='$Z$', NHI='$N_\mathrm{HI}$', aUV=r'$\alpha_\mathrm{UV}$',
                  T='$T$', P='$P$', N='$N_\mathrm{H}$', D='$D$', M='$Mass$')

if 0:
    fig = plt.figure(figsize=(12,12))
    fig.subplots_adjust(left=0.05, bottom=0.05, top=0.94,right=0.94, wspace=1e-4,hspace=1e-4)
    plt.rc('xtick', labelsize=8)
    plt.rc('ytick', labelsize=8)

    names = 'U Z NHI aUV T P N D M'.split()

    #direct = 'U Z NHI aUV'.split()
    axes = triplot(names, vals, sigvals, fig, labels=labels)
    plt.savefig('par.png', dpi=200)

if 1:
    fig = plt.figure(figsize=(8,8))
    fig.subplots_adjust(left=0.095, bottom=0.105, top=0.94,right=0.94, wspace=1e-4,hspace=1e-4)
    plt.rc('xtick', labelsize=9.5)
    plt.rc('ytick', labelsize=9.5)

    names = 'U Z N aUV'.split()
    axes = triplot(names, vals, sigvals, fig, labels=labels, fontsize=16)

    axes['U Z'].set_ylabel('$\log_{10}U$')
    axes['Z U'].set_ylabel('$\log_{10}[Z/Z_\odot]$')
    axes['N U'].set_ylabel('$\log_{10}N_\mathrm{H}$')
    axes['aUV U'].set_ylabel(r'$\log_{10}\alpha_\mathrm{UV}$')
    axes['aUV U'].set_xlabel('$\log_{10}U$')
    axes['aUV Z'].set_xlabel('$\log_{10}[Z/Z_\odot]$')
    axes['aUV N'].set_xlabel('$\log_{10}N_\mathrm{H}$')
    axes['N aUV'].set_xlabel(r'$\log_{10}\alpha_\mathrm{UV}$')
    # special case:
    if os.path.abspath('.') == '/Users/ncrighton/Projects/MPIA_QSO_LBG/Cloudy/J0004_NHI_2/comp1/final':
        for k in ('N U', 'N Z', 'N aUV'):
            axes[k].set_ylim(17.3, 19.2)
        for k in ('U N', 'Z N', 'aUV N'):
            axes[k].set_xlim(17.3, 19.2)

    #plt.savefig('par2.pdf')
    plt.savefig('par2.png',dpi=250)


from pathlib import Path

import pandas as pd

import ibis
import ibis.expr.operations as ops
from ibis.backends.tests.base import BackendTest, RoundHalfToEven

class TestConf(BackendTest, RoundHalfToEven):
    check_names = False
    additional_skipped_operations = frozenset({ops.StringSQLLike})
    supported_to_timestamp_units = BackendTest.supported_to_timestamp_units | {
        'ns'
    }
    supports_divide_by_zero = True
    returned_timestamp_unit = 'ns'

    @staticmethod
    def connect(data_directory: Path) -> ibis.client.Client:
        return ibis.pandas.connect(
            {
                'functional_alltypes': pd.read_csv(
                    str(data_directory / 'functional_alltypes.csv'),
                    index_col=None,
                    dtype={'bool_col': bool, 'string_col': str},
                    parse_dates=['timestamp_col'],
                    encoding='utf-8',
                ),
                'batting': pd.read_csv(str(data_directory / 'batting.csv')),
                'awards_players': pd.read_csv(
                    str(data_directory / 'awards_players.csv')
                ),
            }
        )


"""
=============================
Ledoit-Wolf vs OAS estimation
=============================

The usual covariance maximum likelihood estimate can be regularized
using shrinkage. Ledoit and Wolf proposed a close formula to compute
the asymptotical optimal shrinkage parameter (minimizing a MSE
criterion), yielding the Ledoit-Wolf covariance estimate.

Chen et al. proposed an improvement of the Ledoit-Wolf shrinkage
parameter, the OAS coefficient, whose convergence is significantly
better under the assumption that the data are gaussian.

This example, inspired from Chen's publication [1], shows a comparison
of the estimated MSE of the LW and OAS methods, using gaussian
distributed data.

[1] "Shrinkage Algorithms for MMSE Covariance Estimation"
Chen et al., IEEE Trans. on Sign. Proc., Volume 58, Issue 10, October 2010.

"""
print __doc__

import numpy as np
import pylab as pl
from scipy.linalg import toeplitz, cholesky

from sklearn.covariance import LedoitWolf, OAS

np.random.seed(0)
###############################################################################
n_features = 100
# simulation covariance matrix (AR(1) process)
r = 0.1
real_cov = toeplitz(r ** np.arange(n_features))
coloring_matrix = cholesky(real_cov)

n_samples_range = np.arange(6, 31, 1)
repeat = 100
lw_mse = np.zeros((n_samples_range.size, repeat))
oa_mse = np.zeros((n_samples_range.size, repeat))
lw_shrinkage = np.zeros((n_samples_range.size, repeat))
oa_shrinkage = np.zeros((n_samples_range.size, repeat))
for i, n_samples in enumerate(n_samples_range):
    for j in range(repeat):
        X = np.dot(
            np.random.normal(size=(n_samples, n_features)), coloring_matrix.T)

        lw = LedoitWolf(store_precision=False, assume_centered=True)
        lw.fit(X)
        lw_mse[i, j] = lw.error_norm(real_cov, scaling=False)
        lw_shrinkage[i, j] = lw.shrinkage_

        oa = OAS(store_precision=False, assume_centered=True)
        oa.fit(X)
        oa_mse[i, j] = oa.error_norm(real_cov, scaling=False)
        oa_shrinkage[i, j] = oa.shrinkage_

# plot MSE
pl.subplot(2, 1, 1)
pl.errorbar(n_samples_range, lw_mse.mean(1), yerr=lw_mse.std(1),
            label='Ledoit-Wolf', color='g')
pl.errorbar(n_samples_range, oa_mse.mean(1), yerr=oa_mse.std(1),
            label='OAS', color='r')
pl.ylabel("Squared error")
pl.legend(loc="upper right")
pl.title("Comparison of covariance estimators")
pl.xlim(5, 31)

# plot shrinkage coefficient
pl.subplot(2, 1, 2)
pl.errorbar(n_samples_range, lw_shrinkage.mean(1), yerr=lw_shrinkage.std(1),
            label='Ledoit-Wolf', color='g')
pl.errorbar(n_samples_range, oa_shrinkage.mean(1), yerr=oa_shrinkage.std(1),
            label='OAS', color='r')
pl.xlabel("n_samples")
pl.ylabel("Shrinkage")
pl.legend(loc="lower right")
pl.ylim(pl.ylim()[0], 1. + (pl.ylim()[1] - pl.ylim()[0]) / 10.)
pl.xlim(5, 31)

pl.show()


"""
Test Data Showing Inset Plots
=============================

This example demonstrates the use of a single 3D data cube with time, latitude
and longitude dimensions to plot a temperature series for a single latitude
coordinate, with an inset plot of the data region.

"""

import matplotlib.pyplot as plt
import numpy as np
import iris
import cartopy.crs as ccrs
import iris.quickplot as qplt
import iris.plot as iplt

def main():
    # Load the data
    with iris.FUTURE.context(netcdf_promote=True):
        cube1 = iris.load_cube(iris.sample_data_path('ostia_monthly.nc'))
    # Slice into cube to retrieve data for the inset map showing the
    # data region
    region = cube1[-1, :, :]
    # Average over latitude to reduce cube to 1 dimension
    plot_line = region.collapsed('latitude', iris.analysis.MEAN)

    # Open a window for plotting
    fig = plt.figure()
    # Add a single subplot (axes). Could also use "ax_main = plt.subplot()"
    ax_main = fig.add_subplot(1, 1, 1)
    # Produce a quick plot of the 1D cube
    qplt.plot(plot_line)

    # Set x limits to match the data
    ax_main.set_xlim(0, plot_line.coord('longitude').points.max())
    # Adjust the y limits so that the inset map won't clash with main plot
    ax_main.set_ylim(294, 310)
    ax_main.set_title('Meridional Mean Temperature')
    # Add grid lines
    ax_main.grid()

    # Add a second set of axes specifying the fractional coordinates within
    # the figure with bottom left corner at x=0.55, y=0.58 with width
    # 0.3 and height 0.25.
    # Also specify the projection
    ax_sub = fig.add_axes([0.55, 0.58, 0.3, 0.25],
                          projection=ccrs.Mollweide(central_longitude=180))

    # Use iris.plot (iplt) here so colour bar properties can be specified
    # Also use a sequential colour scheme to reduce confusion for those with
    # colour-blindness
    iplt.pcolormesh(region, cmap='Blues')
    # Manually set the orientation and tick marks on your colour bar
    ticklist = np.linspace(np.min(region.data), np.max(region.data), 4)
    plt.colorbar(orientation='horizontal', ticks=ticklist)
    ax_sub.set_title('Data Region')
    # Add coastlines
    ax_sub.coastlines()
    # request to show entire map, using the colour mesh on the data region only
    ax_sub.set_global()

    qplt.show()

if __name__ == '__main__':
    main()