$$
\begin{equation}
mu^{\prime\prime} + f(u^{\prime}) + s(u) = F(t),\quad u(0)=I,\ u^{\prime}(0)=V,\ t\in (0,T]
\thinspace .
\label{vib:ode2} \tag{1}
\end{equation}
$$
We have also included a possibly nonzero initial value for $u^{\prime}(0)$.
The parameters $m$, $f(u^{\prime})$, $s(u)$, $F(t)$, $I$, $V$, and $T$ are
input data.

There are two main types of damping (friction) forces: linear $f(u^{\prime})=bu$, or
quadratic $f(u^{\prime})=bu^{\prime}|u^{\prime}|$. Spring systems often feature linear
damping, while air resistance usually gives rise to quadratic damping.
Spring forces are often linear: $s(u)=cu$, but nonlinear versions
are also common, the most famous is the gravity force on a pendulum
that acts as a spring with $s(u)\sim \sin(u)$.

## A centered scheme for linear damping

Sampling ([1](#vib:ode2)) at a mesh point $t_n$, replacing
$u^{\prime\prime}(t_n)$ by $[D_tD_tu]^n$, and $u^{\prime}(t_n)$ by $[D_{2t}u]^n$ results
in the discretization

$$
\begin{equation}
[mD_tD_t u + f(D_{2t}u) + s(u) = F]^n,
\label{_auto1} \tag{2}
\end{equation}
$$

which written out means

$$
\begin{equation}
m\frac{u^{n+1}-2u^n + u^{n-1}}{\Delta t^2}
+ f(\frac{u^{n+1}-u^{n-1}}{2\Delta t}) + s(u^n) = F^n,
\label{vib:ode2:step3b} \tag{3}
\end{equation}
$$

where $F^n$ as usual means $F(t)$ evaluated at $t=t_n$.
Solving ([3](#vib:ode2:step3b)) with respect to the unknown
$u^{n+1}$ gives a problem: the $u^{n+1}$ inside the $f$ function
makes the equation *nonlinear* unless $f(u^{\prime})$ is a linear function,
$f(u^{\prime})=bu^{\prime}$. For now we shall assume that $f$ is linear in $u^{\prime}$.
Then

$$
\begin{equation}
m\frac{u^{n+1}-2u^n + u^{n-1}}{\Delta t^2}
+ b\frac{u^{n+1}-u^{n-1}}{2\Delta t} + s(u^n) = F^n,
\label{vib:ode2:step3b2} \tag{4}
\end{equation}
$$

which gives an explicit formula for $u$ at each
new time level:

$$
\begin{equation}
u^{n+1} = (2mu^n + (\frac{b}{2}\Delta t - m)u^{n-1} +
\Delta t^2(F^n - s(u^n)))(m + \frac{b}{2}\Delta t)^{-1}
\label{vib:ode2:step4} \tag{5}
\thinspace .
\end{equation}
$$

For the first time step we need to discretize $u^{\prime}(0)=V$
as $[D_{2t}u = V]^0$ and combine
with ([5](#vib:ode2:step4)) for $n=0$. The discretized initial condition
leads to

$$
\begin{equation}
u^{-1} = u^{1} - 2\Delta t V,
\label{vib:ode2:ic:du} \tag{6}
\end{equation}
$$

which inserted in ([5](#vib:ode2:step4)) for $n=0$ gives an equation
that can be solved for
$u^1$:

$$
\begin{equation}
u^1 = u^0 + \Delta t\, V
+ \frac{\Delta t^2}{2m}(-bV - s(u^0) + F^0)
\thinspace .
\label{vib:ode2:step4b} \tag{7}
\end{equation}
$$

## A centered scheme for quadratic damping

When $f(u^{\prime})=bu^{\prime}|u^{\prime}|$, we get a quadratic equation for $u^{n+1}$
in ([3](#vib:ode2:step3b)). This equation can be straightforwardly
solved by the well-known formula for the roots of a quadratic equation.
However, we can also avoid the nonlinearity by introducing
an approximation with an error of order no higher than what we
already have from replacing derivatives with finite differences.

We start with ([1](#vib:ode2)) and only replace
$u^{\prime\prime}$ by $D_tD_tu$, resulting in

$$
\begin{equation}
[mD_tD_t u + bu^{\prime}|u^{\prime}| + s(u) = F]^n\thinspace .
\label{vib:ode2:quad:idea1} \tag{8}
\end{equation}
$$

Here, $u^{\prime}|u^{\prime}|$ is to be computed at time $t_n$. The idea
is now to introduce
a *geometric mean*, defined by

$$
(w^2)^n \approx w^{n-\frac{1}{2}}w^{n+\frac{1}{2}},
$$

for some quantity $w$ depending on time. The error in the geometric mean
approximation is $\Oof{\Delta t^2}$, the same as in the
approximation $u^{\prime\prime}\approx D_tD_tu$. With $w=u^{\prime}$ it follows
that

$$
[u^{\prime}|u^{\prime}|]^n \approx u^{\prime}(t_{n+\frac{1}{2}})|u^{\prime}(t_{n-\frac{1}{2}})|\thinspace .
$$

The next step is to approximate
$u^{\prime}$ at $t_{n\pm 1/2}$, and fortunately a centered difference
fits perfectly into the formulas since it involves $u$ values at
the mesh points only. With the approximations

$$
\begin{equation}
u^{\prime}(t_{n+1/2})\approx [D_t u]^{n+\frac{1}{2}},\quad u^{\prime}(t_{n-1/2})\approx [D_t u]^{n-\frac{1}{2}},
\label{vib:ode2:quad:idea2} \tag{9}
\end{equation}
$$

we get

$$
\begin{equation}
[u^{\prime}|u^{\prime}|]^n
\approx [D_tu]^{n+\frac{1}{2}}|[D_tu]^{n-\frac{1}{2}}| = \frac{u^{n+1}-u^n}{\Delta t}
\frac{|u^n-u^{n-1}|}{\Delta t}
\thinspace .
\label{_auto2} \tag{10}
\end{equation}
$$

The counterpart to ([3](#vib:ode2:step3b)) is then

$$
\begin{equation}
m\frac{u^{n+1}-2u^n + u^{n-1}}{\Delta t^2}
+ b\frac{u^{n+1}-u^n}{\Delta t}\frac{|u^n-u^{n-1}|}{\Delta t}
+ s(u^n) = F^n,
\label{vib:ode2:step3b:quad} \tag{11}
\end{equation}
$$

which is linear in the unknown $u^{n+1}$. Therefore, we can easily solve
([11](#vib:ode2:step3b:quad))
with respect to $u^{n+1}$ and achieve the explicit updating formula

$$
u^{n+1} =  \left( m + b|u^n-u^{n-1}|\right)^{-1}\times \nonumber
$$

$$
\begin{equation}
 \qquad \left(2m u^n - mu^{n-1} + bu^n|u^n-u^{n-1}| + \Delta t^2 (F^n - s(u^n))
\right)
\thinspace .
\label{vib:ode2:step4:quad} \tag{12}
\end{equation}
$$

In the derivation of a special equation for the first
time step we run into some trouble: inserting ([6](#vib:ode2:ic:du))
in ([12](#vib:ode2:step4:quad)) for $n=0$ results in a complicated nonlinear
equation for $u^1$. By thinking differently about the problem we can
easily get away with the nonlinearity again. We have for $n=0$ that
$b[u^{\prime}|u^{\prime}|]^0 = bV|V|$. Using this value in ([8](#vib:ode2:quad:idea1))
gives

$$
\begin{equation}
[mD_tD_t u + bV|V| + s(u) = F]^0
\thinspace .
\label{_auto3} \tag{13}
\end{equation}
$$

Writing this equation out and using ([6](#vib:ode2:ic:du)) results in the
special equation for the first time step:

$$
\begin{equation}
u^1 = u^0 + \Delta t V + \frac{\Delta t^2}{2m}\left(-bV|V| - s(u^0) + F^0\right)
\thinspace .
\label{vib:ode2:step4b:quad} \tag{14}
\end{equation}
$$

## A forward-backward discretization of the quadratic damping term

The previous section first proposed to discretize the quadratic
damping term $|u^{\prime}|u^{\prime}$ using centered differences:
$[|D_{2t}|D_{2t}u]^n$. As this gives rise to a nonlinearity in
$u^{n+1}$, it was instead proposed to use a geometric mean combined
with centered differences.  But there are other alternatives. To get
rid of the nonlinearity in $[|D_{2t}|D_{2t}u]^n$, one can think
differently: apply a backward difference to $|u^{\prime}|$, such that
the term involves known values, and apply a forward difference to
$u^{\prime}$ to make the term linear in the unknown $u^{n+1}$. With
mathematics,

$$
\begin{equation}
[\beta |u^{\prime}|u^{\prime}]^n \approx \beta |[D_t^-u]^n|[D_t^+ u]^n =
\beta\left\vert\frac{u^n-u^{n-1}}{\Delta t}\right\vert
\frac{u^{n+1}-u^n}{\Delta t}\thinspace .
\label{vib:ode2:nonlin:fbdiff} \tag{15}
\end{equation}
$$

The forward and backward differences both have an error proportional
to $\Delta t$ so one may think the discretization above leads to
a first-order scheme.
However, by looking at the formulas, we realize that the forward-backward
differences in ([15](#vib:ode2:nonlin:fbdiff))
result in exactly the same scheme as in
([11](#vib:ode2:step3b:quad)) where we
used a geometric mean and centered differences and committed errors
of size $\Oof{\Delta t^2}$. Therefore, the forward-backward
differences in ([15](#vib:ode2:nonlin:fbdiff))
act in a symmetric way and actually produce a second-order
accurate discretization of the quadratic damping term.

## Implementation

The algorithm arising from the methods in the sections [A centered scheme for linear damping](#vib:ode2:fdm:flin)
and [A centered scheme for quadratic damping](#vib:ode2:fdm:fquad) is very similar to the undamped case in
the section [vib:ode1:fdm](#vib:ode1:fdm). The difference is
basically a question of different formulas for $u^1$ and
$u^{n+1}$. This is actually quite remarkable. The equation
([1](#vib:ode2)) is normally impossible to solve by pen and paper, but
possible for some special choices of $F$, $s$, and $f$. On the
contrary, the complexity of the
nonlinear generalized model ([1](#vib:ode2)) versus the
simple undamped model is not a big deal when we solve the
problem numerically!

The computational algorithm takes the form

1. $u^0=I$

2. compute $u^1$ from ([7](#vib:ode2:step4b)) if linear
    damping or ([14](#vib:ode2:step4b:quad)) if quadratic damping

3. for $n=1,2,\ldots,N_t-1$:

a. compute $u^{n+1}$ from ([5](#vib:ode2:step4)) if linear
      damping or ([12](#vib:ode2:step4:quad)) if quadratic damping

Modifying the `solver` function for the undamped case is fairly
easy, the big difference being many more terms and if tests on
the type of damping:

```python
def solver(I, V, m, b, s, F, dt, T, damping='linear'):
    """
    Solve m*u'' + f(u') + s(u) = F(t) for t in (0,T],
    u(0)=I and u'(0)=V,
    by a central finite difference method with time step dt.
    If damping is 'linear', f(u')=b*u, while if damping is
    'quadratic', f(u')=b*u'*abs(u').
    F(t) and s(u) are Python functions.
    """
    dt = float(dt); b = float(b); m = float(m) # avoid integer div.
    Nt = int(round(T/dt))
    u = np.zeros(Nt+1)
    t = np.linspace(0, Nt*dt, Nt+1)

    u[0] = I
    if damping == 'linear':
        u[1] = u[0] + dt*V + dt**2/(2*m)*(-b*V - s(u[0]) + F(t[0]))
    elif damping == 'quadratic':
        u[1] = u[0] + dt*V + \
               dt**2/(2*m)*(-b*V*abs(V) - s(u[0]) + F(t[0]))

    for n in range(1, Nt):
        if damping == 'linear':
            u[n+1] = (2*m*u[n] + (b*dt/2 - m)*u[n-1] +
                      dt**2*(F(t[n]) - s(u[n])))/(m + b*dt/2)
        elif damping == 'quadratic':
            u[n+1] = (2*m*u[n] - m*u[n-1] + b*u[n]*abs(u[n] - u[n-1])
                      + dt**2*(F(t[n]) - s(u[n])))/\
                      (m + b*abs(u[n] - u[n-1]))
    return u, t
```

The complete code resides in the file [`vib.py`](${src_vib}/vib.py).

## Verification

### Constant solution

For debugging and initial verification, a constant solution is often
very useful. We choose $\uex(t)=I$, which implies $V=0$.
Inserted in the ODE, we get
$F(t)=s(I)$ for any choice of $f$. Since the discrete derivative
of a constant vanishes (in particular, $[D_{2t}I]^n=0$,
$[D_tI]^n=0$, and $[D_tD_t I]^n=0$), the constant solution also fulfills
the discrete equations. The constant should therefore be reproduced
to machine precision. The function `test_constant` in `vib.py`
implements this test.

### Linear solution

Now we choose a linear solution: $\uex = ct + d$. The initial condition
$u(0)=I$ implies $d=I$, and $u^{\prime}(0)=V$ forces $c$ to be $V$.
Inserting $\uex=Vt+I$ in the ODE with linear damping results in

$$
0 + bV + s(Vt+I) = F(t),
$$

while quadratic damping requires the source term

$$
0 + b|V|V + s(Vt+I) = F(t)\thinspace .
$$

Since the finite difference approximations used to compute $u^{\prime}$ all
are exact for a linear function, it turns out that the linear $\uex$
is also a solution of the discrete equations.
[vib:exer:verify:gen:linear](#vib:exer:verify:gen:linear) asks you to carry out
all the details.

### Quadratic solution

Choosing $\uex = bt^2 + Vt + I$, with $b$ arbitrary,
fulfills the initial conditions and
fits the ODE if $F$ is adjusted properly. The solution also solves
the discrete equations with linear damping. However, this quadratic
polynomial in $t$ does not fulfill the discrete equations in case
of quadratic damping, because the geometric mean used in the approximation
of this term introduces an error.
Doing [vib:exer:verify:gen:linear](#vib:exer:verify:gen:linear) will reveal
the details. One can fit $F^n$ in the discrete equations such that
the quadratic polynomial is reproduced by the numerical method (to
machine precision).

### Catching bugs

How good are the constant and quadratic solutions at catching
bugs in the implementation? Let us check that by introducing some bugs.

 * Use `m` instead of `2*m` in the denominator of `u[1]`: code works for constant
   solution, but fails (as it should) for a quadratic one.

 * Use `b*dt` instead of `b*dt/2` in the updating formula for `u[n+1]`
   in case of linear damping: constant and quadratic both fail.

 * Use `F[n+1]` instead of `F[n]` in case of linear or quadratic damping:
   constant solution works, quadratic fails.

We realize that the constant solution is very useful for catching certain bugs because
of its simplicity (easy to predict what the different terms in the
formula should evaluate to), while the quadratic solution seems
capable of detecting all (?) other kinds of typos in the scheme.
These results demonstrate why we focus so much on exact, simple polynomial
solutions of the numerical schemes in these writings.

## Visualization

The functions for visualizations differ significantly from
those in the undamped case in the `vib_undamped.py` program because,
in the present general case, we do not have an exact solution to
include in the plots. Moreover, we have no good estimate of
the periods of the oscillations as there will be one period
determined by the system parameters, essentially the
approximate frequency $\sqrt{s'(0)/m}$ for linear $s$ and small damping,
and one period dictated by $F(t)$ in case the excitation is periodic.
This is, however,
nothing that the program can depend on or make use of.
Therefore, the user has to specify $T$ and the window width
to get a plot that moves with the graph and shows
the most recent parts of it in long time simulations.

The `vib.py` code
contains several functions for analyzing the time series signal
and for visualizing the solutions.

## User interface

The `main` function is changed substantially from
the `vib_undamped.py` code, since we need to
specify the new data $c$, $s(u)$, and $F(t)$.  In addition, we must
set $T$ and the plot window width (instead of the number of periods we
want to simulate as in `vib_undamped.py`). To figure out whether we
can use one plot for the whole time series or if we should follow the
most recent part of $u$, we can use the `plot_empricial_freq_and_amplitude`
function's estimate of the number of local maxima. This number is now
returned from the function and used in `main` to decide on the
visualization technique.

```python
%matplotlib inline

def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--I', type=float, default=1.0)
    parser.add_argument('--V', type=float, default=0.0)
    parser.add_argument('--m', type=float, default=1.0)
    parser.add_argument('--c', type=float, default=0.0)
    parser.add_argument('--s', type=str, default='u')
    parser.add_argument('--F', type=str, default='0')
    parser.add_argument('--dt', type=float, default=0.05)
    parser.add_argument('--T', type=float, default=140)
    parser.add_argument('--damping', type=str, default='linear')
    parser.add_argument('--window_width', type=float, default=30)
    parser.add_argument('--savefig', action='store_true')
    a = parser.parse_args()
    from scitools.std import StringFunction
    s = StringFunction(a.s, independent_variable='u')
    F = StringFunction(a.F, independent_variable='t')
    I, V, m, c, dt, T, window_width, savefig, damping = \
       a.I, a.V, a.m, a.c, a.dt, a.T, a.window_width, a.savefig, \
       a.damping

    u, t = solver(I, V, m, c, s, F, dt, T)
    num_periods = empirical_freq_and_amplitude(u, t)
    if num_periods <= 15:
        figure()
        visualize(u, t)
    else:
        visualize_front(u, t, window_width, savefig)
    show()
```

The program `vib.py` contains
the above code snippets and can solve the model problem
([1](#vib:ode2)). As a demo of `vib.py`, we consider the case
$I=1$, $V=0$, $m=1$, $c=0.03$, $s(u)=\sin(u)$, $F(t)=3\cos(4t)$,
$\Delta t = 0.05$, and $T=140$. The relevant command to run is

        Terminal> python vib.py --s 'sin(u)' --F '3*cos(4*t)' --c 0.03

This results in a [moving window following the function](${doc_notes}/vib/html//mov-vib/vib_generalized_dt0.05/index.html) on the screen.
[Figure](#vib:ode2:fig:demo) shows a part of the time series.

Damped oscillator excited by a sinusoidal function.

## The Euler-Cromer scheme for the generalized model

The ideas of the Euler-Cromer method from the section [vib:model2x2:EulerCromer](#vib:model2x2:EulerCromer)
carry over to the generalized model. We write ([1](#vib:ode2))
as two equations for $u$ and $v=u^{\prime}$. The first equation is taken as the
one with $v^{\prime}$ on the left-hand side:

$$
\begin{equation}
v^{\prime} = \frac{1}{m}(F(t)-s(u)-f(v)),
\label{vib:ode2:EulerCromer:veq} \tag{16}
\end{equation}
$$

$$
\begin{equation}
u^{\prime} = v\thinspace .
\label{vib:ode2:EulerCromer:ueq} \tag{17}
\end{equation}
$$

Again, the idea is to step ([16](#vib:ode2:EulerCromer:veq)) forward using
a standard Forward Euler method, while we update $u$ from
([17](#vib:ode2:EulerCromer:ueq)) with a Backward Euler method,
utilizing the recent, computed $v^{n+1}$ value. In detail,

$$
\begin{equation}
\frac{v^{n+1}-v^n}{\Delta t} = \frac{1}{m}(F(t_n)-s(u^n)-f(v^n)),
\label{vib:ode2:EulerCromer:dveq0a} \tag{18}
\end{equation}
$$

$$
\begin{equation}
\frac{u^{n+1}-u^n}{\Delta t} = v^{n+1},
\label{vib:ode2:EulerCromer:dueq0a} \tag{19}
\end{equation}
$$

resulting in the explicit scheme

$$
\begin{equation}
v^{n+1} = v^n + \Delta t\frac{1}{m}(F(t_n)-s(u^n)-f(v^n)),
\label{vib:ode2:EulerCromer:dveq} \tag{20}
\end{equation}
$$

$$
\begin{equation}
u^{n+1} = u^n + \Delta t\,v^{n+1}\thinspace .
\label{vib:ode2:EulerCromer:dueq0} \tag{21}
\end{equation}
$$

We immediately note one very favorable feature of this scheme: all the
nonlinearities in $s(u)$ and $f(v)$ are evaluated at a previous time
level. This makes the Euler-Cromer method easier to apply and
hence much more convenient than the centered scheme for the second-order
ODE ([1](#vib:ode2)).

The initial conditions are trivially set as

$$
\begin{equation}
v^0 = V,
\label{_auto4} \tag{22}
\end{equation}
$$

$$
\begin{equation}
u^0 = I\thinspace .
\label{_auto5} \tag{23}
\end{equation}
$$

[hpl 1: odespy for the generalized problem]

## The Stoermer-Verlet algorithm for the generalized model

We can easily apply the ideas from the section [vib:model2x2:StormerVerlet](#vib:model2x2:StormerVerlet) to
extend that method to the generalized model

$$
\begin{align*}
v^{\prime} &= \frac{1}{m}(F(t)-s(u)-f(v)),\\
u^{\prime} &= v\thinspace .
\end{align*}
$$

However, since the scheme is essentially centered differences for
the ODE system on a staggered mesh, we do not go into detail here,
but refer to the section [A staggered Euler-Cromer scheme for a generalized model](#vib:ode2:staggered).

## A staggered Euler-Cromer scheme for a generalized model

The more general model for vibration problems,

$$
\begin{equation}
mu'' + f(u') + s(u) = F(t),\quad u(0)=I,\ u'(0)=V,\ t\in (0,T],
\label{_auto6} \tag{24}
\end{equation}
$$

can be rewritten as a first-order ODE system

$$
\begin{equation}
v' = m^{-1}\left(F(t) - f(v) - s(u)\right),
\label{vib:ode2:staggered:veq} \tag{25}
\end{equation}
$$

$$
\begin{equation}
u' = v\thinspace .
\label{vib:ode2:staggered:ueq} \tag{26}
\end{equation}
$$

mathcal{I}_t is natural to introduce a staggered mesh (see the section [vib:model2x2:staggered](#vib:model2x2:staggered)) and seek $u$ at mesh points $t_n$ (the numerical value is
denoted by $u^n$) and $v$ between mesh points at $t_{n+1/2}$ (the numerical
value is denoted by $v^{n+\frac{1}{2}}$).
A centered difference approximation to ([26](#vib:ode2:staggered:ueq))-([25](#vib:ode2:staggered:veq)) can then be written in operator notation as

$$
\begin{equation}
\lbrack D_tv = m^{-1}\left(F(t) - f(v) - s(u)\right)\rbrack^n,
\label{vib:ode2:staggered:dveq} \tag{27}
\end{equation}
$$

$$
\begin{equation}
\lbrack D_t u = v\rbrack^{n+\frac{1}{2}}\thinspace .
\label{vib:ode2:staggered:dueq} \tag{28}
\end{equation}
$$

Written out,

$$
\begin{equation}
\frac{v^{n+\frac{1}{2}} - v^{n-\frac{1}{2}}}{\Delta t}
= m^{-1}\left(F^n - f(v^n) - s(u^n)\right),
\label{vib:ode2:staggered:dveq2} \tag{29}
\end{equation}
$$

$$
\begin{equation}
\frac{u^n - u^{n-1}}{\Delta t} = v^{n+\frac{1}{2}}\thinspace .
\label{vib:ode2:staggered:dueq2} \tag{30}
\end{equation}
$$

With linear damping, $f(v)=bv$, we can use an arithmetic mean
for $f(v^n)$: $f(v^n)\approx = \frac{1}{2}(f(v^{n-\frac{1}{2}}) +
f(v^{n+\frac{1}{2}}))$. The system
([29](#vib:ode2:staggered:dveq2))-([30](#vib:ode2:staggered:dueq2))
can then be solved with respect to the unknowns $u^n$ and $v^{n+\frac{1}{2}}$:

$$
\begin{equation}
v^{n+\frac{1}{2}} = \left(1 + \frac{b}{2m}\Delta t\right)^{-1}\left(
v^{n-\frac{1}{2}} + {\Delta t}
m^{-1}\left(F^n - {\frac{1}{2}}f(v^{n-\frac{1}{2}}) - s(u^n)\right)\right),
\label{vib:ode2:staggered:v:scheme:lin} \tag{31}
\end{equation}
$$

$$
\begin{equation}
u^n  = u^{n-1} + {\Delta t}v^{n-\frac{1}{2}}\thinspace .
\label{vib:ode2:staggered:u:scheme:lin} \tag{32}
\end{equation}
$$

In case of quadratic damping, $f(v)=b|v|v$, we can use a geometric mean:
$f(v^n)\approx b|v^{n-\frac{1}{2}}|v^{n+\frac{1}{2}}$. Inserting this approximation
in ([29](#vib:ode2:staggered:dveq2))-([30](#vib:ode2:staggered:dueq2)) and
solving for the unknowns $u^n$ and $v^{n+\frac{1}{2}}$ results in

$$
\begin{equation}
v^{n+\frac{1}{2}} = (1 + \frac{b}{m}|v^{n-\frac{1}{2}}|\Delta t)^{-1}\left(
v^{n-\frac{1}{2}} + {\Delta t}
m^{-1}\left(F^n - s(u^n)\right)\right),
\label{vib:ode2:staggered:v:scheme:quad} \tag{33}
\end{equation}
$$

$$
\begin{equation}
u^n  = u^{n-1} + {\Delta t}v^{n-\frac{1}{2}}\thinspace .
\label{vib:ode2:staggered:u:scheme:quad} \tag{34}
\end{equation}
$$

The initial conditions are derived at the end of
the section [vib:model2x2:staggered](#vib:model2x2:staggered):

$$
\begin{equation}
u^0 = I,
\label{vib:ode2:staggered:u02} \tag{35}
\end{equation}
$$

$$
\begin{equation}
v^\frac{1}{2} = V - \frac{1}{2}\Delta t\omega^2I
\label{vib:ode2:staggered:v02} \tag{36}\thinspace .
\end{equation}
$$

[hpl 2: Must deal with two things: how to set initial conditions on $v$, and
how to make phase-space plots with $v$ at integer mesh points. See
article.]

## The PEFRL 4th-order accurate algorithm

A variant of the Euler-Cromer type of algorithm, which provides an
error $\Oof{\Delta t^4}$ if $f(v)=0$, is called PEFRL
[[PEFRL_2002]](#PEFRL_2002). This algorithm is very well suited for integrating
dynamic systems (especially those without damping) over very long time
periods. Define

$$
g(u,v) = \frac{1}{m}(F(t)-s(u)-f(v))\thinspace .
$$

The algorithm is explicit and features these steps:

$$
\begin{equation}
u^{n+1,1} = u^n       + \xi\Delta t v^n,
\label{_auto7} \tag{37}
\end{equation}
$$

$$
\begin{equation}
v^{n+1,1} = v^n       + \frac{1}{2}(1-2\lambda)\Delta t g(u^{n+1,1},v^n),
\label{_auto8} \tag{38}
\end{equation}
$$

$$
\begin{equation}
u^{n+1,2} = u^{n+1,1} + \chi\Delta t v^{n+1,1},
\label{_auto9} \tag{39}
\end{equation}
$$

$$
\begin{equation}
v^{n+1,2} = v^{n+1,1} + \lambda\Delta t g(u^{n+1,2}, v^{n+1,1}),
\label{_auto10} \tag{40}
\end{equation}
$$

$$
\begin{equation}
u^{n+1,3} = u^{n+1,2} + (1-2(\chi + \xi))\Delta t v^{n+1,2},
\label{_auto11} \tag{41}
\end{equation}
$$

$$
\begin{equation}
v^{n+1,3} = v^{n+1,2} + \lambda\Delta t g(u^{n+1,3}, v^{n+1,2}),
\label{_auto12} \tag{42}
\end{equation}
$$

$$
\begin{equation}
u^{n+1,4} = u^{n+1,3} + \chi\Delta t v^{n+1,3},
\label{_auto13} \tag{43}
\end{equation}
$$

$$
\begin{equation}
v^{n+1}   = v^{n+1,3} + \frac{1}{2}(1-2\lambda)\Delta t g(u^{n+1,4},v^{n+1,3}),
\label{_auto14} \tag{44}
\end{equation}
$$

$$
\begin{equation}
u^{n+1}   = u^{n+1,4} + \xi\Delta t v^{n+1}
\label{_auto15} \tag{45}
\end{equation}
$$

The parameters $\xi$, $\lambda$, and $\xi$ have the values

$$
\begin{equation}
\xi = 0.1786178958448091,
\label{_auto16} \tag{46}
\end{equation}
$$

$$
\begin{equation}
\lambda = -0.2123418310626054,
\label{_auto17} \tag{47}
\end{equation}
$$

$$
\begin{equation}
\chi = -0.06626458266981849
\label{_auto18} \tag{48}
\end{equation}
$$

# Exercises and Problems

## Exercise 1: Implement the solver via classes

Reimplement the `vib.py` program using a class `Problem` to hold all
the physical parameters of the problem, a class `Solver` to hold the
numerical parameters and compute the solution, and a class
`Visualizer` to display the solution.

**Hint.**
Use the ideas and examples from Sections 5.5.1 and 5.5.2
in [[Langtangen_decay]](#Langtangen_decay).  More specifically, make a superclass
`Problem` for holding the scalar physical parameters of a problem and
let subclasses implement the $s(u)$ and $F(t)$ functions as methods.
Try to call up as much existing functionality in `vib.py` as possible.

**Solution.**
The complete code looks like this.

```python
# Reimplementation of vib.py using classes

import numpy as np
import scitools.std as plt
import sympy as sym
from vib import solver as vib_solver
from vib import visualize as vib_visualize
from vib import visualize_front as vib_visualize_front
from vib import visualize_front_ascii as vib_visualize_front_ascii
from vib import plot_empirical_freq_and_amplitude as \
                            vib_plot_empirical_freq_and_amplitude

class Vibration:
    '''
    Problem: m*u'' + f(u') + s(u) = F(t) for t in (0,T],
    u(0)=I and u'(0)=V. The problem is solved
    by a central finite difference method with time step dt.
    If damping is 'linear', f(u')=b*u, while if damping is
    'quadratic', f(u')=b*u'*abs(u'). Zero damping is achieved
    with b=0. F(t) and s(u) are Python functions.
    '''
    def __init__(self, I=1, V=0, m=1, b=0, damping='linear'):
        self.I = I; self.V = V; self.m = m; self.b=b;
        self.damping = damping
    def s(self, u):
        return u
    def F(self, t):
        '''Driving force. Zero implies free oscillations'''
        return 0

class Free_vibrations(Vibration):
    '''F(t) = 0'''
    def __init__(self, s=None, I=1, V=0, m=1, b=0, damping='linear'):
        Vibration.__init__(self, I=I, V=V, m=m, b=b, damping=damping)
        if s != None:
            self.s = s

class Forced_vibrations(Vibration):
    '''F(t)! = 0'''
    def __init__(self, F, s=None, I=1, V=0, m=1, b=0,
                 damping='linear'):
        Vibration.__init__(self, I=I, V=V, m=m, b=b,
                           damping=damping)
        if s != None:
            self.s = s
        self.F = F

class Solver:
    def __init__(self, dt=0.05, T=20):
        self.dt = dt; self.T = T

    def solve(self, problem):
        self.u, self.t = vib_solver(
            problem.I, problem.V,
            problem.m, problem.b,
            problem.s, problem.F,
            self.dt, self.T, problem.damping)

class Visualizer:
    def __init__(self, problem, solver, window_width, savefig):
        self.problem = problem; self.solver = solver
        self.window_width = window_width; self.savefig = savefig
    def visualize(self):
        u = self.solver.u; t = self.solver.t   # short forms
        num_periods = vib_plot_empirical_freq_and_amplitude(u, t)
        if num_periods <= 40:
            plt.figure()
            vib_visualize(u, t)
        else:
            vib_visualize_front(u, t, self.window_width, self.savefig)
            vib_visualize_front_ascii(u, t)
        plt.show()

def main():
    # Note: the reading of parameter values would better be done
    # from each relevant class, i.e. class Problem should read I, V,
    #  etc., while class Solver should read dt and T, and so on.
    # Consult, e.g., Langtangen: "A Primer on Scientific Programming",
    # App E.
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--I', type=float, default=1.0)
    parser.add_argument('--V', type=float, default=0.0)
    parser.add_argument('--m', type=float, default=1.0)
    parser.add_argument('--b', type=float, default=0.0)
    parser.add_argument('--s', type=str, default=None)
    parser.add_argument('--F', type=str, default='0')
    parser.add_argument('--dt', type=float, default=0.05)
    parser.add_argument('--T', type=float, default=20)
    parser.add_argument('--window_width', type=float, default=30.,
                        help='Number of periods in a window')
    parser.add_argument('--damping', type=str, default='linear')
    parser.add_argument('--savefig', action='store_true')
    # Hack to allow --SCITOOLS options
    # (scitools.std reads this argument at import)
    parser.add_argument('--SCITOOLS_easyviz_backend',
                        default='matplotlib')
    a = parser.parse_args()

    from scitools.std import StringFunction
    if a.s != None:
        s = StringFunction(a.s, independent_variable='u')
    else:
        s = None
    F = StringFunction(a.F, independent_variable='t')

    if a.F == '0':  # free vibrations
        problem = Free_vibrations(s=s, I=a.I, V=a.V, m=a.m, b=a.b,
                                  damping=a.damping)
    else:   # forced vibrations
        problem = Forced_vibrations(lambda t: np.sin(t),
                                    s=s, I=a.I, V=a.V,
                                    m=a.m, b=a.b, damping=a.damping)

    solver = Solver(dt=a.dt, T=a.T)
    solver.solve(problem)

    visualizer = Visualizer(problem, solver,
                            a.window_width, a.savefig)
    visualizer.visualize()

if __name__ == '__main__':
    main()
```

Filename: `vib_class`.

## Problem 2: Use a backward difference for the damping term

As an alternative to discretizing the damping terms $\beta u^{\prime}$ and
$\beta |u^{\prime}|u^{\prime}$ by centered differences, we may apply
backward differences:

$$
\begin{align*}
[u^{\prime}]^n &\approx [D_t^-u]^n,\\
[|u^{\prime}|u^{\prime}]^n &\approx [|D_t^-u|D_t^-u]^n\\
&= |[D_t^-u]^n|[D_t^-u]^n\thinspace .
\end{align*}
$$

The advantage of the backward difference is that the damping term is
evaluated using known values $u^n$ and $u^{n-1}$ only.
Extend the [`vib.py`](${src_vib}/vib.py) code with a scheme based
on using backward differences in the damping terms. Add statements
to compare the original approach with centered difference and the
new idea launched in this exercise. Perform numerical experiments
to investigate how much accuracy that is lost by using the backward
differences.

**Solution.**
The new discretization approach of the linear and quadratic damping
terms calls for new derivations of the updating formulas (for $u$) in
the solver. Since backward difference approximations will be used for
the damping term, we may also use this approximation for the initial
condition on $u^{\prime}(0)$ without deteriorating the convergence
rate any further. Note that introducing backward difference
approximations for the damping term make our computational schemes
first order, as opposed to the original second order schemes which
used central difference approximations also for the damping terms. The
motivation for also using a backward difference approximation for the
initial condition on $u^{\prime}(0)$, is simply that the computational
schemes get much simpler.

With linear damping, the new discretized form of the equation reads

$$
m\frac{u^{n+1}-2u^n+u^{n-1}}{\Delta t^2} + b\frac{u^n - u^{n-1}}{\Delta t} + s(u^n) = F^n,
\nonumber
$$

which gives us

$$
u^{n+1} = \left(2-\frac{\Delta t b}{m}\right)u^n + \frac{\Delta t^2}{m}\left(F^n - s(u^n)\right) + \left(\frac{\Delta t b}{m} - 1\right)u^{n-1}.
\nonumber
$$

With $n=0$, the updating formula becomes

$$
u^1 = \left(2-\frac{\Delta t b}{m}\right)u^0 + \frac{\Delta t^2}{m}\left(F^0 - s(u^0)\right) + \left(\frac{\Delta t b}{m} - 1\right)u^{-1},
\nonumber
$$

which requires some further elaboration because of the unknown $u^{-1}$. We handle this by discretizing the initial condition $u^{\prime}(0) = V$ by a backward difference approximation as

$$
\frac{u^0 - u^{-1}}{\Delta t} = V,
\nonumber
$$

which implies that

$$
u^{-1} = u^0 - \Delta t V.
\nonumber
$$

Inserting this expression for $u^{-1}$ in the updating formula for $u^{n+1}$, and simplifying,
gives us the following special formula for the first time step:

$$
u^1 = u^0 + \Delta t V + \frac{\Delta t^2}{m}\left(-bV - s(u^0) + F^0\right).
\nonumber
$$

Switching to quadratic damping, the new discretized form of the equations becomes

$$
m\frac{u^{n+1}-2u^n+u^{n-1}}{\Delta t^2} + b|\frac{u^n - u^{n-1}}{\Delta t}|\frac{u^n - u^{n-1}}{\Delta t} + s(u^n) = F^n,
\nonumber
$$

which leads to

$$
u^{n+1} = 2u^n - u^{n-1} - \frac{b}{m}|u^n - u^{n-1}|(u^n - u^{n-1}) + \frac{\Delta t^2}{m}\left(F^n - s(u^n)\right).
\nonumber
$$

With $n=0$, this updating formula becomes

$$
u^1 = 2u^0 - u^{-1} - \frac{b}{m}|u^0 - u^{-1}|(u^0 - u^{-1}) + \frac{\Delta t^2}{m}\left(F^0 - s(u^0)\right).
\nonumber
$$

Again, we handle the unknown $u^{-1}$ via the same expression as above, which be derived from a backward difference approximation to the initial condition on the derivative. Inserting this expression for $u^{-1}$ and simplifying, gives the special updating formula for $u^1$ as

$$
u^1 = u^0 + \Delta t V + \frac{\Delta t^2}{m}\left(-b|V|V - s(u^0) + F^0\right).
\nonumber
$$

We implement these new computational schemes in a new solver function
`solver_bwdamping`, so that the discrete solution for $u$ can be found
by both the original and the new solver. The difference between the
two different solutions is then visualized in the same way as the
original solution in `main`.

The convergence rates computed in `test_mms` demonstrates that our
scheme now is a first order scheme, as $r$ is seen to approach 1.0
with decreasing $\Delta t$.

Both solvers reproduce a constant solution exactly (within machine
precision), whereas sinusoidal and quadratic solutions differ, as
should be expected after comparing the schemes. Pointing out the
"best" approach is difficult: the backward differences yield a much
simpler mathematical problem to be solved, while the more complicated
method converges faster and gives more accuracy for the same cost. On
the other hand, the backward differences can yield any reasonable
accuracy by lowering $\Delta t$, and the results are obtained within a
few seconds on a laptop.

Here is the complete computer code, arising from copying `vib.py` and changing
the functions that have to be changed:

```python
import numpy as np
#import matplotlib.pyplot as plt
import scitools.std as plt

def solver_bwdamping(I, V, m, b, s, F, dt, T, damping='linear'):
    """
    Solve m*u'' + f(u') + s(u) = F(t) for t in (0,T],
    u(0)=I and u'(0)=V. All terms except damping is discretized
    by a central finite difference method with time step dt.
    The damping term is discretized by a backward diff. approx.,
    as is the init.cond. u'(0). If damping is 'linear', f(u')=b*u,
    while if damping is 'quadratic', f(u')=b*u'*abs(u').
    F(t) and s(u) are Python functions.
    """
    dt = float(dt); b = float(b); m = float(m) # avoid integer div.
    Nt = int(round(T/dt))
    u = np.zeros(Nt+1)
    t = np.linspace(0, Nt*dt, Nt+1)

    u_original = np.zeros(Nt+1); u_original[0] = I  # for testing

    u[0] = I
    if damping == 'linear':
        u[1] = u[0] + dt*V + dt**2/m*(-b*V - s(u[0]) + F(t[0]))
    elif damping == 'quadratic':
        u[1] = u[0] + dt*V + \
               dt**2/m*(-b*V*abs(V) - s(u[0]) + F(t[0]))
    for n in range(1, Nt):
        if damping == 'linear':
            u[n+1] = (2 - dt*b/m)*u[n] + dt**2/m*(- s(u[n]) + \
                                F(t[n])) + (dt*b/m - 1)*u[n-1]
        elif damping == 'quadratic':
            u[n+1] = 2*u[n] - u[n-1] - b/m*abs(u[n] - \
            u[n-1])*(u[n] - u[n-1]) + dt**2/m*(-s(u[n]) + F(t[n]))
    return u, t

import sympy as sym

def test_constant():
    """Verify a constant solution."""
    u_exact = lambda t: I
    I = 1.2; V = 0; m = 2; b = 0.9
    w = 1.5
    s = lambda u: w**2*u
    F = lambda t: w**2*u_exact(t)
    dt = 0.2
    T = 2
    #u, t = solver(I, V, m, b, s, F, dt, T, 'linear')
    u, t = solver_bwdamping(I, V, m, b, s, F, dt, T, 'linear')
    difference = np.abs(u_exact(t) - u).max()
    print difference
    tol = 1E-13
    assert difference < tol

    #u, t = solver(I, V, m, b, s, F, dt, T, 'quadratic')
    u, t = solver_bwdamping(I, V, m, b, s, F, dt, T, 'quadratic')
    difference = np.abs(u_exact(t) - u).max()
    print difference
    assert difference < tol

def lhs_eq(t, m, b, s, u, damping='linear'):
    """Return lhs of differential equation as sympy expression."""
    v = sym.diff(u, t)
    if damping == 'linear':
        return m*sym.diff(u, t, t) + b*v + s(u)
    else:
        return m*sym.diff(u, t, t) + b*v*sym.Abs(v) + s(u)

def test_quadratic():
    """Verify a quadratic solution."""
    I = 1.2; V = 3; m = 2; b = 0.9
    s = lambda u: 4*u
    t = sym.Symbol('t')
    dt = 0.2
    T = 2

    q = 2  # arbitrary constant
    u_exact = I + V*t + q*t**2
    F = sym.lambdify(t, lhs_eq(t, m, b, s, u_exact, 'linear'))
    u_exact = sym.lambdify(t, u_exact, modules='numpy')
    #u1, t1 = solver(I, V, m, b, s, F, dt, T, 'linear')
    u1, t1 = solver_bwdamping(I, V, m, b, s, F, dt, T, 'linear')
    diff = np.abs(u_exact(t1) - u1).max()
    print diff
    tol = 1E-13
    #assert diff < tol

    # In the quadratic damping case, u_exact must be linear
    # in order to exactly recover this solution
    u_exact = I + V*t
    F = sym.lambdify(t, lhs_eq(t, m, b, s, u_exact, 'quadratic'))
    u_exact = sym.lambdify(t, u_exact, modules='numpy')
    #u2, t2 = solver(I, V, m, b, s, F, dt, T, 'quadratic')
    u2, t2 = solver_bwdamping(I, V, m, b, s, F, dt, T, 'quadratic')
    diff = np.abs(u_exact(t2) - u2).max()
    print diff
    assert diff < tol

def test_sinusoidal():
    """Verify a numerically exact sinusoidal solution when b=F=0."""
    from math import asin

    def u_exact(t):
        w_numerical = 2/dt*np.arcsin(w*dt/2)
        return I*np.cos(w_numerical*t)

    I = 1.2; V = 0; m = 2; b = 0
    w = 1.5  # fix the frequency
    s = lambda u: m*w**2*u
    F = lambda t: 0
    dt = 0.2
    T = 6
    #u, t = solver(I, V, m, b, s, F, dt, T, 'linear')
    u, t = solver_bwdamping(I, V, m, b, s, F, dt, T, 'linear')
    diff = np.abs(u_exact(t) - u).max()
    print diff
    tol = 1E-14
    #assert diff < tol

    #u, t = solver(I, V, m, b, s, F, dt, T, 'quadratic')
    u, t = solver_bwdamping(I, V, m, b, s, F, dt, T, 'quadratic')
    print diff
    diff = np.abs(u_exact(t) - u).max()
    assert diff < tol

def test_mms():
    """Use method of manufactured solutions."""
    m = 4.; b = 1
    w = 1.5
    t = sym.Symbol('t')
    u_exact = 3*sym.exp(-0.2*t)*sym.cos(1.2*t)
    I = u_exact.subs(t, 0).evalf()
    V = sym.diff(u_exact, t).subs(t, 0).evalf()
    u_exact_py = sym.lambdify(t, u_exact, modules='numpy')
    s = lambda u: u**3
    dt = 0.2
    T = 6
    errors_linear = []
    errors_quadratic = []
    # Run grid refinements and compute exact error
    for i in range(5):
        F_formula = lhs_eq(t, m, b, s, u_exact, 'linear')
        F = sym.lambdify(t, F_formula)
        #u1, t1 = solver(I, V, m, b, s, F, dt, T, 'linear')
        u1, t1 = solver_bwdamping(I, V, m, b, s,
                                  F, dt, T, 'linear')
        error = np.sqrt(np.sum((u_exact_py(t1) - u1)**2)*dt)
        errors_linear.append((dt, error))

        F_formula = lhs_eq(t, m, b, s, u_exact, 'quadratic')
        #print sym.latex(F_formula, mode='plain')
        F = sym.lambdify(t, F_formula)
        #u2, t2 = solver(I, V, m, b, s, F, dt, T, 'quadratic')
        u2, t2 = solver_bwdamping(I, V, m, b, s,
                                  F, dt, T, 'quadratic')
        error = np.sqrt(np.sum((u_exact_py(t2) - u2)**2)*dt)
        errors_quadratic.append((dt, error))
        dt /= 2
    # Estimate convergence rates
    tol = 0.05
    for errors in errors_linear, errors_quadratic:
        for i in range(1, len(errors)):
            dt, error = errors[i]
            dt_1, error_1 = errors[i-1]
            r = np.log(error/error_1)/np.log(dt/dt_1)
        # check r for final simulation with (final and) smallest dt
        # note that the method now is 1st order, i.e. r should
        # approach 1.0
        print r
        assert abs(r - 1.0) < tol

import os, sys
sys.path.insert(0, os.path.join(os.pardir, 'src-vib'))
from vib import (plot_empirical_freq_and_amplitude,
                 visualize_front, visualize_front_ascii,
                 minmax, periods, amplitudes,
                 solver as solver2)

def visualize(list_of_curves, legends, title='', filename='tmp'):
    """Plot list of curves: (u, t)."""
    for u, t in list_of_curves:
        plt.plot(t, u)
        plt.hold('on')
    plt.legend(legends)
    plt.xlabel('t')
    plt.ylabel('u')
    plt.title(title)
    plt.savefig(filename + '.png')
    plt.savefig(filename + '.pdf')
    plt.show()

def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--I', type=float, default=1.0)
    parser.add_argument('--V', type=float, default=0.0)
    parser.add_argument('--m', type=float, default=1.0)
    parser.add_argument('--b', type=float, default=0.0)
    parser.add_argument('--s', type=str, default='4*pi**2*u')
    parser.add_argument('--F', type=str, default='0')
    parser.add_argument('--dt', type=float, default=0.05)
    parser.add_argument('--T', type=float, default=20)
    parser.add_argument('--window_width', type=float, default=30.,
                        help='Number of periods in a window')
    parser.add_argument('--damping', type=str, default='linear')
    parser.add_argument('--savefig', action='store_true')
    # Hack to allow --SCITOOLS options
    # (scitools.std reads this argument at import)
    parser.add_argument('--SCITOOLS_easyviz_backend',
                        default='matplotlib')
    a = parser.parse_args()
    from scitools.std import StringFunction
    s = StringFunction(a.s, independent_variable='u')
    F = StringFunction(a.F, independent_variable='t')
    I, V, m, b, dt, T, window_width, savefig, damping = \
       a.I, a.V, a.m, a.b, a.dt, a.T, a.window_width, a.savefig, \
       a.damping

    # compute u by both methods and then visualize the difference
    u, t = solver2(I, V, m, b, s, F, dt, T, damping)
    u_bw, _ = solver_bwdamping(I, V, m, b, s, F, dt, T, damping)
    u_diff = u - u_bw

    num_periods = plot_empirical_freq_and_amplitude(u_diff, t)
    if num_periods <= 40:
        plt.figure()
        legends = ['1st-2nd order method',
                   '2nd order method',
                   '1st order method']
        visualize([(u_diff, t), (u, t), (u_bw, t)], legends)
    else:
        visualize_front(u_diff, t, window_width, savefig)
        #visualize_front_ascii(u_diff, t)
    plt.show()

if __name__ == '__main__':
    main()
    #test_constant()
    #test_sinusoidal()
    #test_mms()
    #test_quadratic()
    raw_input()
```

Here is a comparison of standard method (2nd order) and backward differences for damping (1st order) for 10 (left) and 40 (right) time steps per period:

Filename: `vib_gen_bwdamping`.

## Exercise 3: Use the forward-backward scheme with quadratic damping

We consider the generalized model with quadratic damping, expressed
as a system of two first-order equations as in the section [A staggered Euler-Cromer scheme for a generalized model](#vib:ode2:staggered):

$$
\begin{align*}
u^{\prime} &= v,\\
v' &= \frac{1}{m}\left( F(t) - \beta |v|v - s(u)\right)\thinspace .
\end{align*}
$$

However, contrary to what is done in the section [A staggered Euler-Cromer scheme for a generalized model](#vib:ode2:staggered),
we want to apply the idea of a forward-backward discretization:
$u$ is marched forward by a one-sided Forward Euler scheme applied
to the first equation, and
thereafter $v$ can be marched forward by a Backward Euler scheme in the
second
% if BOOK == "book":
equation, see in the section [vib:model2x2:EulerCromer](#vib:model2x2:EulerCromer).
% else:
equation.
% endif
Express the idea in operator notation and write out the
scheme. Unfortunately, the backward difference for the $v$ equation
creates a nonlinearity $|v^{n+1}|v^{n+1}$.  To linearize this
nonlinearity, use the known value $v^n$ inside the absolute value
factor, i.e., $|v^{n+1}|v^{n+1}\approx |v^n|v^{n+1}$.  Show that the
resulting scheme is equivalent to the one in the section [A staggered Euler-Cromer scheme for a generalized model](#vib:ode2:staggered) for some time level $n\geq 1$.

What we learn from this exercise is that the first-order differences
and the linearization trick play together in "the right way" such that
the scheme is as good as when we (in the section [A staggered Euler-Cromer scheme for a generalized model](#vib:ode2:staggered))
carefully apply centered differences and a geometric mean on a
staggered mesh to achieve second-order accuracy.
% if BOOK == "book":
There is a
difference in the handling of the initial conditions, though, as
explained at the end of the section [vib:model2x2:EulerCromer](#vib:model2x2:EulerCromer).
% endif
Filename: `vib_gen_bwdamping`.


# 4/ Problem solutions

```python
# helper code needed for running in colab
if 'google.colab' in str(get_ipython()):
    print('Downloading plot_helpers.py to util/ (only neded for colab')
    !mkdir util; wget https://raw.githubusercontent.com/minireference/noBSLAnotebooks/master/util/plot_helpers.py -P util
```

```python
# setup SymPy
from sympy import *
init_printing()

# setup plotting
%matplotlib inline
import matplotlib.pyplot as mpl
from util.plot_helpers import plot_plane, plot_line, plot_vec, plot_vecs

# aliases
Vector = Matrix  # define alias Vector so I don't have to explain this during video
Point = Vector   # define alias Point for Vector since they're the same thing
```

## P4.2

Find the lines of intersection between these pairs of planes:
**a)** $P_1$: $3x-2y-z=2$ and $P_2$: $x+2y+z=0$,
**b)** $P_3$: $2x+y-z=0$ and $P_4$: $x+2y+z=3$.

<!--
		\begin{answer}\textbf{a)}~;
					\textbf{b)}~.\end{answer}
		%	A = Matrix([[1,-2,-1,2],[1,2,1,0]])
		%	A.rref()
		%
-->

```python
# a)         x   y   z  | c
A = Matrix([[3, -2, -1,   2],
            [1,  2,  1,   0]])
A.rref()
```

So $z=s$ is a free variable, and the rest of the equation
can be written as

$$
  \begin{array}{rl}
  x                 &= \frac{1}{2}\\
  y + \frac{1}{2}s  &= -\frac{1}{4}\\
  z                 &= s
  \end{array}
$$

The answer is $(x,y,z) = (\frac{1}{2},-\frac{1}{4},0) + s(0,-\frac{1}{2},1),  \forall s \in \mathbb{R}$.

```python
# b)         x   y   z  | c
B = Matrix([[2,  1, -1,   0],
            [1,  2,  1,   3]])
B.rref()
```

The free variable is $z=t$.
The answer to b) is $\{ (-1,2,0) + t(1,-1,1), \forall t \in \mathbb{R}\}$.

## P4.11

```python
v = Vector([3, 4, 1])
normal = Vector([2, -1, 4])

vPperp = (normal.dot(v)/normal.norm()**2)*normal
print('vPperp =', vPperp)
vP = v - vPperp
print('vP =', vP)

plot_plane(normal, 0)            # plane P
plot_vec(0.2*normal, color='r')  # its normal vec

plot_vecs(v, vPperp, vP)
```

```python
v = Vector([3,   4,  1])
normal = Vector([2,   -1,  4])
D=4
# point on P closest to the origin
alpha = D/normal.norm()**2
p_closest = alpha*normal
# print('len normal', normal.norm())
# print('p_closest', p_closest)
assert p_closest.dot(normal) == 4

vPperp = (normal.dot(v)/normal.norm()**2)*normal
print('vPperp', vPperp)
v_wrong = v - vPperp
print('v_wrong', v_wrong)

plot_plane(normal, D)                          # plane P
plot_vec(0.2*normal, at=p_closest, color='r')  # its normal vec

plot_vecs(v, vPperp, v_wrong)
ax = mpl.gca()
ax.grid(True,which='both')
```

```python
v = Vector([3,   4,  1])
normal = Vector([2,   -1,  4])
D = 4

# some point on P
p0 = Point([2,0,0])

u = v - p0  # vector from p0 to tip of v

uPperp = (normal.dot(u)/normal.norm()**2)*normal
print('uPperp', uPperp)
uInP = u - uPperp

proj_v_on_P = p0 + uInP
print('proj_v_on_P', proj_v_on_P)

plot_plane(normal, D)                          # plane P
plot_vec(0.2*normal, at=p_closest, color='r')  # its normal vec

plot_vec(v)
plot_vec(u, at=p0, color='r')
plot_vec(uPperp, at=p0, color='b')
plot_vec(uInP, at=p0, color='g')
plot_vec(proj_v_on_P, color='y')

ax = mpl.gca()
ax.grid(True,which='both')
```

```python

```


## Cournot Duopoly

Consider two firms competing to maximize their profits by choosing a quantity, $q_i$ to produce for $i = 1, 2$.  Assume the firms have identical costs of production $C_i(q_i) = c q_i$ where $c > 0$.  The two firms face market inverse demand function,

\begin{align}
    p(Q) &= a - b Q \\
    &= a - b(q_1 + q_2)
\end{align}

Each firm must simultaneously choose a strategy $q_i \in [0, \infty)$ given a production choice of their competitor, $q_j$ to maximize their private profits.  The payoff to any choice of firm $i$ is,

\begin{align}
    \pi(q_i | q_j) &= (a - bq_i - bq_j)q_i - cq_i
\end{align}

```python
from sympy import *
from sympy.plotting import plot
init_printing()
```

```python
# Set up the parameters and functions
# for our optimization problem.
a, b, c, q1, q2 = symbols('a b c q_1 q_2',
                          real=True,
                          positive=True,
                          finite=True)

price = a - b * (q1 + q2)
profit1 = price * q1 - c * q1
profit2 = price * q2 - c * q2
```

For the firms to choose a quantity to maximize their profits given the choice of their competitor, each firm will have to find a $q$ that satisfies both a first and second order condition.  Before we start taking derivatives lets check a few properties of the first and second derivatives.

```python
# Take the first and second derivative of either
# profit1 or profit2.
test_deriv2 = profit1.diff(q1).diff(q1)

# What can we tell about the signs of these
# derivatives?
test_deriv2.is_negative
```

    True

So we see that the second derivative of either firm's payoff function is negative, for all $q \in [0, \infty)$.  This means that the function is globally concave in $q$ and the first order condition is both **necessary** and **sufficient** for a maximum.  In other words, the only critical points are maxima.  So if we find one, then we are good to go.

```python
# Take the derivatives
Dprofit1 = profit1.diff(q1)
Dprofit2 = profit2.diff(q2)

# Set the derivatives equal to
# zero and solve for the firm's
# quantity to get their best
# response function.
BR1 = solve(Dprofit1, q1)[0]
BR2 = solve(Dprofit2, q2)[0]

Eq(Symbol('q_1^*'), BR1), Eq(Symbol('q_2^*'), BR2)
```

Lets momentarily assume some parameter values and graphically depict these best response functions.  We will assume the following values:

| Parameter | Value |
| :---: | :---: |
| $a$ | 20 |
| $b$ | 1.5 |
| $c$ | 1 |

```python
# Define generic 'q' symbol
q = Symbol('q')

# Invert q1 BR
BR1_invert = solve(q1 - BR1, q2)[0]

vals = {a: 20,
        b: 1.5,
        c: 1,
        q1: q,
        q2: q}

p1 = plot(BR1_invert.subs(vals),
          BR2.subs(vals),
          (q, 0, 10),
          title="Best Response Functions",
          xlabel="$q_1$",
          ylabel="$q_2$",
          ylim=(0,13))
p1
```

## Nash Equilibrium

Great! So we found the profit maximizing strategy for each firm and have characterized all possible optimal strategies by the firm's best response function.  We can substitute these into each other to obtain the **Nash Equilibrium** values.

```python
# Substitute firm 1's optimal response into
# firm 2's optimal response.
equil_vals_2 = BR2.subs(q1, BR1)

# Now collect q2 on both sides and solve
# for q2 in terms of parameters.
q2_ne = solve(q2 - equil_vals_2, q2)[0]

# Now substitute the Nash equilibrium
# value of q2 into the best response
# of firm 1.
equil_vals_1 = BR1.subs(q2, q2_ne)

# Now collect q1 on both sides and solve
# for q1 in terms of parameters.
q1_ne = solve(q1 - equil_vals_1, q1)[0]

# Display the values
Eq(Symbol('q_1^{ne}'), q1_ne), Eq(Symbol('q_2^{ne}'), q2_ne)
```

```python
# Equilibrium Output
Q_ne = simplify(q1_ne + q2_ne)
P_ne = simplify(price.subs({q1: q1_ne, q2: q2_ne}))

Eq(Symbol('Q^{ne}'), Q_ne), Eq(Symbol('p^{ne}'), P_ne)
```

```python

```


```
"""
Created on Fri Nov 01 17:09:59 2019

@author: andresryes
"""

#Imports
import numpy as np
from numpy import linalg
import scipy as sp
from scipy import optimize
import math
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d
from sympy import symbols, diff

#we receive a 2x2 matrix
A = np.matrix([[2, 1],[1,2]])
B = np.matrix([[1, 0],[0,2]])
C = np.matrix([[5, 1],[1,4]])

A = B
#vector b 2x1
b = np.matrix([[3], [4]])
u,v = symbols('u v', real=True)
print("A:")
print(A)
print("b:")
print(b)

#function
#f = A.A1[3]*v
#print(diff(f, v))

#print(diff(f1, v))
#print(gradientef([1,1]))
#using the linalg library
print("exact solution using numpy:")
print(linalg.solve(A,b))
# F(x) = (1/2)x · Ax − b · x

#Steepest Descent
#this is the function we are going to optimize, based on the original matrix F(x) = 1/2 x^t · Ax − b · x^t
def f(x):
    u = x[0]
    v = x[1]
    return 0.5*(u*(u*A.A1[0] + v*A.A1[2]) + v*(u*A.A1[1] + v*A.A1[3])) - b.A1[0]*u - b.A1[1]*v

#the gradient of the previous function, it indicates the direction
def gradientef(x):
    u,v = symbols('u v', real=True)
    f1 = 0.5*(u*(u*A.A1[0] + v*A.A1[2]) + v*(u*A.A1[1] + v*A.A1[3])) - b.A1[0]*u - b.A1[1]*v
    du = str(diff(f1, u))
    dv = str(diff(f1, v))
    u = x[0]
    v = x[1]
    y=[eval(du),eval(dv)]
    return y

#alpha function of one variable, indicates the step in each iteration/jump
def funcion(alpha,x,s):
    u = x[0]+alpha*s[0]
    v = x[1]+alpha*s[1]
    y1 = 0.5*(u*(u*A.A1[0] + v*A.A1[2]) + v*(u*A.A1[1] + v*A.A1[3])) - b.A1[0]*u - b.A1[1]*v
    return y1

#this is the method itself, receives xk first guess, the function, tolerance
def desc(xk,f,tol):
    #initial error
    error=1
    #this is the array with all the dots to graphic
    X=[]
    #array with all vectors
    S=[]
    #append to the dots array the initial guess
    X.append(xk)
    #make this algorithm while the error is still bigger than the tolerance
    #level
    while(error>tol):
        #calculates the first direction of the steepest descent using the
        #negative of the gradient
        sk=np.dot(-1.0,gradientef(xk))
        #optimizes alpha funciton using Golden Ration optimization of Python
        ak1=sp.optimize.golden(funcion,args=(xk,sk))
        #getting the direction using the step size
        sk_n=np.dot(ak1,sk)
        #append the direction
        S.append(sk_n)
        #calculating our new xk+1
        xk1=xk+np.dot(ak1,sk)
        #append the new dot
        X.append(xk1)
        #calculating the error
        error=np.linalg.norm(xk1-xk)
        #changing xk to the new value
        xk=xk1
    return xk,X,S

y=desc([10, 15],f,0.0000001)
print("the approx solution is: "+str(y[0]))
Xk=y[1]
Sk=y[2]
print("iterations:")
print(len(Sk))

X1=[]
Y1=[]
Z1=[]

#graphic the coordinates
for i in range(len(Xk)):
    X1.append(Xk[i][0])
    Y1.append(Xk[i][1])

#graphic in 3d
fig = plt.figure()
ax = fig.gca(projection="3d")

#graphing the level curves
u, v = np.meshgrid(np.arange(-15, 15.0, 0.05), np.arange(-15.0, 15.0, 0.05))
fmesh = 0.5*(u*(u*A.A1[0] + v*A.A1[2]) + v*(u*A.A1[1] + v*A.A1[3])) - b.A1[0]*u - b.A1[1]*v
ax.plot_surface(u, v, fmesh)
fig2=plt.figure()
plt.axis("equal")
plt.contour(u, v, fmesh)
plt.show()

#graphing the evolution of the solution

#defining the axis
X_1=np.linspace(-10,10,10)
Y_1=np.linspace(-10,10,10)
plt.axvline(x=0, color= 'k', linestyle = '--')
plt.axhline(y=0, color= 'k', linestyle = '--')

#tags
plt.xlabel("X")
plt.ylabel("Y")

#cuadriculate
plt.grid(True)

#calculate the X, Y, Z for meshgrid
X_2,Y_2=np.meshgrid(X_1,Y_1)

#function
Z_1 = f([X_2, Y_2])
A=np.linspace(1,50,2**5)

#to get the tag 'min' in the solution
plt.annotate('Min', xy=(y[0][0]-0.5,y[0][1]+0.1),color='red')

for i in range(len(Sk)):
    plt.arrow(Xk[i][0],Xk[i][1],Sk[i][0],Sk[i][1],head_width=0.05,head_length=0.05)
    plt.scatter(Xk[i][0],Xk[i][1],color='blue')
plt.contour(X_1,Y_1,Z_1,A)
plt.show()
```

```
#Recibimos una matriz 2x2
A = np.matrix([[1, 0],[0,2]])
#un vector resultado b 2x1
b = np.matrix([[3], [4]])
u,v = symbols('u v', real=True)
print(A)
print(b)

#function
f = A.A1[3]*v
from sympy import symbols, diff
print(diff(f, v))

f1 = 0.5*(u*(u*A.A1[0] + v*A.A1[2]) + v*(u*A.A1[1] + v*A.A1[3])) - b.A1[0]*u - b.A1[1]*v
print(diff(f1, v))

#Respuesta
print(linalg.solve(A,b))
# F(x) = (1/2)x · Ax − b · x
```

    [[1 0]
     [0 2]]
    [[3]
     [4]]
    2
    2.0*v - 4
    [[3.]
     [2.]]

```
#Recibimos una matriz 2x2
A = np.matrix([[2, 1],[1,2]])
#un vector resultado b 2x1
b = np.matrix([[4], [6]])
u,v = symbols('u v', real=True)
print(A)
print(b)
def gradientef(x):
    u,v = symbols('u v', real=True)
    f1 = 0.5*(u*(u*A.A1[0] + v*A.A1[2]) + v*(u*A.A1[1] + v*A.A1[3])) - b.A1[0]*u - b.A1[1]*v
    du = str(diff(f1, u))
    dv = str(diff(f1, v))
    print(du)
    u = x[0]
    v = x[1]
    y=[eval(du),eval(dv)]
    return y

gradientef([3, 2])
```

    [[2 1]
     [1 2]]
    [[4]
     [6]]
    2.0*u + 1.0*v - 4

    [4.0, 1.0]


# Cross Correlation

The general form of the cross correlation with integration:
\begin{equation}
{\displaystyle (f\star g)(\tau )\ = \int _{-\infty }^{\infty }{{f^{\star}(t)}}g(t+\tau )\,dt}
\end{equation}
This can be written in discrete form as:
\begin{equation}
{\displaystyle (f\star g)[n]\ = \sum _{m=-\infty }^{\infty }{{f^{\star}[m]}}g[m+n]}
\end{equation}

```python
%pylab inline
pylab.rcParams['savefig.dpi'] = 300
from gps_helper.prn import PRN
from sk_dsp_comm import sigsys as ss
from sk_dsp_comm import digitalcom as dc
from caf_verilog.quantizer import quantize
```

## Test Signals

```python
prn = PRN(10)
prn2 = PRN(20)
fs = 625e3
Ns = fs / 125e3
prn_seq = prn.prn_seq()
prn_seq2 = prn2.prn_seq()
prn_seq,b = ss.NRZ_bits2(array(prn_seq), Ns)
prn_seq2,b2 = ss.NRZ_bits2(array(prn_seq2), Ns)
```

```python
Px,f = psd(prn_seq, 2**12, Fs=fs)
plot(f, 10*np.log10(Px))
```

## Autocorrelation

```python
r, lags = dc.xcorr(prn_seq, prn_seq, 100)
plot(lags, abs(r)) # r -> abs
```

## Time Shifted Signals

```python
r, lags = dc.xcorr(roll(prn_seq, 50), prn_seq, 100)
plot(lags, abs(r))
```

```python
r, lags = dc.xcorr(roll(prn_seq, -50), prn_seq, 100)
plot(lags, abs(r))
```

## No Correlation

```python
r_nc, lags_nc = dc.xcorr(prn_seq, prn_seq2, 100)
plot(lags_nc, abs(r_nc))
ylim([0, 1])
```

## Calculation Space Visualization

```python
from caf_verilog.xcorr import size_visualization
size_visualization(prn_seq[:10], prn_seq[:10], 5)
```

## Simple Cross Correlation

```python
from caf_verilog.xcorr import simple_xcorr
r, lags = simple_xcorr(prn_seq, prn_seq, 100)
plot(lags, r)
```

### Time Shifted Signals

```python
r, lags = simple_xcorr(prn_seq, roll(prn_seq, 50), 100)
plot(lags, abs(array(r)))
```

```python
r, lags = simple_xcorr(prn_seq, roll(prn_seq, -50), 100)
plot(lags, abs(array(r)))
```

### No Correlation

```python
r, lags = simple_xcorr(prn_seq, prn_seq2, 100)
plot(lags, abs(array(r)))
ylim([0, 5000])
```

## Dot Product Method

To ensure the integration time is filled, the secondary or received signal must be twice the length of the reference signal.

```python
from caf_verilog.sim_helper import sim_shift
center = 300
corr_length = 250
shift = 25
ref, rec = sim_shift(prn_seq, center, corr_length, shift=shift, padding=True)
```

```python
f, axarr = plt.subplots(2, sharex=True, gridspec_kw={'hspace': 0})
axarr[0].plot(ref)
axarr[1].plot(rec)
plt.xlabel("Sample Number")
savefig('prn_seq.png')
```

```python
from caf_verilog.xcorr import dot_xcorr
ref, rec = sim_shift(prn_seq, center, corr_length, shift=shift)
rr = dot_xcorr(ref, rec)
rr = array(rr)
```

```python
rxy, lags = dc.xcorr(ref, rec, 1000)
plot(abs(rxy))
xlabel("Center Offset (Samples)")
grid();
savefig('xcorr.png')
```

```python
fig, axs = plt.subplots(2, sharey=True)
axs[0].plot(abs(rr))
axs[0].grid(True)
axs[1].plot(abs(rr))
axs[1].set_xlim([80, 120])
axs[1].set_xlabel('Inverse Center Offset (Samples)')
axs[1].grid(True)
fig.savefig('xcorr_250.png')
```

```python
argmax(rr)
```

```python
from caf_verilog.xcorr import XCorr
xc = XCorr(ref, rec, output_dir='.')
```

```python
xc.gen_tb()
```


# Messtechnik HS2021 - Tutorial 4

## Aufgabe 1: Diskrete Fourier Transformation und Datenverarbeitung
---------
Ein Kohlenstoff NMR Signal ($^{13}{\rm C}$) wurde aufgenommen und mit einer Abtastrate von 15 kHz digitalisiert. Das komplexwertige FID wurde durch ein MATLAB Program digitalisiert, exportiert und ist in `fid.mat` zu finden.

---------

### 1a)
Generieren Sie aus dem komplexwertigen FID das dazugehörige Spektrum und plotten Sie dessen Realteil mit der korrekten Frequenzachse in kHz.

*Hinweis:* Um MATLAB Dateien in Python zu importieren, kann man die ``scipy.io.loadmat`` Funktion brauchen:

````
from scipy.io import loadmat
data = loadmat('fid.mat')
fid = data['fid']
fid = np.squeeze(fid)
````

```python

```

### 1b)
Durch einen Matched Filter lässt sich das Signal-zu-Rausch-Verhältnis (SNR) nachträglich verbessern. Das Signal kann beschrieben werden als
$$ y(t) = \exp(-i\omega_0 t) \exp(-\frac{t}{T_2}) $$
wobei $\omega_0$ die Larmorfrequenz ist und $T_2$ die transversale Relaxationszeit. Bei einem Matched Filter wird das Zeitsignal $ y(t) $ mit der Zerfallsfunktion $ y_e(t) = \exp(-\frac{t}{T_2}) $ multipliziert.
Zeigen Sie die Auswirkungen auf das Spektrum bei der Anwendung von Matched Filtern mit Zerfallszeiten $T_2 = [1,2,5,10,50,100]$.

```python

```

### 1c)
Finden Sie eine passende Zerfallszeit für den Matched Filter, sodass Sie eine gute Balance zwischen Auflösung und SNR erhalten.

```python

```

### 1d)
Nehmen Sie nun an die Aufnahmezeit des FID's ist beschränkt auf 20 ms. Erstellen Sie aus `fid.mat` das entsprechende abgeschnittene Signal, hängen Sie dann Nullen an (Zerofilling), sodass Sie am Ende ein Signal mit insgesamt 2048 Punkte erhalten. Was fällt Ihnen auf, wenn Sie die Spektren des abgeschnittenen und des originalen Signal mit dem vollständigen vergleichen?

```python

```

### 1e)
Programmieren Sie ein Hamming Apodisierungsfenster

$$ w_{\text{hamming}}(t) = 0.54 + 0.46 \cos(\frac{\pi t}{t_{\text{max}}})$$

für das bei 60 ms abgeschnittene FID. Wenden Sie das Hamming Window auf das abgeschnittene Signal an, zerofillen Sie auf 2048 Punkte und vergleichen Sie die Spektren mit und ohne Hamming Window.

```python

```

### 1f)
Zeigen Sie, dass ein Hamming Fenster die Auflösung im Frequenzbereich im Vergleich zu einem "hart" abgeschnittenen FID verschlechtert, jedoch die *ripples* besser unterdrückt. Verwenden Sie Zerofilling au 2048 Punkte. Vergleichen Sie dazu die entsprechende Faltungsfunktion im Frequenzbereich.

```python

```

## Aufgabe 2: Digitalisierung eines Zeitsignals
---------
Ein zeitlich abfallendes oszillierendes Signal hat die Form:
$\begin{equation}
	x(t) =
	\begin{cases}
		 \exp \left( - \frac{t}{T_2} \right) \left\{ \cos(2\pi \nu_1 t) + 3\cos(2\pi \nu_2 t) \right\} & t\geq 0 \\
	 	0 & t < 0
	 \end{cases}
\end{equation}$
mit den beiden Frequenzkomponenten $ \nu_1 =10\,{\rm Hz}$ und $\nu_2 = 150\,{\rm Hz}$ sowie der Zerfallszeit $T_2  = 0.5\,{\rm s}$. Das Signal wird total durch 2048 Punkte aufgenommen und digitalisiert.

---------

### 2a)
Welche Abtastrate schlagen Sie für das Digitalisieren dieses Signals vor? Begründen Sie Ihre Antwort mit Hilfe des Abtasttheorems.
Berechnen Sie anhand der vorgeschlagenen Abtastrate das Zeitsignal und das dazugehörige Spektrum.

```python

```

### 2b)
Vergleichen Sie unterschiedlichen Abtastraten $f_s = 200,320,1000\,{\rm Hz}$ für das obengenannte Signal . Was für Erkenntnisse können Sie aus den entsprechenden Spektren gewinnen?

```python

```


# Quadcopter
## Summary
This notebook outlines a the design of a motion controller for a quadcopter.

## Goals
The ultimate goal is to apply the designed control system to a simulated environment - for this I have chosen Python and specifially [pybullet](https://pybullet.org/) as the 3D physics simulator and [pyglet](http://www.pyglet.org) as the game engine to render the results.

### Control System Requirements
The design criteria for the quadrotor is to reach a step input in 3D space of (1, 1, 1):
* Settling time for x, y, z and yaw (ψ) of less than 5 seconds
* Rise time for x, y, z and yaw (ψ) of less than 2 seconds
* Overshoot of x, y, z and yaw (ψ) less than 5%

## System Description
We will use the following diagram to derive the equations of motion:

TODO - add diagram

* coordinate system and inertial frame is defined with the positive z axis in the opposite direction of gravity
* x, y, and z are the coordinates of the quadcopter centre of mass (CoM) in the inertial frame
* φ, θ, and ψ are the roll, pitch and yaw about the axes x, y and z respectively, with respect to the inertial frame - angles are 0 radians when the quadcopter is hovering, angles are measured CCW when 'looking down' the axis of rotation in the inertial frame
* rotors directly across from eachother rotate in the same direction (this allows the quadcopter to change its yaw angle while keeping its position constant)
* $F_1, F_2, F_3, F_4$ are the forces from the rotors, we assume they can only thrust in the positive z direction, with respect to the quadcopter frame
* $M_1, M_2, M_3, M_4$ are the moments of inertia for the rotors, about the z axis
* $u_1$ and $u_2$ are the motion control system inputs, where $u_1$ controls the throttle, and $u_2$ controls the rotation
* $I$ is the quadcopter moment of inertia (with x, y, and z components)
* $m$ is the mass of the quadcopter, in kg
* $l$ is the distance from the centre of mass to the rotors
* $g$ is gravity
* $R$ is the rotation matrix from the quadcopter frame to the static base frame using the [ZXY Euler angles (Wikipedia)](https://en.wikipedia.org/wiki/Euler_angles#Rotation_matrix) convention
* $\dot{x}, \dot{y}, \dot{z}$ are translational velocities
* $\ddot{x}, \ddot{y}, \ddot{z}$ are translational accelerations
* $\dot{\phi}, \dot{\theta}, \dot{\psi}$ are rotational velocities
* $\ddot{\phi}, \ddot{\theta}, \ddot{\psi}$ are rotational accelerations

## Equations of Motion
We will use the [Newton-Euler equations (Wikipedia)](https://en.wikipedia.org/wiki/Newton%E2%80%93Euler_equations) to define the equations of motion.

Using the [ZXY Euler angles (Wikipedia)](https://en.wikipedia.org/wiki/Euler_angles#Rotation_matrix) convention, we will define the rotation matrix from the quadcopter body frame into the inertial frame as $R$

```python
import sympy as sp
import matplotlib.pyplot as plt
import numpy as np
import math

sp.init_printing()
```

```python
# create our symbols and functions
t, g = sp.symbols('t g')

x = sp.Function('x')(t)
y = sp.Function('y')(t)
z = sp.Function('z')(t)
φ = sp.Function('φ')(t)
θ = sp.Function('θ')(t)
ψ = sp.Function('ψ')(t)
```

```python
# define the rotation matrices about the axes x, y, and z
Rx = sp.Matrix([[1, 0, 0], [0, sp.cos(φ), -sp.sin(φ)], [0, sp.sin(φ), sp.cos(φ)]])
Ry = sp.Matrix([[sp.cos(θ), 0, sp.sin(θ)], [0, 1, 0], [-sp.sin(θ), 0, sp.cos(θ)]])
Rz = sp.Matrix([[sp.cos(ψ), -sp.sin(ψ), 0], [sp.sin(ψ), sp.cos(ψ), 0], [0, 0, 1]])

# create the euler angle Z-X-Y rotation matrix and print it out
R = Rz*Rx*Ry
R
```

$\displaystyle \left[\begin{matrix}- \sin{\left(θ{\left(t \right)} \right)} \sin{\left(φ{\left(t \right)} \right)} \sin{\left(ψ{\left(t \right)} \right)} + \cos{\left(θ{\left(t \right)} \right)} \cos{\left(ψ{\left(t \right)} \right)} & - \sin{\left(ψ{\left(t \right)} \right)} \cos{\left(φ{\left(t \right)} \right)} & \sin{\left(θ{\left(t \right)} \right)} \cos{\left(ψ{\left(t \right)} \right)} + \sin{\left(φ{\left(t \right)} \right)} \sin{\left(ψ{\left(t \right)} \right)} \cos{\left(θ{\left(t \right)} \right)}\\\sin{\left(θ{\left(t \right)} \right)} \sin{\left(φ{\left(t \right)} \right)} \cos{\left(ψ{\left(t \right)} \right)} + \sin{\left(ψ{\left(t \right)} \right)} \cos{\left(θ{\left(t \right)} \right)} & \cos{\left(φ{\left(t \right)} \right)} \cos{\left(ψ{\left(t \right)} \right)} & \sin{\left(θ{\left(t \right)} \right)} \sin{\left(ψ{\left(t \right)} \right)} - \sin{\left(φ{\left(t \right)} \right)} \cos{\left(θ{\left(t \right)} \right)} \cos{\left(ψ{\left(t \right)} \right)}\\- \sin{\left(θ{\left(t \right)} \right)} \cos{\left(φ{\left(t \right)} \right)} & \sin{\left(φ{\left(t \right)} \right)} & \cos{\left(θ{\left(t \right)} \right)} \cos{\left(φ{\left(t \right)} \right)}\end{matrix}\right]$

### Position
Total force $F$ acting on the centre of mass, with respect to the inertial frame.

$ma = F$ where a is acceleration

$m\begin{bmatrix} \ddot{x} \\
\ddot{y} \\
\ddot{z} \end{bmatrix} = \begin{bmatrix} 0 \\
0 \\
-mg \end{bmatrix} + R\begin{bmatrix} 0 \\
0 \\
F_1+F_2+F_3+F_4\end{bmatrix}$

### Rotation
Total torque $\tau$ acting on the centre of mass with respect to the body frame.

$I\alpha = \tau - \omega\times I\omega$ where $\alpha$ is angular acceleration, $\omega$ is angular velocity, and $I$ is the [inertia tensor (Wikipedia)](https://en.wikipedia.org/wiki/Moment_of_inertia#Inertia_tensor) of the quadcopter:

$I = \begin{bmatrix} I_{xx} & I_{xy} & I_{xz} \\
I_{yx} & I_{yy} & I_{yz} \\
I_{zx} & I_{zy} & I_{zz} \end{bmatrix}$

The angular velocity with respect to the body fixed frame is required to fill out this equation.

We can define a time based function that transforms a fixed point $p_b$ in the body frame to the inertial frame as $p_i(t)$:

$p_i(t) = R(t)p_b$

Taking the time derivative gives us the velocity in the inertial frame:

$\dot{p_i(t)} = \dot{R(t)}p_b$

To get the velocity in a body fixed frame, we can multiply by the transpose of the rotation matrix (since the original matrix is converting from body to fixed):

$R^T\dot{p_i(t)} = R^T\dot{R(t)}p_b$

The term $R^T\dot{R(t)}$ can then be used to convert angular velocities in an inertial frame into a body fixed frame.

#### Alternative Calculation
**NOTE:** help needed here...to use this calculation, I also had to negate the result, but not sure why - something to do with skew symmetry?

The angular velocity with respect to the body fixed frame is required to fill out this equation. We can use the rotation matrix from the body frame to inertial frame to calculate the [angular velocity vector (Wikipedia)]((https://en.wikipedia.org/wiki/Rotation_formalisms_in_three_dimensions#Rotation_matrix_%E2%86%94_angular_velocities)) in the body frame:

$\begin{bmatrix} 0 & -\omega_z & \omega_y \\
\omega_z & 0 & -\omega_x \\
-\omega_y & \omega_x & 0 \end{bmatrix} = \dot{A}A^T$

We must first invert/transpose the rotation matrix so we can get the transformation from inertial frame into body frame.

In python we could do the following:
```python
# hack alert! had to negate the value to get it to work...
A = R.transpose()
ω = -sp.simplify(A.diff(t)*A.transpose())
```

```python
# calculate the angular velocity in the body fixed frame
ω = sp.simplify(R.transpose()*R.diff(t))

ω_x = ω[2, 1]
ω_y = ω[0, 2]
ω_z = ω[1, 0]

# print out the equations for the components of omega
ω_x, ω_y, ω_z
```

Rearranging for the angular velocities above:

$\omega = \begin{bmatrix} \omega_x \\
\omega_y \\
\omega_z \end{bmatrix} = \begin{bmatrix} \cos\theta & 0 & -\cos\phi\sin\theta \\
0 & 1 & \sin\phi \\
\sin\theta & 0 & \cos\phi\cos\theta \end{bmatrix} \begin{bmatrix} \dot{\phi} \\
\dot{\theta} \\
\dot{\psi} \end{bmatrix}$

and putting it all together:

$I\begin{bmatrix} \dot{\omega_x} \\
\dot{\omega_y} \\
\dot{\omega_z} \end{bmatrix} = \begin{bmatrix}  l(F_2 - F_4) \\
l(F_3 - F_1) \\
M_1-M_2+M_3-M_4 \end{bmatrix} - \begin{bmatrix} \omega_x \\
\omega_y \\
\omega_z \end{bmatrix} \times I \begin{bmatrix} \omega_x \\
\omega_y \\
\omega_z \end{bmatrix}$

## Approximations for Hover
We will linearize our equations at a stable hover, where we make the following approximations:
* position will be constant and equal to its initial value
* time derivatives of position (e.g. velocity and acceleration) will then be 0
* roll and pitch will be constant and approximately 0
* yaw will be constant and equal to its initial value
* time derivatives of rotation (e.g. angular velocity and acceleration) will also be 0
* the sine of a small value can then be assigned to the small value
* the cosine of a small value can then be assigned to 1
* the non-principal moments of inertia are 0

Or:

$\cos\phi \approx 1 \\
\sin\phi \approx \phi \\
\cos\theta \approx 1 \\
\sin\theta \approx \theta$

The rotation matrix from body frame into inertial frame becomes:

$R = \begin{bmatrix} r11 & r12 & r13 \\
r21 & r22 & r23 \\
r31 & r32 & r33 \end{bmatrix} = \begin{bmatrix} \cos\psi - \sin\psi\phi\theta & -\sin\psi & \cos\psi\theta + \sin\psi\phi \\
\sin\psi + \cos\psi\phi\theta & \cos\psi & \sin\psi\theta - \cos\psi\phi \\
-\theta & \phi & 1 \end{bmatrix}$

### Position

We will define thrust input $u_1 = F_1+F_2+F_3+F_4$ and our force equation becomes:

$m\begin{bmatrix} \ddot{x} \\
\ddot{y} \\
\ddot{z} \end{bmatrix} = \begin{bmatrix} 0 \\
0 \\
-mg \end{bmatrix} + \begin{bmatrix} \cos\psi\theta + \sin\psi\phi \\
\sin\psi\theta - \cos\psi\phi \\
1\end{bmatrix}u_1$

The second derivative of time (acceleration) is proportional to $u_1$.

### Rotation
Assuming the non-principal components of inertia are 0 gives us a revised inertia tensor:

$I = \begin{bmatrix} I_{xx} & 0 & 0 \\
0 & I_{yy} & 0 \\
0 & 0 & I_{zz} \end{bmatrix}$

Linearizing our rotational equations gives us:

$\omega = \begin{bmatrix} \omega_x \\
\omega_y \\
\omega_z \end{bmatrix} = \begin{bmatrix} 1 & 0 & -\theta \\
0 & 1 & \phi \\
\theta & 0 & 1 \end{bmatrix} \begin{bmatrix} \dot{\phi} \\
\dot{\theta} \\
\dot{\psi} \end{bmatrix}$

and also considering that the products of two small numbers is a really small number (the angular velocities are all small), gives us:

$\omega = \begin{bmatrix} \omega_x \\
\omega_y \\
\omega_z \end{bmatrix} = \begin{bmatrix} \dot{\phi} \\
\dot{\theta} \\
\dot{\psi} \end{bmatrix}$

Let $u_2$ be:

$u_2 = \begin{bmatrix} u_{2x} \\
u_{2y} \\
u_{2z} \\
\end{bmatrix} = \begin{bmatrix} l(F_2 - F_4) \\
l(F_3 - F_1) \\
M_1-M_2+M_3-M_4 \end{bmatrix}$

Filling in our torque equation:

$\begin{bmatrix} I_{xx} & 0 & 0 \\
0 & I_{yy} & 0 \\
0 & 0 & I_{zz} \end{bmatrix} \begin{bmatrix} \dot{\omega_x} \\
\dot{\omega_y} \\
\dot{\omega_z} \end{bmatrix} = \begin{bmatrix} u_{2x} \\
u_{2y} \\
u_{2z} \end{bmatrix} - \begin{bmatrix} 0 & -\omega_z & \omega_y \\
\omega_z & 0 & -\omega_x \\
-\omega_y & \omega_x & 0 \end{bmatrix} \begin{bmatrix} I_{xx} & 0 & 0 \\
0 & I_{yy} & 0 \\
0 & 0 & I_{zz} \end{bmatrix} \begin{bmatrix} \omega_x \\
\omega_y \\
\omega_z \end{bmatrix}$

$\begin{bmatrix} I_{xx}\dot{\omega_x} \\
I_{yy}\dot{\omega_y} \\
I_{zz}\dot{\omega_z} \end{bmatrix} = \begin{bmatrix} u_{2x} + I_{yy}\omega_y\omega_z - I_{zz}\omega_z\omega_y \\
u_{2y} - I_{xx}\omega_x\omega_z + I_{zz}\omega_z\omega_x \\
u_{2z} + I_{xx}\omega_x\omega_y - I_{zz}\omega_z\omega_x \end{bmatrix}$

and also approximating that the $\omega_i$ terms multiplied together are approximately zero, we can rearrange and define the following equations:

$I_{xx}\ddot{\phi} = u_{2x} \\
I_{yy}\ddot{\theta} = u_{2y} \\
I_{zz}\ddot{\psi} = u_{2z}$

and finally get equations for the angular accelerations:

$\ddot{\phi} = \dfrac{u_{2x}}{I_{xx}} \\
\ddot{\theta} = \dfrac{u_{2y}}{I_{yy}} \\
\ddot{\psi} = \dfrac{u_{2z}}{I_{zz}}$

by taking one of the linearized equations for position, e.g. $m\ddot{x}$ and then differentiating against time twice, we can substitute in the above equations and observe that the input $u_2$ is proportional to the fourth derivative of time (snap). This means we will desire a minimum snap trajectory to smoothly control a quadcopter (citation needed).

## State Space Representation
This system has 12 states - position, rotation and their associated velocities.

*NOTE: linearized equations for hover*

$x(t) = \begin{bmatrix} x_0 \\
x_1 \\
x_2 \\
x_3 \\
x_4 \\
x_5 \\
x_6 \\
x_7 \\
x_8 \\
x_9 \\
x_{10} \\
x_{11} \end{bmatrix} = \begin{bmatrix} x \\
y \\
z \\
\phi \\
\theta \\
\psi \\
\dot{x} \\
\dot{y} \\
\dot{z} \\
\dot{\phi} \\
\dot{\theta} \\
\dot{\psi} \end{bmatrix}, \dot{x(t)} = \begin{bmatrix} \dot{x_0} \\
\dot{x_1} \\
\dot{x_1} \\
\dot{x_2} \\
\dot{x_3} \\
\dot{x_4} \\
\dot{x_5} \\
\dot{x_6} \\
\dot{x_7} \\
\dot{x_8} \\
\dot{x_9} \\
\dot{x_{10}} \\
\dot{x_{11}} \end{bmatrix} = \begin{bmatrix} x_6 \\
x_7 \\
x_8 \\
x_9 \\
x_{10} \\
x_{11} \\
\ddot{x} \\
\ddot{y} \\
\ddot{z} \\
\ddot{\phi} \\
\ddot{\theta} \\
\ddot{\psi} \end{bmatrix} = \begin{bmatrix} x_6 \\
x_7 \\
x_8 \\
x_9 \\
x_{10} \\
x_{11} \\
\dfrac{(\cos\psi\theta + \sin\psi\phi)u_1}{m} \\
\dfrac{(\sin\psi\theta - \cos\psi\phi)u_1}{m} \\
\dfrac{u_1}{m} - g \\
\dfrac{u_{2x}}{I_{xx}} \\
\dfrac{u_{2y}}{I_{yy}} \\
\dfrac{u_{2z}}{I_{zz}} \end{bmatrix}$

where:

$u(t) = \begin{bmatrix} u_1 \\
u_{2x} \\
u_{2y} \\
u_{2z} \end{bmatrix} = \begin{bmatrix} F_1+F_2+F_3+F_4 \\
l(F_2 - F_4) \\
l(F_3 - F_1) \\
M_1-M_2+M_3-M_4 \end{bmatrix}$

## Control

The following diagram represents the control system:

1. the trajectory generator provides the desired position and yaw to the position controller
1. the position controller reads the current quadcopter position state, determines commanded pitch, roll, and yaw values and sends that to the attitude controller, and simultaneously determines the thrust input ($u_1$) and sends that to the motion controller
1. the attitude controller reads the current quadcopter rotation state, and determines the rotation input ($u_2$) and sends that to the motion controller
1. the motion controller solves for the 4 rotor forces based on the given inputs ($u_1 and u_2$) and updates the quadcopter with the current required forces

UML code provided here in case of edits - image generated on [PlantUML](https://plantuml.com/):
```
@startuml
[Trajectory\nGenerator] -right-> [Position\nController] : desired\nx,y,z,𝜓
[Position\nController] -right-> [Motion\nController] : u1
[Position\nController] -down-> [Attitude\nController] : commanded\n𝜙,𝜃,𝜓
[Attitude\nController] -> [Motion\nController] : u2
[Motion\nController] -right-> [Quadcopter]
[Quadcopter] -down-> [Attitude\nController] : 𝜙,𝜃,𝜓,\nd/dt(𝜙,𝜃,𝜓)
[Quadcopter] -> [Position\nController] : x,y,z,\nd/dt(x,y,z)
@enduml
```

### PID Controllers
We can use [proportional–integral–derivative (PID) controllers (Wikipedia)](https://en.wikipedia.org/wiki/PID_controller) to reach our desired states. In general, to calculate an time based input $u(t)$ based on the time based error $e(t)$ the equation is as follows:

$u(t) = K_pe(t) + K_i\int_0^t \! e(t) + K_d\dot{(e(t)}$

For simplicity, we will ignore the integral term and implement a PD controller to focus on the current error in position and velocity, and ignore the accumulated error (we would want to include this for a real world application, but for an ideal simulation we can ignore it).

#### Position
For position, we will have the following control equations (c = commanded, d = desired):

$\begin{bmatrix} \ddot{x_c} \\
\ddot{y_c} \\
\ddot{z_c}
\end{bmatrix} = \begin{bmatrix} \ddot{x_d} + K_{p,x}(x_d - x) + K_{d,x}(\dot{x_d} - \dot{x}) \\
\ddot{y_d} + K_{p,y}(y_d - y) + K_{d,y}(\dot{y_d} - \dot{y}) \\
\ddot{z_d} + K_{p,z}(z_d - z) + K_{d,z}(\dot{z_d} - \dot{z})
\end{bmatrix}$

We can then calculate $u_1$, which is simply the combined acceleration of commanded acceleration and gravity, multiplied by the quadcopter mass:

$u_1 = m(g+\ddot{z_c}) = m(g + \ddot{z_d} + K_{p,\ddot{z}}(z_d - z) + K_{d,\ddot{z}}(\dot{z_d} - \dot{z}))$

#### Rotation
For rotation, once we have the commanded position accelerations, we can calculate the commanded rotations using trigonometry (c = commanded, d = desired):

$\begin{bmatrix} \phi_c \\
\theta_c \\
\psi_c
\end{bmatrix} = \begin{bmatrix} \dfrac{\ddot{x_c}\sin{\psi_d} - \ddot{y_c}\cos{\psi_d}}{g} \\
\dfrac{\ddot{x_c}\cos{\psi_d} + \ddot{y_c}\sin{\psi_d}}{g} \\
\psi_d
\end{bmatrix}$

*NOTE: We can use the above to calculate the commanded angular velocities and accelerations by taking 2 time derivatives. This will result in a requirement to calculate the commanded jerk and snap for x and y, which we can get from taking 2 time derivatives of our x and y accelerations above. This is why $u_2$ is dependent on the fourth time derivative of position. Commanded jerk and snap for x and y also assumes we can calculate the current jerk and acceleration of the quadcopter.*

and then calculate $u_2$:

$u_2 = \begin{bmatrix} u_{2x} \\
u_{2y} \\
u_{2z}
\end{bmatrix} = \begin{bmatrix} \ddot\phi_c + K_{p,\phi}(\phi_c - \phi) + K_{d,\phi}(\dot{\phi_c} - \dot{\phi}) \\
\ddot\theta_c + K_{p,\theta}(\theta_c - \theta) + K_{d,\theta}(\dot{\theta_c} - \dot{\theta}) \\
\ddot\psi_c + K_{p,\psi}(\psi_c - \psi) + K_{d,\psi}(\dot{\psi_c} - \dot{\psi})
\end{bmatrix}$

#### Motion Control
Now that we have our input $u$ we can then solve for the required rotor forces, and send those to the quadcopter.

*NOTE: in reality we would want to calculate the voltage required to reach a desired rotor speed in order to create the force, but for the purposes of our ideal simulation, we will just use the forces and moments.*

We now have a system of 4 equations and 4 unknowns and can use linear algebra to solve the system:

$\begin{bmatrix} u_1 \\
u_{2x} \\
u_{2y} \\
u_{2z} \end{bmatrix} = \begin{bmatrix} F_1+F_2+F_3+F_4 \\
l(F_2 - F_4) \\
l(F_3 - F_1) \\
M_1-M_2+M_3-M_4 \end{bmatrix}$

We will also need to solve for the moments, $M_i$ but these can be linearly related to the force produced. For a propellor, the force and moment can be represented as a product of some constant and the square of the propellor speed (citation needed):

$F = k_f\omega^2 \\
M = k_m\omega^2$

which means we can represent the moments as:

$M_i = \dfrac{k_m}{k_f}F_i$

and substitute appropriately:

$\begin{bmatrix} u_1 \\
u_{2x} \\
u_{2y} \\
u_{2z} \end{bmatrix} = \begin{bmatrix} F_1+F_2+F_3+F_4 \\
l(F_2 - F_4) \\
l(F_3 - F_1) \\
\dfrac{k_m}{k_f}(F_1-F_2+F_3-F_4) \end{bmatrix}$

we can now represent this system as $Af = u$ where the $f$ vector is the unknown forces, and the $A$ matrix is:

$A = \begin{bmatrix} 1 & 1 & 1 & 1 \\
0 & l & 0 & -l \\
-l & 0 & l & 0 \\
\dfrac{k_m}{k_f} & -\dfrac{k_m}{k_f} & \dfrac{k_m}{k_f} & -\dfrac{k_m}{k_f}
\end{bmatrix}$

We can then solve for $f$:

$f = A^{-1}u$

Since A is all constants, we can calculate the inverse of $A$ once and re-use for each iteration of our control loop.

## Simulation
We now have enough data to simulate the system.

### Constants
I will use a master thesis called [*Modelling, Identification and Control
of a Quadrotor Helicopter*](http://lup.lub.lu.se/luur/download?func=downloadFile&recordOId=8847641&fileOId=8859343) by Tommaso Bresciani as a reference to fill in the constants.

### Constant Control
To ensure our model is doing what we expect, we will first simulate conditions where constant inputs are applied.

```python
# define our constants
constants = dict(
    l = 235e-3,    # 23.5 cm from CoM to rotor
    m = 250e-3,    # 250 g
    Ixx = 8.1e-3,  # Nms^2
    Iyy = 8.1e-3,  # Nms^2
    Izz = 14.2e-3, # Nms^2
    g = 9.81,      # gravity, m/s^2,
    kmkf = 50e-6   # km/kf - the ratio of the moment constant over the force constant
)
```

```python
# define our ivp function using our state space representation
def quadcopter(t, state, forces, constants):
    u1 = sum(forces)
    u2x = constants['l'] * (forces[1] - forces[3])
    u2y = constants['l'] * (forces[2] - forces[0])
    u2z = constants['kmkf'] * (forces[0] - forces[1] + forces[2] - forces[3])
    phi = state[3]
    theta = state[4]
    psi = state[5]
    u1_m = u1 / constants['m']
    x_accel = (math.cos(psi)*theta + math.sin(psi)*phi) * u1_m
    y_accel = (math.sin(psi)*theta - math.cos(psi)*phi) * u1_m
    z_accel = u1_m - constants['g']
    phi_accel = u2x / constants['Ixx']
    theta_accel = u2y / constants['Iyy']
    psi_accel = u2z / constants['Izz']
    updated = np.zeros(12)
    updated[0:6] = state[6:]
    updated[6:] = [x_accel, y_accel, z_accel, phi_accel, theta_accel, psi_accel]
    return updated
```

```python
# define a plotting function to show all the states
def plot_all(title, t, z):
    plt.figure(figsize=(20, 24))

    plt.subplot(4, 3, 1)
    plt.plot(t, z.T[:,0], 'r')
    plt.xlabel('time (s)')
    plt.ylabel('X position (m)')
    plt.title(title)
    plt.grid(True)

    plt.subplot(4, 3, 2)
    plt.plot(t, z.T[:,1], 'b')
    plt.xlabel('time (s)')
    plt.ylabel('Y position (m)')
    plt.grid(True)

    plt.subplot(4, 3, 3)
    plt.plot(t, z.T[:,2], 'g')
    plt.xlabel('time (s)')
    plt.ylabel('Z position (m)')
    plt.grid(True)

    plt.subplot(4, 3, 4)
    plt.plot(t, z.T[:,3], 'y')
    plt.xlabel('time (s)')
    plt.ylabel('Pitch (radians)')
    plt.grid(True)

    plt.subplot(4, 3, 5)
    plt.plot(t, z.T[:,4], 'm')
    plt.xlabel('time (s)')
    plt.ylabel('Roll (radians)')
    plt.grid(True)

    plt.subplot(4, 3, 6)
    plt.plot(t, z.T[:,5], 'c')
    plt.xlabel('time (s)')
    plt.ylabel('Yaw (radians)')
    plt.grid(True)

    plt.subplot(4, 3, 7)
    plt.plot(t, z.T[:,6], 'r')
    plt.xlabel('time (s)')
    plt.ylabel('X velocity (m/s)')
    plt.grid(True)

    plt.subplot(4, 3, 8)
    plt.plot(t, z.T[:,7], 'b')
    plt.xlabel('time (s)')
    plt.ylabel('Y velocity (m/s)')
    plt.grid(True)

    plt.subplot(4, 3, 9)
    plt.plot(t, z.T[:,8], 'g')
    plt.xlabel('time (s)')
    plt.ylabel('Z velocity (m/s)')
    plt.grid(True)

    plt.subplot(4, 3, 10)
    plt.plot(t, z.T[:,9], 'y')
    plt.xlabel('time (s)')
    plt.ylabel('Pitch velocity (radians/s)')
    plt.grid(True)

    plt.subplot(4, 3, 11)
    plt.plot(t, z.T[:,10], 'm')
    plt.xlabel('time (s)')
    plt.ylabel('Roll velocity (radians/s)')
    plt.grid(True)

    plt.subplot(4, 3, 12)
    plt.plot(t, z.T[:,11], 'c')
    plt.xlabel('time (s)')
    plt.ylabel('Yaw velocity (radians/s)')
    plt.grid(True)

    plt.xlabel('time (s)')
    plt.show()
```

```python
from scipy.integrate import solve_ivp

# let the quadcopter fall
forces = [0.0] * 4

# initial values for state vector
y0 = [0.0] * 12

args = (forces, constants)
sol = solve_ivp(quadcopter, [0, 10], y0, args=args, dense_output=True)
```

```python
t = np.linspace(0, 10, 100)
z = sol.sol(t)

# we expect the quadcopter to fall from the sky, z decreasing exponentially, and zdot (velocity) increasing linearly, all other states should be 0
plot_all('Quadcopter Falling', t, z)
```

```python
# determine the force required to hover
force = constants['m'] * constants['g'] / 4
forces = [force] * 4

args = (forces, constants)
sol = solve_ivp(quadcopter, [0, 10], y0, args=args, dense_output=True)

t = np.linspace(0, 10, 100)
z = sol.sol(t)

# we expect the quadcopter to hover so all states should be 0
plot_all('Quadcopter Hovering', t, z)
```

```python
# apply forces such that the quadcopter will change its yaw
force = constants['m'] * constants['g'] / 4
delta_f = force * 0.1
forces = [force] * 4

forces[0] -= delta_f
forces[2] -= delta_f
forces[1] += delta_f
forces[3] += delta_f

args = (forces, constants)
sol = solve_ivp(quadcopter, [0, 10], y0, args=args, dense_output=True)

t = np.linspace(0, 10, 100)
z = sol.sol(t)

# we expect the quadcopter to hover but the yaw should be changing
plot_all('Quadcopter Changing Yaw', t, z)
```

```python
# apply forces such that the quadcopter will start translating and rotating
force = constants['m'] * constants['g'] / 4
delta_f = force * 0.0001
forces = [force] * 4

forces[0] += delta_f
forces[1] += delta_f

args = (forces, constants)
sol = solve_ivp(quadcopter, [0, 10], y0, args=args, dense_output=True)

t = np.linspace(0, 10, 100)
z = sol.sol(t)

# we expect the quadcopter start moving and rotating, but yaw should be 0
plot_all('Quadcopter Movement', t, z)
```

### PD Control
We can now redefine our function to include PD control.

Our trajectory generator will give us the desired x, y, z and yaw positions, velocities, and accelerations, allowing us to calculate the commanded values - which we derived earlier:

$\begin{bmatrix} \ddot{x_c} \\
\ddot{y_c} \\
\ddot{z_c}
\end{bmatrix} = \begin{bmatrix} \ddot{x_d} + K_{p,x}(x_d - x) + K_{d,x}(\dot{x_d} - \dot{x}) \\
\ddot{y_d} + K_{p,y}(y_d - y) + K_{d,y}(\dot{y_d} - \dot{y}) \\
\ddot{z_d} + K_{p,z}(z_d - z) + K_{d,z}(\dot{z_d} - \dot{z})
\end{bmatrix}$

and commanded rotation:

$\begin{bmatrix} \phi_c \\
\theta_c \\
\psi_c
\end{bmatrix} = \begin{bmatrix} \dfrac{\ddot{x_c}\sin{\psi_d} - \ddot{y_c}\cos{\psi_d}}{g} \\
\dfrac{\ddot{x_c}\cos{\psi_d} + \ddot{y_c}\sin{\psi_d}}{g} \\
\psi_d
\end{bmatrix}$

We can then calculate the following input $u$:

$\begin{bmatrix} u_1 \\
u_{2x} \\
u_{2y} \\
u_{2z}
\end{bmatrix} = \begin{bmatrix} m(g+\ddot{z_c}) \\
\ddot\phi_c + K_{p,\phi}(\phi_c - \phi) + K_{d,\phi}(\dot{\phi_c} - \dot{\phi}) \\
\ddot\theta_c + K_{p,\theta}(\theta_c - \theta) + K_{d,\theta}(\dot{\theta_c} - \dot{\theta}) \\
\ddot\psi_c + K_{p,\psi}(\psi_c - \psi) + K_{d,\psi}(\dot{\psi_c} - \dot{\psi})
\end{bmatrix}$

Once we have $u$ we can then solve for the individual rotor forces required.

#### Step Input and Hover
We will simplify the control loop such that we will give the quadcopter a step input - i.e. we will command the quadcopter to go to some point in 3D space with a particular yaw, with the assumption that the quadcopter should hover at this steady state. This is also ideal for tuning our PD controllers.

With a desired step input and hover steady state, we can zero out the commanded velocity and all further time derivatives:

$\begin{bmatrix} \ddot{x_c} \\
\ddot{y_c} \\
\ddot{z_c}
\end{bmatrix} = \begin{bmatrix} K_{p,x}(x_d - x) - K_{d,x}\dot{x} \\
K_{p,y}(y_d - y) - K_{d,y}\dot{y} \\
K_{p,z}(z_d - z) - K_{d,z}\dot{z} \end{bmatrix}$

$\begin{bmatrix} u_1 \\
u_{2x} \\
u_{2y} \\
u_{2z}
\end{bmatrix} = \begin{bmatrix} m(g+\ddot{z_c}) \\
K_{p,\phi}(\phi_c - \phi) - K_{d,\phi}\dot{\phi} \\
K_{p,\theta}(\theta_c-\theta) - K_{d,\theta}\dot{\theta} \\
K_{p,\psi}(\psi_c - \psi) - K_{d,\psi}\dot{\psi}
\end{bmatrix}$

```python
# define our ivp function using our state space representation, and calculating the forces dynamically using PD control
# we will assume a simple trajectory planner that will send a desired hover position - so commanded velocites and accelerations will be 0
def quadcopter_pd(t, state, constants, desired, pd):
    # pull out desired values
    xd, yd, zd, psid = desired[0], desired[1], desired[2], desired[3]

    # pull out the pd tuples
    pdx, pdy, pdz, pdphi, pdtheta, pdpsi = pd[0], pd[1], pd[2], pd[3], pd[4], pd[5]

    # pull out the current state positions...
    x, y, z, phi, theta, psi = state[0], state[1], state[2], state[3], state[4], state[5]

    # ...and state velocities
    xv, yv, zv, phiv, thetav, psiv = state[6], state[7], state[8], state[9], state[10], state[11]

    # calculate commanded values
    xac = pdx[0]*(xd-x) - pdx[1]*xv
    yac = pdy[0]*(yd-y) - pdy[1]*yv
    zac = pdz[0]*(zd-z) - pdz[1]*zv
    phic = (xac*math.sin(psid) - yac*math.cos(psid)) / constants['g']
    thetac  = (xac*math.cos(psid) + yac*math.sin(psid)) / constants['g']
    psic = psid

    # calulate u
    u1 = constants['m'] * (constants['g'] + zac)
    u2x = pdphi[0]*(phic - phi) - pdphi[1]*phiv
    u2y = pdtheta[0]*(thetac - theta) - pdtheta[1]*thetav
    u2z = pdpsi[0]*(psic - psi) - pdpsi[1]*psiv

    # calculate updated accelerations
    u1_m = u1 / constants['m']
    x_accel = (math.cos(psi)*theta + math.sin(psi)*phi) * u1_m
    y_accel = (math.sin(psi)*theta - math.cos(psi)*phi) * u1_m
    z_accel = u1_m - constants['g']
    phi_accel = u2x / constants['Ixx']
    theta_accel = u2y / constants['Iyy']
    psi_accel = u2z / constants['Izz']

    updated = np.zeros(12)
    updated[0:6] = state[6:]
    updated[6:] = [x_accel, y_accel, z_accel, phi_accel, theta_accel, psi_accel]
    return updated
```

```python
# x, y, z, yaw
desired = (1, 1, 1, 0.78)

# each tuple is a PD controller with the values as (P, D)
# we need PD controllers for x, y, z, pitch, roll and yaw
pd = [(3.5, 3.2), # x
      (4.7, 4), # y
      (7, 5), # z
      (10, 0.5), # phi
      (10, 0.5), # theta
      (0.3, 0.25)] # psi

args = (constants, desired, pd)
sol = solve_ivp(quadcopter_pd, [0, 10], y0, args=args, dense_output=True)

t = np.linspace(0, 10, 100)
z = sol.sol(t)

# we expect the quadcopter start moving and rotating, but yaw should be 0
plot_all(f'Quadcopter step input from (0, 0, 0, 0) to {desired}', t, z)
```

#### Putting into code

```python
# define our A matrix once and invert it to save processing time
A = np.array([[1, 1, 1, 1],
              [0, constants['l'], 0, -constants['l']],
              [-constants['l'], 0, constants['l'], 0],
              [constants['kmkf'], -constants['kmkf'], constants['kmkf'], -constants['kmkf']]])
Ainv = np.linalg.inv(A)

def control_loop(state):
    # 1. explode state into values
    # 2. get desired values from trajectory generator
    # 3. calculate commanded values using PID control values
    # 4. calculate u

    # solve for the rotor forces required
    u = np.array([[u1, u2x, u2y, u2z]]).transpose()
    f = Ainv*u

    # lastly, apply the forces to a simulation or convert into a voltage to apply to a motor
```

```python

```


# draw_math_function
Jalankan kode pada sel-sel di bawah ini sehingga identitas Anda tertampilkan dengan benar dan gambarkan beberapa fungsi matematika.

```python
from IPython.core.display import HTML
HTML("")

import sys, os
sys.path.insert(1, os.path.join(sys.path[0], '../..'))

from src.student import Students as stu
nim = 10200999
print(nim, end=' ')
print(stu[nim]['name'], end=' ')
print("https://github.com/", stu[nim]['github'], sep='')
```

```python
import math
import matplotlib.pyplot as plt
```

## function
Sebuah fungsi matematika, misalnya saja

\begin{equation}\tag{1}
f(x) = \left[ x - 2.5 e^{-(x - 2.8)} \right] \sin \tfrac23 \pi (x-1)
\end{equation}

dapat digambarkan dengan dalam rentang $x \in [x_{\rm beg}, x_{\rm end}]$.

```python
# define a function of x
def f(x):
    y1 = (x -  2.5 * math.exp(-(x - 2.8)))
    y2 = math.sin(2 * math.pi * (x - 1) / 3)
    y = y1 * y2
    return y
```

```python
# create data
xbeg = 0
xend = 10
N = 100
dx = (xend - xbeg) / N

xx = []
yy = []

for i in range(0, N + 1):
    x = xbeg + i * dx
    y = f(x)

    xx.append(x)
    yy.append(y)

```

```python
# plot data
plt.ion()
fig, ax = plt.subplots()
ax.set_xlabel("$x$")
ax.set_ylabel("$y$")
plt.plot(xx, yy)

# draw horizontal line
line = plt.Line2D((xbeg, xend), (0, 0), color='#aaa', lw=1, ls='dashed')
plt.gca().add_line(line)

# show plot result
plt.show()
```

## another function
Terdapat fungsi lain berbentuk

\begin{equation}\tag{2}
f(x) = e^{-ax} \sin bx
\end{equation}

dalam rentang $x \in [x_{\rm beg}, x_{\rm end}]$ yang memberikan hasil

dengan nilai-nilai $a$, $b$, $x_{\rm beg}$, $x_{\rm end}$, dan $N$ tertentu.

```python
# set parameters
a = 0.1
b = 4
xbeg = 0
xend = 8
N = 100

# define another function of x
def g(x):
    y = math.exp(-a * x) * math.sin(b * x)
    return y

# create data
dx = (xend - xbeg) / N

xx = []
yy = []

for i in range(0, N + 1):
    x = xbeg + i * dx
    y = g(x)

    xx.append(x)
    yy.append(y)

# plot data
plt.ion()
fig, ax = plt.subplots()
ax.set_xlabel("$x$")
ax.set_ylabel("$y$")
plt.plot(xx, yy)

# draw horizontal line
line = plt.Line2D((xbeg, xend), (0, 0), color='#aaa', lw=1, ls='dashed')
plt.gca().add_line(line)
line = plt.Line2D((xbeg, xend), (0.4, 0.4), color='#aaa', lw=1, ls='dashed')
plt.gca().add_line(line)
line = plt.Line2D((xbeg, xend), (0.8, 0.8), color='#aaa', lw=1, ls='dashed')
plt.gca().add_line(line)
line = plt.Line2D((xbeg, xend), (-0.4, -0.4), color='#aaa', lw=1, ls='dashed')
plt.gca().add_line(line)

# show plot result
plt.show()

# print parameters
print("a = ", a, sep='')
print("b = ", b, sep='')
print("xbeg = ", xbeg, sep='')
print("xend = ", xend, sep='')
print("N = ", N, sep='')
```

Modifikasi kode di atas sehingga gambar berikut dapat diperoleh.