Update forecasting notebook

.github/workflows/ci.yaml

@@ -45,7 +45,7 @@ jobs:
         files: coverage.xml
 
   notebooks:
-      name: notebooks and examples
+      name: notebooks
       runs-on: ubuntu-latest
       needs:
         - test
@@ -64,6 +64,24 @@ jobs:
       - name: Run notebooks
         run: |
           tox -e notebooks
+
+  examples:
+      name: examples
+      runs-on: ubuntu-latest
+      needs:
+        - test
+      strategy:
+        matrix:
+          python-version: [3.9]
+      steps:
+      - uses: actions/checkout@v2
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install dependencies
+        run: |
+          pip install tox codecov
       - name: Run examples
         run: |
           tox -e examples

.pre-commit-config.yaml

@@ -24,14 +24,14 @@ repos:
   hooks:
   - id: black
     args: ["--config=pyproject.toml"]
-    files: "(ramsey|tests)"
+    files: "(ramsey)"
 
 - repo: https://github.com/pycqa/isort
   rev: 5.10.1
   hooks:
   - id: isort
     args: ["--settings-path=pyproject.toml"]
-    files: "(ramsey|tests)"
+    files: "(ramsey)"
 
 - repo: https://github.com/pycqa/bandit
   rev: 1.7.1
@@ -42,7 +42,7 @@ repos:
     types: [python]
     args: ["-c", "pyproject.toml"]
     additional_dependencies: ["toml"]
-    files: "(ramsey|tests)"
+    files: "(ramsey)"
 
 - repo: local
   hooks:
@@ -57,7 +57,7 @@ repos:
   hooks:
   - id: mypy
     args: ["--ignore-missing-imports"]
-    files: "(ramsey|tests)"
+    files: "(ramsey)"
 
 - repo: https://github.com/nbQA-dev/nbQA
   rev: 1.3.1

docs/source/conf.py

@@ -12,25 +12,15 @@
     "sphinx.ext.autodoc",
     "sphinx.ext.autosummary",
     "sphinx.ext.doctest",
-    "sphinx.ext.inheritance_diagram",
     "sphinx.ext.intersphinx",
     "sphinx.ext.mathjax",
     "sphinx.ext.napoleon",
     "sphinx.ext.viewcode",
     "sphinx_autodoc_typehints",
-    "sphinx_gallery.gen_gallery",
     "sphinx_math_dollar",
     "IPython.sphinxext.ipython_console_highlighting",
 ]
 
-sphinx_gallery_conf = {
-    "examples_dirs": ["../../examples"],
-    "gallery_dirs": ["examples"],
-    "filename_pattern": "/plot_",
-    "ignore_pattern": "(__init__)",
-    "min_reported_time": 1,
-}
-
 templates_path = ["_templates"]
 html_static_path = ["_static"]
 

docs/source/index.rst

@@ -1,12 +1,10 @@
-:github_url: https://github.com/dirmeier/ramsey/
+:github_url: https://github.com/ramsey-devs/ramsey/
 
 Ramsey documentation
 ====================
 
 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.
-Ramsey implements (or rather intends to implement) neural and Gaussian process models, normalizing flows,
-and diffusion and score-based models.
 
 .. code-block:: python
 
@@ -82,13 +80,8 @@ Ramsey is licensed under a Apache 2.0 License
     :hidden:
 
     notebooks/neural_process
-
-..  toctree::
-    :maxdepth: 1
-    :caption: Examples
-    :hidden:
-
-    examples/attentive_neural_process
+    notebooks/gaussian_process
+    notebooks/forecasting
 
 ..  toctree::
     :caption: API reference

examples/bayesian_neural_network.py

@@ -22,12 +22,12 @@
 from jax import random
 
 from ramsey.contrib import BNN, BayesianLinear
-from ramsey.data import sample_from_sine_function
+from ramsey.data import sample_from_linear_model
 
 
 def data(key):
