Copyright 2021 Google LLC

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    https://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

## Notebook overview

This notebook provides an interactive demonstration of adaptive dynamics prediction for rigid-body dynamics with data output from a PyBullet simulation of thrown objects. It describes the general structure of an adaptive predictor, mathematically constructs an adaptive predictor for rigid-body dynamics, and then provides an implementation.

```python
!pip install pybullet
!pip install bullet
```

```python
import pybullet as p
import bullet as b
import pybullet_data
import numpy as onp
import jax
from jax import numpy as jnp
from matplotlib import pyplot as plt
import matplotlib as mpl
import seaborn as sns
import functools

from jax.config import config
config.update('jax_enable_x64', True)
```

### Imported Code

```python
""" Quaternion functions """

@jax.jit
def skew(w):
  w = jnp.reshape(w, (3))
  return jnp.array([[0.0, -w[2], w[1]],\
                   [w[2], 0.0, -w[0]],\
                   [-w[1], w[0], 0.0]])

@jax.jit
def to_axis_angle(q):
  angle = 2.0 * jnp.arccos(q[0])
  axis = jnp.where(angle, q[1:] / jnp.sin(angle / 2.0), 0.0)
  return (axis, angle)

@jax.jit
def normalize(q):
  return q / jnp.linalg.norm(q)

@jax.jit
def multiply(q, p):
  return jnp.concatenate((jnp.array([q[0] * p[0] - p[1:] @ q[1:]]),
                         q[0] * p[1:] + p[0] * q[1:] + jnp.cross(q[1:], p[1:])))

@jax.jit
def vector_multiply(v, q):
  """Multiplication between quaternions [0, v0, v1, v2] and q."""
  return multiply(jnp.concatenate((jnp.array([0.0]), v)), q)

@jax.jit
def to_rotation_matrix(q):
  (omega, theta) = to_axis_angle(q)
  omega_sk = skew(omega)
  return jnp.eye(3) + jnp.sin(theta) * omega_sk + (
      1.0 - jnp.cos(theta)) * omega_sk @ omega_sk
```

```python
""" Potential functions for mirror descent """

@jax.jit
def p_norm(x, p):
  return .5*jnp.linalg.norm(x, p)**2

@jax.jit
def grad_p(x, p):
  """Direct calculation of the gradient of p_norm(x, p)."""
  return jnp.linalg.norm(x, p)**(2-p)*jnp.abs(x)**(p-1)*jnp.sign(x)

@jax.jit
def inv_grad_p(x, p):
  """See https://arxiv.org/pdf/1912.13154.pdf, Section 8.1."""
  q = 1/(1 - 1/p)  # 1/q + 1/p == 1
  return grad_p(x, q)

@jax.jit
def hypentropy(x, beta):
  """See https://arxiv.org/pdf/1902.01903.pdf."""
  return jnp.sum(x*jnp.arcsinh(x/beta) - jnp.sqrt(x**2 + beta**2))

@jax.jit
def grad_hyp(x, beta):
  """See https://arxiv.org/pdf/1902.01903.pdf."""
  return jnp.arcsinh(x/beta)

@jax.jit
def inv_grad_hyp(x, beta):
  """See https://arxiv.org/pdf/1902.01903.pdf."""
  return beta*jnp.sinh(x)
```

```python
"""Class for adaptive dynamics prediction."""

from typing import Callable, Dict, Tuple

@functools.partial(jax.jit, static_argnums=(0, 1))
def rollout_prediction(norm_state: Callable[[jnp.ndarray], jnp.ndarray],
                       predict_state: Callable[[jnp.ndarray, jnp.ndarray, int],
                                               jnp.ndarray],
                       params: jnp.ndarray, init_state: jnp.ndarray,
                       t_range: jnp.ndarray) -> jnp.ndarray:
  """Given a measurement, update the parameters and rollout the prediction.

  Args:
    norm_state: A function used to normalize the state after each step.
    predict_state: A function mapping the current state estimate to the state
    params: Current adaptive parameter estimates.
    estimate at the next timestep.
    init_state: Initial state for this rollout.
    t_range: Range of discrete-time indices for the rollout.

  Returns:
    traj: The predicted state trajectory.
  """
  def scan_fn(xhatt: jnp.ndarray, t: int) -> Tuple[jnp.ndarray, jnp.ndarray]:
    xhattp1 = norm_state(predict_state(xhatt, params, t))
    return xhattp1, xhattp1

  _, traj = jax.lax.scan(scan_fn, init_state, t_range)

  return traj

def euler(dynamics, dt):
  def discretized(x, u, t):
    return x + dt*dynamics(x, u, t)
  return discretized

def rk4(dynamics, dt):
  """RK4 discretization."""
  # dynamics(x, u, t)
  def discretized(x, u, t):
    k1 = dynamics(x, u, dt*t)
    k2 = dynamics(x + dt*k1/2, u, dt*t + dt/2)
    k3 = dynamics(x + dt*k2/2, u, dt*t + dt/2)
    k4 = dynamics(x + dt*k3, u, dt*t + dt)
    return x + dt*(k1 + 2*k2 + 2*k3 + k4)/6
  return discretized

def rk4_adaptive(dynamics, dt, sys_dim, inverse_mirror):
  """Discretize the adaptive system, applying the inverse mirror.

  Note: this is basically a Runge-Kutta mirror descent. Because Runge-Kutta
  evaluates the right-hand side at intermediate time points, we need to
  re-apply the inverse mirror map to compute the parameters at intermeditate
  time points.

  Args:
    dynamics: The dynamics to discretize
    dt: The timestep.
    sys_dim: System dimension (without adaptation parameters).
    inverse_mirror: The inverse mirror map to apply to the parameters.

  Returns:
    The Runge-Kutta discretized mirror descent update.
  """
  def discretized(xhat_and_mirror_params, ob, params, t):
    k1 = dynamics(xhat_and_mirror_params, ob, params, dt*t)

    xhat_and_mirror_params1 = xhat_and_mirror_params + dt*k1/2
    p1 = inverse_mirror(xhat_and_mirror_params1[sys_dim:])
    k2 = dynamics(xhat_and_mirror_params1, ob, p1, dt*t + dt/2)

    xhat_and_mirror_params2 = xhat_and_mirror_params + dt*k2/2
    p2 = inverse_mirror(xhat_and_mirror_params2[sys_dim:])
    k3 = dynamics(xhat_and_mirror_params2, ob, p2, dt*t + dt/2)

    xhat_and_mirror_params3 = xhat_and_mirror_params + dt*k3
    p3 = inverse_mirror(xhat_and_mirror_params3[sys_dim:])
    k4 = dynamics(xhat_and_mirror_params3, ob, p3, dt*t + dt)

    return xhat_and_mirror_params + dt*(k1 + 2*k2 + 2*k3 + k4)/6

  return discretized

class AdaptivePredictor(object):
  """An adaptive dynamics predictor."""
  # system info
  sys_dim: int
  unknown_dim: int

  # time info
  # TODO(boffi): Reformulate this just in terms of the horizon of adaptation.
  dt: float
  nsteps: int
  nframes: int
  meas_count: int = 0
  start_ind: int = 0
  prediction_index: int

  # adaptation parameter info
  params: onp.ndarray
  n_features: int
  n_params: int

  # adaptation terms
  grad_psi_rf: Callable[[jnp.ndarray], jnp.ndarray]
  inverse_grad_psi_rf: Callable[[jnp.ndarray], jnp.ndarray]
  grad_psi_phys: Callable[[jnp.ndarray], jnp.ndarray]
  inverse_grad_psi_phys: Callable[[jnp.ndarray], jnp.ndarray]

  # continuous-time formulation
  dynamics: Callable[[jnp.ndarray, jnp.ndarray, float], jnp.ndarray]
  dynamics_with_measurement: Callable[
      [jnp.ndarray, jnp.ndarray, jnp.ndarray, float], jnp.ndarray]

  # discretized updates
  predict_state: Callable[[jnp.ndarray, jnp.ndarray, int], jnp.ndarray]
  predict_state_with_measurement: Callable[
      [jnp.ndarray, jnp.ndarray, jnp.ndarray, int], jnp.ndarray]
  norm_state: Callable[[jnp.ndarray], jnp.ndarray]

  # stores the state prediction over the whole rollout
  state_prediction: onp.ndarray

  # pylint: disable=dangerous-default-value
  def __init__(self,
               predictor_continuous: Callable[[jnp.ndarray, jnp.ndarray, float],
                                              jnp.ndarray],
               measurement_function: Callable[[jnp.ndarray, jnp.ndarray],
                                              jnp.ndarray],
               adapt_continuous: Callable[
                   [jnp.ndarray, jnp.ndarray, float, int], jnp.ndarray],
               norm_state: Callable[[jnp.ndarray], jnp.ndarray],
               ob: onp.ndarray,
               init_params: onp.ndarray,
               n_features: int = 0,
               dt: float = 1e-3,
               nsteps: int = 25,
               nframes: int = 40,
               potentials: Dict[str, str] = {
                   'phys': 'euclid',
                   'rf': 'euclid'
               },
               potential_params: Dict[str, float] = {
                   'phys': 0.0,
                   'rf': 0.0
               },
               prediction_index: int = 40) -> None:
    """Initializes the adaptive predictor.

    Args:
        predictor_continuous: Continuous definition of the predictor.

        measurement_function: Feedback measurement function g(x, hat{x}).

        adapt_continuous: Continuous adaptation dynamics for physical basis
        functions and random features.

        norm_state: An optional state normalization to be applied after each
        timestep, such as a quaternion projection.

        ob: Initial observation of the system state.

        init_params: Initialization for the adaptation parameters (both for
        physical and for random features, physical first).

        n_features: Number of random features.

        dt: Timestep of the predictor.

        nsteps: Number of steps of the predictor between measurements.

        nframes: Define a frame as the number of between-measuremnent periods
        over which the predictor will be active. nframes specifies the number
        of frames so that with initialization, there are 1 + nframes*nsteps
        total integration steps of the predictor.

        potentials: Potential functions for physical and random feature
          adaptation. Currently accepts 'euclid', 'p', and 'hyp' for euclidean,
          p-norm, and hypentropy. Keys should be 'phys' and 'rf'.

        potential_params: Parameters for the potentials (p value in the p norm
        algorithm, beta value in hypentropy). Keys shold be 'phys' and 'rf'.

        prediction_index: Past this point, we clamp the state to the observation
        for maximally accurate predictions.
    """
    self.sys_dim = ob.size

    # initialize adaptation info
    self.n_params = init_params.size
    self.n_features = n_features

    # prepare mirror maps.
    self.setup_mirrors(potentials, potential_params)

    # initialize time info
    self.dt = dt
    self.nsteps = nsteps
    self.nframes = nframes
    self.meas_count = 0
    self.prediction_index = prediction_index

    # construct the continuous formulation of the predictor
    self.dynamics = predictor_continuous
    self.setup_dynamics_with_measurement(measurement_function, adapt_continuous)

    # construct discretized updates
    self.predict_state = jax.jit(rk4(self.dynamics, self.dt))

    self.predict_state_with_measurement = jax.jit(
        rk4_adaptive(self.dynamics_with_measurement, self.dt, self.sys_dim,
                     self.invert_mirrors))
    self.norm_state = norm_state

    # initialize the parameters
    self.params = init_params

    # initialize the predictor state
    self.state_prediction = onp.zeros((nsteps*nframes + 1, self.sys_dim))
    self.state_prediction[0, :] = ob

  def setup_dynamics_with_measurement(
      self,
      measurement_function: Callable[[jnp.ndarray, jnp.ndarray],
                                     jnp.ndarray],
      adapt_continuous: Callable[[jnp.ndarray, jnp.ndarray, float, int],
                                 jnp.ndarray]):
    """Initializes the dynamics_with_measurement() function."""

    def dynamics_with_measurement(xhat_and_mirror_params: jnp.ndarray,
                                  ob: jnp.ndarray,
                                  params: jnp.ndarray, t: float):
      return jnp.concatenate(
          (self.dynamics(xhat_and_mirror_params[:self.sys_dim], params, t) +
           measurement_function(ob, xhat_and_mirror_params[:self.sys_dim]),
           adapt_continuous(ob, xhat_and_mirror_params, t, self.sys_dim)))

    self.dynamics_with_measurement = jax.jit(dynamics_with_measurement)

  def setup_mirrors(self, potentials: Dict[str, str],
                    potential_params: Dict[str, float]):
    """Sets up the forward and inverse mirror maps."""
    # pylint: disable=g-long-lambda
    if potentials['phys'] == 'p':
      self.grad_psi_phys = lambda x: grad_p(
          x, potential_params['phys'])
      self.inverse_grad_psi_phys = lambda x: inv_grad_p(
          x, potential_params['phys'])
    elif potentials['phys'] == 'hyp':
      self.grad_psi_phys = lambda x: grad_hyp(
          x, potential_params['phys'])
      self.inverse_grad_psi_phys = lambda x: inv_grad_hyp(
          x, potential_params['phys'])
    else:
      self.grad_psi_phys = lambda x: x
      self.inverse_grad_psi_phys = lambda x: x

    # set up mirror map for random feature parameters
    if potentials['rf'] == 'p':
      self.grad_psi_rf = lambda x: grad_p(
          x, potential_params['rf'])
      self.inverse_grad_psi_rf = lambda x: inv_grad_p(
          x, potential_params['rf'])
    elif potentials['rf'] == 'hyp':
      self.grad_psi_rf = lambda x: grad_hyp(
          x, potential_params['rf'])
      self.inverse_grad_psi_rf = lambda x: inv_grad_hyp(
          x, potential_params['rf'])
    else:
      self.grad_psi_rf = lambda x: x
      self.inverse_grad_psi_rf = lambda x: x

  def apply_mirrors(self, params: jnp.ndarray):
    """Applies the mirror maps to obtain auxiliary parameters."""
    return jnp.concatenate(
        (self.grad_psi_phys(params[:self.n_params - self.n_features]),
         self.grad_psi_rf(params[self.n_params -
                                 self.n_features:self.n_params]),
         params[self.n_params:]))

  def invert_mirrors(self, mirror_params: jnp.ndarray):
    """Inverts the mirror maps to obtain the parameters."""
    return jnp.concatenate(
        (self.inverse_grad_psi_phys(mirror_params[:self.n_params -
                                                  self.n_features]),
         self.inverse_grad_psi_rf(
             mirror_params[self.n_params - self.n_features:self.n_params]),
         mirror_params[self.n_params:]))

  def rollout_prediction(self, t0: int) -> jnp.ndarray:
    """Given a measurement, update the parameters and rollout the prediction.

    Args:
      t0: The initial time for the rollout.

    Returns:
      traj: The predicted state trajectory.
    """

    # jax will have to re-jit rollout_adaption if the time horizon changes.
    # to fix this, we rollout for nsteps*nframes timesteps every single time,
    # and then only pick off the datapoints that we need.
    return jax.device_get(
        rollout_prediction(
            self.norm_state, self.predict_state,
            jax.device_put(self.params),
            jax.device_put(self.state_prediction[t0, :]),
            jnp.arange(
                t0, t0 +
                self.nframes * self.nsteps)))[:self.nframes * self.nsteps - t0]

  def update_state(self, ob: jnp.ndarray) -> None:
    """Given a measurement, updates the prediction and the parameters.

    Args:
      ob: The state observation.
    """
    self.start_ind = self.meas_count * self.nsteps

    # update the state and the mirrored parameters
    jparams = jax.device_put(self.params)
    init_state_and_mirror_params = jnp.concatenate(
        (jax.device_put(self.state_prediction[self.start_ind, :]),
         self.apply_mirrors(jparams)))

    new_state_and_mirror_params = self.predict_state_with_measurement(
        init_state_and_mirror_params, ob, jparams, self.start_ind)

    # save the updated state. if we are before the place where we want
    # to predict, just let the prediction keep going.
    # otherwise, clamp the estimate to the observation for maximally accurate
    # prediction.
    time_to_predict = self.meas_count >= self.prediction_index
    if time_to_predict:
      self.state_prediction[self.start_ind, :] = ob
    else:
      self.state_prediction[self.start_ind + 1, :] = jax.device_get(
          new_state_and_mirror_params[:self.sys_dim])

    # invert the mirror maps and compute the parameters
    self.params = jax.device_get(
        self.invert_mirrors(new_state_and_mirror_params[self.sys_dim:]))

    if time_to_predict:
      self.state_prediction[self.start_ind + 1:, :] = jax.device_get(
          self.rollout_prediction(self.start_ind))
    else:
      self.state_prediction[self.start_ind + 2:, :] = jax.device_get(
          self.rollout_prediction(self.start_ind + 1))

    self.meas_count += 1

```

