Exploring nested sampling (#807)

* subclass Reparam

* add some examples

* add diagnostic

* update notebook with resampling

* restrict the domain

* compare with mcmc results

* update a working multimodal result

* temp save

* add example and documentation

* fix typo

* fix typos and add to toc

* add contrib.rst

* cleanup TruncatedNormal/Cauchy and Uniform

* make some changes for new jaxns api, not working yet

* fix some typos

* fix typo at Uniform logprob

* make truncated distribution tests pass

* fix lint

* fix failing tests

* update for the new api

* add icdf method to tensorflow

* fix reparam batch logic

* fix event dim of projectednormalreparam

* add missing icdf methods

* enum is working

* add gaussian shell example

* add tests for gaussian shells

* run make license

* add jaxns to the dependency

* add EXPERIMENTAL warning

* fix docs

* temporary save

* add tests for nested sampling

* adjust precision

* fix tests

* make format

* add jaxns to docs requirement

* pump jaxns version

* use jaxns default arguments

* increase threshold of discretehmcgibbs

* add a better docstring for diagnostics

* increase threshold for hmcgibss in x64

* monkeypatch jaxns x64 dtype

* make sure that float32 can be used

* perform optimization

* avoid singularity at 0

* add note about enable x64

* add print summary

* address comments during pair review

* revert change at tfp

* restrict tfp version

Author: fehiepsi

docs/requirements.txt

@@ -3,6 +3,7 @@ flax
 funsor
 jax>=0.1.65
 jaxlib>=0.1.45
+jaxns==0.0.7
 optax==0.0.6
 nbsphinx>=0.8.5
 sphinx-gallery

docs/source/_static/img/examples/gaussian_shells.png

@@ -37,3 +37,12 @@ Inference
    optimizers
    diagnostics
    utilities
+
+Contributed Code
+----------------
+
+.. toctree::
+   :glob:
+   :maxdepth: 1
+
+   contrib

docs/source/api.rst

@@ -0,0 +1,8 @@
+Nested Sampling
+~~~~~~~~~~~~~~~
+
+.. autoclass:: numpyro.contrib.nested_sampling.NestedSampler
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource

docs/source/contrib.rst

@@ -37,6 +37,7 @@ NumPyro documentation
    examples/annotation
    examples/hmm_enum
    examples/capture_recapture
+   examples/gaussian_shells
    tutorials/discrete_imputation
 
 .. nbgallery::

docs/source/index.rst

@@ -0,0 +1,138 @@
+# Copyright Contributors to the Pyro project.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Example: Nested Sampling for Gaussian Shells
+============================================
+
+This example illustrates the usage of the contrib class NestedSampler,
+which is a wrapper of `jaxns` library ([1]) to be used for NumPyro models.
+
+Here we will replicate the Gaussian Shells demo at [2] and compare against
+NUTS sampler.
+
+**References:**
+
+    1. jaxns library: https://github.com/Joshuaalbert/jaxns
+    2. dynesty's Gaussian Shells demo:
+       https://github.com/joshspeagle/dynesty/blob/master/demos/Examples%20--%20Gaussian%20Shells.ipynb
+"""
+
+import argparse
+
+import matplotlib.pyplot as plt
+
+from jax import random
+import jax.numpy as jnp
+
+import numpyro
+from numpyro.contrib.nested_sampling import NestedSampler
+import numpyro.distributions as dist
+from numpyro.infer import MCMC, NUTS, DiscreteHMCGibbs
+
+
+class GaussianShell(dist.Distribution):
+    support = dist.constraints.real_vector
+
+    def __init__(self, loc, radius, width):
+        self.loc, self.radius, self.width = loc, radius, width
+        super().__init__(batch_shape=loc.shape[:-1], event_shape=loc.shape[-1:])
+
+    def sample(self, key, sample_shape=()):
+        return jnp.zeros(
+            sample_shape + self.shape()
+        )  # a dummy sample to initialize the samplers
+
+    def log_prob(self, value):
+        normalizer = (-0.5) * (jnp.log(2.0 * jnp.pi) + 2.0 * jnp.log(self.width))
+        d = jnp.linalg.norm(value - self.loc, axis=-1)
+        return normalizer - 0.5 * ((d - self.radius) / self.width) ** 2
+
+
+def model(center1, center2, radius, width, enum=False):
+    z = numpyro.sample(
+        "z", dist.Bernoulli(0.5), infer={"enumerate": "parallel"} if enum else {}
+    )
+    x = numpyro.sample("x", dist.Uniform(-6.0, 6.0).expand([2]).to_event(1))
+    center = jnp.stack([center1, center2])[z]
+    numpyro.sample("shell", GaussianShell(center, radius, width), obs=x)
+
+
+def run_inference(args, data):
+    print("=== Performing Nested Sampling ===")
+    ns = NestedSampler(model)
+    ns.run(random.PRNGKey(0), **data, enum=args.enum)
+    ns.print_summary()
+    # samples obtained from nested sampler are weighted, so
+    # we need to provide random key to resample from those weighted samples
+    ns_samples = ns.get_samples(random.PRNGKey(1), num_samples=args.num_samples)
+
+    print("\n=== Performing MCMC Sampling ===")
+    if args.enum:
+        mcmc = MCMC(
+            NUTS(model), num_warmup=args.num_warmup, num_samples=args.num_samples
+        )
+    else:
+        mcmc = MCMC(
+            DiscreteHMCGibbs(NUTS(model)),
+            num_warmup=args.num_warmup,
+            num_samples=args.num_samples,
+        )
+    mcmc.run(random.PRNGKey(2), **data)
+    mcmc.print_summary()
+    mcmc_samples = mcmc.get_samples()
+
+    return ns_samples["x"], mcmc_samples["x"]
+
+
+def main(args):
+    data = dict(
+        radius=2.0,
+        width=0.1,
+        center1=jnp.array([-3.5, 0.0]),
+        center2=jnp.array([3.5, 0.0]),
+    )
+    ns_samples, mcmc_samples = run_inference(args, data)
+
+    # plotting
+    fig, (ax1, ax2) = plt.subplots(
+        2, 1, sharex=True, figsize=(8, 8), constrained_layout=True
+    )
+
+    ax1.plot(mcmc_samples[:, 0], mcmc_samples[:, 1], "ro", alpha=0.2)
+    ax1.set(
+        xlim=(-6, 6),
+        ylim=(-2.5, 2.5),
+        ylabel="x[1]",
+        title="Gaussian-shell samples using NUTS",
+    )
+
+    ax2.plot(ns_samples[:, 0], ns_samples[:, 1], "ro", alpha=0.2)
+    ax2.set(
+        xlim=(-6, 6),
+        ylim=(-2.5, 2.5),
+        xlabel="x[0]",
+        ylabel="x[1]",
+        title="Gaussian-shell samples using Nested Sampler",
+    )
+
+    plt.savefig("gaussian_shells_plot.pdf")
+
+
+if __name__ == "__main__":
+    assert numpyro.__version__.startswith("0.6.0")
+    parser = argparse.ArgumentParser(description="Nested sampler for Gaussian shells")
+    parser.add_argument("-n", "--num-samples", nargs="?", default=10000, type=int)
+    parser.add_argument("--num-warmup", nargs="?", default=1000, type=int)
+    parser.add_argument(
+        "--enum",
+        action="store_true",
+        default=False,
+        help="whether to enumerate over the discrete latent variable",
+    )
+    parser.add_argument("--device", default="cpu", type=str, help='use "cpu" or "gpu".')
+    args = parser.parse_args()
+
+    numpyro.set_platform(args.device)
+
+    main(args)