Table of Contents

1  Product Quantization for Model Compression

1.1  Data Preparation and Model

1.2  Product Quantization from Scratch

1.2.1  Learning Code Book

1.2.2  Encode

1.3  Computing Query Distance

1.4  Fasttext Product Quantization

1.5  Product Quantization Recap

2  Reference

```python
# code for loading the format for the notebook
import os

# path : store the current path to convert back to it later
path = os.getcwd()
os.chdir(os.path.join('..', '..', 'notebook_format'))

from formats import load_style
load_style(plot_style=False)
```

```python
os.chdir(path)

# 1. magic for inline plot
# 2. magic to print version
# 3. magic so that the notebook will reload external python modules
# 4. magic to enable retina (high resolution) plots
# https://gist.github.com/minrk/3301035
%matplotlib inline
%load_ext watermark
%load_ext autoreload
%autoreload 2
%config InlineBackend.figure_format='retina'

import time
import fasttext
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

# prevent scientific notations
pd.set_option('display.float_format', lambda x: '%.3f' % x)

%watermark -a 'Ethen' -d -t -v -p numpy,pandas,matplotlib,fasttext,scipy
```

    Ethen 2020-05-21 22:08:36

    CPython 3.6.4
    IPython 7.9.0

    numpy 1.16.5
    pandas 0.25.0
    matplotlib 3.1.1
    fasttext n
    scipy 1.4.1

# Product Quantization for Model Compression

**Product Quantization** or often times PQ for short is an extremely popular algorithm for compressing vectors/embeddings and performing approximate nearest neighborhood search. In this example, we'll be using fasttext to illustrate the concept. The library and the data preparation step has already been introduced in another [documentation](http://ethen8181.github.io/machine-learning/deep_learning/multi_label/fasttext.html), hence we won't be spending too much on them.

## Data Preparation and Model

```python
# download the data and un-tar it under the 'data' folder

# -P or --directory-prefix specifies which directory to download the data to
!wget https://dl.fbaipublicfiles.com/fasttext/data/cooking.stackexchange.tar.gz -P data
# -C specifies the target directory to extract an archive to
!tar xvzf data/cooking.stackexchange.tar.gz -C data
```

    --2020-05-21 22:08:36--  https://dl.fbaipublicfiles.com/fasttext/data/cooking.stackexchange.tar.gz
    Resolving dl.fbaipublicfiles.com (dl.fbaipublicfiles.com)... 104.22.75.142, 104.22.74.142
    Connecting to dl.fbaipublicfiles.com (dl.fbaipublicfiles.com)|104.22.75.142|:443... connected.
    HTTP request sent, awaiting response... 200 OK
    Length: 457609 (447K) [application/x-tar]
    Saving to: ‘data/cooking.stackexchange.tar.gz.3’

    cooking.stackexchan 100%[===================>] 446.88K  --.-KB/s    in 0.1s

    2020-05-21 22:08:36 (4.43 MB/s) - ‘data/cooking.stackexchange.tar.gz.3’ saved [457609/457609]

    x cooking.stackexchange.id
    x cooking.stackexchange.txt
    x readme.txt

```python
!head -n 3 data/cooking.stackexchange.txt
```

    __label__sauce __label__cheese How much does potato starch affect a cheese sauce recipe?
    __label__food-safety __label__acidity Dangerous pathogens capable of growing in acidic environments
    __label__cast-iron __label__stove How do I cover up the white spots on my cast iron stove?

```python
# train/test split
from fasttext_module.split import train_test_split_file
from fasttext_module.utils import prepend_file_name

data_dir = 'data'
test_size = 0.2
input_path = os.path.join(data_dir, 'cooking.stackexchange.txt')
input_path_train = prepend_file_name(input_path, 'train')
input_path_test = prepend_file_name(input_path, 'test')
random_state = 1234
encoding = 'utf-8'

train_test_split_file(input_path, input_path_train, input_path_test,
                      test_size, random_state, encoding)
print('train path: ', input_path_train)
print('test path: ', input_path_test)
```

    train path:  data/train_cooking.stackexchange.txt
    test path:  data/test_cooking.stackexchange.txt

```python
# train the fasttext model
fasttext_params = {
    'input': input_path_train,
    'lr': 0.1,
    'lrUpdateRate': 1000,
    'thread': 8,
    'epoch': 15,
    'wordNgrams': 1,
    'dim': 80,
    'loss': 'ova'
}
model = fasttext.train_supervised(**fasttext_params)

print('vocab size: ', len(model.words))
print('label size: ', len(model.labels))
print('example vocab: ', model.words[:5])
print('example label: ', model.labels[:5])

model_checkpoint = os.path.join('model', 'model.fasttext')
model.save_model(model_checkpoint)
```

    vocab size:  14496
    label size:  733
    example vocab:  ['</s>', 'to', 'a', 'How', 'the']
    example label:  ['__label__baking', '__label__food-safety', '__label__substitutions', '__label__equipment', '__label__bread']

```python
# model.get_input_matrix().shape
print('output matrix shape: ', model.get_output_matrix().shape)
model.get_output_matrix()
```

    output matrix shape:  (733, 80)

    array([[ 0.7556045 , -0.0944011 , -0.21562839, ..., -2.333261  ,
            -0.8223824 ,  1.4007478 ],
           [ 0.8890684 , -0.35554123, -2.706871  , ...,  1.2450662 ,
             0.05990526, -2.0757916 ],
           [ 0.40264592, -1.7393613 ,  1.3168088 , ..., -3.8272636 ,
             1.9295563 ,  1.6614803 ],
           ...,
           [ 1.0975604 , -0.8535318 ,  0.24365285, ..., -0.5385103 ,
            -0.5356003 ,  0.6660056 ],
           [ 1.1728309 , -0.92786735,  0.30251437, ..., -0.60272974,
            -0.5144258 ,  0.7182006 ],
           [ 1.1421866 , -0.89937526,  0.29751846, ..., -0.6316411 ,
            -0.570792  ,  0.7465125 ]], dtype=float32)

## Product Quantization from Scratch

The main goal behind compression is that after training our model, the model size is too large for deployment. e.g. when deploying on edge devices where we have limited memory at hand or would require using a bigger machine which ultimately incur more infrastructure costs.

In this situation, **product quantization** is a compression technique for compressing embeddings. It has been shown that it significantly reduce the size of the model without hurting the performance by much. And without further a due, we'll dive right into how this technique works.

### Learning Code Book

Our first step is to split the a $d$ dimensional vectors/matrix $\mathbf{X}$ into $m$ distinct subvectors $u_j$, $1 \leq j \leq m$ of dimension $d^* = d / m$, where $d$ should ideally be multiple of $m$.

As an example, say our original embedding has a dimension of 1000x80, we can split it into 4 subvectors, each having a dimension of 1000x20.

Next, we'll run a k-means clustering algorithm on each subvectors. Upon computing the cluster centroid for each of $k$ cluster and $m$ subvectors, we would have on our hand, $m$ cluster centroids, each of size $k \times d^*$. In other words, reducing the dimension from $n \times d$ into $m \times k \times d^* = k \times d$

So we'll apply k-means on our 1000x20 subvectors. If $k$ is 10, we would end up having 10 cluster centroids each of size 20. Since we have 4 subvectors, we would have 4 of those.

Information theory nomenclature is often times used to explain these concepts. Where:

- The cluster centroids is referred to as **codebook**.
- The cluster index is called a **code**, a **reproduction value**, a **reconstruction value**.

