Perfect nested sampling functionality

anesthetic/examples/__init__.py

@@ -0,0 +1 @@
+"""Module for generating nested sampling examples."""

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')

anesthetic/examples/perfect_ns.py

@@ -0,0 +1,197 @@
+"""Perfect nested sampling generators."""
+import numpy as np
+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
+
+    Returns
+    -------
+    samples: NestedSamples
+        Nested sampling run
+    """
+
+    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, bounds=None):
+    """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, shape (ndims,)
+        mean of gaussian in parameters. Length of array defines dimensionality
+        of run.
+
+    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 -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.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.uniform(*bounds.T, (nlive, ndims))
+        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
+
+
+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)

anesthetic/examples/utils.py

@@ -0,0 +1,53 @@
+"""Utility functions for nested sampling examples."""
+import numpy as np
+from scipy.stats import special_ortho_group
+from scipy.special import gamma, gammaln
+
+
+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.
+    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.
+    """
+    """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):
+    """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(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

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)