# Series

```python
import pandas as pd
from oeis.sequence import OEIS_Sequence
from matplotlib import pyplot as plt
```

```python
plt.plot(Sequence.terms)
plt.title(Sequence.description)
plt.show()
```

```python
def formula_latex(k, floor=True):
    latex = r"$$\left\lfloor\frac{n^2}{" + str(k) + r"}\right\rfloor$$"
    if not floor:
        latex = latex.replace("floor", "ceil")
    return latex
```

```python
def oeis_md_link(id_):
    return f'[{id_}]({OEIS_URL}{id_})'
```

```python
SEQ_LIST = ['A000290', 'A007590', 'A000212',
            'A002620', 'A118015', 'A056827',
            'A056834', 'A130519', 'A056838',
            'A056865']
```

```python
series_table = pd.DataFrame(columns= ['k',
                             'Secuencia',
                             'Fórmula',
                             'Descripción',
                             'Términos'])
MAX_TERMS = 15

for num, id_ in enumerate(SEQ_LIST):
    Seq = OEIS_Sequence(id_)
    series_table = series_table.append({'k': num + 1,
                                        'Secuencia': oeis_md_link(id_),
                                        'Fórmula': formula_latex(num + 1),
                                        'Descripción': Seq.description,
                                        'Términos': Seq.terms[:MAX_TERMS]
                                       }, ignore_index=True)
```

