Example Notebook explaining changes from 1.0 to 2.0 (#369)

* added examples to standardize transformation documentation

* ran linter

* improve standardize doc further

* rerun linter

* minor cleaning

* added 1.0 to 2.0 example notebook on clean history

* minor cleaning of bayesflow 1.0 to 2.0 notebook

---------

Co-authored-by: Paul-Christian Bürkner <[email protected]>

Author: eodole

examples/From_BayesFlow_1.0_to_2.0.ipynb

@@ -0,0 +1,372 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Moving from Bayesflow 1.0 to 2.0\n",
+    "\n",
+    "_Author: Leona Odole_"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Older users of bayesflow will notice that with the update to version 2.0 many things have changed. This short guide aims to clarify those changes. Users familiar with the previous Quickstart guide will notice that it follows a similar structure, but assumes that users are already familiar with bayesflow. So we omit many of the the mathematical explaination in favor of demonstrating the differences in workflow. For a more detailed explaination of any of the bayesflow framework, users should read, for example, the linear regresion example notebook. \n",
+    "\n",
+    "Additionally to avoid confusion, similarly named objects from _bayesflow1.0_ will have 1.0 after their name, whereas those from _bayesflow2.0_ will not. Finally, a short table with a summary of the function call changes is provided at the end of the guide. "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Keras Framework\n",
+    "\n",
+    "Bayesflow 2.0 looks quite different from 1.0 because the backend has been entirely reformatted in line with `keras` standards.  Previously bayesflow was only compatible with `TensorFlow`, but now users can choose their prefered machine learning framework among `TensorFlow`, `JAX` or `Pytorch`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "\n",
+    "# ensure the backend is set\n",
+    "import os\n",
+    "if \"KERAS_BACKEND\" not in os.environ:\n",
+    "    # set this to \"torch\", \"tensorflow\", or \"jax\"\n",
+    "    os.environ[\"KERAS_BACKEND\"] = \"tensorflow\"\n",
+    "\n",
+    "import keras\n",
+    "import bayesflow as bf\n",
+    "import pandas as pd"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "This version of bayeflow also relies much more heavily on dictionaries since parameters are now named by convention. Many objects now expect a dictionary, so parameters and data are returned as a dictionaries. "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Example Workflow \n",
+    "\n",
+    "### 1. Priors and Likelihood Model\n",
+    "\n",
+    "Any Bayesflow workflow begins with simulated data which is specified with a prior and a corresponding likelihood function. While these two core components are still present, their use and naming conventions within the workflow have changed. \n",
+    "\n",
+    "Previously users would define a prior function, which would then be used by a `Prior1.0` object to sample prior values. The likelihood would then also be specified via function and used by a `Simulator1.0` wrapper to produce observations for a given prior. These were then combined in the `GenerativeModel1.0`.  \n",
+    "\n",
+    "In 2.0 we no longer make use of the  `Prior1.0`, `Simulator1.0` or `GenerativeModel1.0` objects. Instead the roll of the `GenerativeModel1.0` has been renamed to `simulator` which can be invoked as a single function that glues the prior and likelihood functions together to generate samples of both the prior and observations.  "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def theta_prior():\n",
+    "    theta = np.random.normal(size=4)\n",
+    "    # previously: \n",
+    "    # return theta \n",
+    "    return dict(theta=theta) # notice we return a dictionary\n",
+    "    \n",
+    "\n",
+    "def likelihood_model(theta, n_obs):\n",
+    "    x = np.random.normal(loc=theta, size=(n_obs, theta.shape[0]))\n",
+    "    return dict(x=x)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Previously the prior and likelihood were defined as"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Do Not Run\n",
+    "prior_1 = bf.simulation.Prior(prior_fun=theta_prior)\n",
+    "simulator_1 = bf.simulation.Simulator(simulator_fun=likelihood_model)\n",
+    "model_1 = bf.simulation.GenerativeModel(prior=prior_1, simulator=simulator_1)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Whereas the new framework directly uses the likelihood and prior functions directly in the simulator. We also a define a meta function which allows us, for example, to dynamically set the number of observations per simulated dataset. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def meta():\n",
+    "    return dict(n_obs=1)\n",
+    "\n",
+    "simulator = bf.make_simulator([theta_prior, likelihood_model], meta_fn=meta)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We can then generate batches of training samples as follows."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sim_draws = simulator.sample(500)\n",
+    "print(sim_draws[\"x\"].shape)\n",
+    "print(sim_draws[\"theta\"].shape)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 2. Adapter and Data Configuration\n",
+    "\n",
+    "In _bayesflow2.0_ we now need to specify the data configuration. For example we should specify which variables are `summary_variables` meaning observations that will be summarized in the summary network, the `inference_variables` meaning the prior draws on which we're interested in training the posterior network and the `inference_conditions` which specify our number of observations. Previously these things were inferred from the type of network used, but now they should be defined explictly with  the `adapter`. The new approach is much more explicit and extensible. It also makes it easier to change individual settings, while keeping other settings at their defaults."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "adapter = (\n",
+    "    bf.adapters.Adapter()\n",
+    "    .to_array()\n",
+    "    .broadcast(\"n_obs\")\n",
+    "    .convert_dtype(from_dtype=\"float64\", to_dtype=\"float32\")\n",
+    "    .standardize(exclude=[\"n_obs\"])\n",
+    "    .rename(\"x\", \"summary_variables\")\n",
+    "    .rename(\"theta\", \"inference_variables\")\n",
+    "    .rename(\"n_obs\", \"inference_conditions\")\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "In addition the adapter now has built in functions to transform data such as standardization or one-hot encoding. For a full list of the adapter transforms, please see the documentation. "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 3. Summary Network and Inference Network"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "As in _bayesflow1.0_ we still use a summary network, which is still a Deepset model. Nothing has changed in this step of the workflow. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "summary_net = bf.networks.DeepSet(depth=2, summary_dim=10)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "For the inference network there are now several implemented architectures for users to choose from. They are `FlowMatching`, `ConsistencyModel`, `ContinuousConsistencyModel` and `CouplingFlow`.  For this demonstration we use `FlowMatching`, but for further explaination of the different models please see the other examples and documentation. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "inference_net = bf.networks.FlowMatching()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 4. Approximator (Amortizer Posterior)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Previously the actual training and amortization was done in two steps with two different objects the `Amortizer1.0` and `Trainer1.0`. First, users would create an amortizer containing the summary and inference networks."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "### Do Not Run \n",
+    "\n",
+    "# Renamed to Approximator\n",
+    "amortizer = bf.amortizers.AmortizedPosterior(inference_net, summary_net)\n",
+    "\n",
+    "# Defunct\n",
+    "trainer = bf.trainers.Trainer(amortizer=amortizer, generative_model=gen_model)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    " This has been renamed to an `Approximator` and takes the summary network, inference network and the data adapter as arguments. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "approximator = bf.approximators.ContinuousApproximator(\n",
+    "    summary_network=summary_net,\n",
+    "    inference_network=inference_net,\n",
+    "    adapter=adapter\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Whereas previously a  `Trainer1.0` object for training, now users call fit on the `approximator` directly. For additional flexibility in training the `approximator` also has two additional arguments the `learning_rate` and `optimizer`. The optimizer can be any keras optimizer."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "learning_rate = 1e-4\n",
+    "optimizer = keras.optimizers.AdamW(learning_rate=learning_rate, clipnorm=1.0)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Users must then compile the `approximator` with the `optimizer` to make everything ready for training."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "approximator.compile(optimizer=optimizer)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "To train the network, and save output users now need only to call fit on the `approximator`. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "history = approximator.fit(\n",
+    "    epochs=50,\n",
+    "    num_batches=200,\n",
+    "    batch_size=64,\n",
+    "    simulator=simulator\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5.Diagnostics \n",
+    "Another change was made in the model diagnostics, much of the functionality remains the same, but the naming convention has changes. For example previously users would plot losses by using \n",
+    "`bf.diagnostics.plot_losses()`. In *bayesflow2.0*, we instead have all the plotting function grouped together in `bf.diagnostics.plots`. This means, for example, that the loss function is now in `bf.diagnostics.plots.loss()`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "f = bf.diagnostics.plots.loss(\n",
+    "    train_losses=history.history['loss']\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "This was done as we have also added diagnostic metrics such as calibration error, posterior contraction, and root mean squared error. These functions can accordingly be found in `bf.diagnostics.metrics`. For more information please see the documentation."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Summary Change Table \n",
+    "\n",
+    "| Bayesflow 1.0      | Bayesflow 2.0 useage |\n",
+    "| :--------| :---------| \n",
+    "| `Prior`, `Simulator` | Defunct and no longer standalone objects but incorporated into `simulator` | \n",
+    "|`GenerativeModel` | Defunct with it's functionality having been taken over by `simulations.make_simulator` | \n",
+    "| `training.configurator` | Functionality taken over by `Adapter` | \n",
+    "|`Trainer` | Functionality taken over by `fit` method of `Approximator` | \n",
+    "| `AmortizedPosterior`| Renamed to `Approximator` | "
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}