# Neuromatch Academy: Week 2, Day 5, Tutorial 1
# Learning to Predict

__Content creators:__ Marcelo Mattar and Eric DeWitt with help from Matt Krause

__Content reviewers:__ Byron Galbraith and Michael Waskom

---

# Tutorial objectives

In this tutorial, we will learn how to estimate state-value functions in a classical conditioning paradigm using Temporal Difference (TD) learning and examine TD-errors at the presentation of the conditioned and unconditioned stimulus (CS and US) under different CS-US contingencies. These exercises will provide you with an understanding of both how reward prediction errors (RPEs) behave in classical conditioning and what we should expect to see if Dopamine represents a "canonical" model-free RPE.

At the end of this tutorial:
* You will learn to use the standard tapped delay line conditioning model
* You will understand how RPEs move to CS
* You will understand how variability in reward size effects RPEs
* You will understand how differences in US-CS timing effect RPEs

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

```python
#@title Figure settings
import ipywidgets as widgets       # interactive display
%config InlineBackend.figure_format = 'retina'
plt.style.use("https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/nma.mplstyle")
```

```python
# @title Helper functions
from matplotlib import ticker

def plot_value_function(V, ax=None, show=True):
  """Plot V(s), the value function"""
  if not ax:
    fig, ax = plt.subplots()

  ax.stem(V, use_line_collection=True)
  ax.set_ylabel('Value')
  ax.set_xlabel('State')
  ax.set_title("Value function: $V(s)$")

  if show:
    plt.show()

def plot_tde_trace(TDE, ax=None, show=True, skip=400):
  """Plot the TD Error across trials"""
  if not ax:
    fig, ax = plt.subplots()

  indx = np.arange(0, TDE.shape[1], skip)
  im = ax.imshow(TDE[:,indx])
  positions = ax.get_xticks()
  # Avoid warning when setting string tick labels
  ax.xaxis.set_major_locator(ticker.FixedLocator(positions))
  ax.set_xticklabels([f"{int(skip * x)}" for x in positions])
  ax.set_title('TD-error over learning')
  ax.set_ylabel('State')
  ax.set_xlabel('Iterations')
  ax.figure.colorbar(im)
  if show:
    plt.show()

def learning_summary_plot(V, TDE):
  """Summary plot for Ex1"""
  fig, (ax1, ax2) = plt.subplots(nrows = 2, gridspec_kw={'height_ratios': [1, 2]})

  plot_value_function(V, ax=ax1, show=False)
  plot_tde_trace(TDE, ax=ax2, show=False)
  plt.tight_layout()

def reward_guesser_title_hint(r1, r2):
  """"Provide a mildly obfuscated hint for a demo."""
  if (r1==14 and r2==6) or (r1==6 and r2==14):
    return "Technically correct...(the best kind of correct)"

  if  ~(~(r1+r2) ^ 11) - 1 == (6 | 24): # Don't spoil the fun :-)
    return "Congratulations! You solved it!"

  return "Keep trying...."

#@title Default title text
class ClassicalConditioning:

    def __init__(self, n_steps, reward_magnitude, reward_time):

        # Task variables
        self.n_steps = n_steps
        self.n_actions = 0
        self.cs_time = int(n_steps/4) - 1

        # Reward variables
        self.reward_state = [0,0]
        self.reward_magnitude = None
        self.reward_probability = None
        self.reward_time = None

        self.set_reward(reward_magnitude, reward_time)

        # Time step at which the conditioned stimulus is presented

        # Create a state dictionary
        self._create_state_dictionary()

    def set_reward(self, reward_magnitude, reward_time):

        """
        Determine reward state and magnitude of reward
        """
        if reward_time >= self.n_steps - self.cs_time:
            self.reward_magnitude = 0

        else:
            self.reward_magnitude = reward_magnitude
            self.reward_state = [1, reward_time]

    def get_outcome(self, current_state):

        """
        Determine next state and reward
        """
        # Update state
        if current_state < self.n_steps - 1:
            next_state = current_state + 1
        else:
            next_state = 0

        # Check for reward
        if self.reward_state == self.state_dict[current_state]:
            reward = self.reward_magnitude
        else:
            reward = 0

        return next_state, reward

    def _create_state_dictionary(self):

        """
        This dictionary maps number of time steps/ state identities
        in each episode to some useful state attributes:

        state      - 0 1 2 3 4 5 (cs) 6 7 8 9 10 11 12 ...
        is_delay   - 0 0 0 0 0 0 (cs) 1 1 1 1  1  1  1 ...
        t_in_delay - 0 0 0 0 0 0 (cs) 1 2 3 4  5  6  7 ...
        """
        d = 0

        self.state_dict = {}
        for s in range(self.n_steps):
            if s <= self.cs_time:
                self.state_dict[s] = [0,0]
            else:
                d += 1 # Time in delay
                self.state_dict[s] = [1,d]

class MultiRewardCC(ClassicalConditioning):
  """Classical conditioning paradigm, except that one randomly selected reward,
    magnitude, from a list, is delivered of a single fixed reward."""
  def __init__(self, n_steps, reward_magnitudes, reward_time=None):
    """"Build a multi-reward classical conditioning environment
      Args:
        - nsteps: Maximum number of steps
        - reward_magnitudes: LIST of possible reward magnitudes.
        - reward_time: Single fixed reward time
      Uses numpy global random state.
      """
    super().__init__(n_steps, 1, reward_time)
    self.reward_magnitudes = reward_magnitudes

  def get_outcome(self, current_state):
    next_state, reward = super().get_outcome(current_state)
    if reward:
      reward=np.random.choice(self.reward_magnitudes)
    return next_state, reward

class ProbabilisticCC(ClassicalConditioning):
  """Classical conditioning paradigm, except that rewards are stochastically omitted."""
  def __init__(self, n_steps, reward_magnitude, reward_time=None, p_reward=0.75):
    """"Build a multi-reward classical conditioning environment
      Args:
        - nsteps: Maximum number of steps
        - reward_magnitudes: Reward magnitudes.
        - reward_time: Single fixed reward time.
        - p_reward: probability that reward is actually delivered in rewarding state
      Uses numpy global random state.
      """
    super().__init__(n_steps, reward_magnitude, reward_time)
    self.p_reward = p_reward

  def get_outcome(self, current_state):
    next_state, reward = super().get_outcome(current_state)
    if reward:
      reward*= int(np.random.uniform(size=1)[0] < self.p_reward)
    return next_state, reward

```

