Add thumbnail images to examples (#459)

Author: fehiepsi

.gitignore

@@ -2,6 +2,7 @@ numpyro.egg-info
 __pycache__/
 .ipynb_checkpoints/
 build
+tutorials/source/examples/
 
 # built / compiled
 *.pyc

docs/Makefile

@@ -17,4 +17,4 @@ help:
 # Catch-all target: route all unknown targets to Sphinx using the new
 # "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
 %: Makefile
-	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)

examples/README.rst

@@ -0,0 +1,10 @@
+Code Examples
+=============
+
+Examples on specifying models and doing inference in NumPyro.
+
+`View source files on github`__
+
+.. _github: https://github.com/pyro-ppl/numpyro/tree/master/examples
+
+__ github_

examples/baseball.py

@@ -1,18 +1,7 @@
-import argparse
-
-import numpy as onp
-
-import jax.numpy as np
-import jax.random as random
-from jax.scipy.special import logsumexp
-
-import numpyro
-import numpyro.distributions as dist
-from numpyro.examples.datasets import BASEBALL, load_dataset
-from numpyro.infer import MCMC, NUTS, Predictive, log_likelihood
-
-
 """
+Baseball
+========
+
 Original example from Pyro:
 https://github.com/pyro-ppl/pyro/blob/dev/examples/baseball.py
 
@@ -29,28 +18,44 @@
 from our models.
 
 Three models are evaluated:
- - Complete pooling model: The success probability of scoring a hit is shared
-     amongst all players.
- - No pooling model: Each individual player's success probability is distinct and
-     there is no data sharing amongst players.
- - Partial pooling model: A hierarchical model with partial data sharing.
 
+    - Complete pooling model: The success probability of scoring a hit is shared
+      amongst all players.
+    - No pooling model: Each individual player's success probability is distinct and
+      there is no data sharing amongst players.
+    - Partial pooling model: A hierarchical model with partial data sharing.
 
 We recommend Radford Neal's tutorial on HMC ([3]) to users who would like to get a
 more comprehensive understanding of HMC and its variants, and to [4] for details on
 the No U-Turn Sampler, which provides an efficient and automated way (i.e. limited
 hyper-parameters) of running HMC on different problems.
 
-[1] Carpenter B. (2016), ["Hierarchical Partial Pooling for Repeated Binary Trials"]
-    (http://mc-stan.org/users/documentation/case-studies/pool-binary-trials.html).
-[2] Efron B., Morris C. (1975), "Data analysis using Stein's estimator and its
-    generalizations", J. Amer. Statist. Assoc., 70, 311-319.
-[3] Neal, R. (2012), "MCMC using Hamiltonian Dynamics",
-    (https://arxiv.org/pdf/1206.1901.pdf)
-[4] Hoffman, M. D. and Gelman, A. (2014), "The No-U-turn sampler: Adaptively setting
-    path lengths in Hamiltonian Monte Carlo", (https://arxiv.org/abs/1111.4246)
+**References:**
+
+    1. Carpenter B. (2016), `"Hierarchical Partial Pooling for Repeated Binary Trials"
+       <http://mc-stan.org/users/documentation/case-studies/pool-binary-trials.html/>`_.
+    2. Efron B., Morris C. (1975), "Data analysis using Stein's estimator and its
+       generalizations", J. Amer. Statist. Assoc., 70, 311-319.
+    3. Neal, R. (2012), "MCMC using Hamiltonian Dynamics",
+       (https://arxiv.org/pdf/1206.1901.pdf)
+    4. Hoffman, M. D. and Gelman, A. (2014), "The No-U-turn sampler: Adaptively setting
+       path lengths in Hamiltonian Monte Carlo", (https://arxiv.org/abs/1111.4246)
 """
 
