diff --git a/.github/workflows/pytest_and_autopublish.yml b/.github/workflows/pytest_and_autopublish.yml index 7ba0ca7..d6cef97 100644 --- a/.github/workflows/pytest_and_autopublish.yml +++ b/.github/workflows/pytest_and_autopublish.yml @@ -16,8 +16,9 @@ jobs: with: python-version: '3.10' + - run: sudo apt install -y pandoc gsfonts - run: pip --version - - run: pip install -e .[dev,oryx] + - run: pip install -e .[dev,doc,oryx] - run: pip freeze - name: Lint with pylint @@ -26,6 +27,8 @@ jobs: run: pyink --check . - name: Lint with isort run: isort --check . + - name: Build documentation + run: make docs pytest-job: needs: lint diff --git a/.readthedocs.yaml b/.readthedocs.yaml new file mode 100644 index 0000000..856072c --- /dev/null +++ b/.readthedocs.yaml @@ -0,0 +1,17 @@ +version: 2 + +build: + os: ubuntu-22.04 + tools: + python: "3.9" + +sphinx: + configuration: docs/conf.py + +formats: + - pdf + - epub + +python: + install: + - requirements: docs/requirements.txt diff --git a/Makefile b/Makefile index 0c16a97..2f6b0ba 100644 --- a/Makefile +++ b/Makefile @@ -15,4 +15,7 @@ lint: FORCE test: lint FORCE pytest -vv -n auto +docs: FORCE + $(MAKE) -C docs html + FORCE: diff --git a/README.md b/README.md index 1e82492..7f5ae4f 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,7 @@ # coix [![Unittests](https://github.com/jax-ml/coix/actions/workflows/pytest_and_autopublish.yml/badge.svg)](https://github.com/jax-ml/coix/actions/workflows/pytest_and_autopublish.yml) +[![Documentation Status](https://readthedocs.org/projects/coix/badge/?version=latest)](https://coix.readthedocs.io/en/latest/?badge=latest) [![PyPI version](https://badge.fury.io/py/coix.svg)](https://badge.fury.io/py/coix) Inference Combinators in JAX (Coix) is a machine learning framework used to diff --git a/coix/api.py b/coix/api.py index f0f7518..f2e24be 100644 --- a/coix/api.py +++ b/coix/api.py @@ -16,6 +16,7 @@ The implement is pretty much backend-agnostic. We just assume that the core backend supports the following functionality: + + `suffix(p)`: rename latent variables of the program `p`, + `traced_evaluate(p, latents=None)`: execute `p` and collect trace, metrics, optionally we can substitute values in `latents` to `p`, @@ -41,14 +42,14 @@ def compose(q2, q1, suffix=True): - """Executes q2(*q1(...)). + r"""Executes q2(\*q1(...)). Note: We only allow at most one of `q1` or `q2` is weighted. Args: q2: a program q1: a program - suffix: whether to add suffix `_PREV_` to variables in `q1` + suffix: whether to add suffix `\_PREV\_` to variables in `q1` Returns: q: the composed program @@ -62,7 +63,7 @@ def wrapped(*args, **kwargs): def extend(p, f): - """Executes f(*p(...)) with random variables in f marked as auxiliary. + r"""Executes f(\*p(...)) with random variables in f marked as auxiliary. Note: We don't allow recursively marginalize out `p` yet. diff --git a/docs/Makefile b/docs/Makefile new file mode 100644 index 0000000..80bb986 --- /dev/null +++ b/docs/Makefile @@ -0,0 +1,27 @@ +# Minimal makefile for Sphinx documentation +# + +# You can set these variables from the command line. +SPHINXOPTS = -W +SPHINXBUILD = sphinx-build +SPHINXPROJ = coix +SOURCEDIR = . +BUILDDIR = build + +# Put it first so that "make" without argument is like "make help". +help: + @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + +.PHONY: help Makefile + +# Catch-all target: route all unknown targets to Sphinx using the new +# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). +%: Makefile + git clean -dfx build + git clean -dfx examples + git clean -dfx tutorials + git clean -f getting_started.rst + @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + git clean -dfx examples + git clean -dfx tutorials + git clean -f getting_started.rst diff --git a/docs/_templates/breadcrumbs.html b/docs/_templates/breadcrumbs.html new file mode 100644 index 0000000..ea007dd --- /dev/null +++ b/docs/_templates/breadcrumbs.html @@ -0,0 +1,28 @@ +{%- extends "sphinx_rtd_theme/breadcrumbs.html" %} + +{% set display_vcs_links = display_vcs_links if display_vcs_links is defined else True %} + +{% block breadcrumbs_aside %} +
  • + {% if hasdoc(pagename) and display_vcs_links %} + {% if display_github %} + {% if check_meta and 'github_url' in meta %} + + {{ _('Edit on GitHub') }} + {% else %} + {% if 'examples/index' in pagename %} + {{ _('Edit on GitHub') }} + {% elif 'examples/' in pagename %} + {{ _('Edit on GitHub') }} + {% else %} + {{ _('Edit on GitHub') }} + {% endif %} + {% endif %} + {% elif show_source and source_url_prefix %} + {{ _('View page source') }} + {% elif show_source and has_source and sourcename %} + {{ _('View page source') }} + {% endif %} + {% endif %} +
    • +{% endblock %} diff --git a/docs/algo.rst b/docs/algo.rst new file mode 100644 index 0000000..5ed2127 --- /dev/null +++ b/docs/algo.rst @@ -0,0 +1,5 @@ +Inference algorithms +==================== + +.. automodule:: coix.algo + :members: diff --git a/docs/api.rst b/docs/api.rst new file mode 100644 index 0000000..df9e498 --- /dev/null +++ b/docs/api.rst @@ -0,0 +1,5 @@ +Program combinators +=================== + +.. automodule:: coix.api + :members: diff --git a/docs/conf.py b/docs/conf.py new file mode 100644 index 0000000..7da7159 --- /dev/null +++ b/docs/conf.py @@ -0,0 +1,316 @@ +import glob +import os +import shutil +import sys + +import nbsphinx +import sphinx_rtd_theme + +# import pkg_resources + +# -*- coding: utf-8 -*- +# +# Configuration file for the Sphinx documentation builder. +# +# This file does only contain a selection of the most common options. For a +# full list see the documentation: +# http://www.sphinx-doc.org/en/master/config + +# -- Path setup -------------------------------------------------------------- + +# If extensions (or modules to document with autodoc) are in another directory, +# add these directories to sys.path here. If the directory is relative to the +# documentation root, use os.path.abspath to make it absolute, like shown here. +# +sys.path.insert(0, os.path.abspath("..")) + + +os.environ["SPHINX_BUILD"] = "1" + +# -- Project information ----------------------------------------------------- + +project = "coix" +copyright = "The coix Authors" +author = "The coix Authors" + +version = "" + +if "READTHEDOCS" not in os.environ: + # if developing locally, use coix.__version__ as version + from coix import __version__ # noqaE402 + + version = __version__ + + # Add "Edit on GitHub" button on the upper right corner of local docs. + html_context = {"github_version": "main", "display_github": True} + +# release version +release = version + + +# -- General configuration --------------------------------------------------- + +# If your documentation needs a minimal Sphinx version, state it here. +# +# needs_sphinx = '1.0' + +# Add any Sphinx extension module names here, as strings. They can be +# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom +# ones. +extensions = [ + "nbsphinx", + "sphinx.ext.autodoc", + "sphinx.ext.doctest", + "sphinx.ext.imgconverter", + "sphinx.ext.intersphinx", + "sphinx.ext.mathjax", + "sphinx.ext.napoleon", + "sphinx.ext.viewcode", + "sphinx_gallery.gen_gallery", + "sphinx_search.extension", +] + +# Enable documentation inheritance + +autodoc_inherit_docstrings = True + +# autodoc_default_options = { +# 'member-order': 'bysource', +# 'show-inheritance': True, +# 'special-members': True, +# 'undoc-members': True, +# 'exclude-members': '__dict__,__module__,__weakref__', +# } + +# Add any paths that contain templates here, relative to this directory. +templates_path = ["_templates"] + +# The suffix(es) of source filenames. +# You can specify multiple suffix as a list of string: +# +# source_suffix = ['.rst', '.md'] +# NOTE: `.rst` is the default suffix of sphinx, and nbsphinx will +# automatically add support for `.ipynb` suffix. + +# 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" + +# The language for content autogenerated by Sphinx. Refer to documentation +# for a list of supported languages. +# +# This is also used if you do content translation via gettext catalogs. +# Usually you set "language" from the command line for these cases. +language = "en" + +# 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 = [ + ".ipynb_checkpoints", + "examples/*ipynb", + "examples/*py", +] + +# The name of the Pygments (syntax highlighting) style to use. +pygments_style = "sphinx" + + +# do not prepend module name to functions +add_module_names = False + + +# This is processed by Jinja2 and inserted before each notebook +nbsphinx_prolog = r""" +{% set docname = 'notebooks/' + env.doc2path(env.docname, base=None).split('/')[-1] %} +:github_url: https://github.com/jax-ml/coix/blob/main/{{ docname }} + +.. raw:: html + +
    + Interactive online version: + + + Open In Colab + + +
    +""" # noqa: E501 + + +# -- Copy README files + +# replace "# coix" by "# Getting Started with Coix" +with open("../README.md", "rt") as f: + lines = f.readlines() + start_line = 0 + for i, line in enumerate(lines): + start_line = i + if "# coix" == line.rstrip(): + break + lines = lines[start_line:] + lines[0] = "# Getting Started with Coix\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/*.ipynb"): + shutil.copy(src_file, "tutorials/") + +# add index file to `notebooks` path, `:orphan:` is used to +# tell sphinx that this rst file needs not to be appeared in toctree +# with open("../notebooks/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": "__init__", + # not display Total running time of the script because we do not execute it + "min_reported_time": 1, +} + + +# -- Add thumbnails images + +nbsphinx_thumbnails = {} + +for src_file in glob.glob("../notebooks/*.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 Coix logo if not exist png file + if not os.path.exists(png_path): + png_path = "_static/img/coix_logo_wide.png" + nbsphinx_thumbnails[toctree_path + filename] = png_path + + +# -- Options for HTML output ------------------------------------------------- + +# logo +# html_logo = "_static/img/coix_logo_wide.png" + +# logo +# html_favicon = "_static/img/favicon/favicon.ico" + +# The theme to use for HTML and HTML Help pages. See the documentation for +# a list of builtin themes. +# +html_theme = "sphinx_rtd_theme" +html_theme_path = [sphinx_rtd_theme.get_html_theme_path()] + +# Theme options are theme-specific and customize the look and feel of a theme +# further. For a list of options available for each theme, see the +# documentation. +# +# html_theme_options = {} + +# Add any paths that contain custom static files (such as style sheets) here, +# relative to this directory. They are copied after the builtin static files, +# so a file named "default.css" will overwrite the builtin "default.css". +# html_static_path = ["_static"] +# html_style = "css/coix.css" + +# Custom sidebar templates, must be a dictionary that maps document names +# to template names. +# +# The default sidebars (for documents that don't match any pattern) are +# defined by theme itself. Builtin themes are using these templates by +# default: ``['localtoc.html', 'relations.html', 'sourcelink.html', +# 'searchbox.html']``. +# +# html_sidebars = {} + + +# -- Options for HTMLHelp output --------------------------------------------- + +# Output file base name for HTML help builder. +htmlhelp_basename = "coixdoc" + + +# -- Options for LaTeX output ------------------------------------------------ + +latex_elements = { + # The paper size ('letterpaper' or 'a4paper'). + # + # 'papersize': 'letterpaper', + # The font size ('10pt', '11pt' or '12pt'). + # + # 'pointsize': '10pt', + # Additional stuff for the LaTeX preamble. + # + "preamble": r""" + \usepackage{pmboxdraw} + \usepackage{alphabeta} + """, + # Latex figure (float) alignment + # + # 'figure_align': 'htbp', +} + +# Grouping the document tree into LaTeX files. List of tuples +# (source start file, target name, title, +# author, documentclass [howto, manual, or own class]). +latex_documents = [( + master_doc, + "coix.tex", + "Coix Documentation", + "The coix Authors", + "manual", +)] + +# -- Options for manual page output ------------------------------------------ + +# One entry per manual page. List of tuples +# (source start file, name, description, authors, manual section). +man_pages = [(master_doc, "Coix", "Coix Documentation", [author], 1)] + +# -- Options for Texinfo output ---------------------------------------------- + +# Grouping the document tree into Texinfo files. List of tuples +# (source start file, target name, title, author, +# dir menu entry, description, category) +texinfo_documents = [( + master_doc, + "Coix", + "Coix Documentation", + author, + "Coix", + "Inference Combinators in JAX", + "Miscellaneous", +)] + + +# -- Extension configuration ------------------------------------------------- + +# -- Options for intersphinx extension --------------------------------------- + +# Example configuration for intersphinx: refer to the Python standard library. +intersphinx_mapping = { + "python": ("https://docs.python.org/3/", None), + "numpy": ("http://docs.scipy.org/doc/numpy/", None), + "jax": ("https://jax.readthedocs.io/en/latest/", None), + "numpyro": ("http://num.pyro.ai/en/stable/", None), +} diff --git a/docs/core.rst b/docs/core.rst new file mode 100644 index 0000000..95ad347 --- /dev/null +++ b/docs/core.rst @@ -0,0 +1,5 @@ +Program transforms +================== + +.. automodule:: coix.core + :members: diff --git a/docs/index.rst b/docs/index.rst new file mode 100644 index 0000000..232b1ce --- /dev/null +++ b/docs/index.rst @@ -0,0 +1,39 @@ +:github_url: https://github.com/jax-ml/coix + + +Coix Documentation +================== + +.. toctree:: + :maxdepth: 1 + + getting_started + +.. toctree:: + :maxdepth: 1 + :caption: API and Developer Reference + + algo + api + core + loss + util + Change Log + +.. nbgallery:: + :maxdepth: 1 + :caption: Tutorials and Examples + :name: tutorials + + examples/anneal + examples/anneal_oryx + examples/gmm_oryx + examples/dmm_oryx + examples/bmnist + +Indices and tables +================== + +* :ref:`genindex` +* :ref:`modindex` +* :ref:`search` diff --git a/docs/loss.rst b/docs/loss.rst new file mode 100644 index 0000000..be5cbbd --- /dev/null +++ b/docs/loss.rst @@ -0,0 +1,5 @@ +Inference objectives +==================== + +.. automodule:: coix.loss + :members: diff --git a/docs/requirements.txt b/docs/requirements.txt new file mode 100644 index 0000000..da19975 --- /dev/null +++ b/docs/requirements.txt @@ -0,0 +1,15 @@ +ipython +jax +jaxlib +Jinja2 +matplotlib +nbsphinx +numpy +numpyro +pillow +pylab-sdk +pyyaml +readthedocs-sphinx-search +sphinx<5 +sphinx-gallery +sphinx_rtd_theme diff --git a/docs/util.rst b/docs/util.rst new file mode 100644 index 0000000..ad1d451 --- /dev/null +++ b/docs/util.rst @@ -0,0 +1,5 @@ +Utilities +========= + +.. automodule:: coix.util + :members: diff --git a/examples/README.rst b/examples/README.rst new file mode 100644 index 0000000..0acd6e4 --- /dev/null +++ b/examples/README.rst @@ -0,0 +1,10 @@ +Code Examples +============= + +Examples for Coix. + +`View source files on github`__ + +.. _github: https://github.com/jax-ml/coix/tree/main/examples + +__ github_ diff --git a/examples/anneal.py b/examples/anneal.py index b4e675a..6a53a72 100644 --- a/examples/anneal.py +++ b/examples/anneal.py @@ -12,6 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +""" +Example: Anneal example in NumPyro +================================== +""" + import argparse from functools import partial diff --git a/examples/anneal_oryx.py b/examples/anneal_oryx.py index 593dda3..91343c2 100644 --- a/examples/anneal_oryx.py +++ b/examples/anneal_oryx.py @@ -12,6 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +""" +Example: Anneal example in Oryx +=============================== +""" + import argparse from functools import partial diff --git a/examples/bmnist.py b/examples/bmnist.py index 1cdfb3d..6f9ffa3 100644 --- a/examples/bmnist.py +++ b/examples/bmnist.py @@ -12,6 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +""" +BMNIST example in NumPyro +========================= +""" + # TODO: Refactor using numpyro backend. The current code is likely not working yet. import argparse diff --git a/examples/dmm_oryx.py b/examples/dmm_oryx.py index 2b27436..95d2e64 100644 --- a/examples/dmm_oryx.py +++ b/examples/dmm_oryx.py @@ -12,6 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +""" +Example: DMM example in Oryx +============================ +""" + import argparse from functools import partial diff --git a/examples/gmm_oryx.py b/examples/gmm_oryx.py index 68a7b04..5adc609 100644 --- a/examples/gmm_oryx.py +++ b/examples/gmm_oryx.py @@ -1,3 +1,22 @@ +# Copyright 2023 The coix Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Example: GMM example in Oryx +============================ +""" + import argparse from functools import partial diff --git a/pyproject.toml b/pyproject.toml index 8c9bc63..b6b2259 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,7 +32,7 @@ dynamic = ["version"] homepage = "https://github.com/jax-ml/coix" repository = "https://github.com/jax-ml/coix" changelog = "https://github.com/jax-ml/coix/blob/main/CHANGELOG.md" -# TODO(phandu): documentation = "" +documentation = "https://coix.readthedocs.io" [project.optional-dependencies] # Development deps (unittest, linting, formating,...) @@ -48,6 +48,14 @@ dev = [ "pylint>=2.6.0", "pyink", ] +doc = [ + "ipython", + "nbsphinx", + "readthedocs-sphinx-search", + "sphinx<5", + "sphinx_rtd_theme", + "sphinx-gallery", +] oryx = [ "oryx@git+https://github.com/jax-ml/oryx", ]