---
# Section 1: TD-learning

```python
#@title Video 1: Introduction
from IPython.display import YouTubeVideo
video = YouTubeVideo(id="YoNbc9M92YY", width=854, height=480, fs=1)
print("Video available at https://youtu.be/" + video.id)
video
```

    Video available at https://youtu.be/YoNbc9M92YY

__Environment:__

- The agent experiences the environment in episodes or trials.
- Episodes terminate by transitioning to the inter-trial-interval (ITI) state and they are initiated from the ITI state as well. We clamp the value of the terminal/ITI states to zero.
- The classical conditioning environment is composed of a sequence of states that the agent deterministically transitions through. Starting at State 0, the agent moves to State 1 in the first step, from State 1 to State 2 in the second, and so on.  These states represent time in the tapped delay line representation
- Within each episode, the agent is presented a CS and US (reward).
- The CS is always presented at 1/4 of the total duration of the trial. The US (reward) is then delivered after the CS. The interval between the CS and US is specified by `reward_time`.
- The agent's goal is to learn to predict expected rewards from each state in the trial.

**General concepts**

* Return $G_{t}$: future cumulative reward, which can be written in arecursive form
\begin{align}
G_{t} &= \sum \limits_{k = 0}^{\infty} \gamma^{k} r_{t+k+1} \\
&= r_{t+1} + \gamma G_{t+1}
\end{align}
where $\gamma$ is discount factor that controls the importance of future rewards, and $\gamma \in [0, 1]$. $\gamma$ may also be interpreted as probability of continuing the trajectory.
* Value funtion $V_{\pi}(s_t=s)$: expecation of the return
\begin{align}
V_{\pi}(s_t=s) &= \mathbb{E} [ G_{t}\; | \; s_t=s, a_{t:\infty}\sim\pi] \\
& = \mathbb{E} [ r_{t+1} + \gamma G_{t+1}\; | \; s_t=s, a_{t:\infty}\sim\pi]
\end{align}
With an assumption of **Markov process**, we thus have:
\begin{align}
V_{\pi}(s_t=s) &= \mathbb{E} [ r_{t+1} + \gamma V_{\pi}(s_{t+1})\; | \; s_t=s, a_{t:\infty}\sim\pi] \\
&= \sum_a \pi(a|s) \sum_{r, s'}p(s', r)(r + V_{\pi}(s_{t+1}=s'))
\end{align}

**Temporal difference (TD) learning**

* With a Markovian assumption, we can use $V(s_{t+1})$ as an imperfect proxy for the true value $G_{t+1}$ (Monte Carlo bootstrapping), and thus obtain the generalised equation to calculate TD-error:
\begin{align}
\delta_{t} = r_{t+1} + \gamma V(s_{t+1}) - V(s_{t})
\end{align}

* Value updated by using the learning rate constant $\alpha$:
\begin{align}
V(s_{t}) \leftarrow V(s_{t}) + \alpha \delta_{t}
\end{align}

  (Reference: https://web.stanford.edu/group/pdplab/pdphandbook/handbookch10.html)

__Definitions:__

* TD-error:
\begin{align}
\delta_{t} = r_{t+1} + \gamma V(s_{t+1}) - V(s_{t})
\end{align}

* Value updates:
\begin{align}
V(s_{t}) \leftarrow V(s_{t}) + \alpha \delta_{t}
\end{align}

## Exercise 1: TD-learning with guaranteed rewards

Implement TD-learning to estimate the state-value function in the classical-conditioning world with guaranteed rewards, with a fixed magnitude, at a fixed delay after the conditioned stimulus, CS. Save TD-errors over learning (i.e., over trials) so we can visualize them afterwards.

In order to simulate the effect of the CS, you should only update $V(s_{t})$ during the delay period after CS. This period is indicated by the boolean variable `is_delay`. This can be implemented by multiplying the expression for updating the value function by `is_delay`.

Use the provided code to estimate the value function.

```python
def td_learner(env, n_trials, gamma=0.98, alpha=0.001):
  """ Temporal Difference learning

  Args:
    env (object): the environment to be learned
    n_trials (int): the number of trials to run
    gamma (float): temporal discount factor
    alpha (float): learning rate

  Returns:
    ndarray, ndarray: the value function and temporal difference error arrays
  """
  V = np.zeros(env.n_steps) # Array to store values over states (time)
  TDE = np.zeros((env.n_steps, n_trials)) # Array to store TD errors

  for n in range(n_trials):
    state = 0 # Initial state
    for t in range(env.n_steps):
      # Get next state and next reward
      next_state, reward = env.get_outcome(state)
      # Is the current state in the delay period (after CS)?
      is_delay = env.state_dict[state][0]

      # Write an expression to compute the TD-error
      TDE[state, n] = reward + gamma * V[next_state] - V[state]

      # Write an expression to update the value function
      # V[state] += alpha * TDE[state, n]
      V[state] += alpha * TDE[state, n] * is_delay

      # Update state
      state = next_state

  return V, TDE

env = ClassicalConditioning(n_steps=40, reward_magnitude=10, reward_time=10)
V, TDE = td_learner(env, n_trials=20000)
with plt.xkcd():
  learning_summary_plot(V, TDE)
```

## Interactive Demo 1: US to CS Transfer

During classical conditioning, the subject's behavioral response (e.g., salivating) transfers from the unconditioned stimulus (US; like the smell of tasty food) to the conditioned stimulus (CS; like Pavlov ringing his bell) that predicts it. Reward prediction errors play an important role in this process by adjusting the value of states according to their expected, discounted return.

Use the widget below to examine how reward prediction errors change over time. Before training (orange line), only the reward state has high reward prediction error. As training progresses (blue line, slider), the reward prediction errors shift to the conditioned stimulus, where they end up when the trial is complete (green line).

Dopamine neurons, which are thought to carry reward prediction errors _in vivo_, show exactly the same behavior!

```python
#@title

#@markdown Make sure you execute this cell to enable the widget!

n_trials = 20000

@widgets.interact
def plot_tde_by_trial(trial = widgets.IntSlider(value=5000, min=0, max=n_trials-1 , step=1, description="Trial #")):
  if 'TDE' not in globals():
    print("Complete Exercise 1 to enable this interactive demo!")
  else:

    fig, ax = plt.subplots()
    ax.axhline(0, color='k') # Use this + basefmt=' ' to keep the legend clean.
    ax.stem(TDE[:, 0], linefmt='C1-', markerfmt='C1d', basefmt=' ',
            label="Before Learning (Trial 0)",
            use_line_collection=True)
    ax.stem(TDE[:, -1], linefmt='C2-', markerfmt='C2s', basefmt=' ',
            label="After Learning (Trial $\infty$)",
            use_line_collection=True)
    ax.stem(TDE[:, trial], linefmt='C0-', markerfmt='C0o', basefmt=' ',
            label=f"Trial {trial}",
            use_line_collection=True)

    ax.set_xlabel("State in trial")
    ax.set_ylabel("TD Error")
    ax.set_title("Temporal Difference Error by Trial")
    ax.legend()
```

    interactive(children=(IntSlider(value=5000, description='Trial #', max=19999), Output()), _dom_classes=('widge…

## Interactive Demo 2: Learning Rates and Discount Factors

Our TD-learning agent has two parameters that control how it learns: $\alpha$, the learning rate, and $\gamma$, the discount factor. In Exercise 1, we set these parameters to $\alpha=0.001$ and $\gamma=0.98$ for you. Here, you'll investigate how changing these parameters alters the model that TD-learning learns.

Before enabling the interactive demo below, take a moment to think about the functions of these two parameters. $\alpha$ controls the size of the Value function updates produced by each TD-error. In our simple, deterministic world, will this affect the final model we learn? Is a larger $\alpha$ necessarily better in more complex, realistic environments?

The discount rate $\gamma$ applies an exponentially-decaying weight to returns occuring in the future, rather than the present timestep. How does this affect the model we learn? What happens when $\gamma=0$ or $\gamma \geq 1$?

Use the widget to test your hypotheses.

```python
#@title

#@markdown Make sure you execute this cell to enable the widget!

@widgets.interact
def plot_summary_alpha_gamma(alpha = widgets.FloatSlider(value=0.0001, min=0.001, max=0.1, step=0.0001, description="alpha"),
                             gamma = widgets.FloatSlider(value=0.980, min=0, max=1.1, step=0.010, description="gamma")):
  env = ClassicalConditioning(n_steps=40, reward_magnitude=10, reward_time=10)
  try:
    V_params, TDE_params = td_learner(env, n_trials=20000, gamma=gamma, alpha=alpha)
  except NotImplementedError:
    print("Finish Exercise 1 to enable this interactive demo")

  learning_summary_plot(V_params,TDE_params)

```

    interactive(children=(FloatSlider(value=0.001, description='alpha', max=0.1, min=0.001, step=0.0001), FloatSli…

```python
"""
Alpha determines how fast the model learns. In the simple, deterministic world
we're using here, this allows the moodel to quickly converge onto the "true"
model that heavily values the conditioned stimulus. In more complex environments,
however, excessively large values of alpha can slow, or even prevent, learning,
as we'll see later.

Gamma effectively controls how much the model cares about the future: larger values of
gamma cause the model to weight future rewards nearly as much as present ones. At gamma=1,
the model weights all rewards, regardless of when they occur, equally and when greater than one, it
starts to *prefer* rewards in the future, rather than the present (this is rarely good).
When gamma=0, however, the model becomes greedy and only considers rewards that
can be obtained immediately.
 """;
```

---
# Section 2: TD-learning with varying reward magnitudes

In the previous exercise, the environment was as simple as possible. On every trial, the CS predicted the same reward, at the same time, with 100% certainty. In the next few exercises, we will make the environment more progressively more complicated and examine the TD-learner's behavior.

## Interactive Demo 3: Match the Value Functions

First, will replace the environment with one that dispenses one of several rewards, chosen at random. Shown below is the final value function $V$ for a TD learner that was trained in an enviroment where the CS predicted a reward of 6 or 14 units; both rewards were equally likely).

Can you find another pair of rewards that cause the agent to learn the same value function? Assume each reward will be dispensed 50% of the time.

Hints:
* Carefully consider the definition of the value function $V$. This can be solved analytically.
* There is no need to change $\alpha$ or $\gamma$.
* Due to the randomness, there may be a small amount of variation.

```python
#@title

#@markdown Make sure you execute this cell to enable the widget!

n_trials = 20000
np.random.seed(2020)
rng_state = np.random.get_state()
env = MultiRewardCC(40, [6, 14], reward_time=10)
V_multi, TDE_multi = td_learner(env, n_trials, gamma=0.98, alpha=0.001)

@widgets.interact
def reward_guesser_interaction(r1 = widgets.IntText(value=0, min=0, max=50, description="Reward 1"),
                               r2 = widgets.IntText(value=0, min=0, max=50, description="Reward 2")):
  try:
    env2 = MultiRewardCC(40, [r1, r2], reward_time=10)
    V_guess, _ = td_learner(env2, n_trials, gamma=0.98, alpha=0.001)
    fig, ax = plt.subplots()
    m, l, _ = ax.stem(V_multi, linefmt='y-', markerfmt='yo', basefmt=' ', label="Target",
            use_line_collection=True)
    m.set_markersize(15)
    m.set_markerfacecolor('none')
    l.set_linewidth(4)
    m, _, _ = ax.stem(V_guess, linefmt='r', markerfmt='rx', basefmt=' ', label="Guess",
                      use_line_collection=True)
    m.set_markersize(15)

    ax.set_xlabel("State")
    ax.set_ylabel("Value")
    ax.set_title("Guess V(s)\n" + reward_guesser_title_hint(r1, r2))
    ax.legend()
  except NotImplementedError:
    print("Please finish Exercise 1 first!")
```

    interactive(children=(IntText(value=0, description='Reward 1'), IntText(value=0, description='Reward 2'), Outp…

## Section 2.1 Examining the TD Error

Run the cell below to plot the TD errors from our multi-reward environment. A new feature appears in this plot? What is it? Why does it happen?

```python
plot_tde_trace(TDE_multi)
```

```python
"""
The TD trace now takes on negative values because the reward delivered is
sometimes larger than the expected reward and sometimes smaller.
""";
```

---
# Section 3: TD-learning with probabilistic rewards

In this environment, we'll return to delivering a single reward of ten units. However, it will be delivered intermittently: on 20 percent of trials, the CS will be shown but the agent will not receive the usual reward; the remaining 80% will proceed as usual.

 Run the cell below to simulate. How does this compare with the previous experiment?

Earlier in the notebook, we saw that changing $\alpha$ had little effect on learning in a deterministic environment. What happens if you set it to an large value, like 1, in this noisier scenario? Does it seem like it will _ever_ converge?

```python
# np.random.set_state(rng_state) # Resynchronize everyone's notebooks
n_trials = 10000
try:
  env = ProbabilisticCC(n_steps=40, reward_magnitude=10, reward_time=10,
                        p_reward=0.8)
  # V_stochastic, TDE_stochastic = td_learner(env, n_trials*2, alpha=0.001)
  V_stochastic, TDE_stochastic = td_learner(env, n_trials*2, alpha=1)
  learning_summary_plot(V_stochastic, TDE_stochastic)
except NotImplementedError:
  print("Please finish Exercise 1 first")
```

```python
"""
The multi-reward and probabilistic reward enviroments are the same. You
could simulate a probabilistic reward of 10 units, delivered 50% of the time,
by having a mixture of 10 and 0 unit rewards, or vice versa. The takehome message
from these last three exercises is that the *average* or expected reward is
what matters during TD learning.

Large values of alpha prevent the TD Learner from converging. As a result, the
value function seems implausible: one state may have an extremely low value while
the neighboring ones remain high. This pattern persists even if training continues
for hundreds of thousands of trials.
""";
```

---
# Summary

In this notebook, we have developed a simple TD Learner and examined how its state representations and reward prediction errors evolve during training. By manipualting its environment and parameters ($\alpha$, $\gamma$), you developed an intuition for how it behaves.

This simple model closely resembles the behavior of subjects undergoing classical conditioning tasks and the dopamine neurons that may underlie that behavior. You may have implemented TD-reset or used the model to recreate a common experimental error. The update rule used here has been extensively studied for [more than 70 years](https://www.pnas.org/content/108/Supplement_3/15647) as a possible explanation for artificial and biological learning.

However, you may have noticed that something is missing from this notebook. We carefully calculated the value of each state, but did not use it to actually do anything. Using values to plan _**Actions**_ is coming up next!

# Bonus

## Exercise 2: Removing the CS

In Exercise 1, you (should have) included a term that depends on the conditioned stimulus. Remove it and see what happens. Do you understand why?
This phenomena often fools people attempting to train animals--beware!

```python
"""
You should only be updating the V[state] once the conditioned stimulus appears.

If you remove this term the Value Function becomes periodic, dropping towards zero
right after the reward and gradually rising towards the end of the trial. This
behavior is actually correct, because the model is learning the time until the
*next* reward, and State 37 is closer to a reward than State 21 or 22.

In an actual experiment, the animal often just wants rewards; it doesn't care about
/your/ experiment or trial structure!
""";
```


# Recurrent networks

- Once we learn the general equation and properties of the line / hyperplane we can turn around and perform various learning tasks - like regression and classification - using it as a model.

- By the same token now that we have a basic understanding of how to model general ordered data, we can use those models to perform various learning tasks.

- In this Section we use the generic variable order dynamic system introduced in the previously to perform simple learning tasks on time series.  This is often referred to as Recurrent Neural Networks (or RNNs for short).  RNNs can also be used for more sophisticated time series problems - like machine translation - which we discuss in the next Section.

Activate next cell to toggle code on and off

```python
from IPython.display import display
from IPython.display import HTML
import IPython.core.display as di # Example: di.display_html('<h3>%s:</h3>' % str, raw=True)

# This line will hide code by default when the notebook is eåxported as HTML
di.display_html('', raw=True)

# This line will add a button to toggle visibility of code blocks, for use with the HTML export version
di.display_html('''<button onclick="jQuery('.input_area').toggle(); jQuery('.prompt').toggle();">Toggle code</button>''', raw=True)
```

<button onclick="jQuery('.input_area').toggle(); jQuery('.prompt').toggle();">Toggle code</button>

```python
## This code cell will not be shown in the HTML version of this notebook
# imports from custom library
import sys
sys.path.append('../../')

# import custom lib
from mlrefined_libraries import recurrent_library as reclib

# import dataset path
datapath = '../../mlrefined_datasets/recurrent_datasets/'

# import autograd functionality to bulid function's properly for optimizers
import autograd.numpy as np

# this is needed to compensate for %matplotlib notebook's tendancy to blow up images when plotted inline
%matplotlib notebook
from matplotlib import rcParams
rcParams['figure.autolayout'] = True

%load_ext autoreload
%autoreload 2
```

# Using unlimited memory dynamic systems for time series prediction

- Recurrent Neural Networks (RNNs) = the use of [unlimited memory dynamic systems](https://jermwatt.github.io/control-notes/posts/dynamic_systems_unlimited_memory/dynamic_systems_unlimited_memory.html) for learning on time series

- General formula for a simple unlimited memory system

\begin{equation}
h_p = f\left(h_{p-1},x_p \right).
\end{equation}

- If $f$ here is linear we have

\begin{equation}
h_p = w_0 + w_1h_{p-1} + w_2x_p.
\end{equation}

- To use such a model for any sort of prediction task we

choose a function $f$

 train favorite standard ML scheme on input/output pairs

 to predict next future value feed in *only the most recent value* of the series into the trained model

- Making a training set of input/output data for a time series by feeding in the entire series, and "windowing" of ever-increasing size will happen natually

- A common sort of model consisting of a simple unlimited memory dynamic system - often called a *Recurrent Neural Network*  and the function $f$ usually chosen as a standard activation function - adds a final linear combination to the final output of the system

\begin{equation}
\text{model}\left(x_p,\Theta\right) = v_0 + h_pv_1
\end{equation}

- Here the parameter set $\Theta$ contains $v_0, v_1$ the weights of the final linear combination, as well as any parameters internal to $h_p = f\left(h_{p-1},x_p \right)$

- For a regression problem this model generates our predicted output

- Denote by

$$
\,\,\,\,\, \,\,\,   \,\,\,  \,\,\, \,\,\,   \,\,\,   \,\,\,  \,\,    \hat{y}_p:=  \text{model}\left(x_p,\Theta\right)  \\
y_p:= x_p \\
\,\,\,\,\,  \,\,\,\,\,  \,\,\,\,\,  \,\,\,\,\,  \,\,  \ell\left(a,b\right) := \left(a-b \right)^2
$$

- We have a standard Least Squares regression

\begin{equation}
\frac{1}{P-D}\sum_{p=D}^{P}\ell\left(\hat{y}_p,y_p\right)
\end{equation}

- When used for time series prediction we hope hidden states learn interesting and useful *summary statistics* and important *long-term dependencies* - of our entire dataset

- [Deeper systems with additional hidden layers](https://jermwatt.github.io/control-notes/posts/dynamic_systems_unlimited_memory/dynamic_systems_unlimited_memory.html#deeper-systems) can learn higher order summary statistics

- But present a trade-off: the recursion from $x_p$ all the way to the origin at $x_1$ can be very long, which will be performed during each training update

- This can slow down optimization, and can naturally cause 'exploding/vanishing' gradient problems when using gradient descent

- Lots of tricks - from careful selection of activation function, to carefully engineered optimization algorithms, to the construction of more complicated activations like the Long Short Term Memory function

- However while *training* is tough, evaluation of trained models are fast (just plug in the most recent value of the series)

#### Example 1.   Time series prediction using an RNN

- Below we load in a time series dataset we will perform prediction on

- Note: we will first ["contrast normalize"](https://jermwatt.github.io/control-notes/posts/zca_sphereing/ZCA_Sphereing.html) this series - that is we mean center and re-scale by its standard deviation

```python
## This code cell will not be shown in the HTML version of this notebook
# load in driver
csvname = datapath + 'normalized_apple_prices.csv'
x_series = np.loadtxt(csvname,delimiter = ',')[np.newaxis,:]

# standard normalize input series
normalizer,inverse_normalizer = reclib.time_series_lib.normalizers.standard(x_series)
x_series = normalizer(x_series)

# Plot the standard normalized series
reclib.time_series_lib.fixed_order_plotters.plot_sequences(x_series)
```

    <IPython.core.display.Javascript object>

- We perform training / validate our model precisely as we did previously with the autoregressive process

- Here however, we will be able to learn quite a good model using only a single input at-a-time from the true series

- Plotted below is the training fit (in blue) and validation fit (in yellow) resulting from a run of [derivative free/zero order optimization](https://jermwatt.github.io/machine_learning_refined/notes/2_Zero_order_methods/2_6_Coordinate.html) employing the Least Squares cost function

```python
## This code cell will not be shown in the HTML version of this notebook
def step(h_t_prev,x_t_prev,w):
    return w[0] + w[1]*h_t_prev + w[2]*x_t_prev

# exponential average function
def model(x,w):
    # set initial conditions of h to values of x
    h = [x[:,0]]

    # range over x and create h
    for t in range(0,np.size(x)-1):
        # get current point and prior hidden state
        h_t_prev = h[-1]
        x_t_prev = x[:,t]

        # make next element and store
        h_t = step(h_t_prev,x_t_prev,w)
        h.append(h_t)
    return np.array(h).T

# an implementation of the least squares cost function for linear regression
def least_squares(w,x,y):
    # compute cost over batch
    cost = np.sum((model(x,w) - y)**2)
    return cost/float(np.size(y))
```

```python
# This code cell will not be shown in the HTML version of this notebook
# initialize with input/output data
mylib1 = reclib.time_series_lib.rnn_lib.super_setup.Setup(x_series,x_series)

# perform preprocessing step(s) - especially input normalization
mylib1.preprocessing_steps(normalizer = 'standard')

# split into training and validation sets
mylib1.make_train_val_split(train_portion = 0.6)

# choose cost
mylib1.choose_cost(name = 'least_squares',model = model,step = step)

# fit an optimization
w = 0.1*np.random.randn(3,1)
mylib1.fit(max_its = 50,alpha_choice = 'diminishing',optimizer = 'zero_order',w_init = w,verbose = False)

# show cost function history
mylib1.show_histories()
```

    <IPython.core.display.Javascript object>

```python
## This code cell will not be shown in the HTML version of this notebook
# Plot the standard normalized series and its training fit
reclib.time_series_lib.variable_order_plotters.plot_train_val_sequences(mylib1)
```

    <IPython.core.display.Javascript object>

- As with the autoregressive approach, here we too can use our trained model as a generator

```python
## This code cell will not be shown in the HTML version of this notebook
# Plot the standard normalized series
reclib.time_series_lib.variable_order_plotters.plot_train_gen_sequences(mylib1)
```

    <IPython.core.display.Javascript object>


# メモ
数学を Colab で学ぶプロジェクト

いま数学IIの三角関数 trigonometry

# 一般角

**一般角**

平面上で、点 $\mathrm O$ を中心として半直線 $\mathrm{OP}$ を回転させるとき、

この半直線 $\mathrm{OP}$ を動径、半直線 $\mathrm{OX}$ を始線という。

時計の針の回転と逆の向きに図った角度を正の角という。

$\mathrm{OX}$ と $\mathrm{OP}$ のなす角の一つを $\alpha$ とすると、

$\theta = \alpha + 360 \times n$
を動径 $\mathrm{OP}$ の一般角、という。

象限の角 と言う言い方をする。 座標軸に重なる時はどの象限の角でもないとする。

基本例題 108 p.165

次の角の動径を図示せよ。 また、それぞれ第何象限の角か。

* 650°
* 800°
* -630°
* -1280°

```
# sympy の実験 => α β を表す。 * => \times にする。
from  sympy import *
init_printing(latex_printer=lambda *args, **kwargs: latex(*args, mul_symbol='times', **kwargs))
alpha, beta, theta = sympy.symbols('alpha beta theta')
n = symbols('n')
display(Eq(theta, alpha + 360 * n))
```

```
# 参考 角度の記号
%%latex
\angle A = 360^\circ
```

\angle A = 360^\circ

---
弧度法

半径 $r$ の円で、半径に等しい長さの弧 $\mathrm{AB}$ に対する中心角の大きさは、半径 $r$ に関係なく一定である。 この角の大きさを 1 ラジアン (1 弧度) と言い、ラジアンを単位をする角の表し方を弧度法と言う。 直角を 90°とする角度の表し方を度数法と言う。

* 1 ラジアン = $\displaystyle \left ( \frac{180}{\pi} \right )^\circ$
* 180°= $\pi$ ラジアン

* 弧の長さ $\quad l = r \theta$
* 面積 $\quad S = \frac 1 2 r^2 \theta = \frac 1 2 rl$

基本例題 109 p.167

次の角を、度数は弧度に、弧度は度数に、それぞれ書き直せ。
* 72°
* -320°
* $\frac{4}{15} \pi$
* $\frac{-13}{4}\pi$

半径 4、中心角 150°の扇形の弧の長さと面積を求めよ。

```
from sympy import *
display(Rational(72,180)*pi)
display(Rational(-320,180)*pi)
display(Rational(4,15)*180)
display(Rational(-13,4)*180)

```

$\displaystyle \frac{2 \pi}{5}$

$\displaystyle - \frac{16 \pi}{9}$

$\displaystyle 48$

$\displaystyle -585$

```
# 弧の長さは r theta なので
print("{} π".format(4 * 150/180))
# 面積は 1/2 r l なので
print("{} π".format(0.5*4* 4*150/180))
```

    3.3333333333333335 π
    6.666666666666667 π

---
一般角の三角関数の定義

動線 $\mathrm{OP}$ を表す一般角を $\theta$ とする。 $\mathrm{OP} = r,\; \mathrm {P}(x,y)$ とする時、

* 正弦 $\quad \sin \theta = \displaystyle \frac y r $
* 余弦 $\quad \cos \theta = \displaystyle \frac x r $
* 正接 $\quad \tan \theta = \displaystyle \frac y x $

ただし、$\theta = \displaystyle \frac \pi 2 + n\pi$ ($n$ は整数) に対しては $\tan \theta$ の値を定義しない。

三角関数の値域

$-1 \leq \sin \theta \leq 1,\; -1 \leq \cos \theta \leq 1,\; \tan \theta$ は実数全体

---
三角関数の相互関係
* $\tan \theta = \displaystyle \frac {\sin \theta}{\cos \theta}$
* $\sin^2 \theta + \cos^2 \theta = 1$
* $1 + \displaystyle \tan^2 \theta = \frac {1}{\cos^2 \theta}$

---
基本例題 110 p.170

$\theta$ が次の値の時、$\sin \theta, \cos \theta , \tan \theta$ の値を求めよ。

* $\displaystyle \frac{23}{6} \pi$

* $\displaystyle - \frac 5 4 \pi$

```
from sympy.abc import *
from sympy import *
x = Rational(23,6) * pi
display(x)
display(sin(x))
display(cos(x))
display(tan(x))
y = -Rational(5,4) * pi
display(y)
display(sin(y))
display(cos(y))
display(tan(y))

```

$\displaystyle \frac{23 \pi}{6}$

$\displaystyle - \frac{1}{2}$

$\displaystyle \frac{\sqrt{3}}{2}$

$\displaystyle - \frac{\sqrt{3}}{3}$

$\displaystyle - \frac{5 \pi}{4}$

$\displaystyle \frac{\sqrt{2}}{2}$

$\displaystyle - \frac{\sqrt{2}}{2}$

$\displaystyle -1$

---
基本例題 111 p.171 (1)

$\displaystyle \frac 3 2 \pi \lt \theta \lt 2 \pi$ とする。 $\cos \theta = \displaystyle \frac 5 {13}$ のとき、$\sin \theta$ と $\tan \theta$ を求める。

```
from sympy.abc import *
from sympy import *
# 与件 cos(theta) = Rational(5,13)
# 公式から sin(theta)**2 = 1 - cos(theta)**2
display(1 - Rational(5,13)**2)
display(sqrt(1 - Rational(5,13)**2)) # sin(theta)
# 定義より tan(theta) = sin(theta) / cos(theta)
display(Rational(12,13)/Rational(5,13))

```

$\displaystyle \frac{144}{169}$

$\displaystyle \frac{12}{13}$

$\displaystyle \frac{12}{5}$

---
基本例題 111 p.171 (2)

$\displaystyle \pi \lt \theta \lt \frac 3 2 \pi$ とする。 $\tan \theta = 7$ のとき、$\sin \theta$ と $\cos \theta$ を求める。

```
from sympy.abc import *
from sympy import *
# 与件 cos(theta) = 7
# 公式から cos(theta)**2 = 1/(1 + tan(theta)**2)
# 公式から sin(theta) = tan(theta) * cos(theta)
display(S(1) / (1 + 7**2)) # S(1) は 1 の実数表示 Rational(1,1)
display(sqrt(S(1) / (1 + 7**2))) # cos(theta)
display(S(7) * (-sqrt(2)/10))
```

$\displaystyle \frac{1}{50}$

$\displaystyle \frac{\sqrt{2}}{10}$

$\displaystyle - \frac{7 \sqrt{2}}{10}$

---
基本例題 112 p.172 (1)

* 等式 $\;\; \displaystyle \frac {\cos \theta}{1 + \sin \theta} + \tan \theta = \frac 1 {\cos \theta}$ を証明する。

左辺

$\qquad = \displaystyle \frac {\cos \theta}{1 + \sin \theta} + \frac {\sin \theta}{\cos \theta}$

$\qquad = \displaystyle \frac {\cos^2\theta}{(1 + \sin \theta)\cos\theta} + \frac {(1 + \sin \theta)\sin\theta}{(1 + \sin \theta)\cos \theta}$

$\qquad = \displaystyle \frac {\cos^2\theta + \sin^2\theta +\sin\theta}{(1 + \sin \theta)\cos\theta}$

$\qquad = \displaystyle \frac {1 + \sin \theta}{(1 + \sin \theta)\cos\theta}$

$\qquad = \displaystyle \frac {1}{\cos\theta}$

---
基本例題 112 p.172 (2)

* $\cos^2\theta + \sin\theta -\tan\theta (1 - \sin\theta) \cos\theta$ を計算する。

与式

$\qquad = \displaystyle \cos^2\theta + \sin\theta -\frac {\sin\theta}{\cos\theta} (1 - \sin\theta) \cos\theta$

$\qquad = \displaystyle \cos^2\theta + \sin\theta - \sin\theta + \sin^2\theta$

$\qquad = \displaystyle \cos^2\theta + \sin^2\theta$

$\qquad = 1$

---
基本例題 113 p.173

$\sin\theta +\cos\theta = \displaystyle\frac{\sqrt 3}{2}\;\; (\frac \pi 2 \lt \theta \lt \pi)$ の時、次の式の値を求める。

* $\sin\theta\cos\theta$
* $\sin^3\theta+\cos^3\theta$
* $\cos^3\theta-\sin^3\theta$

$\sin\theta +\cos\theta = \displaystyle\frac{\sqrt 3}{2}$

の両辺を 2 乗すると

$\sin^2\theta +\cos^2\theta + 2\sin\theta\cos\theta = \displaystyle\frac{3}{4}$

$1 + 2\sin\theta\cos\theta = \displaystyle\frac{3}{4}$

$2\sin\theta\cos\theta = \displaystyle\frac{3}{4} - 1$

$2\sin\theta\cos\theta = \displaystyle\frac{-1}{4}$

$\sin\theta\cos\theta = \displaystyle\frac{-1}{8}$

(2)

$\sin^3\theta+\cos^3\theta = (\sin\theta + \cos\theta)(\sin^2\theta - \sin\theta\cos\theta + \cos^2\theta)$

$\qquad =  \displaystyle\frac{\sqrt 3}{2} (\sin^2\theta - \frac{-1}{8} + \cos^2\theta)$

$\qquad =  \displaystyle\frac{\sqrt 3}{2} (1 + \frac{1}{8})$

$\qquad =  \displaystyle\frac{9\sqrt 3}{16}$

(3)

$\cos^3\theta-\sin^3\theta = (\cos\theta - \sin\theta)(\sin^2\theta + \sin\theta\cos\theta + \cos^2\theta)$

$\qquad = (\cos\theta - \sin\theta)(1 + \sin\theta\cos\theta)$

$\qquad = (\cos\theta - \sin\theta)(1 + \displaystyle\frac{-1}{8})$

$\qquad = \displaystyle\frac{7}{8}(\cos\theta - \sin\theta)$

ところで

$(\cos\theta - \sin\theta)^2 = 1 - 2\sin\theta\cos\theta$

$\qquad = 1 - 2 \displaystyle\frac{-1}{8} = \frac 5 4$

与えられた象限では $\cos\theta - \sin\theta \lt 0$ なので

$\cos\theta - \sin\theta = - \displaystyle\frac{\sqrt 5}{2}$

よって与式 $\quad = - \displaystyle\frac 7 8 \times \frac {\sqrt 5} 2 = - \frac{7\sqrt 5}{16}$

---
三角関数の性質 p.174

$\theta + 2n\pi$ の三角関数
* $\sin(\theta + 2n\pi) = \sin\theta$
* $\cos(\theta + 2n\pi) = \cos\theta$
* $\tan(\theta + 2n\pi) = \tan\theta$

$-\theta$ の三角関数
* $\sin(-\theta) = - \sin\theta$
* $\cos(-\theta) = \cos\theta$
* $\tan(-\theta) = - \tan\theta$

$-\theta$ の三角関数
* $\sin(-\theta) = - \sin\theta$
* $\cos(-\theta) = \cos\theta$
* $\tan(-\theta) = - \tan\theta$


# Example of an extraction using tikhonov regularisation

TODO: Explain waht tikhonov regularisation is and why we use it?

## Imports

```python
import os
os.environ["OMP_NUM_THREADS"] = "1"

# Import standard packages.
import numpy as np
from astropy.io import fits

# Imports from the extraction.
from SOSS.dms.soss_engine import ExtractionEngine
from SOSS.extract.engine_legacy import ThroughputSOSS, WebbKer

# Imports for plots.
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm
```

### Matplotlib defaults

```python
%matplotlib inline
```

```python
plt.rc('figure', figsize=(13,8))
plt.rc('font', size=16)
plt.rc('image', cmap='inferno')
plt.rc('lines', lw=2)
```

## Read some ref files

```python
# List of orders to consider in the extraction.
soss_orders = [1, 2]

# WaveMap TODO try using DataModels.
wave_maps = []
wave_maps.append(fits.getdata("../extract/Ref_files/wavelengths_m1.fits"))
wave_maps.append(fits.getdata("../extract/Ref_files/wavelengths_m2.fits"))

# SpecProfile
spat_pros = []
spat_pros.append(fits.getdata("../extract/Ref_files/spat_profile_m1.fits").squeeze())
spat_pros.append(fits.getdata("../extract/Ref_files/spat_profile_m2.fits").squeeze())

# Convert data from fits files to float (fits precision is 1e-8)
wave_maps = [wv.astype('float64') for wv in wave_maps]
spat_pros = [p_ord.astype('float64') for p_ord in spat_pros]

# SpecTrace
thrpt_list = [ThroughputSOSS(order) for order in soss_orders]

# SpecKernel
ker_list = [WebbKer(wv_map) for wv_map in wave_maps]

# Put all inputs from reference files in a list
ref_files = [wave_maps, spat_pros, thrpt_list, ker_list]
```

## Load simulation

```python
# Import custom function to read toy simulation.
from SOSS.Fake_data.simu_utils import load_simu

# Load a simulation
simu = load_simu("../Fake_data/phoenix_teff_02300_scale_1.0e+02.fits")
```

```python
# The science array.
scidata = simu["data"]

# Build an associated error array.
bkgd_noise = 20.  # In counts?
scierr = np.sqrt(scidata + bkgd_noise**2)
```

## Extraction

### Extraction Parameters
(Example usage with few inputs parameters.)

```python
params = {}

# Wavelength extraction grid oversampling.
params["n_os"] = 5

# Threshold on the spatial profile.
params["threshold"] = 1e-4
```

### Initiate extraction object
(This needs to be done only once unless the oversampling (`n_os`) changes.)

```python
extract = ExtractionEngine(*ref_files, **params)
```

### Find the best tikhonov factor
This takes some time, so it's better to do it once if the exposures are part of a time series observation, i.e. observations of the same object at similar SNR.

```python
# # Determine which factors to tests.
# factors = np.logspace(-25, -12, 14)

# # Tests all these factors.
# tests = extract.get_tikho_tests(factors, data=scidata, error=scierr)
```

```python
# # Find the best factor.
# best_fac = extract.best_tikho_factor(tests=tests, i_plot=True)
# print(best_fac)
```

```python
# # Refine the grid (span 4 orders of magnitude).
# best_fac = np.log10(best_fac)
# factors = np.logspace(best_fac-2, best_fac+2, 20)

# # Test these new factors.
# tests = extract.get_tikho_tests(factors, data=scidata, error=scierr)
```

```python
# # Find the best factor.
# best_fac = extract.best_tikho_factor(tests=tests, i_plot=True)
# print(best_fac)
```

```python
best_fac = 6.3580702764558024e-24
```

### Extract the oversampled spectrum $f_k$
Can be done in a loop for a timeseries and/or iteratively for different estimates of the reference files.

```python
# Extract the spectrum.
f_k = extract.extract(data=scidata, error=scierr, tikhonov=True, factor=best_fac)
```

```python
# Plot the extracted spectrum.
plt.plot(extract.wave_grid, f_k)

plt.xlabel("Wavelength [$\mu m$]")
plt.ylabel("Oversampled Spectrum $f_k$ [energy$\cdot s^{-1} \cdot \mu m^{-1}$]")
# For now, arbitrairy units, but it should be the flux that hits the detector,
# so energy/time/wavelength

plt.tight_layout()
plt.show()
```

### Bin to pixel native sampling
To get a result comparable to typical extraction methods, we need to integrate the oversampled spectrum ($f_k$) to a grid representative of the native pixel sampling (for each order). This integration is done according to the equation
\begin{equation}
\mathrm{bin}_{i} = \int_{\lambda_{ni}^-}^{\lambda_{ni}^+} T_n(\lambda)\tilde{f}_n(\lambda)\lambda d\lambda \, ,
\end{equation}
where $n$ is a given order, $T_n$ is the throughput of the order and $\tilde{f}_n$ is the underlying flux convolved to the order $n$ resolution. The result of this integral will be in fake counts (it is not directly the sum of the counts so that's why I call it fake).

One could directly extract the integrated flux by setting the throughput to $T_n(\lambda) = 1$ (see second example). The result would then be in flux units instead of counts.

#### Bin in counts

```python
# Save the output in a list for different orders.

flux_bin = []  # Integrated flux.
wave_bin = []  # Wavelength grid.

for i_ord in range(extract.n_orders):

    # Integrate.
    wave_bin_, flux_bin_ = extract.bin_to_pixel(spectrum=f_k, i_order=i_ord)

    # Save.
    wave_bin.append(wave_bin_)
    flux_bin.append(flux_bin_)
```

```python
fig, ax = plt.subplots(2, 1, sharex=True, figsize=(12,6))

for i_ord in range(extract.n_orders):

    label = extract.orders[i_ord]
    ax[i_ord].plot(wave_bin[i_ord], flux_bin[i_ord], label=label)

ax[0].set_ylabel("Extracted signal [counts]")

ax[1].set_xlabel("Wavelength [$\mu m$]")
ax[1].set_ylabel("Extracted signal [counts]")

plt.tight_layout()
plt.show()
```

#### Bin in flux units

```python
# Set throughput to 1
def throughput(x):
    return np.ones_like(x)

flux_bin = []  # Integrated flux.
wave_bin = []  # Wavelength grid.

for i_ord in range(extract.n_orders):

    # Integrate.
    wave_bin_, flux_bin_ = extract.bin_to_pixel(spectrum=f_k, i_order=i_ord, throughput=throughput)

    # Save.
    wave_bin.append(wave_bin_)
    flux_bin.append(flux_bin_)
```

```python
for i_ord in range(extract.n_orders):

    label = extract.orders[i_ord]
    plt.plot(wave_bin[i_ord], flux_bin[i_ord], label=label)

plt.ylabel(r"Convolved flux $\tilde{f_k}$ [energy$\cdot s^{-1} \cdot \mu m^{-1}$]")
plt.xlabel("Wavelength [$\mu m$]")

plt.tight_layout()
plt.legend(title="Order")
plt.show()
```

## Quality estimate

### Rebuild the detector

```python
model_image = extract.rebuild(f_k)
```

```python
plt.subplot(111, aspect=2)

plt.pcolormesh((model_image - scidata)/scierr, vmin=-3, vmax=3)

plt.colorbar(label="Error relative to noise", orientation='horizontal', aspect=40)

plt.tight_layout()
plt.show()
```

We can see that we are very close to the photon noise limit in this case. There are some small structures in the 2nd order in the overlap region, but the extracted spectrum is dominated by the 1st order in this wavelength region anyway, due to the higher throughput.

```python

```