# Homography decomposition
This notebook visualizes the decomposition of a homography into relative poses of cameras viewing the same plane.

**Subjects covered**
1. **Definition of homography and its uses**
2. **Using DLT to compute a homography from point correspondences**
3. **Proof of homography decomposition into poses**
4. **Conclusion**
5. **Sources**

**Why are we interested in these subjects?**
When learning about computer vision, the concept of a homography show up early and often. Homographies appear in many computer vision algorithms, toolsets, and theoretical proofs. They are often used as a tool to map planes to other planes and manipulate quadrilateral shapes. A practical example is the use of a homography to transform a picture of a magazine on a table into a "perspectiveless" frontal view of the magazine.

Similar to the fundamental matrix, a homography enables us to infer relative poses of cameras. If two cameras picture the same plane from different positions, the homography between the two images can be decomposed into the cameras' relative translation and rotation. Knowing the relative pose will allow for a more targeted search of point correspondences between images. It can also serve as an initial pose estimate when using iterative non-linear pose estimation methods (e.g. PnP Levenberg Marquardt in notebook 2).

**Disclaimer**
This notebook is mostly intended as thorough reproduction of the decomposition presented by Faugeras and Lustman in their paper [Motion and structure from motion in a piecewise planar environment](https://hal.inria.fr/file/index/docid/75698/filename/RR-0856.pdf). This notebook does provide definitions and simple examples of homographies. Nevertheless, those unfamiliar with homographies are advised to read more accessible resources on homographies, such as: [this tutorial with code examples](https://docs.opencv.org/4.x/d9/dab/tutorial_homography.html) by OpenCV, [these slides](http://www.cse.psu.edu/~rtc12/CSE486/lecture16.pdf) by Robert Collins, [these slides](http://www.cs.columbia.edu/~allen/F17/NOTES/homography_pka.pdf) collected by Peter Allen, [these linear algebra videos on linear transformations](https://www.3blue1brown.com/topics/linear-algebra) by Grant Sanderson.

## Homography
#### Defintion
A homography is an isomorphism: a structure-preserving mapping between two structures. Projective spaces are the structures that homographies operate on. For a projective space $\mathbb{P}^2$ it preserves concurrency, collinearity, cross-ratio and more (see [Multiple view geometry](https://www.robots.ox.ac.uk/~vgg/hzbook/) Ch.2.4.7). One may be familiar with projective spaces through the use of scale invariant homogenous 3 vectors $\begin{bmatrix} x\\ y\\ w \end{bmatrix}$ to represent quantities usually expressed as Euclidean 2 vectors $\begin{bmatrix} x\\ y \end{bmatrix}$ in computer vision. Here, scale invariance refers to the equivalence relation between scaled versions of the vector $\begin{bmatrix} x\\ y\\ w \end{bmatrix} = k \begin{bmatrix} x\\ y\\ w \end{bmatrix} \hspace{0.5em} \forall k $. A homography $\mathbf{H}$ operates on these projective points in these projective spaces. For example, $\mathbf{H}$ takes a set of points $x_i$ in $\mathbb{P}^2$ and maps them to a corresponding set of points $x'_i$ likewise in $\mathbb{P}^2$. That is,

$$ \begin{align}
    \mathbf{H}&: x \rightarrow x' \\ \\
         x'_1 &= \mathbf{H} \hspace{0.1em} x_1 \end{align}$$

The mapping preserves collinearity by definition, i.e. images of collinear points are themselves collinear i.e. straight lines are mapped to straight lines. A mapping with this property is called a collineation. A collineation is a linear transformation that preserves vector addition and scalar multiplication and can be implemented by a matrix:

$$ x'_i = \begin{bmatrix} h_{1} & h_{2} & h_{3} \\
                          h_{4} & h_{5} & h_{6} \\
                          h_{7} & h_{8} & h_{9}\end{bmatrix} \hspace{0.1em} x_1 $$
#### Degrees of freedom
A homography has $8$ degrees of freedom. While its matrix may suggest $9$, the scale invariance of $\mathbb{P}^2$ implies $\mathbf{H} = k \hspace{0.1em}\mathbf{H} \hspace{0.5em}\forall k $ which deducts a degree of freemdom, similar to how a point in $\mathbb{P}^2$ has $2$ degrees of freedom while it has $3$ parameters.

Besides homographies, there are other collineations with fewer degrees of freedom, each preserving a greater set of properties.

<center>

<br>
Figure taken from Zisserman, Hartley, Multiple view geometry in computer vision
</center>

## Interactive homography
Here we show an interactive widget that visualizes the different types of transformations.

```python
import numpy as np
from bqplot import pyplot as plt
from IPython.display import display
from ipywidgets import FloatSlider, IntSlider, HBox, VBox, Layout, HTML, ToggleButtons, Button, Output, Tab
from numpy import cos, sin
```

```python
# Callback functions.
# NOTE: All functions with parameters `args*` reference global variables.
# I know this is ugly; ipywidgets does not user defined parameters in callbacks.

def set_slider_visibility(desired_sliders, all_sliders):
    """ Updates the sliders exposed to a user after button toggle. """
    for slider in all_sliders:
        if slider in desired_sliders:
            slider.layout.visibility = "visible"
        else:
            slider.layout.visibility = "hidden"

def on_transform_click(*args):
    """ Sets appropriate displays when transformation type toggle button is clicked. """
    transform = transform_buttons.value
    desired_sliders = slider_map[transform]
    set_slider_visibility(desired_sliders, all_sliders)
    html_matrix_1.value = get_symbol_matrix_html(transform, all_sliders)
    html_matrix_2.value = get_symbol_matrix_html(transform, all_sliders, substitute=True)
    on_slider_drag()

def on_slider_drag(*args):
    """ Sets appropriate displays when a value slider is dragged. """
    scatter_orig = transform_fig_1.marks[0]
    lines_orig = transform_fig_1.marks[1]
    scatter_trans = transform_fig_2.marks[0]
    lines_trans = transform_fig_2.marks[1]
    matrix_values = get_matrix_values(transform_buttons.value, all_sliders)

    # Update matrix representation
    html_matrix_2.value = get_symbol_matrix_html(transform_buttons.value, all_sliders, substitute=True)
    html_matrix_3.value = get_value_matrix_html(matrix_values)

    # Update transformation visualization
    with scatter_trans.hold_sync(), lines_trans.hold_sync():
        xs, ys, zs = matrix_values @ np.vstack([scatter_orig.x, scatter_orig.y, np.ones_like(scatter_orig.x)])
        zs += 0.001  # avoid 0 div
        scatter_trans.x, scatter_trans.y = xs/zs, ys/zs

        # xs, ys, zs = matrix_values @ np.vstack([lines_orig.x, lines_orig.y, np.ones_like(lines_orig.x)])
        # zs += 0.001
        lines_trans.x, lines_trans.y = xs/zs, ys/zs

def on_reset_click(*args):
    """ Resets slider values when reset button is clicked. """
    for slider in all_sliders:
        slider.value = slider.default_value
```

```python
# Some verbose formatting functions are placed in notebook_functions.py to save visual space
from notebook_functions import (get_matrix_values,
                                get_symbol_matrix_html,
                                get_value_matrix_html,
                                get_quad_fig, get_html)

#Parameter sliders
rot_angle = FloatSlider(min=-np.pi, max=np.pi, step=0.005, description=r'rotation \(\theta\)')
x_offset = FloatSlider(min=-5, max=5, step=0.005, description=r'x offset  \(t_x\)')
y_offset = FloatSlider(min=-5, max=5, step=0.005,  description=r'y offset \(t_y\)')
scale = FloatSlider(min=-5, max=5, step=0.005, value=1, description=r'scale \(s\)')
x_scale = FloatSlider(min=-5, max=5, step=0.005, value=1, description=r'x scale \(s_x\)')
y_scale = FloatSlider(min=-5, max=5, step=0.005, value=1, description=r'y scale  \(s_y\)')
x_shear = FloatSlider(min=-5, max=5, step=0.005, description=r'x shear \(\alpha_x\)')
y_shear = FloatSlider(min=-5, max=5, step=0.005, description=r'y shear \(\alpha_y\)')
x_perspec = FloatSlider(min=-0.5, max=0.5, step=0.005, description=r'x perspective \(d_x\)', style={'description_width': 'initial'})
y_perspec = FloatSlider(min=-0.5, max=0.5, step=0.005, description=r'y perspective \(d_y\)', style={'description_width': 'initial'})

# Collections of parameter sliders belonging to transformations
euclidean = [rot_angle, x_offset, y_offset]
similarity = euclidean + [scale]
affine = euclidean + [x_shear, y_shear, x_scale, y_scale]
projective = affine + [x_perspec, y_perspec]
all_sliders = [rot_angle, x_offset, y_offset, scale, x_scale, y_scale, x_shear, y_shear, x_perspec, y_perspec]

# Set default values property for easy resetting
for slider in all_sliders:
    slider.default_value = slider.value

# Map from button names to slider sets
slider_map = {'Euclidean': euclidean,
              'Similarity': similarity,
              'Affine': affine,
              'Projective': projective}

# Initialize Widgets
default_transform = 'Projective'
matrix_values = get_matrix_values(default_transform, all_sliders)
html_matrix_1 = HTML(value=get_symbol_matrix_html(default_transform, all_sliders))
html_matrix_2 = HTML(value=get_symbol_matrix_html(default_transform, all_sliders, substitute=True))
html_matrix_3 = HTML(value=get_value_matrix_html(matrix_values))
reset_button = Button(description='Reset', button_style='info')
transform_buttons = ToggleButtons(options=['Projective', 'Affine', 'Similarity', 'Euclidean'])
transform_fig_1 = get_quad_fig('Original')
transform_fig_2 = get_quad_fig('Transformed')
transform_fig_1.background_style = {"fill": '#fafafa'}
set_slider_visibility(slider_map[default_transform], all_sliders)

# Set interactivitiy callbacks
reset_button.on_click(on_reset_click)
transform_buttons.observe(on_transform_click)
for slider in all_sliders:
    slider.observe(on_slider_drag)

# Build UI by stacking ipywidgets
sliders_1 = VBox(all_sliders[:4])
sliders_2 = VBox(all_sliders[4:8])
sliders_3 = VBox(all_sliders[8:])
sliders_ui = HBox([sliders_1, sliders_2, sliders_3], layout=Layout(margin='20px'))
math_ui = HBox([html_matrix_1, html_matrix_2, html_matrix_3])
transform_ui = HBox([transform_fig_1, transform_fig_2])
transform_buttons_ui = HBox([transform_buttons, reset_button])
ui = VBox([VBox([transform_buttons_ui, sliders_ui, math_ui]), transform_ui])
display(ui)
```

    VBox(children=(VBox(children=(HBox(children=(ToggleButtons(options=('Projective', 'Affine', 'Similarity', 'Euc…

## Direct Linear Transform (DLT) - Homography from point correspondences

The Direct Linear Transform (DLT) algorithm is similar to the 8 point algorithm used to estimate the fundamental matrix (in notebook 3).
The relationship between point sets defined by a homography $x' = \mathbf{H} \hspace{0.1em} x$ is used to build linear constraints from point correspondences.
Each point correspondence provides 2 constraints. To solve for the 8 degrees of freedom of a homography we require 4 point correspondences. The approach is as follows:

* Find 4 or more point correspondences and build constraints. Here we will use chessboard points.
* Use an SVD to find the least-squares solution to $\mathbf{H}$ given the constraints $x' = \mathbf{H}\hspace{0.1em} x$. That is, for SVD$= \mathbf {U\Sigma V^{T}}$, take the last column of $\mathbf{V}$ as $\mathbf{H}$. *

*-See the appendix of notebook 3 for an intuitive explanation and proof of why the Singular Value Decomposition can be used to find a least-squares solution.

$\begin{align}
x' &= \mathbf{H} \hspace{0.1em} x  \hspace{0.5em}
\text{For a single piont correspondence }\\ \\
\begin{bmatrix}
x' \\ y' \\ w'
\end{bmatrix} &=
\begin{bmatrix}
h_1 & h_2 & h_3 \\
h_4 & h_5 & h_6 \\
h_7 & h_8 & h_9
\end{bmatrix}
\begin{bmatrix}
x \\ y \\ w
\end{bmatrix} \\ \\
x' &=  \frac{h_1 x + h_2 y + h_3}{h_7 x + h_8 y + h_9}  \hspace{0.5em}
\text{Given } w=1 \\ \\
y' &=  \frac{h_4 x + h_5 y + h_6}{h_7 x + h_8 y + h_9} \\
\\
x'(h_7 x + h_8 y + h_9) - h_1 x - h_2 y - h_3 &= 0 \\ \\
y'(h_4 x + h_5 y + h_6) - h_7 x - h_8 y - h_9 &= 0 \\ \\
\\
\\
\begin{bmatrix}
-x & -y & -1 & 0 & 0 & 0 & xx' & yx' & x' \\
 0 & 0 & 0 & -x & -y & -1 & xy' & yy' & y'
\end{bmatrix}
\begin{bmatrix}
h_1 \\ h_2 \\ h_3 \\ h_4 \\ h_5 \\ h_6 \\ h_7 \\ h_8 \\ h_9
\end{bmatrix} &= \begin{bmatrix} 0 \\ 0 \end{bmatrix}
\end{align}$

```python
  def compute_homography(points_1, points_2):
    constraint_matrix = list()
    for point_1, point_2 in zip(points_1.T, points_2.T):
        x1, y1 = point_1[:2]
        x2, y2 = point_2[:2]
        constraint_1 = [-x1, -y1, -1,   0,   0,  0, x1 * x2, y1 * x2, x2]
        constraint_2 = [  0,   0,  0, -x1, -y1, -1, x1 * y2, y1 * y2, y2]
        constraint_matrix.append(constraint_1)
        constraint_matrix.append(constraint_2)

    constraint_matrix = np.array(constraint_matrix)
    u, s, v_t = np.linalg.svd(constraint_matrix)
    least_squares = v_t[-1] # last column of V
    H = least_squares.reshape((3, 3))
    H = H / H[2, 2]
    return H
```

# Interactive  homography
Here we will interactively plot the homography between two views of a plane

```python
%matplotlib notebook
from notebook_functions import (plot_camera_wireframe,
                                project_points_to_picture,
                                get_rotation_matrix)
import ipyvolume as ipv
import numpy as np
import cv2
import matplotlib.pyplot as plt
import matplotlib
```

```python
class Camera:
    def __init__(self, color='blue'):
        self.extrinsics = np.identity(4)
        self.color = color

    def update_extrinsics(self, x, y, z, rot_x, rot_y, rot_z):
        rot = get_rotation_matrix(rot_x, rot_y, rot_z)
        trans = -np.array([x, y, z])
        self.extrinsics =  np.vstack([np.hstack([rot, trans[:, None]]), [0, 0, 0, 1.0]])

    def change_of_coords(self, coc):
        self.extrinsics = coc @ self.extrinsics

    @property
    def trans(self):
        return self.extrinsics[:3, 3]

    @property
    def rot(self):
        return self.extrinsics[:3, :3]

    @property
    def inv_extrinsics(self):
        return np.linalg.pinv(self.extrinsics)

    @property
    def center(self):
        return (-self.rot.T) @ self.trans

    @property
    def border(self):
        inv_extrinsics = self.inv_extrinsics
        lt = inv_extrinsics @ [-1 , 2/3 , 2/3, 1]
        rt = inv_extrinsics @ [1  , 2/3 , 2/3, 1]
        lb = inv_extrinsics @ [-1 , -2/3, 2/3, 1]
        rb = inv_extrinsics @ [1  , -2/3, 2/3, 1]

        return [lt, rt, lb, rb]

    def project(self, objects):
        projected = self.extrinsics @ objects
        projected /= projected[2]
        return projected[:3]

    def plot_wireframe(self):
        x, y, z = self.center[:3].flatten().astype(float)
        [lt, rt, lb, rb] = self.border

        # Connect camera projective center to wireframe extremities
        self.p1 = ipv.plot([x, lt[0]], [y, lt[1]], [z, lt[2]], color=self.color)
        self.p2 = ipv.plot([x, rt[0]], [y, rt[1]], [z, rt[2]], color=self.color)
        self.p3 = ipv.plot([x, lb[0]], [y, lb[1]], [z, lb[2]], color=self.color)
        self.p4 = ipv.plot([x, rb[0]], [y, rb[1]], [z, rb[2]], color=self.color)

        # Connect wireframe corners with a rectangle
        self.p5 = ipv.plot([lt[0], rt[0]], [lt[1], rt[1]], [lt[2], rt[2]], color=self.color)
        self.p6 = ipv.plot([rt[0], rb[0]], [rt[1], rb[1]], [rt[2], rb[2]], color=self.color)
        self.p7 = ipv.plot([rb[0], lb[0]], [rb[1], lb[1]], [rb[2], lb[2]], color=self.color)
        self.p8 = ipv.plot([lb[0], lt[0]], [lb[1], lt[1]], [lb[2], lt[2]], color=self.color)

        self.p9 = ipv.scatter(np.array([x]), np.array([y]), np.array([z]), size=1, marker="sphere", color=self.color)

    def update_wireframe(self):
        [lt, rt, lb, rb] = self.border

        with self.p1.hold_sync(), self.p2.hold_sync(), self.p3.hold_sync(), self.p4.hold_sync(), self.p5.hold_sync(), self.p6.hold_sync(), self.p7.hold_sync(), self.p8.hold_sync(), self.p9.hold_sync():
            x, y, z = self.center[:3].flatten().astype(float)

            [lt, rt, lb, rb] = self.border
            self.p1.x, self.p1.y, self.p1.z = [x, lt[0]], [y, lt[1]],  [z, lt[2]]
            self.p2.x, self.p2.y, self.p2.z = [x, rt[0]], [y, rt[1]],  [z, rt[2]]
            self.p3.x, self.p3.y, self.p3.z = [x, lb[0]], [y, lb[1]], [z, lb[2]]
            self.p4.x, self.p4.y, self.p4.z = [x, rb[0]], [y, rb[1]], [z, rb[2]]
            self.p5.x, self.p5.y, self.p5.z = [lt[0], rt[0]], [lt[1], rt[1]], [lt[2], rt[2]]
            self.p6.x, self.p6.y, self.p6.z = [rt[0], rb[0]], [rt[1], rb[1]], [rt[2], rb[2]]
            self.p7.x, self.p7.y, self.p7.z = [rb[0], lb[0]], [rb[1], lb[1]], [rb[2], lb[2]]
            self.p8.x, self.p8.y, self.p8.z = [lb[0], lt[0]], [lb[1], lt[1]], [lb[2], lt[2]]
            self.p9.x, self.p9.y, self.p9.z = np.array([x]), np.array([y]), np.array([z])
```

```python
import ipyvolume as ipv
import numpy as np
import matplotlib
from notebook_functions import (plot_camera_wireframe,
                                update_camera_wireframe,
                                project_points_to_picture,
                                get_rotation_matrix,
                                triangulate)
import copy

rot_x_1 = FloatSlider(min=-np.pi, max=np.pi, step=0.005, description=r'\(\theta_x\)')
rot_y_1 = FloatSlider(min=-np.pi, max=np.pi, step=0.005, description=r'\(\theta_y\)')
rot_z_1 = FloatSlider(min=-np.pi, max=np.pi, step=0.005, description=r'\(\theta_z\)')
rot_x_2 = FloatSlider(min=-np.pi, max=np.pi, step=0.005, description=r'\(\theta_x\)')
rot_y_2 = FloatSlider(min=-np.pi, max=np.pi, step=0.005, description=r'\(\theta_y\)')
rot_z_2 = FloatSlider(min=-np.pi, max=np.pi, step=0.005, description=r'\(\theta_z\)')

x_trans_1 = FloatSlider(min=-5, max=5, step=0.005, description=r'x translation  \(t_x\)')
y_trans_1 = FloatSlider(min=-5, max=5, step=0.005,  description=r'y translation \(t_y\)')
z_trans_1 = FloatSlider(min=-5, max=5, step=0.005,  description=r'y translation \(t_y\)')
x_trans_2 = FloatSlider(min=-5, max=5, step=0.005, description=r'x translation  \(t_x\)')
y_trans_2 = FloatSlider(min=-5, max=5, step=0.005,  description=r'y translation \(t_y\)')
z_trans_2 = FloatSlider(min=-5, max=5, step=0.005,  description=r'y translation \(t_y\)')

sliders_1 = [rot_x_1, rot_y_1, rot_z_1, x_trans_1, y_trans_1, z_trans_1]
sliders_2 = [rot_x_2, rot_y_2, rot_z_2, x_trans_2, y_trans_2, z_trans_2]
all_sliders = sliders_1 + sliders_2

camera_view_1 = get_quad_fig('Camera 1 View')
camera_view_2 = get_quad_fig('Camera 2 View')
homography_fig = get_quad_fig('Camera 2 according to homography')

ipv_fig = ipv.figure()
ipv.xlim(-5, 5)
ipv.ylim(-5, 5)
ipv.zlim(-5, 5)

x_coords = np.array([-1, 1, 1, -1]).astype(float)
y_coords = np.array([1, 1, -1, -1]).astype(float)
z_coords = np.array([1, 1, 1, 1]).astype(float)
objects = np.vstack([x_coords, y_coords, z_coords, np.ones_like(x_coords)])
colors = [matplotlib.colors.to_rgba(name) for name in ['salmon', 'lawngreen', 'deepskyblue', 'orange']]
ipv.scatter(x_coords, y_coords, z_coords, marker='sphere', color=colors)
ipv.plot(np.pad(x_coords, 1, 'wrap'), np.pad(y_coords, 1, 'wrap'), np.pad(z_coords, 1, 'wrap'), color='deepskyblue')

cam_1 = Camera('blue')
cam_2 = Camera('green')
cam_1.plot_wireframe()
cam_2.plot_wireframe()

def update_homography():
    """ Updates the 2D scatter plot and lines visualizing the homography """
    scatter = homography_fig.marks[0]
    lines = homography_fig.marks[1]

    points_1 = cam_1.project(objects)
    points_2 = cam_2.project(objects)
    homography = compute_homography(points_1, points_2)
    homography_html.value = get_value_matrix_html(homography)

    points_2_h = homography @ points_1[:3]
    points_2_h = points_2_h / points_2_h[2]
    scatter.x, scatter.y = -points_2_h[0], points_2_h[1]
    lines.x, lines.y = -points_2_h[0], points_2_h[1]

    return homography

homography = None

def on_slider_drag(*args):
    """ Sets appropriate camera position. """
    global homography
    if sliders_tab.selected_index == 0:
        cam = cam_1
        scatter = camera_view_1.marks[0]
        lines = camera_view_1.marks[1]
        sliders = sliders_1
    else:
        cam = cam_2
        scatter = camera_view_2.marks[0]
        lines = camera_view_2.marks[1]
        sliders = sliders_2

    rot_x, rot_y, rot_z, x, y, z = [slider.value for slider in sliders]
    cam.update_extrinsics(x, y, z, rot_x, rot_y, rot_z)
    cam.update_wireframe()

    # project to image
    with scatter.hold_sync(), lines.hold_sync():
        xs, ys, zs = cam.project(objects)
        scatter.x, scatter.y = -xs, ys
        lines.x, lines.y = -xs, ys

    homography = update_homography()

for slider in all_sliders:
    slider.observe(on_slider_drag)

slider_box_1 = VBox(sliders_1[:3])
slider_box_2 = VBox(sliders_1[3:])
slider_box_3 = VBox(sliders_2[:3])
slider_box_4 = VBox(sliders_2[3:])
sliders_ui_1 = HBox([slider_box_1, slider_box_2], layout=Layout(margin='20px'))
sliders_ui_2 = HBox([slider_box_3, slider_box_4], layout=Layout(margin='20px'))
sliders_tab = Tab([sliders_ui_1, sliders_ui_2], _titles={0:'Camera 1', 1:'Camera 2'})
homography_html = HTML(value=get_value_matrix_html(np.zeros([3, 3])))

space_vis = HBox([ipv_fig, camera_view_1, camera_view_2, homography_fig])
ui = VBox([sliders_tab, space_vis, homography_html])
display(ui)
```

    VBox(children=(Tab(children=(HBox(children=(VBox(children=(FloatSlider(value=0.0, description='\\(\\theta_x\\)…

## Homography decomposition

### Symbols
* $\mathbf{n}$ - Normal to the plane
* $\mathbf{R}$ - Rotation that brings cam 1 to cam 2 orientation
* $\mathbf{t}$ - Translation that brings cam 1 to cam 2 position
* $d$ - Distance from camera one to the plane
* $\mathbf{A} = d \mathbf{R} + \mathbf{tn}^t$
* $\mathbf{X_2} = \mathbf{A} \mathbf{X_1}$
* $\mathbf{X_2} = (\mathbf{R} + \frac{\mathbf{tn}^t}{d}) \mathbf{X_1})$
* SVD: $\mathbf{A} = \mathbf{U} \mathbf{\Lambda} \mathbf{V^t} $
* $\mathbf{\Lambda} = d' \mathbf{R'} + \mathbf{t'n'^t}$
* $\mathbf{\Lambda} = \begin{bmatrix} \sigma_1 & 0        & 0 \\
                                      0        & \sigma_2 & 0 \\
                                      0        & 0        & \sigma_3 \end{bmatrix}$

Relation between variables and primes
variables:
* $\mathbf{R} = \mathbf{sUR'V^t}$
* $\mathbf{t} = \mathbf{Ut'}$
* $\mathbf{n} = \mathbf{Vn'}$
* $d = sd'$
* $s = det \hspace{0.1em}\mathbf{U} \hspace{0.2em}det\hspace{0.1em} \mathbf{V}$

primes:
* $\mathbf{R'} = s \mathbf{U^tRV}$
* $\mathbf{t'} = \mathbf{U^t t}$
* $\mathbf{n'} = \mathbf{V^t n}$
* $d' = sd$

### Solutions
solutions for singular values that are all unique (multiplicity $\Lambda$= 3)

$\epsilon_1 \pm 1$
$\epsilon_2 \pm 1$

$n_1 = \epsilon_1 \frac{\sqrt{\sigma_1^2 - \sigma_2^2}}{\sigma_1^2 - \sigma^2_3}$
$n_2 = 0 $
$n_3 = \epsilon_3 \frac{\sqrt{\sigma_2^2 - \sigma_3^2}}{\sigma_1^2 - \sigma^2_3}$

<br>
<br>

Case $d' > 0$:

$\mathbf{R'} = \begin{bmatrix} \cos \theta & 0 & - \sin \theta \\
                               0           & 1 & 0             \\
                               \sin \theta & 0 & \cos \theta
               \end{bmatrix}$

$ \sin \theta = (\sigma_1 - \sigma_3) \frac{n_1 n_3}{\sigma_2} $

$ \cos \theta = \frac{\sigma_1 n^2_3 + \sigma_3 n^2_1}{\sigma_2} $

$\mathbf{t'} = (\sigma_1 - \sigma_3) \begin{bmatrix} n_1 \\ 0 \\ n_3 \end{bmatrix}$

<br>
<br>

Case $d' < 0$:

$\mathbf{R'} = \begin{bmatrix} \cos \theta & 0  & \sin \theta \\
                               0           & -1 & 0           \\
                               \sin \theta & 0  & - \cos \theta
               \end{bmatrix}$

$ \sin \theta = \frac{\sigma_1 + \sigma_3}{\sigma_2}n_1 n_3$

$ \cos \theta = \frac{\sigma_3 n_1^2 - \sigma_1 n_3^2}{\sigma_2}$

$\mathbf{t'} = (\sigma_1 + \sigma_3) \begin{bmatrix} n_1 \\ 0 \\ n_3 \end{bmatrix}$

```python
def decompose_homography(homography):
    u, s, v_t = np.linalg.svd(homography)
    v = v_t.T
    s1, s2, s3 = s

    s = np.linalg.det(u) * np.linalg.det(v)

    if not (s1 != s2  and s2 != s3):
        print('singular values are not unique')
        return None

    n1 = np.sqrt((s1**2 - s2**2) / (s1**2 - s3**2))
    n2 = 0
    n3 = np.sqrt((s2**2 - s3**2) / (s1**2 - s3**2))

    sign_permutations = [(1, 1), (1, -1), (-1, 1), (-1, -1)]

    rots = list()
    transs = list()
    normals = list()
    distances = list()
    change_of_coords = list()

    for d_positive in (True, False):
        for (e1, e3) in sign_permutations:

            n1e = e1 * n1
            n3e = e3 * n3
            normal_prime = np.array([n1e, 0, n3e])

            if d_positive:
                distance = s * s2
                sin_theta = (s1 - s3) * ((n1e * n3e) / s2)
                cos_theta = (s1 * n3e**2 + s3 * n1e**2) / s2

                rot_prime = np.array([[cos_theta, 0, -sin_theta],
                                      [0        , 1, 0         ],
                                      [sin_theta, 0, cos_theta ]])

                trans_prime = (s1 - s3) * np.array([n1e, 0 , n3e])

            else:
                distance =  s * -s2
                sin_theta = ((s1 + s3) / s2) * n1e * n3e
                cos_theta = (s3 * n1e**2 - s1 * n3e**2) / s2

                rot_prime = np.array([[cos_theta, 0 , sin_theta  ],
                                      [0        , -1, 0          ],
                                      [sin_theta, 0 , -cos_theta ]])

                trans_prime = (s1 + s3) * np.array([n1e, 0 , n3e])

            rot = s * u @ rot_prime @ v.T
            trans = u @ trans_prime
            normal = v @ normal_prime

            rots.append(rot)
            transs.append(trans)
            normals.append(normal)
            distances.append(distance)

            coc = np.vstack([np.hstack([rot, trans[:, None]]), [0, 0, 0, 1]])
            change_of_coords.append(coc)

    return change_of_coords, normals, distances
```

# Apply decomposition and test results

```python
change_of_coords, normals, distances = decompose_homography(homography)

test_1_scores = list()
points_1 = cam_1.project(objects)

# Test 1. Does the homography lead to a positive z coordinate
for coc, plane_distance in zip(change_of_coords, distances):
    points_2 = homography @ points_1
    zs = points_2[2] / plane_distance
    visible = (zs > 0).sum()
    test_1_scores.append(visible)

# Sort decompositions highest to lowest.
decompositions = sorted(zip(test_1_scores, change_of_coords, normals, distances), reverse=True, key=lambda x: x[0])
# Keep only the top 4 scoring decompositions
decompositions = decompositions[:4]
_, change_of_coords, normals, distances = list(zip(*decompositions))

ones = np.ones_like(objects[0])[None, :]
points_1 = np.vstack([points_1, ones])
test_2_scores = list()

# Test 2. Are the points observed on the front of the plane? (Does the normal of the plane point in the same direction as the unprojected point)
for coc, normal, plane_distance in zip(change_of_coords, normals, distances):
    vecs = cam_1.inv_extrinsics @ points_1
    dot_products = np.dot(normal, vecs[:3]) / plane_distance
    test_2_scores.append((dot_products > 0).sum())

decompositions = sorted(zip(test_2_scores, change_of_coords, normals, distances), reverse=True, key=lambda x: x[0])
# Keep only the top 2 scoring decompositions
decompositions = decompositions[:2]
test_2_scores, change_of_coords, normals, distances = list(zip(*decompositions))

## If the ratio is good, keep the solution
# ratio = test_2_scores[1] / test_2_scores[0]
# print(ratio)
# # if ratio < 0.9:
# coc[:3, 3] = coc[:3, 3] / np.linalg.norm(coc[:3, 3])
# coc[:3, 3] *= np.linalg.norm(cam_1.trans - cam_2.trans)
trans_cam = Camera('red')
trans_cam.extrinsics = cam_1.extrinsics # Make a copy of cam 1
trans_cam.change_of_coords(coc)
trans_cam.plot_wireframe()
```

# Conclusion

The decomposition is a succes, the pose of ground truth camera 2 is the same as the decomposition.
But notice how the scale of the translation is unknown, as also with the fundamental matrix decomposition (in notebook 3).

## Sources

1. [Zisserman, Richard Hartley Andrew. "Multiple view geometry in computer vision." 2004.](https://www.robots.ox.ac.uk/~vgg/hzbook/)
2. [Elan Dubrofsky, "Homography Estimation." 2009.](http://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.186.4411)

Todo, check this out https://users.cecs.anu.edu.au/~hartley/Papers/CVPR99-tutorial/tutorial.pdf


# Grain Size from Mean Linear Intercept Distance

Håkon Wiik Ånes (hakon.w.anes@ntnu.no)

2019-04-29
___

Following Humphreys *et al.* (2017), for a single-phase material of uniform grain size, the grain size may be related to $\bar{L}$, the mean linear intercept distance (MLI) across a grain as measured on a plane section . $\bar{L}$ is determined by measuring the average number of boundary intercepts ($N_{\mathrm{L}}$) with a line of unit length so that

\begin{equation}
    \bar{L} = \frac{1}{N_{\mathrm{L}}}.
\end{equation}

```python
# Comment to use default next line to use default plotting backend
%matplotlib qt
import numpy as np
import matplotlib.pyplot as plt
plt.rcParams['figure.dpi'] = 200
```

    Warning: Cannot change to a different GUI toolkit: qt. Using qt5 instead.

```python
fname = 'plm_5x.tif'
img = plt.imread(fname)

# Plotting inline just for illustrative purposes, use different backend
# (for example qt or tkAgg) when measuring number of intercepts.
plt.imshow(img)
plt.grid(color='r')
```

```python
scale5x = 0.68  # [um/px]
line_x = img.shape[1] * scale5x  # [um]
line_y = img.shape[0] * scale5x  # [um]

intercepts_x = np.array([29, 32, 26, 30, 26, 23, 31])
mli_x = line_x/np.mean(intercepts_x)  # [um]

print("Grain size from mean linear intercept distance (MLI): {:.2f} um".format(mli_x))
```

    Grain size from mean linear intercept distance (MLI): 61.86 um

## References

Humphreys, Frederick John, Rohrer, Gregory S., Rollett, Anthony D.: **Recrystallization and Related Annealing Phenomena**, 3rd edition, Elsevier, pp. 638, 2017.


# A self-contained tutorial on Lagrangian Neural Networks (LNNs)
Sam Greydanus (this notebook), Miles Cranmer (derivation), Stephan Hoyer (derivation, code)

$\quad$

$\quad$

## Summary

We propose Lagrangian Neural Networks (LNNs), which harness the expressive power of neural networks to learn arbitrary Lagrangians from data. In contrast to [Hamiltonian Neural Networks](https://greydanus.github.io/2019/05/15/hamiltonian-nns/), these models do not require canonical coordinates and perform well in situations where generalized momentum is difficult to compute (eg the [double pendulum](https://diego.assencio.com/?index=e5ac36fcb129ce95a61f8e8ce0572dbf)). This is particularly appealing for use with a learned latent representation, though this application is a work in progress. Unlike [previous work on learning Lagrangians](https://arxiv.org/pdf/1907.04490.pdf), LNNs can learn arbitrary Lagrangians.

|	| Neural Networks  | [Neural ODEs](https://arxiv.org/abs/1806.07366) | [HNNs](https://arxiv.org/abs/1906.01563)  | [DeLaN (ICLR'19)](https://arxiv.org/abs/1907.04490) | LNNs (this work) |
| ------------- |:------------:| :------------:| :------------:| :------------:| :------------:|
| Can model dynamical systems | ✔ | ✔ | ✔ | ✔ | ✔ |
| Learns differential equations | | ✔ | ✔ | ✔ | ✔ |
| Learns exact conservation laws | | | ✔ | ✔ | ✔ |
| Learns from arbitrary coords. |✔ | ✔|| ✔ | ✔ |
| Learns arbitrary Lagrangians | | |  | | ✔ |

## Lagrangians
Following David Morin's [excellent introduction to Lagrangians](http://www.people.fas.harvard.edu/~djmorin/chap6.pdf), we define the Lagrangian as a scalar function $L\equiv T-V$ with $T$ and $V$ as kinetic and potential energies, respectively. It turns out that we can obtain the equations of motion of an arbitrary physical system using the relationship
$$
\frac{d}{dt} \frac{\partial L}{\partial \dot q_j} = \frac{\partial L}{\partial q_j}
$$
We call this the _Euler-Lagrange (E-L) equation_ and we obtain it by requiring that the the quantity
$$S \equiv \int_{t_1}^{t_2} L(q, \dot q, t) dt$$
have a stationary value. In other words, we require that $S$ be at a local minimum, maximum, or saddle point. We call $S$ the "action" of the system, so the E-L relationship is a mathematical tool for describing its "path of stationary action". One way to think of $S$ is as the cost function that Nature minimizes in order to give us the dynamics of the natural world.

## **Step 1: Basic imports**
We will implement this project in [JAX](https://github.com/google/jax).

```python
#!c1.4
%pip install moviepy
```

    Defaulting to user installation because normal site-packages is not writeable
    Collecting moviepy
      Downloading moviepy-1.0.3.tar.gz (388 kB)
         |████████████████████████████████| 388 kB 3.7 MB/s
    Collecting decorator<5.0,>=4.0.2
      Downloading decorator-4.4.2-py2.py3-none-any.whl (9.2 kB)
    Requirement already satisfied: tqdm<5.0,>=4.11.2 in /usr/local/lib/python3.7/dist-packages (from moviepy) (4.50.0)
    Requirement already satisfied: requests<3.0,>=2.8.1 in /kernel/lib/python3.7/site-packages (from moviepy) (2.25.1)
    Collecting proglog<=1.0.0
      Downloading proglog-0.1.9.tar.gz (10 kB)
    Requirement already satisfied: numpy>=1.17.3 in /kernel/lib/python3.7/site-packages (from moviepy) (1.19.4)
    Collecting imageio<3.0,>=2.5
      Downloading imageio-2.9.0-py3-none-any.whl (3.3 MB)
         |████████████████████████████████| 3.3 MB 10.1 MB/s
    Collecting imageio_ffmpeg>=0.2.0
      Downloading imageio_ffmpeg-0.4.3-py3-none-manylinux2010_x86_64.whl (26.9 MB)
         |████████████████████████████████| 26.9 MB 9.0 MB/s
    Requirement already satisfied: pillow in /kernel/lib/python3.7/site-packages (from imageio<3.0,>=2.5->moviepy) (8.2.0)
    Requirement already satisfied: chardet<5,>=3.0.2 in /kernel/lib/python3.7/site-packages (from requests<3.0,>=2.8.1->moviepy) (4.0.0)
    Requirement already satisfied: idna<3,>=2.5 in /kernel/lib/python3.7/site-packages (from requests<3.0,>=2.8.1->moviepy) (2.10)
    Requirement already satisfied: certifi>=2017.4.17 in /kernel/lib/python3.7/site-packages (from requests<3.0,>=2.8.1->moviepy) (2020.12.5)
    Requirement already satisfied: urllib3<1.27,>=1.21.1 in /kernel/lib/python3.7/site-packages (from requests<3.0,>=2.8.1->moviepy) (1.26.4)
    Building wheels for collected packages: moviepy, proglog
      Building wheel for moviepy (setup.py) ... done
      Created wheel for moviepy: filename=moviepy-1.0.3-py3-none-any.whl size=110726 sha256=f3a3096b7fc7c1e1a5ffdef5f4af3ad95427b6bfe23634298027e878b9317732
      Stored in directory: /home/jupyter/.cache/pip/wheels/56/dc/2b/9cd600d483c04af3353d66623056fc03faed76b7518faae4df
      Building wheel for proglog (setup.py) ... done
      Created wheel for proglog: filename=proglog-0.1.9-py3-none-any.whl size=6147 sha256=5a8b78d281d1209bce8598b9a304a02c8c9b7b89a2a0559376575fad73b6a786
      Stored in directory: /home/jupyter/.cache/pip/wheels/12/36/1f/dc61e6ac10781d63cf6fa045eb09fa613a667384e12cb6e6e0
    Successfully built moviepy proglog
    Installing collected packages: proglog, imageio-ffmpeg, imageio, decorator, moviepy
      WARNING: The scripts imageio_download_bin and imageio_remove_bin are installed in '/home/jupyter/.local/bin' which is not on PATH.
      Consider adding this directory to PATH or, if you prefer to suppress this warning, use --no-warn-script-location.
    Successfully installed decorator-4.4.2 imageio-2.9.0 imageio-ffmpeg-0.4.3 moviepy-1.0.3 proglog-0.1.9

```python
import jax
import jax.numpy as jnp
import numpy as np
from jax.experimental.ode import odeint
import matplotlib.pyplot as plt
from functools import partial # reduces arguments to function by making some subset implicit

from jax.experimental import stax
from jax.experimental import optimizers

# visualization
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from moviepy.editor import ImageSequenceClip
from functools import partial
import proglog
from PIL import Image
```

### Make sure we're using a GPU

```python
#!g1.1
from jax.lib import xla_bridge
print(xla_bridge.get_backend().platform)
```

    gpu

## **Step 2: Model the Double Pendulum**
Following [Diego Assencio's derivation](https://diego.assencio.com/?index=1500c66ae7ab27bb0106467c68feebc6), the Lagrangian of the double pendulum is

$$
\begin{align}
L = & \frac{1}{2}(m_1 + m_2) l_1^2 \dot{\theta}_1^2 +
	\frac{1}{2}m_2 l_2^2 \dot{\theta}_2^2 + m_2l_1l_2\dot{\theta}_1\dot{\theta}_2
	\cos(\theta_1 - \theta_2)\nonumber\\[3pt]
     &+ (m_1 + m_2) g l_1 \cos\theta_1 + m_2 g l_2\cos\theta_2
\end{align}
$$

It's worth noting that the first line in the equation above corresponds to kinetic energy (T) and the second line corresponds to negative potential energy (-V).

```python
def lagrangian(q, q_dot, m1, m2, l1, l2, g):
  t1, t2 = q     # theta 1 and theta 2
  w1, w2 = q_dot # omega 1 and omega 2

  # kinetic energy (T)
  T1 = 0.5 * m1 * (l1 * w1)**2
  T2 = 0.5 * m2 * ((l1 * w1)**2 + (l2 * w2)**2 +
                    2 * l1 * l2 * w1 * w2 * jnp.cos(t1 - t2))
  T = T1 + T2

  # potential energy (V)
  y1 = -l1 * jnp.cos(t1)
  y2 = y1 - l2 * jnp.cos(t2)
  V = m1 * g * y1 + m2 * g * y2

  return T - V
```

We'll want to check our results against an analytical solution. We can use equations 19-22 and 28 from [Diego's derivation](https://diego.assencio.com/?index=1500c66ae7ab27bb0106467c68feebc6):

$$
\begin{align}
& \alpha_1(\theta_1,\theta_2) ~:=~ \displaystyle\frac{l_2}{l_1}\left(\frac{m_2}{m_1 + m_2}\right)\cos(\theta_1 - \theta_2)\\
&\alpha_2(\theta_1,\theta_2) ~:=~ \frac{l_1}{l_2}\cos(\theta_1-\theta_2)\\
&\displaystyle f_1(\theta_1, \theta_2, \dot{\theta}_1, \dot{\theta}_2) ~:=~
-\frac{l_2}{l_1}\left(\frac{m_2}{m_1+m_2}\right) \dot{\theta}_2^2\sin(\theta_1 - \theta_2)
- \frac{g}{l_1} \sin\theta_1 \\
&\displaystyle f_2(\theta_1, \theta_2, \dot{\theta}_1, \dot{\theta}_2) ~:=~
\frac{l_1}{l_2}\dot{\theta}_1^2\sin(\theta_1-\theta_2) - \frac{g}{l_2} \sin\theta_2 \\
&g_1 := \displaystyle\frac{f_1 - \alpha_1 f_2}{1 - \alpha_1\alpha_2}
\quad\quad
g_2 := \displaystyle\frac{-\alpha_2 f_1 + f_2}{1 - \alpha_1\alpha_2}
\end{align}
$$

Having defined these quantities, we can express the analytical dynamics of the system in matrix form:

$$
\displaystyle\frac{d}{dt}
\left( \begin{matrix} \theta_1 \\[1pt] \theta_2 \\[1pt] \omega_1 \\[1pt] \omega_1 \end{matrix} \right)
=
\left( \begin{matrix} \omega_1 \\ \omega_2 \\ g_1(\theta_1,\theta_2,\omega_1,\omega_2)
\\ g_2(\theta_1,\theta_2,\omega_1,\omega_2) \end{matrix} \right)
$$

```python
def f_analytical(state, t=0, m1=1, m2=1, l1=1, l2=1, g=9.8):
  t1, t2, w1, w2 = state
  a1 = (l2 / l1) * (m2 / (m1 + m2)) * jnp.cos(t1 - t2)
  a2 = (l1 / l2) * jnp.cos(t1 - t2)
  f1 = -(l2 / l1) * (m2 / (m1 + m2)) * (w2**2) * jnp.sin(t1 - t2) - \
      (g / l1) * jnp.sin(t1)
  f2 = (l1 / l2) * (w1**2) * jnp.sin(t1 - t2) - (g / l2) * jnp.sin(t2)
  g1 = (f1 - a1 * f2) / (1 - a1 * a2)
  g2 = (f2 - a2 * f1) / (1 - a1 * a2)
  return jnp.stack([w1, w2, g1, g2])
```

## **Step 3: Obtain dynamics from the Lagrangian**
We have written our Lagrangian and now we want to use it to obtain the dynamics of the system. Physicists would ordinarily try applying the Euler-Lagrange relationship in order to compute a set of analytical equations of motion. In this situation, though, our Lagrangian is a neural network. Trying to get an analytic solution from a neural network Lagrangian would be like trying to knit with a plate of spaghetti.

A better idea is to obtain the dynamics through numerical integration, and to do this we'll need an expression for $\ddot q$. Let's begin with the E-L constraint and try to solve for $\ddot q$.

$$
\begin{align}
\frac{d}{dt} \frac{\partial \mathcal{L}}{\partial \dot q_j} &= \frac{\partial \mathcal{L}}{\partial q_j} & \text{write down the Euler-Lagrange equation} \quad (1)\\
\frac{d}{dt} \nabla_{\dot q} \mathcal{L} &= \nabla_{q} \mathcal{L} & \text{switch to vector notation} \quad (2)\\
(\nabla_{\dot q}\nabla_{\dot q}^{\top}\mathcal{L})\ddot q + (\nabla_{q}\nabla_{\dot q}^{\top}\mathcal{L}) \dot q &= \nabla_q \mathcal{L} & \text{expand the time derivative }\frac{d}{dt} \quad (3)\\
\ddot q &= (\nabla_{\dot q}\nabla_{\dot q}^{\top}\mathcal{L})^{-1}[\nabla_q \mathcal{L} - (\nabla_{q}\nabla_{\dot q}^{\top}\mathcal{L})\dot q] & \text{use a matrix inverse to solve for } \ddot q \quad (4)\\
\end{align}
$$

**Implementation.** If you look closely at Equation 4, you may notice that it involves both the Hessian and the gradient of a neural network during the forward pass of the LNN. This is not a trivial operation, but modern automatic differentiation makes things surprisingly smooth. Written in JAX, Equation 4 is just a few lines of code
```
q_tt = (jnp.linalg.pinv(jax.hessian(lagrangian, 1)(q, q_t)) @
        (jax.grad(lagrangian, 0)(q, q_t)
          - jax.jacobian(jax.jacobian(lagrangian, 1), 0)(q, q_t) @ q_t))
```

Big props to Stephan Hoyer for figuring this out. In order to get the full dynamics, we also have to integrate the equation of motion (JAX makes this easy with `odeint`).

```python
def equation_of_motion(lagrangian, state, t=None):
  q, q_t = jnp.split(state, 2)
  q_tt = (jnp.linalg.pinv(jax.hessian(lagrangian, 1)(q, q_t))
          @ (jax.grad(lagrangian, 0)(q, q_t)
             - jax.jacobian(jax.jacobian(lagrangian, 1), 0)(q, q_t) @ q_t))
  return jnp.concatenate([q_t, q_tt])

def solve_lagrangian(lagrangian, initial_state, **kwargs):
  # We currently run odeint on CPUs only, because its cost is dominated by
  # control flow, which is slow on GPUs.
  @partial(jax.jit, backend='cpu')
  def f(initial_state):
    return odeint(partial(equation_of_motion, lagrangian),
                  initial_state, **kwargs)
  return f(initial_state)
```

### Helper functions
Here are a few useful helper functions for obtaining dynamics of the system.

```python
# Double pendulum dynamics via the rewritten Euler-Lagrange
@partial(jax.jit, backend='cpu')
def solve_autograd(initial_state, times, m1=1, m2=1, l1=1, l2=1, g=9.8):
  L = partial(lagrangian, m1=m1, m2=m2, l1=l1, l2=l2, g=g)
  return solve_lagrangian(L, initial_state, t=times, rtol=1e-10, atol=1e-10)

# Double pendulum dynamics via analytical forces taken from Diego's blog
@partial(jax.jit, backend='cpu')
def solve_analytical(initial_state, times):
  return odeint(f_analytical, initial_state, t=times, rtol=1e-10, atol=1e-10)

def normalize_dp(state):
  # wrap generalized coordinates to [-pi, pi]
  return jnp.concatenate([(state[:2] + np.pi) % (2 * np.pi) - np.pi, state[2:]])

def rk4_step(f, x, t, h):
  # one step of runge-kutta integration
  k1 = h * f(x, t)
  k2 = h * f(x + k1/2, t + h/2)
  k3 = h * f(x + k2/2, t + h/2)
  k4 = h * f(x + k3, t + h)
  return x + 1/6 * (k1 + 2 * k2 + 2 * k3 + k4)
```

## **Step 4: Compare Autograd vs. Analytic solutions**
Having rewritten the Euler-Lagrange equation so that we can use automatic differentiation to obtain the dynamics of the system, it's time to make sure that this approach matches the analytic solution. Here, we'll choose a random initial state and obtain dynamics via both "autograd" and "analytic" approaches. Then we'll slightly perturb the analytic solution to measure how sensitive it is to initial conditions. Taken together, these measurements should give us an approximate idea of whether our autograd-based approach is doing something reasonable.

```python
# choose an initial state
x0 = np.array([3*np.pi/7, 3*np.pi/4, 0, 0], dtype=np.float32)
noise = np.random.RandomState(0).randn(x0.size)
t = np.linspace(0, 40, num=401, dtype=np.float32)
```

```python
#!g1.1
# compute dynamics analytically
%time x_analytical = jax.device_get(solve_analytical(x0, t))
noise_coeff_1, noise_coeff_2 = 1e-10, 1e-11
x_perturbed_1 = jax.device_get(solve_analytical(x0 + noise_coeff_1 * noise, t))
x_perturbed_2 = jax.device_get(solve_analytical(x0 + noise_coeff_2 * noise, t))
```

    CPU times: user 2.18 s, sys: 1.25 s, total: 3.43 s
    Wall time: 7.51 s

```python
#!g1.1
# compute dynamics via autograd / the rewritten E-L equation
%time x_autograd = jax.device_get(solve_autograd(x0, t))
```

    CPU times: user 3.44 s, sys: 360 ms, total: 3.8 s
    Wall time: 7.62 s

### Plot a comparison

```python
plt.figure(figsize=[10,3], dpi=120) ; plt.xlim(0, 100)
plt.subplot(1,3,1)
plt.title("Analytic vs Perturbed ($\epsilon$={})".format(noise_coeff_1))
plt.xlabel("Time") ; plt.ylabel("State")
plt.plot(t, x_analytical[:, 0], 'g-', alpha=0.5, label='$q$')
plt.plot(t, x_analytical[:, 2], 'c-', alpha=0.5, label='$\dot q$')
plt.plot(t, x_perturbed_1[:, 0], 'g--', alpha=0.5, label='pert. $q$')
plt.plot(t, x_perturbed_1[:, 2], 'c--', alpha=0.5, label='pert. $\dot q$')
plt.legend(fontsize=6)

plt.subplot(1,3,2)
plt.title("Analytic vs Perturbed ($\epsilon$={})".format(noise_coeff_2))
plt.xlabel("Time") ; plt.ylabel("State")
plt.plot(t, x_analytical[:, 0], 'g-', label='$q$')
plt.plot(t, x_analytical[:, 2], 'c-', label='$\dot q$')
plt.plot(t, x_perturbed_2[:, 0], 'g--', label='pert. $q$')
plt.plot(t, x_perturbed_2[:, 2], 'c--', label='pert. $\dot q$')
plt.legend(fontsize=6)

plt.subplot(1,3,3)
plt.title("Analytic vs Autograd")
plt.xlabel("Time") ; plt.ylabel("State")
plt.plot(t, x_analytical[:, 0], 'g-', label='$q$')
plt.plot(t, x_analytical[:, 2], 'c-', label='$\dot q$')
plt.plot(t, x_autograd[:, 0], 'g--', label='autograd $q$')
plt.plot(t, x_autograd[:, 2], 'c--', label='autograd $\dot q$')
plt.legend(fontsize=6)

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

## **Step 5: Generate training data**
Let's generate some training data by using the analytical method to sample double pendulum dynamics.

```python
#!g1.1
time_step = 0.01
N = 1500
analytical_step = jax.jit(jax.vmap(partial(rk4_step, f_analytical, t=0.0, h=time_step)))

# x0 = np.array([-0.3*np.pi, 0.2*np.pi, 0.35*np.pi, 0.5*np.pi], dtype=np.float32)
x0 = np.array([3*np.pi/7, 3*np.pi/4, 0, 0], dtype=np.float32)
t = np.arange(N, dtype=np.float32) # time steps 0 to N
%time x_train = jax.device_get(solve_analytical(x0, t)) # dynamics for first N time steps
%time xt_train = jax.device_get(jax.vmap(f_analytical)(x_train)) # time derivatives of each state
%time y_train = jax.device_get(analytical_step(x_train)) # analytical next step

noise = np.random.RandomState(0).randn(x0.size)
t_test = np.arange(N, 2*N, dtype=np.float32) # time steps N to 2N
%time x_test = jax.device_get(solve_analytical(x0, t_test)) # dynamics for next N time steps
%time xt_test = jax.device_get(jax.vmap(f_analytical)(x_test)) # time derivatives of each state
%time y_test = jax.device_get(analytical_step(x_test)) # analytical next step

```

    CPU times: user 12.8 s, sys: 768 ms, total: 13.6 s
    Wall time: 15.4 s
    CPU times: user 304 ms, sys: 268 ms, total: 572 ms
    Wall time: 1.39 s
    CPU times: user 3.96 s, sys: 72 ms, total: 4.03 s
    Wall time: 5.8 s
    CPU times: user 26.1 s, sys: 20 ms, total: 26.1 s
    Wall time: 26 s
    CPU times: user 44 ms, sys: 0 ns, total: 44 ms
    Wall time: 34.3 ms
    CPU times: user 0 ns, sys: 0 ns, total: 0 ns
    Wall time: 1.22 ms

### Visualize train and test data

```python
# preprocess
train_vis = jax.vmap(normalize_dp)(x_train)
test_vis = jax.vmap(normalize_dp)(x_test)

vel_angle = lambda data:  (np.arctan2(data[:,3], data[:,2]) / np.pi + 1)/2
vel_color = lambda vangle: np.stack( [np.zeros_like(vangle), vangle, 1-vangle]).T
train_colors = vel_color(vel_angle(train_vis))
test_colors = vel_color(vel_angle(test_vis))

# plot
SCALE = 80 ; WIDTH = 0.006
plt.figure(figsize=[8,4], dpi=120)
plt.subplot(1,2,1)
plt.title("Train data") ; plt.xlabel(r'$\theta_1$') ; plt.ylabel(r'$\theta_2$')
plt.quiver(*train_vis.T, color=train_colors, scale=SCALE, width=WIDTH)

plt.subplot(1,2,2)
plt.title("Test data") ; plt.xlabel(r'$\theta_1$') ; plt.ylabel(r'$\theta_2$')
plt.quiver(*test_vis.T, color=test_colors, scale=SCALE, width=WIDTH)

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

## **Step 6: Build a Lagrangian Neural Network**
This section presents a bare-bones approach to approximating the Lagrangian with a neural network

### i) Build model and define the loss

```python
# replace the lagrangian with a parameteric model
def learned_lagrangian(params):
  def lagrangian(q, q_t):
    assert q.shape == (2,)
    state = normalize_dp(jnp.concatenate([q, q_t]))
    return jnp.squeeze(nn_forward_fn(params, state), axis=-1)
  return lagrangian

# define the loss of the model (MSE between predicted q, \dot q and targets)
@jax.jit
def loss(params, batch, time_step=None):
  state, targets = batch
  if time_step is not None:
    f = partial(equation_of_motion, learned_lagrangian(params))
    preds = jax.vmap(partial(rk4_step, f, t=0.0, h=time_step))(state)
  else:
    preds = jax.vmap(partial(equation_of_motion, learned_lagrangian(params)))(state)
  return jnp.mean((preds - targets) ** 2)

# build a neural network model
init_random_params, nn_forward_fn = stax.serial(
    stax.Dense(128),
    stax.Softplus,
    stax.Dense(128),
    stax.Softplus,
    stax.Dense(1),
)
```

### ii) Define optimization and data

```python
#!g1.1
@jax.jit
def update_timestep(i, opt_state, batch):
  params = get_params(opt_state)
  return opt_update(i, jax.grad(loss)(params, batch, time_step), opt_state)

@jax.jit
def update_derivative(i, opt_state, batch):
  params = get_params(opt_state)
  return opt_update(i, jax.grad(loss)(params, batch, None), opt_state)

x_train = jax.device_put(jax.vmap(normalize_dp)(x_train))
y_train = jax.device_put(y_train)

x_test = jax.device_put(jax.vmap(normalize_dp)(x_test))
y_test = jax.device_put(y_test)
```

### iii) Train the model

```python
#!g1.1

rng = jax.random.PRNGKey(0)
_, init_params = init_random_params(rng, (-1, 4))

# numbers in comments denote stephan's settings
batch_size = 100
test_every = 10
num_batches = 1500

train_losses = []
test_losses = []

# adam w learn rate decay
opt_init, opt_update, get_params = optimizers.adam(
    lambda t: jnp.select([t < batch_size*(num_batches//3),
                          t < batch_size*(2*num_batches//3),
                          t > batch_size*(2*num_batches//3)],
                         [1e-4, 5e-5, 1e-5]))
opt_state = opt_init(init_params)

for iteration in range(batch_size*num_batches + 1):
  if iteration % batch_size == 0:
    params = get_params(opt_state)
    train_loss = loss(params, (x_train, xt_train))
    train_losses.append(train_loss)
    test_loss = loss(params, (x_test, xt_test))
    test_losses.append(test_loss)
    if iteration % (batch_size*test_every) == 0:
      print(f"iteration={iteration}, train_loss={train_loss:.6f}, test_loss={test_loss:.6f}")
  opt_state = update_derivative(iteration, opt_state, (x_train, xt_train))

params = get_params(opt_state)
```

    iteration=0, train_loss=3576.855957, test_loss=5191.124512
    iteration=1000, train_loss=115.657616, test_loss=110.434708
    iteration=2000, train_loss=66.741539, test_loss=66.015907
    iteration=3000, train_loss=37.899956, test_loss=38.941505
    iteration=4000, train_loss=30.543676, test_loss=32.286495
    iteration=5000, train_loss=137.313812, test_loss=132.305923
    iteration=6000, train_loss=25.899256, test_loss=29.005585
    iteration=7000, train_loss=15.196664, test_loss=40.832272
    iteration=8000, train_loss=8.597953, test_loss=11.429071
    iteration=9000, train_loss=33.676910, test_loss=32.846699
    iteration=10000, train_loss=3.505478, test_loss=4.494037
    iteration=11000, train_loss=3.358282, test_loss=3.736292
    iteration=12000, train_loss=1.636425, test_loss=2.196020
    iteration=13000, train_loss=1.355138, test_loss=1.822258
    iteration=14000, train_loss=15.706330, test_loss=17.403660
    iteration=15000, train_loss=2.390970, test_loss=2.853291
    iteration=16000, train_loss=1.481209, test_loss=1.862941
    iteration=17000, train_loss=1.496084, test_loss=1.865368
    iteration=18000, train_loss=1.020382, test_loss=1.479395
    iteration=19000, train_loss=1.988720, test_loss=2.119208
    iteration=20000, train_loss=0.726532, test_loss=1.083133
    iteration=21000, train_loss=0.759929, test_loss=1.075109
    iteration=22000, train_loss=0.638381, test_loss=0.939825
    iteration=23000, train_loss=0.813086, test_loss=1.087560
    iteration=24000, train_loss=0.718855, test_loss=0.938150
    iteration=25000, train_loss=0.594811, test_loss=0.861192
    iteration=26000, train_loss=0.659705, test_loss=0.995910
    iteration=27000, train_loss=0.682222, test_loss=0.929229
    iteration=28000, train_loss=0.866563, test_loss=1.038947
    iteration=29000, train_loss=0.562303, test_loss=0.821172
    iteration=30000, train_loss=0.563856, test_loss=0.777511
    iteration=31000, train_loss=0.412002, test_loss=0.611652
    iteration=32000, train_loss=0.468826, test_loss=0.690034
    iteration=33000, train_loss=0.421056, test_loss=0.637357
    iteration=34000, train_loss=0.336975, test_loss=0.517932
    iteration=35000, train_loss=0.557392, test_loss=0.668902
    iteration=36000, train_loss=0.376268, test_loss=0.535015
    iteration=37000, train_loss=0.288204, test_loss=0.465803
    iteration=38000, train_loss=0.257881, test_loss=0.417500
    iteration=39000, train_loss=0.253705, test_loss=0.398200
    iteration=40000, train_loss=0.313284, test_loss=0.480434
    iteration=41000, train_loss=0.639645, test_loss=0.735359
    iteration=42000, train_loss=0.267370, test_loss=0.424087
    iteration=43000, train_loss=0.289885, test_loss=0.422136
    iteration=44000, train_loss=0.359793, test_loss=0.459602
    iteration=45000, train_loss=0.300127, test_loss=0.409058
    iteration=46000, train_loss=0.192382, test_loss=0.314766
    iteration=47000, train_loss=0.442116, test_loss=0.536365
    iteration=48000, train_loss=0.185534, test_loss=0.306417
    iteration=49000, train_loss=0.292655, test_loss=0.405014
    iteration=50000, train_loss=0.478795, test_loss=0.563906
    iteration=51000, train_loss=0.147924, test_loss=0.253956
    iteration=52000, train_loss=0.232359, test_loss=0.313184
    iteration=53000, train_loss=0.155178, test_loss=0.248887
    iteration=54000, train_loss=0.199990, test_loss=0.279796
    iteration=55000, train_loss=0.227304, test_loss=0.299879
    iteration=56000, train_loss=0.144083, test_loss=0.238632
    iteration=57000, train_loss=0.130801, test_loss=0.222120
    iteration=58000, train_loss=0.115036, test_loss=0.210061
    iteration=59000, train_loss=0.188297, test_loss=0.279611
    iteration=60000, train_loss=0.107093, test_loss=0.195451
    iteration=61000, train_loss=0.317091, test_loss=0.383802
    iteration=62000, train_loss=0.100969, test_loss=0.185715
    iteration=63000, train_loss=0.114581, test_loss=0.198864
    iteration=64000, train_loss=0.095644, test_loss=0.182055
    iteration=65000, train_loss=0.381588, test_loss=0.436769
    iteration=66000, train_loss=0.203697, test_loss=0.264329
    iteration=67000, train_loss=0.091463, test_loss=0.173238
    iteration=68000, train_loss=0.359643, test_loss=0.417534
    iteration=69000, train_loss=0.109589, test_loss=0.183835
    iteration=70000, train_loss=0.145655, test_loss=0.224590
    iteration=71000, train_loss=0.080940, test_loss=0.159293
    iteration=72000, train_loss=0.117603, test_loss=0.189496
    iteration=73000, train_loss=0.086342, test_loss=0.158486
    iteration=74000, train_loss=0.105543, test_loss=0.170431
    iteration=75000, train_loss=0.075436, test_loss=0.152614
    iteration=76000, train_loss=0.086185, test_loss=0.156659
    iteration=77000, train_loss=0.130638, test_loss=0.198737
    iteration=78000, train_loss=0.160406, test_loss=0.226429
    iteration=79000, train_loss=0.076000, test_loss=0.141202
    iteration=80000, train_loss=0.086996, test_loss=0.150206
    iteration=81000, train_loss=0.071273, test_loss=0.142989
    iteration=82000, train_loss=0.188655, test_loss=0.232624
    iteration=83000, train_loss=0.136943, test_loss=0.199253
    iteration=84000, train_loss=0.070758, test_loss=0.135860
    iteration=85000, train_loss=0.077590, test_loss=0.135408
    iteration=86000, train_loss=0.073780, test_loss=0.135547
    iteration=87000, train_loss=0.063789, test_loss=0.129135
    iteration=88000, train_loss=0.096395, test_loss=0.159271
    iteration=89000, train_loss=0.306613, test_loss=0.353251
    iteration=90000, train_loss=0.080591, test_loss=0.135561
    iteration=91000, train_loss=0.101489, test_loss=0.151325
    iteration=92000, train_loss=0.092445, test_loss=0.155832
    iteration=93000, train_loss=0.070206, test_loss=0.130379
    iteration=94000, train_loss=0.084526, test_loss=0.150888
    iteration=95000, train_loss=0.053844, test_loss=0.110847
    iteration=96000, train_loss=0.066563, test_loss=0.124486
    iteration=97000, train_loss=0.051100, test_loss=0.113336
    iteration=98000, train_loss=0.066019, test_loss=0.126636
    iteration=99000, train_loss=0.048378, test_loss=0.108911
    iteration=100000, train_loss=0.061389, test_loss=0.114251
    iteration=101000, train_loss=0.044660, test_loss=0.101649
    iteration=102000, train_loss=0.043302, test_loss=0.098634
    iteration=103000, train_loss=0.042150, test_loss=0.096846
    iteration=104000, train_loss=0.040096, test_loss=0.093899
    iteration=105000, train_loss=0.048629, test_loss=0.102880
    iteration=106000, train_loss=0.038267, test_loss=0.089272
    iteration=107000, train_loss=0.037068, test_loss=0.087542
    iteration=108000, train_loss=0.036052, test_loss=0.086194
    iteration=109000, train_loss=0.036346, test_loss=0.086738
    iteration=110000, train_loss=0.034938, test_loss=0.084943
    iteration=111000, train_loss=0.034066, test_loss=0.082635
    iteration=112000, train_loss=0.033279, test_loss=0.080750
    iteration=113000, train_loss=0.036938, test_loss=0.082068
    iteration=114000, train_loss=0.081794, test_loss=0.117171
    iteration=115000, train_loss=0.033241, test_loss=0.076857
    iteration=116000, train_loss=0.032971, test_loss=0.079170
    iteration=117000, train_loss=0.031860, test_loss=0.078451
    iteration=118000, train_loss=0.035184, test_loss=0.080121
    iteration=119000, train_loss=0.029799, test_loss=0.073805
    iteration=120000, train_loss=0.029519, test_loss=0.073477
    iteration=121000, train_loss=0.075845, test_loss=0.117994
    iteration=122000, train_loss=0.042293, test_loss=0.087112
    iteration=123000, train_loss=0.028660, test_loss=0.070426
    iteration=124000, train_loss=0.028276, test_loss=0.069906
    iteration=125000, train_loss=0.027992, test_loss=0.069025
    iteration=126000, train_loss=0.029102, test_loss=0.068678
    iteration=127000, train_loss=0.081024, test_loss=0.119853
    iteration=128000, train_loss=0.036649, test_loss=0.073413
    iteration=129000, train_loss=0.026570, test_loss=0.066678
    iteration=130000, train_loss=0.026323, test_loss=0.066016
    iteration=131000, train_loss=0.026095, test_loss=0.065674
    iteration=132000, train_loss=0.026063, test_loss=0.064942
    iteration=133000, train_loss=0.050393, test_loss=0.089807
    iteration=134000, train_loss=0.025685, test_loss=0.063934
    iteration=135000, train_loss=0.026133, test_loss=0.064237
    iteration=136000, train_loss=0.038952, test_loss=0.069869
    iteration=137000, train_loss=0.035355, test_loss=0.074523
    iteration=138000, train_loss=0.027080, test_loss=0.067015
    iteration=139000, train_loss=0.026190, test_loss=0.061913
    iteration=140000, train_loss=0.024752, test_loss=0.062118
    iteration=141000, train_loss=0.023973, test_loss=0.060721
    iteration=142000, train_loss=0.028209, test_loss=0.062360
    iteration=143000, train_loss=0.023582, test_loss=0.059490
    iteration=144000, train_loss=0.027381, test_loss=0.063406
    iteration=145000, train_loss=0.023635, test_loss=0.058142
    iteration=146000, train_loss=0.029298, test_loss=0.065902
    iteration=147000, train_loss=0.023164, test_loss=0.058751
    iteration=148000, train_loss=0.031361, test_loss=0.069945
    iteration=149000, train_loss=0.023261, test_loss=0.059157
    iteration=150000, train_loss=0.023325, test_loss=0.058530
    CPU times: user 11min 35s, sys: 5min 26s, total: 17min 2s
    Wall time: 15min 41s

    /kernel/lib/python3.7/site-packages/ml_kernel/kernel.py:811: UserWarning: The following variables cannot be serialized: opt_state
      warnings.warn(message)

### iv) Plot losses over training

```python
plt.figure(figsize=(8, 3.5), dpi=120)
plt.plot(train_losses, label='Train loss')
plt.plot(test_losses, label='Test loss')
plt.yscale('log')
plt.ylim(None, 200)
plt.title('Losses over training')
plt.xlabel("Train step") ; plt.ylabel("Mean squared error")
plt.legend() ; plt.show()
```

## **Step 7: Visualize what the model has learned**
In this section we'll run some basic sanity checks to see whether the model has learned something that resembles the Lagrangian of the system

### i) Plot predicted vs actual coordinates (we want a 1:1 mapping)

```python
#!g1.1
xt_pred = jax.vmap(partial(equation_of_motion, learned_lagrangian(params)))(x_test)

fig, axes = plt.subplots(1, 2, figsize=(6, 3), dpi=120)
axes[0].scatter(xt_test[:, 2], xt_pred[:, 2], s=6, alpha=0.2)
axes[0].set_title('Predicting $\dot q$')
axes[0].set_xlabel('$\dot q$ actual')
axes[0].set_ylabel('$\dot q$ predicted')
axes[1].scatter(xt_test[:, 3], xt_pred[:, 3], s=6, alpha=0.2)
axes[1].set_title('Predicting $\ddot q$')
axes[1].set_xlabel('$\ddot q$ actual')
axes[1].set_ylabel('$\ddot q$ predicted')
plt.tight_layout()
```

### ii) Examine the model's performance over the full coordinate space

```python
# make a meshgrid of possible angles
angle_range = -np.pi + 2 * np.pi * np.arange(100)/100
x_scanned = np.stack(np.meshgrid(angle_range, angle_range, [0.0], [0.0], indexing='ij'), axis=-1).reshape(-1, 4)
```

```python
%time xt_actual = jax.vmap(f_analytical)(x_scanned) # actual time derivatives
```

```python
%time xt_predicted = jax.vmap(partial(equation_of_motion, learned_lagrangian(params)))(x_scanned) # predicted time derivatives
```

```python
import xarray
dataarray = xarray.DataArray(np.stack([xt_actual.reshape(100, 100, 4),
                                       xt_predicted.reshape(100, 100, 4)]),
                             dims=['model', 'theta1', 'theta2', 'time_derivative'],
                             coords={'model': ['actual', 'predicted'], 'theta1': angle_range, 'theta2': angle_range})
```

```python
dataarray.isel(time_derivative=2).plot.contourf(levels=31, col='model', x='theta1', y='theta2', vmin=-10, vmax=10)
print("Colorscale is \dot theta1")
plt.show()
```

```python
dataarray.isel(time_derivative=3).plot.contourf(levels=31, col='model', x='theta1', y='theta2', vmin=-10, vmax=10)
print("Colorscale is \dot theta2")
plt.show()
```

## **Step 8: Make predictions with the model**
We should be able to simulate the double pendulum with reasonable degree of accuracy.

```python
# choose an initial state
x1 = np.array([3*np.pi/7, 3*np.pi/4, 0, 0], dtype=np.float32)

t2 = np.linspace(0, 20, num=301)
%time x1_analytical = jax.device_get(solve_analytical(x1, t2))
noise_coeff_3 = 1e-3
%time x1_perturbed = jax.device_get(solve_analytical(x1 + noise_coeff_3 * noise, t2))
%time x1_model = jax.device_get(solve_lagrangian(learned_lagrangian(params), x1, t=t2))
```

```python
plt.figure(figsize=[6,6], dpi=120) ; plt.xlim(0, 20)
plt.subplot(2,1,1)
plt.title("Analytic vs Perturbed ($\epsilon$={})".format(noise_coeff_3))
plt.xlabel("Time") ; plt.ylabel("State")
plt.plot(t2, x1_analytical[:, 2], 'g-', label=r'$\theta_1$ exact')
plt.plot(t2, x1_analytical[:, 3], 'c-', label=r'$\theta_2$ exact')
plt.plot(t2, x1_perturbed[:, 2], 'g--', label=r'$\theta_1$ perturbed')
plt.plot(t2, x1_perturbed[:, 3], 'c--', label=r'$\theta_2$ perturbed')
plt.legend(fontsize=6)

plt.subplot(2,1,2)
plt.title("Analytic vs LNN")
plt.xlabel("Time") ; plt.ylabel("State")
plt.plot(t2, x1_analytical[:, 2], 'g-', label=r'$\theta_1$ exact')
plt.plot(t2, x1_analytical[:, 3], 'c-', label=r'$\theta_2$ exact')
plt.plot(t2, x1_model[:, 2], 'g--', label=r'$\theta_1$ LNN')
plt.plot(t2, x1_model[:, 3], 'c--', label=r'$\theta_2$ LNN')
plt.legend(fontsize=6)

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

## **Step 9: Make videos of the system**

We'll use these to get an intuitive sense of how well our model has learned the dynamics

### i) First, we need to define functions for plotting the pendulum
Sam used [this scipy tutorial](https://scipython.com/blog/the-double-pendulum/) and Stephan Hoyer's [wave equation demo for Jax on TPU](https://colab.research.google.com/github/google/jax/blob/master/cloud_tpu_colabs/Wave_Equation.ipynb) to write the following code.

```python
def make_plot(i, cart_coords, l1, l2, max_trail=30, trail_segments=20, r = 0.05):
    # Plot and save an image of the double pendulum configuration for time step i.
    plt.cla()

    x1, y1, x2, y2 = cart_coords
    ax.plot([0, x1[i], x2[i]], [0, y1[i], y2[i]], lw=2, c='k') # rods
    c0 = Circle((0, 0), r/2, fc='k', zorder=10) # anchor point
    c1 = Circle((x1[i], y1[i]), r, fc='b', ec='b', zorder=10) # mass 1
    c2 = Circle((x2[i], y2[i]), r, fc='r', ec='r', zorder=10) # mass 2
    ax.add_patch(c0)
    ax.add_patch(c1)
    ax.add_patch(c2)

    # plot the pendulum trail (ns = number of segments)
    s = max_trail // trail_segments
    for j in range(trail_segments):
        imin = i - (trail_segments-j)*s
        if imin < 0: continue
        imax = imin + s + 1
        alpha = (j/trail_segments)**2 # fade the trail into alpha
        ax.plot(x2[imin:imax], y2[imin:imax], c='r', solid_capstyle='butt',
                lw=2, alpha=alpha)

    # Center the image on the fixed anchor point. Make axes equal.
    ax.set_xlim(-l1-l2-r, l1+l2+r)
    ax.set_ylim(-l1-l2-r, l1+l2+r)
    ax.set_aspect('equal', adjustable='box')
    plt.axis('off')
    # plt.savefig('./frames/_img{:04d}.png'.format(i//di), dpi=72)

def radial2cartesian(t1, t2, l1, l2):
  # Convert from radial to Cartesian coordinates.
  x1 = l1 * np.sin(t1)
  y1 = -l1 * np.cos(t1)
  x2 = x1 + l2 * np.sin(t2)
  y2 = y1 - l2 * np.cos(t2)
  return x1, y1, x2, y2

def fig2image(fig):
  fig.canvas.draw()
  data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')
  image = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
  return image
```

### ii) Next, we'll make a movie of the pendulum's analytic dynamics

```python
L1, L2 = 1, 1
theta1, theta2 = x1_analytical[:, 0], x1_analytical[:, 1]
cart_coords = radial2cartesian(theta1, theta2, L1, L2)

fig = plt.figure(figsize=(8.3333, 6.25), dpi=72)
ax = fig.add_subplot(111)

import warnings
warnings.filterwarnings("ignore")

images = []
di = 1
N = 300
for i in range(0, N, di):
  print("{}/{}".format(i // di, N // di), end='\n' if i//di%20==0 else ' ')
  make_plot(i, cart_coords, L1, L2)
  images.append( fig2image(fig) )
```

```python
import importlib
importlib.reload(proglog)
print("True (analytical) dynamics of the double pendulum:")
proglog.default_bar_logger = partial(proglog.default_bar_logger, None)
ImageSequenceClip(images, fps=25).ipython_display()
```

### ii) Now let's visualize the LNN's dynamics

```python
L1, L2 = 1, 1
theta1, theta2 = x1_model[:, 0], x1_model[:, 1]
cart_coords = radial2cartesian(theta1, theta2, L1, L2)

fig = plt.figure(figsize=(8.3333, 6.25), dpi=72)
ax = fig.add_subplot(111)

images = []
di = 1
N = 300
for i in range(0, N, di):
  print("{}/{}".format(i // di, N // di), end='\n' if i//di%20==0 else ' ')
  make_plot(i, cart_coords, L1, L2)
  images.append( fig2image(fig) )
```

```python
import importlib
importlib.reload(proglog)
print("LNN-predicted dynamics of the double pendulum:")
proglog.default_bar_logger = partial(proglog.default_bar_logger, None)
ImageSequenceClip(images, fps=25).ipython_display()
```

```python
def save_movie(images, path, duration=100, loop=0, **kwargs):
  images[0].save(path, save_all=True, append_images=images[1:],
                 duration=duration, loop=loop, **kwargs)

# Save GIF.
PIL_images = [Image.fromarray(im, mode='RGB') for im in images]
save_movie(PIL_images, 'double_pend_lnn.gif')
# The movie sometimes takes a second before showing up in the file system.
import time; time.sleep(1)

```

## **Step 10: Conclusions**
We just learned a double pendulum Lagrangian from data! Correction: we learned the double pendulum Lagrangian for a specific case where both arm lengths and both bob masses were 1. Furthermore, all of our train and test data was sampled from a single energy level of the double pendulum. But this is a good first step. These results are exciting because they are better than what an HNN or even a baseline neural ODE could achieve (we also made these comparisons in separate notebooks).

For the purposes of our paper, we learned a model across different energy levels and added better network initializations, training techniques, and more. We also learned the Lagrangians of other physical systems including a 1D wave equation and a particle in a relativistic field. In theory, this approach should generalize well to other systems and to non-canonical coordinate systems such as the latent vector of a VAE.

$\quad$

Thanks for checking out our LNN tutorial, and please contact us with any questions.


#  CS109B Data Science 2: Advanced Topics in Data Science

## Lab 03 - Hierarchical Models in PyMC3

**Harvard University**

**Spring 2022**

**Instructors:** Mark Glickman and Pavlos Protopapas

**Lab instructor and content:** Eleni Angelaki Kaxiras

---

```python
## RUN THIS CELL TO PROPERLY HIGHLIGHT THE EXERCISES
import requests
from IPython.core.display import HTML
styles = requests.get("https://raw.githubusercontent.com/Harvard-IACS/2019-CS109B/master/content/styles/cs109.css").text
HTML(styles)
```

```python
import warnings

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pymc3 as pm
import theano.tensor as tt

from pymc3 import summary
from pymc3 import Model, Normal, HalfNormal, model_to_graphviz, HalfCauchy
from pymc3 import NUTS, sample, find_MAP
from scipy import optimize

warnings.simplefilter(action="ignore", category=FutureWarning)
warnings.filterwarnings('ignore')

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import gridspec
import scipy.stats as stats
import pandas as pd
import seaborn as sns
%matplotlib inline

import warnings
print('Running on PyMC3 v{}'.format(pm.__version__))
```

    Running on PyMC3 v3.8

```javascript
%%javascript
IPython.OutputArea.auto_scroll_threshold = 20000;
```

    <IPython.core.display.Javascript object>

### Learning Objectives

By the end of this lab, you should know how to create Bayesian hierarchical models in PyMC3

**This lab maps to Homework 2.**

## Hierarchical Models

**Gelman et al.** famous radon dataset is a classic for hierarchical modeling. In this dataset the amount of the radioactive gas radon has been measured among different households in all county's of several states. Radon gas is known to be the highest cause of lung cancer in non-smokers. It is believed to be more strongly present in households containing a basement and to differ in amount present among types of soil. Here we'll investigate this differences and try to make predictions of radonlevels in different county's based on the county itself and the presence of a basement.

```python
df = pd.read_csv('../data/radon.csv', index_col=[0])
df['log_radon'] = df['log_radon'].astype('float')
county_names = df.county.unique()
county_idx = df.county_code.values

n_counties = len(df.county.unique())
df.head()
```

      idnum
      state
      state2
      stfips
      zip
      region
      typebldg
      floor
      room
      basement
      ...
      pcterr
      adjwt
      dupflag
      zipflag
      cntyfips
      county
      fips
      Uppm
      county_code
      log_radon

      0
      5081.0
      MN
      MN
      27.0
      55735
      5.0
      1.0
      1.0
      3.0
      N
      ...
      9.7
      1146.499190
      1.0
      0.0
      1.0
      AITKIN
      27001.0
      0.502054
      0
      0.832909

      1
      5082.0
      MN
      MN
      27.0
      55748
      5.0
      1.0
      0.0
      4.0
      Y
      ...
      14.5
      471.366223
      0.0
      0.0
      1.0
      AITKIN
      27001.0
      0.502054
      0
      0.832909

      2
      5083.0
      MN
      MN
      27.0
      55748
      5.0
      1.0
      0.0
      4.0
      Y
      ...
      9.6
      433.316718
      0.0
      0.0
      1.0
      AITKIN
      27001.0
      0.502054
      0
      1.098612

      3
      5084.0
      MN
      MN
      27.0
      56469
      5.0
      1.0
      0.0
      4.0
      Y
      ...
      24.3
      461.623670
      0.0
      0.0
      1.0
      AITKIN
      27001.0
      0.502054
      0
      0.095310

      4
      5085.0
      MN
      MN
      27.0
      55011
      3.0
      1.0
      0.0
      4.0
      Y
      ...
      13.8
      433.316718
      0.0
      0.0
      3.0
      ANOKA
      27003.0
      0.428565
      1
      1.163151

5 rows × 29 columns

Each **row** in the dataframe represents the radon measurements for **one house** in a specific county including whether the house has a basement (floor = 0) or not (floor = 1). We are interested in whether having a basement increases the radon measured in the house.

To keep things simple let's keep only the following three variables: `county`, `log_radon`, and `floor`

```python
# keep only these variables
data = df[['county', 'log_radon', 'floor']]
data.head()
```

      county
      log_radon
      floor

      0
      AITKIN
      0.832909
      1.0

      1
      AITKIN
      0.832909
      0.0

      2
      AITKIN
      1.098612
      0.0

      3
      AITKIN
      0.095310
      0.0

      4
      ANOKA
      1.163151
      0.0

Let's check how many different counties we have. We also notice that they have a different number of houses. Some have a large number of houses measured, some only 1.

```python
data['county'].value_counts().head(5)
```

    ST LOUIS      116
    HENNEPIN      105
    DAKOTA         63
    ANOKA          52
    WASHINGTON     46
    Name: county, dtype: int64

```python
data['county'].value_counts()[-5:]
```

    STEVENS       2
    MILLE LACS    2
    MAHNOMEN      1
    MURRAY        1
    WILKIN        1
    Name: county, dtype: int64

```python
# let's add a column that numbers the counties from 0 to n
# raw_ids = np.unique(data['county'])
# raw2newid = {x:np.where(raw_ids == x)[0][0] for x in raw_ids}

# data['county_id'] = data['county'].map(raw2newid)
# data
```

### 1 - Pooling: Same Linear Regression for all

We can just pool all the data and estimate one big regression to asses the influence of having a basement on radon levels across all counties. Our model would be:

\begin{equation}
y_{i} = \alpha + \beta*floor_{i}
\end{equation}

Where $i$ represents the measurement (house), and floor contains a 0 or 1 if the house has a basement or not. By ignoring the county feature, we do not differenciate on counties.

```python
with pm.Model() as pooled_model:
    # common priors for all
    a = pm.Normal('a', mu=0, sigma=100)
    b = pm.Normal('b', mu=0, sigma=100)

    # radon estimate
    radon_est = a + b*data['floor'].values

    # likelihood after radon observations
    radon_obs = pm.Normal('radon_obs', mu=radon_est,
                           observed=data['log_radon']) # note here we enter the whole dataset
```

```python
from pymc3 import model_to_graphviz
model_to_graphviz(pooled_model)
```

```python
with pooled_model:

    pooled_trace = sample(2000, tune=1000, target_accept=0.9)
    print(f'DONE')
```

    Auto-assigning NUTS sampler...
    Initializing NUTS using jitter+adapt_diag...
    Multiprocess sampling (4 chains in 4 jobs)
    NUTS: [b, a]
    Sampling 4 chains, 0 divergences: 100%|██████████| 12000/12000 [00:03<00:00, 3781.83draws/s]

    DONE

```python
pm.traceplot(pooled_trace);
```

Remember, with the pooled model we have only one intercept, $\alpha$, and only one slope, $\beta$ for all the counties. Let's plot the regression lines.

```python
# plot just a subset of the countries
counties = ['HENNEPIN','AITKIN','WASHINGTON', 'MURRAY', 'YELLOW MEDICINE', 'MAHNOMEN']

plt.figure(figsize=(10,5))
rows = 2
gs = gridspec.GridSpec(rows, len(counties)//rows)

for i, county in enumerate(counties):
    county_data = data.loc[data['county']==county]
    x = np.linspace(-0.2, 1.2)
    radon_est = pooled_trace['a'].mean() + pooled_trace['b'].mean()*x
    subplt = plt.subplot(gs[i])

    subplt.set_ylim(0.,4.)
    subplt.scatter(county_data['floor'], county_data['log_radon'])
    subplt.plot(x, radon_est, c='r', label='pooled line');
    subplt.set_xlabel('floor', fontsize=10)
    subplt.set_ylabel('radon level', fontsize=10)
    subplt.set_title(str(county) + ' County')
    subplt.legend()
    plt.tight_layout()
```

### 2 - Unpooling: Separate Linear Regression for each county

We believe that different counties have different relationships of radon and basements. Our model would be:

\begin{equation}
radon_{i,c} = \alpha_c + \beta_c*floor_{i,c}
\end{equation}

Where $i$ represents the measurement, $c$ the county, and floor contains a 0 or 1 if the house has a basement or not.

Notice we have separate coefficients for each county in $a_c$ and $b_c$. They are totally different, they do not even come from the same distribution.

We will do this for only one county, as an example. We pick HENNEPIN county.

```python
# chose a county
county = 'MEEKER'
county_data = data.loc[data['county']==county]
county_data.head()
```

      county
      log_radon
      floor

      479
      MEEKER
      0.875469
      0.0

      480
      MEEKER
      1.386294
      0.0

      481
      MEEKER
      1.987874
      0.0

      482
      MEEKER
      0.788457
      0.0

      483
      MEEKER
      1.193922
      0.0

```python
#help(pm.Normal)
```

```python
with pm.Model() as unpooled_model:
    mu_a = pm.Normal('mu_a', mu=0., sigma=100)
    sigma_a = pm.HalfNormal('sigma_a', 5.)
    mu_b = pm.Normal('mu_b', mu=0., sigma=100)
    sigma_b = pm.HalfNormal('sigma_b', 5.)

    a = pm.Normal('a', mu=mu_a, sigma=sigma_a)
    b = pm.Normal('b', mu=mu_b, sigma=sigma_b)

    radon_est = a + b*county_data['floor'].values

    radon_obs = pm.Normal('radon_like', mu=radon_est,
                           observed=county_data['log_radon'])
```

```python
model_to_graphviz(unpooled_model)
```

```python
with unpooled_model:

    unpooled_trace = sample(2000, tune=1000, target_accept=0.9)
    print(f'DONE')
```

    Auto-assigning NUTS sampler...
    Initializing NUTS using jitter+adapt_diag...
    Multiprocess sampling (4 chains in 4 jobs)
    NUTS: [b, a, sigma_b, mu_b, sigma_a, mu_a]
    Sampling 4 chains, 465 divergences: 100%|██████████| 12000/12000 [00:58<00:00, 205.25draws/s]
    There were 94 divergences after tuning. Increase `target_accept` or reparameterize.
    There were 184 divergences after tuning. Increase `target_accept` or reparameterize.
    The acceptance probability does not match the target. It is 0.7384761840644513, but should be close to 0.9. Try to increase the number of tuning steps.
    There were 88 divergences after tuning. Increase `target_accept` or reparameterize.
    There were 99 divergences after tuning. Increase `target_accept` or reparameterize.
    The acceptance probability does not match the target. It is 0.8329297259480425, but should be close to 0.9. Try to increase the number of tuning steps.
    The estimated number of effective samples is smaller than 200 for some parameters.

    DONE

```python
pm.traceplot(unpooled_trace);
```

Print the regression line for our chosen county alone.

```python
county = 'MEEKER'
county_data = data.loc[data['county']==county]

x = np.arange(len(county_data['floor'].values))
radon_est_unpooled = unpooled_trace['a'].mean() + unpooled_trace['b'].mean()*county_data['floor'].values

xx = np.linspace(-0.2, 1.2)
radon_est_pooled = pooled_trace['a'].mean() + pooled_trace['b'].mean()*xx

plt.scatter(county_data['floor'], county_data['log_radon'])
plt.xlim(-0.1,1.1)
plt.xlabel('floor', fontsize=10)
plt.ylabel('radon level', fontsize=10)
plt.title(f'{str(county)} county Radon levels')
plt.plot(x, radon_est_unpooled, c='g', label='unpooled line');
plt.plot(xx, radon_est_pooled, c='r', label='pooled line');
plt.legend();
```

### 3 - Partial pooling: Hierarchical Regression (Varying-Coefficients Model)

Merely by the fact that all counties are counties, they share similarities, so there is a middle ground to both of these extremes. Specifically, we may assume that while $\alpha_c$ and $\beta_c$are different for each county as in the unpooled case, the coefficients are all drawn from the same distribution:

\begin{equation}
radon_{i,c} = \alpha_c + \beta_c*floor_{i,c}
\end{equation}

\begin{equation}
a_c \sim  \mathcal{N}(\mu_a,\,\sigma_a^{2})
\end{equation}
\begin{equation}
b_c \sim  \mathcal{N}(\mu_b,\,\sigma_b^{2})
\end{equation}

where the common parameters are:
\begin{eqnarray}
\mu_a \sim  \mathcal{N}(0,\,10) \\
\sigma_a^2 \sim  |\mathcal{N}(0,\,10)| \\
\mu_b \sim  \mathcal{N}(0,\,10) \\
\sigma_b^2 \sim  |\mathcal{N}(0,\,10)|
\end{eqnarray}

Add math for mu and sigma

The different counties are effectively sharing information through the commin priors. We are thus observing what is known as shrinkage; modeling the groups not as independent from each other, neither as a single group but rather as related.

NOTES

We saw that some counties had only one sample, so if that house is a really old with old lead pipes, our prediction will be that all houses in this county have radon. On the other extreme, if we have a newer house with no radon then again we will have missleading results. In one case, you will overestimate the bad quality and in the other underestimate it. Under a hierarchical model, the miss-estimation of one group will be offset by the information provided by the other groups. As always gathering more data helps if this is an option.

### Defining the Model for the Hierarchical Model

```python
with pm.Model() as hierarchical_model:
    # Hyperpriors for group nodes
    mu_a = pm.Normal('mu_a', mu=0., sigma=100)
    sigma_a = pm.HalfNormal('sigma_a', 5.)
    mu_b = pm.Normal('mu_b', mu=0., sigma=100)
    sigma_b = pm.HalfNormal('sigma_b', 5.)

    # Above we just set mu and sd to a fixed value while here we
    # plug in a common group distribution for all a and b (which are
    # vectors of length n_counties).

    # Intercept for each county, distributed around group mean mu_a
    a = pm.Normal('a', mu=mu_a, sigma=sigma_a, shape=n_counties)

    # beta for each county, distributed around group mean mu_b
    b = pm.Normal('b', mu=mu_b, sigma=sigma_b, shape=n_counties)

    # Model error
    #eps = pm.HalfCauchy('eps', 5.)

    radon_est = a[county_idx] + b[county_idx]*data['floor'].values

    # Data likelihood with sigma for random error
#     radon_like = pm.Normal('radon_like', mu=radon_est,
#                            sigma=eps, observed=data['log_radon'])

    # Data likelihood with sigma without random error
    radon_like = pm.Normal('radon_like', mu=radon_est, #sigma=eps,
                           observed=data['log_radon'])
```

```python
model_to_graphviz(hierarchical_model)
```

#### Inference

```python
with hierarchical_model:
    hierarchical_trace = pm.sample(2000, tune=2000, target_accept=.9)
```

    Auto-assigning NUTS sampler...
    Initializing NUTS using jitter+adapt_diag...
    Multiprocess sampling (4 chains in 4 jobs)
    NUTS: [b, a, sigma_b, mu_b, sigma_a, mu_a]
    Sampling 4 chains, 1,822 divergences: 100%|██████████| 16000/16000 [01:02<00:00, 257.47draws/s]
    There were 8 divergences after tuning. Increase `target_accept` or reparameterize.
    There were 49 divergences after tuning. Increase `target_accept` or reparameterize.
    There were 1754 divergences after tuning. Increase `target_accept` or reparameterize.
    The acceptance probability does not match the target. It is 0.1690165871960057, but should be close to 0.9. Try to increase the number of tuning steps.
    There were 11 divergences after tuning. Increase `target_accept` or reparameterize.
    The rhat statistic is larger than 1.4 for some parameters. The sampler did not converge.
    The estimated number of effective samples is smaller than 200 for some parameters.

```python
pm.traceplot(hierarchical_trace, var_names=['mu_a', 'mu_b',
                                            'sigma_a', 'sigma_b']);
```

```python
results = pm.summary(hierarchical_trace)
results[:10]
```

      mean
      sd
      hpd_3%
      hpd_97%
      mcse_mean
      mcse_sd
      ess_mean
      ess_sd
      ess_bulk
      ess_tail
      r_hat

      mu_a
      1.462
      0.048
      1.374
      1.560
      0.002
      0.001
      912.0
      912.0
      681.0
      1786.0
      1.06

      mu_b
      -0.624
      0.089
      -0.824
      -0.476
      0.006
      0.005
      215.0
      193.0
      224.0
      379.0
      1.09

      a[0]
      1.334
      0.210
      0.917
      1.758
      0.003
      0.003
      4753.0
      2289.0
      3168.0
      2972.0
      1.20

      a[1]
      1.043
      0.119
      0.830
      1.276
      0.017
      0.013
      47.0
      44.0
      55.0
      2754.0
      1.07

      a[2]
      1.493
      0.220
      1.027
      1.877
      0.027
      0.019
      68.0
      68.0
      61.0
      2436.0
      1.09

      a[3]
      1.504
      0.190
      1.131
      1.893
      0.004
      0.003
      2658.0
      2658.0
      1129.0
      2511.0
      1.17

      a[4]
      1.382
      0.251
      1.043
      1.853
      0.073
      0.053
      12.0
      12.0
      12.0
      378.0
      1.22

      a[5]
      1.545
      0.249
      1.050
      1.909
      0.057
      0.042
      19.0
      18.0
      21.0
      2916.0
      1.12

      a[6]
      1.746
      0.179
      1.369
      2.046
      0.021
      0.015
      72.0
      72.0
      76.0
      3192.0
      1.08

      a[7]
      1.634
      0.252
      1.126
      1.980
      0.065
      0.048
      15.0
      14.0
      17.0
      1853.0
      1.16

```python
# plot just a subset of the countries
counties = ['HENNEPIN','AITKIN','WASHINGTON', 'LAKE OF THE WOODS', 'YELLOW MEDICINE', 'ANOKA']

plt.figure(figsize=(10,5))
rows = 2
gs = gridspec.GridSpec(rows, len(counties)//rows)

for i, county in enumerate(counties):
    county_data = data.loc[data['county']==county]
    subplt = plt.subplot(gs[i])

    # pooled line (single values coeff for all)
    xx = np.linspace(-0.2, 1.2)
    radon_est = pooled_trace['a'].mean() + pooled_trace['b'].mean()*xx

    radon_est_hier = np.mean(hierarchical_trace['a'][i]) + \
                      np.mean(hierarchical_trace['b'][i])*xx

    # un-pooled (single subject)
    sns.regplot(x='floor', y='log_radon', ci=None, label='unpooled',
                data=county_data).set_title('County '+str(county))

    # hierarchical line
    subplt.set_ylim(0.,4.)
    subplt.scatter(county_data['floor'], county_data['log_radon'])
    subplt.plot(xx, radon_est, c='r', label='pooled');

    # plot the hierarchical, varying coefficient model
    subplt.plot(xx, radon_est_hier, c='g', label='hierarchical');

    subplt.set_xlabel('floor', fontsize=10)
    subplt.set_ylabel('radon level', fontsize=10)
    subplt.set_title(str(county) + ' County')
    subplt.legend()
    plt.tight_layout()
```

This tutorial is modified from PyMC3 docs: https://docs.pymc.io/notebooks/GLM-hierarchical.html