Update documentation page (pyro-ppl#791)

.github/workflows/ci.yml

@@ -24,6 +24,7 @@ jobs:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
       run: |
+        sudo apt install -y pandoc
         python -m pip install --upgrade pip
         pip install https://github.com/pyro-ppl/funsor/archive/master.zip
         pip install jaxlib

.gitignore

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

docs/Makefile

@@ -17,4 +17,5 @@ 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)
+	git clean -dfx source/

docs/requirements.txt

@@ -1,7 +1,9 @@
+dm-haiku
+flax
+funsor
 jax>=0.1.65
 jaxlib>=0.1.45
-git+git://github.com/pyro-ppl/funsor@master#egg=funsor
-tqdm
+nbsphinx
+sphinx-gallery
 tfp_nightly  # TODO: change this to stable release
-flax
-dm-haiku
+tqdm

docs/source/_static/css/pyro.css

@@ -26,4 +26,12 @@
 
 .wy-side-nav-search>a img.logo, .wy-side-nav-search .wy-dropdown>a img.logo {
     max-width: 60%;
-}
+}
+
+.wy-nav-content {
+    max-width: 1000px !important;
+}
+
+.sphx-glr-thumbcontainer {
+    text-align: center !important;
+}

docs/source/api.rst

@@ -0,0 +1,39 @@
+.. currentmodule:: api
+
+API Reference
+=============
+
+Modeling
+--------
+
+.. toctree::
+   :glob:
+   :maxdepth: 1
+
+   primitives
+   handlers
+
+Distributions
+-------------
+
+.. toctree::
+   :glob:
+   :maxdepth: 1
+
+   distributions
+
+Inference
+---------
+
+.. toctree::
+   :glob:
+   :maxdepth: 1
+
+   mcmc
+   svi
+   autoguide
+   reparam
+   funsor
+   optimizers
+   diagnostics
+   utilities

docs/source/conf.py

@@ -1,9 +1,12 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
 
+import glob
 import os
+import shutil
 import sys
 
+import nbsphinx
 import sphinx_rtd_theme
 
 
@@ -41,7 +44,7 @@
 version = ''
 
 if 'READTHEDOCS' not in os.environ:
-    # if developing locally, use pyro.__version__ as version
+    # if developing locally, use numpyro.__version__ as version
     from numpyro import __version__  # noqaE402
     version = __version__
 
@@ -59,11 +62,13 @@
 # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
 # ones.
 extensions = [
+    'nbsphinx',
     'sphinx.ext.autodoc',
     'sphinx.ext.doctest',
     'sphinx.ext.intersphinx',
     'sphinx.ext.mathjax',
     'sphinx.ext.viewcode',
+    'sphinx_gallery.gen_gallery',
 ]
 
 # Enable documentation inheritance
@@ -85,7 +90,13 @@
 # You can specify multiple suffix as a list of string:
 #
 # source_suffix = ['.rst', '.md']
-source_suffix = '.rst'
+source_suffix = ['.rst', '.ipynb']
+
+# do not execute cells
+nbsphinx_execute = 'never'
+
+# Don't add .txt suffix to source files:
+html_sourcelink_suffix = ''
 
 # The master toctree document.
 master_doc = 'index'
@@ -100,7 +111,12 @@
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.
 # This pattern also affects html_static_path and html_extra_path .
-exclude_patterns = []
+exclude_patterns = [
+    '.ipynb_checkpoints',
+    'tutorials/logistic_regression.ipynb',
+    'examples/*ipynb',
+    'examples/*py'
+]
 
 # The name of the Pygments (syntax highlighting) style to use.
 pygments_style = 'sphinx'
@@ -109,6 +125,91 @@
 # do not prepend module name to functions
 add_module_names = False
 
+
+with open('../../numpyro/version.py') as f:
+    for line in f.readlines():
+        if line.startswith("__version__ = "):
+            break
+    pip_version = line.rstrip().split(" = ")[-1].strip("'").strip('"')
+    # resolve the issue: tags for versions before 0.3.0 have the prefix 'v'.
+    tag = "v" + pip_version if version[:3] < "0.3" else pip_version
+
+
+# This is processed by Jinja2 and inserted before each notebook
+nbsphinx_prolog = r"""
+{% set docname = 'notebooks/source/' + env.doc2path(env.docname, base=None).split('/')[-1] %}
+.. raw:: html
+
+    <div class="admonition note">
+      Interactive online version:
+      <span style="white-space: nowrap;">
+        <a href="https://colab.research.google.com/github/pyro-ppl/numpyro/blob/{tag}/{{ docname }}">
+          <img alt="Open In Colab" src="https://colab.research.google.com/assets/colab-badge.svg"
+            style="vertical-align:text-bottom">
+        </a>
+      </span>
+    </div>
+""".replace(r"{tag}", tag)
+
+
+# -- Copy README files
+
+# replace "# NumPyro" by "# Getting Started with NumPyro"
+with open('../../README.md', 'rt') as f:
+    lines = f.readlines()
+    for i, line in enumerate(lines):
+        if "# NumPyro" == line.rstrip():
+            break
+    lines = lines[i:]
+    lines[0] = "# Getting Started with NumPyro\n"
+    text = "\n".join(lines)
+
+with open('getting_started.rst', 'wt') as f:
+    f.write(nbsphinx.markdown2rst(text))
+
+
+# -- Copy notebook files
+
+if not os.path.exists('tutorials'):
+    os.makedirs('tutorials')
+
+for src_file in glob.glob('../../notebooks/source/*.ipynb'):
+    dst_file = os.path.join('tutorials', src_file.split("/")[-1])
+    shutil.copy(src_file, 'tutorials/')
+
+# add index file to `tutorials` path, `:orphan:` is used to
+# tell sphinx that this rst file needs not to be appeared in toctree
+with open('../../notebooks/source/index.rst', 'rt') as f1:
+    with open('tutorials/index.rst', 'wt') as f2:
+        f2.write(":orphan:\n\n")
+        f2.write(f1.read())
+
+
+# -- Convert scripts to notebooks
+
+sphinx_gallery_conf = {
+    'examples_dirs': ['../../examples'],
+    'gallery_dirs': ['examples'],
+    # only execute files beginning with plot_
+    'filename_pattern': '/plot_',
+    'ignore_pattern': '(minipyro|covtype|__init__)',
+}
+
+
+# -- Add thumbnails images
+
+nbsphinx_thumbnails = {}
+
+for src_file in (glob.glob('../../notebooks/source/*.ipynb') + glob.glob('../../examples/*.py')):
+    toctree_path = "tutorials/" if src_file.endswith("ipynb") else "examples/"
+    filename = os.path.splitext(src_file.split("/")[-1])[0]
+    png_path = "_static/img/" + toctree_path + filename + ".png"
+    # use Pyro logo if not exist png file
+    if not os.path.exists(png_path):
+        png_path = "_static/img/pyro_logo_wide.png"
+    nbsphinx_thumbnails[toctree_path + filename] = png_path
+
+
 # -- Options for HTML output -------------------------------------------------
 
 # logo