-    (x_target, y_target), f_target = sample_from_sine_function(
-        key, batch_size=1, num_observations=50
+    (x_target, y_target), f_target = sample_from_linear_model(
+        key, batch_size=1, num_observations=50, noise_scale=1.0
     )
 
     return (x_target.reshape(-1, 1), y_target.reshape(-1, 1)), f_target.reshape(
@@ -37,11 +37,8 @@ def data(key):
 
 def _bayesian_nn(**kwargs):
     layers = [
-        #BayesianLinear(8, with_bias=True),
-        hk.Linear(16, with_bias=True),
-        BayesianLinear(64, with_bias=True),
-        hk.Linear(16, with_bias=True),
-        hk.Linear(1, with_bias=False),
+        BayesianLinear(8, with_bias=True),
+        hk.Linear(2, with_bias=False),
     ]
     bnn = BNN(layers)
     return bnn(**kwargs)
@@ -51,7 +48,7 @@ def train_bnn(
     rng_seq,
     x,  # pylint: disable=invalid-name
     y,  # pylint: disable=invalid-name
-    n_iter=10000,
+    n_iter=20000,
     stepsize=0.001,
 ):
     bnn = hk.transform(_bayesian_nn)

examples/negative_binomial_linear_model.py

@@ -1,3 +1,10 @@
+"""
+Negative binomial regression model
+==================================
+
+This example implements a basic generalized linear model
+"""
+
 import haiku as hk
 import matplotlib.pyplot as plt
 from jax import numpy as jnp

pyproject.toml

@@ -1,7 +1,6 @@
 [build-system]
 requires = ["setuptools", "wheel"]
 
-
 [tool.black]
 line-length = 80
 target-version = ['py38']

ramsey/_src/contrib/bayesian_neural_network/bayesian_neural_network.py

@@ -1,7 +1,6 @@
 import warnings
 from typing import Iterable, Optional
 
-import chex
 import haiku as hk
 from jax import numpy as jnp
 
@@ -40,9 +39,8 @@ def __init__(
         ----------
         layers: Iterable[hk.Module]
             layers of the BNN
-        log_scale_init: Optional[hk.initializers.Initializer]
-            an initializer object from Haiku or None. Used to initialize
-            the log scale of the likelihood
+        family: Family
+            exponential family of the response
         name: Optional[str]
             name of the layer
         kwargs: keyword arguments
@@ -82,7 +80,6 @@ def _loss(self, x, y):
             else:
                 x = layer(x)
 
-        chex.assert_equal_shape([x, y])
         likelihood_fn = self._as_family(x)
         likelihood = jnp.sum(likelihood_fn.log_prob(y))
 
@@ -93,20 +90,4 @@ def _loss(self, x, y):
         return likelihood_fn, -elbo
 
     def _as_family(self, x):
-        param_name, _ = self._family.get_canonical_parameters()
-        return self._family(
-            x, **{param_name: self._get_log_scale_parameter(x.dtype)}
-        )
-
-    def _get_log_scale_parameter(self, dtype):
-        param_name, _ = self._family.get_canonical_parameters()
-        if param_name + "_init" not in self._kwargs:
-            log_scale_init = hk.initializers.TruncatedNormal(
-                jnp.log(0.1), jnp.log(1.0)
-            )
-        else:
-            log_scale_init = self._kwargs[param_name + "_init"]
-        log_scale = hk.get_parameter(
-            "log_scale", [], dtype=dtype, init=log_scale_init
-        )
-        return log_scale
+        return self._family(x)

ramsey/data.py

@@ -160,15 +160,16 @@ def sample_from_negative_binomial_linear_model(
 
 # pylint: disable=too-many-locals,invalid-name
 def sample_from_linear_model(
-    key, batch_size=10, num_observations=100, num_dim=1
+    key, batch_size=10, num_observations=100, num_dim=1, noise_scale=None
 ):
     rng_seq = hk.PRNGSequence(key)
     x = random.normal(next(rng_seq), (num_observations, num_dim))
     ys = []
     fs = []
     for _ in range(batch_size):
         beta = dist.Normal(0.0, 2.0).sample(next(rng_seq), (num_dim,))
-        noise_scale = dist.Gamma(1.0, 10.0).sample(next(rng_seq))
+        if noise_scale is None:
+            noise_scale = dist.Gamma(1.0, 10.0).sample(next(rng_seq))
         f = x @ beta
         y = f + random.normal(next(rng_seq), f.shape) * noise_scale
         fs.append(f.reshape((1, num_observations, 1)))