Contact: mihaelacr@google.com

## Contents

* amortized variational inference (VAEs)
* improving amortized variational inference using KL annleaing
* improving amortized variational inference using constraint optimization

## Tasks
  * [Code] fill in the Elbo for the VAE
  * [Analysis] discuss the difference between sample and reconstruction quality for VAEs, gradient estimation
  * [Code] define the update operation for the KL coefficient
  * [Analysis] discuss the effect of the KL annealing
  * [Code] implement the Lagrangian optimization (GECO)
  * [Analysis] what is the effect of the *learned* Lagrangian on the KL schedule? How would you set the value for the constraint? What are the advantages of using constrained optimization?

## Imports and set up

```
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import math

import tensorflow as tf
import sonnet as snt

import numpy as np

# Plotting library.
from matplotlib import pyplot as plt
import seaborn as sns

# Tensorflow probability utilities
import tensorflow_probability as tfp

tfd = tfp.distributions

```

```
sns.set(rc={"lines.linewidth": 2.8}, font_scale=2)
sns.set_style("whitegrid")
```

```
# Don't forget to select GPU runtime environment in Runtime -> Change runtime type
device_name = tf.test.gpu_device_name()
if device_name != '/device:GPU:0':
  raise SystemError('GPU device not found')
print('Found GPU at: {}'.format(device_name))
```

```
def gallery(array, ncols=10, rescale=False):
    """Data visualization code."""
    if rescale:
      array = (array + 1.) / 2
    nindex, height, width, intensity = array.shape
    nrows = nindex//ncols
    assert nindex == nrows*ncols
    # want result.shape = (height*nrows, width*ncols, intensity)
    result = (array.reshape(nrows, ncols, height, width, intensity)
              .swapaxes(1,2)
              .reshape(height*nrows, width*ncols, intensity))
    return result
```

```
def show_digits(axis, digits, title=''):
  axis.axis('off')
  ncols = int(np.sqrt(digits.shape[0]))
  axis.imshow(gallery(digits, ncols=ncols).squeeze(axis=2),
                 cmap='gray')
  axis.set_title(title, fontsize=15)
```

```
def show_latent_interpolations(generator, prior, session):
  a = np.linspace(0.0, 1.0, BATCH_SIZE)
  a = np.expand_dims(a, axis=1)

  first_latents = prior.sample()[0]
  second_latents = prior.sample()[0]

  # To ensure that the interpolation is still likely under the Gaussian prior,
  # we use Gaussian interpolation - rather than linear interpolation.
  interpolations =  np.sqrt(a) * first_latents + np.sqrt(1 - a) * second_latents

  ncols = int(np.sqrt(BATCH_SIZE))
  samples_from_interpolations = generator(interpolations)
  samples_from_interpolations_np = sess.run(samples_from_interpolations)
  plt.gray()
  axis = plt.gca()
  show_digits(
     axis, samples_from_interpolations_np, title='Latent space interpolations')
```

## Hyperparameters

```
BATCH_SIZE = 64
NUM_LATENTS = 10
TRAINING_STEPS = 10000
```

```
DECODER_VARIABLE_SCOPE = "decoder"
```

```
tf.reset_default_graph()
```

## Get the data

We will use the MNIST dataset. Luckly, TensorFlow comes with a simple way to load it.

```
mnist = tf.contrib.learn.datasets.load_dataset("mnist")
```

```
print(mnist.train.images.shape)
print(type(mnist.train.images))
```

### Transform the data from numpy arrays to in graph tensors.

This allows us to use TensorFlow datasets, which ensure that a new batch from the data is being fed at each session.run. This means that we do not need to use feed_dicts to feed data to each session.

```
def make_tf_data_batch(np_data, shuffle=True):
  # Reshape the data to image size.
  images = np_data.reshape((-1, 28, 28, 1))

  # Create the TF dataset.
  dataset = tf.data.Dataset.from_tensor_slices(images)

  # Shuffle and repeat the dataset for training.
  # This is required because we want to do multiple passes through the entire
  # dataset when training.
  if shuffle:
    dataset = dataset.shuffle(100000).repeat()

  # Batch the data and return the data batch.
  one_shot_iterator = dataset.batch(BATCH_SIZE).make_one_shot_iterator()
  data_batch = one_shot_iterator.get_next()
  return data_batch
```

```
real_data = make_tf_data_batch(mnist.train.images)
print(real_data.shape)
```

# Amortized variational inference (VAEs)

Instead of learning one set of posterior variables per data point, we can use function approximation to learn the distributional variables. Specifically, the posterior parameters for $x_i$ will be the output of a *learned* function $f_\theta(x_i)$, where $\theta$ are parameters shared across all data points. Can you think of why this is useful?

Objective - maximize:
\begin{equation}
 \mathbb{E}_{p^*(x)} \mathbb{E}_{q(z|x)} \log p_{\theta}(x|z)  - \mathbb{E}_{p^*(x)} KL(q(z|x)||p(z))
\end{equation}

For more information, see:
  * https://arxiv.org/abs/1312.6114

## Define the decoder

```
def standard_decoder(z):
  with tf.variable_scope(DECODER_VARIABLE_SCOPE, reuse=tf.AUTO_REUSE):
    h = tf.layers.dense(z, 7 * 7 * 64, activation=tf.nn.relu)
    h = tf.reshape(h, shape=[BATCH_SIZE, 7, 7, 64])
    h = tf.layers.Conv2DTranspose(
          filters=32,
          kernel_size=5,
          strides=2,
          activation=tf.nn.relu,
          padding='same')(h)
    h = tf.layers.Conv2DTranspose(
          filters=1,
          kernel_size=5,
          strides=2,
          activation=None,  # Do not activate the last layer.
          padding='same')(h)
    return tf.distributions.Bernoulli(h)
```

```
def big_decoder(z):
  with tf.variable_scope(DECODER_VARIABLE_SCOPE, reuse=tf.AUTO_REUSE):
    h = tf.layers.dense(z, 4000, activation=tf.nn.relu)
    h = tf.layers.dense(z, 4000, activation=tf.nn.relu)
    h = tf.layers.dense(z, 4000, activation=tf.nn.relu)
    h = tf.layers.dense(h, 7 * 7 * 64, activation=tf.nn.relu)
    h = tf.reshape(h, shape=[BATCH_SIZE, 7, 7, 64])
    h = tf.layers.Conv2DTranspose(
          filters=32,
          kernel_size=5,
          strides=2,
          activation=tf.nn.relu,
          padding='same')(h)
    h = tf.layers.Conv2DTranspose(
          filters=32,
          kernel_size=5,
          strides=1,
          activation=tf.nn.relu,
          padding='same')(h)
    h = tf.layers.Conv2DTranspose(
          filters=32,
          kernel_size=5,
          strides=1,
          activation=tf.nn.relu,
          padding='same')(h)
    h = tf.layers.Conv2DTranspose(
          filters=1,
          kernel_size=5,
          strides=2,
          activation=None,  # Do not activate the last layer.
          padding='same')(h)
    return tf.distributions.Bernoulli(h)
```

```
ENCODER_VARIABLE_SCOPE = 'encoder'
```

```
def encoder(x):
  with tf.variable_scope(ENCODER_VARIABLE_SCOPE, reuse=tf.AUTO_REUSE):

    h = tf.layers.Conv2D(
          filters=8,
          kernel_size=5,
          strides=2,
          activation=tf.nn.relu,
          padding='same')(x)
    h = tf.layers.Conv2D(
          filters=16,
          kernel_size=5,
          strides=2,
          activation=tf.nn.relu,
          padding='same')(h)
    h = tf.layers.Conv2D(
          filters=32,
          kernel_size=5,
          strides=1,
          activation=tf.nn.relu,
          padding='same')(h)

    out_shape = 1
    for s in h.shape.as_list()[1:]:
      out_shape*= s

    h = tf.reshape(h, shape=[BATCH_SIZE, out_shape])
    mean = tf.layers.dense(h, NUM_LATENTS, activation=None)
    scale = tf.layers.dense(h, NUM_LATENTS, activation=None)
    return multi_normal(loc=mean, log_scale=scale)
```

```
def multi_normal(loc, log_scale):
  # We model the latent variables as independent
  return tfd.Independent(
      distribution=tfd.Normal(loc=loc, scale=tf.exp(log_scale)),
      reinterpreted_batch_ndims=1)
```

```
def make_prior():
  # Zero mean, unit variance prior.
  prior_mean = tf.zeros(shape=(BATCH_SIZE, NUM_LATENTS), dtype=tf.float32)
  prior_log_scale = tf.zeros(shape=(BATCH_SIZE, NUM_LATENTS), dtype=tf.float32)

  return multi_normal(prior_mean, prior_log_scale)
```

## Task: define the elbo (variational bound) terms

Compared to the BLR colab, the posterior is over latent variables, and they are conditioned on the input.

You now have to define the two terms of the elbo: the log prob term, and the KL term. Remember that for the KL term you can compute the analytical KL since the prior and the posterior are Gaussian.

```
def bound_terms(data_batch, variational_posterior, decoder_fn):
  latent_sample = None
  # Compute the log probs of the data (log p^*_{\theta}(x|z)).
  # Hint: Use the decoder and look at the shape assertion below!
  all_log_probs = None

  # Since we used broadcasting, log_prob of a Bernoulli will return a log_prob for
  # each dimension.
  all_log_probs.shape.assert_is_compatible_with([BATCH_SIZE, 28, 28, 1])

  # Reduce sum over the data dimensions.
  # This is needed because we used independent distributions for each of the
  # data pixels.
  likelihood_term = None

  # Reduce mean over the batch dimensions
  likelihood_term = tf.reduce_mean(likelihood_term)

  # Compute the KL divergence.
  # Hint: the posterior is a distribution object (check its member functions)!
  kl_term = None
  kl_term.shape.assert_is_compatible_with([BATCH_SIZE])

  # Reduce over the batch dimension.
  kl_term = tf.reduce_mean(kl_term)

  return likelihood_term, kl_term
```

```
real_data = make_tf_data_batch(mnist.train.images)
print(real_data.shape)
```

```
prior = make_prior()
```

```
# The posterior is conditioned on the data: q(z|x).
variational_posterior = encoder(real_data)
decoder = standard_decoder
```

```
# Maximize the data likelihodd and minimize the KL divergence between the prior
# and posterior.
likelihood_term, kl_term = bound_terms(
    real_data, variational_posterior, decoder)
train_elbo = likelihood_term - kl_term
loss =  - train_elbo
```

```
samples = decoder(prior.sample()).mean()
samples.shape.assert_is_compatible_with([BATCH_SIZE, 28, 28, 1])

reconstructions = decoder(variational_posterior.sample()).mean()
```

```
# We now perform joint optimization on the encoder and decoder variables.
optimizer = tf.train.AdamOptimizer(0.001, beta1=0.9, beta2=0.9)
decoder_vars = tf.get_collection(
    tf.GraphKeys.TRAINABLE_VARIABLES, scope=DECODER_VARIABLE_SCOPE)
encoder_vars = tf.get_collection(
    tf.GraphKeys.TRAINABLE_VARIABLES, scope=ENCODER_VARIABLE_SCOPE)
autoencoder_vars = encoder_vars + decoder_vars
update_op = optimizer.minimize(loss, var_list=autoencoder_vars)
```

## Training

```
sess = tf.Session()

# Initialize all variables
sess.run(tf.initialize_all_variables())
```

```
# %hide_pyerr  # - uncomment to interrupt training without a stacktrace
losses = []
kls = []
likelihood_terms = []

for i in xrange(TRAINING_STEPS):
  sess.run(update_op)

  if i % 100 == 0:
    iteration_loss, iteration_likelihood, iteration_kl  = sess.run(
          [loss, likelihood_term, kl_term])
    print('Iteration {}. Loss {}. KL {}'.format(
          i, iteration_loss, iteration_kl))
    losses.append(iteration_loss)
    kls.append(iteration_kl)
    likelihood_terms.append(iteration_likelihood)
```

### Visualize the loss in time

```
fig, axes = plt.subplots(1, 3, figsize=(3*8,5))

axes[0].plot(losses, label='Negative ELBO')
axes[0].set_title('Time', fontsize=15)
axes[0].legend()

axes[1].plot(kls, label='KL')
axes[1].set_title('Time', fontsize=15)
axes[1].legend()

axes[2].plot(likelihood_terms, label='Likelihood Term')
axes[2].set_title('Time', fontsize=15)
axes[2].legend()
```

