Merge branch 'dev' into ssm-wrapper

Author: Chase-Grajeda

.github/workflows/multiversion-docs.yaml

@@ -40,7 +40,7 @@ jobs:
       - name: Create local branches
         run: |
           cd ./dev
-          git branch master remotes/origin/master
+          git branch main remotes/origin/main
       - name: Make the Sphinx docs
         run: |
           cd ./dev/docsrc

.github/workflows/tests.yaml

@@ -5,7 +5,7 @@ on:
   pull_request:
   push:
     branches:
-      - master
+      - main
       - dev
       - streamlined-backend
 

README.md

@@ -4,8 +4,8 @@
 [![DOI](https://img.shields.io/badge/DOI-10.21105%2Fjoss.05702-blue?style=for-the-badge)](https://doi.org/10.21105/joss.05702)
 ![PyPI - License](https://img.shields.io/pypi/l/bayesflow?style=for-the-badge)
 
-BayesFlow is a Python library for simulation-based **Amortized Bayesian Inference** with neural networks.
-It provides users with:
+BayesFlow 2 is a Python library for simulation-based **Amortized Bayesian Inference** with neural networks.
+It provides users and researchers with:
 
 - A user-friendly API for rapid Bayesian workflows
 - A rich collection of neural network architectures
@@ -56,9 +56,8 @@ For an in-depth exposition, check out our walkthrough notebooks below. More tuto
 3. [Two moons starter example](examples/Two_Moons_Starter.ipynb)
 4. [Rapid iteration with point estimators](examples/Lotka_Volterra_point_estimation_and_expert_stats.ipynb)
 5. [SIR model with custom summary network](examples/SIR_Posterior_Estimation.ipynb)
-6. [Hyperparameter optimization](examples/Hyperparameter_Optimization.ipynb)
-7. [Bayesian experimental design](examples/Bayesian_Experimental_Design.ipynb)
-8. [Simple model comparison example](examples/One_Sample_TTest.ipynb)
+6. [Bayesian experimental design](examples/Bayesian_Experimental_Design.ipynb)
+7. [Simple model comparison example](examples/One_Sample_TTest.ipynb)
 
 ## Install
 
@@ -126,9 +125,8 @@ Documentation is available at https://bayesflow.org. Please use the [BayesFlow F
 
 You can cite BayesFlow along the lines of:
 
-- We approximated the posterior with neural posterior estimation and learned summary statistics (NPE; Radev et al., 2020), as implemented in the BayesFlow software for amortized Bayesian workflows (Radev et al., 2023a).
-- We approximated the likelihood with neural likelihood estimation (NLE; Papamakarios et al., 2019) without hand-crafted summary statistics, as implemented in the BayesFlow software for amortized Bayesian workflows (Radev et al., 2023b).
-- We performed simultaneous posterior and likelihood estimation with jointly amortized neural approximation (JANA; Radev et al., 2023a), as implemented in the BayesFlow software for amortized Bayesian workflows (Radev et al., 2023b).
+- We approximated the posterior using neural posterior estimation (NPE) with learned summary statistics (Radev et al., 2020), as implemented in the BayesFlow framework for amortized Bayesian inference (Radev et al., 2023a).
+- We approximated the likelihood using neural likelihood estimation (NLE) without hand-crafted summary statistics (Papamakarios et al., 2019), leveraging its implementation in BayesFlow for efficient and flexible inference.
 
 1. Radev, S. T., Schmitt, M., Schumacher, L., Elsemüller, L., Pratz, V., Schälte, Y., Köthe, U., & Bürkner, P.-C. (2023a). BayesFlow: Amortized Bayesian workflows with neural networks. *The Journal of Open Source Software, 8(89)*, 5702.([arXiv](https://arxiv.org/abs/2306.16015))([JOSS](https://joss.theoj.org/papers/10.21105/joss.05702))
 2. Radev, S. T., Mertens, U. K., Voss, A., Ardizzone, L., Köthe, U. (2020). BayesFlow: Learning complex stochastic models with invertible neural networks. *IEEE Transactions on Neural Networks and Learning Systems, 33(4)*, 1452-1466. ([arXiv](https://arxiv.org/abs/2003.06281))([IEEE TNNLS](https://ieeexplore.ieee.org/document/9298920))
@@ -169,6 +167,10 @@ You can cite BayesFlow along the lines of:
 }
 ```
 
+## Awesome Amortized Inference
+
+If you are interested in a curated list of resources, including reviews, software, papers, and other resources related to amortized inference, feel free to explore our [community-driven list](https://github.com/bayesflow-org/awesome-amortized-inference).
+
 ## Acknowledgments
 
 This project is currently managed by researchers from Rensselaer Polytechnic Institute, TU Dortmund University, and Heidelberg University. It is partially funded by the Deutsche Forschungsgemeinschaft (DFG, German Research Foundation, Project 528702768). The project is further supported by Germany's Excellence Strategy -- EXC-2075 - 390740016 (Stuttgart Cluster of Excellence SimTech) and EXC-2181 - 390900948 (Heidelberg Cluster of Excellence STRUCTURES), as well as the Informatics for Life initiative funded by the Klaus Tschira Foundation.

bayesflow/__init__.py

@@ -33,14 +33,6 @@ def setup():
 
     from bayesflow.utils import logging
 
-    if keras.backend.backend() == "torch":
-        # turn off gradients by default
-        import torch
-
-        torch.autograd.set_grad_enabled(False)
-
-        logging.warning("Disabling gradients by default. Use\nwith torch.enable_grad():\nin custom training loops.")
-
     logging.debug(f"Using backend {keras.backend.backend()!r}")
 
 

bayesflow/workflows/basic_workflow.py

@@ -33,8 +33,8 @@ def __init__(
         checkpoint_name: str = "model",
         save_weights_only: bool = False,
         save_best_only: bool = False,
-        inference_variables: Sequence[str] | str = "theta",
-        inference_conditions: Sequence[str] | str = "x",
+        inference_variables: Sequence[str] | str = None,
+        inference_conditions: Sequence[str] | str = None,
         summary_variables: Sequence[str] | str = None,
         standardize: Sequence[str] | str = "inference_variables",
         **kwargs,
@@ -72,10 +72,10 @@ def __init__(
             as some losses (e.g. flow matching) do not reliably reflect model performance, and outliers in the
             validation data can cause unwanted effects.
         inference_variables : Sequence[str] or str, optional
-            Variables for inference as a sequence of strings or a single string (default is "theta").
+            Variables for inference as a sequence of strings or a single string (default is None).
             Important for automating diagnostics!
         inference_conditions : Sequence[str] or str, optional
-            Variables used as conditions for inference (default is "x").
+            Variables used as conditions for inference (default is None).
         summary_variables : Sequence[str] or str, optional
             Variables for summarizing data, if any (default is None).
         standardize : Sequence[str] or str, optional

docsrc/poly.py

@@ -17,7 +17,7 @@
 root = Git.root(Path(__file__).parent)
 
 #: CodeRegex matching the branches to build docs for
-# BRANCH_REGEX = r"^(master|dev)$"
+# BRANCH_REGEX = r"^(main|dev)$"
 BRANCH_REGEX = r"^(dev)$"
 
 #: Regex matching the tags to build docs for
@@ -76,7 +76,7 @@ def data(driver, rev, env):
     branches, tags = refs_by_type(revisions)
     latest = max(tags or branches)
     for b in branches:
-        if b.name == "master":
+        if b.name == "main":
             latest = b
             break
 
@@ -95,7 +95,7 @@ def root_data(driver):
     branches, tags = refs_by_type(revisions)
     latest = max(tags or branches)
     for b in branches:
-        if b.name == "master":
+        if b.name == "main":
             latest = b
             break
     return {"revisions": revisions, "latest": latest}
@@ -123,8 +123,8 @@ async def selector(rev, keys):
     # map all v1 revisions to one configuration
     if rev.name.startswith("v1."):
         return "v1.1.6"
-    elif rev.name in ["master"]:
-        # special configuration for v1 master branch
+    elif rev.name in ["main"]:
+        # special configuration for v1 main branch
         return rev.name
     elif rev.name == "local":
         return "local"

examples/Bayesian_Experimental_Design.ipynb

@@ -15,10 +15,12 @@
     "In this tutorial, we will:\n",
     "\n",
     "- Introduce the core concepts in static (i.e. non-adaptive) BED\n",
-    "- Show how to use BayesFlow to learn the optimal designs for a simple chemical reaction system\n",
+    "- Show how to use low-level BayesFlow features to learn the optimal designs for a simple chemical reaction system\n",
     "- Demonstrate the multi-backend capabilities of BayesFlow\n",
     "- Use PyTorch components (model, optimisers, etc.) via Keras3 through BayesFlow\n",
-    "- Use a normalizing flow from BayesFlow with a custom loss function that's written in pure PyTorch\n"
+    "- Use a normalizing flow from BayesFlow with a custom loss function that's written in pure PyTorch\n",
+    "\n",
+    "This tutorial is aimed at intermediate / advanced users who want to create custom training loops and capitalize on the stable generative network implementations."
    ]
   },
   {
@@ -90,20 +92,24 @@
     "import torch.distributions as dist\n",
     "import torch.nn as nn\n",
     "\n",
-    "# import keras\n",
-    "from keras.src.backend.common import global_state\n",
-    "\n",
-    "global_state.set_global_attribute(\"torch_device\", \"cpu\")\n",
-    "\n",
     "# for BayesFlow devs: this ensures that the latest dev version can be found\n",
     "# this is not required if you have BayesFlow installed (e.g., via pip)\n",
     "import sys\n",
     "sys.path.append(\"../\")\n",
     "\n",
     "import bayesflow as bf\n",
-    "from cmdstanpy import CmdStanModel\n",
     "\n",
+    "from cmdstanpy import CmdStanModel"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
     "import logging\n",
+    "\n",
     "# disable the long printouts from stan.\n",
     "logger = logging.getLogger(\"cmdstanpy\")\n",
     "logger.addHandler(logging.NullHandler())\n",
@@ -533,10 +539,11 @@
    ],
    "source": [
     "# sample from the amortised posterior\n",
+    "num_samples = 2000\n",
     "amortised_posterior_samples = torch.stack(\n",
     "    [\n",
-    "        posterior_net.sample((2000,), conditions=torch.cat(\n",
-    "                [yy.unsqueeze(0).expand(2000, -1), simulator.designs.expand(2000, -1)], dim=-1\n",
+    "        posterior_net.sample(num_samples, conditions=torch.cat(\n",
+    "                [yy.unsqueeze(0).expand(num_samples, -1), simulator.designs.expand(num_samples, -1)], dim=-1\n",
     "            ),\n",
     "        )\n",
     "        for yy in test_sims[\"scaled_y\"]\n",
@@ -701,7 +708,7 @@
     "designed_amortised_samples = torch.stack(\n",
     "    [\n",
     "        posterior_net_designs.sample(\n",
-    "            (2000,), \n",
+    "            batch_shape=2000, \n",
     "            conditions=torch.cat(\n",
     "                [yy.unsqueeze(0).expand(2000, -1), simulator.designs.expand(2000, -1)], dim=-1\n",
     "            ),\n",

examples/From_ABC_to_BayesFlow.ipynb

@@ -4,8 +4,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# From pyABC to BayesFlow\n",
-    "## Markov Jump Process: Reaction Network\n",
+    "# From ABC to BayesFlow\n",
     "_Author: Jonas Arruda_"
    ]
   },
@@ -17,7 +16,7 @@
    "source": [
     "In the following, we fit stochastic chemical reaction kinetics with Approximate Bayesian Computation (ABC) and show how to transfer the workflow to `BayesFlow`. Any model written in an ABC framework can be easily transferred to `BayesFlow`. The main difference is that `BayesFlow` enables *amortized inference*, allowing us to instantly infer parameters from new data sets without further training.\n",
     "\n",
-    "In ABC, inference is tailored for a single dataset. For both we need to specify a simulator and priors over the parameters. While in ABC we also need to specify a distance function between observation and simulation (and for higher dimensional problems also summary statistics), in `BayesFlow` we only need to specify an adapter that maps the simulator output to the output of a neural `Approximator`. After training, we can validate the computational faithfulness of the trained approximator using various Bayesian metrics, such as simulation-based calibration (SBC). This is typically not feasible with ABC, as we would have to repeat the same expensive approximation loop for every new data set.\n",
+    "In ABC (and non-amortized inference in general), estimation is tailored for a single dataset. For both we need to specify a simulator and priors over the parameters. While in ABC we also need to specify a distance function between observation and simulation (and for higher dimensional problems also summary statistics), in `BayesFlow` we only need to specify an adapter that maps the simulator output to the output of a neural `Approximator`. After training, we can validate the computational faithfulness of the trained approximator using various Bayesian metrics, such as simulation-based calibration (SBC). This is typically not feasible with ABC, as we would have to repeat the same expensive approximation loop for every new data set.\n",
     "\n",
     "For this example, we need to install [pyABC](https://github.com/icb-dcm/pyabc), which is our go-to library for ABC. The example is taken from [the documentation](https://pyabc.readthedocs.io/en/latest/examples/chemical_reaction.html). This tutorial starts with the ABC implementation and then demonstrates how to easily switch to BayesFlow. Readers familiar with pyABC can fast-forward to the second part of the notebook."
    ]

examples/Linear_Regression_Starter.ipynb

@@ -26,31 +26,6 @@
     "At a high level, our architecture consists of a summary network $\\mathbf{h}$ and an inference network $\\mathbf{f}$ which jointly learn to invert a generative model. The summary network transforms input data $\\mathbf{x}$ of potentially variable size to a fixed-length representations. The inference network generates random draws from an approximate posterior $\\mathbf{q}$ via a conditional generative networks (here, an invertible network)."
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Setup\n",
-    "\n",
-    "For this notebook to run, you need to have the latest bayesflow dev version installed,\n",
-    "for example via:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2025-02-14T10:51:27.527747Z",
-     "start_time": "2025-02-14T10:51:27.525961Z"
-    },
-    "id": "JMj_DpOSvFJw"
-   },
-   "outputs": [],
-   "source": [
-    "# !pip install git+https://github.com/bayesflow-org/bayesflow.git@dev"
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {
@@ -77,47 +52,12 @@
     "\n",
     "if \"KERAS_BACKEND\" not in os.environ:\n",
     "    # set this to \"torch\", \"tensorflow\", or \"jax\"\n",
-    "    os.environ[\"KERAS_BACKEND\"] = \"jax\""
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2025-02-14T10:51:27.658247Z",
-     "start_time": "2025-02-14T10:51:27.656660Z"
-    }
-   },
-   "outputs": [],
-   "source": [
-    "# for BayesFlow devs: this ensures that the latest dev version can be found\n",
-    "import sys\n",
-    "sys.path.append('../')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2025-02-14T10:51:29.016951Z",
-     "start_time": "2025-02-14T10:51:27.698656Z"
-    }
-   },
-   "outputs": [],
-   "source": [
-    "import keras\n",
+    "    os.environ[\"KERAS_BACKEND\"] = \"jax\"\n",
+    "\n",
     "import matplotlib.pyplot as plt\n",
-    "import numpy as np"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {},
-   "outputs": [],
-   "source": [
+    "import numpy as np\n",
+    "\n",
+    "import keras\n",
     "import bayesflow as bf"
    ]
   },

examples/Lotka_Volterra_point_estimation_and_expert_stats.ipynb

@@ -5,24 +5,13 @@
    "id": "3cacfcfd-d638-4181-9c16-ba050ab5e367",
    "metadata": {},
    "source": [
-    "# Rapid iteration with point estimation and expert statistics for Lotka-Volterra dynamics\n",
+    "# Rapid Iteration with Point Estimation Lotka-Volterra Dynamics\n",
     "\n",
-    "_Authors: Hans Olischläger_\n",
+    "_Author: Hans Olischläger_\n",
     "\n",
     "In this notebook, we will infer parameters of a famous ecology differential equation with BayesFlow.\n",
     "\n",
-    "We will follow a typical workflow that emphazises rapid iterations early on, before building up towards reliable estimates of the full posterior."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "id": "498100c2-9169-4f58-a37e-684aaf32ea45",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "%load_ext autoreload\n",
-    "%autoreload 2"
+    "We will follow a typical workflow that emphazises rapid iterations early on, before building up towards reliable estimates of the full posterior with end-to-end data embedding."
    ]
   },
   {
@@ -51,25 +40,14 @@
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "import seaborn as sns\n",
+    "\n",
     "from scipy.integrate import odeint\n",
     "\n",
-    "# For BayesFlow devs: this ensures that the latest dev version can be found\n",
-    "import sys\n",
-    "sys.path.append('../')\n",
+    "import keras\n",
     "\n",
     "import bayesflow as bf"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "id": "18b0496b-87e4-46cb-9f51-54ae67b1b9c6",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import keras"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": 4,
@@ -125,7 +103,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "rng = np.random.RandomState(seed=1234) # for reproducibility of the simulations"
+    "rng = np.random.default_rng(seed=1234)"
    ]
   },
   {