# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# + [markdown] id="view-in-github" colab_type="text"
# <a href="https://colab.research.google.com/github/mella30/Deep-Learning-with-Tensorflow-2/blob/main/Course3-Probabilistic_Deep_Learning_with_Tensorflow2/week2_probabilistic_modelling.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>

# + id="zIvqaAUzKaos" colab={"base_uri": "https://localhost:8080/"} outputId="feec3153-a19a-4b83-d176-acc91d59573c"
import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions
tfpl = tfp.layers

print('TF version:', tf.__version__)
print('TFP version:', tfp.__version__)

# + [markdown] id="Nf9HBWGVKaox"
# # Probabilistic layers and Bayesian neural networks

# + [markdown] id="sVYM_IToKaox"
# ## Coding tutorials
# #### [1. The DistributionLambda layer](#coding_tutorial_1)
# #### [2. Probabilistic layers](#coding_tutorial_2)
# #### [3. The DenseVariational layer](#coding_tutorial_3)
# #### [4. Reparameterization layers](#coding_tutorial_4)

# + [markdown] id="QHdmVjG1Kaoy"
# ***
# <a id="coding_tutorial_1"></a>
# ## The `DistributionLambda` layer

# + id="yd8sUI97Kaoy"
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
from tensorflow.keras.optimizers import RMSprop
import numpy as np
import matplotlib.pyplot as plt

# + [markdown] id="rGyQeDUaKao2"
# #### Create a probabilistic model using the `DistributionLambda` layer
#
# Create a model whose first layer represents:
#
# $$
# y = \text{sigmoid}(x) = \frac{1}{1 + \exp(-x)}.
# $$

# + id="F7dgf5p8Kao2" colab={"base_uri": "https://localhost:8080/", "height": 265} outputId="0d5f0389-eb83-4302-cca5-24dc46380c18"
# Create a deterministic and a probabilistic sigmoid model

model_deterministic = Sequential([
    Dense(input_shape=(1,), units=1, activation='sigmoid',
          kernel_initializer=tf.constant_initializer(1),
          bias_initializer=tf.constant_initializer(0)),
])

model_probabilistic = Sequential([
    Dense(input_shape=(1,), units=1, activation='sigmoid',
          kernel_initializer=tf.constant_initializer(1),
          bias_initializer=tf.constant_initializer(0)),
    tfpl.DistributionLambda(lambda t : tfd.Bernoulli(probs=t),
                            convert_to_tensor_fn=tfd.Distribution.sample)  # can also be mean, ...
])

# Plot the function
x_plot = np.linspace(-5, 5, 100)
plt.scatter(x_plot, model_deterministic.predict(x_plot), alpha=0.4)
plt.scatter(x_plot, model_probabilistic.predict(x_plot), alpha=0.4)
plt.plot(x_plot, 1/(1 + np.exp(-x_plot)), color='r', alpha=0.8)
plt.show()


# + id="THWcaknbKao5"
# Create a constant input for this model

x = np.array([[0]])

# + id="P-N5tHQZKao8" colab={"base_uri": "https://localhost:8080/"} outputId="14f36d07-a1b8-4205-ad8b-4a6329706aa4"
# Explore the feedforward object...

y_deterministic = model_deterministic(x)
y_deterministic

# + id="i2g_15Cxkmxd" colab={"base_uri": "https://localhost:8080/"} outputId="a6ffd87e-dd80-44d8-8f7a-aa908054f870"
y_probabilistic = model_probabilistic(x)
y_probabilistic

# + id="ugBgTsxGKao-" colab={"base_uri": "https://localhost:8080/"} outputId="8eed87ae-b9a6-49f7-c715-e8373ef29db8"
# ... and its behaviour under repeated calls

for _ in range(5):
  print(model_deterministic.predict(x))
  print(model_probabilistic.predict(x))
  print("\n")

# + [markdown] id="G6I0aMgpKapD"
# #### Use the forward model to create probabilistic training data

# + id="jSL5TMgVKapD" colab={"base_uri": "https://localhost:8080/", "height": 320} outputId="54e63346-ede7-43e0-fdd0-c0223c767ff5"
# Use the model to create 500 training points

x_train = np.linspace(-5, 5, 500)[:, np.newaxis]
y_train = model_probabilistic.predict(x_train)

# Plot the data and the mean of the distribution
fig, ax = plt.subplots(figsize=(5, 5))
ax.scatter(x_train, y_train, alpha=0.04, color='blue', label='samples')
ax.plot(x_train, model_probabilistic(x_train).mean().numpy().flatten(), 
        color='red', alpha=0.8, label='mean')
ax.legend()
plt.show()

# + [markdown] id="Abtwq7yCKapG"
# #### Create a new probabilistic model with the wrong weights

# + id="gZDJVsXmKapG"
# Create a new version of the model, with the wrong weights

model_untrained = Sequential([
    Dense(input_shape=(1,), units=1, activation='sigmoid',
          kernel_initializer=tf.constant_initializer(2), # change the weight initializer
          bias_initializer=tf.constant_initializer(2)),
    tfpl.DistributionLambda(lambda t : tfd.Bernoulli(probs=t),
                            convert_to_tensor_fn=tfd.Distribution.sample)  # can also be mean, ...
])


