{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1bba0128",
   "metadata": {},
   "source": [
    "# Regression uncertainties\n",
    "\n",
    "In the previous notebook we discussed how to fit data to a neural network. Now, let's try to add uncertainties to it.\n",
    "\n",
    "We are going to use the `torchbnn` module for that, so that we do not have to go into the details of how everything is implemented.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "23feddde",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Tuple\n",
    "\n",
    "# import standard PyTorch modules\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import torchbnn as bnn\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import axes3d\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb1286f0",
   "metadata": {},
   "source": [
    "We start as before: generating the same dataset, which is simple enough that we can visualize the results easily.\n",
    "\n",
    "$f(x) = (3 + \\kappa) x^2 + \\epsilon$, where $\\epsilon \\sim \\mathcal{N}(\\mu=0, \\sigma=10)$ and $\\kappa \\sim \\mathcal{N}(\\mu=0, \\sigma=0.03)$.\n",
    "\n",
    "In this case we do know the true model, so it is interesting to take some time to pinpoint the role of $\\kappa$ and $\\epsilon$. These variables add fluctuation to the results. $\\epsilon$ adds Gaussian noise in a way that is completely independent from $x$ and cannot be traced down to a particular functional dependence. $\\kappa$ changes a specific parameter of the model, in this case the coefficient 3, by around 1%.\n",
    "\n",
    "When fitting a model, the nomenclature *epistemic uncertainty* is often used to refer to uncertainties coming to effects related to different functional models. That is, one can imagine that there are different functions that may fit the data due to the effect of $\\kappa$, such as: $g(x) = 3x^2$ or $h(x) = 2.95x^2$.\n",
    "\n",
    "The nomenclature *aleatoric uncertainty* is used to refer to whichever uncertainty cannot be tracked down to a given model dependence. In this example, different constant factors could be added to the model $g$ to account for the fluctuations in $\\epsilon$.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5d457cd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_data(N: int) -> np.ndarray:\n",
    "    x = 2*np.random.randn(N, 1)\n",
    "    epsilon = 10*np.random.randn(N, 1)\n",
    "    kappa = 0.03*np.random.randn(N, 1)\n",
    "    z = (3 + kappa)*x**2 + epsilon\n",
    "    return np.concatenate((x, z), axis=1).astype(np.float32)\n",
    "\n",
    "train_data = generate_data(N=1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48433f6f",
   "metadata": {},
   "source": [
    "PyTorch allows you to create a class that outputs a single data entry and use that to feed input to your neural network. We will use the following class to feed the data to the neural network. This looks useless if all your data fits in a Numpy array, but notice that if you have a lot of data and cannot load it all in memory, this allows you to read data on demand, as you need it and only the needed samples are stored at a single time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "30205402",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, data: np.ndarray):\n",
    "        self.data = data\n",
    "    def __len__(self) -> int:\n",
    "        \"\"\"How many samples do I have?\"\"\"\n",
    "        return len(self.data)\n",
    "    def __getitem__(self, idx):\n",
    "        # give me item with index idx\n",
    "        return {\"data\": self.data[idx, 0:1], \"target\": self.data[idx, 1:]}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cc0b0774",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1000\n"
     ]
    }
   ],
   "source": [
    "my_dataset = MyDataset(train_data)\n",
    "print(len(my_dataset))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6dccfac6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'data': array([-3.1838348], dtype=float32), 'target': array([27.882494], dtype=float32)}\n"
     ]
    }
   ],
   "source": [
    "print(my_dataset[1])"
   ]
  },
  {
   "attachments": {
    "elbo.png": {
     "image/png": "[encoded data, 27968 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "6bd7c62d",
   "metadata": {},
   "source": [
    "## Let's add uncertainties\n",
    "\n",
    "The method shown before finds the neural network parameters which maximize the log-likelihood of the data. But not all parameters are equally likely and we can estimate an uncertainty for them.\n",
    "\n",
    "With an uncertainty for the parameters, we can propagate the uncertainty through the neural network and obtain an uncertainty on the prediction of the regression output.\n",
    "\n",
    "This can be done assuming each weight in the network function has a given probability distribution and instead of fitting a single value for the weight, we fit the parameters of this probability distribution. For the example shown here, we assume that the probability distribution of the weights is Gaussian and we aim to obtain the mean and variance of the Gaussian.\n",
    "\n",
    "We are going to include the epistemic uncertainty through the variation of the weights. That is, the fact that the weights vary and lead to different effective functions allow us to model different $f(x)$ dependence relationships and this is attributed to the epistemic uncertainty.\n",
    "\n",
    "We additionally assume that the data collected has some aleatoric uncertainty, which means that every point is uncertain by some fixed unknown amount. To model this effect, we assume that the likelihood function $p(\\text{data}|\\theta)$ can be modelled by a Gaussian distribution with a certain standard deviation $\\sigma_a$. This standard deviation will be used to model the aleatoric uncertainty.\n",
    "\n",
    "There are several ways to predict uncertainties in a model in general. The method we use here is called Variational Inference. In this method, we assume that we know to a good approximation some probability distribution for the function $f(x)$, but we don't know its exact probabilities. For example, we assume that the weights of the function are distributed according to a Gaussian and this already restricts how some functions $f(x)$ are more likely than others. The task then becomes simply how to discover which weights are more likely than others, transforming a function- and distribution-finding problem into a weights-finding problem. How do we find then the weights under this approximation?\n",
    "\n",
    "Since we assumed we know that all weights are distributed with a Gaussian distribution, we already know that if we fix the mean value for each weight and the standard deviation of each weight, we can in principle calculate the probability distribution modelled for our network. Let us call this probability distribution $q$. We know what the true probability distribution for the weights should really be, since this comes from the Bayes theorem: \n",
    "\n",
    "$p(\\text{weights}|\\text{data}) \\propto p(\\text{data}|\\text{weights}) p(\\text{weights})$.\n",
    "\n",
    "In this equation, $p(\\text{weights})$ is Gaussian for each weight and $p(\\text{data}|\\text{weights})$ is the likelihood that our prediction matches data with a fixed set of weights, which we shall also assume to be Gaussian and its standard deviation shall model the aleatoric uncertainty. Now, since we know the true form of the weights distribution, given the data, we can minimize some sort of *distance* between $p(\\text{weights}|\\text{data})$ and $q$, which is our approximation. How can we measure distances between probability distributions? There isn't exactly an easy-to-use concept for distances here (but see this if you want to go into more advanced things: http://www.scholarpedia.org/article/Fisher-Rao_metric ), but there is a concept of *divergence* that we can easily take advantage of. Divergences do not necessarily respect the triangular inequality, so they are not interchangeable with distances, but they provide the key element that we need here: smaller divergences between $q$ and $p(\\text{weights}|\\text{data})$ lead to them being more similar to one another.\n",
    "\n",
    "The divergence we shall use here is called the Kullback-Leibler divergence and it is closely related to the concept of Mutual Information and Entropy in Information Theory (which in reality \"borrowed\" these ideas from Statistical Mechanics). The Kullback-Leibler divergence between two distributions is defined as:\n",
    "\n",
    "$KL(q||p) = \\int q(x) \\log \\left(\\frac{p(x)}{q(x)}\\right) dx$.\n",
    "\n",
    "If we plug in this definition of the Kullback-Leibler divergence in the analytical form of our $q$, we obtain a final number that we may choose to minimize, so that our final $p(\\text{weights}|\\text{data})$ is maximized. The details of this derivation can be seen here and in cited sources: https://arxiv.org/pdf/1505.05424.pdf\n",
    "\n",
    "The final loss function obtained is:\n",
    "\n",
    "$\\mathcal{L} = -\\mathbb{E}_{\\text{data}}\\left[\\log p(\\text{data}|\\text{weights})\\right] + \\frac{1}{M} KL(\\text{weights}|\\text{prior})$\n",
    "\n",
    "As mentioned, the first term is assumed to come from a Gaussian with the standard deviation given by the aleatoric uncertainty (assumed to be the same for every data point, but this could be changed to be data-point specific as well!). The second term corresponds to a penalty for varying the weights away from the prior assumption that the weights are Gaussian with a mean zero and standard deviation 0.1. In this equation, $M$ is the number of batches used.\n",
    "\n",
    "A quick sketch from the mentioned paper showing the justification behind this equation is given below:\n",
    "\n",
    "![elbo.png](attachment:elbo.png)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f8d501ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BayesianNetwork(nn.Module):\n",
    "    \"\"\"\n",
    "        A model Bayesian Neural network.\n",
    "        Each weight is represented by a Gaussian with a mean and a standard deviation.\n",
    "        Each evaluation of forward leads to a different choice of the weights, so running\n",
    "        forward several times we can check the effect of the weights variation on the same input.\n",
    "        The nll function implements the negative log likelihood to be used as the first part of the loss\n",
    "        function (the second shall be the Kullback-Leibler divergence).\n",
    "        The negative log-likelihood is simply the negative log likelihood of a Gaussian\n",
    "        between the prediction and the true value. The standard deviation of the Gaussian is left as a\n",
    "        parameter to be fit: sigma.\n",
    "    \"\"\"\n",
    "    def __init__(self, input_dimension: int=1, output_dimension: int=1):\n",
    "        super(BayesianNetwork, self).__init__()\n",
    "        hidden_dimension = 100\n",
    "        self.model = nn.Sequential(\n",
    "                                   bnn.BayesLinear(prior_mu=0,\n",
    "                                                   prior_sigma=1,\n",
    "                                                   in_features=input_dimension,\n",
    "                                                   out_features=hidden_dimension),\n",
    "                                   nn.ReLU(),\n",
    "                                   bnn.BayesLinear(prior_mu=0,\n",
    "                                                   prior_sigma=1,\n",
    "                                                   in_features=hidden_dimension,\n",
    "                                                   out_features=hidden_dimension),\n",
    "                                   nn.ReLU(),\n",
    "                                   bnn.BayesLinear(prior_mu=0,\n",
    "                                                   prior_sigma=1,\n",
    "                                                   in_features=hidden_dimension,\n",
    "                                                   out_features=output_dimension)\n",
    "                                    )\n",
    "        self.log_sigma2 = nn.Parameter(torch.ones(1), requires_grad=True)\n",
    "\n",
    "    def forward(self, x: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Calculate the result f(x) applied on the input x.\n",
    "        \"\"\"\n",
    "        return self.model(x)\n",
    "\n",
    "    def nll(self, prediction: torch.Tensor, target: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Calculate the negative log-likelihood (divided by the batch size, since we take the mean).\n",
    "        \"\"\"\n",
    "        error = prediction - target\n",
    "        squared_error = error**2\n",
    "        sigma2 = torch.exp(self.log_sigma2)[0]\n",
    "        norm_error = 0.5*squared_error/sigma2\n",
    "        norm_term = 0.5*(np.log(2*np.pi) + self.log_sigma2[0])*target.shape[1]\n",
    "        return norm_error.sum(1).mean(0) + norm_term\n",
    "\n",
    "    def aleatoric_uncertainty(self) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "            Get the aleatoric component of the uncertainty.\n",
    "        \"\"\"\n",
    "        return torch.exp(0.5*self.log_sigma2[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "7b9beb21",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create the neural network:\n",
    "b_network = BayesianNetwork()\n",
    "\n",
    "# create the object to load the data:\n",
    "B = 10\n",
    "loader = torch.utils.data.DataLoader(my_dataset, batch_size=B)\n",
    "\n",
    "# create the optimizer to be used \n",
    "optimizer = torch.optim.Adam(b_network.parameters(), lr=0.001)\n",
    "\n",
    "# the Kullback-Leibler divergence should be scaled by 1/number_of_batches\n",
    "# see https://arxiv.org/abs/1505.05424 for more information on this\n",
    "number_of_batches = len(my_dataset)/float(B)\n",
    "weight_kl = 1.0/float(number_of_batches)\n",
    "# divide by B, because we implicitly divide by the number of sample taking the mean in the `nll` function above\n",
    "weight_kl /= float(B)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c68ba2e2",
   "metadata": {},
   "source": [
    "The criteria for finding the optimal weights are based on the Bayes' theorem, on which the posterior probability of the weights is proportional to the likelihood of the data given the weights and to the prior probability of the weights. We assume the prior probability of the weights is Gaussian corresponding to a unit Gaussian centred at zero and with standard deviation 0.1. This prior has a regularizing effect, preventing overtraining.\n",
    "\n",
    "We can translate the Bayes theorem and the assumption that the final posterior distribution is also Gaussian into an optimization procedure to find the posterior mean and variance of the posterior distribution. The function optimized to obtain the mean and variances of the Gaussians for the weights is the sum between the mean-squared-error (corresponding to a Gaussian log-likelihood of the data) and the Kullback-Leibler divergence between the weights distribution and the prior Gaussian."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "fbea6b0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "kl_loss = bnn.BKLLoss(reduction='sum', last_layer_only=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b92ed4b0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0/500  total: 2294.17400, -LL: 703.32788, prior: 0.05234, aleatoric unc.: 1.70432\n",
      "Epoch 1/500  total: 2164.52190, -LL: 4314.84668, prior: 0.05399, aleatoric unc.: 1.75715\n",
      "Epoch 2/500  total: 1671.93444, -LL: 330.02435, prior: 0.05557, aleatoric unc.: 1.80130\n",
      "Epoch 3/500  total: 1633.34416, -LL: 1548.86133, prior: 0.05711, aleatoric unc.: 1.84716\n",
      "Epoch 4/500  total: 1527.28808, -LL: 247.82253, prior: 0.05899, aleatoric unc.: 1.89552\n",
      "Epoch 5/500  total: 1440.56923, -LL: 314.48566, prior: 0.06116, aleatoric unc.: 1.94140\n",
      "Epoch 6/500  total: 1575.21707, -LL: 193.92581, prior: 0.06411, aleatoric unc.: 1.99621\n",
      "Epoch 7/500  total: 1020.99699, -LL: 311.39798, prior: 0.06605, aleatoric unc.: 2.03106\n",
      "Epoch 8/500  total: 889.46376, -LL: 228.17020, prior: 0.06773, aleatoric unc.: 2.06522\n",
      "Epoch 9/500  total: 1075.89125, -LL: 2626.64404, prior: 0.07027, aleatoric unc.: 2.10846\n",
      "Epoch 10/500  total: 783.86553, -LL: 1672.57385, prior: 0.07258, aleatoric unc.: 2.14209\n",
      "Epoch 11/500  total: 836.25104, -LL: 971.67743, prior: 0.07488, aleatoric unc.: 2.17698\n",
      "Epoch 12/500  total: 572.92991, -LL: 691.55060, prior: 0.07698, aleatoric unc.: 2.20801\n",
      "Epoch 13/500  total: 516.95223, -LL: 516.95093, prior: 0.07892, aleatoric unc.: 2.23500\n",
      "Epoch 14/500  total: 785.52413, -LL: 916.61438, prior: 0.08203, aleatoric unc.: 2.27769\n",
      "Epoch 15/500  total: 591.27327, -LL: 184.35178, prior: 0.08498, aleatoric unc.: 2.31321\n",
      "Epoch 16/500  total: 427.04376, -LL: 561.74152, prior: 0.08683, aleatoric unc.: 2.33933\n",
      "Epoch 17/500  total: 524.27719, -LL: 15.85236, prior: 0.08955, aleatoric unc.: 2.37532\n",
      "Epoch 18/500  total: 430.90639, -LL: 273.38846, prior: 0.09209, aleatoric unc.: 2.40437\n",
      "Epoch 19/500  total: 383.77823, -LL: 6.95667, prior: 0.09448, aleatoric unc.: 2.43451\n",
      "Epoch 20/500  total: 378.80883, -LL: 264.90891, prior: 0.09667, aleatoric unc.: 2.46322\n",
      "Epoch 21/500  total: 315.44016, -LL: 72.41981, prior: 0.09888, aleatoric unc.: 2.49201\n",
      "Epoch 22/500  total: 311.58309, -LL: 175.78893, prior: 0.10110, aleatoric unc.: 2.51807\n",
      "Epoch 23/500  total: 329.50581, -LL: 572.85419, prior: 0.10369, aleatoric unc.: 2.55035\n",
      "Epoch 24/500  total: 275.12305, -LL: 307.36102, prior: 0.10603, aleatoric unc.: 2.57908\n",
      "Epoch 25/500  total: 238.38241, -LL: 58.51337, prior: 0.10824, aleatoric unc.: 2.60498\n",
      "Epoch 26/500  total: 254.61200, -LL: 339.08179, prior: 0.11041, aleatoric unc.: 2.63303\n",
      "Epoch 27/500  total: 268.43455, -LL: 1827.35986, prior: 0.11269, aleatoric unc.: 2.66386\n",
      "Epoch 28/500  total: 200.06057, -LL: 172.36951, prior: 0.11525, aleatoric unc.: 2.69245\n",
      "Epoch 29/500  total: 185.10391, -LL: 676.92700, prior: 0.11748, aleatoric unc.: 2.71670\n",
      "Epoch 30/500  total: 185.27358, -LL: 353.72168, prior: 0.11965, aleatoric unc.: 2.74223\n",
      "Epoch 31/500  total: 202.86429, -LL: 56.65427, prior: 0.12221, aleatoric unc.: 2.77346\n",
      "Epoch 32/500  total: 210.36308, -LL: 88.56605, prior: 0.12501, aleatoric unc.: 2.80590\n",
      "Epoch 33/500  total: 190.61076, -LL: 109.12581, prior: 0.12819, aleatoric unc.: 2.83927\n",
      "Epoch 34/500  total: 156.23295, -LL: 87.33363, prior: 0.13093, aleatoric unc.: 2.86657\n",
      "Epoch 35/500  total: 158.04877, -LL: 245.68523, prior: 0.13412, aleatoric unc.: 2.89595\n",
      "Epoch 36/500  total: 119.76969, -LL: 119.74051, prior: 0.13642, aleatoric unc.: 2.92126\n",
      "Epoch 37/500  total: 135.00420, -LL: 33.82713, prior: 0.13874, aleatoric unc.: 2.94849\n",
      "Epoch 38/500  total: 117.64032, -LL: 39.42981, prior: 0.14125, aleatoric unc.: 2.97495\n",
      "Epoch 39/500  total: 131.96143, -LL: 98.90752, prior: 0.14446, aleatoric unc.: 3.00607\n",
      "Epoch 40/500  total: 129.90317, -LL: 234.60854, prior: 0.14738, aleatoric unc.: 3.03759\n",
      "Epoch 41/500  total: 149.99737, -LL: 20.71658, prior: 0.15074, aleatoric unc.: 3.07639\n",
      "Epoch 42/500  total: 103.29358, -LL: 118.18483, prior: 0.15347, aleatoric unc.: 3.10597\n",
      "Epoch 43/500  total: 109.35400, -LL: 66.08629, prior: 0.15591, aleatoric unc.: 3.13605\n",
      "Epoch 44/500  total: 93.15535, -LL: 7.04618, prior: 0.15849, aleatoric unc.: 3.16537\n",
      "Epoch 45/500  total: 83.13112, -LL: 409.44940, prior: 0.16091, aleatoric unc.: 3.19226\n",
      "Epoch 46/500  total: 91.04438, -LL: 73.30114, prior: 0.16382, aleatoric unc.: 3.22404\n",
      "Epoch 47/500  total: 83.55522, -LL: 244.28879, prior: 0.16646, aleatoric unc.: 3.25397\n",
      "Epoch 48/500  total: 73.99347, -LL: 57.34751, prior: 0.16900, aleatoric unc.: 3.28170\n",
      "Epoch 49/500  total: 86.06144, -LL: 97.08743, prior: 0.17213, aleatoric unc.: 3.31545\n",
      "Epoch 50/500  total: 70.17201, -LL: 472.36411, prior: 0.17489, aleatoric unc.: 3.34469\n",
      "Epoch 51/500  total: 61.18556, -LL: 65.52729, prior: 0.17738, aleatoric unc.: 3.37370\n",
      "Epoch 52/500  total: 62.92316, -LL: 35.62979, prior: 0.18012, aleatoric unc.: 3.40398\n",
      "Epoch 53/500  total: 71.60829, -LL: 121.34455, prior: 0.18321, aleatoric unc.: 3.43726\n",
      "Epoch 54/500  total: 53.54997, -LL: 25.85874, prior: 0.18630, aleatoric unc.: 3.46713\n",
      "Epoch 55/500  total: 72.00782, -LL: 46.47833, prior: 0.18960, aleatoric unc.: 3.50607\n",
      "Epoch 56/500  total: 53.76796, -LL: 30.84623, prior: 0.19247, aleatoric unc.: 3.53761\n",
      "Epoch 57/500  total: 48.81183, -LL: 114.88335, prior: 0.19508, aleatoric unc.: 3.56618\n",
      "Epoch 58/500  total: 44.38314, -LL: 53.28208, prior: 0.19776, aleatoric unc.: 3.59475\n",
      "Epoch 59/500  total: 55.60695, -LL: 47.80334, prior: 0.20111, aleatoric unc.: 3.63241\n",
      "Epoch 60/500  total: 48.88861, -LL: 143.25693, prior: 0.20443, aleatoric unc.: 3.66788\n",
      "Epoch 61/500  total: 53.50728, -LL: 4.94818, prior: 0.20783, aleatoric unc.: 3.70812\n",
      "Epoch 62/500  total: 47.01873, -LL: 77.59438, prior: 0.21115, aleatoric unc.: 3.74518\n",
      "Epoch 63/500  total: 46.15298, -LL: 187.30032, prior: 0.21436, aleatoric unc.: 3.78296\n",
      "Epoch 64/500  total: 38.49360, -LL: 192.40350, prior: 0.21705, aleatoric unc.: 3.81615\n",
      "Epoch 65/500  total: 38.48498, -LL: 89.96780, prior: 0.21999, aleatoric unc.: 3.85294\n",
      "Epoch 66/500  total: 33.08836, -LL: 19.20866, prior: 0.22247, aleatoric unc.: 3.88479\n",
      "Epoch 67/500  total: 32.53665, -LL: 10.17999, prior: 0.22546, aleatoric unc.: 3.91701\n",
      "Epoch 68/500  total: 31.42953, -LL: 30.17547, prior: 0.22803, aleatoric unc.: 3.95171\n",
      "Epoch 69/500  total: 29.76948, -LL: 60.30394, prior: 0.23062, aleatoric unc.: 3.98382\n",
      "Epoch 70/500  total: 36.32508, -LL: 22.38429, prior: 0.23436, aleatoric unc.: 4.02776\n",
      "Epoch 71/500  total: 29.30958, -LL: 7.98826, prior: 0.23705, aleatoric unc.: 4.06465\n",
      "Epoch 72/500  total: 28.79093, -LL: 6.86609, prior: 0.24080, aleatoric unc.: 4.10310\n",
      "Epoch 73/500  total: 24.90846, -LL: 39.40199, prior: 0.24306, aleatoric unc.: 4.13660\n",
      "Epoch 74/500  total: 27.69084, -LL: 14.49685, prior: 0.24588, aleatoric unc.: 4.17586\n",
      "Epoch 75/500  total: 30.61155, -LL: 66.61439, prior: 0.24992, aleatoric unc.: 4.22242\n",
      "Epoch 76/500  total: 25.08618, -LL: 42.09279, prior: 0.25323, aleatoric unc.: 4.26264\n",
      "Epoch 77/500  total: 29.45682, -LL: 6.08524, prior: 0.25746, aleatoric unc.: 4.31628\n",
      "Epoch 78/500  total: 21.04458, -LL: 15.61365, prior: 0.26008, aleatoric unc.: 4.35182\n",
      "Epoch 79/500  total: 23.43780, -LL: 14.76187, prior: 0.26322, aleatoric unc.: 4.39626\n",
      "Epoch 80/500  total: 20.39375, -LL: 52.22060, prior: 0.26632, aleatoric unc.: 4.43429\n",
      "Epoch 81/500  total: 20.52036, -LL: 20.75696, prior: 0.26912, aleatoric unc.: 4.47560\n",
      "Epoch 82/500  total: 20.58786, -LL: 23.43789, prior: 0.27223, aleatoric unc.: 4.51860\n",
      "Epoch 83/500  total: 22.81071, -LL: 4.76024, prior: 0.27617, aleatoric unc.: 4.57214\n",
      "Epoch 84/500  total: 21.30295, -LL: 11.03628, prior: 0.27926, aleatoric unc.: 4.62280\n",
      "Epoch 85/500  total: 19.57655, -LL: 25.29999, prior: 0.28254, aleatoric unc.: 4.67018\n",
      "Epoch 86/500  total: 16.27523, -LL: 39.32928, prior: 0.28471, aleatoric unc.: 4.71189\n",
      "Epoch 87/500  total: 17.87385, -LL: 5.49629, prior: 0.28736, aleatoric unc.: 4.76154\n",
      "Epoch 88/500  total: 17.28383, -LL: 18.67742, prior: 0.29044, aleatoric unc.: 4.81090\n",
      "Epoch 89/500  total: 18.66865, -LL: 3.55819, prior: 0.29501, aleatoric unc.: 4.86776\n",
      "Epoch 90/500  total: 15.61203, -LL: 38.41304, prior: 0.29754, aleatoric unc.: 4.91515\n",
      "Epoch 91/500  total: 15.95930, -LL: 49.24847, prior: 0.30078, aleatoric unc.: 4.96588\n",
      "Epoch 92/500  total: 16.55630, -LL: 20.53808, prior: 0.30533, aleatoric unc.: 5.02590\n",
      "Epoch 93/500  total: 12.66083, -LL: 4.94757, prior: 0.30831, aleatoric unc.: 5.07009\n",
      "Epoch 94/500  total: 14.76228, -LL: 71.20815, prior: 0.31009, aleatoric unc.: 5.12316\n",
      "Epoch 95/500  total: 14.62404, -LL: 14.50025, prior: 0.31300, aleatoric unc.: 5.18295\n",
      "Epoch 96/500  total: 14.90159, -LL: 12.78674, prior: 0.31677, aleatoric unc.: 5.24626\n",
      "Epoch 97/500  total: 11.74080, -LL: 7.30341, prior: 0.31915, aleatoric unc.: 5.29628\n",
      "Epoch 98/500  total: 11.48438, -LL: 5.00458, prior: 0.32149, aleatoric unc.: 5.34559\n",
      "Epoch 99/500  total: 11.39528, -LL: 9.96454, prior: 0.32386, aleatoric unc.: 5.39713\n",
      "Epoch 100/500  total: 12.73329, -LL: 7.94399, prior: 0.32652, aleatoric unc.: 5.45732\n",
      "Epoch 101/500  total: 12.82492, -LL: 5.70899, prior: 0.33021, aleatoric unc.: 5.52423\n",
      "Epoch 102/500  total: 10.29261, -LL: 6.06767, prior: 0.33286, aleatoric unc.: 5.57872\n",
      "Epoch 103/500  total: 11.17191, -LL: 14.21713, prior: 0.33638, aleatoric unc.: 5.63869\n",
      "Epoch 104/500  total: 11.81327, -LL: 7.13694, prior: 0.34007, aleatoric unc.: 5.70998\n",
      "Epoch 105/500  total: 9.54222, -LL: 24.48591, prior: 0.34191, aleatoric unc.: 5.76379\n",
      "Epoch 106/500  total: 10.36091, -LL: 6.09785, prior: 0.34401, aleatoric unc.: 5.82894\n",
      "Epoch 107/500  total: 9.04908, -LL: 11.80495, prior: 0.34631, aleatoric unc.: 5.88464\n",
      "Epoch 108/500  total: 11.12780, -LL: 6.29300, prior: 0.35127, aleatoric unc.: 5.96363\n",
      "Epoch 109/500  total: 9.48051, -LL: 9.15429, prior: 0.35349, aleatoric unc.: 6.03229\n",
      "Epoch 110/500  total: 8.67597, -LL: 6.08984, prior: 0.35566, aleatoric unc.: 6.09170\n",
      "Epoch 111/500  total: 9.08319, -LL: 5.30393, prior: 0.35811, aleatoric unc.: 6.16096\n",
      "Epoch 112/500  total: 9.66198, -LL: 25.57287, prior: 0.36034, aleatoric unc.: 6.23762\n",
      "Epoch 113/500  total: 9.07391, -LL: 6.92154, prior: 0.36347, aleatoric unc.: 6.31789\n",
      "Epoch 114/500  total: 9.18264, -LL: 9.77642, prior: 0.36589, aleatoric unc.: 6.39721\n",
      "Epoch 115/500  total: 8.56463, -LL: 4.63342, prior: 0.36976, aleatoric unc.: 6.47501\n",
      "Epoch 116/500  total: 8.41445, -LL: 16.63805, prior: 0.37295, aleatoric unc.: 6.54858\n",
      "Epoch 117/500  total: 8.15860, -LL: 4.07057, prior: 0.37583, aleatoric unc.: 6.62903\n",
      "Epoch 118/500  total: 7.66945, -LL: 12.75606, prior: 0.37780, aleatoric unc.: 6.70255\n",
      "Epoch 119/500  total: 7.49265, -LL: 8.45567, prior: 0.38185, aleatoric unc.: 6.77566\n",
      "Epoch 120/500  total: 7.88782, -LL: 8.02437, prior: 0.38530, aleatoric unc.: 6.86073\n",
      "Epoch 121/500  total: 7.91334, -LL: 9.30178, prior: 0.38813, aleatoric unc.: 6.95149\n",
      "Epoch 122/500  total: 8.09812, -LL: 4.83049, prior: 0.39001, aleatoric unc.: 7.04736\n",
      "Epoch 123/500  total: 7.35547, -LL: 5.80612, prior: 0.39266, aleatoric unc.: 7.13706\n",
      "Epoch 124/500  total: 7.19076, -LL: 7.61347, prior: 0.39384, aleatoric unc.: 7.22395\n",
      "Epoch 125/500  total: 7.22080, -LL: 7.48228, prior: 0.39673, aleatoric unc.: 7.31573\n",
      "Epoch 126/500  total: 6.51365, -LL: 10.89216, prior: 0.39865, aleatoric unc.: 7.39362\n",
      "Epoch 127/500  total: 6.40662, -LL: 6.31307, prior: 0.39938, aleatoric unc.: 7.47254\n",
      "Epoch 128/500  total: 6.93306, -LL: 7.01754, prior: 0.40160, aleatoric unc.: 7.56822\n",
      "Epoch 129/500  total: 6.74271, -LL: 6.63430, prior: 0.40347, aleatoric unc.: 7.66625\n",
      "Epoch 130/500  total: 6.44686, -LL: 11.64886, prior: 0.40589, aleatoric unc.: 7.75812\n",
      "Epoch 131/500  total: 6.72675, -LL: 5.24210, prior: 0.40782, aleatoric unc.: 7.86481\n",
      "Epoch 132/500  total: 6.80228, -LL: 13.50842, prior: 0.41071, aleatoric unc.: 7.97619\n",
      "Epoch 133/500  total: 5.97443, -LL: 4.43141, prior: 0.41181, aleatoric unc.: 8.06552\n",
      "Epoch 134/500  total: 5.92402, -LL: 5.32649, prior: 0.41321, aleatoric unc.: 8.14983\n",
      "Epoch 135/500  total: 6.17365, -LL: 5.18150, prior: 0.41440, aleatoric unc.: 8.25499\n",
      "Epoch 136/500  total: 5.96319, -LL: 5.07157, prior: 0.41622, aleatoric unc.: 8.35583\n",
      "Epoch 137/500  total: 6.23591, -LL: 5.90405, prior: 0.41917, aleatoric unc.: 8.47259\n",
      "Epoch 138/500  total: 6.16629, -LL: 7.00212, prior: 0.42168, aleatoric unc.: 8.59072\n",
      "Epoch 139/500  total: 5.70816, -LL: 4.46014, prior: 0.42198, aleatoric unc.: 8.69135\n",
      "Epoch 140/500  total: 5.89141, -LL: 4.32210, prior: 0.42329, aleatoric unc.: 8.80822\n",
      "Epoch 141/500  total: 5.86564, -LL: 4.79264, prior: 0.42409, aleatoric unc.: 8.92402\n",
      "Epoch 142/500  total: 5.82845, -LL: 7.37854, prior: 0.42666, aleatoric unc.: 9.04643\n",
      "Epoch 143/500  total: 5.52245, -LL: 6.09641, prior: 0.42617, aleatoric unc.: 9.15368\n",
      "Epoch 144/500  total: 5.55177, -LL: 3.60275, prior: 0.42706, aleatoric unc.: 9.26851\n",
      "Epoch 145/500  total: 5.50052, -LL: 7.01150, prior: 0.42671, aleatoric unc.: 9.38153\n",
      "Epoch 146/500  total: 5.57487, -LL: 4.41672, prior: 0.42908, aleatoric unc.: 9.51134\n",
      "Epoch 147/500  total: 5.30800, -LL: 4.13941, prior: 0.43043, aleatoric unc.: 9.61923\n",
      "Epoch 148/500  total: 5.39514, -LL: 3.87754, prior: 0.43000, aleatoric unc.: 9.74195\n",
      "Epoch 149/500  total: 5.25161, -LL: 4.52955, prior: 0.43019, aleatoric unc.: 9.85236\n",
      "Epoch 150/500  total: 5.23710, -LL: 5.29330, prior: 0.43140, aleatoric unc.: 9.96564\n",
      "Epoch 151/500  total: 5.55625, -LL: 12.59121, prior: 0.43435, aleatoric unc.: 10.11595\n",
      "Epoch 152/500  total: 5.40288, -LL: 5.01670, prior: 0.43468, aleatoric unc.: 10.26004\n",
      "Epoch 153/500  total: 5.29786, -LL: 3.69959, prior: 0.43491, aleatoric unc.: 10.40679\n",
      "Epoch 154/500  total: 5.17616, -LL: 5.03672, prior: 0.43485, aleatoric unc.: 10.53531\n",
      "Epoch 155/500  total: 5.16887, -LL: 3.54868, prior: 0.43612, aleatoric unc.: 10.66768\n",
      "Epoch 156/500  total: 5.13514, -LL: 5.77639, prior: 0.43507, aleatoric unc.: 10.79695\n",
      "Epoch 157/500  total: 5.15844, -LL: 5.82699, prior: 0.43707, aleatoric unc.: 10.93950\n",
      "Epoch 158/500  total: 4.89757, -LL: 4.33982, prior: 0.43826, aleatoric unc.: 11.05050\n",
      "Epoch 159/500  total: 5.09657, -LL: 9.34318, prior: 0.43866, aleatoric unc.: 11.18306\n",
      "Epoch 160/500  total: 5.03004, -LL: 4.23301, prior: 0.43979, aleatoric unc.: 11.32898\n",
      "Epoch 161/500  total: 5.19724, -LL: 6.94482, prior: 0.44151, aleatoric unc.: 11.49262\n",
      "Epoch 162/500  total: 4.93899, -LL: 3.74737, prior: 0.44072, aleatoric unc.: 11.62339\n",
      "Epoch 163/500  total: 5.00967, -LL: 4.06411, prior: 0.44184, aleatoric unc.: 11.76584\n",
      "Epoch 164/500  total: 4.92523, -LL: 3.70284, prior: 0.44245, aleatoric unc.: 11.90633\n",
      "Epoch 165/500  total: 4.90263, -LL: 5.19476, prior: 0.44214, aleatoric unc.: 12.03179\n",
      "Epoch 166/500  total: 4.90347, -LL: 5.04848, prior: 0.44011, aleatoric unc.: 12.17206\n",
      "Epoch 167/500  total: 4.85970, -LL: 4.04275, prior: 0.43931, aleatoric unc.: 12.30186\n",
      "Epoch 168/500  total: 4.90966, -LL: 6.62824, prior: 0.43810, aleatoric unc.: 12.45021\n",
      "Epoch 169/500  total: 4.91850, -LL: 4.91183, prior: 0.43710, aleatoric unc.: 12.60517\n",
      "Epoch 170/500  total: 5.03885, -LL: 4.35656, prior: 0.43715, aleatoric unc.: 12.80363\n",
      "Epoch 171/500  total: 4.86051, -LL: 4.03450, prior: 0.43588, aleatoric unc.: 12.94091\n",
      "Epoch 172/500  total: 4.86936, -LL: 4.70055, prior: 0.43449, aleatoric unc.: 13.09483\n",
      "Epoch 173/500  total: 4.83001, -LL: 3.85964, prior: 0.43462, aleatoric unc.: 13.22493\n",
      "Epoch 174/500  total: 4.89690, -LL: 3.95164, prior: 0.43540, aleatoric unc.: 13.39555\n",
      "Epoch 175/500  total: 4.92513, -LL: 5.99314, prior: 0.43664, aleatoric unc.: 13.56865\n",
      "Epoch 176/500  total: 4.89771, -LL: 4.14453, prior: 0.44024, aleatoric unc.: 13.74470\n",
      "Epoch 177/500  total: 4.69510, -LL: 3.80845, prior: 0.43905, aleatoric unc.: 13.83333\n",
      "Epoch 178/500  total: 4.77653, -LL: 3.71582, prior: 0.44094, aleatoric unc.: 13.95735\n",
      "Epoch 179/500  total: 4.76925, -LL: 5.64237, prior: 0.43987, aleatoric unc.: 14.08263\n",
      "Epoch 180/500  total: 4.71789, -LL: 3.96352, prior: 0.43841, aleatoric unc.: 14.19105\n",
      "Epoch 181/500  total: 4.77807, -LL: 3.81962, prior: 0.43777, aleatoric unc.: 14.32319\n",
      "Epoch 182/500  total: 4.72213, -LL: 6.11314, prior: 0.43537, aleatoric unc.: 14.42246\n",
      "Epoch 183/500  total: 4.69993, -LL: 4.08305, prior: 0.43216, aleatoric unc.: 14.51312\n",
      "Epoch 184/500  total: 4.73802, -LL: 4.12818, prior: 0.43229, aleatoric unc.: 14.65313\n",
      "Epoch 185/500  total: 4.67419, -LL: 4.12100, prior: 0.42907, aleatoric unc.: 14.72696\n",
      "Epoch 186/500  total: 4.83301, -LL: 3.82574, prior: 0.43024, aleatoric unc.: 14.91049\n",
      "Epoch 187/500  total: 4.76206, -LL: 3.85872, prior: 0.42845, aleatoric unc.: 15.05879\n",
      "Epoch 188/500  total: 4.72714, -LL: 5.29572, prior: 0.42632, aleatoric unc.: 15.15222\n",
      "Epoch 189/500  total: 4.61022, -LL: 4.08400, prior: 0.42602, aleatoric unc.: 15.21289\n",
      "Epoch 190/500  total: 4.65364, -LL: 4.05610, prior: 0.42413, aleatoric unc.: 15.27519\n",
      "Epoch 191/500  total: 4.67309, -LL: 3.88015, prior: 0.42243, aleatoric unc.: 15.35476\n",
      "Epoch 192/500  total: 4.70079, -LL: 6.96650, prior: 0.41984, aleatoric unc.: 15.43811\n",
      "Epoch 193/500  total: 4.72687, -LL: 3.88274, prior: 0.41902, aleatoric unc.: 15.57714\n",
      "Epoch 194/500  total: 4.76679, -LL: 4.13592, prior: 0.41894, aleatoric unc.: 15.73339\n",
      "Epoch 195/500  total: 4.60989, -LL: 4.12250, prior: 0.41892, aleatoric unc.: 15.76806\n",
      "Epoch 196/500  total: 4.69393, -LL: 3.89446, prior: 0.41817, aleatoric unc.: 15.85876\n",
      "Epoch 197/500  total: 4.60156, -LL: 4.46961, prior: 0.41520, aleatoric unc.: 15.85203\n",
      "Epoch 198/500  total: 4.61889, -LL: 3.85427, prior: 0.41157, aleatoric unc.: 15.88660\n",
      "Epoch 199/500  total: 4.62590, -LL: 3.92039, prior: 0.40973, aleatoric unc.: 15.91067\n",
      "Epoch 200/500  total: 4.60572, -LL: 3.78996, prior: 0.40960, aleatoric unc.: 15.92806\n",
      "Epoch 201/500  total: 4.66548, -LL: 4.11024, prior: 0.40654, aleatoric unc.: 15.99439\n",
      "Epoch 202/500  total: 4.63417, -LL: 3.85801, prior: 0.40403, aleatoric unc.: 16.04211\n",
      "Epoch 203/500  total: 4.71727, -LL: 4.08271, prior: 0.40567, aleatoric unc.: 16.17369\n",
      "Epoch 204/500  total: 4.57863, -LL: 3.82427, prior: 0.40441, aleatoric unc.: 16.14296\n",
      "Epoch 205/500  total: 4.70977, -LL: 4.62687, prior: 0.40592, aleatoric unc.: 16.25426\n",
      "Epoch 206/500  total: 4.59091, -LL: 4.70749, prior: 0.39964, aleatoric unc.: 16.22959\n",
      "Epoch 207/500  total: 4.60897, -LL: 3.91997, prior: 0.39924, aleatoric unc.: 16.23401\n",
      "Epoch 208/500  total: 4.62371, -LL: 4.81561, prior: 0.39723, aleatoric unc.: 16.26201\n",
      "Epoch 209/500  total: 4.59145, -LL: 4.11868, prior: 0.39567, aleatoric unc.: 16.24049\n",
      "Epoch 210/500  total: 4.63973, -LL: 4.25721, prior: 0.39458, aleatoric unc.: 16.29369\n",
      "Epoch 211/500  total: 4.63437, -LL: 3.94333, prior: 0.39565, aleatoric unc.: 16.31607\n",
      "Epoch 212/500  total: 4.62322, -LL: 5.20113, prior: 0.39551, aleatoric unc.: 16.35107\n",
      "Epoch 213/500  total: 4.58316, -LL: 4.30533, prior: 0.39006, aleatoric unc.: 16.33248\n",
      "Epoch 214/500  total: 4.60708, -LL: 4.27092, prior: 0.38732, aleatoric unc.: 16.31506\n",
      "Epoch 215/500  total: 4.55818, -LL: 4.60177, prior: 0.38425, aleatoric unc.: 16.28519\n",
      "Epoch 216/500  total: 4.58447, -LL: 4.49458, prior: 0.38426, aleatoric unc.: 16.26051\n",
      "Epoch 217/500  total: 4.61161, -LL: 4.06498, prior: 0.38396, aleatoric unc.: 16.31507\n",
      "Epoch 218/500  total: 4.54134, -LL: 4.49735, prior: 0.38046, aleatoric unc.: 16.22637\n",
      "Epoch 219/500  total: 4.65291, -LL: 3.99961, prior: 0.37965, aleatoric unc.: 16.31837\n",
      "Epoch 220/500  total: 4.58140, -LL: 4.96251, prior: 0.37568, aleatoric unc.: 16.30339\n",
      "Epoch 221/500  total: 4.59126, -LL: 3.90479, prior: 0.37506, aleatoric unc.: 16.33533\n",
      "Epoch 222/500  total: 4.60168, -LL: 4.12840, prior: 0.37475, aleatoric unc.: 16.35800\n",
      "Epoch 223/500  total: 4.61224, -LL: 4.14133, prior: 0.37542, aleatoric unc.: 16.39504\n",
      "Epoch 224/500  total: 4.56135, -LL: 4.94861, prior: 0.37697, aleatoric unc.: 16.34018\n",
      "Epoch 225/500  total: 4.56561, -LL: 3.90656, prior: 0.37348, aleatoric unc.: 16.31787\n",
      "Epoch 226/500  total: 4.54507, -LL: 3.87136, prior: 0.37132, aleatoric unc.: 16.24625\n",
      "Epoch 227/500  total: 4.57453, -LL: 3.82144, prior: 0.37349, aleatoric unc.: 16.23209\n",
      "Epoch 228/500  total: 4.56564, -LL: 4.09081, prior: 0.37190, aleatoric unc.: 16.20979\n",
      "Epoch 229/500  total: 4.55820, -LL: 4.32381, prior: 0.37007, aleatoric unc.: 16.16996\n",
      "Epoch 230/500  total: 4.53165, -LL: 4.17904, prior: 0.36693, aleatoric unc.: 16.10915\n",
      "Epoch 231/500  total: 4.55167, -LL: 3.88500, prior: 0.36874, aleatoric unc.: 16.08617\n",
      "Epoch 232/500  total: 4.55554, -LL: 4.20630, prior: 0.37019, aleatoric unc.: 16.07415\n",
      "Epoch 233/500  total: 4.64898, -LL: 3.87531, prior: 0.36958, aleatoric unc.: 16.19588\n",
      "Epoch 234/500  total: 4.58516, -LL: 3.82347, prior: 0.37062, aleatoric unc.: 16.23983\n",
      "Epoch 235/500  total: 4.52828, -LL: 3.99777, prior: 0.36784, aleatoric unc.: 16.13410\n",
      "Epoch 236/500  total: 4.56419, -LL: 3.88711, prior: 0.36415, aleatoric unc.: 16.12881\n",
      "Epoch 237/500  total: 4.51364, -LL: 3.85443, prior: 0.36235, aleatoric unc.: 16.06821\n",
      "Epoch 238/500  total: 4.53692, -LL: 4.16873, prior: 0.36272, aleatoric unc.: 16.01422\n",
      "Epoch 239/500  total: 4.52414, -LL: 4.43059, prior: 0.36026, aleatoric unc.: 15.98025\n",
      "Epoch 240/500  total: 4.52974, -LL: 3.82850, prior: 0.35667, aleatoric unc.: 15.96161\n",
      "Epoch 241/500  total: 4.54721, -LL: 3.81967, prior: 0.35717, aleatoric unc.: 15.95436\n",
      "Epoch 242/500  total: 4.52174, -LL: 3.85391, prior: 0.35371, aleatoric unc.: 15.91074\n",
      "Epoch 243/500  total: 4.52171, -LL: 3.84012, prior: 0.35449, aleatoric unc.: 15.87800\n",
      "Epoch 244/500  total: 4.48056, -LL: 4.70426, prior: 0.35078, aleatoric unc.: 15.78308\n",
      "Epoch 245/500  total: 4.53157, -LL: 3.92313, prior: 0.35472, aleatoric unc.: 15.79237\n",
      "Epoch 246/500  total: 4.53108, -LL: 3.85934, prior: 0.35546, aleatoric unc.: 15.77559\n",
      "Epoch 247/500  total: 4.52557, -LL: 4.12003, prior: 0.35167, aleatoric unc.: 15.76002\n",
      "Epoch 248/500  total: 4.50074, -LL: 4.28565, prior: 0.34953, aleatoric unc.: 15.73245\n",
      "Epoch 249/500  total: 4.55166, -LL: 4.17833, prior: 0.34960, aleatoric unc.: 15.77043\n",
      "Epoch 250/500  total: 4.56264, -LL: 4.14792, prior: 0.34951, aleatoric unc.: 15.82986\n",
      "Epoch 251/500  total: 4.54331, -LL: 4.13038, prior: 0.35015, aleatoric unc.: 15.84215\n",
      "Epoch 252/500  total: 4.58222, -LL: 3.93371, prior: 0.34975, aleatoric unc.: 15.92820\n",
      "Epoch 253/500  total: 4.54323, -LL: 4.07649, prior: 0.35027, aleatoric unc.: 15.92687\n",
      "Epoch 254/500  total: 4.55835, -LL: 4.00149, prior: 0.35056, aleatoric unc.: 15.97446\n",
      "Epoch 255/500  total: 4.48282, -LL: 3.77635, prior: 0.35023, aleatoric unc.: 15.89263\n",
      "Epoch 256/500  total: 4.52993, -LL: 4.34475, prior: 0.34669, aleatoric unc.: 15.86973\n",
      "Epoch 257/500  total: 4.48462, -LL: 5.00446, prior: 0.34372, aleatoric unc.: 15.79595\n",
      "Epoch 258/500  total: 4.48687, -LL: 4.10884, prior: 0.34404, aleatoric unc.: 15.74077\n",
      "Epoch 259/500  total: 4.47548, -LL: 4.94737, prior: 0.34176, aleatoric unc.: 15.66432\n",
      "Epoch 260/500  total: 4.45308, -LL: 4.01868, prior: 0.34024, aleatoric unc.: 15.58120\n",
      "Epoch 261/500  total: 4.48031, -LL: 3.95064, prior: 0.33870, aleatoric unc.: 15.53961\n",
      "Epoch 262/500  total: 4.50446, -LL: 3.82153, prior: 0.33506, aleatoric unc.: 15.56419\n",
      "Epoch 263/500  total: 4.46654, -LL: 4.43109, prior: 0.33716, aleatoric unc.: 15.50787\n",
      "Epoch 264/500  total: 4.50047, -LL: 4.56283, prior: 0.33743, aleatoric unc.: 15.50358\n",
      "Epoch 265/500  total: 4.53483, -LL: 4.92259, prior: 0.33880, aleatoric unc.: 15.54195\n",
      "Epoch 266/500  total: 4.48853, -LL: 4.06707, prior: 0.33707, aleatoric unc.: 15.52835\n",
      "Epoch 267/500  total: 4.48643, -LL: 4.31537, prior: 0.33551, aleatoric unc.: 15.52833\n",
      "Epoch 268/500  total: 4.43746, -LL: 3.74320, prior: 0.33580, aleatoric unc.: 15.44481\n",
      "Epoch 269/500  total: 4.45904, -LL: 4.12249, prior: 0.33215, aleatoric unc.: 15.37424\n",
      "Epoch 270/500  total: 4.43447, -LL: 4.48685, prior: 0.32961, aleatoric unc.: 15.30517\n",
      "Epoch 271/500  total: 4.44404, -LL: 4.04688, prior: 0.32919, aleatoric unc.: 15.24123\n",
      "Epoch 272/500  total: 4.47093, -LL: 4.14307, prior: 0.33486, aleatoric unc.: 15.25717\n",
      "Epoch 273/500  total: 4.38182, -LL: 4.15838, prior: 0.33034, aleatoric unc.: 15.08905\n",
      "Epoch 274/500  total: 4.47173, -LL: 3.82219, prior: 0.32999, aleatoric unc.: 15.08135\n",
      "Epoch 275/500  total: 4.36634, -LL: 4.26105, prior: 0.32649, aleatoric unc.: 14.94182\n",
      "Epoch 276/500  total: 4.52364, -LL: 3.95488, prior: 0.32549, aleatoric unc.: 15.03032\n",
      "Epoch 277/500  total: 4.40352, -LL: 4.13583, prior: 0.32405, aleatoric unc.: 14.97387\n",
      "Epoch 278/500  total: 4.40713, -LL: 4.51679, prior: 0.32256, aleatoric unc.: 14.89801\n",
      "Epoch 279/500  total: 4.48028, -LL: 4.06760, prior: 0.32299, aleatoric unc.: 14.96830\n",
      "Epoch 280/500  total: 4.52136, -LL: 5.16841, prior: 0.32424, aleatoric unc.: 15.07113\n",
      "Epoch 281/500  total: 4.37791, -LL: 4.02710, prior: 0.32618, aleatoric unc.: 14.97818\n",
      "Epoch 282/500  total: 4.47312, -LL: 3.96956, prior: 0.32588, aleatoric unc.: 14.99776\n",
      "Epoch 283/500  total: 4.37466, -LL: 3.90545, prior: 0.32387, aleatoric unc.: 14.89272\n",
      "Epoch 284/500  total: 4.45828, -LL: 3.83861, prior: 0.32713, aleatoric unc.: 14.90926\n",
      "Epoch 285/500  total: 4.47821, -LL: 3.77745, prior: 0.32877, aleatoric unc.: 14.95312\n",
      "Epoch 286/500  total: 4.38676, -LL: 3.77643, prior: 0.32696, aleatoric unc.: 14.86059\n",
      "Epoch 287/500  total: 4.41101, -LL: 3.87555, prior: 0.32184, aleatoric unc.: 14.80666\n",
      "Epoch 288/500  total: 4.38039, -LL: 3.79602, prior: 0.31874, aleatoric unc.: 14.73124\n",
      "Epoch 289/500  total: 4.40861, -LL: 3.94453, prior: 0.31796, aleatoric unc.: 14.68562\n",
      "Epoch 290/500  total: 4.42048, -LL: 4.85468, prior: 0.31717, aleatoric unc.: 14.66400\n",
      "Epoch 291/500  total: 4.47149, -LL: 4.16605, prior: 0.32087, aleatoric unc.: 14.74144\n",
      "Epoch 292/500  total: 4.38311, -LL: 4.23254, prior: 0.32335, aleatoric unc.: 14.68072\n",
      "Epoch 293/500  total: 4.43105, -LL: 6.23792, prior: 0.32213, aleatoric unc.: 14.64300\n",
      "Epoch 294/500  total: 4.39989, -LL: 3.77106, prior: 0.32154, aleatoric unc.: 14.64436\n",
      "Epoch 295/500  total: 4.45488, -LL: 4.01945, prior: 0.31872, aleatoric unc.: 14.69023\n",
      "Epoch 296/500  total: 4.38851, -LL: 3.99385, prior: 0.31663, aleatoric unc.: 14.62358\n",
      "Epoch 297/500  total: 4.42410, -LL: 4.41566, prior: 0.31540, aleatoric unc.: 14.64059\n",
      "Epoch 298/500  total: 4.44718, -LL: 4.26388, prior: 0.31464, aleatoric unc.: 14.68472\n",
      "Epoch 299/500  total: 4.38908, -LL: 3.72025, prior: 0.31473, aleatoric unc.: 14.65539\n",
      "Epoch 300/500  total: 4.46657, -LL: 4.89952, prior: 0.31657, aleatoric unc.: 14.68762\n",
      "Epoch 301/500  total: 4.42263, -LL: 4.18494, prior: 0.31660, aleatoric unc.: 14.71025\n",
      "Epoch 302/500  total: 4.42304, -LL: 4.25995, prior: 0.31648, aleatoric unc.: 14.71077\n",
      "Epoch 303/500  total: 4.37865, -LL: 3.74415, prior: 0.31368, aleatoric unc.: 14.63672\n",
      "Epoch 304/500  total: 4.38759, -LL: 4.06919, prior: 0.31460, aleatoric unc.: 14.57741\n",
      "Epoch 305/500  total: 4.38715, -LL: 3.76870, prior: 0.31255, aleatoric unc.: 14.57457\n",
      "Epoch 306/500  total: 4.39209, -LL: 3.94298, prior: 0.31312, aleatoric unc.: 14.54549\n",
      "Epoch 307/500  total: 4.31401, -LL: 3.90953, prior: 0.30980, aleatoric unc.: 14.37811\n",
      "Epoch 308/500  total: 4.36241, -LL: 4.02659, prior: 0.31063, aleatoric unc.: 14.32905\n",
      "Epoch 309/500  total: 4.36258, -LL: 4.23499, prior: 0.31151, aleatoric unc.: 14.27872\n",
      "Epoch 310/500  total: 4.32726, -LL: 3.76714, prior: 0.30940, aleatoric unc.: 14.19101\n",
      "Epoch 311/500  total: 4.38596, -LL: 3.81679, prior: 0.30674, aleatoric unc.: 14.17865\n",
      "Epoch 312/500  total: 4.42068, -LL: 3.75325, prior: 0.31111, aleatoric unc.: 14.24413\n",
      "Epoch 313/500  total: 4.37524, -LL: 3.84334, prior: 0.30775, aleatoric unc.: 14.23629\n",
      "Epoch 314/500  total: 4.27371, -LL: 3.87354, prior: 0.30729, aleatoric unc.: 14.06987\n",
      "Epoch 315/500  total: 4.39870, -LL: 3.89449, prior: 0.30772, aleatoric unc.: 14.08975\n",
      "Epoch 316/500  total: 4.30602, -LL: 3.67755, prior: 0.30566, aleatoric unc.: 14.00856\n",
      "Epoch 317/500  total: 4.33313, -LL: 3.97410, prior: 0.30577, aleatoric unc.: 13.94800\n",
      "Epoch 318/500  total: 4.43432, -LL: 3.71468, prior: 0.30888, aleatoric unc.: 14.04741\n",
      "Epoch 319/500  total: 4.37897, -LL: 3.75499, prior: 0.31079, aleatoric unc.: 14.08446\n",
      "Epoch 320/500  total: 4.30563, -LL: 3.65509, prior: 0.30797, aleatoric unc.: 13.97607\n",
      "Epoch 321/500  total: 4.37533, -LL: 4.24918, prior: 0.30791, aleatoric unc.: 13.98534\n",
      "Epoch 322/500  total: 4.36393, -LL: 3.74137, prior: 0.30698, aleatoric unc.: 13.98914\n",
      "Epoch 323/500  total: 4.35164, -LL: 4.10119, prior: 0.30806, aleatoric unc.: 13.96811\n",
      "Epoch 324/500  total: 4.35739, -LL: 3.90398, prior: 0.30846, aleatoric unc.: 13.93596\n",
      "Epoch 325/500  total: 4.38250, -LL: 4.44499, prior: 0.30673, aleatoric unc.: 13.97261\n",
      "Epoch 326/500  total: 4.35315, -LL: 4.42376, prior: 0.30701, aleatoric unc.: 13.95997\n",
      "Epoch 327/500  total: 4.36795, -LL: 4.17653, prior: 0.30855, aleatoric unc.: 13.97699\n",
      "Epoch 328/500  total: 4.30396, -LL: 3.91943, prior: 0.30560, aleatoric unc.: 13.88542\n",
      "Epoch 329/500  total: 4.36057, -LL: 4.99554, prior: 0.30607, aleatoric unc.: 13.87517\n",
      "Epoch 330/500  total: 4.25497, -LL: 3.72196, prior: 0.30206, aleatoric unc.: 13.75286\n",
      "Epoch 331/500  total: 4.32408, -LL: 3.83675, prior: 0.30096, aleatoric unc.: 13.70025\n",
      "Epoch 332/500  total: 4.41900, -LL: 4.16212, prior: 0.30270, aleatoric unc.: 13.81388\n",
      "Epoch 333/500  total: 4.27495, -LL: 3.77374, prior: 0.30219, aleatoric unc.: 13.72647\n",
      "Epoch 334/500  total: 4.34971, -LL: 3.86083, prior: 0.30085, aleatoric unc.: 13.73270\n",
      "Epoch 335/500  total: 4.37895, -LL: 4.20601, prior: 0.30247, aleatoric unc.: 13.77878\n",
      "Epoch 336/500  total: 4.36633, -LL: 3.74334, prior: 0.30218, aleatoric unc.: 13.81359\n",
      "Epoch 337/500  total: 4.31169, -LL: 3.69489, prior: 0.30070, aleatoric unc.: 13.75855\n",
      "Epoch 338/500  total: 4.29754, -LL: 3.98524, prior: 0.30092, aleatoric unc.: 13.69109\n",
      "Epoch 339/500  total: 4.35188, -LL: 3.76089, prior: 0.29924, aleatoric unc.: 13.71738\n",
      "Epoch 340/500  total: 4.36790, -LL: 6.01456, prior: 0.29973, aleatoric unc.: 13.73557\n",
      "Epoch 341/500  total: 4.32057, -LL: 3.93740, prior: 0.30091, aleatoric unc.: 13.73691\n",
      "Epoch 342/500  total: 4.31353, -LL: 4.73162, prior: 0.29983, aleatoric unc.: 13.69596\n",
      "Epoch 343/500  total: 4.29736, -LL: 3.70710, prior: 0.30202, aleatoric unc.: 13.64569\n",
      "Epoch 344/500  total: 4.30277, -LL: 4.39192, prior: 0.29873, aleatoric unc.: 13.58624\n",
      "Epoch 345/500  total: 4.28730, -LL: 3.70689, prior: 0.30125, aleatoric unc.: 13.53884\n",
      "Epoch 346/500  total: 4.28112, -LL: 3.88174, prior: 0.29967, aleatoric unc.: 13.46685\n",
      "Epoch 347/500  total: 4.29716, -LL: 3.81197, prior: 0.29814, aleatoric unc.: 13.42146\n",
      "Epoch 348/500  total: 4.33209, -LL: 3.68184, prior: 0.29644, aleatoric unc.: 13.45658\n",
      "Epoch 349/500  total: 4.25540, -LL: 4.01308, prior: 0.29434, aleatoric unc.: 13.34073\n",
      "Epoch 350/500  total: 4.33118, -LL: 3.82876, prior: 0.29366, aleatoric unc.: 13.41339\n",
      "Epoch 351/500  total: 4.35809, -LL: 3.98745, prior: 0.29452, aleatoric unc.: 13.47417\n",
      "Epoch 352/500  total: 4.36339, -LL: 4.11955, prior: 0.29371, aleatoric unc.: 13.56014\n",
      "Epoch 353/500  total: 4.27057, -LL: 4.07453, prior: 0.29478, aleatoric unc.: 13.48736\n",
      "Epoch 354/500  total: 4.35561, -LL: 4.00709, prior: 0.29816, aleatoric unc.: 13.54832\n",
      "Epoch 355/500  total: 4.30775, -LL: 3.76553, prior: 0.29959, aleatoric unc.: 13.53639\n",
      "Epoch 356/500  total: 4.30931, -LL: 3.85962, prior: 0.29742, aleatoric unc.: 13.51869\n",
      "Epoch 357/500  total: 4.31901, -LL: 4.00900, prior: 0.29731, aleatoric unc.: 13.50655\n",
      "Epoch 358/500  total: 4.30852, -LL: 3.72225, prior: 0.29563, aleatoric unc.: 13.47658\n",
      "Epoch 359/500  total: 4.28606, -LL: 3.78456, prior: 0.29504, aleatoric unc.: 13.45554\n",
      "Epoch 360/500  total: 4.23826, -LL: 3.60652, prior: 0.29418, aleatoric unc.: 13.32370\n",
      "Epoch 361/500  total: 4.33118, -LL: 4.33109, prior: 0.29590, aleatoric unc.: 13.36398\n",
      "Epoch 362/500  total: 4.30136, -LL: 3.88332, prior: 0.29581, aleatoric unc.: 13.35263\n",
      "Epoch 363/500  total: 4.26135, -LL: 3.99345, prior: 0.29653, aleatoric unc.: 13.29258\n",
      "Epoch 364/500  total: 4.28480, -LL: 4.00098, prior: 0.29697, aleatoric unc.: 13.25982\n",
      "Epoch 365/500  total: 4.29093, -LL: 3.72887, prior: 0.29551, aleatoric unc.: 13.25455\n",
      "Epoch 366/500  total: 4.31470, -LL: 3.67486, prior: 0.29167, aleatoric unc.: 13.27328\n",
      "Epoch 367/500  total: 4.28997, -LL: 3.87706, prior: 0.29307, aleatoric unc.: 13.25669\n",
      "Epoch 368/500  total: 4.25890, -LL: 3.61237, prior: 0.29420, aleatoric unc.: 13.21329\n",
      "Epoch 369/500  total: 4.28513, -LL: 3.66183, prior: 0.29501, aleatoric unc.: 13.19332\n",
      "Epoch 370/500  total: 4.24566, -LL: 3.59994, prior: 0.29461, aleatoric unc.: 13.11074\n",
      "Epoch 371/500  total: 4.28892, -LL: 3.62588, prior: 0.29155, aleatoric unc.: 13.12838\n",
      "Epoch 372/500  total: 4.25191, -LL: 3.96459, prior: 0.29308, aleatoric unc.: 13.06497\n",
      "Epoch 373/500  total: 4.27148, -LL: 3.65914, prior: 0.29409, aleatoric unc.: 13.05644\n",
      "Epoch 374/500  total: 4.22671, -LL: 3.59360, prior: 0.29335, aleatoric unc.: 12.96356\n",
      "Epoch 375/500  total: 4.25709, -LL: 3.75483, prior: 0.29473, aleatoric unc.: 12.92731\n",
      "Epoch 376/500  total: 4.30109, -LL: 3.96422, prior: 0.29433, aleatoric unc.: 12.96784\n",
      "Epoch 377/500  total: 4.27804, -LL: 4.62750, prior: 0.29492, aleatoric unc.: 12.96339\n",
      "Epoch 378/500  total: 4.30291, -LL: 3.66387, prior: 0.29426, aleatoric unc.: 13.02547\n",
      "Epoch 379/500  total: 4.28989, -LL: 3.85442, prior: 0.29308, aleatoric unc.: 13.02152\n",
      "Epoch 380/500  total: 4.28133, -LL: 3.64806, prior: 0.29297, aleatoric unc.: 13.04834\n",
      "Epoch 381/500  total: 4.27531, -LL: 3.91687, prior: 0.29161, aleatoric unc.: 13.04516\n",
      "Epoch 382/500  total: 4.22670, -LL: 3.65751, prior: 0.29002, aleatoric unc.: 12.96334\n",
      "Epoch 383/500  total: 4.25438, -LL: 3.83289, prior: 0.29040, aleatoric unc.: 12.92650\n",
      "Epoch 384/500  total: 4.28284, -LL: 4.12477, prior: 0.29173, aleatoric unc.: 12.93924\n",
      "Epoch 385/500  total: 4.28154, -LL: 5.06327, prior: 0.29164, aleatoric unc.: 12.94179\n",
      "Epoch 386/500  total: 4.29184, -LL: 3.71668, prior: 0.28974, aleatoric unc.: 12.99222\n",
      "Epoch 387/500  total: 4.23449, -LL: 3.61562, prior: 0.29012, aleatoric unc.: 12.93269\n",
      "Epoch 388/500  total: 4.20235, -LL: 3.64429, prior: 0.28787, aleatoric unc.: 12.83775\n",
      "Epoch 389/500  total: 4.23469, -LL: 3.65001, prior: 0.28924, aleatoric unc.: 12.78895\n",
      "Epoch 390/500  total: 4.25743, -LL: 3.70225, prior: 0.28938, aleatoric unc.: 12.79457\n",
      "Epoch 391/500  total: 4.28773, -LL: 3.82099, prior: 0.29032, aleatoric unc.: 12.83519\n",
      "Epoch 392/500  total: 4.24977, -LL: 3.72148, prior: 0.28917, aleatoric unc.: 12.82497\n",
      "Epoch 393/500  total: 4.23988, -LL: 3.75097, prior: 0.29060, aleatoric unc.: 12.79168\n",
      "Epoch 394/500  total: 4.23515, -LL: 4.51935, prior: 0.29009, aleatoric unc.: 12.74808\n",
      "Epoch 395/500  total: 4.25556, -LL: 3.92254, prior: 0.28824, aleatoric unc.: 12.74187\n",
      "Epoch 396/500  total: 4.24138, -LL: 4.08155, prior: 0.28770, aleatoric unc.: 12.73601\n",
      "Epoch 397/500  total: 4.32805, -LL: 3.77373, prior: 0.28914, aleatoric unc.: 12.85976\n",
      "Epoch 398/500  total: 4.29116, -LL: 4.03029, prior: 0.28854, aleatoric unc.: 12.89174\n",
      "Epoch 399/500  total: 4.29906, -LL: 4.08009, prior: 0.29080, aleatoric unc.: 12.94973\n",
      "Epoch 400/500  total: 4.24325, -LL: 4.37855, prior: 0.29110, aleatoric unc.: 12.89844\n",
      "Epoch 401/500  total: 4.21631, -LL: 4.37101, prior: 0.29081, aleatoric unc.: 12.83269\n",
      "Epoch 402/500  total: 4.21490, -LL: 3.70825, prior: 0.28925, aleatoric unc.: 12.77170\n",
      "Epoch 403/500  total: 4.24548, -LL: 3.97789, prior: 0.28979, aleatoric unc.: 12.73532\n",
      "Epoch 404/500  total: 4.25498, -LL: 3.91230, prior: 0.28735, aleatoric unc.: 12.75526\n",
      "Epoch 405/500  total: 4.23092, -LL: 3.64123, prior: 0.28916, aleatoric unc.: 12.73843\n",
      "Epoch 406/500  total: 4.30377, -LL: 3.91702, prior: 0.29196, aleatoric unc.: 12.76713\n",
      "Epoch 407/500  total: 4.25391, -LL: 3.67363, prior: 0.29132, aleatoric unc.: 12.80679\n",
      "Epoch 408/500  total: 4.22299, -LL: 3.62789, prior: 0.28925, aleatoric unc.: 12.74367\n",
      "Epoch 409/500  total: 4.26734, -LL: 3.80965, prior: 0.29048, aleatoric unc.: 12.76100\n",
      "Epoch 410/500  total: 4.25349, -LL: 3.69565, prior: 0.29013, aleatoric unc.: 12.76318\n",
      "Epoch 411/500  total: 4.28336, -LL: 3.69440, prior: 0.29020, aleatoric unc.: 12.79956\n",
      "Epoch 412/500  total: 4.29031, -LL: 4.51302, prior: 0.29134, aleatoric unc.: 12.83943\n",
      "Epoch 413/500  total: 4.19793, -LL: 3.90017, prior: 0.29209, aleatoric unc.: 12.74592\n",
      "Epoch 414/500  total: 4.19218, -LL: 3.57106, prior: 0.29010, aleatoric unc.: 12.65540\n",
      "Epoch 415/500  total: 4.25412, -LL: 3.74360, prior: 0.29186, aleatoric unc.: 12.65665\n",
      "Epoch 416/500  total: 4.26500, -LL: 3.91594, prior: 0.29214, aleatoric unc.: 12.68659\n",
      "Epoch 417/500  total: 4.21341, -LL: 3.72687, prior: 0.29015, aleatoric unc.: 12.63162\n",
      "Epoch 418/500  total: 4.21933, -LL: 3.62715, prior: 0.29108, aleatoric unc.: 12.57077\n",
      "Epoch 419/500  total: 4.25920, -LL: 3.62662, prior: 0.28931, aleatoric unc.: 12.61524\n",
      "Epoch 420/500  total: 4.24498, -LL: 3.89066, prior: 0.29000, aleatoric unc.: 12.60225\n",
      "Epoch 421/500  total: 4.26176, -LL: 4.10444, prior: 0.28994, aleatoric unc.: 12.64584\n",
      "Epoch 422/500  total: 4.19053, -LL: 3.76841, prior: 0.29019, aleatoric unc.: 12.56762\n",
      "Epoch 423/500  total: 4.20613, -LL: 4.07726, prior: 0.28902, aleatoric unc.: 12.51632\n",
      "Epoch 424/500  total: 4.20716, -LL: 3.57636, prior: 0.28839, aleatoric unc.: 12.48092\n",
      "Epoch 425/500  total: 4.18982, -LL: 3.90812, prior: 0.28768, aleatoric unc.: 12.40683\n",
      "Epoch 426/500  total: 4.20946, -LL: 4.42304, prior: 0.28793, aleatoric unc.: 12.37388\n",
      "Epoch 427/500  total: 4.22372, -LL: 4.08125, prior: 0.28901, aleatoric unc.: 12.37556\n",
      "Epoch 428/500  total: 4.23337, -LL: 4.16943, prior: 0.28858, aleatoric unc.: 12.40270\n",
      "Epoch 429/500  total: 4.28320, -LL: 3.92145, prior: 0.29058, aleatoric unc.: 12.47761\n",
      "Epoch 430/500  total: 4.23783, -LL: 3.58313, prior: 0.28923, aleatoric unc.: 12.49446\n",
      "Epoch 431/500  total: 4.20762, -LL: 3.79757, prior: 0.28842, aleatoric unc.: 12.43921\n",
      "Epoch 432/500  total: 4.26053, -LL: 3.61355, prior: 0.29010, aleatoric unc.: 12.51055\n",
      "Epoch 433/500  total: 4.20546, -LL: 3.60644, prior: 0.28897, aleatoric unc.: 12.47210\n",
      "Epoch 434/500  total: 4.17361, -LL: 3.62394, prior: 0.28800, aleatoric unc.: 12.38231\n",
      "Epoch 435/500  total: 4.18867, -LL: 3.65719, prior: 0.28752, aleatoric unc.: 12.32489\n",
      "Epoch 436/500  total: 4.22270, -LL: 4.32778, prior: 0.28681, aleatoric unc.: 12.32984\n",
      "Epoch 437/500  total: 4.16955, -LL: 3.78783, prior: 0.28606, aleatoric unc.: 12.25807\n",
      "Epoch 438/500  total: 4.22591, -LL: 4.07864, prior: 0.28614, aleatoric unc.: 12.28022\n",
      "Epoch 439/500  total: 4.19177, -LL: 4.09317, prior: 0.28660, aleatoric unc.: 12.25075\n",
      "Epoch 440/500  total: 4.27602, -LL: 3.54679, prior: 0.28683, aleatoric unc.: 12.33664\n",
      "Epoch 441/500  total: 4.18954, -LL: 3.72110, prior: 0.28622, aleatoric unc.: 12.31076\n",
      "Epoch 442/500  total: 4.25897, -LL: 3.84474, prior: 0.28746, aleatoric unc.: 12.37183\n",
      "Epoch 443/500  total: 4.19797, -LL: 3.61508, prior: 0.28697, aleatoric unc.: 12.33025\n",
      "Epoch 444/500  total: 4.19945, -LL: 4.23817, prior: 0.28682, aleatoric unc.: 12.30664\n",
      "Epoch 445/500  total: 4.25961, -LL: 3.78511, prior: 0.28722, aleatoric unc.: 12.37745\n",
      "Epoch 446/500  total: 4.25115, -LL: 3.63889, prior: 0.28812, aleatoric unc.: 12.41472\n",
      "Epoch 447/500  total: 4.18577, -LL: 3.74992, prior: 0.28674, aleatoric unc.: 12.35474\n",
      "Epoch 448/500  total: 4.22803, -LL: 3.74013, prior: 0.28695, aleatoric unc.: 12.35672\n",
      "Epoch 449/500  total: 4.24271, -LL: 4.02800, prior: 0.28769, aleatoric unc.: 12.39347\n",
      "Epoch 450/500  total: 4.22130, -LL: 3.77517, prior: 0.28920, aleatoric unc.: 12.40454\n",
      "Epoch 451/500  total: 4.19642, -LL: 4.85910, prior: 0.28584, aleatoric unc.: 12.35776\n",
      "Epoch 452/500  total: 4.17571, -LL: 3.55679, prior: 0.28438, aleatoric unc.: 12.30047\n",
      "Epoch 453/500  total: 4.20105, -LL: 4.47895, prior: 0.28587, aleatoric unc.: 12.26487\n",
      "Epoch 454/500  total: 4.18347, -LL: 3.77700, prior: 0.28563, aleatoric unc.: 12.22977\n",
      "Epoch 455/500  total: 4.23386, -LL: 3.76747, prior: 0.28721, aleatoric unc.: 12.26009\n",
      "Epoch 456/500  total: 4.24766, -LL: 3.71218, prior: 0.28666, aleatoric unc.: 12.31672\n",
      "Epoch 457/500  total: 4.23320, -LL: 3.63747, prior: 0.28566, aleatoric unc.: 12.36460\n",
      "Epoch 458/500  total: 4.21914, -LL: 4.22862, prior: 0.28573, aleatoric unc.: 12.33850\n",
      "Epoch 459/500  total: 4.23816, -LL: 3.77558, prior: 0.28672, aleatoric unc.: 12.39121\n",
      "Epoch 460/500  total: 4.18095, -LL: 3.59544, prior: 0.28580, aleatoric unc.: 12.33295\n",
      "Epoch 461/500  total: 4.19751, -LL: 3.65706, prior: 0.28622, aleatoric unc.: 12.29170\n",
      "Epoch 462/500  total: 4.17444, -LL: 3.58506, prior: 0.28652, aleatoric unc.: 12.24185\n",
      "Epoch 463/500  total: 4.21997, -LL: 3.79871, prior: 0.28532, aleatoric unc.: 12.25640\n",
      "Epoch 464/500  total: 4.17849, -LL: 3.67556, prior: 0.28590, aleatoric unc.: 12.20991\n",
      "Epoch 465/500  total: 4.19469, -LL: 3.65216, prior: 0.28387, aleatoric unc.: 12.18799\n",
      "Epoch 466/500  total: 4.18374, -LL: 3.77503, prior: 0.28464, aleatoric unc.: 12.13527\n",
      "Epoch 467/500  total: 4.15980, -LL: 3.54591, prior: 0.28497, aleatoric unc.: 12.09160\n",
      "Epoch 468/500  total: 4.15528, -LL: 3.64045, prior: 0.28466, aleatoric unc.: 12.03012\n",
      "Epoch 469/500  total: 4.20050, -LL: 3.53782, prior: 0.28442, aleatoric unc.: 12.04823\n",
      "Epoch 470/500  total: 4.18814, -LL: 4.13648, prior: 0.28479, aleatoric unc.: 12.02533\n",
      "Epoch 471/500  total: 4.23799, -LL: 4.06228, prior: 0.28585, aleatoric unc.: 12.09650\n",
      "Epoch 472/500  total: 4.21309, -LL: 3.93179, prior: 0.28526, aleatoric unc.: 12.13297\n",
      "Epoch 473/500  total: 4.22759, -LL: 4.34791, prior: 0.28510, aleatoric unc.: 12.16390\n",
      "Epoch 474/500  total: 4.19927, -LL: 3.58355, prior: 0.28564, aleatoric unc.: 12.18269\n",
      "Epoch 475/500  total: 4.20806, -LL: 3.66223, prior: 0.28625, aleatoric unc.: 12.19320\n",
      "Epoch 476/500  total: 4.18165, -LL: 3.72681, prior: 0.28742, aleatoric unc.: 12.15077\n",
      "Epoch 477/500  total: 4.18182, -LL: 3.96265, prior: 0.28656, aleatoric unc.: 12.11517\n",
      "Epoch 478/500  total: 4.20171, -LL: 3.59803, prior: 0.28639, aleatoric unc.: 12.10981\n",
      "Epoch 479/500  total: 4.20908, -LL: 3.62256, prior: 0.28814, aleatoric unc.: 12.13264\n",
      "Epoch 480/500  total: 4.20871, -LL: 3.70293, prior: 0.28707, aleatoric unc.: 12.14405\n",
      "Epoch 481/500  total: 4.16861, -LL: 4.15197, prior: 0.28590, aleatoric unc.: 12.09377\n",
      "Epoch 482/500  total: 4.19767, -LL: 5.47841, prior: 0.28523, aleatoric unc.: 12.05883\n",
      "Epoch 483/500  total: 4.25352, -LL: 4.24806, prior: 0.28549, aleatoric unc.: 12.17083\n",
      "Epoch 484/500  total: 4.19220, -LL: 3.62151, prior: 0.28546, aleatoric unc.: 12.16249\n",
      "Epoch 485/500  total: 4.20604, -LL: 4.07630, prior: 0.28566, aleatoric unc.: 12.16024\n",
      "Epoch 486/500  total: 4.21192, -LL: 3.81535, prior: 0.28562, aleatoric unc.: 12.18320\n",
      "Epoch 487/500  total: 4.22036, -LL: 3.71840, prior: 0.28568, aleatoric unc.: 12.21167\n",
      "Epoch 488/500  total: 4.17873, -LL: 3.67013, prior: 0.28418, aleatoric unc.: 12.17093\n",
      "Epoch 489/500  total: 4.17590, -LL: 3.84937, prior: 0.28326, aleatoric unc.: 12.13036\n",
      "Epoch 490/500  total: 4.19966, -LL: 3.53921, prior: 0.28432, aleatoric unc.: 12.13181\n",
      "Epoch 491/500  total: 4.15192, -LL: 3.69806, prior: 0.28322, aleatoric unc.: 12.05840\n",
      "Epoch 492/500  total: 4.20871, -LL: 3.99247, prior: 0.28498, aleatoric unc.: 12.09097\n",
      "Epoch 493/500  total: 4.24564, -LL: 3.53812, prior: 0.28562, aleatoric unc.: 12.15556\n",
      "Epoch 494/500  total: 4.21448, -LL: 3.74054, prior: 0.28474, aleatoric unc.: 12.16671\n",
      "Epoch 495/500  total: 4.16589, -LL: 3.60851, prior: 0.28430, aleatoric unc.: 12.13879\n",
      "Epoch 496/500  total: 4.19407, -LL: 3.67642, prior: 0.28451, aleatoric unc.: 12.12040\n",
      "Epoch 497/500  total: 4.20313, -LL: 4.11066, prior: 0.28443, aleatoric unc.: 12.11105\n",
      "Epoch 498/500  total: 4.17024, -LL: 3.79871, prior: 0.28372, aleatoric unc.: 12.08894\n",
      "Epoch 499/500  total: 4.20190, -LL: 3.76114, prior: 0.28452, aleatoric unc.: 12.09669\n"
     ]
    }
   ],
   "source": [
    "epochs = 500\n",
    "# for each epoch\n",
    "for epoch in range(epochs):\n",
    "    losses = list()\n",
    "    # for each mini-batch given by the loader:\n",
    "    for batch in loader:\n",
    "        # get the input in the mini-batch\n",
    "        # this has size (B, C)\n",
    "        # where B is the mini-batch size\n",
    "        # C is the number of features (1 in this case)\n",
    "        features = batch[\"data\"]\n",
    "        # get the targets in the mini-batch (there shall be B of them)\n",
    "        target = batch[\"target\"]\n",
    "        # get the output of the neural network:\n",
    "        prediction = b_network(features)\n",
    "        \n",
    "        # calculate the loss function being minimized\n",
    "        # in this case, it is the mean-squared error between the prediction and the target values added\n",
    "        # to the Kullback-Leibler divergence between the current weight Gaussian and\n",
    "        # the prior Gaussian, set to the unit Normal distribution\n",
    "        nll = b_network.nll(prediction, target)\n",
    "        prior = weight_kl *kl_loss(b_network)\n",
    "        loss = nll + prior\n",
    "\n",
    "        # clean the optimizer temporary gradient storage\n",
    "        optimizer.zero_grad()\n",
    "        # calculate the gradient of the loss function as a function of the gradients\n",
    "        loss.backward()\n",
    "        # ask the Adam optimizer to change the parameters in the direction of - gradient\n",
    "        # Adam scales the gradient by a constant which is adaptively tuned\n",
    "        # take a look at the Adam paper for more details: https://arxiv.org/abs/1412.6980\n",
    "        optimizer.step()\n",
    "        \n",
    "        ale = b_network.aleatoric_uncertainty().detach().numpy()\n",
    "\n",
    "        losses.append(loss.detach().cpu().item())\n",
    "    avg_loss = np.mean(np.array(losses))\n",
    "    print(f\"Epoch {epoch}/{epochs}  total: {avg_loss:.5f}, -LL: {nll.item():.5f}, prior: {prior.item():.5f}, aleatoric unc.: {ale:.5f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15f3f633",
   "metadata": {},
   "source": [
    "After the fit, each time we perform a prediction with the neural network, the weights of the neural network are different: they are randomly chosen from a Gaussian distribution with the fit mean and standard deviation. The logic here is that all these weights match the expected posterior probability optimized, so they should all provide trustworthy results given the data available (but not equally probable!).\n",
    "\n",
    "In order to obtain an uncertainty band for each prediction, we perform each prediction with the neural network fit several times to obtain the results with different set of weights following the posterior probability mentioned. After taking several results for the same input data, we can use the mean prediction as the \"nominal\" prediction and the root-mean-squared-error of the predictions as the uncertainty of the prediction, given the available data.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a3d244e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# let's generate new data!\n",
    "test_data = generate_data(N=1000)\n",
    "\n",
    "b_predicted = list()\n",
    "for k in range(10):\n",
    "    p = b_network(torch.from_numpy(test_data[:,0:1])).detach().numpy()\n",
    "    b_predicted.append(p[:,0])\n",
    "b_predicted = np.stack(b_predicted, axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85fed034",
   "metadata": {},
   "source": [
    "We can now take the average result for each sample, and their root-mean-squared-error as an estimate of the mean and epistemic uncertainty for the results.\n",
    "\n",
    "The aleatoric uncertainty is fitted as an independent parameter. Since we assume the aleatoric uncertainty is independent, we can calculate the total uncertainty as the sum of squares of the epistemic and aleatoric uncertainty."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4a56960d",
   "metadata": {},
   "outputs": [],
   "source": [
    "b_mean = np.mean(b_predicted, axis=1)\n",
    "b_sigma = np.std(b_predicted, axis=1)\n",
    "aleatoric_uncertainty = b_network.aleatoric_uncertainty().detach().numpy()\n",
    "\n",
    "total_uncertainty = (b_sigma**2 + aleatoric_uncertainty**2)**0.5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ef3d430",
   "metadata": {},
   "source": [
    "Let's check how big are those uncertainties found:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "4d01c41f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average epistemic uncertainty:  4.342738\n"
     ]
    }
   ],
   "source": [
    "print(\"Average epistemic uncertainty: \", np.mean(b_sigma))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "67d456a1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Aleatoric uncertainty:  12.096693\n"
     ]
    }
   ],
   "source": [
    "print(\"Aleatoric uncertainty: \", aleatoric_uncertainty)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "364bdcd7",
   "metadata": {},
   "source": [
    "Note that the aleatoric uncertainty is very close to the standard deviation of the $\\epsilon$ component of the model we created in the beginning! Clearly the model could fit the uncertainty coming from that component of the noise.\n",
    "\n",
    "It is not easy to estimate the effect of the epistemic uncertainty, as it is different for every data point (as it is scaled by $x^2$), but we can plot it to take a look at its effect.\n",
    "\n",
    "Note that the uncertainties are the standard deviations of Gaussian models and therefore they correspond to a $1\\sigma$ quantile band, which is a 67% confidence band. The quantile corresponding to $2\\sigma$ corresponds to a 95% confidence band in a Gaussian model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "8b9142e8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 60988 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111)\n",
    "ax.scatter(test_data[:, 0], test_data[:, 1], alpha=0.5, label=\"Test data\")\n",
    "ax.errorbar(test_data[:, 0], b_mean, yerr=2*total_uncertainty, alpha=0.5, fmt='or', label=\"95% band total unc.\")\n",
    "ax.errorbar(test_data[:, 0], b_mean, yerr=2*b_sigma, alpha=0.5, fmt='og', label=\"95% band epistemic unc.\")\n",
    "ax.set(xlabel=\"$x$\", ylabel=\"$f(x)$\")\n",
    "#ax.set_yscale('log')\n",
    "plt.legend(frameon=False)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55635d1f",
   "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
