closed form optimal covariance

Author: AdrienCorenflos

coupled_rejection_sampling/mvn.py

@@ -1,78 +1,47 @@
 import math
-from functools import partial, lru_cache
+from functools import partial
 
-import cvxpy as cp
-import jax.experimental.host_callback as hcb
 import jax.numpy as jnp
 import jax.random
 import jax.scipy.linalg as jlinalg
 import jax.scipy.stats as jstats
 import numpy as np
-import scipy.linalg as linalg
 from jax import numpy as jnp
-from jax._src.scipy.linalg import cho_solve
+from jax.scipy.linalg import cho_solve
 
 from coupled_rejection_sampling.coupled_rejection_sampler import coupled_sampler
 
 _LOG_2PI = math.log(2 * math.pi)
 
 
-@lru_cache
-def _get_cvxpy_pb(d):
-    P_inv_param = cp.Parameter((d, d), PSD=True)
-    Sig_inv_param = cp.Parameter((d, d), PSD=True)
-    z_param = cp.Parameter((d,))
-
-    X = cp.Variable((d, d), PSD=True)
-
-    constraints = [X << P_inv_param, X << Sig_inv_param]
-    objective = cp.log_det(X)
-    problem = cp.Problem(cp.Maximize(objective), constraints)
-
-    return problem, X, P_inv_param, Sig_inv_param, z_param
-
-
-def get_optimal_covariance(m, chol_P, mu, chol_Sig, verbose=False):
+def get_optimal_covariance(chol_P, chol_Sig):
     """
     Get the optimal covariance according to the objective defined in Section 3 of [1].
 
+    The notations roughly follow the ones in the article.
+
     Parameters
     ----------
-    m: jnp.ndarray
-        Mean of X
     chol_P: jnp.ndarray
         Square root of the covariance of X. Lower triangular.
-    mu: jnp.ndarray
-        Mean of Y
     chol_Sig: jnp.ndarray
         Square root of the covariance of Y. Lower triangular.
-    verbose: bool, optional
-        Is the solver verbose. Default is False.
     Returns
     -------
     chol_Q: jnp.ndarray
         Cholesky of the resulting dominating matrix.
     """
-    d = m.shape[0]
+    d = chol_P.shape[0]
     if d == 1:
         return np.maximum(chol_P, chol_Sig)
-    problem, X, P_inv_param, Sig_inv_param, z_param = _get_cvxpy_pb(d)
 
-    eye = np.eye(d)
-    P_inv = linalg.cho_solve((chol_P, True), eye)
-    Sig_inv = linalg.cho_solve((chol_Sig, True), eye)
+    right_Y = jlinalg.solve_triangular(chol_P, chol_Sig, lower=True)  # Y = RY.T RY
+    w_Y, v_Y = jlinalg.eigh(right_Y.T @ right_Y)
+    w_Y = jnp.minimum(w_Y, 1)
+    i_w_Y = 1. / jnp.sqrt(w_Y)
 
-    z_param.value = m - mu
-    P_inv_param.value = P_inv
-    Sig_inv_param.value = Sig_inv
-
-    problem.solve(warm_start=True, verbose=verbose, qcp=True, solver=cp.MOSEK)
-
-    Q_inv = X.value
-    chol_Q_inv = linalg.cholesky(Q_inv, lower=True)
-    Q = linalg.cho_solve((chol_Q_inv, True), eye)
-    chol_Q = np.linalg.cholesky(Q)
-    return chol_Q
+    left_Q = chol_Sig @ (v_Y * i_w_Y[None, :])
+    return jlinalg.cholesky(left_Q @ left_Q.T, lower=True)
 
 
 def coupled_mvns(key, m, chol_P, mu, chol_Sig, N=1, chol_Q=None):
