# Elementary greenhouse models

This notebook is part of [The Climate Laboratory](https://brian-rose.github.io/ClimateLaboratoryBook) by [Brian E. J. Rose](http://www.atmos.albany.edu/facstaff/brose/index.html), University at Albany.

____________

## 1. A single layer atmosphere
____________

We will make our first attempt at quantifying the greenhouse effect in the simplest possible greenhouse model: a single layer of atmosphere that is able to absorb and emit longwave radiation.

### Assumptions

- Atmosphere is a single layer of air at temperature $T_a$
- Atmosphere is **completely transparent to shortwave** solar radiation.
- The **surface** absorbs shortwave radiation $(1-\alpha) Q$
- Atmosphere is **completely opaque to infrared** radiation
- Both surface and atmosphere emit radiation as **blackbodies** ($\sigma T_s^4, \sigma T_a^4$)
- Atmosphere radiates **equally up and down** ($\sigma T_a^4$)
- There are no other heat transfer mechanisms

We can now use the concept of energy balance to ask what the temperature need to be in order to balance the energy budgets at the surface and the atmosphere, i.e. the **radiative equilibrium temperatures**.

### Energy balance at the surface

$$
\begin{align*}
\text{energy in} &= \text{energy out} \\
(1-\alpha) Q + \sigma T_a^4 &= \sigma T_s^4 \\
\end{align*}
$$

The presence of the atmosphere above means there is an additional source term: downwelling infrared radiation from the atmosphere.

We call this the **back radiation**.

### Energy balance for the atmosphere

$$
\begin{align}
\text{energy in} &= \text{energy out} \\
\sigma T_s^4 &= A\uparrow + A\downarrow = 2 \sigma T_a^4 \\
\end{align}
$$

which means that
$$ T_s = 2^\frac{1}{4} T_a \approx 1.2 T_a $$

So we have just determined that, in order to have a purely **radiative equilibrium**, we must have $T_s > T_a$.

*The surface must be warmer than the atmosphere.*

### Solve for the radiative equilibrium surface temperature

Now plug this into the surface equation to find

$$ \frac{1}{2} \sigma T_s^4 = (1-\alpha) Q  $$

and use the definition of the emission temperature $T_e$ to write

$$ (1-\alpha) Q = \sigma T_e^4 $$

*In fact, in this model, $T_e$ is identical to the atmospheric temperature $T_a$, since all the OLR originates from this layer.*

Solve for the surface temperature:

$$ T_s = 2^\frac{1}{4} T_e $$

Putting in observed numbers, $T_e = 255$ K gives a surface temperature of

$$T_s = 303 ~\text{K}$$

This model is one small step closer to reality: surface is warmer than atmosphere, emissions to space generated in the atmosphere, atmosphere heated from below and helping to keep surface warm.

### Why does this model overpredict the surface temperature

Our model now overpredicts the surface temperature by about 15ºC (303 K versus the observed 288 K).

Ideas about why?

Basically we just need to read our **list of assumptions** above and realize that none of them are very good approximations:

- Atmosphere absorbs some solar radiation.
- Atmosphere is NOT a perfect absorber of longwave radiation
- Absorption and emission varies strongly with wavelength *(atmosphere does not behave like a blackbody)*.
- Emissions are not determined by a single temperature $T_a$ but by the detailed *vertical profile* of air temperture.
- Energy is redistributed in the vertical by a variety of dynamical transport mechanisms (e.g. convection and boundary layer turbulence).

____________

## 2. Introducing the two-layer grey gas model
____________

Let's generalize the above model just a little bit to build a slighly more realistic model of longwave radiative transfer.

We will address two shortcomings of our single-layer model:
1. No vertical structure
2. 100% longwave opacity

Relaxing these two assumptions gives us what turns out to be a very useful prototype model for **understanding how the greenhouse effect works**.

### Assumptions

- The atmosphere is **transparent to shortwave radiation** (still)
- Divide the atmosphere up into **two layers of equal mass** (the dividing line is thus at 500 hPa pressure level)
- Each layer **absorbs only a fraction $\epsilon$** of whatever longwave radiation is incident upon it.
- We will call the fraction $\epsilon$ the **absorptivity** of the layer.
- Assume $\epsilon$ is the same in each layer

This is called the **grey gas** model, where **grey** here means the emission and absorption have **no spectral dependence** (same at every wavelength).

We can think of this model informally as a "leaky greenhouse".

Note that the assumption that $\epsilon$ is the same in each layer is appropriate if the absorption is actually carried out by a gas that is **well-mixed** in the atmosphere.

Out of our two most important absorbers:

- CO$_2$ is well mixed
- H$_2$O is not (mostly confined to lower troposphere due to strong temperature dependence of the saturation vapor pressure).

But we will ignore this aspect of reality for now.

### Kirchoff's Law

In order to build our model, we need to introduce one additional piece of physics known as **Kirchoff's Law**:

$$ \text{absorptivity} = \text{emissivity} $$

So **if a layer of atmosphere at temperature $T$ absorbs a fraction $\epsilon$** of incident longwave radiation, it must **emit**

$$ \epsilon ~\sigma ~T^4 $$

both up and down.

### A sketch of the radiative fluxes in the 2-layer atmosphere

- Surface temperature is $T_s$
- Atm. temperatures are $T_0, T_1$ where $T_0$ is closest to the surface.
- absorptivity of atm layers is $\epsilon$
- Surface emission is $\sigma T_s^4$
- Atm emission is $\epsilon \sigma T_0^4, \epsilon \sigma T_1^4$ (up and down)
- Absorptivity = emissivity for atmospheric layers
- a fraction $(1-\epsilon)$ of the longwave beam is **transmitted** through each layer

____________
## 3. Tracing the upwelling beam of longwave radiation
____________

Let's think about the upwelling beam of longwave radiation, which we denote $U$.

### Surface to layer 0

We start at the surface. The upward flux **from the surface to layer 0** is

$$U_0 = \sigma T_s^4$$

(just the emission from the suface).

### Layer 0 to layer 1

Now **following this beam upward**, we first recognize that a fraction $\epsilon$ of this beam is **absorbed** in layer 0.

The upward flux from layer 0 to layer 1 consists of the sum of two parts:

1. The **transmitted part** of whatever is incident from below (i.e. the part that is **not absorbed**)
2. **New upward emissions** from layer 0

We can write this upward flux from layer 0 to layer 1 as:

$$U_1 = (1-\epsilon) \sigma T_s^4 + \epsilon \sigma T_0^4$$

### Beyond layer 1

Continuing to follow the same beam, we follow the same logic! A fraction $\epsilon$ of $U_1$ is absorbed in layer 1, and therefore the transmitted part is $(1-\epsilon) U_1$.

Including new emissions from layer 1, the upwelling flux above layer 1 is

$$U_2 = (1-\epsilon) U_1 + \epsilon \sigma T_1^4$$

### Outgoing Longwave Radiation

Since there is **no more atmosphere above layer 1**, this upwelling beam is our OLR for this model:

$$\text{OLR} = U_2 = (1-\epsilon) U_1 + \epsilon \sigma T_1^4$$

which, plugging in the above expression for $U_1$, works out to

$$\text{OLR} = (1-\epsilon)^2 \sigma T_s^4 + \epsilon(1-\epsilon)\sigma T_0^4 + \epsilon \sigma T_1^4$$

Here the three terms represent **contributions to the total OLR** that **originate from each of the three levels**

### Limits of large and small absorptivity/emissivity

Think about the following two questions:

- What happens to this expression if $\epsilon=1$? *What does this represent physically?*
- What about $\epsilon=0$?

By allowing the atmosphere to partially absorb emissions from other levels, we now see that the Outgoing Longwave Radiation to space **includes emissions from every level** - and therefore **affected by temperature at every level**!

____________

## 4. Tuning the grey gas model to observations
____________

In building our new model we have introduced exactly **one parameter**, the absorptivity $\epsilon$. We need to choose a value for $\epsilon$.

We will tune our model so that it **reproduces the observed global mean OLR** given **observed global mean temperatures**.

### Global mean air temperature observations

To get appropriate temperatures for $T_s, T_0, T_1$, let's revisit the global, annual mean lapse rate plot from NCEP Reanalysis data we first encountered in the [Radiation notes](https://brian-rose.github.io/ClimateLaboratoryBook/courseware/radiation.html).

```python
#  This code is used just to create the skew-T plot of global, annual mean air temperature
%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
import xarray as xr
from metpy.plots import SkewT
ncep_url = "http://www.esrl.noaa.gov/psd/thredds/dodsC/Datasets/ncep.reanalysis.derived/"
ncep_air = xr.open_dataset( ncep_url + "pressure/air.mon.1981-2010.ltm.nc", use_cftime=True)
#  Take global, annual average and convert to Kelvin
coslat = np.cos(np.deg2rad(ncep_air.lat))
weight = coslat / coslat.mean(dim='lat')
Tglobal = (ncep_air.air * weight).mean(dim=('lat','lon','time'))
fig = plt.figure(figsize=(9, 9))
skew = SkewT(fig, rotation=30)
skew.plot(Tglobal.level, Tglobal, color='black', linestyle='-', linewidth=2, label='Observations')
skew.ax.set_ylim(1050, 10)
skew.ax.set_xlim(-75, 45)
# Add the relevant special lines
skew.plot_dry_adiabats(linewidth=0.5)
skew.plot_moist_adiabats(linewidth=0.5)
#skew.plot_mixing_lines()
skew.ax.legend()
skew.ax.set_title('Global, annual mean sounding from NCEP Reanalysis', fontsize = 16);
```

### Target temperatures for our model tuning

First, we set
$$T_s = 288 \text{ K}  $$

From the lapse rate plot, an average temperature for the layer between 1000 and 500 hPa is

$$ T_0 = 275 \text{ K}$$

Defining an average temperature for the layer between 500 and 0 hPa is more ambiguous because of the lapse rate reversal at the tropopause. We will choose

$$ T_1 = 230 \text{ K}$$

From the graph, this is approximately the observed global mean temperature at 275 hPa or about 10 km.

### OLR

From the [observed global energy budget](https://brian-rose.github.io/ClimateLaboratoryBook/courseware/models-budgets-fun.html#the-observed-global-energy-budget) we set

$$\text{OLR} = 238.5 \text{ W m}^{-2}$$

### Solving for $\epsilon$

We wrote down the expression for OLR as a function of temperatures and absorptivity in our model above.

All we need to do is plug the observed values into the above expression for OLR, and solve for $\epsilon$.

It is a **quadratic equation** for the unknown $\epsilon$. We could work out the exact solution using the quadratic formula.

But let's do it **graphically**, using Python!

### Exercise: graphical solution to find the best fit value of $\epsilon$

The OLR formula for the leaky greenhouse that we derived above is

$$\text{OLR} = (1-\epsilon)^2 \sigma T_s^4 + \epsilon(1-\epsilon)\sigma T_0^4 + \epsilon \sigma T_1^4$$

Do the following:

- Write a Python function that implements this formula
- The function should accept **four input parameters**:
    - The three temperatures $T_s, T_0, T_1$
    - The emissivity $\epsilon$
- Using this function, make a **graph of OLR vs. $\epsilon$** for the observed temperature values $T_s = 288, T_0 = 275, T_1 = 230$
- For the graph, $\epsilon$ should range between 0 and 1.
- From your graph, find the approximate value of $\epsilon$ that gives $OLR = 238.5$

```python

```

Note if you solve the quadratic equation algebraically you will get two solutions:

- $\epsilon \approx 0.586$
- $\epsilon \approx 3.93$

(for details, see [the advanced notes here](https://brian-rose.github.io/ClimateLaboratoryBook/courseware/sympy-greenhouse.html))

**Why is the second solution not physically meaningful?**

Hopefully your graph shows that $\epsilon = 0.586$ gives the correct value of OLR.

This is the absorptivity that guarantees that **our model reproduces the observed OLR given the observed temperatures**.

____________

## 5. Level of emission
____________

### Contributions from each level to the outgoing radiation

Now that we have tuned up our model, we can see exactly how strongly each level contributes to the OLR.

The three components of the OLR are

$$
\begin{align*}
\text{OLR}_s &= (1-\epsilon)^2 \sigma T_s^4  \\
\text{OLR}_0 &= \epsilon(1-\epsilon)\sigma T_0^4 \\
\text{OLR}_1 &= \epsilon \sigma T_1^4
\end{align*}
$$

which of course add up to the total OLR we wrote down above.

### Exercise: calculate contributions to OLR

**Write some simple Python code to calculate each term in the OLR using the observed temperatures and the tuned value $\epsilon = 0.586$. Fill out the list below using your calculated numbers.**

```python

```

**Contributions to the OLR originating from each level, in W/m2:**

- Surface:
- Level 0:
- Level 1:

```python
# now sum up the numbers to verify you get something very close to 238.5

```

Notice that the largest single contribution is coming from the top layer.

*This is in spite of the fact that the emissions from this layer are weak, because it is so cold.*

###  Changing the level of emission by adding absorbers

Adding some **extra greenhouse absorbers** will mean that a **greater fraction** of incident longwave radiation is **absorbed in each layer**.

Thus **$\epsilon$ must increase** as we add greenhouse gases.

Suppose we have $\epsilon$ initially, and the absorptivity increases to $\epsilon_2 = \epsilon + \Delta \epsilon$.

Suppose further that this increase happens **abruptly** so that there is no time for the temperatures to respond to this change. **We hold the temperatures fixed** in the column and ask how the radiative fluxes change.

**Question: Do you expect the OLR to increase or decrease?**

### Calculating the change in level of emission

Let's use our two-layer leaky greenhouse model to investigate the answer.

The components of the OLR before the perturbation are

$$
\begin{align*}
\text{OLR}_s &= (1-\epsilon)^2 \sigma T_s^4  \\
\text{OLR}_0 &= \epsilon(1-\epsilon)\sigma T_0^4 \\
\text{OLR}_1 &= \epsilon \sigma T_1^4
\end{align*}
$$

and after the perturbation we have

$$
\begin{align*}
\text{OLR}_s &= (1-\epsilon - \Delta \epsilon)^2 \sigma T_s^4  \\
\text{OLR}_0 &= (\epsilon + \Delta \epsilon)(1-\epsilon - \Delta \epsilon)\sigma T_0^4 \\
\text{OLR}_1 &= (\epsilon + \Delta \epsilon) \sigma T_1^4
\end{align*}
$$

Let's subtract off the original components to get the contributions to the **change in OLR** from each layer:

$$
\begin{align*}
\Delta \text{OLR}_s &= \left[(1-\epsilon - \Delta \epsilon)^2 - (1-\epsilon)^2\right]\sigma T_s^4  \\
\Delta \text{OLR}_0 &= \left[(\epsilon + \Delta \epsilon)(1-\epsilon - \Delta \epsilon) - \epsilon(1-\epsilon)  \right] \sigma T_0^4 \\
\Delta \text{OLR}_1 &= \left[(\epsilon + \Delta \epsilon) - \epsilon \right] \sigma T_1^4
\end{align*}
$$

Now expand this out, but to make things easier to deal with, neglect term in $\Delta \epsilon^2$ (very small - we will be considering changes of less than 10% in $\epsilon$):

$$
\begin{align*}
\Delta \text{OLR}_s &\approx (\Delta \epsilon) \left[ -2(1-\epsilon) \right] \sigma T_s^4  \\
\Delta \text{OLR}_0 &\approx (\Delta \epsilon) (1 - 2 \epsilon) \sigma T_0^4 \\
\Delta \text{OLR}_1 &\approx (\Delta \epsilon) \sigma T_1^4
\end{align*}
$$

Now look at the **sign** of each term. Recall that $0 < \epsilon < 1$. **Which terms in the OLR go up and which go down?**

**THIS IS VERY IMPORTANT, SO STOP AND THINK ABOUT IT.**

The contribution from the **surface** must **decrease**, while the contribution from the **top layer** must **increase**.

**When we add absorbers, the average level of emission goes up!**

____________

## 6. Radiative forcing in the 2-layer grey gas model
____________

### Definition of Radiative Forcing

We now define a very important quantity:

**"Radiative forcing" is the change in total radiative flux at TOA after adding absorbers**

In this model, **only the longwave flux can change**, so we calculate the radiative forcing as

$$ R = - \Delta \text{OLR} $$

(with the minus sign so that $R$ is **positive when the climate system is gaining extra energy**).

### Connection between radiative forcing and level of emission

We just worked out that whenever we **add some extra absorbers**, the **emissions to space** (on average) will originate from **higher levels** in the atmosphere.

What does this mean for OLR? Will it increase or decrease?

To get the answer, we just have to sum up the three contributions we wrote above:

$$
\begin{align*}
R &= -\Delta \text{OLR}_s - \Delta \text{OLR}_0 - \Delta \text{OLR}_1  \\
  &= -\Delta \epsilon \left[  -2(1-\epsilon) \sigma T_s^4 + (1 - 2 \epsilon) \sigma T_0^4  + \sigma T_1^4 \right]
\end{align*}
$$

Is this a positive or negative number? The key point is this:

**It depends on the temperatures, i.e. on the lapse rate.**

### Greenhouse effect for an isothermal atmosphere

Stop and think about this question:

If the **surface and atmosphere are all at the same temperature**, does the OLR go up or down when $\epsilon$ increases (i.e. we add more absorbers)?

Understanding this question is key to understanding how the greenhouse effect works.

#### Let's solve the isothermal case

We will just set $T_s = T_0 = T_1$ in the above expression for the radiative forcing.

What do you get?

```python

```

#### The answer is $R=0$

For an isothermal atmosphere, there is **no change** in OLR when we add extra greenhouse absorbers. Hence, no radiative forcing and no greenhouse effect.

Why?

The level of emission still must go up. But since the temperature at the upper level is the **same** as everywhere else, the emissions are exactly the same.

### The radiative forcing (change in OLR) depends on the lapse rate!

For a more realistic example of radiative forcing due to an increase in greenhouse absorbers, we use our observed temperatures and the tuned value for $\epsilon$.

We'll express the answer in W m$^{-2}$ for a 2% increase in $\epsilon$:

$$ \Delta \epsilon = 0.02 \times 0.58 $$

```python
epsilon = 0.586041150248834
delta_epsilon = 0.02 * epsilon
delta_epsilon
```

    0.011720823004976679

Calculate the three components of the radiative forcing:

```python
sigma = 5.67E-8
Ts = 288.
T0 = 275.
T1 = 230.
```

```python
#  Component originating from the surface
Rs = -delta_epsilon * (-2*(1-epsilon)*sigma * Ts**4)
Rs
```

    3.78528239476332

```python
#  Component originating from level 0
R0 = -delta_epsilon * (1-2*epsilon) * sigma * T0**4
R0
```

    0.6540457761108388

```python
#  Component originating from level 1
R1 = -delta_epsilon * sigma * T1**4
R1
```

    -1.8597411929137297

So just add them up to get the total radiative forcing:

```python
R = Rs + R0 + R1
R
```

    2.579586977960429

So in our example, **the OLR decreases by 2.6 W m$^{-2}$**, or equivalently, the **radiative forcing is +2.6 W m$^{-2}$.**

What we have just calculated is this:

**Given the observed lapse rates, a small increase in absorbers will cause a small decrease in OLR.**

The **greenhouse effect** thus gets **stronger**, and energy will begin to accumulate in the system -- which will eventually **cause temperatures to increase** as the system adjusts to a new equilibrium.

____________

## 7. Summary
____________

### Key physical lessons

- Putting a **layer of longwave absorbers** above the surface keeps the **surface substantially warmer**, because of the **backradiation** from the atmosphere (greenhouse effect).
- The **grey gas** model assumes that each layer absorbs and emits a fraction $\epsilon$ of its blackbody value, independent of wavelength.

- With **incomplete absorption** ($\epsilon < 1$), there are contributions to the OLR from every level and the surface (there is no single **level of emission**)
- Adding more absorbers means that **contributions to the OLR** from **upper levels** go **up**, while contributions from the surface go **down**.
- This upward shift in the weighting of different levels is what we mean when we say the **level of emission goes up**.

- The **radiative forcing** caused by an increase in absorbers **depends on the lapse rate**.
- For an **isothermal atmosphere** the radiative forcing is zero and there is **no greenhouse effect**
- The radiative forcing is positive for our atmosphere **because tropospheric temperatures tend to decrease with height**.

____________

## Credits

This notebook is part of [The Climate Laboratory](https://brian-rose.github.io/ClimateLaboratoryBook), an open-source textbook developed and maintained by [Brian E. J. Rose](http://www.atmos.albany.edu/facstaff/brose/index.html), University at Albany.

It is licensed for free and open consumption under the
[Creative Commons Attribution 4.0 International (CC BY 4.0)](https://creativecommons.org/licenses/by/4.0/) license.

Development of these notes and the [climlab software](https://github.com/brian-rose/climlab) is partially supported by the National Science Foundation under award AGS-1455071 to Brian Rose. Any opinions, findings, conclusions or recommendations expressed here are mine and do not necessarily reflect the views of the National Science Foundation.
____________

```python

```


# Beginner Tutorial for SymPy Mechanics

Once you are at the python command line the first step is to import basic functionality from SymPy and the Mechanics module, otherwise you will only have basic python commands available to work with. We will import the `symbols` function from SymPy core and with the * method
bring in all functionality from the mechanics package.

```
from sympy import symbols
from sympy.physics.mechanics import *
```

You can now see what functions and variables that are available to you with::

```
dir()
```

    ['Dyadic',
     'In',
     'KanesMethod',
     'LagrangesMethod',
     'Lagrangian',
     'MechanicsLatexPrinter',
     'MechanicsPrettyPrinter',
     'MechanicsStrPrinter',
     'Out',
     'Particle',
     'Point',
     'ReferenceFrame',
     'RigidBody',
     'Vector',
     '_',
     '__',
     '___',
     '__builtin__',
     '__builtins__',
     '__name__',
     '__package__',
     '_dh',
     '_i',
     '_i1',
     '_i2',
     '_ih',
     '_ii',
     '_iii',
     '_oh',
     '_sh',
     'angular_momentum',
     'cross',
     'dot',
     'dynamicsymbols',
     'exit',
     'express',
     'get_ipython',
     'help',
     'inertia',
     'inertia_of_point_mass',
     'kinematic_equations',
     'kinetic_energy',
     'linear_momentum',
     'mechanics_printing',
     'mlatex',
     'mpprint',
     'mprint',
     'outer',
     'partial_velocity',
     'potential_energy',
     'quit',
     'symbols']

This is a long list of available functions. Read about the python import statement to learn about better ways to import only
what you need. One good explanation is <http://effbot.org/zone/import-confusion.htm>.

To get started working with vectors we will need to create a reference frame,
as all vectors need to be defined with respect to a reference frame in the mechanics package. If you
know the name of the command that you want to use simply use the builtin help
function to bring up the documentation for the function. In our case we need to
use the ReferenceFrame class.

```
help(ReferenceFrame)
```

    Help on class ReferenceFrame in module sympy.physics.mechanics.essential:

    class ReferenceFrame(__builtin__.object)
     |  A reference frame in classical mechanics.
     |
     |  ReferenceFrame is a class used to represent a reference frame in classical
     |  mechanics. It has a standard basis of three unit vectors in the frame's
     |  x, y, and z directions.
     |
     |  It also can have a rotation relative to a parent frame; this rotation is
     |  defined by a direction cosine matrix relating this frame's basis vectors to
     |  the parent frame's basis vectors.  It can also have an angular velocity
     |  vector, defined in another frame.
     |
     |  Methods defined here:
     |
     |  __getitem__(self, ind)
     |      Returns basis vector for the provided index (index being an str)
     |
     |  __init__(self, name, indices=None, latexs=None)
     |      ReferenceFrame initialization method.
     |
     |      A ReferenceFrame has a set of orthonormal basis vectors, along with
     |      orientations relative to other ReferenceFrames and angular velocities
     |      relative to other ReferenceFrames.
     |
     |      Parameters
     |      ==========
     |
     |      indices : list (of strings)
     |          If custom indices are desired for console, pretty, and LaTeX
     |          printing, supply three as a list. The basis vectors can then be
     |          accessed with the get_item method.
     |      latexs : list (of strings)
     |          If custom names are desired for LaTeX printing of each basis
     |          vector, supply the names here in a list.
     |
     |      Examples
     |      ========
     |
     |      >>> from sympy.physics.mechanics import ReferenceFrame, mlatex
     |      >>> N = ReferenceFrame('N')
     |      >>> N.x
     |      N.x
     |      >>> O = ReferenceFrame('O', ('1', '2', '3'))
     |      >>> O.x
     |      O['1']
     |      >>> O['1']
     |      O['1']
     |      >>> P = ReferenceFrame('P', latexs=('A1', 'A2', 'A3'))
     |      >>> mlatex(P.x)
     |      'A1'
     |
     |  __iter__(self)
     |
     |  __repr__ = __str__(self)
     |
     |  __str__(self)
     |      Returns the name of the frame.
     |
     |  ang_acc_in(self, otherframe)
     |      Returns the angular acceleration Vector of the ReferenceFrame.
     |
     |      Effectively returns the Vector:
     |      ^N alpha ^B
     |      which represent the angular acceleration of B in N, where B is self, and
     |      N is otherframe.
     |
     |      Parameters
     |      ==========
     |
     |      otherframe : ReferenceFrame
     |          The ReferenceFrame which the angular acceleration is returned in.
     |
     |      Examples
     |      ========
     |
     |      >>> from sympy.physics.mechanics import ReferenceFrame, Vector
     |      >>> N = ReferenceFrame('N')
     |      >>> A = ReferenceFrame('A')
     |      >>> V = 10 * N.x
     |      >>> A.set_ang_acc(N, V)
     |      >>> A.ang_acc_in(N)
     |      10*N.x
     |
     |  ang_vel_in(self, otherframe)
     |      Returns the angular velocity Vector of the ReferenceFrame.
     |
     |      Effectively returns the Vector:
     |      ^N omega ^B
     |      which represent the angular velocity of B in N, where B is self, and
     |      N is otherframe.
     |
     |      Parameters
     |      ==========
     |
     |      otherframe : ReferenceFrame
     |          The ReferenceFrame which the angular velocity is returned in.
     |
     |      Examples
     |      ========
     |
     |      >>> from sympy.physics.mechanics import ReferenceFrame, Vector
     |      >>> N = ReferenceFrame('N')
     |      >>> A = ReferenceFrame('A')
     |      >>> V = 10 * N.x
     |      >>> A.set_ang_vel(N, V)
     |      >>> A.ang_vel_in(N)
     |      10*N.x
     |
     |  dcm(self, otherframe)
     |      The direction cosine matrix between frames.
     |
     |      This gives the DCM between this frame and the otherframe.
     |      The format is N.xyz = N.dcm(B) * B.xyz
     |      A SymPy Matrix is returned.
     |
     |      Parameters
     |      ==========
     |
     |      otherframe : ReferenceFrame
     |          The otherframe which the DCM is generated to.
     |
     |      Examples
     |      ========
     |
     |      >>> from sympy.physics.mechanics import ReferenceFrame, Vector
     |      >>> from sympy import symbols
     |      >>> q1 = symbols('q1')
     |      >>> N = ReferenceFrame('N')
     |      >>> A = N.orientnew('A', 'Axis', [q1, N.x])
     |      >>> N.dcm(A)
     |      [1,       0,        0]
     |      [0, cos(q1), -sin(q1)]
     |      [0, sin(q1),  cos(q1)]
     |
     |  orient(self, parent, rot_type, amounts, rot_order='')
     |      Defines the orientation of this frame relative to a parent frame.
     |
     |      Supported orientation types are Body, Space, Quaternion, Axis.
     |      Examples show correct usage.
     |
     |      Parameters
     |      ==========
     |
     |      parent : ReferenceFrame
     |          The frame that this ReferenceFrame will have its orientation matrix
     |          defined in relation to.
     |      rot_type : str
     |          The type of orientation matrix that is being created.
     |      amounts : list OR value
     |          The quantities that the orientation matrix will be defined by.
     |      rot_order : str
     |          If applicable, the order of a series of rotations.
     |
     |      Examples
     |      ========
     |
     |      >>> from sympy.physics.mechanics import ReferenceFrame, Vector
     |      >>> from sympy import symbols
     |      >>> q0, q1, q2, q3, q4 = symbols('q0 q1 q2 q3 q4')
     |      >>> N = ReferenceFrame('N')
     |      >>> B = ReferenceFrame('B')
     |
     |      Now we have a choice of how to implement the orientation. First is
     |      Body. Body orientation takes this reference frame through three
     |      successive simple rotations. Acceptable rotation orders are of length
     |      3, expressed in XYZ or 123, and cannot have a rotation about about an
     |      axis twice in a row.
     |
     |      >>> B.orient(N, 'Body', [q1, q2, q3], '123')
     |      >>> B.orient(N, 'Body', [q1, q2, 0], 'ZXZ')
     |      >>> B.orient(N, 'Body', [0, 0, 0], 'XYX')
     |
     |      Next is Space. Space is like Body, but the rotations are applied in the
     |      opposite order.
     |
     |      >>> B.orient(N, 'Space', [q1, q2, q3], '312')
     |
     |      Next is Quaternion. This orients the new ReferenceFrame with
     |      Quaternions, defined as a finite rotation about lambda, a unit vector,
     |      by some amount theta.
     |      This orientation is described by four parameters:
     |      q0 = cos(theta/2)
     |      q1 = lambda_x sin(theta/2)
     |      q2 = lambda_y sin(theta/2)
     |      q3 = lambda_z sin(theta/2)
     |      Quaternion does not take in a rotation order.
     |
     |      >>> B.orient(N, 'Quaternion', [q0, q1, q2, q3])
     |
     |      Last is Axis. This is a rotation about an arbitrary, non-time-varying
     |      axis by some angle. The axis is supplied as a Vector. This is how
     |      simple rotations are defined.
     |
     |      >>> B.orient(N, 'Axis', [q1, N.x + 2 * N.y])
     |
     |  orientnew(self, newname, rot_type, amounts, rot_order='', indices=None, latexs=None)
     |      Creates a new ReferenceFrame oriented with respect to this Frame.
     |
     |      See ReferenceFrame.orient() for acceptable rotation types, amounts,
     |      and orders. Parent is going to be self.
     |
     |      Parameters
     |      ==========
     |
     |      newname : str
     |          The name for the new ReferenceFrame
     |      rot_type : str
     |          The type of orientation matrix that is being created.
     |      amounts : list OR value
     |          The quantities that the orientation matrix will be defined by.
     |      rot_order : str
     |          If applicable, the order of a series of rotations.
     |
     |
     |      Examples
     |      ========
     |
     |      >>> from sympy.physics.mechanics import ReferenceFrame, Vector
     |      >>> from sympy import symbols
     |      >>> q1 = symbols('q1')
     |      >>> N = ReferenceFrame('N')
     |      >>> A = N.orientnew('A', 'Axis', [q1, N.x])
     |
     |
     |      .orient() documentation:
     |
     |      ========================
     |
     |      Defines the orientation of this frame relative to a parent frame.
     |
     |      Supported orientation types are Body, Space, Quaternion, Axis.
     |      Examples show correct usage.
     |
     |      Parameters
     |      ==========
     |
     |      parent : ReferenceFrame
     |          The frame that this ReferenceFrame will have its orientation matrix
     |          defined in relation to.
     |      rot_type : str
     |          The type of orientation matrix that is being created.
     |      amounts : list OR value
     |          The quantities that the orientation matrix will be defined by.
     |      rot_order : str
     |          If applicable, the order of a series of rotations.
     |
     |      Examples
     |      ========
     |
     |      >>> from sympy.physics.mechanics import ReferenceFrame, Vector
     |      >>> from sympy import symbols
     |      >>> q0, q1, q2, q3, q4 = symbols('q0 q1 q2 q3 q4')
     |      >>> N = ReferenceFrame('N')
     |      >>> B = ReferenceFrame('B')
     |
     |      Now we have a choice of how to implement the orientation. First is
     |      Body. Body orientation takes this reference frame through three
     |      successive simple rotations. Acceptable rotation orders are of length
     |      3, expressed in XYZ or 123, and cannot have a rotation about about an
     |      axis twice in a row.
     |
     |      >>> B.orient(N, 'Body', [q1, q2, q3], '123')
     |      >>> B.orient(N, 'Body', [q1, q2, 0], 'ZXZ')
     |      >>> B.orient(N, 'Body', [0, 0, 0], 'XYX')
     |
     |      Next is Space. Space is like Body, but the rotations are applied in the
     |      opposite order.
     |
     |      >>> B.orient(N, 'Space', [q1, q2, q3], '312')
     |
     |      Next is Quaternion. This orients the new ReferenceFrame with
     |      Quaternions, defined as a finite rotation about lambda, a unit vector,
     |      by some amount theta.
     |      This orientation is described by four parameters:
     |      q0 = cos(theta/2)
     |      q1 = lambda_x sin(theta/2)
     |      q2 = lambda_y sin(theta/2)
     |      q3 = lambda_z sin(theta/2)
     |      Quaternion does not take in a rotation order.
     |
     |      >>> B.orient(N, 'Quaternion', [q0, q1, q2, q3])
     |
     |      Last is Axis. This is a rotation about an arbitrary, non-time-varying
     |      axis by some angle. The axis is supplied as a Vector. This is how
     |      simple rotations are defined.
     |
     |      >>> B.orient(N, 'Axis', [q1, N.x + 2 * N.y])
     |
     |  set_ang_acc(self, otherframe, value)
     |      Define the angular acceleration Vector in a ReferenceFrame.
     |
     |      Defines the angular acceleration of this ReferenceFrame, in another.
     |      Angular acceleration can be defined with respect to multiple different
     |      ReferenceFrames. Care must be taken to not create loops which are
     |      inconsistent.
     |
     |      Parameters
     |      ==========
     |
     |      otherframe : ReferenceFrame
     |          A ReferenceFrame to define the angular acceleration in
     |      value : Vector
     |          The Vector representing angular acceleration
     |
     |      Examples
     |      ========
     |
     |      >>> from sympy.physics.mechanics import ReferenceFrame, Vector
     |      >>> N = ReferenceFrame('N')
     |      >>> A = ReferenceFrame('A')
     |      >>> V = 10 * N.x
     |      >>> A.set_ang_acc(N, V)
     |      >>> A.ang_acc_in(N)
     |      10*N.x
     |
     |  set_ang_vel(self, otherframe, value)
     |      Define the angular velocity vector in a ReferenceFrame.
     |
     |      Defines the angular velocity of this ReferenceFrame, in another.
     |      Angular velocity can be defined with respect to multiple different
     |      ReferenceFrames. Care must be taken to not create loops which are
     |      inconsistent.
     |
     |      Parameters
     |      ==========
     |
     |      otherframe : ReferenceFrame
     |          A ReferenceFrame to define the angular velocity in
     |      value : Vector
     |          The Vector representing angular velocity
     |
     |      Examples
     |      ========
     |
     |      >>> from sympy.physics.mechanics import ReferenceFrame, Vector
     |      >>> N = ReferenceFrame('N')
     |      >>> A = ReferenceFrame('A')
     |      >>> V = 10 * N.x
     |      >>> A.set_ang_vel(N, V)
     |      >>> A.ang_vel_in(N)
     |      10*N.x
     |
     |  ----------------------------------------------------------------------
     |  Data descriptors defined here:
     |
     |  __dict__
     |      dictionary for instance variables (if defined)
     |
     |  __weakref__
     |      list of weak references to the object (if defined)
     |
     |  x
     |      The basis Vector for the ReferenceFrame, in the x direction.
     |
     |  y
     |      The basis Vector for the ReferenceFrame, in the y direction.
     |
     |  z
     |      The basis Vector for the ReferenceFrame, in the z direction.

The `ReferenceFrame` class manages everything about rotations, angular velocities, and angular accelerations with respect to other reference frames. Now create an inertial reference frame called N for "Newtonian" as was described in the the docstring:

```
N = ReferenceFrame('N')
```

Keep in mind that `N` is the variable name of which the reference frame named "N"
is stored. It is important to note that `N` is an object and it has properties
and functions associated with it. To see a list of them type:

```
dir(N)
```

    ['__class__',
     '__delattr__',
     '__dict__',
     '__doc__',
     '__format__',
     '__getattribute__',
     '__getitem__',
     '__hash__',
     '__init__',
     '__iter__',
     '__module__',
     '__new__',
     '__reduce__',
     '__reduce_ex__',
     '__repr__',
     '__setattr__',
     '__sizeof__',
     '__str__',
     '__subclasshook__',
     '__weakref__',
     '_ang_acc_dict',
     '_ang_vel_dict',
     '_cur',
     '_dcm_dict',
     '_dict_list',
     '_dlist',
     '_w_diff_dcm',
     '_x',
     '_y',
     '_z',
     'ang_acc_in',
     'ang_vel_in',
     'dcm',
     'indices',
     'latex_vecs',
     'name',
     'orient',
     'orientnew',
     'pretty_vecs',
     'set_ang_acc',
     'set_ang_vel',
     'str_vecs',
     'x',
     'y',
     'z']

Notice that three of the properties are `x`, `y`, and `z`. These are the
orthonormal unit vectors associated with the reference frame and are the
building blocks for creating more general vectors. We can create a vector by simply
building a linear combination of the unit vectors:

```
v = 1 * N.x + 2 * N.y + 3 * N.z
```

Now a vector expressed in the N reference frame is stored in the variable `v`.
We can print `v` to the screen by typing:

```
print(v)
```

    N.x + 2*N.y + 3*N.z

The vector `v` can be manipulated as expected. You can multiply and divide them by scalars:

```
2 * v
```

    2*N.x + 4*N.y + 6*N.z

```
v / 3.0
```

    0.333333333333333*N.x + 0.666666666666667*N.y + N.z

Note that three is expressed as `3.0` instead of `3`. The python language does
integer division by default. There are ways around this, but for now simply
remember to always declare numbers as floats (i.e. include a decimal).

You can also add and subtract vectors::

```
v + v
```

    2*N.x + 4*N.y + 6*N.z

```
w = 5 * N.x + 7 * N.y
```

```
v - w
```

    - 4*N.x - 5*N.y + 3*N.z

Vectors also have some useful properties:

```
dir(v)
```

    ['__add__',
     '__and__',
     '__class__',
     '__delattr__',
     '__dict__',
     '__div__',
     '__doc__',
     '__eq__',
     '__format__',
     '__getattribute__',
     '__hash__',
     '__init__',
     '__module__',
     '__mul__',
     '__ne__',
     '__neg__',
     '__new__',
     '__or__',
     '__radd__',
     '__rand__',
     '__reduce__',
     '__reduce_ex__',
     '__repr__',
     '__rmul__',
     '__ror__',
     '__rsub__',
     '__setattr__',
     '__sizeof__',
     '__str__',
     '__sub__',
     '__subclasshook__',
     '__truediv__',
     '__weakref__',
     '__xor__',
     '_latex',
     '_pretty',
     '_sympyrepr',
     '_sympystr',
     'args',
     'cross',
     'diff',
     'doit',
     'dot',
     'dt',
     'express',
     'magnitude',
     'normalize',
     'outer',
     'simp',
     'simplify',
     'subs']

You can find the magnitude of a vector by typing:

```
v.magnitude()
```

    sqrt(14)

You can compute a unit vector in the direction of `v`:

```
v.normalize()
```

    sqrt(14)/14*N.x + sqrt(14)/7*N.y + 3*sqrt(14)/14*N.z

You can find the measure numbers and the reference frame the vector was defined in
with::

```
v.args
```

    [([1]
    [2]
    [3], N)]

Dot and cross products of vectors can also be computed::

```
dot(v, w)
```

    19

```
cross(v, w)
```

    - 21*N.x + 15*N.y - 3*N.z

We've only used numbers as our measure numbers so far, but it is just as easy
to use symbols. We will introduce six symbols for our measure numbers with the
SymPy `symbols` [`help(symbols) for the documentation`] function:

```
a1, a2, a3 = symbols('a1 a2 a3')
b1, b2, b3 = symbols('b1 b2 b3')
```

And create two new vectors that are completely symbolic:

```
x = a1 * N.x + a2 * N.y + a3 * N.z
y = b1 * N.x + b2 * N.y + b3 * N.z
```

```
dot(x, y)
```

    a1*b1 + a2*b2 + a3*b3

```
z = cross(x, y)
z
```

    (a2*b3 - a3*b2)*N.x + (-a1*b3 + a3*b1)*N.y + (a1*b2 - a2*b1)*N.z

Numbers and symbols work together seamlessly:

```
dot(v, x)
```

    a1 + 2*a2 + 3*a3

You can also differentiate a vector with respect to a variable in a
reference frame:

```
z.diff(a1, N)
```

    - b3*N.y + b2*N.z

Vectors don't have be defined with respect to just one reference frame. We can
create a new reference frame and orient it with respect to the `N` frame that
has already been created. We will use the `orient` method of the new frame to
do a simple rotation through `alpha` about the `N.x` axis:

```
A = ReferenceFrame('A')
alpha = symbols('alpha')
A.orient(N, 'Axis', [alpha, N.x])
```

Now the direction cosine matrix with of `A` with respect to `N` can be
computed:

```
A.dcm(N)
```

    [1,           0,          0]
    [0,  cos(alpha), sin(alpha)]
    [0, -sin(alpha), cos(alpha)]

Now that SymPy knows that `A` and `N` are oriented with respect to each other
we can express the vectors that we originally wrote in the `A` frame:

```
v.express(A)
```

    A.x + (3*sin(alpha) + 2*cos(alpha))*A.y + (-2*sin(alpha) + 3*cos(alpha))*A.z

```
z.express(A)
```

    (a2*b3 - a3*b2)*A.x + ((a1*b2 - a2*b1)*sin(alpha) + (-a1*b3 + a3*b1)*cos(alpha))*A.y + ((a1*b2 - a2*b1)*cos(alpha) - (-a1*b3 + a3*b1)*sin(alpha))*A.z

In dynamics systems, at least some of the relative orientation of reference
frames and vectors are time varying. The mechanics module provides a way to
specify quantities as time varying. Let's define two variables `beta` and `d` as
variables which are functions of time:

```
beta, d = dynamicsymbols('beta d')
beta, d
```

    (beta(t), d(t))

Now we can create a new reference frame that is oriented with respect to the `A`
frame by `beta` and create a vector in that new frame that is a function of `d`.
This time we will use the `orientnew` method of the `A` frame to create the new
reference frame `B`:

```
B = A.orientnew('B', 'Axis', (beta, A.y))
```

```
vec = d * B.z
```

We can now compute the angular velocity of the reference frame `B` with respect
to other reference frames:

```
B.ang_vel_in(N)
```

    beta'*A.y

This allows us to now differentiate the vector, `vec`, with respect to time and
a reference frame::

```
vecdot = vec.dt(N)
vecdot
```

    d*beta'*B.x + d'*B.z

```
vecdot.express(N)
```

    (d*cos(beta)*beta' + sin(beta)*d')*N.x + (d*sin(alpha)*sin(beta)*beta' - sin(alpha)*cos(beta)*d')*N.y + (-d*sin(beta)*cos(alpha)*beta' + cos(alpha)*cos(beta)*d')*N.z

The `dynamicsymbols` function also allows you to easily create the derivatives of time
varying variables and store them in a variable.

```
theta = dynamicsymbols('theta')
thetad = dynamicsymbols('theta', 1)
theta, thetad
```

    (theta(t), Derivative(theta(t), t))

At this point we now have all the tools needed to setup the kinematics for a
dynamic system. Let's start with the classic mass spring damper system under the influence of gravity. Go to the next notebook to follow along: https://github.com/PythonDynamics/pydy_examples/tree/master/mass_spring_damper.


# Breaking ADAM

In case you haven't heard, one of the top papers at [ICLR 2018](https://iclr.cc/Conferences/2018) (pronounced:
eye-clear, who knew?) was [On the Convergence of Adam and Beyond](https://openreview.net/forum?id=ryQu7f-RZ). In the
paper, the authors determine a flaw in the convergence proof of the ubiquitous ADAM optimizer. They also give an example
of a simple function for which ADAM does not converge to the correct solution. We've seen how torchbearer can be used
for simple function optimisation before and we can do something similar to reproduce the results
from the paper. We should note that this isn't a suggestion that you should necesarily use AMSGrad (the proposed solution from the paper), most work still uses either SGD with momentum or vanilla Adam and some other approaches have been proposed such as [AdamW](https://arxiv.org/abs/1711.05101) which may provide benefits. Our intention here is simply to show how this interesting failure case can be demonstrated with torchbearer.

**Note**: The easiest way to use this tutorial is as a colab notebook, which allows you to dive in with no setup. We recommend you enable a free GPU with

> **Runtime**   →   **Change runtime type**   →   **Hardware Accelerator: GPU**

## Install Torchbearer

First we install torchbearer if needed.

```python
try:
    import torchbearer
except:
    !pip install -q torchbearer
    import torchbearer

print(torchbearer.__version__)
```

    0.4.0.dev

Online Optimization
-----------------------------------

Online learning basically just means learning from one example at a time, in sequence. The function given in the paper
 has a unique minimum at $x=-1$ and is defined as follows:

\begin{equation}
f_t(x) = \begin{cases}1010x, & \text{for } t \; \texttt{mod} \; 101 = 1 \\ -10x, & \text{otherwise}\end{cases}
\end{equation}

We can then write this as a PyTorch model whose forward is a function of its parameters with the following:

```python
import torch.nn as nn

class Online(nn.Module):
    def __init__(self):
        super().__init__()
        self.x = nn.Parameter(torch.zeros(1))

    def forward(self, _, state):
        """
        function to be minimised:
        f(x) = 1010x if t mod 101 = 1, else -10x
        """
        if state[torchbearer.BATCH] % 101 == 1:
            res = 1010 * self.x
        else:
            res = -10 * self.x

        return res
```

We now define a loss (simply return the model output) and a metric which returns the value of our parameter `x`:

```python
def loss(y_pred, _):
    return y_pred

@torchbearer.metrics.to_dict
class est(torchbearer.metrics.Metric):
    def __init__(self):
        super().__init__('est')

    def process(self, state):
        return state[torchbearer.MODEL].x.data.item()
```

In the paper, `x` can only hold values in `[-1, 1]`. We don't strictly need to do anything but we can write
a callback that greedily updates `x` if it is outside of its range as follows:

```python
@torchbearer.callbacks.on_step_training
def greedy_update(state):
    if state[torchbearer.MODEL].x > 1:
        state[torchbearer.MODEL].x.data.fill_(1)
    elif state[torchbearer.MODEL].x < -1:
        state[torchbearer.MODEL].x.data.fill_(-1)
```

Ideally, we would like to generate the graphs from the paper. To do this, we can use the tensorboard callback from torchbearer. However, we might also want to see a live graph directly in the notebook. For that we can use a callback which updates at the end of each step.

```python
try:
  import google.colab
  IN_COLAB = True
except:
  IN_COLAB = False

# Notebook doesn't seem to work right in Colab
if IN_COLAB:
  %matplotlib inline
else:
  %matplotlib notebook

from matplotlib import pyplot as plt

class Plotter():
    def __init__(self):
        self.plots = []
        self.fig = plt.figure(figsize=(5, 5))

    def make_plotter(self, name, c='k', step_size=100):
        idx = len(self.plots)
        self.plots.append(([], [], c, name))

        @torchbearer.callbacks.on_step_training
        @torchbearer.callbacks.only_if(lambda state: state[torchbearer.BATCH] % step_size == 0)
        def store(state):
            self.plots[idx][0].append(len(self.plots[idx][0]) * step_size)
            self.plots[idx][1].append(state[torchbearer.METRICS]['est'])

        @torchbearer.callbacks.on_step_training
        @torchbearer.callbacks.only_if(lambda state: state[torchbearer.BATCH] % (2 * step_size) == 0)
        def plot(state):
            plt.clf()
            for plot in self.plots:
                plt.plot(plot[0], plot[1], plot[2], label=plot[3])
            plt.legend()
            plt.xlabel('Step')
            plt.ylabel('Estimate')
            self.fig.canvas.draw()

        return torchbearer.callbacks.CallbackList([store, plot])
```

Given a figure, the above code will plot the estimate history every given number of steps, although in Colab this will just plot the graph at the end. Finally, we can train this model twice; once with ADAM and once with AMSGrad (included in PyTorch) with just a few
lines (this will take at least a few minutes on a GPU):

```python
import sys
if 'tensorboardX' in sys.modules:
  from torchbearer.callbacks import TensorBoard
else:
  !pip install -q mock
  import mock
  TensorBoard = mock.MagicMock()
import torch

device = 'cuda' if torch.cuda.is_available() else 'cpu'
training_steps = 100000

plt.ion()
plotter = Plotter()
model = Online()
optim = torch.optim.Adam(model.parameters(), lr=0.001, betas=[0.9, 0.99])
tbtrial = torchbearer.Trial(model, optim, loss, [est()],
                            callbacks=[
                                greedy_update,
                                plotter.make_plotter('Adam', c='g', step_size=1000),
                                TensorBoard(comment='adam', write_graph=False, write_batch_metrics=True, write_epoch_metrics=False)
                            ]).to(device)
_ = tbtrial.for_train_steps(training_steps).run(verbose=0)

model = Online()
optim = torch.optim.Adam(model.parameters(), lr=0.001, betas=[0.9, 0.99], amsgrad=True)
tbtrial = torchbearer.Trial(model, optim, loss, [est()],
                            callbacks=[
                                greedy_update,
                                plotter.make_plotter('AMSGrad', c='r', step_size=1000),
                                TensorBoard(comment='amsgrad', write_graph=False, write_batch_metrics=True, write_epoch_metrics=False)
                            ]).to(device)
tbtrial.for_train_steps(training_steps).run(verbose=0)
plt.ioff()
```

    <IPython.core.display.Javascript object>

Note that we have also logged to TensorBoard here (if tensorboardX was installed), which doesn't work on Colab. If you're running locally (for around 6000000 steps) and have tensorboard/tensorboardX installed then after completion, running  `tensorboard --logdir logs` and
navigating to [localhost:6006](http://localhost:6006), you will see a graph like the one in Figure 1 from the paper,
where the top line is with ADAM and the bottom with AMSGrad:

```python
from IPython.display import Image
Image('https://raw.githubusercontent.com/ecs-vlc/torchbearer/master/docs/_static/img/ams_grad_online.png', width=500)
```

Stochastic Optimization
------------------------------

To simulate a stochastic setting, the authors use a slight variant of the function, which changes with some probability:

\begin{equation}
f_t(x) = \begin{cases}1010x, & \text{with probability } 0.01 \\ -10x, & \text{otherwise}\end{cases}
\end{equation}

We can again formulate this as a PyToch model:

```python
import random

class Stochastic(nn.Module):
    def __init__(self):
        super().__init__()
        self.x = nn.Parameter(torch.zeros(1))

    def forward(self, _):
        """
        function to be minimised:
        f(x) = 1010x with probability 0.01, else -10x
        """
        if random.random() <= 0.01:
            res = 1010 * self.x
        else:
            res = -10 * self.x

        return res
```

Using the loss, callback and metric from our previous example, we can train with the following:

```python
import sys
if 'tensorboardX' in sys.modules:
  from torchbearer.callbacks import TensorBoard
else:
  import mock
  TensorBoard = mock.MagicMock()
import torch

device = 'cuda' if torch.cuda.is_available() else 'cpu'
training_steps = 100000

plt.ion()
plotter = Plotter()
model = Stochastic()
optim = torch.optim.Adam(model.parameters(), lr=0.001, betas=[0.9, 0.99])
tbtrial = torchbearer.Trial(model, optim, loss, [est()],
                            callbacks=[
                                greedy_update,
                                plotter.make_plotter('Adam', c='g', step_size=1000),
                                TensorBoard(comment='adam', write_graph=False, write_batch_metrics=True, write_epoch_metrics=False)
                            ]).to(device)
tbtrial.for_train_steps(training_steps).run(verbose=0)

model = Stochastic()
optim = torch.optim.Adam(model.parameters(), lr=0.001, betas=[0.9, 0.99], amsgrad=True)
tbtrial = torchbearer.Trial(model, optim, loss, [est()],
                            callbacks=[
                                greedy_update,
                                plotter.make_plotter('AMSGrad', c='r', step_size=1000),
                                TensorBoard(comment='amsgrad', write_graph=False, write_batch_metrics=True, write_epoch_metrics=False)
                            ]).to(device)
tbtrial.for_train_steps(training_steps).run(verbose=0)
plt.ioff()
```

    <IPython.core.display.Javascript object>

After execution has finished, again running `tensorboard --logdir logs` and navigating to
[localhost:6006](http://localhost:6006) <http://localhost:6006>, we see another graph similar to that of the stochastic setting in Figure 1 of
the paper, where the top line is with ADAM and the bottom with AMSGrad:

```python
from IPython.display import Image
Image('https://raw.githubusercontent.com/ecs-vlc/torchbearer/master/docs/_static/img/ams_grad_stochastic.png', width=500)
```

Conclusions
------------------------------------

So, whatever your thoughts on the AMSGrad optimizer in practice, it's probably the sign of a good paper that you can
re-implement the example and get very similar results without having to try too hard and (thanks to torchbearer) only
writing a small amount of code. The paper includes some more complex, 'real-world' examples, can you re-implement those?

```python

```


# Neuromatch Academy: Week 1, Day 3, Tutorial 4
# Model Fitting: Multiple linear regression and polynomial regression

**Content creators**: Pierre-Étienne Fiquet, Anqi Wu, Alex Hyafil with help from Byron Galbraith, Ella Batty

**Content reviewers**: Lina Teichmann, Saeed Salehi, Patrick Mineault, Michael Waskom

---
#Tutorial Objectives

This is Tutorial 4 of a series on fitting models to data. We start with simple linear regression, using least squares optimization (Tutorial 1) and Maximum Likelihood Estimation (Tutorial 2). We will use bootstrapping to build confidence intervals around the inferred linear model parameters (Tutorial 3). We'll finish our exploration of regression models by generalizing to multiple linear regression and polynomial regression (Tutorial 4). We end by learning how to choose between these various models. We discuss the bias-variance trade-off (Tutorial 5) and Cross Validation for model selection (Tutorial 6).

In this tutorial, we will generalize the regression model to incorporate multiple features.
- Learn how to structure inputs for regression using the 'Design Matrix'
- Generalize the MSE for multiple features using the ordinary least squares estimator
- Visualize data and model fit in multiple dimensions
- Fit polynomial regression models of different complexity
- Plot and evaluate the polynomial regression fits

```python
#@title Video 1: Multiple Linear Regression and Polynomial Regression
from IPython.display import YouTubeVideo
video = YouTubeVideo(id="d4nfTki6Ejc", width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video
```

---
# Setup

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

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

```python
#@title Helper Functions

def plot_fitted_polynomials(x, y, theta_hat):
  """ Plot polynomials of different orders

  Args:
    x (ndarray): input vector of shape (n_samples)
    y (ndarray): vector of measurements of shape (n_samples)
    theta_hat (dict): polynomial regression weights for different orders
  """

  x_grid = np.linspace(x.min() - .5, x.max() + .5)

  plt.figure()

  for order in range(0, max_order + 1):
    X_design = make_design_matrix(x_grid, order)
    plt.plot(x_grid, X_design @ theta_hat[order]);

  plt.ylabel('y')
  plt.xlabel('x')
  plt.plot(x, y, 'C0.');
  plt.legend([f'order {o}' for o in range(max_order + 1)], loc=1)
  plt.title('polynomial fits')
  plt.show()
```

---
# Section 1: Multiple Linear Regression

Now that we have considered the univariate case and how to produce confidence intervals for our estimator, we turn to the general linear regression case, where we can have more than one regressor, or feature, in our input.

Recall that our original univariate linear model was given as

\begin{align}
y = \theta x + \epsilon
\end{align}

where $\theta$ is the slope and $\epsilon$ some noise. We can easily extend this to the multivariate scenario by adding another parameter for each additional feature

\begin{align}
y = \theta_0 + \theta_1 x_1 + \theta_2 x_2 + ... +\theta_d x_d + \epsilon
\end{align}

where $\theta_0$ is the intercept and $d$ is the number of features (it is also the dimensionality of our input).

We can condense this succinctly using vector notation for a single data point

\begin{align}
y_i = \boldsymbol{\theta}^{\top}\mathbf{x}_i + \epsilon
\end{align}

and fully in matrix form

\begin{align}
\mathbf{y} = \mathbf{X}\boldsymbol{\theta} + \mathbf{\epsilon}
\end{align}

where $\mathbf{y}$ is a vector of measurements, $\mathbf{X}$ is a matrix containing the feature values (columns) for each input sample (rows), and $\boldsymbol{\theta}$ is our parameter vector.

This matrix $\mathbf{X}$ is often referred to as the "[design matrix](https://en.wikipedia.org/wiki/Design_matrix)".

For this tutorial we will focus on the two-dimensional case ($d=2$), which allows us to easily visualize our results. As an example, think of a situation where a scientist records the spiking response of a retinal ganglion cell to patterns of light signals that vary in contrast and in orientation. Then contrast and orientation values can be used as features / regressors to predict the cells response.

In this case our model can be writen as:

\begin{align}
y = \theta_0 + \theta_1 x_1 + \theta_2 x_2 + \epsilon
\end{align}

or in matrix form where

\begin{align}
\mathbf{X} =
\begin{bmatrix}
1 & x_{1,1} & x_{1,2} \\
1 & x_{2,1} & x_{2,2} \\
\vdots & \vdots & \vdots \\
1 & x_{n,1} & x_{n,2}
\end{bmatrix},
\boldsymbol{\theta} =
\begin{bmatrix}
\theta_0 \\
\theta_1 \\
\theta_2 \\
\end{bmatrix}
\end{align}

For our actual exploration dataset we shall set $\boldsymbol{\theta}=[0, -2, -3]$ and draw $N=40$ noisy samples from $x \in [-2,2)$. Note that setting the value of $\theta_0 = 0$ effectively ignores the offset term.

```python
#@title
#@markdown Execute this cell to simulate some data

# Set random seed for reproducibility
np.random.seed(1234)

# Set parameters
theta = [0, -2, -3]
n_samples = 40

# Draw x and calculate y
n_regressors = len(theta)
x0 = np.ones((n_samples, 1))
x1 = np.random.uniform(-2, 2, (n_samples, 1))
x2 = np.random.uniform(-2, 2, (n_samples, 1))
X = np.hstack((x0, x1, x2))
noise = np.random.randn(n_samples)
y = X @ theta + noise

ax = plt.subplot(projection='3d')
ax.plot(X[:,1], X[:,2], y, '.')

ax.set(
    xlabel='$x_1$',
    ylabel='$x_2$',
    zlabel='y'
)
plt.tight_layout()
```

Now that we have our dataset, we want to find an optimal vector of paramters $\boldsymbol{\hat\theta}$. Recall our analytic solution to minimizing MSE for a single regressor:

\begin{align}
\hat\theta = \frac{\sum_{i=1}^N x_i y_i}{\sum_{i=1}^N x_i^2}.
\end{align}

The same holds true for the multiple regressor case, only now expressed in matrix form

\begin{align}
\boldsymbol{\hat\theta} = (\mathbf{X}^\top\mathbf{X})^{-1}\mathbf{X}^\top\mathbf{y}.
\end{align}

This is called the [ordinary least squares](https://en.wikipedia.org/wiki/Ordinary_least_squares) (OLS) estimator.

### Exercise 1: Ordinary Least Squares Estimator

In this exercise you will implement the OLS approach to estimating $\boldsymbol{\hat\theta}$ from the design matrix $\mathbf{X}$ and measurement vector $\mathbf{y}$. You can use the `@` symbol for matrix multiplication, `.T` for transpose, and `np.linalg.inv` for matrix inversion.

```python
def ordinary_least_squares(X, y):
  """Ordinary least squares estimator for linear regression.

  Args:
    x (ndarray): design matrix of shape (n_samples, n_regressors)
    y (ndarray): vector of measurements of shape (n_samples)

  Returns:
    ndarray: estimated parameter values of shape (n_regressors)
  """
  ######################################################################
  ## TODO for students: solve for the optimal parameter vector using OLS
  # Fill out function and remove
  raise NotImplementedError("Student exercise: solve for theta_hat vector using OLS")
  ######################################################################

  # Compute theta_hat using OLS
  theta_hat = ...

  return theta_hat

# Uncomment below to test your function
# theta_hat = ordinary_least_squares(X, y)
# print(theta_hat)
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W1D3_ModelFitting/solutions/W1D3_Tutorial4_Solution_2b5abc38.py)

After filling in this function, you should see that $\hat{\theta}$ = [ 0.13861386, -2.09395731, -3.16370742]

Now that we have our $\mathbf{\hat\theta}$, we can obtain $\mathbf{\hat y}$ and thus our mean squared error.

```python
# Compute predicted data
theta_hat = ordinary_least_squares(X, y)
y_hat = X @ theta_hat

# Compute MSE
print(f"MSE = {np.mean((y - y_hat)**2):.2f}")
```

Finally, the following code will plot a geometric visualization of the data points (blue) and fitted plane.

```python
#@title
#@markdown Execute this cell to visualize data and predicted plane

theta_hat = ordinary_least_squares(X, y)
xx, yy = np.mgrid[-2:2:50j, -2:2:50j]
y_hat_grid = np.array([xx.flatten(), yy.flatten()]).T @ theta_hat[1:]
y_hat_grid = y_hat_grid.reshape((50, 50))

ax = plt.subplot(projection='3d')
ax.plot(X[:, 1], X[:, 2], y, '.')
ax.plot_surface(xx, yy, y_hat_grid, linewidth=0, alpha=0.5, color='C1',
                cmap=plt.get_cmap('coolwarm'))

for i in range(len(X)):
  ax.plot((X[i, 1], X[i, 1]),
          (X[i, 2], X[i, 2]),
          (y[i], y_hat[i]),
          'g-', alpha=.5)

ax.set(
    xlabel='$x_1$',
    ylabel='$x_2$',
    zlabel='y'
)
plt.tight_layout()
```

---
# Section 2:  Polynomial Regression

So far today, you learned how to predict outputs from inputs by fitting a linear regression model. We can now model all sort of relationships, including in neuroscience!

One potential problem with this approach is the simplicity of the model. Linear regression, as the name implies, can only capture a linear relationship between the inputs and outputs. Put another way, the predicted outputs are only a weighted sum of the inputs. What if there are more complicated computations happening? Luckily, many more complex models exist (and you will encounter many more over the next 3 weeks). One model that is still very simple to fit and understand, but captures more complex relationships, is **polynomial regression**, an extension of linear regression.

Since polynomial regression is an extension of linear regression, everything you learned so far will come in handy now! The goal is the same: we want to predict the dependent variable $y_{n}$ given the input values $x_{n}$. The key change is the type of relationship between inputs and outputs that the model can capture.

Linear regression models predict the outputs as a weighted sum of the inputs:

$$y_{n}= \theta_0 + \theta x_{n} + \epsilon_{n}$$

With polynomial regression, we model the outputs as a polynomial equation based on the inputs. For example, we can model the outputs as:

$$y_{n}= \theta_0 + \theta_1 x_{n} + \theta_2 x_{n}^2 + \theta_3 x_{n}^3 + \epsilon_{n}$$

We can change how complex a polynomial is fit by changing the order of the polynomial. The order of a polynomial refers to the highest power in the polynomial. The equation above is a third order polynomial because the highest value x is raised to is 3. We could add another term ($+ \theta_4 x_{n}^4$) to model an order 4 polynomial and so on.

First, we will simulate some data to practice fitting polynomial regression models. We will generate random inputs $x$ and then compute y according to $y = x^2 - x - 2 $, with some extra noise both in the input and the output to make the model fitting exercise closer to a real life situation.

```python
#@title
#@markdown Execute this cell to simulate some data

# setting a fixed seed to our random number generator ensures we will always
# get the same psuedorandom number sequence
np.random.seed(121)
n_samples = 30
x = np.random.uniform(-2, 2.5, n_samples)  # inputs uniformly sampled from [-2, 2.5)
y =  x**2 - x - 2   # computing the outputs

output_noise = 1/8 * np.random.randn(n_samples)
y += output_noise  # adding some output noise

input_noise = 1/2 * np.random.randn(n_samples)
x += input_noise  # adding some input noise

fig, ax = plt.subplots()
ax.scatter(x, y)  # produces a scatter plot
ax.set(xlabel='x', ylabel='y');
```

## Section 2.1: Design matrix for polynomial regression

Now we have the basic idea of polynomial regression and some noisy data, let's begin! The key difference between fitting a linear regression model and a polynomial regression model lies in how we structure the input variables.

For linear regression, we used $X = x$ as the input data. To add a constant bias (a y-intercept in a 2-D plot), we use $X = \big[ \boldsymbol 1, x \big]$, where $\boldsymbol 1$ is a column of ones.  When fitting, we learn a weight for each column of this matrix. So we learn a weight that multiples with column 1 - in this case that column is all ones so we gain the bias parameter ($+ \theta_0$). We also learn a weight for every column, or every feature of x, as we learned in Section 1.

This matrix $X$ that we use for our inputs is known as a **design matrix**. We want to create our design matrix so we learn weights for $x^2, x^3,$ etc. Thus, we want to build our design matrix $X$ for polynomial regression of order $k$ as:

$$X = \big[ \boldsymbol 1 , x^1, x^2 , \ldots , x^k \big],$$

where $\boldsymbol{1}$ is the vector the same length as $x$ consisting of of all ones, and $x^p$ is the vector or matrix $x$ with all elements raised to the power $p$. Note that $\boldsymbol{1} = x^0$ and $x^1 = x$

### Exercise 2: Structure design matrix

Create a function (`make_design_matrix`) that structures the design matrix given the input data and the order of the polynomial you wish to fit. We will print part of this design matrix for our data and order 5.

```python
def make_design_matrix(x, order):
  """Create the design matrix of inputs for use in polynomial regression

  Args:
    x (ndarray): input vector of shape (n_samples)
    order (scalar): polynomial regression order

  Returns:
    ndarray: design matrix for polynomial regression of shape (samples, order+1)
  """
  ########################################################################
  ## TODO for students: create the design matrix ##
  # Fill out function and remove
  raise NotImplementedError("Student exercise: create the design matrix")
  ########################################################################

  # Broadcast to shape (n x 1) so dimensions work
  if x.ndim == 1:
    x = x[:, None]

  #if x has more than one feature, we don't want multiple columns of ones so we assign
  # x^0 here
  design_matrix = np.ones((x.shape[0], 1))

  # Loop through rest of degrees and stack columns (hint: np.hstack)
  for degree in range(1, order + 1):
      design_matrix = ...

  return design_matrix

# Uncomment to test your function
order = 5
# X_design = make_design_matrix(x, order)
# print(X_design[0:2, 0:2])
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W1D3_ModelFitting/solutions/W1D3_Tutorial4_Solution_f8576b48.py)

You should see that the printed section of this design matrix is `[[ 1.         -1.51194917]
 [ 1.         -0.35259945]]`

## Section 2.2: Fitting polynomial regression models

Now that we have the inputs structured correctly in our design matrix, fitting a polynomial regression is the same as fitting a linear regression model! All of the polynomial structure we need to learn is contained in how the inputs are structured in the design matrix. We can use the same least squares solution we computed in previous exercises.

### Exercise 3: Fitting polynomial regression models with different orders

Here, we will fit polynomial regression models to find the regression coefficients ($\theta_0, \theta_1, \theta_2,$ ...) by solving the least squares problem. Create a function `solve_poly_reg` that loops over different order polynomials (up to `max_order`), fits that model, and saves out the weights for each. You may invoke the `ordinary_least_squares` function.

We will then qualitatively inspect the quality of our fits for each order by plotting the fitted polynomials on top of the data. In order to see smooth curves, we evaluate the fitted polynomials on a grid of $x$ values (ranging between the largest and smallest of the inputs present in the dataset).

```python
def solve_poly_reg(x, y, max_order):
  """Fit a polynomial regression model for each order 0 through max_order.

  Args:
    x (ndarray): input vector of shape (n_samples)
    y (ndarray): vector of measurements of shape (n_samples)
    max_order (scalar): max order for polynomial fits

  Returns:
    dict: fitted weights for each polynomial model (dict key is order)
  """

  # Create a dictionary with polynomial order as keys,
  # and np array of theta_hat (weights) as the values
  theta_hats = {}

  # Loop over polynomial orders from 0 through max_order
  for order in range(max_order + 1):

    ##################################################################################
    ## TODO for students: Create design matrix and fit polynomial model for this order
    # Fill out function and remove
    raise NotImplementedError("Student exercise: fit a polynomial model")
    ##################################################################################

    # Create design matrix
    X_design = ...

    # Fit polynomial model
    this_theta = ...

    theta_hats[order] = this_theta

  return theta_hats

# Uncomment to test your function
max_order = 5
# theta_hats = solve_poly_reg(x, y, max_order)
# plot_fitted_polynomials(x, y, theta_hats)
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W1D3_ModelFitting/solutions/W1D3_Tutorial4_Solution_9b1eebd9.py)

*Example output:*

## Section 2.4: Evaluating fit quality

As with linear regression, we can compute mean squared error (MSE) to get a sense of how well the model fits the data.

We compute MSE as:

$$ MSE = \frac 1 N ||y - \hat y||^2 = \frac 1 N \sum_{i=1}^N (y_i - \hat y_i)^2 $$

where the predicted values for each model are given by $ \hat y = X \hat \theta$.

*Which model (i.e. which polynomial order) do you think will have the best MSE?*

### Exercise 4: Compute MSE and compare models

We will compare the MSE for different polynomial orders with a bar plot.

```python
mse_list = []
order_list = list(range(max_order + 1))

for order in order_list:

  X_design = make_design_matrix(x, order)

  ############################################################################
  # TODO for students: Compute MSE (fill in ... sections)
  #############################################################################

  # Get prediction for the polynomial regression model of this order
  y_hat = ...

  # Compute the residuals
  residuals = ...

  # Compute the MSE
  mse = ...

  mse_list.append(mse)

fig, ax = plt.subplots()

# Uncomment once above exercise is complete
# ax.bar(order_list, mse_list)

ax.set(title='Comparing Polynomial Fits', xlabel='Polynomial order', ylabel='MSE')
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W1D3_ModelFitting/solutions/W1D3_Tutorial4_Solution_ccea5dd7.py)

*Example output:*

---
# Summary

* Linear regression generalizes naturally to multiple dimensions
* Linear algebra affords us the mathematical tools to reason and solve such problems beyond the two dimensional case

* To change from a linear regression model to a polynomial regression model, we only have to change how the input data is structured

* We can choose the complexity of the model by changing the order of the polynomial model fit

* Higher order polynomial models tend to have lower MSE on the data they're fit with

**Note**: In practice, multidimensional least squares problems can be solved very efficiently (thanks to numerical routines such as LAPACK).

**Suggested readings**

[Introduction to Applied Linear Algebra – Vectors, Matrices, and Least Squares](http://vmls-book.stanford.edu/)
Stephen Boyd and Lieven Vandenberghe