# + [markdown] id="5yENtha1KapJ"
# #### Train the new model with the negative loglikelihood

# + id="VagGrOU7KapJ"
# Define negative loglikelihood, which we will use for training

def nll(y_true, y_pred):
  # .log_prob() take the log probability of the true training data (which is a distribution object)
  return -y_pred.log_prob(y_true)


# + id="8zQaoho3KapL"
# Compile untrained model

model_untrained.compile(loss=nll, optimizer=RMSprop(learning_rate=0.01))

# + id="QnFKk9G5KapO"
# Train model, record weights after each epoch

epochs = [0]
training_weights = [model_untrained.weights[0].numpy()[0, 0]]
training_bias = [model_untrained.weights[1].numpy()[0]]
for epoch in range(100):
    model_untrained.fit(x=x_train, y=y_train, epochs=1, verbose=False)
    epochs.append(epoch)
    training_weights.append(model_untrained.weights[0].numpy()[0, 0])
    training_bias.append(model_untrained.weights[1].numpy()[0])

# + id="OQUK9RCYKapQ" colab={"base_uri": "https://localhost:8080/", "height": 279} outputId="c16f4180-f683-49c3-a4bf-c7bc05195bfe"
# Plot the model weights as they train, converging to the correct values

plt.plot(epochs, training_weights, label='weight')
plt.plot(epochs, training_bias, label='bias')
plt.axhline(y=1, label='true_weight', color='k', linestyle=':')
plt.axhline(y=0, label='true_bias', color='k', linestyle='--')
plt.xlabel('Epochs')
plt.legend()
plt.show()

# epistemic uncertainty prevents the weights to exactly fit the original parameter values
# (data are not perfectly representing the original model)

# + [markdown] id="js3g6c9OKapS"
# ***
# <a id="coding_tutorial_2"></a>
# ## Probabilistic layers

# + id="H07K-uaiKapT"
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
from tensorflow.keras.losses import MeanSquaredError
from tensorflow.keras.optimizers import RMSprop
import numpy as np
import matplotlib.pyplot as plt

# + [markdown] id="osxVhElHKapV"
# #### Create data
#
# The data you'll be working with is artifically created from the following equation:
# $$ y_i = x_i + \frac{3}{10}\epsilon_i$$
# where $\epsilon_i \sim N(0, 1)$ are independent and identically distributed.

# + id="K-CA1wLXKapV" colab={"base_uri": "https://localhost:8080/", "height": 279} outputId="576a0b80-abe2-4676-f21d-6a6f566c13cc"
# Create and plot 100 points of training data

x_train = np.linspace(-1, 1, 100)[:, np.newaxis]
y_train = x_train + 0.3*np.random.randn(100)[:, np.newaxis]

plt.scatter(x_train, y_train, alpha=0.4)
plt.xlabel('x')
plt.ylabel('y')
plt.show()

# + [markdown] id="JXpmtF-uKapX"
# #### Deterministic linear regression with MSE loss

# + id="CF5MF6OPKapY" colab={"base_uri": "https://localhost:8080/", "height": 449} outputId="640228b5-06a2-4c67-9418-b6786dbc3681"
# Create and train deterministic linear model using mean squared error loss

# Create linear regression via Sequential model
model_deterministic = Sequential([
    Dense(units=1, input_shape=(1,))
])
model_deterministic.compile(loss=MeanSquaredError(), optimizer=RMSprop(learning_rate=0.005))
model_deterministic.summary()
model_deterministic.fit(x_train, y_train, epochs=200, verbose=False)

# Plot the data and model
plt.scatter(x_train, y_train, alpha=0.4, label='data')
plt.plot(x_train, model_deterministic.predict(x_train), color='red', alpha=0.8, label='model')
plt.legend()
plt.show()

# + id="jhKQ05n6Kapb" colab={"base_uri": "https://localhost:8080/"} outputId="8833d746-1876-441d-88c9-43f3b5d63a40"
# Examine the model predictions

x = np.array([[0]])
y_deterministic = model_deterministic(x)
y_deterministic

# does not capture the aleatoric uncertainty

# + [markdown] id="SiC14wElKape"
# #### Probabilistic linear regression with both user-defined and learned variance

# + id="jx8f6mmwKape" colab={"base_uri": "https://localhost:8080/"} outputId="c1d9e876-fdc8-43de-8dca-5a34f92df215"
# Create probabilistic regression with normal distribution as final layer
# take the output of the Dense layer to parameterize the mean of the normal distribution

event_shape=1
model_probabilistic = Sequential([
    Dense(units=tfpl.IndependentNormal.params_size(event_shape=1), input_shape=(1,)),
    tfpl.IndependentNormal(event_shape=event_shape)
])

model_probabilistic.summary()  # output is a distribution


# + id="z2NRMucxKapg" colab={"base_uri": "https://localhost:8080/"} outputId="b6a0afb9-60e6-4488-ab98-f06ec04e686a"
# Train model using the negative loglikelihood

def nll(y_true, y_pred):
  return -y_pred.log_prob(y_true)

