Extend documentation and fix API docs

.github/workflows/ci.yaml

@@ -10,8 +10,9 @@ jobs:
   build:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
-      - uses: actions/setup-python@v2
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v3
+      - uses: pre-commit/[email protected]
 
   test:
     name: unit tests
@@ -22,9 +23,9 @@ jobs:
       matrix:
         python-version: [3.9]
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v3
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v3
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
@@ -53,9 +54,9 @@ jobs:
         matrix:
           python-version: [3.9]
       steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v3
         with:
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies
@@ -74,9 +75,9 @@ jobs:
         matrix:
           python-version: [3.9]
       steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v3
         with:
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies

.pre-commit-config.yaml

@@ -24,14 +24,14 @@ repos:
   hooks:
   - id: black
     args: ["--config=pyproject.toml"]
-    files: "(ramsey)"
+    files: "(ramsey|examples)"
 
 - repo: https://github.com/pycqa/isort
-  rev: 5.10.1
+  rev: 5.12.0
   hooks:
   - id: isort
     args: ["--settings-path=pyproject.toml"]
-    files: "(ramsey)"
+    files: "(ramsey|examples)"
 
 - repo: https://github.com/pycqa/bandit
   rev: 1.7.1
@@ -42,22 +42,23 @@ repos:
     types: [python]
     args: ["-c", "pyproject.toml"]
     additional_dependencies: ["toml"]
-    files: "(ramsey)"
+    files: "(ramsey|examples)"
 
-- repo: local
+- repo: https://github.com/PyCQA/flake8
+  rev: 5.0.1
   hooks:
-  - id: pylint
-    name: pylint
-    entry: pylint
-    language: python
-    files: "ramsey"
+   - id: flake8
+     additional_dependencies: [
+       flake8-typing-imports==1.14.0,
+       flake8-pyproject==1.1.0.post0
+     ]
 
 - repo: https://github.com/pre-commit/mirrors-mypy
   rev: v0.910-1
   hooks:
   - id: mypy
     args: ["--ignore-missing-imports"]
-    files: "(ramsey)"
+    files: "(ramsey|examples)"
 
 - repo: https://github.com/nbQA-dev/nbQA
   rev: 1.3.1

.readthedocs.yaml

@@ -1,13 +1,16 @@
 version: 2
 
 build:
-  os: ubuntu-20.04
+  os: "ubuntu-20.04"
   tools:
     python: "3.9"
 
 sphinx:
-   configuration: docs/source/conf.py
+  configuration: docs/conf.py
+  fail_on_warning: false
 
 python:
-   install:
-   - requirements: docs/requirements.txt
+  install:
+    - method: pip
+      path: .
+    - requirements: docs/requirements.txt

Makefile

@@ -0,0 +1,5 @@
+PKG_VERSION=`python setup.py --version`
+
+tag:
+	git tag -a $(PKG_VERSION) -m $(PKG_VERSION)
+	git push --tag

docs/.gitignore

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

docs/Makefile

@@ -5,7 +5,7 @@
 # from the environment for the first two.
 SPHINXOPTS    ?=
 SPHINXBUILD   ?= sphinx-build
-SOURCEDIR     = source
+SOURCEDIR     = ./
 BUILDDIR      = build
 
 # Put it first so that "make" without argument is like "make help".

docs/_static/theme.css

@@ -0,0 +1,27 @@
+html[data-theme="dark"], html[data-theme="light"] {
+  --pst-color-primary: #0048bc;
+}
+
+h1 > code > span {
+    font-family: var(--pst-font-family-monospace);
+    color: #0048bc;
+    font-weight: 700;
+}
+
+nav > li > a > code.literal {
+    padding-top: 0rem;
+    padding-bottom: 0rem;
+    background-color: white;
+    border: 0px;
+}
+
+
+nav.bd-links p.caption {
+    text-transform: uppercase;
+}
+
+code.literal {
+  background-color: white;
+  border: 0px;
+  border-radius: 0px;
+}

docs/source/conf.py → docs/conf.py

@@ -1,15 +1,13 @@
-import glob
-import os
 from datetime import date
 
 project = "Ramsey"
 copyright = f"{date.today().year}, the Ramsey developers"
