ENH: stats.multivariate: introduce Covariance class and subclasses

[mdhaber]

Explore the use of a Covariance class to address the need to accept alternative representations of a multivariate distribution shape matrix and avoid re-processing this matrix each time a frozen distribution method is called.

[mdhaber]

@cached_property is probably not appropriate all over the place, here. The important thing with this Covariance class is to define an interface and provide default documentation (which would not need to be rewritten by subclasses).

[mdhaber]

This name is taken from KDE. Really, I think I'd just replace this method with xTPx, which would do the equivalent of np.sum(np.square(np.dot(dev, prec_U)), axis=-1).

[mdhaber]

It turns out we can get $x^T \Sigma^{-1} x$ without the _J flipping. See scipygh-16987.

[mdhaber]

One advantage of the eigendecomposition is that it supports singular covariance matrices. But we can compute the key $x^T \Sigma^* x$ product using LDL factors, too.

For nonsingular matrices, this is what it looks like:

import numpy as np
from scipy import linalg

n = 40
rng = np.random.default_rng(14)
x = rng.random(size=(n, 1))
A = rng.random(size=(n, n))
A = A @ A.T

# Eigendecomposition route
w, V = linalg.eigh(A)
z = x.T @ (V * w**(-0.5))
xSxT = (z**2).sum()

# LDL route
L, D, perm = linalg.ldl(A)
y = linalg.solve_triangular(L[perm], x[perm], lower=True).T * np.diag(D)**(-0.5)
xSxT2 = (y**2).sum()

print(xSxT, xSxT2)

For singular matrices, we need to mask out the zero eigenvalues, so it looks a little more complicated. When the vector x is within the right subspace, we can perform the calculation as:

import numpy as np
from scipy import linalg

n = 10
rng = np.random.default_rng(1329348965454623456)
x = rng.random(size=(n, 1))
A = rng.random(size=(n, n))

A[0] = A[1] + A[2]  # make it singular
x[0] = x[1] + x[2]  # ensure x is in the subspace
A = A @ A.T

eps = 1e-10

# Eigendecomposition route
w, V = linalg.eigh(A)
mask = np.abs(w) < eps
w_ = w.copy()
w_[mask] = np.inf
z = x.T @ (V * w_**(-0.5))
xSxT = (z**2).sum()

# LDL route
L, D, perm = linalg.ldl(A)
d = np.diag(D).copy()[:, np.newaxis]
mask = np.abs(d) < eps
d_ = d.copy()
d[mask] = 0
d_[mask] = np.inf
y = linalg.solve_triangular(L[perm], x[perm], lower=True) * d_**(-0.5)
xSxT2 = (y**2).sum()

np.testing.assert_allclose(xSxT, xSxT2)

And we can tell it's in the subspace if np.allclose(L@(d**0.5*y), x). Otherwise, the PDF is zero.

[mdhaber]

This needs some work. The pseudodeterminant calculation isn't just the pseudodeterminant of the D matrix (product of positive diagonal elements).

class CovViaLDL(Covariance):

    def __init__(self, ldl):
        L, d, perm = ldl
        d = self._validate_vector(d, 'd')
        perm = self._validate_vector(perm, 'perm')
        L = self._validate_matrix(L, 'L')

        i_zero = d <= 0
        positive_d = np.array(d, dtype=np.float64)
        positive_d[i_zero] = 1  # ones don't affect determinant
        self._log_pdet = np.sum(np.log(positive_d), axis=-1)

        psuedo_reciprocals = 1 / np.sqrt(positive_d)
        psuedo_reciprocals[i_zero] = 0

        self._psuedo_reciprocals = psuedo_reciprocals
        self._d = d
        self._perm = perm
        self._L = L
        self._rank = d.shape[-1] - i_zero.sum(axis=-1)
        self._dimensionality = d.shape[-1]
        self._shape = L.shape
        # This is only used for `_support_mask`, not to decide whether
        # the covariance is singular or not.
        self._eps = 1e-8
        self._allow_singular = True

    def _whiten(self, x):
        L = self.L[self.perm]
        x = self.x[self.perm]
        return linalg.solve_triangular(L, x, lower=True) * self.psuedo_reciprocals

    @cached_property
    def _covariance(self):
        return (self._L * self._d) @ self._L.T

   @staticmethod
    def from_ldl(ldl):
        r"""
        Representation of covariance provided via LDL (aka LDLT) decomposition

        Parameters
        ----------
        ldl : sequence
            A sequence (nominally a tuple) containing the lower factor ``L``,
            the diagonal multipliers ``d``, and row-permutation indices
            ``perm`` as computed by `scipy.linalg.ldl`.

        Notes
        -----
        Let the covariance matrix be :math:`A`, and let :math:`L` be a lower
        triangular matrix and :math:`D` be a diagonal matrix such that
        :math:`L D L^T = A`.

        When all of the eigenvalues are strictly positive, whitening of a
        data point :math:`x` is performed by computing
        :math:`x^T (V W^{-1/2})`, where the inverse square root can be taken
        element-wise.
        :math:`\log\det{A}` is calculated as  :math:`tr(\log{W})`,
        where the :math:`\log` operation is performed element-wise.

        This `Covariance` class supports singular covariance matrices. When
        computing ``_log_pdet``, non-positive eigenvalues are ignored.
        Whitening is not well defined when the point to be whitened
        does not lie in the span of the columns of the covariance matrix. The
        convention taken here is to treat the inverse square root of
        non-positive eigenvalues as zeros.

        Examples
        --------
        Prepare a symmetric positive definite covariance matrix ``A`` and a
        data point ``x``.

        >>> import numpy as np
        >>> from scipy import stats
        >>> rng = np.random.default_rng()
        >>> n = 5
        >>> A = rng.random(size=(n, n))
        >>> A = A @ A.T  # make the covariance symmetric positive definite
        >>> x = rng.random(size=n)

        Perform the eigendecomposition of ``A`` and create the `Covariance`
        object.

        >>> w, v = np.linalg.eigh(A)
        >>> cov = stats.Covariance.from_eigendecomposition((w, v))

        Compare the functionality of the `Covariance` object against
        reference implementations.

        >>> res = cov.whiten(x)
        >>> ref = x @ (v @ np.diag(w**-0.5))
        >>> np.allclose(res, ref)
        True
        >>> res = cov.log_pdet
        >>> ref = np.linalg.slogdet(A)[-1]
        >>> np.allclose(res, ref)
        True

        """
        return CovViaLDL(ldl)