model_probabilistic.compile(loss=nll, optimizer=RMSprop(learning_rate=0.005))
model_probabilistic.fit(x_train, y_train, epochs=200, verbose=False)

# + id="FSJlVYI1Kapi" colab={"base_uri": "https://localhost:8080/"} outputId="943e941a-ff82-405e-f3bd-0476847b4a1d"
# Examine the distribution created as a feedforward value

y_probabilistic = model_probabilistic(x)
y_probabilistic

# + id="YXepFof7Kapk" colab={"base_uri": "https://localhost:8080/", "height": 302} outputId="10e62ff6-fe78-4e47-86b0-5b69b051adda"
# Plot the data and a sample from the model

y_model = model_probabilistic(x_train)
y_sample = y_model.sample()
y_hat = y_model.mean()
y_sd = y_model.stddev()
y_hat_m2sd = y_hat - 2 * y_sd
y_hat_p2sd = y_hat + 2 * y_sd

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5), sharey=True)
ax1.scatter(x_train, y_train, alpha=0.4, label='data')
ax1.scatter(x_train, y_sample, alpha=0.4, color='red', label='model sample')
ax1.legend()
ax2.scatter(x_train, y_train, alpha=0.4, label='data')
ax2.plot(x_train, y_hat, color='red', alpha=0.8, label='model $\mu$')
ax2.plot(x_train, y_hat_m2sd, color='green', alpha=0.8, label='model $\mu \pm 2 \sigma$')
ax2.plot(x_train, y_hat_p2sd, color='green', alpha=0.8)
ax2.legend()
plt.show()


# + [markdown] id="NahVYLPKKapm"
# #### Probabilistic linear regression with nonlinear learned mean & variance

# + [markdown] id="Fj0z3rEtKapn"
# Let's change the data to being nonlinear:
# $$ y_i = x_i^3 + \frac{1}{10}(2 + x_i)\epsilon_i$$
# where $\epsilon_i \sim N(0, 1)$ are independent and identically distributed.

# + id="NOF25FdlKapn" colab={"base_uri": "https://localhost:8080/", "height": 265} outputId="16af8ab0-86e7-4e7f-f908-9c45a93cf75d"
# Create and plot 10000 data points

x_train = np.linspace(-1, 1, 1000)[:, np.newaxis]
y_train = np.power(x_train, 3) + 0.1*(2+x_train)*np.random.randn(1000)[:, np.newaxis]

plt.scatter(x_train, y_train, alpha=0.1)
plt.show()

# + id="DbrN-Dv7Kapp" colab={"base_uri": "https://localhost:8080/"} outputId="9fab3cea-e5b1-47c3-a948-bf328db339cd"
# Create probabilistic regression: normal distribution with fixed variance

model = Sequential([
    Dense(input_shape=(1,), units=8, activation='sigmoid'),
    Dense(tfpl.IndependentNormal.params_size(event_shape=1)),
    tfpl.IndependentNormal(event_shape=1)
])
model.compile(loss=nll, optimizer=RMSprop(learning_rate=0.01))
model.summary()

# + id="LRn7haPjKapr" colab={"base_uri": "https://localhost:8080/"} outputId="0d131f79-3304-4cce-a6d1-e8c34186449d"
# Train model

model.fit(x_train, y_train, epochs=200, verbose=False)
model.evaluate(x_train, y_train)

# + id="GWdceGQPKapu" colab={"base_uri": "https://localhost:8080/", "height": 302} outputId="70c0b320-9b38-4417-9254-795f54c02f3b"
# Plot the data and a sample from the model

y_model = model(x_train)
y_sample = y_model.sample()
y_hat = y_model.mean()
y_sd = y_model.stddev()
y_hat_m2sd = y_hat - 2 * y_sd
y_hat_p2sd = y_hat + 2 * y_sd

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5), sharey=True)
ax1.scatter(x_train, y_train, alpha=0.2, label='data')
ax1.scatter(x_train, y_sample, alpha=0.2, color='red', label='model sample')
ax1.legend()
ax2.scatter(x_train, y_train, alpha=0.2, label='data')
ax2.plot(x_train, y_hat, color='red', alpha=0.8, label='model $\mu$')
ax2.plot(x_train, y_hat_m2sd, color='green', alpha=0.8, label='model $\mu \pm 2 \sigma$')
ax2.plot(x_train, y_hat_p2sd, color='green', alpha=0.8)
ax2.legend()
plt.show()

# + [markdown] id="sk_oEVCmKapw"
# ***
# <a id="coding_tutorial_3"></a>
# ## The `DenseVariational` layer

# + id="QR8mg4KbKapw"
from tensorflow.keras.models import Sequential
from tensorflow.keras.losses import MeanSquaredError
from tensorflow.keras.optimizers import RMSprop
import numpy as np
import matplotlib.pyplot as plt

# + [markdown] id="pFCwsZzVKapy"
# #### Create linear data with Gaussian noise
#
# The data you'll be working with is the same as you used before:
# $$ y_i = x_i + \frac{3}{10}\epsilon_i$$
# where $\epsilon_i \sim N(0, 1)$ are independent and identically distributed. We'll be running a Bayesian linear regression on this data.

