1import numpy as np
 2
 3
 4class QuadraticConvexTransport:
 5    """T(x)=grad(.5*x^T A*x+b^T x), with A symmetric positive definite."""
 6    def __init__(self, A, b=None):
 7        self.A = np.asarray(A, dtype=float)
 8        self.b = np.zeros(self.A.shape[0]) if b is None else np.asarray(b, dtype=float)
 9        if not np.allclose(self.A, self.A.T):
10            raise ValueError("A must be symmetric")
11        if np.min(np.linalg.eigvalsh(self.A)) <= 0:
12            raise ValueError("potential is not strongly convex")
13
14    def map(self, x):
15        return np.asarray(x) @ self.A.T + self.b
16
17    def potential(self, x):
18        x = np.asarray(x)
19        return 0.5 * np.einsum("...i,ij,...j->...", x, self.A, x) + x @ self.b
20
21    def logdet_jacobian(self, x):
22        return np.full(len(np.atleast_2d(x)), np.linalg.slogdet(self.A)[1])
23
24    def gaussian_kl_to_standard(self, mean, covariance):
25        """Exact KL(T#N(mean,cov)||N(0,I))."""
26        mean, covariance = np.asarray(mean), np.asarray(covariance)
27        m = self.A @ mean + self.b
28        c = self.A @ covariance @ self.A.T
29        d = len(mean)
30        return 0.5 * (np.trace(c) + m @ m - d - np.linalg.slogdet(c)[1])
31
32
33def change_of_variables_kl(samples, transport, log_source_density):
34    x = np.asarray(samples)
35    tx = transport.map(x)
36    return float(np.mean(log_source_density(x) - transport.logdet_jacobian(x)
37                         - log_source_density(tx)))
38
39
40def standard_normal_logpdf(x):
41    x = np.asarray(x)
42    return -0.5 * (np.sum(x * x, axis=-1) + x.shape[-1] * np.log(2 * np.pi))