# Source code for MDAnalysis.analysis.encore.covariance

    # -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
    # vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4
    #
    # MDAnalysis --- https://www.mdanalysis.org
    # Copyright (c) 2006-2017 The MDAnalysis Development Team and contributors
    # (see the file AUTHORS for the full list of names)
    #
    # Released under the Lesser GNU Public Licence, v2.1 or any higher version
    #
    # Please cite your use of MDAnalysis in published work:
    #
    # R. J. Gowers, M. Linke, J. Barnoud, T. J. E. Reddy, M. N. Melo, S. L. Seyler,
    # D. L. Dotson, J. Domanski, S. Buchoux, I. M. Kenney, and O. Beckstein.
    # MDAnalysis: A Python package for the rapid analysis of molecular dynamics
    # simulations. In S. Benthall and S. Rostrup editors, Proceedings of the 15th
    # Python in Science Conference, pages 102-109, Austin, TX, 2016. SciPy.
    # doi: 10.25080/majora-629e541a-00e
    #
    # N. Michaud-Agrawal, E. J. Denning, T. B. Woolf, and O. Beckstein.
    # MDAnalysis: A Toolkit for the Analysis of Molecular Dynamics Simulations.
    # J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787
    #
    """
    Covariance calculation --- :mod:`encore.covariance`
    =====================================================================

    The module contains functions to estimate the covariance matrix of
    an ensemble of structures.

    :Author: Matteo Tiberti, Wouter Boomsma, Tone Bengtsen

    .. versionadded:: 0.16.0

    .. deprecated:: 2.8.0
       This module is deprecated in favour of the
       MDAKit `mdaencore <https://mdanalysis.org/mdaencore/>`_ and will be removed
       in MDAnalysis 3.0.0.

    """
    import numpy as np

    def ml_covariance_estimator(coordinates, reference_coordinates=None):
        """
        Standard maximum likelihood estimator of the covariance matrix.

        Parameters
        ----------

        coordinates : numpy.array
            Flattened array of coordiantes

        reference_coordinates : numpy.array
            Optional reference to use instead of mean

        Returns
        -------

        cov_mat : numpy.array
            Estimate of  covariance matrix

        """

        if reference_coordinates is not None:

            # Offset from reference
            coordinates_offset = coordinates - reference_coordinates

        else:
            # Normal covariance calculation: distance to the average
            coordinates_offset = coordinates - np.average(coordinates, axis=0)

        # Calculate covariance manually
        coordinates_cov = np.zeros((coordinates.shape[1], coordinates.shape[1]))
        for frame in coordinates_offset:
            coordinates_cov += np.outer(frame, frame)
        coordinates_cov /= coordinates.shape[0]

        return coordinates_cov

    def shrinkage_covariance_estimator(
        coordinates, reference_coordinates=None, shrinkage_parameter=None
    ):
        """
        Shrinkage estimator of the covariance matrix using the method described in

        Improved Estimation of the Covariance Matrix of Stock Returns With an
        Application to Portfolio Selection. Ledoit, O.; Wolf, M., Journal of
        Empirical Finance, 10, 5, 2003

        This implementation is based on the matlab code made available by Olivier
        Ledoit on his website:
        http://www.ledoit.net/ole2_abstract.htm

        Parameters
        ----------

        coordinates : numpy.array
            Flattened array of coordinates

        reference_coordinates: numpy.array
            Optional reference to use instead of mean

        shrinkage_parameter: None or float
            Optional shrinkage parameter

        Returns
        --------

        cov_mat : nump.array
            Covariance matrix
        """

        x = coordinates
        t = x.shape[0]
        n = x.shape[1]

        mean_x = np.average(x, axis=0)

        # Use provided coordinates as "mean" if provided
        if reference_coordinates is not None:
            mean_x = reference_coordinates

        x = x - mean_x
        xmkt = np.average(x, axis=1)

        # Call maximum likelihood estimator (note the additional column)
        sample = (
            ml_covariance_estimator(np.hstack([x, xmkt[:, np.newaxis]]), 0)
            * (t - 1)
            / float(t)
        )

        # Split covariance matrix into components
        covmkt = sample[0:n, n]
        varmkt = sample[n, n]
        sample = sample[:n, :n]

        # Prior
        prior = np.outer(covmkt, covmkt) / varmkt
        prior[np.ma.make_mask(np.eye(n))] = np.diag(sample)

        # If shrinkage parameter is not set, estimate it
        if shrinkage_parameter is None:

            # Frobenius norm
            c = np.linalg.norm(sample - prior, ord="fro") ** 2

            y = x**2
            p = 1 / float(t) * np.sum(np.dot(np.transpose(y), y)) - np.sum(
                np.sum(sample**2)
            )
            rdiag = 1 / float(t) * np.sum(np.sum(y**2)) - np.sum(
                np.diag(sample) ** 2
            )
            z = x * np.repeat(xmkt[:, np.newaxis], n, axis=1)
            v1 = (
                1 / float(t) * np.dot(np.transpose(y), z)
                - np.repeat(covmkt[:, np.newaxis], n, axis=1) * sample
            )
            roff1 = (
                np.sum(
                    v1 * np.transpose(np.repeat(covmkt[:, np.newaxis], n, axis=1))
                )
                / varmkt
                - np.sum(np.diag(v1) * covmkt) / varmkt
            )
            v3 = 1 / float(t) * np.dot(np.transpose(z), z) - varmkt * sample
            roff3 = (
                np.sum(v3 * np.outer(covmkt, covmkt)) / varmkt**2
                - np.sum(np.diag(v3) * covmkt**2) / varmkt**2
            )
            roff = 2 * roff1 - roff3
            r = rdiag + roff

            # Shrinkage constant
            k = (p - r) / c
            shrinkage_parameter = max(0, min(1, k / float(t)))

        # calculate covariance matrix
        sigma = shrinkage_parameter * prior + (1 - shrinkage_parameter) * sample

        return sigma

    def covariance_matrix(
        ensemble,
        select="name CA",
        estimator=shrinkage_covariance_estimator,
        weights="mass",
        reference=None,
    ):
        """
        Calculates (optionally mass weighted) covariance matrix

        Parameters
        ----------
        ensemble : Universe object
            The structural ensemble
        select : str (optional)
            Atom selection string in the MDAnalysis format.
        estimator : function (optional)
            Function that estimates the covariance matrix. It requires at least
            a "coordinates" numpy array (of shape (N,M,3), where N is the number
            of frames and M the number of atoms). See ml_covariance_estimator and
            shrinkage_covariance_estimator for reference.
        weights : str/array_like (optional)
            specify weights. If ``'mass'`` then chose masses of ensemble atoms, if ``None`` chose uniform weights
        reference : MDAnalysis.Universe object (optional)
            Use the distances to a specific reference structure rather than the
            distance to the mean.

        Returns
        -------
        cov_mat : numpy.array
            Covariance matrix

        """
        # Extract coordinates from ensemble
        coordinates = ensemble.trajectory.timeseries(
            ensemble.select_atoms(select), order="fac"
        )

        # Flatten coordinate matrix into n_frame x n_coordinates
        coordinates = np.reshape(coordinates, (coordinates.shape[0], -1))

        # Extract coordinates from reference structure, if specified
        reference_coordinates = None
        if reference is not None:
            # Select the same atoms in reference structure
            reference_atom_selection = reference.select_atoms(select)
            reference_coordinates = reference_atom_selection.atoms.positions

            # Flatten reference coordinates
            reference_coordinates = reference_coordinates.flatten()

        sigma = estimator(coordinates, reference_coordinates)

        # Optionally correct with weights
        if weights is not None:
            # Calculate mass-weighted covariance matrix
            if (
                not isinstance(weights, (list, tuple, np.ndarray))
                and weights == "mass"
            ):
                if select:
                    weights = ensemble.select_atoms(select).masses
                else:
                    weights = ensemble.atoms.masses
            else:
                if select:
                    req_len = ensemble.select_atoms(select).n_atoms
                else:
                    req_len = ensemble.atoms.n_atoms
                if req_len != len(weights):
                    raise ValueError(
                        "number of weights is unequal to number of "
                        "atoms in ensemble"
                    )

            # broadcast to a (len(weights), 3) array
            weights = np.repeat(weights, 3)

            weight_matrix = np.sqrt(np.identity(len(weights)) * weights)
            sigma = np.dot(weight_matrix, np.dot(sigma, weight_matrix))

        return sigma