# + id="DPbuqlB5Kapz" colab={"base_uri": "https://localhost:8080/", "height": 265} outputId="b09e33b5-88c2-4708-c05d-d5aa10fd8dbf"
# Use the same data as before -- create and plot 100 data points

x_train = np.linspace(-1, 1, 100)[:, np.newaxis]
y_train = x_train + 0.3*np.random.randn(100)[:, np.newaxis]

plt.scatter(x_train, y_train, alpha=0.4)
plt.show()


# + [markdown] id="EjIP1ggBKap1"
# #### Create the prior and posterior distribution for model weights

# + id="vDaiVcweKap1"
# Define the prior weight distribution -- all N(0, 1) -- and not trainable

def prior(kernel_size, bias_size, dtype=None):
  n = kernel_size + bias_size  # total number of paramters in the layer
  prior_model = Sequential([
                            # this model does not have any trainable parameters (hard-coded zeros)
                            # (a trainable prior would generally be possible)
                            tfpl.DistributionLambda(
                                lambda t : tfd.MultivariateNormalDiag(loc=tf.zeros(n),
                                                                scale_diag=tf.ones(n)))                            
  ])
  # returns a distribution
  return prior_model  


# + id="1GdB92laKap5"
# Define variational posterior weight distribution -- multivariate Gaussian

def posterior(kernel_size, bias_size, dtype=None):
  n = kernel_size + bias_size  # total number of paramters in the layer
  posterior_model = Sequential([
                                # this model needs trainable paramters since it should be learned from the data
                                # use a variable layer to learn the parameters of the multivariate normal distribution
                                tfpl.VariableLayer(tfpl.MultivariateNormalTriL.params_size(n), dtype=dtype),  # multivariate gaussian with full covariance matrix
                                tfpl.MultivariateNormalTriL(n)
  ])
  return posterior_model


# + [markdown] id="xnv7LMUXKap6"
# #### Aside: analytical posterior
#
# In this tutorial, we're using a variational posterior because, in most settings, it's not possible to derive an analytical one. However, in this simple setting, it is possible. Specifically, running a Bayesian linear regression on $x_i$ and $y_i$ with $i=1, \ldots, n$ and a unit Gaussian prior on both $\alpha$ and $\beta$:
#
# $$
# y_i = \alpha + \beta x_i + \epsilon_i, \quad 
# \epsilon_i \sim N(0, \sigma^2), \quad 
# \alpha \sim N(0, 1), \quad 
# \beta \sim N(0, 1)
# $$
#
# gives a multivariate Gaussian posterior on $\alpha$ and $\beta$:
#
# $$
# \begin{pmatrix}
# \alpha \\
# \beta
# \end{pmatrix}
# \sim
# N(\mathbf{\mu}, \mathbf{\Sigma})
# $$
# where
# $$ 
# \mathbf{\mu}
# = 
# \mathbf{\Sigma} 
# \begin{pmatrix}
# \hat{n} \bar{y} \\
# \hat{n} \overline{xy}
# \end{pmatrix},
# \quad
# \mathbf{\Sigma} = 
# \frac{1}{(\hat{n} + 1)(\hat{n} \overline{x^2} + 1) - \hat{n}^2 \bar{x}^2}
# \begin{pmatrix}
# \hat{n} \overline{x^2} + 1 & -\hat{n} \bar{x} \\
# -\hat{n} \bar{x} & \hat{n} + 1
# \end{pmatrix}.
# $$
#
# In the above, $\hat{n} = \frac{n}{\sigma^2}$ and $\bar{t} = \frac{1}{n}\sum_{i=1}^n t_i$ for any $t$. In general, however, it's not possible to determine the analytical form for the posterior. For example, in models with a hidden layer with nonlinear activation function, the analytical posterior cannot be determined in general, and variational methods as below are useful.

# + [markdown] id="jGNAZCCYKap7"
# #### Create the model with `DenseVariational` layers

# + id="2Ee6qHUTKap7" colab={"base_uri": "https://localhost:8080/"} outputId="c44c1919-da6e-41e4-d7d6-09d3f463d5b0"
# Create linear regression model with weight uncertainty: weights are
# distributed according to posterior (and, indirectly, prior) distribution

model = Sequential([
                    tfpl.DenseVariational(input_shape=(1,), 
                                          units=1,
                                          make_prior_fn=prior,
                                          make_posterior_fn=posterior,
                                          kl_weight=1/x_train.shape[0],
                                          kl_use_exact=True)  # tfp has stored the analytically calculated KL divergence between many common distributions
])

model.compile(loss=MeanSquaredError(),
              optimizer=RMSprop(learning_rate=0.005))
model.summary()
# usual linear regression has only 2 parameters: the slope and the y-intercept
# here, the weights have distributions, that means:
# - the mean and the variance of the slope
# - the mean and the variance of the y-intercept
# - the covariance between them

# + [markdown] id="yfnMhy6xKap9"
# #### Train model and inspect

# + id="0864YLRGKap-" colab={"base_uri": "https://localhost:8080/"} outputId="d88cb95a-ef55-4afe-d4fd-10e53621b2e0"
# Fit the model, just like a deterministic linear regression