+import argparse
+import os
+
+import numpy as onp
+
+import jax.numpy as np
+import jax.random as random
+from jax.scipy.special import logsumexp
+
+import numpyro
+import numpyro.distributions as dist
+from numpyro.examples.datasets import BASEBALL, load_dataset
+from numpyro.infer import MCMC, NUTS, Predictive, log_likelihood
+
 
 def fully_pooled(at_bats, hits=None):
     r"""
@@ -125,7 +130,8 @@ def partially_pooled_with_logit(at_bats, hits=None):
 
 def run_inference(model, at_bats, hits, rng_key, args):
     kernel = NUTS(model)
-    mcmc = MCMC(kernel, args.num_warmup, args.num_samples, num_chains=args.num_chains)
+    mcmc = MCMC(kernel, args.num_warmup, args.num_samples, num_chains=args.num_chains,
+                progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True)
     mcmc.run(rng_key, at_bats, hits)
     return mcmc.get_samples()
 

examples/bnn.py

@@ -1,9 +1,13 @@
 """
+Bayesian Neural Network
+=======================
+
 We demonstrate how to use NUTS to do inference on a simple (small)
 Bayesian neural network with two hidden layers.
 """
 
 import argparse
+import os
 import time
 
 import matplotlib
@@ -58,7 +62,8 @@ def model(X, Y, D_H):
 def run_inference(model, args, rng_key, X, Y, D_H):
     start = time.time()
     kernel = NUTS(model)
-    mcmc = MCMC(kernel, args.num_warmup, args.num_samples, num_chains=args.num_chains)
+    mcmc = MCMC(kernel, args.num_warmup, args.num_samples, num_chains=args.num_chains,
+                progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True)
     mcmc.run(rng_key, X, Y, D_H)
     mcmc.print_summary()
     print('\nMCMC elapsed time:', time.time() - start)
@@ -124,7 +129,7 @@ def main(args):
     ax.set(xlabel="X", ylabel="Y", title="Mean predictions with 90% CI")
 
     plt.savefig('bnn_plot.pdf')
-    plt.close()
+    plt.tight_layout()
 
 
 if __name__ == "__main__":

examples/funnel.py

@@ -1,22 +1,9 @@
-import argparse
-
-import matplotlib.pyplot as plt
-import seaborn as sns
-
-from jax import random
-import jax.numpy as np
-
-import numpyro
-import numpyro.distributions as dist
-from numpyro.distributions.transforms import AffineTransform
-from numpyro.infer import MCMC, NUTS
-
-sns.set(context='talk')
-
-
 """
+Neal's Funnel
+=============
+
 This example, which is adapted from [1], illustrates how to leverage non-centered
-parameterization using the class `~numpyro.distributions.TransformedDistribution`.
+parameterization using the class :class:`numpyro.distributions.TransformedDistribution`.
 We will examine the difference between two types of parameterizations on the
 10-dimensional Neal's funnel distribution. As we will see, HMC gets trouble at
 the neck of the funnel if centered parameterization is used. On the contrary,
@@ -29,11 +16,26 @@
 inference algorithms know to do reparameterization automatically is to declare
 the random variable as a transformed distribution.
 
