Missing feature: Multivariate normal distribution with a different covariance matrix for each particle

[nchopin]

Ah, sorry, my bad. One way to fix my code :

def multi_dist(xp, beta, gamma, dt):  # xp means X_{t-1}, defines the BM parts
    locS = xp['S'] - beta*xp['S']*xp['I']*dt
    locI =  xp['I'] + (beta*xp['S']*xp['I'] - gamma*xp['I'])*dt
    d = {'S': dists.Normal(loc=locS,
                             scale=(1/10)*np.sqrt(beta*xp['S']*xp['I']*dt)),
           'I': dists.Cond(lambda x: dists.Normal(loc=locI - x['S'] + locS, 
                                        scale = (1/10)*np.sqrt(beta*xp['S']*xp['I']*dt))
          }

where basically I replaced $B_1$ by $S$ minus its expectation.

The current version of MvNormal does not allow for a covariance matrix that varies across the particles. I implemented this for a colleague, but I don't like it so much because there is a loop over $n$, so iit may be slow:

class Generalized_MV_Normal(dists.ProbDist):
    """Multivariate normal, with dim >=2, allowing for a different cov matrix for each
    particle. 
    """
    def __init__(self, loc=None, cov=None):
        self.loc = loc
        self.cov = cov
        self.dim = self.cov.shape[-1]

    def rvs(self, size=1):
        x = np.empty((size, self.dim))
        for n in range(size):
            x[n, :] = random.multivariate_normal(self.loc[n, :], 
                                                 self.cov[n, :, :])
        return x

    def logpdf(self, x):
        N = x.shape[0]
        l = np.empty(N)
        for n in range(N):
            l[n] = stats.multivariate_normal.logpdf(x[n], mean=self.loc[n, :],
                                                    cov=self.cov[n, :, :])
        return l

Since this is the second time someone is asking for something like this, I am going to open an issue and try to think of ways to make the above code more efficient (numba?).

Originally posted by @nchopin in https://github.com/nchopin/particles/issues/54#issuecomment-1152019498

[vwiela]

Thanks for opening this issue. I got one question left about the above Generalized_MV_Normal class as implemented above, especially how to use it for an example as in issue #54.

1. Could you comment how the inputs should look like (arrays, matrices, structured arrays?) Or provide an example use?

Maybe also one additional idea. Often when I used multivariate normal distributions, where the covariance matrix depends on the states of other particles, I encountered problems with the positive definiteness of the matrix. Because due to stochasticity states can encounter for example negative values or values that lead to a non-positive definite covariance matrix in the next step. So when thinking about the implementation of a MvNormal distribution, it might be worth to directly think about implementing truncation options as well.

[AdrienCorenflos]

This should do: the broadcasting should be smart enough to handle everything.

import particles.distributions as dists
import numpy as np
import scipy.stats as stats



def log_det_chol(chol):
    diags = np.diagonal(chol, axis1=-2, axis2=-1)
    return np.sum(np.log(np.abs(diags)), -1)

def mvn_logpdf(x, loc, chol):
    d = loc.shape[-1]
    b = np.broadcast(x, loc)
    diff = np.empty(b.shape)
    diff.flat = [u - v for (u,v) in b]
    z = np.linalg.solve(chol, diff)  # solve_triangular doesn't accept batched matrices
    const = -0.5 * d * np.log(2 * np.pi) - log_det_chol(chol)
    return const -0.5 * np.sum(z * z, -1)


def mvn_sample(loc, chol, size=1):
    # Do some checks on size here
    d = loc.shape[-1]
    broadcast_to = np.broadcast_shapes(loc.shape, (size, d))
    loc = np.broadcast_to(loc, broadcast_to)
    eps = np.random.randn(*broadcast_to)
    return loc + np.einsum("...ij,...j", chol, eps)

class Generalized_MV_Normal(dists.ProbDist):
    """Multivariate normal, with dim >=2, allowing for a different cov matrix for each
    particle. 
    """
    def __init__(self, loc=None, cov=None):
        self.loc = loc
        self.chol = np.linalg.cholesky(cov)
        self.dim = cov.shape[-1]

    def rvs(self, size=1):
        return mvn_sample(self.loc, self.chol, size)

    def logpdf(self, x):
        return mvn_logpdf(x, self.loc, self.chol)