```python
from google.colab import drive
drive.mount('gdrive')
```

    Mounted at gdrive

```python
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.layers import Dense, Input, Flatten, Conv1D, AveragePooling1D, Dropout
from tensorflow.keras.layers import concatenate, Conv2D, MaxPooling2D, MaxPooling1D, BatchNormalization, Softmax
from tensorflow.keras.models import Model, Sequential
import pandas as pd
import numpy as np
%matplotlib inline
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
import seaborn as sns

import tensorflow_probability as tfp
tfd = tfp.distributions
tfpl = tfp.layers
```

# Building Probabilistic Deep Learning Models
- dealing with uncertainty involved in modeling process
- want models to assign higher levels of uncertainty to incorrect predictions
- two main sources of uncertainty: aleatoric and epistemic
  - aleatoric: uncertainty in the data itself (e.g. predicting coin toss; can't do it with certainty but can predict the probability)
  - epistemic: model uncertainty (if there is not enough data, deep learning models may not have the correct parameters)

## Maximum Likelihood estimation
Probability density (or mass, for discrete distributions) function:

$$
P(y | \theta) = \text{Prob} (\text{sampling value $y$ from a distribution with parameter $\theta$})
$$

Multiple samples:

$$
P(y_1, \ldots, y_n | \theta) = \prod_{i=1}^n P(y_i | \theta)
$$

These functions are used for **probability**, where you know the distribution and want to make deductions about possible values sampled from it.

The **likelihood** is the reverse where the samples are fixed and $\theta$ is considered the independent variable.  In this case, you know the samples from data collected but don't know $\theta$.

$$
L(\theta | y_1, \ldots, y_n) = P(y_1, \ldots, y_n | \theta)
$$

## Example: Bernoulli Distribution

$$
\begin{align}
L(\theta | y) &= \begin{cases}
1 - \theta \quad \text{if} \, y = 0 \\
\theta \qquad \, \, \, \text{if} \, y = 1 \\
\end{cases} \\
&= (1 - \theta)^{1 - y} \theta^y \quad y \in \{0,1\}
\end{align}
$$

Assuming samples are independent:

$$
L(\theta | y_1, \ldots, y_n) = \prod_{i=1}^n (1 - \theta)^{1 -y} \theta^y
$$

## Example: Normal Distribution

$$
L(\mu, \sigma | y) = \frac{1}{\sqrt{2 \pi \sigma^2}} \exp \Big( - \frac{1}{2 \sigma^2} (y - \mu)^2 \Big)
$$

Assuming samples are independent:

$$
L(\theta | y_1, \ldots, y_n) = \prod_{i=1}^n \frac{1}{\sqrt{2 \pi \sigma^2}} \exp \Big( - \frac{1}{2 \sigma^2} (y_i - \mu)^2 \Big)
$$

The likelihood is the same as the probability density function except it is a function of $\mu$ and $\sigma$ and the samples are considered constant.

## Maxium Likelihood Extimation

Usually we want to find maximum likelihood estimate $\theta_{MLE}$ which is the value of $\theta$ that maximizes $L(\theta | y_1, \ldots, y_n)$.

## Negative log-likelihood

The likelihood is a product (assuming independent observations) which can be turned into a sum by taking the log of the product.  Since the log function increases with its argument, maximizing the log-likelihood is equivalent to maximizing the likelihood.  Since the loss function must be a minimization problem, we consider the negative log-likelihood:

$$
\theta_{MLE} = \arg \min_{\theta} - \sum_{i=1}^n \log L(y_i | \theta)
$$

This is often used to find the best weights in a neural network where we assume that neural network predictions form part of a distribution that the true label is drawn from.  In the case of binary classifiers, sparse categorical crossentropy is used as the loss function and, in fact, it is exactly the negative log likelihood of a Bernoulli distribution.  In regression, minimizing the sum of squared errors is equivalent to minimizing the negative log-likelihood of a gaussian distribution (assuming Gaussian error term with constant variance).

## DistributionLambda Layer

```python
model = Sequential([
                    Dense(1, input_shape=(2,)),
                    tfpl.DistributionLambda(
                        lambda t: tfd.Normal(loc=t, scale=1), # output of previous layer describes mean of this normal distribution
                        convert_to_tensor_fn=tfd.Distribution.sample) # if you want to use output downstream of model, it
                                                                      # needs to be converted from distribution object into a tensor,
                                                                      # which can be done by sampling from distribution
])                                                                    # (could also be done with mean or mode)
# one required argument for Distribution lambda is a function that takes previous layer's output as input and returns distribution object
# this explicitly captures the uncertainty in the model by converting the output in the dense layer to a distribution

# output is now a distribution object instead of a scaler
# e.g. X_train --> shape (16,2) --> model(X_train) --> Normal dist with batch_shape (16,1) and event shape ()
# batch_shape (16,1) for number of samples (16) and number of outputs from previous dense layer (1)

def nll(y_true, y_pred):
  return - y_pred.log_prob(y_true) # remember that y_pred is now a distribution object

model.compile(loss=nll, optimizer='rmsprop')
# model.fit(X_train, y_train, epochs=10)

# model(X_test).sample() # get predictions by sampling from distribution
# model(X_test).mean()
```

Create a probabilistic model using the `DistributionLambda` layer whose first layer represents:
$$
y = sigmoid(x) = \frac{1}{1 + \exp (-x)}
$$

```python
# deterministic model
model = Sequential([
                    Dense(1, input_shape=(1,), activation='sigmoid',
                          kernel_initializer=tf.constant_initializer(1),
                          bias_initializer=tf.constant_initializer(0))
])
x = np.linspace(-5,5,100)
plt.scatter(x, model.predict(x), alpha=0.4)
plt.plot(x, 1/(1+np.exp(-x)), color='r', alpha=0.8)
plt.show()
```

```python
# deterministic model always returns same prediction for same x value
x = [0]
for _ in range(5):
  print(model.predict(x))
```

    [[0.5]]
    [[0.5]]
    [[0.5]]
    [[0.5]]
    [[0.5]]

```python
# probabilistic model
model = Sequential([
                    Dense(1, input_shape=(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)
])
x = np.linspace(-5,5,100)
plt.scatter(x, model.predict(x), alpha=0.4)
plt.plot(x, 1/(1+np.exp(-x)), color='r', alpha=0.8)
plt.show()
```

```python
x = [0]
for _ in range(5):
  print(model.predict(x))
```

    [[1]]
    [[1]]
    [[0]]
    [[0]]
    [[0]]

```python
X_train = np.linspace(-5,5,500)[:, np.newaxis]
y_train = model.predict(X_train)

fig, ax = plt.subplots(figsize=(5,5))
ax.scatter(X_train, y_train, alpha=0.04, color='b', label='samples')
ax.plot(X_train, model(X_train).mean().numpy().flatten(), color='red', alpha=0.8, label='mean')
ax.legend()
plt.show()
```

```python
# change initial weights/biases so we can train model
model_untrained = Sequential([
                    Dense(1, input_shape=(1,), activation='sigmoid',
                          kernel_initializer=tf.constant_initializer(2),
                          bias_initializer=tf.constant_initializer(2)),
                    tfpl.DistributionLambda(lambda t: tfd.Bernoulli(probs=t),
                                            convert_to_tensor_fn=tfd.Distribution.sample)
])
def nll(y_true, y_pred):
  return -y_pred.log_prob(y_true)

model_untrained.compile(loss=nll, optimizer=tf.keras.optimizers.RMSprop(0.01))

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_train, 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])
```

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

## Probabilistic Layers

```python
model = Sequential([
                    Dense(16, activation='relu', input_shape=(2,)),
                    Dense(2), # (batch_shape, 2)
                    tfpl.DistributionLambda(lambda t: tfd.Independent(tfd.Normal(loc=t[..., :1],
                                                                                scale=tf.math.softplus(t[...,1:]))))
                    # softplus is needed to make sure scale is positive
                    # ... is needed to account for batch_shapes of any rank
])
# alternatively:
model = Sequential([
                    Dense(16, activation='relu', input_shape=(2,)),
                    Dense(2),
                    tfpl.IndependentNormal(event_shape=(1)) # accomplishes exact same thing as above
                    # event shape of 1 means this layer requires two parameters to be defined: mean and std dev
                    # this is why the previous dense layer has two units
                    # it also passes std devs through softplus function to make sure they are positive
])
model = Sequential([
                    Dense(16, activation='relu', input_shape=(2,)),
                    Dense(4), # requires 4 now
                    tfpl.IndependentNormal(event_shape=(2)) # this is specifying 2D multivariate normal distribution so now need 4 inputs
                    # two means and two std devs of a covariance matrix diagonal
])
model = Sequential([
                    Dense(16, activation='relu', input_shape=(2,)),
                    Dense(8), # requires 8 now
                    tfpl.IndependentNormal(event_shape=([2,2]))
])
# potentially tricky to know exactly how many params are needed for a certain event shape
# use params_size method in previous dense layer to have it figure out for you
event_shape = 2
model = Sequential([
                    Dense(16, activation='relu', input_shape=(2,)),
                    Dense(tfpl.IndependentNormal.params_size(event_shape)),
                    tfpl.IndependentNormal(event_shape,
                                           convert_to_tensor_fn=tfd.Distribution.sample)
])
# for multi class classification problems, you can use OneHotCategorical layer instead of Dense layer with softmax at the end
num_classes = 10
model = Sequential([
                    Conv2D(16, 3, activation='relu', input_shape=(32,32,3)),
                    MaxPooling2D(3),
                    Flatten(),
                    Dense(64, activation='relu'),
                    Dense(tfpl.OneHotCategorical.params_size(num_classes)),
                    tfpl.OneHotCategorical(num_classes) # replaces Dense(num_classes, activation='softmax')
])
model.compile(loss=lambda y_true, y_pred: -y_pred.log_prob(y_true)) # y_true is one hot vector --> this returns nll of correct class
```

```python
# create 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)
```

```python
# deterministic linear regression with MSE loss
model = Sequential([
                    Dense(1, input_shape=(1,))
])
model.compile(loss=tf.keras.losses.MeanSquaredError(), optimizer=tf.keras.optimizers.RMSprop(0.005))
model.summary()
model.fit(X_train, y_train, epochs=200, verbose=False)

plt.scatter(X_train, y_train, alpha=0.4)
plt.plot(X_train, model.predict(X_train), color='red', alpha=0.8)
plt.show()
```

Disadvantage is that a deterministic model doesn't capture the uncertainty in predictions

```python
# deterministic linear regression with MSE loss
model = Sequential([
                    Dense(1, input_shape=(1,)),
                    tfpl.DistributionLambda(lambda t: tfd.Independent(tfd.Normal(loc=t, scale=1))) # parameterize mean of normal dist.
])
model.summary()
# same number of trainable params but output is now a distribution object
```

    Model: "sequential_11"
    _________________________________________________________________
    Layer (type)                 Output Shape              Param #
    =================================================================
    dense_17 (Dense)             (None, 1)                 2
    _________________________________________________________________
    distribution_lambda_4 (Distr multiple                  0
    =================================================================
    Total params: 2
    Trainable params: 2
    Non-trainable params: 0
    _________________________________________________________________

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

model.compile(loss=nll, optimizer=tf.keras.optimizers.RMSprop(0.005))
model.fit(X_train, y_train, epochs=200, verbose=False)
```

    <tensorflow.python.keras.callbacks.History at 0x7f06858b9a50>

```python
x = np.array([[0]])
y_model = model(x)
y_model
```

    <tfp.distributions.Independent 'sequential_11_distribution_lambda_4_Independentsequential_11_distribution_lambda_4_Normal' batch_shape=[1] event_shape=[1] dtype=float32>

```python
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.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()
```

Now let's learn mean *AND* standard deviation!

```python
event_shape = 1
model = Sequential([
                    Dense(tfpl.IndependentNormal.params_size(event_shape), input_shape=(1,)), # needs two units now or just use params_size
                    tfpl.IndependentNormal(event_shape) # instead of just mean, this parameterizes mean and std dev
])
model.summary()
```

    Model: "sequential_12"
    _________________________________________________________________
    Layer (type)                 Output Shape              Param #
    =================================================================
    dense_18 (Dense)             (None, 2)                 4
    _________________________________________________________________
    independent_normal_4 (Indepe multiple                  0
    =================================================================
    Total params: 4
    Trainable params: 4
    Non-trainable params: 0
    _________________________________________________________________

```python
model.compile(loss=nll, optimizer=tf.keras.optimizers.RMSprop(0.005))
model.fit(X_train, y_train, epochs=200, verbose=False)
```

    <tensorflow.python.keras.callbacks.History at 0x7f0696987990>

```python
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.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()
```

Much better fit! Now we have information on mean and uncertainty in a model prediction (compared to just mean for deterministic model)

```python
# let's also look at nonlinear models
X_train = np.linspace(-1,1,1000)[:, np.newaxis]
# not only nonlinear trend but the error is also a function of x so it isnt same for all x unlike previous example
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)
```

```python
event_shape = 1
model = Sequential([
                    Dense(8, input_shape=(1,), activation='sigmoid'),
                    Dense(tfpl.IndependentNormal.params_size(event_shape), input_shape=(1,)), # needs two units now or just use params_size
                    tfpl.IndependentNormal(event_shape) # instead of just mean, this parameterizes mean and std dev
])
model.compile(loss=nll, optimizer=tf.keras.optimizers.RMSprop(0.005))
model.summary()
model.fit(X_train, y_train, epochs=200, verbose=False)
```

    Model: "sequential_13"
    _________________________________________________________________
    Layer (type)                 Output Shape              Param #
    =================================================================
    dense_19 (Dense)             (None, 8)                 16
    _________________________________________________________________
    dense_20 (Dense)             (None, 2)                 18
    _________________________________________________________________
    independent_normal_5 (Indepe multiple                  0
    =================================================================
    Total params: 34
    Trainable params: 34
    Non-trainable params: 0
    _________________________________________________________________

    <tensorflow.python.keras.callbacks.History at 0x7f068727e510>

```python
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.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()
```

## Bayes by Backprop
Instead of the weights of each neuron being a single value, you can describe them with distributions. Bayes by Backprop is a method to determine the parameters of the distributions.  Mathematically complex, need to explore further.

## DenseVariational layer

```python
# need to define the prior distribution P(w) - our belief for the distribution of params, w, before seeing any data
# standard assumption is the prior is a spherical gaussian, i.e. an independent normal dist. for each weight and bias
def prior(kernel_size, bias_size, dtype=None): # prior distribution for a dense layer
  n = kernel_size + bias_size
  return lambda t: tfd.Independent(tfd.Normal(loc=tf.zeros(n, dtype=dtype), scale=1), reinterpreted_batch_ndims=1)
  # t is the input to the dense layer that this prior dist describes
  # note that t isnt in the lambda function so the dist is the same for any input
  # the prior does not have any trainable variables so it won't change during optimization

# need to define posterior distribution
def posterior(kernel_size, bias_size, dtype=None):
  n = kernel_size + bias_size
  return Sequential([
                     tfpl.VariableLayer(tfpl.IndependentNormal.params_size(n), dtype=dtype), # VariableLayer --> returns a tf variable
                     tfpl.IndependentNormal(n) # convert_to_tensor_fn is sampling by default (how to get weights for forward pass)
  ])
  # similar to the prior function, this returns a callable object that returns a distribution object but it is a model rather than layer

# rules for defining prior/posterior:
  # 1) inputs are kernel_size, bias_size, dtype
  # 2) they return callable objects that take in tensor as input and return distribution object

N = len(X_train)
model = Sequential([
                    tfpl.DenseVariational(16, posterior, prior, kl_weight=1/N, activation='relu', input_shape=(8,)),
                    tfpl.DenseVariational(2, posterior, prior, kl_weight=1/N),
                    tfpl.IndependentNormal(1)
])
# replace dense layers with DenseVariational layers where everything is the same except the addition of posterior/prior functions
# this layer automatically applies KL_divergence loss which can be weighted according to kl_weight param (1/N is a the best choice)
# objective is to maximize ELBO (evidence lower bound) which is the difference between two terms
  # 1) sum of expected log likelihood over the data (postive of nll loss function I have been using)
  # 2) KL divergence between posterior and prior
           # KL divergence measures similarity btw two functions so it will force the posterior to be similar to prior to minimize the loss
           # kind of like regularizing the posterior
# note that maximizing the ELBO requires a trade-off
# i.e. we want the posterior to maximize log likelihood but at the same time we don't want it to diverge too much from prior

model.compile(loss=lambda y_true, y_pred: -y_pred.log_prob(y_true)) # nll is estimated using a single sample of posteriors to get weights
# the KL_divergence term is automatically included in the DenseVariational layers
```

Note that the posterior is a "variational posterior" to approximate the true posterior because it is generally to difficult to calculate the true posterior using Bayes' Theorem. We tune the parameters of the variational posterior to get as close to the true posterior as possible.

```python
# same data as before with new model that accounts for uncertainty in weights
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)
```

```python
def prior(kernel_size, bias_size, dtype=None):
  n = kernel_size + bias_size
  prior_model = Sequential([
                            tfpl.DistributionLambda(
                                lambda t: tfd.MultivariateNormalDiag(loc=tf.zeros(n), scale_diag=tf.ones(n))
                            ) # prior is just unit gaussian for all parameters of a DenseVariational layer
  ]) # no trainable variables since the zeros and ones are hardcoded in
     # also doesn't use the input so we don't even need to use Sequential but might be convenient since posterior does need it
  return prior_model