-[1] *Stan User's Guide*, https://mc-stan.org/docs/2_19/stan-users-guide/reparameterization-section.html
-[2] Maria I. Gorinova, Dave Moore, Matthew D. Hoffman (2019), "Automatic
-    Reparameterisation of Probabilistic Programs", (https://arxiv.org/abs/1906.03028)
+**References:**
+
+    1. *Stan User's Guide*, https://mc-stan.org/docs/2_19/stan-users-guide/reparameterization-section.html
+    2. Maria I. Gorinova, Dave Moore, Matthew D. Hoffman (2019), "Automatic
+       Reparameterisation of Probabilistic Programs", (https://arxiv.org/abs/1906.03028)
 """
 
+import argparse
+import os
+
+import matplotlib.pyplot as plt
+
+from jax import random
+import jax.numpy as np
+
+import numpyro
+import numpyro.distributions as dist
+from numpyro.distributions.transforms import AffineTransform
+from numpyro.infer import MCMC, NUTS
+
 
 def model(dim=10):
     y = numpyro.sample('y', dist.Normal(0, 3))
@@ -48,41 +50,45 @@ def reparam_model(dim=10):
 
 def run_inference(model, args, rng_key):
     kernel = NUTS(model)
-    mcmc = MCMC(kernel, args.num_warmup, args.num_samples, num_chains=args.num_chains)
+    mcmc = MCMC(kernel, args.num_warmup, args.num_samples, num_chains=args.num_chains,
+                progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True)
     mcmc.run(rng_key)
+    mcmc.print_summary()
     return mcmc.get_samples()
 
 
 def main(args):
     rng_key = random.PRNGKey(0)
 
     # do inference with centered parameterization
+    print("============================= Centered Parameterization ==============================")
     samples = run_inference(model, args, rng_key)
 
     # do inference with non-centered parameterization
+    print("\n=========================== Non-centered Parameterization ============================")
     reparam_samples = run_inference(reparam_model, args, rng_key)
 
     # make plots
-    fig, (ax1, ax2) = plt.subplots(2, 1, sharex=True, figsize=(8, 16))
+    fig, (ax1, ax2) = plt.subplots(2, 1, sharex=True, figsize=(6.4, 6.4))
 
-    sns.scatterplot(samples['x'][:, 0], samples['y'], color='g', alpha=0.3, ax=ax1)
-    ax1.set(xlim=(-20, 20), ylim=(-9, 9), xlabel='x[0]', ylabel='y',
+    ax1.plot(samples['x'][:, 0], samples['y'], "go", alpha=0.3)
+    ax1.set(xlim=(-20, 20), ylim=(-9, 9), ylabel='y',
             title='Funnel samples with centered parameterization')
 
-    sns.scatterplot(reparam_samples['x'][:, 0], reparam_samples['y'], color='g', alpha=0.3, ax=ax2)
+    ax2.plot(reparam_samples['x'][:, 0], reparam_samples['y'], "go", alpha=0.3)
     ax2.set(xlim=(-20, 20), ylim=(-9, 9), xlabel='x[0]', ylabel='y',
             title='Funnel samples with non-centered parameterization')
 
     plt.savefig('funnel_plot.pdf')
-    plt.close()
+    plt.tight_layout()
 
 
 if __name__ == "__main__":
     assert numpyro.__version__.startswith('0.2.1')
     parser = argparse.ArgumentParser(description="Non-centered reparameterization example")
     parser.add_argument("-n", "--num-samples", nargs="?", default=1000, type=int)
     parser.add_argument("--num-warmup", nargs='?', default=1000, type=int)
-    parser.add_argument("--num-chains", nargs='?', default=4, type=int)
+    parser.add_argument("--num-chains", nargs='?', default=1, type=int)
     parser.add_argument("--device", default='cpu', type=str, help='use "cpu" or "gpu".')
     args = parser.parse_args()
 

examples/gp.py

@@ -1,4 +1,13 @@
+"""
+Gaussian Process
+================
+
+In this example we show how to use NUTS to sample from the posterior
+over the hyperparameters of a gaussian process.
+"""
+
 import argparse
+import os
 import time
 
 import matplotlib
@@ -16,11 +25,6 @@
 
 matplotlib.use('Agg')  # noqa: E402
 
-"""
-In this example we show how to use NUTS to sample from the posterior
-over the hyperparameters of a gaussian process.
-"""
-
 
 # squared exponential kernel with diagonal noise term
 def kernel(X, Z, var, length, noise, jitter=1.0e-6, include_noise=True):
@@ -49,7 +53,8 @@ def model(X, Y):
 def run_inference(model, args, rng_key, X, Y):
     start = time.time()
     kernel = NUTS(model)
-    mcmc = MCMC(kernel, args.num_warmup, args.num_samples, num_chains=args.num_chains)
+    mcmc = MCMC(kernel, args.num_warmup, args.num_samples, num_chains=args.num_chains,
+                progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True)
     mcmc.run(rng_key, X, Y)
     mcmc.print_summary()
     print('\nMCMC elapsed time:', time.time() - start)
@@ -117,7 +122,7 @@ def main(args):
     ax.set(xlabel="X", ylabel="Y", title="Mean predictions with 90% CI")
 
     plt.savefig("gp_plot.pdf")
-    plt.close()
+    plt.tight_layout()
 
 
 if __name__ == "__main__":

examples/hmm.py

@@ -1,18 +1,7 @@
-import argparse
-import time
-
-import numpy as onp
-
-from jax import lax, random
-import jax.numpy as np
-from jax.scipy.special import logsumexp
-
-import numpyro
-import numpyro.distributions as dist
-from numpyro.infer import MCMC, NUTS
-
-
 """
