From 7a8ddddff4c84c04a63bce91278f45b2d82137e1 Mon Sep 17 00:00:00 2001
From: Will Handley <wh260@cam.ac.uk>
Date: Fri, 5 Aug 2022 20:10:20 +0100
Subject: [PATCH 01/19] Made utilities for nested sampling generation

---
 anesthetic/examples/__init__.py   |  1 +
 anesthetic/examples/perfect_ns.py |  1 +
 anesthetic/examples/utils.py      | 54 +++++++++++++++++++++++++++++++
 tests/test_examples.py            | 42 ++++++++++++++++++++++++
 4 files changed, 98 insertions(+)
 create mode 100644 anesthetic/examples/__init__.py
 create mode 100644 anesthetic/examples/perfect_ns.py
 create mode 100644 anesthetic/examples/utils.py
 create mode 100644 tests/test_examples.py

diff --git a/anesthetic/examples/__init__.py b/anesthetic/examples/__init__.py
new file mode 100644
index 00000000..cf11d49a
--- /dev/null
+++ b/anesthetic/examples/__init__.py
@@ -0,0 +1 @@
+"""Module for generating nested sampling examples."""
diff --git a/anesthetic/examples/perfect_ns.py b/anesthetic/examples/perfect_ns.py
new file mode 100644
index 00000000..9058cff5
--- /dev/null
+++ b/anesthetic/examples/perfect_ns.py
@@ -0,0 +1 @@
+"""Perfect nested sampling generators."""
diff --git a/anesthetic/examples/utils.py b/anesthetic/examples/utils.py
new file mode 100644
index 00000000..721e4958
--- /dev/null
+++ b/anesthetic/examples/utils.py
@@ -0,0 +1,54 @@
+"""Utility functions for nested sampling examples."""
+import numpy as np
+from scipy.stats import special_ortho_group
+from scipy.special import gamma
+
+
+class random_ellipsoid(object):
+    """Draw a point uniformly in an ellipsoid.
+
+    This is defined so that the volume of the ellipsoid is sqrt(det(cov))*V_n.
+    where V_n is the volume of the unit n ball, and so that its axes have
+    length equal to the square root of the eigenvalues of the covariance
+    matrix.
+
+    Parameters
+    ----------
+    mean: 1d array-like
+        The center of mass of the ellipsoid
+
+    cov: 2d array-like
+        The covariance structure of the ellipsoid. Axes have lengths equal to
+        the square root of the eigenvalues of this matrix.
+    """
+
+    def __init__(self, mean, cov):
+        self.mean = mean
+        self.cov = cov
+
+    def __call__(self, size=None):
+        """Generate samples uniformly from the ellipsoid."""
+        d = len(self.mean)
+        L = np.linalg.cholesky(self.cov)
+        x = np.random.multivariate_normal(np.zeros(d), np.eye(d), size=size)
+        r = np.linalg.norm(x, axis=-1, keepdims=True)
+        u = np.random.power(d, r.shape)
+        return self.mean + (u*x/r) @ L.T
+
+
+def random_covariance(sigmas):
+    """Draw a randomly oriented covariance matrix with axes length sigmas.
+
+    Parameters
+    ----------
+    sigmas, 1d array like
+        Lengths of the axes of the ellipsoid.
+    """
+    d = len(sigmas)
+    R = special_ortho_group.rvs(d)
+    return R @ np.diag(sigmas) @ np.diag(sigmas) @ R.T
+
+
+def volume_n_ball(n, r=1):
+    """Volume of an n dimensional ball, radius r."""
+    return np.pi**(n/2)/gamma(n/2+1)*r**n
diff --git a/tests/test_examples.py b/tests/test_examples.py
new file mode 100644
index 00000000..4f4c27e6
--- /dev/null
+++ b/tests/test_examples.py
@@ -0,0 +1,42 @@
+import numpy as np
+from numpy.testing import assert_allclose
+from anesthetic.examples.utils import (
+    random_ellipsoid, random_covariance, volume_n_ball
+)
+from scipy.spatial import ConvexHull
+
+
+def test_random_covariance():
+    np.random.seed(0)
+    d = 5
+    sigmas = np.random.rand(d)
+    cov = random_covariance(sigmas)
+    assert(cov.shape == (d, d))
+    e = np.linalg.eigvals(cov)
+    assert_allclose(np.sort(e), np.sort(sigmas**2))
+
+
+def test_random_ellipsoid():
+    np.random.seed(0)
+
+    d = 5
+    mean = np.random.rand(d)
+    cov = random_covariance(np.random.rand(d))
+    rvs = random_ellipsoid(mean, cov)
+
+    assert(rvs(10).shape == (10, d))
+    assert(rvs((10,)).shape == (10, d))
+    assert(rvs((10, 20)).shape == (10, 20, d))
+
+    dat = rvs(100000)
+
+    atol = 1/np.sqrt(len(dat))
+    assert_allclose(np.mean(dat, axis=0), mean, atol=atol)
+    assert_allclose(np.cov(dat.T), cov/(d+2), atol=atol)
+
+    cov = random_covariance(np.random.rand(2))
+    rvs = random_ellipsoid([0, 0], cov)
+    dat = rvs(1000000)
+    hull = ConvexHull(dat)
+    vol = hull.volume
+    assert_allclose(vol, np.linalg.det(cov)**0.5*volume_n_ball(2), rtol=1e-3)

From 41b5f35b641804a7c3ccb3452d624b72fcf792dc Mon Sep 17 00:00:00 2001
From: Will Handley <wh260@cam.ac.uk>
Date: Sat, 6 Aug 2022 22:57:16 +0100
Subject: [PATCH 02/19] work in progress

---
 anesthetic/examples/perfect_ns.py | 112 ++++++++++++++++++++++++++++++
 anesthetic/examples/utils.py      |  30 ++++----
 tests/test_examples.py            |  45 +++++++++++-
 3 files changed, 168 insertions(+), 19 deletions(-)

diff --git a/anesthetic/examples/perfect_ns.py b/anesthetic/examples/perfect_ns.py
index 9058cff5..51b96294 100644
--- a/anesthetic/examples/perfect_ns.py
+++ b/anesthetic/examples/perfect_ns.py
@@ -1 +1,113 @@
 """Perfect nested sampling generators."""