# need to assume a distribution for posterior --> let's assume multivariate Gaussian
def posterior(kernel_size, bias_size, dtype=None):
  n = kernel_size + bias_size
  # we want trainable parameters, which we get using VariableLayer which are then passed to next layer
  posterior_model = Sequential([
                             tfpl.VariableLayer(tfpl.MultivariateNormalTriL.params_size(n), dtype=dtype),
                             tfpl.MultivariateNormalTriL(n)
  ])
  return posterior_model
```

```python
model = Sequential([
                    tfpl.DenseVariational(1, posterior, prior, input_shape=(1,), kl_weight=1/len(X_train), kl_use_exact=True)
])
# can use deterministic loss since the output of this model is not a distribution object
model.compile(loss=tf.keras.losses.MeanSquaredError(), optimizer=tf.keras.optimizers.RMSprop(0.005))
model.summary()
```

    Model: "sequential_15"
    _________________________________________________________________
    Layer (type)                 Output Shape              Param #
    =================================================================
    dense_variational_2 (DenseVa (None, 1)                 5
    =================================================================
    Total params: 5
    Trainable params: 5
    Non-trainable params: 0
    _________________________________________________________________

A deterministic linear regression has two trainable params (mean and intercept) but this has 5.  This is the mean and variance of mean and intercept (4) as well as the covariance between them.  Note that the MultivariateNormalTriL scale param is a lower triangular matrix which will have 3 non-zero params for a 2x2 matrix since the upper right corner is 0.  These three params are the variance and covariance described in the second sentence

```python
model.fit(X_train, y_train, epochs=500, verbose=False)
```

    <tensorflow.python.keras.callbacks.History at 0x7f068fdc6b50>

```python
dummy_input = np.array([[0]])
model_prior = model.layers[0]._prior(dummy_input)
model_posterior = model.layers[0]._posterior(dummy_input)
print('prior mean: ', model_prior.mean().numpy()) # did not change during training since it is not trainable
print('prior variance: ', model_prior.variance().numpy()) # did not change during training since it is not trainable
print('posterior mean: ', model_posterior.mean().numpy())
print('posterior covariance: ', model_posterior.covariance().numpy())
```

    prior mean:  [0. 0.]
    prior variance:  [1. 1.]
    posterior mean:  [ 1.0255601  -0.01181691]
    posterior covariance:  [[ 0.01487845 -0.00026619]
     [-0.00026619  0.00594118]]

```python
# plot an ensemble of linear regressions with weights sample from posterior distribution
plt.scatter(X_train, y_train, alpha=0.4)
for _ in range(10):
  y_model = model(X_train) # each time I call model on training data, a different sample is taken from posterior
  plt.plot(X_train, y_model, color='red', alpha=0.8)
plt.show()
```

```python
# let's also look at nonlinear models
X_train = np.linspace(-1,1,1000)[:, np.newaxis]
# not only nonlinear trend but the error is also a function of x so it isnt same for all x unlike previous example
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)
```

```python
# this model accounts for uncertainty in weights and also returns a distribution object
# allows us to model both epistemic and aleatoric uncertainty, respectively
model = Sequential([
                    tfpl.DenseVariational(8, posterior, prior, kl_weight=1/len(X_train), input_shape=(1,), activation='sigmoid'),
                    tfpl.DenseVariational(tfpl.IndependentNormal.params_size(1), posterior, prior, kl_weight=1/len(X_train)),
                    tfpl.IndependentNormal(1)
])
def nll(y_true, y_pred):
  return -y_pred.log_prob(y_true)

model.compile(loss=nll, optimizer=tf.keras.optimizers.RMSprop(0.005))
model.summary()
```

    Model: "sequential_16"
    _________________________________________________________________
    Layer (type)                 Output Shape              Param #
    =================================================================
    dense_variational_3 (DenseVa (None, 8)                 152
    _________________________________________________________________
    dense_variational_4 (DenseVa (None, 2)                 189
    _________________________________________________________________
    independent_normal_7 (Indepe multiple                  0
    =================================================================
    Total params: 341
    Trainable params: 341
    Non-trainable params: 0
    _________________________________________________________________

```python
model.fit(X_train, y_train, epochs=1000, verbose=False)
```

    <tensorflow.python.keras.callbacks.History at 0x7f068f626450>

```python
model.evaluate(X_train, y_train)
```

    32/32 [==============================] - 1s 1ms/step - loss: 0.0142

    0.01424924936145544

```python
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()
```

This allows us to model the uncertainty in the prediction by learning the variation as well as model the uncertainty in the parameters themselves, i.e. both aleatoric and epistemic uncertainty, respectively

## Reparameterization layers
Dense Variational layer is updated version of dense reparameterization layer - need to check if there is an updated version for convolutional layers (there wasn't when the video was filmed)

```python
# reparamaterization layers have same underlying algorithms and theoretical background as dense variational layer
model = Sequential([
                    tfpl.Convolution2DReparameterization(16, 3, activation='relu', input_shape=(28,28,1),
                                                         kernel_posterior_fn=tfpl.default_mean_field_normal_fn(),
                                                         kernel_prior_fn=tfpl.default_multivariate_normal_fn), # replaces Conv2D
                    MaxPooling2D(3),
                    Flatten(),
                    tfpl.DenseReparameterization(tfpl.OneHotCategorical.params_size(10)), # older version of Dense Variational layer
                    tfpl.OneHotCategorical(10)
])
# tfpl.default_mean_field_normal_fn() returns independent normal dist with trainable mean/stddev
# can also define a function to make a posterior/prior
def custom_multivariate_normal_fn(dtype, shape, name, trainable, add_variable_fn): # these are required inputs
  normal = tfd.Normal(loc=tf.zeros(shape, dtype), scale=2*tf.ones(shape, dtype))
  batch_ndims = tf.size(normal.batch_shape_tensor())
  return tfd.Independent(normal, reinterpreted_batch_ndims=batch_ndims) # must return distribution object

model = Sequential([
                    tfpl.Convolution2DReparameterization(16, 3, activation='relu', input_shape=(28,28,1),
                                                         kernel_posterior_fn=tfpl.default_mean_field_normal_fn(),
                                                         kernel_prior_fn=custom_multivariate_normal_fn), # replaces Conv2D
                    MaxPooling2D(3),
                    Flatten(),
                    tfpl.DenseReparameterization(tfpl.OneHotCategorical.params_size(10)), # older version of Dense Variational layer
                    tfpl.OneHotCategorical(10)
])
```

    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/base_layer.py:2191: UserWarning: `layer.add_variable` is deprecated and will be removed in a future version. Please use `layer.add_weight` method instead.
      warnings.warn('`layer.add_variable` is deprecated and '

```python
# load HAR dataset (smartphone accelerometer data while person walks)
def load_HAR_data():
  X_train = np.load('gdrive/MyDrive/Colab Notebooks/Datasets/HAR/x_train.npy')[...,:6]
  y_train = np.load('gdrive/MyDrive/Colab Notebooks/Datasets/HAR/y_train.npy') - 1
  X_test = np.load('gdrive/MyDrive/Colab Notebooks/Datasets/HAR/x_test.npy')[...,:6]
  y_test = np.load('gdrive/MyDrive/Colab Notebooks/Datasets/HAR/y_test.npy') - 1
  return (X_train, y_train), (X_test, y_test)