```python
# timestep of the simulation, and the predictor
dt = 1e-3

# set up the pybullet simulation environment
p.connect(p.DIRECT)
p.resetSimulation()
p.setGravity(0, 0, -9.8)
p.setTimeStep(dt)
useFixedBase = True
flags = p.URDF_INITIALIZE_SAT_FEATURES
plane_pos = [0,0,-0.625]
p.setAdditionalSearchPath(pybullet_data.getDataPath())
plane = p.loadURDF('plane.urdf', plane_pos, flags = flags, useFixedBase=useFixedBase)
sphere = p.loadURDF('sphere_small.urdf', [0, 0, 10], useMaximalCoordinates=True, flags=p.URDF_USE_INERTIA_FROM_FILE)
```

```python
def set_initial_state(object_id, position, orn, linear_velocity, angular_velocity):
  """Set the initial state of the robot."""
  p.resetBasePositionAndOrientation(object_id, position, orn)
  p.resetBaseVelocity(object_id, linear_velocity, angular_velocity)

def get_state(object_id):
  """Converts the pybullet state to a jax array matching the predictor order. """
  position, orientation = p.getBasePositionAndOrientation(object_id)
  linear_velocity, angular_velocity = p.getBaseVelocity(object_id)

  # make pybullet convention match jaxrobotics
  orientation = jnp.array([orientation[3], orientation[0], orientation[1], orientation[2]])
  position = jnp.array(position)
  linear_velocity = jnp.array(linear_velocity)
  angular_velocity = jnp.array(angular_velocity)

  return jnp.concatenate((orientation, position, linear_velocity, angular_velocity))
```

# Adaptive prediction: general formulation
A general adaptive predictor is specified as

\begin{align*}
  \dot{\hat{\mathbf{x}}} &= \mathbf{f}(\hat{\mathbf{x}}, t) + \mathbf{Y}(\hat{\mathbf{x}}, t)_{\text{phys}}\hat{\boldsymbol{\alpha}}_{\text{phys}} + \mathbf{Y}(\hat{\mathbf{x}}, t)_{\text{rf}}\hat{\boldsymbol\alpha}_{\text{rf}} + \mathbf{g}(\mathbf{x}, \hat{\mathbf{x}})
\end{align*}

Above, $\mathbf{f}(\hat{\mathbf{x}}, t)$ represents any part of the dynamics that is actually known. For example, as shown below, for rigid-body dynamics we know the dynamics of the orientation quaternion, and we also know the dynamics for two components of the velocity.

$\mathbf{g}(\mathbf{x}, \hat{\mathbf{x}})$ is a measurement function used to incorporate a measurement when available.

$\mathbf{Y}(\hat{\mathbf{x}}, t)_{\text{phys}}\hat{\boldsymbol{\alpha}}_{\text{phys}}$ represents a physical approximation with physical basis functions $\mathbf{Y}(\hat{\mathbf{x}}, t)_{\text{phys}}$ and corresponding parameters $\hat{\boldsymbol{\alpha}}_{\text{phys}}$.

$\mathbf{Y}(\hat{\mathbf{x}}, t)_{\text{rf}}\hat{\boldsymbol\alpha}_{\text{rf}}$ represents a mathematical approximation with unstructured basis functions $\mathbf{Y}(\hat{\mathbf{x}}, t)_{\text{rf}}$ (for example, here chosen as random features) with corresponding parameters $\hat{\boldsymbol\alpha}_{\text{rf}}$

The two parameter vectors are updated according to the continuous-time algorithm

\begin{align*}
  \frac{d}{dt}\nabla\psi_{\text{phys}}\left(\hat{\boldsymbol{\alpha}}_{\text{phys}}\right) &= - \eta_p \mathbf{Y}\left(\hat{\mathbf{x}}, t\right)_{\text{phys}}^{\mathsf{T}}\mathbf{e}_{\text{phys}}(\mathbf{x}, \hat{\mathbf{x}}),\\
  \frac{d}{dt}\nabla\psi_{\text{rf}}\left(\hat{\boldsymbol{\alpha}}_{\text{rf}}\right) &= - \eta_{\text{rf}} \mathbf{Y}\left(\hat{\mathbf{x}}, t\right)_{\text{rf}}^{\mathsf{T}}\mathbf{e}_{\text{rf}}(\mathbf{x}, \hat{\mathbf{x}}),
\end{align*}

Above, $\psi_i$ represent mirror descent potentials used to regularize the two predictor dynamics, or to improve adaptation based on the geometry of the basis functions.

$\mathbf{e}_i(\cdot, \cdot)$ are error signals used to drive the adaptation

$\eta_i$ are two learning rates for the individual parameter vectos.

Below, we provide a mathematical instantiation of this formulation for rigid-body dynamics, and subsequently implement that mathematical instantation.

```python
#### Adaptive predictor parameters
# mirror descent potentials \psi_i and parameters for those potentials.
potentials       = {'phys': 'euclid', 'rf': 'euclid'}
potential_params = {'phys': 1.0, 'rf': 1.0}

# we define a frame as a sequence of steps between measurements.
# number of predictor steps for each frame (steps between measurements) and number of total
# frames (number of observations we receive over the course of the trajectory)
nsteps = 25
nframes = 40

# learning rates
eta_p = 10.0 * nsteps
eta_rf = 20.0 * nsteps

# measurement feedback gain, used in g(x, \hat{x})
k = 25 * nsteps

# random feature info
n_features = 50
n_avg = 10
input_dim = 6
approx_dim = 3

# initialize the parameters
nparams = 18 + n_features
init_params = onp.zeros(nparams)

# strength of noise added to observations
noise_str = 0.0

# dimension of rigid-body dynamics
sys_dim = 13
```

# Adaptive prediction for rigid-body dynamics
The rotational and translational motion of a rigid-body may be specified in continuous-time as
\begin{align}
   \dot{\mathbf{q}} &= \frac{1}{2}\mathbf{\omega}\circ\boldsymbol{q}\\
  \dot{\mathbf{p}} &= \mathbf{v}\\
  \dot{\mathbf{v}} &= -\mathbf{g}\\
  \dot{\boldsymbol\omega} &= - \mathbf{J}^{-1}\boldsymbol{\omega}\times\mathbf{J}\boldsymbol{\omega}
\end{align}

Above, $\circ$ denotes the quaternion product, $\times$ denotes the cross product, and $\mathbf{J}$ is the moment of inertia matrix. Boldface denotes vectors and capital boldface denotes matrices. We have written the translational dynamics in the world frame, where they decouple from the rotational dynamics. We have written the rotational dynamics in the body frame, where the moment of intertia tensor becomes a constant. For simplicity, in this notebook, we assume that we receive measurements for the translational and rotational motion of the center of mass; measurements from a fixed location in the body can similarly be handled with additional complications.

The unknown physical parameters are the components of $\mathbf{J}$. Because the only unknown parameters are contained in $\mathbf{J}$, we know the dynamics of $\mathbf{q}$, $\mathbf{p}$, and $\mathbf{v}$.

While the term $- \mathbf{J}^{-1}\boldsymbol{\omega}\times\mathbf{J}\boldsymbol{\omega}$ is nonlinear in the unknown inertia matrix $\mathbf{J}$, it can be written as a linear expansion in a set of parameters of larger dimension.