-author = "Ramsey developers"
-release = "0.0.2"
+author = "the Ramsey developers"
 
 extensions = [
     "nbsphinx",
     "sphinx.ext.autodoc",
+    'sphinx_autodoc_typehints',
     "sphinx.ext.autosummary",
     "sphinx.ext.doctest",
     "sphinx.ext.intersphinx",
@@ -19,10 +17,13 @@
     "sphinx_autodoc_typehints",
     "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",

docs/getting_started.rst

@@ -0,0 +1,4 @@
+Philosophy
+==========
+
+Todo

docs/source/index.rst → docs/index.rst

@@ -1,10 +1,17 @@
 :github_url: https://github.com/ramsey-devs/ramsey/
 
-Ramsey documentation
-====================
+Ramsey: probabilistic modelling using Haiku
+===========================================
 
 Ramsey is a library for probabilistic modelling using `Haiku <https://github.com/deepmind/dm-haiku>`_ and `JAX <https://github.com/google/jax>`_.
-It builds upon the same module system that Haiku is using  and is hence fully compatible with Haiku's, NumPyro's API.
+It builds upon the same module system that Haiku is using and is hence fully compatible with its API. Ramsey implements **probabilistic** models, such as neural processes, Gaussian processes,
+Bayesian neural networks, Bayesian timeseries models and state-space-models, and more.
+
+Example
+-------
+
+Ramsey uses to Haiku's module system to construct probabilistic models
+and define parameters. For instance, a simple neural process can be constructed like this:
 
 .. code-block:: python
 
@@ -65,27 +72,40 @@ Contributions in the form of pull requests are more than welcome. A good way to
 
 In order to contribute:
 
-1) Install Ramsey and dev dependencies via :code:`pip install -e '.[dev]'`,
-2) test your contribution/implementation by calling :code:`tox` on the (Unix) command line before submitting a PR.
+1) Clone Ramsey and install it and its dev dependencies via :code:`pip install -e '.[dev]'`,
+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,
+4) test it by calling :code:`tox` on the (Unix) command line,
+5) submit a PR 🙂
 
 License
 -------
 
-Ramsey is licensed under a Apache 2.0 License
+Ramsey is licensed under the Apache 2.0 License.
+
+
+..  toctree::
+    :maxdepth: 1
+    :hidden:
 
+    Home <self>
 
 ..  toctree::
     :caption: Tutorials
     :maxdepth: 1
     :hidden:
 
     notebooks/neural_process
-    notebooks/gaussian_process
     notebooks/forecasting
 
 ..  toctree::
-    :caption: API reference
+    :caption: API
     :maxdepth: 1
     :hidden:
 
-    api
+    ramsey
+    ramsey.attention
+    ramsey.contrib
+    ramsey.family
+    ramsey.kernels
+    ramsey.train

docs/source/notebooks/forecasting.ipynb → docs/notebooks/forecasting.ipynb

@@ -216,7 +216,7 @@
    "outputs": [],
    "source": [
     "from ramsey import GP\n",
-    "from ramsey.covariance_functions import ExponentiatedQuadratic, Periodic\n",
+    "from ramsey.kernels import ExponentiatedQuadratic, Periodic\n",
     "from ramsey.train import train_gaussian_process"
    ]
   },
@@ -620,7 +620,6 @@
     "    n_context,\n",
     "    n_target,\n",
     "):\n",
-    "\n",
     "    key, sample_key = random.split(key, 2)\n",
     "    sample_idxs = random.choice(\n",
     "        sample_key,\n",

docs/source/notebooks/neural_process.ipynb → docs/notebooks/neural_process.ipynb

@@ -129,7 +129,7 @@
     }
    },
    "source": [
-    "## Neural process\n",
+    "# Neural process\n",
     "\n",
     "Having sampled a data set, we can define the model. A neural process model takes a `decoder` argument, a `latent_encoder` argument and an optional `deterministic_encoder` argument.\n",
     "\n",
@@ -334,7 +334,7 @@
     }
    },
    "source": [
-    "## Attentive neural process\n",
+    "# Attentive neural process\n",
     "\n",
     "An attentive neural process model takes a `decoder` argument, a `latent_encoder` argument and a `deterministic_encoder` argument with an attention module. In comparison to before the\n",
     "`deterministic_encoder` is also a tuple in this case. In addition to the network that generates the representation, one is also requires to provide a mechanism for cross-attention.\n",

docs/ramsey.attention.rst

@@ -0,0 +1,18 @@
+``ramsey.attention``
+====================
+
+.. currentmodule:: ramsey.attention
+
+.. automodule:: ramsey.attention
+
+Attention
+---------
+
+.. autosummary::
+    MultiHeadAttention
+
+MultiHeadAttention
+~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: MultiHeadAttention
+   :members:

docs/ramsey.contrib.rst

@@ -0,0 +1,37 @@
+``ramsey.contrib``
+==================
+
+.. currentmodule:: ramsey.contrib
+
+.. automodule:: ramsey.contrib
+
+Modules
+-------
+
+.. autosummary::
+    BayesianLinear
+
+BayesianLinear
+~~~~~~~~~~~~~~
+
+..  autoclass:: BayesianLinear
+    :members:
+
+Models
+------
+
+.. autosummary::
+    BNN
+    RANP
+
+Bayesian neural network
+~~~~~~~~~~~~~~~~~~~~~~~
+
+..  autoclass:: BNN
+    :members:
+
+Recurrent attentive neural process
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+..  autoclass:: RANP
+    :members:

docs/ramsey.family.rst

@@ -0,0 +1,25 @@
+``ramsey.family``
+=================
+
+.. currentmodule:: ramsey.family
+
+.. automodule:: ramsey.family
+
+Exponential families
+--------------------
+
+.. autosummary::
+    Gaussian
+    NegativeBinomial
+
+Gaussian
+~~~~~~~~
+
+..  autoclass:: Gaussian
+    :members:
+
+NegativeBinomial
+~~~~~~~~~~~~~~~~
+
+..  autoclass:: NegativeBinomial
+    :members: