{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "s1i4P0Z-LFnl"
   },
   "source": [
    "# RANS Airfoil Flows with Bayesian Neural Nets\n",
    "\n",
    "## Overview\n",
    "\n",
    "We are now considering the same setup as in the notebook {doc}`supervised-airfoils`: A turbulent airflow around wing profiles, for which we'd like to know the average motion\n",
    "and pressure distribution around this airfoil for different Reynolds numbers and angles of attack. In the earlier notebook, we tackled this by completely bypassing any physical solver and instead training a neural network that learns the quantities of interest. Now, we want to extend this approach to the variational Bayesian Neural Networks (BNNs) of the previous section. In contrast to traditional networks, that learn a single point estimate for each weight value, BNNs aim at learning a _distribution_ over each weight parameter (e.g. a Gaussian with mean $\\mu$ and variance $\\sigma^{2}$). During a forward-pass, each parameter in the network is then sampled from its corresponding approximate posterior distribution $q_{\\phi}(\\theta)$. In that sense, the network parameters themselves are _random variables_ and each forward pass becomes _stochastic_, because for a given input the predictions will vary with every forward-pass. This allows to assess how _uncertain_ the network is: If the predictions vary a lot, we think that the network is uncertain about its output. [[run in colab]](https://colab.research.google.com/github/tum-pbs/pbdl-book/blob/main/bayesian-code.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jzE1wwAjG_QQ"
   },
   "source": [
    "### Read in Data\n",
    "Like in the previous notebook we'll skip the data generation process. This example is adapted from [the Deep-Flow-Prediction codebase](https://github.com/thunil/Deep-Flow-Prediction), which you can check out for details. Here, we'll simply download a small set of training data generated with a Spalart-Almaras RANS simulation in [OpenFOAM](https://openfoam.org/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "JwZudtWauiGa",
    "outputId": "4b8b6f0c-ad97-4527-869c-622ad11c292b"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading training data (300MB), this can take a few minutes the first time...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": []
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded data, 320 training, 80 validation samples\n",
      "Size of the inputs array: (320, 3, 128, 128)\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import os.path, random\n",
    "\n",
    "# get training data: as in the previous supervised example, either download or use gdrive\n",
    "dir = \"./\"\n",
    "if True:\n",
    "    if not os.path.isfile('data-airfoils.npz'):\n",
    "        import requests\n",
    "        print(\"Downloading training data (300MB), this can take a few minutes the first time...\")\n",
    "        with open(\"data-airfoils.npz\", 'wb') as datafile:\n",
    "            resp = requests.get('https://dataserv.ub.tum.de/s/m1615239/download?path=%2F&files=dfp-data-400.npz', verify=False)\n",
    "            datafile.write(resp.content)\n",
    "else: # cf supervised airfoil code:\n",
    "    from google.colab import drive\n",
    "    drive.mount('/content/gdrive')\n",
    "    dir = \"./gdrive/My Drive/\"\n",
    "\n",
    "npfile=np.load(dir+'data-airfoils.npz')\n",
    "print(\"Loaded data, {} training, {} validation samples\".format(len(npfile[\"inputs\"]),len(npfile[\"vinputs\"])))\n",
    "print(\"Size of the inputs array: \"+format(npfile[\"inputs\"].shape))\n",
    "\n",
    "# reshape to channels_last for convencience\n",
    "X_train = np.moveaxis(npfile[\"inputs\"],1,-1)\n",
    "y_train = np.moveaxis(npfile[\"targets\"],1,-1)\n",
    "\n",
    "X_val = np.moveaxis(npfile[\"vinputs\"],1,-1)\n",
    "y_val = np.moveaxis(npfile[\"vtargets\"],1,-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5C_Ooq0JG_Qk"
   },
   "source": [
    "### Look at Data\n",
    "Now we have some training data. We can look at it using the code we also used in the original notebook:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 383
    },
    "id": "1y1zHmKlAWNs",
    "outputId": "2a3223d3-5708-48ee-8186-78d6c403145e"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Here are all 3 inputs are shown at the top (mask,in x, in y) \n",
      "Side by side with the 3 output channels (p,vx,vy) at the bottom:\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 44724 characters]\n",
      "text/plain": [
       "<Figure size 1152x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import pylab\n",
    "from matplotlib import cm\n",
    "\n",
    "# helper to show three target channels: normalized, with colormap, side by side\n",
    "def showSbs(a1,a2, bottom=\"NN Output\", top=\"Reference\", title=None): \n",
    "  c=[]\n",
    "  for i in range(3):\n",
    "    b = np.flipud( np.concatenate((a2[...,i],a1[...,i]),axis=1).transpose())\n",
    "    min, mean, max = np.min(b), np.mean(b), np.max(b); \n",
    "    b -= min; b /= (max-min)\n",
    "    c.append(b)\n",
    "  fig, axes = pylab.subplots(1, 1, figsize=(16, 5))\n",
    "  axes.set_xticks([]); axes.set_yticks([]); \n",
    "  im = axes.imshow(np.concatenate(c,axis=1), origin='upper', cmap='magma')\n",
    "\n",
    "  pylab.colorbar(im); pylab.xlabel('p, ux, uy'); pylab.ylabel('%s           %s'%(bottom,top))\n",
    "  if title is not None: pylab.title(title)\n",
    "\n",
    "NUM=40\n",
    "print(\"\\nHere are all 3 inputs are shown at the top (mask,in x, in y) \\nSide by side with the 3 output channels (p,vx,vy) at the bottom:\")\n",
    "showSbs( X_train[NUM],y_train[NUM], bottom=\"Target Output\", top=\"Inputs\", title=\"Training data sample\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fAd8hWG3atqg"
   },
   "source": [
    "Not surprisingly, the data still looks the same. For details, please check out the description in {doc}`supervised-airfoils`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "C2gdKINAG_Qs"
   },
   "source": [
    "### Neural Network Definition\n",
    "Now let's look at how we can implement BNNs. Instead of PyTorch, we will use TensorFlow, in particular the extension TensorFlow Probability, which has easy-to-implement probabilistic layers. Like in the other notebook, we use a U-Net structure consisting of Convolutional blocks with skip-layer connections. For now, we only want to set up the decoder, i.e. second part of the U-Net as bayesian. For this, we will take advantage of TensorFlows _flipout_ layers (in particular, the convolutional implementation). \n",
    "\n",
    "In a forward pass, those layers automatically sample from the current posterior distribution and store the KL-divergence between prior and posterior in _model.losses_. One can specify the desired divergence measure (typically KL-divergence) and modify the prior and approximate posterior distributions, if other than normal distributions are desired. Other than that, the flipout layers can be used just like regular layers in sequential architectures. The code below implements a single convolutional block of the U-Net:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "9-8O6qJ-G_Qt"
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow_probability.python.distributions as tfd\n",
    "from tensorflow.keras import Sequential\n",
    "from tensorflow.keras.initializers import RandomNormal\n",
    "from tensorflow.keras.layers import Input, Conv2D, Conv2DTranspose,UpSampling2D, BatchNormalization, ReLU, LeakyReLU, SpatialDropout2D, MaxPooling2D\n",
    "from tensorflow_probability.python.layers import Convolution2DFlipout\n",
    "from tensorflow.keras.models import Model\n",
    "\n",
    "def tfBlockUnet(filters=3, transposed=False, kernel_size=4, bn=True, relu=True, pad=\"same\", dropout=0., flipout=False,\n",
    "              kdf=None, name=''):\n",
    "    block = Sequential(name=name)\n",
    "    if relu:\n",
    "        block.add(ReLU())\n",
    "    else:\n",
    "        block.add(LeakyReLU(0.2))\n",
    "    if not transposed:\n",
    "        block.add(Conv2D(filters=filters, kernel_size=kernel_size, padding=pad,\n",
    "                         kernel_initializer=RandomNormal(0.0, 0.02), activation=None,strides=(2,2)))\n",
    "    else:\n",
    "        block.add(UpSampling2D(interpolation = 'bilinear'))\n",
    "        if flipout:\n",
    "            block.add(Convolution2DFlipout(filters=filters, kernel_size=(kernel_size-1), strides=(1, 1), padding=pad,\n",
    "                                           data_format=\"channels_last\", kernel_divergence_fn=kdf,\n",
    "                                           activation=None))\n",
    "        else:\n",
    "            block.add(Conv2D(filters=filters, kernel_size=(kernel_size-1), padding=pad,\n",
    "                         kernel_initializer=RandomNormal(0.0, 0.02), strides=(1,1), activation=None))\n",
    "    \n",
    "    block.add(SpatialDropout2D(rate=dropout))\n",
    "    \n",
    "    if bn:\n",
    "        block.add(BatchNormalization(axis=-1, epsilon=1e-05,momentum=0.9))\n",
    "\n",
    "    return block"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "57AMSMp5sPFZ"
   },
   "source": [
    "Next we define the full network with these blocks - the structure is almost identical to the previous notebook. We manually define the kernel-divergence function as `kdf` and rescale it with a factor called `kl_scaling`. There are two reasons for this: \n",
    "\n",
    "First, we should only apply the kl-divergence once per epoch if we want to use the correct loss (like introduced in {doc}`bayesian-intro`). Since we will use batch-wise training, we need to rescale the Kl-divergence by the number of batches, such that in every parameter update only _kdf / num_batches_ is added to the loss. During one epoch, _num_batches_ parameter updates are performed and the 'full' KL-divergence is used. This batch scaling is computed and passed to the network initialization via `kl_scaling` when instantiating the `Bayes_DfpNet` NN later on.\n",
    "\n",
    "Second, by scaling the KL-divergence part of the loss up or down, we have a way of tuning how much randomness we want to allow in the network: If we neglect the KL-divergence completely, we would just minimize the regular loss (e.g. MSE or MAE), like in a conventional neural network. If we instead neglect the negative-log-likelihood, we would optimize the network such that we obtain random draws from the prior distribution. Balancing those extremes can be done by fine-tuning the scaling of the KL-divergence and is hard in practice. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "pGoGJwqdG_Qv"
   },
   "outputs": [],
   "source": [
    "def Bayes_DfpNet(input_shape=(128,128,3),expo=5,dropout=0.,flipout=False,kl_scaling=10000):\n",
    "    channels = int(2 ** expo + 0.5)\n",
    "    kdf = (lambda q, p, _: tfd.kl_divergence(q, p) / tf.cast(kl_scaling, dtype=tf.float32))\n",
    "\n",
    "    layer1=Sequential(name='layer1')\n",
    "    layer1.add(Conv2D(filters=channels,kernel_size=4,strides=(2,2),padding='same',activation=None,data_format='channels_last'))\n",
    "    layer2=tfBlockUnet(filters=channels*2,transposed=False,bn=True, relu=False,dropout=dropout,name='layer2')\n",
    "    layer3=tfBlockUnet(filters=channels*2,transposed=False,bn=True, relu=False,dropout=dropout,name='layer3')        \n",
    "    layer4=tfBlockUnet(filters=channels*4,transposed=False,bn=True, relu=False,dropout=dropout,name='layer4')        \n",
    "    layer5=tfBlockUnet(filters=channels*8,transposed=False,bn=True, relu=False,dropout=dropout,name='layer5')        \n",
    "    layer6=tfBlockUnet(filters=channels*8,transposed=False,bn=True, relu=False,dropout=dropout,kernel_size=2,pad='valid',name='layer6')        \n",
    "    layer7=tfBlockUnet(filters=channels*8,transposed=False,bn=True, relu=False,dropout=dropout,kernel_size=2,pad='valid',name='layer7')        \n",
    "\n",
    "    # note, kernel size is internally reduced by one for the decoder part\n",
    "    dlayer7=tfBlockUnet(filters=channels*8,transposed=True,bn=True, relu=True,dropout=dropout, flipout=flipout,kdf=kdf, kernel_size=2,pad='valid',name='dlayer7')        \n",
    "    dlayer6=tfBlockUnet(filters=channels*8,transposed=True,bn=True, relu=True,dropout=dropout, flipout=flipout,kdf=kdf, kernel_size=2,pad='valid',name='dlayer6')        \n",
    "    dlayer5=tfBlockUnet(filters=channels*4,transposed=True,bn=True, relu=True,dropout=dropout, flipout=flipout,kdf=kdf,name='dlayer5')        \n",
    "    dlayer4=tfBlockUnet(filters=channels*2,transposed=True,bn=True, relu=True,dropout=dropout, flipout=flipout,kdf=kdf,name='dlayer4')\n",
    "    dlayer3=tfBlockUnet(filters=channels*2,transposed=True,bn=True, relu=True,dropout=dropout, flipout=flipout,kdf=kdf,name='dlayer3')\n",
    "    dlayer2=tfBlockUnet(filters=channels  ,transposed=True,bn=True, relu=True,dropout=dropout, flipout=flipout,kdf=kdf,name='dlayer2')\n",
    "    dlayer1=Sequential(name='outlayer')\n",
    "    dlayer1.add(ReLU())\n",
    "    dlayer1.add(Conv2DTranspose(3,kernel_size=4,strides=(2,2),padding='same'))\n",
    "\n",
    "    # forward pass\n",
    "    inputs=Input(input_shape)\n",
    "    out1 = layer1(inputs)\n",
    "    out2 = layer2(out1)\n",
    "    out3 = layer3(out2)\n",
    "    out4 = layer4(out3)\n",
    "    out5 = layer5(out4)\n",
    "    out6 = layer6(out5)\n",
    "    out7 = layer7(out6)\n",
    "    # ... bottleneck ...\n",
    "    dout6 = dlayer7(out7)\n",
    "    dout6_out6 = tf.concat([dout6,out6],axis=3)\n",
    "    dout6 = dlayer6(dout6_out6)\n",
    "    dout6_out5 = tf.concat([dout6, out5], axis=3)\n",
    "    dout5 = dlayer5(dout6_out5)\n",
    "    dout5_out4 = tf.concat([dout5, out4], axis=3)\n",
    "    dout4 = dlayer4(dout5_out4)\n",
    "    dout4_out3 = tf.concat([dout4, out3], axis=3)\n",
    "    dout3 = dlayer3(dout4_out3)\n",
    "    dout3_out2 = tf.concat([dout3, out2], axis=3)\n",
    "    dout2 = dlayer2(dout3_out2)\n",
    "    dout2_out1 = tf.concat([dout2, out1], axis=3)\n",
    "    dout1 = dlayer1(dout2_out1)\n",
    "    \n",
    "    return Model(inputs=inputs,outputs=dout1)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ETzjWZJzyMEu"
   },
   "source": [
    "Let's define the hyperparameters and create a tensorflow dataset to organize inputs and targets. Since we have 320 observations in the training set, for a batch-size of 10 we should rescale the KL-divergence with a factor of 320/10=32 in order apply the full KL-divergence just once per epoch. We will further scale the KL-divergence down by another factor of `KL_PREF=5000`, which has shown to work well in practice. \n",
    "\n",
    "Furthermore, we will define a function that implements learning rate decay. Intuitively, this allows the optimization to be more precise (by making smaller steps) in later epochs, while still making fast progress (by making bigger steps) in the first epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "H23a0JAJclEN"
   },
   "outputs": [],
   "source": [
    "import math \n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "BATCH_SIZE=10\n",
    "LR=0.001\n",
    "EPOCHS = 120\n",
    "KL_PREF = 5000\n",
    "\n",
    "dataset = tf.data.Dataset.from_tensor_slices((X_train, y_train)).shuffle(len(X_train),\n",
    "    seed=46168531, reshuffle_each_iteration=False).batch(BATCH_SIZE, drop_remainder=False)\n",
    "\n",
    "def compute_lr(i, epochs, minLR, maxLR):\n",
    "  if i < epochs * 0.5:\n",
    "      return maxLR\n",
    "  e = (i / float(epochs) - 0.5) * 2.\n",
    "  # rescale second half to min/max range\n",
    "  fmin = 0.\n",
    "  fmax = 6.\n",
    "  e = fmin + e * (fmax - fmin)\n",
    "  f = math.pow(0.5, e)\n",
    "  return minLR + (maxLR - minLR) * f"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3lUU7A0o1PzV"
   },
   "source": [
    "We can visualize the learning rate decay: We start off with a constant rate and after half of the `EPOCHS` we start to decay it exponentially, until arriving at half of the original learning rate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 296
    },
    "id": "tDXZO0FN02do",
    "outputId": "f19c0484-180c-4780-b7a0-cbc95d2c0fa9"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0, 0.5, 'Learning Rate')"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 13520 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "lrs=[compute_lr(i, EPOCHS, 0.5*LR,LR) for i in range(EPOCHS)]\n",
    "plt.plot(lrs)\n",
    "plt.xlabel('Iteration')\n",
    "plt.ylabel('Learning Rate')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lcwUJ6iH1uTr"
   },
   "source": [
    "Let's initialize the network. Here we're finally computing the `kl_scaling` factor via `KL_PREF` and the batch size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "LGaH2ZW73M1b",
    "outputId": "13c2ada4-79bf-42f7-f2f9-4355c0173fb0"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": []
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The Bayesian U-Net has 846787 parameters.\n"
     ]
    }
   ],
   "source": [
    "from tensorflow.keras.optimizers import RMSprop, Adam\n",
    "\n",
    "model=Bayes_DfpNet(expo=4,flipout=True,kl_scaling=KL_PREF*len(X_train)/BATCH_SIZE)\n",
    "optimizer = Adam(learning_rate=LR, beta_1=0.5,beta_2=0.9999)\n",
    "\n",
    "num_params = np.sum([np.prod(v.get_shape().as_list()) for v in model.trainable_variables])\n",
    "print('The Bayesian U-Net has {} parameters.'.format(num_params))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eF4UTqeB6Y4u"
   },
   "source": [
    "In general, flipout layers come with twice as many parameters as their conventional counterparts, since instead of a single point estimate one has to learn both mean and variance parameters for the Gaussian posterior of the weights. As we only have flipout layers for the decoder part here, the resulting model has 846787 parameters, compared to the 585667 of the conventional NN.\n",
    "\n",
    "## Training\n",
    "\n",
    "Now we are ready to run the training! Note that this might take a while (typically around 4 hours), as the flipout layers are significantly slower to train compared to regular layers.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "VEQuKBegcd1K",
    "outputId": "52d03473-3c95-4df6-9657-9d64c41348c8"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0/120, total loss: 4.265, KL loss: 4.118, MAE loss: 0.1464, MAE loss vali: 0.0872\n",
      "Epoch 1/120, total loss: 4.159, KL loss: 4.089, MAE loss: 0.0706, MAE loss vali: 0.0691\n",
      "Epoch 2/120, total loss: 4.115, KL loss: 4.054, MAE loss: 0.0610, MAE loss vali: 0.0589\n",
      "Epoch 20/120, total loss: 3.344, KL loss: 3.315, MAE loss: 0.0291, MAE loss vali: 0.0271\n",
      "Epoch 40/120, total loss: 2.495, KL loss: 2.471, MAE loss: 0.0245, MAE loss vali: 0.0242\n",
      "Epoch 60/120, total loss: 1.712, KL loss: 1.689, MAE loss: 0.0228, MAE loss vali: 0.0208\n",
      "Epoch 80/120, total loss: 1.190, KL loss: 1.169, MAE loss: 0.0212, MAE loss vali: 0.0200\n",
      "Epoch 100/120, total loss: 0.869, KL loss: 0.848, MAE loss: 0.0208, MAE loss vali: 0.0203\n"
     ]
    }
   ],
   "source": [
    "from tensorflow.keras.losses import mae\n",
    "import math\n",
    "\n",
    "kl_losses=[]\n",
    "mae_losses=[]\n",
    "total_losses=[]\n",
    "mae_losses_vali=[]\n",
    "\n",
    "for epoch in range(EPOCHS):\n",
    "    # compute learning rate - decay is implemented\n",
    "    currLr = compute_lr(epoch,EPOCHS,0.5*LR,LR)\n",
    "    if currLr < LR:\n",
    "            tf.keras.backend.set_value(optimizer.lr, currLr)\n",
    "\n",
    "    # iterate through training data \n",
    "    kl_sum = 0\n",
    "    mae_sum = 0\n",
    "    total_sum=0\n",
    "    for i, traindata in enumerate(dataset, 0):\n",
    "      # forward pass and loss computation\n",
    "        with tf.GradientTape() as tape:\n",
    "            inputs, targets = traindata\n",
    "            prediction = model(inputs, training=True)\n",
    "            loss_mae = tf.reduce_mean(mae(prediction, targets))\n",
    "            kl=sum(model.losses)\n",
    "            loss_value=kl+tf.cast(loss_mae, dtype='float32')\n",
    "        # backpropagate gradients and update parameters \n",
    "        gradients = tape.gradient(loss_value, model.trainable_variables)\n",
    "        optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
    "   \n",
    "        # store losses per batch\n",
    "        kl_sum += kl\n",
    "        mae_sum += tf.reduce_mean(loss_mae)\n",
    "        total_sum+=tf.reduce_mean(loss_value)\n",
    "\n",
    "    # store losses per epoch\n",
    "    kl_losses+=[kl_sum/len(dataset)]\n",
    "    mae_losses+=[mae_sum/len(dataset)]\n",
    "    total_losses+=[total_sum/len(dataset)]\n",
    "\n",
    "    # validation\n",
    "    outputs = model.predict(X_val)\n",
    "    mae_losses_vali += [tf.reduce_mean(mae(y_val, outputs))]\n",
    "\n",
    "    if epoch<3 or epoch%20==0: \n",
    "        print('Epoch {}/{}, total loss: {:.3f}, KL loss: {:.3f}, MAE loss: {:.4f}, MAE loss vali: {:.4f}'.format(epoch, EPOCHS, total_losses[-1], kl_losses[-1], mae_losses[-1], mae_losses_vali[-1]))\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7aM5Ra2C7k1v"
   },
   "source": [
    "The BNN is trained! Let's look at the loss. Since the loss consists of two separate parts, it is helpful to monitor both parts (MAE and KL)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 233
    },
    "id": "3niTOL6CcrHo",
    "outputId": "24acd9cc-c0aa-4848-bd19-e0721eb4b998"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x7f6285c60490>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 40364 characters]\n",
      "text/plain": [
       "<Figure size 1440x288 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig,axs=plt.subplots(ncols=3,nrows=1,figsize=(20,4))\n",
    "\n",
    "axs[0].plot(kl_losses,color='red')\n",
    "axs[0].set_title('KL Loss (Train)')\n",
    "axs[1].plot(mae_losses,color='blue',label='train')\n",
    "axs[1].plot(mae_losses_vali,color='green',label='val')\n",
    "axs[1].set_title('MAE Loss'); axs[1].legend()\n",
    "\n",
    "axs[2].plot(total_losses,label='Total',color='black')\n",
    "axs[2].plot(kl_losses,label='KL',color='red')\n",
    "axs[2].plot(mae_losses,label='MAE',color='blue')\n",
    "axs[2].set_title('Total Train Loss'); axs[2].legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WBGaJng6psg-"
   },
   "source": [
    "This way, we can double-check if minimizing one part of the loss comes at the cost of increasing the other. For our case, we observe that both parts decrease smoothly. In particular, the MAE loss is not increasing for the validation set, indicating that we are not overfitting. \n",
    "\n",
    "It is good practice to double-check how many layers added KL-losses. We can inspect _model.losses_ for that. Since the decoder consists of 6 sequential blocks with flipout layers, we expect 6 entries in _model.losses_."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "LusTjxWRdRjn",
    "outputId": "e4c401d3-c55f-49d2-b780-460bbdf47d0a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "There are 6 entries in model.losses\n",
      "[<tf.Tensor: shape=(), dtype=float32, numpy=0.03536303>, <tf.Tensor: shape=(), dtype=float32, numpy=0.06506929>, <tf.Tensor: shape=(), dtype=float32, numpy=0.2647468>, <tf.Tensor: shape=(), dtype=float32, numpy=0.09337218>, <tf.Tensor: shape=(), dtype=float32, numpy=0.0795429>, <tf.Tensor: shape=(), dtype=float32, numpy=0.075103864>]\n"
     ]
    }
   ],
   "source": [
    "# there should be 6 entries in model.losses since we have 6 blocks with flipout layers in our model\n",
    "print('There are {} entries in model.losses'.format(len(model.losses)))\n",
    "print(model.losses)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "auFzYkx0G_Q6"
   },
   "source": [
    "Now let's visualize how the BNN performs for unseen data from the validation set. Ideally, we would like to integrate out the parameters $\\theta$, i.e. marginalize in order to obtain a prediction. Since this is again hard to realize analytically, one usually approximates the integral via sampling from the posterior:\n",
    "\n",
    "$$\n",
    "  \\hat{y_{i}}=\\int f(x_{i};\\theta)q_{\\phi}(\\theta)d\\theta\\approx\\frac{1}{R}\\sum_{r=1}^{R}f(x_{i};\\theta_{r})$$\n",
    "\n",
    "where each $\\theta_{r}$ is drawn from $q_{\\phi}(\\theta)$. In practice, this just means performing $R$ forward passes for each input $x_{i}$ and computing the average. In the same spirit, one can obtain the standard deviation as a measure of uncertainty:\n",
    "$\\sigma_{i}^{2} = \\frac{1}{R-1}\\sum_{r=1}^{R}(f(x_{i};\\theta)-\\hat{y_{i}})^{2}$.\n",
    "\n",
    "Please note that both $\\hat{y_{i}}$ and $\\sigma_{i}^{2}$ still have shape $128\\times128\\times3$, i.e. the mean and variance computations are performed _per-pixel_ (but might be aggregated to a global measure afterwards)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "Z-gyY4LKG_Q7"
   },
   "outputs": [],
   "source": [
    "REPS=20\n",
    "preds=np.zeros(shape=(REPS,)+X_val.shape)\n",
    "for rep in range(REPS):\n",
    "    preds[rep,:,:,:,:]=model.predict(X_val)\n",
    "preds_mean=np.mean(preds,axis=0)\n",
    "preds_std=np.std(preds,axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mcSUWaiHcVkb"
   },
   "source": [
    "Before inspecting the mean and standard deviation computed in the previous cell, let's visualize one of the outputs of the BNN. In the following plot, the input is shown in the first row, while the second row illustrates the result of a single forward pass. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 321
    },
    "id": "WSEXdGc7G_Q8",
    "outputId": "f0f35cfa-c9d9-408c-8a12-8b2540fe8df0"
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 103980 characters]\n",
      "text/plain": [
       "<Figure size 1152x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "NUM=16\n",
    "# show a single prediction\n",
    "showSbs(y_val[NUM],preds[0][NUM], top=\"Inputs\", bottom=\"Single forward pass\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_b9HoUskeU1B"
   },
   "source": [
    "If you compare this image to one of the outputs from {doc}`supervised-airfoils`, you'll see that it doesn't look to different on first sight. This is a good sign, it seems the network learned to produce the content of the pressure and velocity fields.\n",
    "\n",
    "More importantly, though, we can now visualize the uncertainty over predictions more clearly by inspecting several samples from the posterior distribution  as well as the standard deviation for a given input. Below is code for a function that visualizes precisely that (uncertainty is shown with a different colormap in order to illustrate the differences to previous non-bayesian notebook)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 132
    },
    "id": "gAcKTBCGmohC",
    "outputId": "9fe549dd-c74e-49fa-9c3f-3c1a260bdb00"
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 183116 characters]\n",
      "text/plain": [
       "<Figure size 1440x288 with 8 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# plot repeated samples from posterior for some observations\n",
    "def plot_BNN_predictions(target, preds, pred_mean, pred_std, num_preds=5,channel=0):\n",
    "  if num_preds>len(preds):\n",
    "    print('num_preds was set to {}, but has to be smaller than the length of preds. Setting it to {}'.format(num_preds,len(preds)))\n",
    "    num_preds = len(preds)\n",
    "\n",
    "  # transpose and concatenate the frames that are to plot\n",
    "  to_plot=np.concatenate((target[:,:,channel].transpose().reshape(128,128,1),preds[0:num_preds,:,:,channel].transpose(), \n",
    "                          pred_mean[:,:,channel].transpose().reshape(128,128,1),pred_std[:,:,channel].transpose().reshape(128,128,1)),axis=-1)\n",
    "  fig, axs = plt.subplots(nrows=1,ncols=to_plot.shape[-1],figsize=(20,4))\n",
    "  for i in range(to_plot.shape[-1]):\n",
    "    label='Target' if i==0 else ('Avg Pred' if i == (num_preds+1) else ('Std Dev (normalized)' if i == (num_preds+2) else 'Pred {}'.format(i)))\n",
    "    colmap = cm.viridis if i==to_plot.shape[-1]-1 else cm.magma\n",
    "    frame = np.flipud(to_plot[:,:,i])\n",
    "    min=np.min(frame); max = np.max(frame)\n",
    "    frame -= min; frame /=(max-min)\n",
    "    axs[i].imshow(frame,cmap=colmap)\n",
    "    axs[i].axis('off')\n",
    "    axs[i].set_title(label)\n",
    "\n",
    "OBS_IDX=5\n",
    "plot_BNN_predictions(y_val[OBS_IDX,...],preds[:,OBS_IDX,:,:,:],preds_mean[OBS_IDX,...],preds_std[OBS_IDX,...])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ln7sn-UJpHa7"
   },
   "source": [
    "We are looking at channel 0, i.e. the pressure here. One can observe that the dark and bright regions vary quite a bit across predictions. It is reassuring to note that - at least from visual inspection - the average (i.e. marginal) prediction is closer to the target than most of the single forward passes. \n",
    "\n",
    "It should also be noted that each frame was normalized for the visualization. Therefore, when looking at the uncertainty frame, we can infer where the network is uncertain, but now how uncertain it is in absolute values. \n",
    "\n",
    "In order to assess a global measure of uncertainty we can however compute an average standard deviation over all samples in the validation set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "vQRbK594S2Ak",
    "outputId": "7e3ea005-6d48-4308-8e53-002173f4ebc2"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Average pixel prediction on validation set: \n",
      " pressure: 0.025 +- 0.009, \n",
      " ux: 0.471 +- 0.019,\n",
      " uy: 0.081 +- 0.016\n"
     ]
    }
   ],
   "source": [
    "# Average Prediction with total uncertainty\n",
    "uncertainty_total = np.mean(np.abs(preds_std),axis=(0,1,2))\n",
    "preds_mean_global = np.mean(np.abs(preds),axis=(0,1,2,3))\n",
    "print(\"\\nAverage pixel prediction on validation set: \\n pressure: {} +- {}, \\n ux: {} +- {},\\n uy: {} +- {}\".format(np.round(preds_mean_global[0],3),np.round(uncertainty_total[0],3),np.round(preds_mean_global[1],3),np.round(uncertainty_total[1],3),np.round(preds_mean_global[2],3),np.round(uncertainty_total[2],3)))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-NtAI8ZbMS8a"
   },
   "source": [
    "For a run with standard settings, the uncertainties are on the order of 0.01 for all three fields. As the pressure field has a smaller mean, it's uncertainty is larger in relative terms. This makes sense, as the pressure field is known to be more difficult to predict than the two velocity components."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hcVF9S-JoprL"
   },
   "source": [
    "## Test evaluation\n",
    "\n",
    "Like in the case for a conventional neural network, let's now look at **proper** test samples, i.e. OOD samples, for which in this case we'll use new airfoil shapes. These are shapes that the network never saw in any training samples, and hence it tells us a bit about how well the network generalizes to new shapes.\n",
    "\n",
    "As these samples are at least slightly OOD, we can draw conclusions about how well the network generalizes, which the validation data would not really tell us. In particular, we would like to investigate if the NN is more uncertain when handling OOD data. Like before, we first download the test samples ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ClH53HqtOw3L",
    "outputId": "10ad4fab-48ed-405b-9148-304e4624d29f"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading test data, this should be fast...\n",
      "Loaded 10/10 test samples\n"
     ]
    }
   ],
   "source": [
    "if not os.path.isfile('data-airfoils-test.npz'):\n",
    "  import urllib.request\n",
    "  url=\"https://physicsbaseddeeplearning.org/data/data_test.npz\"\n",
    "  print(\"Downloading test data, this should be fast...\")\n",
    "  urllib.request.urlretrieve(url, 'data-airfoils-test.npz')\n",
    "\n",
    "nptfile=np.load('data-airfoils-test.npz')\n",
    "print(\"Loaded {}/{} test samples\".format(len(nptfile[\"test_inputs\"]),len(nptfile[\"test_targets\"])))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lBRnudybM3Ri"
   },
   "source": [
    "... and then repeat the procedure from above to evaluate the BNN on the test samples, and compute the marginalized prediction and uncertainty."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "TP5D9rhkNlsf",
    "outputId": "5a316880-f5a9-4c4b-8962-59bc6350c788"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Average test error: 0.023046530292471824\n"
     ]
    }
   ],
   "source": [
    "X_test = np.moveaxis(nptfile[\"test_inputs\"],1,-1)\n",
    "y_test = np.moveaxis(nptfile[\"test_targets\"],1,-1)\n",
    "\n",
    "REPS=10\n",
    "preds_test=np.zeros(shape=(REPS,)+X_test.shape)\n",
    "for rep in range(REPS):\n",
    "    preds_test[rep,:,:,:,:]=model.predict(X_test)\n",
    "preds_test_mean=np.mean(preds_test,axis=0)\n",
    "preds_test_std=np.std(preds_test,axis=0)\n",
    "test_loss = tf.reduce_mean(mae(preds_test_mean, y_test))\n",
    "\n",
    "print(\"\\nAverage test error: {}\".format(test_loss))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "kLmWz_kPQ_RE",
    "outputId": "6656bda9-532c-4dce-a4b9-1c4b4c70ecbd"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Average pixel prediction on test set: \n",
      " pressure: 0.03 +- 0.012, \n",
      " ux: 0.466 +- 0.024,\n",
      " uy: 0.091 +- 0.02\n"
     ]
    }
   ],
   "source": [
    "# Average Prediction with total uncertainty\n",
    "uncertainty_test_total = np.mean(np.abs(preds_test_std),axis=(0,1,2))\n",
    "preds_test_mean_global = np.mean(np.abs(preds_test),axis=(0,1,2,3))\n",
    "print(\"\\nAverage pixel prediction on test set: \\n pressure: {} +- {}, \\n ux: {} +- {},\\n uy: {} +- {}\".format(np.round(preds_test_mean_global[0],3),np.round(uncertainty_test_total[0],3),np.round(preds_test_mean_global[1],3),np.round(uncertainty_test_total[1],3),np.round(preds_test_mean_global[2],3),np.round(uncertainty_test_total[2],3)))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is reassuring: The uncertainties on the OOD test set with new shapes are at least slightly higher than on the validation set. \n",
    "\n",
    "### Visualizations\n",
    "\n",
    "The following graph visualizes these measurements: it shows the mean absolute errors for validation and test sets side by side, together with the uncertainties of the predictions as error bars:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 298
    },
    "id": "Sn8g_ACGsVgo",
    "outputId": "26c88b29-1eb9-4316-d5d5-1ab08b3f08f1"
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 8844 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# plot per channel MAE with uncertainty\n",
    "val_loss_c, test_loss_c = [], []\n",
    "for channel in range(3):\n",
    "  val_loss_c.append( tf.reduce_mean(mae(preds_mean[...,channel], y_val[...,channel])) )\n",
    "  test_loss_c.append( tf.reduce_mean(mae(preds_test_mean[...,channel], y_test[...,channel])) )\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ind = np.arange(len(val_loss_c)); width=0.3\n",
    "bars1 = ax.bar(ind - width/2, val_loss_c, width, yerr=uncertainty_total, capsize=4, label=\"validation\")\n",
    "bars2 = ax.bar(ind + width/2, test_loss_c, width, yerr=uncertainty_test_total, capsize=4, label=\"test\")\n",
    "ax.set_ylabel(\"MAE & Uncertainty\")\n",
    "ax.set_xticks(ind); ax.set_xticklabels(('P', 'u_x', 'u_y'))\n",
    "ax.legend(); plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "icgfvAIqoMpE"
   },
   "source": [
    "The mean error is clearly larger, and the slightly larger uncertainties of the predictions are likewise visible via the error bars.\n",
    "\n",
    "In general it is hard to obtain a calibrated uncertainty estimate, but since we are dealing with a fairly simple problem here, the BNN is able to estimate the uncertainty reasonably well.\n",
    "\n",
    "The next graph shows the differences of the BNN predictions for a single case of the test set (using the same style as for the validation sample above):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 132
    },
    "id": "oP1HOciVjh-o",
    "outputId": "932ab62b-f7cc-4626-8848-2d8f4328e532"
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 147356 characters]\n",
      "text/plain": [
       "<Figure size 1440x288 with 8 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "OBS_IDX=5\n",
    "plot_BNN_predictions(y_test[OBS_IDX,...],preds_test[:,OBS_IDX,:,:,:],preds_test_mean[OBS_IDX,...],preds_test_std[OBS_IDX,...])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6gCw2W4z6bYS"
   },
   "source": [
    "We can also visualize several shapes from the test set together with the corresponding marginalized prediction and uncertainty."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 570
    },
    "id": "iJ2y_HmDzkp6",
    "outputId": "6b22347d-9d76-4f6c-e06d-04703cbe1f8b"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'Std. Dev')"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 195896 characters]\n",
      "text/plain": [
       "<Figure size 648x648 with 9 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "IDXS = [1,3,8]\n",
    "CHANNEL = 0\n",
    "fig, axs = plt.subplots(nrows=len(IDXS),ncols=3,sharex=True, sharey = True, figsize = (9,len(IDXS)*3))\n",
    "for i, idx in enumerate(IDXS):\n",
    "  axs[i][0].imshow(np.flipud(X_test[idx,:,:,CHANNEL].transpose()), cmap=cm.magma)\n",
    "  axs[i][1].imshow(np.flipud(preds_test_mean[idx,:,:,CHANNEL].transpose()), cmap=cm.magma)\n",
    "  axs[i][2].imshow(np.flipud(preds_test_std[idx,:,:,CHANNEL].transpose()), cmap=cm.viridis)\n",
    "axs[0][0].set_title('Shape')\n",
    "axs[0][1].set_title('Avg Pred')\n",
    "axs[0][2].set_title('Std. Dev')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3to3b5OS6p_9"
   },
   "source": [
    "As we can see, the shapes from the test set differ quite a bit from another. Nevertheless, the uncertainty estimate is reasonably distributed. It is especially high in the boundary layer around the airfoil, and in the regions of the low pressure pocket."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "trAG3xoQG_Q-"
   },
   "source": [
    "### Discussion\n",
    "\n",
    "Despite these promising results, there are still several issues with Bayesian Neural Nets, limiting their use in many practical applications. One serious drawback is the need for additional scaling of the KL-loss and the fact that there is no convincing argument on why it is necessary yet (read e.g. [here](http://proceedings.mlr.press/v119/wenzel20a/wenzel20a.pdf) or [here](https://arxiv.org/abs/2008.05912)).\n",
    "Furthermore, some people think that assuming independent normal distributions as variational approximations to the posterior is an oversimplification since in practice the weights are actually highly correlated ([paper](https://arxiv.org/abs/1909.00719)). Other people instead argue that this might not be an issue, as long as the networks in use are deep enough ([paper](https://arxiv.org/abs/2002.03704)). On top of that, there is research on different (e.g. heavy-tailed) priors other than normals and many other aspects of BNNs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "h_z2_i_VA1HP"
   },
   "source": [
    "## Next steps \n",
    "\n",
    "But now it's time to experiment with BNNs yourself. \n",
    "\n",
    "* One interesting thing to look at is how the behavior of our BNN changes, if we adjust the KL-prefactor. In the training loop above we set it to 5000 without further justification. You can check out what happens, if you use a value of 1, as it is suggested by the theory, instead of 5000. According to our implementation, this should make the network 'more bayesian', since we assign larger importance to the KL-divergence than before. \n",
    "\n",
    "* So far, we have only worked with variational BNNs, implemented via TensorFlows probabilistic layers. Recall that there is a simpler way of getting uncertainty estimates: Using dropout not only at training, but also at inference time. You can check out how the outputs change for that case. In order to do so, you can, for instance, just pass a non-zero dropout rate to the network specification and change the prediction phase in the above implementation from _model.predict(...)_ to _model(..., training=True)_. Setting the _training=True_ flag will tell TensorFlow to forward the input as if it were training data and hence, it will apply dropout. Please note that the _training=True_ flag can also affect other features of the network. Batch normalization, for instance, works differently in training and prediction mode. As long as we don't deal with overly different data and use sufficiently large batch-sizes, this should not introduce large errors, though. Sensible dropout rates to start experimenting with are e.g. around 0.1."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "bayesian-code.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
