# Neuromatch Academy: Week 2, Day 5, Tutorial 3
# Learning to Act: Q-Learning

__Content creators:__ Marcelo Mattar and Eric DeWitt with help from Byron Galbraith

__Content reviewers:__ Matt Krause and Michael Waskom

---

# Tutorial Objectives

In this tutorial you will learn how to act in the more realistic setting of sequential decisions, formalized by Markov Decision Processes (MDPs). In a sequential decision problem, the actions executed in one state not only may lead to immediate rewards (as in a bandit problem), but may also affect the states experienced next (unlike a bandit problem). Each individual action may therefore affect affect all future rewards. Thus, making decisions in this setting requires considering each action in terms of their expected **cumulative** future reward.

We will consider here the example of spatial navigation, where actions (movements) in one state (location) affect the states experienced next, and an agent might need to execute a whole sequence of actions before a reward is obtained.

By the end of this tutorial, you will learn
* what grid worlds are and how they help in evaluating simple reinforcement learning agents
* the basics of the Q-learning algorithm for estimating action values
* how the concept of exploration and exploitation, reviewed in the bandit case, also applies to the sequential decision setting

---
# Setup

```python
# Imports
import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import convolve as conv
```

```python
#@title Figure settings
%config InlineBackend.figure_format = 'retina'
plt.style.use("/share/dataset/COMMON/nma.mplstyle.txt")
```

```python
#@title Helper functions
def epsilon_greedy(q, epsilon):
  """Epsilon-greedy policy: selects the maximum value action with probabilty
  (1-epsilon) and selects randomly with epsilon probability.

  Args:
    q (ndarray): an array of action values
    epsilon (float): probability of selecting an action randomly

  Returns:
    int: the chosen action
  """
  if np.random.random() > epsilon:
    action = np.argmax(q)
  else:
    action = np.random.choice(len(q))

  return action

class CliffWorld:
  """
  World: Cliff world.
  40 states (4-by-10 grid world).
  The mapping from state to the grids are as follows:
  30 31 32 ... 39
  20 21 22 ... 29
  10 11 12 ... 19
  0  1  2  ...  9
  0 is the starting state (S) and 9 is the goal state (G).
  Actions 0, 1, 2, 3 correspond to right, up, left, down.
  Moving anywhere from state 9 (goal state) will end the session.
  Taking action down at state 11-18 will go back to state 0 and incur a
      reward of -100.
  Landing in any states other than the goal state will incur a reward of -1.
  Going towards the border when already at the border will stay in the same
      place.
  """
  def __init__(self):
    self.name = "cliff_world"
    self.n_states = 40
    self.n_actions = 4
    self.dim_x = 10
    self.dim_y = 4
    self.init_state = 0

  def get_outcome(self, state, action):
    if state == 9:  # goal state
      reward = 0
      next_state = None
      return next_state, reward
    reward = -1  # default reward value
    if action == 0:  # move right
      next_state = state + 1
      if state % 10 == 9:  # right border
        next_state = state
      elif state == 0:  # start state (next state is cliff)
        next_state = None
        reward = -100
    elif action == 1:  # move up
      next_state = state + 10
      if state >= 30:  # top border
        next_state = state
    elif action == 2:  # move left
      next_state = state - 1
      if state % 10 == 0:  # left border
        next_state = state
    elif action == 3:  # move down
      next_state = state - 10
      if state >= 11 and state <= 18:  # next is cliff
        next_state = None
        reward = -100
      elif state <= 9:  # bottom border
        next_state = state
    else:
      print("Action must be between 0 and 3.")
      next_state = None
      reward = None
    return int(next_state) if next_state is not None else None, reward

  def get_all_outcomes(self):
    outcomes = {}
    for state in range(self.n_states):
      for action in range(self.n_actions):
        next_state, reward = self.get_outcome(state, action)
        outcomes[state, action] = [(1, next_state, reward)]
    return outcomes

def learn_environment(env, learning_rule, params, max_steps, n_episodes):
  # Start with a uniform value function
  value = np.ones((env.n_states, env.n_actions))

  # Run learning
  reward_sums = np.zeros(n_episodes)

  # Loop over episodes
  for episode in range(n_episodes):
    state = env.init_state  # initialize state
    reward_sum = 0

    for t in range(max_steps):
      # choose next action
      action = epsilon_greedy(value[state], params['epsilon'])

      # observe outcome of action on environment
      next_state, reward = env.get_outcome(state, action)

      # update value function
      value = learning_rule(state, action, reward, next_state, value, params)

      # sum rewards obtained
      reward_sum += reward

      if next_state is None:
          break  # episode ends
      state = next_state

      reward_sums[episode] = reward_sum

  return value, reward_sums

def plot_state_action_values(env, value, ax=None):
  """
  Generate plot showing value of each action at each state.
  """
  if ax is None:
    fig, ax = plt.subplots()

  for a in range(env.n_actions):
    ax.plot(range(env.n_states), value[:, a], marker='o', linestyle='--')
  ax.set(xlabel='States', ylabel='Values')
  ax.legend(['R','U','L','D'], loc='lower right')

def plot_quiver_max_action(env, value, ax=None):
  """
  Generate plot showing action of maximum value or maximum probability at
    each state (not for n-armed bandit or cheese_world).
  """
  if ax is None:
    fig, ax = plt.subplots()

  X = np.tile(np.arange(env.dim_x), [env.dim_y,1]) + 0.5
  Y = np.tile(np.arange(env.dim_y)[::-1][:,np.newaxis], [1,env.dim_x]) + 0.5
  which_max = np.reshape(value.argmax(axis=1), (env.dim_y,env.dim_x))
  which_max = which_max[::-1,:]
  U = np.zeros(X.shape)
  V = np.zeros(X.shape)
  U[which_max == 0] = 1
  V[which_max == 1] = 1
  U[which_max == 2] = -1
  V[which_max == 3] = -1

  ax.quiver(X, Y, U, V)
  ax.set(
      title='Maximum value/probability actions',
      xlim=[-0.5, env.dim_x+0.5],
      ylim=[-0.5, env.dim_y+0.5],
  )
  ax.set_xticks(np.linspace(0.5, env.dim_x-0.5, num=env.dim_x))
  ax.set_xticklabels(["%d" % x for x in np.arange(env.dim_x)])
  ax.set_xticks(np.arange(env.dim_x+1), minor=True)
  ax.set_yticks(np.linspace(0.5, env.dim_y-0.5, num=env.dim_y))
  ax.set_yticklabels(["%d" % y for y in np.arange(0, env.dim_y*env.dim_x,
                                                  env.dim_x)])
  ax.set_yticks(np.arange(env.dim_y+1), minor=True)
  ax.grid(which='minor',linestyle='-')

def plot_heatmap_max_val(env, value, ax=None):
  """
  Generate heatmap showing maximum value at each state
  """
  if ax is None:
    fig, ax = plt.subplots()

  if value.ndim == 1:
      value_max = np.reshape(value, (env.dim_y,env.dim_x))
  else:
      value_max = np.reshape(value.max(axis=1), (env.dim_y,env.dim_x))
  value_max = value_max[::-1,:]

  im = ax.imshow(value_max, aspect='auto', interpolation='none', cmap='afmhot')
  ax.set(title='Maximum value per state')
  ax.set_xticks(np.linspace(0, env.dim_x-1, num=env.dim_x))
  ax.set_xticklabels(["%d" % x for x in np.arange(env.dim_x)])
  ax.set_yticks(np.linspace(0, env.dim_y-1, num=env.dim_y))
  if env.name != 'windy_cliff_grid':
      ax.set_yticklabels(
          ["%d" % y for y in np.arange(
              0, env.dim_y*env.dim_x, env.dim_x)][::-1])
  return im

def plot_rewards(n_episodes, rewards, average_range=10, ax=None):
  """
  Generate plot showing total reward accumulated in each episode.
  """
  if ax is None:
    fig, ax = plt.subplots()

  smoothed_rewards = (conv(rewards, np.ones(average_range), mode='same')
                      / average_range)

  ax.plot(range(0, n_episodes, average_range),
          smoothed_rewards[0:n_episodes:average_range],
          marker='o', linestyle='--')
  ax.set(xlabel='Episodes', ylabel='Total reward')

def plot_performance(env, value, reward_sums):
  fig, axes = plt.subplots(nrows=2, ncols=2, figsize=(16, 12))
  plot_state_action_values(env, value, ax=axes[0,0])
  plot_quiver_max_action(env, value, ax=axes[0,1])
  plot_rewards(n_episodes, reward_sums, ax=axes[1,0])
  im = plot_heatmap_max_val(env, value, ax=axes[1,1])
  fig.colorbar(im)
```

---
# Section 1: Markov Decision Processes

```python
#@title Video 1: MDPs and Q-learning
# Insert the ID of the corresponding youtube video
from IPython.display import IFrame
class BiliVideo(IFrame):
    def __init__(self, id, page=1, width=400, height=300, **kwargs):
        self.id=id
        src = "https://player.bilibili.com/player.html?bvid={0}&page={1}".format(id, page)
        super(BiliVideo, self).__init__(src, width, height, **kwargs)

video = BiliVideo(id='BV1ft4y1Q7bX', width=854, height=480, fs=1)
print("Video available at https://www.bilibili.com/video/{0}".format(video.id))
video
```

    Video available at https://www.bilibili.com/video/BV1ft4y1Q7bX

## Section 1.1: Grid Worlds

As pointed out, bandits only have a single state and immediate rewards for our actions. Many problems we are interested in have multiple states and delayed rewards, i.e. we won't know if the choices we made will pay off over time, or which actions we took contributed to the outcomes we observed.

In order to explore these ideas, we turn the a common problem setting: the grid world. Grid worlds are simple environments where each state corresponds to a tile on a 2D grid, and the only actions the agent can take are to move up, down, left, or right across the grid tiles. The agent's job is almost always to find a way to a goal tile in the most direct way possible while overcoming some maze or other obstacles, either static or dynamic.

For our discussion we will be looking at the classic Cliff World, or Cliff Walker, environment. This is a 4x10 grid with a starting position in the lower-left and the goal position in the lower-right. Every tile between these two is the "cliff", and should the agent enter the cliff, they will receive a -100 reward and be sent back to the starting position. Every tile other than the cliff produces a -1 reward when entered. The goal tile ends the episode after taking any action from it.

Given these conditions, the maximum achievable reward is -11 (1 up, 9 right, 1 down). Using negative rewards is a common technique to encourage the agent to move and seek out the goal state as fast as possible.

---
# Section 2: Q-Learning

Now that we have our environment, how can we solve it?

One of the most famous algorithms for estimating action values (aka Q-values) is the Temporal Differences (TD) **control** algorithm known as *Q-learning* (Watkins, 1989).

\begin{align}
Q(s_t,a_t) \leftarrow Q(s_t,a_t) + \alpha \big(r_t + \gamma\max_\limits{a} Q(s_{t+1},a_{t+1}) - Q(s_t,a_t)\big)
\end{align}

where $Q(s,a)$ is the value function for action $a$ at state $s$, $\alpha$ is the learning rate, $r$ is the reward, and $\gamma$ is the temporal discount rate.

The expression $r_t + \gamma\max_\limits{a} Q(s_{t+1},a_{t+1})$ is referred to as the TD target while the full expression
\begin{align}
r_t + \gamma\max_\limits{a} Q(s_{t+1},a_{t+1}) - Q(s_t,a_t),
\end{align}
i.e. the difference between the TD target and the current Q-value, is referred to as the TD error, or reward prediction error.

Because of the max operator used to select the optimal Q-value in the TD target, Q-learning directly estimates the optimal action value, i.e. the cumulative future reward that would be obtained if the agent behaved optimally, regardless of the policy currently followed by the agent. For this reason, Q-learning is referred to as an **off-policy** method.

## Exercise 1: Implement the Q-learning algorithm

In this exercise you will implement the Q-learning update rule described above. It takes in as arguments the previous state $s_t$, the action $a_t$ taken, the reward received $r_t$, the current state $s_{t+1}$, the Q-value table, and a dictionary of parameters that contain the learning rate $\alpha$ and discount factor $\gamma$. The method returns the updated Q-value table. For the parameter dictionary, $\alpha$: `params['alpha']` and $\gamma$: `params['gamma']`.

```python
def q_learning(state, action, reward, next_state, value, params):
  """Q-learning: updates the value function and returns it.

  Args:
    state (int): the current state identifier
    action (int): the action taken
    reward (float): the reward received
    next_state (int): the transitioned to state identifier
    value (ndarray): current value function of shape (n_states, n_actions)
    params (dict): a dictionary containing the default parameters

  Returns:
    ndarray: the updated value function of shape (n_states, n_actions)
  """
  # Q-value of current state-action pair
  q = value[state, action]

  ##########################################################
  ## TODO for students: implement the Q-learning update rule
  # Fill out function and remove
  raise NotImplementedError("Student excercise: implement the Q-learning update rule")
  ##########################################################

  # write an expression for finding the maximum Q-value at the current state
  if next_state is None:
    max_next_q = 0
  else:
    max_next_q = ...

  # write the expression to compute the TD error
  td_error = ...
  # write the expression that updates the Q-value for the state-action pair
  value[state, action] = ...

  return value
```