We'll be using [scipy vector quantization module's k-means clustering](https://docs.scipy.org/doc/scipy/reference/cluster.vq.html) implementation in the example code.

```python
from scipy.cluster.vq import kmeans2, vq

def compute_code_books(vectors, sub_size=2, n_cluster=128, n_iter=20, minit='points', seed=123):
    n_rows, n_cols = vectors.shape
    n_sub_cols = n_cols // sub_size

    np.random.seed(seed)
    code_books = np.zeros((sub_size, n_cluster, n_sub_cols), dtype=np.float32)
    for subspace in range(sub_size):
        sub_vectors = vectors[:, subspace * n_sub_cols:(subspace + 1) * n_sub_cols]
        centroid, label = kmeans2(sub_vectors, n_cluster, n_iter, minit=minit)
        code_books[subspace] = centroid

    return code_books
```

```python
sub_size = 2  # m
n_cluster = 64  # k

# learning the cluster centroids / code books for our output matrix/embedding
code_books = compute_code_books(model.get_output_matrix(), sub_size, n_cluster)
print('code book size: ', code_books.shape)
```

    code book size:  (2, 64, 40)

### Encode

The cluster centroids for each of the subvectors represents the average/common pattern of each subvectors, and what we're going to do is to replace the original vector with the cluster centroid of each subvectors. The effect of doing so is instead of storing the original floating values for every record in our dataset, we'll be replacing it with the closest cluster centroid id within each subvectors.

i.e. At the end, we'll be "compressing" our original vector of size $n \times d$ into a size of $n \times m$

```python
def encode(vectors, code_books):
    n_rows, n_cols = vectors.shape
    sub_size = code_books.shape[0]
    n_sub_cols = n_cols // sub_size

    codes = np.zeros((n_rows, sub_size), dtype=np.int32)
    for subspace in range(sub_size):
        sub_vectors = vectors[:, subspace * n_sub_cols:(subspace + 1) * n_sub_cols]
        code, dist = vq(sub_vectors, code_books[subspace])
        codes[:, subspace] = code

    return codes
```

```python
# our original embedding now becomes the cluster centroid for each subspace
vector_codes = encode(model.get_output_matrix(), code_books)
print('encoded vector codes size: ', vector_codes.shape)
vector_codes
```

    encoded vector codes size:  (733, 2)

    array([[34, 61],
           [29, 19],
           [ 7,  4],
           ...,
           [42, 52],
           [42, 47],
           [42, 52]], dtype=int32)

We can calculate the potential size/memory savings if we were to go from storing the original vector into storing the encoded codes and the code book.

```python
(vector_codes.nbytes + code_books.nbytes) / model.get_output_matrix().nbytes
```

    0.11231241473396998

Instead of directly compressing the original embedding, we can also learn the codebooks and compress all new incoming embeddings on the fly.

## Computing Query Distance

We can also compute nearest neighbors using the compressed vectors.

```python
# we'll get one of the labels to find its nearest neighbors
label_id = 0
print(model.labels[label_id])

query = model.get_output_matrix()[label_id]
query.shape
```

    __label__baking

    (80,)

```python
# printing out the shape of the code book to hopefully make it easier
code_books.shape
```

    (2, 64, 40)

To do so, we'll be computing the distance between each subspace of the query with the cluster centroid of each subspace, giving us a $m \times k$ distance table, where each one denotes the squared Euclidean distance between the $m_{th}$ subvector of the query and the $k_{th}$ code/cluster centroid for that $m_{th}$ subvector.

```python
def query_dist_table(query, code_books):
    sub_size, n_cluster, n_sub_cols = code_books.shape

    dist_table = np.zeros((sub_size, n_cluster))
    for subspace in range(sub_size):
        sub_query = query[subspace * n_sub_cols:(subspace + 1) * n_sub_cols]

        diff = code_books[subspace] - sub_query.reshape(1, -1)
        diff = np.sum(diff ** 2, axis=1)
        dist_table[subspace, :] = diff

    return dist_table
```

```python
dist_table = query_dist_table(query, code_books)
print(dist_table.shape)
dist_table[:, :5]
```

    (2, 64)

    array([[ 99.99280548,  67.29737854, 127.93981934, 110.84047699,
            128.75978088],
           [108.44404602, 209.63134766,  87.0813446 , 129.2036438 ,
            158.81433105]])

Then assuming for original vector is already encoded in advance, we can lookup the distances for each cluster centroid and add them up.

```python
# lookup the distance
dists = np.sum(dist_table[range(sub_size), vector_codes], axis=1)
dists[:5]
```

    array([ 42.44870567, 408.77243042, 256.50379944, 247.31576538,
           172.66897583])

```python
# the numpy indexing trick is equivalent to the following loop approach
n_rows = vector_codes.shape[0]
dists = np.zeros(n_rows).astype(np.float32)
for n in range(n_rows):
    for m in range(sub_size):
        dists[n] += dist_table[m][vector_codes[n][m]]

dists[:5]
```

    array([ 42.448708, 408.77243 , 256.50378 , 247.31577 , 172.66898 ],
          dtype=float32)

```python
# find the nearest neighbors and "translate" it to the original labels
k = 5
nearest = np.argsort(dists)[:k]
[model.labels[label] for label in nearest]
```

    ['__label__baking',
     '__label__cake',
     '__label__baking-powder',
     '__label__baking-soda',
     '__label__cookies']

The approach illustrated here is more of a naive approach as it still involves calculating the distances to all the vectors, which can still be inefficient for large $n$ (number of data points). We won't be discussing how to speed up the nearest neighborhood search process for product quantization as this documentation is more focused on the compression aspect of it.

## Fasttext Product Quantization

Fasttext comes with built-in capabilities for doing model compression using product quantization. We'll experiment with different options/parameter and measure the model performance and model size. i.e. compression ratio v.s. model performance dip.

The next couple of code chunks defines the functions to measure the model performance and model file size.

```python
from typing import Dict

def score(input_path_train: str,
          input_path_test: str,
          model: fasttext.FastText._FastText,
          k: int,
          round_digits: int=3) -> Dict[str, float]:

    file_path_dict = {
        'train': input_path_train,
        'test': input_path_test
    }

    result = {}
    for group, file_path in file_path_dict.items():
        num_records, precision_at_k, recall_at_k = model.test(file_path, k)
        f1_at_k = 2 * (precision_at_k * recall_at_k) / (precision_at_k + recall_at_k)
        metric = {
            f'{group}_precision@{k}': round(precision_at_k, round_digits),
            f'{group}_recall@{k}': round(recall_at_k, round_digits),
            f'{group}_f1@{k}': round(f1_at_k, round_digits)
        }
        result.update(metric)

    return result
```

```python
k = 1
result = score(input_path_train, input_path_test, model, k)
result
```

    {'train_precision@1': 0.638,
     'train_recall@1': 0.277,
     'train_f1@1': 0.386,
     'test_precision@1': 0.489,
     'test_recall@1': 0.211,
     'test_f1@1': 0.295}

```python
def compute_file_size(file_path: str) -> str:
    """
    Calculate the file size and format it into a human readable string.

    References
    ----------
    https://stackoverflow.com/questions/2104080/how-can-i-check-file-size-in-python
    """
    file_size = compute_raw_file_size(file_path)
    file_size_str = convert_bytes(file_size)
    return file_size_str

def compute_raw_file_size(file_path: str) -> int:
    """Calculate the file size in bytes."""
    file_info = os.stat(file_path)
    return file_info.st_size

def convert_bytes(num: int) -> str:
    """Convert bytes into more human readable MB, GB, etc."""
    for unit in ['bytes', 'KB', 'MB', 'GB', 'TB']:
        if num < 1024.0:
            return "%3.1f %s" % (num, unit)
        num /= 1024.0
```

```python
compute_file_size(model_checkpoint)
```

    '4.9 MB'

For this part of the experiment, we'll tweak the parameter, dimension of subvector, `dsub`. Remember that this is one of main parameter that controls the tradeoff between the compression ratio and amount of distortion (deviation from the original vector).

```python
dsubs = [-1, 2, 4, 8]

results = []
for dsub in dsubs:
    # ensure we are always loading from the original model,
    # i.e. do not over-ride the model_checkpoint variable
    fasttext_model = fasttext.load_model(model_checkpoint)
    if dsub > 0:
        dir_name = os.path.dirname(model_checkpoint)
        model_path = os.path.join(dir_name, f'model_quantized_dsub{dsub}.fasttext')

        # qnorm, normalized the vector and quantize it
        fasttext_model.quantize(dsub=dsub, qnorm=True)
        fasttext_model.save_model(model_path)
    else:
        model_path = model_checkpoint

    result = score(input_path_train, input_path_test, fasttext_model, k)
    result['dsub'] = dsub
    result['file_size'] = compute_raw_file_size(model_path)
    results.append(result)

df_results = pd.DataFrame.from_dict(results)
df_results
```

      train_precision@1
      train_recall@1
      train_f1@1
      test_precision@1
      test_recall@1
      test_f1@1
      dsub
      file_size

      0
      0.638
      0.277
      0.386
      0.489
      0.211
      0.295
      -1
      5143093

      1
      0.638
      0.277
      0.386
      0.487
      0.210
      0.294
      2
      1181690

      2
      0.630
      0.274
      0.381
      0.483
      0.209
      0.291
      4
      891770

      3
      0.587
      0.255
      0.355
      0.470
      0.203
      0.283
      8
      746810

We can visualize the table results. Our main observation is that setting `dsub` to 2 seems to give the most memory reduction while preserving most of the model's performance.

```python
# change default style figure and font size
plt.rcParams['figure.figsize'] = 15, 6
plt.rcParams['font.size'] = 12

fig, (ax1, ax2) = plt.subplots(1, 2)
fig.suptitle('Fasttext Quantization Experiments')

ax1.plot(df_results['dsub'], df_results['file_size'])
ax1.set_title('dsub versus file size')
ax1.set_xlabel('dsub')
ax1.set_ylabel('file size (bytes)')

ax2.plot(df_results['dsub'], df_results['test_precision@1'])
ax2.set_title('dsub versus test precision@1')
ax2.set_xlabel('dsub')
ax2.set_ylabel('precision@1')

plt.show()
```

The `.quantize` method also provide other options that we did not use here such as:

- Whether to quantize the output matrix, `qout`
- Whether to retrain the model's vector after doing the quantization, `retrain`.

Here we also list down the other compression tidbits from the original paper.

- **Quantize the Norm** fasttext has a parameter `qnorm` to normalize of the vector and also quantize that norm. This often times leads to a lesser drop in accuracy.
- **Retrain after Quantization:** This suggests a bottom-up learning strategy where we first quantize the input matrix, then retrain and quantize the output matrix (the input matrix being frozen).
- **Vocabulary Pruning:** Upon training the model, we can remove features that do not play a large role in the model. For each document, we verify if it is already covered by a retained feature and, if not, we add the feature with the highest norm to our set of retained features. If the number of features is below some user specified threshold, we add the features with the highest norm that have not yet been picked.
- **Choice of subvectors:** We observe in practice that using $k = d/2$ subvectors, i.e., half of the components of the embeddings, works well in practice. Using less subquantizers significantly decreases the performance for a small memory gain.

## Product Quantization Recap

To wrap up, we'll do a recap of product quantization, this time using a bit more notation.

The idea behind product quantization is to compress our original matrix into compact codes, where the comparison of the compact codes approximates the comparison in the original space. In the process of doing so, we are essentially retaining the most useful information within our matrix while discarding the less-relevant ones.

We have an original matrix $\mathbf{X} = [\mathbf{x}^1, \mathbf{x}^2, ..., \mathbf{x}^d]$, where $\mathbf{x}^i \in \mathbb{R}^n$. The input matrix will first be splitted into $m$ distinct sub-matrix, $\mathbf{U}^j$, $1 \leq j \leq m$, each of dimension $d^* = d / m$, where $d$ should ideally be multiple of $m$.

\begin{align}
\underbrace{\mathbf{x}^1, ..., \mathbf{x}^{d^*}}_{\mathbf{U}^1}, ..., \underbrace{\mathbf{x}^{d - d^* + 1}, ..., \mathbf{x}^{d}}_{\mathbf{U}^m}
\end{align}

Then a product quantizer function, $q(\cdot)$, is defined as a concatenation of sub-quantizer.

\begin{align}
q(\mathbf{X}) = \big[q^1(\mathbf{U}^1), q^2(\mathbf{U}^2), ..., q^M(\mathbf{U}^m)\big]
\end{align}

Where each sub-quantizer, $q^j$, is usually a low complexity quantizer/clustering algorithm, such as k-means. In other words, each sub-quantizer learns a sub-codebook/sub-cluster centroids $C^j$ comprises of $k$ centroids each of size $d^*$. Then the quantizer would map an input into its respective code/centroid under each sub-matrix and the final representation would be the concatenation of $m$ centroids.

\begin{align}
\mathbf{c} = [\mathbf{c}^1, \mathbf{c}^2, ..., \mathbf{c}^m] \in C = C^1 \times C^2, ..., C^m
\end{align}

Typically, the number of cluster centroids, $k$, is set to 256 so that each code/centroid can be represented by 8 bits.

# Reference

- [Github: Nano Product Quantization (nanopq)](https://github.com/matsui528/nanopq)
- [Blog: Product Quantizers for k-NN Tutorial Part 1](http://mccormickml.com/2017/10/13/product-quantizer-tutorial-part-1)
- [Paper: A. Joulin, E. Grave, P. Bojanowski, M. Douze, H. Jegou, T. Mikolov - FastText.zip: Compressing text classification models (2016)](https://arxiv.org/abs/1612.03651)
- [Paper: H. Jegou, M. Douze, and C. Schmid. - Product quantization for nearest neighbor search (2011)](https://lear.inrialpes.fr/pubs/2011/JDS11/jegou_searching_with_quantization.pdf)


```python
# This cell is mandatory in all Dymos documentation notebooks.
missing_packages = []
try:
    import openmdao.api as om
except ImportError:
    if 'google.colab' in str(get_ipython()):
        !python -m pip install openmdao[notebooks]
    else:
        missing_packages.append('openmdao')
try:
    import dymos as dm
except ImportError:
    if 'google.colab' in str(get_ipython()):
        !python -m pip install dymos
    else:
        missing_packages.append('dymos')
try:
    import pyoptsparse
except ImportError:
    if 'google.colab' in str(get_ipython()):
        !pip install -q condacolab
        import condacolab
        condacolab.install_miniconda()
        !conda install -c conda-forge pyoptsparse
    else:
        missing_packages.append('pyoptsparse')
if missing_packages:
    raise EnvironmentError('This notebook requires the following packages '
                           'please install them and restart this notebook\'s runtime: {",".join(missing_packages)}')
```

# SSTO Earth Launch

This example is based on the _Time-Optimal Launch of a Titan II_
example given in Appendix B of Longuski {cite}`longuski2014optimal`.
It finds the pitch profile for a single-stage-to-orbit launch vehicle that minimizes the time
required to reach orbit insertion under constant thrust.

The vehicle dynamics are given by

\begin{align}
  \frac{dx}{dt} &= v_x \\
  \frac{dy}{dt} &= v_y \\
  \frac{dv_x}{dt} &= \frac{1}{m} (T \cos \theta - D \cos \gamma) \\
  \frac{dv_y}{dt} &= \frac{1}{m} (T \sin \theta - D \sin \gamma) - g \\
  \frac{dm}{dt} &= \frac{T}{g I_{sp}}
\end{align}

The initial conditions are

\begin{align}
  x_0 &= 0 \\
  y_0 &= 0 \\
  v_{x0} &= 0 \\
  v_{y0} &= 0 \\
  m_0 &= 117000 \rm{\,kg}
\end{align}

and the final conditions are

\begin{align}
  x_f &= \rm{free} \\
  y_f &= 185 \rm{\,km} \\
  v_{xf} &= V_{circ} \\
  v_{yf} &= 0 \\
  m_f &= \rm{free}
\end{align}

## Defining the ODE

Generally, one could define the ODE system as a composite group of multile components.
The atmosphere component computes density ($\rho$).
The eom component computes the state rates.
Decomposing the ODE into smaller calculations makes it easier to derive the analytic derivatives.

However, for this example we will demonstrate the use of complex-step differentiation and define the ODE as a single component.
This saves time up front in the deevlopment of the ODE at a minor cost in execution time.

The unconnected inputs to the EOM at the top of the diagram are provided by the Dymos phase as states, controls, or time values.
The outputs, including the state rates, are shown on the right side of the diagram.
The Dymos phases use state rate values to ensure that the integration technique satisfies the dynamics of the system.

```python
import openmdao.api as om
import numpy as np

class LaunchVehicleODE(om.ExplicitComponent):

    def initialize(self):
        self.options.declare('num_nodes', types=int,
                             desc='Number of nodes to be evaluated in the RHS')

        self.options.declare('g', types=float, default=9.80665,
                             desc='Gravitational acceleration, m/s**2')

        self.options.declare('rho_ref', types=float, default=1.225,
                             desc='Reference atmospheric density, kg/m**3')

        self.options.declare('h_scale', types=float, default=8.44E3,
                             desc='Reference altitude, m')

        self.options.declare('CD', types=float, default=0.5,
                             desc='coefficient of drag')

        self.options.declare('S', types=float, default=7.069,
                             desc='aerodynamic reference area (m**2)')

    def setup(self):
        nn = self.options['num_nodes']

        self.add_input('y',
                       val=np.zeros(nn),
                       desc='altitude',
                       units='m')

        self.add_input('vx',
                       val=np.zeros(nn),
                       desc='x velocity',
                       units='m/s')

        self.add_input('vy',
                       val=np.zeros(nn),
                       desc='y velocity',
                       units='m/s')

        self.add_input('m',
                       val=np.zeros(nn),
                       desc='mass',
                       units='kg')

        self.add_input('theta',
                       val=np.zeros(nn),
                       desc='pitch angle',
                       units='rad')

        self.add_input('thrust',
                       val=2100000 * np.ones(nn),
                       desc='thrust',
                       units='N')

        self.add_input('Isp',
                       val=265.2 * np.ones(nn),
                       desc='specific impulse',
                       units='s')
        # Outputs
        self.add_output('xdot',
                        val=np.zeros(nn),
                        desc='velocity component in x',
                        units='m/s')

        self.add_output('ydot',
                        val=np.zeros(nn),
                        desc='velocity component in y',
                        units='m/s')

        self.add_output('vxdot',
                        val=np.zeros(nn),
                        desc='x acceleration magnitude',
                        units='m/s**2')

        self.add_output('vydot',
                        val=np.zeros(nn),
                        desc='y acceleration magnitude',
                        units='m/s**2')

        self.add_output('mdot',
                        val=np.zeros(nn),
                        desc='mass rate of change',
                        units='kg/s')

        self.add_output('rho',
                        val=np.zeros(nn),
                        desc='density',
                        units='kg/m**3')

        # Setup partials
        # Complex-step derivatives
        self.declare_coloring(wrt='*', method='cs', show_sparsity=True)

    def compute(self, inputs, outputs):

        theta = inputs['theta']
        cos_theta = np.cos(theta)
        sin_theta = np.sin(theta)
        vx = inputs['vx']
        vy = inputs['vy']
        m = inputs['m']
        F_T = inputs['thrust']
        Isp = inputs['Isp']
        y = inputs['y']

        g = self.options['g']
        rho_ref = self.options['rho_ref']
        h_scale = self.options['h_scale']

        CDA = self.options['CD'] * self.options['S']

        outputs['rho'] = rho_ref * np.exp(-y / h_scale)
        outputs['xdot'] = vx
        outputs['ydot'] = vy
        outputs['vxdot'] = (F_T * cos_theta - 0.5 * CDA * outputs['rho'] * vx**2) / m
        outputs['vydot'] = (F_T * sin_theta - 0.5 * CDA * outputs['rho'] * vy**2) / m - g
        outputs['mdot'] = -F_T / (g * Isp)

```

## Solving the problem

```python
import matplotlib.pyplot as plt
import openmdao.api as om
import dymos as dm

#
# Setup and solve the optimal control problem
#
p = om.Problem(model=om.Group())
p.driver = om.pyOptSparseDriver()
p.driver.declare_coloring(tol=1.0E-12)

from dymos.examples.ssto.launch_vehicle_ode import LaunchVehicleODE

#
# Initialize our Trajectory and Phase
#
traj = dm.Trajectory()

phase = dm.Phase(ode_class=LaunchVehicleODE,
                 transcription=dm.GaussLobatto(num_segments=12, order=3, compressed=False))

traj.add_phase('phase0', phase)
p.model.add_subsystem('traj', traj)

#
# Set the options for the variables
#
phase.set_time_options(fix_initial=True, duration_bounds=(10, 500))

phase.add_state('x', fix_initial=True, ref=1.0E5, defect_ref=10000.0,
                rate_source='xdot')
phase.add_state('y', fix_initial=True, ref=1.0E5, defect_ref=10000.0,
                rate_source='ydot')
phase.add_state('vx', fix_initial=True, ref=1.0E3, defect_ref=1000.0,
                rate_source='vxdot')
phase.add_state('vy', fix_initial=True, ref=1.0E3, defect_ref=1000.0,
                rate_source='vydot')
phase.add_state('m', fix_initial=True, ref=1.0E3, defect_ref=100.0,
                rate_source='mdot')

phase.add_control('theta', units='rad', lower=-1.57, upper=1.57, targets=['theta'])
phase.add_parameter('thrust', units='N', opt=False, val=2100000.0, targets=['thrust'])

#
# Set the options for our constraints and objective
#
phase.add_boundary_constraint('y', loc='final', equals=1.85E5, linear=True)
phase.add_boundary_constraint('vx', loc='final', equals=7796.6961)
phase.add_boundary_constraint('vy', loc='final', equals=0)

phase.add_objective('time', loc='final', scaler=0.01)

p.model.linear_solver = om.DirectSolver()

#
# Setup and set initial values
#
p.setup(check=True)

p.set_val('traj.phase0.t_initial', 0.0)
p.set_val('traj.phase0.t_duration', 150.0)
p.set_val('traj.phase0.states:x', phase.interp('x', [0, 1.15E5]))
p.set_val('traj.phase0.states:y', phase.interp('y', [0, 1.85E5]))
p.set_val('traj.phase0.states:vy', phase.interp('vx', [1.0E-6, 0]))
p.set_val('traj.phase0.states:m', phase.interp('vy', [117000, 1163]))
p.set_val('traj.phase0.controls:theta', phase.interp('theta', [1.5, -0.76]))
p.set_val('traj.phase0.parameters:thrust', 2.1, units='MN')

#
# Solve the Problem
#
dm.run_problem(p)

#
# Get the explicitly simulated results
#
exp_out = traj.simulate()

#
# Plot the results
#
fig, axes = plt.subplots(nrows=2, ncols=1, figsize=(10, 8))

axes[0].plot(p.get_val('traj.phase0.timeseries.states:x'),
             p.get_val('traj.phase0.timeseries.states:y'),
             marker='o',
             ms=4,
             linestyle='None',
             label='solution')

axes[0].plot(exp_out.get_val('traj.phase0.timeseries.states:x'),
             exp_out.get_val('traj.phase0.timeseries.states:y'),
             marker=None,
             linestyle='-',
             label='simulation')

axes[0].set_xlabel('range (m)')
axes[0].set_ylabel('altitude (m)')
axes[0].set_aspect('equal')

axes[1].plot(p.get_val('traj.phase0.timeseries.time'),
             p.get_val('traj.phase0.timeseries.controls:theta'),
             marker='o',
             ms=4,
             linestyle='None')

axes[1].plot(exp_out.get_val('traj.phase0.timeseries.time'),
             exp_out.get_val('traj.phase0.timeseries.controls:theta'),
             linestyle='-',
             marker=None)

axes[1].set_xlabel('time (s)')
axes[1].set_ylabel('theta (deg)')

plt.suptitle('Single Stage to Orbit Solution Using Linear Tangent Guidance')
fig.legend(loc='lower center', ncol=2)

plt.show()
```

```python
from openmdao.utils.assert_utils import assert_near_equal

assert_near_equal(p.get_val('traj.phase0.timeseries.time')[-1], 143, tolerance=0.05)
assert_near_equal(p.get_val('traj.phase0.timeseries.states:y')[-1], 1.85E5, 1e-4)
assert_near_equal(p.get_val('traj.phase0.timeseries.states:vx')[-1], 7796.6961, 1e-4)
assert_near_equal(p.get_val('traj.phase0.timeseries.states:vy')[-1], 0, 1e-4)
```

## References

```{bibliography}
:filter: docname in docnames
```


# PHY321: Classical Mechanics 1

**Homework 4, due Monday  February 10**

Date: **Feb 8, 2021**

### Practicalities about  homeworks and projects

1. You can work in groups (optimal groups are often 2-3 people) or by yourself. If you work as a group you can hand in one answer only if you wish. **Remember to write your name(s)**!

2. Homeworks are available Wednesday/Thursday the week before the deadline. The deadline is at the Friday lecture.

3. How do I(we)  hand in?  You can hand in the paper and pencil exercises as a hand-written document. For this homework this applies to exercises 1-5. Alternatively, you can hand in everyhting (if you are ok with typing mathematical formulae using say Latex) as a jupyter notebook at D2L. The numerical exercise(s) (exercise 6 here) should always be handed in as a jupyter notebook by the deadline at D2L.

### Introduction to homework 4

This week's sets of classical pen and paper and computational
exercises deal with simple motion problems and conservation laws; energy, momentum and angular momentum. These conservation laws are central in Physics and understanding them properly lays the foundation for understanding and analyzing more complicated physics problems.
The relevant reading background is
1. chapters 3, 4.1, 4.2 and 4.3 of Taylor (there are many good examples there) and

2. chapters 10-13 of Malthe-Sørenssen.

In both textbooks there are many nice worked out examples. Malthe-Sørenssen's text contains also several coding examples you may find useful.

The numerical homework focuses on another motion problem where you can
use the code you developed in homework 3, almost entirely. Please take
a look at the posted solution (jupyter-notebook) for homework 3. You
need only to change the forces at play. The problem at hand is a
classic, the gravitational force acting between the Sun and the
Earth. Here you will notice also that the standard Euler-integration
algorithm is not the best choice and we will introduce the so-called
Euler-Cromer method and the Velocity-Verlet method. These methods will
give much more stable numerical results with only few additions to
your code.

The code you develop here will also be reused when we analyze energy
conservation in homework set 5. And for those of you doing the honors project, it serves as a starting point for the solar system variant.

### Exercise 1 (10 pt), Conservation laws, Energy and momentum

* 1a (2pt) How do we define a conservative force?

A conservative force is a force whose property is that the total work
done in moving an object between two points is independent of the
taken path. This means that the work on an object under the influence
of a conservative force, is independent on the path of the object. It
depends only on the spatial degrees of freedom and it is possible to
assign a numerical value for the potential at any point. It leads to
conservation of energy. The gravitational force is an example of a
conservative force.

If you wish to read more about conservative forces or not, Feyman's lectures from 1963 are quite interesting.
He states for example that **All fundamental forces in nature appear to be conservative**.
This statement was made while developing his argument that *there are no nonconservative forces*.
You may enjoy the link to [Feynman's lecture](http://www.feynmanlectures.caltech.edu/I_14.html).

An important condition for the final work to be independent of the path is that the **curl** of the force is zero, that
is

$$
\boldsymbol{\nabla} \times \boldsymbol{F}=0
$$

* 1b (4pt) Use the work-energy theorem to show that energy is conserved with a conservative force.

The work-energy theorem states that the work done $W$ by a force $\boldsymbol{F}$ that moves an object from a position $\boldsymbol{r}_0$ to a new position $\boldsymbol{r}_1$

$$
W=\int_{\boldsymbol{r}_0}^{\boldsymbol{r}_1}\boldsymbol{F}\boldsymbol{dr}=\frac{1}{2}mv_1^2-\frac{1}{2}mv_0^2,
$$

where $v_1^2$ is the velocity squared at a time $t_1$ and $v_0^2$ the corresponding quantity at a time $t_0$.
The work done is thus the difference in kinetic energies. We can rewrite the above equation as

$$
\frac{1}{2}mv_1^2=\int_{\boldsymbol{r}_0}^{\boldsymbol{r}_1}\boldsymbol{F}\boldsymbol{dr}+\frac{1}{2}mv_0^2,
$$

that is the final kinetic energy is equal to the initial kinetic energy plus the work done by the force over a given path from a  position $\boldsymbol{r}_0$ at time $t_0$ to a final position position $\boldsymbol{r}_1$ at a later time $t_1$.

* 1c (4pt) Assume that you have only internal two-body forces acting on $N$ objects in an isolated system. The force from object $i$ on object $j$ is $\boldsymbol{f}_{ij}$. Show that the linear momentum is conserved.

Here we use Newton's third law and assume that our system is only
affected by so-called internal forces.  This means that the force
$\boldsymbol{f}_{ij}$ from object $i$ acting on object $j$ is equal to the
force acting on object $j$ from object $i$ but with opposite sign,
that is $\boldsymbol{f}_{ij}=-\boldsymbol{f}_{ji}$.

The total linear momentum is defined as

$$
\boldsymbol{P}=\sum_{i=1}^N\boldsymbol{p}_i=\sum_{i=1}^Nm_i\boldsymbol{v}_i,
$$

where $i$ runs over all objects, $m_i$ is the mass of object $i$ and $\boldsymbol{v}_i$ its corresponding velocity.

The force acting on object $i$ from all the other objects is (lower
case letters for individual objects and upper case letters for total
quantities)

$$
\boldsymbol{f}_i=\sum_{j=1}^N\boldsymbol{f}_{ji}.
$$

Summing over all objects the net force is

$$
\sum_{i=1}^N\boldsymbol{f}_i=\sum_{i=1}^N\sum_{j=1;j\ne i}^N\boldsymbol{f}_{ji}.
$$

We are summing freely over all objects with the constraint that $i\ne j$ (no self-interactions).
We can now manipulate the double sum as

$$
\sum_{i=1}^N\sum_{j=1;j\ne i}^N\boldsymbol{f}_{ji}=\sum_{i=1}^N\sum_{j>i}^N(\boldsymbol{f}_{ji}+\boldsymbol{f}_{ij}).
$$

Convince yourself about this by setting $N=2$ and $N=3$. Nweton's third law says
$\boldsymbol{f}_{ij}=-\boldsymbol{f}_{ji}$, which means we have

$$
\sum_{i=1}^N\sum_{j=1;j\ne i}^N\boldsymbol{f}_{ji}=\sum_{i=1}^N\sum_{j>i}^N(\boldsymbol{f}_{ji}-\boldsymbol{f}_{ji})=0.
$$

The total force due to internal degrees of freedom only is thus $0$.
If we then use the definition that

$$
\sum_{i=1}^N\boldsymbol{f}_i=\sum_{i=1}^Nm_i\frac{d\boldsymbol{v}_i}{dt}=\sum_{i=1}^N\frac{d\boldsymbol{p}_i}{dt}=\frac{d \boldsymbol{P}}{dt}=0,
$$

where we assumed that $m_i$ is independent of time, we see that time derivative of the total momentum is zero.
We say then that the linear momentum is a constant of the motion. It is conserved.

### Exercise 2 (10 pt), Conservation of angular momentum

* 2a (2pt) Define angular momentum and the torque for a single object with external forces only.

The angular moment $\boldsymbol{l}_i$ for a given object $i$ is defined as

$$
\boldsymbol{l}_i = \boldsymbol{r}_i \times \boldsymbol{p}_i,
$$

where $\boldsymbol{p}_i=m_i\boldsymbol{v}_i$. With external forces only defining the acceleration and the mass being time independent, the momentum is the integral over the external force as function of time, that is

$$
\boldsymbol{p}_i(t)=\boldsymbol{p}_i(t_0)+\int_{t_0}^t \boldsymbol{f}_i^{\mathrm{ext}}(t')dt'.
$$

The torque for one object is

$$
\boldsymbol{\tau}_i=\frac{d\boldsymbol{l}_i}{dt} = \frac{dt(\boldsymbol{r}_i \times \boldsymbol{p}_i)}{dt}=\boldsymbol{r}_i \times \frac{d\boldsymbol{p}_i}{dt}=\boldsymbol{r}_i \times \boldsymbol{f}_i,
$$

* 2b (4pt) Define angular momentum and the torque for a system with $N$ objects/particles  with external and internal forces. The force from object $i$ on object $j$ is $\boldsymbol{F}_{ij}$.

The total angular momentum $\boldsymbol{L}$ is defined as

$$
\boldsymbol{L}=\sum_{i=1}^N\boldsymbol{l}_i = \sum_{i=1}^N\boldsymbol{r}_i \times \boldsymbol{p}_i.
$$

and the total torque is (using the expression for one object from 2a)

$$
\boldsymbol{\tau}=\sum_{i=1}^N\frac{d\boldsymbol{l}_i}{dt} = \sum_{i=1}^N\boldsymbol{r}_i \times \boldsymbol{f}_i.
$$

The force acting on one object is $\boldsymbol{f}_i=\boldsymbol{f}_i^{\mathrm{ext}}+\sum_{j=1}^N\boldsymbol{f}_{ji}$.

* 2c (4pt) With internal forces only, what is the mathematical form of the forces that allows for angular momentum to be conserved?

Using the results from 1c, we can rewrite without external forces our torque as

$$
\boldsymbol{\tau}=\sum_{i=1}^N\frac{\boldsymbol{l}_i}{dt} = \sum_{i=1}^N\boldsymbol{r}_i \times \boldsymbol{f}_i=\sum_{i=1}^N(\boldsymbol{r}_i \times \sum_{j=1}^N\boldsymbol{f}_{ji}),
$$

which gives

$$
\boldsymbol{\tau}=\sum_{i=1}^N\sum_{j=1;j\ne i}^N(\boldsymbol{r}_i \times \boldsymbol{f}_{ji}).
$$

We can rewrite this as (convince yourself again about this)

$$
\boldsymbol{\tau}=\sum_{i=1}^N\sum_{j>i}^N(\boldsymbol{r}_i \times \boldsymbol{f}_{ji}+\boldsymbol{r}_j \times \boldsymbol{f}_{ij}),
$$

and using Newton's third law we have

$$
\boldsymbol{\tau}=\sum_{i=1}^N\sum_{j>i}^N(\boldsymbol{r}_i -\boldsymbol{r}_j) \times \boldsymbol{f}_{ji}.
$$

If the force is proportional to $\boldsymbol{r}_i -\boldsymbol{r}_j$ then angular momentum is conserved since the cross-product of a vector with itself is zero. We say thus that angular momentum is a constant of the motion.

### Exsercise 3 (10pt), Example of potential

Consider a particle of mass $m$ moving according to the potential

$$
V(x,y,z)=A\exp\left\{-\frac{x^2+z^2}{2a^2}\right\}.
$$

* 3a (2pt) Is energy conserved? If so, why?

In this exercise $A$ and $a$ are constants. The force is given by the derivative of $V$ with respect to the spatial degrees of freedom and since the potential depends only of these degrees of freedom force is conservative and energy is conserved.

* 3b (4pt) Which of  the quantities, $p_x,p_y,p_z$ are conserved?

Taking the derivatives with respect to time shows that only $p_y$ is conserved.

* 3c (4pt) Which of  the quantities, $L_x,L_y,L_z$ are conserved?

Only $L_y$ is conserved.

### Exercise 4 (10pt), Angular momentum case

At $t=0$ we have a single object with position $\boldsymbol{r}_0=x_0\boldsymbol{e}_x+y_0\boldsymbol{e}_y$. We add also a force in the $x$-direction at $t=0$. We assume that the object is at rest at $t=0$.

$$
\boldsymbol{F} = F\boldsymbol{e}_x.
$$

* 4a (3pt) Find the velocity and momentum at a given time $t$ by integrating over time with the above initial conditions.

There is no velocity in the $x$- and $y$-directions at $t=0$, thus $\boldsymbol{v}_0=0$. The force is constant and acting only in the $x$-direction. We have then (dropping vector symbols and setting $t_0=0$)

$$
v_x(t) = \int_0^t a(t')dt'=\int_0^t\frac{F}{m}dt'=\frac{F}{m}t.
$$

* 4b (3pt) Find also the position at a time $t$.

In the $x$-direction we have then

$$
x(t) = \int_0^t v_x(t')dt'=x_0+\frac{F}{2m}t^2,
$$

resulting in

$$
\boldsymbol{r}(t)=(x_0+\frac{F}{2m}t^2)\boldsymbol{e}_x+y_0\boldsymbol{e}_y.
$$

* 4c (4pt) Use the position and the momentum to find the angular momentum and the torque. Is angular momentum conserved?

Velocity and position are defined in the $xy$-plane only which means that only angular momentum in the $z$-direction is non-zero. The angular momentum is

$$
\boldsymbol{l} = (x(t)v_y(t)-y(t)v_x(t))\boldsymbol{e}_z=-y_0\frac{F}{m}t\boldsymbol{e}_z,
$$

which results in a torque $\boldsymbol{\tau}=-y_0\frac{F}{m}\boldsymbol{e}_z$, which is not zero. Thus, angular momentum is not conserved.

### Exercise 5 (10pt), forces  and potentials

A particle of mass $m$ has velocity $v=\alpha/x$, where $x$ is its displacement.

* 5a (3pt) Find the force $F(x)$ responsible for the motion.

Here, since the force is assumed to be conservative (only dependence on $x$), we can use energy conservation.
Assuming that the total energy at $t=0$ is $E_0$, we have

$$
E_0=V(x)+\frac{1}{2}mv^2=V(x)+\frac{1}{2}m\frac{\alpha^2}{x^2}.
$$

Taking the derivative wrt $x$ we have

$$
\frac{dV}{dx}-m\frac{\alpha^2}{x^3}=0,
$$

and since $F(x)=-dV/dx$ we have

$$
F(x)=-m\frac{\alpha^2}{x^3}.
$$

A particle is thereafter under the influence of a force $F=-kx+kx^3/\alpha^2$, where $k$ and $\alpha$ are constants and $k$ is positive.

* 5b (3pt) Determine $V(x)$  and discuss the motion. It can be convenient here to make a sketch/plot of the potential as function of $x$.

We assume that the potential is zero at say $x=0$. Integrating the force from zero to $x$ gives

$$
V(x) = \int_0^x F(x')dx'=\frac{kx^2}{2}-\frac{kx^4}{4\alpha^2}.
$$

The following code plots the potential. We have chosen values of $\alpha=k=1.0$. Feel free to experiment with other values. We plot $V(x)$ for a domain of $x\in [-2,2]$.

```python
%matplotlib inline

import numpy as np
import matplotlib.pyplot as plt
import math

x0= -2.0
xn = 2.1
Deltax = 0.1
alpha = 1.0
k = 1.0
#set up arrays
x = np.arange(x0,xn,Deltax)
n = np.size(x)
V = np.zeros(n)
V = 0.5*k*x*x-0.25*k*(x**4)/(alpha*alpha)
plt.plot(x, V)
plt.xlabel("x")
plt.ylabel("V")
plt.show()
```

From the plot here (with the chosen parameters)
1. we see that with a given initial velocity we can overcome the potential energy barrier

and leave the potential well for good.
1. If the initial velocity is smaller (see next exercise) than a certain value, it will remain trapped in the potential well and oscillate back and forth around $x=0$. This is where the potential has its minimum value.

2. If the kinetic energy at $x=0$ equals the maximum potential energy, the object will oscillate back and forth between the minimum potential energy at $x=0$ and the turning points where the kinetic energy turns zero. These are the so-called non-equilibrium points.

* 5c (4pt)  What happens when the energy of the particle is $E=(1/4)k\alpha^2$? Hint: what is the maximum value of the potential energy?

From the figure we see that
the potential has a minimum at at $x=0$ then rises until $x=\alpha$ before falling off again. The maximum
potential, $V(x\pm \alpha) = k\alpha^2/4$. If the energy is higher, the particle cannot be contained in the
well. The turning points are thus defined by $x=\pm \alpha$. And from the previous plot you can easily see that this is the case ($\alpha=1$ in the abovementioned Python code).

### Exercise 6 (10pt) and 7 (40pt)

The aim of this exercise (as well as the next) is to study the motion
of objects under the influence of the gravitational force.  We will
limit ourselves to the Earth-Sun system. Here we will scale the
equations and sketch our first algorithm for solving the equations,
namely using Euler's method again, as we did in homework 3.  This part
together with the numerical part forms also the entry point for the
solar system honors project. Furthermore, we will reuse parts of these
results when analyzing energy conservation in homework 5.

We will limit ourselves (in order to test the algorithm) to a
hypothetical solar system with the Earth only orbiting around the sun.
The only force in the problem is gravity. Newton's law of gravitation
is given by a force $F_G$

$$
F_G=\frac{GM_{\odot}M_{\mathrm{Earth}}}{r^2},
$$

where $M_{\odot}$ is the mass of the Sun and $M_{\mathrm{Earth}}$ is
the mass of the Earth. The gravitational constant is $G$ and $r$ is
the distance between the Earth and the Sun.  The Sun
has a mass which is much larger than that of the Earth. We can
therefore safely neglect the motion of the Sun in this problem.

We assume that the orbit of the Earth around the Sun
is co-planar, and we take this to be the $xy$-plane.
Using Newton's second law of motion we get the following equations

$$
\frac{d^2x}{dt^2}=\frac{F_{G,x}}{M_{\mathrm{Earth}}},
$$

and

$$
\frac{d^2y}{dt^2}=\frac{F_{G,y}}{M_{\mathrm{Earth}}},
$$

where $F_{G,x}$ and $F_{G,y}$ are the $x$ and $y$ components of the
gravitational force.

We will use so-called astronomical units when rewriting our equations.
Using astronomical units (AU as abbreviation)it means that one
astronomical unit of length, known as 1 AU, is the average distance
between the Sun and Earth, that is $1$ AU = $1.5\times 10^{11}$ m.  It
can also be convenient to use years instead of seconds since years
match better the time evolution of the solar system. The mass of the
Sun is $M_{\mathrm{sun}}=M_{\odot}=2\times 10^{30}$ kg. The masses of
all relevant planets and their distances from the sun are listed in
the table here in kg and AU.

 Planet	                  Mass in kg                	 Distance to  sun in AU

   Earth      	    $M_{\mathrm{Earth}}=6\times 10^{24}$ kg        	    1AU

   Jupiter    	    $M_{\mathrm{Jupiter}}=1.9\times 10^{27}$ kg    	    5.20 AU

   Mars       	    $M_{\mathrm{Mars}}=6.6\times 10^{23}$ kg       	    1.52 AU

   Venus      	    $M_{\mathrm{Venus}}=4.9\times 10^{24}$ kg      	    0.72 AU

   Saturn     	    $M_{\mathrm{Saturn}}=5.5\times 10^{26}$ kg     	    9.54 AU

   Mercury    	    $M_{\mathrm{Mercury}}=3.3\times 10^{23}$ kg    	    0.39 AU

   Uranus     	    $M_{\mathrm{Uranus}}=8.8\times 10^{25}$ kg     	    19.19 AU

   Neptun     	    $M_{\mathrm{Neptun}}=1.03\times 10^{26}$ kg    	    30.06 AU

   Pluto      	    $M_{\mathrm{Pluto}}=1.31\times 10^{22}$ kg     	    39.53 AU

In setting up the equations we limit ourselves to a co-planar
motion and use only the $x$ and $y$ coordinates. But you should feel
free to extend your equations to three dimensions, it is not very
difficult and the data from NASA are all in three dimensions.

[NASA](http://www.nasa.gov/index.html) has an excellent site at <http://ssd.jpl.nasa.gov/horizons.cgi#top>.
From there you can extract initial conditions in order to start your differential equation solver.
At the above website you need to change from **OBSERVER** to **VECTOR** and then write in the planet you are interested in.
The generated data contain the $x$, $y$ and $z$ values as well as their corresponding velocities. The velocities are in units of AU per day.
Alternatively they can be obtained in terms of km and km/s.

For the system below involving only the Earth and the Sun, you
could just initialize the position with say $x=1$ AU and $y=0$ AU.

We assume that mass units can be obtained by using the fact that Earth's orbit is almost circular around the Sun.

For circular motion we know that the force must obey the following relation

$$
F_G= \frac{M_{\mathrm{Earth}}v^2}{r}=\frac{GM_{\odot}M_{\mathrm{Earth}}}{r^2},
$$

where $v$ is the velocity of Earth.
The latter equation can be used to show that

$$
v^2r=GM_{\odot}=4\pi^2\mathrm{AU}^3/\mathrm{yr}^2.
$$

* 6a (5pt) Show how to derive the last equation and use this to scale the differential equations, getting thus rid of the constant $G$ and the two masses. Split the differential equations for the motion in the $x$ and $y$ directions in terms of four coupled differential equations.

* 6b (5pt)  Discretize the above differential equations and set up an algorithm for solving these equations using Euler's forward algorithm and the so-called velocity Verlet method [discussed in the lecture notes](https://mhjensen.github.io/Physics321/doc/pub/energyconserv/html/energyconserv.html). Here you can reuse what you did in homework 3, exercises 6 and 7.

### Exercise 7 (40pt), Numerical elements, solving exercise 6 numerically

* 7a (20pt)  Write then a program which solves the above differential equations for the Earth-Sun system using Euler's  method and the velocity Verlet method.  Find out which initial value for the velocity that gives a circular orbit and test the stability of your algorithm as function of different time steps $\Delta t$.  Make a plot of the results you obtain for the position of the Earth (plot the $x$ and $y$ values and/or if you prefer to use three dimensions the $z$-value as well) orbiting  the Sun. Discuss eventual differences between the Verlet algorithm and the Euler algorithm.

* 7b (20pt) Consider then a planet which begins at a distance of 1 AU from the sun. Find out by trial and error what the initial velocity must be in order for the planet to escape from the sun.  Can you find an exact answer?  How does that match your numerical results?

### Answers

We start with a simpler case first, the Earth-Sun system  in two dimensions only.  The gravitational force $F_G$ on the earth from the sun is

$$
\boldsymbol{F}_G=-\frac{GM_{\odot}M_E}{r^3}\boldsymbol{r},
$$

where $G$ is the gravitational constant,

$$
M_E=6\times 10^{24}\mathrm{Kg},
$$

the mass of Earth,

$$
M_{\odot}=2\times 10^{30}\mathrm{Kg},
$$

the mass of the Sun and

$$
r=1.5\times 10^{11}\mathrm{m},
$$

is the distance between Earth and the Sun. The latter defines what we call an astronomical unit **AU**.
From Newton's second law we have then for the $x$ direction

$$
\frac{d^2x}{dt^2}=-\frac{F_{x}}{M_E},
$$

and

$$
\frac{d^2y}{dt^2}=-\frac{F_{y}}{M_E},
$$

for the $y$ direction.

Here we will use  that  $x=r\cos{(\theta)}$, $y=r\sin{(\theta)}$ and

$$
r = \sqrt{x^2+y^2}.
$$

We can rewrite

$$
F_{x}=-\frac{GM_{\odot}M_E}{r^2}\cos{(\theta)}=-\frac{GM_{\odot}M_E}{r^3}x,
$$

and

$$
F_{y}=-\frac{GM_{\odot}M_E}{r^2}\sin{(\theta)}=-\frac{GM_{\odot}M_E}{r^3}y,
$$

for the $y$ direction.

We can rewrite these two equations

$$
F_{x}=-\frac{GM_{\odot}M_E}{r^2}\cos{(\theta)}=-\frac{GM_{\odot}M_E}{r^3}x,
$$

and

$$
F_{y}=-\frac{GM_{\odot}M_E}{r^2}\sin{(\theta)}=-\frac{GM_{\odot}M_E}{r^3}y,
$$

as four first-order coupled differential equations

$$
\frac{dv_x}{dt}=-\frac{GM_{\odot}}{r^3}x,
$$

and

$$
\frac{dx}{dt}=v_x,
$$

and

$$
\frac{dv_y}{dt}=-\frac{GM_{\odot}}{r^3}y,
$$

and

$$
\frac{dy}{dt}=v_y.
$$

The four coupled differential equations

$$
\frac{dv_x}{dt}=-\frac{GM_{\odot}}{r^3}x,
$$

and

$$
\frac{dx}{dt}=v_x,
$$

and

$$
\frac{dv_y}{dt}=-\frac{GM_{\odot}}{r^3}y,
$$

and

$$
\frac{dy}{dt}=v_y,
$$

can be turned into dimensionless equations or we can introduce astronomical units with $1$ AU = $1.5\times 10^{11}$.

Using the equations from circular motion (with $r =1\mathrm{AU}$)

$$
\frac{M_E v^2}{r} = F = \frac{GM_{\odot}M_E}{r^2},
$$

we have

$$
GM_{\odot}=v^2r,
$$

and using that the velocity of Earth (assuming circular motion) is
$v = 2\pi r/\mathrm{yr}=2\pi\mathrm{AU}/\mathrm{yr}$, we have

$$
GM_{\odot}= v^2r = 4\pi^2 \frac{(\mathrm{AU})^3}{\mathrm{yr}^2}.
$$

The four coupled differential equations can then be discretized using Euler's method as (with step length $h$)

$$
v_{x,i+1}=v_{x,i}-h\frac{4\pi^2}{r_i^3}x_i,
$$

and

$$
x_{i+1}=x_i+hv_{x,i},
$$

and

$$
v_{y,i+1}=v_{y,i}-h\frac{4\pi^2}{r_i^3}y_i,
$$

and

$$
y_{i+1}=y_i+hv_{y,i},
$$

The code here implements Euler's method for the Earth-Sun system using a more compact way of representing the vectors. Alternatively, you could have spelled out all the variables $v_x$, $v_y$, $x$ and $y$ as one-dimensional arrays.

```python
# Common imports
import numpy as np
import pandas as pd
from math import *
import matplotlib.pyplot as plt
import os

# Where to save the figures and data files
PROJECT_ROOT_DIR = "Results"
FIGURE_ID = "Results/FigureFiles"
DATA_ID = "DataFiles/"

if not os.path.exists(PROJECT_ROOT_DIR):
    os.mkdir(PROJECT_ROOT_DIR)

if not os.path.exists(FIGURE_ID):
    os.makedirs(FIGURE_ID)

if not os.path.exists(DATA_ID):
    os.makedirs(DATA_ID)

def image_path(fig_id):
    return os.path.join(FIGURE_ID, fig_id)

def data_path(dat_id):
    return os.path.join(DATA_ID, dat_id)

def save_fig(fig_id):
    plt.savefig(image_path(fig_id) + ".png", format='png')

DeltaT = 0.01
#set up arrays
tfinal = 10 # in years
n = ceil(tfinal/DeltaT)
# set up arrays for t, a, v, and x
t = np.zeros(n)
v = np.zeros((n,2))
r = np.zeros((n,2))
# Initial conditions as compact 2-dimensional arrays
r0 = np.array([1.0,0.0])
v0 = np.array([0.0,2*pi])
r[0] = r0
v[0] = v0
Fourpi2 = 4*pi*pi
# Start integrating using Euler's method
for i in range(n-1):
    # Set up the acceleration
    # Here you could have defined your own function for this
    rabs = sqrt(sum(r[i]*r[i]))
    a =  -Fourpi2*r[i]/(rabs**3)
    # update velocity, time and position using Euler's forward method
    v[i+1] = v[i] + DeltaT*a
    r[i+1] = r[i] + DeltaT*v[i]
    t[i+1] = t[i] + DeltaT
# Plot position as function of time
fig, ax = plt.subplots()
#ax.set_xlim(0, tfinal)
ax.set_ylabel('x[m]')
ax.set_xlabel('y[m]')
ax.plot(r[:,0], r[:,1])
fig.tight_layout()
save_fig("EarthSunEuler")
plt.show()
```

We notice here that Euler's method doesn't give a stable orbit with for example $\Delta t =0.01$. It
means that we cannot trust Euler's method. Euler's method does not conserve energy. It is an
example of an integrator which is not
[symplectic](https://en.wikipedia.org/wiki/Symplectic_integrator).

Here we present thus two methods, which with simple changes allow us
to avoid these pitfalls. The simplest possible extension is the
so-called Euler-Cromer method.  The changes we need to make to our
code are indeed marginal here.  We need simply to replace

```python
    r[i+1] = r[i] + DeltaT*v[i]
```

in the above code with the velocity at the new time $t_{i+1}$

```python
    r[i+1] = r[i] + DeltaT*v[i+1]
```

By this simple caveat we get stable orbits.  Below we derive the
Euler-Cromer method as well as one of the most utlized algorithms for
solving the above type of problems, the so-called Velocity-Verlet
method.

Let us repeat Euler's method.
We have a differential equation

$$
\begin{equation}
  y'(t_i)=f(t_i,y_i)
\label{_auto1} \tag{1}
\end{equation}
$$

and if we truncate at the first derivative, we have from the Taylor expansion

$$
\begin{equation}
   y_{i+1}=y(t_i) + (\Delta t) f(t_i,y_i) + O(\Delta t^2), \label{eq:euler} \tag{2}
\end{equation}
$$

which when complemented with $t_{i+1}=t_i+\Delta t$ forms
the algorithm for the well-known Euler method.
Note that at every step we make an approximation error
of the order of $O(\Delta t^2)$, however the total error is the sum over all
steps $N=(b-a)/(\Delta t)$ for $t\in [a,b]$, yielding thus a global error which goes like
$NO(\Delta t^2)\approx O(\Delta t)$.

To make Euler's method more precise we can obviously
decrease $\Delta t$ (increase $N$), but this can lead to loss of numerical precision.
Euler's method is not recommended for precision calculation,
although it is handy to use in order to get a first
view on how a solution may look like.

Euler's method is asymmetric in time, since it uses information about the derivative at the beginning
of the time interval. This means that we evaluate the position at $y_1$ using the velocity
at $v_0$. A simple variation is to determine $x_{n+1}$ using the velocity at
$v_{n+1}$, that is (in a slightly more generalized form)

$$
\begin{equation}
   y_{n+1}=y_{n}+ v_{n+1}+O(\Delta t^2)
\label{_auto2} \tag{3}
\end{equation}
$$

and

$$
\begin{equation}
   v_{n+1}=v_{n}+(\Delta t) a_{n}+O(\Delta t^2).
\label{_auto3} \tag{4}
\end{equation}
$$

The acceleration $a_n$ is a function of $a_n(y_n, v_n, t_n)$ and needs to be evaluated
as well. This is the Euler-Cromer method. It is easy to change the above code and see that with the same
time step we get stable results.

Let us stay with $x$ (position) and $v$ (velocity) as the quantities we are interested in.

We have the Taylor expansion for the position given by

$$
x_{i+1} = x_i+(\Delta t)v_i+\frac{(\Delta t)^2}{2}a_i+O((\Delta t)^3).
$$

The corresponding expansion for the velocity is

$$
v_{i+1} = v_i+(\Delta t)a_i+\frac{(\Delta t)^2}{2}v^{(2)}_i+O((\Delta t)^3).
$$

Via Newton's second law we have normally an analytical expression for the derivative of the velocity, namely

$$
a_i= \frac{d^2 x}{dt^2}\vert_{i}=\frac{d v}{dt}\vert_{i}= \frac{F(x_i,v_i,t_i)}{m}.
$$

If we add to this the corresponding expansion for the derivative of the velocity

$$
v^{(1)}_{i+1} = a_{i+1}= a_i+(\Delta t)v^{(2)}_i+O((\Delta t)^2)=a_i+(\Delta t)v^{(2)}_i+O((\Delta t)^2),
$$

and retain only terms up to the second derivative of the velocity since our error goes as $O(h^3)$, we have

$$
(\Delta t)v^{(2)}_i\approx a_{i+1}-a_i.
$$

We can then rewrite the Taylor expansion for the velocity as

$$
v_{i+1} = v_i+\frac{(\Delta t)}{2}\left( a_{i+1}+a_{i}\right)+O((\Delta t)^3).
$$

Our final equations for the position and the velocity become then

$$
x_{i+1} = x_i+(\Delta t)v_i+\frac{(\Delta t)^2}{2}a_{i}+O((\Delta t)^3),
$$

and

$$
v_{i+1} = v_i+\frac{(\Delta t)}{2}\left(a_{i+1}+a_{i}\right)+O((\Delta t)^3).
$$

Note well that the term $a_{i+1}$ depends on the position at $x_{i+1}$. This means that you need to calculate
the position at the updated time $t_{i+1}$ before the computing the next velocity.  Note also that the derivative of the velocity at the time
$t_i$ used in the updating of the position can be reused in the calculation of the velocity update as well.

We can now easily add the Verlet method to our original code as

```python
DeltaT = 0.01
#set up arrays
tfinal = 10
n = ceil(tfinal/DeltaT)
# set up arrays for t, a, v, and x
t = np.zeros(n)
v = np.zeros((n,2))
r = np.zeros((n,2))
# Initial conditions as compact 2-dimensional arrays
r0 = np.array([1.0,0.0])
v0 = np.array([0.0,2*pi])
r[0] = r0
v[0] = v0
Fourpi2 = 4*pi*pi
# Start integrating using the Velocity-Verlet  method
for i in range(n-1):
    # Set up forces, air resistance FD, note now that we need the norm of the vecto
    # Here you could have defined your own function for this
    rabs = sqrt(sum(r[i]*r[i]))
    a =  -Fourpi2*r[i]/(rabs**3)
    # update velocity, time and position using the Velocity-Verlet method
    r[i+1] = r[i] + DeltaT*v[i]+0.5*(DeltaT**2)*a
    rabs = sqrt(sum(r[i+1]*r[i+1]))
    anew = -4*(pi**2)*r[i+1]/(rabs**3)
    v[i+1] = v[i] + 0.5*DeltaT*(a+anew)
    t[i+1] = t[i] + DeltaT
# Plot position as function of time
fig, ax = plt.subplots()
ax.set_ylabel('x[m]')
ax.set_xlabel('y[m]')
ax.plot(r[:,0], r[:,1])
fig.tight_layout()
save_fig("EarthSunVV")
plt.show()
```

You can easily generalize the calculation of the forces by defining a function
which takes in as input the various variables. We leave this as a challenge to you.

Running the above code for various time steps we see that the Velocity-Verlet is fully stable for various time steps.

We can also play around with different initial conditions in order to find the escape velocity from an orbit around the sun with distance one astronomical unit, 1 AU. The theoretical value for the escape velocity, is given by

$$
v = \sqrt{8\pi^2}{r},
$$

and with $r=1$ AU, this means that the escape velocity is $2\pi\sqrt{2}$ AU/yr. To obtain this we required that the kinetic energy of Earth equals the potential energy given by the gravitational force.

Setting

$$
\frac{1}{2}M_{\mathrm{Earth}}v^2=\frac{G\M_{\odot}}{r},
$$

and with $G\M_{\odot}=4\pi^2$ we obtain the above relation for the velocity. Setting an initial velocity say equal to $9$ in the above code, yields a planet (Earth) which escapes a stable orbit around the sun, as seen by running the ode here.

```python
DeltaT = 0.01
#set up arrays
tfinal = 100
n = ceil(tfinal/DeltaT)
# set up arrays for t, a, v, and x
t = np.zeros(n)
v = np.zeros((n,2))
r = np.zeros((n,2))
# Initial conditions as compact 2-dimensional arrays
r0 = np.array([1.0,0.0])
# setting initial velocity larger than escape velocity
v0 = np.array([0.0,9.0])
r[0] = r0
v[0] = v0
Fourpi2 = 4*pi*pi
# Start integrating using the Velocity-Verlet  method
for i in range(n-1):
    # Set up forces, air resistance FD, note now that we need the norm of the vecto
    # Here you could have defined your own function for this
    rabs = sqrt(sum(r[i]*r[i]))
    a =  -Fourpi2*r[i]/(rabs**3)
    # update velocity, time and position using the Velocity-Verlet method
    r[i+1] = r[i] + DeltaT*v[i]+0.5*(DeltaT**2)*a
    rabs = sqrt(sum(r[i+1]*r[i+1]))
    anew = -4*(pi**2)*r[i+1]/(rabs**3)
    v[i+1] = v[i] + 0.5*DeltaT*(a+anew)
    t[i+1] = t[i] + DeltaT
# Plot position as function of time
fig, ax = plt.subplots()
ax.set_ylabel('x[m]')
ax.set_xlabel('y[m]')
ax.plot(r[:,0], r[:,1])
fig.tight_layout()
save_fig("EscapeEarthSunVV")
plt.show()
```