### Generate samples and latent interpolations

```
real_data_vals, final_samples_vals, data_reconstructions_vals = sess.run(
      [real_data, samples, reconstructions])
```

```
fig, axes = plt.subplots(1, 3, figsize=(3*4,4))

show_digits(axes[0], real_data_vals, 'Data')
show_digits(axes[1], data_reconstructions_vals, 'Reconstructions')
show_digits(axes[2], final_samples_vals, 'Samples')
```

```
show_latent_interpolations(lambda x: decoder(x).mean(), prior, sess)
```

## Questions about amortized variational inference

* What do you observe about sample quality and reconstruction quality?
* What do you observe about the ELBO and KL term?
* Which gradient estimation method is used to compute the gradients with respect to the encoder parameters?

# KL annealing

Objective - maximize:
\begin{equation}
 \mathbb{E}_{p^*(x)} \mathbb{E}_{q(z|x)} \log p_\theta(x|z)  - \alpha \mathbb{E}_{p^*(x)} KL(q(z|x)||p(z))
\end{equation}

Where $\alpha$ changes during training, to weigh in the KL term more.

```
tf.reset_default_graph()
```

```
real_data = make_tf_data_batch(mnist.train.images)
prior = make_prior()
decoder = standard_decoder
```

```
kl_coefficient = tf.Variable(
    tf.zeros(shape=(1), dtype=tf.float32),
    trainable=False,
    name='kl_coeff')
```

## Task: define the update schedule for the coefficient of the KL

We want the kl coefficient (`kl_coefficient`) to increase linearly by `kl_step` at each iteration. To achieve this, we need to define an update operation in TensorFlow, which we will run manually in the training loop (see below).

```
kl_step = 1.0 / TRAINING_STEPS
update_kl_coeff = None
```

```
variational_posterior = encoder(real_data)
# Maximize the data likelihodd and minimize the KL divergence between the prior
# and posterior.
likelihood_term, kl_term = bound_terms(
    real_data, variational_posterior, decoder)
train_elbo = likelihood_term - kl_coefficient * kl_term
loss =  - train_elbo
```

```
samples = decoder(prior.sample()).mean()
samples.shape.assert_is_compatible_with([BATCH_SIZE, 28, 28, 1])

reconstructions = decoder(variational_posterior.sample()).mean()
```

```
# We now perform joint optimization on the encoder and decoder variables.
optimizer = tf.train.AdamOptimizer(0.001, beta1=0.9, beta2=0.9)
decoder_vars = tf.get_collection(
    tf.GraphKeys.TRAINABLE_VARIABLES, scope=DECODER_VARIABLE_SCOPE)
encoder_vars = tf.get_collection(
    tf.GraphKeys.TRAINABLE_VARIABLES, scope=ENCODER_VARIABLE_SCOPE)
autoencoder_vars = encoder_vars + decoder_vars
variables_update_op = optimizer.minimize(loss, var_list=autoencoder_vars)

# Ensure that a variable update is followed by an update in the KL coefficient.
with tf.control_dependencies([variables_update_op]):
  update_op = tf.identity(update_kl_coeff)
```

## Training

```
sess = tf.Session()

# Initialize all variables
sess.run(tf.initialize_all_variables())
```

```
# %hide_pyerr  # - uncomment to interrupt training without a stacktrace
losses = []
kls = []
likelihood_terms = []

for i in xrange(TRAINING_STEPS):
  sess.run(update_op)

  if i % 100 == 0:
    iteration_loss, iteration_likelihood, iteration_kl  = sess.run(
          [loss, likelihood_term, kl_term])
    print('Iteration {}. Loss {}. KL {}'.format(
          i, iteration_loss, iteration_kl))
    losses.append(iteration_loss)
    kls.append(iteration_kl)
    likelihood_terms.append(iteration_likelihood)
```

```
fig, axes = plt.subplots(1, 3, figsize=(3*8,5))

axes[0].plot(losses, label='Negative ELBO')
axes[0].set_title('Time', fontsize=15)
axes[0].legend()

axes[1].plot(kls, label='KL')
axes[1].set_title('Time', fontsize=15)
axes[1].legend()

axes[2].plot(likelihood_terms, label='Likelihood Term')
axes[2].set_title('Time', fontsize=15)
axes[2].legend()
```

```
real_data_vals, final_samples_vals, data_reconstructions_vals = sess.run(
      [real_data, samples, reconstructions])
```

```
fig, axes = plt.subplots(1, 3, figsize=(3*4,4))

show_digits(axes[0], real_data_vals, 'Data')
show_digits(axes[1], data_reconstructions_vals, 'Reconstructions')
show_digits(axes[2], final_samples_vals, 'Samples')
```

```
show_latent_interpolations(lambda x: decoder(x).mean(), prior, sess)
```

## Questions about KL annealing

* What do you observe about the KL behaviour throughout training as opposed to amortized variational inference without any KL annealing?

# Constrained optimization

Instead of using KL annealing, constrained optimization can be used to automatically tuned the relative weight of the likelihood and kl terms. This removes the need to manually create an optimization schedule, which can be problem specific.

The objective now becomes:

\begin{equation}
  \text{minimize } \mathbb{E}_{p^*(x)} KL(q(z|x)||p(z)) \text{ such that }  \mathbb{E}_{p^*(x)} \mathbb{E}_{q(z|x) \log p_\theta(x|z)} > \alpha
\end{equation}

This can be solved using the use of Lagrange multipliers. The objective becomes:

\begin{equation}
  \text{minimize }  \mathbb{E}_{p^*(x)} KL(q(z|x)||p(z)) + \lambda  (\mathbb{E}_{p^*(x)} \mathbb{E}_{q(z|x)} (\alpha - \log p_\theta(x|z)))
\end{equation}

The difference compared to the KL annealing is that:

   * $\lambda$ is a learned parameter - it will be learned using stochastic gradient descent, like the network parameters.  The difference is that the lagrangian has to solve a maximization problem. You can see this intuitively: the graadient with respect to $\lambda$ in the objective above is $\mathbb{E}_{p^*(x)} \mathbb{E}_{q(z|x)} (\alpha - \log p_\theta(x|z))$. If $ \mathbb{E}_{p^*(x)} \mathbb{E}_{q(z|x)} (\alpha - \log p_\theta(x|z))> 0$, the constraint is not being satisfied, so the value of the lagrangian needs to increase. This will be done by doing gradient ascent, instead of gradient descent. Note that for $\lambda$ to be a valid lagranian in a minimization problem, it has to be positive.
   * The practicioner has to specify the hyperparameter $\alpha$, which determines the reoncstruction quality of the model.
   * the coefficient is in front of the likelihood term, not the KL term. This is mainly for convenience, as it is easier to specify the hyperparameter $\alpha$ for the likelihood (reconstruction loss).

For more assumptions made by this method, see the Karush–Kuhn–Tucker conditions.

For more information, see:
  * http://bayesiandeeplearning.org/2018/papers/33.pdf

```
tf.reset_default_graph()
```

```
real_data = make_tf_data_batch(mnist.train.images)
prior = make_prior()
decoder = standard_decoder
```

## Task: define the lagrangian multiplier

The lagarngian multiplier is a learned variable, and always has to be positive. We ideally would like to have the coefficient initialized around 1.

```
# Unlike in the KL annealing case, we learn the coefficient.
lagrangian_var = None
# Ensure that the lagrangian is positive and has stable dynamics.
lagrangian = None
```

```
# How good do we want the reconstruction loss to be?
# We can look at previous runs to get an idea what a reasonable value would be.
reconstruction_target = - 90
```

```
variational_posterior = encoder(real_data)
likelihood_term, kl_term = bound_terms(
    real_data, variational_posterior, decoder)
# Note: now the elbo is a different quanitity to what we optimize.
train_elbo = likelihood_term - kl_term
loss =  kl_term + lagrangian * (reconstruction_target - likelihood_term )
```

```
samples = decoder(prior.sample()).mean()
samples.shape.assert_is_compatible_with([BATCH_SIZE, 28, 28, 1])

reconstructions = decoder(variational_posterior.sample()).mean()
```

```
tf.trainable_variables()
```

```
# We now perform joint optimization on the encoder and decoder variables.
optimizer = tf.train.AdamOptimizer(0.001, beta1=0.9, beta2=0.9)
decoder_vars = tf.get_collection(
    tf.GraphKeys.TRAINABLE_VARIABLES, scope=DECODER_VARIABLE_SCOPE)
encoder_vars = tf.get_collection(
    tf.GraphKeys.TRAINABLE_VARIABLES, scope=ENCODER_VARIABLE_SCOPE)
autoencoder_vars = encoder_vars + decoder_vars
autoencoder_variables_update_op = optimizer.minimize(
      loss, var_list=autoencoder_vars)
```

## Task: optimize the lagrangian multiplier

The lagrange multiplier needs to maximize the loss (not minimise it). Because of this, we use a separate optimizer which minimises the negative loss.

```
lagrangian_optimizer = tf.train.GradientDescentOptimizer(0.001)

# Ensure that a variable update is followed by an update to the Lagrangian.
with tf.control_dependencies([autoencoder_variables_update_op]):
  update_op = None
```

## Training

```
sess = tf.Session()

# Initialize all variables
sess.run(tf.initialize_all_variables())
```

```
# %hide_pyerr  # - uncomment to interrupt training without a stacktrace
losses = []
kls = []
likelihood_terms = []
lagrangian_values = []

for i in xrange(TRAINING_STEPS):
  sess.run(update_op)

  if i % 100 == 0:
    iteration_loss, iteration_likelihood, iteration_kl, lag_val  = sess.run(
          [loss, likelihood_term, kl_term, lagrangian])
    print('Iteration {}. Loss {}. KL {}. Lagrangian {}'.format(
          i, iteration_loss, iteration_kl, lag_val))
    losses.append(iteration_loss)
    kls.append(iteration_kl)
    likelihood_terms.append(iteration_likelihood)
    lagrangian_values.append(lag_val)
```

```
fig, axes = plt.subplots(2, 2, figsize=(2*8, 2* 5))

axes[0, 0].plot(losses, label='Negative ELBO')
axes[0, 0].set_title('Time', fontsize=15)
axes[0, 0].legend()

axes[0, 1].plot(kls, label='KL')
axes[0, 1].set_title('Time', fontsize=15)
axes[0, 1].legend()

axes[1, 0].plot(likelihood_terms, label='Likelihood Term')
axes[1, 0].set_title('Time', fontsize=15)
axes[1, 0].legend()

axes[1, 1].plot(lagrangian_values, label='Lagrangian Values')
axes[1, 1].set_title('Time', fontsize=15)
axes[1, 1].legend()
```

### Generate samples and latent interpolations

```
real_data_vals, final_samples_vals, data_reconstructions_vals = sess.run(
      [real_data, samples, reconstructions])
```

```
fig, axes = plt.subplots(1, 3, figsize=(3*4,4))

show_digits(axes[0], real_data_vals, 'Data')
show_digits(axes[1], data_reconstructions_vals, 'Reconstructions')
show_digits(axes[2], final_samples_vals, 'Samples')
```

```
show_latent_interpolations(lambda x: decoder(x).mean(), prior, sess)
```

## Discussion about constrained optimization

* What do you observe about the behaviour of the likelihood and KL term throughout training? How is it different than in Amortized Variational inference with and without KL annealing?
* What do you notice about the behaviour of the lagrangian during training?
* How would you choose the hyperparameters of the constrained optimization problem?


**Brief Honor Code**. Do the homework on your own. You may discuss ideas with your classmates, but DO NOT copy the solutions from someone else or the Internet. If stuck, discuss with TA.

```python
%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
```

```python
import warnings
warnings.filterwarnings(action="ignore", module="scipy", message="^internal gelsd")
```

**1**. (20 points)

Consider the linear transformation $f(x)$ on $\mathbb{R}^3$ that takes the standard basis $\left\{e_1,e_2,e_3\right\}$ to $\left\{v_1,v_2,v_3\right\}$ where

$$v_1=\left(\begin{matrix}10\\-10\\16\end{matrix}\right), v_2=\left(\begin{matrix}2\\-5\\20\end{matrix}\right) \textrm {and } v_3=\left(\begin{matrix}1\\-4\\13\end{matrix}\right)$$