[*Click for solution*](https://github.com/erlichlab/course-content/tree/master//tutorials/W2D5_ReinforcementLearning/solutions/W2D5_Tutorial3_Solution_910598ca.py)

Now that we have our Q-learning algorithm, let's see how it handles learning to solve the Cliff World environment.

You will recall from the previous tutorial that a major part of reinforcement learning algorithms are their ability to balance exploitation and exploration. For our Q-learning agent, we again turn to the epsilon-greedy strategy. At each step, the agent will decide with probability $1 - \epsilon$ to use the best action for the state it is currently in by looking at the value function, otherwise just make a random choice.

The process by which our the agent will interact with and learn about the environment is handled for you in the helper function `learn_environment`. This implements the entire learning episode lifecycle of stepping through the state observation, action selection (epsilon-greedy) and execution, reward, and state transition. Feel free to review that code later to see how it all fits together, but for now let's test out our agent.

```python
# set for reproducibility, comment out / change seed value for different results
np.random.seed(1)

# parameters needed by our policy and learning rule
params = {
  'epsilon': 0.1,  # epsilon-greedy policy
  'alpha': 0.1,  # learning rate
  'gamma': 1.0,  # discount factor
}

# episodes/trials
n_episodes = 500
max_steps = 1000

# environment initialization
env = CliffWorld()

# solve Cliff World using Q-learning
results = learn_environment(env, q_learning, params, max_steps, n_episodes)
value_qlearning, reward_sums_qlearning = results

# Plot results
plot_performance(env, value_qlearning, reward_sums_qlearning)
```

If all went well, we should see four plots that show different aspects on our agent's learning and progress.

* The top left is a representation of the Q-table itself, showing the values for different actions in different states. Notably, going right from the starting state or down when above the cliff is clearly very bad.
* The top right figure shows the greedy policy based on the Q-table, i.e. what action would the agent take if it only took its best guess in that state.
* The bottom right is the same as the top, only instead of showing the action, it's showing a representation of the maximum Q-value at a particular state.
* The bottom left is the actual proof of learning, as we see the total reward steadily increasing after each episode until asymptoting at the maximum possible reward of -11.

Feel free to try changing the parameters or random seed and see how the agent's behavior changes.

---
# Summary

In this tutorial you implemented a reinforcement learning agent based on Q-learning to solve the Cliff World environment. Q-learning combined the epsilon-greedy approach to exploration-expoitation with a table-based value function to learn the expected future rewards for each state.

---
# Bonus

## SARSA

An alternative to Q-learning, the SARSA algorithm also estimates action values. However, rather than estimating the optimal (off-policy) values, SARSA estimates the **on-policy** action value, i.e. the cumulative future reward that would be obtained if the agent behaved according to its current beliefs.

\begin{align}
Q(s_t,a_t) \leftarrow Q(s_t,a_t) + \alpha \big(r_t + \gamma Q(s_{t+1},a_{t+1}) - Q(s_t,a_t)\big)
\end{align}

where, once again, $Q(s,a)$ is the value function for action $a$ at state $s$, $\alpha$ is the learning rate, $r$ is the reward, and $\gamma$ is the temporal discount rate.

In fact, you will notices that the *only* difference between Q-learning and SARSA is the TD target calculation uses the policy to select the next action (in our case epsilon-greedy) rather than using the action that maximizes the Q-value.

### Exercise 2: Implement the SARSA algorithm

In this exercise you will implement the SARSA update rule described above. Just like Q-learning, it takes in as arguments the previous state $s_t$, the action $a_t$ taken, the reward received $r_t$, the current state $s_{t+1}$, the Q-value table, and a dictionary of parameters that contain the learning rate $\alpha$ and discount factor $\gamma$. The method returns the updated Q-value table. You may use the `epsilon_greedy` function to acquire the next action. For the parameter dictionary, $\alpha$: `params['alpha']`, $\gamma$: `params['gamma']`, and $\epsilon$: `params['epsilon']`.

```python
def sarsa(state, action, reward, next_state, value, params):
  """SARSA: updates the value function and returns it.

  Args:
    state (int): the current state identifier
    action (int): the action taken
    reward (float): the reward received
    next_state (int): the transitioned to state identifier
    value (ndarray): current value function of shape (n_states, n_actions)
    params (dict): a dictionary containing the default parameters

  Returns:
    ndarray: the updated value function of shape (n_states, n_actions)
  """
  # value of previous state-action pair
  q = value[state, action]

  ##########################################################
  ## TODO for students: implement the SARSA update rule
  # Fill out function and remove
  raise NotImplementedError("Student excercise: implement the SARSA update rule")
  ##########################################################

  # select the expected value at current state based on our policy by sampling
  # from it
  if next_state is None:
    policy_next_q = 0
  else:
    # write an expression for selecting an action using epsilon-greedy
    policy_action = ...
    # write an expression for obtaining the value of the policy action at the
    # current state
    policy_next_q = ...

  # write the expression to compute the TD error
  td_error = ...
  # write the expression that updates the Q-value for the state-action pair
  value[state, action] = ...

  return value
```

```python
def sarsa(state, action, reward, next_state, value, params):
  """SARSA: updates the value function and returns it.

  Args:
    state (int): the current state identifier
    action (int): the action taken
    reward (float): the reward received
    next_state (int): the transitioned to state identifier
    value (ndarray): current value function of shape (n_states, n_actions)
    params (dict): a dictionary containing the default parameters

  Returns:
    ndarray: the updated value function of shape (n_states, n_actions)
  """
  # value of previous state-action pair
  q = value[state, action]

  # select the expected value at current state based on our policy by sampling
  # from it
  if next_state is None:
    policy_next_q = 0
  else:
    # write an expression for selecting an action using epsilon-greedy
    policy_action = epsilon_greedy(value[next_state], params['epsilon'])
    # write an expression for obtaining the value of the policy action at the
    # current state
    policy_next_q = value[next_state, policy_action]

  # write the expression to compute the TD error
  td_error = reward + params['gamma'] * policy_next_q - q
  # write the expression that updates the Q-value for the state-action pair
  value[state, action] = q + params['alpha'] * td_error

  return value
```

Now that we have an implementation for SARSA, let's see how it tackles Cliff World. We will again use the same setup we tried with Q-learning.

```python
# set for reproducibility, comment out / change seed value for different results
np.random.seed(1)

# parameters needed by our policy and learning rule
params = {
  'epsilon': 0.1,  # epsilon-greedy policy
  'alpha': 0.1,  # learning rate
  'gamma': 1.0,  # discount factor
}

# episodes/trials
n_episodes = 500
max_steps = 1000

# environment initialization
env = CliffWorld()

# learn Cliff World using Sarsa
results = learn_environment(env, sarsa, params, max_steps, n_episodes)
value_sarsa, reward_sums_sarsa = results

# Plot results
plot_performance(env, value_sarsa, reward_sums_sarsa)
```

We should see that SARSA also solves the task with similar looking outcomes to Q-learning. One notable difference is that SARSA seems to be skittsh around the cliff edge and often goes further away before coming back down to the goal.

Again, feel free to try changing the parameters or random seed and see how the agent's behavior changes.

## On-Policy vs Off-Policy

We have now seen an example of both on- and off-policy learning algorithms. Let's compare both Q-learning and SARSA reward results again, side-by-side, to see how they stack up.

```python
# parameters needed by our policy and learning rule
params = {
  'epsilon': 0.1,  # epsilon-greedy policy
  'alpha': 0.1,  # learning rate
  'gamma': 1.0,  # discount factor
}

# episodes/trials
n_episodes = 500
max_steps = 1000

# environment initialization
env = CliffWorld()

# learn Cliff World using Sarsa
np.random.seed(1)
results = learn_environment(env, q_learning, params, max_steps, n_episodes)
value_qlearning, reward_sums_qlearning = results
np.random.seed(1)
results = learn_environment(env, sarsa, params, max_steps, n_episodes)
value_sarsa, reward_sums_sarsa = results
```

```python
fig, ax = plt.subplots()
ax.plot(reward_sums_qlearning, label='Q-learning')
ax.plot(reward_sums_sarsa, label='SARSA')
ax.set(xlabel='Episodes', ylabel='Total reward')
plt.legend(loc='lower right');
```

On this simple Cliff World task, Q-learning and SARSA are almost indistinguisable from a performance standpoint, but we can see that Q-learning has a slight-edge within the 500 episode time horizon. Let's look at the illustrated "greedy policy" plots again.

```python
fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=(16, 6))
plot_quiver_max_action(env, value_qlearning, ax=ax1)
ax1.set(title='Q-learning maximum value/probability actions')
plot_quiver_max_action(env, value_sarsa, ax=ax2)
ax2.set(title='SARSA maximum value/probability actions');
```

What should immediately jump out is that Q-learning learned to go up, then immediately go to the right, skirting the cliff edge, until it hits the wall and goes down to the goal. The policy further away from the cliff is less certain.

SARSA, on the other hand, appears to avoid the cliff edge, going up one more tile before starting over to the goal side. This also clearly solves the challenge of getting to the goal, but does so at an additional -2 cost over the truly optimal route.

Why do you think these behaviors emerged the way they did?


```python
from IPython.display import Image
Image('../../Python_probability_statistics_machine_learning_2E.png',width=200)
```

```python
import numpy as np
np.random.seed(1234)
```

The estimation problem starts with the desire to infer something meaningful
from
data.  For parametric estimation, the strategy is to postulate a model for
the
data and then use the data to fit model parameters.  This leads to two
fundamental questions: where to get the model and how to estimate the
parameters? The first question is best answered by the maxim: *all models are
wrong, some are useful*. In other words, choosing a model depends as much on
the
application as on the model itself. Think about models as building
different
telescopes to view the sky. No one would ever claim that the
telescope generates
the sky! It is same with data models. Models give us
multiple perspectives on
the data that themselves are proxies for some deeper
underlying phenomenon.
Some
categories of data may be more commonly studied using certain types of
models,
but this is usually very domain-specific and ultimately depends on the
aims of
the analysis. In some cases, there may be strong physical reasons
behind
choosing a model.  For example, one could postulate that the model is
linear
with some noise as in the following:

$$
Y = a X + \epsilon
$$

 which basically
says that you, as the experimenter, dial in some
value for $X$
and then read off
something directly proportional to $X$ as the
measurement,
$Y$, plus some
additive noise that you attribute to jitter in the
apparatus.
Then, the next
step is to estimate the paramater $a$ in the model,
given some
postulated claim
about the nature of $\epsilon$. How to compute the
model
parameters depends on
the particular methodology. The two broad rubrics
are
parametric and non-
parametric estimation. In the former, we assume we know
the
density function of
the data and then try to derive the embedded parameters
for
it. In the latter,
we claim only to know that the density function is a
member
of a broad class of
density functions and then use the data
to characterize a
member of that class.
Broadly speaking, the former consumes
less data than the
latter, because there
are fewer unknowns to compute from
the data.

Let's
concentrate on parametric
estimation for now. The tradition is to denote
the
unknown parameter to be
estimated as $\theta$ which is a member of a large
space
of alternates,
$\Theta$. To judge between potential $\theta$ values, we
need an
objective
function, known as a *risk* function,
$L(\theta,\hat{\theta})$, where
$\hat{\theta}(\mathbf{x})$ is an
estimate for the unknown $\theta$ that is
derived from the available
data $\mathbf{x}$. The most common and useful risk
function is the
squared error loss,

$$
L(\theta,\hat{\theta}) =
(\theta-\hat{\theta})^2
$$

 Although neat, this is not practical because we
need to know the
unknown
$\theta$ to compute it. The other problem is because
$\hat{\theta}$ is
a
function of the observed data, it is also a random variable
with its own
probability density function.  This leads to the notion of the
*expected risk*
function,

$$
R(\theta,\hat{\theta}) =
\mathbb{E}_\theta(L(\theta,\hat{\theta})) = \int
L(\theta,\hat{\theta}(\mathbf{x})) f(\mathbf{x};\theta) d \mathbf{x}
$$

 In
other words, given a fixed $\theta$, integrate over the
probability density
function of the data, $f(\mathbf{x})$,  to compute the
risk. Plugging in for the
squared error loss,  we compute the
mean squared error,

$$
\mathbb{E}_\theta(\theta-\hat{\theta})^2 =\int (\theta-\hat{\theta})^2
f(\mathbf{x};\theta) d \mathbf{x}
$$

 This has the important factorization into
the *bias*,

$$
\texttt{bias} = \mathbb{E}_\theta(\hat{\theta})-\theta
$$
with the corresponding variance, $\mathbb{V}_\theta(\hat{\theta})$ as
in the
following *mean squared error* (MSE):

$$
\mathbb{E}_\theta(\theta-\hat{\theta})^2=
\texttt{bias}^2+\mathbb{V}_\theta(\hat{\theta})
$$

 This is an important
trade-off that we will return to repeatedly. The
idea is
the bias is nonzero
when the estimator $\hat{\theta}$, integrated
over all
possible data,
$f(\mathbf{x})$, does not equal the underlying target
parameter
$\theta$. In
some sense, the estimator misses the target, no matter
how much
data is used.
When the bias equals zero, the estimated is *unbiased*.
For fixed
MSE, low bias
implies high variance and vice-versa. This trade-off
was once not
emphasized and
instead much attention was paid to the smallest
variance of
unbiased estimators
(see Cramer-Rao bounds). In practice,
understanding and
exploiting the trade-off
between bias and variance and
reducing the MSE is more
important.

With all this
set up, we can now ask how bad can bad get by
examining *minimax* risk,

$$
R_{\texttt{mmx}} = \inf_{\hat{\theta}} \sup_\theta R(\theta,\hat{\theta})
$$
where the $\inf$ is take over all estimators.  Intuitively, this
means if we
found the worst possible $\theta$ and swept over all possible
parameter
estimators $\hat{\theta}$, and then took the smallest possible risk
we could
find, we would have the minimax risk. Thus, an estimator,
$\hat{\theta}_{\texttt{mmx}}$, is a *minimax estimator* if it achieves this
feat,

$$
\sup_\theta R(\theta,\hat{\theta}_{\texttt{mmx}})
=\inf_{\hat{\theta}}
\sup_\theta R(\theta,\hat{\theta})
$$

 In other words,
even in the face of the worst $\theta$ (i.e., the
$\sup_\theta$),
$\hat{\theta}_{\texttt{mmx}}$ still achieves the minimax
risk. There is a
greater theory that revolves around minimax estimators of
various kinds, but
this is far beyond our scope here. The main thing to focus
on
is that under
certain technical but easily satisfiable conditions, the
maximum
likelihood
estimator is approximately minimax. Maximum likelihood is
the subject
of the
next section.  Let's get started with the simplest
application: coin-
flipping.
## Setting up the Coin Flipping Experiment

Suppose we have coin and
want to
estimate the probability of heads ($p$) for
it. We model the
distribution of
heads and tails as a Bernoulli distribution
with the following
probability mass
function:

$$
\phi(x)= p^x (1-p)^{(1-x)}
$$

 where $x$ is the outcome, *1* for
heads and *0* for tails. Note that
maximum
likelihood is a parametric method
that requires the specification of a
particular model for which we will compute
embedded parameters. For $n$
independent flips, we have the joint density as the
product of $n$ of
these
functions as in,

$$
\phi(\mathbf{x})=\prod_{i=1}^n
p^x_i (1-p)^{(1-x_i)}
$$

 The following is the *likelihood function*,

$$
\mathcal{L}(p ; \mathbf{x})= \prod_{i=1}^n p^{ x_i }(1-p)^{1-x_i}
$$

 This is
basically notation. We have just renamed the
previous equation to
emphasize the
$p$ parameter, which is what
we want to estimate.

The principle
of *maximum
likelihood* is to maximize the likelihood as the
function of $p$
after plugging
in all of the $x_i$ data. We then call this
maximizer $\hat{p}$
which is a
function of the observed $x_i$ data, and as
such, is a random
variable with its
own distribution. This method therefore
ingests data and an
assumed model for
the probability density, and produces a
function that
estimates the embedded
parameter in the assumed probability
density.  Thus,
maximum likelihood
generates the *functions* of data that we
need in order to
get at the underlying
parameters of the model. Note that there
is no limit to
the ways we can
functionally manipulate the data we have
collected. The maximum
likelihood
principle gives us a systematic method for
constructing these
functions subject
to the assumed model. This is a point
worth emphasizing: the
maximum likelihood
principle yields functions as
solutions the same way solving
differential
equations yields functions as
solutions. It is very, very much
harder to produce
a function than to produce a
value as a solution, even with
the assumption of a
convenient probability
density.  Thus, the power of the
principle is that you
can construct such
functions subject to the model
assumptions.

### Simulating the Experiment

We need the following code to
simulate coin flipping.

```python
from scipy.stats import bernoulli
p_true=1/2.0         # estimate this!
fp=bernoulli(p_true) # create bernoulli random variate
xs = fp.rvs(100)     # generate some samples
print (xs[:30])        # see first 30 samples
```

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

Now, we can write out the likelihood function using Sympy.  Note
that we give
the Sympy variables the `positive=True` attribute upon
construction because this
eases Sympy's internal simplification algorithms.

```python
import sympy
x,p,z=sympy.symbols('x p z', positive=True)
phi=p**x*(1-p)**(1-x) # distribution function
L=np.prod([phi.subs(x,i) for i in xs]) # likelihood function
print (L) # approx 0.5?
```

    p**57*(1 - p)**43

Note that, once we plug in the data, the likelihood function is
solely a
function of the unknown parameter ($p$ in this case).  The following
code uses
calculus to find the extrema of the likelihood function.  Note that
taking the
`log` of $L$ makes the maximization problem tractable but doesn't
change the
extrema.

```python
logL=sympy.expand_log(sympy.log(L))
sol,=sympy.solve(sympy.diff(logL,p),p)
print (sol)
```

    57/100

**Programming Tip.**

Note that `sol,=sympy.solve` statement includes
a comma
after the `sol` variable. This is because the `solve`
function returns a list
containing a single element. Using
this assignment unpacks that single element
into the `sol` variable
directly. This is another one of the many small
elegancies of Python.

The following code generates
[Figure](#fig:Maximum_likelihood_10_2).

```python

from matplotlib.pylab import subplots
fig,ax=subplots()
x=np.linspace(0.001,1-.001,100)
ax.plot(x,list(map(sympy.lambdify(p,logL,'numpy'),x)),'k-',lw=3)
ax.plot(sol,logL.subs(p,sol),'o',
        color='gray',ms=15,label='Estimated')
ax.plot(p_true,logL.subs(p,p_true),'s',
        color='k',ms=15,label='Actual')
ax.set_xlabel('$p$',fontsize=18)
ax.set_ylabel('Likelihood',fontsize=18)
ax.set_title('Estimate not equal to true value',fontsize=18)
ax.legend(loc=0)
```

    <matplotlib.legend.Legend at 0x7f644ae93cc0>

**Programming Tip.**

In the prior code, we use the `lambdify` function in
`lambdify(p,logJ,'numpy')` to
take a Sympy expression and convert it into a
Numpy version that is easier to
compute.  The `lambdify` function has an extra
argument where you can specify
the function space that it should use to convert
the expression. In the above
this is set to Numpy.

Maximum likelihood estimate vs. true
parameter. Note that the estimate is slightly off from the true value. This is a
consequence of the fact that the estimator is a function of the data and lacks
knowledge of the true underlying value.

[Figure](#fig:Maximum_likelihood_10_2) shows  that our estimator $\hat{p}$
(circle) is not equal to the true value of $p$ (square), despite being
the
maximum of the likelihood function. This may sound disturbing, but keep in
mind
this estimate is a function of the random data; and since that data can
change,
the ultimate estimate can likewise change.
Remember that the estimator is a
*function* of the data and is thus also a
*random variable*, just like the data
is. This means it has its own probability
distribution with corresponding mean
and variance. So, what we are observing is
a consequence of that variance.

Histogram of
maximum likelihood estimates. The title shows the estimated mean and standard
deviation of the samples.

[Figure](#fig:Maximum_likelihood_30_2) shows what happens when you run many
thousands of coin experiments and compute the maximum likelihood
estimate for
each experiment, given a particular number of samples
per experiment. This
simulation gives us a histogram of the maximum likelihood
estimates, which is an
approximation of the probability distribution of the
$\hat{p}$ estimator itself.
This figure shows that the sample mean
of the estimator ($\mu = \frac{1}{n}\sum
\hat{p}_i$) is pretty close to the
true value, but looks can be deceiving. The
only way to know for sure is to
check if the estimator is unbiased, namely, if
$$
\mathbb{E}(\hat{p}) = p
$$

 Because this problem is simple, we can solve for
this in general
noting that
the terms above are either $p$, if $x_i=1$ or $1-p$
if $x_i=0$.
This means that
we can write

$$
\mathcal{L}(p\vert \mathbf{x})=
p^{\sum_{i=1}^n x_i}(1-p)^{n-\sum_{i=1}^n
x_i}
$$

 with corresponding logarithm
as

$$
J=\log(\mathcal{L}(p\vert \mathbf{x})) =  \log(p)  \sum_{i=1}^n x_i +
\log(1-p) \left(n-\sum_{i=1}^n x_i\right)
$$

 Taking the derivative of this
gives:

$$
\frac{dJ}{dp} = \frac{1}{p}\sum_{i=1}^n x_i + \frac{(n-\sum_{i=1}^n
x_i)}{p-1}
$$

 and solving this for $p$ leads to

$$
\hat{p} = \frac{1}{ n}
\sum_{i=1}^n x_i
$$

This is our *estimator* for $p$. Up until now, we have been
using Sympy to
solve
for this based on the data $x_i$ but now that we have it
analytically we
don't
have to solve for it each time. To check if this estimator
is biased, we
compute
its expectation:

$$
\mathbb{E}\left(\hat{p}\right)
=\frac{1}{n}\sum_i^n \mathbb{E}(x_i) =
\frac{1}{n} n \mathbb{E}(x_i)
$$

 by
linearity of the expectation and where

$$
\mathbb{E}(x_i)  = p
$$

 Therefore,
$$
\mathbb{E}\left(\hat{p}\right) =p
$$

 This means that the estimator is
*unbiased*. Similarly,

$$
\mathbb{E}\left(\hat{p}^2\right) = \frac{1}{n^2}
\mathbb{E}\left[\left(
\sum_{i=1}^n x_i \right)^2 \right]
$$

 and where

$$
\mathbb{E}\left(x_i^2\right) =p
$$

 and by the independence assumption,

$$
\mathbb{E}\left(x_i x_j\right) =\mathbb{E}(x_i)\mathbb{E}(x_j) =p^2
$$

 Thus,
$$
\mathbb{E}\left(\hat{p}^2\right) =\left(\frac{1}{n^2}\right) n \left[
p+(n-1)p^2 \right]
$$

 So, the variance of the estimator, $\hat{p}$, is the
following:

$$
\mathbb{V}(\hat{p}) = \mathbb{E}\left(\hat{p}^2\right)-
\mathbb{E}\left(\hat{p}\right)^2  = \frac{p(1-p)}{n}
$$

 Note that the $n$ in
the denominator means that the variance
asymptotically
goes to zero as $n$
increases (i.e., we consider more and
more samples). This is
good news because
it means that more and
more coin flips lead to a better
estimate of the
underlying $p$.

Unfortunately, this formula for the variance is
practically
useless because we
need $p$ to compute it and $p$ is the parameter
we are trying
to estimate in
the first place! However, this is where the *plug-
in* principle
[^invariance-property]
saves the day.  It turns out in this
situation, you can
simply substitute the maximum likelihood estimator,
$\hat{p}$, for the $p$ in
the above equation to obtain the asymptotic variance
for $\mathbb{V}(\hat{p})$.
The fact that this works is guaranteed by the
asymptotic theory of maximum
likelihood estimators.

[^invariance-property]:
This is also known as the
*invariance property*
of maximum likelihood
estimators. It basically states that
the
maximum likelihood estimator of any
function, say, $h(\theta)$, is
the same
$h$ with the maximum likelihood
estimator for $\theta$ substituted
in for
$\theta$; namely, $h(\theta_{ML})$.
Nevertheless, looking at
$\mathbb{V}(\hat{p})^2$, we can immediately notice
that
if $p=0$, then there is
no estimator variance because the outcomes are
guaranteed to be tails. Also, for
any $n$,  the maximum of this variance
happens
at $p=1/2$. This is our worst
case scenario and the only way to
compensate is
with larger $n$.

All we have
computed is the mean and variance of the
estimator. In general,
this is
insufficient to characterize the underlying
probability density of
$\hat{p}$,
except if we somehow knew that $\hat{p}$ were
normally distributed.
This is
where the powerful *Central Limit Theorem* we
discussed in the section
[ch:stats:sec:limit](#ch:stats:sec:limit) comes in. The
form of the estimator,
which is just a
sample mean, implies that we can apply
this theorem and conclude
that $\hat{p}$
is asymptotically normally distributed.
However, it doesn't
quantify how many
samples $n$ we need. In our simulation
this is no problem
because we can
generate as much data as we like, but in the
real world, with a
costly
experiment, each sample may be precious [^edgeworth].
In the following,
we won't apply the Central Limit Theorem and instead proceed
analytically.
[^edgeworth]: It turns out that the central limit theorem
augmented with an
Edgeworth expansion tells us that convergence is regulated by
the skewness
of
the distribution
[[feller1950introduction]](#feller1950introduction). In other
words, the
more
symmetric the distribution, the faster it converges to the
normal
distribution
according to the central limit theorem.

### Probability Density for the Estimator

To write out the full density for $\hat{p}$, we first
have to ask
what is
the probability that the estimator will equal a specific
value and the
tally up
all the ways that could happen with their corresponding
probabilities.
For
example, what is the probability that

$$
\hat{p} =
\frac{1}{n}\sum_{i=1}^n x_i  = 0
$$

  This can only happen one way: when $x_i=0
\hspace{0.5em} \forall i$. The
probability of this happening can be computed
from the density

$$
f(\mathbf{x},p)= \prod_{i=1}^n \left(p^{x_i} (1-p)^{1-x_i}
\right)
$$

$$
f\left(\sum_{i=1}^n x_i  = 0,p\right)= \left(1-p\right)^n
$$
Likewise, if $\lbrace x_i \rbrace$ has only one nonzero element, then

$$
f\left(\sum_{i=1}^n x_i  = 1,p\right)= n p \prod_{i=1}^{n-1} \left(1-p\right)
$$
where the $n$ comes from the $n$ ways to pick one element
from the $n$ elements
$x_i$. Continuing this way, we can construct the
entire density as

$$
f\left(\sum_{i=1}^n x_i  = k,p\right)= \binom{n}{k} p^k  (1-p)^{n-k}
$$

 where
the first term on the right is the binomial coefficient of $n$ things
taken $k$
at a time. This is the binomial distribution and it's not the
density
for
$\hat{p}$, but rather for $n\hat{p}$. We'll leave this as-is
because it's
easier
to work with below. We just have to remember to keep
track of the $n$
factor.
**Confidence Intervals**

Now that we have the full density for
$\hat{p}$, we
are ready to ask some
meaningful questions. For example, what is
the probability
the estimator is within
$\epsilon$ fraction of the true value
of $p$?

$$
\mathbb{P}\left( \vert  \hat{p}-p \vert  \le \epsilon p \right)
$$

 More
concretely, we want to know how often the
estimated $\hat{p}$ is trapped
within
$\epsilon$ of the actual value.  That is,
suppose we ran the experiment
1000
times to generate 1000 different estimates
of $\hat{p}$. What percentage of
the
1000 so-computed values are trapped within
$\epsilon$ of the underlying
value.
Rewriting the above equation as the
following,

$$
\mathbb{P}\left(p-\epsilon p
< \hat{p} < p + \epsilon p \right) =
\mathbb{P}\left(  n p - n \epsilon p <
\sum_{i=1}^n x_i < n p + n \epsilon p
\right)
$$

 Let's plug in some live
numbers here for our worst case
scenario (i.e., highest
variance scenario) where
$p=1/2$. Then, if
$\epsilon = 1/100$, we have

$$
\mathbb{P}\left( \frac{99
n}{100} < \sum_{i=1}^n x_i   < \frac{101 n}{100}
\right)
$$

 Since the sum in
integer-valued, we need $n> 100$ to even compute this.
Thus,
if $n=101$ we have,
$$
\begin{eqnarray*}
\mathbb{P}\left(\frac{9999}{200} < \sum_{i=1}^{101} x_i <
\frac{10201}{200} \right) = f\left(\sum_{i=1}^{101} x_i = 50,p\right) &  \ldots
\\\
= \binom{101}{50} (1/2)^{50}  (1-1/2)^{101-50} & = & 0.079
\end{eqnarray*}
$$

 This means that in the worst-case scenario for $p=1/2$, given $n=101$
trials,
we will only get within 1\% of the actual $p=1/2$ about 8\% of the
time.
If you
feel disappointed, it is because you've been paying attention.
What if
the coin
was really heavy and it was hard work to repeat this 101 times?

Let's
come at
this another way: given I could only flip the coin 100
times, how close
could I
come to the true underlying value with high
probability (say, 95\%)? In
this
case, instead of picking a value for
$\epsilon$, we are solving for
$\epsilon$.
Plugging in gives,

$$
\mathbb{P}\left(50 - 50\epsilon <
\sum_{i=1}^{100} x_i < 50 + 50 \epsilon
\right) = 0.95
$$

 which we have to
solve for $\epsilon$. Fortunately, all the tools we
need to
solve for this are
already in Scipy.

```python
from scipy.stats import binom
# n=100, p = 0.5, distribution of the estimator phat
b=binom(100,.5)
# symmetric sum the probability around the mean
g = lambda i:b.pmf(np.arange(-i,i)+50).sum()
print (g(10)) # approx 0.95
```

    0.9539559330706295

```python
%matplotlib inline

from matplotlib.pylab import subplots, arange
fig,ax= subplots()
fig.set_size_inches((10,5))
# here is the density of the sum of x_i
_=ax.stem(arange(0,101),b.pmf(arange(0,101)),
        linefmt='k-', markerfmt='ko',use_line_collection=True)
_=ax.vlines( [50+10,50-10],0 ,ax.get_ylim()[1] ,color='k',lw=3.)
_=ax.axis(xmin=30,xmax=70)
_=ax.tick_params(labelsize=18)
fig.savefig('fig-statistics/Maximum_likelihood_20_2.png')
fig.tight_layout()
```

Probability
mass function for $\hat{p}$. The two vertical lines form the confidence
interval.

 The two vertical lines in
[Figure](#fig:Maximum_likelihood_20_2)
show how far out from the mean we have to
go to accumulate 95\% of the
probability. Now, we can solve this as

$$
50+50\epsilon=60
$$

 which makes $\epsilon=1/5$ or 20\%. So, flipping 100 times
means I can
only get
within 20\% of the real $p$ 95\% of the time in the worst
case
scenario (i.e.,
$p=1/2$). The following code verifies the situation.

```python
from scipy.stats import bernoulli
b=bernoulli(0.5) # coin distribution
xs = b.rvs(100) # flip it 100 times
phat = np.mean(xs) # estimated p
print (abs(phat-0.5) < 0.5*0.20) # make it w/in interval?
```

    True

Let's keep doing this and see if we can get within this interval 95\% of
the
time.

```python
out=[]
b=bernoulli(0.5) # coin distribution
for i in range(500):   # number of tries
    xs = b.rvs(100)    # flip it 100 times
    phat = np.mean(xs) # estimated p
    out.append(abs(phat-0.5) < 0.5*0.20 ) # within 20% ?

# percentage of tries w/in 20% interval
print (100*np.mean(out))
```

    97.39999999999999

Well, that seems to work! Now we have a way to get at the quality of
the
estimator, $\hat{p}$.

**Maximum Likelihood Estimator Without Calculus**

The
prior example showed how we can use calculus to compute the maximum
likelihood
estimator. It's important to emphasize that the maximum likelihood
principle
does *not* depend on calculus and extends to more general situations
where
calculus is impossible. For example, let $X$ be uniformly distributed in
the
interval $[0,\theta]$.  Given $n$ measurements of $X$, the likelihood
function
is the following:

$$
L(\theta) = \prod_{i=1}^n \frac{1}{\theta} =
\frac{1}{\theta^n}
$$

 where each $x_i \in [0,\theta]$. Note that the slope of
this function
is not
zero anywhere so the usual calculus approach is not going
to work here.
Because
the likelihood is the product of the individual uniform
densities, if
any of the
$x_i$ values were outside of the proposed $[0,\theta]$
interval,
then the
likelihood would go to zero, because the uniform density is
zero
outside of the
$[0,\theta]$. This is no good for maximization. Thus,
observing
that the
likelihood function is strictly decreasing with increasing
$\theta$,
we conclude
that the value for $\theta$ that maximizes the likelihood
is the
maximum of the
$x_i$ values. To summarize, the maximum likelihood
estimator is
the following:

$$
\theta_{ML} = \max_i x_i
$$

 As always, we want
the distribution of this estimator to judge its
performance.
In this case, this
is pretty straightforward. The cumulative
density function
for the $\max$
function is the following:

$$
\mathbb{P} \left( \hat{\theta}_{ML} < v \right) =
\mathbb{P}( x_0 \leq v
\wedge x_1 \leq v \ldots \wedge x_n \leq v)
$$

 and
since all the $x_i$ are uniformly distributed in $[0,\theta]$, we have

$$
\mathbb{P} \left( \hat{\theta}_{ML} < v  \right) =
\left(\frac{v}{\theta}\right)^n
$$

 So, the probability density function is
then,

$$
f_{\hat{\theta}_{ML}}(\theta_{ML}) =  n \theta_{ML}^{ n-1 } \theta^{
-n }
$$

 Then, we can compute the $\mathbb{E}(\theta_{ML}) = (\theta n)/(n+1)$
with
corresponding variance as $\mathbb{V}(\theta_{ML}) = (\theta^2
n)/(n+1)^2/(n+2)$.

For a quick sanity check, we can write the following
simulation for $\theta =1$
as in the following:

```python
from scipy import stats
rv = stats.uniform(0,1)  # define uniform random variable
mle=rv.rvs((100,500)).max(0) # max along row-dimension
print (np.mean(mle)) # approx n/(n+1) = 100/101 ~= 0.99
print (np.var(mle)) #approx n/(n+1)**2/(n+2) ~= 9.61E-5
9.95762009884e-05
```

    0.9902508350190666
    9.414736602783561e-05

    9.95762009884e-05

**Programming Tip.**

The `max(0)` suffix on for the `mle` computation takes
the
maximum of the so-computed array along the row (`axis=0`)
dimension.

 You can
also plot `hist(mle)` to see the histogram of the simulated
maximum likelihood
estimates and match it up against the probability density
function we derived
above.

In this section, we explored the concept of maximum
likelihood
estimation using a coin flipping experiment both analytically and
numerically
with the scientific Python stack. We also explored the case when
calculus is not
workable for maximum likelihood estimation.  There are two key
points to
remember. First, maximum likelihood estimation produces a function of
the data
that is itself a random variable, with its own probability
distribution. We can
get at the quality of the so-derived estimators by
examining the confidence
intervals around the estimated values using the
probability distributions
associated with the estimators themselves.
Second, maximum likelihood
estimation applies even in situations
where using basic calculus is not
applicable [[wasserman2004all]](#wasserman2004all).

## Delta Method

Sometimes we want to characterize the distribution
of a *function* of a random
variable. In order to extend and generalize the
Central Limit Theorem in this
way, we need the Taylor series expansion. Recall
that the Taylor series
expansion is an approximation of a function of the
following form,

$$
T_r(x) =\sum_{i=0}^r \frac{g^{(i)}(a)}{i!}(x-a)^i
$$

 this
basically says that a function $g$ can be adequately
approximated about a
point
$a$ using a polynomial based on its derivatives
evaluated at $a$. Before
we
state the general theorem, let's examine
an example to understand how the
mechanics work.

**Example.**  Suppose that $X$ is a random variable with
$\mathbb{E}(X)=\mu\neq 0$.  Furthermore, supposedly have a suitable
function $g$
and we want the distribution of $g(X)$. Applying the
Taylor series expansion, we
obtain the following,

$$
g(X) \approx g(\mu)+ g^{\prime}(\mu)(X-\mu)
$$

 If we
use $g(X)$  as an estimator for $g(\mu)$, then we can say that
we
approximately
have the following

$$
\begin{align*}
\mathbb{E}(g(X)) &=g(\mu) \\\
\mathbb{V}(g(X))
&=(g^{\prime}(\mu))^2 \mathbb{V}(X) \\\
\end{align*}
$$
Concretely, suppose we want to estimate the odds, $\frac{p}{1-p}$.
For example,
if $p=2/3$, then we say that the odds is `2:1` meaning that the
odds of the one
outcome are twice as likely as the odds of the other outcome.
Thus, we have
$g(p)=\frac{p}{1-p}$ and we want to find
$\mathbb{V}(g(\hat{p}))$.  In our coin-
flipping problem, we have the
estimator $\hat{p}=\frac{1}{n}\sum X_k$ from the
Bernoulli-distributed data
$X_k$ individual coin-flips. Thus,

$$
\begin{align*}
\mathbb{E}(\hat{p}) &= p  \\\
\mathbb{V}(\hat{p}) &=
\frac{p(1-p)}{n}  \\\
\end{align*}
$$

  Now, $g^\prime(p)=1/(1-p)^2$, so we have,

$$
\begin{align*}
\mathbb{V}(g(\hat{p}))&=(g^\prime(p))^2 \mathbb{V}(\hat{p})
\\\
&=\left(\frac{1}{(1-p)^2}\right)^2 \frac{p(1-p)}{n}
\\\
                      &=
\frac{p}{n(1-p)^3}  \\\
\end{align*}
$$

 which is an approximation of the
variance of the estimator
$g(\hat{p})$. Let's
simulate this and see how it
agrees.

```python
from scipy import stats
# compute MLE estimates
d=stats.bernoulli(0.1).rvs((10,5000)).mean(0)
# avoid divide-by-zero
d=d[np.logical_not(np.isclose(d,1))]
# compute odds ratio
odds = d/(1-d)
print ('odds ratio=',np.mean(odds),'var=',np.var(odds))
```

    odds ratio= 0.12363809523809527 var= 0.017607461164021166

The first number above is the mean of the simulated odds
ratio and the second
is
the variance of the estimate.  According to
the variance estimate above, we
have
$\mathbb{V}(g(1/10))\approx
0.0137$, which is not too bad for this
approximation. Recall we want
to estimate the odds from $\hat{p}$. The code
above takes $5000$
estimates of the $\hat{p}$ to estimate $\mathbb{V}(g)$. The
odds ratio
for $p=1/10$ is $1/9\approx 0.111$.

**Programming Tip.**

The code
above uses the `np.isclose` function to identify the ones from
the simulation
and the `np.logical_not` removes these elements from the
data because the odds
ratio has a zero in the denominator
for these values.

Let's try this again
with a probability of heads of `0.5` instead of
`0.3`.

```python
from scipy import stats
d=stats.bernoulli(.5).rvs((10,5000)).mean(0)
d=d[np.logical_not(np.isclose(d,1))]
print( 'odds ratio=',np.mean(d),'var=',np.var(d))
```

    odds ratio= 0.4984584584584585 var= 0.024323949976002027

The odds ratio is this case is equal to one, which
is not close to what was
reported. According to our
approximation, we should have $\mathbb{V}(g)=0.4$,
which does not
look like what our simulation just reported.  This is
because the
approximation is best when the odds ratio is
nearly linear and worse otherwise
(see [Figure](#fig:Maximum_likelihood_0001)).

The odds
ratio is close to linear for small values but becomes unbounded as $p$
approaches one. The delta method is more effective for small underlying values
of $p$, where the linear approximation is better.


# 轨道旋转 MP2 方法 (OO-MP2) 简单理解

> 创建日期：2021-01-09

这篇文档会尝试简单介绍轨道旋转 MP2 方法 (Orbital-Optimized Second-Order Møller–Plesset Perturbation, OO-MP2 or OMP2) 的基础概念与 PySCF 上的程序实现和理解。

这篇文档的编写并没有翻阅很多文献，并作测评上的认识。为数不多的文献与参考资料是

> Sun, Chan, et al. [^Sun-Chan.JCP.2020] (PySCF 进展文章)
>
> PySCF 并没有一个完整或独立的 OO-MP2 模块。实现 OO-MP2 可以通过仿 CASSCF 的方式实现。之后使用到的 `MP2AsFCISolver` class 就是直接从该文章中截取的演示代码。

> Psi4NumPy 演示文档 [10a_orbital-optimized-mp2.ipynb](https://github.com/psi4/psi4numpy/blob/master/Tutorials/10_Orbital_Optimized_Methods/10a_orbital-optimized-mp2.ipynb)
>
>   这是一份比较不错的基于 Psi4 的程序简要文档，使用的算法与技巧也不复杂。

需要指出，这里的 OO-MP2 程序实现完全是基于 Post-HF 的闭壳层、无冻结轨道 MP2 实现的。更复杂的开壳层、冻结轨道、双杂化泛函方法，都不予以考虑。

```python
import numpy as np
import scipy
from pyscf import gto, mcscf, fci, mp, scf
from functools import partial

np.random.seed(0)
np.einsum = partial(np.einsum, optimize=True)
np.set_printoptions(precision=4, linewidth=150, suppress=True)
```

这篇文档的程序理解部分，我们都会使用下述水分子。但文档末尾，我们会用氢分子的例子，说明 OO-MP2 的能量未必要比 MP2 能量要低。

```python
mol = gto.Mole()
mol.atom = """
O  0. 0. 0.
H  0. 0. 1.
H  0. 1. 0.
"""
mol.basis = "6-31G"
mol.verbose = 0
mol.build()
```

    <pyscf.gto.mole.Mole at 0x7efdd0fbc8b0>

## PySCF 程序实现：高效方式

这段程序 `MP2AsFCISolver` class 是直接从 Sun 的 JCP 文章截取的。通过在 CASSCF 中，将活性空间更改为全轨道、更改约化密度矩阵 (1-RDM, 2-RDM) 的生成方式为 MP2 的约化密度矩阵、并且允许活性空间的轨道旋转，就可以实现 OO-MP2。

```python
class MP2AsFCISolver:
    def kernel(self, h1, h2, norb, nelec, ci0=None, ecore=0, **kwargs):
        # Kernel takes the set of integrals from the current set of orbitals
        fakemol = gto.M(verbose=0)
        fakemol.nelectron = sum(nelec)
        fake_hf = fakemol.RHF()
        fake_hf._eri = h2
        fake_hf.get_hcore = lambda *args: h1
        fake_hf.get_ovlp = lambda *args: np.eye(norb)

        # Build an SCF object fake_hf without SCF iterations to perform MP2
        fake_hf.mo_coeff = np.eye(norb)
        fake_hf.mo_occ = np.zeros(norb)
        fake_hf.mo_occ[:fakemol.nelectron//2] = 2
        self.mp2 = fake_hf.MP2().run()
        return self.mp2.e_tot + ecore, self.mp2.t2

    def make_rdm12(self, t2, norb, nelec):
        dm1 = self.mp2.make_rdm1(t2)
        dm2 = self.mp2.make_rdm2(t2)
        return dm1, dm2
```

`mf_rhf` 为 RHF 实例：

```python
mf_rhf = mol.RHF().run()
mf_rhf.e_tot
```

    -75.9697009626036

`mf_mp2` 为 MP2 实例：

```python
mf_mp2 = mp.MP2(mf_rhf).run()
mf_mp2.e_tot
```

    -76.1040356515777

```python
mf_mp2.e_corr
```

    -0.13433468897410067

`mf_cas` 在这里是指 OO-MP2 实例：

```python
mf_cas = mcscf.CASSCF(mf_rhf, mol.nao, mol.nelectron)
mf_cas.fcisolver = MP2AsFCISolver()
mf_cas.internal_rotation = True
cas_result = mf_cas.kernel()
cas_result[0]
```

    -76.10510419427318

## PySCF 程序实现：大体思路拆解

在这一段中，我们不会使用 PySCF 的 `CASSCF` class，而是从 RHF 与 MP2 的结果，了解 OO-MP2 的大体思路。

从结果上，这种实现方式与 PySCF 会相同。但 PySCF 的 `CASSCF` class 一般会使用二阶方法 (即使用 Orbital Hessian) 加速收敛，而我们这里只使用一阶梯度下降方法 (Orbital Gradient) 进行收敛；一阶收敛方法显然会慢一些，但公式与程序会简单一些。

首先对一些基础变量作声明：

- `nocc` $n_\mathrm{occ}$ 占据轨道数，`nvir` $n_\mathrm{vir}$ 非占轨道数；

- `nmo` $n_\mathrm{MO}$ 分子轨道数，一般与原子轨道数相等；

- `so` $[0:n_\mathrm{occ}]$ 占据轨道分割 (slice)，`sv` $[n_\mathrm{occ}:n_\mathrm{MO}]$ 非占轨道分割 (slice)；

- `mo_occ` PySCF 中用于表示轨道占据数的变量。

```python
nocc, nmo = mol.nelec[0], mol.nao
nvir = nmo - nocc
so, sv = slice(0, nocc), slice(nocc, nmo)
mo_occ = mf_rhf.mo_occ
```

OO-MP2 的大体过程可以拆分为如下循环：

1. 代入分子轨道系数 $C_{\mu i}$，得到该系数下 MP2 的激发张量 $t_{ij}^{ab}$；

2. 进而生成该情形下的 1-RDM $\gamma_{pq}$ 与 2-RDM $\Gamma_{pr}^{qs}$；

3. 进而生成广义 Fock 矩阵 $F_{pq}$ 与轨道梯度 $x_{pq} = F_{pq} - F_{qp}$；

4. 最后更新分子轨道系数 $C_{\mu i}$。

最终的收敛条件判据是 $F_{pq} - F_{qp} = 0$，即广义 Fock 矩阵 $F_{pq}$ 为对称矩阵。

```python
def oomp2_cycle(C):
    # Generate Psuedo objects, and therefore t_iajb
    mf_prhf = scf.RHF(mol)
    mf_prhf.mo_occ, mf_prhf.mo_coeff = mo_occ, C
    mf_pmp2 = mp.MP2(mf_prhf).run()                                                         # Step 1
    # Generate 1-RDM, 2-RDM and orbital gradient from generalized Fock matrix
    rdm1, rdm2 = mf_pmp2.make_rdm1(), mf_pmp2.make_rdm2()                                   # Step 2
    gfock_grad = mf_cas.unpack_uniq_var(mf_cas.get_grad(C, (rdm1, rdm2), mf_cas.ao2mo(C)))  # Step 3
    # Returned value: Updated MO Coefficient; OO-MP2 Energy (in current cycle); orbital gradient error
    return update_C(C, gfock_grad), mf_pmp2.e_tot, np.linalg.norm(gfock_grad)               # Step 4
```

而更新轨道系数是通过下述过程实现：

$$
\begin{gather}
X_{ai} = - X_{ia} = \frac{x_{ai}}{- \varepsilon_a + \varepsilon_i} = \frac{F_{ai} - F_{ia}}{- \varepsilon_a + \varepsilon_i} \\
X_{ij} = 0, \; = X_{ab} = 0 \\
\mathbf{C} \leftarrow \mathbf{C} \exp(\lambda \mathbf{X})
\end{gather}
$$

其中 $\lambda$ 是梯度下降率，对应于机器学习，它与梯度下降算法的学习率是类似的。这里取为 $\lambda = 0.5$。

```python
def update_C(C, gfock_grad):
    # Generate anti-symmetrized rotation matrix
    D = mf_rhf.make_rdm1(C, mo_occ)
    e = (C.T @ mf_rhf.get_fock(dm=D) @ C).diagonal()
    X = np.zeros_like(C)
    X[sv, so] = gfock_grad[sv, so] / (- e[sv, None] + e[None, so])
    X[so, sv] = gfock_grad[so, sv] / (- e[None, sv] + e[so, None])
    # Control rotation by factor
    X *= 0.5
    # Generate rotated MO coefficient
    C_new = C @ scipy.linalg.expm(X)
    return C_new
```

如果将 RHF 的分子轨道系数 `mf_rhf.mo_coeff` 作为分子轨道系数的初猜，那么收敛过程可以用下述迭代代码给出：

```python
C_oo = np.copy(mf_rhf.mo_coeff)
```

```python
print("Cycle | OO-MP2 Energy | G-Fock Gradient Norm")
for i in range(15):
    C_oo, eng, err = oomp2_cycle(C_oo)
    print("{:5d} | {:<13.8f} |  {:<16.8e}".format(i, eng, err))
```

    Cycle | OO-MP2 Energy | G-Fock Gradient Norm
        0 | -76.10403565  |  7.90255467e-02
        1 | -76.10503066  |  2.20049872e-02
        2 | -76.10509490  |  7.36831750e-03
        3 | -76.10510186  |  4.69833400e-03
        4 | -76.10510336  |  2.78455388e-03
        5 | -76.10510386  |  1.90302779e-03
        6 | -76.10510405  |  1.22381869e-03
        7 | -76.10510413  |  8.22563315e-04
        8 | -76.10510417  |  5.38178673e-04
        9 | -76.10510418  |  3.59079689e-04
       10 | -76.10510419  |  2.36428274e-04
       11 | -76.10510419  |  1.57182522e-04
       12 | -76.10510419  |  1.03794108e-04
       13 | -76.10510419  |  6.88759621e-05
       14 | -76.10510419  |  4.55466994e-05

:::{admonition} 记号区别

在这份文档中，RHF 的 Fock 矩阵记号定义为 $f_{pq}$，而 Post-HF 方法的 Fock 矩阵记号定义为 $F_{pq}$。这两者并非相同，并且非轨道优化的方法下，广义 Fock 矩阵 $F_{pq}$ 矩阵一般是非对称的。

:::

## PySCF 程序实现：理解与分解

我们会对上面程序中的重要步骤进行说明。

### 原子轨道电子积分定义

- `h` $h_{\mu \nu}$，维度 $(\mu, \nu)$，原子轨道基的 Hamiltonian Core 矩阵，即动能与原子核-电子势能积分；

- `S` $S_{\mu \nu}$，维度 $(\mu, \nu)$，原子轨道基的重叠积分；

- `eri` $(\mu \nu | \kappa \lambda)$，维度 $(\mu, \nu, \kappa, \lambda)$，原子轨道基的双电子积分。

```python
h = mol.intor("int1e_kin") + mol.intor("int1e_nuc")
S = mol.intor("int1e_ovlp")
eri = mol.intor("int2e")
```

### Canonical MP2

我们先简单回顾一下在 Canonical RHF 下，MP2 的激发系数 $t_{ij}^{ab}$ 与能量 $E_\mathrm{corr}^\mathsf{MP2}$ 的导出方式。我们留意到 PySCF 的自洽场过程给出的是 Canonical 情况，即分子轨道的 Fock 矩阵 $f_{pq}$ 是对角矩阵。

- `C` $C_{\mu p}$ 为分子轨道系数，`e` $e_p$ 为轨道能量；

- `D_iajb` $D_{ij}^{ab}$ MP2 分母项，维度 $(i, a, j, b)$：

    $$
    D_{ij}^{ab} = \varepsilon_i - \varepsilon_a + \varepsilon_j - \varepsilon_b
    $$

- `eri_mo` $(pq|rs)$ 分子轨道基下的双电子积分，维度 $(p, q, r, s)$：

    $$
    (pq|rs) = C_{\mu p} C_{\nu q} (\mu \nu | \kappa \lambda) C_{\kappa r} C_{\lambda s}
    $$

- `t_iajb` $t_{ij}^{ab}$ MP2 激发系数：

    $$
    t_{ij}^{ab} = \frac{(ia|jb)}{D_{ij}^{ab}}
    $$

```python
C, e = mf_rhf.mo_coeff, mf_rhf.mo_energy
D_iajb = e[so, None, None, None] - e[None, sv, None, None] + e[None, None, so, None] - e[None, None, None, sv]
eri_mo = np.einsum("up, vq, uvkl, kr, ls -> pqrs", C, C, eri, C, C)
t_iajb = eri_mo[so, sv, so, sv] / D_iajb
```

因此，MP2 相关能可以写为 (参考值为 -0.134335 a.u.)

$$
E_\mathrm{corr}^\mathsf{MP2} = \big( 2 t_{ij}^{ab} - t_{ij}^{ba} \big) (ia|jb)
$$

```python
((2 * t_iajb - t_iajb.swapaxes(-1, -3)) * eri_mo[so, sv, so, sv]).sum()
```

    -0.13433468897410067

### Non-Canonical MP2：PySCF 程序

但对于 OO-MP2 而言，由于产生了轨道旋转，我们需要考察 Non-Canonical RHF 的 MP2。

Non-Canonical 意指 RHF 的 Fock 矩阵 $f_{pq}$ 是分块对角化的，即占据-非占和非占-占据分块 $f_{ia}$、$f_{ai}$ 均为零；而占据和非占分块 $f_{ij}$、$f_{ab}$ 的矩阵并非是对角化的。

为了构造这样一个 Non-Canonical RHF 的情形，我们可以对 Canonical RHF 分子轨道系数矩阵 `C_rhf` 作下述变换，得到 Non-Canonical RHF 分子轨道系数矩阵 `C_rot`：

$$
\mathbf{C} \leftarrow \mathbf{C} \mathbf{U}
$$

上述的 $\mathbf{U}$ 矩阵是分块对角化的正交矩阵。为了构造这样的正交矩阵，我们可以生成一个分块对角化、且反对称的 `X` $\mathbf{X}$ 矩阵，并令 $\mathbf{U} = \exp(\mathbf{X})$。

```python
C_rhf = mf_rhf.mo_coeff
```

```python
X = np.random.randn(nmo, nmo)
X[sv, so] = X[so, sv] = 0
X -= X.T
X *= 0.02
C_rot = C_rhf @ scipy.linalg.expm(X)
```

由此构建出的 Non-Canonical 分子轨道 Fock 矩阵 $f_{pq}$ 是分块对角化的，即不一定要求 $f_{ij} = \delta_{ij} \varepsilon_i$ 与 $f_{ab} = \delta_{ab} \varepsilon_a$：

```python
fock_rot = np.einsum("up, uv, vq -> pq", C_rot, mf_rhf.get_fock(), C_rot)
fock_rot
```

    array([[-20.4748,  -0.0683,  -0.3396,  -1.0108,  -0.9682,  -0.    ,  -0.    ,  -0.    ,   0.    ,   0.    ,  -0.    ,  -0.    ,   0.    ],
           [ -0.0683,  -1.3485,  -0.0071,  -0.0425,  -0.0204,   0.    ,  -0.    ,   0.    ,   0.    ,   0.    ,   0.    ,   0.    ,  -0.    ],
           [ -0.3396,  -0.0071,  -0.6668,  -0.0222,  -0.0172,  -0.    ,   0.    ,   0.    ,   0.    ,  -0.    ,  -0.    ,   0.    ,  -0.    ],
           [ -1.0108,  -0.0425,  -0.0222,  -0.6362,  -0.0523,  -0.    ,   0.    ,  -0.    ,   0.    ,   0.    ,  -0.    ,   0.    ,   0.    ],
           [ -0.9682,  -0.0204,  -0.0172,  -0.0523,  -0.5539,  -0.    ,  -0.    ,  -0.    ,  -0.    ,  -0.    ,   0.    ,   0.    ,  -0.    ],
           [ -0.    ,   0.    ,  -0.    ,  -0.    ,  -0.    ,   0.1967,   0.0008,  -0.0182,   0.0501,  -0.002 ,   0.0269,  -0.0106,   0.0776],
           [ -0.    ,  -0.    ,   0.    ,   0.    ,  -0.    ,   0.0008,   0.2872,  -0.0032,   0.011 ,   0.0263,   0.0326,  -0.0324,   0.0407],
           [ -0.    ,   0.    ,   0.    ,  -0.    ,  -0.    ,  -0.0182,  -0.0032,   0.9833,   0.0038,  -0.0097,   0.0066,   0.0088,  -0.0122],
           [  0.    ,   0.    ,   0.    ,   0.    ,  -0.    ,   0.0501,   0.011 ,   0.0038,   1.1572,  -0.0006,  -0.0001,   0.0048,  -0.0293],
           [  0.    ,   0.    ,  -0.    ,   0.    ,  -0.    ,  -0.002 ,   0.0263,  -0.0097,  -0.0006,   1.1616,  -0.0056,  -0.0042,   0.0036],
           [ -0.    ,   0.    ,  -0.    ,  -0.    ,   0.    ,   0.0269,   0.0326,   0.0066,  -0.0001,  -0.0056,   1.2463,  -0.0016,  -0.0138],
           [ -0.    ,   0.    ,   0.    ,   0.    ,   0.    ,  -0.0106,  -0.0324,   0.0088,   0.0048,  -0.0042,  -0.0016,   1.3518,  -0.0047],
           [  0.    ,  -0.    ,  -0.    ,   0.    ,  -0.    ,   0.0776,   0.0407,  -0.0122,  -0.0293,   0.0036,  -0.0138,  -0.0047,   1.7381]])

对于这样的分子轨道系数矩阵 `C_rot`，PySCF 的程序照样可以给出正确的 MP2 相关能量 -0.134335 a.u. (其中 `mf_prhf` 是指虚假的 (Pseudo) RHF 实例)：

```python
mf_prhf = scf.RHF(mol)
mf_prhf.mo_occ, mf_prhf.mo_coeff = mo_occ, C_rot
mf_pmp2 = mp.MP2(mf_prhf).run()
mf_pmp2.e_corr
```

    -0.13433467530628806

### Non-Canonical MP2：激发系数 $t_{ij}^{ab}$ 迭代更新方式

首先为程序与公式说明，定义一些变量：

- RHF Fock 对角线占据部分记为 `eo` $\varepsilon_i = f_{ii}$；

- RHF Fock 对角线非占部分记为 `ev` $\varepsilon_a = f_{aa}$；

- RHF Fock 去除对角线的占据分块记为 `fock_oo` $f'_{ij} = (1 - \delta_{ij}) f_{ij}$；

- RHF Fock 去除对角线的非占分块记为 `fock_vv` $f'_{ab} = (1 - \delta_{ab}) f_{ab}$；

- 双电子积分 `eri_mo` $(pq|rs)$；

- 只包含占据-非占分块的双电子积分 `eri_iajb` $(ia|jb)$

- MP2 分母 `D_iajb` $D_{ij}^{ab}$。

```python
eo, ev = fock_rot.diagonal()[so], fock_rot.diagonal()[sv]
fock_oo, fock_vv = fock_rot[so, so], fock_rot[sv, sv]
fock_oop, fock_vvp = fock_oo - np.diag(eo), fock_vv - np.diag(ev)
eri_mo = np.einsum("up, vq, uvkl, kr, ls -> pqrs", C_rot, C_rot, eri, C_rot, C_rot)
eri_iajb = eri_mo[so, sv, so, sv]
D_iajb = eo[:, None, None, None] - ev[None, :, None, None] + eo[None, None, :, None] - ev[None, None, None, :]
```

:::{caution}

**变量重新定义**

上面代码块中 `eo`, `ev`, `eri_mo`, `D_iajb` 就在 Non-Canonical 的系数矩阵 `C_rot` 下给出；但我们曾经也在 Canonical 系数矩阵下给出过类似的变量。

由于我们会经常切换各种系数矩阵的旋转方式 (非旋转、Non-Canonical、Non-HF)，因此一些变量也会被复用与复写，也暂不区分旋转后与旋转前的分子轨道角标。这可能会对阅读造成困惑。

:::

依据不同的微扰理论定义方式，Non-Canonical RHF 的 MP2 相关能可能与 Canonical RHF 的 MP2 相关能不同。因此这里采用两个相关能相同的定义。此时激发系数 $t_{ij}^{ab}$ 应当满足

$$
(ia|jb) = t_{kj}^{ab} f_{ki} + t_{ik}^{ab} f_{kj} - t_{ij}^{cb} f_{ca} - t_{ij}^{ac} f_{cb}
$$

上式是对等式右的 $k$ 进行求和的。如果现在用 $f_{ij} = f'_{ij} + \delta_{ij} \varepsilon_i$，$f_{ab} = f'_{ab} + \delta_{ab} \varepsilon_a$ 展开，那么上式可以写为

$$
(ia|jb) = t_{ij}^{ab} D_{ij}^{ab} + t_{kj}^{ab} f'_{ki} + t_{ik}^{ab} f'_{kj} - t_{ij}^{cb} f'_{ca} - t_{ij}^{ac} f'_{cb}
$$

整理上式，就可以得到迭代关系

$$
t_{ij}^{ab} \leftarrow \frac{(ia|jb) - t_{kj}^{ab} f'_{ki} - t_{ik}^{ab} f'_{kj} + t_{ij}^{cb} f'_{ca} + t_{ij}^{ac} f'_{cb}}{D_{ij}^{ab}}
$$

一般来说，如果轨道的旋转并不是很剧烈，那么 $f'_{ij}$, $f'_{ab}$ 两者的贡献较小，因此 $t_{ij}^{ab} \simeq (ia|jb) / D_{ij}^{ab}$ 会是一个比较好的近似。

在此情形下，Non-Canonical MP2 的能量计算方式如下：

$$
E_\mathrm{corr}^\mathsf{MP2} = \big( 2 t_{ij}^{ab} - t_{ij}^{ba} \big) (ia|jb)
$$

下面的程序就是实现 Non-Canonical MP2 的流程。

- `update_t_iajb` 即使用迭代关系，更新 $t_{ij}^{ab}$；

- `calculate_noncan_mp2` 即计算 Non-Canonical MP2 相关能的函数。

```python
def update_t_iajb(t_iajb):
    t_iajb_new = np.zeros_like(t_iajb)
    t_iajb_new += np.einsum("icjb, ca -> iajb", t_iajb, fock_vvp)
    t_iajb_new += np.einsum("iajc, cb -> iajb", t_iajb, fock_vvp)
    t_iajb_new -= np.einsum("iakb, kj -> iajb", t_iajb, fock_oop)
    t_iajb_new -= np.einsum("kajb, ki -> iajb", t_iajb, fock_oop)
    t_iajb_new += eri_iajb
    t_iajb_new /= D_iajb
    return t_iajb_new
```

```python
def calculate_noncan_mp2(t_iajb):
    return ((2 * t_iajb - t_iajb.swapaxes(-1, -3)) * eri_iajb).sum()
```

随后声明初猜 $t_{ij}^{ab} \simeq (ia|jb) / D_{ij}^{ab}$ 并以此迭代更新；并以 Canonical MP2 的相关能加以验证。在 5 次循环后，几乎收敛到了正确能量。

```python
t_iajb = eri_mo[so, sv, so, sv] / D_iajb
```

```python
for i in range(10):
    print("Error: {:16.8e}".format(calculate_noncan_mp2(t_iajb) - mf_mp2.e_corr))
    t_iajb = update_t_iajb(t_iajb)
```

    Error:   3.41981239e-03
    Error:   2.09994114e-03
    Error:   9.08474334e-05
    Error:   9.06169148e-05
    Error:   3.54576068e-06
    Error:   5.43397725e-06
    Error:   6.95752296e-08
    Error:   4.42378672e-07
    Error:  -2.06561550e-08
    Error:   4.46581002e-08

事实上，在 PySCF 中，包含占据-非占轨道旋转的 Non-RHF 下的 MP2 方法，也是通过上述过程进行计算的。

### MP2 1-RDM

对于一阶约化密度 1-RDM $\gamma_{pq}$，其需要通过分块的方式生成：

$$
\begin{align}
\gamma_{ij}^\mathsf{RHF} &= 2 \delta_{ij} \\
\gamma_{ab}^\mathsf{RHF} &= \gamma_{ia}^\mathsf{RHF} = \gamma_{ai}^\mathsf{RHF} = 0 \\
\gamma_{ij}^\mathrm{corr} &= - 4 t_{ik}^{ab} t_{jk}^{ab} + 2 t_{ik}^{ba} t_{jk}^{ab} \\
\gamma_{ab}^\mathrm{corr} &= 4 t_{ij}^{ac} t_{ij}^{bc} - 2 t_{ij}^{ca} t_{ij}^{bc} \\
\gamma_{ia}^\mathrm{corr} &= \gamma_{ai}^\mathrm{corr} = 0 \\
\gamma_{pq} &= \gamma_{pq}^\mathsf{RHF} + \gamma_{pq}^\mathrm{corr}
\end{align}
$$

这种生成方式无关乎方法是否是 Canonical 的。

首先生成 RHF 的 1-RDM `rdm1_rhf` $\gamma_{pq}^\mathsf{RHF}$：

```python
rdm1_rhf = np.zeros((nmo, nmo))
np.fill_diagonal(rdm1_rhf[so, so], 2)
```

随后给出 MP2 相关部分所给出的 1-RDM `rdm1_corr` $\gamma_{pq}^\mathrm{corr}$：

```python
rdm1_corr = np.zeros((nmo, nmo))
rdm1_corr[so, so] = - 4 * np.einsum("iakb, jakb -> ij", t_iajb, t_iajb) + 2 * np.einsum("ibka, jakb -> ij", t_iajb, t_iajb)
rdm1_corr[sv, sv] = 4 * np.einsum("iajc, ibjc -> ab", t_iajb, t_iajb) - 2 * np.einsum("icja, ibjc -> ab", t_iajb, t_iajb)
```

总 1-RDM `rdm1` $\gamma_{pq}$ 可以通过简单相加获得：

```python
rdm1 = rdm1_rhf + rdm1_corr
np.allclose(rdm1, mf_pmp2.make_rdm1())
```

    True

### MP2 2-RDM

对于二阶约化密度 2-RDM `rdm2` $\Gamma_{pr}^{qs}$ (维度 $(p, q, r, s)$)，其也需要通过分块生成。首先生成 $\Gamma_{ia}^{jb}$, $\Gamma_{ai}^{bj}$, $\Gamma_{ik}^{jl}$, $\Gamma_{ac}^{bd}$ 部分：

$$
\Gamma_{pr}^{qs} = \left( \gamma_{pq} \gamma_{rs} - \frac{1}{2} \gamma_{ps} \gamma_{rq} \right) - \left( \gamma_{pq}^\mathrm{corr} \gamma_{rs}^\mathrm{corr} - \frac{1}{2} \gamma_{ps}^\mathrm{corr} \gamma_{rq}^\mathrm{corr} \right)
$$

其余的部分是 $\Gamma_{ij}^{ab}$ 与 $\Gamma_{ab}^{ij}$：

$$
\Gamma_{ij}^{ab} = \Gamma_{ab}^{ij} = 4 t_{ij}^{ab} - 2 t_{ij}^{ba}
$$

```python
rdm2 = np.zeros((nmo, nmo, nmo, nmo))
rdm2 = np.einsum("pq, rs -> pqrs", rdm1, rdm1) - 0.5 * np.einsum("ps, rq -> pqrs", rdm1, rdm1)
rdm2 -= np.einsum("pq, rs -> pqrs", rdm1_corr, rdm1_corr) - 0.5 * np.einsum("ps, rq -> pqrs", rdm1_corr, rdm1_corr)
rdm2[so, sv, so, sv] = 4 * np.einsum("iajb -> iajb", t_iajb) - 2 * np.einsum("ibja -> iajb", t_iajb)
rdm2[sv, so, sv, so] = 4 * np.einsum("iajb -> aibj", t_iajb) - 2 * np.einsum("ibja -> aibj", t_iajb)
np.allclose(rdm2, mf_pmp2.make_rdm2(), atol=1e-7)
```

    False

由此，我们可以通过 1-RDM $\gamma_{pq}$ 与 2-RDM $\Gamma_{pr}^{qs}$ 验证 MP2 总能量 -76.104036 a.u.：

$$
E_\mathrm{tot}^\mathsf{MP2} = h_{pq} \gamma_{pq} + \frac{1}{2} (pq|rs) \Gamma_{pr}^{qs} + E_\mathrm{nuc}
$$

但这里的单电子积分 $h_{pq}$ 与双电子积分 $(pq|rs)$ 都是在旋转过后的系数轨道矩阵 `C_rot` $\mathbf{C}$ 为基给出，因此需要重新生成一下。

```python
h_mo = np.einsum("up, uv, vq -> pq", C_rot, h, C_rot)
eri_mo = np.einsum("up, vq, uvkl, kr, ls -> pqrs", C_rot, C_rot, eri, C_rot, C_rot)
```

```python
(
    + np.einsum("pq, pq ->", h_mo, rdm1)
    + 0.5 * np.einsum("pqrs, pqrs ->", eri_mo, rdm2)
    + mol.energy_nuc()
)
```

    -76.10403565383504

### 生成广义 Fock 矩阵

广义 Fock 矩阵 `gfock` $F_{pq}$ 区别于 RHF 的 Fock 矩阵 $f_{pq}$。其定义为

$$
F_{pq} = h_{pm} \gamma_{mq} + (pm|rs) \Gamma_{mr}^{qs}
$$

```python
gfock = np.einsum("pr, rq -> pq", h_mo, rdm1) + np.einsum("pmrs, mqrs -> pq", eri_mo, rdm2)
```

事实上，RHF 的 Fock 矩阵中，占据轨道部分也可以用类似的方法定义：

$$
\begin{align}
2 f_{ij} &= h_{im} \gamma_{mj}^\mathsf{RHF} + (im|rs) \Gamma_{mr}^{js, \mathsf{RHF}} \\
\Gamma_{pr}^{qs, \mathsf{RHF}} &= \gamma_{pq}^\mathsf{RHF} \gamma_{rs}^\mathsf{RHF} - \frac{1}{2} \gamma_{ps}^\mathsf{RHF} \gamma_{rq}^\mathsf{RHF}
\end{align}
$$

```python
rdm2_rhf = np.einsum("pq, rs -> pqrs", rdm1_rhf, rdm1_rhf) - 0.5 * np.einsum("ps, rq -> pqrs", rdm1_rhf, rdm1_rhf)
np.allclose(
    (np.einsum("pr, rq -> pq", h_mo, rdm1_rhf) + np.einsum("pmrs, mqrs -> pq", eri_mo, rdm2_rhf))[so, so],
    2 * fock_rot[so, so],
)
```

    True

但在 PySCF 的 CASSCF 模块中，似乎没有直接生成广义 Fock 矩阵的方式。但其有广义 Fock 的导数量，被称为轨道梯度 (Orbital Gradient) `gfock_grad` $x_{pq}$：

$$
x_{pq} = F_{pq} - F_{qp}
$$

```python
gfock_grad = gfock - gfock.T
np.allclose(
    mf_cas.unpack_uniq_var(mf_cas.get_grad(C_rot, (rdm1, rdm2), mf_cas.ao2mo(C_rot))),
    gfock_grad
)
```

    True

至此，所有生成 OO-MP2 所需要的单步复杂计算都已经涵盖到了。

## 轨道旋转的意义

讨论到现在，我们仅仅知道了 OO-MP2 的程序实现是如何进行的；但对其根源的合理性问题，我们在这里才开始说明。

出于一般性，我们现在考虑 Non-HF 形式的轨道系数，即相对于 RHF 系数已经一定程度的旋转。该 Non-HF 轨道系数称为 `C_base` $C_{\mu p}^\mathrm{base}$。我们之后的讨论都基于该 Non-HF 轨道系数开始。

```python
X = np.random.randn(nmo, nmo)
X = (X - X.T) * 0.02
C_base = C_rhf @ scipy.linalg.expm(X)
```

首先需要说明，轨道的旋转矩阵必须是正交矩阵 (酉矩阵)。这是因为轨道系数必须满足

$$
\mathbf{C}^\dagger \mathbf{S} \mathbf{C} = \mathbf{I}
$$

旋转矩阵 $\mathbf{U}$ 通过下式定义：$\mathbf{C} = \mathbf{C}^\mathrm{base} \mathbf{U}$。因此，

$$
\mathbf{C}^\dagger \mathbf{S} \mathbf{C} = \mathbf{U}^\dagger \mathbf{C}^\dagger \mathbf{S} \mathbf{C} \mathbf{U} = \mathbf{U}^\dagger \mathbf{I} \mathbf{U} = \mathbf{U}^\dagger \mathbf{U} = \mathbf{I}
$$

而任何正交矩阵都可以通过反对称矩阵 $\mathbf{X} = \mathbf{X}^\dagger$ 的幂次给出 $\mathbf{U} = \exp(\mathbf{X})$。

现在考察在微扰下，能量随轨道系数的变化情况。令一般情况下轨道系数 $C_{\mu p}$ 为关于反对称矩阵 $X_{pq}$ 的函数：

$$
\mathbf{C} = \mathbf{C}^\mathrm{base} \exp (\mathbf{X})
$$

而 $C_{\mu p}$ 对应的 MP2 能量写作关于 $X_{pq}$ 的函数 $E_\mathrm{tot}^\mathsf{MP2} (\mathbf{X})$。下面的代码 `eng_mp2_pert` 即是代入反对称矩阵 $X_{pq}$，生成 MP2 能量的函数。

```python
def eng_mp2_pert(X):
    C_rot = C_base @ scipy.linalg.expm(X)
    mf_prhf = scf.RHF(mol)
    mf_prhf.mo_occ, mf_prhf.mo_coeff = mo_occ, C_rot
    mf_pmp2 = mp.MP2(mf_prhf).run()
    return mf_pmp2.e_tot
```

由此，能量关于旋转矩阵的导数关系可以写为矩阵 `dX` ${\mathrm{d} \mathbf{X}}$，其维度为 $(p, q)$：

$$
{\mathrm{d}X}_{pq} = \frac{\partial E_\mathrm{tot}^\mathsf{MP2}}{\partial X_{pq}}
$$

这种导数可以写成三点差分的数值微分的形式：

$$
{\mathrm{d}X}_{pq} \simeq \frac{E_\mathrm{tot}^\mathsf{MP2} (d_{pq}) - E_\mathrm{tot}^\mathsf{MP2} (- d_{pq})}{2 d_{pq}}
$$

$E_\mathrm{tot}^\mathsf{MP2} (d_{pq})$ 的意义是，反对称矩阵 $\mathbf{X}$ 仅在第 $p$ 行、第 $q$ 列上，$X_{pq} = d_{pq}$；且在第 $q$ 行、第 $p$ 列上，$X_{qp} = - d_{pq}$；其它位置上，$\mathbf{X}$ 均取到零值。如果 $p = q$，那么 $\mathbf{X} = \mathbf{0}$。生成这种反对称矩阵的函数 `gen_pert_X` 如下所示：

```python
def gen_pert_X(p, q, interval):
    X = np.zeros((nmo, nmo))
    X[p, q] = interval
    X -= X.T
    return X
```

那么依据上述反对称矩阵，所求出的 MP2 能量随 $X_{pq}$ 变化的数值导数 ${\mathrm{d}X}_{pq}$ 的生成函数如下：

```python
def eng_mp2_numdiff(p, q, interval):
    X_positive = gen_pert_X(p, q, interval)
    X_negative = gen_pert_X(p, q, -interval)
    return (eng_mp2_pert(X_positive) - eng_mp2_pert(X_negative)) / (2 * interval)
```

对角标 $p, q$ 循环，我们就能求出完整的导数矩阵 `dX` ${\mathrm{d} \mathbf{X}}$ (这里选取的数值微分的间隙值 `interval` 为 $10^{-4}$ a.u.)：

```python
dX = np.zeros((nmo, nmo))
for a in range(nmo):
    for i in range(nmo):
        dX[a, i] = eng_mp2_numdiff(a, i, 1e-4)
dX
```

    array([[ 0.    , -0.    , -0.    , -0.    ,  0.    ,  0.243 , -0.6191,  0.7465, -1.7459,  1.0327, -0.5983,  1.3028, -1.8584],
           [ 0.    ,  0.    , -0.    ,  0.    ,  0.    , -0.044 , -0.0219,  0.3123, -0.1693, -0.1755,  0.1931, -0.0509,  0.033 ],
           [ 0.    ,  0.    ,  0.    , -0.    ,  0.    , -0.0894,  0.0924,  0.2349,  0.1175,  0.0443, -0.3922,  0.1505, -0.4868],
           [ 0.    , -0.    ,  0.    ,  0.    , -0.    ,  0.0648, -0.0899, -0.0568, -0.0668, -0.137 ,  0.2291, -0.0017, -0.1029],
           [-0.    , -0.    , -0.    ,  0.    ,  0.    , -0.0091,  0.0252,  0.1021, -0.0093, -0.0796,  0.0327, -0.067 , -0.0761],
           [-0.243 ,  0.044 ,  0.0894, -0.0648,  0.0091,  0.    , -0.    ,  0.    , -0.    ,  0.    ,  0.    ,  0.    , -0.    ],
           [ 0.6191,  0.0219, -0.0924,  0.0899, -0.0252,  0.    ,  0.    ,  0.    , -0.    ,  0.    ,  0.    ,  0.    , -0.    ],
           [-0.7465, -0.3123, -0.2349,  0.0568, -0.1021, -0.    , -0.    ,  0.    ,  0.    ,  0.    ,  0.    ,  0.    , -0.    ],
           [ 1.7459,  0.1693, -0.1175,  0.0668,  0.0093,  0.    ,  0.    , -0.    ,  0.    ,  0.    ,  0.    ,  0.    , -0.    ],
           [-1.0327,  0.1755, -0.0443,  0.137 ,  0.0796, -0.    , -0.    , -0.    , -0.    ,  0.    , -0.    ,  0.    ,  0.    ],
           [ 0.5983, -0.1931,  0.3922, -0.2291, -0.0327, -0.    , -0.    , -0.    , -0.    ,  0.    ,  0.    , -0.    ,  0.    ],
           [-1.3028,  0.0509, -0.1505,  0.0017,  0.067 , -0.    , -0.    , -0.    , -0.    , -0.    ,  0.    ,  0.    , -0.    ],
           [ 1.8584, -0.033 ,  0.4868,  0.1029,  0.0761,  0.    ,  0.    ,  0.    ,  0.    , -0.    , -0.    ,  0.    ,  0.    ]])

注意到这是一个反对称且分块的矩阵；在占据与非占分块值完全为零，有值处仅有 $\mathrm{d} X_{ai} = - \mathrm{d} X_{ia}$。这实际上近乎等于 2 倍的轨道梯度矩阵 `2 * gfock_grad`：

$$
\mathrm{d} X_{pq} = 2 x_{pq} = 2 (F_{pq} - F_{qp})
$$

```python
mf_prhf = scf.RHF(mol)
mf_prhf.mo_occ, mf_prhf.mo_coeff = mo_occ, C_base
mf_pmp2 = mp.MP2(mf_prhf).run()
rdm1, rdm2 = mf_pmp2.make_rdm1(), mf_pmp2.make_rdm2()
gfock_grad = mf_cas.unpack_uniq_var(mf_cas.get_grad(C_base, (rdm1, rdm2), mf_cas.ao2mo(C_base)))
```

```python
np.allclose(2 * gfock_grad, dX, atol=5e-6)
```

    True

因此，可以说 OO-MP2 的意义是，找到一个合适的 $\mathbf{C}^\mathrm{base}$，使得对于任意的很小的、用于旋转的反对称矩阵 $\mathbf{X}$，有 $E_\mathrm{tot}^\mathsf{MP2} (\mathbf{X})$ 不会更改。

## OO-MP2 能量并非一定比 MP2 低

在文档最后，我们会指出，OO-MP2 能量并非 MP2 的下界。尽管 OO-MP2 看起来对轨道进行变分式的优化，但其变分的对象应当认为是 Hylleraas 泛函，而非总 MP2 能量。

对于下述拉长的氢分子，就是一个 OO-MP2 能量比 MP2 能量高的例子。

```python
mol = gto.Mole()
mol.atom = """
H  0. 0. 0.
H  0. 0. 15.
"""
mol.basis = "6-31G"
mol.verbose = 0
mol.build()
```

    <pyscf.gto.mole.Mole at 0x7efd968f4310>

其 MP2 能量为

```python
mol.RHF().run().MP2().run().e_tot
```

    -1.7458592201255043

而其 OO-MP2 能量为

```python
mf_cas = mcscf.CASSCF(mol.RHF().run(), mol.nao, mol.nelectron)
mf_cas.fcisolver = MP2AsFCISolver()
mf_cas.internal_rotation = True
cas_result = mf_cas.kernel()
cas_result[0]
```

    -1.7280760742391805

但即使 OO-MP2 的能量比 MP2 高，它仍然无法解决 MP2 方法在解离两个氢原子所产生的相当大的解离误差。

[^Sun-Chan.JCP.2020]: Recent Developments in the PySCF Program Package. *J. Chem. Phys.* **2020**, *153* (2), 24109. doi: [10.1063/5.0006074](https://doi.org/10.1063/5.0006074).


# "Revisiting AlexNet"

- toc: true
- branch: master
- badges: true
- comments: true
- categories: [fastpages, jupyter]
- image: images/chef_bepita.jpg
- hide: false
- search_exclude: true
- metadata_key1: metadata_value1
- metadata_key2: metadata_value2

## Architecture Tricks

### ReLU Nonlinearity

The first innovation that the paper brings to the table is the now most common ReLU ($ f(x) = max(0,x) $) activation function. Till then the most common way of applying non linearity to neural networks was the tanh ( $ f(x) = (1 + e^{-x})^{-1} $) activation function. Actually ReLUs were first ??? introduced by Nair and Hinton[???] but it was in the context of restricted Boltzman machines ???. The advantage of ReLU over tanh is that ReLU does not saturate and even if some fraction of the current batch will produce a positive input learning will occur in that pass. Therefore training a deep neural network with ReLUs is considerably faster.

### Training on multiple GPUs

Back in the old days of 2011-2012 RTX2080 Ti were not available. The GPU they had was GTX 580 with 3GB of memory. But AlexNet implemntation taks advantage of the ability of GPUs to cross parrallelize the training. So AlexNet was trained on a pair of GTX 580 and the network architecture was designed for that purpase

### Local Response Normalization

The use of ReLUs makes input normalization less crucial because because of their non saturating nature. As long as there are positive inputs in a batch learning will happen. That said, the authors used normalization to make training faster. The normalization used is not common nowaddays and called `Local Response Normalization` and is defined as follows:

\begin{equation}
b_{x,y}^{i} = \frac{a_{x,y}^{i}}{\left( k + \alpha \sum_{j=max(0,i-n/2)}^{min(N-1,i+n/2) } (a_{x,y}^{j})^{2} \right)^{\beta}}
\end{equation}

where $k, n, \alpha, \beta$ hyperparameters. The convolution filters are ordered by index before the training begins and the normalization is done among the n "neighbour" filters.  The authors call this `brightness normalization` because the mean of activations is not subtracted.

### Overlapping Pooling

Max pooling is applied to the output of the first, second and fifth layers. The pooling kernals overlap in a sense that the stride is smaller than the kernel width.the authors mention that it helps to reduce overfitting.

## Reduce Overfitting Tricks

### Data Augmentation

Data augmentation is done in two forms:

Transformations - Randomly take a 224X224 patch from our 256X256 images and their horisontal reflections.

RGB Intensities - Preform PCA on the ImageNET training set and add multiples of the found princples. Specifically to each pixel the values
\begin{equation}
[P_{1},P_{2},P_{3}][\alpha_{1}\lambda_{1},\alpha_{2}\lambda_{2},\alpha_{3}\lambda_{3}]^{T}
\end{equation}

are added to each pixel values where $P_{i}$ and $\alpha_{i}$ are the $i$-th eigenvalues and eigenvectors of 3X3 covariance matrix of RGB pixel values and $\alpha_{i}$ are drawn per image per use in train from a zero mean Gaussian distribution with standart deviation 0.1.

### Dropout

This models is one of the first ones that used dropout. In this thchnique every neuron's output is zeroed in a probability of 0.5 and does not contribute to the forward pass and does not participate in the backward pass. In this way we ensure that complex dependencies between neurons do not replace a single neuron contribution and simplicity is encouraged.
At test time we multiply all neaurons by 0.5 to keep the overall sum of weights.

## Bibiliography


# Multigrid Lecture and Homework Example

### Aron Ahmadia (US Army ERDC) and David Ketcheson (KAUST)

### Teaching Numerical Methods with IPython Notebooks, SciPy 2014

This example by Aron Ahmadia and David Ketcheson is licensed under a Creative Commons Attribution 4.0 International License.  All code examples are also licensed under the [MIT license](http://opensource.org/licenses/MIT).

```
from IPython.core.display import HTML
css_file = './example.css'
HTML(open(css_file, "r").read())
```

<link href='http://fonts.googleapis.com/css?family=Alegreya+Sans:100,300,400,500,700,800,900,100italic,300italic,400italic,500italic,700italic,800italic,900italic' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=Arvo:400,700,400italic' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=PT+Mono' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=Shadows+Into+Light' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=Philosopher:400,700,400italic,700italic' rel='stylesheet' type='text/css'>

# Multigrid

Multigrid is one of the great algorithms for numerically solving PDEs.  It is also one of the few essentially optimal algorithms for solving a linear system of equations, since it computes the solution of an $N\times N$ system in ${\mathcal O}(N\log(N))$ -- or even just ${\mathcal O}(N)$ -- operations.

This notebook is meant to accompany a reading of [Section 4.6 of Randall LeVeque's text on finite difference methods](http://0-epubs.siam.org.library.kaust.edu.sa/doi/abs/10.1137/1.9780898717839.ch4).  Other good resources are cited there.

```
%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
from clawpack.visclaw.JSAnimation import IPython_display
```

## A two-point boundary value problem
Let's use Jacobi's method to solve the one-dimensional boundary value problem

\begin{align}
u''(x) & = f(x) & 0<x<1
\end{align}

with boundary conditions $u(0) = 1, u(1) = 3$ where

\begin{align}
f(x) & = -20 + \frac{1}{2} \phi''(x) \cos(\phi(x)) - \frac{1}{2} (\phi'(x))^2\sin(\phi(x)) \\ \\
\phi(x) & = 20\pi x^3.
\end{align}

This problem appears on page 103 of the text. In the box below, we set up the grid, the right hand side,  and the boundary values.

```
m=2**8-1                           # Number of grid points
h=1./(m+1)                         # Mesh width
x=np.linspace(0,1,m+2); x=x[1:-1]  # grid
alpha=1.; beta=3.                  # boundary values
phi = lambda x: 20.* np.pi * x**3
f = lambda x: -20 + 0.5*120*np.pi*x * np.cos(phi(x)) - \
                0.5*(60*np.pi*x**2)**2 * np.sin(phi(x)) # RHS
plt.plot(x,f(x)); plt.xlabel('x'); plt.ylabel('f(x)')
```

## Jacobi's Method
A centered-difference method for this problem gives us the system of equations

\begin{align}
\frac{U_{i-1} - 2 U_i + U_{i+1}}{h^2} & = f(x_i).
\end{align}

Solving for $U_i$ gives
\begin{align}
U_i & = \frac{1}{2} ( U_{i+1} + U_{i-1} ) - \frac{h^2}{2} f(x_i).
\end{align}

Jacobi's method consists of using the equation above as a fixed-point iteration:

\begin{align}
U_i^{[k+1]} & = \frac{1}{2} ( U_{i+1}^{[k]} + U_{i-1}^{[k]} ) - \frac{h^2}{2} f(x_i).
\end{align}

In the cell below, implement Jacobi's method for the BVP.  Don't forget about the boundary conditions.  Take 100 iterations with it and plot the result.

```
U=np.linspace(alpha,beta,m)       # Just use a straight line as initial guess
UU=[U]
F=0.5*h**2*f(x)                   # Construct the RHS
F[0]-=alpha/2.; F[-1]-=beta/2.    # boundary conditions
e=np.ones(m-1)
G=0.5*(np.diag(e,-1)+np.diag(e,1)) # Jacobi matrix (average neighbors)

for i in range(100):
    U=np.dot(G,U) - F
    UU.append(U)
```

If you can, try storing the solution after each Jacobi iteration in a list, and make an animated plot of the solution after each iteration.

```
xf = np.linspace(0,1,1000);
uf = 1.+12.*xf-10.*xf**2 + 0.5*np.sin(phi(xf))  # Exact solution on fine grid
u = 1.+12.*x-10.*x**2 + 0.5*np.sin(phi(x))  # Exact solution on computational grid

fig = plt.figure()
ax = plt.axes(xlim=(0, 1), ylim=(-3, 6))
ax.plot(xf,uf)
line1, = ax.plot([],[],'-')
line2, = ax.plot([],[],'-o')
plt.legend(['Exact','Jacobi','Error'],loc=2)

def fplot(i):
    line1.set_data(x,UU[i])
    line2.set_data(x,UU[i]-u)
    ax.set_title('# of iterations: '+str(i))
    return line1,

animation.FuncAnimation(fig, fplot, frames=range(len(UU)))
```

## Convergence of Jacobi

Notice how the high-frequency components of the error are damped very quickly, but the low-frequency error barely changes.  This is because the matrix $G$ appearing in Jacobi's method is tridiagonal, so in a single iteration each solution point only 'feels' the values of its immediate neighbors.  It takes very many iterations for low-frequency information to propagate.  As we'll see, part of the power of multigrid lies in enabling information to propagate rapidly between distant points in the solution.

At each Jacobi iteration, the error in the $p$th Fourier mode
$$u^p = \sin(p\pi x)$$
is reduced by a factor of

$$|\hat{\gamma}_p|= |\cos(p\pi h)|.$$

Which modes of the error converge most slowly?

## Under-relaxation
The *under-relaxed Jacobi method* is obtained by taking, at each iteration, a weighted average of the previous iterate and the new iterate:

\begin{align}
U_i^{[k+1]} & = (1-\omega)U^{[k]}_i + \omega\left(\frac{1}{2} ( U_{i+1}^{[k]} + U_{i-1}^{[k]} ) - \frac{h^2}{2} f(x_i)\right).
\end{align}

Now at each iteration the $p$th Fourier mode is reduced by a factor of

$$|\hat{\gamma}_p| := |1-\omega + \omega\cos(p\pi h)| = 1 + \omega(\cos(p \pi h) - 1)|.$$

The code below plots this range of values. Try changing $\omega$ to $\frac{2}{3}$ and notice how the second half of the eigenvalues (for $p\ge m/2$) are all shifted to have smaller absolute value.

```
M=30; H=1./(M+1)
omegas=np.linspace(0.,1.)
p=np.arange(1,M+1)

fig = plt.figure()
ax = plt.axes(xlim=(0, M), ylim=(0,1))
ax.set_xlabel('p', fontsize=20)
ax.set_ylabel('$\widehat{\gamma}_p$', fontsize=20)
line1, = ax.plot([],[],'o')

def fplot(i):
    omega = omegas[i]
    gamma = abs(1.+omega*(np.cos(p*np.pi*H)-1))
    line1.set_data(p,gamma)
    ax.set_title('$\omega = %02.2f $' % omega, fontsize=20)
    return line1,

animation.FuncAnimation(fig, fplot, frames=range(50))
```

In the cell below, implement under-relaxed Jacobi and run 100 iterations for the BVP with $\omega=2/3$.  Does it make the results any better?  Why or why not?

```
omega = 2./3.
U=np.linspace(alpha,beta,m)  # Just use a straight line as initial guess
UU=[U]
F=0.5*h**2*f(x)
F[0]-=alpha/2.; F[-1]-=beta/2.    # Construct the RHS, including boundary conditions
e=np.ones(m-1)
G=0.5*(np.diag(e,-1)+np.diag(e,1))
niter=100

for i in range(niter):
    U=(1.-omega)*U+omega*(np.dot(G,U)-F)
    UU.append(U)

xf = np.linspace(0,1,1000);
uf = 1.+12.*xf-10.*xf**2 + 0.5*np.sin(phi(xf))  # Exact solution on fine grid
u = 1.+12.*x-10.*x**2 + 0.5*np.sin(phi(x))  # Exact solution on computational grid

fig = plt.figure()
ax = plt.axes(xlim=(0, 1), ylim=(-3, 6))
ax.plot(xf,uf)
line1, = ax.plot([],[],'-')
line2, = ax.plot([],[],'-o')
plt.legend(['Exact','Jacobi','Error'],loc=2)

def fplot(i):
    line1.set_data(x,UU[i])
    line2.set_data(x,UU[i]-u)
    ax.set_title('# of iterations: '+str(i))
    return line1,

animation.FuncAnimation(fig, fplot, frames=range(niter))
```

## Grid representations of a given frequency (*aliasing*)
The code below plots a mode of a certain frequency, and then shows how it is represented on grids with decreasing numbers of points.  Notice how a *low frequency* mode (i.e., one with many points per wavelength) on a fine grid becomes a *high frequency* mode (one with very few points per wavelength).

Eventually, on grids with very few points, the mode is *aliased* to a completely different mode.

```
from matplotlib import animation
import matplotlib.pyplot as plt
from clawpack.visclaw.JSAnimation import IPython_display

fig = plt.figure()
ax = plt.axes(xlim=(0, 1), ylim=(-1.2, 1.2))

p=5
xf=np.linspace(0,1,1000) # fine grid

line1, = ax.plot(xf, np.sin(p*np.pi*xf), '-r', lw=2)
line2, = ax.plot([],[],'o-',lw=2)

def fplot(m):
    x=np.linspace(0,1,m+2); # grid
    line2.set_data(x,np.sin(p*np.pi*x))
    ax.set_title('m='+str(m))
    return line2,

animation.FuncAnimation(fig, fplot, frames=(63,31,15,7,3,1))
```

## Restriction and Prolongation
The multigrid idea is to solve successively on different grids, so that every mode has a small damping factor $\hat{\gamma}$ on some grid.

In order to implement multigrid, we need a way to take a function with values on a fine grid and restrict it to its values on a coarser grid (restriction). We also need a way to take a function with values on a coarse grid and interpolate onto a fine grid (prolongation). The functions below do just this, using linear interpolation in the case of the prolongation operator. Make sure you understand what they are doing. You may find it helpful to look up slicing notation in the NumPy help online.

```
def coarsen(f):
    return f[1::2]  # This slicing just takes the odd-numbered points

def interpolate(f,alpha,beta):
    m_coarse=len(f)
    m_fine  =2*m_coarse+1
    f_interp = np.zeros(m_fine)
    f_interp[1::2]=f                          #Set the values of the odd numbered points
    f_interp[2:-1:2]=0.5*(f[:-1]+f[1:])       #Set the values of the (interior) even numbered points
    f_interp[0]=0.5*(f_interp[1]+alpha)     #Set the values of the endpoints
    f_interp[-1]=0.5*(f_interp[-2]+beta)
    return f_interp
```

## Correction using a second grid

Now we compute the residual $$r = f-AU$$ on the fine grid, coarsen it, and solve $$Ae=-r$$ for the approximate error on the coarse grid. Then interpolate and subtract the interpolated error from the fine grid solution.

```
#compute the residual
A=2./h**2*(G-np.eye(m))
F=f(x); F[0]-=alpha/h**2; F[-1]-=beta/h**2
r=F-np.dot(A,U)
print max(abs(r))
omega=2./3.
m2=(m-1)/2.               # Number of points in coarse grid
h2=1./(m2+1)              # coarse mesh width
x2=coarsen(x)             # coarse grid points
r2=coarsen(r)             # residual restricted to coarse grid
err2=np.zeros(m2)         # initial guess for the error
e2=np.ones(m2-1); G2=0.5*(np.diag(e2,-1)+np.diag(e2,1))
for i in range(3):                        #Solve Ae=-r by Jacobi iteration
    err2=(1.-omega)*err2+omega*(np.dot(G2,err2)+0.5*h2**2*r2)
err=interpolate(err2,0,0)     # interpolate the error
print err[0]
U=U-err                   # and use it to correct our solution
plt.clf()
plt.plot(x,u,x,U,x,U-u)
plt.legend(['Exact','Jacobi','Error'],loc='best')
```

Finally, in order to keep our multigrid code clean, we'll write a function to set up and take a fixed number of under-relaxed Jacobi iterations. The function returns the approximate solution and the residual.

```
def Jacobi(U,f,alpha,beta,m,nu):
    """Perform nu Jacobi iterations on a grid with m points, with initial guess U, right hand side function f and
       Dirichlet boundary conditions with values alpha and beta.  The function returns both the approximate
       solution and the residual."""
    omega=2./3.
    h=1./(m+1)
    F=0.5*h**2*f.copy()
    F[0]-=alpha/2.; F[-1]-=beta/2.
    e=np.ones(m-1)
    G=0.5*(np.diag(e,-1)+np.diag(e,1))
    for i in range(nu):
        U=(1.-omega)*U + omega*(np.dot(G,U)-F)
    A=2./h**2*(G-np.eye(m))
    FF=f.copy(); FF[0]-=alpha/h**2; FF[-1]-=beta/h**2
    rr=FF-np.dot(A,U)
    return U,rr
```

##Multigrid V-cycle
Now let's do a full V-cycle. Look carefully through the code below until you understand what it does.

```
k=8;
m=2**k-1
rdep=k-1  # Recursion depth; this is how many grids down we want to go
          # rdep=k-1 gives a full V-cycle
nu=3      # Number of Jacobi iterations to take at each step
U=np.linspace(alpha,beta,m)  # Initial guess
x=np.linspace(0,1,m+2); x=x[1:-1]  # grid
phi = lambda x: 20.* np.pi * x**3
u = 1.+12.*x-10.*x**2 + 0.5*np.sin(phi(x))
f = lambda x: -20 + 0.5*120*np.pi*x * np.cos(phi(x)) - 0.5*(60*np.pi*x**2)**2 * np.sin(phi(x))
alpha=1.; beta=3.
F=f(x)
r=[None]*(rdep+1); error=[None]*(rdep+1)        # This just initializes these lists to have the right length
U,rr=Jacobi(U,F,alpha,beta,m,nu)  # Initial iteration on fine grid
for i in range(1,rdep+1): # Going down
    m=(m-1)/2 #  = 2**(k-i) - 1
    r[i]=coarsen(rr)             # residual restricted to next coarser grid
    error[i],rr=Jacobi(np.zeros(m),-r[i],0.,0.,m,nu)
for i in range(1,rdep): # Coming up
    m=2*m+1
    err=error[rdep-i]-interpolate(error[rdep+1-i],0,0)     # Interpolate and subtract the correction
    error[-i-1],rr=Jacobi(err,-r[rdep-i],0.,0.,m,nu)
m=2*m+1
U=U-interpolate(error[1],0,0)                   # final solution correction
U,rr=Jacobi(U,F,alpha,beta,m,nu)                #Final iterations on original grid
plt.clf()
plt.plot(x,u,x,U,'o-',x,U-u)
plt.legend(['Exact','Jacobi','Error'],loc='best')
```

Since we have saved the errors on each grid, we can plot them to see how they represent the accumulation of errors at different scales.

```
fig = plt.figure()
ax = plt.axes(xlim=(0, 1), ylim=(-5,5))
line1, = ax.plot([],[],'-ok',lw=2)

def fplot(i):
    y = error[-i]
    m = len(y)
    x=np.linspace(0,1,m+2); x=x[1:-1]
    line1.set_data(x,y)
    return line1,

animation.FuncAnimation(fig, fplot, frames=range(1,len(error)))
```

## Conclusion

Now you have the basic ideas of multigrid.  Of course, for this 1-dimensional problem (where the system to be solved is tridiagonal) it does not provide a big advantage over more naive methods.  The power of multigrid is that it can be used in multiple dimensions, where the algebraic system is much more challenging for other methods.  Furthermore, multigrid can be applied to nonlinear and time-dependent problems.

# Exercise #11:
Modify the V-cycle code above to answer the following questions. Try to explain your results.
(a) How does the accuracy change as we change the number of Jacobi iterations performed at each step?

(b) Is it better to use a finer grid, or more Jacobi iterations if we want to improve the solution accuracy?

(c) What happens if we don't perform any Jacobi iterations in the "up" part of the V-cycle?

(d) What happens if we don't recurse all the way down to the 1-point grid?

(e) What happens if we use the original Jacobi method, or some other value of $\omega$?

```

```

# Exercise #12
Use the V-cycle code above to implement the full multigrid algorithm discussed in the text. You will probably find it helpful to write a Vcycle() function first.

```

```

## A Multigrid interact

```
from IPython.display import display, HTML
from IPython.html import widgets
from IPython.html.widgets import interact, interactive

def Jacobi(U,f,alpha,beta,m,nu,omega=2./3):
    """Perform nu Jacobi iterations on a grid with m points, with initial guess U, right hand side function f and
       Dirichlet boundary conditions with values alpha and beta.  The function returns both the approximate
       solution and the residual."""
    h=1./(m+1)
    F=0.5*h**2*f.copy()
    F[0]-=alpha/2.; F[-1]-=beta/2.
    e=np.ones(m-1)
    G=0.5*(np.diag(e,-1)+np.diag(e,1))
    for i in range(nu):
        U=(1.-omega)*U + omega*(np.dot(G,U)-F)
    A=2./h**2*(G-np.eye(m))
    FF=f.copy(); FF[0]-=alpha/h**2; FF[-1]-=beta/h**2
    rr=FF-np.dot(A,U)
    return U,rr

def Vcycle(k=8,nu=3,omega=2./3,smooth_after_interp=True):
    m=2**k-1
    rdep=k-1  # Recursion depth; this is how many grids down we want to go
              # rdep=k-1 gives a full V-cycle
    alpha=1.; beta=3.
    U=np.linspace(alpha,beta,m)  # Initial guess
    x=np.linspace(0,1,m+2); x=x[1:-1]  # grid
    u = 1.+12.*x-10.*x**2 + 0.5*np.sin(phi(x))

    F=f(x)
    r=[None]*(rdep+1); error=[None]*(rdep+1) # This just initializes these lists to have the right length
    U,rr=Jacobi(U,F,alpha,beta,m,nu,omega)  # Initial iteration on fine grid
    for i in range(1,rdep+1): # Going down
        m=(m-1)/2 #  = 2**(k-i) - 1
        r[i]=coarsen(rr)             # residual restricted to next coarser grid
        error[i],rr=Jacobi(np.zeros(m),-r[i],0.,0.,m,nu,omega)
    for i in range(1,rdep): # Coming up
        m=2*m+1
        err=error[rdep-i]-interpolate(error[rdep+1-i],0,0)     # Interpolate and subtract the correction
        if smooth_after_interp:
            error[-i-1],rr=Jacobi(err,-r[rdep-i],0.,0.,m,nu,omega)
        else:
            error[-i-1] = err

    m=2*m+1
    U=U-interpolate(error[1],0,0)                   # final solution correction
    U,rr=Jacobi(U,F,alpha,beta,m,nu,omega)                #Final iterations on original grid

    fig, ax = plt.subplots(figsize=(8,6),subplot_kw={'axisbg':'#EEEEEE',
                                       'axisbelow':True})
    ax = plt.axes(xlim=(0, 1), ylim=(-5,6))
    ax.plot(x,u,x,U,'o-',x,U-u)
    plt.legend(['Exact','Jacobi','Error'],loc='best')

interact(Vcycle, k=(4,10,1), nu=(1,4,1), omega=(0,2,1./6))
```

```

```