\begin{align}
  \dot{\hat{\mathbf{q}}} &= \frac{1}{2}\hat{\mathbf{\omega}}\circ\hat{\boldsymbol{q}} + k(\mathbf{q} - \hat{\mathbf{q}})\\
  \dot{\hat{\mathbf{p}}} &= \hat{\mathbf{v}} + k(\mathbf{p} - \hat{\mathbf{p}})\\
  \dot{\hat{\mathbf{v}}} &= -\mathbf{g} + k\left(\mathbf{v} - \hat{\mathbf{v}}\right)\\
  \dot{\hat{\boldsymbol\omega}} &= \mathbf{Y}(\hat{\boldsymbol\omega})\hat{\boldsymbol\alpha}_{\text{phys}} + k(\boldsymbol\omega - \hat{\boldsymbol\omega})
\end{align}

We have chosen a simple linear feedback term, $\mathbf{g}(\mathbf{x}, \hat{\mathbf{x}}) = k(\mathbf{x} - \hat{\mathbf{x}})$.

The physical parameters are upated according to

\begin{equation*}
  \frac{d}{dt}\nabla\psi_{\text{phys}}\left(\hat{\boldsymbol{\alpha}}_{\text{phys}}\right) = - \eta_p \mathbf{Y}\left(\hat{\mathbf{x}}, t\right)_{\text{phys}}^{\mathsf{T}}\left(\hat{\boldsymbol\omega} - \boldsymbol\omega\right)
\end{equation*}

The error signal $\mathbf{e}_{\text{phys}}(\mathbf{x}, \hat{\mathbf{x}}) = \hat{\boldsymbol\omega} - \boldsymbol{\omega}$ because the parameter estimates $\hat{\boldsymbol{\alpha}}_{\text{phys}}$ only appear in the $\dot{\hat{\boldsymbol\omega}}$ dynamics.

We may also add random features to the entire predictor. The parameters for this unstructured approximation would then be updated according to

\begin{equation*}
  \frac{d}{dt}\nabla\psi_{\text{rf}}\left(\hat{\boldsymbol{\alpha}}_{\text{rf}}\right) = - \eta_{\text{rf}} \mathbf{Y}\left(\hat{\mathbf{x}}, t\right)_{\text{rf}}^{\mathsf{T}}\left(\hat{\mathbf{x}} - \mathbf{x}\right)
\end{equation*}

Here, the error signal is given by the total error $\mathbf{\hat{x}} - \mathbf{x}$ because we add the random feature approximation to the entire predictor. In general, if we plan to estimate disturbances in only a few components, the random feature adaptation can be driven by the error in those components and does not have to be a function of the full predictor state.

```python
grav = jnp.array([0., 0., -9.8])

@jax.jit
def world_to_body(q: jnp.ndarray, vec: jnp.ndarray) -> jnp.ndarray:
  """ Takes vec in the world frame and maps it to the body frame. """
  R = to_rotation_matrix(q).T
  return  R @ vec

@jax.jit
def phys_error_signal(x: jnp.ndarray, xhat: jnp.ndarray) -> jnp.ndarray:
  """Returns the error used in the adaptation for the inertial params.
  Assumes that the measurements all come in in world coordinates."""
  q = x[:4]
  w_world = x[10:]
#  w = world_to_body(q, w_world)
  w = w_world
  what = xhat[10:]
  return what - w

@jax.jit
def norm_state(xhat: jnp.ndarray) -> jnp.ndarray:
  """ State normalization via quaternion projection. """
  return jnp.concatenate((normalize(xhat[:4]), xhat[4:]))

@jax.jit
def inertial_basis(xhat: jnp.ndarray, t: float) -> jnp.ndarray:
  """ Overparameterized inertial basis matrix. """
  x, y, z = xhat[-3:]
  inert_mat = jnp.array([
      [x*z, y*z, x*y, x*x, y*y, z*z, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
      [0, 0, 0, 0, 0, 0, x*y, x*z, z*y, y*y, x*x, z*z, 0, 0, 0, 0, 0, 0],
      [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, x*z, z*y, x*y, x*x, y*y, z*z]])
  return inert_mat

@jax.jit
def measurement_function(x: jnp.ndarray, xhat: jnp.ndarray) -> jnp.ndarray:
  """Measurement function."""
  q = x[:4]
  w_world = x[10:]
#  w = world_to_body(q, w_world)
  w = w_world
  obs = jnp.concatenate((q, x[4:10], w))
  return k*(obs - xhat)

def init_random_features(key: jnp.ndarray, approx_dim: int, input_dim: int,
                         n_features: int, n_avg: int):
  """Constructs the random basis functions."""
  key, skey = jax.random.split(key)
  w = jax.random.normal(skey, shape=(approx_dim, n_features, n_avg, input_dim))
  key, skey = jax.random.split(key)
  b = jax.random.uniform(skey, shape=(approx_dim, n_features, n_avg))
  return lambda x: jnp.sum(jnp.cos(w @ x + b), axis=2)/jnp.sqrt(n_avg)

# construct the random feature approximators
key = jax.random.PRNGKey(0)
key, skey = jax.random.split(key)
rf_basis = init_random_features(skey, approx_dim, input_dim, n_features, n_avg)

@jax.jit
def rf_error_signal(x: jnp.ndarray, xhat: jnp.ndarray):
  """ Error signal for random feature adaptation. """
  v, vhat = x[7:10], xhat[7:10]
  return vhat - v

@functools.partial(jax.jit, static_argnums=3)
def adapt_continuous(x: jnp.ndarray, xhat_and_mirror_params: jnp.ndarray, t: float, sys_dim) -> jnp.ndarray:
  """Continuous definition of the adaptive predictor."""
  xhat = xhat_and_mirror_params[:sys_dim]
  vhat_and_what = xhat[7:]
  phys_adapt = -eta_p* inertial_basis(xhat, t).T @ phys_error_signal(x, xhat)
  rf_adapt = -eta_rf * rf_basis(vhat_and_what).T @ rf_error_signal(x, xhat)
  return jnp.concatenate((phys_adapt, rf_adapt))

@jax.jit
def predictor_continuous(xhat: jnp.ndarray, params: jnp.ndarray, t: float) -> jnp.ndarray:
  """Returns the dynamics predictor in continuous-time without measurement.

  Assumes that the state is given by:
              (quaternion, position, velocity, angular velocity)

  Args:
    xhat: Predicted state vector.
    params: Current parameter estimates.
    t: The current time.

  Returns:
    The continuous-time dynamics.
  """
  qhat = xhat[0:4]
  vhat = xhat[7:10]
  what = xhat[10:]
  vhat_and_what = xhat[7:]

  qhat_dot = .5 * vector_multiply(what, qhat)
  phat_dot = vhat

  # add random features to velocity for estimation of magnus effect.
  # magnus and lift depend on translational and angular velocity.
  vhat_dot = grav + rf_basis(vhat_and_what) @ params[nparams - n_features:]
  what_dot = inertial_basis(xhat, t) @ params[:nparams - n_features]

  return jnp.concatenate((qhat_dot, phat_dot, vhat_dot, what_dot))
```

The following code cell demonstrates how to construct an ```AdaptivePredictor``` object using the above implementation and how to use the predictor for rigid-body dynamics with PyBullet.

```python
# set object for experiments
object_id = sphere

# set initial velocity and grab the initial measurement
set_initial_state(object_id, [0, 0, 10.], [0, 0, 0, 1.], [1, 3, 2.], [3., 2., 1.])
x_init = get_state(object_id)

# construct the adaptive predictor
predictor = AdaptivePredictor(predictor_continuous, measurement_function, adapt_continuous, norm_state,
                              x_init, init_params, n_features, dt, nsteps, nframes, potentials,
                              potential_params, prediction_index=nframes)

# store the state of the predictor and the parameters over the entire trajectory.
predictor_states = []
params_over_time = onp.zeros((nframes, nparams))
params_over_time[0, :] = init_params

# store the pybullet state over time
state_observations = onp.zeros((1+nsteps*nframes, 13))
state_observations[0, :] = x_init

# rollout the simulation
state_observation = x_init
for i in range(nsteps*nframes):
  p.stepSimulation()

  # incorporate the measurement and update the parameters
  new_measurement = (i % nsteps) == 0
  if new_measurement:
    ind = int(i / nsteps)
    predictor.update_state(state_observation + 2*noise_str*onp.random.uniform(size=13) - noise_str)
    predictor_states.append(predictor.state_prediction.copy())
    params_over_time[ind, :] = predictor.params

  # save the state observations over time for direct comparison
  state_observation = get_state(object_id)
  state_observations[i+1, :] = state_observation
```

The following two code cells provide ```matplotlib``` defaults for nice-looking plots, as well as convenience functions for visualizing the performance of the predictor.

```python
### figure parameters
%matplotlib inline
sns.set_style("white")
mpl.rcParams['axes.grid'] = True
mpl.rcParams['axes.grid.which'] = 'both'
mpl.rcParams['xtick.minor.visible'] = True
mpl.rcParams['ytick.minor.visible'] = True
mpl.rcParams['grid.color'] = '0.8'
mpl.rcParams['axes.facecolor'] = 'white'
mpl.rcParams['text.usetex'] = False
mpl.rcParams['font.family'] = 'serif'

cmap3 = sns.cubehelix_palette(n_colors=3, start=.5, rot=-.75, light=.6, dark=.15, hue=1, gamma=.95, reverse=False)
cmap4 = sns.cubehelix_palette(n_colors=4, start=.5, rot=-.75, light=.6, dark=.15, hue=1, gamma=.95, reverse=False)
cmap_ps = sns.cubehelix_palette(n_colors=nparams, start=.5, rot=-.75, light=.6, dark=.15, hue=1, gamma=.95, reverse=False)

fs         = (8, 5)
lab_size   = 20
title_size = 20
ticksize   = 20
leg_size   = 20
lw         = 2
leg_alpha  = 0
```

