{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1f01028e-f89f-48cb-9a39-f1b6148eb80b",
   "metadata": {},
   "source": [
    "# Neural network architecture\n",
    "\n",
    "Dingo employs different network architectures for posterior estimation, e.g. [Neural posterior estimation](https://arxiv.org/abs/1605.06376) and [Flow Matching Posterior estimation](https://arxiv.org/pdf/2305.17161), see [here](sbi.md) for an introduction. A central object is the conditional neural density estimator, a deep neural network trained to represent the Bayesian posterior. This section describes the neural network architecture developed in {cite:p}`Dax:2021tsq`, and subsequently used in {cite:p}`Dax:2021myb`, {cite:p}`Dax:2022pxd` and {cite:p}`Wildberger:2022agw`.\n",
    "\n",
    "## Neural spline flow with SVD compression\n",
    "The NPE architecture consists of two compenents, the embedding network which compresses the high-dimensionl data to a lower dimensional feature vector, and the conditional normalizing flow which estimates the Bayesian posterior based on this feature vector. Both components are trained jointly and end-to-end with the objective descriped [here](sbi.md). The network can be build with"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "12202bc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dingo.core.nn.nsf import create_nsf_with_rb_projection_embedding_net"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93dfe641",
   "metadata": {},
   "source": [
    "### Embedding network\n",
    "The embedding network compresses the high-dimensional conditioning information (consisting of frequency domain strain and PSD data). The first layer of this network is initialized with an [SVD](https://en.wikipedia.org/wiki/Singular_value_decomposition) matrix from a reduced basis built with non-noisy waveforms. This projection filters out the noise that is orthogonal to the signal manifold, and significantly simplifies the task for the neural network.\n",
    "\n",
    "The initial compression layer is followed by a sequence of residual blocks consisting of dense layers for further compression. Example kwargs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77c8469c",
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_kwargs = {\n",
    "    \"input_dims\": (2, 3, 8033),\n",
    "    \"output_dim\": 128,\n",
    "    \"hidden_dims\": [\n",
    "        1024, 1024, 1024, 1024, 1024, 1024, \\\n",
    "        512, 512, 512, 512, 512, 512, \\\n",
    "        256, 256, 256, 256, 256, 256, \\\n",
    "        128, 128, 128, 128, 128, 128\n",
    "    ],\n",
    "    \"activation\": \"elu\",\n",
    "    \"dropout\": 0.0,\n",
    "    \"batch_norm\": True,\n",
    "    \"svd\": {\n",
    "        \"num_training_samples\": 50000,\n",
    "        \"num_validation_samples\": 5000,\n",
    "        \"size\": 200,\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f1261fa",
   "metadata": {},
   "source": [
    "Here, `input_dims=(2, 3, 8033)` refers to the input dimension, for frequency domain data with `8033` frequency bins and `3` channels (real part, complex part, ASD) in `2` detectors. The embedding network compresses this to `output_dim=128` components. The SVD initialization is controlled with the `svd` argument, and the residual blocks are specified with `hidden_dims`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18525738",
   "metadata": {},
   "source": [
    "```{note}\n",
    "Not all of these arguments have to be set in the configuration file when training dingo. For example, the `input_dims` argument is automatically filled in based on the specified domain information and number of detectors. Similarly, the `context_dim` of the flow (see below) is filled in based on the `output_dim` of the embedding network and the number of [GNPE](gnpe.md) proxies. See the [Dingo examples](https://github.com/dingo-gw/dingo/tree/main/examples) for the corresponding configuration files and training commands.\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "051cddb1",
   "metadata": {},
   "source": [
    "### Discrete Normalizing Flow\n",
    "We use the [neural spline flow](https://arxiv.org/abs/1906.04032) as a density estimator. This takes the output of the embedding network as context information and estimates the Bayesian posterior distribution. Example kwargs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "e0a5eb25",
   "metadata": {},
   "outputs": [],
   "source": [
    "posterior_kwargs = {\n",
    "    \"input_dim\": 15,\n",
    "    \"context_dim\": 129,\n",
    "    \"num_flow_steps\": 30,\n",
    "    \"base_transform_kwargs\": {\n",
    "        \"hidden_dim\": 512,\n",
    "        \"num_transform_blocks\": 5,\n",
    "        \"activation\": \"elu\",\n",
    "        \"dropout_probability\": 0.0,\n",
    "        \"batch_norm\": True,\n",
    "        \"num_bins\": 8,\n",
    "        \"base_transform_type\": \"rq-coupling\",\n",
    "    },\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc24eb0d",
   "metadata": {},
   "source": [
    "This creates a neural spline flow with `input_dim=15` parameters, conditioned on a `129` dimensional context vector, corresponding to the `128` dimensional output of the embedding network and one [GNPE](gnpe.md) proxy variable. The neural spline flow consists of `num_flow_steps=30` layers, for which the transformation is specified with `base_transform_kwargs`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "3455f052",
   "metadata": {},
   "outputs": [],
   "source": [
    "nde = create_nsf_with_rb_projection_embedding_net(posterior_kwargs, embedding_kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbf2abd6",
   "metadata": {},
   "source": [
    "### Continuous Normalizing Flow\n",
    "\n",
    "Flow Matching Posterior Estimation (FMPE) utilizes continuous normalizing flows to represent posterior distributions. Instead of discrete mappings, FMPE models a velocity field that governs transformations over time, mapping a base distribution (e.g., standard normal) to the posterior.\n",
    "\n",
    "We use a dense residual network to parameterize the velocity field, which takes as input the parameters $ \\theta $, time $ t $, and context $ d $. The model is trained by minimizing the flow matching loss, which ensures that the learned velocity field matches the optimal transport-inspired target velocity field.\n",
    "\n",
    "Example kwargs for an FMPE model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a79ef39",
   "metadata": {},
   "outputs": [],
   "source": [
    "posterior_kwargs = {\n",
    "    \"activation\": \"gelu\",\n",
    "    \"batch_norm\": True,\n",
    "    \"context_with_glu\": False,\n",
    "    \"dropout\": 0.0,\n",
    "    \"hidden_dims\": [\n",
    "        1024, 1024, 1024, 1024, 1024, 1024,\n",
    "        512, 512, 512, 512, 512, 512, 512, 512\n",
    "    ],\n",
    "    \"sigma_min\": 0.001,\n",
    "    \"theta_embedding_kwargs\": {\n",
    "        \"embedding_net\": {\n",
    "            \"activation\": \"gelu\",\n",
    "            \"hidden_dims\": [16, 32, 64, 128, 256],\n",
    "            \"output_dim\": 256,\n",
    "            \"type\": \"DenseResidualNet\",\n",
    "        },\n",
    "        \"encoding\": {\n",
    "            \"encode_all\": False,\n",
    "            \"frequencies\": 0,\n",
    "        },\n",
    "    },\n",
    "    \"theta_with_glu\": True,\n",
    "    \"time_prior_exponent\": 1,\n",
    "    \"type\": \"DenseResidualNet\",\n",
    "}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2a2e230",
   "metadata": {},
   "source": [
    "### Explanation of Key Settings\n",
    "\n",
    "- **`theta_embedding_kwargs`**: \n",
    "  - Specifies how the parameter \\(\\theta\\) is embedded for input into the velocity network.\n",
    "  - Contains settings for:\n",
    "    - `embedding_net`: Defines the embedding network, such as a `DenseResidualNet`. This network transforms the input \\(\\theta\\) into a higher-dimensional representation for downstream processing. The `hidden_dims` control the architecture depth and complexity, while the `output_dim` specifies the dimensionality of the embedding.\n",
    "    - `encoding`: Configures the positional encoding for \\(\\theta\\). The `encode_all` flag determines whether all components of \\(\\theta\\) or only a subset (e.g., the first component) are encoded. `frequencies` adjusts the sinusoidal encoding settings.\n",
    "\n",
    "- **`hidden_dims`**:\n",
    "  - Defines the layer sizes of the velocity network. The dimensions decrease progressively to capture hierarchical features and allow for efficient computation. Larger dimensions can better approximate complex velocity fields but increase computational costs.\n",
    "\n",
    "- **`sigma_min`**:\n",
    "  - Controls the minimum noise level in the interpolation path between the base distribution and the target posterior distribution. Too low values of sigma_min can lead to sharp trajectories that are harder for the velocity network to model accurately, potentially resulting in unstable training.\n",
    "\n",
    "- **`theta_with_glu`**:\n",
    "  - If `True`, applies Gated Linear Units (GLUs) for processing \\(\\theta\\). GLUs introduce an additional nonlinear gating mechanism, which can improve the network’s ability to model complex relationships between inputs.\n",
    "\n",
    "- **`context_with_glu`**:\n",
    "  - Similar to `theta_with_glu`, but applies GLUs to the context \\(d\\). Useful for enhancing feature extraction when the context data has intricate dependencies.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26c7e178",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}