model.fit(x_train, y_train, epochs=500, verbose=False)

# + id="NzCT1byTKaqA" colab={"base_uri": "https://localhost:8080/"} outputId="e8220740-b601-450e-b01d-eb659b9bb6e5"
# Check out the parameters of the prior and posterior distribution

dummy_input = np.array([[0]])
model_prior = model.layers[0]._prior(dummy_input)
model_posterior = model.layers[0]._posterior(dummy_input)
# the variables of the prior are set to unit gaussians and are not trainable
print('prior mean:           ', model_prior.mean().numpy())
print('prior variance:       ', model_prior.variance().numpy())
# learned mean for the slope and y-intercept (were 1. and 0.)
print('posterior mean:       ', model_posterior.mean().numpy())
# learned variance for the slope and the y-intercept (uncertainty) are on the diagnoal of the covariance matrix
print('posterior covariance: ', model_posterior.covariance().numpy()[0])
# learned covariance between slope and y-intercept are on the off-diagonal (are equal)
print('                      ', model_posterior.covariance().numpy()[1])

# + id="7RWg-xARKaqC" colab={"base_uri": "https://localhost:8080/", "height": 265} outputId="e743a0e3-be72-4dbc-d3c3-7135103bcaed"
# Plot an ensemble of linear regressions, with weights sampled from the posterior distribution
# this shows the epistemic uncertainty: we don't actually know which model produced the data

plt.scatter(x_train, y_train, alpha=0.4, label='data')
for _ in range(10):
    y_model = model(x_train)
    if _ == 0:
        plt.plot(x_train, y_model, color='red', alpha=0.8, label='model')
    else:
        plt.plot(x_train, y_model, color='red', alpha=0.8)        
plt.legend()
plt.show()

# + [markdown] id="8x7I47G6KaqE"
# #### Explore the effect of sample size

# + id="9B3Sqy4aKaqF" colab={"base_uri": "https://localhost:8080/", "height": 265} outputId="295c4be7-5aa6-4183-bddc-21c080bec951"
# epistemic uncertainty can be explained away with more data

# Create two datasets, one with 1000 points, another with 100

x_train_1000 = np.linspace(-1, 1, 1000)[:, np.newaxis]
y_train_1000 = x_train_1000 + 0.3*np.random.randn(1000)[:, np.newaxis]

x_train_100 = np.linspace(-1, 1, 100)[:, np.newaxis]
y_train_100 = x_train_100 + 0.3*np.random.randn(100)[:, np.newaxis]

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4), sharex=True, sharey=True)
ax1.scatter(x_train_1000, y_train_1000, alpha=0.1)
ax2.scatter(x_train_100, y_train_100, alpha=0.4)
plt.show()

# + id="vDpDiW14KaqH" colab={"base_uri": "https://localhost:8080/"} outputId="95c24a1f-576b-4bec-bcb0-6ce2c158f0d8"
# Train a model on each dataset

model_1000 = Sequential([tfpl.DenseVariational(input_shape=(1,), 
                                               units=1,
                                               make_prior_fn=prior, 
                                               make_posterior_fn=posterior,
                                               kl_weight=1/1000)])

model_100 = Sequential([tfpl.DenseVariational(input_shape=(1,), 
                                              units=1,
                                              make_prior_fn=prior, 
                                              make_posterior_fn=posterior,
                                              kl_weight=1/100)])

model_1000.compile(loss=MeanSquaredError(), optimizer=RMSprop(learning_rate=0.005))
model_100.compile(loss=MeanSquaredError(), optimizer=RMSprop(learning_rate=0.005))

model_1000.fit(x_train_1000, y_train_1000, epochs=50, verbose=False)
model_100.fit(x_train_100, y_train_100, epochs=500, verbose=False)

# + id="v_q1wfCwKaqJ" colab={"base_uri": "https://localhost:8080/", "height": 265} outputId="ab15ca1f-4411-4ff7-8b2a-6dc8656d6aac"
# Plot an ensemble of linear regressions from each model

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4), sharex=True, sharey=True)
for _ in range(10):
    y_model_1000 = model_1000(x_train_1000)
    ax1.scatter(x_train_1000, y_train_1000, color='C0', alpha=0.02)
    ax1.plot(x_train_1000, y_model_1000, color='red', alpha=0.8)
    y_model_100 = model_100(x_train_100)
    ax2.scatter(x_train_100, y_train_100, color='C0', alpha=0.05)
    ax2.plot(x_train_100, y_model_100, color='red', alpha=0.8)
plt.show()

# the variance (= uncertainty) in the outputs of the right model (with fewer samples) 
# is higher as in the left model (with more samples)

# + [markdown] id="ooeqV6GxKaqL"
# #### Put it all together: nonlinear probabilistic regression with weight uncertainty

# + [markdown] id="CZ-gVKXVKaqM"
# Let's change the data to being nonlinear:
# $$ y_i = x_i^3 + \frac{1}{10}(2 + x_i)\epsilon_i$$
# where $\epsilon_i \sim N(0, 1)$ are independent and identically distributed.

# + id="GNopZXPcKaqM" colab={"base_uri": "https://localhost:8080/", "height": 265} outputId="ac559932-6688-4765-a672-588da12c05ea"
# Create and plot 1000 data points

x_train = np.linspace(-1, 1, 1000)[:, np.newaxis]
y_train = np.power(x_train, 3) + 0.1*(2+x_train)*np.random.randn(1000)[:, np.newaxis]

plt.scatter(x_train, y_train, alpha=0.1)
plt.show()

# + id="WF9yxmDiKaqO" colab={"base_uri": "https://localhost:8080/"} outputId="530cf22d-a0d9-47fb-a9b1-3cd3501c5c90"
# Create probabilistic regression with one hidden layer, weight uncertainty

model = Sequential([
    # epistemic uncertainty modelling
    tfpl.DenseVariational(units=8,
                          input_shape=(1,),
                          make_prior_fn=prior,
                          make_posterior_fn=posterior,
                          kl_weight=1/x_train.shape[0],
                          activation='sigmoid'),
    tfpl.DenseVariational(units=tfpl.IndependentNormal.params_size(1),
                          make_prior_fn=prior,
                          make_posterior_fn=posterior,
                          kl_weight=1/x_train.shape[0]),
    # aleatoric uncertainty modelling
    tfpl.IndependentNormal(1)
])

def nll(y_true, y_pred):
    return -y_pred.log_prob(y_true)

model.compile(loss=nll, optimizer=RMSprop(learning_rate=0.005))
model.summary()

# + id="tGlYg-jnKaqP" colab={"base_uri": "https://localhost:8080/"} outputId="d7c73656-c2d1-48b8-a3fa-58ff9b7b72e2"
# Train the model

model.fit(x_train, y_train, epochs=1000, verbose=False)
model.evaluate(x_train, y_train)

# + id="QdeHFDzXKaqR" colab={"base_uri": "https://localhost:8080/", "height": 265} outputId="652cdb98-c776-4ee5-c7b4-5d8add3fab0d"
# Plot an ensemble of trained probabilistic regressions

plt.scatter(x_train, y_train, marker='.', alpha=0.2, label='data')
for _ in range(5):
    y_model = model(x_train)
    y_hat = y_model.mean()
    y_hat_m2sd = y_hat - 2 * y_model.stddev()
    y_hat_p2sd = y_hat + 2 * y_model.stddev()
    if _ == 0:
        plt.plot(x_train, y_hat, color='red', alpha=0.8, label='model $\mu$')
        plt.plot(x_train, y_hat_m2sd, color='green', alpha=0.8, label='model $\mu \pm 2 \sigma$')
        plt.plot(x_train, y_hat_p2sd, color='green', alpha=0.8)
    else:
        plt.plot(x_train, y_hat, color='red', alpha=0.8)
        plt.plot(x_train, y_hat_m2sd, color='green', alpha=0.8)
        plt.plot(x_train, y_hat_p2sd, color='green', alpha=0.8)        
plt.legend()
plt.show()

# epistemic uncertainty: possible models vary
# aleatoric uncertainty: points don't lie exactly on the lines

# + [markdown] id="AkinPW72KaqU"
# ***
# <a id="coding_tutorial_4"></a>
# ## Reparameterization layers

# + id="OKiqoElQKaqV"
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Conv1D, MaxPooling1D, Flatten
from tensorflow.keras.losses import SparseCategoricalCrossentropy
from tensorflow.keras.optimizers import RMSprop
import os
import numpy as np
import matplotlib.pyplot as plt

# in the newer tfp versions, there are implementations of the convolutional variational layers:
# https://www.tensorflow.org/probability/api_docs/python/tfp/layers/conv_variational

# + [markdown] id="UVagTkYgKaqX"
# #### Load in the HAR dataset

# + [markdown] id="icl9TqZAKaqY"
# You'll be working with the [Human Activity Recognition (HAR) Using Smartphones](https://archive.ics.uci.edu/ml/datasets/Human+Activity+Recognition+Using+Smartphones) dataset. It consists of the readings from an accelerometer (which measures acceleration) carried by a human doing different activities. The six activities are walking horizontally, walking upstairs, walking downstairs, sitting, standing and laying down. The accelerometer is inside a smartphone, and, every 0.02 seconds (50 times per second), it takes six readings: linear and gyroscopic acceleration in the x, y and z directions. See [this link](https://archive.ics.uci.edu/ml/datasets/Human+Activity+Recognition+Using+Smartphones) for details and download. If you use it in your own research, please cite the following paper:
#
# - <NAME>, <NAME>, <NAME>, <NAME> and <NAME>. A Public Domain Dataset for Human Activity Recognition Using Smartphones. 21th European Symposium on Artificial Neural Networks, Computational Intelligence and Machine Learning, ESANN 2013. Bruges, Belgium 24-26 April 2013. 
#
# The goal is to use the accelerometer data to predict the activity.

# + [markdown] id="sLpIPEaEKe25"
# #### Import the data
#
# The dataset required for this coding tutorial can be downloaded from the following link:
#
# https://drive.google.com/file/d/1U_O3bhvuSAzQDKHGWcBIAO80iV9FVmDo/view?usp=sharing
#
# You should store this file in Drive for use in this Colab notebook.

# + id="15FeLLY8KfBU" colab={"base_uri": "https://localhost:8080/"} outputId="20ad1610-799f-47b4-d3fa-4648140ef743"
# Run this cell to connect to your Drive folder

from google.colab import drive
drive.mount('/content/gdrive')


# + id="Y_p1wpaTKaqY"
# Load the HAR dataset and create some data processing functions

# Function to load the data from file
def load_HAR_data():
    data_dir = '/content/gdrive/MyDrive/HAR.zip (Unzipped Files)/HAR'
    x_train = np.load(os.path.join(data_dir, 'x_train.npy'))[..., :6]
    y_train = np.load(os.path.join(data_dir, 'y_train.npy')) - 1
    x_test  = np.load(os.path.join(data_dir, 'x_test.npy'))[..., :6]
    y_test  = np.load(os.path.join(data_dir, 'y_test.npy')) - 1
    return (x_train, y_train), (x_test, y_test)

# Dictionary containing the labels and the associated activities
label_to_activity = {0: 'walking horizontally', 1: 'walking upstairs', 2: 'walking downstairs',
                     3: 'sitting', 4: 'standing', 5: 'laying'}

# Function to change integer labels to one-hot labels
def integer_to_onehot(data_integer):
    data_onehot = np.zeros(shape=(data_integer.shape[0], data_integer.max()+1))
    for row in range(data_integer.shape[0]):
        integer = int(data_integer[row])
        data_onehot[row, integer] = 1
    return data_onehot

# Load the data
(x_train, y_train), (x_test, y_test) = load_HAR_data()
y_train_oh = integer_to_onehot(y_train)
y_test_oh = integer_to_onehot(y_test)


# + id="UUybzk8lKaqa" colab={"base_uri": "https://localhost:8080/", "height": 353} outputId="04491115-5d3e-4c4e-83df-1d99f007cd7e"
# Inspect some of the data by making plots

def make_plots(num_examples_per_category):
    for label in range(6):
        x_label = x_train[y_train[:, 0] == label]
        for i in range(num_examples_per_category):
            fig, ax = plt.subplots(figsize=(10, 1))
            ax.imshow(x_label[100*i].T, cmap='Greys', vmin=-1, vmax=1)
            ax.axis('off')
            if i == 0:
                ax.set_title(label_to_activity[label])
            plt.show()
        
make_plots(1)

# + [markdown] id="XxhEIraGKaqb"
# #### 1D deterministic convolutional neural network

# + id="2CV7Sxy6Kaqc" colab={"base_uri": "https://localhost:8080/"} outputId="c2f5eafc-3522-43bc-8650-2c6bf1cb72df"
# Create standard deterministic model with:
# - Conv1D
# - MaxPooling
# - Flatten
# - Dense with Softmax

model_deterministic = Sequential([
    Conv1D(input_shape=(128, 6), filters=8, kernel_size=16, activation='relu'),
    MaxPooling1D(pool_size=16),
    Flatten(),
    Dense(units=6, activation='softmax')
])
model_deterministic.summary()

# + [markdown] id="rsUT-G8YKaqd"
# #### Probabilistic 1D convolutional neural network, with both weight and output uncertainty

# + id="euajafSYKaqe" colab={"base_uri": "https://localhost:8080/"} outputId="255bd278-a81d-4a37-84b7-d050997b45e7"
# Create probablistic model with the following layers:
#  - Conv1D (epistemic uncertainty modelling)
#  - MaxPooling
#  - Flatten
#  - Dense (epistemic uncertainty modelling)
#  - OneHotCategorical (aleatoric uncertainty modelling) 

# for that, tfp know the analytical solution
divergence_fn = lambda q, p, _ : tfd.kl_divergence(q, p) / x_train.shape[0]  # make it unbiased
# the reparameterization layers allow other loss functions as well

model_probabilistic = Sequential([
    # probabilistic version of 1D convolution layer
    tfpl.Convolution1DReparameterization(
        input_shape=(128, 6),  # size of the sensor outputs
        filters=8, kernel_size=16, activation='relu',  # same as previously
        kernel_prior_fn=tfpl.default_multivariate_normal_fn,  # has fixed parameters
        kernel_posterior_fn=tfpl.default_mean_field_normal_fn(is_singular=False),  # so it doesn't collapse into a deterministic (singular) parameter with just a mean an no (0) variance
        kernel_divergence_fn=divergence_fn,
        bias_prior_fn=tfpl.default_multivariate_normal_fn,  # has fixed parameters
        bias_posterior_fn=tfpl.default_mean_field_normal_fn(is_singular=False),  # so it doesn't collapse into a deterministic (singular) parameter with just a mean an no (0) variance
        bias_divergence_fn=divergence_fn
    ),
    # MaxPooling and Flatten don't have parameters, thus no probabilistic versions
    MaxPooling1D(pool_size=16),
    Flatten(),
    # probabilistic version of Dense layer
    tfpl.DenseReparameterization(
        units=tfpl.OneHotCategorical.params_size(6),  # number of categories
        activation=None,  # same as previously
        kernel_prior_fn=tfpl.default_multivariate_normal_fn,  # has fixed parameters
        kernel_posterior_fn=tfpl.default_mean_field_normal_fn(is_singular=False),  # so it doesn't collapse into a deterministic (singular) parameter with just a mean an no (0) variance
        kernel_divergence_fn=divergence_fn,
        bias_prior_fn=tfpl.default_multivariate_normal_fn,  # has fixed parameters
        bias_posterior_fn=tfpl.default_mean_field_normal_fn(is_singular=False),  # so it doesn't collapse into a deterministic (singular) parameter with just a mean an no (0) variance
        bias_divergence_fn=divergence_fn
    ),
    # probabilistic version of one hot encoded classification
    tfpl.OneHotCategorical(6)
])

