$\newcommand{\calf}{{\cal F}}
\newcommand{\dnu}{d \nu}
\newcommand{\mf}{{\bf F}}
\newcommand{\vu}{{\bf u}}
\newcommand{\ve}{{\bf e}}
\newcommand{\mg}{{\bf G}}
\newcommand{\ml}{{\bf L}}
\newcommand{\mg}{{\bf G}}
\newcommand{\mi}{{\bf I}}
\newcommand{\diag}{\mbox{diag}}
\newcommand{\begeq}{{\begin{equation}}}
\newcommand{\endeq}{{\end{equation}}}
$

```julia
include("mprnote.jl")
```

# Newton's Method in Multiple Precision: C. T. Kelley

This notebook documents the results in
C. T. Kelley, ***Newton's Method in Mixed Precision***, 2020.(Kel20)

As an example we will solve the Chandrasekhar H-equation (Cha60). This equation, which we describe in detail in the Example section, has a fast $O(N \log(N))$ function evaluation, a Jacobian evaluation that is $O(N^2)$ work analytically and $O(N^2 \log(N))$ with a finite difference. This means that most of the work, if you do things right, is in the LU factorization of the Jacobian.

The difference between double, single, and half precision will be clear in the results from the examples. This notebook has no half-precision computations. Julia does half-precsion in software and that is very slow.

## Contents