+import numpy as np
+from scipy.stats import multivariate_normal
+from anesthetic.examples.utils import random_ellipsoid
+from anesthetic import NestedSamples
+from anesthetic.samples import merge_nested_samples
+
+
+def gaussian(nlive, ndims, sigma=0.1, R=1, logLmin=-1e-2):
+    """Perfect nested sampling run for a spherical Gaussian & prior.
+
+    Up to normalisation this is identical to the example in John Skilling's
+    original paper https://doi.org/10.1214/06-BA127 . Both spherical Gaussian
+    width sigma and spherical uniform prior width R are centered on zero
+    
+    Parameters
+    ----------
+    nlive: int
+        number of live points
+
+    ndims: int
+        dimensionality of gaussian
+
+    sigma: float
+        width of gaussian likelihood
+
+    R: float
+        radius of gaussian prior
+
+    logLmin: float
+        loglikelihood at which to terminate
+    """
+
+    def logLike(x):
+        return  -(x**2).sum(axis=-1)/2/sigma**2
+
+    def random_sphere(n):
+        return random_ellipsoid(np.zeros(ndims), np.eye(ndims), n)
+
+    samples = []
+    r = R
+    logL_birth = np.ones(nlive) * -np.inf
+    logL = logL_birth.copy()
+    while logL.min() < logLmin:
+        points = r * random_sphere(nlive)
+        logL = logLike(points)
+        samples.append(NestedSamples(points, logL=logL, logL_birth=logL_birth))
+        logL_birth = logL.copy()
+        r = (points**2).sum(axis=-1, keepdims=True)**0.5
+
+    samples = merge_nested_samples(samples) 
+    samples.logL
+    logLend = samples[samples.nlive >= nlive].logL.max()
+    return samples[samples.logL_birth < logLend].recompute()
+
+
+def correlated_gaussian(nlive, mean, cov):
+    """Perfect nested sampling run for a correlated gaussian likelihood.
+
+    This produces a perfect nested sampling run with a uniform prior over the
+    unit hypercube, with a likelihood gaussian in the parameters normalised so
+    that the evidence is unity. The algorithm proceeds by simultaneously
+    rejection sampling from the prior and sampling exactly and uniformly from
+    the known ellipsoidal contours.
+
+    This can produce perfect runs in up to around d~15 dimensions. Beyond
+    this rejection sampling from a truncated gaussian in the early stage
+    becomes too inefficient.
+
+    Parameters
+    ----------
+    nlive: int
+        minimum number of live points across the run
+
+    mean: 1d array-like
+        mean of gaussian in parameters. Length of array defines dimensionality
+        of run.
+
+    cov: 2d array-like
+        covariance of gaussian in parameters
+
+    Returns
+    ------
+    samples: NestedSamples
+        Nested sampling run
+    """
+
+    def logLike(x):
+        return multivariate_normal(mean, cov).logpdf(x)
+
+    logLmax = logLike(mean)
+
+    points = np.random.rand(2*nlive, len(mean))
+    samples = NestedSamples(points, logL=logLike(points), logL_birth=-np.inf)
+
+    while (1/samples.nlive.iloc[:-nlive]).sum() < samples.D()*2:
+        logLs = samples.logL.iloc[-nlive]
+
+        # Cube round
+        points = np.random.rand(nlive, len(mean))
+        logL = logLike(points)
+        i = logL > logLs
+        samps_1 = NestedSamples(points[i], logL=logL[i], logL_birth=logLs)
+
+        # Ellipsoidal round
+        points = random_ellipsoid(mean, cov*2*(logLmax - logLs), nlive)
+        logL = logLike(points)
+        i = ((points > 0) & (points < 1)).all(axis=1)
+        samps_2 = NestedSamples(points[i], logL=logL[i], logL_birth=logLs)
+
+        samples = merge_nested_samples([samples, samps_1, samps_2])
+
+    return samples
diff --git a/anesthetic/examples/utils.py b/anesthetic/examples/utils.py
index 721e4958..6e02c7a6 100644
--- a/anesthetic/examples/utils.py
+++ b/anesthetic/examples/utils.py
@@ -1,10 +1,10 @@
 """Utility functions for nested sampling examples."""
 import numpy as np
 from scipy.stats import special_ortho_group
-from scipy.special import gamma
+from scipy.special import gamma, gammaln
 
 
-class random_ellipsoid(object):
+def random_ellipsoid(mean, cov, size=None):
     """Draw a point uniformly in an ellipsoid.
 
     This is defined so that the volume of the ellipsoid is sqrt(det(cov))*V_n.
@@ -21,19 +21,13 @@ class random_ellipsoid(object):
         The covariance structure of the ellipsoid. Axes have lengths equal to
         the square root of the eigenvalues of this matrix.
     """
-
-    def __init__(self, mean, cov):
-        self.mean = mean
-        self.cov = cov
-
-    def __call__(self, size=None):
-        """Generate samples uniformly from the ellipsoid."""
-        d = len(self.mean)
-        L = np.linalg.cholesky(self.cov)
-        x = np.random.multivariate_normal(np.zeros(d), np.eye(d), size=size)
-        r = np.linalg.norm(x, axis=-1, keepdims=True)
-        u = np.random.power(d, r.shape)
-        return self.mean + (u*x/r) @ L.T
+    """Generate samples uniformly from the ellipsoid."""
+    d = len(mean)
+    L = np.linalg.cholesky(cov)
+    x = np.random.multivariate_normal(np.zeros(d), np.eye(d), size=size)
+    r = np.linalg.norm(x, axis=-1, keepdims=True)
+    u = np.random.power(d, r.shape)
+    return mean + (u*x/r) @ L.T
 
 
 def random_covariance(sigmas):
@@ -51,4 +45,8 @@ def random_covariance(sigmas):
 
 def volume_n_ball(n, r=1):
     """Volume of an n dimensional ball, radius r."""
-    return np.pi**(n/2)/gamma(n/2+1)*r**n
+    return np.pi**(n/2)/gamma(1+n/2)*r**n
+
+def log_volume_n_ball(n, r=1):
+    """Log-volume of an n dimensional ball, radius r."""
+    return np.log(np.pi)*(n/2)- gammaln(1+n/2) + np.log(r)*n
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 4f4c27e6..20844109 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -1,9 +1,13 @@
 import numpy as np
 from numpy.testing import assert_allclose
 from anesthetic.examples.utils import (
-    random_ellipsoid, random_covariance, volume_n_ball
+    random_ellipsoid, random_covariance, volume_n_ball, log_volume_n_ball
+)
+from anesthetic.examples.perfect_ns import (
+    gaussian, correlated_gaussian
 )
 from scipy.spatial import ConvexHull
+from scipy.special import gamma, gammaln
 
 
 def test_random_covariance():
@@ -22,7 +26,9 @@ def test_random_ellipsoid():
     d = 5
     mean = np.random.rand(d)
     cov = random_covariance(np.random.rand(d))
-    rvs = random_ellipsoid(mean, cov)
+
+    def rvs(shape=None):
+        return random_ellipsoid(mean, cov, shape)
 
     assert(rvs(10).shape == (10, d))
     assert(rvs((10,)).shape == (10, d))
@@ -35,8 +41,41 @@ def test_random_ellipsoid():
     assert_allclose(np.cov(dat.T), cov/(d+2), atol=atol)
 
     cov = random_covariance(np.random.rand(2))
-    rvs = random_ellipsoid([0, 0], cov)
+    mean = [0, 0]
     dat = rvs(1000000)
     hull = ConvexHull(dat)
     vol = hull.volume
     assert_allclose(vol, np.linalg.det(cov)**0.5*volume_n_ball(2), rtol=1e-3)
+
+
+def test_perfect_ns_gaussian():
+    np.random.seed(0)
+    nlive = 500
+    ndims = 10
+    sigma = 0.1
+    R = 1
+    D = ndims/2*np.log(2*np.pi*np.e*sigma**2) - log_volume_n_ball(ndims)
+    logZ = gammaln(ndims/2) + log_volume_n_ball(ndims-1) - log_volume_n_ball(ndims, R)  + np.log(sigma)*ndims + (ndims/2-1)*np.log(2)
+    plt.hist(samples.logZ(1000))
+    samples = gaussian(nlive, ndims)
+    samples
+    samples.gui()
+    gaussian(nlive, ndims).logZ()
+    samples.gui()
+    plt.plot(samples.logL, samples.nlive)
+    samples.nlive.plot()
+    samples.gui()
+
+
+def test_perfect_ns_correlated_gaussian():
+    np.random.seed(0)
+    nlive = 500
+    d = 5
+    mean = 0.5*np.ones(d)
+    cov = random_covariance(np.random.rand(d)*0.1)
+    samples = correlated_gaussian(nlive, mean, cov)
+    samples.nlive.plot()
+    assert_allclose(samples.logZ(), 0, atol=3*samples.logZ(12).std())
+    assert (samples[:-nlive].nlive >= nlive).all()
+    assert_allclose(samples[[0, 1, 2, 3, 4]].mean(), mean, atol=1e-2)
+    assert_allclose(samples[[0, 1, 2, 3, 4]].cov(), cov/(d+2), atol=1e-2)

From cecb85c35e7bb57f24b067dc788eae7642da8d7d Mon Sep 17 00:00:00 2001
From: Will Handley <wh260@cam.ac.uk>
Date: Sun, 7 Aug 2022 04:33:05 +0100
Subject: [PATCH 03/19] perfect gaussian implemented

---
 anesthetic/examples/perfect_ns.py | 13 +++++++----
 anesthetic/examples/utils.py      |  3 ++-
 tests/test_examples.py            | 38 ++++++++++++++++---------------
 3 files changed, 31 insertions(+), 23 deletions(-)