label_to_activity = {0:'walking horizontally', 1:'walking upstairs', 2:'walking downstairs',
                     3:'sitting', 4:'standing', 5:'laying'}

# change label to one-hot encoding

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

(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)
```

```python
# plot data
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)
      if i == 0:
        ax.set_title(label_to_activity[label])
      plt.show()

make_plots(1)
```

```python
# standard deterministic model:
model = Sequential([
                    Conv1D(input_shape=(128,6), filters=8, kernel_size=16, activation='relu'),
                    MaxPooling1D(16),
                    Flatten(),
                    Dense(6, activation='softmax')
])
model.summary()
```

    Model: "sequential_19"
    _________________________________________________________________
    Layer (type)                 Output Shape              Param #
    =================================================================
    conv1d (Conv1D)              (None, 113, 8)            776
    _________________________________________________________________
    max_pooling1d (MaxPooling1D) (None, 7, 8)              0
    _________________________________________________________________
    flatten_3 (Flatten)          (None, 56)                0
    _________________________________________________________________
    dense_21 (Dense)             (None, 6)                 342
    =================================================================
    Total params: 1,118
    Trainable params: 1,118
    Non-trainable params: 0
    _________________________________________________________________

```python
# convert to probabilistic model

divergence_fn = lambda q, p, _: tfd.kl_divergence(q,p) / len(X_train) # only works if analytical solution is known by tensorflow

model = Sequential([
    tfpl.Convolution1DReparameterization(
        input_shape=(128,6), filters=8, kernel_size=16, activation='relu',
        kernel_prior_fn=tfpl.default_multivariate_normal_fn,
        kernel_posterior_fn=tfpl.default_mean_field_normal_fn(is_singular=False),
        kernel_divergence_fn=divergence_fn,
        bias_prior_fn=tfpl.default_multivariate_normal_fn,
        bias_posterior_fn=tfpl.default_mean_field_normal_fn(is_singular=False),
        bias_divergence_fn=divergence_fn,
    ),
    MaxPooling1D(16),
    Flatten(),
    tfpl.DenseReparameterization(
        units=tfpl.OneHotCategorical.params_size(6), activation=None,
        kernel_prior_fn=tfpl.default_multivariate_normal_fn,
        kernel_posterior_fn=tfpl.default_mean_field_normal_fn(is_singular=False),
        kernel_divergence_fn=divergence_fn,
        bias_prior_fn=tfpl.default_multivariate_normal_fn,
        bias_posterior_fn=tfpl.default_mean_field_normal_fn(is_singular=False),
        bias_divergence_fn=divergence_fn,
    ),
    tfpl.OneHotCategorical(event_size=6)
])
model.summary() # twice as many trainable params since each param is replaced with mean/stddev
```

    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/base_layer.py:2191: UserWarning: `layer.add_variable` is deprecated and will be removed in a future version. Please use `layer.add_weight` method instead.
      warnings.warn('`layer.add_variable` is deprecated and '

    Model: "sequential_21"
    _________________________________________________________________
    Layer (type)                 Output Shape              Param #
    =================================================================
    conv1d_reparameterization_3  (None, 113, 8)            1552
    _________________________________________________________________
    max_pooling1d_3 (MaxPooling1 (None, 7, 8)              0
    _________________________________________________________________
    flatten_6 (Flatten)          (None, 56)                0
    _________________________________________________________________
    dense_reparameterization_4 ( (None, 6)                 684
    _________________________________________________________________
    one_hot_categorical_5 (OneHo multiple                  0
    =================================================================
    Total params: 2,236
    Trainable params: 2,236
    Non-trainable params: 0
    _________________________________________________________________

```python
def kl_approx(q, p, q_tensor): # Monte Carlo approximation of KL divergence
  return tf.reduce_mean(q.log_prob(q_tensor) - p.log_prob(q_tensor))

divergence_fn = lambda q, p, q_tensor: k1_approx(q, p, q_tensor) / len(X_train) # need to approximate k1_div if analytical solution unknown
```

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

model.compile(loss=nll, optimizer=tf.keras.optimizers.RMSprop(0.005), metrics='accuracy', experimental_run_tf_function=False)
# need experimental_run_tf_function=False when using reparameterization layers
```

```python
model.fit(X_train, y_train_oh, epochs=20, verbose=False)
model.evaluate(X_train, y_train_oh)
model.evaluate(X_test, y_test_oh)
```

    230/230 [==============================] - 1s 2ms/step - loss: 0.7120 - accuracy: 0.6878
    93/93 [==============================] - 0s 3ms/step - loss: 0.9826 - accuracy: 0.6607

    [0.9825744032859802, 0.6606718897819519]

```python
def analyze_model_predictions(image_num):
  # show accelerometer data
  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('True activity: ', label_to_activity[y_test[image_num, 0]])
  print('')

  # print probability model assigns
  print('Model estimated probs:')
  predicted_probs = np.empty(shape=(200,6))
  for i in range(200):
    predicted_probs[i] = model(X_test[image_num][np.newaxis, ...]).mean().numpy()[0]
  pct_2p5 = np.array([np.percentile(predicted_probs[:,i], 2.5) for i in range(6)])
  pct_97p5 = np.array([np.percentile(predicted_probs[:,i], 97.5) for i in range(6)])

  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, 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()
```

```python
analyze_model_predictions(21)
```

```python
analyze_model_predictions(1137)
```

```python

```


# Equações diferenciais com atrasos
Muito fenômenos naturais apresentam atrasos na relação entre causa e efeito. Se formos estritos preceberemos que poucas relações de causa e efeito em biologia podem ser consideradas instantâneas, Por exemplo: um predador leva tempo para digerir sua presa e convertê-la em novos predadores. Sob esta ótica, os modelos baseados em EDOs que vimos até agora, são meras aproximações às dinâmicas reais.

Técnicamente falando, os modelos EDO, partem do princípio que apena o conhecimento do presente e sua derivada é suficiente para determinar o estado imediatamente futuro do sistema. Equações diferencias com atrasos, incluem a dependência de outros pontos na história do sistema para a determinação dos estados futuros. Desta maneira os sistema pode ser escrito da seguinte forma:
$$\frac{dx}{dt} = F(t, x(t), x(t-\tau))$$

onde $\tau$ é o parâmetro de atraso.

## Equação Logística com atraso discreto
Como nosso primeiro exemplo. vamos examinar a equação de Hutchinson (1948):
$$\frac{dx(t)}{dt}=r x(t)\left[1-\frac{x(t-\tau)}{K}\right]$$