1. Write a matrix $A$ that represents the same linear transformation. (4 points)

2. Compute the rank of $A$ using two different methods (do not use `matrix_rank`!). (4 points)

3. Find the eigenvalues and eigenvectors of $A$. (4 points)

4. What is the matrix representation of $f$ with respect to the eigenbasis? (8 points)

```python
A = np.array([[10,-10,16], [2, -5, 20], [1, -4, 13]]).T
A
```

    array([[ 10,   2,   1],
           [-10,  -5,  -4],
           [ 16,  20,  13]])

```python
import scipy.linalg as la
```

Number of pivots after LU

```python
P, L, U = la.lu(A)
np.diag(U)
```

    array([ 16.        , -10.5       ,  -0.96428571])

Number of singular values greater than machine precision

```python
U, s, V = la.svd(A)
len(s)
```

    3

Number of non-zero eigenvalues

```python
np.real_if_close(la.eigvals(A))
```

    array([9., 3., 6.])

Determinant is not zero so must be full rank

```python
la.det(A)
```

    161.9999999999999

Find eigenvalues and eigenvectors

```python
vals, vecs = la.eig(A)
```

Matrix representation of f with respect to the eigenbasis

```python
np.diag(np.real_if_close(vals))
```

    array([[9., 0., 0.],
           [0., 3., 0.],
           [0., 0., 6.]])

**2**. (20 points)

You are given the following x-y coordinates (first column is x, second is y)

```
array([[  0.        ,   4.12306991],
       [  3.        , -15.47355729],
       [  4.        , -11.68725507],
       [  3.        , -20.33756693],
       [  5.        ,  -6.06401989],
       [  6.        ,  32.79353057],
       [  8.        ,  82.48658405],
       [  9.        ,  84.02971858],
       [  4.        ,  -1.30587276],
       [  8.        ,  68.59409878]])
```

- Find the coefficients $(a, b, c)$ of the least-squares fit of a quadratic function $y = a + bx + cx^2$ to the data.
- Plot the data and fitted curve using `matplotlib`.

```python
xs = np.array([
    [  0.        ,   4.12306991],
    [  3.        , -15.47355729],
    [  4.        , -11.68725507],
    [  3.        , -20.33756693],
    [  5.        ,  -6.06401989],
    [  6.        ,  32.79353057],
    [  8.        ,  82.48658405],
    [  9.        ,  84.02971858],
    [  4.        ,  -1.30587276],
    [  8.        ,  68.59409878]])
```

```python
X = xs[:, 0].reshape(-1,1) ** np.array([0,1,2])
y = xs[:, 1].reshape(-1,1)
```

```python
from scipy.linalg import lstsq
```

```python
(a, b, c), res, rank, s = lstsq(X, y)
```

```python
a, b, c
```

    (array([-0.35762896]), array([-11.78531232]), array([2.53125199]))

```python
xp = np.linspace(xs[:,0].min(), xs[:,0].max(), 100)
plt.plot(xs[:, 0], xs[:, 1], 'o')
plt.plot(xp, a + b*xp + c*xp**2, '-')
pass
```

**3**. (20 points)

Use the `svd` function to solve the least squares problem above, and repeat the same plot. Calculate the residual error $\lvert y - X\beta \rvert$.

```python
from scipy.linalg import svd, norm
```

```python
U, S, V = svd(X, )
```

```python
U.shape, S.shape, V.shape
```

    ((10, 10), (3,), (3, 3))

```python
S
```

    array([132.86424774,   4.91317461,   0.97590399])

```python
G = np.zeros((V.shape[0], U.shape[0]))
for i, s in enumerate(S):
    G[i,i] = 1/s
betas = V.T @ G @ U.T @ y
```

```python
res = norm(y - X @ betas)
res
```

    5.3868440487073235

```python
a, b, c = betas
xp = np.linspace(xs[:,0].min(), xs[:,0].max(), 100)
plt.plot(xs[:, 0], xs[:, 1], 'o')
plt.plot(xp, a + b*xp + c*xp**2, '-')
pass
```

**4**. (20 points)

Avoiding catastrophic cancellation.

