```python
from IPython.display import Image
```

# CNTK 203: Reinforcement Learning Basics

Reinforcement learning (RL) is an area of machine learning inspired by behaviorist psychology, concerned with how [software agents](https://en.wikipedia.org/wiki/Software_agent) ought to take [actions](https://en.wikipedia.org/wiki/Action_selection) in an environment so as to maximize some notion of cumulative reward. In machine learning, the environment is typically formulated as a [Markov decision process](https://en.wikipedia.org/wiki/Markov_decision_process) (MDP) as many reinforcement learning algorithms for this context utilize [dynamic programming](https://en.wikipedia.org/wiki/Dynamic_programming) techniques.

In some machine learning settings, we do not have immediate access to labels, so we cannot rely on supervised learning techniques. If, however, there is something we can interact with and thereby get some feedback that tells us occasionally, whether our previous behavior was good or not, we can use RL to learn how to improve our behavior.

Unlike in supervised learning, in RL, labeled correct input/output pairs are never presented and sub-optimal actions are never explicitly corrected. This mimics many of the online learning paradigms which involves finding a balance between exploration (of conditions or actions never learnt before) and exploitation (of already learnt conditions or actions from previous encounters). Multi-arm bandit problems is one of the category of RL algorithms where exploration vs. exploitation trade-off have been thoroughly studied. See figure below for [reference](http://www.simongrant.org/pubs/thesis/3/2.html).

```python
# Figure 1
Image(url="https://cntk.ai/jup/polecart.gif", width=300, height=300)
```

**Problem**

We will use the [CartPole](https://gym.openai.com/envs/CartPole-v0) environment from OpenAI's [gym](https://github.com/openai/gym) simulator to teach a cart to balance a pole. As described in the link above, in the CartPole example, a pole is attached by an un-actuated joint to a cart, which moves along a frictionless track. The system is controlled by applying a force of +1 or -1 to the cart. A reward of +1 is provided for every timestep that the pole remains upright. The episode ends when the pole is more than 15 degrees from vertical, or the cart moves more than 2.4 units from the center. See figure below for reference.

**Goal**
Our goal is to prevent the pole from falling over as the cart moves with the pole in upright position (perpendicular to the cart) as the starting state. More specifically if the pole is less than 15 degrees from vertical while the cart is within 2.4 units of the center we will collect reward. In this tutorial, we will train till we learn a set of actions (policies) that lead to an average reward of 200 or more over last 50 batches.

In, RL terminology, the goal is to find _policies_ $a$, that maximize the _reward_ $r$ (feedback) through interaction with some environment (in this case the pole being balanced on the cart). So given a series of experiences $$s \xrightarrow{a} r, s'$$ we then can learn how to choose action $a$ in a given state $s$ to maximize the accumulated reward $r$ over time:
\begin{align}
Q(s,a) &= r_0 + \gamma r_1 + \gamma^2 r_2 + \ldots \newline
&= r_0 + \gamma \max_a Q^*(s',a)
\end{align}
where $\gamma \in [0,1)$ is the discount factor that controls how much we should value reward that is further away. This is called the [*Bellmann*-equation](https://en.wikipedia.org/wiki/Bellman_equation).

In this tutorial we will show how to model the state space, how to use the received reward to figure out which action yields the highest future reward.

We present two different popular approaches here:

**Deep Q-Networks**: DQNs have become famous in 2015 when they were successfully used to train how to play Atari just form raw pixels. We train neural network to learn the $Q(s,a)$ values (thus _Q-Network _). From these $Q$ functions values we choose the best action.

**Policy gradient**: This method directly estimates the policy (set of actions) in the network. The outcome is a learning of an ordered set of actions which leads to maximize reward by probabilistically choosing a subset of actions. In this tutorial, we learn the actions using a gradient descent approach to learn the policies.

In this tutorial, we focus how to implement RL in CNTK. We choose a straight forward shallow network. One can extend the approaches by replacing our shallow model with deeper networks that are introduced in other CNTK tutorials.

Additionally, this tutorial is in its early stages and will be evolving in future updates.

## Before we start...
Please run the following cell from the menu above or select the cell below and hit `Shift + Enter` to ensure the environment is ready. Verify that the following imports work in your notebook.

```python
from __future__ import print_function
from __future__ import division
import matplotlib.pyplot as plt
from matplotlib import style
import numpy as np
import pandas as pd
import seaborn as sns

style.use('ggplot')
%matplotlib inline
```

We use the following construct to install the OpenAI gym package if it is not installed. For users new to Jupyter environment, this construct can be used to install any python package.

```python
try:
    import gym
except:
    !pip install gym
    import gym
```

### Select the notebook run mode

There are two run modes:
- *Fast mode*: `isFast` is set to `True`. This is the default mode for the notebooks, which means we train for fewer iterations or train / test on limited data. This ensures functional correctness of the notebook though the models produced are far from what a completed training would produce.

- *Slow mode*: We recommend the user to set this flag to `False` once the user has gained familiarity with the notebook content and wants to gain insight from running the notebooks for a longer period with different parameters for training.

```python
isFast = True
```

# CartPole: Data and Environment

We will use the [CartPole](https://gym.openai.com/envs/CartPole-v0) environment from OpenAI's [gym](https://github.com/openai/gym) simulator to teach a cart to balance a pole. Please follow the links to get more details.

In every time step, the agent
 * gets an observation $(x, \dot{x}, \theta, \dot{\theta})$, corresponding to *cart position*, *cart velocity*, *pole angle with the vertical*, *pole angular velocity*,
 * performs an action `LEFT` or `RIGHT`, and
 * receives
  * a reward of +1 for having survived another time step, and
  * a new state $(x', \dot{x}', \theta', \dot{\theta}')$

The episode ends, if
 * the pole is more than 15 degrees from vertical and/or
 * the cart is moving more than 2.4 units from center.

The task is considered done, if
 * the agent achieved and averaged reward of 200 over the last 50 episodes (if you manage to get a reward of 200 averaged over the last 100 episode you can consider submitting it to OpenAI).

In fast mode these targets are relaxed.

# Part 1: DQN

After a transition $(s,a,r,s')$, we are trying to move our value function $Q(s,a)$ closer to our target $r+\gamma \max_{a'}Q(s',a')$, where $\gamma$ is a discount factor for future rewards and ranges in value between 0 and 1.

DQNs
 * learn the _Q-function_ that maps observation (state, action) to a `score`
 * use memory replay (previously recorded $Q$ values corresponding to different $(s,a)$ to decorrelate experiences (sequence state transitions)
 * use a second network to stabilize learning (*not* part of this tutorial)

### Setting up the model
\begin{equation}
l_1 = relu( x W_1 + b_1) \\
Q(s,a) = l_1 W_2 + b_2 \\
\end{equation}

We will start with a slightly modified version for Keras, https://github.com/jaara/AI-blog/blob/master/CartPole-basic.py, published by Jaromír Janisch in his [AI blog](https://jaromiru.com/2016/09/27/lets-make-a-dqn-theory/), and will then incrementally convert it to use CNTK.

We use a simple two-layer densely connected network, for simpler illustrations. More advance networks can be substituted.

**CNTK** concepts: The commented out code is meant to be an illustration of the similarity of concepts between CNTK API/abstractions against Keras.

```python
import random, numpy, math, os

#from keras.models import Sequential
#from keras.layers import *
#from keras.optimizers import *
from cntk import *
from cntk.layers import *
from cntk.ops.sequence import input
# Select the right target device when this notebook is being tested:
if 'TEST_DEVICE' in os.environ:
    if os.environ['TEST_DEVICE'] == 'cpu':
        device.try_set_default_device(device.cpu())
    else:
        device.try_set_default_device(device.gpu(0))
```

STATE_COUNT = 4 (corresponding to $(x, \dot{x}, \theta, \dot{\theta})$),

ACTION_COUNT = 2 (corresponding to `LEFT` or `RIGHT`)

```python
env = gym.make('CartPole-v0')

STATE_COUNT  = env.observation_space.shape[0]
ACTION_COUNT = env.action_space.n

STATE_COUNT, ACTION_COUNT
```

    [2017-03-29 10:30:29,176] Making new env: CartPole-v0

    (4, 2)

Note: in the cell below we highlight how one would do it in Keras. And a marked similarity with CNTK. While CNTK allows for more compact representation, we present a slightly verbose illustration for ease of learning.

Additionally, you will note that, CNTK model doesn't need to be compiled explicitly and is implicitly done when data is processed during training.

CNTK effectively uses available memory on the system between minibatch execution. Thus the learning rates are stated as **rates per sample** instead of **rates per minibatch** (as with other toolkits).

```python
# Targetted reward
REWARD_TARGET = 30 if isFast else 200
# Averaged over these these many episodes
BATCH_SIZE_BASELINE = 20 if isFast else 50

H = 64 # hidden layer size

class Brain:
    def __init__(self):
        self.params = {}
        self.model, self.trainer, self.loss = self._create()
        # self.model.load_weights("cartpole-basic.h5")

    def _create(self):
        observation = input(STATE_COUNT, np.float32, name="s")
        q_target = input(ACTION_COUNT, np.float32, name="q")

        # model = Sequential()
        # model.add(Dense(output_dim=64, activation='relu', input_dim=STATE_COUNT))
        # model.add(Dense(output_dim=ACTION_COUNT, activation='linear'))

        # Following a style similar to Keras
        l1 = Dense(H, activation=relu)
        l2 = Dense(ACTION_COUNT)
        unbound_model = Sequential([l1, l2])
        model = unbound_model(observation)

        self.params = dict(W1=l1.W, b1=l1.b, W2=l2.W, b2=l2.b)

        lr = 0.00025
        # opt = RMSprop(lr=0.00025)
        # model.compile(loss='mse', optimizer=opt)

        # loss='mse'
        loss = reduce_mean(square(model - q_target), axis=0)
        meas = reduce_mean(square(model - q_target), axis=0)

        # optimizer=opt
        lr_schedule = learning_rate_schedule(lr, UnitType.minibatch)
        learner = sgd(model.parameters, lr_schedule, gradient_clipping_threshold_per_sample=10)
        trainer = Trainer(model, (loss, meas), learner)

        # CNTK: return trainer and loss as well
        return model, trainer, loss

    def train(self, x, y, epoch=1, verbose=0):
        #self.model.fit(x, y, batch_size=64, nb_epoch=epoch, verbose=verbose)
        arguments = dict(zip(self.loss.arguments, [x,y]))
        updated, results =self.trainer.train_minibatch(arguments, outputs=[self.loss.output])

    def predict(self, s):
        return self.model.eval([s])
```

The `Memory` class stores the different states, actions and rewards.

```python
class Memory:   # stored as ( s, a, r, s_ )
    samples = []

    def __init__(self, capacity):
        self.capacity = capacity

    def add(self, sample):
        self.samples.append(sample)

        if len(self.samples) > self.capacity:
            self.samples.pop(0)

    def sample(self, n):
        n = min(n, len(self.samples))
        return random.sample(self.samples, n)
```

The `Agent` uses the `Brain` and `Memory` to replay the past actions to choose optimal set of actions that maximize the rewards.

```python
MEMORY_CAPACITY = 100000
BATCH_SIZE = 64

GAMMA = 0.99 # discount factor

MAX_EPSILON = 1
MIN_EPSILON = 0.01 # stay a bit curious even when getting old
LAMBDA = 0.0001    # speed of decay

class Agent:
    steps = 0
    epsilon = MAX_EPSILON

    def __init__(self):
        self.brain = Brain()
        self.memory = Memory(MEMORY_CAPACITY)

    def act(self, s):
        if random.random() < self.epsilon:
            return random.randint(0, ACTION_COUNT-1)
        else:
            return numpy.argmax(self.brain.predict(s))

    def observe(self, sample):  # in (s, a, r, s_) format
        self.memory.add(sample)

        # slowly decrease Epsilon based on our eperience
        self.steps += 1
        self.epsilon = MIN_EPSILON + (MAX_EPSILON - MIN_EPSILON) * math.exp(-LAMBDA * self.steps)

    def replay(self):
        batch = self.memory.sample(BATCH_SIZE)
        batchLen = len(batch)

        no_state = numpy.zeros(STATE_COUNT)

        # CNTK: explicitly setting to float32
        states = numpy.array([ o[0] for o in batch ], dtype=np.float32)
        states_ = numpy.array([(no_state if o[3] is None else o[3]) for o in batch ], dtype=np.float32)

        p = agent.brain.predict(states)
        p_ = agent.brain.predict(states_)

        # CNTK: explicitly setting to float32
        x = numpy.zeros((batchLen, STATE_COUNT)).astype(np.float32)
        y = numpy.zeros((batchLen, ACTION_COUNT)).astype(np.float32)

        for i in range(batchLen):
            s, a, r, s_ = batch[i]

            # CNTK: [0] because of sequence dimension
            t = p[0][i]
            if s_ is None:
                t[a] = r
            else:
                t[a] = r + GAMMA * numpy.amax(p_[0][i])

            x[i] = s
            y[i] = t

        self.brain.train(x, y)
```

### Brain surgery

As any learning experiences, we expect to see the initial state of actions to be wild exploratory and over the iterations the system learns the range of actions that yield longer runs and collect more rewards. The tutorial below implements the [$\epsilon$-greedy](https://en.wikipedia.org/wiki/Reinforcement_learning) approach.

```python
def plot_weights(weights, figsize=(7,5)):
    '''Heat map of weights to see which neurons play which role'''
    sns.set(style="white")
    f, ax = plt.subplots(len(weights), figsize=figsize)
    cmap = sns.diverging_palette(220, 10, as_cmap=True)

    for i, data in enumerate(weights):
        axi = ax if len(weights)==1 else ax[i]
        if isinstance(data, tuple):
            w, title = data
            axi.set_title(title)
        else:
            w = data

        sns.heatmap(w.asarray(), cmap=cmap, square=True, center=True, #annot=True,
                    linewidths=.5, cbar_kws={"shrink": .25}, ax=axi)
```

### Exploration - exploitation trade-off

Note that the initial $\epsilon$ is set to 1 which implies we are entirely exploring but as steps increase we reduce exploration and start leveraging the learnt space to collect rewards (a.k.a. exploitation) as well.

```python
def epsilon(steps):
    return MIN_EPSILON + (MAX_EPSILON - MIN_EPSILON) * np.exp(-LAMBDA * steps)
plt.plot(range(10000), [epsilon(x) for x in range(10000)], 'r')
plt.xlabel('step');plt.ylabel('$\epsilon$')
```

We are now ready to train our agent using **DQN**. Note this would take anywhere between 2-10 min and we stop whenever the learner hits the average reward of 200 over past 50 batches. One would get better results if they could train the learner until say one hits a reward of 200 or higher for say larger number of runs. This is left as an exercise.

```python
TOTAL_EPISODES = 2000 if isFast else 3000

def run(agent):
    s = env.reset()
    R = 0

    while True:
        # Uncomment the line below to visualize the cartpole
        # env.render()

        # CNTK: explicitly setting to float32
        a = agent.act(s.astype(np.float32))

        s_, r, done, info = env.step(a)

        if done: # terminal state
            s_ = None

        agent.observe((s, a, r, s_))
        agent.replay()

        s = s_
        R += r

        if done:
            return R

agent = Agent()

episode_number = 0
reward_sum = 0
while episode_number < TOTAL_EPISODES:
    reward_sum += run(agent)
    episode_number += 1
    if episode_number % BATCH_SIZE_BASELINE == 0:
        print('Episode: %d, Average reward for episode %f.' % (episode_number,
                                                               reward_sum / BATCH_SIZE_BASELINE))
        if episode_number%200==0:
            plot_weights([(agent.brain.params['W1'], 'Episode %i $W_1$'%episode_number)], figsize=(14,5))
        if reward_sum / BATCH_SIZE_BASELINE > REWARD_TARGET:
            print('Task solved in %d episodes' % episode_number)
            plot_weights([(agent.brain.params['W1'], 'Episode %i $W_1$'%episode_number)], figsize=(14,5))
            break
        reward_sum = 0
agent.brain.model.save('dqn.mod')
```

If you run it, you should see something like
```
Episode: 20, Average reward for episode 20.700000.
Episode: 40, Average reward for episode 20.150000.
Episode: 60, Average reward for episode 21.100000.
...
Episode: 960, Average reward for episode 20.150000.
Episode: 980, Average reward for episode 26.700000.
Episode: 1000, Average reward for episode 32.900000.
Task solved in 1000 episodes
```

#### Task 1.1
Rewrite the model without using the layer lib.
#### Task 1.2
Play with different [learners](https://cntk.ai/pythondocs/cntk.learner.html#module-cntk.learner). Which one works better? Worse? Think about which parameters you would need to adapt when switching from one learner to the other.

### Running the DQN model

```python
import cntk as C
env = gym.make('CartPole-v0')

num_episodes = 10  # number of episodes to run

modelPath = 'dqn.mod'
root = C.load_model(modelPath)

for i_episode in range(num_episodes):
    print(i_episode)
    observation = env.reset()  # reset environment for new episode
    done = False
    while not done:
        if not 'TEST_DEVICE' in os.environ:
            env.render()
        action = np.argmax(root.eval([observation.astype(np.float32)]))
        observation, reward, done, info  = env.step(action)
```

    [2017-03-29 10:35:19,016] Making new env: CartPole-v0

    0
    1
    2
    3
    4
    5
    6
    7
    8
    9

# Part 2: Policy gradient
**Goal:**
\begin{equation}\text{maximize } E [R | \pi_\theta]
\end{equation}

**Approach:**
1. Collect experience (sample a bunch of trajectories through $(s,a)$ space)
2. Update the policy so that _good_ experiences become more probable

**Difference to DQN: **
 * we don't consider single $(s,a,r,s')$ transitions, but rather use whole episodes for the gradient updates
 * our parameters directly model the policy (output is an action probability), whereas in DQN they model the value function (output is raw score)

#### Rewards:
Remember, we get +1 reward for every time step, in which we still were in the game.

The problem: we normally do not know, which action led to a continuation of the game, and which was actually a bad one. Our simple heuristic: actions in the beginning of the episode are good, and those towards the end are likely bad (they led to losing the game after all).

```python
def discount_rewards(r, gamma=0.999):
    """Take 1D float array of rewards and compute discounted reward """
    discounted_r = np.zeros_like(r)
    running_add = 0
    for t in reversed(range(0, r.size)):
        running_add = running_add * gamma + r[t]
        discounted_r[t] = running_add
    return discounted_r
```

```python
discounted_epr = discount_rewards(np.ones(10))
f, ax = plt.subplots(1, figsize=(5,2))
sns.barplot(list(range(10)), discounted_epr, color="steelblue")
```

We normalize the rewards so that they tank below zero towards the end. gamma controls how late the rewards tank.

```python
discounted_epr_cent = discounted_epr - np.mean(discounted_epr)
discounted_epr_norm = discounted_epr_cent/np.std(discounted_epr_cent)
f, ax = plt.subplots(1, figsize=(5,2))
sns.barplot(list(range(10)), discounted_epr_norm, color="steelblue")
```

```python
discounted_epr = discount_rewards(np.ones(10), gamma=0.5)
discounted_epr_cent = discounted_epr - np.mean(discounted_epr)
discounted_epr_norm = discounted_epr_cent/np.std(discounted_epr_cent)
f, ax = plt.subplots(2, figsize=(5,3))
sns.barplot(list(range(10)), discounted_epr, color="steelblue", ax=ax[0])
sns.barplot(list(range(10)), discounted_epr_norm, color="steelblue", ax=ax[1])
```

### Setting up the model
\begin{equation}
l_1 = relu( x W_1 + b_1) \\
l_2 = l_1 W_2 + b_2 \\
\pi(a|s) = sigmoid(l_2)
\end{equation}

Note: in policy gradient approach, the output of the dense layer is mapped into to a 0-1 range via the sigmoid function.

```python
import cntk as C

TOTAL_EPISODES = 2000 if isFast else 10000

D = 4  # input dimensionality
H = 10 # number of hidden layer neurons

observations = input(STATE_COUNT, np.float32, name="obs")

W1 = C.parameter(shape=(STATE_COUNT, H), init=C.glorot_uniform(), name="W1")
b1 = C.parameter(shape=H, name="b1")
layer1 = C.relu(C.times(observations, W1) + b1)

W2 = C.parameter(shape=(H, ACTION_COUNT), init=C.glorot_uniform(), name="W2")
b2 = C.parameter(shape=ACTION_COUNT, name="b2")
score = C.times(layer1, W2) + b2
# Until here it was similar to DQN

probability = C.sigmoid(score, name="prob")
```

**Policy Search**: The optimal policy search can be carried out with either gradient free approaches or by computing gradients over the policy space ($\pi_\theta$) which is parameterized by $\theta$. In this tutorial, we use the classic forward (`loss.forward`) and back (`loss.backward`) propagation of errors over the parameterized space $\theta$. In this case, $\theta = \{W_1, b_1, W_2, b_2\}$, our model parameters.

```python
input_y = input(1, np.float32, name="input_y")
advantages = input(1, np.float32, name="advt")

loss = -C.reduce_mean(C.log(C.square(input_y - probability) + 1e-4) * advantages, axis=0, name='loss')

lr = 0.001
lr_schedule = learning_rate_schedule(lr, UnitType.sample)
sgd = C.sgd([W1, W2], lr_schedule)

gradBuffer = dict((var.name, np.zeros(shape=var.shape)) for var in loss.parameters if var.name in ['W1', 'W2', 'b1', 'b2'])

xs, hs, label, drs = [], [], [], []
running_reward = None
reward_sum = 0
episode_number = 1

observation = env.reset()

while episode_number <= TOTAL_EPISODES:
    x = np.reshape(observation, [1, STATE_COUNT]).astype(np.float32)

    # Run the policy network and get an action to take.
    prob = probability.eval(arguments={observations: x})[0][0][0]
    action = 1 if np.random.uniform() < prob else 0

    xs.append(x)  # observation
    # grad that encourages the action that was taken to be taken

    y = 1 if action == 0 else 0  # a "fake label"
    label.append(y)

    # step the environment and get new measurements
    observation, reward, done, info = env.step(action)
    reward_sum += float(reward)

    # Record reward (has to be done after we call step() to get reward for previous action)
    drs.append(float(reward))

    if done:
        # Stack together all inputs, hidden states, action gradients, and rewards for this episode
        epx = np.vstack(xs)
        epl = np.vstack(label).astype(np.float32)
        epr = np.vstack(drs).astype(np.float32)
        xs, label, drs = [], [], []  # reset array memory

        # Compute the discounted reward backwards through time.
        discounted_epr = discount_rewards(epr)
        # Size the rewards to be unit normal (helps control the gradient estimator variance)
        discounted_epr -= np.mean(discounted_epr)
        discounted_epr /= np.std(discounted_epr)

        # Forward pass
        arguments = {observations: epx, input_y: epl, advantages: discounted_epr}
        state, outputs_map = loss.forward(arguments, outputs=loss.outputs,
                                          keep_for_backward=loss.outputs)

        # Backward psas
        root_gradients = {v: np.ones_like(o) for v, o in outputs_map.items()}
        vargrads_map = loss.backward(state, root_gradients, variables=set([W1, W2]))

        for var, grad in vargrads_map.items():
            gradBuffer[var.name] += grad

        # Wait for some batches to finish to reduce noise
        if episode_number % BATCH_SIZE_BASELINE == 0:
            grads = {W1: gradBuffer['W1'].astype(np.float32),
                     W2: gradBuffer['W2'].astype(np.float32)}
            updated = sgd.update(grads, BATCH_SIZE_BASELINE)

            # reset the gradBuffer
            gradBuffer = dict((var.name, np.zeros(shape=var.shape))
                              for var in loss.parameters if var.name in ['W1', 'W2', 'b1', 'b2'])

            print('Episode: %d. Average reward for episode %f.' % (episode_number, reward_sum / BATCH_SIZE_BASELINE))

            if reward_sum / BATCH_SIZE_BASELINE > REWARD_TARGET:
                print('Task solved in: %d ' % episode_number)
                break

            reward_sum = 0

        observation = env.reset()  # reset env
        episode_number += 1
probability.save('pg.mod')
```

    Episode: 20. Average reward for episode 16.250000.
    Episode: 40. Average reward for episode 16.200000.
    Episode: 60. Average reward for episode 20.450000.
    Episode: 80. Average reward for episode 17.600000.
    Episode: 100. Average reward for episode 17.200000.
    Episode: 120. Average reward for episode 18.550000.
    Episode: 140. Average reward for episode 19.250000.
    Episode: 160. Average reward for episode 19.550000.
    Episode: 180. Average reward for episode 18.600000.
    Episode: 200. Average reward for episode 21.600000.
    Episode: 220. Average reward for episode 16.750000.
    Episode: 240. Average reward for episode 17.800000.
    Episode: 260. Average reward for episode 18.450000.
    Episode: 280. Average reward for episode 17.800000.
    Episode: 300. Average reward for episode 21.800000.
    Episode: 320. Average reward for episode 20.150000.
    Episode: 340. Average reward for episode 19.450000.
    Episode: 360. Average reward for episode 20.500000.
    Episode: 380. Average reward for episode 20.200000.
    Episode: 400. Average reward for episode 20.800000.
    Episode: 420. Average reward for episode 19.200000.
    Episode: 440. Average reward for episode 21.600000.
    Episode: 460. Average reward for episode 20.150000.
    Episode: 480. Average reward for episode 23.850000.
    Episode: 500. Average reward for episode 19.400000.
    Episode: 520. Average reward for episode 20.250000.
    Episode: 540. Average reward for episode 19.950000.
    Episode: 560. Average reward for episode 22.250000.
    Episode: 580. Average reward for episode 21.800000.
    Episode: 600. Average reward for episode 17.000000.
    Episode: 620. Average reward for episode 17.600000.
    Episode: 640. Average reward for episode 19.700000.
    Episode: 660. Average reward for episode 21.200000.
    Episode: 680. Average reward for episode 23.650000.
    Episode: 700. Average reward for episode 23.700000.
    Episode: 720. Average reward for episode 24.350000.
    Episode: 740. Average reward for episode 19.650000.
    Episode: 760. Average reward for episode 27.050000.
    Episode: 780. Average reward for episode 25.700000.
    Episode: 800. Average reward for episode 22.300000.
    Episode: 820. Average reward for episode 17.800000.
    Episode: 800. Average reward for episode 24.500000.
    Episode: 820. Average reward for episode 28.050000.
    Episode: 880. Average reward for episode 19.800000.
    Episode: 900. Average reward for episode 18.500000.
    Episode: 920. Average reward for episode 21.200000.
    Episode: 840. Average reward for episode 31.550000.
    Episode: 960. Average reward for episode 25.150000.
    Episode: 1000. Average reward for episode 21.000000.
    Episode: 1100. Average reward for episode 25.000000.
    Episode: 1120. Average reward for episode 26.650000.
    Episode: 1140. Average reward for episode 22.450000.
    Episode: 1160. Average reward for episode 22.500000.
    Episode: 1300. Average reward for episode 26.300000.
    Episode: 1320. Average reward for episode 32.000000.
    Task solved in: 840

# Solutions
#### Solution 1.1

```python
observation = input(STATE_COUNT, np.float32, name="s")

W1 = parameter(shape=(STATE_COUNT, H), init=glorot_uniform(), name="W1")
b1 = parameter(shape=H, name="b1")
layer1 = relu(times(observation, W1) + b1)
W2 = parameter(shape=(H, ACTION_COUNT), init=glorot_uniform(), name="W2")
b2 = parameter(shape=ACTION_COUNT, name="b2")
model = times(layer1, W2) + b2
W1.shape, b1.shape, W2.shape, b2.shape, model.shape
```

    ((4, 10), (10,), (10, 2), (2,), (2,))


# Start-to-Finish Example: `GiRaFFE_NRPy` 1D tests

### Authors: Patrick Nelson & Terrence Pierre Jacques

### Adapted from [Start-to-Finish Example: Head-On Black Hole Collision](../Tutorial-Start_to_Finish-BSSNCurvilinear-Two_BHs_Collide.ipynb)

## This module compiles and runs code tests for all 1D initial data options available in GiRaFFE-NRPy+, evolving one-dimensional GRFFE waves.

### NRPy+ Source Code for this module:

* Main python module for all 1D initial data: [GiRaFFEfood_NRPy/GiRaFFEfood_NRPy_1D_tests.py](../../edit/in_progress/GiRaFFEfood_NRPy/GiRaFFEfood_NRPy_1D_tests.py) __Options:__
    1. [Fast Wave](Tutorial-GiRaFFEfood_NRPy_1D_tests-fast_wave.ipynb)
    1. [Alfven Wave](Tutorial-GiRaFFEfood_NRPy_1D_alfven_wave.ipynb)
    1. [Degenerate Alfven Wave](Tutorial-GiRaFFEfood_NRPy_1D_tests-degen_Alfven_wave.ipynb)
    1. [Three Alfven Waves](Tutorial-GiRaFFEfood_NRPy_1D_tests-three_waves.ipynb)
    1. [FFE Breakdown](Tutorial-GiRaFFEfood_NRPy_1D_tests-FFE_breakdown.ipynb)
* [GiRaFFE_NRPy/GiRaFFE_NRPy_staggered_Afield_flux.py](../../edit/in_progress/GiRaFFE_NRPy/GiRaFFE_NRPy_staggered_Afield_flux.py) [\[**tutorial**\]](Tutorial-GiRaFFE_NRPy_staggered-Afield_flux.ipynb) Generates the expressions to find the flux term of the induction equation.
* [GiRaFFE_NRPy/GiRaFFE_NRPy_staggered_A2B.py](../../edit/in_progress/GiRaFFE_NRPy/GiRaFFE_NRPy_staggered_A2B.py) [\[**tutorial**\]](Tutorial-GiRaFFE_NRPy_staggered-A2B.ipynb) Generates the driver to compute the magnetic field from the vector potential/
* [GiRaFFE_NRPy/GiRaFFE_NRPy_BCs.py](../../edit/in_progress/GiRaFFE_NRPy/GiRaFFE_NRPy_BCs.py) [\[**tutorial**\]](Tutorial-GiRaFFE_NRPy-BCs.ipynb) Generates the code to apply boundary conditions to the vector potential, scalar potential, and three-velocity.
* [GiRaFFE_NRPy/GiRaFFE_NRPy_C2P_P2C.py](../../edit/in_progress/GiRaFFE_NRPy/GiRaFFE_NRPy_C2P_P2C.py) [\[**tutorial**\]](Tutorial-GiRaFFE_NRPy-C2P_P2C.ipynb) Generates the conservative-to-primitive and primitive-to-conservative solvers.
* [GiRaFFE_NRPy/GiRaFFE_NRPy_Metric_Face_Values.py](../../edit/in_progress/GiRaFFE_NRPy/GiRaFFE_NRPy_Metric_Face_Values.py) [\[**tutorial**\]](Tutorial-GiRaFFE_NRPy-Metric_Face_Values.ipynb) Generates code to interpolate metric gridfunctions to cell faces.
* [GiRaFFE_NRPy/GiRaFFE_NRPy_PPM.py](../../edit/in_progress/GiRaFFE_NRPy/GiRaFFE_NRPy_PPM.py) [\[**tutorial**\]](Tutorial-GiRaFFE_NRPy-PPM.ipynb) Genearates code to reconstruct primitive variables on cell faces.
* [GiRaFFE_NRPy/GiRaFFE_NRPy_staggered_Source_Terms.py](../../edit/in_progress/GiRaFFE_NRPy/GiRaFFE_NRPy_staggered_Source_Terms.py) [\[**tutorial**\]](Tutorial-GiRaFFE_NRPy_staggered-Source_Terms.ipynb) Generates the expressions to find the flux term of the Poynting flux evolution equation.
* [GiRaFFE_NRPy/Stilde_flux.py](../../edit/in_progress/GiRaFFE_NRPy/Stilde_flux.py) [\[**tutorial**\]](Tutorial-GiRaFFE_NRPy-Stilde_flux.ipynb) Generates the expressions to find the flux term of the Poynting flux evolution equation.
* [../GRFFE/equations.py](../../edit/GRFFE/equations.py) [\[**tutorial**\]](../Tutorial-GRFFE_Equations-Cartesian.ipynb) Generates code necessary to compute the source terms.
* [../GRHD/equations.py](../../edit/GRHD/equations.py) [\[**tutorial**\]](../Tutorial-GRHD_Equations-Cartesian.ipynb) Generates code necessary to compute the source terms.

Here we use NRPy+ to generate the C source code necessary to set up initial data for an Alfvén wave (see [the original GiRaFFE paper](https://arxiv.org/pdf/1704.00599.pdf)). Then we use it to generate the RHS expressions for [Method of Lines](https://reference.wolfram.com/language/tutorial/NDSolveMethodOfLines.html) time integration based on the [explicit Runge-Kutta fourth-order scheme](https://en.wikipedia.org/wiki/Runge%E2%80%93Kutta_methods) (RK4).

# Table of Contents
$$\label{toc}$$

This notebook is organized as follows

1. [Step 1](#initializenrpy): Set core NRPy+ parameters for numerical grids
1. [Step 2](#grffe): Output C code for GRFFE evolution
    1. [Step 2.a](#mol): Output macros for Method of Lines timestepping
1. [Step 3](#gf_id): Import `GiRaFFEfood_NRPy` initial data modules
1. [Step 4](#cparams): Output C codes needed for declaring and setting Cparameters; also set `free_parameters.h`
1. [Step 5](#mainc): `GiRaFFE_NRPy_standalone.c`: The Main C Code
1. [Step 6](#compileexec): Compile and execute C codes
1. [Step 7](#plots): Data Visualization
1. [Step 8](#latex_pdf_output): Output this notebook to $\LaTeX$-formatted PDF file

# Step 1: Set up core functions and parameters for solving  GRFFE equations \[Back to [top](#toc)\]
$$\label{setup}$$

```python
import shutil, os, sys           # Standard Python modules for multiplatform OS-level functions
# First, we'll add the parent directory to the list of directories Python will check for modules.
nrpy_dir_path = os.path.join("..")
if nrpy_dir_path not in sys.path:
    sys.path.append(nrpy_dir_path)

# Step P1: Import needed NRPy+ core modules:
from outputC import outCfunction, lhrh # NRPy+: Core C code output module
import sympy as sp               # SymPy: The Python computer algebra package upon which NRPy+ depends
import finite_difference as fin  # NRPy+: Finite difference C code generation module
import NRPy_param_funcs as par   # NRPy+: Parameter interface
import grid as gri               # NRPy+: Functions having to do with numerical grids
import indexedexp as ixp         # NRPy+: Symbolic indexed expression (e.g., tensors, vectors, etc.) support
import cmdline_helper as cmd     # NRPy+: Multi-platform Python command-line interface

# Step P2: Create C code output directory:
Ccodesdir = os.path.join("GiRaFFE_staggered_1D_Tests_standalone_Ccodes/")
# First remove C code output directory if it exists
# Courtesy https://stackoverflow.com/questions/303200/how-do-i-remove-delete-a-folder-that-is-not-empty
# !rm -r ScalarWaveCurvilinear_Playground_Ccodes
shutil.rmtree(Ccodesdir, ignore_errors=True)
# Then create a fresh directory
cmd.mkdir(Ccodesdir)

# Step P3: Create executable output directory:
outdir = os.path.join(Ccodesdir,"output/")
cmd.mkdir(outdir)

# Step P5: Set timestepping algorithm (we adopt the Method of Lines)
REAL      = "double"    # Best to use double here.
default_CFL_FACTOR= 0.5 # (GETS OVERWRITTEN WHEN EXECUTED.) In pure axisymmetry (symmetry_axes = 2 below) 1.0 works fine. Otherwise 0.5 or lower.

# Step P6: Set the finite differencing order to 2.
par.set_parval_from_str("finite_difference::FD_CENTDERIVS_ORDER",4)

thismodule = "Start_to_Finish-GiRaFFE_NRPy-1D_tests"
TINYDOUBLE = par.Cparameters("REAL", thismodule, "TINYDOUBLE", 1e-100)

import GiRaFFE_NRPy.GiRaFFE_NRPy_Main_Driver_staggered as md
# par.set_paramsvals_value("GiRaFFE_NRPy.GiRaFFE_NRPy_C2P_P2C::enforce_speed_limit_StildeD = False")
par.set_paramsvals_value("GiRaFFE_NRPy.GiRaFFE_NRPy_C2P_P2C::enforce_current_sheet_prescription = False")
```

# Step 2: Output C code for GRFFE evolution \[Back to [top](#toc)\]
$$\label{grffe}$$

We will first write the C codes needed for GRFFE evolution. We have already written a module to generate all these codes and call the functions in the appropriate order, so we will import that here. We will take the slightly unusual step of doing this before we generate the initial data functions because the main driver module will register all the gridfunctions we need. It will also generate functions that, in addition to their normal spot in the MoL timestepping, will need to be called during the initial data step to make sure all the variables are appropriately filled in.

All of this is handled with a single call to `GiRaFFE_NRPy_Main_Driver_generate_all()`, which will register gridfunctions, write all the C code kernels, and write the C code functions to call those.

```python
md.GiRaFFE_NRPy_Main_Driver_generate_all(Ccodesdir)
```

    Output C function calculate_StildeD0_source_term() to file GiRaFFE_staggered_1D_Tests_standalone_Ccodes/RHSs/calculate_StildeD0_source_term.h
    Output C function calculate_StildeD1_source_term() to file GiRaFFE_staggered_1D_Tests_standalone_Ccodes/RHSs/calculate_StildeD1_source_term.h
    Output C function calculate_StildeD2_source_term() to file GiRaFFE_staggered_1D_Tests_standalone_Ccodes/RHSs/calculate_StildeD2_source_term.h
    Output C function calculate_Stilde_rhsD() to file GiRaFFE_staggered_1D_Tests_standalone_Ccodes/RHSs/calculate_Stilde_rhsD.h
    Output C function GiRaFFE_NRPy_cons_to_prims() to file GiRaFFE_staggered_1D_Tests_standalone_Ccodes/C2P/GiRaFFE_NRPy_cons_to_prims.h
    Output C function GiRaFFE_NRPy_prims_to_cons() to file GiRaFFE_staggered_1D_Tests_standalone_Ccodes/C2P/GiRaFFE_NRPy_prims_to_cons.h

## Step 2.a: Output macros for Method of Lines timestepping \[Back to [top](#toc)\]
$$\label{mol}$$

Now, we generate the code to implement the method of lines using the fourth-order Runge-Kutta algorithm.

```python
RK_method = "RK4"

# Step 3: Generate Runge-Kutta-based (RK-based) timestepping code.
#       As described above the Table of Contents, this is a 3-step process:
#       3.A: Evaluate RHSs (RHS_string)
#       3.B: Apply boundary conditions (post_RHS_string, pt 1)
import MoLtimestepping.C_Code_Generation as MoL
from MoLtimestepping.RK_Butcher_Table_Dictionary import Butcher_dict
RK_order  = Butcher_dict[RK_method][1]
cmd.mkdir(os.path.join(Ccodesdir,"MoLtimestepping/"))
MoL.MoL_C_Code_Generation(RK_method,
    RHS_string      = """
GiRaFFE_NRPy_RHSs(&params,auxevol_gfs,RK_INPUT_GFS,RK_OUTPUT_GFS);""",
    post_RHS_string = """
GiRaFFE_NRPy_post_step(&params,xx,auxevol_gfs,RK_OUTPUT_GFS,n+1);\n""",
    outdir = os.path.join(Ccodesdir,"MoLtimestepping/"))
```

# Step 3: Import `GiRaFFEfood_NRPy` initial data modules \[Back to [top](#toc)\]
$$\label{gf_id}$$

With the preliminaries out of the way, we will write the C functions to set up initial data. There are two categories of initial data that must be set: the spacetime metric variables, and the GRFFE plasma variables. We will set up the spacetime first, namely the Minkowski spacetime.

```python
gammaDD = ixp.zerorank2(DIM=3)
for i in range(3):
    for j in range(3):
        if i==j:
            gammaDD[i][j] = sp.sympify(1) # else: leave as zero
betaU = ixp.zerorank1() # All should be 0
alpha = sp.sympify(1)

# Description and options for this initial data
desc = "Generate a flat spacetime metric."
loopopts_id ="AllPoints" # we don't need to read coordinates for flat spacetime.

# For testing: Also set inverse metric:
gammaUU, unused_gammaDET = ixp.symm_matrix_inverter3x3(gammaDD)

name = "set_initial_spacetime_metric_data"
values_to_print = [
                   lhrh(lhs=gri.gfaccess("auxevol_gfs","gammaDD00"),rhs=gammaDD[0][0]),
                   lhrh(lhs=gri.gfaccess("auxevol_gfs","gammaDD01"),rhs=gammaDD[0][1]),
                   lhrh(lhs=gri.gfaccess("auxevol_gfs","gammaDD02"),rhs=gammaDD[0][2]),
                   lhrh(lhs=gri.gfaccess("auxevol_gfs","gammaDD11"),rhs=gammaDD[1][1]),
                   lhrh(lhs=gri.gfaccess("auxevol_gfs","gammaDD12"),rhs=gammaDD[1][2]),
                   lhrh(lhs=gri.gfaccess("auxevol_gfs","gammaDD22"),rhs=gammaDD[2][2]),
                   lhrh(lhs=gri.gfaccess("auxevol_gfs","betaU0"),rhs=betaU[0]),
                   lhrh(lhs=gri.gfaccess("auxevol_gfs","betaU1"),rhs=betaU[1]),
                   lhrh(lhs=gri.gfaccess("auxevol_gfs","betaU2"),rhs=betaU[2]),
                   lhrh(lhs=gri.gfaccess("auxevol_gfs","alpha"),rhs=alpha)
                  ]

outCfunction(
    outfile  = os.path.join(Ccodesdir,name+".h"), desc=desc, name=name,
    params   ="const paramstruct *params,REAL *xx[3],REAL *auxevol_gfs",
    body     = fin.FD_outputC("returnstring",values_to_print,params="outCverbose=False"),
    loopopts = loopopts_id)

```

    Output C function set_initial_spacetime_metric_data() to file GiRaFFE_staggered_1D_Tests_standalone_Ccodes/set_initial_spacetime_metric_data.h

Now, we will write out the initials data function for the GRFFE variables.

```python
initial_data_dir = os.path.join(Ccodesdir,"InitialData/")
cmd.mkdir(initial_data_dir)

ID_opts = ["AlfvenWave", "ThreeAlfvenWaves", "DegenAlfvenWave", "FastWave", "FFEBD"]

import GiRaFFEfood_NRPy.GiRaFFEfood_NRPy as gid
for initial_data in ID_opts:
    if initial_data=="AlfvenWave":
        gid.GiRaFFEfood_NRPy_generate_initial_data(ID_type = initial_data, stagger_enable = True)
        desc = "Generate Alfven wave 1D initial data for GiRaFFEfood_NRPy."
    elif initial_data=="ThreeAlfvenWaves":
        gid.GiRaFFEfood_NRPy_generate_initial_data(ID_type = "ThreeWaves", stagger_enable = True)
        desc = "Generate three Alfven wave 1D initial data for GiRaFFEfood_NRPy."
    elif initial_data=="DegenAlfvenWave":
        gid.GiRaFFEfood_NRPy_generate_initial_data(ID_type = initial_data, stagger_enable = True)
        desc = "Generate degenerate Alfven wave 1D initial data for GiRaFFEfood_NRPy."
    elif initial_data=="FastWave":
        gid.GiRaFFEfood_NRPy_generate_initial_data(ID_type = initial_data, stagger_enable = True)
        desc = "Generate fast wave 1D initial data for GiRaFFEfood_NRPy."
    elif initial_data=="FFEBD":
        gid.GiRaFFEfood_NRPy_generate_initial_data(ID_type = "FFE_Breakdown", stagger_enable = True)
        desc = "Generate FFE breakdown 1D initial data for GiRaFFEfood_NRPy."

    name = initial_data

    values_to_print = [
                       lhrh(lhs=gri.gfaccess("out_gfs","AD0"),rhs=gid.AD[0]),
                       lhrh(lhs=gri.gfaccess("out_gfs","AD1"),rhs=gid.AD[1]),
                       lhrh(lhs=gri.gfaccess("out_gfs","AD2"),rhs=gid.AD[2]),
                       lhrh(lhs=gri.gfaccess("auxevol_gfs","ValenciavU0"),rhs=gid.ValenciavU[0]),
                       lhrh(lhs=gri.gfaccess("auxevol_gfs","ValenciavU1"),rhs=gid.ValenciavU[1]),
                       lhrh(lhs=gri.gfaccess("auxevol_gfs","ValenciavU2"),rhs=gid.ValenciavU[2]),
#                        lhrh(lhs=gri.gfaccess("auxevol_gfs","BU0"),rhs=gid.BU[0]),
#                        lhrh(lhs=gri.gfaccess("auxevol_gfs","BU1"),rhs=gid.BU[1]),
#                        lhrh(lhs=gri.gfaccess("auxevol_gfs","BU2"),rhs=gid.BU[2]),
                       lhrh(lhs=gri.gfaccess("out_gfs","psi6Phi"),rhs=sp.sympify(0))
                      ]

    outCfunction(
        outfile  = os.path.join(initial_data_dir,name+".c"), desc=desc, name=name,
        params   ="const paramstruct *params, REAL *xx[3], REAL *auxevol_gfs, REAL *out_gfs",
        body     = fin.FD_outputC("returnstring",values_to_print,params="outCverbose=False"),
        rel_path_to_Cparams='../',
        loopopts ="AllPoints,Read_xxs")

inital_data_body = """

const char *option1 = "AlfvenWave";
const char *option2 = "ThreeAlfvenWaves";
const char *option3 = "DegenAlfvenWave";
const char *option4 = "FastWave";
const char *option5 = "FFEBD";

if (strcmp(initial_data_option, option1) == 0) {
    AlfvenWave(params, xx, auxevol_gfs, out_gfs);
}

else if (strcmp(initial_data_option, option2) == 0) {
    ThreeAlfvenWaves(params, xx, auxevol_gfs, out_gfs);
}

else if (strcmp(initial_data_option, option3) == 0) {
    DegenAlfvenWave(params, xx, auxevol_gfs, out_gfs);
}

else if (strcmp(initial_data_option, option4) == 0) {
    FastWave(params, xx, auxevol_gfs, out_gfs);
}

else if (strcmp(initial_data_option, option5) == 0) {
    FFEBD(params, xx, auxevol_gfs, out_gfs);
}

else {
    printf("ERROR: Invalid choice of initial data.");
    exit(1);
}

"""
name = "initial_data"
desc = "Main initial data function."
includes = ["AlfvenWave.c", "ThreeAlfvenWaves.c", "DegenAlfvenWave.c", "FastWave.c", "FFEBD.c"]
outCfunction(
    outfile  = os.path.join(initial_data_dir,name+".h"), desc=desc, name=name,
    params   ="const char *initial_data_option, const paramstruct *restrict params,REAL *xx[3],REAL *restrict auxevol_gfs,REAL *restrict out_gfs",
    body     = inital_data_body,
    includes = includes,
    prefunc="#include <string.h>",
    rel_path_to_Cparams='../',
    loopopts ="")

```

    Output C function AlfvenWave() to file GiRaFFE_staggered_1D_Tests_standalone_Ccodes/InitialData/AlfvenWave.c
    Output C function ThreeAlfvenWaves() to file GiRaFFE_staggered_1D_Tests_standalone_Ccodes/InitialData/ThreeAlfvenWaves.c
    Output C function DegenAlfvenWave() to file GiRaFFE_staggered_1D_Tests_standalone_Ccodes/InitialData/DegenAlfvenWave.c
    Output C function FastWave() to file GiRaFFE_staggered_1D_Tests_standalone_Ccodes/InitialData/FastWave.c
    Output C function FFEBD() to file GiRaFFE_staggered_1D_Tests_standalone_Ccodes/InitialData/FFEBD.c
    Output C function initial_data() to file GiRaFFE_staggered_1D_Tests_standalone_Ccodes/InitialData/initial_data.h

# Step 4: Output C codes needed for declaring and setting Cparameters; also set `free_parameters.h` \[Back to [top](#toc)\]
$$\label{cparams}$$

Based on declared NRPy+ Cparameters, first we generate `declare_Cparameters_struct.h`, `set_Cparameters_default.h`, and `set_Cparameters[-SIMD].h`.

Then we output `free_parameters.h`, which sets initial data parameters, as well as grid domain & reference metric parameters, applying `domain_size` and `sinh_width`/`SymTP_bScale` (if applicable) as set above

```python
# Step 3.e: Output C codes needed for declaring and setting Cparameters; also set free_parameters.h
# Step 3.e.i: Generate declare_Cparameters_struct.h, set_Cparameters_default.h, and set_Cparameters[-SIMD].h
par.generate_Cparameters_Ccodes(os.path.join(Ccodesdir))

# Step 3.e.ii: Set free_parameters.h
with open(os.path.join(Ccodesdir,"free_parameters.h"),"w") as file:
    file.write("""// Override parameter defaults with values based on command line arguments and NGHOSTS.
params.Nxx0 = atoi(argv[1]);
params.Nxx1 = atoi(argv[2]);
params.Nxx2 = atoi(argv[3]);
params.Nxx_plus_2NGHOSTS0 = params.Nxx0 + 2*NGHOSTS;
params.Nxx_plus_2NGHOSTS1 = params.Nxx1 + 2*NGHOSTS;
params.Nxx_plus_2NGHOSTS2 = params.Nxx2 + 2*NGHOSTS;
// Step 0d: Set up space and time coordinates
// Step 0d.i: Declare \Delta x^i=dxx{0,1,2} and invdxx{0,1,2}, as well as xxmin[3] and xxmax[3]:
const REAL xxmin[3] = {-1.3255,-0.085,-0.085};
const REAL xxmax[3] = { 1.6745, 0.115, 0.115};

params.dxx0 = (xxmax[0] - xxmin[0]) / ((REAL)params.Nxx0+1);
params.dxx1 = (xxmax[1] - xxmin[1]) / ((REAL)params.Nxx1+1);
params.dxx2 = (xxmax[2] - xxmin[2]) / ((REAL)params.Nxx2+1);
printf("dxx0,dxx1,dxx2 = %.5e,%.5e,%.5e\\n",params.dxx0,params.dxx1,params.dxx2);
params.invdx0 = 1.0 / params.dxx0;
params.invdx1 = 1.0 / params.dxx1;
params.invdx2 = 1.0 / params.dxx2;

const int poison_grids = 0;
// Standard GRFFE parameters:
params.GAMMA_SPEED_LIMIT = 2000.0;
params.diss_strength = 0.1;
""")

```

# Step 4: Set up boundary condition functions for chosen singular, curvilinear coordinate system \[Back to [top](#toc)\]
$$\label{bc_functs}$$

Next apply singular, curvilinear coordinate boundary conditions [as documented in the corresponding NRPy+ tutorial notebook](Tutorial-Start_to_Finish-Curvilinear_BCs.ipynb)

...But, for the moment, we're actually just using this because it writes the file `gridfunction_defines.h`.

```python
import CurviBoundaryConditions.CurviBoundaryConditions as cbcs
cbcs.Set_up_CurviBoundaryConditions(os.path.join(Ccodesdir,"boundary_conditions/"),Cparamspath=os.path.join("../"),enable_copy_of_static_Ccodes=False)
```

    Wrote to file "GiRaFFE_staggered_1D_Tests_standalone_Ccodes/boundary_conditions/parity_conditions_symbolic_dot_products.h"
    Evolved parity: ( AD0:1, AD1:2, AD2:3, StildeD0:1, StildeD1:2, StildeD2:3,
        psi6Phi:0 )

    AuxEvol parity: ( BU0:1, BU1:2, BU2:3, B_lU0:1, B_lU1:2, B_lU2:3, B_rU0:1,
        B_rU1:2, B_rU2:3, BstaggerU0:1, BstaggerU1:2, BstaggerU2:3,
        Bstagger_lU0:1, Bstagger_lU1:2, Bstagger_lU2:3, Bstagger_rU0:1,
        Bstagger_rU1:2, Bstagger_rU2:3, Stilde_flux_HLLED0:1,
        Stilde_flux_HLLED1:2, Stilde_flux_HLLED2:3, ValenciavU0:1,
        ValenciavU1:2, ValenciavU2:3, Valenciav_lU0:1, Valenciav_lU1:2,
        Valenciav_lU2:3, Valenciav_llU0:1, Valenciav_llU1:2, Valenciav_llU2:3,
        Valenciav_lrU0:1, Valenciav_lrU1:2, Valenciav_lrU2:3, Valenciav_rU0:1,
        Valenciav_rU1:2, Valenciav_rU2:3, Valenciav_rlU0:1, Valenciav_rlU1:2,
        Valenciav_rlU2:3, Valenciav_rrU0:1, Valenciav_rrU1:2, Valenciav_rrU2:3,
        alpha:0, alpha_face:0, betaU0:1, betaU1:2, betaU2:3, beta_faceU0:1,
        beta_faceU1:2, beta_faceU2:3, cmax_x:0, cmax_y:0, cmax_z:0, cmin_x:0,
        cmin_y:0, cmin_z:0, gammaDD00:4, gammaDD01:5, gammaDD02:6, gammaDD11:7,
        gammaDD12:8, gammaDD22:9, gamma_faceDD00:4, gamma_faceDD01:5,
        gamma_faceDD02:6, gamma_faceDD11:7, gamma_faceDD12:8, gamma_faceDD22:9,
        psi6_temp:0, psi6center:0 )
    Wrote to file "GiRaFFE_staggered_1D_Tests_standalone_Ccodes/boundary_conditions/EigenCoord_Cart_to_xx.h"

# Step 5: `GiRaFFE_NRPy_standalone.c`: The Main C Code \[Back to [top](#toc)\]
$$\label{mainc}$$

```python
# Part P0: Define REAL, set the number of ghost cells NGHOSTS (from NRPy+'s FD_CENTDERIVS_ORDER),
#          and set the CFL_FACTOR (which can be overwritten at the command line)

with open(os.path.join(Ccodesdir,"GiRaFFE_NRPy_REAL__NGHOSTS__CFL_FACTOR.h"), "w") as file:
    file.write("""
// Part P0.a: Set the number of ghost cells, from NRPy+'s FD_CENTDERIVS_ORDER
#define NGHOSTS """+str(3)+"""
#define NGHOSTS_A2B """+str(2)+"""
// Part P0.b: Set the numerical precision (REAL) to double, ensuring all floating point
//            numbers are stored to at least ~16 significant digits
#define REAL """+REAL+"""
// Part P0.c: Set the CFL Factor. Can be overwritten at command line.
REAL CFL_FACTOR = """+str(default_CFL_FACTOR)+";")
```

```python
%%writefile $Ccodesdir/GiRaFFE_NRPy_standalone.c
// Step P0: Define REAL and NGHOSTS; and declare CFL_FACTOR. This header is generated in NRPy+.
#include "GiRaFFE_NRPy_REAL__NGHOSTS__CFL_FACTOR.h"

#include "declare_Cparameters_struct.h"

const int NSKIP_1D_OUTPUT = 1;

// Step P1: Import needed header files
#include "stdio.h"
#include "stdlib.h"
#include "math.h"
#include "time.h"
#include "stdint.h" // Needed for Windows GCC 6.x compatibility
#ifndef M_PI
#define M_PI 3.141592653589793238462643383279502884L
#endif
#ifndef M_SQRT1_2
#define M_SQRT1_2 0.707106781186547524400844362104849039L
#endif

// Step P2: Declare the IDX4S(gf,i,j,k) macro, which enables us to store 4-dimensions of
//           data in a 1D array. In this case, consecutive values of "i"
//           (all other indices held to a fixed value) are consecutive in memory, where
//           consecutive values of "j" (fixing all other indices) are separated by
//           Nxx_plus_2NGHOSTS0 elements in memory. Similarly, consecutive values of
//           "k" are separated by Nxx_plus_2NGHOSTS0*Nxx_plus_2NGHOSTS1 in memory, etc.
#define IDX4S(g,i,j,k) \
( (i) + Nxx_plus_2NGHOSTS0 * ( (j) + Nxx_plus_2NGHOSTS1 * ( (k) + Nxx_plus_2NGHOSTS2 * (g) ) ) )
#define IDX4ptS(g,idx) ( (idx) + (Nxx_plus_2NGHOSTS0*Nxx_plus_2NGHOSTS1*Nxx_plus_2NGHOSTS2) * (g) )
#define IDX3S(i,j,k) ( (i) + Nxx_plus_2NGHOSTS0 * ( (j) + Nxx_plus_2NGHOSTS1 * ( (k) ) ) )
#define LOOP_REGION(i0min,i0max, i1min,i1max, i2min,i2max) \
  for(int i2=i2min;i2<i2max;i2++) for(int i1=i1min;i1<i1max;i1++) for(int i0=i0min;i0<i0max;i0++)
#define LOOP_ALL_GFS_GPS(ii) _Pragma("omp parallel for") \
  for(int (ii)=0;(ii)<Nxx_plus_2NGHOSTS_tot*NUM_EVOL_GFS;(ii)++)

// Step P3: Set gridfunction macros
#include "boundary_conditions/gridfunction_defines.h"

// Step P4: Include the RHS, BC, and primitive recovery functions
#include "GiRaFFE_NRPy_Main_Driver.h"

// Step P5: Include the initial data functions
#include "set_initial_spacetime_metric_data.h"
#include "InitialData/initial_data.h"
// main() function:
// Step 0: Read command-line input, set up grid structure, allocate memory for gridfunctions, set up coordinates
// Step 1: Set up scalar wave initial data
// Step 2: Evolve scalar wave initial data forward in time using Method of Lines with RK4 algorithm,
//         applying quadratic extrapolation outer boundary conditions.
// Step 3: Output relative error between numerical and exact solution.
// Step 4: Free all allocated memory
int main(int argc, const char *argv[]) {
    paramstruct params;
#include "set_Cparameters_default.h"
    // Step 0a: Read command-line input, error out if nonconformant
    if(argc != 5 || atoi(argv[1]) < NGHOSTS || atoi(argv[2]) < NGHOSTS || atoi(argv[3]) < NGHOSTS) {
      printf("Error: Expected three command-line arguments: ./GiRaFFE_NRPy_standalone [Nx] [Ny] [Nz],\n");
      printf("where Nx is the number of grid points in the x direction, and so forth.\n");
      printf("Nx,Ny,Nz MUST BE larger than NGHOSTS (= %d)\n",NGHOSTS);
      exit(1);
    }
    // Step 0c: Set free parameters, overwriting Cparameters defaults
    //          by hand or with command-line input, as desired.
#include "free_parameters.h"
#include "set_Cparameters-nopointer.h"

    //          ... and then set up the numerical grid structure in time:
    const REAL t_final = 2.0;
    const REAL CFL_FACTOR = 0.5; // Set the CFL Factor

    // Step 0c: Allocate memory for gridfunctions
    const int Nxx_plus_2NGHOSTS_tot = Nxx_plus_2NGHOSTS0*Nxx_plus_2NGHOSTS1*Nxx_plus_2NGHOSTS2;
    // Step 0k: Allocate memory for gridfunctions
#include "MoLtimestepping/RK_Allocate_Memory.h"
    REAL *restrict auxevol_gfs = (REAL *)malloc(sizeof(REAL) * NUM_AUXEVOL_GFS * Nxx_plus_2NGHOSTS_tot);
    REAL *evol_gfs_exact = (REAL *)malloc(sizeof(REAL) * NUM_EVOL_GFS * Nxx_plus_2NGHOSTS_tot);
    REAL *auxevol_gfs_exact  = (REAL *)malloc(sizeof(REAL) * NUM_AUXEVOL_GFS * Nxx_plus_2NGHOSTS_tot);

    // For debugging, it can be useful to set everything to NaN initially.
    if(poison_grids) {
        for(int ii=0;ii<NUM_EVOL_GFS * Nxx_plus_2NGHOSTS_tot;ii++) {
            y_n_gfs[ii] = 1.0/0.0;
            y_nplus1_running_total_gfs[ii] = 1.0/0.0;
            //k_odd_gfs[ii] = 1.0/0.0;
            //k_even_gfs[ii] = 1.0/0.0;
            diagnostic_output_gfs[ii] = 1.0/0.0;
            evol_gfs_exact[ii] = 1.0/0.0;
        }
        for(int ii=0;ii<NUM_AUXEVOL_GFS * Nxx_plus_2NGHOSTS_tot;ii++) {
            auxevol_gfs[ii] = 1.0/0.0;
            auxevol_gfs_exact[ii] = 1.0/0.0;
        }
    }

    // Step 0d: Set up coordinates: Set dx, and then dt based on dx_min and CFL condition
    // This is probably already defined above, but just in case...
    #ifndef MIN
    #define MIN(A, B) ( ((A) < (B)) ? (A) : (B) )
    #endif
    REAL dt = CFL_FACTOR * MIN(dxx0,MIN(dxx1,dxx2)); // CFL condition
    int Nt = (int)(t_final / dt + 0.5); // The number of points in time.
                                        //Add 0.5 to account for C rounding down integers.

    // Step 0e: Set up cell-centered Cartesian coordinate grids
    REAL *xx[3];
    xx[0] = (REAL *)malloc(sizeof(REAL)*Nxx_plus_2NGHOSTS0);
    xx[1] = (REAL *)malloc(sizeof(REAL)*Nxx_plus_2NGHOSTS1);
    xx[2] = (REAL *)malloc(sizeof(REAL)*Nxx_plus_2NGHOSTS2);
    for(int j=0;j<Nxx_plus_2NGHOSTS0;j++) xx[0][j] = xxmin[0] + (j-NGHOSTS+1)*dxx0;
    for(int j=0;j<Nxx_plus_2NGHOSTS1;j++) xx[1][j] = xxmin[1] + (j-NGHOSTS+1)*dxx1;
    for(int j=0;j<Nxx_plus_2NGHOSTS2;j++) xx[2][j] = xxmin[2] + (j-NGHOSTS+1)*dxx2;

    // Step 1: Set up initial data to be exact solution at time=0:
    REAL time = 0.0;
    set_initial_spacetime_metric_data(&params, xx, auxevol_gfs);
    const char *initial_data_option = argv[4];
    initial_data(initial_data_option, &params, xx, auxevol_gfs, y_n_gfs);

    // Fill in  the remaining quantities
    GiRaFFE_compute_B_and_Bstagger_from_A(&params,
                                          auxevol_gfs+Nxx_plus_2NGHOSTS_tot*GAMMADD00GF,
                                          auxevol_gfs+Nxx_plus_2NGHOSTS_tot*GAMMADD01GF,
                                          auxevol_gfs+Nxx_plus_2NGHOSTS_tot*GAMMADD02GF,
                                          auxevol_gfs+Nxx_plus_2NGHOSTS_tot*GAMMADD11GF,
                                          auxevol_gfs+Nxx_plus_2NGHOSTS_tot*GAMMADD12GF,
                                          auxevol_gfs+Nxx_plus_2NGHOSTS_tot*GAMMADD22GF,
                                          auxevol_gfs+Nxx_plus_2NGHOSTS_tot*PSI6_TEMPGF, /* Temporary storage,overwritten */
                                          y_n_gfs+Nxx_plus_2NGHOSTS_tot*AD0GF,
                                          y_n_gfs+Nxx_plus_2NGHOSTS_tot*AD1GF,
                                          y_n_gfs+Nxx_plus_2NGHOSTS_tot*AD2GF,
                                          auxevol_gfs+Nxx_plus_2NGHOSTS_tot*BU0GF,
                                          auxevol_gfs+Nxx_plus_2NGHOSTS_tot*BU1GF,
                                          auxevol_gfs+Nxx_plus_2NGHOSTS_tot*BU2GF,
                                          auxevol_gfs+Nxx_plus_2NGHOSTS_tot*BSTAGGERU0GF,
                                          auxevol_gfs+Nxx_plus_2NGHOSTS_tot*BSTAGGERU1GF,
                                          auxevol_gfs+Nxx_plus_2NGHOSTS_tot*BSTAGGERU2GF);
    //override_BU_with_old_GiRaFFE(&params,auxevol_gfs,0);
    GiRaFFE_NRPy_prims_to_cons(&params,auxevol_gfs,y_n_gfs);
    // Extra stack, useful for debugging:
    GiRaFFE_NRPy_cons_to_prims(&params,xx,auxevol_gfs,y_n_gfs);

    for(int n=0;n<=Nt;n++) { // Main loop to progress forward in time.
    //for(int n=0;n<=1;n++) { // Main loop to progress forward in time.
        // Step 1a: Set current time to correct value & compute exact solution
        time = ((REAL)n)*dt;

        /* Step 2: Validation: Output relative error between numerical and exact solution, */
        if(time == 0.0 || time == 0.5 || time == 1.0 || time == 2.0 || time == 0.02 || time == 0.56) {
            // Step 2c: Output relative error between exact & numerical at center of grid.
            const int i0mid=Nxx_plus_2NGHOSTS0/2;
            const int i1mid=Nxx_plus_2NGHOSTS1/2;
            const int i2mid=Nxx_plus_2NGHOSTS2/2;
            char filename[100];
            sprintf(filename,"out%d__%s-%08d.txt", Nxx0, initial_data_option, n);
            FILE *out2D = fopen(filename, "w");
            for(int i0=0;i0<Nxx_plus_2NGHOSTS0;i0++) {
                const int idx = IDX3S(i0,i1mid,i2mid);
                fprintf(out2D,"%.16e %.16e %.16e %.16e %.16e %.16e %.16e %.16e %.16e %.16e %.16e %.16e %.16e %.16e %.16e\n",
                        xx[0][i0],
                        auxevol_gfs[IDX4ptS(BU0GF,idx)],auxevol_gfs[IDX4ptS(BU1GF,idx)],auxevol_gfs[IDX4ptS(BU2GF,idx)],
                        y_n_gfs[IDX4ptS(AD0GF,idx)],y_n_gfs[IDX4ptS(AD1GF,idx)],y_n_gfs[IDX4ptS(AD2GF,idx)],
                        y_n_gfs[IDX4ptS(STILDED0GF,idx)],y_n_gfs[IDX4ptS(STILDED1GF,idx)],y_n_gfs[IDX4ptS(STILDED2GF,idx)],
                        auxevol_gfs[IDX4ptS(VALENCIAVU0GF,idx)],auxevol_gfs[IDX4ptS(VALENCIAVU1GF,idx)],auxevol_gfs[IDX4ptS(VALENCIAVU2GF,idx)],
                        y_n_gfs[IDX4ptS(PSI6PHIGF,idx)], time);
            }
            fclose(out2D);
            // For convergence testing, we'll shift the grid x -> x-1 and output initial data again, giving the exact solution.
            LOOP_REGION(0,Nxx_plus_2NGHOSTS0,0,1,0,1) {
                xx[0][i0] += -mu_AW*time;
                //xx[0][i0] += -time;
            }
            set_initial_spacetime_metric_data(&params,xx,auxevol_gfs_exact);
            initial_data(initial_data_option, &params,xx,auxevol_gfs_exact,evol_gfs_exact);
            // Fill in  the remaining quantities
            //driver_A_to_B(&params,evol_gfs_exact,auxevol_gfs_exact);
            GiRaFFE_NRPy_prims_to_cons(&params,auxevol_gfs_exact,evol_gfs_exact);
            // And now, we'll set the grid back to rights.
            LOOP_REGION(0,Nxx_plus_2NGHOSTS0,0,1,0,1) {
                xx[0][i0] -= -mu_AW*time;
                //xx[0][i0] -= -time;
            }
            sprintf(filename,"out%d-%08d_exact.txt",Nxx0,n);
            FILE *out2D_exact = fopen(filename, "w");
            for(int i0=0;i0<Nxx_plus_2NGHOSTS0;i0++) {
                const int idx = IDX3S(i0,i1mid,i2mid);
                fprintf(out2D_exact,"%.16e %.16e %.16e %.16e %.16e %.16e %.16e %.16e %.16e %.16e %.16e %.16e %.16e %.16e\n",
                        xx[0][i0],
                        auxevol_gfs_exact[IDX4ptS(BU0GF,idx)],auxevol_gfs_exact[IDX4ptS(BU1GF,idx)],auxevol_gfs_exact[IDX4ptS(BU2GF,idx)],
                        evol_gfs_exact[IDX4ptS(AD0GF,idx)],evol_gfs_exact[IDX4ptS(AD1GF,idx)],evol_gfs_exact[IDX4ptS(AD2GF,idx)],
                        evol_gfs_exact[IDX4ptS(STILDED0GF,idx)],evol_gfs_exact[IDX4ptS(STILDED1GF,idx)],evol_gfs_exact[IDX4ptS(STILDED2GF,idx)],
                        auxevol_gfs_exact[IDX4ptS(VALENCIAVU0GF,idx)],auxevol_gfs_exact[IDX4ptS(VALENCIAVU1GF,idx)],auxevol_gfs_exact[IDX4ptS(VALENCIAVU2GF,idx)],
                        evol_gfs_exact[IDX4ptS(PSI6PHIGF,idx)]);
            }
            fclose(out2D_exact);

        }

        // Step 3: Evolve scalar wave initial data forward in time using Method of Lines with RK4 algorithm,
        //         applying quadratic extrapolation outer boundary conditions.
        // Step 3.b: Step forward one timestep (t -> t+dt) in time using
        //           chosen RK-like MoL timestepping algorithm
#include "MoLtimestepping/RK_MoL.h"
    } // End main loop to progress forward in time.

    // Step 4: Free all allocated memory
#include "MoLtimestepping/RK_Free_Memory.h"
    free(auxevol_gfs);
    free(auxevol_gfs_exact);
    free(evol_gfs_exact);
    for(int i=0;i<3;i++) free(xx[i]);
    return 0;
}
```

    Writing GiRaFFE_staggered_1D_Tests_standalone_Ccodes//GiRaFFE_NRPy_standalone.c

# Step 6: Compile generated C codes & perform GRFFE simulations \[Back to [top](#toc)\]
$$\label{compileexec}$$

To aid in the cross-platform-compatible (with Windows, MacOS, & Linux) compilation and execution, we make use of `cmdline_helper` [(**Tutorial**)](Tutorial-cmdline_helper.ipynb).

```python
cmd.C_compile(os.path.join(Ccodesdir,"GiRaFFE_NRPy_standalone.c"),
              os.path.join(Ccodesdir,"output","GiRaFFE_NRPy_standalone"),compile_mode="optimized")

# Change to output directory
os.chdir(outdir)
# Clean up existing output files
cmd.delete_existing_files("out*.txt")
cmd.delete_existing_files("out*.png")

# ID options are: "AlfvenWave", "ThreeAlfvenWaves", "DegenAlfvenWave", "FastWave", "FFEBD"
for opt in ID_opts:
    cmd.Execute("GiRaFFE_NRPy_standalone", "299 4 4 "+opt, "out_298"+opt+".txt")
#     cmd.Execute("GiRaFFE_NRPy_standalone", "1280 9 9 "+opt, "out_1280"+opt+".txt")
#     cmd.Execute("GiRaFFE_NRPy_standalone", "1280 32 32 "+opt, "out_"+opt+".txt")

# cmd.Execute("GiRaFFE_NRPy_standalone", "149 9 9 AlfvenWave","out149.txt")
# Return to root directory
os.chdir(os.path.join("../../"))
```

    Compiling executable...
    (EXEC): Executing `gcc -std=gnu99 -Ofast -fopenmp -march=native -funroll-loops GiRaFFE_staggered_1D_Tests_standalone_Ccodes/GiRaFFE_NRPy_standalone.c -o GiRaFFE_staggered_1D_Tests_standalone_Ccodes/output/GiRaFFE_NRPy_standalone -lm`...
    (BENCH): Finished executing in 3.0177741050720215 seconds.
    Finished compilation.
    (EXEC): Executing `taskset -c 0,1,2,3 ./GiRaFFE_NRPy_standalone 299 4 4 AlfvenWave`...
    (BENCH): Finished executing in 15.846468687057495 seconds.
    (EXEC): Executing `taskset -c 0,1,2,3 ./GiRaFFE_NRPy_standalone 299 4 4 ThreeAlfvenWaves`...
    (BENCH): Finished executing in 15.84378170967102 seconds.
    (EXEC): Executing `taskset -c 0,1,2,3 ./GiRaFFE_NRPy_standalone 299 4 4 DegenAlfvenWave`...
    (BENCH): Finished executing in 15.844474077224731 seconds.
    (EXEC): Executing `taskset -c 0,1,2,3 ./GiRaFFE_NRPy_standalone 299 4 4 FastWave`...
    (BENCH): Finished executing in 15.850006580352783 seconds.
    (EXEC): Executing `taskset -c 0,1,2,3 ./GiRaFFE_NRPy_standalone 299 4 4 FFEBD`...
    (BENCH): Finished executing in 16.047165870666504 seconds.

# Step 7: Data Visualization \[Back to [top](#toc)\]
$$\label{plots}$$

Now we plot the data and recreate figure 1 from the [GiRaFFE paper](https://arxiv.org/pdf/1704.00599.pdf). We reconstruct the electric field via

$$
E_i = -\epsilon_{ijk}v^j B^k
$$

the `calc_E` function below. We also calculate the FFE condition $B^2 - E^2$ below using the `calc_Bsquared_minus_Esquared` function.

```python
eDDD = ixp.LeviCivitaSymbol_dim3_rank3()
def calc_E(data):
    VU0 = data[:, 10]
    VU1 = data[:, 11]
    VU2 = data[:, 12]
    BU0 = data[:, 1]
    BU1 = data[:, 2]
    BU2 = data[:, 3]

    VU = [VU0, VU1, VU2]
    BU = [BU0, BU1, BU2]
    ED = np.zeros((VU0.size, 3))
    for i in range(3):
        for j in range(3):
            for k in range(3):
                ED[:,i] = ED[:,i] - eDDD[i][j][k]*VU[j]*BU[k]
    return ED

def calc_Bsquared_minus_Esquared(data):
    EU = calc_E(data)
    BU0 = data[:, 1]
    BU1 = data[:, 2]
    BU2 = data[:, 3]
    return (BU0**2 + BU1**2 + BU2**2) - (EU[:,0]**2 + EU[:,1]**2 + EU[:,2]**2)

```

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

#          0   1     2     3     4     5     6     7          8          9          10            11            12             13
labels = ["x","BU0","BU1","BU2","AD0","AD1","AD2","StildeD0","StildeD1","StildeD2","ValenciavU0","ValenciavU1","ValenciavU2", "psi6Phi"]
fig = plt.figure(figsize=(6, 15))

# spec = mpl.gridspec.GridSpec(ncols=6, nrows=2,wspace=0.65, hspace=0.4) # 6 columns evenly divides both 2 & 3
# ax1 = fig.add_subplot(spec[0,0:2]) # row 0 with axes spanning 2 cols on evens
# ax2 = fig.add_subplot(spec[0,2:4])
# ax3 = fig.add_subplot(spec[0,4:])
# ax4 = fig.add_subplot(spec[1,1:3]) # row 0 with axes spanning 2 cols on odds
# ax5 = fig.add_subplot(spec[1,3:5])

gs = gridspec.GridSpec(nrows=5, ncols=1, hspace=0.5)
ax1 = fig.add_subplot(gs[0, 0])
ax2 = fig.add_subplot(gs[1, 0])
ax3 = fig.add_subplot(gs[2, 0])
ax4 = fig.add_subplot(gs[3, 0])
ax5 = fig.add_subplot(gs[4, 0])

Data_num_Fast_A = np.loadtxt(os.path.join(Ccodesdir,"output","out299__FastWave-00000000.txt"))
Data_num_Fast_B = np.loadtxt(os.path.join(Ccodesdir,"output","out299__FastWave-00000100.txt"))

E_Fast_A = calc_E(Data_num_Fast_A)
E_Fast_B = calc_E(Data_num_Fast_B)

ax1.scatter(Data_num_Fast_A[:,0], np.abs(E_Fast_A[:,2]), s=1,label = 't = 0')
ax1.plot(Data_num_Fast_B[:,0], np.abs(E_Fast_B[:,2]), 'k-', label = 't = 0.5')
ax1.set_xlim(-0.5, 1.5)
ax1.set_ylim(0.6)
ax1.text(0.95, 0.01, 'Fast Wave',
        verticalalignment='bottom', horizontalalignment='right',
        transform=ax1.transAxes,
        color='black', fontsize=14)
ax1.set_xlabel('x')
ax1.set_ylabel(r'$|E^z|$')
ax1.legend()

Data_num_Alf_A = np.loadtxt(os.path.join(Ccodesdir,"output","out299__AlfvenWave-00000000.txt"))
Data_num_Alf_B = np.loadtxt(os.path.join(Ccodesdir,"output","out299__AlfvenWave-00000400.txt"))

ax2.scatter(Data_num_Alf_A[:,0], Data_num_Alf_A[:,3], s=1, label = 't = 0')
ax2.plot(Data_num_Alf_B[:,0], Data_num_Alf_B[:,3], 'k-', label = 't = 2.0')
ax2.set_xlim(-1.5, 1.5)
ax2.set_ylim(1.1)
ax2.text(0.95, 0.01, 'Alfven Wave',
        verticalalignment='bottom', horizontalalignment='right',
        transform=ax2.transAxes,
        color='black', fontsize=14)
ax2.set_xlabel('x')
ax2.set_ylabel(r'$B^z$')
ax2.legend(loc='center right')

Data_num_DegenAlf_A = np.loadtxt(os.path.join(Ccodesdir,"output","out299__DegenAlfvenWave-00000000.txt"))
Data_num_DegenAlf_B = np.loadtxt(os.path.join(Ccodesdir,"output","out299__DegenAlfvenWave-00000200.txt"))

E_DegenAlf_A = calc_E(Data_num_DegenAlf_A)
E_DegenAlf_B = calc_E(Data_num_DegenAlf_B)

ax3.scatter(Data_num_DegenAlf_A[:,0], E_DegenAlf_A[:,1], s=1, label = 't = 0')
ax3.plot(Data_num_DegenAlf_B[:,0], E_DegenAlf_B[:,1], 'k-', label = 't = 1.0')
ax3.set_xlim(-1.5, 1.5)
ax3.set_ylim(-1.35)
ax3.text(0.95, 0.01, 'Deg. Alfven Wave',
        verticalalignment='bottom', horizontalalignment='right',
        transform=ax3.transAxes,
        color='black', fontsize=14)
ax3.set_xlabel('x')
ax3.set_ylabel(r'$E^y$')
ax3.legend()

# Data_num_ThreeAlf_A = np.loadtxt(os.path.join(Ccodesdir,"output","out149__ThreeAlfvenWaves-00000000.txt"))
Data_num_ThreeAlf_B = np.loadtxt(os.path.join(Ccodesdir,"output","out299__ThreeAlfvenWaves-00000112.txt"))

# ax2.plot(Data_num_ThreeAlf_A[:,0], Data_num_ThreeAlf_A[:,2], 'k-')
ax4.scatter(Data_num_ThreeAlf_B[:,0], Data_num_ThreeAlf_B[:,2], s=1, label = 't = 0.56')
ax4.set_xlim(-1.0, 1.0)
# ax4.set_ylim()
ax4.text(0.95, 0.01, 'Three Waves',
        verticalalignment='bottom', horizontalalignment='right',
        transform=ax4.transAxes,
        color='black', fontsize=14)
ax4.set_xlabel('x')
ax4.set_ylabel(r'$B^y$')
ax4.legend(loc='center')

Data_num_FFEBD_A = np.loadtxt(os.path.join(Ccodesdir,"output","out299__FFEBD-00000000.txt"))
Data_num_FFEBD_B = np.loadtxt(os.path.join(Ccodesdir,"output","out299__FFEBD-00000004.txt"))

B2mE2_A = calc_Bsquared_minus_Esquared(Data_num_FFEBD_A)
B2mE2_B = calc_Bsquared_minus_Esquared(Data_num_FFEBD_B)

ax5.scatter(Data_num_FFEBD_A[:,0], B2mE2_A, s=1, label = 't = 0')
ax5.plot(Data_num_FFEBD_B[:,0], B2mE2_B, 'k-', label = 't = 0.02')
ax5.set_xlim(-0.4, 0.6)
ax5.text(0.95, 0.01, 'FFE Breakdown',
        verticalalignment='bottom', horizontalalignment='right',
        transform=ax5.transAxes,
        color='black', fontsize=14)
ax5.set_xlabel('x')
ax5.set_ylabel(r'$B^2 - E^2$')
ax5.legend()

plt.savefig(os.path.join(Ccodesdir,"output","NRPy-GiRaFFE"), dpi=800, bbox_inches="tight")
plt.close(fig)

```

```python
img1 = plt.imread(os.path.join(Ccodesdir,"output","NRPy-GiRaFFE.png"))
img2 =  plt.imread(os.path.join("GiRaFFE_NRPy/example_par_files/figure1_GiRaFFE_paper.png"))

NUM_ROWS = 1
IMGs_IN_ROW = 2
f, ax = plt.subplots(NUM_ROWS, IMGs_IN_ROW, figsize=(28,18))
plt.subplots_adjust(wspace=0.05)
plt.axis('off')

ax[0].imshow(img1)
ax[1].imshow(img2)

ax[0].set_title('image 1')
ax[1].set_title('image 2')

# title = 'side by side view of images'
# f.suptitle(title, fontsize=16)
plt.tight_layout()
# plt.xticks([])
# plt.yticks([])
plt.show()
```

# Step 8: Output this notebook to $\LaTeX$-formatted PDF file \[Back to [top](#toc)\]
$$\label{latex_pdf_output}$$

The following code cell converts this Jupyter notebook into a proper, clickable $\LaTeX$-formatted PDF file. After the cell is successfully run, the generated PDF may be found in the root NRPy+ tutorial directory, with filename
[Tutorial-Start_to_Finish-GiRaFFE_NRPy-1D_tests-staggered.pdf](Tutorial-Start_to_Finish-GiRaFFE_NRPy-1D_tests-staggered.pdf) (Note that clicking on this link may not work; you may need to open the PDF file through another means.)

```python
import cmdline_helper as cmd    # NRPy+: Multi-platform Python command-line interface
cmd.output_Jupyter_notebook_to_LaTeXed_PDF("Tutorial-Start_to_Finish-GiRaFFE_NRPy-1D_tests-staggered",location_of_template_file=os.path.join(".."))
```

    Created Tutorial-Start_to_Finish-GiRaFFE_NRPy-1D_tests-staggered.tex, and
        compiled LaTeX file to PDF file Tutorial-Start_to_Finish-
        GiRaFFE_NRPy-1D_tests-staggered.pdf


##### Copyright 2018 The TensorFlow Probability Authors.

Licensed under the Apache License, Version 2.0 (the "License");

```
#@title Licensed under the Apache License, Version 2.0 (the "License"); { display-mode: "form" }
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
```

# Gaussian Process Regression in TensorFlow Probability

    View on TensorFlow.org

    Run in Google Colab

    View source on GitHub

    Download notebook

In this colab, we explore Gaussian process regression using
TensorFlow and TensorFlow Probability. We generate some noisy observations from
some known functions and fit GP models to those data. We then sample from the GP
posterior and plot the sampled function values over grids in their domains.

## Background
Let $\mathcal{X}$ be any set. A *Gaussian process*
(GP) is a collection of random variables indexed by $\mathcal{X}$ such that if
$\{X_1, \ldots, X_n\} \subset \mathcal{X}$ is any finite subset, the marginal density
$p(X_1 = x_1, \ldots, X_n = x_n)$ is multivariate Gaussian. Any Gaussian
distribution is completely specified by its first and second central moments
(mean and covariance), and GP's are no exception. We can specify a GP completely
in terms of its mean function $\mu : \mathcal{X} \to \mathbb{R}$ and covariance function
$k : \mathcal{X} \times \mathcal{X} \to \mathbb{R}$. Most of the expressive power of GP's is encapsulated
in the choice of covariance function. For various reasons, the covariance
function is also referred to as a *kernel function*. It is required only to be
symmetric and positive-definite (see [Ch. 4 of Rasmussen & Williams](
http://www.gaussianprocess.org/gpml/chapters/RW4.pdf)). Below we make use of the
ExponentiatedQuadratic covariance kernel. Its form is

$$
k(x, x') := \sigma^2 \exp \left( \frac{\|x - x'\|^2}{\lambda^2} \right)
$$

where $\sigma^2$ is called the 'amplitude' and $\lambda$ the *length scale*.
The kernel parameters can be selected via a maximum likelihood optimization
procedure.

A full sample from a GP comprises a real-valued function over the entire space
$\mathcal{X}$ and is in practice impractical to realize; often one chooses a set of
points at which to observe a sample and draws function values at these points.
This is achieved by sampling from an appropriate (finite-dimensional)
multi-variate Gaussian.

Note that, according to the above definition, any finite-dimensional
multivariate Gaussian distribution is also a Gaussian process. Usually, when
one refers to a GP, it is implicit that the index set is some $\mathbb{R}^n$
and we will indeed make this assumption here.

A common application of Gaussian processes in machine learning is Gaussian
process regression. The idea is that we wish to estimate an unknown function
given noisy observations $\{y_1, \ldots, y_N\}$ of the function at a finite
number of points $\{x_1, \ldots x_N\}.$ We imagine a generative process

$$
\begin{align}
f \sim \: & \textsf{GaussianProcess}\left(
    \text{mean_fn}=\mu(x),
    \text{covariance_fn}=k(x, x')\right) \\
y_i \sim \: & \textsf{Normal}\left(
    \text{loc}=f(x_i),
    \text{scale}=\sigma\right), i = 1, \ldots, N
\end{align}
$$

As noted above, the sampled function is impossible to compute, since we would
require its value at an infinite number of points. Instead, one considers a
finite sample from a multivariate Gaussian.

$$
  \begin{gather}
    \begin{bmatrix}
      f(x_1) \\
      \vdots \\
      f(x_N)
    \end{bmatrix}
    \sim
    \textsf{MultivariateNormal} \left( \:
      \text{loc}=
      \begin{bmatrix}
        \mu(x_1) \\
        \vdots \\
        \mu(x_N)
      \end{bmatrix} \:,\:
      \text{scale}=
      \begin{bmatrix}
        k(x_1, x_1) & \cdots & k(x_1, x_N) \\
        \vdots & \ddots & \vdots \\
        k(x_N, x_1) & \cdots & k(x_N, x_N) \\
      \end{bmatrix}^{1/2}
    \: \right)
  \end{gather} \\
  y_i \sim \textsf{Normal} \left(
      \text{loc}=f(x_i),
      \text{scale}=\sigma
  \right)
$$

Note the exponent $\frac{1}{2}$ on the covariance matrix: this denotes a
Cholesky decomposition. Comptuing the Cholesky is necessary because the MVN is
a location-scale family distribution. Unfortunately the Cholesky decomposition
is computationally expensive, taking $O(N^3)$ time and $O(N^2)$ space. Much of
the GP literature is focused on dealing with this seemingly innocuous little
exponent.

It is common to take the prior mean function to be constant, often zero. Also,
some notational conventions are convenient. One often writes $\mathbf{f}$ for the
finite vector of sampled function values. A number of interesting notations are
used for the covariance matrix resulting from the application of $k$ to pairs of
inputs. Following [(Quiñonero-Candela, 2005)][QuinoneroCandela2005], we note
that the components of the matrix are covariances of function values at
particular input points. Thus we can denote the covariance matrix as $K_{AB}$
where $A$ and $B$ are some indicators of the collection of function values along
the given matrix dimensions.

[QuinoneroCandela2005]: http://www.jmlr.org/papers/volume6/quinonero-candela05a/quinonero-candela05a.pdf

For example, given observed data $(\mathbf{x}, \mathbf{y})$ with implied latent function
values $\mathbf{f}$, we can write

$$
K_{\mathbf{f},\mathbf{f}} = \begin{bmatrix}
  k(x_1, x_1) & \cdots & k(x_1, x_N) \\
  \vdots & \ddots & \vdots \\
  k(x_N, x_1) & \cdots & k(x_N, x_N) \\
\end{bmatrix}
$$

Similarly, we can mix sets of inputs, as in

$$
K_{\mathbf{f},*} = \begin{bmatrix}
  k(x_1, x^*_1) & \cdots & k(x_1, x^*_T) \\
  \vdots & \ddots & \vdots \\
  k(x_N, x^*_1) & \cdots & k(x_N, x^*_T) \\
\end{bmatrix}
$$

where we suppose there are $N$ training inputs, and $T$ test inputs. The above
generative process may then be written compactly as

$$
\begin{align}
\mathbf{f} \sim \: & \textsf{MultivariateNormal} \left(
        \text{loc}=\mathbf{0},
        \text{scale}=K_{\mathbf{f},\mathbf{f}}^{1/2}
    \right) \\
y_i \sim \: & \textsf{Normal} \left(
    \text{loc}=f_i,
    \text{scale}=\sigma \right), i = 1, \ldots, N
\end{align}
$$

The sampling operation in the first line yields a finite set of $N$ function
values from a multivariate Gaussian -- *not an entire function as in the above
GP draw notation*. The second line describes a collection of $N$ draws from
*univariate* Gaussians centered at the various function values, with fixed
observation noise $\sigma^2$.

With the above generative model in place, we can proceed to consider the
posterior inference problem. This yields a posterior distribution over function
values at a new set of test points, conditioned on the observed noisy data from
the process above.

With the above notation in place, we can compactly write the posterior
predictive distribution over future (noisy) observations conditional on
corresponding inputs and training data as follows (for more details, see §2.2 of
[Rasmussen & Williams](http://www.gaussianprocess.org/gpml/)).

$$
\mathbf{y}^* \mid \mathbf{x}^*, \mathbf{x}, \mathbf{y} \sim \textsf{Normal} \left(
    \text{loc}=\mathbf{\mu}^*,
    \text{scale}=(\Sigma^*)^{1/2}
\right),
$$

where

$$
\mathbf{\mu}^* = K_{*,\mathbf{f}}\left(K_{\mathbf{f},\mathbf{f}} + \sigma^2 I \right)^{-1} \mathbf{y}
$$

and

$$
\Sigma^* = K_{*,*} - K_{*,\mathbf{f}}
    \left(K_{\mathbf{f},\mathbf{f}} + \sigma^2 I \right)^{-1} K_{\mathbf{f},*}
$$

## Imports

```
import time

import numpy as np
import matplotlib.pyplot as plt
import tensorflow.compat.v2 as tf
import tensorflow_probability as tfp
tfb = tfp.bijectors
tfd = tfp.distributions
tfk = tfp.math.psd_kernels
tf.enable_v2_behavior()

from mpl_toolkits.mplot3d import Axes3D
%pylab inline
# Configure plot defaults
plt.rcParams['axes.facecolor'] = 'white'
plt.rcParams['grid.color'] = '#666666'
%config InlineBackend.figure_format = 'png'
```

    Populating the interactive namespace from numpy and matplotlib

## Example: Exact GP Regression on Noisy Sinusoidal Data
Here we generate training data from a noisy sinusoid, then sample a bunch of
curves from the posterior of the GP regression model. We use
[Adam](https://arxiv.org/abs/1412.6980) to optimize the kernel hyperparameters
(we minimize the negative log likelihood of the data under the prior). We
plot the training curve, followed by the true function and the posterior
samples.

```
def sinusoid(x):
  return np.sin(3 * np.pi * x[..., 0])

def generate_1d_data(num_training_points, observation_noise_variance):
  """Generate noisy sinusoidal observations at a random set of points.

  Returns:
     observation_index_points, observations
  """
  index_points_ = np.random.uniform(-1., 1., (num_training_points, 1))
  index_points_ = index_points_.astype(np.float64)
  # y = f(x) + noise
  observations_ = (sinusoid(index_points_) +
                   np.random.normal(loc=0,
                                    scale=np.sqrt(observation_noise_variance),
                                    size=(num_training_points)))
  return index_points_, observations_
```

```
# Generate training data with a known noise level (we'll later try to recover
# this value from the data).
NUM_TRAINING_POINTS = 100
observation_index_points_, observations_ = generate_1d_data(
    num_training_points=NUM_TRAINING_POINTS,
    observation_noise_variance=.1)
```

We'll put priors on the kernel hyperparameters, and write the joint distribution of the hyperparameters and observed data using `tfd.JointDistributionNamed`.

```
def build_gp(amplitude, length_scale, observation_noise_variance):
  """Defines the conditional dist. of GP outputs, given kernel parameters."""

  # Create the covariance kernel, which will be shared between the prior (which we
  # use for maximum likelihood training) and the posterior (which we use for
  # posterior predictive sampling)
  kernel = tfk.ExponentiatedQuadratic(amplitude, length_scale)

  # Create the GP prior distribution, which we will use to train the model
  # parameters.
  return tfd.GaussianProcess(
      kernel=kernel,
      index_points=observation_index_points_,
      observation_noise_variance=observation_noise_variance)

gp_joint_model = tfd.JointDistributionNamed({
    'amplitude': tfd.LogNormal(loc=0., scale=np.float64(1.)),
    'length_scale': tfd.LogNormal(loc=0., scale=np.float64(1.)),
    'observation_noise_variance': tfd.LogNormal(loc=0., scale=np.float64(1.)),
    'observations': build_gp,
})
```

We can sanity-check our implementation by verifying that we can sample from the prior, and compute the log-density of a sample.

```
x = gp_joint_model.sample()
lp = gp_joint_model.log_prob(x)

print("sampled {}".format(x))
print("log_prob of sample: {}".format(lp))
```

    sampled {'length_scale': <tf.Tensor: shape=(), dtype=float64, numpy=9.526471247059396>, 'observation_noise_variance': <tf.Tensor: shape=(), dtype=float64, numpy=0.7266366332947352>, 'amplitude': <tf.Tensor: shape=(), dtype=float64, numpy=1.6714059577238705>, 'observations': <tf.Tensor: shape=(100,), dtype=float64, numpy=
    array([3.28671007, 3.68127692, 1.46988415, 1.77062709, 1.67811476,
           2.55931267, 1.65328635, 1.73481701, 2.30216517, 1.29280539,
           2.48049681, 2.06504851, 1.24508666, 1.76258146, 1.1567751 ,
           1.12563896, 0.88050047, 2.62427676, 3.35178458, 0.89685013,
           3.22490858, 2.41177716, 2.09414069, 1.77904813, 0.16945602,
           2.53614107, 1.81245418, 1.35642751, 3.68498843, 1.82001509,
           2.03604088, 3.81705144, 3.18142879, 0.57016071, 1.68918439,
           2.85699517, 2.14177352, 3.72928207, 2.17608584, 3.28719282,
           0.6005666 , 1.96646054, 2.13131158, 1.95124967, 0.59090051,
           3.2427046 , 0.83849964, 1.35907388, 2.86403843, 1.56229808,
           1.47042658, 1.67109072, 3.20026542, 3.5325695 , 3.5645868 ,
           2.62192175, 2.68818167, 3.11055677, 2.67670315, 1.3779872 ,
           1.06653007, 2.15629192, 2.29109232, 2.48774572, 2.49993369,
           1.87633097, 4.14883802, 1.28336995, 0.62767486, 2.3908827 ,
           2.59999376, 1.84284938, 2.52603837, 1.25639193, 2.48792358,
           0.95560965, 2.59174573, 3.06116313, 2.30181874, 1.64854234,
           2.20463203, 0.86536671, 1.18708743, 1.88104063, 2.55473466,
           1.94555342, 1.64555061, 2.55028422, 0.86947597, 2.3485391 ,
           0.92084282, 1.51403219, 2.1829704 , 2.02841574, 1.0644413 ,
           3.26609039, 2.32084387, 2.74834809, 2.5470367 , 2.51504817])>}
    log_prob of sample: -137.270757398

Now let's optimize to find the parameter values with highest posterior probability. We'll define a variable for each parameter, and constrain their values to be positive.

```
# Create the trainable model parameters, which we'll subsequently optimize.
# Note that we constrain them to be strictly positive.

constrain_positive = tfb.Shift(np.finfo(np.float64).tiny)(tfb.Exp())

amplitude_var = tfp.util.TransformedVariable(
    initial_value=1.,
    bijector=constrain_positive,
    name='amplitude',
    dtype=np.float64)

length_scale_var = tfp.util.TransformedVariable(
    initial_value=1.,
    bijector=constrain_positive,
    name='length_scale',
    dtype=np.float64)

observation_noise_variance_var = tfp.util.TransformedVariable(
    initial_value=1.,
    bijector=constrain_positive,
    name='observation_noise_variance_var',
    dtype=np.float64)

trainable_variables = [v.trainable_variables[0] for v in
                       [amplitude_var,
                       length_scale_var,
                       observation_noise_variance_var]]
```

To condition the model on our observed data, we'll define a `target_log_prob` function, which takes the (still to be inferred) kernel hyperparameters.

```
# Use `tf.function` to trace the loss for more efficient evaluation.
@tf.function(autograph=False, jit_compile=False)
def target_log_prob(amplitude, length_scale, observation_noise_variance):
  return gp_joint_model.log_prob({
      'amplitude': amplitude,
      'length_scale': length_scale,
      'observation_noise_variance': observation_noise_variance,
      'observations': observations_
  })
```

```
# Now we optimize the model parameters.
num_iters = 1000
optimizer = tf.optimizers.Adam(learning_rate=.01)

# Store the likelihood values during training, so we can plot the progress
lls_ = np.zeros(num_iters, np.float64)
for i in range(num_iters):
  with tf.GradientTape() as tape:
    loss = -target_log_prob(amplitude_var, length_scale_var,
                            observation_noise_variance_var)
  grads = tape.gradient(loss, trainable_variables)
  optimizer.apply_gradients(zip(grads, trainable_variables))
  lls_[i] = loss

print('Trained parameters:')
print('amplitude: {}'.format(amplitude_var._value().numpy()))
print('length_scale: {}'.format(length_scale_var._value().numpy()))
print('observation_noise_variance: {}'.format(observation_noise_variance_var._value().numpy()))
```

    Trained parameters:
    amplitude: 0.875065718641
    length_scale: 0.176252877463
    observation_noise_variance: 0.0971537576792

```
# Plot the loss evolution
plt.figure(figsize=(12, 4))
plt.plot(lls_)
plt.xlabel("Training iteration")
plt.ylabel("Log marginal likelihood")
plt.show()
```

```
# Having trained the model, we'd like to sample from the posterior conditioned
# on observations. We'd like the samples to be at points other than the training
# inputs.
predictive_index_points_ = np.linspace(-1.2, 1.2, 200, dtype=np.float64)
# Reshape to [200, 1] -- 1 is the dimensionality of the feature space.
predictive_index_points_ = predictive_index_points_[..., np.newaxis]

optimized_kernel = tfk.ExponentiatedQuadratic(amplitude_var, length_scale_var)
gprm = tfd.GaussianProcessRegressionModel(
    kernel=optimized_kernel,
    index_points=predictive_index_points_,
    observation_index_points=observation_index_points_,
    observations=observations_,
    observation_noise_variance=observation_noise_variance_var,
    predictive_noise_variance=0.)

# Create op to draw  50 independent samples, each of which is a *joint* draw
# from the posterior at the predictive_index_points_. Since we have 200 input
# locations as defined above, this posterior distribution over corresponding
# function values is a 200-dimensional multivariate Gaussian distribution!
num_samples = 50
samples = gprm.sample(num_samples)
```

```
# Plot the true function, observations, and posterior samples.
plt.figure(figsize=(12, 4))
plt.plot(predictive_index_points_, sinusoid(predictive_index_points_),
         label='True fn')
plt.scatter(observation_index_points_[:, 0], observations_,
            label='Observations')
for i in range(num_samples):
  plt.plot(predictive_index_points_, samples[i, :], c='r', alpha=.1,
           label='Posterior Sample' if i == 0 else None)
leg = plt.legend(loc='upper right')
for lh in leg.legendHandles:
    lh.set_alpha(1)
plt.xlabel(r"Index points ($\mathbb{R}^1$)")
plt.ylabel("Observation space")
plt.show()
```

*Note: if you run the above code several times, sometimes it looks great and
other times it looks terrible! The maximum likelihood training of the parameters
is quite sensitive and sometimes converges to poor models. The best approach
is to use MCMC to marginalize the model hyperparameters.*

## Marginalizing hyperparameters with HMC

Instead of optimizing the hyperparameters, let's try integrating them out with Hamiltonian Monte Carlo. We'll first define and run a sampler to approximately draw from the posterior distribution over kernel hyperparameters, given the observations.

```
num_results = 100
num_burnin_steps = 50

sampler = tfp.mcmc.TransformedTransitionKernel(
    tfp.mcmc.HamiltonianMonteCarlo(
        target_log_prob_fn=target_log_prob,
        step_size=tf.cast(0.1, tf.float64),
        num_leapfrog_steps=8),
    bijector=[constrain_positive, constrain_positive, constrain_positive])

adaptive_sampler = tfp.mcmc.DualAveragingStepSizeAdaptation(
    inner_kernel=sampler,
    num_adaptation_steps=int(0.8 * num_burnin_steps),
    target_accept_prob=tf.cast(0.75, tf.float64))

initial_state = [tf.cast(x, tf.float64) for x in [1., 1., 1.]]
```

```
# Speed up sampling by tracing with `tf.function`.
@tf.function(autograph=False, jit_compile=False)
def do_sampling():
  return tfp.mcmc.sample_chain(
      kernel=adaptive_sampler,
      current_state=initial_state,
      num_results=num_results,
      num_burnin_steps=num_burnin_steps,
      trace_fn=lambda current_state, kernel_results: kernel_results)

t0 = time.time()
samples, kernel_results = do_sampling()
t1 = time.time()
print("Inference ran in {:.2f}s.".format(t1-t0))
```

    Inference ran in 48.94s.

Let's sanity-check the sampler by examining the hyperparameter traces.

```
(amplitude_samples,
 length_scale_samples,
 observation_noise_variance_samples) = samples

f = plt.figure(figsize=[15, 3])
for i, s in enumerate(samples):
  ax = f.add_subplot(1, len(samples) + 1, i + 1)
  ax.plot(s)
```

Now instead of constructing a single GP with the optimized hyperparameters, we construct the *posterior predictive distribution* as a mixture of GPs, each defined by a sample from the posterior distribution over hyperparameters. This approximately integrates over the posterior parameters via Monte Carlo sampling to compute the marginal predictive distribution at unobserved locations.

```
# The sampled hyperparams have a leading batch dimension, `[num_results, ...]`,
# so they construct a *batch* of kernels.
batch_of_posterior_kernels = tfk.ExponentiatedQuadratic(
    amplitude_samples, length_scale_samples)

# The batch of kernels creates a batch of GP predictive models, one for each
# posterior sample.
batch_gprm = tfd.GaussianProcessRegressionModel(
    kernel=batch_of_posterior_kernels,
    index_points=predictive_index_points_,
    observation_index_points=observation_index_points_,
    observations=observations_,
    observation_noise_variance=observation_noise_variance_samples,
    predictive_noise_variance=0.)

# To construct the marginal predictive distribution, we average with uniform
# weight over the posterior samples.
predictive_gprm = tfd.MixtureSameFamily(
    mixture_distribution=tfd.Categorical(logits=tf.zeros([num_results])),
    components_distribution=batch_gprm)

num_samples = 50
samples = predictive_gprm.sample(num_samples)
```

```
# Plot the true function, observations, and posterior samples.
plt.figure(figsize=(12, 4))
plt.plot(predictive_index_points_, sinusoid(predictive_index_points_),
         label='True fn')
plt.scatter(observation_index_points_[:, 0], observations_,
            label='Observations')
for i in range(num_samples):
  plt.plot(predictive_index_points_, samples[i, :], c='r', alpha=.1,
           label='Posterior Sample' if i == 0 else None)
leg = plt.legend(loc='upper right')
for lh in leg.legendHandles:
    lh.set_alpha(1)
plt.xlabel(r"Index points ($\mathbb{R}^1$)")
plt.ylabel("Observation space")
plt.show()
```

Although the differences are subtle in this case, in general, we would expect the posterior predictive distribution to generalize better (give higher likelihood to held-out data) than just using the most likely parameters as we did above.