docs/source/index.rst

@@ -5,41 +5,58 @@ NumPyro documentation
 =====================
 
 .. toctree::
-   :glob:
-   :maxdepth: 2
-   :caption: Modeling:
-
-   primitives
-   handlers
+   :maxdepth: 1
 
+   getting_started
 
 .. toctree::
-   :glob:
-   :maxdepth: 2
-   :caption: Distributions:
-
-   distributions
-
-
-.. toctree::
-   :glob:
-   :maxdepth: 2
-   :caption: Inference:
-
-   mcmc
-   svi
-   autoguide
-   reparam
-   funsor
-   optimizers
-   diagnostics
-   utilities
+   :maxdepth: 1
+   :caption: API and Developer Reference
+
+   api
+   Change Log <https://github.com/pyro-ppl/numpyro/releases>
+
+.. nbgallery::
+   :maxdepth: 1
+   :caption: Introductory Tutorials
+   :name: introductory-tutorials
+
+   tutorials/bayesian_regression
+   tutorials/bayesian_hierarchical_linear_regression
+   examples/baseball
+   examples/vae
+   examples/funnel
+   examples/stochastic_volatility
+
+.. nbgallery::
+   :maxdepth: 1
+   :caption: Discrete Latent Variables
+   :name: discrete-latent-variables
+
+   examples/annotation
+   examples/hmm_enum
+   tutorials/discrete_imputation
+
+.. nbgallery::
+   :maxdepth: 1
+   :caption: Applications
+   :name: applications
+
+   tutorials/time_series_forecasting
+   tutorials/ordinal_regression
+   tutorials/bayesian_imputation
+   examples/gp
+   examples/bnn
+   examples/sparse_regression
+   examples/proportion_test
+   examples/ucbadmit
+   examples/ode
+   examples/neutra
 
 
 Indices and tables
 ==================
 
 * :ref:`genindex`
+* :ref:`modindex`
 * :ref:`search`
-
-.. * :ref:`modindex`

examples/annotation.py

@@ -2,8 +2,8 @@
 # SPDX-License-Identifier: Apache-2.0
 
 """
-Bayesian Models of Annotation
-=============================
+Example: Bayesian Models of Annotation
+======================================
 
 In this example, we run MCMC for various crowdsourced annotation models in [1].
 

examples/baseball.py

@@ -2,8 +2,8 @@
 # SPDX-License-Identifier: Apache-2.0
 
 """
-Baseball
-========
+Example: Baseball Batting Average
+=================================
 
 Original example from Pyro:
 https://github.com/pyro-ppl/pyro/blob/dev/examples/baseball.py
@@ -197,7 +197,7 @@ def main(args):
 
 if __name__ == "__main__":
     assert numpyro.__version__.startswith('0.4.1')
-    parser = argparse.ArgumentParser(description="Baseball batting average using HMC")
+    parser = argparse.ArgumentParser(description="Baseball batting average using MCMC")
     parser.add_argument("-n", "--num-samples", nargs="?", default=3000, type=int)
     parser.add_argument("--num-warmup", nargs='?', default=1500, type=int)
     parser.add_argument("--num-chains", nargs='?', default=1, type=int)

examples/bnn.py

@@ -2,8 +2,8 @@
 # SPDX-License-Identifier: Apache-2.0
 
 """
-Bayesian Neural Network
-=======================
+Example: Bayesian Neural Network
+================================
 
 We demonstrate how to use NUTS to do inference on a simple (small)
 Bayesian neural network with two hidden layers.

examples/funnel.py

@@ -2,8 +2,8 @@
 # SPDX-License-Identifier: Apache-2.0
 
 """
-Neal's Funnel
-=============
+Example: Neal's Funnel
+======================
 
 This example, which is adapted from [1], illustrates how to leverage non-centered
 parameterization using the class :class:`numpyro.distributions.TransformedDistribution`.

examples/gp.py

@@ -2,8 +2,8 @@
 # SPDX-License-Identifier: Apache-2.0
 
 """
-Gaussian Process
-================
+Example: Gaussian Process
+=========================
 
 In this example we show how to use NUTS to sample from the posterior
 over the hyperparameters of a gaussian process.