+Hidden Markov Model
+===================
+
 In this example, we will follow [1] to construct a semi-supervised Hidden Markov
 Model for a generative model with observations are words and latent variables
 are categories. Instead of automatically marginalizing all discrete latent
@@ -26,12 +15,30 @@
 JAX's `lax.scan` primitive. The primitive will greatly improve compiling for the
 model.
 
-[1] https://mc-stan.org/docs/2_19/stan-users-guide/hmms-section.html
-[2] http://pyro.ai/examples/hmm.html
-[3] https://en.wikipedia.org/wiki/Forward_algorithm
-[4] https://discourse.pymc.io/t/how-to-marginalized-markov-chain-with-categorical/2230
+**References:**
+
+    1. https://mc-stan.org/docs/2_19/stan-users-guide/hmms-section.html
+    2. http://pyro.ai/examples/hmm.html
+    3. https://en.wikipedia.org/wiki/Forward_algorithm
+    4. https://discourse.pymc.io/t/how-to-marginalized-markov-chain-with-categorical/2230
 """
 
+import argparse
+import os
+import time
+
+import matplotlib.pyplot as plt
+import numpy as onp
+from scipy.stats import gaussian_kde
+
+from jax import lax, random
+import jax.numpy as np
+from jax.scipy.special import logsumexp
+
+import numpyro
+import numpyro.distributions as dist
+from numpyro.infer import MCMC, NUTS
+
 
 def simulate_data(rng_key, num_categories, num_words, num_supervised_data, num_unsupervised_data):
     rng_key, rng_key_transition, rng_key_emission = random.split(rng_key, 3)
@@ -148,13 +155,29 @@ def main(args):
     rng_key = random.PRNGKey(2)
     start = time.time()
     kernel = NUTS(semi_supervised_hmm)
-    mcmc = MCMC(kernel, args.num_warmup, args.num_samples)
+    mcmc = MCMC(kernel, args.num_warmup, args.num_samples,
+                progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True)
     mcmc.run(rng_key, transition_prior, emission_prior, supervised_categories,
              supervised_words, unsupervised_words)
     samples = mcmc.get_samples()
     print_results(samples, transition_prob, emission_prob)
     print('\nMCMC elapsed time:', time.time() - start)
 
+    # make plots
+    fig, ax = plt.subplots(1, 1)
+
+    x = onp.linspace(0, 1, 101)
+    for i in range(transition_prob.shape[0]):
+        for j in range(transition_prob.shape[1]):
+            ax.plot(x, gaussian_kde(samples['transition_prob'][:, i, j])(x),
+                    label="transition_prob[{}, {}], true value = {:.2f}"
+                    .format(i, j, transition_prob[i, j]))
+    ax.set(xlabel="Probability", ylabel="Frequency",
+           title="Transition probability posterior")
+
+    plt.savefig("hmm_plot.pdf")
+    plt.tight_layout()
+
 
 if __name__ == '__main__':
     assert numpyro.__version__.startswith('0.2.1')