diff --git a/anesthetic/examples/perfect_ns.py b/anesthetic/examples/perfect_ns.py
index 51b96294..e4217542 100644
--- a/anesthetic/examples/perfect_ns.py
+++ b/anesthetic/examples/perfect_ns.py
@@ -12,7 +12,7 @@ def gaussian(nlive, ndims, sigma=0.1, R=1, logLmin=-1e-2):
     Up to normalisation this is identical to the example in John Skilling's
     original paper https://doi.org/10.1214/06-BA127 . Both spherical Gaussian
     width sigma and spherical uniform prior width R are centered on zero
-    
+
     Parameters
     ----------
     nlive: int
@@ -29,10 +29,15 @@ def gaussian(nlive, ndims, sigma=0.1, R=1, logLmin=-1e-2):
 
     logLmin: float
         loglikelihood at which to terminate
+
+    Returns
+    -------
+    samples: NestedSamples
+        Nested sampling run
     """
 
     def logLike(x):
-        return  -(x**2).sum(axis=-1)/2/sigma**2
+        return -(x**2).sum(axis=-1)/2/sigma**2
 
     def random_sphere(n):
         return random_ellipsoid(np.zeros(ndims), np.eye(ndims), n)
@@ -48,7 +53,7 @@ def random_sphere(n):
         logL_birth = logL.copy()
         r = (points**2).sum(axis=-1, keepdims=True)**0.5
 
-    samples = merge_nested_samples(samples) 
+    samples = merge_nested_samples(samples)
     samples.logL
     logLend = samples[samples.nlive >= nlive].logL.max()
     return samples[samples.logL_birth < logLend].recompute()
@@ -80,7 +85,7 @@ def correlated_gaussian(nlive, mean, cov):
         covariance of gaussian in parameters
 
     Returns
-    ------
+    -------
     samples: NestedSamples
         Nested sampling run
     """
diff --git a/anesthetic/examples/utils.py b/anesthetic/examples/utils.py
index 6e02c7a6..c0e47ec8 100644
--- a/anesthetic/examples/utils.py
+++ b/anesthetic/examples/utils.py
@@ -47,6 +47,7 @@ def volume_n_ball(n, r=1):
     """Volume of an n dimensional ball, radius r."""
     return np.pi**(n/2)/gamma(1+n/2)*r**n
 
+
 def log_volume_n_ball(n, r=1):
     """Log-volume of an n dimensional ball, radius r."""
-    return np.log(np.pi)*(n/2)- gammaln(1+n/2) + np.log(r)*n
+    return np.log(np.pi)*n/2 - gammaln(1+n/2) + np.log(r)*n
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 20844109..3303178b 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -7,7 +7,6 @@
     gaussian, correlated_gaussian
 )
 from scipy.spatial import ConvexHull
-from scipy.special import gamma, gammaln
 
 
 def test_random_covariance():
@@ -54,28 +53,31 @@ def test_perfect_ns_gaussian():
     ndims = 10
     sigma = 0.1
     R = 1
-    D = ndims/2*np.log(2*np.pi*np.e*sigma**2) - log_volume_n_ball(ndims)
-    logZ = gammaln(ndims/2) + log_volume_n_ball(ndims-1) - log_volume_n_ball(ndims, R)  + np.log(sigma)*ndims + (ndims/2-1)*np.log(2)
-    plt.hist(samples.logZ(1000))
-    samples = gaussian(nlive, ndims)
-    samples
-    samples.gui()
-    gaussian(nlive, ndims).logZ()
-    samples.gui()
-    plt.plot(samples.logL, samples.nlive)
-    samples.nlive.plot()
-    samples.gui()
+
+    samples = gaussian(nlive, ndims, sigma, R)
+
+    assert (samples[:-nlive].nlive == nlive).all()
+
+    mean = samples[np.arange(ndims)].mean()
+    assert_allclose(mean, 0, atol=1e-2)
+    cov = samples[np.arange(ndims)].cov()
+    assert_allclose(cov, np.eye(ndims)*sigma**2, atol=1e-3)
+
+    D = log_volume_n_ball(ndims) - ndims/2*np.log(2*np.pi*np.e*sigma**2)
+    assert_allclose(samples.D(), D, atol=3*samples.D(12).std())
+
+    logZ = ndims/2 * np.log(2*np.pi*sigma**2) - log_volume_n_ball(ndims, R)
+    assert_allclose(samples.logZ(), logZ, atol=3*samples.logZ(12).std())
 
 
 def test_perfect_ns_correlated_gaussian():
     np.random.seed(0)
     nlive = 500
-    d = 5
-    mean = 0.5*np.ones(d)
-    cov = random_covariance(np.random.rand(d)*0.1)
+    ndims = 5
+    mean = 0.5*np.ones(ndims)
+    cov = random_covariance(np.random.rand(ndims)*0.1)
     samples = correlated_gaussian(nlive, mean, cov)
-    samples.nlive.plot()
     assert_allclose(samples.logZ(), 0, atol=3*samples.logZ(12).std())
     assert (samples[:-nlive].nlive >= nlive).all()
-    assert_allclose(samples[[0, 1, 2, 3, 4]].mean(), mean, atol=1e-2)
-    assert_allclose(samples[[0, 1, 2, 3, 4]].cov(), cov/(d+2), atol=1e-2)
+    assert_allclose(samples[np.arange(ndims)].mean(), mean, atol=1e-2)
+    assert_allclose(samples[np.arange(ndims)].cov(), cov/(ndims+2), atol=1e-2)

From 09685efed26d7ba6f78a5f94708a73860bf0af7a Mon Sep 17 00:00:00 2001
From: Will Handley <wh260@cam.ac.uk>
Date: Tue, 9 Aug 2022 02:23:07 +0100
Subject: [PATCH 04/19] Moved wedding cake to examples and adjusted tests

---
 anesthetic/examples/_matplotlib_agg.py |  3 +
 anesthetic/examples/perfect_ns.py      | 70 +++++++++++++++++++++
 tests/matplotlib_agg.py                |  2 -
 tests/test_examples.py                 | 24 ++++++-
 tests/test_gui.py                      |  2 +-
 tests/test_plot.py                     |  2 +-
 tests/test_reader.py                   |  2 +-
 tests/test_samples.py                  | 15 +----
 tests/wedding_cake.py                  | 87 --------------------------
 9 files changed, 100 insertions(+), 107 deletions(-)
 create mode 100644 anesthetic/examples/_matplotlib_agg.py
 delete mode 100644 tests/matplotlib_agg.py
 delete mode 100644 tests/wedding_cake.py

diff --git a/anesthetic/examples/_matplotlib_agg.py b/anesthetic/examples/_matplotlib_agg.py
new file mode 100644
index 00000000..a140116a
--- /dev/null
+++ b/anesthetic/examples/_matplotlib_agg.py
@@ -0,0 +1,3 @@
+"""When imported forces matplotlib to use Agg backend. Useful for tests."""
+import matplotlib
+matplotlib.use('Agg')
diff --git a/anesthetic/examples/perfect_ns.py b/anesthetic/examples/perfect_ns.py
index e4217542..406491d7 100644
--- a/anesthetic/examples/perfect_ns.py
+++ b/anesthetic/examples/perfect_ns.py
@@ -116,3 +116,73 @@ def logLike(x):
         samples = merge_nested_samples([samples, samps_1, samps_2])
 
     return samples
