00_PyData_Global_2021_nb_full.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import tensorflow_probability as tfp\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "plt.style.use('fivethirtyeight')\n",
    "\n",
    "tfd = tfp.distributions\n",
    "tfpl = tfp.layers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "TFP docs are [here](https://www.tensorflow.org/probability/api_docs/python/tfp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Modeling aleatoric and epistemic uncertainty using TensorFlow and TensorFlow Probability\n",
    "\n",
    "## PyData Global 2021\n",
    "\n",
    "Modeling uncertainty is incredibly useful when we want to understand how sure the model is about its own predictions. In the talk we’ll discuss differences between two types of uncertainty (aleatoric and epistemic) and show how to model them using TensorFlow and TensorFlow Probability. At the end of the talk, you’ll have practical understanding how to apply uncertainty modeling in your own project.\n",
    "\n",
    "\n",
    "Prerequisites:\n",
    "\n",
    "* Have good understanding of deep learning basics\n",
    "* Have good understanding of Bayes’ theorem\n",
    "* Have practical experience in a contemporary Python deep learning framework (TensorFlow recommended)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. TFP Components: `tfp.distributions`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's start with Normal(100, 15)\n",
    "normal = tfd.Normal(loc=100., scale=15.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "normal"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1 Methods -  `.sample()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Draw 2000 samples form `normal`\n",
    "samples = normal.sample(2000)\n",
    "samples.numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot a histogram of `samples`\n",
    "plt.figure(figsize=(6, 4))\n",
    "plt.hist(samples.numpy(), density=True, bins=np.sqrt(2000).astype(int))\n",
    "plt.title('A histogram of 2000 samples\\nfrom $N(100, 15)$')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 Methods -  `.prob()` & `.log_prob()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's compare PDF values of a couple of values given `normal`\n",
    "values = [70., 85., 100., 115., 130.]\n",
    "\n",
    "for v in values:\n",
    "    print(f'PDF of {v:05.1f} == {normal.prob(v):.4f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now, let's do the same for log_probs\n",
    "values = [70., 85., 100., 115., 130.]\n",
    "\n",
    "for v in values:\n",
    "    print(f'Log PDF {v:05.1f} == {normal.log_prob(v):.4f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Modeling aleatoric uncertainty - `tfp.layers`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 `tfpl.MultivariateNormal` layer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.1.1 Generate data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, let's generate some data! \n",
    "\n",
    "Our data generating process will be defined by the following formula:\n",
    "\n",
    "$$\\large y_i = 1.5x_i + 3 + 0.35\\epsilon$$\n",
    "\n",
    "where:\n",
    "\n",
    "* $\\epsilon \\sim N(0, 1)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate the data\n",
    "x_train = np.linspace(-1, 1, 200)[:, np.newaxis]\n",
    "y_train = 1.5*x_train + 3 + 0.35*np.random.randn(200)[:, np.newaxis]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the data\n",
    "plt.figure(figsize=(6, 4))\n",
    "plt.scatter(x_train, y_train, alpha=.5)\n",
    "plt.xlabel('$x$')\n",
    "plt.ylabel('$y$')\n",
    "plt.title('Training data')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.1.2 Simple model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's build a simple linear regression model\n",
    "model = tf.keras.Sequential([\n",
    "    tf.keras.layers.Dense(1, input_shape=(1,))\n",
    "])\n",
    "\n",
    "# Compile \n",
    "model.compile(loss='mse', optimizer='sgd')\n",
    "\n",
    "# Fit\n",
    "history = model.fit(x_train, y_train, epochs=100, verbose=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Predict\n",
    "y_pred = model(x_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the data and a trained regression line\n",
    "plt.figure(figsize=(6, 4))\n",
    "plt.scatter(x_train, y_train, alpha=.5)\n",
    "plt.plot(x_train, y_pred, label='Fitted reg. line', c='r')\n",
    "plt.xlabel('$x$')\n",
    "plt.ylabel('$y$')\n",
    "plt.title('Training data & fitted line')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.1.3 Model with probabilistic output - aleatoric uncertainty (learn mean & std)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's build a model with probabilistic output (mean & std) - a simple way\n",
    "event_shape = 1\n",
    "\n",
    "model = tf.keras.Sequential([\n",
    "    tf.keras.layers.Dense(2, input_shape=(1,)),\n",
    "    tfpl.IndependentNormal(event_shape=event_shape)\n",
    "])\n",
    "\n",
    "# Define neg. loglik. loss function\n",
    "def neg_loglik(y_true, y_pred):\n",
    "    return -y_pred.log_prob(y_true)\n",
    "\n",
    "# Compile \n",
    "model.compile(loss=neg_loglik, optimizer='sgd')\n",
    "\n",
    "# Fit\n",
    "history = model.fit(x_train, y_train, epochs=300, verbose=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's compare distributions - y_train & y_model\n",
    "y_model = model(x_train).sample().numpy()\n",
    "\n",
    "# Plot \n",
    "plt.figure(figsize=(6, 4))\n",
    "plt.scatter(x_train, y_train, alpha=.5, label='Training data')\n",
    "plt.scatter(x_train, y_model, alpha=.5, label='Learned dist. samples')\n",
    "plt.xlabel('$x$')\n",
    "plt.ylabel('$y$')\n",
    "plt.legend()\n",
    "plt.title('Training data vs. learned dist.')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute conf ints\n",
    "y_sample = model(x_train).sample()\n",
    "y_hat = model(x_train).mean()\n",
    "y_sd = model(x_train).stddev()\n",
    "y_hat_lower = y_hat - 2 * y_sd\n",
    "y_hat_upper = y_hat + 2 * y_sd\n",
    "\n",
    "# Plot conf ints\n",
    "plt.figure(figsize=(6, 4))\n",
    "plt.scatter(x_train, y_train, alpha=.5)\n",
    "plt.plot(x_train, y_hat, label='Fitted reg. line', c='r')\n",
    "plt.fill_between(np.squeeze(x_train), np.squeeze(y_hat_lower), np.squeeze(y_hat_upper), alpha=.1, label='$+/- 2SD$')\n",
    "plt.xlabel('$x$')\n",
    "plt.ylabel('$y$')\n",
    "plt.title('Training data & fitted line\\n+ 95% CIs')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Modeling epistemic uncertainty - `tfpl.DenseVariational`\n",
    "\n",
    "Now, we're going to model weight's uncertainty. \n",
    "\n",
    "This means, that each weight in our network will now be represented by a **distribution**, and **not** just a **point estimate**.\n",
    "\n",
    "<br><br>\n",
    "\n",
    "<img src=\"https://sanjaykthakur.files.wordpress.com/2018/12/bayes_nn.png\">\n",
    "<p style=\"text-align: center\"><sup>Image from <a href=\"https://sanjaykthakur.com/2018/12/05/the-very-basics-of-bayesian-neural-networks/\">https://sanjaykthakur.com/2018/12/05/the-very-basics-of-bayesian-neural-networks/</a></sup></p>\n",
    "\n",
    "<br><br>\n",
    "\n",
    "To do this we'll use **Bayes by Backprop** method introduced in a Blundell et al. paper [Weight uncertainty in neural networks](https://arxiv.org/pdf/1505.05424.pdf) (2015).\n",
    "\n",
    "<br>\n",
    "\n",
    "We'll try to estimate weight distribution parameters $\\theta$, given data $D$:\n",
    "\n",
    "\n",
    "$$\\large P(\\theta | D) = \\frac{P(D | \\theta) P(\\theta)}{P(D)} $$\n",
    "\n",
    "\n",
    "We'll follow a three-step formula to achieve this:\n",
    "\n",
    "1. Pick a prior density over weights $P(\\theta)$\n",
    "2. Use training data $D$ to determine the likelihood $P(D | \\theta)$\n",
    "3. Estimate the posterior density over weights $P(\\theta | D)$ \n",
    "\n",
    "<br>\n",
    "\n",
    "Keywords: [***KL divergence***](https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence), [***ELBO***](https://en.wikipedia.org/wiki/Evidence_lower_bound), [***reparametrization trick***](https://arxiv.org/pdf/1505.05424.pdf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.1 Epistemic uncertainty - linear case"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's start with our linear data!\n",
    "x_train_100 = np.linspace(-1, 1, 100)[:, np.newaxis]\n",
    "x_train_1000 = np.linspace(-1, 1, 1000)[:, np.newaxis]\n",
    "\n",
    "y_train_100 = 1.5*x_train_100 + 3 + 0.35*np.random.randn(100)[:, np.newaxis]\n",
    "y_train_1000 = 1.5*x_train_1000 + 3 + 0.35*np.random.randn(1000)[:, np.newaxis]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the data\n",
    "plt.figure(figsize=(15, 4))\n",
    "\n",
    "plt.subplot(121)\n",
    "plt.scatter(x_train_100, y_train_100, alpha=.5)\n",
    "plt.xlabel('$x$')\n",
    "plt.ylabel('$y$')\n",
    "plt.title('$n=100$', fontsize=14)\n",
    "\n",
    "plt.subplot(122)\n",
    "plt.scatter(x_train_1000, y_train_1000, alpha=.5)\n",
    "plt.xlabel('$x$')\n",
    "plt.ylabel('$y$')\n",
    "plt.title('$n=1000$', fontsize=14)\n",
    "\n",
    "plt.suptitle('Training data')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 3.1.1 Define prior and posterior"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`prior()` and `posterior()` functions are necessary to parametrize `tfpl.DenseVariational` layer.\n",
    "\n",
    "They both need to take `kernel_size`, `bias_size` and `dtype` arguments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prior - diagonal MVN ~ N(0, 1)\n",
    "def prior(kernel_size, bias_size, dtype=None):\n",
    "    \n",
    "    n = kernel_size + bias_size\n",
    "    \n",
    "    prior_model = tf.keras.Sequential([\n",
    "        \n",
    "        tfpl.DistributionLambda(\n",
    "            # Non-trianable distribution\n",
    "            lambda t: tfd.MultivariateNormalDiag(loc=tf.zeros(n), scale_diag=tf.ones(n)))\n",
    "    ])\n",
    "    \n",
    "    return prior_model\n",
    "\n",
    "\n",
    "# Posterior\n",
    "def posterior(kernel_size, bias_size, dtype=None):\n",
    "    \n",
    "    n = kernel_size + bias_size\n",
    "    \n",
    "    posterior_model = tf.keras.Sequential([\n",
    "        \n",
    "        tfpl.VariableLayer(tfpl.MultivariateNormalTriL.params_size(n), dtype=dtype),\n",
    "        tfpl.MultivariateNormalTriL(n)\n",
    "    ])\n",
    "    \n",
    "    return posterior_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 3.1.2 Aleatoric uncertainty - define a model with `DenseVariational` layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the model\n",
    "def get_aleatoric_model(x_train_shape):\n",
    "    \n",
    "    model = tf.keras.Sequential([\n",
    "        tfpl.DenseVariational(\n",
    "            1,\n",
    "            input_shape=(1,),\n",
    "            make_prior_fn=prior,\n",
    "            make_posterior_fn=posterior,\n",
    "            kl_weight=1/x_train_shape, # Normalizing to scale the D_KL term in ELBO properly when using minibatches.\n",
    "            kl_use_exact=False) # could be `True` in this case, but we go for estimated value\n",
    "    ])\n",
    "\n",
    "    # Compile\n",
    "    model.compile(loss='mse', optimizer='sgd')\n",
    "    \n",
    "    return model\n",
    "\n",
    "# Initialize\n",
    "model = get_aleatoric_model(100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fit the model on two datasets \n",
    "print('Fitting `model_100`...')\n",
    "model_100 = get_aleatoric_model(100)\n",
    "history_100 = model_100.fit(x_train_100, y_train_100, epochs=500, verbose=False)\n",
    "\n",
    "print('Fitting `model_1000`...')\n",
    "model_1000 = get_aleatoric_model(1000)\n",
    "history_1000 = model_1000.fit(x_train_1000, y_train_1000, epochs=500, verbose=False)\n",
    "\n",
    "print('Done.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot loss\n",
    "plt.figure(figsize=(6, 4))\n",
    "plt.plot(history_100.history['loss'], label='model_100', alpha=.5)\n",
    "plt.plot(history_1000.history['loss'], label='model_1000', alpha=.5)\n",
    "plt.title('Loss')\n",
    "plt.xlabel('No. of epochs')\n",
    "plt.ylabel('Loss')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot fitted regression lines\n",
    "N_ITERS = 15\n",
    "\n",
    "plt.figure(figsize=(15, 4))\n",
    "\n",
    "#------------------------------\n",
    "\n",
    "plt.subplot(121)\n",
    "plt.scatter(x_train_100, y_train_100, alpha=.5, label='Training data')\n",
    "for _ in range(N_ITERS):\n",
    "    y_model_100 = model_100(x_train_100)\n",
    "    if _ == 0:\n",
    "        plt.plot(x_train_100, y_model_100, color='red', alpha=0.5, label='Fitted reg. lines')\n",
    "    else:\n",
    "        plt.plot(x_train_100, y_model_100, color='red', alpha=0.5)        \n",
    "\n",
    "plt.xlabel('$x$')\n",
    "plt.ylabel('$y$')\n",
    "plt.title('$n=100$', fontsize=14)\n",
    "plt.legend()\n",
    "\n",
    "#------------------------------\n",
    "\n",
    "plt.subplot(122)\n",
    "plt.scatter(x_train_1000, y_train_1000, alpha=.5, label='Training data')\n",
    "for _ in range(N_ITERS):\n",
    "    y_model_1000 = model_1000(x_train_1000)\n",
    "    if _ == 0:\n",
    "        plt.plot(x_train_1000, y_model_1000, color='red', alpha=0.5, label='Fitted reg. lines')\n",
    "    else:\n",
    "        plt.plot(x_train_1000, y_model_1000, color='red', alpha=0.5)        \n",
    "\n",
    "plt.xlabel('$x$')\n",
    "plt.ylabel('$y$')\n",
    "plt.title('$n=1000$', fontsize=14)\n",
    "plt.legend()\n",
    "\n",
    "#------------------------------\n",
    "\n",
    "plt.suptitle('Training data and fitted lines')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Going fully probabilistic: aleatoric & epistemic uncertainty"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.1 Fully probabilistic non-linear case "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create non-linear data\n",
    "n_samples = int(5e3)\n",
    "x_train = np.linspace(-1, 1, n_samples)[:, np.newaxis]\n",
    "y_train = x_train**3 + 0.1*(1.5 + x_train)*np.random.randn(n_samples)[:, np.newaxis]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the data\n",
    "plt.figure(figsize=(6, 4))\n",
    "plt.scatter(x_train, y_train, alpha=0.05)\n",
    "plt.title('Non-linear training data')\n",
    "plt.xlabel('$x$')\n",
    "plt.ylabel('$y$')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prior - diagonal MVN ~ N(0, 1)\n",
    "def prior(kernel_size, bias_size, dtype=None):\n",
    "    \n",
    "    n = kernel_size + bias_size\n",
    "    \n",
    "    prior_model = tf.keras.Sequential([\n",
    "        \n",
    "        tfpl.DistributionLambda(\n",
    "            # Non-trianable distribution\n",
    "            lambda t: tfd.MultivariateNormalDiag(loc=tf.zeros(n), scale_diag=tf.ones(n)))\n",
    "    ])\n",
    "    \n",
    "    return prior_model\n",
    "\n",
    "\n",
    "# Posterior\n",
    "def posterior(kernel_size, bias_size, dtype=None):\n",
    "    \n",
    "    n = kernel_size + bias_size\n",
    "    \n",
    "    posterior_model = tf.keras.Sequential([\n",
    "        tfpl.VariableLayer(tfpl.MultivariateNormalTriL.params_size(n), dtype=dtype),\n",
    "        tfpl.MultivariateNormalTriL(n)\n",
    "    ])\n",
    "    \n",
    "    return posterior_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the model\n",
    "def get_full_model(x_train_shape):\n",
    "    \n",
    "    model = tf.keras.Sequential([\n",
    "\n",
    "        # Epistemic uncertainty\n",
    "        tfpl.DenseVariational(units=8,\n",
    "                              input_shape=(1,),\n",
    "                              make_prior_fn=prior,\n",
    "                              make_posterior_fn=posterior,\n",
    "                              kl_weight=1/x_train_shape,\n",
    "                              kl_use_exact=False,\n",
    "                              activation='sigmoid'),\n",
    "        \n",
    "        tfpl.DenseVariational(units=tfpl.IndependentNormal.params_size(1),\n",
    "                              make_prior_fn=prior,\n",
    "                              make_posterior_fn=posterior,\n",
    "                              kl_use_exact=False,\n",
    "                              kl_weight=1/x_train_shape),\n",
    "\n",
    "        # Aleatoric uncertainty\n",
    "        tfpl.IndependentNormal(1)\n",
    "    ])\n",
    "\n",
    "    def neg_loglik(y_true, y_pred):\n",
    "        return -y_pred.log_prob(y_true)\n",
    "\n",
    "    model.compile(loss=neg_loglik, optimizer='rmsprop')\n",
    "    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the model\n",
    "model_full = get_full_model(n_samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# Initialize the model\n",
    "model_full = get_full_model(n_samples)\n",
    "\n",
    "# Fit\n",
    "history = model_full.fit(x_train, y_train, epochs=1000, verbose=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot fitted regression lines\n",
    "N_ITERS = 15\n",
    "\n",
    "plt.figure(figsize=(9, 6))\n",
    "\n",
    "#------------------------------\n",
    "\n",
    "plt.scatter(x_train, y_train, alpha=.1, label='Training data')\n",
    "for _ in range(N_ITERS):\n",
    "    \n",
    "    # Compute conf ints\n",
    "    y_sample = model_full(x_train).sample()\n",
    "    y_hat = model_full(x_train).mean()\n",
    "    y_sd = model_full(x_train).stddev()\n",
    "    y_hat_lower = y_hat - 2 * y_sd\n",
    "    y_hat_upper = y_hat + 2 * y_sd\n",
    "    \n",
    "\n",
    "    if _ == 0:\n",
    "        plt.plot(x_train, y_hat, color='red', alpha=.5, lw=1, label='Fitted reg. lines')\n",
    "        plt.plot(x_train, y_hat_lower, c='g', alpha=.5, lw=1, label='95% CI')\n",
    "        plt.plot(x_train, y_hat_upper, c='g', alpha=.5, lw=1)\n",
    "    else:\n",
    "        plt.plot(x_train, y_hat, color='red', alpha=.5, lw=1)    \n",
    "        plt.plot(x_train, y_hat_lower, c='g', alpha=.5, lw=1)\n",
    "        plt.plot(x_train, y_hat_upper, c='g', alpha=.5, lw=1)\n",
    "\n",
    "plt.xlabel('$x$')\n",
    "plt.ylabel('$y$')\n",
    "plt.legend()\n",
    "\n",
    "#------------------------------\n",
    "\n",
    "plt.title(f'Training data + fitted lines + 95% CIs\\nn={n_samples}')\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:tf-proba-env]",
   "language": "python",
   "name": "conda-env-tf-proba-env-py"
  },
  "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}