```python
series_table
```

      k
      Secuencia
      Fórmula
      Descripción
      Términos

      0
      1
      [A000290](http://oeis.org/A000290)
      $$\left\lfloor\frac{n^2}{1}\right\rfloor$$
      The squares: a(n) = n^2.
      [0, 1, 4, 9, 16, 25, 36, 49, 64, 81, 100, 121,...

      1
      2
      [A007590](http://oeis.org/A007590)
      $$\left\lfloor\frac{n^2}{2}\right\rfloor$$
      a(n) = floor(n^2/2).
      [0, 0, 2, 4, 8, 12, 18, 24, 32, 40, 50, 60, 72...

      2
      3
      [A000212](http://oeis.org/A000212)
      $$\left\lfloor\frac{n^2}{3}\right\rfloor$$
      a(n) = floor(n^2/3).
      [0, 0, 1, 3, 5, 8, 12, 16, 21, 27, 33, 40, 48,...

      3
      4
      [A002620](http://oeis.org/A002620)
      $$\left\lfloor\frac{n^2}{4}\right\rfloor$$
      Quarter-squares: floor(n/2)*ceiling(n/2). Equi...
      [0, 0, 1, 2, 4, 6, 9, 12, 16, 20, 25, 30, 36, ...

      4
      5
      [A118015](http://oeis.org/A118015)
      $$\left\lfloor\frac{n^2}{5}\right\rfloor$$
      a(n) = floor(n^2/5).
      [0, 0, 0, 1, 3, 5, 7, 9, 12, 16, 20, 24, 28, 3...

      5
      6
      [A056827](http://oeis.org/A056827)
      $$\left\lfloor\frac{n^2}{6}\right\rfloor$$
      a(n) = floor(n^2/6).
      [0, 0, 0, 1, 2, 4, 6, 8, 10, 13, 16, 20, 24, 2...

      6
      7
      [A056834](http://oeis.org/A056834)
      $$\left\lfloor\frac{n^2}{7}\right\rfloor$$
      a(n) = floor(n^2/7).
      [0, 0, 0, 1, 2, 3, 5, 7, 9, 11, 14, 17, 20, 24...

      7
      8
      [A130519](http://oeis.org/A130519)
      $$\left\lfloor\frac{n^2}{8}\right\rfloor$$
      a(n) = Sum_{k=0..n} floor(k/4). (Partial sums ...
      [0, 0, 0, 0, 1, 2, 3, 4, 6, 8, 10, 12, 15, 18,...

      8
      9
      [A056838](http://oeis.org/A056838)
      $$\left\lfloor\frac{n^2}{9}\right\rfloor$$
      a(n) = floor(n^2/9).
      [0, 0, 0, 1, 1, 2, 4, 5, 7, 9, 11, 13, 16, 18,...

      9
      10
      [A056865](http://oeis.org/A056865)
      $$\left\lfloor\frac{n^2}{10}\right\rfloor$$
      a(n) = floor(n^2/10).
      [0, 0, 0, 0, 1, 2, 3, 4, 6, 8, 10, 12, 14, 16,...

```python
# Tabla en markdown para incluir en el capítulo
print(series_table.to_markdown(index=False))
```

    |   k | Secuencia                          | Fórmula                                     | Descripción                                                           | Términos                                                      |
    |----:|:-----------------------------------|:--------------------------------------------|:----------------------------------------------------------------------|:--------------------------------------------------------------|
    |   1 | [A000290](http://oeis.org/A000290) | $$\left\lfloor\frac{n^2}{1}\right\rfloor$$  | The squares: a(n) = n^2.                                              | [0, 1, 4, 9, 16, 25, 36, 49, 64, 81, 100, 121, 144, 169, 196] |
    |   2 | [A007590](http://oeis.org/A007590) | $$\left\lfloor\frac{n^2}{2}\right\rfloor$$  | a(n) = floor(n^2/2).                                                  | [0, 0, 2, 4, 8, 12, 18, 24, 32, 40, 50, 60, 72, 84, 98]       |
    |   3 | [A000212](http://oeis.org/A000212) | $$\left\lfloor\frac{n^2}{3}\right\rfloor$$  | a(n) = floor(n^2/3).                                                  | [0, 0, 1, 3, 5, 8, 12, 16, 21, 27, 33, 40, 48, 56, 65]        |
    |   4 | [A002620](http://oeis.org/A002620) | $$\left\lfloor\frac{n^2}{4}\right\rfloor$$  | Quarter-squares: floor(n/2)*ceiling(n/2). Equivalently, floor(n^2/4). | [0, 0, 1, 2, 4, 6, 9, 12, 16, 20, 25, 30, 36, 42, 49]         |
    |   5 | [A118015](http://oeis.org/A118015) | $$\left\lfloor\frac{n^2}{5}\right\rfloor$$  | a(n) = floor(n^2/5).                                                  | [0, 0, 0, 1, 3, 5, 7, 9, 12, 16, 20, 24, 28, 33, 39]          |
    |   6 | [A056827](http://oeis.org/A056827) | $$\left\lfloor\frac{n^2}{6}\right\rfloor$$  | a(n) = floor(n^2/6).                                                  | [0, 0, 0, 1, 2, 4, 6, 8, 10, 13, 16, 20, 24, 28, 32]          |
    |   7 | [A056834](http://oeis.org/A056834) | $$\left\lfloor\frac{n^2}{7}\right\rfloor$$  | a(n) = floor(n^2/7).                                                  | [0, 0, 0, 1, 2, 3, 5, 7, 9, 11, 14, 17, 20, 24, 28]           |
    |   8 | [A130519](http://oeis.org/A130519) | $$\left\lfloor\frac{n^2}{8}\right\rfloor$$  | a(n) = Sum_{k=0..n} floor(k/4). (Partial sums of A002265.)            | [0, 0, 0, 0, 1, 2, 3, 4, 6, 8, 10, 12, 15, 18, 21]            |
    |   9 | [A056838](http://oeis.org/A056838) | $$\left\lfloor\frac{n^2}{9}\right\rfloor$$  | a(n) = floor(n^2/9).                                                  | [0, 0, 0, 1, 1, 2, 4, 5, 7, 9, 11, 13, 16, 18, 21]            |
    |  10 | [A056865](http://oeis.org/A056865) | $$\left\lfloor\frac{n^2}{10}\right\rfloor$$ | a(n) = floor(n^2/10).                                                 | [0, 0, 0, 0, 1, 2, 3, 4, 6, 8, 10, 12, 14, 16, 19]            |

```python
for k, seq in enumerate(SEQ_LIST):
    lst = SEQ_LIST.copy()
    del lst[k]
    print(k, seq, lst)
```

    0 A000290 ['A007590', 'A000212', 'A002620', 'A118015', 'A056827', 'A056834', 'A130519', 'A056838', 'A056865']
    1 A007590 ['A000290', 'A000212', 'A002620', 'A118015', 'A056827', 'A056834', 'A130519', 'A056838', 'A056865']
    2 A000212 ['A000290', 'A007590', 'A002620', 'A118015', 'A056827', 'A056834', 'A130519', 'A056838', 'A056865']
    3 A002620 ['A000290', 'A007590', 'A000212', 'A118015', 'A056827', 'A056834', 'A130519', 'A056838', 'A056865']
    4 A118015 ['A000290', 'A007590', 'A000212', 'A002620', 'A056827', 'A056834', 'A130519', 'A056838', 'A056865']
    5 A056827 ['A000290', 'A007590', 'A000212', 'A002620', 'A118015', 'A056834', 'A130519', 'A056838', 'A056865']
    6 A056834 ['A000290', 'A007590', 'A000212', 'A002620', 'A118015', 'A056827', 'A130519', 'A056838', 'A056865']
    7 A130519 ['A000290', 'A007590', 'A000212', 'A002620', 'A118015', 'A056827', 'A056834', 'A056838', 'A056865']
    8 A056838 ['A000290', 'A007590', 'A000212', 'A002620', 'A118015', 'A056827', 'A056834', 'A130519', 'A056865']
    9 A056865 ['A000290', 'A007590', 'A000212', 'A002620', 'A118015', 'A056827', 'A056834', 'A130519', 'A056838']

```python
V = ['A000290', 'A007590', 'A000212', 'A002620', 'A118015', 'A056827', 'A056834', 'A130519', 'A056865']
```

```python
txt = ""

for x in V:
    txt = txt + ", " + x
txt
```

    ', A000290, A007590, A000212, A002620, A118015, A056827, A056834, A130519, A056865'

```python
from math import floor

def f(n, k):
    return floor((n ** 2) / k)
```

```python
lst = []
acc = 0
for n in range(20):
    acc += f(n, 2)
    lst.append(acc)
```

```python
print(lst)
```

    [0, 0, 2, 6, 14, 26, 44, 68, 100, 140, 190, 250, 322, 406, 504, 616, 744, 888, 1050, 1230]

```python
lista = []

for n in range(20):
    lista.append(floor((2 * (n ** 3) + 3 * (n ** 2) - 2 * n)/12))
```

```python
print(lista)
```

    [0, 0, 2, 6, 14, 26, 44, 68, 100, 140, 190, 250, 322, 406, 504, 616, 744, 888, 1050, 1230]

```python
from sympy import Sum, symbols, simplify
```

```python
i, k, n = symbols('i k n', integer=True)
simplify(Sum((i ** 2) / k , (i, 1, n)).doit())
```

$\displaystyle \frac{n \left(2 n^{2} + 3 n + 1\right)}{6 k}$

```python

```


$\newcommand{\mb}[1]{\mathbf{ #1 }}$
$\newcommand{\bb}[1]{\mathbb{ #1 }}$
$\newcommand{\bs}[1]{\boldsymbol{ #1 }}$
$\newcommand{\norm}[1]{\left\Vert #1 \right\Vert}$

$\newcommand{\der}[2]{\frac{\mathrm{d} #1 }{\mathrm{d} #2 }}$
$\newcommand{\derp}[2]{\frac{\partial #1 }{\partial #2 }}$

$\newcommand{\R}{\bb{R}}$

# Double Inverted Pendulum

Consider a frictionless two-link pendulum. Torque $\tau_1$ is applied at a fixed base, where the first link is attached. Torque $\tau_2$ is applied at the opposite end of the first link, where the second link is attached.

The first and second links are modeled as massless, with point masses $m_1$ and $m_2$ at lengths of $l_1$ and $l_2$ from the torque applications, respectively.

Assume the angle of the first link from upright $\theta_1$ and the angle of the second link from the first link $\theta_2$ are measured, along with the corresponding angular rates $\dot{\theta}_1$ and $\dot{\theta}_2$. Let $\mb{q} = (\theta_1, \theta_2)$ and $\dot{\mb{q}} = (\dot{\theta}_1, \dot{\theta}_2)$.

## Dynamics

Let $\mathcal{I} = (\mb{i}_1, \mb{i}_2, \mb{i}_3)$ be an inertial reference frame and let $\mathcal{A} = (\mb{a}_1, \mb{a}_2, \mb{a}_3)$ and $\mathcal{B} = (\mb{b}_1, \mb{b}_2, \mb{b}_3)$ be body reference frames such that $\mb{i}_3 = \mb{a}_3 = \mb{b}_3$. Let $\mb{e}_1, \mb{e}_2, \mb{e}_3 \in \R^3$ be elementary basis vectors.

Let $\mb{f}_1^\mathcal{I} = -m_1 g \mb{e}_1$ and $\mb{f}_2^{\mathcal{I}} = -m_2 g \mb{e}_1$ denote graviational forces on the first and second link, respectively, both resolved in the inertial reference frame.

Let $\mb{r}_{10}^\mathcal{A} = l_1 \mb{e}_1$ denote the position of the joint relative to the base resolved in the $\mathcal{A}$ frame. Let $\mb{r}_{21}^\mathcal{B} = l_2 \mb{e}_1$ denote the position of the end of the pendulum relative to the joint resolved in the $\mathcal{B}$ frame.

Let $\mb{R}: \R \to \R^{3 \times 3}$ be a rotation matrix defined as:

\begin{equation}
    \mb{R}(\theta) = \begin{bmatrix} \cos{\theta} & -\sin{\theta} & 0 \\ \sin{\theta} & \cos{\theta} & 0 \\ 0 & 0 & 1 \end{bmatrix},
\end{equation}

for all $\theta \in \R$.

The direction cosine matrix (DCM) relating $\mathcal{A}$ frame coordinates to inertial frame coordinates is $\mb{C}^{\mathcal{A}\mathcal{I}} = \mb{R}(\theta_1)^\top$, while the DCM relating $\mathcal{B}$ frame coordinates to $\mathcal{A}$ frame coordinates is $\mb{C}^{\mathcal{B}\mathcal{A}} = \mb{R}(\theta_2)^\top$.

Since $\mb{R}(\theta)$ is orthogonal for all $\theta \in \R$, $\mb{C}^{\mathcal{I}\mathcal{A}} = \mb{R}(\theta_1)$ and $\mb{C}^{\mathcal{A}\mathcal{B}} = \mb{R}(\theta_2)$.

Additionally, the DCM relating $\mathcal{B}$ frame coordinates to $\mathcal{I}$ frame coordinates is $\mb{C}^{\mathcal{B}\mathcal{I}} = \mb{C}^{\mathcal{B}\mathcal{A}}\mb{C}^{\mathcal{A}\mathcal{I}}$, and $\mb{C}^{\mathcal{I}\mathcal{B}} = \mb{C}^{\mathcal{I}\mathcal{A}}\mb{C}^{\mathcal{A}\mathcal{B}} = \mb{R}(\theta_1)\mb{R}(\theta_2)$.

The position of the joint relative to the base resolved in the inertial frame is:

\begin{align}
    \mb{r}_{10}^\mathcal{I} &= \mb{C}^{\mathcal{I}\mathcal{A}}\mb{r}_{10}^{\mathcal{A}}\\
    &= l_1 \mb{R}(\theta_1)\mb{e}_1\\
    &= \begin{bmatrix} l_1 \cos{\theta} \\ l_1 \sin{\theta} \\ 0 \end{bmatrix},
\end{align}

while the position of the end of the pendulum relative to the base resolved in the inertial frame is:

\begin{align}
    \mb{r}_{20}^\mathcal{I} &= \mb{r}_{21}^\mathcal{I} + \mb{r}_{10}^\mathcal{I}\\
    &= \mb{C}^{\mathcal{IB}}\mb{r}_{21}^\mathcal{B} + \mb{C}^{\mathcal{I}\mathcal{A}}\mb{r}_{10}^\mathcal{A}\\
    &= l_2 \mb{R}(\theta_1)\mb{R}(\theta_2) \mb{e}_1 + l_1 \mb{R}(\theta_1) \mb{e}_1\\
    &= \mb{R}(\theta_1) \begin{bmatrix} l_2 \cos{\theta_2} \\ l_2 \sin{\theta_2} \\ 0 \end{bmatrix} + \begin{bmatrix} l_1 \cos{\theta_1} \\ l_1 \sin{\theta_1} \\ 0 \end{bmatrix}\\
    &= \begin{bmatrix} l_2 \cos{\theta_1} \cos{\theta_2} - l_2 \sin{\theta_1} \sin{\theta_2} + l_1 \cos{\theta_1} \\ l_2 \sin{\theta_1} \cos{\theta_2} + l_2 \cos{\theta_1} \sin{\theta_2} + l_1 \sin{\theta_1} \\ 0 \end{bmatrix}\\
    &= \begin{bmatrix} l_2 \cos{(\theta_1 + \theta_2)} + l_1 \cos{\theta_1} \\ l_2 \sin{(\theta_1 + \theta_2)} + l_1 \sin{\theta_1} \\ 0 \end{bmatrix},
\end{align}

making use of the facts that $\cos{(\theta_1 + \theta_2)} = \cos{\theta_1}\cos{\theta_2} - \sin{\theta_1}\sin{\theta_2}$ and $\sin{(\theta_1 + \theta_2)} = \sin{\theta_1}\cos{\theta_2} + \cos{\theta_1}\sin{\theta_2}$.

The corresponding velocities are:

\begin{equation}
    \dot{\mb{r}}_{10}^\mathcal{I} = \begin{bmatrix} -l_1 \sin{\theta_1} & 0 \\ l_1 \cos{\theta_1} & 0 \\ 0 & 0 \end{bmatrix} \begin{bmatrix} \dot{\theta}_1 \\ \dot{\theta}_2 \end{bmatrix},
\end{equation}

as well as:

\begin{equation}
    \dot{\mb{r}}_{20}^\mathcal{I} = \begin{bmatrix} -l_2 \sin{(\theta_1 + \theta_2)} -l_1 \sin{\theta_1} & -l_2 \sin{(\theta_1 + \theta_2)} \\ l_2 \cos{(\theta_1 + \theta_2)} + l_1 \cos{\theta_1} & l_2 \cos{(\theta_1 + \theta_2)} \\ 0 & 0 \end{bmatrix} \begin{bmatrix} \dot{\theta}_1 \\ \dot{\theta}_2 \end{bmatrix},
\end{equation}

The corresponding kinetic energy terms are:

\begin{equation}
    \frac{1}{2}m_1\norm{ \dot{\mb{r}}_{10}^\mathcal{I} }_2^2 = \frac{1}{2} m_1 \dot{\mb{q}}^\top \begin{bmatrix} l_1^2 & 0 \\ 0 & 0 \end{bmatrix} \dot{\mb{q}},
\end{equation}

as well as:

\begin{align}
    \frac{1}{2}m_2\norm{ \dot{\mb{r}}_{20}^\mathcal{I} }_2^2 &= \frac{1}{2} m_2 \dot{\mb{q}}^\top \begin{bmatrix} l_1^2 + 2 l_1 l_2 (\sin{(\theta_1 + \theta_2)}\sin{\theta_1} + \cos{(\theta_1 + \theta_2)}\cos{\theta_1})  + l_2^2 & l_1 l_2 (\sin{(\theta_1 + \theta_2)\sin{\theta_1}} + \cos{(\theta_1 + \theta_2)}\cos{\theta_1}) + l_2^2 \\ l_1 l_2 (\sin{(\theta_1 + \theta_2)\sin{\theta_1}} + \cos{(\theta_1 + \theta_2)}\cos{\theta_1}) + l_2^2 & l_2^2 \end{bmatrix} \dot{\mb{q}}\\
    &= \frac{1}{2}m_2\dot{\mb{q}}^\top \begin{bmatrix} l_1^2 + 2 l_1 l_2 \cos{\theta_2} + l_2^2 & l_1 l_2 \cos{\theta_2} + l_2^2 \\ l_1 l_2 \cos{\theta_2} + l_2^2 & l_2^2 \end{bmatrix} \dot{\mb{q}},
\end{align}

using the fact that $\cos{\theta_2} = \cos{((\theta_1 + \theta_2) - \theta_1)} = \cos{(\theta_1 + \theta_2)}\cos{\theta_1} + \sin{(\theta_1 + \theta_2)}\sin{\theta_1}$.

The total kinetic energy of the system can now be described as:

\begin{equation}
    T(\mb{q}, \dot{\mb{q}}) = \frac{1}{2} \dot{\mb{q}}^\top \underbrace{ \begin{bmatrix} (m_1 + m_2)l_1^2 + 2 m_2 l_1 l_2 \cos{\theta_2} + m_2 l_2^2 & m_2 l_1 l_2 \cos{\theta_2} + m_2 l_2^2 \\ m_2 l_1 l_2 \cos{\theta_2} + m_2 l_2^2 & m_2 l_2^2 \end{bmatrix} }_{\mb{D}(\mb{q})} \dot{\mb{q}}.
\end{equation}

The Coriolis terms matrix can be expressed in terms of the inertia matrix:

\begin{equation}
    \mb{C}(\mb{q}, \dot{\mb{q}}) = \derp{\mb{D}}{\mb{q}}(\cdot, \dot{\mb{q}}, \cdot) - \frac{1}{2}\left( \derp{\mb{D}}{\mb{q}}(\cdot, \dot{\mb{q}}, \cdot) \right)^\top.
\end{equation}

Here, $\derp{\mb{D}}{\mb{q}} \in \R^{2 \times 2 \times 2}$ is represented as a multilinear operator for all $\mb{q} \in \R^2$, with:

\begin{equation}
    \derp{\mb{D}}{\mb{q}}(\mb{e}_i, \mb{e}_j, \mb{e}_k) = \derp{D_{ij}(\mb{q})}{\theta_k},
\end{equation}

for all $i, j, k \in \{1, 2\}$ and elementary basis vectors $\mb{e}_i, \mb{e}_j, \mb{e}_k \in \R^2$.

We write:

\begin{equation}
    \derp{\mb{D}}{\mb{q}} = \sum_{i=1}^2 \sum_{j=1}^2 \sum_{k=1}^2 \derp{D_{ij}(\mb{q})}{\theta_k} \mb{e}_i \otimes \mb{e}_j \otimes \mb{e}_k,
\end{equation}

and contract along the second argument as:

\begin{align}
    \derp{\mb{D}}{\mb{q}} &= \sum_{i=1}^2 \sum_{j=1}^2 \sum_{k=1}^2 \derp{D_{ij}(\mb{q})}{\theta_k} \dot{\theta}_j \mb{e}_i \otimes \mb{e}_k\\
    &= \sum_{i=1}^2 \sum_{k=1}^2 \left( \sum_{j=1}^2 \derp{D_{ij}(\mb{q})}{\theta_k} \dot{\theta}_j \right) \mb{e}_i \otimes \mb{e}_k\\
    &= \begin{bmatrix} 0 & -2 m_2 l_1 l_2 \dot{\theta}_1 \sin{\theta_2} - m_2 l_1 l_2 \dot{\theta}_2 \sin{\theta_2} \\ 0 & -m_2 l_1 l_2 \dot{\theta}_1 \sin{\theta_2} \end{bmatrix}.
\end{align}

Therefore, the Coriolis terms matrix is:

\begin{equation}
    \mb{C}(\mb{q}, \dot{\mb{q}}) = \begin{bmatrix} 0 & -m_2 l_1 l_2 (2 \dot{\theta}_1 + \dot{\theta}_2) \sin{\theta}_2 \\ \frac{1}{2} m_2 l_1 l_2 (2 \dot{\theta}_1 + \dot{\theta}_2) \sin{\theta}_2 & -\frac{1}{2} m_2 l_1 l_2 \dot{\theta}_1 \sin{\theta}_2 \end{bmatrix}.
\end{equation}

The potential energy is:

\begin{align}
    U(\mb{q}) &= -(\mb{f}_1^\mathcal{I})^\top\mb{r}_{10}^\mathcal{I} - (\mb{f}_2^\mathcal{I})^\top\mb{r}_{20}^\mathcal{I}\\
    &= m_1 g l_1 \cos{\theta_1} + m_2 g (l_2 \cos{(\theta_1 + \theta_2)} + l_1 \cos{\theta_1})\\
    &= (m_1 + m_2) g l_1 \cos{\theta_1} + m_2 g l_2 \cos{(\theta_1 + \theta_2)}.
\end{align}

The gradient of the potential energy is:

\begin{align}
    \mb{G}(\mb{q}) &= \left(\derp{U}{\mb{q}}\right)^\top\\
    &= \begin{bmatrix} -(m_1 + m_2) g l_1 \sin{\theta_1} - m_2 g l_2 \sin{(\theta_1 + \theta_2)} \\ -m_2 g l_2 \sin{(\theta_1 + \theta_2)} \end{bmatrix}.
\end{align}

```python
# %load core/systems/double_inverted_pendulum.py
from numpy import array, cos, identity, sin
from core.dynamics import RoboticDynamics, SystemDynamics

class DoubleInvertedPendulum(SystemDynamics, RoboticDynamics):
    def __init__(self, m_1, m_2, l_1, l_2, g=9.81):
        SystemDynamics.__init__(self, 4, 2)
        RoboticDynamics.__init__(self, identity(2))
        self.params = m_1, m_2, l_1, l_2, g

    def D(self, q):
        m_1, m_2, l_1, l_2, _ = self.params
        _, theta_2 = q
        D_11 = (m_1 + m_2) * (l_1 ** 2) + 2 * m_2 * l_1 * l_2 * cos(theta_2) + m_2 * (l_2 ** 2)
        D_12 = m_2 * l_1 * l_2 * cos(theta_2) + m_2 * (l_2 ** 2)
        D_21 = D_12
        D_22 = m_2 * (l_2 ** 2)
        return array([[D_11, D_12], [D_21, D_22]])

    def C(self, q, q_dot):
        _, m_2, l_1, l_2, _ = self.params
        _, theta_2 = q
        theta_1_dot, theta_2_dot = q_dot
        C_11 = 0
        C_12 = -m_2 * l_1 * l_2 * (2 * theta_1_dot + theta_2_dot) * sin(theta_2)
        C_21 = -C_12 / 2
        C_22 = -m_2 * l_1 * l_2 * theta_1_dot * sin(theta_2) / 2
        return array([[C_11, C_12], [C_21, C_22]])

    def U(self, q):
        m_1, m_2, l_1, l_2, g = self.params
        theta_1, theta_2 = q
        return (m_1 + m_2) * g * l_1 * cos(theta_1) + m_2 * g * l_2 * cos(theta_1 + theta_2)

    def G(self, q):
        m_1, m_2, l_1, l_2, g = self.params
        theta_1, theta_2 = q
        G_1 = -(m_1 + m_2) * g * l_1 * sin(theta_1) - m_2 * g * l_2 * sin(theta_1 + theta_2)
        G_2 = -m_2 * g * l_2 * sin(theta_1 + theta_2)
        return array([G_1, G_2])
```

```python
m_1, m_2 = 0.25, 0.25
l_1, l_2 = 0.5, 0.5
dip = DoubleInvertedPendulum(m_1, m_2, l_1, l_2)
```

## Feedback Linearization

```python
from numpy import identity
from core.controllers import FBLinController, LQRController

Q, R = 10 * identity(4), identity(2)
lqr = LQRController.build(dip, Q, R)
fb_lin = FBLinController(dip, lqr)
```

```python
from numpy import linspace

x_0 = array([1, 0, 0, 0])
ts = linspace(0, 10, 1000 + 1)
fb_lin_data = dip.simulate(x_0, fb_lin, ts)
```

```python
from matplotlib.pyplot import figure, grid, plot, show, subplot, xlabel, ylabel
```

```python
xs, us = fb_lin_data

figure(figsize=(12, 8))

subplot(221)
plot(ts, xs[:, 0], linewidth=3)
grid()
xlabel('Time (sec)', fontsize=16)
ylabel('$\\theta_1$', fontsize=16)

subplot(222)
plot(ts, xs[:, 1], linewidth=3)
grid()
xlabel('Time (sec)', fontsize=16)
ylabel('$\\theta_2$', fontsize=16)

subplot(223)
plot(ts[:-1], us[:, 0], linewidth=3)
grid()
xlabel('Time (sec)', fontsize=16)
ylabel('$\\tau_1$', fontsize=16)

subplot(224)
plot(ts[:-1], us[:, 1], linewidth=3)
grid()
xlabel('Time (sec)', fontsize=16)
ylabel('$\\tau_2$', fontsize=16)

show()
```

## Energy-Based Controller

```python
from core.controllers import EnergyController

K_p = identity(2)
K_d = identity(2)
energy = EnergyController(dip, K_p, K_d)
```

```python
energy_data = dip.simulate(x_0, energy, ts)
```

```python
from matplotlib.pyplot import legend

figure(figsize=(12, 8))

subplot(221)
xs, _ = fb_lin_data
plot(ts, xs[:, 0], linewidth=3, label='FBL')
xs, _ = energy_data
plot(ts, xs[:, 0], linewidth=3, label='Energy')
grid()
legend(fontsize=16)
xlabel('Time (sec)', fontsize=16)
ylabel('$\\theta_1$', fontsize=16)

subplot(222)
xs, _ = fb_lin_data
plot(ts, xs[:, 1], linewidth=3, label='FBL')
xs, _ = energy_data
plot(ts, xs[:, 1], linewidth=3, label='Energy')
grid()
legend(fontsize=16)
xlabel('Time (sec)', fontsize=16)
ylabel('$\\theta_2$', fontsize=16)

subplot(223)
_, us = fb_lin_data
plot(ts[:-1], us[:, 0], linewidth=3, label='FBL')
_, us = energy_data
plot(ts[:-1], us[:, 0], linewidth=3, label='Energy')
grid()
legend(fontsize=16)
xlabel('Time (sec)', fontsize=16)
ylabel('$\\tau_1$', fontsize=16)

subplot(224)
_, us = fb_lin_data
plot(ts[:-1], us[:, 1], linewidth=3, label='FBL')
_, us = energy_data
plot(ts[:-1], us[:, 1], linewidth=3, label='Energy')
grid()
legend(fontsize=16)
xlabel('Time (sec)', fontsize=16)
ylabel('$\\tau_2$', fontsize=16)

show()
```

## Optimization-Based Controller

```python
from core.controllers import QPController

qp = QPController.build_care(dip, Q, R)
qp.add_regularizer(fb_lin, 25)
```

```python
qp_data = dip.simulate(x_0, qp, ts)
```

```python
figure(figsize=(12, 8))

subplot(221)
xs, _ = fb_lin_data
plot(ts, xs[:, 0], linewidth=3, label='FBL')
xs, _ = energy_data
plot(ts, xs[:, 0], linewidth=3, label='Energy')
xs, _ = qp_data
plot(ts, xs[:, 0], linewidth=3, label='QP')
grid()
legend(fontsize=16)
xlabel('Time (sec)', fontsize=16)
ylabel('$\\theta_1$', fontsize=16)

subplot(222)
xs, _ = fb_lin_data
plot(ts, xs[:, 1], linewidth=3, label='FBL')
xs, _ = energy_data
plot(ts, xs[:, 1], linewidth=3, label='Energy')
xs, _ = qp_data
plot(ts, xs[:, 1], linewidth=3, label='QP')
grid()
legend(fontsize=16)
xlabel('Time (sec)', fontsize=16)
ylabel('$\\theta_2$', fontsize=16)

subplot(223)
_, us = fb_lin_data
plot(ts[:-1], us[:, 0], linewidth=3, label='FBL')
_, us = energy_data
plot(ts[:-1], us[:, 0], linewidth=3, label='Energy')
_, us = qp_data
plot(ts[:-1], us[:, 0], linewidth=3, label='QP')
grid()
legend(fontsize=16)
xlabel('Time (sec)', fontsize=16)
ylabel('$\\tau_1$', fontsize=16)

subplot(224)
_, us = fb_lin_data
plot(ts[:-1], us[:, 1], linewidth=3, label='FBL')
_, us = energy_data
plot(ts[:-1], us[:, 1], linewidth=3, label='Energy')
_, us = qp_data
plot(ts[:-1], us[:, 1], linewidth=3, label='QP')
grid()
legend(fontsize=16)
xlabel('Time (sec)', fontsize=16)
ylabel('$\\tau_2$', fontsize=16)

show()
```

# Safe Controller

Consider the task of moving the pendulum from configuration $\mb{q}_0 = (1, 0)$ at rest to $\mb{q}_d = (-1, 0)$ while evading an obstacle above the base of the pendulum. Suppose the position of the obstacle relative to the base of the pendulum is resolved in the inertial frame as $\mb{p}^{\mathcal{I}} = R \mb{e}_1$, with $0 < R \leq l_1 + l_2$. Additionally, suppose we must keep the end at least $\delta R$ m away from the obstacle, for $0 < \delta < 1$.

Consider the function:

\begin{equation}
    h(\mb{q}) = \frac{1}{2} \norm{ \mb{r}_{20}^{\mathcal{I}} - \mb{p}^{\mathcal{I}} }_2^2 - \frac{1}{2} \delta^2 R^2.
\end{equation}

This function is positive for any configuration satisfying $\norm{ \mb{r}_{20}^{\mathcal{I}} - \mb{p}^{\mathcal{I}} }_2 \geq \delta R$. Denote the $0$-superlevel set of this function as $\mathcal{C}$.

The corresponding time derivative is:

\begin{equation}
    \dot{h}(\mb{q}, \dot{\mb{q}}) = ( \mb{r}_{20}^{\mathcal{I}} - \mb{p}^{\mathcal{I}} )^\top \dot{\mb{r}}_{20}^\mathcal{I}.
\end{equation}

The second time derivative is:

\begin{equation}
    \ddot{h}(\mb{q}, \dot{\mb{q}}, \mb{u}) = (\dot{\mb{r}}_{20}^{\mathcal{I}})^\top\dot{\mb{r}}_{20}^\mathcal{I} + (\mb{r}_{20}^\mathcal{I} - \mb{p}^\mathcal{I})^\top \ddot{\mb{r}}_{20}^\mathcal{I}
\end{equation}

Note that $\ddot{\mb{r}}_{20}^\mathcal{I}$ can be expressed as:

\begin{equation}
    \ddot{\mb{r}}_{20}^\mathcal{I} = \begin{bmatrix} -l_2 (\dot{\theta}_1 + \dot{\theta}_2) \cos{(\theta_1 + \theta_2)} - l_1 \dot{\theta}_1 \cos{\theta_1} & -l_2 (\dot{\theta}_1 + \dot{\theta}_2) \cos{(\theta_1 + \theta_2)} \\ -l_2 (\dot{\theta}_1 + \dot{\theta}_2) \sin{(\theta_1 + \theta_2)} - l_1 \dot{\theta}_1 \sin{\theta_1} & -l_2 (\dot{\theta}_1 + \dot{\theta}_2) \sin{(\theta_1 + \theta_2)} \\ 0 & 0 \end{bmatrix} \begin{bmatrix} \dot{\theta}_1 \\ \dot{\theta}_2 \end{bmatrix} + \begin{bmatrix} -l_2 \sin{(\theta_1 + \theta_2)} -l_1 \sin{\theta_1} & -l_2 \sin{(\theta_1 + \theta_2)} \\ l_2 \cos{(\theta_1 + \theta_2)} + l_1 \cos{\theta_1} & l_2 \cos{(\theta_1 + \theta_2)} \\ 0 & 0 \end{bmatrix} \begin{bmatrix} \ddot{\theta}_1 \\ \ddot{\theta}_2 \end{bmatrix},
\end{equation}

that is, $\ddot{\mb{r}}$ is an affine expression in $\ddot{\mb{q}}$. This means $\ddot{h}$ is an affine expression in $\mb{u}$ and can therefore be used in an inequality constraint of an optimization problem.

Pick $\alpha > 0$. Since the pendulum starts from rest, $\dot{\mb{q}}_0 = \mb{0}_2$, and since $\dot{\mb{r}}_{20}^\mathcal{I}$ is linear in $\dot{\mb{q}}$, initially $\dot{h}(\mb{q}_0, \dot{\mb{q}}_0) = 0$. If $h(\mb{q}_0) \geq 0$, that is, the system is initially safe, then:

\begin{equation}
    \dot{h}(\mb{q}_0, \dot{\mb{q}}_0) \geq -\alpha h(\mb{q}_0).
\end{equation}

Picking actions $\mb{u}$ such that:

\begin{align}
    \der{}{t}(\dot{h}(\mb{q}, \dot{\mb{q}}) + \alpha h(\mb{q})) &= \ddot{h}(\mb{q}, \dot{\mb{q}}, \mb{u}) + \alpha \dot{h}(\mb{q}, \dot{\mb{q}})\\
    &\geq -\alpha ( \dot{h}(\mb{q}, \dot{\mb{q}}) + \alpha h(\mb{q}) )\\
    &= -\alpha \dot{h}(\mb{q}, \dot{\mb{q}}) - \alpha^2 h(\mb{q})
\end{align}

for all $\mb{q} \in \mathcal{C}$ and $\dot{\mb{q}} \in \R^2$ ensures the inequality $\dot{h}(\mb{q}, \dot{\mb{q}}) \geq -\alpha h(\mb{q})$ continues to be satisfied by Nagumo's theorem. It it is always possible to pick such actions, then $\dot{h} + \alpha h$ is called a control safety function (CSF).

Applying this reasoning once more, if $\dot{h}(\mb{q}, \dot{\mb{q}}) \geq -\alpha h(\mb{q})$ for all $\mb{q} \in \mathcal{C}$ and $\dot{\mb{q}} \in \R^2$, then $h(\mb{q}) \geq 0$ continues to be satisfied by Nagumo's theorem. This means the actions maintain the safety of system by ensuring $\mb{q} \in \mathcal{C}$, or that $\mathcal{C}$ is a safe set.

Let $\bs{\Lambda} \in \bb{S}_{++}^{2}$ be an action cost, let $\mb{k}: \mathcal{Q} \times \R^2 \to \R^2$ be a baseline controller, let $\rho \geq 0$ be a regularization coefficient. Consider the optimization-based controller solving:

\begin{align}
    \inf_{\substack{\mb{u} \in \R^2 \\ \mb{x} \in \mathcal{Q} \times \R^2}} ~&~ \mb{u}^\top\bs{\Lambda}\mb{u} + \rho \norm{ \mb{u} - \mb{k}(\mb{x}) }_2^2\\
    \mathrm{s.t.} ~&~ \mb{x} = (\mb{q} - \mb{q}_d, \dot{\mb{q}})\\
    &~ \dot{V}(\mb{x}, \mb{u}) \leq -\alpha' V(\mb{x})\\
    &~ \ddot{h}(\mb{q}, \dot{\mb{q}}, \mb{u}) + \alpha \dot{h}(\mb{q}, \dot{\mb{q}}) \geq -\alpha (\dot{h}(\mb{q}, \dot{\mb{q}}) +\alpha h(\mb{q})),
\end{align}

for all $\mb{q} \in \mathcal{Q}$ and $\mb{\dot{\mb{q}}} \in \R^2$. Here, $\mb{x}$ is denoted as a decision variable only for convenience of notation.

The additional constraint may render the optimization problem infeasible. The stability constraint can be relaxed to find safe solutions. Let $\sigma > 0$ be a slack coefficient, and consider the optimization-based controller solving:

\begin{align}
    \inf_{\substack{\mb{u} \in \R^2 \\ \mb{x} \in \mathcal{Q} \times \R^2 \\ \delta \in \R}} ~&~ \mb{u}^\top\bs{\Lambda}\mb{u} + \rho \norm{ \mb{u} - \mb{k}(\mb{x}) }_2^2 + \sigma \delta^2\\
    \mathrm{s.t.} ~&~ \mb{x} = (\mb{q} - \mb{q}_d, \dot{\mb{q}})\\
    &~ \dot{V}(\mb{x}, \mb{u}) \leq -\alpha' V(\mb{x}) + \delta\\
    &~ \ddot{h}(\mb{q}, \dot{\mb{q}}, \mb{u}) + \alpha \dot{h}(\mb{q}, \dot{\mb{q}}) \geq -\alpha (\dot{h}(\mb{q}, \dot{\mb{q}}) +\alpha h(\mb{q})),
\end{align}

```python
from numpy import concatenate, reshape

class SetPointOutput(RoboticDynamics):
    def __init__(self, robotic_dynamics, q_d):
        RoboticDynamics.__init__(self, robotic_dynamics.B)
        self.robotic_dynamics = robotic_dynamics
        self.q_d = q_d

    def eval(self, x, t):
        q, q_dot = reshape(x, (2, -1))
        return concatenate([q - self.q_d, q_dot])

    def drift(self, x, t):
        return self.robotic_dynamics.drift(x, t)

    def act(self, x, t):
        return self.robotic_dynamics.act(x, t)
```

```python
from core.dynamics import AffineQuadCLF
from numpy.linalg import eigvalsh

q_d = array([-1, 0])
set_point = SetPointOutput(dip, q_d)
lyap = AffineQuadCLF.build_care(set_point, Q, R)
slack_coeff = 1e3
alpha = min(eigvalsh(Q)) / max(eigvalsh(lyap.P))
comp = lambda r: alpha * r

lqr = LQRController.build(set_point, Q, R)
fb_lin = FBLinController(set_point, lqr)

qp = QPController(dip, 2)
qp.add_regularizer(fb_lin, 1)
qp.add_stability_constraint(lyap, comp, True, slack_coeff)
```

```python
qp_data = dip.simulate(x_0, qp, ts)
```

```python
xs, us = qp_data

figure(figsize=(12, 8))

subplot(221)
plot(ts, xs[:, 0], linewidth=3)
grid()
xlabel('Time (sec)', fontsize=16)
ylabel('$\\theta_1$', fontsize=16)

subplot(222)
plot(ts, xs[:, 1], linewidth=3)
grid()
xlabel('Time (sec)', fontsize=16)
ylabel('$\\theta_2$', fontsize=16)

subplot(223)
plot(ts[:-1], us[:, 0], linewidth=3)
grid()
xlabel('Time (sec)', fontsize=16)
ylabel('$\\tau_1$', fontsize=16)

subplot(224)
plot(ts[:-1], us[:, 1], linewidth=3)
grid()
xlabel('Time (sec)', fontsize=16)
ylabel('$\\tau_2$', fontsize=16)

show()
```

```python
from numpy import dot
from numpy.linalg import norm
from core.dynamics import AffineDynamics, ScalarDynamics

class SafetyFunction(AffineDynamics, ScalarDynamics):
    def __init__(self, double_inverted_pendulum, height, prop, alpha):
        self.dynamics = double_inverted_pendulum
        self.height = height
        self.prop = prop
        self.alpha = alpha
        self.p = array([self.height, 0, 0])

    def r(self, q):
        _, _, l_1, l_2, _ = self.dynamics.params
        theta_1, theta_2 = q
        r_1 = l_2 * cos(theta_1 + theta_2) + l_1 * cos(theta_1)
        r_2 = l_2 * sin(theta_1 + theta_2) + l_1 * sin(theta_1)
        r_3 = 0
        return array([r_1, r_2, r_3])

    def r_dot(self, q, q_dot):
        return dot(self.r_ddot_act(q), q_dot)

    def r_ddot_drift(self, q, q_dot):
        _, _, l_1, l_2, _ = self.dynamics.params
        theta_1, theta_2 = q
        theta_1_dot, theta_2_dot = q_dot
        drift_1 = -l_2 * ((theta_1_dot + theta_2_dot) ** 2) * cos(theta_1 + theta_2) - l_1 * (theta_1_dot ** 2) * cos(theta_1)
        drift_2 = -l_2 * ((theta_1_dot + theta_2_dot) ** 2) * sin(theta_1 + theta_2) - l_1 * (theta_1_dot ** 2) * sin(theta_1)
        drift_3 = 0
        return array([drift_1, drift_2, drift_3])

    def r_ddot_act(self, q):
        _, _, l_1, l_2, _ = self.dynamics.params
        theta_1, theta_2 = q
        act_1 = array([-l_2 * sin(theta_1 + theta_2) - l_1 * sin(theta_1), -l_2 * sin(theta_1 + theta_2)])
        act_2 = array([l_2 * cos(theta_1 + theta_2) + l_1 * cos(theta_1), l_2 * cos(theta_1 + theta_2)])
        act_3 = array([0, 0])
        return array([act_1, act_2, act_3])

    def h(self, q):
        return (norm(self.r(q) - self.p) ** 2 - (self.prop ** 2) * (self.height ** 2)) / 2

    def h_dot(self, q, q_dot):
        return dot(self.r(q) - self.p, self.r_dot(q, q_dot))

    def h_ddot_drift(self, q, q_dot):
        r_dot = self.r_dot(q, q_dot)
        return dot(r_dot, r_dot) + dot(self.r(q) - self.p, self.r_ddot_drift(q, q_dot))

    def h_ddot_act(self, q):
        return dot(self.r(q) - self.p, self.r_ddot_act(q))

    def eval(self, x, t):
        q, q_dot = reshape(x, (2, -1))
        return self.h_dot(q, q_dot) + self.alpha * self.h(q)

    def drift(self, x, t):
        q, q_dot = reshape(x, (2, -1))
        return self.h_ddot_drift(q, q_dot) + self.alpha * self.h_dot(q, q_dot)

    def act(self, x, t):
        q, q_dot = reshape(x, (2, -1))
        return self.h_ddot_act(q)
```

```python
height = l_1 + l_2
prop = 0.2
alpha = 2.5
safety = SafetyFunction(dip, height, prop, alpha)

alpha_stability = min(eigvalsh(Q)) / max(eigvalsh(lyap.P))
comp_stability = lambda r: alpha_stability * r
comp_safety = lambda r: 5 * r

qp = QPController(dip, 2)
qp.add_regularizer(fb_lin, 1)
qp.add_stability_constraint(lyap, comp_stability, True, slack_coeff)
qp.add_safety_constraint(safety, comp_safety)

qp_data = dip.simulate(x_0, qp, ts)
```

```python
from matplotlib.pyplot import axis

xs, us = qp_data
rs = array([safety.r(q) for q in xs[:, :2]])
r_joints = array([array([l_1 * cos(theta_1), l_1 * sin(theta_1), 0]) for theta_1 in xs[:, 0]])
hs = array([safety.h(q) for q in xs[:, :2]])

thetas = linspace(0, 2 * 3.14, 1000)
unsafe = array([array([prop * height * cos(theta), height + prop * height * sin(theta)]) for theta in thetas])
base = array([0, 0])
theta_1_d, theta_2_d = q_d
target = array([l_2 * sin(theta_1_d + theta_2_d) + l_1 * sin(theta_1_d), l_2 * cos(theta_1_d + theta_2_d) + l_1 * cos(theta_1_d)])

figure(figsize=(12, 4))

subplot(121)
plot(rs[:, 1], rs[:, 0], linewidth=3, label='End')
plot(r_joints[:, 1], r_joints[:, 0], linewidth=3, label='Joint')
plot(unsafe[:, 0], unsafe[:, 1], 'r', linewidth=3)
ylabel('$x_1$ (m)', fontsize=16)
xlabel('$x_2$ (m)', fontsize=16)
grid()
legend(fontsize=16)
axis('equal')
axis([-1, 1, -0.2, 1.2])

subplot(122)
plot(ts, hs, linewidth=3)
xlabel('$t$ (sec)', fontsize=16)
ylabel('$h$', fontsize=16)
grid()
```

```python
from IPython.display import HTML
from matplotlib.animation import FuncAnimation

fig = figure()

ax = fig.add_subplot(111)
ax.fill(unsafe[:, 0], unsafe[:, 1], 'r')
ax.plot(target[0], target[1], 'x', markersize=20)
ax.axis('equal')
ax.grid('on')
ax.set_xlim([-1, 1])
ax.set_ylim([-0.2, 1.2])

link1, = ax.plot([], [], linewidth=3)
link2, = ax.plot([], [], linewidth=3)

def init():
    link1.set_data([], [])
    link2.set_data([], [])
    return link1, link2

def update(idx):
    base = array([0, 0])
    joint = r_joints[idx]
    end = rs[idx]
    link1.set_data([base[1], joint[1]], [base[0], joint[0]])
    link2.set_data([joint[1], end[1]], [joint[0], end[0]])
    return link1, link2

anim = FuncAnimation(fig, update, init_func=init, frames=500, interval=20, blit=True)
```

```python
HTML(anim.to_html5_video())
```


```python
import sympy
import matplotlib.pyplot as plt
p = sympy.Symbol("p")
class supplyDemand:
    def __init__(self,demandEquation,supplyEquation):
        self.demandEquation = demandEquation
        self.supplyEquation = supplyEquation
        self.priceEnd = sympy.solve(demandEquation)[0]
        self.prices = []
        self.demandQ = []
        self.supplyQ = []
        for price in range(0,self.priceEnd+1):
            self.prices += [price]
            self.demandQ += [demandEquation.subs(p,price)]
            self.supplyQ += [supplyEquation.subs(p,price)]
        self.startingQ = sympy.solve(demandEquation-supplyEquation)[0]
        self.startingP = demandEquation.subs(p,self.startingQ)
        self.mode = "Normal"

    def plot(self):
        plt.plot(self.demandQ,self.prices,'k')
        plt.plot(self.supplyQ,self.prices,'k')
        plt.plot(self.startingQ,self.startingP, 'bo')
        plt.xlabel("Supply and Demand Quantity")
        plt.ylabel("Price")
        plt.show()
```

```python
import sympy
import matplotlib.pyplot as plt
p = sympy.Symbol("p")
class supplyDemand:
    def __init__(self,demandEquation,supplyEquation):
        self.demandEquation = demandEquation
        self.supplyEquation = supplyEquation
        self.priceEnd = sympy.solve(demandEquation)[0]
        self.prices = []
        self.demandQ = []
        self.supplyQ = []
        for price in range(0,self.priceEnd+1):
            self.prices += [price]
            self.demandQ += [demandEquation.subs(p,price)]
            self.supplyQ += [supplyEquation.subs(p,price)]
        self.startingQ = sympy.solve(demandEquation-supplyEquation)[0]
        self.startingP = demandEquation.subs(p,self.startingQ)
        self.mode = "Normal"

    def plot(self):
        plt.plot(self.demandQ,self.prices,'k')
        plt.plot(self.supplyQ,self.prices,'k')
        plt.plot(self.startingQ,self.startingP, 'bo')
        plt.xlabel("Supply and Demand Quantity")
        plt.ylabel("Price")
        plt.show()
    def internationalTrade(self,worldPrice):
        self.mode = "International"
        self.worldPrice = worldPrice
        if worldPrice>self.startingP:
            self.export = True
            self.domesticQ = self.demandEquation.subs(p,worldPrice)
            self.internationalQ = self.supplyEquation.subs(p,worldPrice)-self.demandEquation.subs(p,worldPrice)
            self.totalQ = self.domesticQ+self.internationalQ
        else:
            self.export = False
            self.domesticQ = self.supplyEquation.subs(p,worldPrice)
            self.internationalQ = self.demandEquation.subs(p,worldPrice)-self.domesticQ
            self.totalQ = self.domesticQ+self.internationalQ
```

```python
economy1 = supplyDemand(10-p,p)
economy1.plot()
economy1.internationalTrade(6)
print(economy1.startingQ)
print(economy1.export)
print(economy1.domesticQ)
print(economy1.totalQ)
```

```python
import sympy
import matplotlib.pyplot as plt
p = sympy.Symbol("p")
class supplyDemand:
    def __init__(self,demandEquation,supplyEquation):
        self.demandEquation = demandEquation
        self.supplyEquation = supplyEquation
        self.priceEnd = sympy.solve(demandEquation)[0]
        self.prices = []
        self.demandQ = []
        self.supplyQ = []
        for price in range(0,self.priceEnd+1):
            self.prices += [price]
            self.demandQ += [demandEquation.subs(p,price)]
            self.supplyQ += [supplyEquation.subs(p,price)]
        self.startingQ = sympy.solve(demandEquation-supplyEquation)[0]
        self.startingP = demandEquation.subs(p,self.startingQ)
        self.mode = "Normal"
    def internationalTrade(self,worldPrice):
        self.mode = "International"
        self.worldPrice = worldPrice
        if worldPrice>self.startingP:
            self.export = True
            self.domesticQ = self.demandEquation.subs(p,worldPrice)
            self.internationalQ = self.supplyEquation.subs(p,worldPrice)-self.demandEquation.subs(p,worldPrice)
            self.totalQ = self.domesticQ+self.internationalQ
        else:
            self.export = False
            self.domesticQ = self.supplyEquation.subs(p,worldPrice)
            self.internationalQ = self.demandEquation.subs(p,worldPrice)-self.domesticQ
            self.totalQ = self.domesticQ+self.internationalQ

    def plot(self):
        plt.plot(self.demandQ,self.prices,'k')
        plt.plot(self.supplyQ,self.prices,'k')
        plt.plot(self.startingQ,self.startingP, 'bo')
        plt.xlabel("Supply and Demand Quantity")
        plt.ylabel("Price")
        if self.mode=="International":
            plt.plot([self.totalQ,self.priceEnd],[self.worldPrice,self.worldPrice],'k')
            if self.startingP!=self.worldPrice:
                plt.plot(self.totalQ,self.worldPrice, 'bo')
                plt.plot(self.domesticQ,self.worldPrice,'bo')
                plt.plot([0,self.domesticQ],[self.worldPrice,self.worldPrice],label='Domestic Quantity')
                plt.plot([self.domesticQ,self.totalQ],[self.worldPrice,self.worldPrice], label='International Quantity')
                plt.legend()
        plt.show()
economy1 = supplyDemand(10-p,p)
economy1.internationalTrade(6)
economy1.plot()
```

```python
import sympy
import matplotlib.pyplot as plt
p = sympy.Symbol("p")
class supplyDemand:
    def __init__(self,demandEquation,supplyEquation):
        self.demandEquation = demandEquation
        self.supplyEquation = supplyEquation
        self.priceEnd = sympy.solve(demandEquation)[0]
        self.prices = []
        self.demandQ = []
        self.supplyQ = []
        for price in range(0,self.priceEnd+1):
            self.prices += [price]
            self.demandQ += [demandEquation.subs(p,price)]
            self.supplyQ += [supplyEquation.subs(p,price)]
        self.startingQ = sympy.solve(demandEquation-supplyEquation)[0]
        self.startingP = demandEquation.subs(p,self.startingQ)
        self.mode = "Normal"
    def internationalTrade(self,worldPrice):
        self.mode = "International"
        self.worldPrice = worldPrice
        if worldPrice>self.startingP:
            self.export = True
            self.domesticQ = self.demandEquation.subs(p,worldPrice)
            self.internationalQ = self.supplyEquation.subs(p,worldPrice)-self.demandEquation.subs(p,worldPrice)
            self.totalQ = self.domesticQ+self.internationalQ
        else:
            self.export = False
            self.domesticQ = self.supplyEquation.subs(p,worldPrice)
            self.internationalQ = self.demandEquation.subs(p,worldPrice)-self.domesticQ
            self.totalQ = self.domesticQ+self.internationalQ

    def plot(self):
        plt.plot(self.demandQ,self.prices,'k')
        plt.plot(self.supplyQ,self.prices,'k')
        plt.plot(self.startingQ,self.startingP, 'bo')
        plt.xlabel("Supply and Demand Quantity")
        plt.ylabel("Price")
        if self.mode=="International":
            plt.plot([self.totalQ,self.priceEnd],[self.worldPrice,self.worldPrice],'k')
            if self.startingP!=self.worldPrice:
                plt.plot(self.totalQ,self.worldPrice, 'bo')
                plt.plot(self.domesticQ,self.worldPrice,'bo')
                plt.plot([0,self.domesticQ],[self.worldPrice,self.worldPrice],label='Domestic Quantity')
                plt.plot([self.domesticQ,self.totalQ],[self.worldPrice,self.worldPrice], label='International Quantity')
                plt.legend()
                if self.export:
                    print("The country exports during international trade, and the price moves from "+str(self.startingP)+" to "+str(self.worldPrice)+".")
                    print("The domestic producers go from producing "+str(self.startingQ)+" units to "+str(self.totalQ)+" units.")
                    print("The domestic consumers go from consuming "+str(self.startingQ)+" units to "+str(self.domesticQ)+" units.")
                    print("Total domestic producer revenue is "+str(self.totalQ*self.worldPrice))
                else:
                    print("The country imports during international trade, and the price moves from "+str(self.startingP)+" to "+str(self.worldPrice)+".")
                    print("The domestic producers go from producing "+str(self.startingQ)+" units to "+str(self.domesticQ)+" units.")
                    print("The domestic consumers go from consuming "+str(self.startingQ)+" units to "+str(self.totalQ)+" units.")
                    print("Total domestic producer revenue is "+str(self.domesticQ*self.worldPrice))
        plt.show()
economy1 = supplyDemand(10-p,p)
economy1.internationalTrade(6)
economy1.plot()
```

```python
import sympy
import matplotlib.pyplot as plt
import matplotlib.patches as patches
p = sympy.Symbol("p")
class supplyDemand:
    def __init__(self,demandEquation,supplyEquation):
        self.demandEquation = demandEquation
        self.supplyEquation = supplyEquation
        self.priceEnd = sympy.solve(demandEquation)[0]
        self.prices = []
        self.demandQ = []
        self.supplyQ = []
        for price in range(0,self.priceEnd+1):
            self.prices += [price]
            self.demandQ += [demandEquation.subs(p,price)]
            self.supplyQ += [supplyEquation.subs(p,price)]
        self.startingQ = sympy.solve(demandEquation-supplyEquation)[0]
        self.startingP = demandEquation.subs(p,self.startingQ)
        self.mode = "Normal"
    def internationalTrade(self,worldPrice):
        self.mode = "International"
        self.worldPrice = worldPrice
        if worldPrice>self.startingP:
            self.export = True
            self.domesticQ = self.demandEquation.subs(p,worldPrice)
            self.internationalQ = self.supplyEquation.subs(p,worldPrice)-self.demandEquation.subs(p,worldPrice)
            self.totalQ = self.domesticQ+self.internationalQ
        else:
            self.export = False
            self.domesticQ = self.supplyEquation.subs(p,worldPrice)
            self.internationalQ = self.demandEquation.subs(p,worldPrice)-self.domesticQ
            self.totalQ = self.domesticQ+self.internationalQ
        self.domesticQ
    def drawSurplus1(self):
        if self.export:
            triangle1 = patches.Polygon([[self.domesticQ,self.worldPrice],[0,self.worldPrice],[0,self.priceEnd]],True,label="Consumer Surplus")
            triangle2 = patches.Polygon([[0,0],[0,self.startingP],[self.startingQ,self.startingP]],True,label="Original Producer Surplus",color="red")
            rect1 = patches.Polygon([[0,self.startingP],[0,self.worldPrice],[self.domesticQ,self.worldPrice],[self.startingQ,self.startingP]],True,label="Transferred Surplus",color="orange")
            triangle3 = patches.Polygon([[self.domesticQ,self.worldPrice],[self.totalQ,self.worldPrice],[self.startingQ,self.startingP]],True,label="Surplus from International",color="yellow")
            currentAxis = plt.gca()
            currentAxis.add_patch(triangle1)
            currentAxis.add_patch(triangle2)
            currentAxis.add_patch(rect1)
            currentAxis.add_patch(triangle3)
        else:
            triangle1 = patches.Polygon([[self.startingQ,self.startingP],[0,self.startingP],[0,self.priceEnd]],True,label="Original Consumer Surplus", color="red")
            triangle2 = patches.Polygon([[0,0],[0,self.worldPrice],[self.domesticQ,self.worldPrice]],True,label="Producer Surplus")
            rect1 = patches.Polygon([[0,self.startingP],[0,self.worldPrice],[self.domesticQ,self.worldPrice],[self.startingQ,self.startingP]],True,label="Transferred Surplus",color="orange")
            triangle3 = patches.Polygon([[self.domesticQ,self.worldPrice],[self.totalQ,self.worldPrice],[self.startingQ,self.startingP]],True,label="Surplus from International",color="yellow")
            currentAxis = plt.gca()
            currentAxis.add_patch(triangle1)
            currentAxis.add_patch(triangle2)
            currentAxis.add_patch(rect1)
            currentAxis.add_patch(triangle3)

    def plot(self):
        plt.plot(self.demandQ,self.prices,'k')
        plt.plot(self.supplyQ,self.prices,'k')
        plt.plot(self.startingQ,self.startingP, 'bo')
        plt.xlabel("Supply and Demand Quantity")
        plt.ylabel("Price")
        if self.mode=="International":
            plt.plot([0,self.priceEnd],[self.worldPrice,self.worldPrice],"k")
            if self.startingP!=self.worldPrice:
                plt.plot(self.totalQ,self.worldPrice, 'bo')
                plt.plot(self.domesticQ,self.worldPrice,'bo')
                plt.plot([0,self.domesticQ],[self.worldPrice,self.worldPrice],"m",label='Domestic Quantity')
                plt.plot([self.domesticQ,self.totalQ],[self.worldPrice,self.worldPrice],"g", label='International Quantity')

                if self.export:
                    print("The country exports during international trade, and the price moves from "+str(self.startingP)+" to "+str(self.worldPrice)+".")
                    print("The domestic producers go from producing "+str(self.startingQ)+" units to "+str(self.totalQ)+" units.")
                    print("The domestic consumers go from consuming "+str(self.startingQ)+" units to "+str(self.domesticQ)+" units.")
                    print("Total domestic producer revenue is "+str(self.totalQ*self.worldPrice))
                else:
                    print("The country imports during international trade, and the price moves from "+str(self.startingP)+" to "+str(self.worldPrice)+".")
                    print("The domestic producers go from producing "+str(self.startingQ)+" units to "+str(self.domesticQ)+" units.")
                    print("The domestic consumers go from consuming "+str(self.startingQ)+" units to "+str(self.totalQ)+" units.")
                    print("Total domestic producer revenue is "+str(self.domesticQ*self.worldPrice))
                plt.legend()
        plt.show()
```

```python
economy1 = supplyDemand(10-p,p)
economy1.internationalTrade(4)
economy1.drawSurplus1()
economy1.plot()
```

```python
import sympy
import matplotlib.pyplot as plt
import matplotlib.patches as patches
p = sympy.Symbol("p")
class supplyDemand:
    def __init__(self,demandEquation,supplyEquation):
        self.demandEquation = demandEquation
        self.supplyEquation = supplyEquation
        self.priceEnd = sympy.solve(demandEquation)[0]
        self.prices = []
        self.demandQ = []
        self.supplyQ = []
        for price in range(0,self.priceEnd+1):
            self.prices += [price]
            self.demandQ += [demandEquation.subs(p,price)]
            self.supplyQ += [supplyEquation.subs(p,price)]
        self.startingQ = sympy.solve(demandEquation-supplyEquation)[0]
        self.startingP = demandEquation.subs(p,self.startingQ)
        self.mode = "Normal"
    def internationalTrade(self,worldPrice):
        self.mode = "International"
        self.worldPrice = worldPrice
        if worldPrice>self.startingP:
            self.export = True
            self.domesticQ = self.demandEquation.subs(p,worldPrice)
            self.internationalQ = self.supplyEquation.subs(p,worldPrice)-self.demandEquation.subs(p,worldPrice)
            self.totalQ = self.domesticQ+self.internationalQ
        else:
            self.export = False
            self.domesticQ = self.supplyEquation.subs(p,worldPrice)
            self.internationalQ = self.demandEquation.subs(p,worldPrice)-self.domesticQ
            self.totalQ = self.domesticQ+self.internationalQ
        self.domesticQ
    def drawSurplus1(self):
        if self.export:
            triangle1 = patches.Polygon([[self.domesticQ,self.worldPrice],[0,self.worldPrice],[0,self.priceEnd]],True,label="Consumer Surplus")
            triangle2 = patches.Polygon([[0,0],[0,self.startingP],[self.startingQ,self.startingP]],True,label="Original Producer Surplus",color="red")
            rect1 = patches.Polygon([[0,self.startingP],[0,self.worldPrice],[self.domesticQ,self.worldPrice],[self.startingQ,self.startingP]],True,label="Transferred Surplus",color="orange")
            triangle3 = patches.Polygon([[self.domesticQ,self.worldPrice],[self.totalQ,self.worldPrice],[self.startingQ,self.startingP]],True,label="Surplus from International",color="yellow")
            currentAxis = plt.gca()
            currentAxis.add_patch(triangle1)
            currentAxis.add_patch(triangle2)
            currentAxis.add_patch(rect1)
            currentAxis.add_patch(triangle3)
            print("The consumer surplus goes from "+str((self.priceEnd-self.startingP)*self.startingQ*.5)+" to "+str((self.priceEnd-self.worldPrice)*self.domesticQ*.5))
            print("The producer surplus goes from "+str(self.startingP*self.startingQ*.5)+" to "+str(self.worldPrice*self.totalQ*.5))
        else:
            triangle1 = patches.Polygon([[self.startingQ,self.startingP],[0,self.startingP],[0,self.priceEnd]],True,label="Original Consumer Surplus", color="red")
            triangle2 = patches.Polygon([[0,0],[0,self.worldPrice],[self.domesticQ,self.worldPrice]],True,label="Producer Surplus")
            rect1 = patches.Polygon([[0,self.startingP],[0,self.worldPrice],[self.domesticQ,self.worldPrice],[self.startingQ,self.startingP]],True,label="Transferred Surplus",color="orange")
            triangle3 = patches.Polygon([[self.domesticQ,self.worldPrice],[self.totalQ,self.worldPrice],[self.startingQ,self.startingP]],True,label="Surplus from International",color="yellow")
            currentAxis = plt.gca()
            currentAxis.add_patch(triangle1)
            currentAxis.add_patch(triangle2)
            currentAxis.add_patch(rect1)
            currentAxis.add_patch(triangle3)
            print("The producer surplus goes from "+str(self.startingP*self.startingQ*.5)+" to "+str(self.worldPrice*self.domesticQ*.5))
            print("The consumer surplus goes from "+str((self.priceEnd-self.startingP)*self.startingQ*.5)+" to "+str((self.priceEnd-self.worldPrice)*self.domesticQ*.5))

    def plot(self):
        plt.plot(self.demandQ,self.prices,'k')
        plt.plot(self.supplyQ,self.prices,'k')
        plt.plot(self.startingQ,self.startingP, 'bo')
        plt.xlabel("Supply and Demand Quantity")
        plt.ylabel("Price")
        if self.mode=="International":
            plt.plot([0,self.priceEnd],[self.worldPrice,self.worldPrice],"k")
            if self.startingP!=self.worldPrice:
                plt.plot(self.totalQ,self.worldPrice, 'bo')
                plt.plot(self.domesticQ,self.worldPrice,'bo')
                plt.plot([0,self.domesticQ],[self.worldPrice,self.worldPrice],"m",label='Domestic Quantity')
                plt.plot([self.domesticQ,self.totalQ],[self.worldPrice,self.worldPrice],"g", label='International Quantity')

                if self.export:
                    print("The country exports during international trade, and the price moves from "+str(self.startingP)+" to "+str(self.worldPrice)+".")
                    print("The domestic producers go from producing "+str(self.startingQ)+" units to "+str(self.totalQ)+" units.")
                    print("The domestic consumers go from consuming "+str(self.startingQ)+" units to "+str(self.domesticQ)+" units.")
                    print("Total domestic producer revenue is "+str(self.totalQ*self.worldPrice))
                else:
                    print("The country imports during international trade, and the price moves from "+str(self.startingP)+" to "+str(self.worldPrice)+".")
                    print("The domestic producers go from producing "+str(self.startingQ)+" units to "+str(self.domesticQ)+" units.")
                    print("The domestic consumers go from consuming "+str(self.startingQ)+" units to "+str(self.totalQ)+" units.")
                    print("Total domestic producer revenue is "+str(self.domesticQ*self.worldPrice))
                plt.legend()
        plt.show()
```

```python
economy1 = supplyDemand(10-p,p)
economy1.internationalTrade(6)
economy1.drawSurplus1()
economy1.plot()
```

FinanceAndPython.com-EconomicFoundations-master/11 Imports and Exports.ipynb


```python
from scipy.stats import binom, poisson, norm, lognorm
import matplotlib.pyplot as plt
from iminuit import Minuit        # The actual fitting tool, better than scipy's
from scipy import stats
import sympy as sy
import numpy as np
import math
import sys

sys.path.append(r'/home/saim/External_Functions')
from ExternalFunctions import Chi2Regression, BinnedLH, UnbinnedLH
from ExternalFunctions import nice_string_output, add_text_to_ax
```

```python
# Plotting stuff

plt.rcParams['font.size'] = 18
plt.style.use(['science', 'notebook', 'grid'])

pink = '#e377c2'
blue = '#1f77b4'
golden = '#ff7f0e'
green = '#2ca02c'
red = '#d62728'
purple = '#9467bd'
light_blue = '#17becf'
```

```python
r = np.random
r.seed(42)
```

```python
def chi2_eval(fitted_object, Npoints, Nparams):

    Chi2_value = fitted_object.fval
    Ndof = Npoints - Nparams   # Number of degrees of freedom
    Chi2_prob = stats.chi2.sf(Chi2_value, Ndof)

    return Chi2_value, Ndof, Chi2_prob
```

```python
# Turning histogram data into x, y, and sigma_y values for all non-zero entries (not considered in Chi2 fit):
def hist_data(data, Nbins, mini, maxi):

    counts, bin_edges = np.histogram(data,
                                     bins = Nbins,
                                     range = (mini, maxi),
                                     density = False)

    bin_centers = (bin_edges[1:] + bin_edges[:-1]) / 2
    x = bin_centers[counts > 0]

    y = counts[counts > 0]
    sy = np.sqrt(y)

    return x, y, sy
```

```python
def draw_chi2fit(Nparams, x_values, x_min, x_max, PDF,
                 fitted_dist, Nbins, x_bin, y_bin, sigma):

    # Produce the points for drawing the fit:
    x_axis = np.linspace(x_min, x_max, Nbins)
    y_axis = PDF(x_axis, *fitted_dist.values[:]
                      )

# Produce figure with histogram (with error bars) and fit overlayed:
    fig, ax = plt.subplots(figsize=(14, 6))

    ax.errorbar(x_bin, y_bin, sigma, fmt = '.', color = '#1f77b4', label = 'Data')
    ax.plot(x_axis, y_axis, '-', color = golden, label = 'Fit')

    ax.set(xlabel = "Value",
           ylabel = "Frequency",
           title = "")

    ax.legend(loc = 'lower right',
              fontsize=14);

    # Fitting results
    chi2_value = fitted_dist.fval
    Ndof = Nbins - fitted_dist.nfit
    chi2_prob = stats.chi2.sf(chi2_value, Ndof)

# Define figure text
    d = {'Entries': len(x_values),
         'Chi2':       chi2_value,
         'ndf':        Ndof,
         'Prob':       chi2_prob,
    }

    for name in fitted_dist.parameters:
        d[name] = [fitted_dist.values[name], fitted_dist.errors[name]]

    text = nice_string_output(d, extra_spacing = 2, decimals = 3)
    add_text_to_ax(0.69, 0.95, text, ax, fontsize = 15)
    fig.tight_layout()
```

## Gaussian chi2 fit

```python
Npointz = 10000     # Number of random points produced
x_all = r.normal(loc = 0.2,
                 scale = 1.1,
                 size = Npointz)

Nbinz = 100
xmin, xmax = np.min(x_all), np.max(x_all)
binwidth_gauss = np.ptp(x_all) / Nbinz
#binwidth = (xmax - xmin) / Nbins

# Fitting function which is NOT normalised but has normalisation constants "N" in,
# and includes the bin width:
def func_gauss_norm(x, N, mu, sigma) :
    norm = binwidth_gauss * N / np.sqrt(2.0 * np.pi) / sigma
    z = (x - mu) / sigma
    return norm * np.exp(-0.5 * (z**2))

def func_gaussian_alt(x, N, mu, sigma) :
    return binwidth_gauss * N * norm.pdf(x, mu, sigma)
```

```python
x1, y1, sy1 = hist_data(x_all, Nbinz, xmin, xmax)
```

```python
# Fitting
chi2_gaussian = Chi2Regression(func_gauss_norm, x1, y1, sy1)     # Fitting object
chi2_gaussian.errordef = Minuit.LEAST_SQUARES

minuit_gaussian = Minuit(chi2_gaussian,
                         N = Npointz,
                         mu = 0,
                         sigma = 0.05)

minuit_gaussian.migrad()       # Perform the actual fit
```

         FCN = 75.35
         Nfcn = 196 (196 total)

         EDM = 2.32e-07 (Goal: 0.0002)

         Valid Minimum
         Valid Parameters
         No Parameters at limit

         Below EDM threshold (goal x 10)
         Below call limit

         Hesse ok
         Has Covariance
         Accurate
         Pos. def.
         Not forced

         Name
         Value
         Hesse Error
         Minos Error-
         Minos Error+
         Limit-
         Limit+
         Fixed

         0
         N
         9.93e3
         0.10e3

         1
         mu
         0.199
         0.011

         2
         sigma
         1.094
         0.008

```python
draw_chi2fit(3, x_all, xmin, xmax, func_gauss_norm,
             minuit_gaussian, Nbinz, x1, y1, sy1)
```

## Linear chi2 fit

```python
# Fitting function
def func_linear(x, alpha0, alpha1):
        return alpha0 + alpha1*x
```

```python
# Parameters
alpha0 = 3.6
alpha1 = 0.3
sigma_y = 0.5
```

```python
lin_Npoints = 50                        # Number of random points produced
lin_x = np.arange(lin_Npoints)               # Generate points in array
#exLin = np.zeros_like(lin_x)

lin_y = alpha0 + alpha1 * lin_x + r.normal(0, sigma_y, lin_Npoints)    # linear function + gaussian errors
error_lin_y = sigma_y * np.ones_like(lin_x)
```

```python
# Fitting
chi2_linear = Chi2Regression(func_linear, lin_x, lin_y, error_lin_y)   # Fitting object
chi2_linear.errordef = Minuit.LEAST_SQUARES

# Give fitting function, its parameters their starting fitting values
minuit_linear = Minuit(chi2_linear,
                         alpha0 = 2,
                         alpha1 = 0.1)

minuit_linear.migrad()  # perform the actual fit
```

         FCN = 37.08
         Nfcn = 32 (32 total)

         EDM = 1.82e-21 (Goal: 0.0002)

         Valid Minimum
         Valid Parameters
         No Parameters at limit

         Below EDM threshold (goal x 10)
         Below call limit

         Hesse ok
         Has Covariance
         Accurate
         Pos. def.
         Not forced

         Name
         Value
         Hesse Error
         Minos Error-
         Minos Error+
         Limit-
         Limit+
         Fixed

         0
         alpha0
         3.63
         0.14

         1
         alpha1
         0.301
         0.005

```python
chi2_linear, Ndof_linear, pval_linear = chi2_eval(minuit_linear, len(lin_x), 2)
```

```python
figLin, axLin = plt.subplots(figsize=(16, 8))

axLin.errorbar(lin_x,
               lin_y,
               error_lin_y,
               fmt = 'ro',
               ecolor = 'k',
               elinewidth = 1,
               capsize = 1,
               capthick = 1)

axLin.plot(lin_x,
           func_linear(lin_x, *minuit_linear.values[:]),
           '-r',
           color = blue)

d = {'Intercept':[minuit_linear.values['alpha0'], minuit_linear.errors['alpha0']],
     'Slope':    [minuit_linear.values['alpha1'], minuit_linear.errors['alpha1']],
     'Chi2':     chi2_linear,
     'ndf':      Ndof_linear,
     'Prob':     pval_linear,
    }

text = nice_string_output(d, extra_spacing=2, decimals=3)
add_text_to_ax(0.04, 0.95, text, axLin, fontsize=20)
figLin.tight_layout()
```

## Monte Carlo Simulation and Fitting

```python
N_points = 10000
N_bins = 100
# inverse integrated function added to itself 4 times
exp_inv = sum(-0.8*np.log(r.uniform(size = N_points)) for i in range(4))

# Function given in problem statement which is summed 4 times
def exp_func(x):
    return sum(r.exponential(0.8, N_points) for i in range(4))

xmin_exp = 0
xmax_exp = 20

x_axis_exp = np.linspace(start = xmin_exp,
                      stop = xmax_exp,
                      num = 10000)
y_axis_exp = exp_func(x_axis_exp)
```

```python
# Init plot object
fig, ax = plt.subplots(figsize=(15, 9))

# Plot generated data
ax.hist(exp_inv,
        bins = N_bins,
        range = (xmin_exp, xmax_exp),
        color = blue,
        histtype = 'step'
       )
# Plot labels
ax.set(xlabel = "x - following f(x)",
       ylabel = "Frequency",
       xlim = (xmin_exp -1.0 , xmax_exp+1.0))

# Define figure text
textstr = '\n'.join((
    r'$\mathrm{Entries}=%.2f$' % (len(exp_inv), ),
    r'$\mathrm{Mean}=%.2f$' % (exp_inv.mean(), ),
    r'$\mathrm{Std}=%.2f$' % (exp_inv.std(ddof=1), )))

# Plot figure text
props = dict(boxstyle = 'round',
             facecolor = 'white',
             edgecolor = 'black',
             alpha=0.5)

# place a text box in upper left in axes coords
ax.text(0.86,
        0.95,
        textstr,
        transform = ax.transAxes,
        fontsize = 14,
        verticalalignment='top',
        bbox = props)

fig.tight_layout()
```

```python
# Binning the data
x3, y3, sigma_y3 = hist_data(exp_inv, 100, 0, 20)
```

```python
# Fitting
chi2_MC_Gauss = Chi2Regression(func_gauss_norm, x3, y3, sigma_y3)     # Fitting object
chi2_MC_Gauss.errordef = Minuit.LEAST_SQUARES

minuit_MC_Gauss = Minuit(chi2_MC_Gauss,
                         N = N_points,
                         mu = 3,
                         sigma = 1.6)

minuit_MC_Gauss.migrad()       # Perform the actual fit
```

         FCN = 1148
         Nfcn = 65 (65 total)

         EDM = 2.98e-05 (Goal: 0.0002)

         Valid Minimum
         Valid Parameters
         No Parameters at limit

         Below EDM threshold (goal x 10)
         Below call limit

         Hesse ok
         Has Covariance
         Accurate
         Pos. def.
         Not forced

         Name
         Value
         Hesse Error
         Minos Error-
         Minos Error+
         Limit-
         Limit+
         Fixed

         0
         N
         20.82e3
         0.22e3

         1
         mu
         3.126
         0.019

         2
         sigma
         1.349
         0.015

```python
draw_chi2fit(3, exp_inv, xmin_exp, xmax_exp,
             func_gaussian_alt, minuit_MC_Gauss, N_bins, x3, y3, sigma_y3)
```

## Exponential Fit

```python
N_exp = 10000     # Number of random points produced
x_exp = r.exponential(np.e, N_exp)

exp_bins = 100
binwidth_exp = np.ptp(x_exp) / exp_bins
exp_min, exp_max = np.min(x_exp), np.max(x_exp)
```

```python
def exp_pdf(x, N, tau):
    return N * binwidth_exp / tau * np.exp(-x/tau)
```

```python
# Binning data
x4, y4, sy4 = hist_data(x_exp, exp_bins, exp_min, exp_max)
```

```python
# Fitting
chi2_exp = Chi2Regression(exp_pdf, x4, y4, sy4)     # Fitting object
chi2_exp.errordef = Minuit.LEAST_SQUARES

minuit_exp = Minuit(chi2_exp,
                         N = 10000,
                         tau = 2)

minuit_exp.migrad()       # Perform the actual fit
```

         FCN = 85.13
         Nfcn = 38 (38 total)

         EDM = 7.49e-07 (Goal: 0.0002)

         Valid Minimum
         Valid Parameters
         No Parameters at limit

         Below EDM threshold (goal x 10)
         Below call limit

         Hesse ok
         Has Covariance
         Accurate
         Pos. def.
         Not forced

         Name
         Value
         Hesse Error
         Minos Error-
         Minos Error+
         Limit-
         Limit+
         Fixed

         0
         N
         9.93e3
         0.10e3

         1
         tau
         2.693
         0.027

```python
draw_chi2fit(2, x_exp, exp_min, exp_max, exp_pdf,
             minuit_exp, exp_bins, x4, y4, sy4)
```

## Power Law Fit

```python
N_pow = 10000     # Number of random points produced
x_pow = r.power(a = 15,
                size = N_pow)

pow_bins = 100
binwidth_pow = np.ptp(x_pow) / pow_bins
pow_min, pow_max = np.min(x_pow), np.max(x_pow)
```

```python
def power_pdf(x, N, a, b):
    return N * binwidth_pow / a * np.power(x, b)
```

```python
# Binning data
x5, y5, sy5 = hist_data(x_pow, pow_bins, pow_min, pow_max)
```

```python
# Fitting
chi2_pow = Chi2Regression(power_pdf, x5, y5, sy5)     # Fitting object
chi2_pow.errordef = Minuit.LEAST_SQUARES

minuit_pow = Minuit(chi2_pow,
                         N = 10000,
                         a = 4,
                         b = 1)

minuit_pow.migrad()       # Perform the actual fit
```

         FCN = 59.57
         Nfcn = 129 (129 total)

         EDM = 1.17e-05 (Goal: 0.0002)

         Valid Minimum
         Valid Parameters
         No Parameters at limit

         Below EDM threshold (goal x 10)
         Below call limit

         Hesse ok
         Has Covariance
         Accurate
         Pos. def.
         Not forced

         Name
         Value
         Hesse Error
         Minos Error-
         Minos Error+
         Limit-
         Limit+
         Fixed

         0
         N
         0.11e6
         0.12e6

         1
         a
         0.7
         0.8

         2
         b
         14.24
         0.16

```python
draw_chi2fit(3, x_pow, pow_min, pow_max, power_pdf,
             minuit_pow, pow_bins, x5, y5, sy5)
```

```python
stats.anderson(x_exp, dist='expon')
```

    AndersonResult(statistic=0.8263378287992964, critical_values=array([0.922, 1.078, 1.341, 1.606, 1.957]), significance_level=array([15. , 10. ,  5. ,  2.5,  1. ]))

```python
# Null is accepted at all significance levels as test statistic is lower than critical values
```


### Prelude

```python
import numpy as np

from sympy import *
import scipy.integrate as integrate
import math

import matplotlib.pyplot as plt
%matplotlib inline

def gaussian(x,mu,sigma):
    return (1/(sigma * sqrt(2*pi)) * exp(-(mu-x)**2 / (2*sigma**2)))
```

### Input

```python
hours_estimated = 6
plus_minus = 3
```

### Computation

```python
x = np.linspace(hours_estimated - 4*plus_minus, hours_estimated + 4*plus_minus,100)
y = np.array([gaussian(v,hours_estimated, plus_minus) for v in x],dtype='float')
plt.grid(True)
plt.title('Expected distribution (normal)')
plt.xlabel('effective time (hours)')
plt.ylabel('probability distribution')
plt.plot(x,y,color='gray')
plt.fill_between(x,y,0,color='#c0f0c0')
plt.show()
```

```python
x = np.linspace(hours_estimated - 4*plus_minus, hours_estimated + 4*plus_minus,100)
y = np.array([gaussian(v,hours_estimated, plus_minus) for v in x],dtype='float')
plt.grid(True)
plt.title('The impossible cases')
plt.xlabel('effective time (hours)')
plt.ylabel('probability distribution')
plt.plot(x,y,color='gray')
plt.fill_between(x,y,where=x>0,color='#c0f0c0')
plt.fill_between(x,y,where=x<=0, color='#ff0000')
plt.show()
```

### summing up the error

```python
red_surface = integrate.quad(lambda x: gaussian(x, hours_estimated, plus_minus), -(math.inf), 0)
red_surface[0]
```

    0.022750131948179212

```python
corrected_mean = 0
for x_val in x:
    integral_so_far = integrate.quad(lambda x: gaussian(x, hours_estimated, plus_minus), hours_estimated, x_val)
    if(integral_so_far[0] > red_surface[0]):
        corrected_mean = x_val
        break

def belowzero(x):
    if (x<=0):
        return true
    else:
        return false

def zerotomean(x):
    if (x>0 and x<= hours_estimated):
        return true
    else:
        return false

def meantocorrectedmean(x):
    if (x>hours_estimated and x<= corrected_mean):
        return true
    else:
        return false

def abovecorrected(x):
    if (x>corrected_mean):
        return true
    else:
        return false
```

```python

x = np.linspace(hours_estimated - 4*plus_minus, hours_estimated + 4*plus_minus,100)
y = np.array([gaussian(v,hours_estimated, plus_minus) for v in x],dtype='float')
plt.grid(True)
plt.title('Correction on mean')
plt.xlabel('effective time (hours)')
plt.ylabel('probability distribution')
plt.plot(x,y,color='gray')
plt.fill_between(x,y,where=(x<=0), color='#ff0000')
plt.fill_between(x,y,where=((x>0) & (x<=hours_estimated)), color='#c0f0c0')
plt.fill_between(x,y,where=((x>hours_estimated) & (x<=corrected_mean)), color='#ff0000')
plt.fill_between(x,y,where=(x>corrected_mean), color='#c0f0c0')
plt.show()
```

### Error (%)

```python
err = (corrected_mean - hours_estimated) / hours_estimated * 100
print("A task estimated at %.2f hours (+/- %.2f hours) is likely to turn out %.2f percent above estimation" % (hours_estimated, plus_minus, err))
```

    A task estimated at 6.00 hours (+/- 3.00 hours) is likely to turn out 6.06 percent above estimation


```python
import transforms3d.quaternions as tf
import sympy
sympy.init_printing()
```

```python
a = sympy.symbols('a_1:3')
b = sympy.symbols('b_1:3')
c = sympy.symbols('c_1:3')
d = sympy.symbols('d_1:3')

q0 = sympy.Matrix([a[0], b[0], c[0], d[0]])
q1 = sympy.Matrix([a[1], b[1], c[1], d[1]])
```

```python
def qmult(q0, q1):
    res = sympy.Matrix(tf.qmult(q0, q1))
    res.simplify()
    return res

def qconj(q):
    return sympy.Matrix([q[0], -q[1], -q[2], -q[3]])

def qnorm(q):
    return sympy.sqrt(qmult(q, qconj(q))[0])
```

```python
qmult(q0, q1)
```

```python
qconj(q1).T
```

```python
qnorm(q1)
```

```python
qmult(q0, qconj(q0)/qnorm(q0)**2)
```

```python
qmult(q0, qconj(q0)/qnorm(q0))
```

```python
def qinv(q):
    res = qconj(q)/qnorm(q)**2
    res.simplify()
    return res
```

```python
qmult(q0, qinv(q1))
```