vpet-abc 0.1.2

Fast likelihood-free PET kinetic modelling in JAX

vPET-ABC-JAX

Fast, likelihood‑free PET kinetic modelling implemented in JAX

---

1 · What is this repository?

vpetabc is a pure‑Python re‑implementation of the vPET‑ABC framework (Grazian et al., 2021) for large‑scale dynamic PET kinetic modelling, written from the ground up in JAX.

Compared with the earlier CuPy version, the JAX rewrite

• removes CUDA‑specific boiler‑plate – the same code runs on CPU, multi‑GPU, or TPU via XLA;
• exposes a clean, PyTorch‑like API centred on an abstract KineticModel;
• relies on vectorised primitives (vmap, lax.scan) so that even > 40 M‑voxel datasets fit into a single JIT‑compiled graph;
• delivers further speed‑ups;
• depends only on jax, pandas, and tqdm – no CuPy, no manual builds.

---

2 · Repository layout

.
├── data/
│   ├── sample_2TCM.csv
│   └── sample_lpntPET.csv
├── dist/                     # wheels/sdists created by `python -m build`
├── example_usage.ipynb
├── pyproject.toml
├── README.md
├── src/
    └── vpetabc/
        ├── __init__.py       # package namespace
        ├── engine.py         # ABC engine + helpers
        ├── models.py         # TwoTissueModel, lpntPETModel, …
        ├── priors.py         # prior samplers
        └── utilities.py      # I/O + posterior utilities

Module	Description
engine.py	ABCRejection, the fully vectorised, JIT‑compiled rejection‑ABC driver
models.py	KineticModel base‑class + TwoTissueModel, lpntPETModel implementations
priors.py	Uniform × Bernoulli prior samplers (TwoTissuePrior, lpntPETPrior)
utilities.py	preprocess_table, get_conditional_posterior_mean, misc. helpers

---

3 · Installation

TL;DR pip install vpet-abc[cuda]

3.1  Stable release from PyPI

# create & activate a virtual environment
python -m venv .venv && source .venv/bin/activate

# CPU‑only:
pip install vpet-abc

# NVIDIA GPUs
pip install vpet-abc[cuda]

jax[cuda] wheels already bundle matching CUDA/cuDNN libraries; you only need a driver on Linux / Windows. For TPU, Metal (macOS), or ROCm see the official
JAX installation guide.

3.2  Tested environments

OS	Python	jax / jaxlib	Accelerator
macOS 15.5 (arm64)	3.11.12	0.6.1 (CPU)	Apple M2
Rocky Linux 8.10	3.9.2	0.4.30 [cuda12]	NVIDIA V100

---

4 · Quick start

See example_usage.ipynb for an executable walkthrough, or run:

import jax
import jax.numpy as jnp
import jax.random as jr
import pandas as pd
from vpetabc import *

df = pd.read_csv("data/sample_2TCM.csv", index_col=0)
Cp_fine, A, TACs, _ = preprocess_table(df)

engine = ABCRejection(
    TwoTissueModel(),
    prior_sampler = TwoTissuePrior,
    num_sims      = 200_000,
    accept_frac   = 0.005,
)

post = engine.run(jr.PRNGKey(0), TACs, Cp_fine, A, batch_size=50_000)
means, chosen = get_conditional_posterior_mean(post)
print(means[:3])            # first 3 voxels

---

5 · Extending the framework

1. Define your kinetic model

class MyModel(KineticModel):
    @partial(jax.jit, static_argnums=(0,))
    def simulate(self, θ, Cp, dt):
        # return Ct(t) as a (T_fine,) array
        ...

2. Write a prior

def MyPrior(key, n, lows, highs):
    return jr.uniform(key, (n, P), lows, highs)

3. Pass MyModel and MyPrior to ABCRejection.
   Batching, GPU kernels, and distance evaluation are handled automatically.

---

6 · Benchmarks

TBD

---

7 · Citation

TBA soon.

---

8 · Licence

vpetabc is released under the MIT Licence (see LICENCE).
The sample dataset is provided for non‑commercial research use only.