+
+
+def wedding_cake(nlive, ndims, sigma=0.01, alpha=0.5):
+    """Perfect nested sampling run for a wedding cake likelihood.
+
+    This is a likelihood with nested hypercuboidal plateau regions of constant
+    likelihood centered on 0.5, with geometrically decreasing volume by a
+    factor of alpha. The value of the likelihood in these plateau regions has a
+    gaussian profile with width sigma.
+
+    logL = - alpha^(2 floor(D*log_alpha(2|x-0.5|_infinity))/D) / (8 sigma^2)
+
+    Parameters
+    ----------
+    nlive: int
+        number of live points
+
+    ndims: int
+        dimensionality of the likelihood
+
+    sigma: float
+        width of gaussian profile
+
+    alpha: float
+        volume compression between plateau regions
+    """
+
+    def i(x):
+        """Plateau number of a parameter point."""
+        r = np.max(abs(x-0.5), axis=-1)
+        return np.floor(ndims*np.log(2*r)/np.log(alpha))
+
+    def logL(x):
+        """Gaussian log-likelihood."""
+        ri = alpha**(i(x)/ndims)/2
+        return - ri**2/2/sigma**2
+
+    points = np.zeros((0, ndims))
+    death_likes = np.zeros(0)
+    birth_likes = np.zeros(0)
+
+    live_points = np.random.rand(nlive, ndims)
+    live_likes = logL(live_points)
+    live_birth_likes = np.ones(nlive)*-np.inf
+
+    while True:
+        logL_ = live_likes.min()
+        j = live_likes == logL_
+
+        death_likes = np.concatenate([death_likes, live_likes[j]])
+        birth_likes = np.concatenate([birth_likes, live_birth_likes[j]])
+        points = np.concatenate([points, live_points[j]])
+
+        i_ = i(live_points[j][0])+1
+        r_ = alpha**(i_/ndims)/2
+        x_ = np.random.uniform(0.5-r_, 0.5+r_, size=(j.sum(), ndims))
+        live_birth_likes[j] = logL_
+        live_points[j] = x_
+        live_likes[j] = logL(x_)
+
+        samps = NestedSamples(points, logL=death_likes, logL_birth=birth_likes)
+
+        if samps.iloc[-nlive:].weights.sum()/samps.weights.sum() < 0.001:
+            break
+
+    death_likes = np.concatenate([death_likes, live_likes])
+    birth_likes = np.concatenate([birth_likes, live_birth_likes])
+    points = np.concatenate([points, live_points])
+
+    return NestedSamples(points, logL=death_likes, logL_birth=birth_likes)
diff --git a/tests/matplotlib_agg.py b/tests/matplotlib_agg.py
deleted file mode 100644
index 60b27b09..00000000
--- a/tests/matplotlib_agg.py
+++ /dev/null
@@ -1,2 +0,0 @@
-import matplotlib
-matplotlib.use('Agg')
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 3303178b..3086bc97 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -4,9 +4,10 @@
     random_ellipsoid, random_covariance, volume_n_ball, log_volume_n_ball
 )
 from anesthetic.examples.perfect_ns import (
-    gaussian, correlated_gaussian
+    gaussian, correlated_gaussian, wedding_cake
 )
 from scipy.spatial import ConvexHull
+from scipy.special import logsumexp
 
 
 def test_random_covariance():
@@ -81,3 +82,24 @@ def test_perfect_ns_correlated_gaussian():
     assert (samples[:-nlive].nlive >= nlive).all()
     assert_allclose(samples[np.arange(ndims)].mean(), mean, atol=1e-2)
     assert_allclose(samples[np.arange(ndims)].cov(), cov/(ndims+2), atol=1e-2)
+
+
+def test_wedding_cake():
+    np.random.seed(0)
+    nlive = 500
+    ndims = 4
+    sigma = 0.01
+    alpha = 0.5
+    samples = wedding_cake(nlive, ndims, sigma, alpha)
+
+    assert samples.nlive.iloc[0] == nlive
+    assert samples.nlive.iloc[-1] == 1
+    assert (samples.nlive <= nlive).all()
+
+    mean = np.sqrt(ndims/2) * (np.log(4*ndims*sigma**2)-1) / np.log(alpha)
+    std = -np.sqrt(ndims/2)/np.log(alpha)
+    i = np.arange(mean + std*10)
+    logZ = logsumexp(-alpha**(2*i/ndims)/8/sigma**2
+                     + i*np.log(alpha) + np.log(1-alpha))
+
+    assert_allclose(samples.logZ(), logZ, atol=3*samples.logZ(12).std())
diff --git a/tests/test_gui.py b/tests/test_gui.py
index c458a592..edcbc09f 100644
--- a/tests/test_gui.py
+++ b/tests/test_gui.py
@@ -1,4 +1,4 @@
-import matplotlib_agg  # noqa: F401
+import anesthetic.examples._matplotlib_agg  # noqa: F401
 from packaging import version
 from matplotlib import __version__ as mpl_version
 from anesthetic import NestedSamples
diff --git a/tests/test_plot.py b/tests/test_plot.py
index b7055a8e..6d5d50e7 100644
--- a/tests/test_plot.py
+++ b/tests/test_plot.py
@@ -1,4 +1,4 @@
-import matplotlib_agg  # noqa: F401
+import anesthetic.examples._matplotlib_agg  # noqa: F401
 import pytest
 import numpy as np
 import sys
diff --git a/tests/test_reader.py b/tests/test_reader.py
index 14e71143..f7cf1cef 100644
--- a/tests/test_reader.py
+++ b/tests/test_reader.py
@@ -1,4 +1,4 @@
-import matplotlib_agg  # noqa: F401
+import anesthetic.examples._matplotlib_agg  # noqa: F401
 import os
 import sys
 import pytest
diff --git a/tests/test_samples.py b/tests/test_samples.py
index 82f54bbc..043e528f 100644
--- a/tests/test_samples.py
+++ b/tests/test_samples.py
@@ -1,5 +1,5 @@
 import warnings
-import matplotlib_agg  # noqa: F401
+import anesthetic.examples._matplotlib_agg  # noqa: F401
 import sys
 import pytest
 import numpy as np
@@ -14,7 +14,6 @@
 from pandas.testing import assert_frame_equal
 from matplotlib.colors import to_hex
 from scipy.stats import ks_2samp, kstest, norm
-from wedding_cake import WeddingCake
 try:
     import montepython  # noqa: F401
 except ImportError:
@@ -790,18 +789,6 @@ def test_MCMCSamples_importance_sample():
     assert mc3.limits is not mc0.limits
 
 
-def test_wedding_cake():
-    np.random.seed(3)
-    wc = WeddingCake(4, 0.5, 0.01)
-    nlive = 500
-    samples = wc.sample(nlive)
-    assert samples.nlive.iloc[0] == nlive
-    assert samples.nlive.iloc[-1] == 1
-    assert (samples.nlive <= nlive).all()
-    out = samples.logZ(100)
-    assert abs(out.mean()-wc.logZ()) < out.std()*3
-
-
 def test_logzero_mask_prior_level():
     np.random.seed(3)
     ns0 = NestedSamples(root='./tests/example_data/pc')
