Add docs folder

README.md

@@ -1,18 +1,16 @@
 # sbijax
 
-[![status](http://www.repostatus.org/badges/latest/concept.svg)](http://www.repostatus.org/#concept)
+[![active](https://www.repostatus.org/badges/latest/active.svg)](https://www.repostatus.org/#active)
 [![ci](https://github.com/dirmeier/sbijax/actions/workflows/ci.yaml/badge.svg)](https://github.com/dirmeier/sbijax/actions/workflows/ci.yaml)
 [![version](https://img.shields.io/pypi/v/sbijax.svg?colorB=black&style=flat)](https://pypi.org/project/sbijax/)
 
 > Simulation-based inference in JAX
 
 ## About
 
-`sbijax` implements several algorithms for simulation-based inference using
-[JAX](https://github.com/google/jax), [Haiku](https://github.com/deepmind/dm-haiku),
-[Distrax](https://github.com/deepmind/distrax) and [BlackJAX](https://github.com/blackjax-devs/blackjax).
-
-`sbijax` so far implements
+`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, `sbijax` implements
 
 - [Sequential Monte Carlo ABC](https://www.routledge.com/Handbook-of-Approximate-Bayesian-Computation/Sisson-Fan-Beaumont/p/book/9780367733728) (`SMCABC`),
 - [Neural Likelihood Estimation](https://arxiv.org/abs/1805.07226) (`SNL`)
@@ -26,7 +24,11 @@ where the acronyms in parentheses denote the names of the methods in `sbijax`.
 
 ## Examples
 
-You can find several self-contained examples on how to use the algorithms in `examples`.
+You can find several self-contained examples on how to use the algorithms in [examples](https://github.com/dirmeier/sbijax/tree/main/examples).
+
+## Documentation
+
+Documentation can be found [here](https://sbijax.readthedocs.io/en/latest/).
 
 ## Installation
 
@@ -45,6 +47,11 @@ To install the latest GitHub <RELEASE>, use:
 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
+feature-complete and user-friendly, and better documented.
+
 ## Author
 
 Simon Dirmeier <a href="mailto:sfyrbnd @ pm me">sfyrbnd @ pm me</a>

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,54 @@
+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"
+]
+
+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 🚀"

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 package draws significant inspiration from 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: 1
+    :hidden:
+
+    sbijax
+    sbijax.nn