Add a test that runs everything (#12)

.pre-commit-config.yaml

@@ -40,6 +40,8 @@ repos:
     language: python
     language_version: python3
     types: [python]
+    args: ["-c", "pyproject.toml"]
+    additional_dependencies: ["toml"]
     files: "(sbijax|examples)"
 
 - repo: https://github.com/PyCQA/flake8
@@ -58,6 +60,17 @@ repos:
     args: ["--ignore-missing-imports"]
     files: "(sbijax|examples)"
 
+- repo: https://github.com/nbQA-dev/nbQA
+  rev: 1.6.3
+  hooks:
+  - id: nbqa-black
+  - id: nbqa-pyupgrade
+    args: [--py39-plus]
+  - id: nbqa-isort
+    args: ['--profile=black']
+  - id: nbqa-flake8
+    args: ['--ignore=E501,E203,E302,E402,E731,W503']
+
 - repo: https://github.com/jorisroovers/gitlint
   rev: v0.18.0
   hooks:

README.md

@@ -2,13 +2,14 @@
 
 [![status](http://www.repostatus.org/badges/latest/concept.svg)](http://www.repostatus.org/#concept)
 [![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
-[BlackJAX](https://github.com/blackjax-devs/blackjax), [Haiku](https://github.com/deepmind/dm-haiku) and [JAX](https://github.com/google/jax).
+[JAX](https://github.com/google/jax), [Haiku](https://github.com/deepmind/dm-haiku) and [BlackJAX](https://github.com/blackjax-devs/blackjax).
 
 SbiJAX so far implements
 
@@ -37,29 +38,6 @@ To install the latest GitHub <RELEASE>, use:
 pip install git+https://github.com/dirmeier/sbijax@<RELEASE>
 ```
 
-## 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) Fork the repository and install `hatch` and `pre-commit`
-
-```bash
-pip install hatch pre-commit
-pre-commit install
-```
-
-2) Create a new branch in your fork and implement your contribution
-
-3) Test your contribution/implementation by calling `hatch run test` on the (Unix) command line before submitting a PR
-
-```bash
-hatch run test:lint
-hatch run test:test
-```
-
-4) Submit a pull request :slightly_smiling_face:
-
 ## Author
 
 Simon Dirmeier <a href="mailto:sfyrbnd @ pm me">sfyrbnd @ pm me</a>

pyproject.toml

@@ -25,12 +25,16 @@ dependencies = [
     "dm-haiku>=0.0.9",
     "flax>=0.6.3",
     "optax>=0.1.3",
+    "surjectors@git+https://[email protected]/dirmeier/[email protected]",
 ]
 dynamic = ["version"]
 
 [project.urls]
 homepage = "https://github.com/dirmeier/sbijax"
 
+[tool.hatch.metadata]
+allow-direct-references = true
+
 [tool.hatch.version]
 path = "sbijax/__init__.py"
 
@@ -50,7 +54,7 @@ dependencies = [
 
 [tool.hatch.envs.test.scripts]
 lint = 'pylint sbijax'
-test = 'pytest -v --doctest-modules --cov=./sbi --cov-report=xml sbijax'
+test = 'pytest -v --doctest-modules --cov=./sbijax --cov-report=xml sbijax'
 
 
 [tool.black]

sbijax/__init__.py

@@ -2,7 +2,7 @@
 sbijax: Simulation-based inference in JAX
 """
 
-__version__ = "0.0.10"
+__version__ = "0.0.11"
 
 
 from sbijax.abc.rejection_abc import RejectionABC

sbijax/snl.py

@@ -5,6 +5,8 @@
 import numpy as np
 import optax
 from absl import logging
+
+# TODO(simon): this is a bit an annoying dependency to have
 from flax.training.early_stopping import EarlyStopping
 from jax import numpy as jnp
 

sbijax/snl_test.py

@@ -1,6 +1,80 @@
 # pylint: skip-file
-import chex
+
+import distrax
+import haiku as hk
+import optax
+from jax import numpy as jnp
+from surjectors import Chain, MaskedCoupling, TransformedDistribution
+from surjectors.conditioners import mlp_conditioner
+from surjectors.util import make_alternating_binary_mask
+
+from sbijax import SNL
+
+
+def prior_model_fns():
+    p = distrax.Independent(
+        distrax.Uniform(jnp.full(2, -3.0), jnp.full(2, 3.0)), 1
+    )
+    return p.sample, p.log_prob
+
+
+def simulator_fn(seed, theta):
+    p = distrax.MultivariateNormalDiag(theta, 0.1 * jnp.ones_like(theta))
+    y = p.sample(seed=seed)
+    return y
+
+
+def log_density_fn(theta, y):
+    prior = distrax.Uniform(jnp.full(2, -3.0), jnp.full(2, 3.0))
+    likelihood = distrax.MultivariateNormalDiag(
+        theta, 0.1 * jnp.ones_like(theta)
+    )
+
+    lp = jnp.sum(prior.log_prob(theta)) + jnp.sum(likelihood.log_prob(y))
+    return lp
+
+
+def make_model(dim):
+    def _bijector_fn(params):
+        means, log_scales = jnp.split(params, 2, -1)
+        return distrax.ScalarAffine(means, jnp.exp(log_scales))
+
+    def _flow(method, **kwargs):
+        layers = []
+        for i in range(2):
+            mask = make_alternating_binary_mask(dim, i % 2 == 0)
+            layer = MaskedCoupling(
+                mask=mask,
+                bijector=_bijector_fn,
+                conditioner=mlp_conditioner([8, 8, dim * 2]),
+            )
+            layers.append(layer)
+        chain = Chain(layers)
+        base_distribution = distrax.Independent(
+            distrax.Normal(jnp.zeros(dim), jnp.ones(dim)),
+            1,
+        )
+        td = TransformedDistribution(base_distribution, chain)
+        return td(method, **kwargs)
+
+    td = hk.transform(_flow)
+    td = hk.without_apply_rng(td)
+    return td
 
 
 def test_snl():
-    chex.assert_equal(1, 1)
+    rng_seq = hk.PRNGSequence(0)
+    y_observed = jnp.array([-1.0, 1.0])
+
+    prior_simulator_fn, prior_logdensity_fn = prior_model_fns()
+    fns = (prior_simulator_fn, prior_logdensity_fn), simulator_fn
+
+    snl = SNL(fns, make_model(2))
+    params, info = snl.fit(
+        next(rng_seq),
+        y_observed,
+        n_rounds=1,
+        optimizer=optax.adam(1e-4),
+        sampler="slice",
+    )
+    _ = snl.sample_posterior(params, 2, 100, 50, sampler="slice")