Move to UV and Google docstyle

.github/workflows/ci.yaml

@@ -14,95 +14,51 @@ jobs:
       - uses: actions/setup-python@v3
       - uses: pre-commit/[email protected]
 
-  build:
-    runs-on: ubuntu-latest
-    needs:
-      - precommit
-    strategy:
-      matrix:
-        python-version: [3.11]
-    steps:
-    - uses: actions/checkout@v3
-    - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v3
-      with:
-        python-version: ${{ matrix.python-version }}
-    - name: Install dependencies
-      run: |
-        pip install hatch
-    - name: Build package
-      run: |
-        hatch build
-
   lints:
     runs-on: ubuntu-latest
     needs:
       - precommit
     strategy:
       matrix:
-        python-version: [3.11]
+        python-version: [ 3.11, 3.12 ]
     steps:
       - uses: actions/checkout@v3
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v3
         with:
           python-version: ${{ matrix.python-version }}
+      - uses: astral-sh/setup-uv@v5
+        with:
+          version: "latest"
       - name: Install dependencies
         run: |
-          pip install hatch
+          uv sync --dev
       - name: Run lints
         run: |
-          hatch run test:lints
+          make lints
 
   tests:
     runs-on: ubuntu-latest
     needs:
       - lints
     strategy:
       matrix:
-        python-version: [3.11]
+        python-version: [3.11, 3.12]
     steps:
       - uses: actions/checkout@v3
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v3
         with:
           python-version: ${{ matrix.python-version }}
-      - name: Install dependencies
-        run: |
-          pip install hatch
-      - name: Build package
-        run: |
-          pip install jaxlib==0.4.24 jax==0.4.24
-      - name: Run tests
-        run: |
-          hatch run test:tests
-      - name: Upload coverage reports to Codecov
-        uses: codecov/codecov-action@v3
-        env:
-          CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
-
-  examples:
-    runs-on: ubuntu-latest
-    needs:
-      - tests
-    strategy:
-      matrix:
-        python-version: [3.11]
-    steps:
-      - uses: actions/checkout@v3
-      - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v3
+      - uses: astral-sh/setup-uv@v5
         with:
-          python-version: ${{ matrix.python-version }}
+          version: "latest"
       - name: Install dependencies
         run: |
-          pip install hatch
-      - name: Build package
-        run: |
-          pip install jaxlib==0.4.24 jax==0.4.24
-      - name: Run examples
+          uv sync --dev
+      - name: Run tests
         run: |
-          hatch run test:examples
+          make tests
       - name: Upload coverage reports to Codecov
         uses: codecov/codecov-action@v3
         env:

.github/workflows/examples.yaml

@@ -0,0 +1,39 @@
+name: examples
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  precommit:
+    name: Pre-commit checks
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+
+  examples:
+    runs-on: ubuntu-latest
+    needs:
+      - precommit
+    strategy:
+      matrix:
+        python-version: [ 3.11, 3.12 ]
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v3
+        with:
+          python-version: ${{ matrix.python-version }}
+      - uses: astral-sh/setup-uv@v5
+        with:
+          version: "latest"
+      - name: Install dependencies
+        run: |
+          uv sync --all-extras
+      - name: Run tests
+        run: |
+          uv run python examples/neural_processes.py -n 10 -m NP
+          uv run python examples/neural_processes.py -n 10 -m ANP
+          uv run python examples/neural_processes.py -n 10 -m DANP

.pre-commit-config.yaml

@@ -14,7 +14,7 @@ repos:
     - id: trailing-whitespace
 
 - repo: https://github.com/pre-commit/mirrors-mypy
-  rev: v0.910-1
+  rev: v1.15.0
   hooks:
   - id: mypy
     args: ["--ignore-missing-imports"]

Makefile

@@ -1,22 +1,14 @@
-.PHONY: tag, tests, lints, docs, format, examples
-
-PKG_VERSION=`hatch version`
-
-tag:
-	 git tag -a v${PKG_VERSION} -m v${PKG_VERSION}
-	 git push --tag
+.PHONY: tests, lints, docs, format
 
 tests:
-	hatch run test:tests
+	uv run pytest
 
 lints:
-	hatch run test:lints
+	uv run ruff check ramsey examples
 
 format:
-	hatch run test:format
+	uv run ruff check --select I --fix ramsey examples
+	uv run ruff format ramsey examples
 
 docs:
 	cd docs && make html
-
-examples:
-	hatch run test:examples

README.md

@@ -11,49 +11,65 @@
 
 ## About
 