model_probabilistic.summary()
# has the double number of parameters than the deterministic one (mean and stddev for every weight)

# + id="NbS5vk0_Kaqg"
# Replace analytical Kullback-Leibler divergence with approximated one
# (in case that tfp doesn't know the analytical solution for the chosen prior and posterior)

def kl_approx(q, p, q_tensor):
  # this equals to the Monte Carlo approximation of the KL divergence (just the mean of the integrand)
    return tf.reduce_mean(q.log_prob(q_tensor) - p.log_prob(q_tensor))

divergence_fn = lambda q, p, q_tensor : kl_approx(q, p, q_tensor) / x_train.shape[0]


# + id="qkwLmGHaKaqh"
# Compile the model using the negative loglikelihood

def nll(y_true, y_pred):
    return -y_pred.log_prob(y_true)

model_probabilistic.compile(loss=nll,
              optimizer=RMSprop(learning_rate=0.005),
              metrics=['accuracy'],
              experimental_run_tf_function=False)  # has to do with eager exec

# + id="ibkhBetoKaqj" colab={"base_uri": "https://localhost:8080/"} outputId="ad763a83-c58e-43d8-9164-bd9c91851e1c"
# Train the model

model_probabilistic.fit(x_train, y_train_oh, epochs=20, verbose=False)
model_probabilistic.evaluate(x_train, y_train_oh)
model_probabilistic.evaluate(x_test, y_test_oh)


# + [markdown] id="wtBXmPJIKaql"
# #### Inspect model performance

# + id="tBnw9p_tKaqm"
# Define function to analyse model predictions versus true labels

def analyse_model_predictions(image_num):

    # Show the accelerometer data
    print('------------------------------')
    print('Accelerometer data:')
    fig, ax = plt.subplots(figsize=(10, 1))
    ax.imshow(x_test[image_num].T, cmap='Greys', vmin=-1, vmax=1)
    ax.axis('off')
    plt.show()

    # Print the true activity
    print('------------------------------')
    print('True activity:', label_to_activity[y_test[image_num, 0]])
    print('')

    # Print the probabilities the model assigns
    print('------------------------------')
    print('Model estimated probabilities:')
    # Create ensemble of predicted probabilities
    predicted_probabilities = np.empty(shape=(200, 6))
    for i in range(200):
        predicted_probabilities[i] = model_probabilistic(x_test[image_num][np.newaxis, ...]).mean().numpy()[0]
    pct_2p5 = np.array([np.percentile(predicted_probabilities[:, i], 2.5) for i in range(6)])
    pct_97p5 = np.array([np.percentile(predicted_probabilities[:, i], 97.5) for i in range(6)])
    # Make the plots
    fig, ax = plt.subplots(figsize=(9, 3))
    bar = ax.bar(np.arange(6), pct_97p5, color='red')
    bar[y_test[image_num, 0]].set_color('green')
    bar = ax.bar(np.arange(6), pct_2p5-0.02, color='white', linewidth=1, edgecolor='white')
    ax.set_xticklabels([''] + [activity for activity in label_to_activity.values()],
                       rotation=45, horizontalalignment='right')
    ax.set_ylim([0, 1])
    ax.set_ylabel('Probability')
    plt.show()


# + id="ibhqPDcJKaqo" colab={"base_uri": "https://localhost:8080/", "height": 452} outputId="b77a7a16-ba4b-44c2-e64f-133092128774"
# call the analyze function with certain sample
analyse_model_predictions(image_num=79)

# here, this model is totally confident and also correct on the classification

# + id="dwg676VGKaqr" colab={"base_uri": "https://localhost:8080/", "height": 452} outputId="5e5c9a45-1530-4462-8334-f79c837484f2"
# call the analyze function with an uncertain sample
analyse_model_predictions(image_num=633)

# from the data it's hard to classify the activity (very little information)
# model is very uncertain about its classification 
# (assigns very different probabilities throughout the runs and also classifies different activities)

# + id="rL9SAg6bKaqu" colab={"base_uri": "https://localhost:8080/", "height": 452} outputId="a05ca3bd-a482-4e70-8e7d-33855501c89c"
# call the analyze function with an very uncertain sample
analyse_model_predictions(image_num=1137)

# model is extreme uncertain about its classification 
# (from 0-100% probability for multiple actions)
# -> it knows what it doesn't know (needs more data on what it sees)

# + id="ZFLYbjPabPUd"


StarcoderdataJupyterScriptsDedupFiltered