diff --git a/.github/workflows/draft.yaml b/.github/workflows/draft.yaml new file mode 100644 index 0000000..6bae4d7 --- /dev/null +++ b/.github/workflows/draft.yaml @@ -0,0 +1,26 @@ +name: draft + +on: + push: + branches: [ main ] + pull_request: + branches: [ main ] + + +jobs: + paper: + runs-on: ubuntu-latest + name: Paper Draft + steps: + - name: Checkout + uses: actions/checkout@v4 + - name: Build draft PDF + uses: openjournals/openjournals-draft-action@master + with: + journal: joss + paper-path: paper/paper.md + - name: Upload + uses: actions/upload-artifact@v4 + with: + name: paper + path: paper/paper.pdf diff --git a/manuscript+supplement/.gitignore b/manuscript+supplement/.gitignore deleted file mode 100644 index 17c4ba8..0000000 --- a/manuscript+supplement/.gitignore +++ /dev/null @@ -1,7 +0,0 @@ -venv/ -venv/* - -figs/ -figs/* - -eeg_data.pkl diff --git a/manuscript+supplement/README.md b/manuscript+supplement/README.md deleted file mode 100644 index d6f9379..0000000 --- a/manuscript+supplement/README.md +++ /dev/null @@ -1,49 +0,0 @@ -# Reproducing results of the `sbijax` manuscript - -In order to reproduce the results of the paper, please follow the steps below. - -Install a new Python 3.11 virtual environment and activate it using: - -```shell -python3.11 -m venv venv -source venv/bin/activate -``` - -Install `sbijax` and all dependencies using: -```shell -pip install -r requirements.txt -``` - -## Usage with Jupyter notebooks - -For each section that contains code and figures, we provide a separate Jupyter notebook. -All notebooks contain the results and figures shown in the main manuscript. - - To run all experiments using Jupyter, first install a Jupyter kernel: -```shell -python -m ipykernel install --name sbi-dev --user -``` - -Then call -```shell -cd experimental_code -jupyter lab -``` - -This opens Jupyter on your web browser. You can now run any of the four notebooks. - -## Usage with Python files - -For each section that contains code and figures, we also provide separate Python files which -are exported from the notebooks. - -To execute each and reproduce the results of a section, call: -```shell -cd experimental_code/scripts -python chp3-the_sbi_package.py -python chp4-examples.py -python chp5-real_data_example.py -python chpx-appendix.py -``` - -Calling a Python script creates files in a folder called `figs`. diff --git a/manuscript+supplement/experimental_code/convert.sh b/manuscript+supplement/experimental_code/convert.sh deleted file mode 100755 index 1a69654..0000000 --- a/manuscript+supplement/experimental_code/convert.sh +++ /dev/null @@ -1,8 +0,0 @@ -#!/usr/bin/env bash - -set -euxo pipefail - -jupytext --to py *.ipynb -sed -i '' -e 's/%matplotlib inline//g' *.py -sed -i '' -e 's/plt.show()//g' *.py -mv *py scripts diff --git a/manuscript+supplement/requirements.txt b/manuscript+supplement/requirements.txt deleted file mode 100644 index e749459..0000000 --- a/manuscript+supplement/requirements.txt +++ /dev/null @@ -1,189 +0,0 @@ -absl-py==2.1.0 -anyio==4.7.0 -appnope==0.1.4 -argon2-cffi==23.1.0 -argon2-cffi-bindings==21.2.0 -arrow==1.3.0 -arviz==0.19.0 -asttokens==3.0.0 -astunparse==1.6.3 -async-lru==2.0.4 -attrs==24.3.0 -babel==2.16.0 -beautifulsoup4==4.12.3 -blackjax-nightly==1.1.1.post7 -bleach==6.2.0 -certifi==2024.12.14 -cffi==1.17.1 -charset-normalizer==3.4.1 -chex==0.1.88 -click==8.1.8 -cloudpickle==3.1.0 -comm==0.2.2 -contourpy==1.3.1 -coverage==7.6.10 -cycler==0.12.1 -debugpy==1.8.11 -decorator==5.1.1 -defusedxml==0.7.1 -distrax==0.1.5 -dm-haiku==0.0.12 -dm-tree==0.1.8 -docstring_parser==0.16 -edfio==0.4.5 -EDFlib-Python==1.0.8 -einops==0.8.0 -etils==1.11.0 -executing==2.1.0 -fastjsonschema==2.21.1 -fastprogress==1.0.3 -flatbuffers==24.12.23 -flax==0.10.2 -fonttools==4.55.3 -fqdn==1.5.1 -fsspec==2024.12.0 -gast==0.6.0 -google-pasta==0.2.0 -googleapis-common-protos==1.66.0 -grpcio==1.68.1 -h11==0.14.0 -h5netcdf==1.4.1 -h5py==3.12.1 -httpcore==1.0.7 -httpx==0.28.1 -humanize==4.11.0 -idna==3.10 -immutabledict==4.2.1 -importlib_resources==6.5.2 -iniconfig==2.0.0 -ipykernel==6.29.5 -ipython==8.31.0 -isoduration==20.11.0 -jax==0.4.31 -jaxlib==0.4.31 -jaxopt==0.8.3 -jedi==0.19.2 -Jinja2==3.1.5 -jmp==0.0.4 -joblib==1.4.2 -jrnmm @ git+https://github.com/dirmeier/jrnmm@9489dc72e2ac8f3a82db3fa12ae179fb161f744c -json5==0.10.0 -jsonpointer==3.0.0 -jsonschema==4.23.0 -jsonschema-specifications==2024.10.1 -jupyter-events==0.11.0 -jupyter-lsp==2.2.5 -jupyter_client==8.6.3 -jupyter_core==5.7.2 -jupyter_server==2.15.0 -jupyter_server_terminals==0.5.3 -jupyterlab==4.3.1 -jupyterlab_pygments==0.3.0 -jupyterlab_server==2.27.3 -keras==3.7.0 -kiwisolver==1.4.8 -lazy_loader==0.4 -libclang==18.1.1 -Markdown==3.7 -markdown-it-py==3.0.0 -MarkupSafe==3.0.2 -matplotlib==3.9.2 -matplotlib-inline==0.1.7 -mdurl==0.1.2 -memory-profiler==0.61.0 -mistune==3.1.0 -ml-dtypes==0.4.1 -mne==1.8.0 -mne-bids==0.14 -moabb==1.1.1 -msgpack==1.1.0 -namex==0.0.8 -nbclient==0.10.2 -nbconvert==7.16.5 -nbformat==5.10.4 -nest-asyncio==1.6.0 -notebook_shim==0.2.4 -numpy==1.26.4 -opt_einsum==3.4.0 -optax==0.2.4 -optree==0.13.1 -orbax-checkpoint==0.6.4 -overrides==7.7.0 -packaging==24.2 -pandas==1.5.3 -pandocfilters==1.5.1 -parso==0.8.4 -pexpect==4.9.0 -pillow==11.1.0 -platformdirs==4.3.6 -pluggy==1.5.0 -pooch==1.8.2 -prometheus_client==0.21.1 -promise==2.3 -prompt_toolkit==3.0.48 -protobuf==5.29.2 -psutil==6.1.1 -ptyprocess==0.7.0 -pure_eval==0.2.3 -pyarrow==18.1.0 -pycparser==2.22 -Pygments==2.18.0 -pyparsing==3.2.1 -pyriemann==0.6 -pytest==7.4.4 -python-dateutil==2.9.0.post0 -python-json-logger==3.2.1 -pytz==2024.2 -PyYAML==6.0.2 -pyzmq==26.2.0 -referencing==0.35.1 -requests==2.32.3 -rfc3339-validator==0.1.4 -rfc3986-validator==0.1.1 -rich==13.9.4 -rpds-py==0.22.3 -sbijax @ git+https://github.com/dirmeier/sbijax@708ebb69e9e7271ffa946e468e6e209ac9cc215d -scikit-learn==1.6.0 -scipy==1.15.0 -seaborn==0.12.2 -Send2Trash==1.8.3 -session-info==1.0.0 -simple-parsing==0.1.6 -six==1.17.0 -sniffio==1.3.1 -soupsieve==2.6 -stack-data==0.6.3 -stdlib-list==0.11.0 -surjectors==0.3.3 -tabulate==0.9.0 -tensorboard==2.18.0 -tensorboard-data-server==0.7.2 -tensorflow==2.18.0 -tensorflow-datasets==4.9.7 -tensorflow-io-gcs-filesystem==0.37.1 -tensorflow-metadata==1.16.1 -tensorflow-probability==0.25.0 -tensorstore==0.1.71 -termcolor==2.5.0 -terminado==0.18.1 -tfp-nightly==0.25.0.dev20240815 -threadpoolctl==3.5.0 -tinycss2==1.4.0 -toml==0.10.2 -toolz==1.0.0 -tornado==6.4.2 -tqdm==4.67.1 -traitlets==5.14.3 -types-python-dateutil==2.9.0.20241206 -typing_extensions==4.12.2 -uri-template==1.3.0 -urllib3==1.26.20 -wcwidth==0.2.13 -webcolors==24.11.1 -webencodings==0.5.1 -websocket-client==1.8.0 -Werkzeug==3.1.3 -wrapt==1.17.0 -xarray==2024.3.0 -xarray-einstats==0.8.0 -zipp==3.21.0 diff --git a/paper/paper.bib b/paper/paper.bib new file mode 100644 index 0000000..3efc9e1 --- /dev/null +++ b/paper/paper.bib @@ -0,0 +1,614 @@ +@article{dirmeier2025causal, + title={Causal posterior estimation}, + author={Dirmeier, Simon and Mira, Antonietta}, + journal={arXiv preprint arXiv:2505.21468}, + year={2025} +} + +@article{ulzega2025shedding, + title={Shedding light on the solar dynamo using data-driven Bayesian parameter inference}, + author={Ulzega, Simone and Beer, J{\"u}rg and Ferriz-Mas, Antonio and Dirmeier, Simon and Albert, Carlo}, + journal={The Astrophysical Journal}, + volume={992}, + number={1}, + pages={61}, + year={2025}, + publisher={The American Astronomical Society} +} + + +@inproceedings{ + sharrock2024sequential, + title={Sequential Neural Score Estimation: Likelihood-Free Inference with Conditional Score Based Diffusion Models}, + author={Louis Sharrock and Jack Simons and Song Liu and Mark Beaumont}, + booktitle={Forty-first International Conference on Machine Learning}, + year={2024} +} + +@inproceedings{ + gloeckler2024allinone, + title={All-in-one simulation-based inference}, + author={Manuel Gloeckler and Michael Deistler and Christian Dietrich Weilbach and Frank Wood and Jakob H. Macke}, + booktitle={Forty-first International Conference on Machine Learning}, + year={2024} +} + +@article{albert2025simulated, + title={Simulated Annealing ABC with multiple summary statistics}, + author={Albert, Carlo and Ulzega, Simone and Dirmeier, Simon and Scheidegger, Andreas and Bassi, Alberto and Mira, Antonietta}, + journal={arXiv preprint arXiv:2505.23261}, + year={2025} +} + + +@inproceedings{dirmeier2025simulationbased, + title={Simulation-based Inference for High-dimensional Data using Surjective Sequential Neural Likelihood Estimation}, + author={Simon Dirmeier and Carlo Albert and Fernando Perez-Cruz}, + booktitle={The 41st Conference on Uncertainty in Artificial Intelligence}, + year={2025} +} + +@article{dirmeier2024surjectors, + author = {Simon Dirmeier}, + title = {Surjectors: surjection layers for density estimation with normalizing flows}, + year = {2024}, + journal = {Journal of Open Source Software}, + publisher = {The Open Journal}, + volume = {9}, + number = {94}, + pages = {6188} +} + +@InProceedings{papamakarios2019sequential, + title = {Sequential Neural Likelihood: Fast Likelihood-free Inference with Autoregressive Flows}, + author = {Papamakarios, George and Sterratt, David and Murray, Iain}, + booktitle = {Proceedings of the 22nd International Conference on Artificial Intelligence and Statistics}, + year = {2019} +} + +@article{papamakarios2021normalizing, + title={Normalizing flows for probabilistic modeling and inference}, + author={Papamakarios, George and Nalisnick, Eric and Rezende, Danilo Jimenez and Mohamed, Shakir and Lakshminarayanan, Balaji}, + journal={The Journal of Machine Learning Research}, + volume={22}, + number={1}, + pages={2617--2680}, + year={2021}, + publisher={JMLRORG} +} + + +@inproceedings{ + wildberger2023flow, + title={Flow Matching for Scalable Simulation-Based Inference}, + author={Jonas Bernhard Wildberger and Maximilian Dax and Simon Buchholz and Stephen R Green and Jakob H. Macke and Bernhard Sch{\"o}lkopf}, + booktitle={Advances in Neural Information Processing Systems}, + year={2023}, +} + +@inproceedings{ + chen2021neural, + title={Neural Approximate Sufficient Statistics for Implicit Models}, + author={Yanzhi Chen and Dinghuai Zhang and Michael U. Gutmann and Aaron Courville and Zhanxing Zhu}, + booktitle={International Conference on Learning Representations}, + year={2021} +} + +@article{schmitt2023consistency, + title={Consistency Models for Scalable and Fast Simulation-Based Inference}, + author={Schmitt, Marvin and Pratz, Valentin and K{\"o}the, Ullrich and B{\"u}rkner, Paul-Christian and Radev, Stefan T}, + journal={arXiv preprint arXiv:2312.05440}, + year={2023} +} + +@inproceedings{chen2023learning, + title={Is learning summary statistics necessary for likelihood-free inference?}, + author={Chen, Yanzhi and Gutmann, Michael U and Weller, Adrian}, + booktitle={Proceedings of the 40th International Conference on Machine +Learning}, + year={2023} +} + + +@inproceedings{greenberg2019automatic, + title = {Automatic Posterior Transformation for Likelihood-Free Inference}, + author = {Greenberg, David and Nonnenmacher, Marcel and Macke, Jakob}, + booktitle = {Proceedings of the 36th International Conference on Machine Learning}, + year = {2019}, +} + +@inproceedings{miller2022contrastive, + title={Contrastive Neural Ratio Estimation}, + author={Benjamin Kurt Miller and Christoph Weniger and Patrick Forr{\'e}}, + booktitle={Advances in Neural Information Processing Systems}, + year={2022}, +} + +@software{deepmind2020jax, + title = {The {D}eep{M}ind {JAX} {E}cosystem}, + author = {DeepMind and Babuschkin, Igor and Baumli, Kate and Bell, Alison and Bhupatiraju, Surya and Bruce, Jake and Buchlovsky, Peter and Budden, David and Cai, Trevor and Clark, Aidan and Danihelka, Ivo and Dedieu, Antoine and Fantacci, Claudio and Godwin, Jonathan and Jones, Chris and Hemsley, Ross and Hennigan, Tom and Hessel, Matteo and Hou, Shaobo and Kapturowski, Steven and Keck, Thomas and Kemaev, Iurii and King, Michael and Kunesch, Markus and Martens, Lena and Merzic, Hamza and Mikulik, Vladimir and Norman, Tamara and Papamakarios, George and Quan, John and Ring, Roman and Ruiz, Francisco and Sanchez, Alvaro and Sartran, Laurent and Schneider, Rosalia and Sezener, Eren and Spencer, Stephen and Srinivasan, Srivatsan and Stanojevi\'{c}, Milo\v{s} and Stokowiec, Wojciech and Wang, Luyu and Zhou, Guangyao and Viola, Fabio}, + url = {http://github.com/deepmind}, + year = {2020}, +} + +@software{jax2018github, + author = {James Bradbury and Roy Frostig and Peter Hawkins and Matthew James Johnson and Chris Leary and Dougal Maclaurin and George Necula and Adam Paszke and Jake Vander{P}las and Skye Wanderman-{M}ilne and Qiao Zhang}, + title = {{JAX}: composable transformations of {P}ython+{N}um{P}y programs}, + url = {http://github.com/google/jax}, + version = {0.3.13}, + year = {2018}, +} + +@article{cabezas2024blackjax, + title={BlackJAX: Composable {B}ayesian inference in {JAX}}, + author={Alberto Cabezas and Adrien Corenflos and Junpeng Lao and Rémi Louf}, + year={2024}, + journal={arXiv preprint arXiv:2402.10797}, +} + +@article{dillon2017tensorflow, + title={Tensorflow distributions}, + author={Dillon, Joshua V and Langmore, Ian and Tran, Dustin and Brevdo, Eugene and Vasudevan, Srinivas and Moore, Dave and Patton, Brian and Alemi, Alex and Hoffman, Matt and Saurous, Rif A}, + journal={arXiv preprint arXiv:1711.10604}, + year={2017} +} + +@article{kumar2019arviz, + year = {2019}, + volume = {4}, + number = {33}, + pages = {1143}, + author = {Ravin Kumar and Colin Carroll and Ari Hartikainen and Osvaldo Martin}, + title = {ArviZ a unified library for exploratory analysis of Bayesian models in Python}, + journal = {Journal of Open Source Software} +} + +@software{haiku2020github, + author = {Tom Hennigan and Trevor Cai and Tamara Norman and Lena Martens and Igor Babuschkin}, + title = {{H}aiku: {S}onnet for {JAX}}, + url = {http://github.com/deepmind/dm-haiku}, + version = {0.0.10}, + year = {2020}, +} + +@software{dirmeier2024jrnmm, + author = {Simon Dirmeier}, + title = {jrnmm: The Jansen-Rit neural mass model SDE in JAX}, + year = {2024}, + version = {0.1.0.post1}, + url = {http://github.com/dirmeier/jrnmm} +} + +@article{bishop1994mixture, + title={Mixture density networks}, + author={Bishop, Christopher M}, + year={1994}, + publisher={Aston University} +} + +@article{hoffman2014no, + title={The No-U-Turn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo.}, + author={Hoffman, Matthew D and Gelman, Andrew and others}, + journal={Journal of Machine Learning Research}, + volume={15}, + number={1}, + pages={1593--1623}, + year={2014} +} + +@article{neal2003slice, + title={Slice sampling}, + author={Neal, Radford M}, + journal={The Annals of Statistics}, + volume={31}, + number={3}, + pages={705--767}, + year={2003} +} + +@article{blei2017variational, + title={Variational inference: A review for statisticians}, + author={Blei, David M and Kucukelbir, Alp and McAuliffe, Jon D}, + journal={Journal of the American statistical Association}, + volume={112}, + number={518}, + pages={859--877}, + year={2017}, + publisher={Taylor \& Francis} +} + +@inproceedings{lipman2023flow, + title={Flow Matching for Generative Modeling}, + author={Yaron Lipman and Ricky T Q Chen and Heli Ben-Hamu and Maximilian Nickel and Matthew Le}, + booktitle={The Eleventh International Conference on Learning Representations}, + year={2023} +} + +@inproceedings{chen2018neurao, + author = {Chen, Ricky T Q and Rubanova, Yulia and Bettencourt, Jesse and Duvenaud, David K}, + booktitle = {Advances in Neural Information Processing Systems}, + editor = {S. Bengio and H. Wallach and H. Larochelle and K. Grauman and N. Cesa-Bianchi and R. Garnett}, + title = {Neural Ordinary Differential Equations}, + year = {2018} +} + +@book{sugiyama2012density, + title={Density ratio estimation in machine learning}, + author={Sugiyama, Masashi and Suzuki, Taiji and Kanamori, Takafumi}, + year={2012}, + publisher={Cambridge University Press} +} + +@inproceedings{paszke2019pytorch, + author = {Paszke, Adam and Gross, Sam and Massa, Francisco and Lerer, Adam and Bradbury, James and Chanan, Gregory and Killeen, Trevor and Lin, Zeming and Gimelshein, Natalia and Antiga, Luca and Desmaison, Alban and Kopf, Andreas and Yang, Edward and DeVito, Zachary and Raison, Martin and Tejani, Alykhan and Chilamkurthy, Sasank and Steiner, Benoit and Fang, Lu and Bai, Junjie and Chintala, Soumith}, + booktitle = {Advances in Neural Information Processing Systems}, + title = {PyTorch: An Imperative Style, High-Performance Deep Learning Library}, + year = {2019}, +} + +@article{tejero-cantero2020sbi, + year = {2020}, + publisher = {The Open Journal}, + volume = {5}, number = {52}, + pages = {2505}, + author = {Alvaro Tejero-Cantero and Jan Boelts and Michael Deistler and Jan-Matthis Lueckmann and Conor Durkan and Pedro J. Gonçalves and David S. Greenberg and Jakob H. Macke}, + title = {sbi: A toolkit for simulation-based inference}, + journal = {Journal of Open Source Software} +} + +@book{sisson2018handbook, + title={Handbook of approximate Bayesian computation}, + author={Sisson, Scott A and Fan, Yanan and Beaumont, Mark}, + year={2018}, + publisher={CRC Press} +} + +@book{brooks2011mcmc, + author = {Brooks, Steve and Gelman, Andrew and Jones, Galin and Meng, Xiao-Li}, + publisher = {CRC press}, + title = {Handbook of Markov Chain Monte Carlo}, + year = {2011}, +} + +@article{albert2015simulated, + title={A simulated annealing approach to approximate {B}ayes computations}, + author={Albert, Carlo and K{\"u}nsch, Hans R and Scheidegger, Andreas}, + journal={Statistics and Computing}, + volume={25}, + pages={1217--1232}, + year={2015}, + publisher={Springer} +} + +@article{beaumont2009adaptive, + title={Adaptive approximate Bayesian computation}, + author={Beaumont, Mark A and Cornuet, Jean-Marie and Marin, Jean-Michel and Robert, Christian P}, + journal={Biometrika}, + volume={96}, + number={4}, + pages={983--990}, + year={2009}, + publisher={Oxford University Press} +} + +@article{del2012adaptive, + title={An adaptive sequential Monte Carlo method for approximate Bayesian computation}, + author={Del Moral, Pierre and Doucet, Arnaud and Jasra, Ajay}, + journal={Statistics and Computing}, + volume={22}, + pages={1009--1020}, + year={2012}, + publisher={Springer} +} + +@article{vehtari2021split, + author = {Aki Vehtari and Andrew Gelman and Daniel Simpson and Bob Carpenter and Paul-Christian B{\"u}rkner}, + title = {{Rank-Normalization, Folding, and Localization: An Improved $\widehat{R}$ for Assessing Convergence of MCMC (with Discussion)}}, + volume = {16}, + journal = {Bayesian Analysis}, + number = {2}, + publisher = {International Society for Bayesian Analysis}, + pages = {667 -- 718}, + year = {2021}, +} + +@article{hunter2007mat, + Author = {Hunter, J. D.}, + Title = {Matplotlib: A 2D graphics environment}, + Journal = {Computing in Science \& Engineering}, + Volume = {9}, + Number = {3}, + Pages = {90--95}, + publisher = {IEEE COMPUTER SOC}, + year = {2006} +} + + + +@article{besag1994comments, + title={Comments on “Representations of knowledge in complex systems” by U. Grenander and MI Miller}, + author={Besag, Julian}, + journal = {Journal of the Royal Statistical Society Series B: Statistical Methodology}, + volume={56}, + number={591-592}, + pages={4}, + year={1994} +} + +@inproceedings{he2016deep, + author = {He, Kaiming and Zhang, Xiangyu and Ren, Shaoqing and Sun, Jian}, + title = {Deep Residual Learning for Image Recognition}, + booktitle = {Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR)}, + year = {2016} +} + +@inproceedings{durkan2019neural, + title={Neural spline flows}, + author={Durkan, Conor and Bekasov, Artur and Murray, Iain and Papamakarios, George}, + booktitle={Advances in Neural Information Processing Systems}, + year={2019} +} + +@inproceedings{papamakarios2017masked, + title={Masked autoregressive flow for density estimation}, + author={Papamakarios, George and Pavlakou, Theo and Murray, Iain}, + booktitle={Advances in Neural Information Processing Systems}, + year={2017} +} + +@inproceedings{dinh2017density, + title={Density estimation using Real {NVP}}, + author={Laurent Dinh and Jascha Sohl-Dickstein and Samy Bengio}, + booktitle={International Conference on Learning Representations}, + year={2017} +} + +@inproceedings{song2023consistency, + title={Consistency Models}, + author={Yang Song and Prafulla Dhariwal and Mark Chen and Ilya Sutskever}, + year={2023}, + booktitle={Proceedings of the 40th International Conference on Machine Learning}, +} + +@article{cranmer2020frontier, + title={The frontier of simulation-based inference}, + author={Cranmer, Kyle and Brehmer, Johann and Louppe, Gilles}, + journal={Proceedings of the National Academy of Sciences}, + volume={117}, + number={48}, + pages={30055--30062}, + year={2020}, + publisher={National Acadademy of Sciences} +} + +@article{harris2020array, + title={Array programming with NumPy}, + author={Harris, Charles R and Millman, K Jarrod and Van Der Walt, St{\'e}fan J and Gommers, Ralf and Virtanen, Pauli and Cournapeau, David and Wieser, Eric and Taylor, Julian and Berg, Sebastian and Smith, Nathaniel J and others}, + journal={Nature}, + volume={585}, + number={7825}, + pages={357--362}, + year={2020}, + publisher={Nature Publishing Group UK London} +} + +@misc{gabry2024bayesplot, + title = {bayesplot: Plotting for Bayesian Models}, + author = {Jonah Gabry and Tristan Mahr}, + year = {2024}, + note = {R package version 1.11.0}, + url = {https://mc-stan.org/bayesplot/}, +} + +@article{gabry2019visualization, + title = {Visualization in Bayesian workflow}, + author = {Jonah Gabry and Daniel Simpson and Aki Vehtari and Michael Betancourt and Andrew Gelman}, + year = {2019}, + journal = {J. R. Stat. Soc. A}, + volume = {182}, + issue = {2}, + pages = {389-402}, +} + +@inproceedings{hermans2020likelihood, + title = {Likelihood-free {MCMC} with Amortized Approximate Ratio Estimators}, + author = {Hermans, Joeri and Begy, Volodimir and Louppe, Gilles}, + booktitle = {Proceedings of the 37th International Conference on Machine Learning}, + year = {2020}, +} + + +@article{hoyer2017xarray, + title = {xarray: {N-D} labeled arrays and datasets in {Python}}, + author = {Hoyer, S. and J. Hamman}, + journal = {Journal of Open Research Software}, + volume = {5}, + number = {1}, + year = {2017} +} + +@inproceedings{yao2023discriminative, + author = {Yao, Yuling and Domke, Justin}, + booktitle = {Advances in Neural Information Processing Systems}, + title = {Discriminative Calibration: Check Bayesian Computation from Simulations and Flexible Classifier}, + year = {2023} +} + +@inproceedings{linhart2023lc2st, + author = {Linhart, Julia and Gramfort, Alexandre and Rodrigues, Pedro}, + booktitle = {Advances in Neural Information Processing Systems}, + title = {L-C2ST: Local Diagnostics for Posterior Approximations in Simulation-Based Inference}, + year = {2023} +} + +@inproceedings{zhao2022comparing, + title={Comparing Distributions by Measuring Differences that Affect Decision Making}, + author={Shengjia Zhao and Abhishek Sinha and Yutong He and Aidan Perreault and Jiaming Song and Stefano Ermon}, + booktitle={International Conference on Learning Representations}, + year={2022} +} + +@article{sisson2007sequential, + title={Sequential Monte Carlo without likelihoods}, + author={Sisson, Scott A and Fan, Yanan and Tanaka, Mark M}, + journal={Proceedings of the National Academy of Sciences}, + volume={104}, + number={6}, + pages={1760--1765}, + year={2007}, +} + +@article{charbonneau2005fluctuations, + title={Fluctuations in {B}abcock-{L}eighton dynamos. I. period doubling and transition to chaos}, + author={Charbonneau, Paul and St-Jean, C{\'e}dric and Zacharias, Pia}, + journal={The Astrophysical Journal}, + volume={619}, + number={1}, + pages={613}, + year={2005}, + publisher={IOP Publishing} +} + +@inproceedings{ + liu2023flow, + title={Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow}, + author={Xingchao Liu and Chengyue Gong and qiang liu}, + booktitle={The Eleventh International Conference on Learning Representations }, + year={2023} +} + +@article{phan2019composable, + title={Composable effects for flexible and accelerated probabilistic programming in NumPyro}, + author={Phan, Du and Pradhan, Neeraj and Jankowiak, Martin}, + journal={arXiv preprint arXiv:1912.11554}, + year={2019} +} + +@InProceedings{gloeckler2024allinone, + title = {All-in-one simulation-based inference}, + author = {Gloeckler, Manuel and Deistler, Michael and Weilbach, Christian Dietrich and Wood, Frank and Macke, Jakob H.}, + booktitle = {Proceedings of the 41st International Conference on Machine Learning}, + year = {2024}, +} + +@InProceedings{h2024quantile, + title = {Simulation-Based Inference with Quantile Regression}, + author = {Jia, He}, + booktitle = {Proceedings of the 41st International Conference on Machine Learning}, + year = {2024}, +} + +@article{tierney1994markov, + title={Markov chains for exploring posterior distributions}, + author={Tierney, Luke}, + journal={The Annals of Statistics}, + pages={1701--1728}, + year={1994} +} + +@book{compbayes2019, + title={Computational Bayesian Statistics: An Introduction}, + publisher={Cambridge University Press}, + author={Amaral Turkman, M. Antónia and Paulino, Carlos Daniel and Müller, Peter}, + year={2019} +} + +@article{schaelte2022, + title = {pyABC: Efficient and robust easy-to-use approximate Bayesian computation}, + author = {Yannik Schälte and Emmanuel Klinger and Emad Alamoudi and Jan Hasenauer}, + year = 2022, + journal = {Journal of Open Source Software}, + publisher = {The Open Journal}, + volume = 7, + number = 74, + pages = 4304, +} + + +@article{dutta2021abcpy, + title={ABCpy: A High-Performance Computing Perspective to Approximate Bayesian Computation}, + volume={100}, + number={7}, + journal={Journal of Statistical Software}, + author={Dutta, Ritabrata and Schoengens, Marcel and Pacchiardi, Lorenzo and Ummadisingu, Avinash and Widmer, Nicole and Künzli, Pierre and Onnela, Jukka-Pekka and Mira, Antonietta}, + year={2021}, + pages={1–38} +} + +@article{approxbayescomp, + title={Approximate Bayesian Computations to fit and compare insurance loss models}, + author={Goffard, Pierre-Olivier and Laub, Patrick J}, + journal={Insurance: Mathematics and Economics}, + volume={100}, + pages={350--371}, + year={2021} +} + +@article{rousset2017summary, + title={The summary-likelihood method and its implementation in the Infusion package}, + author={Rousset, Francois and Gouy, Alexandre and Martinez-Almoyna, Camille and Courtiol, Alexandre}, + journal={Molecular ecology resources}, + volume={17}, + number={1}, + pages={110--119}, + year={2017}, + publisher={Wiley Online Library} +} + +@article{rodrigues2021hnpe, + title={HNPE: Leveraging global parameters for neural posterior estimation}, + author={Rodrigues, Pedro and Moreau, Thomas and Louppe, Gilles and Gramfort, Alexandre}, + journal={Advances in Neural Information Processing Systems}, + volume={34}, + pages={13432--13443}, + year={2021} +} + +@article{ableidinger2017stochastic, + author = {Ableidinger, Markus and Buckwar, Evelyn and Hinterleitner, Harald}, + journal = {The Journal of Mathematical Neuroscience}, + number = {1}, + pages = {8}, + title = {A Stochastic Version of the {J}ansen and {R}it Neural Mass Model: Analysis and Numerics}, + volume = {7}, + year = {2017} +} + +@article{buckwar2020spectral, + title={Spectral density-based and measure-preserving {ABC} for partially observed diffusion processes. An illustration on {H}amiltonian {SDE}s}, + author={Buckwar, Evelyn and Tamborrino, Massimiliano and Tubikanec, Irene}, + journal={Statistics and Computing}, + volume={30}, + number={3}, + pages={627--648}, + year={2020} +} + +@dataset{cattan2018, + author = {Grégoire Cattan and + Pedro L. C. Rodrigues and + Marco Congedo}, + title = {EEG Alpha Waves dataset}, + year = 2018, + publisher = {Zenodo}, + doi = {10.5281/zenodo.2348892} +} + +@article{virtanen2020scipy, + title={SciPy 1.0: fundamental algorithms for scientific computing in Python}, + author={Virtanen, Pauli and Gommers, Ralf and Oliphant, Travis E and Haberland, Matt and Reddy, Tyler and Cournapeau, David and Burovski, Evgeni and Peterson, Pearu and Weckesser, Warren and Bright, Jonathan and others}, + journal={Nature methods}, + volume={17}, + number={3}, + pages={261--272}, + year={2020}, + publisher={Nature Publishing Group} +} + +@inproceedings{ + sharrock2024sequential, + title={Sequential Neural Score Estimation: Likelihood-Free Inference with Conditional Score Based Diffusion Models}, + author={Louis Sharrock and Jack Simons and Song Liu and Mark Beaumont}, + booktitle={Forty-first International Conference on Machine Learning}, + year={2024} +} diff --git a/paper/paper.md b/paper/paper.md new file mode 100644 index 0000000..2eaaca9 --- /dev/null +++ b/paper/paper.md @@ -0,0 +1,125 @@ +--- +title: 'Simulation-based Inference with the Python Package sbijax' +tags: + - Python + - Simulation-based Inference + - Approximate Bayesian Computation + - Machine learning +authors: + - name: Simon Dirmeier + affiliation: "1, 2" + - name: Antonietta Mira + affiliation: "3, 4" + - name: Carlo Albert + affiliation: "5" +affiliations: + - name: Swiss Data Science Center, Zurich, Switzerland + index: 1 + - name: ETH Zurich, Zurich, Switzerland + index: 2 + - name: Università della Svizzera italiana, Switzerland + index: 3 + - name: University of Insubria, Italy + index: 4 + - name: Swiss Federal Institute of Aquatic Science and Technology, Switzerland + index: 5 +date: 19 March 2026 +bibliography: paper.bib +--- + +# Summary + +Neural simulation-based inference (SBI) describes an emerging family of methods for Bayesian inference for simulator models that use neural networks as surrogate models. +Here we introduce `sbijax`, a Python package that implements a wide variety of state-of-the-art methods in neural simulation-based inference using a user-friendly +programming interface. sbijax offers high-level functionality to quickly construct SBI estimators, and compute and visualize posterior distributions with only a few lines of code. +In addition, the package provides functionality for conventional approximate Bayesian computation, to compute model diagnostics, and to automatically estimate summary +statistics. By virtue of being entirely written in `JAX`, sbijax is extremely computationally efficient, allowing rapid training of neural networks and executing code automatically in parallel on both CPU and GPU. + +# Statement of Need + +Modern approaches to neural simulation-based inference (SBI) utilize recent developments in neural density estimation or score-based generative modelling to build surrogate models to approximate Bayesian posterior distributions. +Similarly to conventional methods, such as approximate Bayesian computation (ABC) and its sequential (SMC-ABC) and annealing-based (e.g., SABC) variants, neural SBI methods infer this posterior distribution by first simulating synthetic +data and then numerically constructing an appropriate approximation to this pseudo data set. SBI methods are attractive for a couple of reasons. +On the one hand, this family of methods has been shown to be more computationally efficient and often more accurate than ABC methods, in particular for smaller simulation budgets. On the other hand, +SBI allows to easily amortize inference, i.e., to infer the posterior distribution for multiple different observations once a neural model has been trained. + +Here we propose `sbijax`, a Python package implementing state-of-the-art methodology of neural simulation-based inference. +While the main focus of the package is the implementation of recent algorithms to make them available to practitioners, e.g., @wildberger2023flow or @schmitt2023consistency, +`sbijax` also implements common methods from approximate Bayesian computation, e.g., SMC-ABC [@beaumont2009adaptive], to have the entire SBI toolbox in one efficient package (see Table \ref{tbl-methods} for an overview). +In addition, `sbijax` provides functionality for model diagnostics, posterior visualization and Markov Chain Monte Carlo (MCMC) sampling. + +The package uses the high-performance computing framework `JAX` as a backend [@jax2018github]. +Using `JAX` has several advantages, including a) that it uses the same syntax as `numpy` [@harris2020array] which enables a seamless transition for applied scientists who already are familiar with it, +and b) that empirical evaluations have shown that `JAX` can be significantly faster than `PyTorch` (see, e.g., @phan2019composable). Our package heavily builds on libraries from +the `JAX`-verse and common Bayesian inference tools. Specifically, we use `Haiku` [@haiku2020github] to construct and train neural networks, `ArviZ` [@kumar2019arviz] for various model visualizations, `surjectors` for normalizing-flow based +density estimation [@dirmeier2024surjectors], TensorFlow Probability [@dillon2017tensorflow] to define statistical distributions, and `BlackJAX` [@cabezas2024blackjax] for posterior sampling using Markov Chain Monte Carlo. + +| **Model** | **Class name** | **Reference** | +|------------------------------------------------|------------|-----------------------------| +| Sequential Monte Carlo ABC | `SMCABC` | @beaumont2009adaptive | +| Neural likelihood estimation | `NLE` | @papamakarios2019sequential | +| Surjective neural likelihood estimation | `SNLE` | @dirmeier2025simulationbased | +| Automatic posterior transformation | `NPE` | @greenberg2019automatic | +| Contrastive neural ratio estimation | `NRE` | @miller2022contrastive | +| Flow matching posterior estimation | `FMPE` | @wildberger2023flow | +| Posterior Score Estimation | `NPSE` | @sharrock2024sequential | +| All-In-One Posterior Estimation | `AIO` | @gloeckler2024allinone | +| Consistency model posterior estimation | `CMPE` | @schmitt2023consistency | +| Neural approximate sufficient statistics | `NASS` | @chen2021neural | +| Neural approximate slice sufficient statistics | `NASSS` | @chen2023learning | + +:Implemented SBI methods in `sbijax` \label{tbl-methods}. + +# State of the field + +While a plethora of different models has been proposed in the recent literature, the development of adequate software packages has not followed at the same pace, +and only few packages exist that allow modelers to use these methods. Most prominently, the Python package `sbi` [@tejero-cantero2020sbi] +implements several approaches for neural simulation-based inference, such as a neural posterior, likelihood-ratio, and likelihood estimation [@cranmer2020frontier] utilizing a `PyTorch` backend [@paszke2019pytorch]. +The package additionally provides an API for model diagnostics, such as posterior predictive checks, effective sample size computations and simulation-based calibration. +However, the package lacks implementations of recent developments which pose the state-of-the-art in the field, such as by @chen2023learning, @dirmeier2025simulationbased or @schmitt2023consistency. +Also, by virtue of being developed in `PyTorch` it is potentially restrictive to practitioners that do not have experience with it. +For approximate Bayesian computation, several `Python` packages are available. In particular `abcpy` [@dutta2021abcpy] implements a multitude of different ABC algorithms. +However, none of these packages implement modern (neural) SBI methods. + +# Software design + +`Sbijax` is designed as a modular and extensible toolbox for SBI in `JAX`, +combining a functional programming philosophy with a flexible interface for both +expert users and domain scientists. Its design is guided by the following principles: + +1) Faithful alignment with `JAX`’s functional paradigm. +Sbijax adheres closely to the low-level, functional programming model of `JAX`. +While SBI methods are implemented using an object-oriented structure, all member functions return +immutable objects rather than modifying internal state. +This avoids hidden side effects, facilitates composability, and ensures compatibility with `JAX` transformations such as `jit`, `vmap`, and `pmap`. + +2) Separation of inference methods and neural network implementations. +Sbijax focuses on implementing SBI algorithms, while neural network components are defined using `Haiku`. +This allows users to construct models directly with `Haiku`, and to incorporate probabilistic building blocks from `Distrax` + [@deepmind2020jax] and `surjectors` [@dirmeier2024surjectors]. + As a result, model definitions remain flexible and low-level, while seamlessly interoperating with the broader `JAX` ecosystem. + +3) Support for extensibility and research. +Sbijax is structured to facilitate experimentation with new SBI methods. +Its modular design allows components such as neural architectures, training objectives, and sampling strategies to be easily replaced or extended, making it suitable as both a research framework and a practical toolbox. + +4) Accessibility for domain scientists. +In addition to its flexibility, Sbijax includes pre-implemented models with sensible defaults, enabling use without in-depth expertise in deep learning. +In the simplest case, users only need to define a prior and a simulator to run inference workflows—often in as few as five lines of code. + +# Research impact statement + +`sbijax` has already been used extensively in the Machine Learning literature. @dirmeier2025simulationbased and @dirmeier2025causal proposed novel SBI methods for +surjective neural likelihood estimation and posterior estimation using causal constraints, respectively, where they used `sbijax` heavily for their experimental evaluations. +@ulzega2025shedding used `sbijax` to infer the posterior distribution of a complicated Bayesian model from the astrophysics literature. +@albert2025simulated developed a novel ABC method that uses `sbijax` for model evaluation. + +# AI usage disclosure + +No GenAI or other AI tools have been used in writing the software or this manuscript. + +# Acknowledgements + +This research was supported by the Swiss National Science Foundation (Grant No. $200021\_208249$). + +# References diff --git a/uv.lock b/uv.lock index 5a3810f..49d483c 100644 --- a/uv.lock +++ b/uv.lock @@ -789,7 +789,7 @@ wheels = [ [[package]] name = "distrax" -version = "0.1.5" +version = "0.1.7" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "absl-py" }, @@ -797,12 +797,11 @@ dependencies = [ { name = "jax" }, { name = "jaxlib" }, { name = "numpy" }, - { name = "setuptools" }, - { name = "tensorflow-probability" }, + { name = "tfp-nightly" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/81/a7/d87f5c5caf2ae9e28a9f506274870c25eb7e38c1ce0a1b7b3696afc73ac9/distrax-0.1.5.tar.gz", hash = "sha256:ec41522d389af69efedc8d475a7e6d8f229429c00f2140dcd641feacf7e21948", size = 190162, upload-time = "2023-11-22T12:53:44.836Z" } +sdist = { url = "https://files.pythonhosted.org/packages/21/a1/eb2086d2fc27ddf627f47af392d67e7e023a153d1b00087d3d12eab0465a/distrax-0.1.7.tar.gz", hash = "sha256:78deec181894e811d416b34e6b500f5e8f6d39bc850bd6f3152ee9e2164c9bef", size = 183672, upload-time = "2025-09-01T23:06:00.108Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/d8/42/efa7985eefd516f273bc5557550539ebd687cc6a3a2a176d5e1f5b0398ca/distrax-0.1.5-py3-none-any.whl", hash = "sha256:5020f4b53a9a480d019c12e44292fbacb7de857cce478bc594dacf29519c61b7", size = 319716, upload-time = "2023-11-22T12:53:42.449Z" }, + { url = "https://files.pythonhosted.org/packages/a2/3e/e67008a232c37f98f1563469e206b1f2820f01c9cce1118c02215f1bf361/distrax-0.1.7-py3-none-any.whl", hash = "sha256:cf502d64022c7a682882ae6f07aea35206c6a3d8ddfb1e3beb006723baf283fd", size = 312680, upload-time = "2025-09-01T23:05:58.175Z" }, ] [[package]] @@ -3572,24 +3571,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/99/55/db07de81b5c630da5cbf5c7df646580ca26dfaefa593667fc6f2fe016d2e/tabulate-0.10.0-py3-none-any.whl", hash = "sha256:f0b0622e567335c8fabaaa659f1b33bcb6ddfe2e496071b743aa113f8774f2d3", size = 39814, upload-time = "2026-03-04T18:55:31.284Z" }, ] -[[package]] -name = "tensorflow-probability" -version = "0.23.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "absl-py" }, - { name = "cloudpickle" }, - { name = "decorator" }, - { name = "dm-tree", version = "0.1.8", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" }, - { name = "dm-tree", version = "0.1.9", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13'" }, - { name = "gast" }, - { name = "numpy" }, - { name = "six" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/1e/90/4e8c686f2e691f48e40e16a539c61a6e9880743733d8c4dc3f275d12268e/tensorflow_probability-0.23.0-py2.py3-none-any.whl", hash = "sha256:dda5cacfe50cb19ecd96f3ce81e6ff8680d84213bcfe94ca0aaf6e5f51c88061", size = 6915575, upload-time = "2023-11-20T23:21:06.624Z" }, -] - [[package]] name = "terminado" version = "0.18.1"