-Ramsey is a library for probabilistic modelling using [JAX](https://github.com/google/jax),
+Ramsey is a library for probabilistic deep learning using [JAX](https://github.com/google/jax),
 [Flax](https://github.com/google/flax) and [NumPyro](https://github.com/pyro-ppl/numpyro).
-It offers high quality implementations of neural processes, Gaussian processes, Bayesian time series and state-space models, clustering processes,
-and everything else Bayesian.
 
-Ramsey makes use of
+Ramsey's scope covers
 
-- Flax`s module system for models with trainable parameters (such as neural or Gaussian processes),
-- NumPyro for models where parameters are endowed with prior distributions (such as Gaussian processes, Bayesian neural networks, ARMA models)
-
-and is hence aimed at being fully compatible with both of them.
+- neural processes (vanilla, attentive, Markovian, convolutional, ...),
+- neural Laplace and Fourier operator models,
+- flow matching and denoising diffusion models,
+- etc.
 
 ## Example usage
 
 You can, for instance, construct a simple neural process like this:
 
 ```python
-from jax import random as jr
+from flax import nnx
 
 from ramsey import NP
 from ramsey.nn import MLP
-from ramsey.data import sample_from_sine_function
 
-def get_neural_process():
-    dim = 128
-    np = NP(
-        decoder=MLP([dim] * 3 + [2]),
-        latent_encoder=(
-            MLP([dim] * 3), MLP([dim, dim * 2])
-        )
+def get_neural_process(in_features, out_features):
+  dim = 128
+  np = NP(
+    decoder=MLP(in_features, [dim, dim, out_features * 2], rngs=nnx.Rngs(0)),\
+    latent_encoder=(
+      MLP(in_features, [dim, dim], rngs=nnx.Rngs(1)),
+      MLP(dim, [dim, dim * 2], rngs=nnx.Rngs(2))
     )
-    return np
+  )
+  return np
+
+neural_process = get_neural_process(1, 1)
+```
+
+The neural process takes a decoder and a set of two latent encoders as arguments. All of these are typically `flax.nnx` MLPs, but
+Ramsey is flexible enough that you can change them, for instance, to CNNs or RNNs. Once the model is defined, you can train
+it by accessing the ELBO given input-output pairs via
 
-key = jr.PRNGKey(23)
+```python
+from jax import random as jr
+from ramsey.data import sample_from_sine_function
+
+key = jr.PRNGKey(0)
 data = sample_from_sine_function(key)
 
-neural_process = get_neural_process()
-params = neural_process.init(key, x_context=data.x, y_context=data.y, x_target=data.x)
+x_context, y_context = data.x[:, :20, :],  data.y[:, :20, :]
+x_target, y_target = data.x, data.y
+loss = neural_process.loss(
+  x_context=x_context,
+  y_context=y_context,
+  x_target=x_target,
+  y_target=y_target
+)
 ```
 
-The neural process takes a decoder and a set of two latent encoders as arguments. All of these are typically MLPs, but
-Ramsey is flexible enough that you can change them, for instance, to CNNs or RNNs. Once the model is defined, you can initialize
-its parameters just like in Flax.
+Making predictions can be done like this:
+```python
+pred = neural_process(x_context=x_context, y_context=y_context, x_target=x_target)
+```
 
 ## Installation
 
@@ -70,8 +86,7 @@ command line:
 pip install git+https://github.com/ramsey-devs/ramsey@<RELEASE>
 ```
 
-See also the installation instructions for [JAX](https://github.com/google/jax), if
-you plan to use Ramsey on GPU/TPU.
+See also the installation instructions for [JAX](https://github.com/google/jax), if you plan to use Ramsey on GPU/TPU.
 
 ## Contributing
 
@@ -80,11 +95,12 @@ Contributions in the form of pull requests are more than welcome. A good way to
 
 In order to contribute:
 
-1) Clone Ramsey and install the package manager `hatch` via `pip install hatch`,
+1) Clone Ramsey and install  `uv` from [here](https://github.com/astral-sh/uv),
 2) create a new branch locally `git checkout -b feature/my-new-feature` or `git checkout -b issue/fixes-bug`,
-3) implement your contribution and ideally a test case,
-4) test it by calling `make format`, `make lints` and `make tests` on the (Unix) command line,
-5) submit a PR 🙂
+3) install all dependencies via `uv sync --all-extras`,
+4) implement your contribution and ideally a test case,
+5) test it by calling `make format`, `make lints` and `make tests` on the (Unix) command line,
+6) submit a PR 🙂
 
 ## Why Ramsey
 

docs/conf.py

@@ -13,6 +13,7 @@
     "sphinx.ext.intersphinx",
     "sphinx.ext.mathjax",
     "sphinx.ext.napoleon",
+    'sphinxcontrib.bibtex',
     "sphinx.ext.viewcode",
     "sphinx_autodoc_typehints",
     "sphinx_copybutton",
@@ -41,14 +42,32 @@
     "examples/*py",
 ]
 
-html_theme = "sphinx_book_theme"
+autosummary_generate = True
+autodoc_typehints = 'none'
+typehints_fully_qualified = True
+always_document_param_types = True
 
+html_theme = "sphinx_book_theme"
 html_theme_options = {
     "repository_url": "https://github.com/ramsey-devs/ramsey",
     "use_repository_button": True,
     "use_download_button": False,
     "use_fullscreen_button": False,
     "launch_buttons": {"colab_url": "https://colab.research.google.com"},
 }
-
 html_title = "Ramsey"
+
+
+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']
+bibtex_default_style = 'plain'
+bibtex_reference_style = 'label'