```python
### Visualization convenience functions.
def make_prediction_comparison_plot(n_total_steps: int, dt: float, state_observations: jnp.ndarray, predictor_states: list, measurement_no: int, variable_name: str):
  """Produces a plot of the prediction and the prediction error, produced after receiving measurement measurement_no, for the state variables variable_name."""
  state_predictions = predictor_states[measurement_no]
  times = dt*onp.arange(n_total_steps)
  if variable_name == 'orientation':
    fig, ax = plt.subplots(figsize=fs)
    plt.plot(times, state_observations[:, 0], label=r"$q_w$", color=cmap4[0])
    plt.plot(times, state_observations[:, 1], label=r"$q_x$", color=cmap4[1])
    plt.plot(times, state_observations[:, 2], label=r"$q_y$", color=cmap4[2])
    plt.plot(times, state_observations[:, 3], label=r"$q_z$", color=cmap4[3])
    plt.plot(times, state_predictions[:, 0], label=r"$\hat{q}_w$", color=cmap4[0], linestyle='--')
    plt.plot(times, state_predictions[:, 1], label=r"$\hat{q}_x$", color=cmap4[1], linestyle='--')
    plt.plot(times, state_predictions[:, 2], label=r"$\hat{q}_y$", color=cmap4[2], linestyle='--')
    plt.plot(times, state_predictions[:, 3], label=r"$\hat{q}_z$", color=cmap4[3], linestyle='--')
    plt.axvline(x = dt*nsteps*measurement_no)
    plt.ylabel(variable_name, fontsize=lab_size)
    plt.xlabel('time', fontsize=lab_size)
    plt.tick_params(axis='both', labelsize=ticksize)
    ax.minorticks_on()
    plt.tight_layout()
    ax.grid(True, which='both')
    plt.legend(fontsize=lab_size, ncol=4, framealpha=leg_alpha)

    fig, ax = plt.subplots(figsize=fs)
    plt.plot(times, state_predictions[:, 0] - state_observations[:, 0], label=r"$\tilde{q}_w$", color=cmap4[0])
    plt.plot(times, state_predictions[:, 1] - state_observations[:, 1], label=r"$\tilde{q}_x$", color=cmap4[1])
    plt.plot(times, state_predictions[:, 2] - state_observations[:, 2], label=r"$\tilde{q}_y$", color=cmap4[2])
    plt.plot(times, state_predictions[:, 3] - state_observations[:, 3], label=r"$\tilde{q}_z$", color=cmap4[3])
    plt.axvline(x = dt*nsteps*measurement_no)
    plt.ylabel(variable_name + ' error', fontsize=lab_size)
    plt.xlabel('time', fontsize=lab_size)
    plt.tick_params(axis='both', labelsize=ticksize)
    ax.minorticks_on()
    plt.tight_layout()
    ax.grid(True, which='both')
    plt.legend(fontsize=lab_size, ncol=2, framealpha=leg_alpha)
  elif variable_name == 'position':
    fig, ax = plt.subplots(figsize=fs)
    plt.plot(times, state_observations[:, 4], label=r"$x$", color=cmap3[0])
    plt.plot(times, state_observations[:, 5], label=r"$y$", color=cmap3[1])
    plt.plot(times, state_observations[:, 6], label=r"$z$", color=cmap3[2])
    plt.plot(times, state_predictions[:, 4], label=r"$\hat{x}$", color=cmap3[0], linestyle='--')
    plt.plot(times, state_predictions[:, 5], label=r"$\hat{y}$", color=cmap3[1], linestyle='--')
    plt.plot(times, onp.maximum(0, state_predictions[:, 6]), label=r"$\hat{z}$", color=cmap3[2], linestyle='--')
    plt.axvline(x = dt*nsteps*measurement_no)
    plt.ylabel(variable_name, fontsize=lab_size)
    plt.xlabel('time', fontsize=lab_size)
    plt.tick_params(axis='both', labelsize=ticksize)
    ax.minorticks_on()
    plt.tight_layout()
    ax.grid(True, which='both')
    plt.legend(fontsize=lab_size, ncol=2, framealpha=leg_alpha)

    fig, ax = plt.subplots(figsize=fs)
    plt.plot(times, state_predictions[:, 4] - state_observations[:, 4], label=r"$\tilde{x}$", color=cmap3[0])
    plt.plot(times, state_predictions[:, 5] - state_observations[:, 5], label=r"$\tilde{y}$", color=cmap3[1])
    plt.plot(times, onp.maximum(0, state_predictions[:, 6]) - state_observations[:, 6], label=r"$\tilde{z}$", color=cmap3[2])
    plt.axvline(x = dt*nsteps*measurement_no)
    plt.ylabel(variable_name + ' error', fontsize=lab_size)
    plt.xlabel('time', fontsize=lab_size)
    plt.tick_params(axis='both', labelsize=ticksize)
    ax.minorticks_on()
    plt.tight_layout()
    ax.grid(True, which='both')
    plt.legend(fontsize=lab_size, framealpha=leg_alpha)
  elif variable_name == 'linear velocity':
    fig, ax = plt.subplots(figsize=fs)
    plt.plot(times, state_observations[:, 7], label=r"$v_x$",      color=cmap3[0])
    plt.plot(times, state_observations[:, 8], label=r"$v_y$",      color=cmap3[1])
    plt.plot(times, state_observations[:, 9], label=r"$v_z$",      color=cmap3[2])
    plt.plot(times, state_predictions[:, 7], label=r"$\hat{v}_x$", color=cmap3[0], linestyle='--')
    plt.plot(times, state_predictions[:, 8], label=r"$\hat{v}_y$", color=cmap3[1], linestyle='--')
    plt.plot(times, state_predictions[:, 9], label=r"$\hat{v}_z$", color=cmap3[2], linestyle='--')
    plt.axvline(x = dt*nsteps*measurement_no)
    plt.ylabel(variable_name, fontsize=lab_size)
    plt.xlabel('time', fontsize=lab_size)
    plt.tick_params(axis='both', labelsize=ticksize)
    ax.minorticks_on()
    plt.tight_layout()
    ax.grid(True, which='both')
    plt.legend(fontsize=lab_size, ncol=2, framealpha=leg_alpha)

    fig, ax = plt.subplots(figsize=fs)
    plt.plot(times, state_predictions[:, 7] - state_observations[:, 7], label=r"$\tilde{v}_x$", color=cmap3[0])
    plt.plot(times, state_predictions[:, 8] - state_observations[:, 8], label=r"$\tilde{v}_y$", color=cmap3[1])
    plt.plot(times, state_predictions[:, 9] - state_observations[:, 9], label=r"$\tilde{v}_z$", color=cmap3[2])
    plt.ylabel(variable_name + ' error', fontsize=lab_size)
    plt.axvline(x = dt*nsteps*measurement_no)
    plt.xlabel('time', fontsize=lab_size)
    plt.tick_params(axis='both', labelsize=ticksize)
    ax.minorticks_on()
    plt.tight_layout()
    ax.grid(True, which='both')
    plt.legend(fontsize=lab_size, framealpha=leg_alpha)
  elif variable_name == 'angular velocity':
    fig, ax = plt.subplots(figsize=fs)
    plt.plot(times, state_observations[:, 10], label=r"$w_x$",      color=cmap3[0])
    plt.plot(times, state_observations[:, 11], label=r"$w_y$",      color=cmap3[1])
    plt.plot(times, state_observations[:, 12], label=r"$w_z$",      color=cmap3[2])
    plt.plot(times, state_predictions[:, 10], label=r"$\hat{w}_x$", color=cmap3[0], linestyle='--')
    plt.plot(times, state_predictions[:, 11], label=r"$\hat{w}_y$", color=cmap3[1], linestyle='--')
    plt.plot(times, state_predictions[:, 12], label=r"$\hat{w}_z$", color=cmap3[2], linestyle='--')
    plt.axvline(x = dt*nsteps*measurement_no)
    plt.ylabel(variable_name, fontsize=lab_size)
    plt.xlabel('time', fontsize=lab_size)
    plt.tick_params(axis='both', labelsize=ticksize)
    ax.minorticks_on()
    plt.tight_layout()
    ax.grid(True, which='both')
    plt.legend(fontsize=lab_size, ncol=3, framealpha=leg_alpha)

    fig, ax = plt.subplots(figsize=fs)
    plt.plot(times, state_predictions[:, 10] - state_observations[:, 10], label=r"$\tilde{w}_x$", color=cmap3[0])
    plt.plot(times, state_predictions[:, 11] - state_observations[:, 11], label=r"$\tilde{w}_y$", color=cmap3[1])
    plt.plot(times, state_predictions[:, 12] - state_observations[:, 12], label=r"$\tilde{w}_z$", color=cmap3[2])
    plt.axvline(x = dt*nsteps*measurement_no)
    plt.ylabel(variable_name + ' error', fontsize=lab_size)
    plt.xlabel('time', fontsize=lab_size)
    plt.tick_params(axis='both', labelsize=ticksize)
    ax.minorticks_on()
    plt.tight_layout()
    ax.grid(True, which='both')
    plt.legend(fontsize=lab_size, ncol=3, framealpha=leg_alpha)

def make_parameter_plot(nframes: int, nsteps: int, dt: float, parameters: jnp.ndarray, measurement_no: int):
  """ Produces a plot of the parameters over time. """
  times = dt*nsteps*onp.arange(nframes+1)
  fig, ax = plt.subplots(figsize=fs)
  plt.ylabel('parameters', fontsize=title_size)
  plt.xlabel('time', fontsize=lab_size)
  for ii in range(parameters.shape[1]):
    plt.plot(times[:-1], parameters[:, ii], color=cmap_ps[ii])
  plt.tick_params(axis='both', labelsize=ticksize)
  ax.minorticks_on()
  ax.grid(True, which='both')
  plt.axvline(x = dt*nsteps*measurement_no)
  plt.tight_layout()
```

```python
### Visualize the results.
measurement_no = 33
make_prediction_comparison_plot(nsteps*nframes+1, dt, state_observations, predictor_states, measurement_no, 'linear velocity')
make_parameter_plot(nframes, nsteps, dt, params_over_time, measurement_no)
```

## Dynamics prediction with rotational dynamics

Here we consider performing dynamics prediction for a system with nontrivial rotational dynamics. PyBullet makes it a bit difficult to do this, and also has a different integration system than our predictor. To simplify the problem, and to avoid these integration differences, we implement the rotational dynamics ourselves.

We will implement the "ground truth" rotational dynamics in the world frame, meaning that we have to keep track of a changing moment of inertia tensor using the parallel axis theorem. We will then take observations from a rollout of this system. The predictor will implement the translational dynamics in the world frame and the rotational dynamics in the body frame.

```python
# moment of inertia matrix for an ellipse.
J_def = jnp.array([[5.0, 1.0, 0],
                   [1.0, 2, 0],
                   [0, 0, 0.5]])

# mass of the object
m_def = 1.0

# drag and magnus terms.
Cd = 0.5
Cl = 0.5

def rotational_dynamics_body(x, u, t, J=J_def):
  """ Implements rotational dynamics w.r.t the center of mass in the body frame."""
  q, p, v, w = x[:4], x[4:7], x[7:10], x[10:]
  qdot = .5 * vector_multiply(w, q)
  pdot = v
  R = to_rotation_matrix(q)
  w_world = R @ w
  vdot = grav - Cd*jnp.linalg.norm(v)*v + Cl*jnp.cross(w_world, v)
  wdot = -jnp.linalg.inv(J) @ jnp.cross(w, J @ w)

  return jnp.concatenate((qdot, pdot, vdot, wdot))

#disc_rot_dyn = jax.jit(rk4(rotational_dynamics, dt=dt))
disc_rot_dyn = jax.jit(rk4(rotational_dynamics_body, dt=dt))
```

