```python
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
plt.style.use('classic')

pd.plotting.register_matplotlib_converters()
%matplotlib inline
```

# Discussion: Week 4

## Exercise: The Perpetual Inventory Method

The stock of physical capital is a fundamental component of many macroeconomic models but measuring, and therefore acquiring data on, the value of a country's capital stock is hard. The challenge arise because it is hard to estimate the value of something that is not changing ownership in a transaction. What is value of a high-rise apartment building that has had the same owner for 50 years?

One way to overcome this challenge is to use the **perpetual inventory method**. The perpetual inventory method is based on the law of motion of the capital stock:

\begin{align}
K_{t+1} & = I_t + (1-\delta) K_t
\end{align}

where $K_t$ denotes physical capital, $I_t$ denotes investment (purchases of new capital), and $\delta$ denotes the rate of depreciation. Give $K_0$ and data on investment $I_0, I_1, \ldots,$, the the capital stock in periods $1, 2, \ldots$ can be computed iteratively. To obtain $K_0$, the method assumes that the economy is in  a steady state in the initial year. Recall that in the steady state:

\begin{align}
I & = \delta K
\end{align}

So, set $K_0 = \delta^{-1} I_0$.

Steps:

1. Obtain real investment data (https://fred.stlouisfed.org/series/GPDICA)
2. Select a value for the depreciation rate $\delta$. 0.1 is a standard choice.
3. Initialize an array of zeros (or an empty array) to store capital values
4. Set the first value of the captial array equal to the $\delta^{-1} I_0$
5. Iterate over the values in the investment data to compute the remaining values of capital.

Read Timothy Kehoe's notes (http://users.econ.umn.edu/~tkehoe/classes/GrowthAccountingNotes.pdf) for more background on the perpetual inventory method if you're interested.

```python
# Import real investment data from FRED into a DataFrame called 'data' (PROVIDED)
data = pd.read_csv('https://fred.stlouisfed.org/data/GPDICA.txt',sep='\s+',skiprows=13,index_col=0,parse_dates=True)

# Set the columns attribute of data to ['Investment']
data.columns = ['Investment']

# Set the index.name attribute of data to 'Date' (i.e., data.index.name)
data.index.name='Date'

# Create a variable called 'delta' that stores the value for delta as 0.1
delta = 0.1

# Create a variable called 'capital' equal to an array of zeros with the same length as the variable 'data'
capital = np.arange(len(data))

# Set the initial value of the 'capital' variable equal to the initial value of investment divided by delta
capital[0] = data['Investment'].iloc[0]/delta

# Iterate (i.e., use a for loop) over t in range(len(capital)-1) to construct values in the capital array
for t in range(len(capital)-1):
    capital[t+1] = data['Investment'].iloc[t] + (1-delta)*capital[t]

# Add a 'Capital' column to data
data['Capital'] = capital

# Divide data by 1000 so that the values are in units are in trillions of $ (instead of billions)
data = data/1000

# Print the first 5 rows of data
print(data.head())
```

                Investment  Capital
    Date
    1929-01-01    0.120381    1.203
    1930-01-01    0.081988    1.203
    1931-01-01    0.053297    1.164
    1932-01-01    0.019658    1.100
    1933-01-01    0.026589    1.009

Now construct a figure that contains a plot of capital in the top panel and investment in the bottom panel.

1. The figure should have dimensions $12\times 8$.
2. Y-axis labels for investment and capital should both be "trillions of chained dollars"
3. Each axis should have an appropriate title.

```python
# Create figure
fig = plt.figure(figsize=(12,8))

# Construct the plot of capital
ax1 = fig.add_subplot(2,1,1)
ax1.plot(data.index,data.Capital,lw=4,alpha=0.75)
ax1.set_ylabel('Trillions of chained Dollars')
ax1.set_title('Capital')
ax1.grid()

# Construct the plot of investment
ax2 = fig.add_subplot(2,1,2)
ax2.plot(data.index,data.Investment,lw=4,alpha=0.75)
ax2.set_ylabel('Trillions of chained Dollars')
ax2.set_xlabel('Date')
ax2.set_title('Investment')
ax2.grid()

# Use fig.tight_layout() to adjust whitespace
fig.tight_layout()
```

**Question**

1. Why does the computed capital series fluctuate less than the investment data?
2. Given that there are about 325 million people in the US, does the computed value of capital in the last year look large, small, or about right on a per capita basis?

**Answer**

1. Because 10 percent rate of capital accumulation implies an autoregressive coefficent on capital of 0.9 so changes in capital are highly persistent.<!-- answer -->
2. 26 trillion divided by 325 million is about 80,000 which seems to be in the right ballpark. Note that that the amount of capital in the US *per worker* will be an order of magnitude higher since there is a stubstantial share of people not in the workforce: e.g., children, retired persons, institutionalized persons.<!-- answer -->


# Task 3: Expectation values

There are two main quantities that we wish to compute for the ground state in this project.
They are the full many-body ground state energy $E$ from the general Hartree-Fock method, and the particle density (also known as the one-body density and the electron density) $\rho(x)$.

## The general Hartree-Fock energy

We wish to compute the ground state energy from the Hartree-Fock ansatz, $| \Psi \rangle = |\phi_1\dots\phi_n\rangle$.
This is done by
\begin{align}
    E &= \langle \Psi | \hat{H} | \Psi \rangle
    = \sum_{i = 1}^{n} \langle \phi_i | \hat{h} | \phi_i \rangle
    + \frac{1}{2} \sum_{i, j = 1}^{n}
    \langle \phi_i \phi_j | \hat{u} | \phi_i \phi_j \rangle_{AS},
\end{align}
where the molecular orbitals are the optimized Hartree-Fock orbitals.
Inserting the basis expansion the energy can be found as a function of the coefficient matrices and the atomic orbital matrix elements.

## The particle density

The particle density can be computed from the single-particle functions evaluated on a grid, and the one-body density matrix.
We have already found the one-body density matrix when constructing the Fock matrix, it is given by $D_{\nu \mu} = \sum_{i = 1}^{n} C^{*}_{\mu i} C_{\nu i}$.
We can then compute the particle density from
\begin{align}
    \rho(x) = \sum_{\mu, \nu = 1}^{l} \psi^{*}_{\mu}(x) D_{\nu \mu} \psi_{\nu}(x)
\end{align}


# Demo - Helmholtz equation in polar coordinates

**Mikael Mortensen** (email: `mikaem@math.uio.no`), Department of Mathematics, University of Oslo.

Date: **April 8, 2020**

**Summary.** This is a demonstration of how the Python module [shenfun](https://github.com/spectralDNS/shenfun) can be used to solve the
Helmholtz equation on a circular disc, using polar coordinates. This demo is implemented in
a single Python file [unitdisc_helmholtz.py](https://github.com/spectralDNS/shenfun/blob/master/demo/unitdisc_helmholtz.py),
and the numerical method is described in more detail by J. Shen [[shen3]](#shen3).

Helmholtz on the unit disc.

## Helmholtz equation

The Helmholtz equation is given as

$$
\begin{equation}
-\nabla^2 u(\boldsymbol{x}) + \alpha u(\boldsymbol{x}) = f(\boldsymbol{x}) \quad \text{for }\, \boldsymbol{x}=(x, y) \in \Omega, \label{eq:helmholtz} \tag{1}
\end{equation}
$$

$$
\begin{equation}
u =0 \text{ on } \partial \Omega,
\label{_auto1} \tag{2}
\end{equation}
$$

where $u(\boldsymbol{x})$ is the solution, $f(\boldsymbol{x})$ is a function and $\alpha$ a constant.
The domain is a circular disc $\Omega = \{(x, y): x^2+y^2 < a^2\}$ with radius $a$.
We use polar coordinates $(\theta, r)$, defined as

$$
\begin{equation}
 x = r \cos \theta,
\label{_auto2} \tag{3}
\end{equation}
$$

$$
\begin{equation}
 y = r \sin \theta,
\label{_auto3} \tag{4}
\end{equation}
$$

which leads to a Cartesian product mesh $(\theta, r) \in [0, 2\pi) \times [0, a]$
suitable for numerical implementations. Note that the
two directions are ordered with $\theta$ first and then $r$, which is less common
than $(r, \theta)$. This has to do with the fact that we will need to
solve linear equation systems along the radial direction, but not
the $\theta$-direction, since Fourier matrices are diagonal. When
the radial direction is placed last, the data in the radial direction
will be contigeous in a row-major C memory, leading to faster memory
access where it is needed the most. Note that it takes very few
changes in `shenfun` to switch the directions to $(r, \theta)$ if this
is still desired.

We will use Chebyshev
or Legendre basis functions $\psi_j(r)$ for the radial direction and
a periodic Fourier expansion in $\exp(\imath k \theta)$ for the
azimuthal direction. The polar basis functions are as such

$$
\begin{equation}
v_{kj}(\theta, r) = \exp(\imath k \theta) \psi_j(r),
\label{_auto4} \tag{5}
\end{equation}
$$

and we look for solutions

$$
\begin{equation}
u(\mathbf{x}) = \tilde{u}(\theta, r) = \sum_{k} \sum_{j} \hat{u}_{kj} v_{kj}(\theta, r).
\label{_auto5} \tag{6}
\end{equation}
$$

Note that $\tilde{u}$ is the function $u$ mapped to computational space.
From now on we will simply use $u(\theta, r)$ without the tilde, and assume that
the proper version of the function is understood from its arguments.

A discrete Fourier approximation space with $N$ basis functions is then

$$
\begin{equation}
V_F^N = \text{span} \{\exp(\imath k \theta)\}, \text{ for } k \in K,
\label{_auto6} \tag{7}
\end{equation}
$$

where $K = \{-N/2, -N/2+1, \ldots, N/2-1\}$. Since the solution $u(\theta, r)$
is real, there is Hermitian symmetry and $\hat{u}_{k,j} = \hat{u}_{k,-j}^*$
(with $*$ denoting a complex conjugate).
For this reason we use only $k \in K=\{0, 1, \ldots, N/2\}$ in solving for
$\hat{u}_{kj}$, and then use Hermitian symmetry to get the remaining
unknowns.

The radial basis is more tricky, because there is a nontrivial 'boundary'
condition (pole condition) that needs to be applied at the center of the disc $(r=0)$

$$
\begin{equation}
\frac{\partial u(\theta, 0)}{\partial \theta} = 0.
\label{_auto7} \tag{8}
\end{equation}
$$

To apply this condition we split the solution into Fourier
coefficients with wavenumber 0 and $K\backslash \{0\}$,
remembering that the Fourier basis function with $k=0$ is
simply 1

$$
\begin{equation}
u(\theta, r) = \sum_{j} \left( \hat{u}_{0j} \psi_{j}(r) + \sum_{k=1}^{N/2} \hat{u}_{kj} \exp(\imath k \theta) \psi_j(r) \right).
\label{_auto8} \tag{9}
\end{equation}
$$

We then apply a different radial basis for the two $\psi$'s in
the above equation (renaming the first $\overline{\psi}$)

$$
\begin{equation}
u(\theta, r) = \sum_{j} \left( \hat{u}_{0j} \overline{\psi}_{j}(r) + \sum_{k=1}^{N/2} \hat{u}_{kj} \exp(\imath k \theta) \psi_j(r) \right).
\label{_auto9} \tag{10}
\end{equation}
$$

Note that the first term $\sum_{j} \hat{u}_{0j} \overline{\psi}_{j}(r)$ is independent
of $\theta$. Now, to enforce conditions

$$
\begin{equation}
u(\theta, a) = 0,
\label{_auto10} \tag{11}
\end{equation}
$$

$$
\begin{equation}
\frac{\partial u(\theta, 0)}{\partial \theta} = 0,
\label{_auto11} \tag{12}
\end{equation}
$$

it is sufficient for the two bases ($\overline{\psi}$ and $\psi$) to
satisfy

$$
\begin{equation}
\overline{\psi}_j(a) = 0,
\label{_auto12} \tag{13}
\end{equation}
$$

$$
\begin{equation}
\psi_j(a) = 0,
\label{_auto13} \tag{14}
\end{equation}
$$

$$
\begin{equation}
\psi_j(0) = 0.
\label{_auto14} \tag{15}
\end{equation}
$$

Bases that satisfy these conditions can be found both with Legendre and
Chebyshev polynomials.
If $\phi_j(x)$ is used for either the Legendre polynomial $L_j(x)$ or the
Chebyshev polynomial of the first kind $T_j(x)$, we can have

$$
\begin{equation}
\overline{\psi}_j(r) = \phi_j(2r/a-1) - \phi_{j+1}(2r/a-1), \text{ for } j \in 0, 1, \ldots N-1,
\label{_auto15} \tag{16}
\end{equation}
$$

$$
\begin{equation}
\psi_j(r) = \phi_j(2r/a-1) - \phi_{j+2}(2r/a-1), \text{ for } j \in 0, 1, \ldots N-2.
\label{eq:psi} \tag{17}
\end{equation}
$$

Define the following approximation spaces for the radial direction

$$
\begin{equation}
V_D^N = \text{span} \{\psi_j\}_{j=0}^{N-3}
\label{_auto16} \tag{18}
\end{equation}
$$

$$
\begin{equation}
V_U^N = \text{span} \{\overline{\psi}_j\}_{j=0}^{N-2}
\label{_auto17} \tag{19}
\end{equation}
$$

$$
\begin{equation}
\label{_auto18} \tag{20}
\end{equation}
$$

and split the function space for the azimuthal direction into

$$
\begin{equation}
V_F^0 =  \text{span}\{1\},
\label{_auto19} \tag{21}
\end{equation}
$$

$$
\begin{equation}
V_F^{1} = \text{span} \{\exp(\imath k \theta)\}, \text{ for } k \in K \backslash \{0\}.
\label{_auto20} \tag{22}
\end{equation}
$$

We then look for solutions

$$
\begin{equation}
u(\theta, r) = u^0(r) + u^1(\theta, r),
\label{_auto21} \tag{23}
\end{equation}
$$

where

$$
\begin{equation}
u^0(r) = \sum_{j=0}^{N-2} \hat{u}^0_j \overline{\psi}_j(r),
\label{_auto22} \tag{24}
\end{equation}
$$

$$
\begin{equation}
u^1(\theta, r) = \sum_{j=0}^{N-3}\sum_{k=1}^{N/2} \hat{u}^1_{kj} \exp(\imath k \theta) \psi_j(r) .
\label{_auto23} \tag{25}
\end{equation}
$$

As such the Helmholtz problem is split in two smaller problems.
The two problems read with the spectral Galerkin method:

Find $u^0 \in V_F^0 \otimes V_U^N$ such that

$$
\begin{equation}
   \int_{\Omega} (-\nabla^2 u^0 + \alpha u^0) v^0 w d\sigma = \int_{\Omega} f v^0 w d\sigma, \quad \forall \, v^0 \in V_F^0 \otimes V_U^N.
\label{eq:u0} \tag{26}
   \end{equation}
$$

Find $u^1 \in V_F^1 \otimes V_D^N$ such that

$$
\begin{equation}
   \int_{\Omega} (-\nabla^2 u^1 + \alpha u^1) v^1 w d\sigma = \int_{\Omega} f v^1 w d\sigma, \quad \forall \, v^1 \in V_F^1 \otimes V_D^N.
\label{eq:u1} \tag{27}
   \end{equation}
$$

Note that integration over the domain is done using
polar coordinates with an integral measure of $d\sigma=rdrd\theta$.
However, the integral in the radial direction needs to be mapped
to $t=2r/a-1$, where $t \in [-1, 1]$, which suits the basis functions used,
see ([17](#eq:psi)). This leads to a measure of $0.5(t+1)adtd\theta$.
Furthermore, the weight $w(t)$ will be unity for the Legendre basis and
$(1-t^2)^{-0.5}$ for the Chebyshev bases.

## Implementation

A complete implementation is found in the file [unitdisc_helmholtz.py](https://github.com/spectralDNS/shenfun/blob/master/demo/unitdisc_helmholtz.py).
Here we give a brief explanation for the implementation. Start by
importing all functionality from [shenfun](https://github.com/spectralDNS/shenfun)
and [sympy](https://sympy.org), where Sympy is required for handeling the
polar coordinates. Also, we choose to work with covariant
basis vectors.

```
from shenfun import *
import sympy as sp
config['basisvectors'] = 'covariant'

# Define polar coordinates using angle along first axis and radius second
theta, r = psi = sp.symbols('x,y', real=True, positive=True)
rv = (r*sp.cos(theta), r*sp.sin(theta)) # Map to Cartesian (x, y)
```

Note that Sympy symbols are both positive and real, $\theta$ is
chosen to be along the first axis and $r$ second. This has to agree with
the next step, which is the creation of tensorproductspaces
$V_F^0 \otimes V_U^N$ and $V_F^1 \otimes V_D^N$. We use
`domain=(0, 1)` for the radial direction to get a unit disc, whereas
the default domain for the Fourier bases is already the
required $(0, 2\pi)$.

```
N = 32
F = FunctionSpace(N, 'F', dtype='d')
F0 = FunctionSpace(1, 'F', dtype='d')
L = FunctionSpace(N, 'L', bc=(0, 0), domain=(0, 1))
L0 = FunctionSpace(N, 'L', bc=(None, 0), domain=(0, 1))
T = TensorProductSpace(comm, (F, L), axes=(1, 0), coordinates=(psi, rv))
T0 = TensorProductSpace(MPI.COMM_SELF, (F0, L0), axes=(1, 0), coordinates=(psi, rv))
```

Note that since `F0` only has one component we could actually use
`L0` without creating `T0`. But the code turns out to be simpler
if we use `T0`, much because the additional $\theta$-direction is
required for the polar coordinates to apply. Using one single basis
function for the $\theta$ direction is as such a generic way to handle
polar 1D problems (i.e., problems that are only functions of the
radial direction, but still using polar coordinates).
Also note that `F` is created using the entire range of wavenumbers
even though it should not include wavenumber 0.
As such we need to make sure that the coefficient created for
$k=0$ (i.e., $\hat{u}^1_{0,j}$) will be exactly zero.
Finally, note that
`T0` is not distributed with MPI, which is accomplished using
`MPI.COMM_SELF` instead of `comm` (which equals `MPI.COMM_WORLD`).
The purely radial problem ([26](#eq:u0)) is only solved on the one
processor with rank = 0.

Polar coordinates are ensured by feeding `coordinates=(psi, rv)`
to [TensorProductSpace](https://shenfun.readthedocs.io/en/latest/shenfun.html#shenfun.tensorproductspace.TensorProductSpace). Operators like [div()](https://shenfun.readthedocs.io/en/latest/shenfun.forms.html#shenfun.forms.operators.div)
[grad()](https://shenfun.readthedocs.io/en/latest/shenfun.forms.html#shenfun.forms.operators.grad) and  [curl()](https://shenfun.readthedocs.io/en/latest/shenfun.forms.html#shenfun.forms.operators.curl) will now work on
items of [Function](https://shenfun.readthedocs.io/en/latest/shenfun.forms.html#shenfun.forms.arguments.Function), [TestFunction](https://shenfun.readthedocs.io/en/latest/shenfun.forms.html#shenfun.forms.arguments.TestFunction) and
[TrialFunction](https://shenfun.readthedocs.io/en/latest/shenfun.forms.html#shenfun.forms.arguments.TrialFunction) using a polar coordinate system.

To define the equations ([26](#eq:u0)) and ([27](#eq:u1)) we first declare
these test- and trialfunctions, and then use code that
is remarkably similar to the mathematics.

```
v = TestFunction(T)
u = TrialFunction(T)
v0 = TestFunction(T0)
u0 = TrialFunction(T0)
alpha = 1

mats = inner(v, -div(grad(u))+alpha*u)
if comm.Get_rank() == 0:
    mats0 = inner(v0, -div(grad(u0))+alpha*u0)
```

Here `mats` and `mats0` will contain several tensor product
matrices in the form of
[TPMatrix](https://shenfun.readthedocs.io/en/latest/shenfun.html#shenfun.matrixbase.TPMatrix). Since there is only one non-periodic direction
the matrices can be easily solved using [SolverGeneric1ND](https://shenfun.readthedocs.io/en/latest/shenfun.html#shenfun.la.SolverGeneric1ND).
But first we need to define the function $f(\theta, r)$.
To this end we use sympy and the method of
manufactured solution to define a possible solution `ue`,
and then compute `f` exactly using exact differentiation

```

# Manufactured solution
ue = (r*(1-r))**2*sp.cos(8*theta)-0.1*(r-1)
#f = -ue.diff(r, 2) - (1/r)*ue.diff(r, 1) - (1/r**2)*ue.diff(theta, 2) + alpha*ue
f = (-div(grad(u))+alpha*u).tosympy(basis=ue, psi=psi)

# Compute the right hand side on the quadrature mesh
fj = Array(T, buffer=f)

# Take scalar product
f_hat = Function(T)
f_hat = inner(v, fj, output_array=f_hat)
if T.local_slice(True)[0].start == 0: # The processor that owns k=0
    f_hat[0] = 0

# For k=0 we solve only a 1D equation. Do the scalar product for Fourier
# coefficient 0 by hand (or sympy).
if comm.Get_rank() == 0:
    f0_hat = Function(T0)
    gt = sp.lambdify(r, sp.integrate(f, (theta, 0, 2*sp.pi))/2/sp.pi)(L0.mesh())
    f0_hat = T0.scalar_product(gt, f0_hat)
```

Note that for $u^0$ we perform the interal in the $\theta$ direction
exactly using sympy. This is necessary since one Fourier coefficient
is not sufficient to do this integral numerically. For the $u^1$
case we do the integral numerically as part of the [inner()](https://shenfun.readthedocs.io/en/latest/shenfun.forms.html#shenfun.forms.inner.inner) product.
With the correct right hand side assembled we can solve the
linear system of equations

```
u_hat = Function(T)
Sol1 = la.SolverGeneric1ND(mats)
u_hat = Sol1(f_hat, u_hat)

# case k = 0
u0_hat = Function(T0)
if comm.Get_rank() == 0:
    Sol0 = la.SolverGeneric1ND(mats0)
    u0_hat = Sol0(f0_hat, u0_hat)
comm.Bcast(u0_hat, root=0)
```

Having found the solution in spectral space all that is
left is to transform it back to real space.

```
# Transform back to real space. Broadcast 1D solution
sl = T.local_slice(False)
uj = u_hat.backward() + u0_hat.backward()[:, sl[1]]
```

## Postprocessing
The solution can now be compared with the exact solution
through

```
uq = Array(T, buffer=ue)
print('Error =', np.linalg.norm(uj-uq))
```

We can also get the gradient of the solution. For this we need
a space without boundary conditions, and a vector space

```
TT = T.get_orthogonal()
V = VectorSpace(TT)
```

Notice that we do not have the solution in one single space
in spectral space, since it is a combination of `u_hat` and
`u0_hat`. For this reason we first transform the solution from
real space `uj` to the new orthogonal space `TT`

```
ua = Array(TT, buffer=uj)
uh = ua.forward()
```

With the solution as a [Function](https://shenfun.readthedocs.io/en/latest/shenfun.forms.html#shenfun.forms.arguments.Function) we can simply project
the gradient to `V`

```
dv = project(grad(uh), V)
du = dv.backward()
```

Note that the gradient `du` now contains the contravariant components
of the covariant basis vector `b`. The basis vector `b` is not normalized
(it's length is not unity), because we have set
`config['basisvectors']='covariant'`. The basisvectors can
be seen as

```
from IPython.display import Math
Math(T.coors.latex_basis_vectors(symbol_names={theta: '\\theta', r: 'r'}))
```

and we see that they are given in terms of the Cartesian unit vectors.
The gradient we have computed is (and yes, it should be $r^2$ because we
do not have unit vectors)

$$
\begin{equation}
\nabla u = \underbrace{\frac{1}{r^2}\frac{\partial u}{\partial \theta}}_{du[0]}\mathbf{b}_{\theta} + \underbrace{\frac{\partial u}{\partial r}}_{du[1]} \mathbf{b}_{r}
\label{eq:gradu} \tag{28}
\end{equation}
$$

Now it makes sense to plot the solution and its gradient in Cartesian
instead of computational coordinates. To this end we need to
project the gradient to a Cartesian basis

$$
\begin{align*}
\frac{\partial u}{\partial x} &= \nabla u \cdot \mathbf{i},\\
\frac{\partial u}{\partial y} &= \nabla u \cdot \mathbf{j}.
\end{align*}
$$

We compute the Cartesian gradient by assembling ([28](#eq:gradu))
on the computational grid

```
ui, vi = TT.local_mesh(True)
b = T.coors.get_covariant_basis()
bij = np.array(sp.lambdify(psi, b)(ui, vi))
gradu = du[0]*bij[0] + du[1]*bij[1]
```

Because of the way the vectors are stored, `gradu[0]` will now
contain $\nabla u \cdot \mathbf{i}$ and
`gradu[1]` will contain $\nabla u \cdot \mathbf{j}$.

To validate the gradient we compute the $L^2$ error norm

$$
\sqrt{\int_{\Omega} |\nabla u - \nabla u_e|^2 d\sigma}
  = \sqrt{\int_{\theta=0}^{2\pi}\int_{r=0}^{1} \left(\left(\frac{1}{r^2}\frac{\partial u-u_e}{\partial \theta}\right)^2\mathbf{b}_{\theta}\cdot \mathbf{b}_{\theta} + \left(\frac{\partial u-u_e}{\partial r}\right)^2\mathbf{b}_{r}\cdot \mathbf{b}_{r} \right)rd\theta dr}
$$

implemented as

```
gradue = Array(V, buffer=grad(u).tosympy(basis=ue, psi=psi))
gij = T.coors.get_covariant_metric_tensor()
ui, vi = TT.local_mesh(True, uniform=True)
# Evaluate metric on computational mesh
g = np.array(sp.lambdify(psi, gij)(ui, vi), dtype=object)
# Compute L2 error
errorg = inner(1, (du[0]-gradue[0])**2*g[0, 0]+ (du[1]-gradue[1])**2*g[1, 1])
print('Error gradient', np.sqrt(errorg))
```

We now refine the solution to make it look better,
and plot on the unit disc.

```
%matplotlib inline

u_hat2 = u_hat.refine([N*3, N*3])
u0_hat2 = u0_hat.refine([1, N*3])
sl = u_hat2.function_space().local_slice(False)
ur = u_hat2.backward() + u0_hat2.backward()[:, sl[1]]

# Wrap periodic plot around since it looks nicer
xx, yy = u_hat2.function_space().local_cartesian_mesh()
xp = np.vstack([xx, xx[0]])
yp = np.vstack([yy, yy[0]])
up = np.vstack([ur, ur[0]])
# For vector no need to wrap around and no need to refine:
xi, yi = TT.local_cartesian_mesh()

# plot
import matplotlib.pyplot as plt
plt.figure()
plt.contourf(xp, yp, up)
plt.quiver(xi, yi, gradu[0], gradu[1], scale=40, pivot='mid', color='white')
plt.colorbar()
plt.title('Helmholtz - unitdisc')
plt.xticks([])
plt.yticks([])
plt.axis('off')
plt.show()
```

1.  **J. Shen**.
    Efficient Spectral-Galerkin Methods III: Polar and Cylindrical Geometries,
    *SIAM Journal on Scientific Computing*,
    18(6),
    pp. 1583-1604,
    [doi: 10.1137/S1064827595295301](http://dx.doi.org/10.1137/S1064827595295301),
    1997.


# Linearized Dynamics Near Steady-State

We start from the linearized state-equations of the passively mode-locked laser [ref1] [ref1]

[ref1]: https://github.com/adrianschlatter/notebooks/blob/master/Dynamics%20of%20Passively%20Mode-Locked%20Lasers.ipynb

```
import sympy as sym
sym.init_printing(use_latex='mathjax')

Pst, gst, EsatL, dqPdEP = sym.symbols("P_{st} g_{st} E_{satL} q'_P")
TR, tauL, etaP, Toc, PP0 = sym.symbols('T_R tau_L eta_P T_{oc} P_{P0}')
alpha, beta, gamma, epsilon = sym.symbols('alpha beta gamma epsilon')
w0 = sym.symbols('omega_0')

Pst_ = sym.Eq(Pst, -EsatL / tauL + etaP * PP0 / gst)

x = sym.MatrixSymbol('x', 2, 1)
xdot = sym.MatrixSymbol('xdot', 2, 1)
y = sym.MatrixSymbol('y', 1, 1)
u = sym.MatrixSymbol('u', 1, 1)

A = sym.Matrix([[-Pst * dqPdEP, Pst / TR],
                [-gst / EsatL, -1 / tauL - Pst / EsatL]])
B = sym.Matrix([[0], [etaP / EsatL]])
C = sym.Matrix([[Toc, 0]])

sym.Eq(xdot, A * x + B * u)
```

$$\dot{x} = \left[\begin{matrix}0\\\frac{\eta_{P}}{E_{{satL}}}\end{matrix}\right] u + \left[\begin{matrix}- P_{{st}} q'_{P} & \frac{P_{{st}}}{T_{R}}\\- \frac{g_{{st}}}{E_{{satL}}} & - \frac{1}{\tau_{L}} - \frac{P_{{st}}}{E_{{satL}}}\end{matrix}\right] x$$

```
sym.Eq(y, C * x)
```

$$y = \left[\begin{matrix}T_{{oc}} & 0\end{matrix}\right] x$$

where $\vec{x} = \left[ \delta P, \delta g \right]^T$, $\vec u = \left[\delta P_P \right]$, and $\vec y = \left[ \delta P_{out} \right]$.

Now, we rewrite this in new state variables $\vec x' = \left[ \delta \dot{P} / \omega_0, \delta P \right]$, with

$\omega_0 = \sqrt{ \frac{P_{st} g_{st}}{T_R E_{satL}} + P_{st} q'_P \left( 1 / \tau_L + P_{st} / E_{satL} \right)}$.

From the state equation above, we immediately see the transformation matrix $M$:

```
M = sym.Matrix([[-Pst * dqPdEP / w0, Pst / TR / w0],
                [1, 0]])
w0_ = sym.Eq(w0, sym.sqrt(-A[0, 1] * A[1, 0] + A[0, 0] * A[1, 1]))
xp = sym.MatrixSymbol("x'", 2, 1)
sym.simplify(sym.Eq(xp, M * x))
```

$$x' = \left[\begin{matrix}- \frac{P_{{st}} q'_{P}}{\omega_{0}} & \frac{P_{{st}}}{T_{R} \omega_{0}}\\1 & 0\end{matrix}\right] x$$

The transformed state equation is:

$ \begin{align*}
\dot{x'} &= M A M^{-1} x' + M B u \\
y &= C \cdot x = C M^{-1} x'
\end{align*}$

```
M = M.subs(w0_.lhs, w0_.rhs)

Bp = M * B
Ap = M * A * M**-1
sym.Eq(sym.MatrixSymbol("A'", 2, 2), sym.simplify(Ap))
```

$$A' = \left[\begin{matrix}- P_{{st}} q'_{P} - \frac{1}{\tau_{L}} - \frac{P_{{st}}}{E_{{satL}}} & - \sqrt{\frac{P_{{st}} q'_{P}}{\tau_{L}} + \frac{P_{{st}}^{2} q'_{P}}{E_{{satL}}} + \frac{P_{{st}} g_{{st}}}{E_{{satL}} T_{R}}}\\\sqrt{\frac{P_{{st}} q'_{P}}{\tau_{L}} + \frac{P_{{st}}^{2} q'_{P}}{E_{{satL}}} + \frac{P_{{st}} g_{{st}}}{E_{{satL}} T_{R}}} & 0\end{matrix}\right]$$

```
sym.Eq(sym.MatrixSymbol("B'", 2, 1), sym.simplify(Bp))
```

$$B' = \left[\begin{matrix}\frac{P_{{st}} \eta_{P}}{E_{{satL}} T_{R} \sqrt{\frac{P_{{st}}}{E_{{satL}} T_{R} \tau_{L}} \left(T_{R} q'_{P} \left(E_{{satL}} + P_{{st}} \tau_{L}\right) + g_{{st}} \tau_{L}\right)}}\\0\end{matrix}\right]$$

We recognize that the off-diagonal elements of $A$ are $\omega_0$ and $-\omega_0$. Furthermore, we define

$ \begin{align*}
\zeta &= \frac{1}{2 \omega_0} \left(1 / \tau_L + P_{st} \cdot \left(q'_P + 1 / E_{satL} \right) \right) \\
\rho &= \frac{P_{st} \eta_P}{E_{satL} T_R \omega_0^2}
\end{align*} $

```
zeta = sym.Symbol('zeta', real=True)
zeta_ = sym.Eq(zeta, (1 / tauL + Pst * (dqPdEP + 1 / EsatL)) / 2 / w0)
rho = sym.Symbol('rho', real=True)
rho_ = sym.Eq(rho, Pst * etaP / (EsatL * TR * w0) / w0)

Bpn = sym.Matrix(Bp)
Apn = sym.Matrix(Ap)

Bpn[0, 0] = Bp[0, 0] / (w0_.lhs / w0_.rhs) * rho_.lhs / rho_.rhs

Apn[0, 0] = Ap[0, 0] * zeta_.lhs / zeta_.rhs
Apn[0, 1] = Ap[0, 1] * w0_.lhs / w0_.rhs
Apn[1, 0] = Ap[1, 0] * w0_.lhs / w0_.rhs

sym.Eq(sym.MatrixSymbol("x'dot", 2, 1), sym.simplify(Apn * xp + Bpn * u))
```

$$\dot{x'} = \left[\begin{matrix}\omega_{0} \rho\\0\end{matrix}\right] u + \left[\begin{matrix}- 2 \omega_{0} \zeta & - \omega_{0}\\\omega_{0} & 0\end{matrix}\right] x'$$

Looking at the eigenvalues (see below), we see that the system

- is stable if $\zeta > 0$
- is unstable if $\zeta < 0$ (always)
- has oscillatory modes for $-1 < \zeta < 1$ (radicand is negative)
- is critically damped for $\zeta = 1$
- is over-critically damped for $\zeta > 1$ (radicand is positive)

Of course, this analysis is only true if $\omega_0$ is real.

```
Apn.eigenvals()
```

$$\left \{ \omega_{0} \left(- \zeta - \sqrt{\left(\zeta - 1\right) \left(\zeta + 1\right)}\right) : 1, \quad \omega_{0} \left(- \zeta + \sqrt{\left(\zeta - 1\right) \left(\zeta + 1\right)}\right) : 1\right \}$$

Now, we apply a unit-step to the input of the system and determine where the output will stabilize. To do that, we have to solve the following equation:

```
eq = sym.Eq(sym.Matrix([[0], [0]]), Apn * xp + Bpn)
sym.simplify(eq)
```

$$\left[\begin{matrix}0\\0\end{matrix}\right] = \left[\begin{matrix}\omega_{0} \rho\\0\end{matrix}\right] + \left[\begin{matrix}- 2 \omega_{0} \zeta & - \omega_{0}\\\omega_{0} & 0\end{matrix}\right] x'$$

From the resulting $x'$ we obtain for the steady-state $y$ response:

```
sym.simplify(C * M**-1 * -Apn**-1 * Bpn)[0]
```

$$T_{{oc}} \rho$$

The DC-response to small pump-power changes is the slope efficiency.

```
sym.simplify((Toc * rho_.rhs).subs([(w0_.lhs, w0_.rhs),
                                    (Pst_.lhs, Pst_.rhs)]))
```

$$\frac{T_{{oc}} \eta_{P} g_{{st}}}{P_{{P0}} T_{R} \eta_{P} q'_{P} + g_{{st}}^{2}}$$

For $q'_P = 0$ (no saturable absorber), this simplifies to $\eta_P \cdot \frac{T_{oc}}{g_{st}}$, i.e. the slope efficiency depends only on the pump-absorption efficiency and the ratio of output coupling to total losses.

Passively Mode-Locked Laser - Steady-State Dynamics.ipynb


# Rearranging the system matrices from  (Lin, Parker, 1999)

[1] Lin, J., & Parker, R. G. (1999). Analytical characterization of the unique properties of planetary gear free vibration. Journal of Vibration and Acoustics, Transactions of the ASME, 121(3), 316–321. http://doi.org/10.1115/1.2893982 \
[2] Cooley, C. G., & Parker, R. G. (2012). Vibration Properties of High-Speed Planetary Gears With Gyroscopic Effects. Journal of Vibration and Acoustics, 134(6). http://doi.org/10.1115/1.4006646

```python
from sympy import *
init_printing()

def symb(x,y):
    return symbols('{0}_{1}'.format(x,y), type = float)
```


# McIntyre Problem 10.20
For an infinite square well subject to the perturbation,
$$
H' = \gamma x(L-x)
$$
calculate the first order energy corrections.

```python
# load the sympy library
from sympy import *
```

```python
# declare some symbols
# being as specific as possible at this step with help things simplify below
x, γ = symbols("x γ", real=True)
n, m = symbols("n m", integer=True, positive=True) # state labels
ħ, M, L = symbols("ħ M L", real=True, positive=True) # constants (big M is mass)
π = pi # use a pretty pi
```

```python
# define the unperturbed eigenenergies as functions of n
# "lambda" is a shortcut to define simple functions
# the syntax is lambda (variables): some one-line expression
E0 = lambda n: n**2*π**2*ħ**2/(2*M*L**2)
```

```python
# check the definition
E0(n)
```

$\displaystyle \frac{\pi^{2} n^{2} ħ^{2}}{2 L^{2} M}$

```python
# define the unperturbed eigenstates as functions of n and x
# note that we're not checking the domain, so the user must ensure 0 < x < L
ϕ0 = lambda n,x: sqrt(2/L)*sin(n*π*x/L)
```

```python
# check the definition
ϕ0(n,x)
```

$\displaystyle \frac{\sqrt{2} \sin{\left(\frac{\pi n x}{L} \right)}}{\sqrt{L}}$

```python
# define the perturbation as a function of x
Hp = lambda x: γ*x*(L-x)
```

```python
# check
Hp(x)
```

$\displaystyle x γ \left(L - x\right)$

## First-order energy corrections
I'll define a function to setup a matrix element integral
$$
\langle n \vert H' \vert m \rangle = \int\limits_0^L \phi^{(0)}_n(x)^*\,H'\,\phi^{(0)}_m(x)\,dx
$$

```python
# define a matrix element integral of the operator Hp w.r.t. the unperturbed eigenstates |n> and |m>
# Note, I could have also used the lambda notation like above
def HpMatrixElement(n,m):
    # because the eigenfunctions are real functions, I'll skip the complex conjugate
    # the form of the Integral command is integrate(function, (integration variable, lower limit, upper limit))
    # note that the limits may include infinity, written as oo (two letters small o)
    return Integral(ϕ0(n,x)*Hp(x)*ϕ0(m,x),(x,0,L))
```

```python
# check our definition
HpMatrixElement(n,m)
```

$\displaystyle \int\limits_{0}^{L} \frac{2 x γ \left(L - x\right) \sin{\left(\frac{\pi m x}{L} \right)} \sin{\left(\frac{\pi n x}{L} \right)}}{L}\, dx$

```python
# Calculate the 1st order energy correction
# we perform the integral using the "doit()" method
E1 = lambda n: HpMatrixElement(n,n).doit()
```

```python
# display the result
E1(n)
```

$\displaystyle \frac{2 γ \left(\frac{L^{3}}{12} + \frac{L^{3}}{4 \pi^{2} n^{2}}\right)}{L}$

## First-order state corrections
The first order state corrections require an infinite sum.
$$
\vert n^{(1)} \rangle = \sum\limits_{m\neq n} \frac{\langle m^{(0)}\vert H'\vert n^{(0)}\rangle}{E^{(0)}_n-E^{(0)}_m}\,\vert m^{(0)} \rangle
$$

```python
# Let's try the matrix element first
HpMatrixElement(n,m).doit()
```

$\displaystyle \begin{cases} - \frac{2 γ \left(\frac{2 \pi L^{3} m^{3} n}{\pi^{3} m^{6} - 3 \pi^{3} m^{4} n^{2} + 3 \pi^{3} m^{2} n^{4} - \pi^{3} n^{6}} - \frac{2 \pi L^{3} m n^{3}}{\pi^{3} m^{6} - 3 \pi^{3} m^{4} n^{2} + 3 \pi^{3} m^{2} n^{4} - \pi^{3} n^{6}}\right)}{L} + \frac{2 γ \left(- \frac{2 \left(-1\right)^{m} \left(-1\right)^{n} \pi L^{3} m^{3} n}{\pi^{3} m^{6} - 3 \pi^{3} m^{4} n^{2} + 3 \pi^{3} m^{2} n^{4} - \pi^{3} n^{6}} + \frac{2 \left(-1\right)^{m} \left(-1\right)^{n} \pi L^{3} m n^{3}}{\pi^{3} m^{6} - 3 \pi^{3} m^{4} n^{2} + 3 \pi^{3} m^{2} n^{4} - \pi^{3} n^{6}}\right)}{L} & \text{for}\: m \neq n \\\frac{L^{2} γ}{6} + \frac{L^{2} γ}{2 \pi^{2} n^{2}} & \text{otherwise} \end{cases}$

We get two answers, one for $m=n$ (which we don't need here) and one for $m\neq n$. We can extract the useful one using a hack.  Piecewise functions are actually stored in memory as a nested list called `args`.  The piece we want is the first element of the first sublist, which we specify as `[0][0]`.  The code looks like

```python
HpMatrixElement(n,m).doit().args[0][0]
```

$\displaystyle - \frac{2 γ \left(\frac{2 \pi L^{3} m^{3} n}{\pi^{3} m^{6} - 3 \pi^{3} m^{4} n^{2} + 3 \pi^{3} m^{2} n^{4} - \pi^{3} n^{6}} - \frac{2 \pi L^{3} m n^{3}}{\pi^{3} m^{6} - 3 \pi^{3} m^{4} n^{2} + 3 \pi^{3} m^{2} n^{4} - \pi^{3} n^{6}}\right)}{L} + \frac{2 γ \left(- \frac{2 \left(-1\right)^{m} \left(-1\right)^{n} \pi L^{3} m^{3} n}{\pi^{3} m^{6} - 3 \pi^{3} m^{4} n^{2} + 3 \pi^{3} m^{2} n^{4} - \pi^{3} n^{6}} + \frac{2 \left(-1\right)^{m} \left(-1\right)^{n} \pi L^{3} m n^{3}}{\pi^{3} m^{6} - 3 \pi^{3} m^{4} n^{2} + 3 \pi^{3} m^{2} n^{4} - \pi^{3} n^{6}}\right)}{L}$

```python
# Let's try to simplify
_.simplify()
```

$\displaystyle - \frac{4 L^{2} m n γ \left(1 - \left(-1\right)^{m + n + 1}\right)}{\pi^{2} \left(m^{4} - 2 m^{2} n^{2} + n^{4}\right)}$

Note that this is equal to zero if m+n is odd. we can use that fact to simplify the sum later

```python
# print this as code so that we can copy and paste it into a new function
str(_)
```

    '-4*L**2*m*n*γ*(1 - (-1)**(m + n + 1))/(pi**2*(m**4 - 2*m**2*n**2 + n**4))'

```python
# copy and paste from above
HpOffDiagonal = lambda n, m: -4*L**2*m*n*γ*(1 - (-1)**(m + n + 1))/(pi**2*(m**4 - 2*m**2*n**2 + n**4))
```

I'm going to attempt to evaluate the sums inside $\vert n^{(1)}\rangle$.  I know these don't simplify, but I just want to show that if you *really* need to do this in python you can.  This is very much a brute-force solution made by trial and error.  There is probably a more eligant way to code it.  You may skip down to **Results** below if you just want to see the first few terms of the sums.

To actually do the sum we have to break it into pieces $m<n$ and $m>m$.  This gives us special cases when $n=1$ (no < sum) and $n=2$ (only one term in the < sum).  I also break the function into n=odd and n=even cases.

```python
k = symbols("k", integer=True) # new summation variable
```

```python
def ϕ1(n,x,mmax=oo):
    # mmax is an optional parameter to terminate the sum at a maximum value of m
    # calculate the kmax values for below
    if mmax==oo:
        kmax = oo
    elif ask(Q.odd(mmax)):
        kmax = (mmax-1)/Integer(2)
    else:
        kmax = mmax/Integer(2) # FIXME for odd n this gives one extra term

    if n==1:
        # we only need to sum the odd terms, so I'll sub m=2k+1.  The sum will start at k = 1
        return Sum(HpOffDiagonal(n,2*k+1)/(E0(n)-E0(2*k+1))*ϕ0(2*k+1,x),(k,1,kmax))
    elif n==2:
        # here we only need to sum the even terms, so I'll sub m=2k.  The sum will start from k=2
        return Sum(HpOffDiagonal(n,2*k)/(E0(n)-E0(2*k))*ϕ0(2*k,x),(k,2,kmax))
    elif ask(Q.odd(n)): # odd values of n>1
        return Sum(HpOffDiagonal(n,2*k+1)/(E0(n)-E0(2*k+1))*ϕ0(2*k+1,x),(k,0,(n-3)/Integer(2)))+Sum(HpOffDiagonal(n,2*k+1)/(E0(n)-E0(2*k+1))*ϕ0(2*k+1,x),(k,(n+1)/Integer(2),kmax))
    else: # even values of n>2
        return Sum(HpOffDiagonal(n,2*k)/(E0(n)-E0(2*k))*ϕ0(2*k,x),(k,1,(n-2)/Integer(2)))+Sum(HpOffDiagonal(n,2*k)/(E0(n)-E0(2*k))*ϕ0(2*k,x),(k,(n+2)/Integer(2),kmax))
```

```python
ϕ1(1,x).simplify()
```

$\displaystyle \frac{\sqrt{2} L^{\frac{7}{2}} M γ \sum_{k=1}^{\infty} \frac{\left(2 + \frac{1}{k}\right) \sin{\left(\frac{\pi x \left(2 k + 1\right)}{L} \right)}}{k^{2} \left(k^{3} + 3 k^{2} + 3 k + 1\right)}}{4 \pi^{4} ħ^{2}}$

```python
# The result doesn't simplify (as expected), but we see the key terms
ϕ1(2,x).simplify()
```

$\displaystyle - \frac{\sqrt{2} L^{\frac{7}{2}} M γ \sum_{k=2}^{\infty} \frac{k \sin{\left(\frac{2 \pi k x}{L} \right)}}{- k^{6} + 3 k^{4} - 3 k^{2} + 1}}{\pi^{4} ħ^{2}}$

```python
ϕ1(3,x).simplify()
```

$\displaystyle \frac{3 \sqrt{2} L^{\frac{7}{2}} M γ \left(- \frac{\sin{\left(\frac{\pi x}{L} \right)}}{8} + \sum_{k=2}^{\infty} \frac{\left(2 k + 1\right) \sin{\left(\frac{\pi x \left(2 k + 1\right)}{L} \right)}}{k^{6} + 3 k^{5} - 3 k^{4} - 11 k^{3} + 6 k^{2} + 12 k - 8}\right)}{4 \pi^{4} ħ^{2}}$

```python
ϕ1(4,x).simplify()
```

$\displaystyle \frac{2 \sqrt{2} L^{\frac{7}{2}} M γ \left(- \frac{\sin{\left(\frac{2 \pi x}{L} \right)}}{27} + \sum_{k=3}^{\infty} \frac{k \sin{\left(\frac{2 \pi k x}{L} \right)}}{k^{6} - 12 k^{4} + 48 k^{2} - 64}\right)}{\pi^{4} ħ^{2}}$

By specifying `mmax` we can start to get approximate functions we can evaluate.

*Question:* We could in theory estimate the size of the 2nd order corrections. Then we only need a few terms in the sum where *m* is near *n* to have a numerical result consistent with the order of the approximation.  How many terms do we need?

### Results

```python
ϕ1(1,x,10-1).doit()
```

$\displaystyle - \frac{8 \sqrt{2} L^{\frac{3}{2}} γ \left(- \frac{3 L^{2} M \sin{\left(\frac{3 \pi x}{L} \right)}}{256 \pi^{2} ħ^{2}} - \frac{5 L^{2} M \sin{\left(\frac{5 \pi x}{L} \right)}}{6912 \pi^{2} ħ^{2}} - \frac{7 L^{2} M \sin{\left(\frac{7 \pi x}{L} \right)}}{55296 \pi^{2} ħ^{2}} - \frac{9 L^{2} M \sin{\left(\frac{9 \pi x}{L} \right)}}{256000 \pi^{2} ħ^{2}}\right)}{\pi^{2}}$

```python
# now add back in the unperturbed state to get the final result
(ϕ0(1,x)+ϕ1(1,x,10-1).doit()).simplify()
```

$\displaystyle \frac{\sqrt{2} \left(\frac{L^{4} M γ \left(81000 \sin{\left(\frac{3 \pi x}{L} \right)} + 5000 \sin{\left(\frac{5 \pi x}{L} \right)} + 875 \sin{\left(\frac{7 \pi x}{L} \right)} + 243 \sin{\left(\frac{9 \pi x}{L} \right)}\right)}{864000} + \pi^{4} ħ^{2} \sin{\left(\frac{\pi x}{L} \right)}\right)}{\pi^{4} \sqrt{L} ħ^{2}}$

```python
# sub in some numbers and get a function
values = [(ħ,1),(M,1),(L,1),(γ,20)] # values for our constants
N((ϕ0(1,x)+ϕ1(1,x,10-1).doit()).subs(values)) # The N command returns a decimal approximation
```

$\displaystyle 1.4142135623731 \sin{\left(\pi x \right)} + 0.0272217962543552 \sin{\left(3 \pi x \right)} + 0.00168035779347872 \sin{\left(5 \pi x \right)} + 0.000294062613858775 \sin{\left(7 \pi x \right)} + 8.16653887630656 \cdot 10^{-5} \sin{\left(9 \pi x \right)}$

Notice how small the added terms are, even with a large value of $\gamma$.

```python
# sub in some numbers and plot it
p=plot(ϕ0(1,x).doit().subs(values),(x,0,1),line_color="blue",show=False,label="unpert.")
p.extend(plot((ϕ0(1,x)+ϕ1(1,x,10).doit()).subs(values),(x,0,1),line_color="red",show=False,label="1st ord."))
p.legend = True
p.ylabel = "ϕ(x)"
p.show()
```

```python
# compare the energies for these values of the constants
N(E0(1).subs(values))
```

$\displaystyle 4.93480220054468$

```python
N((E0(1)+E1(1)).subs(values))
```

$\displaystyle 9.28134737030139$

Notice that the energies change by a lot even though the wavefunction does not.

```python

```


```python
from sympy import *
init_printing()
```

```python
def skew(l):
    l1, l2, l3 = l
    return Matrix([
        [0, -l3, l2],
        [l3, 0, -l1],
        [-l2, l1, 0]
    ])
```

```python
# define state variables
x, y, z, phi, theta, psi, u, v, w, p, q, r = symbols('x y z phi theta psi u v w p q r', real=True)
s = Matrix([x, y, z, phi, theta, psi, u, v, w, p, q, r])

# position and orientation
eta = Matrix([x, y, z, phi, theta, psi])
nu = Matrix([u, v, w, p, q, r])

# centre of gravity
xg, yg, zg = symbols('xg yg zg', real=True)
rg = Matrix([xg, yg, zg])

# centre of bouyancy
xb, yb, zb = symbols('xb yb zb', real=True)
rb = Matrix([xb, yb, zb])

# center of pressure
xcp, ycp, zcp = symbols('xcp ycp zcp', real=True)
rcp = Matrix([xcp, ycp, zcp])

# mass matrix
m = symbols('m', real=True, positive=True)
Ixx, Iyy, Izz = symbols('Ixx Iyy Izz')
I0 = diag(Ixx, Iyy, Izz)
M = BlockMatrix([
    [m*eye(3), -m*skew(rg)],
    [m*skew(rg), I0]
])
M = Matrix(M)
# M = simplify(M)

# Coriolis and centripetal matrix
nu1 = Matrix([u, v, w])
nu2 = Matrix([p, q, r])
crb = BlockMatrix([
    [zeros(3), -m*skew(nu1)-m*skew(nu2)*skew(rg)],
    [-m*skew(nu1)+m*skew(rg)*skew(nu2), -skew(I0*nu2)]
])
crb = Matrix(crb)
# crb = simplify(crb)

# damping matrix
Xuu, Yvv, Zww, Kpp, Mqq, Nrr = symbols(
    'Xuu Yvv Zww Kpp Mqq Nrr', real=True
)
D = Matrix([
    [Xuu*abs(u), 0, 0, 0, 0, 0],
    [0, Yvv*abs(v), 0, 0, 0, 0],
    [0, 0, Zww*abs(w), 0, 0, 0],
    [0, -zcp*Yvv*abs(v), ycp*Zww*abs(w), Kpp*abs(p), 0, 0],
    [zcp*Xuu*abs(u), 0, -xcp*Zww*abs(w), 0, Mqq*abs(q), 0],
    [-ycp*Xuu*abs(u), xcp*Yvv*abs(v), 0, 0, 0, Nrr*abs(r)]
])

# rotational transform between body and NED quaternions
# Tq = Rational(1,2)*Matrix([
#     [-eps1, -eps2, -eps3],
#     [eta0, -eps3, eps2],
#     [eps3, eta0, -eps1],
#     [-eps2, eps1, eta0]
# ])
T_euler = Matrix([
    [1, sin(phi)*tan(theta), cos(phi)*tan(theta)],
    [0, cos(phi), -sin(phi)],
    [0, sin(phi)/cos(theta), cos(phi)/cos(theta)]
])

# Rq = Matrix([
#     [1-2*(eps2**2+eps3**2), 2*(eps1*eps2-eps3*eta0), 2*(eps1*eps3+eps2*eta0)],
#     [2*(eps1*eps2+eps3*eta0), 1-2*(eps1**2+eps3**2), 2*(eps2*eps3-eps1*eta0)],
#     [2*(eps1*eps3-eps2*eta0), 2*(eps2*eps3+eps1*eta0), 1-2*(eps1**2+eps2**2)]
# ])
R_euler = Matrix([
    [cos(psi)*cos(theta), -sin(psi)*cos(phi)+cos(psi)*sin(theta)*sin(phi), sin(psi)*sin(phi)+cos(psi)*cos(phi)*sin(theta)],
    [sin(psi)*cos(theta), cos(psi)*cos(phi)+sin(phi)*sin(theta)*sin(psi), -cos(psi)*sin(phi)+sin(theta)*sin(psi)*cos(phi)],
    [-sin(theta), cos(theta)*sin(phi), cos(theta)*cos(phi)]
])

# Jeta = BlockMatrix([
#     [Rq, zeros(3)],
#     [zeros(4,3), Tq]
# ])
Jeta = BlockMatrix([
    [R_euler, zeros(3)],
    [zeros(3), T_euler]
])

Jeta = Matrix(Jeta)
# Jeta = simplify(Jeta)

# bouyancy in quaternions
# W, B = symbols('W B', real=True)
# fg = Matrix([0, 0, W])
# fb = Matrix([0, 0, -B])
# Rqinv = Rq.inv()
# geta = Matrix([
#     Rqinv*(fg+fb),
#     skew(rg)*Rqinv*fg + skew(rb)*Rqinv*fb
# ])
W, B = symbols('W B', real=True)
fg = Matrix([0, 0, W])
fb = Matrix([0, 0, -B])
R_euler_inv = R_euler.inv()
geta = Matrix([
    R_euler_inv*(fg+fb),
    skew(rg)*R_euler_inv*fg + skew(rb)*R_euler_inv*fb
])
# geta = simplify(geta)
```

```python
Jeta
```

$\displaystyle \left[\begin{matrix}\cos{\left(\psi \right)} \cos{\left(\theta \right)} & \sin{\left(\phi \right)} \sin{\left(\theta \right)} \cos{\left(\psi \right)} - \sin{\left(\psi \right)} \cos{\left(\phi \right)} & \sin{\left(\phi \right)} \sin{\left(\psi \right)} + \sin{\left(\theta \right)} \cos{\left(\phi \right)} \cos{\left(\psi \right)} & 0 & 0 & 0\\\sin{\left(\psi \right)} \cos{\left(\theta \right)} & \sin{\left(\phi \right)} \sin{\left(\psi \right)} \sin{\left(\theta \right)} + \cos{\left(\phi \right)} \cos{\left(\psi \right)} & - \sin{\left(\phi \right)} \cos{\left(\psi \right)} + \sin{\left(\psi \right)} \sin{\left(\theta \right)} \cos{\left(\phi \right)} & 0 & 0 & 0\\- \sin{\left(\theta \right)} & \sin{\left(\phi \right)} \cos{\left(\theta \right)} & \cos{\left(\phi \right)} \cos{\left(\theta \right)} & 0 & 0 & 0\\0 & 0 & 0 & 1 & \sin{\left(\phi \right)} \tan{\left(\theta \right)} & \cos{\left(\phi \right)} \tan{\left(\theta \right)}\\0 & 0 & 0 & 0 & \cos{\left(\phi \right)} & - \sin{\left(\phi \right)}\\0 & 0 & 0 & 0 & \frac{\sin{\left(\phi \right)}}{\cos{\left(\theta \right)}} & \frac{\cos{\left(\phi \right)}}{\cos{\left(\theta \right)}}\end{matrix}\right]$

```python
# thrust model
Kt0, Kt1 = symbols('Kt0 Kt1', real=True)
Kt = Matrix([Kt0, Kt1])
Qt0, Qt1 = symbols('Qt0 Qt1', real=True)
Qt = Matrix([Qt0, Qt1])

# control inputs
rpm0, rpm1 = symbols('rpm0 rpm1', real=True)
rpm = Matrix([rpm0, rpm1])
de, dr = symbols('de dr', real=True)
control_vector = Matrix([rpm0, rpm1, de, dr])

# control force vector
Ft = Kt.dot(rpm)
Mt = Qt.dot(rpm)

# coefficient for each element in cost function
tauc = Matrix([
    Ft*cos(de)*cos(dr),
    -Ft*sin(dr),
    Ft*sin(de)*cos(dr),
    Mt*cos(de)*cos(dr),
    -Mt*sin(dr),
    Mt*sin(de)*cos(dr)
])
```

```python
etadot = Jeta*nu
nudot = M.inv()*(tauc - (crb + D)*nu - geta)
```

```python
sdot = Matrix([
   etadot,
   nudot
])
```

```python

```

```python
print(cse(sdot))
```

    ([(x0, cos(psi)), (x1, cos(theta)), (x2, u*x1), (x3, sin(phi)), (x4, sin(psi)), (x5, x3*x4), (x6, sin(theta)), (x7, cos(phi)), (x8, x0*x7), (x9, x4*x7), (x10, x0*x3), (x11, x1*x3), (x12, x1*x7), (x13, tan(theta)), (x14, q*x3), (x15, r*x7), (x16, 1/x1), (x17, Yvv*Abs(v)), (x18, m*xg), (x19, p*x18), (x20, m*zg), (x21, r*x20), (x22, m*u), (x23, m*yg), (x24, r*x23), (x25, sin(dr)), (x26, Kt0*rpm0), (x27, Kt1*rpm1), (x28, m*w), (x29, -x28), (x30, p*x23), (x31, -B + W), (x32, x3**2), (x33, x6**2), (x34, x1**2), (x35, x7**2), (x36, 1/(x32*x33 + x32*x34 + x33*x35 + x34*x35)), (x37, x11*x36), (x38, -p*(x29 - x30) - q*(x19 + x21) - r*(x22 - x24) - v*x17 + x25*(-x26 - x27) - x31*x37), (x39, m**2), (x40, xg**4), (x41, x39*x40), (x42, yg**2), (x43, xg**2), (x44, x39*x43), (x45, x42*x44), (x46, zg**2), (x47, x44*x46), (x48, Iyy*Izz), (x49, Ixx*x48), (x50, Ixx*Izz), (x51, m*x50), (x52, m*x48), (x53, -x42*x52 - x43*x51 + x49), (x54, Ixx*Iyy), (x55, m*x54), (x56, -x43*x55 - x46*x52), (x57, Ixx*x41 + Iyy*x45 + Izz*x47 + x53 + x56), (x58, yg**4), (x59, x39*x58), (x60, x42*x46), (x61, x39*x60), (x62, -x42*x55 - x46*x51), (x63, Ixx*x45 + Iyy*x59 + Izz*x61 + x62), (x64, zg**4), (x65, x39*x64), (x66, Ixx*x47 + Iyy*x61 + Izz*x65), (x67, 1/(x57 + x63 + x66)), (x68, x54*yg), (x69, xg**3), (x70, Ixx*x69), (x71, yg**3), (x72, Iyy*m), (x73, x71*x72), (x74, Izz*x46), (x75, x23*x74), (x76, x67*(x23*x70 - x68*xg + x73*xg + x75*xg)), (x77, x50*zg), (x78, zg**3), (x79, Izz*x78), (x80, x20*x42), (x81, Iyy*x80), (x82, x18*x79 + x20*x70 - x77*xg + x81*xg), (x83, Zww*Abs(w)), (x84, m*v), (x85, p*x20), (x86, q*x23), (x87, sin(de)), (x88, cos(dr)), (x89, x88*(x26 + x27)), (x90, -x22), (x91, q*x20), (x92, x12*x36), (x93, -p*(x84 - x85) - q*(x90 - x91) - r*(x19 + x86) - w*x83 - x31*x92 + x87*x89), (x94, x67*x93), (x95, Xuu*Abs(u)), (x96, q*x18), (x97, cos(de)), (x98, -x84), (x99, r*x18), (x100, x6/(x33 + x34)), (x101, -p*(x21 + x86) - q*(x28 - x96) - r*(x98 - x99) - u*x95 + x100*x31 + x89*x97), (x102, m**3), (x103, Ixx*x102), (x104, Iyy*x102), (x105, Izz*x102), (x106, x39*x42), (x107, x39*x46), (x108, x103*x43), (x109, 1/(m*x49 + x103*x40 + x104*x42*x43 + x104*x58 + x104*x60 + x105*x43*x46 + x105*x60 + x105*x64 - x106*x48 - x106*x54 - x107*x48 - x107*x50 + x108*x42 + x108*x46 - x44*x50 - x44*x54)), (x110, Ixx*x43), (x111, x110*x23), (x112, Iyy*x23), (x113, -x111 - x112*x46 + x68 - x73), (x114, p*q), (x115, Qt0*rpm0), (x116, Qt1*rpm1), (x117, x88*(x115 + x116)), (x118, -x19), (x119, -x86), (x120, x67*(B*(x100*yb + x37*xb) + Ixx*x114 - Iyy*x114 - Nrr*r*Abs(r) - W*(x100*yg + x37*xg) - u*(x84 - x95*ycp + x99) - v*(x17*xcp + x24 + x90) - w*(x118 + x119) + x117*x87)), (x121, Izz*m), (x122, x121*x78), (x123, Izz*x80 + x110*x20 + x122 - x77), (x124, p*r), (x125, -x21), (x126, x67*(B*(-x100*zb - x92*xb) - Ixx*x124 + Izz*x124 - Mqq*q*Abs(q) - W*(-x100*zg - x92*xg) - u*(x29 + x95*zcp + x96) - v*(x118 + x125) - w*(x22 - x83*xcp + x91) + x25*(-x115 - x116))), (x127, x23*xg), (x128, Izz*x127), (x129, x128*zg), (x130, x112*xg*zg), (x131, x129 - x130), (x132, q*r), (x133, x67*(B*(-x37*zb + x92*yb) + Iyy*x132 - Izz*x132 - Kpp*p*Abs(p) - W*(-x37*zg + x92*yg) - u*(x119 + x125) - v*(-x17*zcp + x28 + x30) - w*(x83*ycp + x85 + x98) + x117*x97)), (x134, x48*zg), (x135, Iyy*x20), (x136, x111*zg - x134*yg + x135*x71 + x23*x79), (x137, Ixx*zg), (x138, x127*x137), (x139, -x129 + x138), (x140, Ixx*m), (x141, x140*x69), (x142, Ixx*x18), (x143, x42*x72), (x144, x141 + x142*x46 + x143*xg - x54*xg), (x145, -Izz*x20*x43 - x122 + x134 - x81), (x146, x101*x67), (x147, x38*x67), (x148, x130 - x138), (x149, -x141 - x142*x42 - x18*x74 + x50*xg), (x150, x112*x43 - x48*yg + x73 + x75), (x151, x39*xg), (x152, x39*x69), (x153, -x135*xg + x151*x42*zg + x151*x78 + x152*zg), (x154, -x128 + x151*x46*yg + x151*x71 + x152*yg), (x155, -x121*x46 + x45), (x156, -x143 + x47), (x157, -x137*x23 + x39*x71*zg + x39*x78*yg + x44*yg*zg), (x158, -x140*x43 + x61)], [Matrix([
    [                                                                        v*(x10*x6 - x9) + w*(x5 + x6*x8) + x0*x2],
    [                                                                       v*(x5*x6 + x8) + w*(-x10 + x6*x9) + x2*x4],
    [                                                                                -u*x6 + v*x11 + w*x12],
    [                                                                                p + x13*x14 + x13*x15],
    [                                                                                q*x7 - r*x3],
    [                                                                                x14*x16 + x15*x16],
    [                                           x101*x109*x57 + x113*x120 + x123*x126 + x131*x133 + x38*x76 + x82*x94],
    [                                  x101*x76 + x109*x38*(x53 + x63) + x120*x144 + x126*x139 + x133*x145 + x136*x94],
    [                     x109*x93*(x49 + x56 + x62 + x66) + x120*x148 + x126*x149 + x133*x150 + x136*x147 + x146*x82],
    [ x120*x153 + x126*x154 + x131*x146 + x133*(-x121*x43 + x155 + x156 + x41 - x43*x72 + x48) + x145*x147 + x150*x94],
    [x120*x157 + x123*x146 + x126*(-x121*x42 - x140*x42 + x155 + x158 + x50 + x59) + x133*x154 + x139*x147 + x149*x94],
    [ x113*x146 + x120*(-x140*x46 + x156 + x158 - x46*x72 + x54 + x65) + x126*x157 + x133*x153 + x144*x147 + x148*x94]])])

```python

```