Add example gallery

Author: jessegrabowski

doc/conf.py

@@ -21,6 +21,8 @@
 import sys
 import pytensor
 
+sys.path.insert(0, os.path.abspath(os.path.join("..", "scripts")))
+
 # General configuration
 # ---------------------
 
@@ -34,7 +36,9 @@
     "sphinx.ext.linkcode",
     "sphinx.ext.mathjax",
     "sphinx_design",
-    "sphinx.ext.intersphinx"
+    "sphinx.ext.intersphinx",
+    "myst_nb",
+    "generate_gallery",
 ]
 
 intersphinx_mapping = {
@@ -295,3 +299,30 @@ def find_source():
 
 # If false, no module index is generated.
 # latex_use_modindex = True
+
+
+# -- MyST config  -------------------------------------------------
+myst_enable_extensions = [
+    "colon_fence",
+    "deflist",
+    "dollarmath",
+    "amsmath",
+    "substitution",
+]
+myst_dmath_double_inline = True
+
+myst_substitutions = {
+    "pip_dependencies": "{{ extra_dependencies }}",
+    "conda_dependencies": "{{ extra_dependencies }}",
+    "extra_install_notes": "",
+}
+
+nb_execution_mode = "off"
+nbsphinx_execute = "never"
+nbsphinx_allow_errors = True
+
+
+# -- Bibtex config  -------------------------------------------------
+bibtex_bibfiles = ["references.bib"]
+bibtex_default_style = "unsrt"
+bibtex_reference_style = "author_year"

doc/gallery/introduction/what_is_pytensor.ipynb

@@ -80,6 +80,7 @@ Community
    introduction
    user_guide
    API <library/index>
+   Examples <gallery/gallery>
    Contributing <dev_start_guide>
 
 .. _Theano: https://github.com/Theano/Theano

doc/index.rst

@@ -43,6 +43,10 @@ dependencies:
   - ipython
   - pymc-sphinx-theme
   - sphinx-design
+  - myst-nb
+  - matplotlib
+  - watermark
+
   # code style
   - ruff
   # developer tools

environment.yml

@@ -0,0 +1,184 @@
+"""
+Sphinx plugin to run generate a gallery for notebooks
+
+Modified from the pymc project, which modified the seaborn project, which modified the mpld3 project.
+"""
+
+import base64
+import json
+import os
+import shutil
+from pathlib import Path
+
+import matplotlib
+
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import sphinx
+from matplotlib import image
+
+
+logger = sphinx.util.logging.getLogger(__name__)
+
+DOC_SRC = Path(__file__).resolve().parent
+# DEFAULT_IMG_LOC = os.path.join(os.path.dirname(DOC_SRC), "_static", "PyMC.png")
+
+DEFAULT_IMG_LOC = None
+external_nbs = {}
+
+HEAD = """
+Example Gallery
+===============
+
+.. toctree::
+   :hidden:
+
+"""
+
+SECTION_TEMPLATE = """
+.. _{section_id}:
+
+{section_title}
+{underlines}
+
+.. grid:: 1 2 3 3
+   :gutter: 4
+
+"""
+
+ITEM_TEMPLATE = """
+   .. grid-item-card:: :doc:`{doc_name}`
+      :img-top: {image}
+      :link: {doc_reference}
+      :link-type: {link_type}
+      :shadow: none
+"""
+
+folder_title_map = {
+    "introduction": "Introduction",
+}
+
+
+def create_thumbnail(infile, width=275, height=275, cx=0.5, cy=0.5, border=4):
+    """Overwrites `infile` with a new file of the given size"""
+    im = image.imread(infile)
+    rows, cols = im.shape[:2]
+    size = min(rows, cols)
+    if size == cols:
+        xslice = slice(0, size)
+        ymin = min(max(0, int(cx * rows - size // 2)), rows - size)
+        yslice = slice(ymin, ymin + size)
+    else:
+        yslice = slice(0, size)
+        xmin = min(max(0, int(cx * cols - size // 2)), cols - size)
+        xslice = slice(xmin, xmin + size)
+    thumb = im[yslice, xslice]
+    thumb[:border, :, :3] = thumb[-border:, :, :3] = 0
+    thumb[:, :border, :3] = thumb[:, -border:, :3] = 0
+
+    dpi = 100
+    fig = plt.figure(figsize=(width / dpi, height / dpi), dpi=dpi)
+
+    ax = fig.add_axes([0, 0, 1, 1], aspect="auto", frameon=False, xticks=[], yticks=[])
+    ax.imshow(thumb, aspect="auto", resample=True, interpolation="bilinear")
+    fig.savefig(infile, dpi=dpi)
+    plt.close(fig)
+    return fig
+
+
+class NotebookGenerator:
+    """Tools for generating an example page from a file"""
+
+    def __init__(self, filename, root_dir, folder):
+        self.folder = folder
+
+        self.basename = Path(filename).name
+        self.stripped_name = Path(filename).stem
+        self.image_dir = Path(root_dir) / "_thumbnails" / folder
+        self.png_path = self.image_dir / f"{self.stripped_name}.png"
+
+        with filename.open(encoding="utf-8") as fid:
+            self.json_source = json.load(fid)
+        self.default_image_loc = DEFAULT_IMG_LOC
+
+    def extract_preview_pic(self):
+        """By default, just uses the last image in the notebook."""
+        pic = None
+        for cell in self.json_source["cells"]:
+            for output in cell.get("outputs", []):
+                if "image/png" in output.get("data", []):
+                    pic = output["data"]["image/png"]
+        if pic is not None:
+            return base64.b64decode(pic)
+        return None
+
+    def gen_previews(self):
+        preview = self.extract_preview_pic()
+        if preview is not None:
+            with self.png_path.open("wb") as buff:
+                buff.write(preview)
+        else:
+            logger.warning(
+                f"Didn't find any pictures in {self.basename}",
+                type="thumbnail_extractor",
+            )
+            shutil.copy(self.default_image_loc, self.png_path)
+        create_thumbnail(self.png_path)
+
+
+def main(app):
+    logger.info("Starting thumbnail extractor.")
+
+    working_dir = Path.getcwd()
+    os.chdir(app.builder.srcdir)
+
+    file = [HEAD]
+
+    for folder, title in folder_title_map.items():
+        file.append(
+            SECTION_TEMPLATE.format(
+                section_title=title, section_id=folder, underlines="-" * len(title)
+            )
+        )
+
+        thumbnail_dir = Path("..") / "_thumbnails" / folder
+        if not thumbnail_dir.exists():
+            Path.mkdir(thumbnail_dir, parents=True)
+
+        if folder in external_nbs.keys():
+            file += [
+                ITEM_TEMPLATE.format(
+                    doc_name=descr["doc_name"],
+                    image=descr["image"],
+                    doc_reference=descr["doc_reference"],
+                    link_type=descr["link_type"],
+                )
+                for descr in external_nbs[folder]
+            ]
+
+        nb_paths = sorted(Path.glob(f"gallery/{folder}/*.ipynb"))
+
+        for nb_path in nb_paths:
+            nbg = NotebookGenerator(
+                filename=nb_path, root_dir=Path(".."), folder=folder
+            )
+            nbg.gen_previews()
+
+            file.append(
+                ITEM_TEMPLATE.format(
+                    doc_name=Path(folder) / nbg.stripped_name,
+                    image="/" + str(nbg.png_path),
+                    doc_reference=Path(folder) / nbg.stripped_name,
+                    link_type="doc",
+                )
+            )
+
+    with Path("gallery", "gallery.rst").open("w", encoding="utf-8") as f:
+        f.write("\n".join(file))
+
+    os.chdir(working_dir)
+
+
+def setup(app):
+    app.connect("builder-inited", main)