Read the Wikipedia entry on [loss of significance](https://en.wikipedia.org/wiki/Loss_of_significance). Then answer the following problem:

The tail of the standard logistic distribution is given by $1 - F(t) = 1 - (1+e^{-t})^{-1}$.

- Define a function `f1` to calculate the tail probability of the logistic distribution using the formula given above
- Use [`sympy`](http://docs.sympy.org/latest/index.html) to find the exact value of the tail distribution (using the same symbolic formula) to 20 decimal digits
- Calculate the *relative error* of `f1` when $t = 25$ (The relative error is given by `abs(exact - approximate)/exact`)
- Rewrite the expression for the tail of the logistic distribution using simple algebra so that there is no risk of cancellation, and write a function `f2` using this formula. Calculate the *relative error* of `f2` when $t = 25$.
- How much more accurate is `f2` compared with `f1` in terms of the relative error?

```python
def f1(t):
    """Calculates tail probabilty of the logistic distribution."""
    return 1 - 1.0/(1 + np.exp(-t))

def fsymb(t, n=20):
    """Exact value to n decimal digits using symbolic algebra."""
    from sympy import exp
    return (1 - 1/(1 + exp(-t))).evalf(n=n)

def f2(t):
    """Calculates tail probabilty of the logistic distribution - no cancellation."""
    return 1/(1 + np.exp(t))

r1 = abs(fsymb(25) - f1(25))/fsymb(25)
r2 = abs(fsymb(25) - f2(25))/fsymb(25)

print("Relative error of f1:\t%.16f" % r1)
print("Relative error of f2\t%.16f" % r2)
print("f2 improvieemnt over f1\t%g" % (r1/r2))
```

    Relative error of f1:	0.0000041759147666
    Relative error of f2	0.0000000000000001
    f2 improvieemnt over f1	3.66247e+10

**5**. (50 points)

Read in `figs/elephant.jpg` as a gray-scale image. The image has $1066 \times 1600$ values. Using SVD, recreate the image with a relative error of less than 0.5%. What is the relative size of the compressed image as a percentage?

```python
from skimage import io
img = io.imread('figs/elephant.jpg', as_grey=True)
```

```python
img.shape
```

    (1066, 1600)

```python
plt.imshow(img, cmap='gray')
plt.gcf().set_frameon(False)
plt.gca().set_axis_off()
pass
```

```python
U, S, V = svd(img)
```

```python
U.shape, S.shape, V.shape
```

    ((1066, 1066), (1066,), (1600, 1600))

```python
k = np.sum(np.cumsum((S**2)/(S**2).sum()) <= 0.995)
```

```python
i = U[:,:k] @ np.diag(S[:k])@ V[:k, :]
```

```python
plt.imshow(i, cmap='gray')
plt.gcf().set_frameon(False)
plt.gca().set_axis_off()
pass
```

```python
space = U.shape[0]*k + k + k*V.shape[1]
100.0 * space / np.prod(img.shape)
```

    3.12734521575985


# Markov决策过程

## Markov决策过程模型

### 离散时间Markov决策过程

> 在离散时间智能体/环境接口中，智能体与环境交互的时刻为{0,1,2,3,...}。在时刻t，发生如下事情：
> * 智能体观测时间得$O_t$
> * 智能体根据观测环境决定做出动作$A_t$
> * 环境根据智能体动作，给予智能体奖励$R_{t+1}$，并进入下一步状态$S_{t+1}$

由此得：
* S：状态空间，状态取值的综合
* O：观测空间，观测取值的集合
* A：动作集合
  * 为了分析方便，往往采用一个包括所有可能动作的更大集合，使得每一步动作集合在数学上以同样的字母表示
* R：奖励空间，是实数集的子集

离散时间不一定是间隔相同或间隔预先设定好的时间，这里的离散时间指标t只是表示决策和动作的指标

【**轨道**】$S_0,O_0,A_0,R_1,S_1,O_1,A_1,R_2,...$

对于回合制任务，可能存在终止状态$s_{终止}$。
* 达到终止状态时，回合结束，不再有任何观测或动作
* *状态空间S中不包括终止状态*
* 强调含有终止状态的状态空间为$S^+$
* 回合制任务的轨道形式：$S_0,O_0,A_0,R_1,S_1,O_1,A_1,R_2,...,S_T = s_{终止}$
* 回合制度任务中一个回合的步数$T$：随机变量，随机过程中可以看做一个停时

> 部分不完全可观测问题：建模为部分可观测的Markov决策过程（POMDP）进行求解

【Markov决策过程模型】
* **概率**：在时间t，从状态$S_t = s$和动作$A_t = a$跳转到下一状态$S_{t+1} = s'$和动作$R_{t+1} = r$的概率为：$Pr[S_{t+1} = s',R_{t+1} = r|S_t = s,A_t = a]$
* **Markov性**：认为奖励$R_{t+1}$和下一状态$S_{t+1}$仅仅依赖于当前状态$S_t$和动作$A_t$，而不依赖于更早的状态和动作
  * 对状态的额外约束，状态必须含义可能对未来产生影响的所有过去信息

> 智能体/环境接口没有假设状态满足Markov性。Markov性是满足Markov决策过程的特点
> 可以从不满足Markov性的观测中构造满足Markov性的状态，或去学习Markov性

### 环境与动力

对于有限Markov过程，可定义：
* **动力**：$p(s',r|s,a) = Pr[S_{t+1} = s',R_{t+1} = r|S_t = s,A_t = a]$
* **状态转移概率**：$p(s'|s,a) = Pr[S_{t+1} = s'|S_t = s,A_t = a] = \sum_{r \in R} p(s',r|s,a)$
* **给定“状态-动作”的期望奖励**：$ r(s,a) = E[R_{t+1} | S_t = s,A_t = a] = \sum_{r \in R} r \sum_{s' \in S} p(s',r|s,a)$
* **给定“状态-动作-下一状态”的期望奖励**：$ r(s,a,s') = E[R_{t+1} | S_t = s,A_t = a,S_{t+1} = s'] = \sum_{r \in R} r \frac{p(s',r|s,a)}{p(s'|s,a)} $

注：$E[R_{t+1} | S_t = s,A_t = a,S_{t+1} = s'] = \sum_{r \in R} r p(r|s,s',a) = \sum_{r \in R} r \frac{p(r,s,s',a)}{p(s,s',a)} = \sum_{r \in R} r \frac{p(s',r|s,a)}{p(s'|s,a)}$

非有限Markov决策过程：定义时使用概率分布函数

### 智能体与策略

【**策略**】：从状态到动作的转移概率

对于有限Markov决策过程：$\pi: S \times A \to [0,1]$可定义为：$\pi(a|s) = Pr[A_t=a|S_t=s],s \in S, a\in A$

动作集为连续，则采用概率分布进行定义

【**确定性策略**】若某个策略$\pi$对状态集内任意状态，均有$a \in A$满足$\pi(a'|s) = 0, a' \ne a$，则称此策略为**确定性策略**。

### 奖励、回报和价值函数

【**回报**】
* 对于*回合制任务*，假设某一回合在第T步达到终止状态，则从步骤t以后的回报$G_t$可定义为未来奖励之和：$G_t = R_{t+1} + R_{t+2} +\dots +R_{T} $
* t为确定性变量，T为随机变量，因此在$G_t$定义中，不仅每一项是随机变量，而且含有的项数也是随机变量
* 对*连续性任务*，若只是简单求和，会趋于无穷大。由此引入**折扣**，此时定义回报：$G_t = R_{t+1} + \gamma R_{t+2} + \gamma^2 R_{t+3} +\dots = \sum_{\tau = 0}^{+\infty}\gamma^\tau R_{t+\tau+1} $
  其中$\gamma \in [0,1]$
* 折扣因子决定了如何在最近的奖励和未来的奖励之间进行折中。未来$\tau$步后得到的1单位奖励相当于现在的$\gamma^\tau$单位奖励。
  * $\tau = 0$：只考虑眼前利益，无视未来利益——**贪心算法**
  * $\tau = 1$：未来1单位利益和眼前1单位利益同等重要
  一般设定$\gamma \in (0,1)$，若未来每一步奖励有界，则回报也有界。
* 也可为*连续性任务*定义**平均奖励**：$\bar{R} = \lim_{t \to \infty} E[\frac1t \sum_{\tau=1}^t R_\tau]$

【**状态价值函数**$v_\pi(s)$】
* 从状态s开始采用策略$\pi$的期望回报
* $v_\pi(s) = E_\pi[G_t|S_t = s]$

【**动作价值函数**$q_\pi(s,a)$】
* 状态s采用动作a后，采用策略$\pi$的期望回报
* $q_\pi(s,a) = E_\pi[G_t|S_t=s,A_t=a]$

对于终止状态：$v_pi(s_{终止}) = 0$；$q_pi(s_{终止},a) = 0$

## Bellman期望方程

【**策略评估**】试图求解给定策略的价值函数。Bellman方程常用于策略评估

【**状态价值和动作价值的互相表示**】
* 用t时刻的动作价值函数表示t时刻的状态价值函数：$v_{pi} (s) = \sum_a \pi(a|s)q_\pi(s,a), s \in S $，推导略。
  * 空心圆圈代表状态，实心圆圈代表状态动作对，则动作价值函数表示状态价值函数的过程可用**备份图**表示
* 用t+1时刻的状态价值函数表示t时刻的动作价值函数：$q_\pi(s,a) = \sum_{s',r}p(s',r|s,a)[r+\gamma v_\pi(s'')]= r(s,a) + \gamma\sum_{s'} p(s'|s,a)v_\pi(s'),s \in S, a \in A $

以上推导可灵活运用全概率公式、贝叶斯公式、概率乘法公式等内容灵活计算。

### Bellman方程求解

具体例子和参数设定参考原书第二章的表2-1和图2-1，此时可通过[`sympy`](https://docs.sympy.org/latest/index.html)求解线性方程组，得价值策略

```plaintext
import sympy
import gym
import numpy as np
import scipy
from sympy import symbols
sympy.init_printing()
```

求解Bellman期望方程：

```plaintext
v_hungry, v_full = symbols('v_hungry v_full')
q_hungry_eat, q_hungry_none, q_full_eat, q_full_none = \
        symbols('q_hungry_eat q_hungry_none q_full_eat q_full_none')
alpha, beta, gamma = symbols('alpha beta gamma')
x, y = symbols('x y')   # 相关变量
```

此方程组为含有字母的线性方程组，可利用`solve_linear_system()`求解

```plaintext
system = sympy.Matrix((
        (1, 0, x-1, -x, 0, 0, 0),
        (0, 1, 0, 0, -y, y-1, 0),
        (-gamma, 0, 1, 0, 0, 0, -2),
        ((alpha-1)*gamma, -alpha*gamma, 0, 1, 0, 0, 4*alpha-3),
        (-beta*gamma, (beta-1)*gamma, 0, 0, 1, 0, -4*beta+2),
        (0, -gamma, 0, 0, 0, 1, 1) ))   # 增广矩阵
sympy.solve_linear_system(system,v_hungry, v_full,
        q_hungry_none, q_hungry_eat, q_full_none, q_full_eat)
```

【**总结**】
用`sympy`求解线性方程组：
1. `sympy.init_printing()`可使得最后结果以合适形式打印
2. 声明相关量为symbol
3. 构造`Matrix`对象，具体为线性方程组的增广矩阵
4. 调用函数求解

## 最优策略及其性质

定义一个偏序关系：对于两个策略$\pi$和$\pi'$，若对于任意$s \in S$都满足$v_\pi(s) \le v_{\pi'}(s)$，则称策略$\pi$小于等于$\pi'$，记作$\pi \le \pi'$

### 最优策略和最优价值函数

【**最优策略**】
* 对于一个动力，若状态空间$A(s)$是闭集，那么就存在一个策略$\pi_*$，使得所有的策略都小于等于这个策略，则称$\pi_*$为最优策略
* **最优价值函数**：最优策略的价值函数
  * **最优状态价值函数**：$v_* = \max_{\pi} v_\pi(s),s \in S$
  * **最优动作价值函数**：$q_*(s,a) = \max_{\pi} q_\pi(s,a)$

最优策略可能存在多个，但其价值函数是相同的，任取一个最优策略来考察也不是一般性。其中一种选取方法是选择这样的一种**确定性策略**：

$$\pi_*(s) = \mathop{\arg\max}_{a \in A} q_*(s,a),s \in S $$

若有多个动作a使得$q_*(s,a)$取得最大值，任选一个动作即可。

### Bellman最优方程

*最优价值函数的重要性质：Bellman最优方程*

最优策略的性质：
$$
    \pi_*(a|s) = \begin{cases}
        1, &a = \mathop{\arg\max}_{a \in A} q_*(s,a')\\
        0,&others
    \end{cases}
$$

* 最优动作价值函数表示最优状态价值函数：$v_*(s) = \max_{a\in A} q_*(s,a)$
* 最优状态价值函数表示最优动作价值函数：$q_*(s,a) = r(s,a) + \gamma\sum_{s'}p(s'|s,a)v_*(s,a) = \sum_{s',r}p(s',r|s,a)[r+\gamma v_*(s')]$

由此得**Bellman最优方程**：
* 最优状态价值函数表示：$v_*(s) = \max_{a \in A} \left[ r(s,a) + \gamma\sum_{s'}p(s'|s,a)v_*(s,a) \right]$
* 最优动作价值函数表示：$q_*(s,a) = r(s,a) + \gamma\sum_{s'}p(s'|s,a)\max_{a'} q_*(s',a')$

用`sympy`求解。由于有max，为非线性方程，可通过分类化为线性方程

```plaintext
xy_tuples = ((0, 0), (1, 0), (0, 1), (1, 1))
for x, y in xy_tuples:  # 分类讨论
    system = sympy.Matrix((
            (1, 0, x-1, -x, 0, 0, 0),
            (0, 1, 0, 0, -y, y-1, 0),
            (-gamma, 0, 1, 0, 0, 0, -2),
            ((alpha-1)*gamma, -alpha*gamma, 0, 1, 0, 0, 4*alpha-3),
            (-beta*gamma, (beta-1)*gamma, 0, 0, 1, 0, -4*beta+2),
            (0, -gamma, 0, 0, 0, 1, 1) ))
    result = sympy.solve_linear_system(system,
            v_hungry, v_full,
            q_hungry_none, q_hungry_eat, q_full_none, q_full_eat, simplification=True)
    msgx = 'v(饿) = q(饿,{}吃)'.format('' if x else '不')
    msgy = 'v(饱) = q(饱,{}吃)'.format('不' if y else '')
    print('==== {}, {} ==== x = {}, y = {} ===='.format(msgx, msgy, x, y))
    display(result)
```

对于**给定数值**情况，常常将$v_*(s) = \max_{a\in A} q_*(s,a)$松弛为$v_*(s) \ge q_*(s,a)$，并消去q_*(s,a)以减少决策变量，得新的线性规划：
$$
    \begin{aligned}
        &\min &\sum_{s \in S} c(s)v(s)\\
        &\mathrm{over}&v(s),s \in S\\
        &\mathrm{s.t.}&v(s) \ge r(s,a) + \gamma\sum_{s'}p(s'|s,a)v_*(s,a)
    \end{aligned}
$$

$c(s)$为一组任意取值的整实数。Bellman最优方程的解显然会在线性规划的可行域内。此线性规划最优解满足Bellman最优方程。

### 用Bellman最优方程解求解最优策略

**最优策略**：对于每个状态$s$，总是选择让$q_*(s,a)$最大的动作$a$

**实际困难**：
* 难以列出Bellman最优方程
* 难以求解Bellman最优方程

## 案例：悬崖寻路

使用gym库中的悬崖寻路问题（CliffWalking-v0）作为案例分析，该环境是一个有限Markov决策过程。具体解释见原书2.4节

### 环境使用

* 离散Markov决策过程，状态为{0,1,...,46}的整数值（带上终止状态则加上47），动作为{0,1,2,3}的int值，分别表示上右下左。奖励取自{-1,-100}。

```plaintext
env = gym.make('CliffWalking-v0')
env.seed(0)
print('观测空间 = {}'.format(env.observation_space))
print('动作空间 = {}'.format(env.action_space))
print('状态数量 = {}, 动作数量 = {}'.format(env.nS, env.nA))
print('地图大小 = {}'.format(env.shape))
```

    观测空间 = Discrete(48)
    动作空间 = Discrete(4)
    状态数量 = 48, 动作数量 = 4
    地图大小 = (4, 12)

运行一回合：

```plaintext
def play_once(env,policy):
    total_reward = 0
    state = env.reset()
    while True:
        loc = np.unravel_index(state, env.shape)
        print('状态 = {}, 位置 = {}'.format(state, loc), end=' ')
        action = np.random.choice(env.nA, p=policy[state])
        next_state, reward, done, _ = env.step(action)
        print('动作 = {}, 奖励 = {}'.format(action, reward))
        total_reward += reward
        if done:
            break
        state = next_state
    return total_reward
```

最优策略运行一回合

```plaintext
actions = np.ones(env.shape,dtype=int)
actions[-1, :] = 0
actions[:, -1] = 2
optimal_policy = np.eye(4)[actions.reshape(-1)]

total_reward = play_once(env, optimal_policy)
print('回合奖励 = {}'.format(total_reward))
```

    状态 = 36, 位置 = (3, 0) 动作 = 0, 奖励 = -1
    状态 = 24, 位置 = (2, 0) 动作 = 1, 奖励 = -1
    状态 = 25, 位置 = (2, 1) 动作 = 1, 奖励 = -1
    状态 = 26, 位置 = (2, 2) 动作 = 1, 奖励 = -1
    状态 = 27, 位置 = (2, 3) 动作 = 1, 奖励 = -1
    状态 = 28, 位置 = (2, 4) 动作 = 1, 奖励 = -1
    状态 = 29, 位置 = (2, 5) 动作 = 1, 奖励 = -1
    状态 = 30, 位置 = (2, 6) 动作 = 1, 奖励 = -1
    状态 = 31, 位置 = (2, 7) 动作 = 1, 奖励 = -1
    状态 = 32, 位置 = (2, 8) 动作 = 1, 奖励 = -1
    状态 = 33, 位置 = (2, 9) 动作 = 1, 奖励 = -1
    状态 = 34, 位置 = (2, 10) 动作 = 1, 奖励 = -1
    状态 = 35, 位置 = (2, 11) 动作 = 2, 奖励 = -1
    回合奖励 = -13

求解Bellman期望方程：

```plaintext
def evaluate_bellman(env, policy, gamma=1.):
    a, b = np.eye(env.nS), np.zeros((env.nS))
    for state in range(env.nS - 1):
        for action in range(env.nA):
            pi = policy[state][action]
            for p, next_state, reward, done in env.P[state][action]:
                a[state, next_state] -= (pi * gamma * p)
                b[state] += (pi * reward * p)
    v = np.linalg.solve(a, b)
    q = np.zeros((env.nS, env.nA))
    for state in range(env.nS - 1):
        for action in range(env.nA):
            for p, next_state, reward, done in env.P[state][action]:
                q[state][action] += ((reward + gamma * v[next_state]) * p)
    return v, q
```

评估随机策略的均值

```plaintext
policy = np.random.uniform(size=(env.nS, env.nA))
policy = policy / np.sum(policy, axis=1)[:, np.newaxis]

state_values, action_values = evaluate_bellman(env, policy)
print('状态价值 = {}'.format(state_values))
print('动作价值 = {}'.format(action_values))
```

    状态价值 = [-227182.96580657 -227026.10344732 -226687.624366   -226415.66297675
     -223753.30004204 -220264.99683447 -211898.74623424 -195221.65576656
     -177271.79127495 -122831.66464187 -106585.18274491 -104349.01792373
     -227189.99058898 -227065.19921615 -226538.68719358 -225761.38463778
     -223869.03416183 -221225.07098338 -219071.20535663 -208545.40124377
     -188944.95053559 -118275.59108578 -100330.22350904  -92682.09377381
     -227356.56076217 -227367.00749714 -226919.77958566 -226624.73047364
     -225492.40168571 -224403.24209048 -224790.21582123 -216857.22745747
     -209439.77779116 -173721.42436915  -94878.87815452  -52646.12961402
     -227877.02015437 -227882.43625112 -227579.76453949 -227799.72026599
     -227365.26619454 -227902.10974978 -227687.75321092 -223890.90474862
     -225382.28337005 -214120.53897423 -163388.85618291       0.        ]
    动作价值 = [[-2.27183966e+05 -2.27027103e+05 -2.27190991e+05 -2.27183966e+05]
     [-2.27027103e+05 -2.26688624e+05 -2.27066199e+05 -2.27183966e+05]
     [-2.26688624e+05 -2.26416663e+05 -2.26539687e+05 -2.27027103e+05]
     [-2.26416663e+05 -2.23754300e+05 -2.25762385e+05 -2.26688624e+05]
     [-2.23754300e+05 -2.20265997e+05 -2.23870034e+05 -2.26416663e+05]
     [-2.20265997e+05 -2.11899746e+05 -2.21226071e+05 -2.23754300e+05]
     [-2.11899746e+05 -1.95222656e+05 -2.19072205e+05 -2.20265997e+05]
     [-1.95222656e+05 -1.77272791e+05 -2.08546401e+05 -2.11899746e+05]
     [-1.77272791e+05 -1.22832665e+05 -1.88945951e+05 -1.95222656e+05]
     [-1.22832665e+05 -1.06586183e+05 -1.18276591e+05 -1.77272791e+05]
     [-1.06586183e+05 -1.04350018e+05 -1.00331224e+05 -1.22832665e+05]
     [-1.04350018e+05 -1.04350018e+05 -9.26830938e+04 -1.06586183e+05]
     [-2.27183966e+05 -2.27066199e+05 -2.27357561e+05 -2.27190991e+05]
     [-2.27027103e+05 -2.26539687e+05 -2.27368007e+05 -2.27190991e+05]
     [-2.26688624e+05 -2.25762385e+05 -2.26920780e+05 -2.27066199e+05]
     [-2.26416663e+05 -2.23870034e+05 -2.26625730e+05 -2.26539687e+05]
     [-2.23754300e+05 -2.21226071e+05 -2.25493402e+05 -2.25762385e+05]
     [-2.20265997e+05 -2.19072205e+05 -2.24404242e+05 -2.23870034e+05]
     [-2.11899746e+05 -2.08546401e+05 -2.24791216e+05 -2.21226071e+05]
     [-1.95222656e+05 -1.88945951e+05 -2.16858227e+05 -2.19072205e+05]
     [-1.77272791e+05 -1.18276591e+05 -2.09440778e+05 -2.08546401e+05]
     [-1.22832665e+05 -1.00331224e+05 -1.73722424e+05 -1.88945951e+05]
     [-1.06586183e+05 -9.26830938e+04 -9.48798782e+04 -1.18276591e+05]
     [-1.04350018e+05 -9.26830938e+04 -5.26471296e+04 -1.00331224e+05]
     [-2.27190991e+05 -2.27368007e+05 -2.27878020e+05 -2.27357561e+05]
     [-2.27066199e+05 -2.26920780e+05 -2.27977020e+05 -2.27357561e+05]
     [-2.26539687e+05 -2.26625730e+05 -2.27977020e+05 -2.27368007e+05]
     [-2.25762385e+05 -2.25493402e+05 -2.27977020e+05 -2.26920780e+05]
     [-2.23870034e+05 -2.24404242e+05 -2.27977020e+05 -2.26625730e+05]
     [-2.21226071e+05 -2.24791216e+05 -2.27977020e+05 -2.25493402e+05]
     [-2.19072205e+05 -2.16858227e+05 -2.27977020e+05 -2.24404242e+05]
     [-2.08546401e+05 -2.09440778e+05 -2.27977020e+05 -2.24791216e+05]
     [-1.88945951e+05 -1.73722424e+05 -2.27977020e+05 -2.16858227e+05]
     [-1.18276591e+05 -9.48798782e+04 -2.27977020e+05 -2.09440778e+05]
     [-1.00331224e+05 -5.26471296e+04 -2.27977020e+05 -1.73722424e+05]
     [-9.26830938e+04 -5.26471296e+04 -1.00000000e+00 -9.48798782e+04]
     [-2.27357561e+05 -2.27977020e+05 -2.27878020e+05 -2.27878020e+05]
     [-2.27368007e+05 -2.27977020e+05 -2.27977020e+05 -2.27878020e+05]
     [-2.26920780e+05 -2.27977020e+05 -2.27977020e+05 -2.27977020e+05]
     [-2.26625730e+05 -2.27977020e+05 -2.27977020e+05 -2.27977020e+05]
     [-2.25493402e+05 -2.27977020e+05 -2.27977020e+05 -2.27977020e+05]
     [-2.24404242e+05 -2.27977020e+05 -2.27977020e+05 -2.27977020e+05]
     [-2.24791216e+05 -2.27977020e+05 -2.27977020e+05 -2.27977020e+05]
     [-2.16858227e+05 -2.27977020e+05 -2.27977020e+05 -2.27977020e+05]
     [-2.09440778e+05 -2.27977020e+05 -2.27977020e+05 -2.27977020e+05]
     [-1.73722424e+05 -2.27977020e+05 -2.27977020e+05 -2.27977020e+05]
     [-9.48798782e+04 -1.00000000e+00 -2.27977020e+05 -2.27977020e+05]
     [ 0.00000000e+00  0.00000000e+00  0.00000000e+00  0.00000000e+00]]

评估最优策略均值

```plaintext
optimal_state_values, optimal_action_values = evaluate_bellman(env, optimal_policy)
print('最优状态价值 = {}'.format(optimal_state_values))
print('最优动作价值 = {}'.format(optimal_action_values))
```

    最优状态价值 = [-14. -13. -12. -11. -10.  -9.  -8.  -7.  -6.  -5.  -4.  -3. -13. -12.
     -11. -10.  -9.  -8.  -7.  -6.  -5.  -4.  -3.  -2. -12. -11. -10.  -9.
      -8.  -7.  -6.  -5.  -4.  -3.  -2.  -1. -13. -12. -11. -10.  -9.  -8.
      -7.  -6.  -5.  -4.  -3.   0.]
    最优动作价值 = [[ -15.  -14.  -14.  -15.]
     [ -14.  -13.  -13.  -15.]
     [ -13.  -12.  -12.  -14.]
     [ -12.  -11.  -11.  -13.]
     [ -11.  -10.  -10.  -12.]
     [ -10.   -9.   -9.  -11.]
     [  -9.   -8.   -8.  -10.]
     [  -8.   -7.   -7.   -9.]
     [  -7.   -6.   -6.   -8.]
     [  -6.   -5.   -5.   -7.]
     [  -5.   -4.   -4.   -6.]
     [  -4.   -4.   -3.   -5.]
     [ -15.  -13.  -13.  -14.]
     [ -14.  -12.  -12.  -14.]
     [ -13.  -11.  -11.  -13.]
     [ -12.  -10.  -10.  -12.]
     [ -11.   -9.   -9.  -11.]
     [ -10.   -8.   -8.  -10.]
     [  -9.   -7.   -7.   -9.]
     [  -8.   -6.   -6.   -8.]
     [  -7.   -5.   -5.   -7.]
     [  -6.   -4.   -4.   -6.]
     [  -5.   -3.   -3.   -5.]
     [  -4.   -3.   -2.   -4.]
     [ -14.  -12.  -14.  -13.]
     [ -13.  -11. -113.  -13.]
     [ -12.  -10. -113.  -12.]
     [ -11.   -9. -113.  -11.]
     [ -10.   -8. -113.  -10.]
     [  -9.   -7. -113.   -9.]
     [  -8.   -6. -113.   -8.]
     [  -7.   -5. -113.   -7.]
     [  -6.   -4. -113.   -6.]
     [  -5.   -3. -113.   -5.]
     [  -4.   -2. -113.   -4.]
     [  -3.   -2.   -1.   -3.]
     [ -13. -113.  -14.  -14.]
     [ -12. -113. -113.  -14.]
     [ -11. -113. -113. -113.]
     [ -10. -113. -113. -113.]
     [  -9. -113. -113. -113.]
     [  -8. -113. -113. -113.]
     [  -7. -113. -113. -113.]
     [  -6. -113. -113. -113.]
     [  -5. -113. -113. -113.]
     [  -4. -113. -113. -113.]
     [  -3.   -1. -113. -113.]
     [   0.    0.    0.    0.]]

求解Bellman最优方程

```plaintext
def optimal_bellman(env, gamma=1.):
    p = np.zeros((env.nS, env.nA, env.nS))
    r = np.zeros((env.nS, env.nA))
    for state in range(env.nS - 1):
        for action in range(env.nA):
            for prob, next_state, reward, done in env.P[state][action]:
                p[state, action, next_state] += prob
                r[state, action] += (reward * prob)
    c = np.ones(env.nS)
    a_ub = gamma * p.reshape(-1, env.nS) - \
            np.repeat(np.eye(env.nS), env.nA, axis=0)
    b_ub = -r.reshape(-1)
    a_eq = np.zeros((0, env.nS))
    b_eq = np.zeros(0)
    bounds = [(None, None),] * env.nS
    res = scipy.optimize.linprog(c, a_ub, b_ub, bounds=bounds,
            method='interior-point')
    v = res.x
    q = r + gamma * np.dot(p, v)
    return v, q
```

```plaintext
optimal_state_values, optimal_action_values = optimal_bellman(env)
print('最优状态价值 = {}'.format(optimal_state_values))
print('最优动作价值 = {}'.format(optimal_action_values))
optimal_actions = optimal_action_values.argmax(axis=1)
print('最优策略 = {}'.format(optimal_actions))
```

    最优状态价值 = [-1.40000000e+01 -1.30000000e+01 -1.20000000e+01 -1.10000000e+01
     -1.00000000e+01 -9.00000000e+00 -8.00000000e+00 -7.00000000e+00
     -6.00000000e+00 -5.00000000e+00 -4.00000000e+00 -3.00000000e+00
     -1.30000000e+01 -1.20000000e+01 -1.10000000e+01 -1.00000000e+01
     -9.00000000e+00 -8.00000000e+00 -7.00000000e+00 -6.00000000e+00
     -5.00000000e+00 -4.00000000e+00 -3.00000000e+00 -2.00000000e+00
     -1.20000000e+01 -1.10000000e+01 -1.00000000e+01 -9.00000000e+00
     -8.00000000e+00 -7.00000000e+00 -6.00000000e+00 -5.00000000e+00
     -4.00000000e+00 -3.00000000e+00 -2.00000000e+00 -1.00000000e+00
     -1.30000000e+01 -1.20000000e+01 -1.10000000e+01 -1.00000000e+01
     -9.00000000e+00 -8.00000000e+00 -7.00000000e+00 -6.00000000e+00
     -5.00000000e+00 -4.00000000e+00 -9.99999999e-01  1.82319061e-11]
    最优动作价值 = [[ -14.99999999  -13.99999999  -13.99999999  -14.99999999]
     [ -13.99999999  -13.          -13.          -14.99999999]
     [ -13.          -12.          -12.          -13.99999999]
     [ -12.          -11.          -11.          -13.        ]
     [ -11.          -10.          -10.          -12.        ]
     [ -10.           -9.           -9.          -11.        ]
     [  -9.           -8.           -8.          -10.        ]
     [  -8.           -7.           -7.           -9.        ]
     [  -7.           -6.           -6.           -8.        ]
     [  -6.           -5.           -5.           -7.        ]
     [  -5.           -4.           -4.           -6.        ]
     [  -4.           -4.           -3.           -5.        ]
     [ -14.99999999  -13.          -13.          -13.99999999]
     [ -13.99999999  -12.          -12.          -13.99999999]
     [ -13.          -11.          -11.          -13.        ]
     [ -12.          -10.          -10.          -12.        ]
     [ -11.           -9.           -9.          -11.        ]
     [ -10.           -8.           -8.          -10.        ]
     [  -9.           -7.           -7.           -9.        ]
     [  -8.           -6.           -6.           -8.        ]
     [  -7.           -5.           -5.           -7.        ]
     [  -6.           -4.           -4.           -6.        ]
     [  -5.           -3.           -3.           -5.        ]
     [  -4.           -3.           -2.           -4.        ]
     [ -13.99999999  -12.          -14.          -13.        ]
     [ -13.          -11.         -113.          -13.        ]
     [ -12.          -10.         -113.          -12.        ]
     [ -11.           -9.         -113.          -11.        ]
     [ -10.           -8.         -113.          -10.        ]
     [  -9.           -7.         -113.           -9.        ]
     [  -8.           -6.         -113.           -8.        ]
     [  -7.           -5.         -113.           -7.        ]
     [  -6.           -4.         -113.           -6.        ]
     [  -5.           -3.         -113.           -5.        ]
     [  -4.           -2.         -113.           -4.        ]
     [  -3.           -2.           -1.           -3.        ]
     [ -13.         -113.          -14.          -14.        ]
     [ -12.         -113.         -113.          -14.        ]
     [ -11.         -113.         -113.         -113.        ]
     [ -10.         -113.         -113.         -113.        ]
     [  -9.         -113.         -113.         -113.        ]
     [  -8.         -113.         -113.         -113.        ]
     [  -7.         -113.         -113.         -113.        ]
     [  -6.         -113.         -113.         -113.        ]
     [  -5.         -113.         -113.         -113.        ]
     [  -4.         -113.         -113.         -113.        ]
     [  -3.           -1.         -113.         -113.        ]
     [   0.            0.            0.            0.        ]]
    最优策略 = [2 1 1 1 1 1 1 1 1 1 1 2 1 1 1 1 1 1 1 1 1 1 1 2 1 1 1 1 1 1 1 1 1 1 1 2 0
     0 0 0 0 0 0 0 0 0 1 0]


# Iteration and Stability

**ENGSCI233: Computational Techniques and Computer Systems**

*Department of Engineering Science, University of Auckland*

# 0 What's in this Notebook?

Differential equations describe how physical systems change over time. If we want to predict those changes into the future, we need to solve the equations. Most can't be solved analytically.

The next best option is a numerical method. You have already seen Euler's method - its not very accurate, although the philosophy is spot on. Estimate the derivative, use it to take a jump forward, re-evaluate the derivative. Here, we'll explore some more complex methods, and how to measure their performance.

You need to know:
- How to use jumps back and forth - predictor and corrector steps - when implementing an improved version of Euler's method. This concept generalises to higher order Runge-Kutta methods.
- The accuracy of these methods depends on how large a step we take. We can use a convergence test to decide how small those steps need to be for a decent result.
- Sometimes these methods are unstable. That is both a drawback of the method and but also depends on the difficulty of the problem it has been applied to (stiffness).

Note for user: **equations format poorly** in this notebook. If they are running off the screen "re-run" the cell containing the equation (click on the cell, hit Enter, then Ctrl+Enter).

```python
# imports and environment: this cell must be executed before any other in the notebook
%matplotlib inline
from matplotlib import pyplot as plt
```

## 1 ODE solution methods

***Different methods to approximately solve Ordinary Differential Equations***

Consider the initial value problem $\frac{dy}{dx} = (1 + xy)^{2}, y(0) = 1$.  This equation is **non-separable** and **nonlinear** which makes it difficult to solve.  An approximate solution can still be generated using a numerical method.

In this module, we will consider **first-order** ODEs of the form $\frac{dy}{dx} = f(x,y)$.  Higher-order ODEs can be written as a system of first-order ODEs so this approach can be applied to almost any ODE. Furthermore, we will restrict our interest to **initial value problems**, i.e., cases where information about $y$ (and its derivatives) is known at some starting **reference point** $x^{(0)}$.

The numerical methods we will consider are all **step-by-step** methods.  This means that we will start from the **known solution** $y^{(0)}$ at $x^{(0)}$ and proceed **stepwise**.  First, we take a step of length $h$ and **estimate a solution** $y^{(1)}$ at $x^{(1)} = x^{(0)} + h$.  Then, we use the value $y^{(1)}$ as a **pseudo-initial value** to estimate a solution $y^{(2)}$ at $x^{(2)} = x^{(1)} + h = x^{(0)} + 2h$.  Proceeding in this manner we will be able to construct an approximate solution for $y$ over a specified range of $x$.

### 1.1 Euler method

All the methods in this module match the derivatives in a **Taylor's expansion** over a finite region to some order. The omission of **higher-order** Taylor's series terms in the approximation means that each method will have an associated **error**.  The **Euler method** is the simplest in that it only approximates the first derivative of the function. It does this using a first-order **finite difference** approximation to the derivative.

The finite difference approximation to the derivative is:

\begin{equation}
  \frac{dy}{dx} \thickapprox \dfrac{\Delta y}{\Delta x} = \dfrac{y^{(k+1)} - y^{(k)}}{h} = f\left(x^{(k)},y^{(k)}\right)
\end{equation}

This can be rearranged to give

\begin{equation}
  y^{(k+1)} = y^{(k)} + h\,f\left(x^{(k)},y^{(k)}\right)
\end{equation}

where $f$ is a known function for the derivative.

***Execute the cell below for a demonstration of the Euler method for the ODE $dy/dx = (1+xy)^2$***

```python
# import functions
from ipywidgets import interact
from iteration233 import plot_euler_elements

def euler_figure(steps=0.5, h = 0.1):
    # create figure
    f,(ax1, ax2) = plt.subplots(1,2)
    f.set_size_inches(12,5)

    # display plot
    plot_euler_elements(ax1, ax2, steps, h)

# run the interactive figure environment
interact(euler_figure, steps = (0.5,10,0.5), h = (0.02,0.2,0.02));
```

***How does error in the approximation change with the step-size, $h$? Describe any trade-offs.***

> *~ your answer here ~*

***What type of error is introduced by truncating the Taylor series to a finite number of terms?***

> *~ your answer here ~*

The exercise below attempts to solve the ODE

\begin{equation}
\frac{dy}{dx} = \sin\left(a\sin(x)\sqrt{x}+\frac{\cos(bx)}{x+1}\right)
\end{equation}

However, we should be mindful of the tradeoff between **efficiency and accuracy** (or effort and error).

***Execute the cell below. Experiment with the input boxes.***

***In general, how many steps are required to get the error below 5%?***

```python
# import functions
from ipywidgets import fixed
from iteration233 import plot_tradeoffs_elements, tradeoffs_setup

# function parameters
p = [8, 8.5]
def tradeoffs_figure(steps, predict_value, p):
    # create figure
    f,ax = plt.subplots(1,1)
    f.set_size_inches([16, 6])

    # display plot
    plot_tradeoffs_elements(ax, steps, predict_value, p)

# run the interactive figure environment
box1,box2 = tradeoffs_setup()
interact(tradeoffs_figure, steps = box1, predict_value = box2, p = fixed(p));
```

### 1.2 Improved Euler method

The Euler method uses the slope (i.e., the value of $y'$) at $y^{(k)}$ to predict $y^{(k+1)}$.  If the slope is **changing significantly** over the step length then the Euler method will give an **inaccurate** solution.  An improved method can be devised by **averaging** the slopes at $y^{(k)}$ and $y^{(k+1)}$.

The slope at $y^{(k)}$ is $f\left(x^{(k)}, y^{(k)}\right)$. A **prediction** of the solution at $x^{(k+1)}$ is $y^{(k+1)}_p=y^{(k)}+h\,f\left(x^{(k)}, y^{(k)}\right)$. The slope at $y^{(k+1)}$ is $f\left(x^{(k+1)}, y^{(k+1)}_p\right)$.

Using the average slope in the previous definition of the Euler method yields the **Improved Euler method**:

\begin{equation}
\begin{split}
y^{(k+1)} &= y^{(k)} + h \times \text{average slope} \\
&= y^{(k)} + \dfrac{h}{2} \left[f\left(x^{(k)},y^{(k)}\right) + f\left(x^{(k)}+h,y^{(k)} + h\,f\left(x^{(k)},y^{(k)}\right)\right)\right].
\end{split}
\end{equation}

The Improved Euler Method is an example of a **predictor-corrector** method. First, the value $y^{(k+1)}_p$ is predicted using an Euler step. This allows extra slope information to be determined which can be used to make a corrected estimate of $y^{(k+1)}$.

The Improved Euler Method is a **second-order** method, i.e., truncation error $\sim\mathcal{O}(h^{3})$. The method is computationally **more expensive** than the Euler Method but it provides **more accurate** solutions.

***Execute the cell below for a demonstration of the Improved Euler method for the ODE $dy/dx = (1+xy)^2$***

```python
# import functions
from iteration233 import plot_improved_elements

def improved_euler_figure(steps=6.25, compare_euler=False):
    # create figure
    f,ax = plt.subplots(1,1)
    f.set_size_inches([10,4])

    # display plot
    plot_improved_elements(ax, steps, compare_euler)

# run the interactive figure environment
interact(improved_euler_figure, steps = (0.25,9,0.25), compare_euler = False);
```

### 1.3 Runge-Kutta methods

Could you design a better method than improved Euler, that builds complexity in the same manner Euler -> improved Euler? One thing we could do is evaluate a derivative at the mid-point of the step and use that information to compute a **weighted** derivative, e.g.,

\begin{equation}
\begin{split}
f_0 & = f\left(x^{(k)}, y^{(k)}\right) \\
f_1 & = f\left(x^{(k)}+\frac{h}{2}, y^{(k)}+\frac{h}{2}f_0\right) \\
f_2 & = f\left(x^{(k)}+h, y^{(k)}-hf_0+2hf_1\right) \\
y^{(k+1)} &= y^{(k)} + \dfrac{h}{6} \left(f_0 + 4f_1 + f_2\right).
\end{split}
\end{equation}

This is a third order scheme with truncation error $\sim\mathcal{O}(h^{4})$.

The schemes get more complicated, although its diminishing returns beyond 4th order. However, they all fall under a broad category called Runge-Kutta methods. A particularly popular implementation is called **RK45**.

## 2 Accuracy

***Measuring how close a numerical solution is to the true solution.***

Numerical computation provides a number of opportunities for **error** to creep in. Improvements to algorithms can **minimize** this error - usually at the **cost** of execution time - but can **never eliminate** it entirely. It is therefore useful to **estimate** the error that accompanies a particular solution method.

### 2.1 Euler method

If $y(x^{(k)}+h)$ is the exact value at $x^{(k)}+h$ then, by Taylor's series:

\begin{equation}
  \begin{split}
    y\left(x^{(k)}+h\right) &= y\left(x^{(k)}\right) + h\,\frac{dy}{dx}\left(x^{(k)}\right) + \dfrac{h^{2}}{2!} \frac{d^2y}{dx^2}\left(x^{(k)}\right) +
    \dfrac{h^{3}}{3!} \frac{d^3y}{dx^3}\left(x^{(k)}\right)+ \ldots \\
    &= \text{Euler's Method estimate} + \mathcal{O}(h^{2})
  \end{split}
\end{equation}

and thus the **truncation error** at each step involves terms that are proportional to $h^{2}$ or higher powers.

If we halve the step length then the error per step is proportional to $\left(\dfrac{h}{2}\right)^{2} = \dfrac{h^{2}}{4}$. However, to get to the same point we must do twice as many steps. Thus the total error is $E \varpropto 2 \left(\dfrac{h^{2}}{4}\right)= \dfrac{h^{2}}{2}$. Thus, by **halving the step length we halve the error** (and by doubling the step length we double the error). Because error for Euler's method is proportional to step length, it is termed a **first-order method**. In general, an $n$-th order method will have truncation error of $\mathcal{O}(h^{n+1})$.

### 2.1 Improved Euler method

The Improved Euler method is **second order accurate**, i.e., truncation error $\sim \mathcal{O}(h^{3})$, but this improvement over the Euler method comes at the cost of **greater computational expense**.

From the formulation of the Improved Euler method, the derivative estimated at the **corrector step** can be expanded (applying a first order Taylor expansion, first for $x$, then for $y$)

\begin{equation}
\begin{split}
f\left(x^{(k)}+h,y^{(k)} + h\,f\left(x^{(k)},y^{(k)}\right)\right) &=&\,f\left(x^{(k)},y^{(k)} + h\,f\left(x^{(k)},y^{(k)}\right)\right) + h \frac{\partial f}{\partial x}\left(x^{(k)},y^{(k)} + h\,f\left(x^{(k)},y^{(k)}\right)\right) \\
&=&\,f\left(x^{(k)},y^{(k)}\right) +  h\,f\left(x^{(k)},y^{(k)}\right)\,\frac{\partial f}{\partial y}\left(x^{(k)},y^{(k)}\right) + h \left(\frac{\partial f}{\partial x}\left(x^{(k)},y^{(k)}\right) + h\,f\left(x^{(k)},y^{(k)}\right)\,\frac{\partial^2 f}{\partial x\partial y}\left(x^{(k)},y^{(k)}\right)\right)  \\
&=&\,f\left(x^{(k)},y^{(k)}\right) + h \frac{df}{dx}\left(x^{(k)}, y^{(k)}\right) + h^2\,f\left(x^{(k)},y^{(k)}\right)\,\frac{\partial^2 f}{\partial x\partial y}\left(x^{(k)},y^{(k)}\right),
\end{split}
\end{equation}

where the total derivative is expressed $\frac{df}{dx}=\frac{\partial f}{\partial x}+\frac{\partial y}{\partial x}\frac{\partial f}{\partial y}$. Substituting this expression into the Improved Euler formula we get

\begin{equation}
\begin{split}
y^{(k+1)} &=&\,y^{(k)} + h\,f\left(x^{(k)},y^{(k)}\right) + \frac{h^2}{2}\frac{df}{dx}\left(x^{(k)}, y^{(k)}\right) + \mathcal{O}(h^3)\\
&=&\,\text{Improved Euler method estimate} + \mathcal{O}(h^3).
\end{split}
\end{equation}

### 2.3 Runge-Kutta methods

We won't bother with a full derivation of truncation error for the higher order methods. It is more important to appreciate the concept: if a method is $n^{th}$-order, then it has truncation error $\sim\mathcal{O}(h^{n+1})$. Halving the step size means that error goes from $h^{n+1} \rightarrow h^{n+1}/2^{n+1}$, but we need to take twice the number of steps, so the total error is $h^{n+1}/2^n$. So the accuracy is improved, but that comes at a cost of having to evaluate more derivatives at ever more midpoints.

***Execute the cell below to compare accuracy of the different schemes for different step sizes.***

```python
# import functions
from iteration233 import plot_accuracy_elements

# function parameters
def accuracy_figure(more_steps=0):
    # create figure
    f,(ax,ax2) = plt.subplots(1,2)
    f.set_size_inches([12,4])

    # display plot
    plot_accuracy_elements(ax, ax2, more_steps)

# run the interactive figure environment
interact(accuracy_figure, more_steps = (0,15,1));
```

***What are the slopes of the curves on the RIGHTHAND plot?***

> *~ your answer here ~*

***Which method is the most accurate?***

> *~ your answer here ~*

## 3 Stability

***Eliminating the wiggles and wobbles in the numerical solution.***

### 3.1 Euler method

Consider the simple ODE $y' = ay$, which we know has an exponential solution. For the case that $a<0$, the solution should decay to zero.

The Euler update step is

\begin{equation}
\begin{split}
y^{(k+1)} &=&\,y^{(k)} + h\,a\,y^{(k)} \\
&=&\,\left(1+ha\right)y^{(k)} \\
&=&\,\left(1+ha\right)^k y^{(0)}.
\end{split}
\end{equation}

If $|1+ha|>1$ then $y^{(k)}$ will keep increasing in value rather than decaying to 0. This defines a **stability criterion** for Euler's method applied to this problem: $-1 < 1+ha < 1 \Rightarrow 1 + ha > -1 \Rightarrow ha > -2 \Rightarrow h < - \dfrac{2}{a}$.  Therefore the stability criterion for the Euler method to give stable solutions for this problem is $h < -\dfrac{2}{a}$.

### 3.2 Improved Euler method

The Improved Euler method is a second-order accurate method and therefore we would expect it to provide **more accurate** solutions than the Euler method for the same choice of step size. However, the stability of the method is **not necessarily improved** as the solution still relies on the Euler estimate of $y_p^{(k+1)}$.

The Improved Euler update step is

\begin{equation}
\begin{split}
y^{(k+1)} &=&\,y^{(k)} + \frac{h}{2}\,\left(a\,y^{(k)}+a(1+ha)y^{(k)}\right) \\
&=&\,\left(1+ha+\frac{h^2}{2}a^2\right)y^{(k)} \\
&=&\,\left(1+ha+\frac{h^2}{2}a^2\right)^k y^{(0)}.
\end{split}
\end{equation}

For $a<0$ we require $y^{(k)}$ to decrease and not oscillate in sign, which implies $0 < 1 + ha + \tfrac{h^2}{2} a^2 < 1$. Thus, if

\begin{equation}
1 + ha + \tfrac{h^2}{2} a^2 < 1, \quad\Rightarrow\quad ha ( 1 + \tfrac{h}{2} a ) < 0 \quad\Rightarrow\quad h < -\tfrac{2}{a},
\end{equation}

which is the **same criterion** for the Euler method.

Similarly, if

\begin{equation}
1 + ha + \tfrac{h^2}{2} a^2 > 0 \quad\Rightarrow\quad 2 + 2ha + h^2 a^2 > 0 \quad\Rightarrow\quad ( 1 + ha ) ^2 > -1
\end{equation}

which is **always true** since the left-hand-side is positive and the right-hand-side is negative for any choice of $h$. Therefore the criterion for the Improved Euler method to give **stable, non-oscillating** solutions for this problem is $h < -\tfrac{2}{a}$.

***Execute the cell below to see how stability depends on method type and step size***

```python
# import function
from iteration233 import plot_stability_elements

# function parameters
def stability_figure(steps=28, method='Euler'):
    # create figure
    f,ax = plt.subplots(1,1)
    f.set_size_inches([10,4])

    # display plot
    plot_stability_elements(ax, steps, method)

# run the interactive figure environment
interact(stability_figure, steps = (3,15,1), method = ['Euler','Improved Euler','RK45']);
```

***Set STEPS to 15. Which method is the most accurate?***

> *~ your answer here ~*

***For the equation above, $a$=-10. For what step size does the Euler method exhibit oscillation?***

> *~ your answer here ~*

***Set STEPS to 6. Which method is the most accurate?***

> *~ your answer here ~*

***For what step size does the Euler method exhibit instability?***

> *~ your answer here ~*

### 3.1 Stiffness

The example above illustrates a quality of some differential equations called **stiffness**. While the definition of stiffness is a little nebulous, the following characteristics can be used to diagnose stiff solutions:
- Method stability is **constrained** by step size.
- Derivatives depend strongly on, or are large compared to, the solution.
- Unnecessarily **expensive** with explicit methods like Euler, improved Euler, RK45. Performs better with **implicit** methods (next year).

Stiffness is a **property of the solution** whereas stability is a property of the **solution method**. Consider the example below

\begin{equation}
\frac{dy}{dx}=a(x)y,\quad\quad\begin{Bmatrix} & 0.1 \quad & x < 20 \\ a(x) = & 1\quad & 20\leq x < 40 \\ & 0.1 \quad & x \geq 40 \end{Bmatrix},\quad\quad y(0)=1
\end{equation}

***Execute the cell below for a demonstration of solution stiffness.***

```python
# import functions
from iteration233 import plot_stiffness_elements

# function parameters
def stiffness_figure(step_size=0.5, log_yscale=False):
    # create figure
    f,(ax,ax2) = plt.subplots(1,2)
    f.set_size_inches([12,5])

    # display plot
    plot_stiffness_elements(ax, ax2, step_size)
    if log_yscale: ax.set_yscale('symlog',linthreshy=1.e-5)

# run the interactive figure environment
interact(stiffness_figure, step_size = (0.75,2.5,0.25), log_yscale=False);
```

***Set the STEP_SIZE to 0.75. Does the Euler method to a reasonable job solving this problem?***

> *~ your answer here ~*

***Set the STEP_SIZE to 1.75. Where have solution oscillations started to appear? Use the righthand plot to explain the appearance of solution oscillations in terms of the step size and value of $a$.***

> *~ your answer here ~*

***Set the STEP_SIZE to 2.75. Where is the solution exhibiting stiffness? Use the righthand plot to explain how ONLY PART of the solution can be stiff. Explain the stiffness in terms of the step size and value of $a$.***

> *~ your answer here ~*

## 4 Convergence

***Determining when a numerical solution is close enough to the true solution.***

Most ODEs do not have an **analytical solution**. How should we determine that a solution is **sufficiently accurate** when the true solution can never be known? One way is to solve the ODE for a range of step sizes, $h$, and look for **convergent behaviour**.

For example, take the ODE in Section 3 with $a=10$, $y(0) = 1$ and say we wish to estimate $y(0.1)$

***Execute the cell below for a demonstration of ODE convergence.***

```python
# import functions
from iteration233 import plot_convergence_elements

# function parameters
def convergence_figure(more_steps=0, zoom=False):
    # create figure
    f,(ax,ax2) = plt.subplots(1,2)
    f.set_size_inches([12,4])

    # display plot
    plot_convergence_elements(ax, ax2, more_steps, zoom)

# run the interactive figure environment
interact(convergence_figure, more_steps = (0,15,1), zoom = False);
```

***Set MORE_STEPS to 3. How has the ODE solution at $x=0.2$ changed by taking smaller steps?***

> *~ your answer here ~*

***What does the horizontal axis on the righthand plot represent?***

> *~ your answer here ~*

***Set MORE_STEPS to 15. What value is the solution converging toward? (Use the ZOOM if you need to.)***

> *~ your answer here ~*

***How should we quantify and monitor for convergent behaviour?***

> *~ your answer here ~*

***Calculate the true solutions for $y(0.2)$ - how do the estimates above compare?***

> *~ your answer here ~*

***The [Bulirsch-Stoer](https://en.wikipedia.org/wiki/Bulirsch%E2%80%93Stoer_algorithm) algorithm fits a function to the points on the RIGHTHAND plot and extrapolates the result at $h=0$. What value does this correspond to on the horizontal axis?***

> *~ your answer here ~*


```python
%matplotlib inline
%config InlineBackend.figure_format = 'svg'

import numpy as np
import qutip as qt
import scqubits as scq
```

# Composite Hilbert Spaces, QuTiP Interface

The `HilbertSpace` class provides data structures and methods for handling composite Hilbert spaces which may consist of multiple qubits and/or oscillators coupled to each other. To harness the power of QuTiP (a toolbox for simulating stationary and dynamical properties of closed and open quantum systems), `HilbertSpace` provides a convenient interface: it generates `qutip.qobj` objects which can then be handed over directly to QuTiP.

In the following, the basic functionality of the `HilbertSpace` class is illustrated for the example of two transmons coupled to a harmonic mode.

Transmon qubits can be capacitively coupled to a common harmonic mode, realized by an LC oscillator or a transmission-line resonator. The Hamiltonian describing such a composite system is given by:
\begin{equation}
H=H_\text{tmon,1} + H_\text{tmon,2} + \omega_r a^\dagger a + \sum_{j=1,2}g_j n_j(a+a^\dagger),
\end{equation}
where $j=1,2$ enumerates the two transmon qubits, $\omega_r$ is the (angular) frequency of the resonator. Furthermore, $n_j$ is the charge number operator for qubit $j$, and $g_j$ is the coupling strength between qubit $j$ and the resonator.

## Creating a HilbertSpace instance

The first step consists of creating the objects describing the individual building blocks of the full Hilbert space. Here, these will be the two transmons and one oscillator:

```python
tmon1 = scq.Transmon(
    EJ=40.0,
    EC=0.2,
    ng=0.3,
    ncut=40,
    truncated_dim=4     # after diagonalization, we will keep 3 levels
)

tmon2 = scq.Transmon(
    EJ=15.0,
    EC=0.15,
    ng=0.0,
    ncut=30,
    truncated_dim=4
)

resonator = scq.Oscillator(
    E_osc=4.5,
    truncated_dim=4  # up to 3 photons (0,1,2,3)
)
```

The `HilbertSpace` object is now created in one of two ways. The first is by utilizing the GUI:

```python
g = 0.1
hilbertspace = scq.HilbertSpace.create()
```

    VBox(children=(HBox(children=(VBox(children=(HBox(children=(Label(value='Select subsystems (Ctrl-Click)'), But…

    Output()

The second way of creating a `HilbertSpace` object is through regular programmatic creation: provide a list of all subsystems and then specify individual interaction terms (see the next subsection for the latter):

```python
hilbertspace = scq.HilbertSpace([tmon1, tmon2, resonator])
print(hilbertspace)
```

    HilbertSpace:  subsystems
    -------------------------

    Transmon------------|
                        | EJ: 40.0
                        | EC: 0.2
                        | ng: 0.3
                        | ncut: 40
                        | truncated_dim: 4
                        |
                        | dim: 81

    Transmon------------|
                        | EJ: 15.0
                        | EC: 0.15
                        | ng: 0.0
                        | ncut: 30
                        | truncated_dim: 4
                        |
                        | dim: 61

    Oscillator----------|
                        | E_osc: 4.5
                        | l_osc: None
                        | truncated_dim: 4
                        |
                        | dim: 4

## Specifying interactions between subsystems

.. note::
   **The code interface for specifying interactions has changed with version 1.4.** The old interface (see [here](./hilbertspace_legacy.ipynb)) is
   deprecated and will cease to be supported in the future.

Interaction terms describing the coupling between subsystems can be specified in three different ways.

### Option 1: operator-product based interface
Interaction terms involving multiple subsystems $S=1,2,3,\ldots$ are often of the form

$V = g A_1 A_2 A_3 \cdots \qquad \text{or} \qquad
V = g B_1 B_2 B_3 \cdots + \text{h.c.}$

where the operators $A_j$, $B_j$ act on subsystem $j$. In the first case, the operators $A_j$ are expected to be hermitean.

This structure is captured by using the `add_interaction` method in the following way:

```python
g1 = 0.1  # coupling resonator-tmon1 (without charge matrix elements)
operator1 = tmon1.n_operator()
operator2 = resonator.creation_operator() + resonator.annihilation_operator()

hilbertspace.add_interaction(
    g=g1,
    op1=(operator1, tmon1),
    op2=(operator2, resonator)
)

g2 = 0.2  # coupling resonator-CPB2 (without charge matrix elements)
hilbertspace.add_interaction(
    g=g2,
    op1=tmon2.n_operator,
    op2=resonator.creation_operator,
    add_hc=True
)
```

In this instance, `op1` and `op2` take either

1. `(<array>, <subsystem>)`: a tuple with an array or sparse matrix in the first position and the subsystem in the second position, or
2. `<callable>`: the operator as a callable function (which will automatically yield the subsystem the operator function is bound to).

These two choices can be mixed and matched.

Note:

* Interactions based on only one operator are enabled (simply drop the `op2` entry). One use case of this is to create a higher order non-linearity $a^\dagger a^\dagger a^\dagger a a a$ in a Kerr oscillator.
* Interactions with more than two operators are built by specifying `op3`, `op4`, etc.

The option `add_hc = True` adds the hermitian conjugate to the Hamiltonian, and is available in the present operator-product based interface, as well as in the one described next as option 2.

### Option 2: string- based interface

The `add_interaction` method can also be used to define the interaction in string form, providing an expression that can be evaluated by the Python interpreter (after operator name substitution).

```python
hilbertspace_example = scq.HilbertSpace([tmon1, tmon2, resonator])

g3 = 0.1

hilbertspace_example.add_interaction(
    expr="g3 * cos(n) * adag",
    op1=("n", tmon1.n_operator),
    op2=("adag", resonator.creation_operator),
    add_hc=True
)
```

`expr` is a string used to define the interaction as a Python expression. It should use variables that are already defined globally, and operators given by  the names provided in `op1`, `op2`, `op3`, etc. When using `expr`, each argument `op1`, `op2`, ... should be a tuple of the form

`(<operator name as string>,  <operator as callable>)` or

`(<operator name as string>,  <operator as array | sparse matrix>, <subsystem>)`.

Beyond built-in Python functions and operations, acceptable functions to use in `expr` are:

|  Function             |   Function translation          |
|-----------------------|---------------------------------|
| `'cos'`               | `'Qobj.cosm'`                   |
| `'sin'`               | `'Qobj.sinm'`                   |
| `'exp'`               | `'Qobj.expm'`                   |
| `'sqrt'`              | `'Qobj.sqrtm'`                  |

### Option 3: *Qobj* interface

Finally `add_interaction` can also been used to directly add a `Qobj` object that has already been properly identity wrapped:

```python
# Generate a Qobj
g = 0.1
a = qt.destroy(4)
kerr = a.dag() * a.dag() * a * a
id = qt.qeye(4)
V = g * qt.tensor(id, id, kerr)

hilbertspace_example.add_interaction(qobj = V)
```

In this case, the `add_interaction` method only requires the argument `qobj` which must be an object of the type `Qobj`.

.. note:
    By default, `add_interaction` will perform a validity check that internally constructs the interaction Hamiltonian. For a large Hilbert space, this check itself may require noticable runtime. Validity checking can be disabled by adding the keyword argument `check_validity=False`.

## Obtaining the Hamiltonian

With the interactions specified, the full Hamiltonian of the coupled system can be obtained via the method `.hamiltonian()`. This conveniently results in a `qubit.Qobj` operator:

```python
dressed_hamiltonian = hilbertspace.hamiltonian()
dressed_hamiltonian
```

Quantum object: dims = [[4, 4, 4], [4, 4, 4]], shape = (64, 64), type = oper, isherm = True\begin{equation*}\left(\begin{array}{*{11}c}-48.968 & 0.030 & 0.0 & 0.0 & 0.0 & \cdots & 0.0 & 0.0 & 0.0 & 0.0 & 0.0\\0.030 & -44.468 & 0.042 & 0.0 & 0.261 & \cdots & 0.0 & 0.0 & 0.0 & 0.0 & 0.0\\0.0 & 0.042 & -39.968 & 0.052 & 0.0 & \cdots & 0.0 & 0.0 & 0.0 & 0.0 & 0.0\\0.0 & 0.0 & 0.052 & -35.468 & 0.0 & \cdots & 0.0 & 0.0 & 0.0 & 0.0 & 0.0\\0.0 & 0.261 & 0.0 & 0.0 & -44.881 & \cdots & 0.0 & 0.0 & 0.0 & 0.0 & 0.0\\\vdots & \vdots & \vdots & \vdots & \vdots & \ddots & \vdots & \vdots & \vdots & \vdots & \vdots\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & \cdots & -4.723 & 0.0 & 0.0 & -0.749 & 0.0\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & \cdots & 0.0 & -14.477 & 0.030 & 0.0 & 0.0\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & \cdots & 0.0 & 0.030 & -9.977 & 0.042 & 0.0\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & \cdots & -0.749 & 0.0 & 0.042 & -5.477 & 0.052\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & \cdots & 0.0 & 0.0 & 0.0 & 0.052 & -0.977\\\end{array}\right)\end{equation*}

This matrix is the representation of the Hamiltonian in the bare product basis.

The non-interacting ("bare") Hamiltonian and the interaction Hamiltonian can be accessed similarly by using the methods `.bare_hamiltonian()` and `.interaction_hamiltonian()`.

### Obtaining the eigenspectrum via QuTiP

Since the Hamiltonian obtained this way is a proper `qutip.qobj`, all QuTiP routines are now available. In the first case, we are still making use of the scqubit `HilbertSpace.eigensys()` method. In the second, case, we use QuTiP's method `.eigenenergies()`:

```python
evals, evecs = hilbertspace.eigensys(evals_count=4)
print(evals)
```

    [-48.97770317 -45.02707241 -44.36656205 -41.18438832]

```python
dressed_hamiltonian = hilbertspace.hamiltonian()
dressed_hamiltonian.eigenenergies()
```

    array([-48.97770317, -45.02707241, -44.36656205, -41.18438832,
           -41.1776098 , -40.46448065, -39.76202396, -37.45533167,
           -37.22705488, -36.65156212, -36.56694139, -35.88617024,
           -35.14255357, -33.58960343, -33.38439485, -33.12061816,
           -32.66494366, -32.04065582, -31.96284558, -30.93536847,
           -29.65538466, -29.63912664, -28.97946785, -28.85287223,
           -28.64748897, -28.09427075, -27.35745603, -26.97461415,
           -26.21586637, -25.79648462, -25.3216198 , -25.07747135,
           -24.37567532, -24.24760531, -23.66618527, -23.15199748,
           -22.2580349 , -22.06752989, -21.57303944, -21.26626732,
           -20.85288717, -20.51511596, -19.78521899, -19.1938092 ,
           -18.41228666, -17.73463904, -17.66080329, -16.93618597,
           -16.66699841, -15.89010361, -15.58181323, -14.68288545,
           -13.84725844, -13.27037775, -13.08388811, -12.34366698,
           -11.62655644, -10.308913  ,  -9.23168892,  -8.32791939,
            -8.13883089,  -5.82301882,  -4.18180351,  -0.88957289])

## Spectrum lookup: converting between bare and dressed indices

We can refer to eigenstates and eigenenergies of the interacting quantum systems in two ways:

1. **dressed indices**: enumerating eigenenergies and -states as `j=0,1,2,...` starting from the ground state
2. **bare product state indices**: for a Hilbert space composed of three subsystems, specify tuples of the form `(l,m,n)` where `l`, `m`, `n` denote the excitation levels of each bare subsystem.

To use lookup functions converting back and forth between these indexing schemes, first generate the spectrum lookup table via:

```python
hilbertspace.generate_lookup()
```

This makes lookup functions available via `hilbertspace.lookup.<lookup method>` where

+-----------------------+---------------------------------------------------------------------------+
|  lookup method        | purpose                                                                   |
+=======================+===========================================================================+
| `bare_index`          | yields bare product state index for given dressed index                   |
+-----------------------+---------------------------------------------------------------------------+
| `dressed_index`       | yields dressed index for given bare product state index                   |
+-----------------------+---------------------------------------------------------------------------+
|                       |                                                                           |
+-----------------------+---------------------------------------------------------------------------+
| `energy_dressed_index`| return eigenenergy for given dressed index                                |
+-----------------------+---------------------------------------------------------------------------+
| `energy_bare_index`   | return eigenenergy for given bare product-state index                     |
+-----------------------+---------------------------------------------------------------------------+
|                       |                                                                           |
+-----------------------+---------------------------------------------------------------------------+
| `bare_eigenenergies`  | return bare eigenenergies for specified subsystem                         |
+-----------------------+---------------------------------------------------------------------------+
| `bare_eigenstates`    | return bare eigenstates for specified subsystem                           |
+-----------------------+---------------------------------------------------------------------------+
| `bare_productstate`   | return bare product state for given bare index                            |
+-----------------------+---------------------------------------------------------------------------+

Here are the bare energies of the first Cooper pair box system:

```python
hilbertspace.lookup.bare_eigenenergies(tmon1)
```

    array([-36.05064983, -28.25601136, -20.67410141, -13.31487799])

The dressed state with index j=8 corresponds to following bare product state index:

```python
hilbertspace.lookup.bare_index(8)
```

    (1, 1, 0)

.. note:
    In cases of strong hybridization between bare product states, this conversion can fail in which case it returns `None`. Conversion between the two different index types is based on maximum overlap between bare and dressed states.

The bare product state (2,1,0) most closely matches the following dressed state:

```python
hilbertspace.lookup.dressed_index((2, 1, 0))
```

    21

This is the eigenenergy for dressed index j=10:

```python
hilbertspace.lookup.energy_dressed_index(10)
```

    -36.56694138527015

## Helper function:  identity wrapping

The abbreviated notation commonplace in physics usually omits tensor products with identities. As an example, in a composite Hilbert space consisting of three subsystems $j=1,2,3$ with individual Hamiltonians $H_j$, we typically use the notation

$H=H_1+H_2+H_3$

which is meant to stand for

$H = H_1\otimes\mathbb{I}\otimes \mathbb{I} + \mathbb{I}\otimes H_2\otimes \mathbb{I} + \mathbb{I}\otimes\mathbb{I}\otimes H_3$

In most places of the `HilbertSpace` class, scqubits manages this "identity wrapping" automatically. In some use cases, it may become necessary to invoke this functionality manually. For that purpose, scqubits implements the helper routine `identity_wrap`. For instance, the term `\mathbb{I}\otimes H_2\otimes \mathbb{I}` would be generated via:

```python
scq.identity_wrap(H2, subsystem2, [subsystem1, subsystem2, subsystem3]
```

The complete interface is:

```scq.identity_wrap(operator, subsystem, subsys_list, op_in_eigenbasis=False, evecs=None)```

Wrap given operator in subspace subsystem in identity operators to form full Hilbert-space operator.

**Parameters**

- *operator*: operator acting in Hilbert space of subsystem; if str, then this should be an operator name in the subsystem, typically not in eigenbasis

- *subsystem (QuantumSys)*: subsystem to which *operator* belongs

- *subsys_list*: list of all subsystems (in fixed order!)

- *op_in_eigenbasis* (bool, optional): whether operator is given in the subsystem eigenbasis; otherwise, the internal QuantumSys basis is assumed

- *evecs* (ndarray, optional): internal QuantumSys eigenstates, used to convert operator into eigenbasis

**Return type**: Qobj

```python

```