Add sphinx documentation (#22)

README.md

@@ -49,7 +49,7 @@ pip install git+https://github.com/dirmeier/sbijax@<RELEASE>
 
 ## Acknowledgements
 
-> 📝 The package draws significant inspiration from the excellent Pytorch-based [`sbi`](https://github.com/sbi-dev/sbi) package which is substantially more
+> 📝 The API of the package is heavily inspired by the excellent Pytorch-based [`sbi`](https://github.com/sbi-dev/sbi) package which is substantially more
 feature-complete and user-friendly, and better documented.
 
 ## Author

docs/.gitignore

@@ -0,0 +1,6 @@
+source/examples/
+source/examples/*
+build/
+build/*
+_autosummary/
+_autosummary/*

docs/Makefile

@@ -0,0 +1,22 @@
+# Minimal makefile for Sphinx documentation
+#
+
+# You can set these variables from the command line, and also
+# from the environment for the first two.
+SPHINXOPTS    ?=
+SPHINXBUILD   ?= sphinx-build
+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
+	rm -rf build
+	rm -rf source/examples
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)

docs/_static/theme.css

@@ -0,0 +1,48 @@
+html[data-theme="light"] {
+    --pst-color-primary: rgb(121, 40, 161);
+    --pst-color-primary-bg: #ffe9dd;
+    --pst-color-secondary: #b26679;
+    --pst-color-inline-code-links: #b26679;
+}
+
+h1 > code > span {
+    font-family: var(--pst-font-family-monospace);
+    font-weight: 700;
+}
+
+nav > li > a > code.literal {
+    padding-top: 0;
+    padding-bottom: 0;
+    background-color: white;
+    border: 0;
+}
+
+nav.bd-links p.caption {
+    text-transform: uppercase;
+}
+
+code.literal {
+  background-color: white;
+  border: 0;
+  border-radius: 0;
+}
+a:hover {
+      text-decoration-thickness: 1px !important;
+}
+
+
+ul.bd-breadcrumbs li.breadcrumb-item a:hover {
+  text-decoration-thickness: 1px;
+}
+
+nav.bd-links li > a:hover {
+  text-decoration-thickness: 1px;
+}
+
+.prev-next-area a p.prev-next-title {
+  text-decoration: none !important;
+}
+
+button.theme-switch-button {
+    display: none !important;
+}

docs/conf.py

@@ -0,0 +1,69 @@
+from datetime import date
+
+project = "sbijax"
+copyright = f"{date.today().year}, the sbijax developers"
+author = "the sbijax developers"
+
+extensions = [
+    "nbsphinx",
+    "sphinx.ext.autodoc",
+    'sphinx_autodoc_typehints',
+    "sphinx.ext.autosummary",
+    "sphinx.ext.doctest",
+    "sphinx.ext.intersphinx",
+    "sphinx.ext.mathjax",
+    "sphinx.ext.napoleon",
+    "sphinx.ext.viewcode",
+    "sphinx_autodoc_typehints",
+    "sphinx_copybutton",
+    "sphinx_math_dollar",
+    "IPython.sphinxext.ipython_console_highlighting",
+    'sphinx_design',
+]
+
+
+templates_path = ["_templates"]
+html_static_path = ["_static"]
+html_css_files = ['theme.css']
+
+autodoc_default_options = {
+    "member-order": "bysource",
+    "special-members": True,
+    "exclude-members": "__repr__, __str__, __weakref__",
+}
+
+exclude_patterns = [
+    "_build",
+    "build",
+    "Thumbs.db",
+    ".DS_Store",
+    "notebooks/.ipynb_checkpoints",
+    "examples/*ipynb",
+    "examples/*py"
+]
+
+autodoc_typehints = "both"
+
+html_theme = "sphinx_book_theme"
+
+html_theme_options = {
+    "repository_url": "https://github.com/dirmeier/sbijax",
+    "use_repository_button": True,
+    "use_download_button": False,
+    "extra_navbar": ""
+}
+
+html_title = "sbijax 🚀"
+
+
+def skip(app, what, name, obj, would_skip, options):
+    if name == "__init__":
+        return True
+    return would_skip
+
+
+def setup(app):
+    app.connect("autodoc-skip-member", skip)
+
+
+bibtex_bibfiles = ['references.bib']

docs/examples.rst

@@ -0,0 +1,7 @@
+More examples
+=============
+
+.. note::
+
+    Self-contained example code can be found on GitHub in `examples <https://github.com/dirmeier/sbijax/tree/main/examples>`_.
+    The examples are executable from the command line, so forking/cloning the code suffices to run them.

docs/index.rst

@@ -0,0 +1,89 @@
+:github_url: https://github.com/dirmeier/sbijax
+
+👋 Welcome to :code:`sbijax`!
+=============================
+
+.. div:: sd-text-left sd-font-italic
+
+    Simulation-based inference in JAX
+
+----
+
+:code:`sbijax` implements several algorithms for simulation-based inference in
+`JAX <https://github.com/google/jax>`_ using `Haiku <https://github.com/deepmind/dm-haiku>`_,
+`Distrax <https://github.com/deepmind/distrax>`_ and `BlackJAX <https://github.com/blackjax-devs/blackjax>`_. Specifically, :code:`sbijax` implements
+
+- `Sequential Monte Carlo ABC <https://www.routledge.com/Handbook-of-Approximate-Bayesian-Computation/Sisson-Fan-Beaumont/p/book/9780367733728>`_ (:code:`SMCABC`),
+- `Neural Likelihood Estimation <https://arxiv.org/abs/1805.07226>`_  (:code:`SNL`)
+- `Surjective Neural Likelihood Estimation <https://arxiv.org/abs/2308.01054>`_ (:code:`SSNL`)
+- `Neural Posterior Estimation C <https://arxiv.org/abs/1905.07488>`_ (short :code:`SNP`)
+- `Contrastive Neural Ratio Estimation <https://arxiv.org/abs/2210.06170>`_ (short :code:`SNR`)
+- `Neural Approximate Sufficient Statistics <https://arxiv.org/abs/2010.10079>`_ (:code:`SNASS`)
+- `Neural Approximate Slice Sufficient Statistics <https://openreview.net/forum?id=jjzJ768iV1>`_ (:code:`SNASSS`)
+
+Installation
+------------
+
+To install from PyPI, call:
+
+.. code-block:: bash
+
+    pip install sbijax
+
+To install the latest GitHub <RELEASE>, just call the following on the
+command line:
+
+.. code-block:: bash
+
+    pip install git+https://github.com/dirmeier/sbijax@<RELEASE>
+
+See also the installation instructions for `JAX <https://github.com/google/jax>`_, if
+you plan to use :code:`sbijax` on GPU/TPU.
+
+Contributing
+------------
+
+Contributions in the form of pull requests are more than welcome. A good way to start is to check out issues labelled
+`"good first issue" <https://github.com/dirmeier/sbijax/issues?q=is%3Aissue+is%3Aopen+label%3A%22good+first+issue%22>`_.
+
+In order to contribute:
+
+1) Clone :code:`sbijax` and install :code:`hatch` via :code:`pip install hatch`,
+2) create a new branch locally :code:`git checkout -b feature/my-new-feature` or :code:`git checkout -b issue/fixes-bug`,
+3) implement your contribution and ideally a test case,
+4) test it by calling :code:`hatch run test` on the (Unix) command line,
+5) submit a PR 🙂
+
+Acknowledgements
+----------------
+
+.. note::
+
+    📝 The API of the package is heavily inspired by the excellent Pytorch-based `sbi <https://github.com/sbi-dev/sbi>`_ package which is
+    substantially more feature-complete and user-friendly.
+
+License
+-------
+
+:code:`sbijax` is licensed under the Apache 2.0 License.
+
+..  toctree::
+    :maxdepth: 1
+    :hidden:
+
+    🏠 Home <self>
+
+..  toctree::
+    :caption: 🎓 Examples
+    :maxdepth: 1
+    :hidden:
+
+    Self-contained scripts <examples>
+
+..  toctree::
+    :caption: 🧱 API
+    :maxdepth: 2
+    :hidden:
+
+    sbijax
+    sbijax.nn

docs/requirements.txt

@@ -0,0 +1,20 @@
+-e .
+ipykernel
+ipython
+matplotlib
+nbsphinx
+pandas
+scikit-learn
+seaborn
+session_info
+sphinx
+sphinx-autobuild
+sphinx-book-theme>=1.1.0
+sphinx-copybutton
+sphinx-math-dollar
+sphinx_autodoc_typehints
+sphinx_design
+sphinx_fontawesome
+sphinx_gallery
+sphinxcontrib-fulltoc
+tqdm

docs/sbijax.nn.rst

@@ -0,0 +1,27 @@
+``sbijax.nn``
+=============
+
+.. currentmodule:: sbijax.nn
+
+----
+
+``sbijax.nn`` contains utility functions and classes to construct neural
+networks and normalizing flows.
+
+.. autosummary::
+    make_affine_maf
+    make_surjective_affine_maf
+    make_resnet
+    make_snass_net
+    make_snasss_net
+
+
+.. autofunction:: make_affine_maf
+
+.. autofunction:: make_surjective_affine_maf
+
+.. autofunction:: make_resnet
+
+.. autofunction:: make_snass_net
+
+.. autofunction:: make_snasss_net

docs/sbijax.rst

@@ -0,0 +1,55 @@
+``sbijax``
+==========
+
+.. currentmodule:: sbijax
+
+----
+
+:code:`sbijax` contains the implemented methods for simulation-based inference.
+
+Methods
+-------
+
+.. autosummary::
+    SMCABC
+    SNL
+    SNP
+    SNR
+    SNASS
+    SNASSS
+
+SMCABC
+~~~~~~
+
+..  autoclass:: SMCABC
+    :members: fit, simulate_data_and_possibly_append, sample_posterior
+
+SNL+SSNL
+~~~~~~~~
+
+..  autoclass:: SNL
+    :members: fit, simulate_data_and_possibly_append, sample_posterior
+
+SNP
+~~~
+
+..  autoclass:: SNP
+    :members: fit, simulate_data_and_possibly_append, sample_posterior
+
+SNR
+~~~
+
+..  autoclass:: SNR
+    :members: fit, simulate_data_and_possibly_append, sample_posterior
+
+SNASS
+~~~~~
+
+..  autoclass:: SNASS
+    :members: fit, simulate_data_and_possibly_append, sample_posterior
+
+SNASSS
+~~~~~~
+
+..  autoclass:: SNASSS
+    :members: fit, simulate_data_and_possibly_append, sample_posterior

examples/bivariate_gaussian_snl.py

@@ -1,5 +1,5 @@
 """
-Example using sequential neural likelihood estimation  on a bivariate Gaussian
+Example using sequential neural likelihood estimation on a bivariate Gaussian
 """
 
 from functools import partial

examples/bivariate_gaussian_snp.py

@@ -1,5 +1,5 @@
 """
-Example using sequential neural posterior estimation  on a bivariate Gaussian
+Example using sequential neural posterior estimation on a bivariate Gaussian
 """
 
 import distrax

examples/bivariate_gaussian_snr.py

@@ -1,5 +1,5 @@
 """
-Example using sequential neural ratio estimation  on a bivariate Gaussian
+Example using sequential neural ratio estimation on a bivariate Gaussian
 """
 
 import distrax

pyproject.toml

@@ -62,7 +62,7 @@ dependencies = [
 
 [tool.hatch.envs.test.scripts]
 lint = 'pylint sbijax'
-test = 'pytest -v --doctest-modules --cov=./sbijax --cov-report=xml sbijax'
+test = 'pytest -v --cov=./sbijax --cov-report=xml sbijax'
 
 [tool.black]
 line-length = 80

sbijax/__init__.py

@@ -2,7 +2,7 @@
 sbijax: Simulation-based inference in JAX
 """
 
-__version__ = "0.1.6"
+__version__ = "0.1.7"
 
 
 from sbijax._src.abc.smc_abc import SMCABC