```python
# initialize with nontrivial rotational dynamics.
q_init = jnp.array([0, 0, 0, 1])
p_init = jnp.array([0, 0, 10.0])
v_init = jnp.array([1.5, 0.0, 2.5])
w_init = jnp.array([3.0, 2.0, 1.0])
x_init = jnp.concatenate((q_init, p_init, v_init, w_init))

# construct the adaptive predictor
predictor = AdaptivePredictor(predictor_continuous, measurement_function, adapt_continuous, norm_state,
                              x_init, init_params, n_features, dt, nsteps, nframes, potentials,
                              potential_params, prediction_index=nframes)

# store the state of the predictor and the parameters over the entire trajectory.
predictor_states = []
params_over_time = onp.zeros((nframes, nparams))
params_over_time[0, :] = init_params

# store the pybullet state over time
state_observations = onp.zeros((1+nsteps*nframes, 13))
state_observations[0, :] = x_init

# rollout the simulation
state_observation = x_init
xt = x_init
for t in range(nsteps*nframes):
  xt = disc_rot_dyn(xt, 0, t)

  # incorporate the measurement and update the parameters
  new_measurement = (t % nsteps) == 0
  if new_measurement:
    ind = int(t / nsteps)
    predictor.update_state(state_observation)
    predictor_states.append(predictor.state_prediction.copy())
    params_over_time[ind, :] = predictor.params

  # save the state observations over time for direct comparison
  state_observation = xt
  state_observations[t+1, :] = xt

  print("Finished iteration %d" % t)
```

    Finished iteration 0
    Finished iteration 1
    Finished iteration 2
    Finished iteration 3
    Finished iteration 4
    Finished iteration 5
    Finished iteration 6
    Finished iteration 7
    Finished iteration 8
    Finished iteration 9
    Finished iteration 10
    Finished iteration 11
    Finished iteration 12
    Finished iteration 13
    Finished iteration 14
    Finished iteration 15
    Finished iteration 16
    Finished iteration 17
    Finished iteration 18
    Finished iteration 19
    Finished iteration 20
    Finished iteration 21
    Finished iteration 22
    Finished iteration 23
    Finished iteration 24
    Finished iteration 25
    Finished iteration 26
    Finished iteration 27
    Finished iteration 28
    Finished iteration 29
    Finished iteration 30
    Finished iteration 31
    Finished iteration 32
    Finished iteration 33
    Finished iteration 34
    Finished iteration 35
    Finished iteration 36
    Finished iteration 37
    Finished iteration 38
    Finished iteration 39
    Finished iteration 40
    Finished iteration 41
    Finished iteration 42
    Finished iteration 43
    Finished iteration 44
    Finished iteration 45
    Finished iteration 46
    Finished iteration 47
    Finished iteration 48
    Finished iteration 49
    Finished iteration 50
    Finished iteration 51
    Finished iteration 52
    Finished iteration 53
    Finished iteration 54
    Finished iteration 55
    Finished iteration 56
    Finished iteration 57
    Finished iteration 58
    Finished iteration 59
    Finished iteration 60
    Finished iteration 61
    Finished iteration 62
    Finished iteration 63
    Finished iteration 64
    Finished iteration 65
    Finished iteration 66
    Finished iteration 67
    Finished iteration 68
    Finished iteration 69
    Finished iteration 70
    Finished iteration 71
    Finished iteration 72
    Finished iteration 73
    Finished iteration 74
    Finished iteration 75
    Finished iteration 76
    Finished iteration 77
    Finished iteration 78
    Finished iteration 79
    Finished iteration 80
    Finished iteration 81
    Finished iteration 82
    Finished iteration 83
    Finished iteration 84
    Finished iteration 85
    Finished iteration 86
    Finished iteration 87
    Finished iteration 88
    Finished iteration 89
    Finished iteration 90
    Finished iteration 91
    Finished iteration 92
    Finished iteration 93
    Finished iteration 94
    Finished iteration 95
    Finished iteration 96
    Finished iteration 97
    Finished iteration 98
    Finished iteration 99
    Finished iteration 100
    Finished iteration 101
    Finished iteration 102
    Finished iteration 103
    Finished iteration 104
    Finished iteration 105
    Finished iteration 106
    Finished iteration 107
    Finished iteration 108
    Finished iteration 109
    Finished iteration 110
    Finished iteration 111
    Finished iteration 112
    Finished iteration 113
    Finished iteration 114
    Finished iteration 115
    Finished iteration 116
    Finished iteration 117
    Finished iteration 118
    Finished iteration 119
    Finished iteration 120
    Finished iteration 121
    Finished iteration 122
    Finished iteration 123
    Finished iteration 124
    Finished iteration 125
    Finished iteration 126
    Finished iteration 127
    Finished iteration 128
    Finished iteration 129
    Finished iteration 130
    Finished iteration 131
    Finished iteration 132
    Finished iteration 133
    Finished iteration 134
    Finished iteration 135
    Finished iteration 136
    Finished iteration 137
    Finished iteration 138
    Finished iteration 139
    Finished iteration 140
    Finished iteration 141
    Finished iteration 142
    Finished iteration 143
    Finished iteration 144
    Finished iteration 145
    Finished iteration 146
    Finished iteration 147
    Finished iteration 148
    Finished iteration 149
    Finished iteration 150
    Finished iteration 151
    Finished iteration 152
    Finished iteration 153
    Finished iteration 154
    Finished iteration 155
    Finished iteration 156
    Finished iteration 157
    Finished iteration 158
    Finished iteration 159
    Finished iteration 160
    Finished iteration 161
    Finished iteration 162
    Finished iteration 163
    Finished iteration 164
    Finished iteration 165
    Finished iteration 166
    Finished iteration 167
    Finished iteration 168
    Finished iteration 169
    Finished iteration 170
    Finished iteration 171
    Finished iteration 172
    Finished iteration 173
    Finished iteration 174
    Finished iteration 175
    Finished iteration 176
    Finished iteration 177
    Finished iteration 178
    Finished iteration 179
    Finished iteration 180
    Finished iteration 181
    Finished iteration 182
    Finished iteration 183
    Finished iteration 184
    Finished iteration 185
    Finished iteration 186
    Finished iteration 187
    Finished iteration 188
    Finished iteration 189
    Finished iteration 190
    Finished iteration 191
    Finished iteration 192
    Finished iteration 193
    Finished iteration 194
    Finished iteration 195
    Finished iteration 196
    Finished iteration 197
    Finished iteration 198
    Finished iteration 199
    Finished iteration 200
    Finished iteration 201
    Finished iteration 202
    Finished iteration 203
    Finished iteration 204
    Finished iteration 205
    Finished iteration 206
    Finished iteration 207
    Finished iteration 208
    Finished iteration 209
    Finished iteration 210
    Finished iteration 211
    Finished iteration 212
    Finished iteration 213
    Finished iteration 214
    Finished iteration 215
    Finished iteration 216
    Finished iteration 217
    Finished iteration 218
    Finished iteration 219
    Finished iteration 220
    Finished iteration 221
    Finished iteration 222
    Finished iteration 223
    Finished iteration 224
    Finished iteration 225
    Finished iteration 226
    Finished iteration 227
    Finished iteration 228
    Finished iteration 229
    Finished iteration 230
    Finished iteration 231
    Finished iteration 232
    Finished iteration 233
    Finished iteration 234
    Finished iteration 235
    Finished iteration 236
    Finished iteration 237
    Finished iteration 238
    Finished iteration 239
    Finished iteration 240
    Finished iteration 241
    Finished iteration 242
    Finished iteration 243
    Finished iteration 244
    Finished iteration 245
    Finished iteration 246
    Finished iteration 247
    Finished iteration 248
    Finished iteration 249
    Finished iteration 250
    Finished iteration 251
    Finished iteration 252
    Finished iteration 253
    Finished iteration 254
    Finished iteration 255
    Finished iteration 256
    Finished iteration 257
    Finished iteration 258
    Finished iteration 259
    Finished iteration 260
    Finished iteration 261
    Finished iteration 262
    Finished iteration 263
    Finished iteration 264
    Finished iteration 265
    Finished iteration 266
    Finished iteration 267
    Finished iteration 268
    Finished iteration 269
    Finished iteration 270
    Finished iteration 271
    Finished iteration 272
    Finished iteration 273
    Finished iteration 274
    Finished iteration 275
    Finished iteration 276
    Finished iteration 277
    Finished iteration 278
    Finished iteration 279
    Finished iteration 280
    Finished iteration 281
    Finished iteration 282
    Finished iteration 283
    Finished iteration 284
    Finished iteration 285
    Finished iteration 286
    Finished iteration 287
    Finished iteration 288
    Finished iteration 289
    Finished iteration 290
    Finished iteration 291
    Finished iteration 292
    Finished iteration 293
    Finished iteration 294
    Finished iteration 295
    Finished iteration 296
    Finished iteration 297
    Finished iteration 298
    Finished iteration 299
    Finished iteration 300
    Finished iteration 301
    Finished iteration 302
    Finished iteration 303
    Finished iteration 304
    Finished iteration 305
    Finished iteration 306
    Finished iteration 307
    Finished iteration 308
    Finished iteration 309
    Finished iteration 310
    Finished iteration 311
    Finished iteration 312
    Finished iteration 313
    Finished iteration 314
    Finished iteration 315
    Finished iteration 316
    Finished iteration 317
    Finished iteration 318
    Finished iteration 319
    Finished iteration 320
    Finished iteration 321
    Finished iteration 322
    Finished iteration 323
    Finished iteration 324
    Finished iteration 325
    Finished iteration 326
    Finished iteration 327
    Finished iteration 328
    Finished iteration 329
    Finished iteration 330
    Finished iteration 331
    Finished iteration 332
    Finished iteration 333
    Finished iteration 334
    Finished iteration 335
    Finished iteration 336
    Finished iteration 337
    Finished iteration 338
    Finished iteration 339
    Finished iteration 340
    Finished iteration 341
    Finished iteration 342
    Finished iteration 343
    Finished iteration 344
    Finished iteration 345
    Finished iteration 346
    Finished iteration 347
    Finished iteration 348
    Finished iteration 349
    Finished iteration 350
    Finished iteration 351
    Finished iteration 352
    Finished iteration 353
    Finished iteration 354
    Finished iteration 355
    Finished iteration 356
    Finished iteration 357
    Finished iteration 358
    Finished iteration 359
    Finished iteration 360
    Finished iteration 361
    Finished iteration 362
    Finished iteration 363
    Finished iteration 364
    Finished iteration 365
    Finished iteration 366
    Finished iteration 367
    Finished iteration 368
    Finished iteration 369
    Finished iteration 370
    Finished iteration 371
    Finished iteration 372
    Finished iteration 373
    Finished iteration 374
    Finished iteration 375
    Finished iteration 376
    Finished iteration 377
    Finished iteration 378
    Finished iteration 379
    Finished iteration 380
    Finished iteration 381
    Finished iteration 382
    Finished iteration 383
    Finished iteration 384
    Finished iteration 385
    Finished iteration 386
    Finished iteration 387
    Finished iteration 388
    Finished iteration 389
    Finished iteration 390
    Finished iteration 391
    Finished iteration 392
    Finished iteration 393
    Finished iteration 394
    Finished iteration 395
    Finished iteration 396
    Finished iteration 397
    Finished iteration 398
    Finished iteration 399
    Finished iteration 400
    Finished iteration 401
    Finished iteration 402
    Finished iteration 403
    Finished iteration 404
    Finished iteration 405
    Finished iteration 406
    Finished iteration 407
    Finished iteration 408
    Finished iteration 409
    Finished iteration 410
    Finished iteration 411
    Finished iteration 412
    Finished iteration 413
    Finished iteration 414
    Finished iteration 415
    Finished iteration 416
    Finished iteration 417
    Finished iteration 418
    Finished iteration 419
    Finished iteration 420
    Finished iteration 421
    Finished iteration 422
    Finished iteration 423
    Finished iteration 424
    Finished iteration 425
    Finished iteration 426
    Finished iteration 427
    Finished iteration 428
    Finished iteration 429
    Finished iteration 430
    Finished iteration 431
    Finished iteration 432
    Finished iteration 433
    Finished iteration 434
    Finished iteration 435
    Finished iteration 436
    Finished iteration 437
    Finished iteration 438
    Finished iteration 439
    Finished iteration 440
    Finished iteration 441
    Finished iteration 442
    Finished iteration 443
    Finished iteration 444
    Finished iteration 445
    Finished iteration 446
    Finished iteration 447
    Finished iteration 448
    Finished iteration 449
    Finished iteration 450
    Finished iteration 451
    Finished iteration 452
    Finished iteration 453
    Finished iteration 454
    Finished iteration 455
    Finished iteration 456
    Finished iteration 457
    Finished iteration 458
    Finished iteration 459
    Finished iteration 460
    Finished iteration 461
    Finished iteration 462
    Finished iteration 463
    Finished iteration 464
    Finished iteration 465
    Finished iteration 466
    Finished iteration 467
    Finished iteration 468
    Finished iteration 469
    Finished iteration 470
    Finished iteration 471
    Finished iteration 472
    Finished iteration 473
    Finished iteration 474
    Finished iteration 475
    Finished iteration 476
    Finished iteration 477
    Finished iteration 478
    Finished iteration 479
    Finished iteration 480
    Finished iteration 481
    Finished iteration 482
    Finished iteration 483
    Finished iteration 484
    Finished iteration 485
    Finished iteration 486
    Finished iteration 487
    Finished iteration 488
    Finished iteration 489
    Finished iteration 490
    Finished iteration 491
    Finished iteration 492
    Finished iteration 493
    Finished iteration 494
    Finished iteration 495
    Finished iteration 496
    Finished iteration 497
    Finished iteration 498
    Finished iteration 499
    Finished iteration 500
    Finished iteration 501
    Finished iteration 502
    Finished iteration 503
    Finished iteration 504
    Finished iteration 505
    Finished iteration 506
    Finished iteration 507
    Finished iteration 508
    Finished iteration 509
    Finished iteration 510
    Finished iteration 511
    Finished iteration 512
    Finished iteration 513
    Finished iteration 514
    Finished iteration 515
    Finished iteration 516
    Finished iteration 517
    Finished iteration 518
    Finished iteration 519
    Finished iteration 520
    Finished iteration 521
    Finished iteration 522
    Finished iteration 523
    Finished iteration 524
    Finished iteration 525
    Finished iteration 526
    Finished iteration 527
    Finished iteration 528
    Finished iteration 529
    Finished iteration 530
    Finished iteration 531
    Finished iteration 532
    Finished iteration 533
    Finished iteration 534
    Finished iteration 535
    Finished iteration 536
    Finished iteration 537
    Finished iteration 538
    Finished iteration 539
    Finished iteration 540
    Finished iteration 541
    Finished iteration 542
    Finished iteration 543
    Finished iteration 544
    Finished iteration 545
    Finished iteration 546
    Finished iteration 547
    Finished iteration 548
    Finished iteration 549
    Finished iteration 550
    Finished iteration 551
    Finished iteration 552
    Finished iteration 553
    Finished iteration 554
    Finished iteration 555
    Finished iteration 556
    Finished iteration 557
    Finished iteration 558
    Finished iteration 559
    Finished iteration 560
    Finished iteration 561
    Finished iteration 562
    Finished iteration 563
    Finished iteration 564
    Finished iteration 565
    Finished iteration 566
    Finished iteration 567
    Finished iteration 568
    Finished iteration 569
    Finished iteration 570
    Finished iteration 571
    Finished iteration 572
    Finished iteration 573
    Finished iteration 574
    Finished iteration 575
    Finished iteration 576
    Finished iteration 577
    Finished iteration 578
    Finished iteration 579
    Finished iteration 580
    Finished iteration 581
    Finished iteration 582
    Finished iteration 583
    Finished iteration 584
    Finished iteration 585
    Finished iteration 586
    Finished iteration 587
    Finished iteration 588
    Finished iteration 589
    Finished iteration 590
    Finished iteration 591
    Finished iteration 592
    Finished iteration 593
    Finished iteration 594
    Finished iteration 595
    Finished iteration 596
    Finished iteration 597
    Finished iteration 598
    Finished iteration 599
    Finished iteration 600
    Finished iteration 601
    Finished iteration 602
    Finished iteration 603
    Finished iteration 604
    Finished iteration 605
    Finished iteration 606
    Finished iteration 607
    Finished iteration 608
    Finished iteration 609
    Finished iteration 610
    Finished iteration 611
    Finished iteration 612
    Finished iteration 613
    Finished iteration 614
    Finished iteration 615
    Finished iteration 616
    Finished iteration 617
    Finished iteration 618
    Finished iteration 619
    Finished iteration 620
    Finished iteration 621
    Finished iteration 622
    Finished iteration 623
    Finished iteration 624
    Finished iteration 625
    Finished iteration 626
    Finished iteration 627
    Finished iteration 628
    Finished iteration 629
    Finished iteration 630
    Finished iteration 631
    Finished iteration 632
    Finished iteration 633
    Finished iteration 634
    Finished iteration 635
    Finished iteration 636
    Finished iteration 637
    Finished iteration 638
    Finished iteration 639
    Finished iteration 640
    Finished iteration 641
    Finished iteration 642
    Finished iteration 643
    Finished iteration 644
    Finished iteration 645
    Finished iteration 646
    Finished iteration 647
    Finished iteration 648
    Finished iteration 649
    Finished iteration 650
    Finished iteration 651
    Finished iteration 652
    Finished iteration 653
    Finished iteration 654
    Finished iteration 655
    Finished iteration 656
    Finished iteration 657
    Finished iteration 658
    Finished iteration 659
    Finished iteration 660
    Finished iteration 661
    Finished iteration 662
    Finished iteration 663
    Finished iteration 664
    Finished iteration 665
    Finished iteration 666
    Finished iteration 667
    Finished iteration 668
    Finished iteration 669
    Finished iteration 670
    Finished iteration 671
    Finished iteration 672
    Finished iteration 673
    Finished iteration 674
    Finished iteration 675
    Finished iteration 676
    Finished iteration 677
    Finished iteration 678
    Finished iteration 679
    Finished iteration 680
    Finished iteration 681
    Finished iteration 682
    Finished iteration 683
    Finished iteration 684
    Finished iteration 685
    Finished iteration 686
    Finished iteration 687
    Finished iteration 688
    Finished iteration 689
    Finished iteration 690
    Finished iteration 691
    Finished iteration 692
    Finished iteration 693
    Finished iteration 694
    Finished iteration 695
    Finished iteration 696
    Finished iteration 697
    Finished iteration 698
    Finished iteration 699
    Finished iteration 700
    Finished iteration 701
    Finished iteration 702
    Finished iteration 703
    Finished iteration 704
    Finished iteration 705
    Finished iteration 706
    Finished iteration 707
    Finished iteration 708
    Finished iteration 709
    Finished iteration 710
    Finished iteration 711
    Finished iteration 712
    Finished iteration 713
    Finished iteration 714
    Finished iteration 715
    Finished iteration 716
    Finished iteration 717
    Finished iteration 718
    Finished iteration 719
    Finished iteration 720
    Finished iteration 721
    Finished iteration 722
    Finished iteration 723
    Finished iteration 724
    Finished iteration 725
    Finished iteration 726
    Finished iteration 727
    Finished iteration 728
    Finished iteration 729
    Finished iteration 730
    Finished iteration 731
    Finished iteration 732
    Finished iteration 733
    Finished iteration 734
    Finished iteration 735
    Finished iteration 736
    Finished iteration 737
    Finished iteration 738
    Finished iteration 739
    Finished iteration 740
    Finished iteration 741
    Finished iteration 742
    Finished iteration 743
    Finished iteration 744
    Finished iteration 745
    Finished iteration 746
    Finished iteration 747
    Finished iteration 748
    Finished iteration 749
    Finished iteration 750
    Finished iteration 751
    Finished iteration 752
    Finished iteration 753
    Finished iteration 754
    Finished iteration 755
    Finished iteration 756
    Finished iteration 757
    Finished iteration 758
    Finished iteration 759
    Finished iteration 760
    Finished iteration 761
    Finished iteration 762
    Finished iteration 763
    Finished iteration 764
    Finished iteration 765
    Finished iteration 766
    Finished iteration 767
    Finished iteration 768
    Finished iteration 769
    Finished iteration 770
    Finished iteration 771
    Finished iteration 772
    Finished iteration 773
    Finished iteration 774
    Finished iteration 775
    Finished iteration 776
    Finished iteration 777
    Finished iteration 778
    Finished iteration 779
    Finished iteration 780
    Finished iteration 781
    Finished iteration 782
    Finished iteration 783
    Finished iteration 784
    Finished iteration 785
    Finished iteration 786
    Finished iteration 787
    Finished iteration 788
    Finished iteration 789
    Finished iteration 790
    Finished iteration 791
    Finished iteration 792
    Finished iteration 793
    Finished iteration 794
    Finished iteration 795
    Finished iteration 796
    Finished iteration 797
    Finished iteration 798
    Finished iteration 799
    Finished iteration 800
    Finished iteration 801
    Finished iteration 802
    Finished iteration 803
    Finished iteration 804
    Finished iteration 805
    Finished iteration 806
    Finished iteration 807
    Finished iteration 808
    Finished iteration 809
    Finished iteration 810
    Finished iteration 811
    Finished iteration 812
    Finished iteration 813
    Finished iteration 814
    Finished iteration 815
    Finished iteration 816
    Finished iteration 817
    Finished iteration 818
    Finished iteration 819
    Finished iteration 820
    Finished iteration 821
    Finished iteration 822
    Finished iteration 823
    Finished iteration 824
    Finished iteration 825
    Finished iteration 826
    Finished iteration 827
    Finished iteration 828
    Finished iteration 829
    Finished iteration 830
    Finished iteration 831
    Finished iteration 832
    Finished iteration 833
    Finished iteration 834
    Finished iteration 835
    Finished iteration 836
    Finished iteration 837
    Finished iteration 838
    Finished iteration 839
    Finished iteration 840
    Finished iteration 841
    Finished iteration 842
    Finished iteration 843
    Finished iteration 844
    Finished iteration 845
    Finished iteration 846
    Finished iteration 847
    Finished iteration 848
    Finished iteration 849
    Finished iteration 850
    Finished iteration 851
    Finished iteration 852
    Finished iteration 853
    Finished iteration 854
    Finished iteration 855
    Finished iteration 856
    Finished iteration 857
    Finished iteration 858
    Finished iteration 859
    Finished iteration 860
    Finished iteration 861
    Finished iteration 862
    Finished iteration 863
    Finished iteration 864
    Finished iteration 865
    Finished iteration 866
    Finished iteration 867
    Finished iteration 868
    Finished iteration 869
    Finished iteration 870
    Finished iteration 871
    Finished iteration 872
    Finished iteration 873
    Finished iteration 874
    Finished iteration 875
    Finished iteration 876
    Finished iteration 877
    Finished iteration 878
    Finished iteration 879
    Finished iteration 880
    Finished iteration 881
    Finished iteration 882
    Finished iteration 883
    Finished iteration 884
    Finished iteration 885
    Finished iteration 886
    Finished iteration 887
    Finished iteration 888
    Finished iteration 889
    Finished iteration 890
    Finished iteration 891
    Finished iteration 892
    Finished iteration 893
    Finished iteration 894
    Finished iteration 895
    Finished iteration 896
    Finished iteration 897
    Finished iteration 898
    Finished iteration 899
    Finished iteration 900
    Finished iteration 901
    Finished iteration 902
    Finished iteration 903
    Finished iteration 904
    Finished iteration 905
    Finished iteration 906
    Finished iteration 907
    Finished iteration 908
    Finished iteration 909
    Finished iteration 910
    Finished iteration 911
    Finished iteration 912
    Finished iteration 913
    Finished iteration 914
    Finished iteration 915
    Finished iteration 916
    Finished iteration 917
    Finished iteration 918
    Finished iteration 919
    Finished iteration 920
    Finished iteration 921
    Finished iteration 922
    Finished iteration 923
    Finished iteration 924
    Finished iteration 925
    Finished iteration 926
    Finished iteration 927
    Finished iteration 928
    Finished iteration 929
    Finished iteration 930
    Finished iteration 931
    Finished iteration 932
    Finished iteration 933
    Finished iteration 934
    Finished iteration 935
    Finished iteration 936
    Finished iteration 937
    Finished iteration 938
    Finished iteration 939
    Finished iteration 940
    Finished iteration 941
    Finished iteration 942
    Finished iteration 943
    Finished iteration 944
    Finished iteration 945
    Finished iteration 946
    Finished iteration 947
    Finished iteration 948
    Finished iteration 949
    Finished iteration 950
    Finished iteration 951
    Finished iteration 952
    Finished iteration 953
    Finished iteration 954
    Finished iteration 955
    Finished iteration 956
    Finished iteration 957
    Finished iteration 958
    Finished iteration 959
    Finished iteration 960
    Finished iteration 961
    Finished iteration 962
    Finished iteration 963
    Finished iteration 964
    Finished iteration 965
    Finished iteration 966
    Finished iteration 967
    Finished iteration 968
    Finished iteration 969
    Finished iteration 970
    Finished iteration 971
    Finished iteration 972
    Finished iteration 973
    Finished iteration 974
    Finished iteration 975
    Finished iteration 976
    Finished iteration 977
    Finished iteration 978
    Finished iteration 979
    Finished iteration 980
    Finished iteration 981
    Finished iteration 982
    Finished iteration 983
    Finished iteration 984
    Finished iteration 985
    Finished iteration 986
    Finished iteration 987
    Finished iteration 988
    Finished iteration 989
    Finished iteration 990
    Finished iteration 991
    Finished iteration 992
    Finished iteration 993
    Finished iteration 994
    Finished iteration 995
    Finished iteration 996
    Finished iteration 997
    Finished iteration 998
    Finished iteration 999