diff --git a/tests/wedding_cake.py b/tests/wedding_cake.py
deleted file mode 100644
index 46b755f9..00000000
--- a/tests/wedding_cake.py
+++ /dev/null
@@ -1,87 +0,0 @@
-import numpy as np
-from scipy.special import logsumexp
-from anesthetic import NestedSamples
-
-
-class WeddingCake():
-    """Class for generating samples from a wedding cake 'likelihood'.
-
-    This is a likelihood with nested hypercuboidal plateau regions of constant
-    likelihood centered on 0.5, with geometrically decreasing volume by a
-    factor of alpha. The value of the likelihood in these plateau regions has a
-    gaussian profile with width sigma.
-
-    logL = - alpha^(2 floor(D*log_alpha(2|x-0.5|_infinity))/D) / (8 sigma^2)
-
-    Parameters
-    ----------
-    D: int
-        dimensionality (number of parameters) of the likelihood
-
-    alpha: float
-        volume compression between plateau regions
-
-    sigma: float
-        width of gaussian profile
-    """
-    def __init__(self, D=4, alpha=0.5, sigma=0.01):
-        self.D = D
-        self.alpha = alpha
-        self.sigma = sigma
-
-    def logZ(self):
-        """Numerically compute the true evidence."""
-        mean = np.sqrt(self.D/2.) * (np.log(4*self.D*self.sigma**2)-1
-                                     ) / np.log(self.alpha)
-        std = -np.sqrt(self.D/2.)/np.log(self.alpha)
-        i = np.arange(mean + std*10)
-
-        return logsumexp(-self.alpha**(2*i/self.D)/8/self.sigma**2
-                         + i*np.log(self.alpha) + np.log(1-self.alpha))
-
-    def i(self, x):
-        """Plateau number of a parameter point."""
-        r = np.max(abs(x-0.5), axis=-1)
-        return np.floor(self.D*np.log(2*r)/np.log(self.alpha))
-
-    def logL(self, x):
-        """Gaussian log-likelihood."""
-        ri = self.alpha**(self.i(x)/self.D)/2
-        return - ri**2/2/self.sigma**2
-
-    def sample(self, nlive=500):
-        """Generate samples from a perfect nested sampling run."""
-        points = np.zeros((0, self.D))
-        death_likes = np.zeros(0)
-        birth_likes = np.zeros(0)
-
-        live_points = np.random.rand(nlive, self.D)
-        live_likes = self.logL(live_points)
-        live_birth_likes = np.ones(nlive)*-np.inf
-
-        while True:
-            logL_ = live_likes.min()
-            j = live_likes == logL_
-
-            death_likes = np.concatenate([death_likes, live_likes[j]])
-            birth_likes = np.concatenate([birth_likes, live_birth_likes[j]])
-            points = np.concatenate([points, live_points[j]])
-
-            i_ = self.i(live_points[j][0])+1
-            r_ = self.alpha**(i_/self.D)/2
-            x_ = np.random.uniform(0.5-r_, 0.5+r_, size=(j.sum(), self.D))
-            live_birth_likes[j] = logL_
-            live_points[j] = x_
-            live_likes[j] = self.logL(x_)
-
-            samps = NestedSamples(points, logL=death_likes,
-                                  logL_birth=birth_likes)
-
-            if samps.iloc[-nlive:].weights.sum()/samps.weights.sum() < 0.001:
-                break
-
-        death_likes = np.concatenate([death_likes, live_likes])
-        birth_likes = np.concatenate([birth_likes, live_birth_likes])
-        points = np.concatenate([points, live_points])
-
-        return NestedSamples(points, logL=death_likes, logL_birth=birth_likes)

From 120205583648856c1037748b90e31f0ac6b70d81 Mon Sep 17 00:00:00 2001
From: Will Handley <wh260@cam.ac.uk>
Date: Tue, 9 Aug 2022 02:35:42 +0100
Subject: [PATCH 05/19] updated plot command

---
 bin/plot.py | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/bin/plot.py b/bin/plot.py
index cee7c3a9..85c86e4b 100644
--- a/bin/plot.py
+++ b/bin/plot.py
@@ -9,7 +9,7 @@
 a, b, m = 2., 4., 0.5         # x4 parameters
 
 n = 1000
-ls = 'k--'
+ls = 'k'
 
 x = np.linspace(-0.4,0.4,n)
 p = np.exp(-x**2/sigma0**2/2)/np.sqrt(2*np.pi)/sigma0
@@ -39,14 +39,20 @@
 axes['x0']['x3'].plot(x0, x3, ls)
 axes['x0']['x3'].plot(-x0, x3, ls)
 axes['x0']['x3'].plot(-2*x0, x3, ls)
+axes['x0']['x3'].plot(np.linspace(-2*sigma0, 2*sigma0, n), np.ones(n), ls)
+axes['x0']['x3'].plot(np.linspace(-2*sigma0, 2*sigma0, n), np.zeros(n), ls)
 
 axes['x1']['x3'].plot(2*x0, x3, ls)
 axes['x1']['x3'].plot(x0, x3, ls)
 axes['x1']['x3'].plot(-x0, x3, ls)
 axes['x1']['x3'].plot(-2*x0, x3, ls)
+axes['x1']['x3'].plot(np.linspace(-2*sigma1, 2*sigma1, n), np.ones(n), ls)
+axes['x1']['x3'].plot(np.linspace(-2*sigma1, 2*sigma1, n), np.zeros(n), ls)
 
-for p in [0.66, 0.95]:
+for p in [0, 0.66, 0.95]:
     axes['x2']['x3'].plot(-np.log(1-p)*x0, x3, ls)
+axes['x2']['x3'].plot(np.linspace(0, -np.log(1-p)*sigma0, n), np.ones(n), ls)
+axes['x2']['x3'].plot(np.linspace(0, -np.log(1-p)*sigma0, n), np.zeros(n), ls)
 
 from scipy.optimize import root
 from scipy.special import erf
@@ -57,6 +63,9 @@
     y = k - x**2/2
     axes['x0']['x2'].plot(x*sigma0, y*sigma2, ls)
     axes['x1']['x2'].plot(x*sigma1, y*sigma2, ls)
+    axes['x0']['x2'].plot(x*sigma0, 0*x, ls)
+    axes['x1']['x2'].plot(x*sigma1, 0*x, ls)
+
 
 t = np.linspace(0, 2*np.pi, n)
 x0 = sigma1*eps*np.cos(t) + sigma0 * np.sin(t)

From 3545c7cc3ebea53d7544bc862b68d052115ab015 Mon Sep 17 00:00:00 2001
From: Will Handley <wh260@cam.ac.uk>
Date: Tue, 9 Aug 2022 06:53:33 +0100
Subject: [PATCH 06/19] Equipped correlated gaussian with bounds

---
 anesthetic/examples/perfect_ns.py | 24 +++++++++++++++++-------
 tests/test_examples.py            |  4 +++-
 2 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/anesthetic/examples/perfect_ns.py b/anesthetic/examples/perfect_ns.py
index 406491d7..12b8a5dc 100644
--- a/anesthetic/examples/perfect_ns.py
+++ b/anesthetic/examples/perfect_ns.py
@@ -1,6 +1,5 @@
 """Perfect nested sampling generators."""
 import numpy as np
-from scipy.stats import multivariate_normal
 from anesthetic.examples.utils import random_ellipsoid
 from anesthetic import NestedSamples
 from anesthetic.samples import merge_nested_samples
@@ -59,7 +58,7 @@ def random_sphere(n):
     return samples[samples.logL_birth < logLend].recompute()
 
 
-def correlated_gaussian(nlive, mean, cov):
+def correlated_gaussian(nlive, mean, cov, bounds=None):
     """Perfect nested sampling run for a correlated gaussian likelihood.
 
     This produces a perfect nested sampling run with a uniform prior over the
@@ -77,32 +76,43 @@ def correlated_gaussian(nlive, mean, cov):
     nlive: int
         minimum number of live points across the run
 
-    mean: 1d array-like
+    mean: 1d array-like, shape (ndims,)
         mean of gaussian in parameters. Length of array defines dimensionality
         of run.
 
-    cov: 2d array-like
+    cov: 2d array-like, shape (ndims,ndims)
         covariance of gaussian in parameters
 
+    bounds: 2d array-like, shape (ndims, 2)
+        bounds of a gaussian, default [[0, 1]]*ndims
+
     Returns
     -------
     samples: NestedSamples
         Nested sampling run
     """
 
+    invcov = np.linalg.inv(cov)
+
     def logLike(x):
-        return multivariate_normal(mean, cov).logpdf(x)
+        return -0.5 * ((x-mean) @ invcov * (x-mean)).sum(axis=-1)
+
+    ndims = len(mean)
+
+    if bounds is None:
+        bounds = [[0, 1]]*ndims
 
+    bounds = np.array(bounds)
     logLmax = logLike(mean)
 
-    points = np.random.rand(2*nlive, len(mean))
+    points = np.random.uniform(*bounds.T, (2*nlive, ndims))
     samples = NestedSamples(points, logL=logLike(points), logL_birth=-np.inf)
 
     while (1/samples.nlive.iloc[:-nlive]).sum() < samples.D()*2:
         logLs = samples.logL.iloc[-nlive]
 
         # Cube round
-        points = np.random.rand(nlive, len(mean))
+        points = np.random.uniform(*bounds.T, (nlive, ndims))
         logL = logLike(points)
         i = logL > logLs
         samps_1 = NestedSamples(points[i], logL=logL[i], logL_birth=logLs)
