docs/tutorials/simulators.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Simulators and datasets\n",
    "\n",
    "In `lampe`, a simulator can be any Python callable that takes (a vector of) parameters $\\theta$ as input, in the form of a NumPy array or a PyTorch tensor, and returns a stochastic observation $x$ as output. As such, a simulator implicitly defines a likelihood distribution $p(x | \\theta)$. The prior distribution $p(\\theta)$ is restricted to be a [PyTorch distribution](torch.distributions) implementing the `sample` and `log_prob` methods. Together, the prior and the simulator form a joint distribution $p(\\theta, x) = p(\\theta) p(x | \\theta)$ from which parameters-observation pairs $(\\theta, x)$ can be drawn.\n",
    "\n",
    "This notebook walks you through setting up a prior and a simulator, sampling pairs from the joint, saving pairs on disk and, finally, loading pairs from disk."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import lampe\n",
    "import torch\n",
    "import zuko\n",
    "\n",
    "from itertools import islice\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prior\n",
    "\n",
    "In Bayesian inference problems, it is quite common to have a uniform prior distribution $p(\\theta)$. We recommend to use the [`BoxUniform`](zuko.distributions.BoxUniform) class from the [`zuko`](zuko) package to create a multivariate uniform distribution between a lower and an upper bound.\n",
    "\n",
    "For instance, let's say $\\theta$ is uniform over the domain $[-1, 1]^3$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BoxUniform(low: torch.Size([3]), high: torch.Size([3]))"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "LOWER = -torch.ones(3)\n",
    "UPPER = torch.ones(3)\n",
    "\n",
    "prior = zuko.distributions.BoxUniform(LOWER, UPPER)\n",
    "prior"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the `sample` method, parameters $\\theta$ can be sampled from the newly created prior distribution. The log-density can also be evaluated for given parameters using the `log_prob` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-0.7064,  0.0166, -0.3484])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "theta = prior.sample()\n",
    "theta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(-2.0794)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prior.log_prob(theta)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulator\n",
    "\n",
    "As mentioned previously, the only constraints for the simulator are that the vector of parameters $\\theta$ and the observation $x$ must either be NumPy arrays or PyTorch tensors. The callable is otherwise unconstrained, which allows for a large variety of simulators, including already existing ones and ones implemented in other programming languages than Python.\n",
    "\n",
    "Let's consider the following generative process as our simulator:\n",
    "\n",
    "$$ x = \\begin{pmatrix} \\theta_1 + \\theta_2 \\times \\theta_3 \\\\ \\theta_1 \\times \\theta_2 + \\theta_3 \\end{pmatrix} + 0.05 \\times \\varepsilon $$\n",
    "\n",
    "where $\\varepsilon$ is a 2-d standard Gaussian random variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([ 0.9121, -0.5785,  0.7630])\n",
      "tensor([0.4579, 0.3166])\n"
     ]
    }
   ],
   "source": [
    "def simulator(theta: torch.Tensor) -> torch.Tensor:\n",
    "    x = torch.stack([\n",
    "        theta[..., 0] + theta[..., 1] * theta[..., 2],\n",
    "        theta[..., 0] * theta[..., 1] + theta[..., 2],\n",
    "    ], dim=-1)\n",
    "\n",
    "    return x + 0.05 * torch.randn_like(x)\n",
    "\n",
    "theta = prior.sample()\n",
    "x = simulator(theta)\n",
    "\n",
    "print(theta, x, sep='\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This implementation is vectorized, which means that it can process a batch of parameters at once."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 0.9290, -0.1757, -0.7820],\n",
      "        [ 0.5494, -0.7134, -0.3816],\n",
      "        [-0.2231, -0.6708,  0.0089],\n",
      "        [-0.3692,  0.8460,  0.4389]])\n",
      "tensor([[ 1.1476, -1.0494],\n",
      "        [ 0.9043, -0.7716],\n",
      "        [-0.3210,  0.2096],\n",
      "        [-0.0251,  0.1417]])\n"
     ]
    }
   ],
   "source": [
    "theta = prior.sample((4,))\n",
    "x = simulator(theta)\n",
    "\n",
    "print(theta, x, sep='\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sampling\n",
    "\n",
    "Given a prior and a simulator, the [`JointLoader`](lampe.data.JointLoader) class provided by the [`lampe.data`](lampe.data) module creates an iterable [`DataLoader`](torch.utils.data.DataLoader) of batched parameters-observation pairs $(\\theta, x) \\sim p(\\theta, x)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 256/256 [00:03<00:00, 84.28it/s]\n"
     ]
    }
   ],
   "source": [
    "loader = lampe.data.JointLoader(prior, simulator, batch_size=256)\n",
    "\n",
    "for theta, x in tqdm(islice(loader, 256), total=256):\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the simulator takes a NumPy array as input, it should be indicated with `numpy=True`. Additionally, if the simulator is vectorized, it is recommended to indicate it with `vectorized=True` as it can improve sampling performances significantly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 256/256 [00:00<00:00, 10390.68it/s]\n"
     ]
    }
   ],
   "source": [
    "loader = lampe.data.JointLoader(prior, simulator, batch_size=256, vectorized=True)\n",
    "\n",
    "for theta, x in tqdm(islice(loader, 256), total=256):\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading in memory\n",
    "\n",
    "If the simulator is fast or inexpensive, it is reasonable to generate pairs $(\\theta, x)$ on demand. Otherwise, the pairs have to be generated ahead of time. The [`lampe.data`](lampe.data) module provides the [`JointDataset`](lampe.data.JointDataset) class to interact with in-memory pairs $(\\theta, x)$. This is ideal when your data fits in RAM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([ 0.6757, -0.3787,  0.9581])\n",
      "tensor([0.3410, 0.7707])\n"
     ]
    }
   ],
   "source": [
    "theta = prior.sample((1024,))\n",
    "x = simulator(theta)\n",
    "\n",
    "dataset = lampe.data.JointDataset(theta, x)\n",
    "\n",
    "print(*dataset[42], sep='\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1024/1024 [00:00<00:00, 282991.85it/s]\n"
     ]
    }
   ],
   "source": [
    "for theta, x in tqdm(dataset):\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`JointDataset` can be wrapped in a [`DataLoader`](torch.utils.data.DataLoader) to enable batching and shuffling."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Saving on disk\n",
    "\n",
    "If your data does not fit in RAM or you need to reuse it later, you may want to store it on disk. The [HDF5](https://en.wikipedia.org/wiki/Hierarchical_Data_Format) file format is commonly used for this purpose, as it was specifically designed to hold large amounts of numerical data. The [`lampe.data`](lampe.data) module provides the [`H5Dataset`](lampe.data.H5Dataset) class to help load and store pairs $(\\theta, x)$ in HDF5 files. The [`H5Dataset.store`](lampe.data.H5Dataset.store) function takes an iterable of batched pairs $(\\theta, x)$ as input and stores them into a new HDF5 file. The iterable can be a precomputed list, a custom generator or even a `JointLoader` instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 65536/65536 [00:00<00:00, 518158.51sample/s]\n"
     ]
    }
   ],
   "source": [
    "pairs = []\n",
    "\n",
    "for _ in range(256):\n",
    "    theta = prior.sample((256,))\n",
    "    x = simulator(theta)\n",
    "\n",
    "    pairs.append((theta, x))\n",
    "\n",
    "lampe.data.H5Dataset.store(pairs, 'data_0.h5', size=2**16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 65536/65536 [00:00<00:00, 375452.15sample/s]\n"
     ]
    }
   ],
   "source": [
    "def generate():\n",
    "    while True:\n",
    "        theta = prior.sample((256,))\n",
    "        x = simulator(theta)\n",
    "\n",
    "        yield theta, x\n",
    "\n",
    "lampe.data.H5Dataset.store(generate(), 'data_1.h5', size=2**16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 65536/65536 [00:00<00:00, 369042.89sample/s]\n"
     ]
    }
   ],
   "source": [
    "lampe.data.H5Dataset.store(loader, 'data_2.h5', size=2**16)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading from disk\n",
    "\n",
    "Now that the pairs are stored, we need to load them. The `H5Dataset` class creates an [`IterableDataset`](torch.utils.data.IterableDataset) of pairs $(\\theta, x)$ from HDF5 files. The pairs are dynamically loaded, meaning they are read from disk on demand instead of being cached in memory. This allows for very large datasets that do not even fit in memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 65536/65536 [00:53<00:00, 1227.49it/s]\n"
     ]
    }
   ],
   "source": [
    "dataset = lampe.data.H5Dataset('data_0.h5')\n",
    "\n",
    "for i in tqdm(range(len(dataset))):\n",
    "    theta, x = dataset[i]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, as it can be slow to read pairs from disk one by one, `H5Dataset` implements a custom `__iter__` method which loads pairs by chunks to reduce the number of disk reads."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 65536/65536 [00:00<00:00, 188557.83it/s]\n"
     ]
    }
   ],
   "source": [
    "for theta, x in tqdm(dataset):\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alternatively, if your data fits in memory, you can load it at once with the `to_memory` method, which returns a `JointDataset`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = lampe.data.H5Dataset('data_0.h5').to_memory()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Batching and shuffling\n",
    "\n",
    "As `DataLoader` is not able to shuffle the elements of an `IterableDataset` instance, `H5Dataset` provides a `shuffle` argument to enable (disabled by default) shuffling the pairs when iterating. A `batch_size` argument is also available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = lampe.data.H5Dataset('data_0.h5', batch_size=4, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-0.5102, -0.1597,  0.4225],\n",
      "        [ 0.6263, -0.5490, -0.5424],\n",
      "        [ 0.3914,  0.8860,  0.7201],\n",
      "        [-0.6238,  0.9327, -0.6588]])\n",
      "tensor([[-0.5276,  0.5434],\n",
      "        [ 0.8239, -0.9316],\n",
      "        [ 0.9018,  1.0367],\n",
      "        [-1.3076, -1.2636]])\n"
     ]
    }
   ],
   "source": [
    "for theta, x in dataset:\n",
    "    print(theta, x, sep='\\n')\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 0.3951, -0.0103, -0.9181],\n",
      "        [-0.5488, -0.9719, -0.5787],\n",
      "        [-0.2422, -0.0027,  0.8907],\n",
      "        [ 0.6688,  0.7993, -0.6854]])\n",
      "tensor([[ 0.2745, -0.8634],\n",
      "        [-0.0400, -0.0773],\n",
      "        [-0.2196,  0.9108],\n",
      "        [ 0.1359, -0.1790]])\n"
     ]
    }
   ],
   "source": [
    "for theta, x in dataset:\n",
    "    print(theta, x, sep='\\n')\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Merging datasets\n",
    "\n",
    "Another feature of `H5Dataset` is that it can load data from any number of HDF5 files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 196608/196608 [00:00<00:00, 206486.69it/s]\n"
     ]
    }
   ],
   "source": [
    "dataset = lampe.data.H5Dataset('data_0.h5', 'data_1.h5', 'data_2.h5')\n",
    "\n",
    "for theta, x in tqdm(dataset):\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This feature is very useful when aggregating data from multiple sources, such as data generated on several machines or at different times."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 196608/196608 [00:00<00:00, 423420.86sample/s]\n"
     ]
    }
   ],
   "source": [
    "dataset = lampe.data.H5Dataset('data_0.h5', 'data_1.h5', 'data_2.h5', batch_size=256)\n",
    "\n",
    "lampe.data.H5Dataset.store(dataset, 'data_all.h5', size=len(dataset))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 196608/196608 [00:00<00:00, 226463.55it/s]\n"
     ]
    }
   ],
   "source": [
    "dataset = lampe.data.H5Dataset('data_all.h5')\n",
    "\n",
    "for theta, x in tqdm(dataset):\n",
    "    pass"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "lampe",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  },
  "vscode": {
   "interpreter": {
    "hash": "269a9aa2480715f25fe9d59d402ce545c19ea9afade4a738b9b8ab1d764662d4"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}