- [Example: The Chandrasekhar H-Equation](#The-Chandrasekhar-H-Equation): Integral equations example

- [Setting up the notebook](#Setting-up): Install the application.

- [First run of the solver](#Running-the-solver): First solver teset

- [How to use the solver](#NSOL): Using nsol.jl.

- [The results in the paper](#The-results-in-the-paper): Running the codes that generated the results

## The Chandrasekhar H-Equation

The example is the mid-point rule discretization of the Chandrasekhar H-equation (Cha60).

   \begin{equation}
    {\calf}(H)(\mu) = H(\mu) -
\left(
1 - \frac{c}{2} \int_0^1 \frac{\mu H(\mu)}{\mu + \nu} \dnu
\right)^{-1} = 0.
    \end{equation}

The nonlinear operator $\calf$ is defined on $C[0,1]$, the space of
continuous functions on $[0,1]$.

The equation has a well-understood dependence on the parameter $c$
(Mul68),
    (DK80).
The equation has unique solutions at $c=0$
and $c=1$ and two solutions for $0 < c < 1$. There is a simple fold
singularity
        (Kel87)
        at $c=1$. Only one
        (Cha60),
        (Bus60)
            of the two solutions for $0 < c < 1$ is of physical interest
and that is the one easiest to find numerically. One must do
a continuation computation to find the other one.

The structure of the singularity is preserved if one discretizes
the integral with any rule that integrates constants exactly. For
the purposes of this paper the composite midpoint rule will suffice.
The $N$-point composite midpoint rule is
\begin{equation}
\int_0^1 f(\nu) \dnu \approx \frac{1}{N} \sum_{j=1}^N f(\nu_j)
\end{equation}
where $\nu_j = (j - 1/2)/N$ for $1 \le j \le N$. This rule is
second-order accurate for sufficiently smooth functions $f$. The
solution of the integral equation is, however, not smooth enough. $H'(\mu)$
has a logarithmic singularity at $\mu=0$.

The discrete problem is

\begin{equation}
\mf(\vu)_i \equiv
u_i - \left(
1  - \frac{c}{2N} \sum_{j=1}^N \frac{u_j \mu_i}{\mu_j + \mu_i}
\right)^{-1}
=0.
\end{equation}

One can simplify the approximate integral operator
and expose some useful structure. Since

\begin{equation}
\frac{c}{2N} \sum_{j=1}^N \frac{u_j \mu_i}{\mu_j + \mu_i}
= \frac{c (i - 1/2) }{2N} \sum_{j=1}^N \frac{u_j}{i+j -1}.
\end{equation}

hence the approximate integral operator is
the product of a diagonal matrix and a Hankel matrix and
one can use an FFT to evaluate that operator with $O(N \log(N))$
work
(GV96).

We can express the approximation of the integral operator in matrix
form
\begin{equation}
(\ml \vu)_i = \frac{c (i - 1/2) }{2N} \sum_{j=1}^N \frac{u_j}{i+j -1}
\end{equation}
and compute the Jacobian analytically as
\begin{equation}
\mf'(\vu) = \mi - \diag(\mg(\vu))^2 \ml
\end{equation}
where
\begin{equation}
\mg(\vu)_i = \left(
1  - \frac{c}{2N} \sum_{j=1}^N \frac{u_j \mu_i}{\mu_j + \mu_i}
\right)^{-1}.
\end{equation}
Hence the data for the Jacobian is already available after
one computes $\mf(\vu) = \vu - \mg(\vu)$ and the Jacobian can
be computed with $O(N^2)$ work.
We do that in this example and therefore the only $O(N^3)$
part of the solve is the matrix factorization.

One could also approximate the Jacobian with forward differences.
In this case one approximates the $j$th column $\mf'(\vu)_j$
of the Jacobian with
\begin{equation}
\frac{\mf(\vu + h {\tilde \ve}_j) - \mf(\vu)}{h}
\end{equation}
where ${\tilde \ve}_j$ is a unit vector in the $j$th coordinate
direction and $h$ is a suitable difference increment. If one computes
$\mf$ in double precision with unit roundoff $u_d$, then
$h =O(\| \vu \| \sqrt{u_d})$ is a reasonable choice
(Kel95).
Then
the error in the Jacobian is $O(\sqrt{u_d}) = O(u_s)$ where $u_s$ is
unit roundoff in single precision. The cost of a finite difference Jacobian
in this example is $O(N^2 \log(N))$ work.

The analysis in  (Kel20) suggests that there
is no significant difference in the nonlinear iteration
from either the choice of analytic or finite difference Jacobians
or the choice of single or double precision for the linear solver. This notebook has the data used in that paper
to support that assertion. You will be able to duplicate the results and play with the codes.

Half precision is another story and we have those codes for you, too.

## Setting up

You need to install these packages with __Pkg__. I assume you know how to do that.

- SIAMFANLEquations
- PyPlot
- LinearAlgebra
- Printf
- IJulia (You must have done this already or you would not be looking at this notebook.)

To render the LaTeX you'll need to run  the first markdown cell in the notebook. That sets up the commands you need to render the LaTeX correctly.

The directory is a Julia project. So all you should need to do to get going is to run the first code cell in this notebook. That cell has one line

```
include("mprnote.jl")
```

Then you can do a simple solve and test that you did it right by typing
```Julia
hout=heqtest()
```
which I will do in the next code cell. Now ...

**Make absolutely sure that you are in the MPResults directory**. Then ...

## Running the solver

The codes solve the H-equation and plot/tabulate the results in various ways. __heqtest__ prints on column of the tables in Chandrasekhar's book. It calls __nsold.jl__ . I've shown you the output from the solver but that is not important for now. You get the details on the solver in the next section.

```julia
heqtest()
```

    0.00e+00      1.00000e+00
    5.00e-02      1.04424e+00
    1.00e-01      1.07236e+00
    1.50e-01      1.09470e+00
    2.00e-01      1.11346e+00
    2.50e-01      1.12965e+00
    3.00e-01      1.14389e+00
    3.50e-01      1.15657e+00
    4.00e-01      1.16797e+00
    4.50e-01      1.17830e+00
    5.00e-01      1.18773e+00
    5.50e-01      1.19638e+00
    6.00e-01      1.20435e+00
    6.50e-01      1.21172e+00
    7.00e-01      1.21856e+00
    7.50e-01      1.22493e+00
    8.00e-01      1.23088e+00
    8.50e-01      1.23646e+00
    9.00e-01      1.24169e+00
    9.50e-01      1.24662e+00
    1.00e+00      1.25126e+00

    [1.54457e+00, 7.94167e-03, 1.55209e-07, 2.67377e-15, 1.53837e-15, 1.35064e-15, 1.52226e-15, 1.42178e-15, 1.47288e-15, 1.42178e-15, 1.33227e-15]

    (exactout = (solution = [1.00707e+00, 1.01754e+00, 1.02622e+00, 1.03392e+00, 1.04094e+00, 1.04744e+00, 1.05352e+00, 1.05925e+00, 1.06469e+00, 1.06986e+00  …  1.24220e+00, 1.24320e+00, 1.24419e+00, 1.24517e+00, 1.24614e+00, 1.24709e+00, 1.24804e+00, 1.24897e+00, 1.24989e+00, 1.25081e+00], functionval = [0.00000e+00, 0.00000e+00, 0.00000e+00, 0.00000e+00, 0.00000e+00, 0.00000e+00, 0.00000e+00, 0.00000e+00, 0.00000e+00, 0.00000e+00  …  0.00000e+00, 0.00000e+00, 4.44089e-16, 0.00000e+00, 0.00000e+00, -2.22045e-16, 0.00000e+00, 0.00000e+00, 0.00000e+00, 0.00000e+00], history = [1.54457e+00, 7.94167e-03, 1.55209e-07, 2.67377e-15, 1.53837e-15, 1.35064e-15, 1.52226e-15, 1.42178e-15, 1.47288e-15, 1.42178e-15, 1.33227e-15], stats = (ifun = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], ijac = [0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], iarm = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]), idid = false, errcode = 10), fdout = (solution = [1.00707e+00, 1.01754e+00, 1.02622e+00, 1.03392e+00, 1.04094e+00, 1.04744e+00, 1.05352e+00, 1.05925e+00, 1.06469e+00, 1.06986e+00  …  1.24220e+00, 1.24320e+00, 1.24419e+00, 1.24517e+00, 1.24614e+00, 1.24709e+00, 1.24804e+00, 1.24897e+00, 1.24989e+00, 1.25081e+00], functionval = [0.00000e+00, 0.00000e+00, 0.00000e+00, 0.00000e+00, 0.00000e+00, 0.00000e+00, 0.00000e+00, 0.00000e+00, 0.00000e+00, 0.00000e+00  …  0.00000e+00, 0.00000e+00, 4.44089e-16, 0.00000e+00, 0.00000e+00, -2.22045e-16, 0.00000e+00, 0.00000e+00, 0.00000e+00, 0.00000e+00], history = [1.54457e+00, 7.94166e-03, 1.55239e-07, 2.24254e-15, 1.58572e-15, 1.53837e-15, 1.47288e-15, 1.53837e-15, 1.67640e-15, 1.57009e-15, 1.23629e-15], stats = (ifun = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], ijac = [0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], iarm = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]), idid = false, errcode = 10))

heqtest.jl, calls the solver and harvests some iteration statistics. The two columns of numbers are the reults from (Cha60) (page 125). The iteration statistics are from nsold.jl, the solver.

## NSOL

The solver is ```nsol``` from my pacakge [SIAMFANLEquations.jl](#https://github.com/ctkelley/SIAMFANLEquations.jl)
(Kel20d). That pacackage and an
[IJulia Notebook](#https://github.com/ctkelley/NotebookSIAMFANL)
(Kel20b)
support my upcoming book __Solving Nonlinear Equations with Iterative Methods:
Solvers and Examples in Julia__
(Kel20c)
The solver and the H-equation example are both from that package. When you run the first code cell you are set up.

## Using nsol.jl

At the level of this notebook, it's pretty simple. Remember that Julia hates to allocate memory. So your function and Jacobian evaluation routines should expect the calling function to **preallocate** the storage for both the function and Jacobian. Your functions will then use __.=__ to put the function and Jacobian where they are supposed to be.

It's worthwhile to look at the help screen.

```julia
?nsol
```

    search: nsol nsolsc transcode transpose Transpose transpose! DimensionMismatch

```
nsol(F!, x0, FS, FPS, J!=diffjac!; rtol=1.e-6, atol=1.e-12,
           maxit=20, solver="newton", sham=5, armmax=10, resdec=.1,
           dx = 1.e-7, armfix=false,
           pdata = nothing, jfact = klfact,
           printerr = true, keepsolhist = false, stagnationok=false)
```

)

C. T. Kelley, 2020

Julia versions of the nonlinear solvers from my SIAM books.  Herewith: nsol

You must allocate storage for the function and Jacobian in advance –> in the calling program <– ie. in FS and FPS

Inputs:

  * F!: function evaluation, the ! indicates that F! overwrites FS, your   preallocated storage for the function.

    So FS=F!(FS,x) or FS=F!(FS,x,pdata) returns FS=F(x)

  * x0: initial iterate

  * FS: Preallocated storage for function. It is an N x 1 column vector

  * FPS: preallocated storage for Jacobian. It is an N x N matrix

  * J!: Jacobian evaluation, the ! indicates that J! overwrites FPS, your   preallocated storage for the Jacobian. If you leave this out the   default is a finite difference Jacobian.

    So, FP=J!(FP,FS,x) or FP=J!(FP,FS,x,pdata) returns FP=F'(x).

    (FP,FS, x) must be the argument list, even if FP does not need FS.   One reason for this is that the finite-difference Jacobian   does and that is the default in the solver.
  * Precision: Lemme tell ya 'bout precision. I designed this code for    full precision functions and linear algebra in any precision you want.    You can declare   FPS as Float64, Float32, or Float16 and nsol will do the right thing if   YOU do not destroy the declaration in your J! function. I'm amazed   that this works so easily. If the Jacobian is reasonably well    conditioned, you can cut the cost of Jacobian factorization and   storage in half with no loss. For large dense Jacobians and inexpensive   functions, this is a good deal.

    BUT ... There is very limited support for direct sparse solvers in   anything other than Float64. I recommend that you only use Float64   with direct sparse solvers unless you really know what you're doing. I   have a couple examples in the notebook, but watch out.

---

Keyword Arguments (kwargs):

  * rtol and atol: relative and absolute error tolerances

  * maxit: limit on nonlinear iterations

solver: default = "newton"

Your choices are "newton" or "chord". However, you have sham at your disposal only if you chose newton. "chord" will keep using the initial derivative until the iterate converges, uses the iteration budget, or the line search fails. It is not the same as sham=Inf, which is smarter.

sham: default = 5 (ie Newton)

This is the Shamanskii method. If sham=1, you have Newton. The iteration updates the derivative every sham iterations. The convergence rate has local q-order sham+1 if you only count iterations where you update the derivative. You need not provide your own derivative function to use this option. sham=Inf is chord only if chord is converging well.

I made sham=1 the default for scalar equations. For systems I'm more aggressive and want to invest as little energy in linear algebra as possible. So the default is sham=5.

armmax: upper bound on step size reductions in line search

resdec: default = .1

This is the target value for residual reduction. The default value is .1. In the old MATLAB codes it was .5. I only turn Shamanskii on if the residuals are decreasing rapidly, at least a factor of resdec, and the line search is quiescent. If you want to eliminate resdec from the method ( you don't ) then set resdec = 1.0 and you will never hear from it again.

dx: default = 1.e-7

difference increment in finite-difference derivatives       h=dx*norm(x,Inf)+1.e-8

armfix: default = false

The default is a parabolic line search (ie false). Set to true and the step size will be fixed at .5. Don't do this unless you are doing experiments for research.

pdata:

precomputed data for the function/Jacobian.  Things will go better if you use this rather than hide the data  in global variables within the module for your function/Jacobian

jfact: default = klfact (tries to figure out best choice)

If your Jacobian has any special structure, please set jfact to the correct choice for a factorization.

I use jfact when I call PrepareJac! to evaluate the Jacobian (using your J!) and factor it. The default is to use klfact (an internal function) to do something reasonable. For general matrices, klfact picks lu! to compute an LU factorization and share storage with the Jacobian.  You may change LU to something else by, for example, setting jfact = cholseky! if your Jacobian is spd.

klfact knows about banded matrices and picks qr. You should, however RTFM, allocate the extra two upper bands, and use jfact=qr! to override klfact.

If you give me something that klfact does not know how to dispatch on, then nothing happens. I just return the original Jacobian matrix and  nsol will use backslash to compute the Newton step. I know that this is probably not optimal in your situation, so it is  good to pick something else, like jfact = lu.

Please do not mess with the line that calls PrepareJac!.

```
    FPF = PrepareJac!(FPS, FS, x, ItRules)
```

FPF is not the same as FPS (the storage you allocate for the Jacobian) for a reason. FPF and FPS do not have the same type, even though they share storage. So, FPS=PrepareJac!(FPS, FS, ...) will break things.

printerr: default = true

I print a helpful message when the solver fails. To suppress that message set printerr to false.

keepsolhist: default = false

Set this to true to get the history of the iteration in the output tuple. This is on by default for scalar equations and off for systems. Only turn it on if you have use for the data, which can get REALLY LARGE.

stagnationok: default = false

Set this to true if you want to disable the line search and either observe divergence or stagnation. This is only useful for research or writing a book.

Output:

A named tuple (solution, functionval, history, stats, idid,                errcode, solhist) where

solution = converged result functionval = F(solution) history = the vector of residual norms (||F(x)||) for the iteration stats = named tuple of the history of (ifun, ijac, iarm), the number of functions/derivatives/steplength reductions at each iteration.

I do not count the function values for a finite-difference derivative because they count toward a Jacobian evaluation.

idid=true if the iteration succeeded and false if not.

errcode = 0 if if the iteration succeeded         = -1 if the initial iterate satisfies the termination criteria         = 10 if no convergence after maxit iterations         = 1  if the line search failed

solhist:

This is the entire history of the iteration if you've set keepsolhist=true

solhist is an N x K array where N is the length of x and K is the number of iteration + 1. So, for scalar equations, it's a row vector.

---

# Examples

#### World's easiest problem example. Test 64 and 32 bit Jacobians. No meaningful difference in the residual histories or the converged solutions.

```jldoctest
 julia> function f!(fv,x)
       fv[1]=x[1] + sin(x[2])
       fv[2]=cos(x[1]+x[2])
       end
f (generic function with 1 method)

julia> x=ones(2,); fv=zeros(2,); jv=zeros(2,2); jv32=zeros(Float32,2,2);
julia> nout=nsol(f!,x,fv,jv; sham=1);
julia> nout32=nsol(f!,x,fv,jv32; sham=1);
julia> [nout.history nout32.history]
5×2 Array{Float64,2}:
 1.88791e+00  1.88791e+00
 2.43119e-01  2.43119e-01
 1.19231e-02  1.19231e-02
 1.03266e-05  1.03262e-05
 1.46416e-11  1.43548e-11

julia> [nout.solution nout.solution - nout32.solution]
2×2 Array{Float64,2}:
 -7.39085e-01  -5.42899e-14
  2.30988e+00   3.49498e-13
```

#### H-equation example. I'm taking the sham=5 default here, so the convergence is not quadratic. The good news is that we evaluate the Jacobian only once.

```jldoctest
julia> n=16; x0=ones(n,); FV=ones(n,); JV=ones(n,n);
julia> hdata=heqinit(x0, .5);
julia> hout=nsol(heqf!,x0,FV,JV;pdata=hdata);
julia> hout.history
4-element Array{Float64,1}:
 6.17376e-01
 3.17810e-03
 2.75227e-05
 2.35817e-07
```

## How nsol.jl controls the precision of the Jacobian

You can control the precision of the Jacobian by simply allocating FPS in your favorite precision. So if I have a problem with N=256 unknowns I will declare FP as zeros(N,1) and may declare FPS as zeros(N,N) or Float32.(zeros(N,N)).

Note the __.__ between Float32 and the paren. This, as is standard Julia practice, applies the conversion to Float32 to everyelement in the array. If you forget the __.__ Julia will complain.

# The results in the paper

The paper has plots for double, single, and half precsion computations for c=.5, .99, and 1.0. The half precision results take a very long time to get. On my computer (2019 iMac; 8 cores, 64GB of memory) the half precision compute time was over two weeks. Kids, don't try this at home.

The data for the paper are in the cleverly named directory __Data_From_Paper__

__cd to the directory MPResults__ and __from that directory__ run

```Julia
data_harvest()
```

at the julia prompt you will generate all the tables and plots in the paper.

If you have the time and patience you can also generate the data with __data_populate.jl__. This creates binary files with the iteration histories and you can see for yourself how long it takes. You can reduce the number of grid levels and __turn half precision off__. I will turn half precision off in the example in this notebook. That means that the code will run in a reasonable amount of time instead of the __two weeks__ it needs for the half precision results.

```data_populate(c; half=false,level=p)``` does double and single preicsion solves for $1024 \times 2^k$ point grids for k=0, ... p-1 . Set ```half=true``` to make the computations take much longer.

Here is a simple example of using data_populate and the plotter code plot_nsold.jl. I'm only using the 1024, 2048 and 4096 point grids. The plot in the paper uses more levels. This is part of Figure 1 in the paper.

Look at the source to __data_populate.jl__ and __PlotData.jl__ and you'll see how I did this. These codes only mange files, plots, and tables. There is nothing really exciting here. You don't need to know much Julia to understand this, but you do need to know something and I can't help with that.

To begin with, I will create a directory called Data_Test to put all this stuff. I will cd to that directory.

```julia
cd(MPRdir)
Home4mp="Data_Test"
try (mkdir(Home4mp))
catch
end
cd(Home4mp)
pwd()
```

    "/Users/ctk/Dropbox/Julia/dev/MPResults/Data_Test"

Now I'll run __data_populate__ to create a subdirectory with the data. cd to that directory and run __PlotData__. That's how I created the plots in the paper with all values of *c* and all the problem sizes. The half precision computation took two weeks.

__PlotData(.5, level=3)__ makes the four-plot with only double and single (no half).

Even with this small problem, __data_populate__ takes a while. Be patient. Once it's done the plots will appear
pretty rapidly.

```julia
using PyPlot
data_populate(.5; level=3)
PlotData(.5; level=3);

```

Finally, I will duplicate the tables and plots in the paper with the precomputed data in the Complete_Data directory. I'll cd to that directory and make the plots with __data_harvest.jl__.

```julia
cd(MPRdir)
cd("Complete_Data")
data_harvest()
```

The three LaTeX tables are pretty vivid demonstrations that using a half-precision Jacobian is a poor idea. The columns are ratios of successive residual norms. Those ratios are supposed to go to zero if the convergence is q-superlinear. In the half precision case, it is not. You can also see that from the plots at the bottom. I make the tables with three calls to __MakeTable.jl__ which is contained in the file __TableData.jl__.

```julia
MakeTable(.5)
```

    \begin{tabular}{llllll}
            n &     1024 &     2048 &     4096 &     8192 &    16384 \\
    \hline
    1 & 1.56706e-01 & 1.56708e-01 & 1.56708e-01 & 1.56705e-01 & 1.56706e-01   \\
    2 & 1.53569e-01 & 1.53573e-01 & 1.53579e-01 & 1.53578e-01 & 1.53576e-01   \\
    3 & 1.52949e-01 & 1.52944e-01 & 1.52946e-01 & 1.52948e-01 & 1.52949e-01   \\
    4 & 1.52853e-01 & 1.52848e-01 & 1.52844e-01 & 1.52847e-01 & 1.52843e-01   \\
    5 & 1.52831e-01 & 1.52829e-01 & 1.52832e-01 & 1.52830e-01 & 1.52830e-01   \\
    6 & 1.52828e-01 & 1.52825e-01 & 1.52826e-01 & 1.52830e-01 & 1.52827e-01   \\
    7 & 1.52830e-01 & 1.52824e-01 & 1.52827e-01 & 1.52826e-01 & 1.52825e-01   \\
    8 & 1.52832e-01 & 1.52832e-01 & 1.52824e-01 & 1.52825e-01 & 1.52828e-01   \\
    9 & 1.52838e-01 & 1.52830e-01 & 1.52831e-01 & 1.52830e-01 & 1.52826e-01   \\
    10 & 1.52828e-01 & 1.52828e-01 & 1.52827e-01 & 1.52829e-01 & 1.52827e-01   \\
    \hline
    \end{tabular}

```julia
MakeTable(.99)
```

    \begin{tabular}{llllll}
            n &     1024 &     2048 &     4096 &     8192 &    16384 \\
    \hline
    1 & 2.63294e-01 & 4.88603e-01 & 5.06480e-01 & 5.06480e-01 & 5.06480e-01   \\
    2 & 1.39099e-01 & 4.50348e-01 & 5.83959e-01 & 5.83962e-01 & 5.83961e-01   \\
    3 & 1.44278e-01 & 4.13186e-01 & 6.38900e-01 & 6.38900e-01 & 6.38898e-01   \\
    4 & 1.01754e-01 & 3.60960e-01 & 6.64343e-01 & 6.77977e-01 & 6.77976e-01   \\
    5 & 9.42216e-02 & 3.76718e-01 & 6.78547e-01 & 7.06190e-01 & 7.06192e-01   \\
    6 & 1.03855e-01 & 3.36916e-01 & 7.12325e-01 & 7.26959e-01 & 7.26957e-01   \\
    7 & 8.71357e-02 & 4.10487e-01 & 6.98753e-01 & 7.42506e-01 & 7.42507e-01   \\
    8 & 9.85171e-02 & 3.63814e-01 & 7.53512e-01 & 7.54294e-01 & 7.54296e-01   \\
    9 & 9.81771e-02 & 3.73218e-01 & 7.13812e-01 & 7.63337e-01 & 7.63337e-01   \\
    10 & 7.82248e-02 & 3.60341e-01 & 7.51307e-01 & 7.70318e-01 & 7.70319e-01   \\
    \hline
    \end{tabular}

```julia
MakeTable(1.0)
```

    \begin{tabular}{llllll}
            n &     1024 &     2048 &     4096 &     8192 &    16384 \\
    \hline
    1 & 2.89271e-01 & 4.97487e-01 & 5.18347e-01 & 5.18347e-01 & 5.18347e-01   \\
    2 & 2.02306e-01 & 4.37133e-01 & 6.02756e-01 & 6.02754e-01 & 6.02755e-01   \\
    3 & 3.04726e-01 & 4.60824e-01 & 6.62228e-01 & 6.64944e-01 & 6.64946e-01   \\
    4 & 3.28286e-01 & 4.37880e-01 & 6.87719e-01 & 7.11345e-01 & 7.11344e-01   \\
    5 & 4.09534e-01 & 5.29573e-01 & 7.12647e-01 & 7.46800e-01 & 7.46800e-01   \\
    6 & 4.67802e-01 & 5.59574e-01 & 7.46514e-01 & 7.74644e-01 & 7.74642e-01   \\
    7 & 5.30915e-01 & 5.87828e-01 & 7.65099e-01 & 7.97048e-01 & 7.97048e-01   \\
    8 & 5.86394e-01 & 6.75211e-01 & 7.87571e-01 & 8.15444e-01 & 8.15444e-01   \\
    9 & 6.53255e-01 & 6.77888e-01 & 8.02109e-01 & 8.30813e-01 & 8.30815e-01   \\
    10 & 6.76681e-01 & 7.53966e-01 & 8.09954e-01 & 8.43842e-01 & 8.43838e-01   \\
    11 & 7.25254e-01 & 7.27053e-01 & 8.24948e-01 & 8.55014e-01 & 8.55015e-01   \\
    12 & 7.42052e-01 & 7.91974e-01 & 8.66721e-01 & 8.64711e-01 & 8.64711e-01   \\
    13 & 7.82260e-01 & 7.94773e-01 & 7.82944e-01 & 8.73197e-01 & 8.73196e-01   \\
    14 & 7.97188e-01 & 7.84908e-01 & 8.21089e-01 & 8.80684e-01 & 8.80685e-01   \\
    15 & 7.92356e-01 & 8.61411e-01 & 8.78150e-01 & 8.87342e-01 & 8.87343e-01   \\
    16 & 8.54814e-01 & 7.99366e-01 & 7.81216e-01 & 8.93304e-01 & 8.93303e-01   \\
    17 & 8.05580e-01 & 8.94867e-01 & 9.34862e-01 & 8.98662e-01 & 8.98662e-01   \\
    18 & 8.74513e-01 & 8.01814e-01 & 7.57971e-01 & 9.03515e-01 & 9.03515e-01   \\
    19 & 8.35880e-01 & 9.29478e-01 & 8.87219e-01 & 9.07927e-01 & 9.07925e-01   \\
    20 & 8.96103e-01 & 8.30895e-01 & 9.29619e-01 & 9.11949e-01 & 9.11952e-01   \\
    21 & 8.44739e-01 & 9.25185e-01 & 8.49444e-01 & 9.15640e-01 & 9.15637e-01   \\
    22 & 9.13096e-01 & 8.25132e-01 & 9.18586e-01 & 9.19032e-01 & 9.19033e-01   \\
    23 & 8.53739e-01 & 9.61787e-01 & 8.64937e-01 & 9.22164e-01 & 9.22166e-01   \\
    24 & 9.27608e-01 & 8.63339e-01 & 9.44576e-01 & 9.25064e-01 & 9.25063e-01   \\
    25 & 8.67368e-01 & 9.27716e-01 & 9.49682e-01 & 9.27758e-01 & 9.27757e-01   \\
    26 & 9.37645e-01 & 9.02116e-01 & 8.63996e-01 & 9.30262e-01 & 9.30264e-01   \\
    27 & 8.64444e-01 & 8.84138e-01 & 9.57200e-01 & 9.32602e-01 & 9.32599e-01   \\
    28 & 9.50195e-01 & 9.80767e-01 & 9.72447e-01 & 9.34786e-01 & 9.34786e-01   \\
    29 & 8.76453e-01 & 8.51197e-01 & 8.62460e-01 & 9.36834e-01 & 9.36837e-01   \\
    30 & 9.42673e-01 & 9.90469e-01 & 9.56378e-01 & 9.38760e-01 & 9.38762e-01   \\
    \hline
    \end{tabular}

Finally, I will cd to the directory that contains the notebook.

```julia
cd(MPRdir)
```


# The Harmonic Oscillator Strikes Back

*Note:* Much of this is adapted/copied from https://flothesof.github.io/harmonic-oscillator-three-methods-solution.html

This week we continue our adventures with the harmonic oscillator.

The harmonic oscillator is a system that, when displaced from its equilibrium position, experiences a restoring force F proportional to the displacement x:

$$F=-kx$$

The potential energy of this system is

$$V = {1 \over 2}k{x^2}$$

These are sometime rewritten as

$$ F=- \omega_0^2 m x, \text{ }  V(x) = {1 \over 2} m \omega_0^2 {x^2}$$

Where $\omega_0  = \sqrt {{k \over m}} $

If the equilibrium value of the harmonic oscillator is not zero, then

$$ F=- \omega_0^2 m (x-x_{eq}), \text{ }  V(x) = {1 \over 2} m \omega_0^2 (x-x_{eq})^2$$

## 1. Harmonic oscillator from last time (with some better defined conditions)

Applying the harmonic oscillator force to Newton's second law leads to the following second order differential equation

$$ F = m a $$

$$ F= -m \omega_0^2 (x-x_{eq}) $$

$$ a = - \omega_0^2 (x-x_{eq}) $$

$$ x(t)'' = - \omega_0^2 (x-x_{eq}) $$

The final expression can be rearranged into a second order homogenous differential equation, and can be solved using the methods we used above

This is already solved to remind you how we found these values

```python
import sympy as sym
sym.init_printing()
```

**Note** that this time we define some of the properties of the symbols. Namely, that the frequency is always positive and real and that the positions are always real

```python
omega0,t=sym.symbols("omega_0,t",positive=True,nonnegative=True,real=True)
xeq=sym.symbols("x_{eq}",real=True)
x=sym.Function("x",real=True)
x(t),omega0
```

```python
dfeq=sym.Derivative(x(t),t,2)+omega0**2*(x(t)-xeq)
dfeq
```

```python
sol = sym.dsolve(dfeq)
sol
```

```python
sol,sol.args[0],sol.args[1]
```

**Note** this time we define the initial positions and velocities as real

```python
x0,v0=sym.symbols("x_0,v_0",real=True)
ics=[sym.Eq(sol.args[1].subs(t, 0), x0),
     sym.Eq(sol.args[1].diff(t).subs(t, 0), v0)]
ics
```

```python
solved_ics=sym.solve(ics)
solved_ics
```

### 1.1 Equation of motion for $x(t)$

```python
full_sol = sol.subs(solved_ics[0])
full_sol
```

### 1.2 Equation of motion for $p(t)$

```python
m=sym.symbols("m",positive=True,nonnegative=True,real=True)
p=sym.Function("p")
sym.Eq(p(t),m*sol.args[1].subs(solved_ics[0]).diff(t))
```

## 2. Time average values for a harmonic oscillator

If we want to understand the average value of a time dependent observable, we need to solve the following integral

$${\left\langle {A(t)} \right\rangle}_t  = \begin{array}{*{20}{c}}
{\lim }\\
{\tau  \to 0}
\end{array}\frac{1}{\tau }\int\limits_0^\tau  {A(t)dt} $$

### 2.1 Average position ${\left\langle {x} \right\rangle}_t$ for a harmonic oscillator

```python
tau=sym.symbols("tau",nonnegative=True,real=True)
xfunc=full_sol.args[1]
xavet=(xfunc.integrate((t,0,tau))/tau).limit(tau,sym.oo)
xavet
```

The computer does not always make the best choices the first time. If you treat each sum individually this is not a hard limit to do by hand. The computer is not smart. We can help it by inseting an `expand()` function in the statement

```python
xavet=(xfunc.integrate((t,0,tau))/tau).expand().limit(tau,sym.oo)
xavet
```

### 2.2 Excercise: Calculate the average momenta  ${\left\langle {p} \right\rangle}_t$ for a harmonic oscillator

```python
tau=sym.symbols("tau",nonnegative=True,real=True)
p=sym.Eq(p(t),m*sol.args[1].subs(solved_ics[0]).diff(t))
pfunc=p.args[1]
pavet=(pfunc.integrate((t,0,tau))/tau).limit(tau,sym.oo)
pavet
```

### 2.3 Exercise: Calculate the average kinetic energy of a harmonic oscillator

```python
ke=pfunc**2/2*m
ke
```

```python
keavet=(ke.integrate((t,0,tau))/tau).limit(tau,sym.oo)
keavet
```

```python
keavet=(ke.integrate((t,0,tau))/tau).expand().limit(tau,sym.oo)
keavet
```

## 3.  Ensemble (Thermodynamic) Average values for a harmonic oscillator

If we want to understand the thermodynamics ensemble average value of an observable, we need to solve the following integral.

$${\left\langle {A(t)} \right\rangle}_{T}  = \frac{\int{A e^{-\beta H}dqdp}}{\int{e^{-\beta H}dqdp} } $$

You can think of this as a Temperature average instead of a time average.

Here $\beta=\frac{1}{k_B T}$ and the classical Hamiltonian, $H$ is

$$ H = \frac{p^2}{2 m} + V(q)$$

**Note** that the factors of $1/h$ found in the classical partition function cancel out when calculating average values

### 3.1 Average position ${\left\langle {x} \right\rangle}_t$ for a harmonic oscillator

For a harmonic oscillator with equilibrium value $x_{eq}$, the Hamiltonian is
$$ H = \frac{p^2}{2 m} + \frac{1}{2} m \omega_0 (x-x_{eq})^2 $$

First we will calculate the partition function $\int{e^{-\beta H}dqdp}$

```python
k,T=sym.symbols("k,T",positive=True,nonnegative=True,real=True)
xT,pT=sym.symbols("x_T,p_T",real=True)
ham=sym.Rational(1,2)*(pT)**2/m + sym.Rational(1,2)*m*omega0**2*(xT-xeq)**2
beta=1/(k*T)
bolz=sym.exp(-beta*ham)
z=sym.integrate(bolz,(xT,-sym.oo,sym.oo),(pT,-sym.oo,sym.oo))
z
```

Then we can calculate the numerator $\int{A e^{-\beta H}dqdp}$

```python
numx=sym.integrate(xT*bolz,(xT,-sym.oo,sym.oo),(pT,-sym.oo,sym.oo))
numx
```

And now the average value

```python
xaveT=numx/z
xaveT
```

### 3.2 Exercise: Calculate the average momenta  ${\left\langle {p} \right\rangle}_t$ for a harmonic oscillator

After calculating the value, explain why you think you got this number

```python
nump=sym.integrate(pT*bolz,(xT,-sym.oo,sym.oo),(pT,-sym.oo,sym.oo))
nump
```

```python
pavg=nump/z
pavg
```

### 3.3 Exercise: Calculate the average kinetic energy

The answer you get here is a well known result related to the energy equipartition theorem

```python
KE=pT**2/2*m
numKE=sym.integrate(KE*bolz,(xT,-sym.oo,sym.oo),(pT,-sym.oo,sym.oo))
numKE
```

```python
KEavg=numKE/z
KEavg
```

# Back to the lecture

## 4. Exercise Verlet integrators

In this exercise we will write a routine to solve for the equations of motion for a hamonic oscillator.

Plot the positions and momenta (seprate plots) of the harmonic oscillator as a functions of time.

Calculaate trajectories using the following methods:
1. Exact solution
2. Simple taylor series expansion
3. Predictor-corrector method
4. Verlet algorithm
5. Leapfrog algorithm
6. Velocity Verlet algorithm

```python
#Exact solution for x

fuller_sol = sym.simplify(full_sol.subs({x0:10, xeq:0 , v0:10, omega0:1}))
sym.plot(fuller_sol.rhs,(t,-10,10))
```

```python
#Exact solution for p

m=sym.symbols("m",positive=True,nonnegative=True,real=True)
p=sym.Function("p")
sym.Eq(p(t),m*sol.args[1].subs(solved_ics[0]).diff(t))
Momenta=sym.Eq(p(t),m*sol.args[1].subs(solved_ics[0]).diff(t))
Momentaa = sym.simplify(Momenta.subs({x0:10, xeq:0 , v0:10, omega0:1, m:1}))
sym.plot(Momentaa.rhs,(t,-10,10))
```

```python
#Simple taylor series expansion

import matplotlib.pyplot as plt
xt0=.1
t=.1
vt0=.1
xlist=[]
for i in range(0,1000):
    vt=vt0+(1/2)*t**2*(-(xt0))
    vt0=vt
    xt= xt0+(vt0)*(t)
    xt0=xt
    xlist.append(xt)
plt.plot(xlist)
plt.grid(True)
plt.xlabel('x')
plt.ylabel('y')
```

```python
# Verlet algorithm for position
#r(t+dt)=2r(t)-r(t-dt)+dt**2*a
#v(t)=(r(t+dt)-r(t-dt))/2dt

xt1=.1
xt2=.1
t=.1
a=.1
plist=[]
for i in range(0,100):
    vt=(xt2-xt1)/2*t
    xt1=xt2
    xt3=vt
    plist.append(vt)
plt.plot(xlist)
plt.grid(True)
plt.xlabel('x')
plt.ylabel('y')
```

```python
# Verlet algorithm for momentum
#r(t+dt)=2r(t)-r(t-dt)+dt**2*a
#v(t)=(r(t+dt)-r(t-dt))/2dt

xt1=.1
xt2=.1
t=.1
a=.1
xlist=[]
for i in range(0,100):
    xt1=2*xt2-xt3+t**2*-xt2
    xt3=xt2
    xt2=xt1
    xlist.append(xt3)
plt.plot(xlist)
plt.grid(True)
plt.xlabel('x')
plt.ylabel('y')
```

```python
#Leapfrog for position
#r(t+dt)=r(t)+dtv(t+1/2dt)
#v(t+1/2dt)=v(t-1/2dt)+dta(t)
#v(t)=1/2(v(t+1/2dt)+v(t-1/2dt))
xt1=1
t=.2
xtwohalf=.2
xlist=[]
for i in range(0,100):
    xtwhalf=xtwohalf+t*-xt1
    xtwohalf=xtwhalf
    xt2=xt1+t*xtwhalf
    xt1=xt2
    xlist.append(xt2)
plt.plot(xlist)
plt.grid(True)
plt.xlabel('x')
plt.ylabel('y')
```

```python
#Leapfrog for velocity
#r(t+dt)=r(t)+dtv(t+1/2dt)
#v(t+1/2dt)=v(t-1/2dt)+dta(t)
#v(t)=1/2(v(t+1/2dt)+v(t-1/2dt))
vt1=2
t=.2
vtwohalf=.2
vtwhalf=2
vlist=[]
for i in range(0,100):
    vt11=(1/2)*((vtwhalf)+(vtwohalf))
    vtwohalf=vt11
    vtwhalf=vtwohalf+t*-vt1
    vtwohalf=vtwhalf
    vt2=vt1+t*vtwhalf
    vt1=vt2
    vlist.append(vt11)
plt.plot(vlist)
plt.grid(True)
plt.xlabel('x')
plt.ylabel('y')
```

```python
#velocity verlet for position
#r(t+dt)=r(t)+dtv(t)+1/2dt^2a(t)
#v(t+dt)=v(t)+1/2dt(a(t)+a(t+dt))
#v(t+1/2dt)=v(t)+1/2dta(t)
#v(t+dt)=v(t+1/2dt)+1/2dta(t+dt)

dt=0.01
x1=0
v1=1
x=[]
for i in range(0,1000):
    x2=x1+v1*dt+1/2*-x1*dt**2
    x1=x2
    v2=v1+1/2*(-x2-x1)*dt
    v1=v2
    x.append(x2)
plt.plot(x)
plt.grid(True)
plt.xlabel('x')
plt.ylabel('y')
```

```python
#velocity verlet for Velocity
#r(t+dt)=r(t)+dtv(t)+1/2dt^2a(t)
#v(t+dt)=v(t)+1/2dt(a(t)+a(t+dt))
#v(t+1/2dt)=v(t)+1/2dta(t)
#v(t+dt)=v(t+1/2dt)+1/2dta(t+dt)

xt1=1
dt=2
vt1=1
vell=[]
for i in range(0,10):
    vwhalf=vt1+1/2*dt*-xt1
    xt1=vwhalf
    vell.append(vwhalf)
plt.plot(vell)
plt.grid(True)
plt.xlabel('x')
plt.ylabel('y')
```


```python
%matplotlib inline

import numpy as np
import pylab as pl
import math
from sympy import *
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from mpl_toolkits.mplot3d import Axes3D
```

```python
from NN import NN
```

コブ・ダクラス型生産関数と課題文で例に出された関数を用いる。

いずれも定義域は0≤x≤1である。

コブ・ダグラス型生産関数は以下の通りである。

z = x_1**0.5*x_2*0.5

```python
def example1(x_1, x_2):
    z = x_1**0.5*x_2*0.5
    return z

fig = pl.figure()
ax = Axes3D(fig)

X = np.arange(0, 1, 0.1)
Y = np.arange(0, 1, 0.1)

X, Y = np.meshgrid(X, Y)
Z = example1(X, Y)

ax.plot_surface(X, Y, Z, rstride=1, cstride=1)

pl.show()
```

課題の例で使われた関数は以下の通りである。

z = (1+np.sin(4*math.pi*x_1))*x_2*1/2

NNのクラスはすでにNN.pyからimportしてある。

```python
nn = NN()
```

以下に使い方を説明する。

初めに、このコブ・ダグラス型生産関数を用いる。

```python
x_1 = Symbol('x_1')
x_2 = Symbol('x_2')
f = x_1**0.5*x_2*0.5
```

入力層、中間層、出力層を作る関数を実行する。引数には層の数を用いる。

```python
nn.set_input_layer(2)
```

    array([ 0.,  0.])

```python
nn.set_hidden_layer(2)
```

    (array([ 0.,  0.]), array([ 0.,  0.]))

```python
nn.set_output_layer(2)
```

    (array([ 0.,  0.]), array([ 0.,  0.]), array([ 0.,  0.]))

nn.set_hidden_layer()は同時にシグモイド関数で変換する前の中間層も作る。

set_output_layer()は同時にシグモイド関数で変換する前の出力層、さらに教師データを入れる配列も作る。

nn.setup()で入力層ー中間層、中間層ー出力層間の重みを入れる配列を作成する。

nn.initialize()で重みを初期化する。重みは-1/√d ≤ w ≤ 1/√d (dは入力層及び中間層の数)の範囲で一様分布から決定される。

```python
nn.setup()
nn.initialize()
```

nn.supervised_function(f, idata)は教師データを作成する。引数は関数とサンプルデータをとる。

```python
idata = [1, 2]
```

```python
nn.supervised_function(f, idata)
```

nn.simulate(N, eta)は引数に更新回数と学習率をとる。普通はN=1で行うべきかもしれないが、工夫として作成してみた。N回学習した後に出力層を返す。

```python
nn.simulate(1, 0.1)
```

    array([ 0.64956737,  0.63013034])

nn.calculation()は学習せずに入力層から出力層の計算を行う。nn.simulate()内にも用いられている。

次に実際に学習を行う。サンプルデータは、

```python
X = np.arange(0, 1, 0.2)
Y = np.arange(0, 1, 0.2)
print X, Y
```

    [ 0.   0.2  0.4  0.6  0.8] [ 0.   0.2  0.4  0.6  0.8]

の組み合わせである。

```python
X = np.arange(0, 1, 0.2)
Y = np.arange(0, 1, 0.2)

a = np.array([])
b = np.array([])
c = np.array([])

nn = NN()
nn.set_network()

for x in X:
    for y in Y:
        a = np.append(a, x)
        b = np.append(b, y)

for i in range(100):
        l = np.random.choice([i for i in range(len(a))])
        m = nn.main(1, f, [a[l], b[l]], 0.5)

for x in X:
    for y in Y:
        idata = [x, y]
        c = np.append(c, nn.realize(f, idata))
```

```python
a
```

    array([ 0. ,  0. ,  0. ,  0. ,  0. ,  0.2,  0.2,  0.2,  0.2,  0.2,  0.4,
            0.4,  0.4,  0.4,  0.4,  0.6,  0.6,  0.6,  0.6,  0.6,  0.8,  0.8,
            0.8,  0.8,  0.8])

```python
b
```

    array([ 0. ,  0.2,  0.4,  0.6,  0.8,  0. ,  0.2,  0.4,  0.6,  0.8,  0. ,
            0.2,  0.4,  0.6,  0.8,  0. ,  0.2,  0.4,  0.6,  0.8,  0. ,  0.2,
            0.4,  0.6,  0.8])

```python
c
```

    array([ 0.29071791,  0.2888693 ,  0.28691894,  0.28477467,  0.28237209,
            0.27339473,  0.27172975,  0.26999122,  0.26809561,  0.26598798,
            0.25736673,  0.25605663,  0.25468786,  0.253184  ,  0.25149742,
            0.24310494,  0.24225539,  0.2413458 ,  0.24030653,  0.23909651,
            0.2308551 ,  0.23050565,  0.23008039,  0.22951698,  0.22878105])

例えば(0, 0)を入力すると0.52328635を返している（つまりa[0]とb[0]を入力して、c[0]の値を返している）。

```python
fig = pl.figure()
ax = Axes3D(fig)

ax.scatter(a, b, c)

pl.show()
```

確率的勾配降下法を100回繰り返したが見た感じから近づいている。回数を10000回に増やしてみる。

```python
X = np.arange(0, 1, 0.2)
Y = np.arange(0, 1, 0.2)

a = np.array([])
b = np.array([])
c = np.array([])

nn = NN()
nn.set_network()

for x in X:
    for y in Y:
        a = np.append(a, x)
        b = np.append(b, y)

for i in range(10000):
        l = np.random.choice([i for i in range(len(a))])
        m = nn.main(1, f, [a[l], b[l]], 0.5)

for x in X:
    for y in Y:
        idata = [x, y]
        c = np.append(c, nn.realize(f, idata))
```

```python
fig = pl.figure()
ax = Axes3D(fig)

ax.scatter(a, b, c)

pl.show()
```

見た感じ随分近づいているように見える。

```python
ここで、
```

```python
nn.hidden_layer
```

    array([ 0.19685327,  0.48830241])

同様のことを課題の例で使われた関数でも試してみる。

```python
x_1 = Symbol('x_1')
x_2 = Symbol('x_2')
f = (1+sin(4*math.pi*x_1))*x_2*1/2
```

```python
X = np.arange(0, 1, 0.2)
Y = np.arange(0, 1, 0.2)

a = np.array([])
b = np.array([])
c = np.array([])

nn = NN()
nn.set_network()

for x in X:
    for y in Y:
        a = np.append(a, x)
        b = np.append(b, y)

for i in range(1000):
        l = np.random.choice([i for i in range(len(a))])
        m = nn.main(1, f, [a[l], b[l]], 0.5)

for x in X:
    for y in Y:
        idata = [x, y]
        c = np.append(c, nn.realize(f, idata))
```

```python
fig = pl.figure()
ax = Axes3D(fig)

ax.scatter(a, b, c)

pl.show()
```

上手く近似できないので中間層の数を変えてみる。5層にしてみる。

```python
X = np.arange(0, 1, 0.2)
Y = np.arange(0, 1, 0.2)

a = np.array([])
b = np.array([])
c = np.array([])

nn = NN()
nn.set_network(h=5)

for x in X:
    for y in Y:
        a = np.append(a, x)
        b = np.append(b, y)

for i in range(1000):
        l = np.random.choice([i for i in range(len(a))])
        m = nn.main(1, f, [a[l], b[l]], 0.5)

for x in X:
    for y in Y:
        idata = [x, y]
        c = np.append(c, nn.realize(f, idata))
```

```python
fig = pl.figure()
ax = Axes3D(fig)

ax.scatter(a, b, c)

pl.show()
```

目標と比べると大きく異なる。

サンプル数を、

```python
X = np.arange(-1, 1, 0.1)
Y = np.arange(-1, 1, 0.1)
print X, Y
```

で取り、学習の際にランダムに一個選ばれたサンプルを何十回も学習させてみた。

```python
fig = pl.figure()
ax = Axes3D(fig)

f = (1+sin(4*math.pi*x_1))*x_2*1/2

X = np.arange(-1, 1, 0.1)
Y = np.arange(-1, 1, 0.1)

a = np.array([])
b = np.array([])
c = np.array([])

fig = plt.figure()

nn = NN()
nn.set_network()

for x in X:
    for y in Y:
        a = np.append(a, x)
        b = np.append(b, y)
        c = np.append(c, nn.main2(50, f, [x, y], 0.8))

for i in range(50):
    l = np.random.choice([i for i in range(len(a))])
    m = nn.main2(20, f, [a[l], b[l]], 0.5)
    c[l] = m

a = np.array([])
b = np.array([])
c = np.array([])

for x in X:
    for y in Y:
        a = np.append(a, x)
        b = np.append(b, y)
        c = np.append(c, nn.realize(f, [x, y]))

ax.scatter(a, b, c)
ax.set_zlim(0, 1)

pl.show()
```

本来ならば下のような形になるべきであるので上手くいっているとは言い難い。

```python
def example2(x_1, x_2):
    z = (1+np.sin(4*math.pi*x_1))*x_2*1/2
    return z

fig = pl.figure()
ax = Axes3D(fig)

X = np.arange(-1, 1, 0.1)
Y = np.arange(-1, 1, 0.1)

X, Y = np.meshgrid(X, Y)
Z = example2(X, Y)

ax.plot_surface(X, Y, Z, rstride=1, cstride=1)
ax.set_zlim(-1, 1)

pl.show()
```

同じ方法でコブ・ダグラス型生産関数を学習させた様子をアニメーションにしてみた。この方法が何の意味を持つかは分からないが学習はまあまよくできていた。

mp4ファイルで添付した。

[考察]ここにはコードを書けなかったが、サンプル数が多くなった時は、ミニバッチ式の学習を試すべきだと思った。

最後に交差検定を行う。

初めに学習回数が極めて少ないNNである。

```python
X = np.arange(0, 1, 0.2)
Y = np.arange(0, 1, 0.2)

a = np.array([])
b = np.array([])
c = np.array([])

for x in X:
    for y in Y:
        a = np.append(a, x)
        b = np.append(b, y)

evl = np.array([])

for i in range(len(a)):
    nn = NN()
    nn.set_network()
    for j in range(1):
        l = np.random.choice([i for i in range(len(a))])
        if l != i:
            m = nn.main(1, f, [a[l], b[l]], 0.5)
    idata = [a[i], b[i]]
    est = nn.realize(f, idata)
    evl = np.append(evl, math.fabs(est - nn.supervised_data))
```

```python
np.average(evl)
```

    0.31083460645411465

次に十分大きく(1000回に)してみる。

```python
X = np.arange(0, 1, 0.2)
Y = np.arange(0, 1, 0.2)

a = np.array([])
b = np.array([])
c = np.array([])

nn = NN()
nn.set_network(h=7)

for x in X:
    for y in Y:
        a = np.append(a, x)
        b = np.append(b, y)

evl = np.array([])

for i in range(len(a)):
    for j in range(10000):
        nn = NN()
        nn.set_network()
        l = np.random.choice([i for i in range(len(a))])
        if l != i:
            m = nn.main(1, f, [a[l], b[l]], 0.5)
    idata = [a[i], b[i]]
    evl = np.append(evl, math.fabs(nn.realize(f, idata) - nn.supervised_data))
```

```python
evl
```

    array([ 0.17422959,  0.37774617,  0.29845184,  0.4247869 ,  0.42139651,
            0.28095943,  0.48263009,  0.12094186,  0.35090278,  0.07039276,
            0.60270475,  0.38029955,  0.47076667,  0.27789737,  0.26027828,
            0.61710335,  0.51095281,  0.29832648,  0.51007664,  0.53242457,
            0.50885681,  0.48114824,  0.53944263,  0.36336567,  0.37925241])

```python
np.average(evl)
```

    0.28091957975738507

誤差の平均であるので小さい方よい。
学習回数を大幅に増やした結果、精度が上がった。

次に回数を100回にして、中間層の数を2と5で精度を比較する。

```python
X = np.arange(0, 1, 0.2)
Y = np.arange(0, 1, 0.2)

a = np.array([])
b = np.array([])
c = np.array([])

for x in X:
    for y in Y:
        a = np.append(a, x)
        b = np.append(b, y)

evl = np.array([])

for i in range(len(a)):
    for j in range(100):
        nn = NN()
        nn.set_network()
        l = np.random.choice([i for i in range(len(a))])
        if l != i:
            m = nn.main(1, f, [a[l], b[l]], 0.5)
    idata = [a[i], b[i]]
    est = nn.realize(f, idata)
    evl = np.append(evl, math.fabs(est - nn.supervised_data))
```

```python
np.average(evl)
```

    0.3090473511199906

```python
X = np.arange(0, 1, 0.2)
Y = np.arange(0, 1, 0.2)

a = np.array([])
b = np.array([])
c = np.array([])

for x in X:
    for y in Y:
        a = np.append(a, x)
        b = np.append(b, y)

evl = np.array([])

for i in range(len(a)):
    for j in range(100):
        nn = NN()
        nn.set_network(h=5)
        l = np.random.choice([i for i in range(len(a))])
        if l != i:
            m = nn.main(1, f, [a[l], b[l]], 0.5)
    idata = [a[i], b[i]]
    est = nn.realize(f, idata)
    evl = np.append(evl, math.fabs(est - nn.supervised_data))
```

```python
np.average(evl)
```

    0.28516243998794422

コブ・ダグラス型の関数が比較的単純であるから、意味はないかと思ったが、精度はかなり良くなった。

```python

```


# Volatility Smile
In the Black-Scholes model, note that the *volatility* $\sigma$ is the only unobservable parameter, we need to consider an estimate of the volatility. We'll some covered material in Chap_04.
1. **Historic volatility**: use historical data for a period of the *same length as the time to maturity* to find the volatility. However, to price, the input value of $\sigma$ should be the *forecast volatility* of the value *over the remaining life of the option* rather than an estimated volatility value from the *past market data* of the *asset price* (which is the **historic volatility**).
2. **Implied volatility**: the value of $\sigma$ that matches the *theoretical price from BS equation* with the *observed market price* of the option. Mathematically,$\\[0.5em]$
$$
% color
% Aquamarine, black, blue, brown, cyan, darkgray, gray, green, lightgray, lime, magenta, olive, orange, pink, brown, red, teal, violet, white, yellow
\DeclareMathOperator*{\argmin}{argmin}
\DeclareMathOperator*{\argmax}{argmax}
\DeclareMathOperator*{\plim}{plim}
\newcommand{\ffrac}{\displaystyle \frac}
\newcommand{\d}[1]{\displaystyle{#1}}
\newcommand{\space}{\text{ }}
\newcommand{\bspace}{\;\;\;\;}
\newcommand{\bbspace}{\;\;\;\;\;\;\;\;}
\newcommand{\QQQ}{\boxed{?\:}}
\newcommand{\void}{\left.\right.}
\newcommand{\Tran}[1]{{#1}^{\mathrm{T}}}
\newcommand{\CB}[1]{\left\{ #1 \right\}}
\newcommand{\SB}[1]{\left[ #1 \right]}
\newcommand{\P}[1]{\left( #1 \right)}
\newcommand{\abs}[1]{\left| #1 \right|}
\newcommand{\norm}[1]{\left\| #1 \right\|}
\newcommand{\given}[1]{\left. #1 \right|}
\newcommand{\using}[1]{\stackrel{\mathrm{#1}}{=}}
\newcommand{\asim}{\overset{\text{a}}{\sim}}
\newcommand{\RR}{\mathbb{R}}
\newcommand{\EE}{\mathbb{E}}
\newcommand{\II}{\mathbb{I}}
\newcommand{\NN}{\mathbb{N}}
\newcommand{\ZZ}{\mathbb{Z}}
\newcommand{\QQ}{\mathbb{Q}}
\newcommand{\PP}{\mathbb{P}}
\newcommand{\AcA}{\mathcal{A}}
\newcommand{\FcF}{\mathcal{F}}
\newcommand{\AsA}{\mathscr{A}}
\newcommand{\FsF}{\mathscr{F}}
\newcommand{\dd}{\mathrm{d}}
\newcommand{\I}[1]{\mathrm{I}\left( #1 \right)}
\newcommand{\N}[1]{\mathcal{N}\left( #1 \right)}
\newcommand{\Exp}[1]{\mathrm{E}\left[ #1 \right]}
\newcommand{\Var}[1]{\mathrm{Var}\left[ #1 \right]}
\newcommand{\Avar}[1]{\mathrm{Avar}\left[ #1 \right]}
\newcommand{\Cov}[1]{\mathrm{Cov}\left( #1 \right)}
\newcommand{\Corr}[1]{\mathrm{Corr}\left( #1 \right)}
\newcommand{\ExpH}{\mathrm{E}}
\newcommand{\VarH}{\mathrm{Var}}
\newcommand{\AVarH}{\mathrm{Avar}}
\newcommand{\CovH}{\mathrm{Cov}}
\newcommand{\CorrH}{\mathrm{Corr}}
\newcommand{\ow}{\text{otherwise}}
\newcommand{\wp}{\text{with probability }}
\newcommand{\FSD}{\text{FSD}}
\newcommand{\SSD}{\text{SSD}}v_{\text{mkt}} = v_{\text{BS}}\P{t,S_t,T,r,\sigma_{\text{imp}}, K}$$

    and thus $\sigma_{\text{imp}}$ will depend on strik price $K$ and maturity $T$.

## Volatility Smiles
A plot of the **implied volatility** $\sigma_{\text{imp}}$ of an option as a function of its **strike price** $K$ is known as a ***volatility smile***.

### For Calls and Puts
First, the parity: $p + S_0 \color{brown}{e^{-qT}} = c + Ke^{-rT}$. then

$$\left.\begin{align}
p_{\text{BS}} + S_0 \color{brown}{e^{-qT}} &= c_{\text{BS}} + Ke^{-rT}\\
p_{\text{imp}} + S_0 \color{brown}{e^{-qT}} &= c_{\text{imp}} + Ke^{-rT}
\end{align}\right\} \Rightarrow p_{\text{BS}} - p_{\text{imp}} = c_{\text{BS}} - c_{\text{imp}}$$

and this implies that $\sigma_{\text{imp}}$ are the same for both: suppose that the **implied volatility** of the put option is $22\%$, then $p_{\text{BS}}\P{t,S_t,T,r,22\%,K} = p_{\text{mkt}}$ and thus $c_{\text{BS}}\P{t,S_t,T,r,22\%,K} - c_{\text{mkt}} = 0$. Namely, the **implied volatility** of the call option is also $22\%$.

### Risk neutral density function
Let $g\P{T,S_T,t,S_t}$ (or simply $g\P{S_T}$) denote the transition density function of the asset price, from $T$ to $t$. The time-$t$ price of a *European call* with **maturity date** $T$ and **strike price** $K$ is given by

$$c\P{t,S_t,T,K} = e^{-r\P{T-t}}\int_K^\infty \P{S_T-K} g\P{S_T}\;\dd S_T$$

If we differentiate $c$ with respect to $K$, twise, we obtain

$$\ffrac{\partial c}{\partial K} = -e^{-r\P{T-t}}\int_K^\infty g\P{S_T}\;\dd S_T, \bspace \ffrac{\partial^2 c}{\partial K^2} = e^{-r\P{T-t}} g\P{K}$$

$$\Rightarrow\color{brown}{g\P{T,K,t,S_t} = g\P{K} = e^{r\P{T-t}}\ffrac{\partial^2 c}{\partial K^2}}$$

We introduce this function to declare that:

$\bspace$Suppose that market European option prices *at all strikes* are available, the **risk neutral density function** can be inferred *completely* from the *market prices* of options with *the same **maturity*** and
*different **strikes***, *without* knowing the **volatility** function.

**e.g.**

Assume that $c_1$, $c_2$ and $c_3$ are the prices of $T$-year *European call options* with **strike** $K-\delta$, $K$ and $K+\delta$, respectively. Assuming $\delta$ is small, then an estimate of $g\P K$ is

$$g\P{T,K,t,S_t} = g\P K = e^{r\P{T-t}} \ffrac{\partial^2 c}{\partial K^2} \approx e^{r\P{T-t}} \ffrac{c_1 + c_3 - 2c_2}{\delta^2}$$

$Remark$

>$$\ffrac{\partial f\P x}{\partial x} \approx \ffrac{f\P{x+\Delta x} - f\P x}{\Delta x} \Rightarrow \ffrac{\partial^2 f\P x}{\partial x^2} \approx \ffrac{f\P{x+\Delta x} - f\P x - \P{f\P{x} - f\P{x-\Delta x}}}{\Delta x^2}$$

### Volatility Smile for Foreign Currency Options
skipped...
### Implied and Lognormal Distribution for Foreign Currency Options
skipped...
### Empirical Results
skipped...

## Equity (Stock) Options
skipped...
## Alternative Ways of Characterizing the Volatility Smile
skipped...

***

FinMath/Models and Pricing of Financial Derivatives/Chap_06_Volatility_Smile.ipynb


```python
import sympy
from phasor.utilities.ipynb.displays import *
from phasor.utilities.ipynb.ipy_sympy import *
import scipy.linalg
```

    /usr/lib64/python3.5/site-packages/matplotlib/font_manager.py:273: UserWarning: Matplotlib is building the font cache using fc-list. This may take a moment.
      warnings.warn('Matplotlib is building the font cache using fc-list. This may take a moment.')
    /usr/lib64/python3.5/site-packages/matplotlib/font_manager.py:273: UserWarning: Matplotlib is building the font cache using fc-list. This may take a moment.
      warnings.warn('Matplotlib is building the font cache using fc-list. This may take a moment.')

    Populating the interactive namespace from numpy and matplotlib
    Sympy version:  1.0

```python
a, b, c= sympy.var('a, b, c')
ac, bc, cc = a.conjugate(), b.conjugate(), c.conjugate()
X = sympy.ps_In * sympy.Matrix([
[0,        0,        1 * bc, 0,        1 * c,   0],
[0,        0,        1 * ac, 1 * c,   0,        0],
[2 * b,   2 * a,   0,       0,        0,        0],
[0,        -1 * cc, 0,       0,        0,        -1 * b],
[-1 * cc, 0,        0,       0,        0,        -1 * a],
[0,        0,        0,       -2 * bc, -2 * ac, 0],
])

X = sympy.ps_In * sympy.Matrix([
[0,       1 * ac, 1 * c,   0],
[2 * a,   0,      0,       0],
[-1 * cc, 0,      0,       -1 * a],
[0,       0,      -2 * ac, 0],
])

RX = sympy.Matrix([a,c,ac,cc])

G = np.matrix([
    [0,  0,  0,  1, 0, 0],
    [0,  0,  0,  0, 1, 0],
    [0,  0,  0,  0, 0, 1],
    [-1, 0,  0,  0, 0, 0],
    [0,  -1, 0,  0, 0, 0],
    [0,  0,  -1, 0, 0, 0],
])
```

```python
dX = X * RX
tau = .1
X2 = X.subs([(a, a + tau * dX[0]/4), (c, c + tau * dX[1]/4)])
```

```python
R = np.matrix([.1,0,1,.1,0,1]).T
subs = [(a, 1), (b, 0.02003302j), (c, 1j)]
T = np.matrix(X2.subs(subs).evalf()).astype(complex)
R = np.matrix(RX.subs(subs).evalf()).astype(complex)
#R = np.matrix([1,1j,1,-1j]).T
```

```python
T
```

    matrix([[ 0.00+0.j    ,  0.00+0.9475j, -1.05+0.j    ,  0.00+0.j    ],
            [ 0.00+1.895j ,  0.00+0.j    ,  0.00+0.j    ,  0.00+0.j    ],
            [-1.05+0.j    ,  0.00+0.j    ,  0.00+0.j    ,  0.00-0.9475j],
            [ 0.00+0.j    ,  0.00+0.j    ,  0.00-1.895j ,  0.00+0.j    ]])

```python
V = np.matrix(scipy.linalg.expm(0.1 * T))
V2 = np.matrix(scipy.linalg.expm(.2 * T))
```

```python
abs(V[0,0])**2
```

```python
abs(V)
```

    matrix([[  9.96528699e-01,   9.46404435e-02,   1.04564768e-01,
               4.96406339e-03],
            [  1.89280887e-01,   9.91027626e-01,   9.92812678e-03,
               3.13823833e-04],
            [  1.04564768e-01,   4.96406339e-03,   9.96528699e-01,
               9.46404435e-02],
            [  9.92812678e-03,   3.13823833e-04,   1.89280887e-01,
               9.91027626e-01]])

```python
Y = V * R
print(Y.transpose().conjugate() * Y)
print(R.transpose().conjugate() * R)
Y
```

    [[ 4.02838049+0.j]]
    [[ 4.+0.j]]

    matrix([[ 0.80228755+0.j        ],
            [ 0.00000000+1.17069421j],
            [ 0.80228755+0.j        ],
            [ 0.00000000-1.17069421j]])

```python
Y
```

    matrix([[ 0.80228755+0.j        ],
            [ 0.00000000+1.17069421j],
            [ 0.80228755+0.j        ],
            [ 0.00000000-1.17069421j]])

```python
(V*(V.transpose().conjugate())).real
```

    matrix([[ 1.01298469,  0.        , -0.20934319,  0.        ],
            [ 0.        ,  1.01806168,  0.        ,  0.00313639],
            [-0.20934319,  0.        ,  1.01298469,  0.        ],
            [ 0.        ,  0.00313639,  0.        ,  1.01806168]])

```python
V[:3,:3]
```

    matrix([[ 0.99500417+0.j        ,  0.00000000+0.09983342j,  0.00000000+0.j        ],
            [ 0.00000000+0.09983342j,  0.99500417+0.j        ,  0.00000000+0.j        ],
            [ 0.00000000+0.j        ,  0.00000000+0.j        ,  0.99500417+0.j        ]])

```python
V[:3,3:6]
```

    matrix([[ 0.00499583+0.j        ],
            [ 0.00000000+0.00016658j],
            [ 0.00000000-0.09999992j]])

```python
V[3:6,:3]
```

    matrix([[ 0.+0.j        ,  0.+0.j        ,  0.-0.09983342j]])

```python
V[3:6,3:6]
```

    matrix([[ 0.99500417+0.j]])

```python
(V*G*V.transpose()).real
```

```python
X**2 + X + X**3
```

```python

```


# Calculating the steady solution for $\eta$ and $\theta$ domain
\begin{equation}\label{p1introhpa} \Psi _{\eta \eta \tau }  = \frac{1}{{R_s }}\Psi _{\eta \eta \eta \eta }  + \Psi _{\eta \eta \eta } \Psi _\theta   - \Psi _{\eta \eta \theta } \Psi _\eta\\
-1 \le \eta \le1, \quad 0 \le \theta \le \pi
\end{equation}

The numerical solution found in this notebook is only stable for $0 < R_s \lesssim 55$
Iterating without any time loop, this simply iterates on $\psi$ with the $\psi_{\eta\eta t}=0$ the method is highly unstable with only a few loops required to make the solution unbounded.

```julia
using DifferentialEquations
using Plots
using BenchmarkTools
using DelimitedFiles
using LaTeXStrings
using SparseArrays
using LinearAlgebra
```

In this code the vector $\Psi_j$ is used to store the solution at the current time step and $\Psi_{jnext}$ is psi value at the next time step

```julia
pyplot()
function Ψ_smallr(θ::Float64, η::Float64 ,r::Float64)
    return(-η*(η^2-1)*sin(θ)/2.0 + r*η*(η^2-1)^2.0*(2.0+η^2)*sin(2*θ)/560.0 +
        r^2.0*η*(η^2-1)^2*((-591 -2294*η^2 + 161*η^4 + 1428*η^6)*sin(θ) + 3.0*(423 + 166*η^2 - 553*η^4 - 84*η^6)sin(3*θ))/62092800)
end

function Ψ_smallr_d1(θ::Float64, η::Float64 ,r::Float64)
    return(-(3.0*η^2-1)*sin(θ)/2.0 + r*(7*η^6-9*η^2+2)*sin(2*θ)/560.0 +
        r^2.0*(3*(-197 - 1112*η^2 + 6930*η^4 - 2772*η^6 - 8085*η^8 + 5236*η^10)*sin(θ) + (1269 - 6120*η^2 - 6930*η^4 + 24948*η^6 - 10395*η^8 - 2772*η^10)*sin(3*θ))/62092800.0)
end

function  Ψ_smallr_d2(θ::Float64, η::Float64 ,r::Float64)
    return(-3.0*η*sin(θ) + r*(42*η^5-18*η)*sin(2*θ)/560.0 +
        r^2.0*(3*( - 2224*η + 27720*η^3 - 16632*η^5 - 64680*η^7 + 52360*η^9)*sin(θ) + (- 12240*η - 27720*η^3 + 149688*η^5 - 83160*η^7 - 27720*η^9)*sin(3*θ))/62092800.0)
end

function  Ψ_smallr_d3(θ::Float64, η::Float64 ,r::Float64)
    return(-3.0*sin(θ) + r*(210*η^4-18)*sin(2*θ)/560.0 +
        r^2.0*(3*( - 2224 + 83160*η^2 - 83160*η^4 - 452760*η^6 + 471240*η^8)*sin(θ) + (- 12240 - 83160*η^2 + 748440*η^4 - 582120*η^6 - 249480*η^8)*sin(3*θ))/62092800.0)
end

function Ψ_bigr(θ::Float64, η::Float64)
    return(sin(pi*η)*sin(θ)/pi)
end

function Ψ_bigr_d1(θ::Float64, η::Float64)
    return(cos(pi*η)*sin(θ))
end

function  Ψ_bigr_d2(θ::Float64, η::Float64)
    return(-sin(pi*η)*sin(θ)*pi)
end

function  Ψ_bigr_d3(θ::Float64, η::Float64)
     return(-cos(pi*η)*sin(θ)*pi^2)
end

```

    Ψ_bigr_d3 (generic function with 1 method)

```julia
#Change these to constants after
const R = Float64(200.0) # Reynolds number
const N = 200  # Total number of intervals (steps-1) in the eta direction
const M = 256  # Total number of intervals (steps-1) theta direction
const J = 2    # Total number of time steps
ψ_current = zeros(N+1,M+1)
ψ_next =zeros(N+1,M+1)
ψ_j = ψ_current
ψ_jnext = ψ_next
const Δη = 2.0/N
const Δθ = pi/M
#const Δt = 1.0
const α = 0
const c = 1.0
const d3 =       - c
const d2 =       4*c
const d1 =      -6*c

A =spdiagm(-2=>ones(N-3)*d3 ,-1 => ones(N-2)*d2, 0 => ones(N-1)*d1, 1 => ones(N-2)*d2,2 => ones(N-3)*d3)

A[1,1] = d1 + d3 #Boundary condition psi@i=2 = 0, psi_eta@i=2 = -sin(theta) add -2d3\delta\eta\sine(theta) to RHS
A[end,end] = d1 + d3 #Boundary condition psi@i=N+1 = 0, psi_eta@i=N+1 = -sin(theta) add 2d3\delta\eta\sine(theta) to RHS
LUA = lu(A)
varinfo(r"A")
```

| name                    |       size | summary                                       |
|:----------------------- | ----------:|:--------------------------------------------- |
| A                       | 17.172 KiB | 199×199 SparseMatrixCSC{Float64, Int64}       |
| LUA                     | 17.195 KiB | SuiteSparse.UMFPACK.UmfpackLU{Float64, Int64} |
| matrix_A_at_boundaries! |    0 bytes | typeof(matrix_A_at_boundaries!)               |
| matrix_A_for_coreflow!  |    0 bytes | typeof(matrix_A_for_coreflow!)                |

```julia
ψ_current[:,1] .=0.0
ψ_current[:,M+1] .=0.0
ψ_current[1,:] .= 0.0
ψ_current[N+1,:] .= 0.0
for n in 2:N
    for m in 2:M
        #ψ_current[n,m] = Ψ_smallr(Δθ*(m-1),-1.0+Δη*(n-1),R)
        ψ_current[n,m] = 0
    end
end
ψ_current
#n=1 is ficticous point at eta=-1-deltaeta, n=2 is at eta=-1 and n=N+2 and n=N+3 is ficticious at eta=1+deltaeta

```

    201×257 Matrix{Float64}:
     0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  …  0.0  0.0  0.0  0.0  0.0  0.0  0.0
     0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0     0.0  0.0  0.0  0.0  0.0  0.0  0.0
     0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0     0.0  0.0  0.0  0.0  0.0  0.0  0.0
     0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0     0.0  0.0  0.0  0.0  0.0  0.0  0.0
     0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0     0.0  0.0  0.0  0.0  0.0  0.0  0.0
     0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  …  0.0  0.0  0.0  0.0  0.0  0.0  0.0
     0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0     0.0  0.0  0.0  0.0  0.0  0.0  0.0
     0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0     0.0  0.0  0.0  0.0  0.0  0.0  0.0
     0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0     0.0  0.0  0.0  0.0  0.0  0.0  0.0
     0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0     0.0  0.0  0.0  0.0  0.0  0.0  0.0
     0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  …  0.0  0.0  0.0  0.0  0.0  0.0  0.0
     0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0     0.0  0.0  0.0  0.0  0.0  0.0  0.0
     0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0     0.0  0.0  0.0  0.0  0.0  0.0  0.0
     ⋮                        ⋮              ⋱  ⋮                        ⋮
     0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0     0.0  0.0  0.0  0.0  0.0  0.0  0.0
     0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  …  0.0  0.0  0.0  0.0  0.0  0.0  0.0
     0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0     0.0  0.0  0.0  0.0  0.0  0.0  0.0
     0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0     0.0  0.0  0.0  0.0  0.0  0.0  0.0
     0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0     0.0  0.0  0.0  0.0  0.0  0.0  0.0
     0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0     0.0  0.0  0.0  0.0  0.0  0.0  0.0
     0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  …  0.0  0.0  0.0  0.0  0.0  0.0  0.0
     0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0     0.0  0.0  0.0  0.0  0.0  0.0  0.0
     0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0     0.0  0.0  0.0  0.0  0.0  0.0  0.0
     0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0     0.0  0.0  0.0  0.0  0.0  0.0  0.0
     0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0     0.0  0.0  0.0  0.0  0.0  0.0  0.0
     0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  …  0.0  0.0  0.0  0.0  0.0  0.0  0.0

```julia
#wireframe(range(0,stop = pi,length = M+1),range(-1,stop = 1,length = N+1),ψ_j)

```

```julia
function make_ψ_jnext_on_boundaries!(ψ_j, ψ_jnext, parameters)
    (Δθ, Δη, R, N, M, rc, rd1, rd2, rd3, d1, d2, d3, A, LUA) = parameters
    for m = 2:M-2:M #For the first theta profile and last profile use second order difference formula for theta near boundaries

    #at n=1 psi[1,:] = 0 and we are only calculating the interior points from n=2 to n=N
    n=2
    bc=2*Δη*sin(Δθ*(m-1))
    bcp1=2*Δη*sin(Δθ*m)
    bcm1=2*Δη*sin(Δθ*(m-2))
    rhs = @view ψ_jnext[2:N,m]

    #Remember to fix the RHS so that the bc terms are still included

    rhs[1] = rd3*(ψ_j[2,m]+bc)    +    rd2*ψ_j[n-1,m] + rd1*ψ_j[n,m]+rd2*ψ_j[n+1,m] + rd3*ψ_j[n+2,m]+
                 rc*(-bc/2)*(-ψ_j[n,m-1]/2 + ψ_j[n,m+1]/2)-
                 rc*(-(-(ψ_j[2,m-1]+bcm1)/12 +   4*ψ_j[n-1,m-1]/3     -5*ψ_j[n,m-1]/2 + 4*ψ_j[n+1,m-1]/3 -ψ_j[n+2,m-1]/12)/2+
                      (-(ψ_j[2,m+1]+bcp1)/12 +   4*ψ_j[n-1,m+1]/3     -5*ψ_j[n,m+1]/2 + 4*ψ_j[n+1,m+1]/3 -ψ_j[n+2,m+1]/12)/2)*
                       (bc/12)- d3*bc

        for n = 3:N-1
            rhs[n-1] = rd3*ψ_j[n-2,m]    +    rd2*ψ_j[n-1,m] + rd1*ψ_j[n,m]+rd2*ψ_j[n+1,m] + rd3*ψ_j[n+2,m]
        end#End of the n loop

    #at n=N psi[N+1,:] = 0
    n=N
    rhs[N-1] = rd3*ψ_j[n-2,m]    +    rd2*ψ_j[n-1,m] + rd1*ψ_j[n,m]+rd2*ψ_j[n+1,m] + rd3*(ψ_j[N,m]-bc)+
                 rc*(-bc/2)*(-ψ_j[n,m-1]/2 + ψ_j[n,m+1]/2)-
                 rc*(-(-ψ_j[n-2,m-1]/12 + 4*ψ_j[n-1,m-1]/3  -5*ψ_j[n,m-1]/2 + 4*ψ_j[n+1,m-1]/3 -(ψ_j[N,m-1]-bcm1)/12)/2+
                      (-ψ_j[n-2,m+1]/12 + 4*ψ_j[n-1,m+1]/3  -5*ψ_j[n,m+1]/2 + 4*ψ_j[n+1,m+1]/3 -(ψ_j[N,m+1]-bcp1)/12)/2)*
                    (bc/12) + d3*bc

    matrix_A_at_boundaries!(ψ_j, m, parameters)
    lu!(LUA,A)
    ldiv!(LUA,rhs)
    end#End of the m loop
end

function make_ψ_jnext_in_core_region!(ψ_j, ψ_jnext, parameters)
    (Δθ, Δη, R, N, M, rc, rd1, rd2, rd3, d1, d2, d3, A, LUA) = parameters
    for m = 3:M-1

    #at n=1 psi[1,:] = 0 and we are only calculating the interior points from n=2 to n=N
    n=2
    bc=2*Δη*sin(Δθ*(m-1))
    bcp1=2*Δη*sin(Δθ*m)
    bcp2=2*Δη*sin(Δθ*(m+1))
    bcm1=2*Δη*sin(Δθ*(m-2))
    bcm2=2*Δη*sin(Δθ*(m-3))

    rhs = @view ψ_jnext[2:N,m]

    #Remember to fix the RHS so that the bc terms are still included

    rhs[1] = rd3*(ψ_j[2,m]+bc)    +    rd2*ψ_j[n-1,m] + rd1*ψ_j[n,m]+rd2*ψ_j[n+1,m] + rd3*ψ_j[n+2,m]+
             rc*(-bc/2)*( ψ_j[n,m-2]/12 -  2*ψ_j[n,m-1]/3  +2*ψ_j[n,m+1]/3 - ψ_j[n,m+2]/12)-
             rc*((-(ψ_j[2,m-2]+bcm2)/12 +  4*ψ_j[n-1,m-2]/3 -5*ψ_j[n,m-2]/2 + 4*ψ_j[n+1,m-2]/3 -ψ_j[n+2,m-2]/12)/12-
               2*(-(ψ_j[2,m-1]+bcm1)/12 +  4*ψ_j[n-1,m-1]/3 -5*ψ_j[n,m-1]/2 + 4*ψ_j[n+1,m-1]/3 -ψ_j[n+2,m-1]/12)/3+
               2*(-(ψ_j[2,m+1]+bcp1)/12 +  4*ψ_j[n-1,m+1]/3 -5*ψ_j[n,m+1]/2 + 4*ψ_j[n+1,m+1]/3 -ψ_j[n+2,m+1]/12)/3-
                 (-(ψ_j[2,m+2]+bcp2)/12 +  4*ψ_j[n-1,m+2]/3 -5*ψ_j[n,m+2]/2 + 4*ψ_j[n+1,m+2]/3 -ψ_j[n+2,m+2]/12)/12)*
                (bc/12) - d3*bc

        for n = 3:N-1
            rhs[n-1] = rd3*ψ_j[n-2,m] + rd2*ψ_j[n-1,m] + rd1*ψ_j[n,m]+rd2*ψ_j[n+1,m] + rd3*ψ_j[n+2,m]
        end#End of the n loop

    #at n=N psi[N+1,:] = 0
    n=N
    rhs[N-1] = rd3*ψ_j[n-2,m]    +    rd2*ψ_j[n-1,m] + rd1*ψ_j[n,m]+rd2*ψ_j[n+1,m] + rd3*(ψ_j[N,m]-bc)+
            rc*(-bc/2)*( ψ_j[n,m-2]/12 -  2*ψ_j[n,m-1]/3  +2*ψ_j[n,m+1]/3 - ψ_j[n,m+2]/12)-
            rc*((-ψ_j[n-2,m-2]/12 + 4*ψ_j[n-1,m-2]/3 -5*ψ_j[n,m-2]/2 + 4*ψ_j[n+1,m-2]/3 -(ψ_j[N,m-2]-bcm2)/12)/12-
              2*(-ψ_j[n-2,m-1]/12 + 4*ψ_j[n-1,m-1]/3 -5*ψ_j[n,m-1]/2 + 4*ψ_j[n+1,m-1]/3 -(ψ_j[N,m-1]-bcm1)/12)/3+
              2*(-ψ_j[n-2,m+1]/12 + 4*ψ_j[n-1,m+1]/3 -5*ψ_j[n,m+1]/2 + 4*ψ_j[n+1,m+1]/3 -(ψ_j[N,m+1]-bcp1)/12)/3-
                (-ψ_j[n-2,m+2]/12 + 4*ψ_j[n-1,m+2]/3 -5*ψ_j[n,m+2]/2 + 4*ψ_j[n+1,m+2]/3 -(ψ_j[N,m+2]-bcp2)/12)/12)*
                (bc/12) + d3*bc

    matrix_A_for_coreflow!(ψ_j, m, parameters)
    lu!(LUA,A)
    ldiv!(LUA,rhs)
    end#End of the m loop
end

function matrix_A_at_boundaries!( ψ_j, m, parameters)
    (Δθ, Δη, R, N, M, rc, rd1, rd2, rd3, d1, d2, d3, A, LUA) = parameters
    #This code should only be used to calculate matrix A for m= 2 or m = M because one sided theta derivatives are used at the boundaries
    bc=2*Δη*sin(Δθ*(m-1))
    bcp1=2*Δη*sin(Δθ*m)
    bcm1=2*Δη*sin(Δθ*(m-2))

    n=2
    ψ_θ = -ψ_j[n,m-1]/2 + ψ_j[n,m+1]/2
    ψ_ηηθ =  -(-(ψ_j[2,m-1]+bcm1)/12 +   4*ψ_j[n-1,m-1]/3  -5*ψ_j[n,m-1]/2 + 4*ψ_j[n+1,m-1]/3 -ψ_j[n+2,m-1]/12)/2+
              (-(ψ_j[2,m+1]+bcp1)/12 +   4*ψ_j[n-1,m+1]/3  -5*ψ_j[n,m+1]/2 + 4*ψ_j[n+1,m+1]/3 -ψ_j[n+2,m+1]/12)/2

    A[1,1]= d1+d3 -rc*(-ψ_θ/2 -  ψ_ηηθ/12)#-ve
    A[1,2]= d2    -rc*(-ψ_θ   -2*ψ_ηηθ/3)#-ve
    A[1,3]= d3    -rc*( ψ_θ/2 +  ψ_ηηθ/12)#-ve

    n=3
    ψ_θ = -ψ_j[n,m-1]/2 + ψ_j[n,m+1]/2
    ψ_ηηθ =  -(-ψ_j[n-2,m-1]/12 +   4*ψ_j[n-1,m-1]/3  -5*ψ_j[n,m-1]/2 + 4*ψ_j[n+1,m-1]/3 -ψ_j[n+2,m-1]/12)/2+
              (-ψ_j[n-2,m+1]/12 +   4*ψ_j[n-1,m+1]/3  -5*ψ_j[n,m+1]/2 + 4*ψ_j[n+1,m+1]/3 -ψ_j[n+2,m+1]/12)/2

    A[2,1]= d2 -rc*( ψ_θ   + 2*ψ_ηηθ/3)#-ve
    A[2,2]= d1
    A[2,3]= d2 -rc*(-ψ_θ   - 2*ψ_ηηθ/3)#-ve
    A[2,4]= d3 -rc*( ψ_θ/2 +   ψ_ηηθ/12)#-ve

    for n = 4:N-2
        ψ_θ = -ψ_j[n,m-1]/2 + ψ_j[n,m+1]/2
        ψ_ηηθ =  -(-ψ_j[n-2,m-1]/12 +  4*ψ_j[n-1,m-1]/3  -5*ψ_j[n,m-1]/2 + 4*ψ_j[n+1,m-1]/3 -ψ_j[n+2,m-1]/12)/2+
                  (-ψ_j[n-2,m+1]/12 +  4*ψ_j[n-1,m+1]/3  -5*ψ_j[n,m+1]/2 + 4*ψ_j[n+1,m+1]/3 -ψ_j[n+2,m+1]/12)/2

        A[n-1,n-3] = d3 -rc*(-ψ_θ/2 -  ψ_ηηθ/12)#-ve
        A[n-1,n-2] = d2 -rc*( ψ_θ   + 2*ψ_ηηθ/3)#-ve
        A[n-1,n-1] = d1
        A[n-1,n]   = d2 -rc*(-ψ_θ   - 2*ψ_ηηθ/3)#-ve
        A[n-1,n+1] = d3 -rc*( ψ_θ/2 +  ψ_ηηθ/12)#-ve

    end

    n=N-1
    ψ_θ = -ψ_j[n,m-1]/2 + ψ_j[n,m+1]/2
    ψ_ηηθ =  -(-ψ_j[n-2,m-1]/12 +  4*ψ_j[n-1,m-1]/3  -5*ψ_j[n,m-1]/2 + 4*ψ_j[n+1,m-1]/3 -ψ_j[n+2,m-1]/12)/2+
              (-ψ_j[n-2,m+1]/12 +  4*ψ_j[n-1,m+1]/3  -5*ψ_j[n,m+1]/2 + 4*ψ_j[n+1,m+1]/3 -ψ_j[n+2,m+1]/12)/2

    A[N-2,N-4]= d3 -   rc*(-ψ_θ/2 -  ψ_ηηθ/12)#-ve
    A[N-2,N-3]= d2    -rc*( ψ_θ   + 2*ψ_ηηθ/3)#-ve
    A[N-2,N-2]= d1
    A[N-2,N-1]= d2 -   rc*(-ψ_θ   - 2*ψ_ηηθ/3)#-ve

    n=N
    ψ_θ = -ψ_j[n,m-1]/2 + ψ_j[n,m+1]/2
    ψ_ηηθ =  -(-ψ_j[n-2,m-1]/12 + 4*ψ_j[n-1,m-1]/3  -5*ψ_j[n,m-1]/2 + 4*ψ_j[n+1,m-1]/3 -(-ψ_j[N,m-1]-bcm1)/12)/2+
              (-ψ_j[n-2,m+1]/12 + 4*ψ_j[n-1,m+1]/3  -5*ψ_j[n,m+1]/2 + 4*ψ_j[n+1,m+1]/3 -(-ψ_j[N,m+1]-bcp1)/12)/2

    A[N-1,N-3]= d3    -   rc*(-ψ_θ/2 -  ψ_ηηθ/12)#-ve
    A[N-1,N-2]= d2    -   rc*( ψ_θ   +2*ψ_ηηθ/3)#-ve
    A[N-1,N-1]= d1+d3 -   rc*( ψ_θ/2   +ψ_ηηθ/12)#-ve

end

function matrix_A_for_coreflow!(ψ_j, m, parameters)
    (Δθ, Δη, R, N, M, rc, rd1, rd2, rd3, d1, d2, d3, A, LUA) = parameters

    bc  =2*Δη*sin(Δθ*(m-1))
    bcp1=2*Δη*sin(Δθ*m)
    bcp2=2*Δη*sin(Δθ*(m+1))
    bcm1=2*Δη*sin(Δθ*(m-2))
    bcm2=2*Δη*sin(Δθ*(m-3))

    n=2
    ψ_θ =  ψ_j[n,m-2]/12 -  2*ψ_j[n,m-1]/3  +2*ψ_j[n,m+1]/3 - ψ_j[n,m+2]/12
    ψ_ηηθ =  (-(ψ_j[2,m-2]+bcm2)/12 +  4*ψ_j[n-1,m-2]/3  -5*ψ_j[n,m-2]/2 + 4*ψ_j[n+1,m-2]/3 -ψ_j[n+2,m-2]/12)/12-
           2*(-(ψ_j[2,m-1]+bcm1)/12 +  4*ψ_j[n-1,m-1]/3  -5*ψ_j[n,m-1]/2 + 4*ψ_j[n+1,m-1]/3 -ψ_j[n+2,m-1]/12)/3+
           2*(-(ψ_j[2,m+1]+bcp1)/12 +  4*ψ_j[n-1,m+1]/3  -5*ψ_j[n,m+1]/2 + 4*ψ_j[n+1,m+1]/3 -ψ_j[n+2,m+1]/12)/3-
             (-(ψ_j[2,m+2]+bcp2)/12 +  4*ψ_j[n-1,m+2]/3  -5*ψ_j[n,m+2]/2 + 4*ψ_j[n+1,m+2]/3 -ψ_j[n+2,m+2]/12)/12

    A[1,1]= d1+d3 -rc*(-ψ_θ/2 -  ψ_ηηθ/12)#-ve
    A[1,2]= d2    -rc*(-ψ_θ   -2*ψ_ηηθ/3)#-ve
    A[1,3]= d3    -rc*( ψ_θ/2 +  ψ_ηηθ/12)#-ve

    n=3
    ψ_θ =  ψ_j[n,m-2]/12 -  2*ψ_j[n,m-1]/3  +2*ψ_j[n,m+1]/3 - ψ_j[n,m+2]/12
    ψ_ηηθ =  (-ψ_j[n-2,m-2]/12 +  4*ψ_j[n-1,m-2]/3  -5*ψ_j[n,m-2]/2 + 4*ψ_j[n+1,m-2]/3 -ψ_j[n+2,m-2]/12)/12-
           2*(-ψ_j[n-2,m-1]/12 +  4*ψ_j[n-1,m-1]/3  -5*ψ_j[n,m-1]/2 + 4*ψ_j[n+1,m-1]/3 -ψ_j[n+2,m-1]/12)/3+
           2*(-ψ_j[n-2,m+1]/12 +  4*ψ_j[n-1,m+1]/3  -5*ψ_j[n,m+1]/2 + 4*ψ_j[n+1,m+1]/3 -ψ_j[n+2,m+1]/12)/3-
             (-ψ_j[n-2,m+2]/12 +  4*ψ_j[n-1,m+2]/3  -5*ψ_j[n,m+2]/2 + 4*ψ_j[n+1,m+2]/3 -ψ_j[n+2,m+2]/12)/12

    A[2,1]= d2 -rc*( ψ_θ   + 2*ψ_ηηθ/3)#-ve
    A[2,2]= d1
    A[2,3]= d2 -rc*(-ψ_θ   - 2*ψ_ηηθ/3)#-ve
    A[2,4]= d3 -rc*( ψ_θ/2 +   ψ_ηηθ/12)#-ve

    for n = 4:N-2
        ψ_θ = ψ_j[n,m-2]/12 -  2*ψ_j[n,m-1]/3  +2*ψ_j[n,m+1]/3 - ψ_j[n,m+2]/12
        ψ_ηηθ =  (-ψ_j[n-2,m-2]/12 +   4*ψ_j[n-1,m-2]/3  -5*ψ_j[n,m-2]/2 + 4*ψ_j[n+1,m-2]/3 -ψ_j[n+2,m-2]/12)/12-
               2*(-ψ_j[n-2,m-1]/12 +   4*ψ_j[n-1,m-1]/3  -5*ψ_j[n,m-1]/2 + 4*ψ_j[n+1,m-1]/3 -ψ_j[n+2,m-1]/12)/3+
               2*(-ψ_j[n-2,m+1]/12 +   4*ψ_j[n-1,m+1]/3  -5*ψ_j[n,m+1]/2 + 4*ψ_j[n+1,m+1]/3 -ψ_j[n+2,m+1]/12)/3-
                 (-ψ_j[n-2,m+2]/12 +   4*ψ_j[n-1,m+2]/3  -5*ψ_j[n,m+2]/2 + 4*ψ_j[n+1,m+2]/3 -ψ_j[n+2,m+2]/12)/12

        A[n-1,n-3]= d3 -rc*(-ψ_θ/2 -  ψ_ηηθ/12)#-ve
        A[n-1,n-2]= d2 -rc*( ψ_θ   + 2*ψ_ηηθ/3)#-ve
        A[n-1,n-1]= d1
        A[n-1,n]  = d2 -rc*(-ψ_θ   - 2*ψ_ηηθ/3)#-ve
        A[n-1,n+1]= d3 -rc*( ψ_θ/2 +  ψ_ηηθ/12)#-ve

    end

    n=N-1
    ψ_θ =  ψ_j[n,m-2]/12 -  2*ψ_j[n,m-1]/3  +2*ψ_j[n,m+1]/3 - ψ_j[n,m+2]/12
    ψ_ηηθ =  (-ψ_j[n-2,m-2]/12 +  4*ψ_j[n-1,m-2]/3  -5*ψ_j[n,m-2]/2 + 4*ψ_j[n+1,m-2]/3 -ψ_j[n+2,m-2]/12)/12-
           2*(-ψ_j[n-2,m-1]/12 +  4*ψ_j[n-1,m-1]/3  -5*ψ_j[n,m-1]/2 + 4*ψ_j[n+1,m-1]/3 -ψ_j[n+2,m-1]/12)/3+
           2*(-ψ_j[n-2,m+1]/12 +  4*ψ_j[n-1,m+1]/3  -5*ψ_j[n,m+1]/2 + 4*ψ_j[n+1,m+1]/3 -ψ_j[n+2,m+1]/12)/3-
             (-ψ_j[n-2,m+2]/12 +  4*ψ_j[n-1,m+2]/3  -5*ψ_j[n,m+2]/2 + 4*ψ_j[n+1,m+2]/3 -ψ_j[n+2,m+2]/12)/12

    A[N-2,N-4]= d3 -   rc*(-ψ_θ/2 -  ψ_ηηθ/12)#-ve
    A[N-2,N-3]= d2    -rc*( ψ_θ   + 2*ψ_ηηθ/3)#-ve
    A[N-2,N-2]= d1
    A[N-2,N-1]= d2 -   rc*(-ψ_θ   - 2*ψ_ηηθ/3)#-ve

    n=N
    ψ_θ =  ψ_j[n,m-2]/12 -  2*ψ_j[n,m-1]/3  +2*ψ_j[n,m+1]/3 - ψ_j[n,m+2]/12
    ψ_ηηθ =  (-ψ_j[n-2,m-2]/12 + 4*ψ_j[n-1,m-2]/3 -5*ψ_j[n,m-2]/2 + 4*ψ_j[n+1,m-2]/3 -(ψ_j[N,m-2]-bcm2)/12)/12-
           2*(-ψ_j[n-2,m-1]/12 + 4*ψ_j[n-1,m-1]/3 -5*ψ_j[n,m-1]/2 + 4*ψ_j[n+1,m-1]/3 -(ψ_j[N,m-1]-bcm1)/12)/3+
           2*(-ψ_j[n-2,m+1]/12 + 4*ψ_j[n-1,m+1]/3 -5*ψ_j[n,m+1]/2 + 4*ψ_j[n+1,m+1]/3 -(ψ_j[N,m+1]-bcp1)/12)/3-
             (-ψ_j[n-2,m+2]/12 + 4*ψ_j[n-1,m+2]/3 -5*ψ_j[n,m+2]/2 + 4*ψ_j[n+1,m+2]/3 -(ψ_j[N,m+2]-bcp2)/12)/12

    A[N-1,N-3]= d3    -   rc*(-ψ_θ/2 -  ψ_ηηθ/12)#-ve
    A[N-1,N-2]= d2    -   rc*( ψ_θ   +2*ψ_ηηθ/3)#-ve
    A[N-1,N-1]= d1+d3 -   rc*( ψ_θ/2   +ψ_ηηθ/12)#-ve

end
```

    matrix_A_for_coreflow! (generic function with 1 method)

```julia
const rd3 =      0.0
const rd2 =      0.0
const rd1 =      0.0
const rc    =   R*Δη/Δθ

parameters =  (Δθ, Δη, R, N, M, rc, rd1, rd2, rd3, d1, d2, d3, A, LUA);

```

    WARNING: redefinition of constant rc. This may fail, cause incorrect answers, or produce other errors.

```julia
#Forming the righthand side vector

percent = 5.0
old_norm = 0.0

#Only theta profiles from m =2 to m = M need to be calculated with one sided theta derivatives second order.
#11111111111111111111111111111111111111111111111111111111
@time for j =1:J

    make_ψ_jnext_on_boundaries!(ψ_j, ψ_jnext, parameters)
    make_ψ_jnext_in_core_region!(ψ_j, ψ_jnext, parameters)

    if( mod(j,2) == 1)
        ψ_j    = ψ_next
        ψ_jnext=ψ_current
    else
        ψ_j    = ψ_current
        ψ_jnext=ψ_next
    end #End of if statement

    if( j> round(percent*J/100))
        #Calculate a simple norm of psi
        psi_norm = norm(ψ_j[:,4])
        println(string(percent)*"% compelete T="*string(round(j; digits=2))*" Core flow norm = "*string(round(psi_norm; digits=6)))
        percent+=5
        if isnan(psi_norm )
           println("Psi is no longer a number, exiting... ")
           exit()
        elseif(abs(psi_norm-old_norm)<0.00001)
            println("Psi has converged, exiting... ")
            break
       end
    end

end #End of time loop j
println(string(percent)*"% compelete      Core flow norm = "*string(round(norm(ψ_j[:,4]); digits=6)))
println("Reynolds number= "*string(R))
#At this point inverse of A should be applied to the rhs vector to produce the next iterate of psi at the j+1 time step

```

    5.0% compelete T=1.0 Core flow norm = 0.071721
    10.0% compelete T=2.0 Core flow norm = 0.14169
      0.372373 seconds (898.57 k allocations: 142.587 MiB, 9.12% gc time, 75.40% compilation time)
    15.0% compelete      Core flow norm = 0.14169
    Reynolds number= 200.0

```julia
open("psi_solution.txt", "w") do io
           writedlm(io, ψ_current)
       end
#wireframe(range(0,stop = pi,length = M+1),range(-1,stop = 1,length = N+1),ψ_j)
```

```julia
m=120
MatrixPsi=readdlm("psifin.m",Float64,);
xinterval=range(-1,stop=1,length=101)
size(MatrixPsi)
eta=range(-1, stop=1,length= N+1)
plot(eta,ψ_j[:,m],linecolor=:orange, label="Julia", linestyle=:solid)
plot!(x->Ψ_smallr(Δθ*(m-1),x,R),-1,1, label="Small R analytical", linestyle=:dot, linecolor=:blue,xlabel = L"\eta", ylabel = L"\Psi",size=(800,400))
plot!(xinterval,MatrixPsi[2:102,m], label="Spectral", linestyle=:dash, linecolor=:black,xlabel = L"\eta", ylabel = L"\Psi")
plot!(x->Ψ_bigr((m-1)*Δθ,x),-1,1, label="Large R analytical",  linestyle=:dash, linecolor=:red,xlabel = L"\eta", ylabel = L"\Psi",size=(800,400),yguidefontrotation=-90, legend=:outertopright)
```

```julia
eta_profile =0.6
eta_steps = size(MatrixPsi)[1]
delta_eta_steps = 2/(eta_steps-3)
theta=range(0, stop=pi,length= M+1)
plot(theta,ψ_j[round(Int,(eta_profile+1)/Δη)+1,:],linecolor=:orange, label="Julia", linestyle=:solid)
plot!(x->Ψ_smallr(x,eta_profile,R),0,pi, label="Small R analytical", linestyle=:dot, linecolor=:blue,xlabel = L"\theta", ylabel = L"\Psi",size=(800,400))
#plot!(x->Ψ_bigr(x,eta_profile),0,pi, label="Large R analytical",  linestyle=:dash, linecolor=:red,xlabel = L"\theta", ylabel = L"\Psi",size=(800,400),yguidefontrotation=-90, legend=:outertopright)
plot!(theta,MatrixPsi[round(Int,(eta_profile+1)/delta_eta_steps)+2,1:256])
```

```julia
round(Int,(eta_profile+1)/Δη)+1
```

    161

```julia
for n =2:N
    for m =2:M
            η = -1 + Δη*(n-1)
            θ = -1 + Δθ*(m-1)
            ψ_jnext[n,m] = ψ_j[n, m] - Ψ_smallr(θ,η,R)
    end
end

```

```julia
wireframe(range(0,stop = pi,length = M+1),range(-1,stop = 1,length = N+1),ψ_jnext)
```

```julia

```


#### _Speech Processing Labs 2020: Module 1_

```python
## Run this first!
%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np

import cmath
from math import floor
from matplotlib.animation import FuncAnimation
from IPython.display import HTML
plt.style.use('ggplot')

from dspMisc import *
```

# 5 Interpreting the  Discrete Fourier Transform

### Learning Outcomes
* Understand how sampling rate effects the DFT output
* Understand what the DFT leakage is.

### Need to know
* Topic Videos: Fourier Analysis, Frequency Domain
* [Digital Signals: Sampling sinusoids](./sp-m1-3-sampling-sinusoids.ipynb)
* [The Discrete Fourier Transform](./sp-m1-4-discrete-fourier-transform.ipynb) _(desirable)_

Equation alert: If you're viewing this on github, please note that the equation rendering is not always perfect. You should view the notebooks through a jupyter notebook server for an accurate view.

## 5.1 A Very Quick DFT recap
The previous worksheet went through the mechanics of the Discrete Fourier Transform (DFT).  The DFT input and output are broadly:

* **Input:**  $N$ amplitude samples over time
    * $x[n]$, for $n=1..N$ (i.e. a time series of $N$ samples)

* **Output:** the dot product (i.e., the similiarity) between the input and $N$ sinusoids with different frequencies
    * DFT[k] $= Me^{j\phi}$, i.e. a complex number with **magnitude** $M$ and **phase** angle $\phi$

The outputs are calculated using the following formula for $k=0,...N-1$.

$$
\begin{align}
DFT[k] &= \sum_{n=0}^{N-1} x[n] e^{-j \frac{2\pi nk}{N}} \\
&= \sum_{n=0}^{N-1} x[n]\big[\cos\big(\frac{2\pi nk}{N} \big) - j \sin\big(\frac{2\pi nk}{N} \big) \big]
\end{align}
$$

We saw how the complex numbers in this equation $e^{-j2\pi nk/N}$ for $k=1...N-1$ represent sinusoids with specific frequencies as projections of phasors. You can just think of a phasor as an analogue clock with one hand (vector) ticking (rotating) around a clockface (i.e. a circle), where the length of the hand is the amplitude of that wave, and how fast it goes around the clock is it's frequency.

Each DFT[k] output essentially tells us how strongly the input is repeats itself (i.e. is periodic) at the frequency of the $k$th DFT phasor. So, we talk about the DFT outputs as providing a **frequency response**.

Since the the DFT outputs are complex numbers, we can talk about them in terms of magnitude and phase angle: The magnitude of DFT[k] tells us how much we'd weight the $k$-th phasor if we were to try to reconstruct the original input by adding all the DFT phasors together.  The phase of DFT[k] tells use whether we need to shift that wave along the time axis.

## 5.2 The DFT Frequency Response:  Which Frequencies?

You will have noticed in the previous worksheet that we only talked about relative frequencies above, e.g. the $e^{-j2\pi nk/N}$ phasor has $k$ times the frequency of the $e^{-j2\pi n/N}$ phasor.

Since the frequencies of the DFT phasors are all multiples of the $k=1$ phasor, we just have to figure out the frequency of that phasor and we'll know all the rest. Now, the $k=1$ phasor takes $N$ steps to complete a full cycle of the unit circle.  So, it's frequency depends on how long each of these steps takes.  That is, the **sampling rate**.

Assume we are sampling at a rate of $f_s$ samples per second. Then, we can calculate:

* How long between each sample (i.e. the **sampling time**):
    * $t_s = 1/f_s$  (seconds)

* How long does it take to take $N$ samples?
    * $t_s \times N$  (seconds x samples = seconds)

* How long will it take the $k=1$ phasor to make 1 complete cycle?
    * $T_1 = t_s \times N$  (seconds)
    * This is the **period** or **wavelength** of the phasor

* What is the **frequency** of the $k=1$ phasor?
    * $f_{min} = 1/T_1 = 1/(t_s N) $ (cycles/second)

* What is the **frequency** of the $k$th phasor?
    * $kf_{min}$ (cycles/second)

So, the important thing to remember is that the DFT outputs depend on:

* The **number of samples** in the input sequence, $N$
* The **sampling rate**, $f_s$ samples/second

### Example
Assume we have a sampling rate of $f_s = 32$ samples/second, and an input length of $N=16$.

We can then calculate the $k=1$ DFT output frequency by noting the following:

* The sampling time is: $t_s = 1/f_s = 1/32$ seconds
* So, the time it takes to complete one cycle is: $T_1 = t_sN = 1/32 \times 16 = 1/2$ seconds
* So, the frequency of DFT[1] is $f_{min} = 1/1/2 = 2$ cycles/second

### Exercise

* What's the frequency of the $k=5$ phasor in the previous example?

### Notes

$kf_{min} = (5)(2) = 10$ cycles per second

### Exercise

Try running this code which sets the input to be a sine wave with the same frequency as the the 1st DFT component. That is, `f_in = 1/wavlen_in`.

```python
N=64

#sampling rate:
f_s = 128

## sample time
t_s = 1/f_s

## Notice that the DFT component frequencies as completely determined by the number of samples N
## and the sampling frequency f_s.
## So if we wanted the input frequency to exactly match the frequency of the 1st DFT component we
## could just do this:

wavlen_in = t_s*N
f_in = 1/wavlen_in

print("f_s: %f\nt_s: %f\nwavlen: %f\nf_in %f" % (f_s, t_s, wavlen_in, f_in))

## indices of input sequence of size N
nsteps = np.array(range(N))

samples_per_cycle = wavlen_in/t_s
print("samples per cycle %f" % samples_per_cycle)

## Generate a sine wave with f_in matching the frequency of the k=1 DFT phasor
## You can look at the code for this function in the file XXX.py

x, time_steps = gen_sinusoid(frequency=f_in, phase=0, amplitude=1, sample_rate=f_s, seq_length=N, gen_function=np.sin)

## Plot the input
fig, timedom = plt.subplots(figsize=(16, 4))
timedom.scatter(time_steps, x, color='magenta')
timedom.plot(time_steps, x, color='magenta')
timedom.set_xlabel("time (s)")
timedom.set_ylabel("Amplitude")
timedom.set_title("Input: $x[n]$ in the time domain")
```

```python
## Calculate the frequencies associated with each DFT output
def get_dft_freqs_all(sample_rate, seq_len):
    ## Get the minimum frequency
    f_min = sample_rate/seq_len

    ## All the other frequencies are just integer multiples of the minimum frequency
    ## We have 0,...,seq_len-1 outputs
    dft_outs = np.arange(seq_len)

    return f_min * dft_outs
```

```python
## Find the DFT output frequencies for our current sample rate and input length
dft_freqs = get_dft_freqs_all(sample_rate=f_s, seq_len=N)
print(dft_freqs)

```

    [  0.   2.   4.   6.   8.  10.  12.  14.  16.  18.  20.  22.  24.  26.
      28.  30.  32.  34.  36.  38.  40.  42.  44.  46.  48.  50.  52.  54.
      56.  58.  60.  62.  64.  66.  68.  70.  72.  74.  76.  78.  80.  82.
      84.  86.  88.  90.  92.  94.  96.  98. 100. 102. 104. 106. 108. 110.
     112. 114. 116. 118. 120. 122. 124. 126.]

```python
## Apply the DFT to the input and return the magnitudes and phases for each DFT output
mags, phases = get_dft_mag_phase(x, seq_len=N)

## Plot the magnitudes
fig, fdom = plt.subplots(figsize=(16, 4))
fdom.scatter(dft_freqs, mags)
fdom.set_xlabel("Frequency (Hz)")
fdom.set_ylabel("Magnitude")
fdom.set_title("Magnitude response")

#print(mags)

## Plot the phases
fig, fdom = plt.subplots(figsize=(16, 4))
fdom.scatter(dft_freqs, phases)
fdom.set_xlabel("Frequency (Hz)")
fdom.set_ylabel("Phase angle (radians)")
fdom.set_title("Phase response")

#print(phases)

```

### Exercise

* Is the output what you'd expect given the input?
* Why is the phase non-zero for the DFT output's 1 and 63?
* What happens if you change the `amplitude` of the input sine wave? What about the `phase`?
* What happens to the range of DFT output frequencies if you increase the sampling rate `f_s` but keep the input sequence length `N` fixed?
* What happens to the range of DFT output frequencies for $k > N/2$? Can the DFT actually model those frequencies?

### Notes

```python

```

## 5.3 DFT on a compound waveform

Now, let's see what happens when we use a compound waverform as input to the DFT.  First we need to set the length of the sequence and the sample rate.

```python
## Set the number of samples N, sampling rate f_s
N=64

#sampling rate:
f_s = 64

## sample time
t_s = 1/f_s

print("N: %d\nf_s: %f\nt_s: %f" % (N, f_s, t_s))
```

    N: 64
    f_s: 64.000000
    t_s: 0.015625

Let's create a compound waveform as input with frequency components of 4Hz and 12Hz:

```python
## Now let's create a compound waveform as input
## We're setting gen_function=np.cos, meaning that each of our sinusoids is a cosine wave
x1, time_steps = gen_sinusoid(frequency=4, phase=0, amplitude=1, sample_rate=f_s, seq_length=N, gen_function=np.cos)

## Let's shift the phase of this one by 90 degrees (pi/2 rads)
x2, time_steps = gen_sinusoid(frequency=12, phase=np.pi/2, amplitude=1, sample_rate=f_s, seq_length=N, gen_function=np.cos)

x_compound = x1 + x2

## Plot the input
fig, timedom = plt.subplots(figsize=(16, 4))
timedom.scatter(time_steps, x_compound, color='magenta')
timedom.plot(time_steps, x_compound, color='magenta')
timedom.set_xlabel("time (s)")
timedom.set_ylabel("Amplitude")
timedom.set_title("A compound waveform: $x_compound$ in the time domain")
```

You should see a somewhat complicated waveform whose overall pattern repeats 4 times in a second.  Now we do the DFT and look at the magnitude and phase values for the different DFT outputs.

```python
## Do the DFT on the compound waveform as above:
mags, phases = get_dft_mag_phase(x_compound, seq_len=N)
dft_freqs = get_dft_freqs_all(sample_rate=f_s, seq_len=N)

## Plot the magnitudes
fig, fdom = plt.subplots(figsize=(16, 4))
#timedom.scatter(dft_freqs[:round(N)], mags[:round(N)])
fdom.scatter(dft_freqs[:round(N/2)], mags[:round(N/2)])
fdom.set_xlabel("Frequency (Hz)")
fdom.set_ylabel("Magnitude")
fdom.set_title("Magnitude response")
#print(mags)

## Plot the phases - let's just skip this for the moment
fig, fdom = plt.subplots(figsize=(16, 4))
fdom.scatter(dft_freqs[:round(N/2)], phases[:round(N/2)])
fdom.set_xlabel("Frequency (Hz)")
fdom.set_ylabel("Phase (rad)")
fdom.set_title("Phase response")
#print(phases)

```

### Exercise

* Do you see what you would expect given the input?
* Why does the number of non-zero phases differ from the number of non-zero magnitudes?

### Notes

```python

```

### Exercise

Explore the outputs of the DFT when you change the following parameters
* sampling rate `f_s`
* sequence length `N`
* frequency of the input sequence(s) `frequency` (in the function `gen_sinusoid`)

**Question**
* What happens if you set the frequency of one of the input components to be greater than the Nyquist frequency?
* What happens when the frequency of the input components doesn't match any of the DFT component frequencies?
(e.g. for the example above what if one of the input frequencies is not a whole number?)

### Notes

```python

```

## 5.4  Leakage

One of the main things to remember about the DFT is that you're calculating the correlation between the input and phasors with specific frequencies.  If your input exactly matches one of those phasor frequencies the magnitude response will show a positive magnitude for that phasor and zero for everything else.  However, if the input frequency falls between phasor frequencies (i.e. doesn't correspond to a specific DFT output), then you'll see spectral **leakage**.  The DFT outputs close to the input frequency will also get positive magnitudes, with the DFT output closest to the input frequency getting the highest magnitude.

The following code gives an example

```python
N=64
f_s = 64
f_in1 = 4.5 ## In between DFT output frequencies
f_in2 = 12  ## One of the DFT outputs

x1, time_steps = gen_sinusoid(frequency=f_in1, phase=0, amplitude=1, sample_rate=f_s, seq_length=N, gen_function=np.cos)
x2, time_steps = gen_sinusoid(frequency=f_in2, phase=np.pi/2, amplitude=1, sample_rate=f_s, seq_length=N, gen_function=np.cos)
x_compound = x1  + x2

## Plot the input
fig, timedom = plt.subplots(figsize=(16, 4))
timedom.scatter(time_steps, x_compound, color='magenta')
timedom.plot(time_steps, x_compound, color='magenta')
timedom.set_xlabel("time (s)")
timedom.set_ylabel("Amplitude")
timedom.set_title("Leakage example (time domain)")

```

```python
## Do the DFT on the compound waveform as above:
mags, phases = get_dft_mag_phase(x_compound, seq_len=N)
dft_freqs = get_dft_freqs_all(sample_rate=f_s, seq_len=N)

## Plot the magnitudes
fig, fdom = plt.subplots(figsize=(16, 4))

## Just plot the first N/2 frequencies since we know that they are the mirrored for k>N/2
fdom.scatter(dft_freqs[:round(N/2)], mags[:round(N/2)])
fdom.set_xlabel("Frequency (Hz)")
fdom.set_ylabel("Magnitude")
fdom.set_title("Leakage example: Magnitude response")
#print(mags)

```

### Leakage as the normalized sinc function

Leakage makes the DFT harder to interpret. However, we can derive the shape that leakage will have from the the DFT equation and some algebra about rectangular functions.  It turns out that leakage for a particular frequency follows the normalized **sinc** function:

$$
\begin{align}
X(m) &= \Big|\frac{AN}{2} \cdot \mathop{sinc}(c-m)\Big|\\
&= \Big|\frac{AN}{2} \cdot \frac{\sin(\pi(c-m))}{2\pi(c-m))}\Big|\\
\end{align}
$$

Where $A$ is the peak amplitude of the input, $N$ is the input sequence length, $c$ is the number of cycles completed in the input sequence time.  If $c$ is a whole number we just get the usual DFT response, but if $c$ is not a whole number, we get a spread across output frequency bins.

Let's check whether the sinc function matches what we get in the DFT. First we write a function to evaluate the leakage function in between our DFT outputs.

```python
## Calculate the approximated leakage as the sinc function
def calc_leakage(freq,  sample_rate, seqlen, amplitude=1):
    sequence_time = (1/sample_rate)*seqlen

    ## number of cycles in input for component
    c = freq * sequence_time
    print("c=", c)

    ## Interpolate between DFT ouput indices
    ms = np.arange(0, seqlen, 0.05)

    ## Approximated response - we could actually just return a function here, but
    ## let's just keep things concrete for now.
    leakage = np.abs(amplitude * seqlen * 0.5 * np.sinc((c-ms)))

    return leakage, ms * (sample_rate/seqlen)

```

Now let's plot the leakage predicted for our two input components separately (top) and added together (bottom)

```python
## Calculate the leakage function for our know input wave components
leakage1, ms = calc_leakage(f_in1, f_s, N)
leakage2, ms = calc_leakage(f_in2, f_s, N)

## Plot the magnitude response and the leakage function for each of our input components
fig, fdom = plt.subplots(figsize=(16, 4))
fdom.set(xlim=(-1, N/2), ylim=(-1, N))
fdom.scatter(dft_freqs, mags)

fdom.plot(ms, leakage1)
fdom.plot(ms, leakage2)
fdom.set_xlabel("Frequency (Hz)")
fdom.set_ylabel("Magnitude")
fdom.set_title("Leakage function for each input component frequency")

## Plot the magnitude response and the sum of the leakage functions
fig, fdom = plt.subplots(figsize=(16, 4))
fdom.set(xlim=(-1, N/2), ylim=(-1, N))
fdom.scatter(dft_freqs, mags)
fdom.plot(ms, leakage1 + leakage2, color='C5')
fdom.set_xlabel("Frequency (Hz)")
fdom.set_ylabel("Magnitude")
fdom.set_title("Sum of leakage functions for input components")

## It fits, though not perfectly!
```

In the top figure, you should see that peaks (**main lobes**) of each leakage function are aligned with our input component frequencies. The peaks are at the same points as the DFT outputs when the sinusoidal component has a frequency matching the DFT output frequency (i.e. 12 above Hz).  Otherwise we see the spread of leakage around the input component frequency (i.e. around 4.5 Hz).

You'll also notice that our DFT magnitudes points don't always line up perfectly with our sinc curves.  Part of this is because the leakage function is an _approximation_.  Nevertheless, it's a very good approximation!

### Exercise

* What happens if the frequencies of the two components making the compound waveform are very close together?
    * e.g. make `f_in2=6`
* What if one of the components has a relatively small amplitude?
    * e.g. change `amplitude` of the second input to 0.4

### Notes

```python

```

### Shaping the lobes
The leakage sinc function has a peak around a specific frequency. If we want our DFT to be better able to distinguish between close frequencies, we need that peak, the **main lobe** to be narrower.  We also want the other peaks, the **side lobes** to be flatter.  We can achieve this using different **windowing methods** on our input.  This is why you see 'Hanning' as the default option for window method in the advanced spectrogram settings in praat.

We'll see some more examples of this when we look at different types of filters later.  But for now the main thing to observe is that leakage might give you the impression that specific frequency components are present in your waveform, when what's actually happening is that your waveform has frequencies that don't match the DFT phasors.  Another thing that can happen is that the peak for a low amplitude component gets subsumed into the main lobe of a nearby frequency component.  This might make you miss frequency components in your input!

```python

```

## 5.6 Extra:

Here's the composition and decomposition for the compound waveform one you saw in Notebook 2!

```python
N=1028
f_s = 1028
f_in1 = 8 ## In between DFT output frequencies
f_in2 = 20  ## One of the DFT outputs
f_in3 = 36  ## One of the DFT outputs

x1, time_steps = gen_sinusoid(frequency=f_in1, phase=0, amplitude=0.5, sample_rate=f_s, seq_length=N, gen_function=np.cos)
x2, time_steps = gen_sinusoid(frequency=f_in2, phase=np.pi/2, amplitude=1, sample_rate=f_s, seq_length=N, gen_function=np.cos)
x3, time_steps = gen_sinusoid(frequency=f_in3, phase=0, amplitude=0.3, sample_rate=f_s, seq_length=N, gen_function=np.cos)

x_compound = x1  + x2 + x3

## Plot the input
fig, timedom = plt.subplots(figsize=(16, 4))
#timedom.scatter(time_steps, x_compound, color='magenta')
#timedom.set(xlim=(0, 0.2))
timedom.plot(time_steps, x_compound, color='magenta')
```

```python
## Plot the input

fig = plt.figure(figsize=(15,15))
gs = fig.add_gridspec(3,2)

ymax=2
timedom = fig.add_subplot(gs[0, 0])
timedom.set(xlim=(-0.1, 1), ylim=(-ymax,ymax))
timedom.plot(time_steps, x_compound, color='magenta')
timedom.set_title("waveform made from adding 3 sine waves")

s1 = fig.add_subplot(gs[0, 1])
s1.set(xlim=(-0.1, 1), ylim=(-ymax,ymax))
s1.plot(time_steps, x1)
s1.set_title("component 1: 8 Hz")

s2 = fig.add_subplot(gs[1, 1])
s2.set(xlim=(-0.1, 1), ylim=(-ymax,ymax))
s2.plot(time_steps, x2)
s2.set_title("component 2: 20 Hz")

s3 = fig.add_subplot(gs[2, 1])
s3.set(xlim=(-0.1, 1), ylim=(-ymax,ymax))
s3.plot(time_steps, x3)
s3.set_title("component 3: 36 Hz")

fig.savefig("./fig/compound_waveform.png")
```

```python
## Do the DFT on the compound waveform as above:
mags, phases = get_dft_mag_phase(x_compound, seq_len=N)
dft_freqs = get_dft_freqs_all(sample_rate=f_s, seq_len=N)

## Plot the magnitudes
fig,fdom = plt.subplots(figsize=(16, 4))
fdom.set(xlim=(-1, N/2))
## Just plot the first N/2 frequencies since we know that they are the mirrored for k>N/2
fdom.scatter(dft_freqs[:round(N/2)], mags[:round(N/2)])
fdom.set_xlabel("Frequency (Hz)")
fdom.set_ylabel("Magnitude")
fdom.set_title("Magnitude Response")
#print(mags)

```

```python

```

```python

```