diff --git a/tests/test_examples.py b/tests/test_examples.py
index 3086bc97..ca87c07c 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -78,7 +78,9 @@ def test_perfect_ns_correlated_gaussian():
     mean = 0.5*np.ones(ndims)
     cov = random_covariance(np.random.rand(ndims)*0.1)
     samples = correlated_gaussian(nlive, mean, cov)
-    assert_allclose(samples.logZ(), 0, atol=3*samples.logZ(12).std())
+    samples.gui()
+    logZ = np.linalg.slogdet(2*np.pi*cov)[1]/2
+    assert_allclose(samples.logZ(), logZ, atol=3*samples.logZ(12).std())
     assert (samples[:-nlive].nlive >= nlive).all()
     assert_allclose(samples[np.arange(ndims)].mean(), mean, atol=1e-2)
     assert_allclose(samples[np.arange(ndims)].cov(), cov/(ndims+2), atol=1e-2)

From df30a6cf2552bbc90b7a281bde5f4b945dc88724 Mon Sep 17 00:00:00 2001
From: Will Handley <wh260@cam.ac.uk>
Date: Tue, 9 Aug 2022 19:21:01 +0100
Subject: [PATCH 07/19] Fixed to failing lint

---
 anesthetic/examples/perfect_ns.py | 1 -
 tests/test_examples.py            | 8 ++++----
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/anesthetic/examples/perfect_ns.py b/anesthetic/examples/perfect_ns.py
index 12b8a5dc..12bd8b82 100644
--- a/anesthetic/examples/perfect_ns.py
+++ b/anesthetic/examples/perfect_ns.py
@@ -91,7 +91,6 @@ def correlated_gaussian(nlive, mean, cov, bounds=None):
     samples: NestedSamples
         Nested sampling run
     """
-
     invcov = np.linalg.inv(cov)
 
     def logLike(x):
diff --git a/tests/test_examples.py b/tests/test_examples.py
index ca87c07c..dedb44c5 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -15,7 +15,7 @@ def test_random_covariance():
     d = 5
     sigmas = np.random.rand(d)
     cov = random_covariance(sigmas)
-    assert(cov.shape == (d, d))
+    assert cov.shape == (d, d) 
     e = np.linalg.eigvals(cov)
     assert_allclose(np.sort(e), np.sort(sigmas**2))
 
@@ -30,9 +30,9 @@ def test_random_ellipsoid():
     def rvs(shape=None):
         return random_ellipsoid(mean, cov, shape)
 
-    assert(rvs(10).shape == (10, d))
-    assert(rvs((10,)).shape == (10, d))
-    assert(rvs((10, 20)).shape == (10, 20, d))
+    assert rvs(10).shape == (10, d)
+     ssert(rvs((10,)).shape == (10, d)
+     ssert(rvs((10, 20)).shape == (10, 20, d)
 
     dat = rvs(100000)
 

From a07d2d31b4f3311a2833cb904bd76b3e1968f5ff Mon Sep 17 00:00:00 2001
From: Will Handley <wh260@cam.ac.uk>
Date: Wed, 10 Aug 2022 01:25:27 +0100
Subject: [PATCH 08/19] lint corrections

---
 tests/test_examples.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/test_examples.py b/tests/test_examples.py
index dedb44c5..6e8ece34 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -15,7 +15,7 @@ def test_random_covariance():
     d = 5
     sigmas = np.random.rand(d)
     cov = random_covariance(sigmas)
-    assert cov.shape == (d, d) 
+    assert cov.shape == (d, d)
     e = np.linalg.eigvals(cov)
     assert_allclose(np.sort(e), np.sort(sigmas**2))
 
@@ -31,8 +31,8 @@ def rvs(shape=None):
         return random_ellipsoid(mean, cov, shape)
 
     assert rvs(10).shape == (10, d)
-     ssert(rvs((10,)).shape == (10, d)
-     ssert(rvs((10, 20)).shape == (10, 20, d)
+    assert rvs((10,)).shape == (10, d)
+    assert rvs((10, 20)).shape == (10, 20, d)
 
     dat = rvs(100000)
 

From cc5f931a705b5d76f70183f727900245d8f48b24 Mon Sep 17 00:00:00 2001
From: Lukas Hergt <lthergt@phas.ubc.ca>
Date: Thu, 1 Dec 2022 20:22:15 -0800
Subject: [PATCH 09/19] Update test_samples.py

---
 tests/test_samples.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tests/test_samples.py b/tests/test_samples.py
index 7149b79e..29a0509a 100644
--- a/tests/test_samples.py
+++ b/tests/test_samples.py
@@ -1,4 +1,3 @@
-import warnings
 import anesthetic.examples._matplotlib_agg  # noqa: F401
 
 import sys

From 4a3621cb177c9df560416afb86951f714c9d0f9f Mon Sep 17 00:00:00 2001
From: Lukas Hergt <lthergt@phas.ubc.ca>
Date: Thu, 1 Dec 2022 20:23:35 -0800
Subject: [PATCH 10/19] Update test_gui.py

---
 tests/test_gui.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/tests/test_gui.py b/tests/test_gui.py
index d77333fe..f5111e13 100644
--- a/tests/test_gui.py
+++ b/tests/test_gui.py
@@ -1,6 +1,4 @@
 import anesthetic.examples._matplotlib_agg  # noqa: F401
-from packaging import version
-from matplotlib import __version__ as mpl_version
 from anesthetic import read_chains
 
 

From 1a958dd708a395e82ca42dba65c834f1b82f1378 Mon Sep 17 00:00:00 2001
From: Will Handley <wh260@cam.ac.uk>
Date: Mon, 16 Jan 2023 18:12:26 +0000
Subject: [PATCH 11/19] Removed erroneous line

---
 anesthetic/examples/perfect_ns.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/anesthetic/examples/perfect_ns.py b/anesthetic/examples/perfect_ns.py
index 712d948c..43adc6db 100644
--- a/anesthetic/examples/perfect_ns.py
+++ b/anesthetic/examples/perfect_ns.py
@@ -53,7 +53,6 @@ def random_sphere(n):
         r = (points**2).sum(axis=-1, keepdims=True)**0.5
 
     samples = merge_nested_samples(samples)
-    samples.logL
     logLend = samples[samples.nlive >= nlive].logL.max()
     return samples[samples.logL_birth < logLend].recompute()
 

From c2decc702b99832748c0539d404b3b5e5e1a4c64 Mon Sep 17 00:00:00 2001
From: Will Handley <wh260@cam.ac.uk>
Date: Mon, 16 Jan 2023 18:18:53 +0000
Subject: [PATCH 12/19] Updated docstrings in utility functions

---
 anesthetic/examples/utils.py | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/anesthetic/examples/utils.py b/anesthetic/examples/utils.py
index c0e47ec8..2ace401a 100644
--- a/anesthetic/examples/utils.py
+++ b/anesthetic/examples/utils.py
@@ -20,8 +20,13 @@ def random_ellipsoid(mean, cov, size=None):
     cov: 2d array-like
         The covariance structure of the ellipsoid. Axes have lengths equal to
         the square root of the eigenvalues of this matrix.
+
+    Returns
+    -------
+    points: array-like
+        Points drawn uniformly from the ellipsoid.
+        shape (*size, len(mean)).
     """
-    """Generate samples uniformly from the ellipsoid."""
     d = len(mean)
     L = np.linalg.cholesky(cov)
     x = np.random.multivariate_normal(np.zeros(d), np.eye(d), size=size)
@@ -37,6 +42,11 @@ def random_covariance(sigmas):
     ----------
     sigmas, 1d array like
         Lengths of the axes of the ellipsoid.