@@ -109,9 +78,7 @@ def coupled_mvns(key, m, chol_P, mu, chol_Sig, N=1, chol_Q=None):
     """
 
     if chol_Q is None:
-        # This is gonna be slow, but there is no real JAX version...
-        chol_Q = hcb.call(lambda args: get_optimal_covariance(*args), (m, chol_P, mu, chol_Sig), result_shape=chol_P)
-
+        chol_Q = get_optimal_covariance(chol_P, chol_Sig)
     log_det_chol_P = tril_log_det(chol_P)
     log_det_chol_Sig = tril_log_det(chol_Sig)
     log_det_chol_Q = tril_log_det(chol_Q)

examples/optimal_covariance.py

@@ -1,14 +1,13 @@
 """
 This corresponds to the Application 1 of [1]. 
 """
+import time
 from functools import partial
 
-import jax.experimental.host_callback as hcb
 import jax.numpy as jnp
 import jax.random
 import matplotlib.pyplot as plt
 import numpy as np
-import tikzplotlib
 import tqdm.auto as tqdm
 from jax.scipy import linalg
 
@@ -19,15 +18,15 @@
 save_path = "out/optimal.npz"
 
 jax_key = jax.random.PRNGKey(42)
-D = 10
+D = 15
 P = jnp.diag(jnp.arange(1, D + 1))
 
 chol_P = jnp.linalg.cholesky(P)
 m = mu = jnp.zeros((D,))
 
 B = 200  # number of covariances
 M = 500  # number of RS experiments
-NS = [1, 8, 64, 512]
+NS = [4 ** k for k in range(6)]
 
 
 def rejection_experiment():
@@ -39,15 +38,15 @@ def get_chol_covs(k):
         orth, _ = linalg.qr(gauss)
         Sigma = orth @ P @ orth.T
         chol_Sigma = jnp.linalg.cholesky(Sigma)
-        chol_Q = hcb.call(lambda args: get_optimal_covariance(*args), (m, chol_P, mu, chol_Sigma),
-                          result_shape=chol_P)
+        chol_Q = get_optimal_covariance(chol_P, chol_Sigma)
         return chol_Sigma, chol_Q
 
     res_shape = B, len(NS), 2
     res_mean_coupled = np.zeros(res_shape)
     res_var_coupled = np.zeros(res_shape)
     res_mean_trials = np.zeros(res_shape)
     res_var_trials = np.zeros(res_shape)
+    res_runtime = np.zeros(res_shape)
 
     res_coupling_bounds = np.zeros((B, 2))
     res_trials_bounds = np.zeros(res_shape)
@@ -79,24 +78,31 @@ def test_fun(op_key, chol_Sigma, chol_Q, N):
             res_trials_bounds[b, i, 0] = (1 + (n - 1) / trials_opt) / (n / trials_opt)
             res_trials_bounds[b, i, 1] = (1 + (n - 1) / trials_max) / (n / trials_max)
 
+            tic = time.time()
             (res_mean_coupled[b, i, 0], res_var_coupled[b, i, 0], res_mean_trials[b, i, 0],
              res_var_trials[b, i, 0]) = test_fun(rs_key, chol_Sigma, chol_Q, n)
+            res_runtime[b, i, 0] = time.time() - tic
 
+            tic = time.time()
             (res_mean_coupled[b, i, 1], res_var_coupled[b, i, 1], res_mean_trials[b, i, 1],
              res_var_trials[b, i, 1]) = test_fun(rs_key, chol_Sigma, chol_Q_max, n)
-
-    return res_mean_coupled, res_var_coupled, res_mean_trials, res_var_trials, res_coupling_bounds, res_trials_bounds
+            res_runtime[b, i, 1] = time.time() - tic
+    return res_mean_coupled, res_var_coupled, res_mean_trials, res_var_trials, res_coupling_bounds, res_trials_bounds, res_runtime
 
 
 if RUN:
-    (coupled_mean, coupled_var, n_trials_mean, n_trials_var, coupling_bounds, n_trials_bounds) = rejection_experiment()
+    (coupled_mean, coupled_var, n_trials_mean, n_trials_var, coupling_bounds, n_trials_bounds,
+     runtime) = rejection_experiment()
     np.savez(save_path, coupled_mean=coupled_mean, coupled_var=coupled_var,
              n_trials_mean=n_trials_mean, n_trials_var=n_trials_var, coupling_bounds=coupling_bounds,
-             n_trials_bounds=n_trials_bounds)
+             n_trials_bounds=n_trials_bounds, runtime=runtime)
 
 if PLOT:
     data = np.load(save_path)
-    fig, axes = plt.subplots(nrows=2, ncols=2, figsize=(12, 6), sharex=True, sharey=True)
+    fig, axes = plt.subplots(nrows=2, ncols=3, figsize=(15, 7), sharex=True, sharey=True)
+
+    print(np.mean(data["runtime"][1:], 0))
+    print(np.std(data["runtime"][1:], 0))
 
     for i, n in enumerate(NS):
         ax = axes.flatten()[i]
@@ -108,16 +114,16 @@ def test_fun(op_key, chol_Sigma, chol_Q, N):
         ax.scatter(range(B), data["coupled_mean"][arg_sort, i, 1], label=f"Empirical MAX",
                    color="tab:orange",
                    alpha=0.75)
-
-        ax.plot(range(B), data["coupling_bounds"][arg_sort, 0], label="Optimised bound",
-                color="tab:blue")
-        ax.plot(range(B), data["coupling_bounds"][arg_sort, 1], label="MAX bound",
-                color="tab:orange")
+        twinx = ax.twinx()
+        twinx.semilogy(range(B), data["coupling_bounds"][arg_sort, 0], label="Optimised bound",
+                   color="tab:blue")
+        twinx.semilogy(range(B), data["coupling_bounds"][arg_sort, 1], label="MAX bound",
+                   color="tab:orange")
     axes[0, 0].legend(loc="upper left")
-    # plt.show()
-    tikzplotlib.save("out/gaussian_opt_coupling.tikz")
+    plt.show()
+    # tikzplotlib.save("out/gaussian_opt_coupling.tikz")
 
-    fig, axes = plt.subplots(nrows=2, ncols=2, figsize=(12, 6), sharex=True, sharey=True)
+    fig, axes = plt.subplots(nrows=2, ncols=3, figsize=(15, 7), sharex=True, sharey=True)
     for i, n in enumerate(NS):
         ax = axes.flatten()[i]
         ax.set_title(f"$N={n}$")
@@ -131,6 +137,6 @@ def test_fun(op_key, chol_Sigma, chol_Q, N):
                    alpha=0.75)
         ax.plot(range(B), 1 / data["n_trials_bounds"][arg_sort, i, 1], label="MAX bound", color="tab:orange")
     axes[0, 0].legend(loc="upper left")
-    # plt.show()
+    plt.show()
 
-    tikzplotlib.save("out/gaussian_opt_acceptance.tikz")
+    # tikzplotlib.save("out/gaussian_opt_acceptance.tikz")

tests/test_mvn.py

@@ -1,3 +1,4 @@
+import cvxpy as cp
 import jax.random
 import numpy as np
 import numpy.testing as np_test
@@ -17,6 +18,25 @@ def is_inverse_less(chol_A, chol_B):
     assert np.max(np.real(linalg.eigvals(A_inv - B_inv))) > -1e-7
 
 
+def cvxpy_get_optimal_cov(chol_P, chol_Sig):
+    d = chol_P.shape[0]
+    eye = np.eye(d)
+    P_inv = linalg.cho_solve((chol_P, True), eye)
+    Sig_inv = linalg.cho_solve((chol_Sig, True), eye)
+
+    X = cp.Variable((d, d), PSD=True)
+
+    constraints = [X << P_inv, X << Sig_inv]
+    objective = cp.log_det(X)
+    problem = cp.Problem(cp.Maximize(objective), constraints)
+    problem.solve(warm_start=True, qcp=True, solver=cp.SCS)
+    Q_inv = X.value
+    chol_Q_inv = linalg.cholesky(Q_inv, lower=True)
+    Q = linalg.cho_solve((chol_Q_inv, True), eye)
+    chol_Q = np.linalg.cholesky(Q)
+    return chol_Q
+
+
 def test_ordering():
     chol_P = np.array([
         [1.0, 0.0, 0.0],
@@ -25,15 +45,14 @@ def test_ordering():
     ])
 
     chol_Sig = np.array([
-        [1.3, 0.0, 0.0],
+        [0.2, 0.0, 0.0],
         [0.5, 0.9, 0.0],
-        [0.1, 1.7, 0.2]
+        [0.1, 1.7, 1.5]
     ])
 
-    a = np.array([0., 1., -1.])
-    b = np.array([1., -1., 0.])
-
-    chol_Q = get_optimal_covariance(a, chol_P, b, chol_Sig, verbose=True)
+    chol_Q = get_optimal_covariance(chol_P, chol_Sig)
+    cvxpy_chol_Q = cvxpy_get_optimal_cov(chol_P, chol_Sig)
+    np.testing.assert_allclose(chol_Q, cvxpy_chol_Q, rtol=1e-4, atol=1e-5)
 
     is_inverse_less(chol_P, chol_Q)
     is_inverse_less(chol_Sig, chol_Q)
@@ -108,22 +127,7 @@ def test_mvns_different_cov(d, M, mocker: MockerFixture):
     np_test.assert_allclose(np.cov(ys, rowvar=False), chol_Sigma @ chol_Sigma.T, atol=1e-2, rtol=1e-2)
 
 
-
 def test_lower_bound():
-    # A trivial test
-    m = np.array([1.])
-    mu = np.array([1.])
-    Sig = np.array([[1.]])
-    P = np.array([[1.]])
-    Q = np.array([[1.]])
-
-    chol_P = linalg.cholesky(P, lower=True)
-    chol_Sig = linalg.cholesky(Sig, lower=True)
-    chol_Q = linalg.cholesky(Q, lower=True)
-    K = coupled_mvns.lower_bound(m, chol_P, mu, chol_Sig, chol_Q)
-
-    assert K == pytest.approx(1., rel=1e-5, abs=1e-5)
-
     # Test against reflection coupling
     m = np.array([1., 2.])
     mu = np.array([1.5, 1.1])