v0.1.0

.github/actions/setup/action.yml

@@ -0,0 +1,10 @@
+name: Setup project
+description: Checkout code, install dependencies
+
+runs:
+  using: "composite"
+  steps:
+    - name: Install dependencies
+      shell: bash
+      run: |
+        pip install .[dev,docs]

.github/workflows/build-docs.yml

@@ -0,0 +1,57 @@
+name: documentation
+
+on:
+  push:
+    branches: ["main", "dev"]
+
+jobs:
+  build_docs:
+    runs-on: ubuntu-latest
+
+    steps:
+    - name: Checkout repository
+      uses: actions/checkout@v4
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: "3.10"
+    - uses: ./.github/actions/setup
+
+    - name: Generate .rst files from package
+      run: |
+        sphinx-apidoc -o docs/source/api diabayes
+
+    - name: Build Sphinx HTML documentation
+      run: |
+        cd docs
+        make clean && make html
+
+    - name: Upload documentation artifact
+      uses: actions/upload-artifact@v4
+      with:
+        name: sphinx-html
+        path: docs/build/html
+
+    # Publish built docs to gh-pages branch
+    - name: Commit documentation changes
+      run: |
+        git clone https://${{ secrets.GH_TOKEN }}@github.com/martijnende/diabayes.git --branch gh-pages --single-branch gh-pages
+        cd gh-pages
+        git clean -fdx  # Remove old untracked files and directories
+        cp -r ../docs/build/html/* .  # Copy new HTML files to gh-pages
+        touch .nojekyll
+        git config --local user.email "action@github.com"
+        git config --local user.name "GitHub Action"
+        git add .
+        git commit -m "Update documentation" -a || true
+        # The above command will fail if no changes were present, so we ignore
+        # that.
+      env:
+        GH_PAT: ${{ secrets.GH_TOKEN }}
+
+    - name: Push changes
+      uses: ad-m/github-push-action@master
+      with:
+        branch: gh-pages
+        directory: gh-pages
+        GH_PAT: ${{ secrets.GH_TOKEN }}

.github/workflows/python-test.yml

@@ -0,0 +1,26 @@
+name: tests
+
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+
+permissions:
+  contents: read
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python 3.10
+      uses: actions/setup-python@v5
+      with:
+        python-version: "3.10"
+    - uses: ./.github/actions/setup
+    - name: Test with pytest
+      run: |
+        pytest

.gitignore

@@ -1,3 +1,6 @@
+scratch/
+.vscode/
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[codz]
@@ -70,6 +73,8 @@ instance/
 
 # Sphinx documentation
 docs/_build/
+docs/build/
+docs/source/_autosummary
 
 # PyBuilder
 .pybuilder/

README.md

@@ -1 +1,32 @@
-# diabayes
+# DiaBayes: rock friction inversion tools
+
+[![GitHub Release](https://img.shields.io/github/release/martijnende/diabase.svg?style=flat)]() 
+[![tests](https://github.com/martijnende/diabayes/actions/workflows/python-test.yml/badge.svg)](https://github.com/martijnende/diabayes/actions/workflows/python-test.yml)
+[![documentation](https://github.com/martijnende/diabayes/actions/workflows/build-docs.yml/badge.svg)](https://martijnende.github.io/diabayes)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+<!-- [![PyPi Version](https://img.shields.io/pypi/v/diabayes.svg)](https://pypi.python.org/pypi/diabayes/) -->
+
+[Documentation](https://martijnende.github.io/diabayes) | [Example usage](#example-usage) | [Installation](#installation) | [How to cite?](#how-to-cite)
+
+## Example usage
+
+See `examples/simple_example.ipynb` for a self-contained Jupyter notebook that illustrates forward and inverse modelling.
+
+## Installation
+
+Installing the latest version with CPU support from GitHub:
+```bash
+pip install git@github.com:martijnende/diabayes.git
+```
+Installing from a local repository with Nvidia GPU support (CUDA version 12):
+```bash
+cd /path/to/diabase && pip install .[gpu]
+```
+If you plan to contribute to the development of this package, please include the development tools:
+```bash
+pip install .[gpu,dev]
+```
+
+## How to cite?
+
+A publication describing this software package is underway. Until then, feel free to refer to this repository.

diabayes/SVI.py

@@ -0,0 +1,108 @@
+from typing import Callable, TypeAlias
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+from jaxtyping import Array, Float
+
+from diabayes.typedefs import Variables, _Params
+
+ParticleArray = Float[Array, "N M"]
+GradientArray: TypeAlias = ParticleArray
+
+
+@eqx.filter_jit
+def _distance(x1: ParticleArray, x2: ParticleArray) -> Float[Array, "N N"]:
+    """Euclidean distance between two particles `x1` and `x2`"""
+    return jnp.square(x1 - x2).sum(axis=-1)
+
+
+@eqx.filter_jit
+def _exponential_kernel(
+    x1: ParticleArray, x2: ParticleArray, h: float
+) -> Float[Array, "N N"]:
+    """Radial basis distance between two particles `x1` and `x2`"""
+    dist_sq = _distance(x1, x2)
+    return jnp.exp(-dist_sq / h)
+
+
+@eqx.filter_jit
+def _median_trick_h(x: ParticleArray) -> Float:
+    """
+    Compute the scaling factor `h` proportional to the
+    squared median of the RMS distance between all pairs
+    of particles.
+    """
+    dist_sq = jnp.square(x[:, None] - x[None]).sum(axis=-1)
+    # Replace this one with jnp.nanmedian to avoid spreading NaNs?
+    med_sq = jnp.median(jnp.sqrt(dist_sq))
+    h = med_sq**2 / jnp.log(len(x) + 1)
+    return h
+
+
+@eqx.filter_jit
+def compute_phi(
+    x: ParticleArray, gradp: GradientArray, gradq: GradientArray
+) -> GradientArray:
+    """
+    Compute the Stein variational gradients for a set of
+    particles (``x``) and the gradients of the log-likelihood
+    function (``gradp``) and the log-prior (``gradq``).
+
+    Parameters
+    ----------
+    x : ParticleArray
+        The set of invertible parameters ("particles")
+        of shape ``(Nparticles, Ndimensions)``
+    gradp, gradq : GradientArray
+        The gradients of the log-likelihood (``gradp``) and the
+        log-prior (``gradq``) with respect to the invertible parameters.
+        Has a shape ``(Nparticles, Ndimensions)``
+
+    Returns
+    -------
+    grad_x : Gradients
+        The directional gradients for the particle updates, e.g.
+        ``new_x = x + step_size * grad_x``. Has the same
+        shape as ``x`` and ``gradp, gradq``.
+    """
+    h = _median_trick_h(x)
+    map_over_a = (0, None, None)
+    map_over_b = (None, 0, None)
+
+    map1 = jax.vmap(_exponential_kernel, map_over_a, 0)
+    kernel = jax.vmap(map1, map_over_b, 1)
+    K = kernel(x, x, h)
+
+    map1 = jax.vmap(jax.grad(_exponential_kernel), map_over_a, 0)
+    grad_kernel = jax.vmap(map1, map_over_b, 1)
+    grad_K = grad_kernel(x, x, h).sum(axis=0)
+
+    grad_x = (jnp.matmul(K, gradp) + jnp.matmul(K, gradq) + grad_K) / len(gradp)
+    return -grad_x
+
+
+@eqx.filter_value_and_grad
+def _log_likelihood(
+    log_params: _Params,
+    mu_obs: Float[Array, "Nt"],
+    noise_std: Float,
+    forward_fn: Callable[[_Params], Variables],
+) -> Float:
+    # Transform the log_params by taking exponential
+    params = type(log_params).from_array(jnp.exp(log_params.to_array()))
+    # Forward pass to get friction curve
+    mu_hat = jnp.squeeze(forward_fn(params).mu)
+    # Log-likelihood of residuals
+    p = -(jnp.square(mu_obs - mu_hat).mean() / (2 * noise_std**2))
+    return p
+
+
+"""
+vmap _log_likelihood along particle axis. Since _log_likelihood is 
+decorated with eqx.filter_value_and_grad, it returns the log-
+likelihood and its gradient (hence it needs 2 values for the out_axes).
+"""
+mapped_log_likelihood = jax.vmap(
+    _log_likelihood, in_axes=(0, None, None, None), out_axes=(0, 0)
+)

diabayes/__init__.py

@@ -0,0 +1,7 @@
+from diabayes.typedefs import (
+    BayesianSolution,
+    RSFConstants,
+    RSFParams,
+    SpringBlockConstants,
+    Variables,
+)