+
+    Returns
+    -------
+    Covariance matrix: 2d np.array
+        shape (len(sigmas, len(sigmas)).
     """
     d = len(sigmas)
     R = special_ortho_group.rvs(d)

From ba4707de8d3bfba78898494feda4396ba151a4d0 Mon Sep 17 00:00:00 2001
From: Will Handley <wh260@cam.ac.uk>
Date: Mon, 16 Jan 2023 18:20:49 +0000
Subject: [PATCH 13/19] Updated module docstring

---
 anesthetic/examples/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/anesthetic/examples/__init__.py b/anesthetic/examples/__init__.py
index cf11d49a..fa0a5121 100644
--- a/anesthetic/examples/__init__.py
+++ b/anesthetic/examples/__init__.py
@@ -1 +1 @@
-"""Module for generating nested sampling examples."""
+"""Module for generating example data of nested sampling."""

From 57e35e733479e61a25292bbb0d4a39bee2c496b0 Mon Sep 17 00:00:00 2001
From: Will Handley <wh260@cam.ac.uk>
Date: Sun, 22 Jan 2023 18:35:43 +0000
Subject: [PATCH 14/19] Updated bounds treatment

---
 anesthetic/examples/perfect_ns.py | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/anesthetic/examples/perfect_ns.py b/anesthetic/examples/perfect_ns.py
index 712d948c..18afbd17 100644
--- a/anesthetic/examples/perfect_ns.py
+++ b/anesthetic/examples/perfect_ns.py
@@ -91,6 +91,8 @@ def correlated_gaussian(nlive, mean, cov, bounds=None):
     samples: NestedSamples
         Nested sampling run
     """
+    mean = np.array(mean, dtype=float)
+    cov = np.array(cov, dtype=float)
     invcov = np.linalg.inv(cov)
 
     def logLike(x):
@@ -101,7 +103,8 @@ def logLike(x):
     if bounds is None:
         bounds = [[0, 1]]*ndims
 
-    bounds = np.array(bounds)
+    bounds = np.array(bounds, dtype=float)
+
     logLmax = logLike(mean)
 
     points = np.random.uniform(*bounds.T, (2*nlive, ndims))
@@ -110,6 +113,10 @@ def logLike(x):
     while (1/samples.nlive.iloc[:-nlive]).sum() < samples.D_KL()*2:
         logLs = samples.logL.iloc[-nlive]
 
+        sig = np.diag(cov*2*(logLmax - logLs))**0.5
+        bounds[:, 0] = np.max([bounds[:, 0], mean - sig], axis=0)
+        bounds[:, 1] = np.min([bounds[:, 1], mean + sig], axis=0)
+
         # Cube round
         points = np.random.uniform(*bounds.T, (nlive, ndims))
         logL = logLike(points)
@@ -119,8 +126,9 @@ def logLike(x):
         # Ellipsoidal round
         points = random_ellipsoid(mean, cov*2*(logLmax - logLs), nlive)
         logL = logLike(points)
-        i = ((points > 0) & (points < 1)).all(axis=1)
+        i = ((points > bounds.T[0]) & (points < bounds.T[1])).all(axis=1)
         samps_2 = NestedSamples(points[i], logL=logL[i], logL_birth=logLs)
+        print(len(samps_1), len(samps_2))
 
         samples = merge_nested_samples([samples, samps_1, samps_2])
 

From e6c5ad0863bd68313809bade744f6614ea92b30c Mon Sep 17 00:00:00 2001
From: Will Handley <wh260@cam.ac.uk>
Date: Mon, 23 Jan 2023 13:48:13 +0000
Subject: [PATCH 15/19] Added a planck example

---
 anesthetic/examples/perfect_ns.py | 54 +++++++++++++++++++++++++++++++
 anesthetic/samples.py             |  4 +--
 tests/test_examples.py            | 44 ++++++++++++++++++++++++-
 3 files changed, 99 insertions(+), 3 deletions(-)

diff --git a/anesthetic/examples/perfect_ns.py b/anesthetic/examples/perfect_ns.py
index a70f24b6..07a80bfe 100644
--- a/anesthetic/examples/perfect_ns.py
+++ b/anesthetic/examples/perfect_ns.py
@@ -202,3 +202,57 @@ def logL(x):
     points = np.concatenate([points, live_points])
 
     return NestedSamples(points, logL=death_likes, logL_birth=birth_likes)
+
+
+def planck_gaussian(nlive=500):
+    """A Gaussian likelihood with Planck-like posterior.
+
+    This is a gaussian likelihood with the same mean, parameter covariance and
+    average loglikelihood as the Planck 2018 legacy chains
+    base/plikHM_TTTEEE_lowl_lowE_lensing
+
+    Parameters
+    ----------
+    nlive : int, optional
+        number of live points
+
+    Returns
+    -------
+    samples: NestedSamples
+        Nested sampling run
+    """
+
+    columns = ['omegabh2', 'omegach2', 'theta', 'tau', 'logA', 'ns']
+    labels = [r'$\Omega_b h^2$', r'$\Omega_c h^2$', r'$100\theta_{MC}$',
+              r'$\tau$', r'${\rm{ln}}(10^{10} A_s)$', r'$n_s$']
+
+    cov = np.array([
+        [2.12e-08, -9.03e-08, 1.76e-08, 2.96e-07, 4.97e-07, 2.38e-07],
+        [-9.03e-08, 1.39e-06, -1.26e-07, -3.41e-06, -4.15e-06, -3.28e-06],
+        [1.76e-08, -1.26e-07, 9.71e-08, 4.30e-07, 7.41e-07, 4.13e-07],
+        [2.96e-07, -3.41e-06, 4.30e-07, 5.33e-05, 9.53e-05, 1.05e-05],
+        [4.97e-07, -4.15e-06, 7.41e-07, 9.53e-05, 2.00e-04, 1.35e-05],
+        [2.38e-07, -3.28e-06, 4.13e-07, 1.05e-05, 1.35e-05, 1.73e-05]])
+
+    mean = np.array([0.02237, 0.1200, 1.04092, 0.0544, 3.044, 0.9649])
+
+    bounds = np.array([
+        [5.00e-03, 1.00e-01],
+        [1.00e-03, 9.90e-01],
+        [5.00e-01, 1.00e+01],
+        [1.00e-02, 8.00e-01],
+        [1.61e+00, 3.91e+00],
+        [8.00e-01, 1.20e+00]])
+
+    logL_mean = -1400.35
+
+    samples = correlated_gaussian(nlive, mean, cov, bounds)
+
+    data = samples.iloc[:, :len(columns)].to_numpy()
+    logL = samples['logL'].to_numpy()
+    logL_birth = samples['logL_birth'].to_numpy()
+    logL_birth += logL_mean - samples.logL.mean()
+    logL += logL_mean - samples.logL.mean()
+    samples = NestedSamples(data=data, columns=columns, labels=labels,
+                            logL=logL, logL_birth=logL_birth)
+    return samples
diff --git a/anesthetic/samples.py b/anesthetic/samples.py
index c205ea0d..feb3e384 100644
--- a/anesthetic/samples.py
+++ b/anesthetic/samples.py
@@ -615,7 +615,7 @@ def stats(self, nsamples=None, beta=None):
         samples['D_KL'] = (S*w).sum()
         samples.set_label('D_KL', r'$\mathcal{D}_\mathrm{KL}$')
 
-        samples['d_G'] = ((S-samples.D_KL)**2*w).sum()
+        samples['d_G'] = ((S-samples.D_KL)**2*w).sum()*2
         samples.set_label('d_G', r'$d_\mathrm{G}$')
 
         samples['logL_P'] = samples['logZ'] + samples['D_KL']
@@ -819,7 +819,7 @@ def d_G(self, nsamples=None, beta=None):
         S = (logw*0).add(betalogL, axis=0) - logZ
         w = np.exp(logw-logZ)
         D_KL = (S*w).sum()
-        d_G = ((S-D_KL)**2*w).sum()
+        d_G = ((S-D_KL)**2*w).sum()*2
         if np.isscalar(d_G):
             return d_G
         else:
diff --git a/tests/test_examples.py b/tests/test_examples.py
index e1e2cdb8..f0399816 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -4,7 +4,7 @@
     random_ellipsoid, random_covariance, volume_n_ball, log_volume_n_ball
 )
 from anesthetic.examples.perfect_ns import (
-    gaussian, correlated_gaussian, wedding_cake
+    gaussian, correlated_gaussian, wedding_cake, planck_gaussian
 )
 from scipy.spatial import ConvexHull
 from scipy.special import logsumexp