```python
### Visualize the results.
measurement_no = 34
make_prediction_comparison_plot(nsteps*nframes+1, dt, state_observations, predictor_states, measurement_no, 'linear velocity')
```

```python
make_parameter_plot(nframes, nsteps, dt, params_over_time, measurement_no)
```


###### Content under Creative Commons Attribution license CC-BY 4.0, code under MIT license (c)2014 L.A. Barba, G.F. Forsyth, I. Hawke. Partly based on content by David Ketcheson, also under CC-BY.

# Full phugoid model

This is the third Jupyter Notebook of the series on the _phugoid model of glider flight_, our first learning module of the course ["Practical Numerical Methods with Python"](https://openedx.seas.gwu.edu/courses/course-v1:MAE+MAE6286+2017/about). In the [first notebook](https://nbviewer.jupyter.org/github/numerical-mooc/numerical-mooc/blob/master/lessons/01_phugoid/01_01_Phugoid_Theory.ipynb), we described the physics of the trajectories known as phugoids obtained from an exchange of potential and kinetic energy in an idealized motion with no drag. We gave you a neat little code to play with and plot various phugoid curves.

In the second notebook, we looked at the equation representing small perturbations on the straight-line phugoid, resulting in simple harmonic motion. This is a second-order ordinary differential equation, and we solved it numerically using **Euler's method**: the simplest numerical method of all. We learned about convergence and calculated the error of the numerical solution, comparing with an analytical solution. That is a good foundation!

Now, let's go back to the dynamical model, and take away the idealization of no-drag. Let's remind ourselves of the forces affecting an aircraft, considering now that it may be accelerating, with an instantaneous upward trajectory. We use the designation $\theta$ for the angle, and consider it positive upwards.

#### Figure 1. Forces with a positive trajectory angle.

In Figure 1, $L$ is the lift, $W$ is the weight, $D$ is the drag, and $\theta$ the positive angle of the trajectory, instantaneously.

In [Lesson 1](https://nbviewer.jupyter.org/github/numerical-mooc/numerical-mooc/blob/master/lessons/01_phugoid/01_01_Phugoid_Theory.ipynb), we wrote the force balance in the directions perpendicular and parallel to the trajectory for a glider in _equilibrium_. What if the forces are _not_ in balance? Well, there will be acceleration terms in the equations of motion, and we would have in that case:

$$
\begin{align}
m \frac{dv}{dt} & = - W \sin\theta - D \\
m v \, \frac{d\theta}{dt} & = - W \cos\theta + L
\end{align}
$$

We can use a few little tricks to make these equations more pleasing. First, use primes to denote the time derivatives and divide through by the weight:

$$
\begin{align}
 \frac{v'}{g} & = - \sin\theta - D/W \\
\frac{v}{g} \, \theta' & = - \cos\theta + L/W
\end{align}
$$

Recall, from our first lesson, that the ratio of lift to weight is known from the trim conditions—$L/W=v^2/v_t^2$— and also from the definitions of lift and drag,

$$
\begin{eqnarray}
L &=& C_L S \times \frac{1}{2} \rho v^2 \\
D &=& C_D S \times \frac{1}{2} \rho v^2
\end{eqnarray}
$$

we see that $L/D=C_L/C_D$. The system of equations can be re-written:

$$
\begin{align}
v' & = - g\, \sin\theta - \frac{C_D}{C_L} \frac{g}{v_t^2} v^2 \\
\theta' & = - \frac{g}{v}\,\cos\theta + \frac{g}{v_t^2}\, v
\end{align}
$$

It is very interesting that the first equation has the factor $C_D/C_L$, which is the inverse of a measure of the aerodynamic efficiency of the aircraft. It turns out, this is the term that contributes damping to the phugoid model: if drag is zero, there is no damping. Drag is never zero in real life, but as engineers design more aerodynamically efficient aircraft, they make the phugoid mode more weakly damped. At altitude, this is nothing but a slight bother, but vertical oscillations are unsafe during final approach to land, so this is something to watch out for!

## The initial value problem

If we want to visualize the flight trajectories predicted by this model, we are going to need to integrate the spatial coordinates, which depend on both the forward velocity (tangential to the trajectory) and the trajectory angle. The position of the glider on a vertical plane will be designated by coordinates $(x, y)$ with respect to an inertial frame of reference, and are obtained from:

$$
\begin{align}
x'(t) & = v \cos(\theta) \\
y'(t) & = v \sin(\theta)
\end{align}
$$

Augmenting our original two differential equations by the two equations above, we have a system of four first-order differential equations to solve. We will use a time-stepping approach, like in the previous lesson. To do so, we do need *initial values* for every unknown:

$$
v(0) = v_0 \quad \text{and} \quad \theta(0) = \theta_0 \\
x(0) = x_0 \quad \text{and} \quad y(0) = y_0
$$

## Solve with Euler's method

We know how to apply Euler's method from the previous lesson. We replace each of the time derivatives by an approximation of the form:

$$
v'(t) \approx \frac{v^{n+1} - v^n}{\Delta t}
$$

where we are now using a superscript $n$ to indicate the $n$-th value in the time iterations. The first differential equation, for example, gives:

$$
\frac{v^{n+1} - v^n}{\Delta t} = - g\, \sin\theta^n - \frac{C_D}{C_L} \frac{g}{v_t^2} (v^n)^2
$$

Alright, we know where this is going. At each time iteration $t^n$, we want to evaluate all the known data of our system to obtain the state at $t^{n+1}$—the next time step. We say that we are _stepping in time_ or _time marching_.

The full system of equations discretized with Euler's method is:

$$
\begin{align}
v^{n+1} & = v^n + \Delta t \left(- g\, \sin\theta^n - \frac{C_D}{C_L} \frac{g}{v_t^2} (v^n)^2 \right) \\
\theta^{n+1} & = \theta^n + \Delta t \left(- \frac{g}{v^n}\,\cos\theta^n + \frac{g}{v_t^2}\, v^n \right) \\
x^{n+1} & = x^n + \Delta t \, v^n \cos\theta^n \\
y^{n+1} & = y^n + \Delta t \, v^n \sin\theta^n.
\end{align}
$$

As we've learned before, the system of differential equations can also be written as a vector equation:

$$
u'(t) = f(u)
$$

where

$$
\begin{align}
u & = \begin{pmatrix} v \\ \theta \\ x \\ y \end{pmatrix} & f(u) & = \begin{pmatrix} - g\, \sin\theta - \frac{C_D}{C_L} \frac{g}{v_t^2} v^2 \\ - \frac{g}{v}\,\cos\theta + \frac{g}{v_t^2}\, v \\ v\cos\theta \\ v\sin\theta \end{pmatrix}
\end{align}
$$

It's a bit tricky to code the solution using a NumPy array holding all your independent variables. But if you do, a function for the Euler step can be written that takes any number of simultaneous equations. It simply steps in time using the same line of code:

```Python
def euler_step(u, f, dt):
    return u + dt * f(u)
```

This function can take a NumPy array `u` with any number of components. All we need to do is create an appropriate function `f(u)` describing our system of differential equations. Notice how we are passing a _function_ as part of the arguments list to `euler_step()`. Neat!

### And solve!

As always, we start by loading the modules and libraries that we need for this problem. We'll need a few transcendental functions, including the $\log$ for a convergence study later on. And remember: the line `%matplotlib inline` is a magic function that tells Matplotlib to give us the plots in the notebook (the default behavior of Matplotlib is to open a pop-up window).

```python
import math
import numpy
from matplotlib import pyplot
%matplotlib inline
```

In addition, we modify some entries of the `rcParams` dictionary of `pyplot` to define notebook-wide plotting parameters: font family and font size.
Here we go!

```python
# Set the font family and size to use for Matplotlib figures.
pyplot.rcParams['font.family'] = 'serif'
pyplot.rcParams['font.size'] = 16
```

Next, we need to set things up to start our numerical solution: the parameter
values and the _initial values_. You know what the acceleration of gravity is: 9.81 m/s$^2$, but what are good values for $C_D/C_L$, the inverse of the aerodynamic efficiency? Some possible values are given on a table in the Wikipedia entry for [lift-to-drag ratio](http://en.wikipedia.org/wiki/Lift-to-drag_ratio): a modern sailplane can have $L/D$ of 40 to 60, depending on span (and, in case you're interested, a flying squirrel has $L/D$ close to 2).

For the _trim velocity_, the speed range for typical sailplanes is between 65 and 280 km/hr, according to Wikipedia (it must be right!). Let's convert that to meters per second: 18 to 78 m/s. We'll pick a value somewhere in the middle of this range.

Here's a possible set of parameters for the simulation, but be sure to come back and change some of these, and see what happens!

```python
# Set parameters.
g = 9.81  # gravitational acceleration (m.s^{-2})
vt = 30.0  # trim velocity (m.s)
CD = 1.0 / 40  # drag coefficient
CL = 1.0  # lift coefficient

# Set initial conditions.
v0 = vt  # start at the trim velocity
theta0 = 0.0  # trajectory angle
x0 = 0.0  # horizontal position
y0 = 1000.0  # vertical position (altitude)
```

We'll define a function `rhs_phugoid()` to match the right-hand side of Equation (15), the full differential system in vector form. This function assumes that we have available the parameters defined above. If you re-execute the cell above with different parameter values, you can just run the solution without re-executing the function definition.

```python
def rhs_phugoid(u, CL, CD, g, vt):
    """
    Returns the right-hand side of the phugoid system of equations.

    Parameters
    ----------
    u : list or numpy.ndarray
        Solution at the previous time step
        as a list or 1D array of four floats.
    CL : float
        Lift coefficient.
    CD : float
        Drag coefficient.
    g : float
        Gravitational acceleration.
    vt : float
        Trim velocity.

    Returns
    -------
    rhs : numpy.ndarray
        The right-hand side of the system
        as a 1D array of four floats.
    """
    v, theta, x, y = u
    rhs = numpy.array([-g * math.sin(theta) - CD / CL * g / vt**2 * v**2,
                       -g * math.cos(theta) / v + g / vt**2 * v,
                       v * math.cos(theta),
                       v * math.sin(theta)])
    return rhs
```

Compare the code defining function `rhs_phugoid()` with the differential equations, and convince yourself that it's right!

$$
\begin{align}
u & = \begin{pmatrix} v \\ \theta \\ x \\ y \end{pmatrix} & f(u) & = \begin{pmatrix} - g\, \sin\theta - \frac{C_D}{C_L} \frac{g}{v_t^2} v^2 \\ - \frac{g}{v}\,\cos\theta + \frac{g}{v_t^2}\, v \\ v\cos\theta \\ v\sin\theta \end{pmatrix} \nonumber
\end{align}
$$

Now, Euler's method is implemented in a simple function `euler_step()`:

```python
def euler_step(u, f, dt, *args):
    """
    Returns the solution at the next time step using Euler's method.

    Parameters
    ----------
    u : numpy.ndarray
        Solution at the previous time step
        as a 1D array of floats.
    f : function
        Function to compute the right-hand side of the system.
    dt : float
        Time-step size.
    args : tuple, optional
        Positional arguments to pass to the function f.

    Returns
    -------
    u_new : numpy.ndarray
        The solution at the next time step
        as a 1D array of floats.
    """
    u_new = u + dt * f(u, *args)
    return u_new
```

**Note**—We use an optional input to the function `euler_step()`, named `*args`. It passes to the function `f()` an arbitrary number of arguments. Doing so, `euler_step()` can take any function `f()`, regardless of the number of arguments this function needs. Sweet! (Read the Python documentation about [Arbitrary Argument Lists](https://docs.python.org/3/tutorial/controlflow.html#arbitrary-argument-lists) for more explanations.)

After defining a final time for the solution, and the time step $\Delta t$, we can construct the grid in time using the NumPy function [`linspace()`](http://docs.scipy.org/doc/numpy/reference/generated/numpy.linspace.html). Make sure you study the decisions we made here to build the time grid: why do we add 1 to the definition of `N`, for example?

Look at the code below, and make sure you understand the following aspects of it.

* The NumPy array `u` contains the solution at every time-step,  consisting of the velocity, angle and location of the glider.
* The first element of the array `u` is set to contain the initial conditions.
* In the `for`-loop, the function `euler_step()` is called to get the solution at time-step $n+1$.

```python
T = 100.0  # length of the time interval
dt = 0.1  # time-step size
N = int(T / dt) + 1  # number of time steps

# Create array to store the solution at each time step.
u = numpy.empty((N, 4))
# Set the initial conditions.
u[0] = numpy.array([v0, theta0, x0, y0])

# Time integration with Euler's method.
for n in range(N - 1):
    u[n + 1] = euler_step(u[n], rhs_phugoid, dt, CL, CD, g, vt)
```

## Plot the trajectory

In order to plot the path of the glider, we need the location (`x`, `y`) with respect to time. That information is already contained in our NumPy array containing the solution; we just need to pluck it out.

Make sure you understand the indices to `u`, below, and the use of the colon notation. If any of it is confusing, read the Python documentation on [Indexing](https://docs.scipy.org/doc/numpy/reference/arrays.indexing.html).

```python
# Get the glider's position over the time.
x = u[:, 2]
y = u[:, 3]
```

Time to plot the path of the glider and get the distance travelled!

```python
# Plot the path of the glider.
pyplot.figure(figsize=(9.0, 4.0))
pyplot.title('Path of the glider (flight time = {})'.format(T))
pyplot.xlabel('x')
pyplot.ylabel('y')
pyplot.grid()
pyplot.plot(x, y, color='C0', linestyle='-', linewidth=2);
```

## Grid convergence

Let's study the convergence of Euler's method for the phugoid model. In the previous lesson, when we studied the straight-line phugoid under a small perturbation, we looked at convergence by comparing the numerical solution with the exact solution. Unfortunately, most problems don't have an exact solution (that's why we compute in the first place!). But here's a neat thing: we can use numerical solutions computed on different grids to study the convergence of the method, even without an analytical solution.

We need to be careful, though, and make sure that the fine-grid solution is resolving all of the features in the mathematical model. How can we know this? We'll have a look at that in a bit. Let's see how this works first.

You need a sequence of numerical solutions of the same problem, each with a different number of time grid points.

Let's create a list of floats called `dt_values` that contains the time-step size of each grid to be solved on. For each element of `dt_values`, we will compute the solution `u` of the glider model using Euler's method and add it to the list `u_values` (initially empty).
If we want to use five different values of $\Delta t$, we'll have five elements in the list `u_values`, each element being a Numpy array. We'll have a list of Numpy arrays! How meta is that?

Read the code below carefully, and remember: you can get a help panel on any function by entering a question mark followed by the function name. For example, add a new code cell below and type:  `?numpy.empty`.

```python
# Set the list of time-step sizes.
dt_values = [0.1, 0.05, 0.01, 0.005, 0.001]

# Create an empty list that will contain the solution of each grid.
u_values = []

for dt in dt_values:
    N = int(T / dt) + 1  # number of time-steps
    # Create array to store the solution at each time step.
    u = numpy.empty((N, 4))
    # Set the initial conditions.
    u[0] = numpy.array([v0, theta0, x0, y0])
    # Temporal integration using Euler's method.
    for n in range(N - 1):
        u[n + 1] = euler_step(u[n], rhs_phugoid, dt, CL, CD, g, vt)
    # Store the solution for the present time-step size
    u_values.append(u)
```

In [Lesson 2](https://nbviewer.jupyter.org/github/numerical-mooc/numerical-mooc/blob/master/lessons/01_phugoid/01_02_Phugoid_Oscillation.ipynb), we compared our numerical result to an analytical solution, but now we will instead compare numerical results from different grids.

For each solution, we'll compute the difference relative to the finest grid. You will be tempted to call this an _"error"_, but be careful: the solution at the finest grid is _not the exact_ solution, it is just a reference value that we can use to estimate grid convergence.

To calculate the difference between one solution `u_current` and the solution at the finest grid, `u_finest`, we'll use the $L_1$-norm, but any norm will do.

There is a small problem with this, though.  The coarsest grid, where $\Delta t = 0.1$, has 1001 grid points, while the finest grid, with $\Delta t = 0.001$ has 100001 grid points.  How do we know which grid points correspond to the same location in two numerical solutions, in order to compare them?

If we had time grids of 10 and 100 steps, respectively, this would be relatively simple to calculate.  Each element in our 10-step grid would span ten elements in our 100-step grid.

Calculating the _ratio_ of the two grid sizes will tell us how many elements in our fine-grid will span over one element in our coarser grid.

Recall that we can _slice_ a NumPy array and grab a subset of values from it.  The syntax for that is

```Python
my_array[3:8]
```

An additional slicing trick that we can take advantage of is the "slice step size."  We add an additional `:` to the slice range and then specify how many steps to take between elements. For example, this code

```Python
my_array[3:8:2]
```

will return the values of `my_array[3]`, `my_array[5]` and `my_array[7]`

With that, we can write a function to obtain the differences between coarser and finest grids. Here we go ...

```python
def l1_diff(u_coarse, u_fine, dt):
    """
    Returns the difference in the L1-norm between the solution on
    a coarse grid and the solution on a fine grid.

    Parameters
    ----------
    u_coarse : numpy.ndarray
        Solution on the coarse grid as an array of floats.
    u_fine : numpy.ndarray
        Solution on the fine grid as an array of floats.
    dt : float
        Time-step size.

    Returns
    -------
    diff : float
        The difference between the two solutions in the L1-norm
        scaled by the time-step size.
    """
    N_coarse = len(u_coarse)
    N_fine = len(u_fine)
    ratio = math.ceil(N_fine / N_coarse)
    diff = dt * numpy.sum(numpy.abs(u_coarse - u_fine[::ratio]))
    return diff
```

Now that the function has been defined, let's compute the grid differences for each solution, relative to the fine-grid solution. Call the function `l1_diff()` with two solutions, one of which is always the one at the finest grid. Here's a neat Python trick: you can use negative indexing in Python!  If you have an array called `my_array` you access the _first_ element with

`my_array[0]`

But you can also access the _last_ element with

`my_array[-1]`

and the next to last element with

`my_array[-2]`

and so on.

```python
# Create an empty list to store the difference in the solution
# between two consecutive grids.
diff_values = []

for i, dt in enumerate(dt_values[:-1]):
    diff = l1_diff(u_values[i][:, 2], u_values[-1][:, 2], dt)
    diff_values.append(diff)
```

Time to create a plot of the results! We'll create a *log-log* plot with the Matplotlib function [`loglog()`](https://matplotlib.org/api/pyplot_api.html?highlight=loglog#matplotlib.pyplot.loglog). Remember to skip the difference of the finest-grid solution with itself, which is zero.

```python
# Plot the difference versus the time-step size.
pyplot.figure(figsize=(6.0, 6.0))
pyplot.title('L1-norm difference vs. time-step size')  # set the title
pyplot.xlabel('$\Delta t$')  # set the x-axis label
pyplot.ylabel('Difference')  # set the y-axis label
pyplot.grid()
pyplot.loglog(dt_values[:-1], diff_values,
              color='C0', linestyle='--', marker='o')  # log-log plot
pyplot.axis('equal');  # make axes scale equally
```

### Order of convergence

The order of convergence is the rate at which the numerical solution approaches the exact one as the mesh is refined. Considering that we're not comparing with an exact solution, we use 3 grid resolutions that are refined at a constant ratio $r$ to find the *observed order of convergence* ($p$), which is given by:

$$
\begin{equation}
p = \frac{\log \left(\frac{f_3-f_2}{f_2-f_1} \right) }{\log(r)}
\end{equation}
$$

where $f_1$ is the finest mesh solution, and $f_3$ the coarsest.

```python
r = 2  # refinement ratio for the time-step size
h = 0.001  # base grid size

dt_values2 = [h, r * h, r**2 * h]
u_values2 = []

for dt in dt_values2:
    N = int(T / dt) + 1  # number of time steps
    # Create array to store the solution at each time step.
    u = numpy.empty((N, 4))
    # Set initial conditions.
    u[0] = numpy.array([v0, theta0, x0, y0])
    # Time integration using Euler's method.
    for n in range(N - 1):
        u[n + 1] = euler_step(u[n], rhs_phugoid, dt, CL, CD, g, vt)
    # Store the solution.
    u_values2.append(u)

# Calculate f2 - f1.
f2_f1 = l1_diff(u_values2[1][:, 2], u_values2[0][:, 2], dt_values2[1])
# Calculate f3 - f2.
f3_f2 = l1_diff(u_values2[2][:, 2], u_values2[1][:, 2], dt_values2[2])
# Calculate the observed order of convergence.
p = math.log(f3_f2 / f2_f1) / math.log(r)
print('Observed order of convergence: p = {:.3f}'.format(p))
```

    Observed order of convergence: p = 1.014

See how the observed order of convergence is close to 1? This means that the rate at which the grid differences decrease match the mesh-refinement ratio. We say that Euler's method is of *first order*, and this result is a consequence of that.

# Paper airplane challenge

Suppose you wanted to participate in a paper-airplane competition, and you want to use what you know about the phugoid model to improve your chances. For a given value of $L/D$ that you can obtain in your design, you want to know what is the best initial velocity and launch angle to fly the longest distance from a given height.

Using the phugoid model, write a new code to analyze the flight of a paper airplane, with the following conditions:

*  Assume $L/D$ of 5.0 (a value close to measurements in Feng et al. 2009)
*  For the trim velocity, let's take an average value of 4.9 m/s.
*  Find a combination of launch angle and velocity that gives the best distance.
*  Think about how you will know when the flight needs to stop ... this will influence how you organize the code.
*  How can you check if your answer is realistic?

## References

* Feng, N. B. et al. _"On the aerodynamics of paper airplanes"_, AIAA paper 2009-3958, 27th AIAA Applied Aerodynamics Conference, San Antonio, TX. [PDF](http://arc.aiaa.org/doi/abs/10.2514/6.2009-3958)

* Simanca, S. R. and Sutherland, S. _"Mathematical problem-solving with computers,"_ 2002 course notes, Stony Brook University, chapter 3: [The Art of Phugoid](https://www.math.sunysb.edu/~scott/Book331/Art_Phugoid.html). (Note that there is an error in the figure: sine and cosine are switched.)

---

###### The cell below loads the style of the notebook.

```python
from IPython.core.display import HTML
css_file = '../../styles/numericalmoocstyle.css'
HTML(open(css_file, 'r').read())
```

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

@font-face {
    font-family: "Computer Modern";
    src: url('http://mirrors.ctan.org/fonts/cm-unicode/fonts/otf/cmunss.otf');
}

#notebook_panel { /* main background */
    background: rgb(245,245,245);
}

div.cell { /* set cell width */
    width: 750px;
}

div #notebook { /* centre the content */
    background: #fff; /* white background for content */
    width: 1000px;
    margin: auto;
    padding-left: 0em;
}

#notebook li { /* More space between bullet points */
    margin-top:0.8em;
}

/* draw border around running cells */
div.cell.border-box-sizing.code_cell.running {
    border: 1px solid #111;
}

/* Put a solid color box around each cell and its output, visually linking them*/
div.cell.code_cell {
    background-color: rgb(256,256,256);
    border-radius: 0px;
    padding: 0.5em;
    margin-left:1em;
    margin-top: 1em;
}

div.text_cell_render{
    font-family: 'Alegreya Sans' sans-serif;
    line-height: 140%;
    font-size: 125%;
    font-weight: 400;
    width:600px;
    margin-left:auto;
    margin-right:auto;
}

/* Formatting for header cells */
.text_cell_render h1 {
    font-family: 'Nixie One', serif;
    font-style:regular;
    font-weight: 400;
    font-size: 45pt;
    line-height: 100%;
    color: rgb(0,51,102);
    margin-bottom: 0.5em;
    margin-top: 0.5em;
    display: block;
}

.text_cell_render h2 {
    font-family: 'Nixie One', serif;
    font-weight: 400;
    font-size: 30pt;
    line-height: 100%;
    color: rgb(0,51,102);
    margin-bottom: 0.1em;
    margin-top: 0.3em;
    display: block;
}

.text_cell_render h3 {
    font-family: 'Nixie One', serif;
    margin-top:16px;
    font-size: 22pt;
    font-weight: 600;
    margin-bottom: 3px;
    font-style: regular;
    color: rgb(102,102,0);
}

.text_cell_render h4 {    /*Use this for captions*/
    font-family: 'Nixie One', serif;
    font-size: 14pt;
    text-align: center;
    margin-top: 0em;
    margin-bottom: 2em;
    font-style: regular;
}

.text_cell_render h5 {  /*Use this for small titles*/
    font-family: 'Nixie One', sans-serif;
    font-weight: 400;
    font-size: 16pt;
    color: rgb(163,0,0);
    font-style: italic;
    margin-bottom: .1em;
    margin-top: 0.8em;
    display: block;
}

.text_cell_render h6 { /*use this for copyright note*/
    font-family: 'PT Mono', sans-serif;
    font-weight: 300;
    font-size: 9pt;
    line-height: 100%;
    color: grey;
    margin-bottom: 1px;
    margin-top: 1px;
}

.CodeMirror{
    font-family: "Source Code Pro";
    font-size: 90%;
}

.alert-box {
    padding:10px 10px 10px 36px;
    margin:5px;
}

.success {
    color:#666600;
    background:rgb(240,242,229);
}
</style>