Gaussian tails experiment (#4)

* c(x) and right solver implmemented

* left exp exp solver

* fixed equation solvers and implemented various other parts

* tp and tq sampling

* Gamma_hat and other yet untested parts added

* p and q sampling

* Converted while loops to while_loops

* First version that appears to work

* Unit tests and fixes

* Experiment draft

* coupled truncated exponentials in functional paradigm + logspace

* Made gauss tails residual using quick solver.

* I think this works just fine. Let's just design a bunch of unittest, document and merge.

* Make experiment use the "new" code.

* Visibly some issue with the second marginal of the exponential coupling.

* Finished refactoring and changed experiment

* Finalise Gaussian tails experiment.

* Some fixes on sampling tails + better test using p-value.

* truncnorm was failing.

* Fixed the Gaussian tails.

* Started the reflected max example

* Reflected mh (#3)

* need to commit before holidays

* This should be the reflected mh implemented.

* Finalised MMALA example

* add some replication

* Update README.md

Co-authored-by: Simo Särkkä <simo.sarkka@aalto.fi>

Author: AdrienCorenflos

README.md

@@ -14,7 +14,7 @@ Otherwise, simply download the repository and run `pip .`.
 
 ### Example
 
-See the `examples` and `tests` folders. 
+See the `examples` and `tests` folders. The examples implement the paper illustrations, while the tests test single functions behaviours and are a bit more low level.
 
 
 ### References

coupled_rejection_sampling/gauss_tails.py

@@ -0,0 +1,234 @@
+# Coupled sampling from 2 different tails x > mu and x > eta of N(x | 0,1) Gaussian
+# by using a maximal coupling of translated exponential proposals hatp(x) = alpha exp(-alpha (x - mu))
+# and hatq(x) = beta exp(-beta (x - eta)) similarly to Robert (1995), but with the coupling.
+#
+# Robert, C. P. (1995). Simulation of truncated normal variables. Statistics and Computing, Volume 5, pages 121–125.
+import chex as chex
+import jax.lax
+import jax.numpy as jnp
+import jax.random
+import tensorflow_probability.substrates.jax as tfp
+from jax.experimental.host_callback import id_print
+
+from coupled_rejection_sampling.utils import logsubexp, log1mexp
+
+
+@jax.jit
+def coupled_gaussian_tails(key: chex.PRNGKey,
+                           mu: chex.Numeric, eta: chex.Numeric):
+    """
+    A coupled version of Robert's truncated normal sampling algorithm.
+    We want to sample from 2 different tails x > mu and x > eta of a N(x | 0,1) Gaussian
+
+    Parameters
+    ----------
+    key: jnp.ndarray
+        JAX random key
+    mu, eta: float
+        The tails of the Gaussians we want to sample from
+
+    Returns
+    -------
+    X: chex.Numeric
+        The sample from x > mu
+    Y: chex.Numeric
+        The sample from x > eta
+    is_coupled: jnp.ndarray
+        Coupling flag
+    """
+    alpha_mu = get_alpha(mu)
+    alpha_eta = get_alpha(eta)
+
+    p = lambda k: _robert_sampler(k, mu, alpha_mu)
+    q = lambda k: _robert_sampler(k, eta, alpha_eta)
+
+    log_w_p = lambda x: -0.5 * (x - alpha_mu) ** 2
+    log_w_q = lambda x: -0.5 * (x - alpha_eta) ** 2
+
+    def cond(carry):
+        accept_X, accept_Y, *_ = carry
+        return ~accept_X & ~accept_Y
+
+    def body(carry):
+        *_, i, curr_key = carry
+        next_key, sample_key, accept_key = jax.random.split(curr_key, 3)
+        X_hat, Y_hat, are_coupled = coupled_exponentials(sample_key, mu, alpha_mu, eta, alpha_eta)
+        log_w_X = log_w_p(X_hat)
+        log_w_Y = log_w_q(Y_hat)
+
+        log_u = jnp.log(jax.random.uniform(accept_key))
+        accept_X = log_u < log_w_X
+        accept_Y = log_u < log_w_Y
+
+        return accept_X, accept_Y, X_hat, Y_hat, are_coupled, i + 1, next_key
+
+    # initialisation
+    residual_key, loop_key = jax.random.split(key)
+
+    output = jax.lax.while_loop(cond,
+                                lambda carry: body(carry),
+                                (False, False, 0., 0., False, 0, loop_key))
+
+    is_X_accepted, is_Y_accepted, X, Y, is_coupled, n_trials, _ = output
+
+    X = jax.lax.cond(is_X_accepted, lambda _: X, p, residual_key)
+    Y = jax.lax.cond(is_Y_accepted, lambda _: Y, q, residual_key)
+
+    is_coupled = is_coupled & is_X_accepted & is_Y_accepted
+
+    return X, Y, is_coupled
+
+
+@jax.jit
+def coupled_exponentials(key:chex.PRNGKey, mu:chex.Numeric, alpha_mu:chex.Numeric, eta:chex.Numeric, alpha_eta:chex.Numeric):
+    """
+    Sampling from a maximal coupling of shifted exponentials.
+    p(x) = exp(-alpha (x - m)) / m  for x >= m, 0 otherwise
+
+    It assumes that eta > mu and alpha_eta > alpha_mu
+
+    Parameters
+    ----------
+    key: chex.PRNGKey
+        JAX random key
+    mu, eta: chex.Numeric
+        The shift of the exponentials
+    alpha_mu, alpha_eta: chex.Numeric
+        The rate parameters
+
+    Returns
+    -------
+    X: chex.Numeric
+        The sample from x > mu
+    Y: chex.Numeric
+        The sample from x > eta
+    is_coupled: chex.Numeric
+        Coupling flag
+    """
+
+    gamma = get_gamma(mu, eta, alpha_mu, alpha_eta)
+
+    eta_mu = -alpha_mu * (eta - mu)
+    gamma_mu = -alpha_mu * (gamma - mu)
+    gamma_eta = -alpha_eta * (gamma - eta)
+
+    log_max_coupling_proba = logsubexp(eta_mu, jnp.logaddexp(gamma_mu, gamma_eta))
+
+    subkey1, subkey2 = jax.random.split(key)
+
+    log_u = jnp.log(jax.random.uniform(subkey1, shape=()))
+    are_coupled = (log_u <= log_max_coupling_proba)
+
+    def if_coupled(k):
+        x = _sampled_from_coupled_exponentials(k, mu, eta, alpha_mu, alpha_eta, eta_mu, gamma_mu, gamma_eta, gamma)
+        return x, x
+
+    def otherwise(k):
+        x = _sample_from_first_marginal(k, mu, eta, alpha_mu, alpha_eta, eta_mu, gamma_mu, gamma_eta, gamma)
+        y = _sample_from_second_marginal(k, mu, eta, alpha_mu, alpha_eta, eta_mu, gamma_mu, gamma_eta, gamma)
+        return x, y
+
+    x_out, y_out = jax.lax.cond(are_coupled, if_coupled, otherwise, subkey2)
+    return x_out, y_out, are_coupled
+
+
+def _sampled_from_coupled_exponentials(key, mu, _eta, alpha_mu, alpha_eta, eta_mu, gamma_mu, gamma_eta, gamma):
+    def C1_inv(log_u):
+        return mu - logsubexp(eta_mu, log_u + logsubexp(eta_mu, gamma_mu)) / alpha_mu
+
+    def C2_inv(log_u):
+        return gamma - log_u / alpha_eta
+
+    log_p1 = logsubexp(eta_mu, gamma_mu)
+    log_p2 = gamma_eta
+    log_p = log_p1 - jnp.logaddexp(log_p1, log_p2)
+
+    log_u1, log_u2 = jnp.log(jax.random.uniform(key, shape=(2,)))
+
+    res = jax.lax.cond(log_u1 < log_p, C1_inv, C2_inv, log_u2)
+    return res
+
+
+def _sample_from_first_marginal(key, mu, eta, alpha_mu, alpha_eta, eta_mu, gamma_mu, gamma_eta, gamma):
+    key1, key2 = jax.random.split(key, 2)
+    log_u1 = jnp.log(jax.random.uniform(key1))
+
+    log_p1 = logsubexp(gamma_mu, gamma_eta)  # This has the same value as $\log(\tilde{Z})$
+    log_p2 = log1mexp(eta_mu)
+    log_p = log_p1 - jnp.logaddexp(log_p1, log_p2)
+
+    def _sample_from_tail(log_u):
+        return mu - log1mexp(log_u + log1mexp(eta_mu)) / alpha_mu
+
+    def _sample_from_overlap(log_u):
+
+        def log_f(x):
+            return logsubexp(-alpha_mu * (x - mu), -alpha_eta * (x - eta)) - log_p1 - log_u
+
+        # upper bound for the solution is given by a lower bounding of the density
+        def upper_bound_loop(carry):
+
+            curr_upper_bound, _ = carry
+            curr_upper_bound = 1.5 * curr_upper_bound
+            obj = log_f(curr_upper_bound)
+            return curr_upper_bound, obj >= 0
+
+        upper_bound, _ = jax.lax.while_loop(lambda carry: carry[-1], upper_bound_loop, (gamma, True))
+
+        res, objective_at_estimated_root, *_ = tfp.math.find_root_chandrupatla(log_f, gamma, upper_bound,
+                                                                               position_tolerance=1e-6,
+                                                                               value_tolerance=1e-6)
+
+        return res
+
+    return jax.lax.cond(log_u1 < log_p, _sample_from_overlap, _sample_from_tail, jnp.log(jax.random.uniform(key2)))
+
+
+def _sample_from_second_marginal(key, mu, eta, alpha_mu, alpha_eta, eta_mu, gamma_mu, gamma_eta, gamma):
+    log_Zq_1 = jnp.logaddexp(0, gamma_mu)  # log(1 + exp(gamma_mu))
+    log_Zq_2 = jnp.logaddexp(eta_mu, gamma_eta)
+    log_Zq = logsubexp(log_Zq_1, log_Zq_2)
+    log_u = jnp.log(jax.random.uniform(key))
+
+    def log_f(x):
+        res_1 = jnp.logaddexp(0, -alpha_mu * (x - mu))  # log(1 + exp(...))
+        res_2 = jnp.logaddexp(eta_mu, -alpha_eta * (x - eta))
+        res = logsubexp(res_1, res_2)
+        res = res - log_Zq - log_u
+        return res
+
+    out, objective_at_estimated_root, *_ = tfp.math.find_root_chandrupatla(log_f, eta, gamma, position_tolerance=1e-6,
+                                                                           value_tolerance=1e-6)
+
+    return out
+
+
+def _robert_sampler(key, mu, alpha):
+    def body(carry):
+        curr_k, *_ = carry
+        curr_k, subkey = jax.random.split(curr_k, 2)
+
+        u1, u2 = jax.random.uniform(subkey, shape=(2,))
+
+        x = mu - jnp.log(1 - u1) / alpha
+        accepted = u2 <= jnp.exp(-0.5 * (x - alpha) ** 2)
+
+        return curr_k, x, accepted
+
+    _, x_out, _ = jax.lax.while_loop(lambda carry: ~carry[-1], body, (key, 0., False))
+    return x_out
+
+
+def get_alpha(mu):
+    """ Compute the optimal alpha as per Robert (1995) """
+    return 0.5 * (mu + jnp.sqrt(mu ** 2 + 4))
+
+
+def get_gamma(mu, eta, alpha, beta):
+    """ Threshold when hatp(x) = hatq(x) """
+    return (jnp.log(beta) - jnp.log(alpha) + beta * eta - alpha * mu) / (beta - alpha)
+
+
+def texp_logpdf(x, mu, alpha):
+    """ Translated exponential density """
+    return jnp.where(x < mu, -jnp.inf, jnp.log(alpha) - alpha * (x - mu))

examples/gaussian_tail_coupling.py

@@ -0,0 +1,103 @@
+"""
+Example of coupling of tails of N(0,1) Gaussian distributions.
+"""
+import time
+
+import jax.lax
+import jax.random
+import matplotlib.pyplot as plt
+import numpy as np
+import scipy.integrate
+import tikzplotlib
+import tqdm
+from scipy.stats import truncnorm
+import seaborn as sns
+from coupled_rejection_sampling.gauss_tails import coupled_gaussian_tails
+
+RUN = False
+PLOT = True
+
+save_path = "out/gauss_tails.npz"
+
+
+if RUN:
+    JAX_KEY = jax.random.PRNGKey(0)
+
+    # Tails >5 is considered "hard" and the closer these are, the worse Thorisson's method is
+    mu = 6.
+
+    M = 100_000  # Samples to draw
+    DELTAS = np.linspace(1e-6, 0.5, num=200)
+
+    # Compute maximal coupling probability (note that this can be inaccurate)
+    p = lambda x: truncnorm.pdf(x, mu, np.inf)
+
+    def experiment():
+        pxy_list = np.empty((len(DELTAS), 2))
+        x_samples = np.empty((len(DELTAS), M))
+        y_samples = np.empty((len(DELTAS), M))
+        runtimes = np.empty((len(DELTAS),))
+
+        keys = jax.random.split(JAX_KEY, M)
+        sampler = jax.jit(jax.vmap(coupled_gaussian_tails, in_axes=[0, None, None]))
+        # compilation run
+        *_, acc = sampler(keys, mu, mu + DELTAS[0])
+        acc.block_until_ready()
+
+        for n in tqdm.trange(len(DELTAS)):
+            delta = DELTAS[n]
+            eta = mu + delta
+
+            q = lambda x: truncnorm.pdf(x, eta, np.inf)
+            mpq = lambda x: np.minimum(p(x), q(x))
+            true_pxy = scipy.integrate.quad(mpq, 0, np.inf)[0]
+
+            tic = time.time()
+            x_samples[n], y_samples[n], acc = sampler(keys, mu, eta)
+            acc.block_until_ready()
+            toc = time.time()
+            runtimes[n] = toc - tic
+            pxy = np.mean(acc)
+            pxy_list[n] = pxy, true_pxy
+
+        return x_samples, y_samples, pxy_list, runtimes
+
+
+    x_samples, y_samples, pxy_list, runtimes = experiment()
+    np.savez(save_path, x_samples=x_samples, y_samples=y_samples, pxy_list=pxy_list, M=M, etas=mu + DELTAS, mu=mu, runtimes=runtimes)
+
+if PLOT:
+    data = np.load(save_path)
+
+    x_samples = data["x_samples"]
+    y_samples = data["y_samples"]
+    pxy_list = data["pxy_list"]
+    M = data["M"]
+
+    mu = data["mu"]
+    etas = data["etas"]
+    runtimes = data["runtimes"]
+    xs = np.linspace(mu, 10, 1000)
+
+    fig = plt.figure(figsize=(10, 10))
+    g = sns.jointplot(x=x_samples[10], y=y_samples[10], s=10)
+
+    plt.show()
+
+    fig, ax = plt.subplots()
+    plt.title("Coupling probability as function of $\eta$")
+    ax.semilogy(etas, pxy_list[:, 0], label='Actual coupling probability', linestyle="-", color="k")
+    ax.semilogy(etas, pxy_list[:, 1], label='True coupling probability', linestyle="--", color="k")
+    plt.xlabel("$\eta$")
+    ax.legend()
+    # plt.show()
+    tikzplotlib.save("out/gaussian_tails_coupling.tikz")
+
+
+    fig, ax = plt.subplots()
+    plt.title("Run time as function of $\eta$")
+    ax.plot(etas, runtimes, linestyle="-", color="k")
+    plt.xlabel("$\eta$")
+    plt.ylabel("Run time (s)")
+    ax.legend()
+    tikzplotlib.save("out/gaussian_tails_runtime.tikz")

requirements.txt

@@ -1,3 +1,4 @@
+chex
 cvxpy
 jax[cpu]
 jupyter
@@ -9,5 +10,6 @@ pytest
 pytest-mock
 scipy
 seaborn
+tensorflow_probability
 tikzplotlib
 tqdm