Modelos com atraso são usados muito menos frequentemente, devido à dificuldade na sua solução analítica. Felizmente dispomos de uma [biblioteca](https://github.com/Zulko/ddeint) em Python para sua solução numérica.

```python
%display typeset
import warnings
warnings.filterwarnings('always')
```

```python
from ddeint import ddeint
import numpy as np
```

Wright demonstrou em 1955, que se $r\tau<\frac{37}{24}$ e $x(0)>0$, $x(t) \rightarrow K$ quando $t\rightarrow \infty$. Mas se $r\tau > \frac{pi}{2}$, então a equação logística com atraso, apresenta uma solução periódica (ciclo limite), que oscila em torno de K.

Para resolver uma equação diferencial com atrasos, precisamos definir uma função de história que retorna o comportamento do sistema no intervalo $[-\tau,0]$, assumindo que o período da dinâmica em que estamos interessado começa em $t=0$.

```python
def logdelay(y,t,*pars):
    x=y(t)
    r,K,tau = pars
    xd = y(t-tau)
    return np.array([
        r*x*(1-(xd/K))
    ])
@interact
def plot_logistic(r=slider(1,3,0.1,pi/2),tau=slider(0,2,.2,1)):
    g = lambda t: 1 # valores de x(t) quando x pertence ao intervalo [-tau,0]
    tsteps = np.linspace(0,50,1000)
    show("r tau > pi/2? ", (r*tau)>(pi/2).n())
    K = 1000

    res=ddeint(logdelay,g, tsteps, (r,K,tau))
    P = list_plot(list(zip(tsteps,np.array(res))), legend_label='$x(t)$')
    l = plot(1000, (x,0,50), legend_label='K', color='green', linestyle=':')
    D = list_plot(list(zip(res[2:],res[:-2])), axes_labels=['$x(t)$','$x(t-2)$'])
#     show(D)
    show(graphics_array([(P+l,),(D,)]), figsize=[8,10])
```

    Interactive function <function plot_logistic at 0x7fbb44ec5048> with 2 widgets
      r: TransformFloatSlider(value…

## Análise de Estabilidade Linear

Os equilíbrios (x*) de um sistema dinâmico com atraso, de forma similar a EDOs, são pontos que satisfazem a seguinte condição:
$$f(x^*)=0$$

Para a análise da estabilidade dos equilíbrios, uma diferença é que em DDEs perturbações do sistema na vizinhaça destes equilíbrio levam a deslocamentos que são funções do tempo $\delta x(t)$. Seja $x_{\tau}=x(t-\tau)$,

```python

```

## Modelo predador-presa

Considere o seguinte modelo de Lotka-Volterra com atrasos:

\begin{align}
\frac{dx}{dt} &= x(r_1-a_{11}x_{\tau}-a_{12}y_{\tau}) \\
\frac{dy}{dt} &= y (-r_2 +a_{21}x_{\tau}-a_{22}y_{\tau})
\end{align}
onde $r_1$, $r_2$, $a_{11}$, $a_{12}$, $a_{21}$ e $a_{22}$ são constantes e o atraso $\tau \geq 0$, corresponte ao período de gestação do predador.

```python
def pred_presa(S,t,*pars):
    x,y = S(t)
    r1,r2,a11,a12,a21,a22,tau = pars
    xd,yd = S(t-tau)
    return np.array([
        x*(r1-a11*xd-a12*yd),
        y*(-r2 +a21*xd -a22*yd)
    ])

@interact
def show_sol(tau=slider(0,2.5,0.5,1)):
    g = lambda t: np.array([1,1])
    r1=1.1
    r2=1
    a11=1
    a12=1
    a21=2
    a22=1
    tsteps = np.linspace(0,100,1000)
    res = ddeint(pred_presa, g, tsteps,(r1,r2,a11,a12,a21,a22,tau))
    P = list_plot(list(zip(res[:,0],res[:,1])), plotjoined=True, axes_labels=['$x(t)$','$y(t)$'])
    show(P);
```

    Interactive function <function show_sol at 0x7fbb44d17e18> with 1 widget
      tau: TransformFloatSlider(value=1.0…

## Hematopoiesis
Considere o modelo de Mackey-Glass de uma dimensão.
$$\frac{dx(t)}{dt}=\frac{\beta x(t-\tau)}{1+x(t-\tau)^n}-\delta x(t)$$
onde $x(t)$ é o número de células sanguíneas no tempo $t$, $\tau$ é um atraso constante e $n$ uma constante. O primeiro termo da equação representa a produção atrasada de células e $delta$ a taxa de mortalidade das células.

```python
import pylab as plt
```

```python
def hematopoiesis(S,t,*pars):
    x = S(t)
    beta, n,delta, tau = pars
    xd = S(t-tau)
    return np.array([
        (beta*xd)/(1+xd**n)-delta*x
    ])

@interact
def show_sol(beta=2,n=10, tau=2,delta=0.7):
    g = lambda t: 0.5
    tsteps = np.linspace(0,512,4096)
    res = ddeint(hematopoiesis, g, tsteps,(beta,n,delta,tau))
    res = res[200:]
    P = list_plot(list(zip(res[tau:],res[:-tau])), plotjoined=True, axes_labels=['$x(t)$',r'$x_{\tau}$'])
    Q = list_plot(list(zip(tsteps[-256:],res[-256:])), plotjoined=True, axes_labels=['$t$',r'$x(t)$'])

    show(graphics_array([[P],[Q]]), figsize=[8,10]);
```

    Interactive function <function show_sol at 0x7fbb3eaab8c8> with 4 widgets
      beta: IntSlider(value=2, descripti…

```python
g = lambda t: 0.5
tsteps = np.linspace(0,512,2048)
res = ddeint(hematopoiesis, g, tsteps,(2,10,0.5,2))
```

```python
plt.psd(res, noverlap=128);
```

```python
def drawbif(l,u):
    g = lambda t: 0.5
    tsteps = np.linspace(0,256,768)

    pts = []
    for delta in np.linspace(l,u,50):
        res = ddeint(hematopoiesis, g, tsteps,(2,10,delta,2))
        res = np.array([res[0]]+[i[0] for i in res[1:]])
        dres = np.diff(res)
        pts += [(delta,x) for i,x in enumerate(res[-128:]) if dres[i] <=1e-15]

    show(points(pts),axes_labels=[r'$\delta$','$x(t)$'],gridlines=True, figsize=[10,8])

drawbif(0,1.8)
```

```python

```

```python

```

Planilhas Sage/Modelos Dinamicos com atrasos.ipynb


```python
%matplotlib inline

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import sklearn.linear_model as linear_model

np.random.seed(1337)
from timeit import timeit

import matplotlib
import matplotlib.pyplot as plt
```

### Well-conditioned linear regression

```python
%matplotlib inline

import warnings
warnings.filterwarnings('ignore')

import numpy as np
import scipy as sp
import matplotlib.pyplot as plt
import random
from scipy import stats
from scipy.optimize import fmin

from matplotlib import rc
#rc('font',**{'family':'sans-serif','sans-serif':['Helvetica']})
## for Palatino and other serif fonts use:
rc('font',**{'family':'serif','serif':['Palatino']})
rc('text', usetex=True)

from PIL import Image

import random
from numpy import linalg as la
from scipy.sparse.linalg import svds, eigs
import math

p = 1000  # Ambient dimension
n = 1000  # Number of samples

# Generate a p-dimensional zero vector
x_star = np.random.randn(p)
# Normalize
x_star = (1 / la.norm(x_star, 2)) * x_star

# Well-conditioned matrix A
A = np.random.randn(n, p)
[U, S, V] = la.svd(A)
S = np.linspace(1, 5, n)
S = np.diag(S)
A = U.dot(S.dot(V.T))

y = A.dot(x_star)

# Print condition number of A
print(la.cond(A))

# Compute largest singular value of A
U,S,V = la.svd(A)
L = max(S)**2
mu = min(S)**2

# eta = 1/s**2
# print(eta)
```

    5.000000000000003

```python
def GD(y, A, eta, iters, epsilon, verbose, x_star):
    p = A.shape[1]   # Length of original signal
    n = A.shape[0]   # Length of measurement vector

    x_new = np.zeros(p)    # Initial estimate
    At = np.transpose(A)   # Transpose of A

    x_list, f_list = [1], [f(y, A, x_new)]

    for i in range(iters):
        x_old = x_new

        # Compute gradient
        grad = -At @ (y - A @ x_new)

        # Perform gradient step
        x_new = x_old - eta * grad

        if (la.norm(x_new - x_old, 2) / la.norm(x_new, 2)) < epsilon:
            break

        # Keep track of solutions and objective values
        x_list.append(la.norm(x_new - x_star, 2))
        f_list.append(f(y, A, x_new))

        if verbose:
            print("iter# = "+ str(i) + ", ||x_new - x_old||_2 = " + str(la.norm(x_new - x_old, 2)))

    print("Number of steps:", len(f_list))
    return x_new, x_list, f_list

def f(y, A, x):
    return 0.5 * math.pow(la.norm(y - A @ x, 2), 2)
```

```python
def HB(y, A, eta, beta, iters, epsilon, verbose, x_star):
    p = A.shape[1]   # Length of original signal
    n = A.shape[0]   # Length of measurement vector

    x_new = np.zeros(p)    # Initial estimate
    x_prev = np.zeros(p)   # Initial estimate
    x_cur = np.zeros(p)    # Initial estimate
    At = np.transpose(A)   # Transpose of A

    x_list, f_list = [1], [f(y, A, x_new)]

    for i in range(iters):
        x_prev = x_cur
        x_cur = x_new

        # Compute gradient
        grad = -At @ (y - A @ x_cur)

        # Perform gradient step
        x_new = x_cur - eta * grad + beta * (x_cur - x_prev)

        if (la.norm(x_new - x_cur, 2) / la.norm(x_new, 2)) < epsilon:
            break

        # Keep track of solutions and objective values
        x_list.append(la.norm(x_new - x_star, 2))
        f_list.append(f(y, A, x_new))

        if verbose:
            print("iter# = "+ str(i) + ", ||x_new - x_old||_2 = " + str(la.norm(x_new - x_cur, 2)))

    print("Number of steps:", len(f_list))
    return x_new, x_list, f_list
```

```python
# Run algorithm
epsilon = 1e-6                # Precision parameter
iters = 1000

# # Setting 1
# eta_GD = 2/(L + mu)
# eta_HB = 4/(np.sqrt(L) + np.sqrt(mu))**2
# beta = max(np.abs(1 - np.sqrt(eta_HB * mu)), np.abs(1 - np.sqrt(eta_HB * L)))**2

# Setting 2
eta_GD = 1/L
eta_HB = 1/L
beta = max(np.abs(1 - np.sqrt(eta_HB * mu)), np.abs(1 - np.sqrt(eta_HB * L)))**2

# #Setting 3
# eta_GD = 1/L
# eta_HB = 1/L
# beta = 1

# #Setting 4
# eta_GD = 0.1
# eta_HB = 0.1
# beta = 0.9

x_GD, x_list_GD, f_list_GD = GD(y, A, eta_GD, iters, epsilon, True, x_star)
x_HB, x_list_HB, f_list_HB = HB(y, A, eta_HB, beta, iters, epsilon, True, x_star)

# Plot
plt.rc('text', usetex=True)
plt.rc('font', family='serif')

xs_GD = range(len(x_list_GD))
xs_HB = range(len(x_list_HB))

plt.plot(xs_GD, x_list_GD, '-o', color = '#3399FF', linewidth = 4, alpha = 0.7, markerfacecolor = 'b')
plt.plot(xs_HB, x_list_HB, '-o', color = '#FF6666', linewidth = 4, alpha = 0.7, markerfacecolor = 'r')
plt.yscale('log')
plt.xlabel('Iterations')
plt.ylabel(r"$\|x^\star - \widehat{x}\|_2$")

# Make room for the ridiculously large title.
plt.subplots_adjust(top=0.8)
plt.show()

print(beta)
```

### Nesterov acceleration

+ Acceleration #1:

\begin{align}
\beta = \frac{\theta_t - 1}{\theta_{t+1}}, ~~\text{where}~~ \theta_0 = 1, ~\theta_{t+1} = \frac{1 + \sqrt{1 + 4\theta_t^2}}{2}
\end{align}

```python
def Nesterov1(y, A, eta, iters, epsilon, verbose, x_star):
    p = A.shape[1]   # Length of original signal
    n = A.shape[0]   # Length of measurement vector

    x_prev = np.zeros(p)   # Initial estimate
    x_cur = np.zeros(p)    # Initial estimate
    y_cur = np.zeros(p)
    At = np.transpose(A)   # Transpose of A

    theta_cur = 1

    x_list, f_list = [1], [f(y, A, x_cur)]

    for i in range(iters):
        x_prev = x_cur

        # Compute gradient
        grad = -At @ (y - A @ y_cur)

        # Perform gradient step
        x_cur = y_cur - eta * grad

        # Perform the extrapolation step
        theta_prev = theta_cur
        theta_cur = (1 + np.sqrt(1 + 4*theta_prev**2))/2
        beta = (theta_prev - 1)/theta_cur

        y_cur = x_cur + beta * (x_cur - x_prev)

        if (la.norm(x_cur - x_prev, 2) / la.norm(x_cur, 2)) < epsilon:
            break

        # Keep track of solutions and objective values
        x_list.append(la.norm(x_cur - x_star, 2))
        f_list.append(f(y, A, x_cur))

        if verbose:
            print("iter# = "+ str(i) + ", ||x_new - x_old||_2 = " + str(la.norm(x_cur - x_prev, 2)))

    print("Number of steps:", len(f_list))
    return x_cur, x_list, f_list
```

```python
# Run algorithm
epsilon = 1e-6                # Precision parameter
iters = 1000

# # Setting 1
# eta_GD = 1/L
# eta_HB = 4/(np.sqrt(L) + np.sqrt(mu))**2
# beta = max(np.abs(1 - np.sqrt(eta_HB * mu)), np.abs(1 - np.sqrt(eta_HB * L)))**2
# eta_Nest = 1/L

# # Setting 2
# eta_GD = 1/L
# eta_HB = 1/L
# eta_Nest = 1/L
# beta = max(np.abs(1 - np.sqrt(eta_HB * mu)), np.abs(1 - np.sqrt(eta_HB * L)))**2

# Setting 3
eta_GD = 1/L
eta_HB = 1/L
eta_Nest = 1/L
beta = 0.9

# # Setting 4
# eta_GD = 0.1
# eta_HB = 0.1
# eta_Nest = 0.1
# beta = 0.9

x_GD, x_list_GD, f_list_GD = GD(y, A, eta_GD, iters, epsilon, True, x_star)
x_HB, x_list_HB, f_list_HB = HB(y, A, eta_HB, beta, iters, epsilon, True, x_star)
x_Nest1, x_list_Nest1, f_list_Nest1 = Nesterov1(y, A, eta_Nest, iters, epsilon, True, x_star)

# Plot
plt.rc('text', usetex=True)
plt.rc('font', family='serif')

xs_GD = range(len(x_list_GD))
xs_HB = range(len(x_list_HB))
xs_Nest1 = range(len(x_list_Nest1))

plt.plot(xs_GD, x_list_GD, '-', color = '#3399FF', linewidth = 2, alpha = 0.7)
plt.plot(xs_HB, x_list_HB, '-', color = '#FF6666', linewidth = 2, alpha = 0.7)
plt.plot(xs_Nest1, x_list_Nest1, '-', color = '#669900', linewidth = 2, alpha = 0.7)
plt.yscale('log')
plt.xlabel('Iterations')
plt.ylabel(r"$\|x^\star - \widehat{x}\|_2$")

# Make room for the ridiculously large title.
plt.subplots_adjust(top=0.8)
plt.show()

print(beta)
```

#### Some observations

+ The curve has ripples: Actually, Nesterov's method is not a descent method (i.e., it does not guarantee that *per iteration* the distance to the optimum (and the function values) should drop)
+ Nesterov's method is hyper-parameter free: you don't have to worry about how to set up the momentum parameter.
+ It actually works (even in theory) beyond strongly convex functions: we presented a theorem for the Heavy Ball method that shows acceleration, but assumes strong convexity (as well as knowing many parameters such as Lipschitz gradient continuity, strong convexity). Nesterov's method requires only the knowledge of the Lipschitz gradient continuity (which is the easiest to learn - why?)

```python
def Nesterov2(y, A, eta, iters, epsilon, verbose, x_star):
    p = A.shape[1]   # Length of original signal
    n = A.shape[0]   # Length of measurement vector

    x_prev = np.zeros(p)   # Initial estimate
    x_cur = np.zeros(p)    # Initial estimate
    y_cur = np.zeros(p)
    At = np.transpose(A)   # Transpose of A

    x_list, f_list = [1], [f(y, A, x_cur)]

    for i in range(iters):
        x_prev = x_cur

        # Compute gradient
        grad = -At @ (y - A @ y_cur)

        # Perform gradient step
        x_cur = y_cur - eta * grad

        # Perform the extrapolation step
        beta = i / (i + 3)

        y_cur = x_cur + beta * (x_cur - x_prev)

        if (la.norm(x_cur - x_prev, 2) / la.norm(x_cur, 2)) < epsilon:
            break

        # Keep track of solutions and objective values
        x_list.append(la.norm(x_cur - x_star, 2))
        f_list.append(f(y, A, x_cur))

        if verbose:
            print("iter# = "+ str(i) + ", ||x_new - x_old||_2 = " + str(la.norm(x_cur - x_prev, 2)))

    print("Number of steps:", len(f_list))
    return x_cur, x_list, f_list
```

```python
# Run algorithm
epsilon = 1e-6                # Precision parameter
iters = 1000

# Setting 1
eta_GD = 2/(L + mu)
eta_HB = 4/(np.sqrt(L) + np.sqrt(mu))**2
beta = max(np.abs(1 - np.sqrt(eta_HB * mu)), np.abs(1 - np.sqrt(eta_HB * L)))**2
eta_Nest = 1/L

# Setting 2
eta_GD = 1/L
eta_HB = 1/L
eta_Nest = 1/L
beta = max(np.abs(1 - np.sqrt(eta_HB * mu)), np.abs(1 - np.sqrt(eta_HB * L)))**2

# # Setting 3
# eta_GD = 1/L
# eta_HB = 1/L
# eta_Nest = 1/L
# beta = 0.9

# # Setting 4
# eta_GD = 0.1
# eta_HB = 0.1
# eta_Nest = 0.1
# beta = 0.9

x_GD, x_list_GD, f_list_GD = GD(y, A, eta_GD, iters, epsilon, True, x_star)
x_HB, x_list_HB, f_list_HB = HB(y, A, eta_HB, beta, iters, epsilon, True, x_star)
eta_GD = 1/L
x_Nest1, x_list_Nest1, f_list_Nest1 = Nesterov1(y, A, eta_Nest, iters, epsilon, True, x_star)
x_Nest2, x_list_Nest2, f_list_Nest2 = Nesterov2(y, A, eta_Nest, iters, epsilon, True, x_star)

# Plot
plt.rc('text', usetex=True)
plt.rc('font', family='serif')

xs_GD = range(len(x_list_GD))
xs_HB = range(len(x_list_HB))
xs_Nest1 = range(len(x_list_Nest1))
xs_Nest2 = range(len(x_list_Nest2))

plt.plot(xs_GD, x_list_GD, '-', color = '#3399FF', linewidth = 2, alpha = 0.7)
plt.plot(xs_HB, x_list_HB, '-', color = '#FF6666', linewidth = 2, alpha = 0.7)
plt.plot(xs_Nest1, x_list_Nest1, '-', color = '#669900', linewidth = 2, alpha = 0.7)
plt.plot(xs_Nest2, x_list_Nest2, '-', color = '#660099', linewidth = 2, alpha = 0.7)
plt.yscale('log')
plt.xlabel('Iterations')
plt.ylabel(r"$\|x^\star - \widehat{x}\|_2$")

# Make room for the ridiculously large title.
plt.subplots_adjust(top=0.8)
plt.show()

print(beta)
```

```python
def Nesterov3(y, A, eta, L, mu, iters, epsilon, verbose, x_star):
    p = A.shape[1]   # Length of original signal
    n = A.shape[0]   # Length of measurement vector

    x_prev = np.zeros(p)   # Initial estimate
    x_cur = np.zeros(p)    # Initial estimate
    y_cur = np.zeros(p)
    At = np.transpose(A)   # Transpose of A

    beta = (np.sqrt(L) - np.sqrt(mu))/(np.sqrt(L) + np.sqrt(mu))

    x_list, f_list = [1], [f(y, A, x_cur)]

    for i in range(iters):
        x_prev = x_cur

        # Compute gradient
        grad = -At @ (y - A @ y_cur)

        # Perform gradient step
        x_cur = y_cur - eta * grad

        # Perform the extrapolation step
        y_cur = x_cur + beta * (x_cur - x_prev)

        if (la.norm(x_cur - x_prev, 2) / la.norm(x_cur, 2)) < epsilon:
            break

        # Keep track of solutions and objective values
        x_list.append(la.norm(x_cur - x_star, 2))
        f_list.append(f(y, A, x_cur))

        if verbose:
            print("iter# = "+ str(i) + ", ||x_new - x_old||_2 = " + str(la.norm(x_cur - x_prev, 2)))

    print("Number of steps:", len(f_list))
    return x_cur, x_list, f_list
```

```python
# Run algorithm
epsilon = 1e-6                # Precision parameter
iters = 1000

# Setting 1
eta_GD = 2/(L + mu)
eta_HB = 4/(np.sqrt(L) + np.sqrt(mu))**2
beta = max(np.abs(1 - np.sqrt(eta_HB * mu)), np.abs(1 - np.sqrt(eta_HB * L)))**2
eta_Nest = 2/(L + mu)

# # Setting 2
# eta_GD = 1/L
# eta_HB = 1/L
# eta_Nest = 1/L
# beta = max(np.abs(1 - np.sqrt(eta_HB * mu)), np.abs(1 - np.sqrt(eta_HB * L)))**2

# # Setting 3
# eta_GD = 1/L
# eta_HB = 1/L
# eta_Nest = 1/L
# beta = 0.9

# # Setting 4
# eta_GD = 0.1
# eta_HB = 0.1
# eta_Nest = 0.1
# beta = 0.9

x_GD, x_list_GD, f_list_GD = GD(y, A, eta_GD, iters, epsilon, True, x_star)
x_HB, x_list_HB, f_list_HB = HB(y, A, eta_HB, beta, iters, epsilon, True, x_star)
eta_GD = 1/L
x_Nest1, x_list_Nest1, f_list_Nest1 = Nesterov1(y, A, eta_Nest, iters, epsilon, True, x_star)
x_Nest2, x_list_Nest2, f_list_Nest2 = Nesterov2(y, A, eta_Nest, iters, epsilon, True, x_star)
x_Nest3, x_list_Nest3, f_list_Nest3 = Nesterov3(y, A, eta_Nest, L, mu, iters, epsilon, True, x_star)

# Plot
plt.rc('text', usetex=True)
plt.rc('font', family='serif')

xs_GD = range(len(x_list_GD))
xs_HB = range(len(x_list_HB))
xs_Nest1 = range(len(x_list_Nest1))
xs_Nest2 = range(len(x_list_Nest2))
xs_Nest3 = range(len(x_list_Nest3))

plt.plot(xs_GD, x_list_GD, '-', color = '#3399FF', linewidth = 2, alpha = 0.7)
plt.plot(xs_HB, x_list_HB, '-', color = '#FF6666', linewidth = 2, alpha = 0.7)
plt.plot(xs_Nest1, x_list_Nest1, '-', color = '#669900', linewidth = 2, alpha = 0.7)
plt.plot(xs_Nest2, x_list_Nest2, '-', color = '#660099', linewidth = 2, alpha = 0.7)
plt.plot(xs_Nest3, x_list_Nest3, '-', color = '#666600', linewidth = 2, alpha = 0.7)
plt.yscale('log')
plt.xlabel('Iterations')
plt.ylabel(r"$\|x^\star - \widehat{x}\|_2$")

# Make room for the ridiculously large title.
plt.subplots_adjust(top=0.8)
plt.show()

print(beta)
```

### Noise tolerance in acceleration

Let us first describe a very simple case (actually we have already described a very similar case above) on top of which we will build some intuition for the result we discussed just now.
Consider the minimization of a smooth convex function
$f\colon \mathbb{R}^n\rightarrow\mathbb{R}$ defined as:

$$
\displaystyle f(x) = \frac{1}{2} x^\top A x - b^\top x.
$$
Here, $A$ is a positive semidefinite matrix with $n \times n$ dimensions, and $b \in \mathbb{R}^n$, both given.
One can easy compute that:
$$
\nabla f(x) = Ax - b
$$

For the purpose of this example, we will consider the case of a Laplacian cycle graph:
$$
A = \begin{bmatrix} 2 & -1 & 0 & 0 & 0 & \cdots & -1 \\
                   -1 & 2 & -1 & 0 & 0 & \cdots & 0 \\
                   0 & -1 & 2 & -1 & 0 & \cdots & 0 \\
                   \vdots & & & \ddots & & \vdots & \vdots
    \end{bmatrix}
$$

By definition, since $A$ is PSD, the problem is convex.
Moreover, by definition of $A$ it is not hard to see that the function satisfies Lipschitz gradient continuity with constaint $L = 4$.

Let us do the following experiment then:

```python
import matplotlib.pyplot as plt
import numpy as np

def GD(init, grad, L = 1, iters = 100):
    x_cur = init
    x_list = [init]
    for t in range(0, iters):
        x_cur = x_cur - (1/L)*grad(x_cur)
        x_list.append(x_cur)
    return x_list

def AccGD(init, grad, L = 1, iters = 100):
    x_cur = init
    y_cur = init
    x_list = [init]
    theta = 1
    for t in range(0, iters):
        theta = 0.5 * (1.0 + np.sqrt(1+4.0*theta**2))
        theta2 = 0.5 * (1.0 + np.sqrt(1+4.0*theta**2))
        beta = (theta - 1)/theta2
        y2 = x_cur - (1/L) * grad(x_cur)
        x_cur = y2 + beta * (y2 - y_cur)
        y_cur = y2
        x_list.append(x_cur)
    return x_list
```

```python
n = 100
B = np.zeros((n,n))
for i in range(1,n-1):
    B[i,i+1] = 1
    B[i,i-1] = 1
B[0,1] = 1
B[n-1,n-2] = 1
A = 2.0*np.eye(n) - B
b = np.zeros(n)
b[0] = 1
opt = np.dot(np.linalg.pinv(A),b)

def f(x):
    return 0.5*np.dot(x, np.dot(A,x)) - np.dot(x,b)

def grad(x):
    return np.dot(A,x) - b

f_star = f(opt)

iters = 5000

xs3 = GD(np.zeros(n), grad, 4, iters)
ys3 = [ abs(f(xs3[i])-f_star) for i in range(0, iters) ]
xs3acc = AccGD(np.zeros(n), grad, 4, iters)
ys3acc = [ abs(f(xs3acc[i])-f_star) for i in range(0, iters) ]

plt.yscale('log')
plt.plot(range(0,iters), ys3, range(0,iters), ys3acc)
```

```python
def noisygrad(x):
    return np.dot(A,x) - b + np.random.normal(0,1,(n))

iters = 500

xs3 = GD(np.zeros(n), noisygrad, 4, iters)
ys3 = [ abs(f(xs3[i])-f_star) for i in range(0, iters) ]
xs3acc = AccGD(np.zeros(n), noisygrad, 4, iters)
ys3acc = [ abs(f(xs3acc[i])-f_star) for i in range(0, iters) ]

plt.yscale('linear')
plt.plot(range(0, iters), ys3, range(0, iters), ys3acc)
```

```python

```


```python
"""
3D version of approx2D.py.
Numerical integration only.
"""
import sympy as sym
import numpy as np
import scipy.integrate

def least_squares(f, psi, Omega):
    """
    Given a function f(x,y,z) on a rectangular domain
    Omega=[[xmin,xmax],[ymin,ymax],[zmin,zmax]],
    return the best approximation to f in the space V
    spanned by the functions in the list psi.
    f and psi are symbolic (sympy) expressions, but will
    be converted to numeric functions for faster integration.
    """
    N = len(psi) - 1
    A = np.zeros((N+1, N+1))
    b = np.zeros(N+1)
    x, y, z = sym.symbols('x y z')
    f = sym.lambdify([x, y, z], f, modules='numpy')
    psi_sym = psi[:]  # take a copy, needed for forming u later
    psi = [sym.lambdify([x, y, z], psi[i]) for i in range(len(psi))]

    print('...evaluating matrix...')
    for i in range(N+1):
        for j in range(i, N+1):
            print(('(%d,%d)' % (i, j)))

            integrand = lambda x, y, z: psi[i](x,y,z)*psi[j](x,y,z)
            I, err = scipy.integrate.nquad(
                integrand,
                [[Omega[0][0], Omega[0][1]],
                 [Omega[1][0], Omega[1][1]],
                 [Omega[2][0], Omega[2][1]]])
            A[i,j] = A[j,i] = I
        integrand = lambda x, y, z: psi[i](x,y,z)*f(x,y,z)
        I, err = scipy.integrate.nquad(
            integrand,
            [[Omega[0][0], Omega[0][1]],
             [Omega[1][0], Omega[1][1]],
             [Omega[2][0], Omega[2][1]]])
        b[i] = I
    print()
    c = np.linalg.solve(A, b)
    if N <= 10:
        print(('A:\n', A, '\nb:\n', b))
        print(('coeff:', c))
    u = sum(c[i]*psi_sym[i] for i in range(len(psi_sym)))
    print(('approximation:', u))
    return u, c

def sine_basis(Nx, Ny, Nz):
    """
    Compute basis sin((p+1)*pi*x)*sin((q+1)*pi*y)*sin((r+1)*pi*z),
    p=0,...,Nx, q=0,...,Ny, r=0,...,Nz.
    """
    x, y, z = sym.symbols('x y z')
    psi = []
    for r in range(0, Nz+1):
        for q in range(0, Ny+1):
            for p in range(0, Nx+1):
                s = sym.sin((p+1)*sym.pi*x)*\
                    sym.sin((q+1)*sym.pi*y)*sym.sin((r+1)*sym.pi*z)
                psi.append(s)
    return psi

def test_least_squares():
    # Use sine functions
    x, y, z = sym.symbols('x y z')
    N = 1  # (N+1)**3 = 8 basis functions
    psi = sine_basis(N, N, N)
    f_coeff = [0]*len(psi)
    f_coeff[3] = 2
    f_coeff[4] = 3
    f = sum(f_coeff[i]*psi[i] for i in range(len(psi)))
    # Check that u exactly reproduces f
    u, c = least_squares(f, psi, Omega=[[0,1], [0,1], [0,1]])
    diff = np.abs(np.array(c) - np.array(f_coeff)).max()
    print(('diff:', diff))
    tol = 1E-15
    assert diff < tol

if __name__ == '__main__':
    import time
    t0 = time.clock()
    test_least_squares()

```

    C:\Users\TURBULENCE\anaconda3\lib\site-packages\ipykernel_launcher.py:87: DeprecationWarning: time.clock has been deprecated in Python 3.3 and will be removed from Python 3.8: use time.perf_counter or time.process_time instead

    ...evaluating matrix...
    (0,0)
    (0,1)
    (0,2)
    (0,3)
    (0,4)
    (0,5)
    (0,6)
    (0,7)
    (1,1)
    (1,2)
    (1,3)
    (1,4)
    (1,5)
    (1,6)
    (1,7)
    (2,2)
    (2,3)
    (2,4)
    (2,5)
    (2,6)
    (2,7)
    (3,3)
    (3,4)
    (3,5)
    (3,6)
    (3,7)
    (4,4)
    (4,5)
    (4,6)
    (4,7)
    (5,5)
    (5,6)
    (5,7)
    (6,6)
    (6,7)
    (7,7)

    ('A:\n', array([[ 1.25000000e-01,  6.10673003e-18,  6.59062874e-18,
             4.63225289e-20,  1.03621141e-17,  1.64513355e-20,
             6.65864441e-19,  6.55298598e-20],
           [ 6.10673003e-18,  1.25000000e-01, -1.24709424e-19,
             7.58940250e-18, -1.82466179e-20,  4.45915727e-18,
            -6.56260751e-21, -2.43715642e-19],
           [ 6.59062874e-18, -1.24709424e-19,  1.25000000e-01,
             6.66621203e-18,  1.58156740e-18, -6.56260751e-21,
             6.30489643e-18,  3.89654534e-20],
           [ 4.63225289e-20,  7.58940250e-18,  6.66621203e-18,
             1.25000000e-01, -7.56664053e-20,  1.45155411e-19,
             1.42073238e-19,  7.52993970e-18],
           [ 1.03621141e-17, -1.82466179e-20,  1.58156740e-18,
            -7.56664053e-20,  1.25000000e-01,  6.28106179e-18,
             6.76605095e-18,  6.02078188e-19],
           [ 1.64513355e-20,  4.45915727e-18, -6.56260751e-21,
             1.45155411e-19,  6.28106179e-18,  1.25000000e-01,
            -5.25927377e-19,  9.25789840e-18],
           [ 6.65864441e-19, -6.56260751e-21,  6.30489643e-18,
             1.42073238e-19,  6.76605095e-18, -5.25927377e-19,
             1.25000000e-01,  6.44716881e-18],
           [ 6.55298598e-20, -2.43715642e-19,  3.89654534e-20,
             7.52993970e-18,  6.02078188e-19,  9.25789840e-18,
             6.44716881e-18,  1.25000000e-01]]), '\nb:\n', array([1.88803776e-17, 1.36664989e-17, 1.12513463e-17, 2.50000000e-01,
           3.75000000e-01, 2.46497132e-17, 1.97267708e-17, 1.63522547e-17]))
    ('coeff:', array([-9.83888788e-17, -1.16605297e-17, -5.46062399e-17,  2.00000000e+00,
            3.00000000e+00,  4.41297363e-17, -6.84422781e-18, -4.11087412e-18]))
    ('approximation:', -9.83888788340761e-17*sin(pi*x)*sin(pi*y)*sin(pi*z) + 3.0*sin(pi*x)*sin(pi*y)*sin(2*pi*z) - 5.46062398956636e-17*sin(pi*x)*sin(2*pi*y)*sin(pi*z) - 6.84422780958398e-18*sin(pi*x)*sin(2*pi*y)*sin(2*pi*z) - 1.16605296707883e-17*sin(2*pi*x)*sin(pi*y)*sin(pi*z) + 4.41297363317943e-17*sin(2*pi*x)*sin(pi*y)*sin(2*pi*z) + 2.0*sin(2*pi*x)*sin(2*pi*y)*sin(pi*z) - 4.1108741176374e-18*sin(2*pi*x)*sin(2*pi*y)*sin(2*pi*z))
    ('diff:', 4.440892098500626e-16)

```python

```

Data Science and Machine Learning/Machine-Learning-In-Python-THOROUGH/EXAMPLES/FINITE_ELEMENTS/INTRO/EXERCICES/02_APPROX_3D.ipynb


## 10 Sept 2019

RULES

Date: Level 2 heading ##

Example Heading: Level 3 heading ###

Method Heading: Level 4 heading ####

### Imports

```python
import numpy as np
import math
import sympy
import matplotlib.pyplot as plt
import matplotlib as mpl
from mpl_toolkits.mplot3d import Axes3D
from IPython.display import display
from IPython.display import display_latex
from sympy import latex
```

### Setup

```python
sympy.init_printing()
np.set_printoptions(precision=3)
```

### Functions

```python
# Usage: display_equation('u_x', x)
def display_equation(idx, symObj):
    if(isinstance(idx, str)):
        eqn = '\\[' + idx + ' = ' + latex(symObj) + '\\]'
        display_latex(eqn, raw=True)
    else:
        eqn = '\\[' + latex(idx) + ' = ' + latex(symObj) + '\\]'
        display_latex(eqn, raw=True)
    return
```

```python
def rotation(theta):
    theta = (np.pi / 180) * theta
    R = np.zeros((3, 3))
    R[0][0] = np.cos(theta)
    R[0][1] = -np.sin(theta)
    R[0][2] = 0
    R[1][0] = np.sin(theta)
    R[1][1] = np.cos(theta)
    R[1][2] = 0
    R[2][0] = 0
    R[2][1] = 0
    R[2][2] = 1
    return R
```

```python
def createEquivRot(theta, K):
    R = np.zeros((3, 3))
    theta = (np.pi / 180) * theta

    R[0][0] = K[0][0] * K[0][0] * (1 - math.cos(theta)) + math.cos(theta)
    R[0][1] = K[0][0] * K[1][0] * (1 - math.cos(theta)) - K[2][0] * math.sin(theta)
    R[0][2] = K[0][0] * K[2][0] * (1 - math.cos(theta)) + K[1][0] * math.sin(theta)

    R[1][0] = K[0][0] * K[1][0] * (1 - math.cos(theta)) + K[2][0] * math.sin(theta)
    R[1][1] = K[1][0] * K[1][0] * (1 - math.cos(theta)) + math.cos(theta)
    R[1][2] = K[1][0] * K[2][0] * (1 - math.cos(theta)) - K[0][0] * math.sin(theta)

    R[2][0] = K[2][0] * K[2][0] * (1 - math.cos(theta)) - K[1][0] * math.sin(theta)
    R[2][1] = K[1][0] * K[2][0] * (1 - math.cos(theta)) + K[0][0] * math.sin(theta)
    R[2][2] = K[2][0] * K[2][0] * (1 - math.cos(theta)) + math.cos(theta)
    return R
```

```python
def vplot2d(O, a, origin_label='O', tip_label='P'):
    dx = a[0][0] - O[0][0]
    dy = a[1][0] - O[1][0]

    head_length = 0.3
    vec_ab = [dx,dy]

    vec_ab_magnitude = math.sqrt(dx**2+dy**2)

    dx = dx / vec_ab_magnitude
    dy = dy / vec_ab_magnitude

    vec_ab_magnitude = vec_ab_magnitude - head_length

    ax = plt.axes()
    ax.set_aspect('equal', 'box')
    #ax.axis([-range, range, -range, range]) Sub for xlim and ylim

    ax.arrow(O[0][0], O[1][0], vec_ab_magnitude*dx, vec_ab_magnitude*dy, head_width=0.1, head_length=head_length, fc='lightblue', ec='black')
    plt.scatter(O[0][0],O[1][0],color='black')
    plt.scatter(a[0][0],a[1][0],color='black')

    ax.annotate(origin_label, (O[0][0]-0.4,O[1][0]),fontsize=14)
    ax.annotate(tip_label, (a[0][0]+0.3,a[1][0]),fontsize=14)
    return
```

```python
def vplot3d(O, P, origin_label='O', tip_label='P'):
    mpl.rcParams['legend.fontsize'] = 10

    ax = fig.gca(projection='3d')
    x = np.linspace(O[0][0], P[0][0], 100)
    y = np.linspace(O[1][0], P[1][0], 100)
    z = np.linspace(O[2][0], P[2][0], 100)

    ax.plot(x, y, z, label=tip_label)
    ax.scatter(P[0][0], P[1][0], P[2][0], marker='x')
    ax.axis('equal')
    ax.legend()
    return
```

```python
# plot a frame
def fplot3d(B, T = np.identity(4), origin_label='O', tip_label='P', x_range = [-2, 2], y_range = [-2, 2], z_range=[-2, 2]):
    mpl.rcParams['legend.fontsize'] = 10
    m = 1

    ax = fig.gca(projection='3d')
    R = T[:3,:3]

    B = np.expand_dims(T[:3,3], axis=0).T

    Bx = homTrans(T, np.array([[m, 0, 0]]).T)
    By = homTrans(T, np.array([[0, m, 0]]).T)
    Bz = homTrans(T, np.array([[0, 0, m]]).T)

    vplot3d(B, Bx, origin_label='B', tip_label='x')
    vplot3d(B, By, origin_label='B', tip_label='y')
    vplot3d(B, Bz, origin_label='B', tip_label='z')

    ax.scatter(B[0][0], B[1][0], B[2][0], marker='x')
    ax.axis('equal')
    ax.set_aspect('equal')
    ax.auto_scale_xyz(x_range, y_range, z_range)
    #ax.pbaspect= [20, 2, 2]
    ax.legend()
    return
```

```python
def homTrans(T, A):
    R = T[:3,:3]
    return (R @ A) + np.expand_dims(T[:3,3], axis=0).T
```

```python
#tVec (translation vector) is a column vector, e.g. tMat = np.array([[10],[5],[0]])
# r is either a 3x3 rotation transformation matrix or an angle in degrees
def createHomTrans(r, tVec):
    tVec = np.vstack((tVec, [1]))
    T = np.zeros((4, 4))

    if (isinstance(r, np.ndarray)):
        T = np.vstack((r, np.array([0, 0, 0])))
    else:
        T = np.vstack((rotation(r), np.array([0, 0, 0])))

    T = np.append(T, tVec, axis=1)
    return T
```

### Example 2.1

```python
R = rotation(30)
BP = np.array([[0.0], [2.0], [0.0]])
# print(R.dot(BP))
AP = R @ BP # New operator since python 3.5
print(AP)
```

    [[-1.   ]
     [ 1.732]
     [ 0.   ]]

```python
range = 3
O = [[0],[0], [0]]
vplot2d(O, AP, tip_label='A')
vplot2d(O, BP, tip_label='B')

plt.grid()

plt.xlim(-range,range)
plt.ylim(-range,range)

plt.title('VECTOR',fontsize=10)
plt.show()
```

FIGURE: 2D plot

```python
O = [[0],[0], [0]]
fig = plt.figure()
vplot3d(O, AP, tip_label='A')
vplot3d(O, BP, tip_label='B')
plt.show()
```

FIGURE: 3D plot

## 11 Sept 2019

### Example 2.2

```python
O = [[0],[0], [0]]
tVec = np.array([[10],[5],[0]])
A = np.array([[3], [7], [0]])
T = createHomTrans(30, tVec)
result = homTrans(T, A)
print(result)
```

    [[ 9.098]
     [12.562]
     [ 0.   ]]

```python
range = 20
vplot2d(O, A, tip_label='A', origin_label='')
vplot2d(O, result, tip_label='B', origin_label='')

plt.grid()

plt.xlim(-range,range)
plt.ylim(-range,range)

plt.title('VECTOR',fontsize=10)
plt.show()
```

```python
O = [[0],[0], [0]]
fig = plt.figure()
vplot3d(O, A, tip_label='A')
vplot3d(O, result, tip_label='B')
plt.show()
```

### Example 2.5

#### Method 1

```python
# Inverting the 4x4 matrix
O = [[0],[0], [0]]
tVec = np.array([[4],[3],[0]])
ATB = createHomTrans(30, tVec)
BTA = np.linalg.inv(T)
print(BTA)
```

    [[  0.866   0.5     0.    -11.16 ]
     [ -0.5     0.866   0.      0.67 ]
     [  0.      0.      1.      0.   ]
     [  0.      0.      0.      1.   ]]

#### Method 2

```python
# Using the method in the book
R = rotation(30)
tVec = np.array([[4],[3],[0]])
R = R.T
tVec = -R @ tVec
BTA = createHomTrans(R, tVec)
print(BTA)
```

    [[ 0.866  0.5    0.    -4.964]
     [-0.5    0.866  0.    -0.598]
     [ 0.     0.     1.     0.   ]
     [ 0.     0.     0.     1.   ]]

## 12 Sept 2019

### X-Y-Z Fixed Angles Representation Derivation

```python
c_a, c_b, c_y = sympy.symbols('c_alpha c_beta c_gamma')
s_a, s_b, s_y = sympy.symbols('s_alpha s_beta s_gamma')

R_z = sympy.Matrix([[c_a, -s_a, 0], [s_a, c_a, 0], [0, 0, 1]])
R_y = sympy.Matrix([[c_b, 0, s_b], [0, 1, 0], [-s_b, 0, c_b]])
R_x = sympy.Matrix([[1, 0, 0], [0, c_y, -s_y], [0, s_y, c_y]])

R = R_z * R_y * R_x
display_equation('R_z', R_z)
display_equation('R_y', R_y)
display_equation('R_x', R_x)
display_equation('R = R_z R_y R_x', R)

```

\[R_z = \left[\begin{matrix}c_{\alpha} & - s_{\alpha} & 0\\s_{\alpha} & c_{\alpha} & 0\\0 & 0 & 1\end{matrix}\right]\]

\[R_y = \left[\begin{matrix}c_{\beta} & 0 & s_{\beta}\\0 & 1 & 0\\- s_{\beta} & 0 & c_{\beta}\end{matrix}\right]\]

\[R_x = \left[\begin{matrix}1 & 0 & 0\\0 & c_{\gamma} & - s_{\gamma}\\0 & s_{\gamma} & c_{\gamma}\end{matrix}\right]\]

\[R = R_z R_y R_x = \left[\begin{matrix}c_{\alpha} c_{\beta} & c_{\alpha} s_{\beta} s_{\gamma} - c_{\gamma} s_{\alpha} & c_{\alpha} c_{\gamma} s_{\beta} + s_{\alpha} s_{\gamma}\\c_{\beta} s_{\alpha} & c_{\alpha} c_{\gamma} + s_{\alpha} s_{\beta} s_{\gamma} & - c_{\alpha} s_{\gamma} + c_{\gamma} s_{\alpha} s_{\beta}\\- s_{\beta} & c_{\beta} s_{\gamma} & c_{\beta} c_{\gamma}\end{matrix}\right]\]

### Equivalent angle—axis representation

```python
#k_x, k_y, k_z, c, s = sympy.symbols('k_x k_y k_z c_theta s_theta')
#R_theta = sympy.symbols('R_theta')
#R_x = Matrix([[1, 1, 0], [0, c, -s], [0, s, c]])
#R_y = Matrix([[c, 0, s], [0, 1, 0], [-s, 0, c]])
#R_z = Matrix([[c, -s, 0], [s, c, 0], [0, 0, 1]])

# To be completed
```

### Example 2.8

```python
O = np.zeros((3, 1))
R = np.zeros((3, 3))
K = np.array([[0.7070, 0.7070, 0]]).T
tVec = np.array([[1.0, 2.0, 3.0]]).T

R = createEquivRot(30, K)
T = createHomTrans(R, O)

print(T)
```

    [[ 0.933  0.067  0.353  0.   ]
     [ 0.067  0.933 -0.353  0.   ]
     [-0.353  0.353  0.866  0.   ]
     [ 0.     0.     0.     1.   ]]

```python
m = 2 # This is a magnitude scalar to make K visible in the plot
A = np.zeros((3, 1))
fig = plt.figure()
# Origin frame
fplot3d(A, tip_label='B')
vplot3d(A, K * m, origin_label='A', tip_label='K')
fplot3d(A, T, tip_label='B',
        x_range = [-2, 2], y_range = [-2, 2], z_range = [-2, 2])
plt.show()
```

## 13 Sept 2019

```python
O = np.zeros((3, 1))
R = np.zeros((3, 3))
K = np.array([[0.7070, 0.7070, 0]]).T
tVec = np.array([[1.0, 2.0, 3.0]]).T
P = np.array([[1.0, 2.0, 3.0]]).T

R = createEquivRot(30, K)
BTB = createHomTrans(0, -tVec)
ATA = createHomTrans(0, tVec)
ATB = createHomTrans(R, O)

T = ATA @ ATB @ BTB
print(T)

m = 2 # This is a magnitude scalar to make K visible in the plot
A = np.zeros((3, 1))
P *= 0.5

S = P + K

fig = plt.figure()
ax = fig.gca(projection='3d')
ax.scatter(P[0][0], P[1][0], P[2][0], marker='P')

fplot3d(A, tip_label='B')
vplot3d(P, S, origin_label='A', tip_label='K')
fplot3d(A, T, tip_label='B',
        x_range = [-2, 2], y_range = [-2, 2], z_range = [-2, 2])
plt.show()
```

### How to show two frames where one is the origin, and the other is transformed

```python
A = np.zeros((3, 1))
R = rotation(90)
tVec = np.ones((3, 1))
T = createHomTrans(R, tVec)
#print(T)
fig = plt.figure()
# Origin frame
fplot3d(A, tip_label='B')
# B vector
vplot3d(A, tVec, origin_label='B', tip_label='x')
# B vector frame
fplot3d(A, T, tip_label='B',
        x_range = [-2, 2], y_range = [-2, 2], z_range = [-2, 2])
plt.show()
```

```python

```


# Random Signals

*This jupyter notebook is part of a [collection of notebooks](../index.ipynb) on various topics of Digital Signal Processing. Please direct questions and suggestions to [Sascha.Spors@uni-rostock.de](mailto:Sascha.Spors@uni-rostock.de).*

## Power Spectral Density

The (auto-) [power spectral density](https://en.wikipedia.org/wiki/Spectral_density#Power_spectral_density) (PSD) is defined as the Fourier transformation of the [auto-correlation function](correlation_functions.ipynb) (ACF).

### Definition

For a continuous-amplitude real-valued wide-sense stationary (WSS) random signal $x[k]$ the PSD is given as

\begin{equation}
\Phi_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega}) = \mathcal{F}_* \{ \varphi_{xx}[\kappa] \}
\end{equation}

where $\mathcal{F}_* \{ \cdot \}$ denotes the [discrete-time Fourier transformation](https://en.wikipedia.org/wiki/Discrete-time_Fourier_transform) (DTFT) and $\varphi_{xx}[\kappa]$ the ACF of $x[k]$. Note, the DTFT is performed with respect to $\kappa$. The ACF of a random signal of finite length $N$ can be expressed by way of a linear convolution

\begin{equation}
\varphi_{xx}[\kappa] = \frac{1}{N} \cdot x_N[k] * x_N[-k]
\end{equation}

Taking the DTFT of the left- and right-hand side results in

\begin{equation}
\Phi_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega}) = \frac{1}{N} \, X_N(\mathrm{e}^{\,\mathrm{j}\,\Omega})\, X_N(\mathrm{e}^{-\,\mathrm{j}\,\Omega}) =
\frac{1}{N} \, | X_N(\mathrm{e}^{\,\mathrm{j}\,\Omega}) |^2
\end{equation}

The last equality results from the definition of the magnitude and the symmetry of the DTFT for real-valued signals. The spectrum $X_N(\mathrm{e}^{\,\mathrm{j}\,\Omega})$ quantifies the amplitude density of the signal $x_N[k]$. It can be concluded from above result that the PSD quantifies the squared amplitude or power density of a random signal. This explains the term power spectral density.

### Properties

The properties of the PSD can be deduced from the properties of the ACF and the DTFT as

1. From the link between the PSD $\Phi_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega})$ and the spectrum $X_N(\mathrm{e}^{\,\mathrm{j}\,\Omega})$ derived above it can be concluded that the PSD real valued

    $$\Phi_{xx}(\mathrm{e}^{\,\mathrm{j}\,\Omega}) \in \mathbb{R}$$