@@ -105,3 +105,45 @@ def test_wedding_cake():
                      + i*np.log(alpha) + np.log(1-alpha))
 
     assert_allclose(samples.logZ(), logZ, atol=3*samples.logZ(12).std())
+
+
+def test_planck():
+    np.random.seed(0)
+    nlive = 50
+    cov = np.array([
+        [2.12e-08, -9.03e-08, 1.76e-08, 2.96e-07, 4.97e-07, 2.38e-07],
+        [-9.03e-08, 1.39e-06, -1.26e-07, -3.41e-06, -4.15e-06, -3.28e-06],
+        [1.76e-08, -1.26e-07, 9.71e-08, 4.30e-07, 7.41e-07, 4.13e-07],
+        [2.96e-07, -3.41e-06, 4.30e-07, 5.33e-05, 9.53e-05, 1.05e-05],
+        [4.97e-07, -4.15e-06, 7.41e-07, 9.53e-05, 2.00e-04, 1.35e-05],
+        [2.38e-07, -3.28e-06, 4.13e-07, 1.05e-05, 1.35e-05, 1.73e-05]])
+
+    mean = np.array([0.02237, 0.1200, 1.04092, 0.0544, 3.044, 0.9649])
+
+    bounds = np.array([
+        [5.00e-03, 1.00e-01],
+        [1.00e-03, 9.90e-01],
+        [5.00e-01, 1.00e+01],
+        [1.00e-02, 8.00e-01],
+        [1.61e+00, 3.91e+00],
+        [8.00e-01, 1.20e+00]])
+
+    logL_mean = -1400.35
+    d = 6
+    logdetroottwopicov = np.linalg.slogdet(2*np.pi*cov)[1]/2
+    logZ = logL_mean + logdetroottwopicov + d/2
+    logV = np.log(np.diff(bounds)).sum()
+    D_KL = logV - logdetroottwopicov - d/2
+    planck = planck_gaussian(nlive)
+
+    params = ['omegabh2', 'omegach2', 'theta', 'tau', 'logA', 'ns']
+
+    assert (planck[params] > bounds.T[0]).all().all()
+    assert (planck[params] < bounds.T[1]).all().all()
+    assert_allclose(planck[params].mean(), mean, atol=1e-3)
+    assert_allclose(planck[params].cov(), cov, atol=1e-3)
+    assert_allclose(planck.logL.mean(), logL_mean, atol=1e-3)
+    assert_allclose(planck.logZ(), logZ, atol=3*planck.logZ(12).std())
+    assert_allclose(planck.D_KL(), D_KL, atol=3*planck.D_KL(12).std())
+    assert_allclose(planck.d_G(), len(params), atol=3*planck.d_G(12).std())
+    assert_allclose(planck.logL_P(), logL_mean, atol=3*planck.logL_P(12).std())

From 6c713bff366af6b97fc05e1bf258ede671d0e8e8 Mon Sep 17 00:00:00 2001
From: Will Handley <wh260@cam.ac.uk>
Date: Mon, 23 Jan 2023 14:52:11 +0000
Subject: [PATCH 16/19] Fixed lint

---
 anesthetic/examples/perfect_ns.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/anesthetic/examples/perfect_ns.py b/anesthetic/examples/perfect_ns.py
index 07a80bfe..75f229cb 100644
--- a/anesthetic/examples/perfect_ns.py
+++ b/anesthetic/examples/perfect_ns.py
@@ -205,7 +205,7 @@ def logL(x):
 
 
 def planck_gaussian(nlive=500):
-    """A Gaussian likelihood with Planck-like posterior.
+    """Gaussian likelihood with Planck-like posterior.
 
     This is a gaussian likelihood with the same mean, parameter covariance and
     average loglikelihood as the Planck 2018 legacy chains
@@ -221,7 +221,6 @@ def planck_gaussian(nlive=500):
     samples: NestedSamples
         Nested sampling run
     """
-
     columns = ['omegabh2', 'omegach2', 'theta', 'tau', 'logA', 'ns']
     labels = [r'$\Omega_b h^2$', r'$\Omega_c h^2$', r'$100\theta_{MC}$',
               r'$\tau$', r'${\rm{ln}}(10^{10} A_s)$', r'$n_s$']

From ee1513debeb0829e4e09c79c08eeb6b67da01048 Mon Sep 17 00:00:00 2001
From: Will Handley <wh260@cam.ac.uk>
Date: Wed, 25 Jan 2023 10:15:41 +0000
Subject: [PATCH 17/19] d_G check for perfect gaussian case

---
 tests/test_examples.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/test_examples.py b/tests/test_examples.py
index f0399816..506388fd 100644
--- a/tests/test_examples.py
+++ b/tests/test_examples.py
@@ -67,6 +67,8 @@ def test_perfect_ns_gaussian():
     D_KL = log_volume_n_ball(ndims) - ndims/2*np.log(2*np.pi*np.e*sigma**2)
     assert_allclose(samples.D_KL(), D_KL, atol=3*samples.D_KL(12).std())
 
+    assert_allclose(samples.d_G(), ndims, atol=3*samples.d_G(12).std())
+
     logZ = ndims/2 * np.log(2*np.pi*sigma**2) - log_volume_n_ball(ndims, R)
     assert_allclose(samples.logZ(), logZ, atol=3*samples.logZ(12).std())
 

From 80d68062bdc0675c9920c8be087660224117bb47 Mon Sep 17 00:00:00 2001
From: Will Handley <wh260@cam.ac.uk>
Date: Wed, 25 Jan 2023 10:17:09 +0000
Subject: [PATCH 18/19] Corrected parenthesis

---
 anesthetic/examples/utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/anesthetic/examples/utils.py b/anesthetic/examples/utils.py
index 2ace401a..4eab80af 100644
--- a/anesthetic/examples/utils.py
+++ b/anesthetic/examples/utils.py
@@ -46,7 +46,7 @@ def random_covariance(sigmas):
     Returns
     -------
     Covariance matrix: 2d np.array
-        shape (len(sigmas, len(sigmas)).
+        shape (len(sigmas), len(sigmas)).
     """
     d = len(sigmas)
     R = special_ortho_group.rvs(d)

From 62d2c17ebf77be2916b66d34ed592bf6b58c2967 Mon Sep 17 00:00:00 2001
From: Will Handley <wh260@cam.ac.uk>
Date: Wed, 25 Jan 2023 10:27:08 +0000
Subject: [PATCH 19/19] Updated ellipsoid docstring to mirror multivariate
 normal

---
 anesthetic/examples/utils.py | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/anesthetic/examples/utils.py b/anesthetic/examples/utils.py
index 4eab80af..6615a36d 100644
--- a/anesthetic/examples/utils.py
+++ b/anesthetic/examples/utils.py
@@ -21,11 +21,19 @@ def random_ellipsoid(mean, cov, size=None):
         The covariance structure of the ellipsoid. Axes have lengths equal to
         the square root of the eigenvalues of this matrix.
 
+    size: int or tuple of ints, optional
+        Given a shape of, for example, (m,n,k), m*n*k samples are generated,
+        and packed in an m-by-n-by-k arrangement. Because each sample is
+        N-dimensional, the output shape is (m,n,k,N). If no shape is specified,
+        a single (N-D) sample is returned.
+
     Returns
     -------
     points: array-like
-        Points drawn uniformly from the ellipsoid.
-        shape (*size, len(mean)).
+        The drawn samples, of shape size, if that was provided. If not, the
+        shape is (N,).
+        In other words, each entry out[i,j,...,:] is an N-dimensional value
+        drawn uniformly from the ellipsoid.
     """
     d = len(mean)
     L = np.linalg.cholesky(cov)