2. From the even symmetry $\varphi_{xx}[\kappa] = \varphi_{xx}[-\kappa]$ of the ACF it follows that

    $$ \Phi_{xx}(\mathrm{e}^{\,\mathrm{j} \, \Omega}) = \Phi_{xx}(\mathrm{e}^{\,-\mathrm{j}\, \Omega}) $$

3. The PSD of an uncorrelated random signal is given as

    $$ \Phi_{xx}(\mathrm{e}^{\,\mathrm{j} \, \Omega}) = \sigma_x^2 + \mu_x^2 \cdot {\bot \!\! \bot \!\! \bot}\left( \frac{\Omega}{2 \pi} \right) $$

    which can be deduced from the [ACF of an uncorrelated signal](correlation_functions.ipynb#Properties).

4. The quadratic mean of a random signal is given as

    $$ E\{ x[k]^2 \} = \varphi_{xx}[0] = \frac{1}{2\pi} \int\limits_{-\pi}^{\pi} \Phi_{xx}(\mathrm{e}^{\,\mathrm{j}\, \Omega}) \,\mathrm{d} \Omega $$

    The last relation can be found by expressing the ACF by the inverse DTFT.

### Example - Power Spectral Density of a Speech Signal

In this example the PSD $\Phi_{xx}(\mathrm{e}^{\,\mathrm{j} \,\Omega})$ of a speech signal $x[k]$ is estimated by applying a discrete Fourier transformation (DFT) to its ACF. For a better interpretation of the PSD, the frequency axis $f = \frac{\Omega}{2 \pi} \cdot f_s$ has been chosen for illustration, where $f_s$ denotes the sampling frequency of the signal.

In Python the ACF is stored in a vector with indexes $0, 1, ..., 2N -1$ where the indexes correspond to the lags $\kappa = -N+1,-N+2,....,N-1$. When computing the discrete Fourier transform (DFT) of the ACF numerically by the fast Fourier transform (FFT) one has to take this shift into account. For instance, by multiplying the DFT $\Phi_{xx}[\mu]$ by $e^{j \mu \frac{2 \pi}{2N - 1} (N-1)}$ where $N$ denotes the length of the signal $N$.

```python
%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
from scipy.io import wavfile

# read audio file
fs, x = wavfile.read('../data/speech_8k.wav')
x = np.asarray(x, dtype=float)
N = len(x)

# compute ACF
acf = 1/len(x) * np.correlate(x, x, mode='full')
# compute PSD
psd = np.fft.fft(acf)
psd = psd * np.exp(1j*np.arange(2*N-1)*2*np.pi*(N-1)/(2*N-1))
f = np.fft.fftfreq(2*N-1, d=1/fs)

# plot PSD
plt.figure(figsize = (10, 8))
plt.plot(f, np.real(psd))
plt.title('Estimated power spectral density')
plt.ylabel(r'$\hat{\Phi}_{xx}(e^{j \Omega})$')
plt.xlabel(r'$f$')
plt.axis([0, 2000, 0, 1.1*max(np.abs(psd))]);
plt.grid()
```

**Exercise**

* What does the PSD tell you about the spectral contents of a speech signal?

Solution: It can be concluded from the shown PSD that the main power of a speech signal is contained in the frequency range below 500 Hz. The speech signal exhibits furthermore a harmonic structure with a dominant fundamental frequency and a number of harmonics.

## Cross-Power Spectral Density

The cross-power spectral density is defined as the Fourier transformation of the [cross-correlation function](correlation_functions.ipynb#Cross-Correlation-Function) (CCF).

### Definition

For two continuous-amplitude real-valued wide-sense stationary (WSS) random signals $x[k]$ and $y[k]$ the cross-power spectral density is given as

\begin{equation}
\Phi_{xy}(\mathrm{e}^{\,\mathrm{j} \, \Omega}) = \mathcal{F}_* \{ \varphi_{xy}[\kappa] \}
\end{equation}

where $\varphi_{xy}[\kappa]$ denotes the CCF of $x[k]$ and $y[k]$. Note, the DTFT is performed with respect to $\kappa$. The CCF of two random signals of finite lengths $N$ and $M$ can be expressed by way of a linear convolution

\begin{equation}
\varphi_{xy}[\kappa] = \frac{1}{N} \cdot x_N[k] * y_M[-k]
\end{equation}

Taking the DTFT of the left- and right-hand side results in

\begin{equation}
\Phi_{xy}(\mathrm{e}^{\,\mathrm{j}\,\Omega}) = \frac{1}{N} \, X_N(\mathrm{e}^{\,\mathrm{j}\,\Omega})\, Y_M(\mathrm{e}^{-\,\mathrm{j}\,\Omega})
\end{equation}

### Properties

1. The symmetries of $\Phi_{xy}(\mathrm{e}^{\,\mathrm{j}\, \Omega})$ can be derived from the symmetries of the CCF and the DTFT as

    $$ \underbrace {\Phi_{xy}(\mathrm{e}^{\,\mathrm{j}\, \Omega}) = \Phi_{xy}^*(\mathrm{e}^{-\,\mathrm{j}\, \Omega})}_{\varphi_{xy}[\kappa] \in \mathbb{R}} =
\underbrace {\Phi_{yx}(\mathrm{e}^{\,- \mathrm{j}\, \Omega}) = \Phi_{yx}^*(\mathrm{e}^{\,\mathrm{j}\, \Omega})}_{\varphi_{yx}[-\kappa] \in \mathbb{R}}$$

    from which can be concluded that $|\Phi_{xy}(\mathrm{e}^{\,\mathrm{j}\, \Omega})| = |\Phi_{yx}(\mathrm{e}^{\,\mathrm{j}\, \Omega})|$

2. The cross PSD of two uncorrelated random signals is given as

    $$ \Phi_{xy}(\mathrm{e}^{\,\mathrm{j} \, \Omega}) = \mu_x^2 \mu_y^2 \cdot {\bot \!\! \bot \!\! \bot}\left( \frac{\Omega}{2 \pi} \right) $$

    which can be deduced from the CCF of an uncorrelated signal.

### Example - Cross Power Spectral Density

The following example estimates and plots the cross PSD $\Phi_{xy}(\mathrm{e}^{\,\mathrm{j}\, \Omega})$ of two random signals $x_N[k]$ and $y_M[k]$ of finite lengths $N$ and $M = 2 N$. The estimated cross PSD is calculated from the valid part of the CCF $\varphi_{xy}[\kappa]$ by an DFT in order to exclude boundary effects.

```python
N = 1024  # length of random signal x

# generate two random signals
np.random.seed(2)
x = 2 + np.random.normal(size=N)
y = 1 + np.random.normal(size=2*N)

# compute cross PSD via CCF
acf = 1/N * np.correlate(x, y, mode='valid')
psd = np.fft.fft(acf)
psd = psd * np.exp(1j*np.arange(N+1)*2*np.pi*(N-1)/(2*N-1))

# plot results
f = np.fft.fftfreq(len(psd), d=1/2)
plt.figure(figsize = (10, 4))
plt.stem(f, np.real(psd))
plt.title('Estimated cross power spectral density')
plt.ylabel(r'$\hat{\Phi}_{xy}(e^{j \Omega})$')
plt.xlabel(r'$\Omega/ \pi$')
plt.grid()
```

**Exercise**

* What does the cross PSD $\Phi_{xy}(\mathrm{e}^{\,\mathrm{j} \, \Omega})$ tell you about the statistical properties of the two random signals?

Solution: The cross PSD $\Phi_{xy}(\mathrm{e}^{\,\mathrm{j} \, \Omega})$ is essential only non-zero for $\Omega=0$. It hence can be concluded that the two random signals are not mean-free and uncorrelated to each other.

**Copyright**

This notebook is provided as [Open Educational Resource](https://en.wikipedia.org/wiki/Open_educational_resources). Feel free to use the notebook for your own purposes. The text is licensed under [Creative Commons Attribution 4.0](https://creativecommons.org/licenses/by/4.0/), the code of the IPython examples under the [MIT license](https://opensource.org/licenses/MIT). Please attribute the work as follows: *Sascha Spors, Digital Signal Processing - Lecture notes featuring computational examples, 2016-2018*.


# Taylor problem 2.43

last revised: 08-Jan-2019 by Dick Furnstahl [furnstahl.1@osu.edu]

**Go through and fill in the blanks where ### appears.**

The goal of this problem is to plot the trajectory of a basketball with specified parameters, with and without air resistance.  From the plot we are to find how far the basketball travels horizontally in the two cases before hitting the floor.

The differential equations to solve are equation (2.61) in Taylor:

$
\begin{align}
    m \ddot x &= - c \sqrt{\dot x^2 + \dot y^2}\, \dot x  \\
    m \ddot y &= - mg - c \sqrt{\dot x^2 + \dot y^2}\, \dot y
\end{align}
$

where the initial conditions we are given (in meters) are:

$
\begin{align}
   x(0) &= 0 \;, & y(0) &= 2  \\
   \dot x(0) &= 15 \cos \pi/4 \;, & \dot y(0) &= 15\sin \pi/4
\end{align}
$

and $c = \gamma D^2$ with $\gamma = 0.25$ and $D = 0.24\,$m.

### Gameplan:

We'll follow the treatment at the end of Taylor_problem_1.39.ipynb.  So we rewrite the coupled 2nd-order differential equations $x(t)$ and $y(t)$ as four coupled 1st-order differential equations for $x(t)$, $y(t)$, $v_x(t)$, and $v_y(t)$:

$
\begin{align}
   \frac{d}{dt}\left(\begin{array}{c}
                          \mathbf{x} \\
                          \mathbf{v}
                      \end{array}\right)
               = \left(\begin{array}{c}
                          \mathbf{v} \\
                          \mathbf{F}/m
                       \end{array}\right)
      \qquad \Longrightarrow \qquad
   \frac{d}{dt}\left(\begin{array}{c}
                          x \\
                          y \\
                          v_x \\
                          v_y
                      \end{array}\right)
               = \left(\begin{array}{c}
                          v_x \\
                          v_y \\
                          F_x/m \\
                          F_y/m
                       \end{array}\right)
\end{align}
$

```python
### Import what you need to make arrays, to make plots, and to use odeint.
```

```python
def ode_rhs(u_vec, t, *params):
    """
    Right-hand side (rhs) of the differential equation, with
    u_vec = [x, y, v_x, v_y] and params = [m, g, c].
    """
    x, y, v_x, v_y = u_vec
    m, g, c = params
    return [v_x,
            , ### fill in the 2nd line of the return vector
            , ### fill in the 3rd line of the return vector
            -g - (c/m) * np.sqrt(v_x**2 + v_y**2) * v_y]
```

```python
m = 0.6      # units: kg
g = 9.8      # units: m/s^2
gamma = 0.25 # units: N s^2/m^4
D = 0.24     # units: m
c = gamma * D**2

theta = np.pi/4.  # angle of launch
v0 = 15.          # launch speed

# set the initial conditions
u0_vec = [0, 2, v0*np.cos(theta), ]  ### complete the initial conditions

t_max = 3.  # integration time
t_pts = np.arange()  ### create a suitable array of t values

# absolute and relative tolerances for ode solver
abserr = 1.0e-8
relerr = 1.0e-6

# Integrate the differential equation
#  Note the use of T here to take the *transpose* to change a
#  (# points) by 4 array to a 4 by (# points) array.
x, y, v_x, v_y = odeint(ode_rhs, u0_vec, t_pts, args=(m, g, c),
                        atol=abserr, rtol=relerr).T
```

Check the shape (i.e., dimensions) of the output from odeint:

```python
(odeint(ode_rhs, u0_vec, t_pts, args=(m, g, c),
                        atol=abserr, rtol=relerr)).shape
```

```python
(odeint(ode_rhs, u0_vec, t_pts, args=(m, g, c),
                        atol=abserr, rtol=relerr).T).shape
```

Define a function for the trajectory in vacuum, using the analytic formula:

```python
### Just check that this makes sense
def vacuum_trajectory(u0_vec, t_pts, g=9.8):
    """Given the initial condition vector and an array of time values,
       return the trajectory (x and y arrays) in the vacuum, using the
       exact solution.
    """
    x0, y0, v0_x, v0_y = u0_vec
    return v0_x * t_pts, \
           y0 + v0_y * t_pts - g * t_pts**2/2.

# create arrays for the vaccum trajectory (using default for g)
x_vac, y_vac = vacuum_trajectory(u0_vec, t_pts)
```

Make a plot with and without drag:

```python
fig = plt.figure(figsize=(6,4))
ax = fig.add_subplot(1,1,1)
ax.plot(x, y, 'b-', label='air resistance', )
ax.plot(x_vac, y_vac, 'r--', label='vacuum', )
ax.set_ylim()    ### set appropriate limits on the y-axis
ax.set_xlim()    ### set appropriate limits on the x-axis
ax.set_xlabel('x (m)')
ax.set_ylabel('y (m)')

ax.legend()
### add a line to save this figure
```

We can read off rough values for the distances where the ball hits the ground.  Or we can find the closest values to zero in the arrays for `y` and `y_vac` and look up the value of `x` and `x_vac` at the same array index:

```python
# Note the use of fabs(y), which is |y|, to get closest to
#  zero instead of the most negative value of y.
index_y_min = (np.fabs(y)).argmin() # finds nearest index to minimum of |y|
print(f'distance with drag = {x[index_y_min]:.2f}' )

index_y_vac_min = (np.fabs(y_vac)).argmin() # finds nearest index to minimum of |y|
print(f'distance in vacuum = {x_vac[index_y_vac_min]:.2f}' )

```

**What change could you make to get a better estimate with the same method?**

```python

```