```python
import numpy as np
from numba import jit, float64, int64
from numba.experimental import jitclass
import matplotlib.pyplot as plt
%matplotlib inline
np.set_printoptions(precision=16, suppress=True)

from naginterfaces.library import machine, roots
from naginterfaces.base.utils import NagAlgorithmicWarning
eps = machine.precision()

from numba.core.errors import NumbaDeprecationWarning, NumbaPendingDeprecationWarning, NumbaWarning
import warnings
warnings.simplefilter('ignore', category=NumbaDeprecationWarning)
warnings.simplefilter('ignore', category=NumbaPendingDeprecationWarning)
warnings.simplefilter('ignore', category=NumbaWarning)
warnings.simplefilter('ignore', category=NagAlgorithmicWarning)
```

# Using Anderson Acceleration to speed up fixed-point iterations

##  Accelerating the fixed point solution of cos(x) = x

Type any real number into a calculator and repeatedly press the cosine button. The result eventually converges to about 0.739085 which is the solution of the equation $x = \cos(x)$.  This is an example of a fixed-point iteration.  The general case is any computation that looks like this

\begin{align}
x_{n+1} = f(x_n)
\end{align}

which is repeated until a convergence criterion is reached. Such computations arise naturally in many areas of finance, physics and engineering.

We could code the cosine fixed-point computation in Python as follows

```python
def fcos(x):
    """Performs one iteration of the cosine fixed point iteration
    """
    return np.cos(x)

def fp_driver(x0, func, tol):
    x = x0 # starting guess
    err = 1000000 # Huge error to start things off
    iterations = 0
    while err > tol:
        nextx = func(x)
        err = np.sqrt((x - nextx)**2)
        x = nextx
        iterations += 1
    return iterations, x

fp_driver(3, fcos, 1e-15)
```

    (88, 0.7390851332151603)

Demonstrating that we get convergence to a tolerance of 1e-15 in 88 iterations.  A natural question to ask is 'Can we accelerate the convergence somehow?' and the answer is 'Yes', using a technique called [Anderson Acceleration](https://epubs.siam.org/doi/abs/10.1137/10078356X).  Anderson Acceleration has been known in the chemistry community since the 1960s but until relatively recently it wasn't well known in other fields. The [2011 paper by Walker and Ni](https://epubs.siam.org/doi/abs/10.1137/10078356X) describing the technique has been cited over 245 times as of late 2019.

NAG recently introduced two Anderson Acceleration routines into the NAG Library following the success of the technique in the evaluation of Nearest correlation matrices (NCMs) (the accelerated nearest correlation matrix routine is implemented in `naginterfaces.library.correg.corrmat_fixed`). Here, we will demonstrate the use of NAG's standard Anderson Acceleration routine `roots.sys_func_aa`.

`roots.sys_func_aa` takes the role of `fp_driver` in the code above.  Here's how to use it for the cosine example

```python
def nag_fcos(x):
    """Performs one iteration of the cosine fixed point iteration
    """
    # Note that NAG requires f(x) - x to be returned instead of just f(x)
    return np.cos(x)-x

m = 0 # When m=0, Acceleration is turned off so this is identical to running a basic fixed-point calculation
tol = 1e-10
nag_cos = roots.sys_func_aa(nag_fcos, 3.0, tol, eps, m).x
nag_cos
```

    array([0.7390851331723561])

We turn on acceleration by setting `m` to something more than 0.

```python
nag_cos = roots.sys_func_aa(nag_fcos, 3.0, tol, eps, m=1).x
nag_cos
```

    array([0.7390851331642223])

Unfortunately, the NAG routine does not return the number of iterations required for convergence so we have no idea what the benefit of acceleration is unless we count them ourselves.  We do this with a helper class which we pass to the `nag_fcos` function.

```python
class fpinfo:
    """A class used to get information to/from the NAG routine
    """
    def __init__(self):
        self.iterations = 0

def nag_fcos_withdata(x, data):
    """Performs one iteration of the cosine fixed point iteration
    """
    data.iterations += 1
    # Note that NAG requires f(x) - x to be returned instead of just f(x)
    return np.cos(x)-x

tol = 1e-15
data = fpinfo()
nag_cos = roots.sys_func_aa(nag_fcos_withdata, 3.0, tol, eps, m=0, data=data).x
print(f"Solution found was {nag_cos[0]} in {data.iterations} iterations with acceleration switched off")
data = fpinfo()
nag_cos = roots.sys_func_aa(nag_fcos_withdata, 3.0, tol, eps, m=1, data=data).x
print(f"Solution found was {nag_cos[0]} in {data.iterations} iterations with acceleration switched on")
```

    Solution found was 0.7390851332151611 in 88 iterations with acceleration switched off
    Solution found was 0.7390851332151607 in 10 iterations with acceleration switched on

An almost 9x improvement in iteration count simply by switching on acceleration is extremely useful.  We move on to consider a more challenging example of fixed-point iteration.

## Solving the Poisson Equation Using Fixed-Point Iterations

[Poisson's equation](https://en.wikipedia.org/wiki/Poisson%27s_equation) is an elliptic Partial Differential Equation (PDE) which in two dimensions takes the form

\begin{align}
\nabla^2 u & = f(x,y)\nonumber \\
\left( \frac{\partial^2}{\partial x^2} + \frac{\partial^2}{\partial y^2} \right)u(x,y) &= f(x,y)\nonumber
\end{align}

The special case where $f(x)=0$ is referred to as Laplace's equation.  Both Poisson's and Laplace's equations arise in many areas of physics and engineering.

One example which requires the numerical solution of Poisson's equation is finding the steady state of the potential field of two parallel lines of electric charge .  One with total charge $+1$ and the other with total charge $-1$.  The Dirichlet boundary condition $u(x,y)=0$ is applied at the domain boundary.

```python
def source(N):
    #  Produces a matrix corresponding to a source term
    #  Want two line sources with total charge of +/- 1 on each line
    #  This is our function f(x, y) in the poisson equation
    x0 = np.zeros((N, N))
    h = 1/(N-1)
    source_rows = range((N//4), (3*N//4))
    source_col1 = N//4
    source_col2 = 3*N//4

    x0[source_rows, source_col1] = x0[source_rows, source_col1] + 1/(N*0.5*h**2)
    x0[source_rows, source_col2] = x0[source_rows, source_col2] - 1/(N*0.5*h**2)
    return x0

def init_problem(N=50):
    # Initialises an N x N problem domain with an initialised source
    x0 = np.zeros((N, N))
    x0 += source(N)
    return x0

x0 = init_problem(50)
plt.title('Initial conditions of our example problem on a 50 x 50 grid')
_ = plt.imshow(x0)
```

Many introductory courses on the solution of PDEs will demonstrate three methods for the solution of systems such as this:

* The Jacobi Method
* The Gauss-Seidel Method
* Successive Over-relaxation

A paper that discusses these methods in depth using similar notation to this notebook is [Parallel S.O.R. iterative methods](https://www.sciencedirect.com/science/article/pii/S0167819184903806) by D.J. Evans.

We now look at implementing the three solution methods in turn

### The Jacobi Method

Starting from the <i>n</i>th configuration of the discretized grid $u^{n}$, we get the next configuration by applying the following formula to every $(i,j)$th grid point other than the very edges where we apply the boundary conditions.

\begin{align}
u^{n+1}_{j,i} = \frac{1}{4} \left(u^{n}_{j+1,i} +u^{n}_{j-1,i} +u^{n}_{j,i+1} +u^{n}_{j,i-1}  \right)+\frac{h^2}{4}f_{j,i}
\end{align}

An implementation of this iteration in Python is

```python
@jit
def jacobi(x, source, solverinfo):
    """Performs one iteration of the jacobi method

    Arguments:
    x - N x N Matrix containing the previous iteration of the solution
    source - N x N Matrix containing our source function evaluated on the grid
    solverinfo - Takes a solverinfo object to get info in/out of the solver
    """
    N = x.shape[0]
    nextx = np.zeros((N, N))
    h = 1/(N-1)
    # loop over the grid
    # Only iterate over interior points thus keeping the edges untouched and hence enforcing
    # the boundary condition that x = 0 at the edges.
    for i in range(1, N - 1):
        for j in range(1, N - 1):
            nextx[j, i] = 0.25 * (x[j+1, i] + x[j, i+1] + x[j-1, i] + x[j, i-1]) + 0.25*h**2*source[j, i]
    solverinfo.iterations += 1
    return nextx
```

Next we need a driver for the Jacobi iteration function.  Something that is general enough to drive subsequent methods of solving our problem.

```python
# This spec is required for Numba jit compilation
spec = [
    ('iterations', int64),
    ('w', float64)
]

@jitclass(spec)
class solverinfo:
    """A class used to get information to/from a solver
    """
    def __init__(self, _w=1):
        self.iterations = 0
        self.w = 1 # Used for SOR and ignored for everything else

def solve_poisson(x, tol, data, method):
    """Solves Poisson's equation

    x - N x N Matrix containing the initial configuration of the problem space
    tol - Convergence tolerance
    data - A solverinfo object to get arbitrary info in/out of the solver
    method - A function that performs one solution iteration
    """
    N = int(np.sqrt(x.size))
    err = 1 # Big error to start things off with
    while err > tol:
        nextx = method(x, source(N), data)
        err = np.sqrt(np.sum((x - nextx)**2))
        x = np.copy(nextx)
    return x
```

Finally, we can set up and solve an instance of our problem using the Jacobi method.

```python
N = 100
u = init_problem(N)
jacobi_info = solverinfo()
tol = 1e-9
jacobi_sol = solve_poisson(u, tol, jacobi_info, jacobi)
print(f"Gauss-Seidel on a {N} by {N} grid")
print(f"Solution found in {jacobi_info.iterations} iterations")
_ = plt.imshow(jacobi_sol)
```

Poisson's equation gives us the potential.  For completeness, let's find the electric field by taking the gradient

```python
Ex, Ey = np.gradient(jacobi_sol)
E = np.sqrt(Ex**2+Ey**2) # Magnitude of Electric field
plt.imshow(E)
_ = plt.title('Electric Field')
```

28,734 iterations is a lot and we would hope that we could do better with more sophisticated techniques. Fortunately, we can and 'more sophisticated' isn't quite as sophisticated as you might fear.

### Gauss--Seidel Method

In the Jacobi method, we kept two independent copies of the solution space.  Each element of the **new** solution is computed from the surrounding elements of the **old** solution.  In Gauss-Seidel we make use of new grid elements as soon as they become available.  That is, when we reach position $(j,i)$ in the grid we have already updated $u_{j-1,i}$ and $u_{j,i-1}$ so we can use those instead of the old values.

\begin{align}
u^{n+1}_{j,i} = \frac{1}{4} \left(u^{n}_{j+1,i} +u^{n+1}_{j-1,i} +u^{n}_{j,i+1} +u^{n+1}_{j,i-1}  \right)+\frac{h^2}{4}f_{j,i}
\end{align}

The Gauss-Seidel iteration, therefore, looks like this

```python
@jit
def gauss_seidel(x, source, solverinfo):
    N = x.shape[0]
    nextx = np.copy(x)
    h = 1/(N-1)
    for i in range(1, N - 1):
        for j in range(1, N - 1):
            nextx[j, i] = 0.25 * (nextx[j+1, i] + nextx[j, i+1] + nextx[j-1, i] + nextx[j, i-1]) + 0.25*h**2*source[j, i]
    solverinfo.iterations += 1
    return nextx
```

```python
N = 100
u = init_problem(N)
gs_info = solverinfo()
tol = 1e-9
gs_sol = solve_poisson(u, tol, gs_info, gauss_seidel)

print(f"Gauss--Seidel on a {N} by {N} grid")
print(f"Solution found in {gs_info.iterations} iterations")
fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(10, 6))
axes[0].imshow(gs_sol)
axes[0].set_title('Potential')
Ex, Ey = np.gradient(gs_sol)
E = np.sqrt(Ex**2+Ey**2) # Magnitude of Electric field
axes[1].imshow(E)
_ = axes[1].set_title('Electric Field')
```

The number of iterations has almost halved which is a fantastic return on such a small investment.  The standard treatment of this problem has yet another trick up its sleeve with the Successive Over Relaxation method (SOR) which is an extension of the Gauss--Seidel method.

### Successive Over Relaxation Method (SOR)

Take the Gauss--Seidel update formula above (3) and add and subtract $u_{i,j}^n$ to the right-hand side

\begin{align}
u^{n+1}_{j,i} & =  \frac{1}{4} \left(u^{n}_{j+1,i} +u^{n+1}_{j-1,i} +u^{n}_{j,i+1} +u^{n+1}_{j,i-1}  +h^2 f_{j,i} \right) \nonumber \\
& = u_{i,j}^n + \frac{1}{4} \left(u^{n}_{j+1,i} +u^{n+1}_{j-1,i} +u^{n}_{j,i+1} +u^{n+1}_{j,i-1} + h^2 f_{j,i} - 4 u_{i,j}^n \right) \nonumber \\
& = u_{i,j}^n + r_{i,j}
\end{align}

So $r_{i,j}$ is the amount of change of $u_{i,j}$ for one iteration of Gauss_Seidel.  If this amount of change is good, maybe more of the same is better and it turns out that the convergence of Gauss--Seidel can be accelerated by making a larger change as follows

\begin{align}
u^{n+1}_{j,i} & = u_{i,j}^n + \omega r_{i,j}
\end{align}

where $\omega$ is positive constant called the acceleration factor that in practice lies between 1 and 2.  Note that when $\omega = 1$ the SOR reduces to standard Gauss-Seidel.

The Python looks like this

```python
@jit
def SOR(x, source, solverinfo):
    N = x.shape[0]
    nextx = np.copy(x)
    h = 1/(N-1)
    w = solverinfo.w
    for i in range(1, N - 1):
        for j in range(1, N - 1):
            new = 0.25 * (nextx[j-1, i] + nextx[j+1, i] + nextx[j, i-1]+ nextx[j, i+1]) + 0.25*h**2*source[j, i]
            nextx[j, i] += w * (new - nextx[j, i])
    solverinfo.iterations += 1
    return nextx
```

Using this to solve our problem with $\omega=1.94$  we have

```python
N = 100
u = init_problem(N)
SOR_info = solverinfo()
SOR_info.w = 1.94
tol = 1e-9
SOR_sol = solve_poisson(u, tol, SOR_info, SOR)

print(f"SOR on a {N} by {N} grid")
print(f"Solution found in {SOR_info.iterations} iterations")
fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(10, 6))
axes[0].imshow(SOR_sol)
axes[0].set_title('Potential')
Ex, Ey = np.gradient(SOR_sol)
E = np.sqrt(Ex**2+Ey**2) # Magnitude of Electric field
axes[1].imshow(E)
_ = axes[1].set_title('Electric Field')
```

486 iterations is over 59x faster than our first attempt using the Jacobi method which needed 28,734 iterations to converge.

Choose the SOR $\omega$ parameter badly, however, and you can get performance that's more than 10x slower than when we chose well!  It's still several times faster than Jacobi or Gauss--Seidel however.

```python
N = 100
u = init_problem(N)
SOR_info = solverinfo()
SOR_info.w = 1.5
tol = 1e-9
SOR_sol = solve_poisson(u, tol, SOR_info, SOR)

print(f"SOR on a {N} by {N} grid")
print(f"Solution found in {SOR_info.iterations} iterations")
fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(10, 6))
axes[0].imshow(SOR_sol)
axes[0].set_title('Potential')
Ex, Ey = np.gradient(SOR_sol)
E = np.sqrt(Ex**2+Ey**2) # Magnitude of Electric field
axes[1].imshow(E)
_ = axes[1].set_title('Electric Field')
```

There are analytic methods available to estimate the optimum $\omega$ but it's often necessary to find it empirically.  The optimum value of $\omega$ depends on both the PDE being solved and on the grid resolution.  More advanced extensions of SOR exist that have even better convergence properties (Chebyshev acceleration for example).  I could also consider parallelisation schemes for these methods such as Red-Black ordering but I need to stop somewhere if I am going to get onto Anderson Acceleration before the deadline of this article!

Before moving on, let's look at the number of iterations taken by SOR for this problem as a function of $\omega$

```python
def SOR_iterations(w):
    N = 100
    u = init_problem(N)
    SOR_info = solverinfo()
    SOR_info.w = w
    tol = 1e-9
    _ = solve_poisson(u, tol, SOR_info, SOR)
    return SOR_info.iterations

w = np.arange(1.0, 1.99, 0.01)
iterations = [SOR_iterations(omega) for omega in w]
plt.plot(w, iterations)
plt.title('Number of iterations vs SOR parameter')
plt.xlabel('w')
_ = plt.ylabel('Iterations')
```

## Applying NAG's Anderson Acceleration to these methods

### Code adaptations required to call Jacobi iterations from NAG

I started on this journey because I noted that all of these methods are fixed-point and wondered if NAG's Anderson acceleration would be useful.  Let's go back to Jacobi and recall the definition of each iteration

```python
@jit
def jacobi(x, source, solverinfo): # pylint: disable=function-redefined
    """Performs one iteration of the jacobi method

    Arguments:
    x - N x N Matrix containing the previous iteration of the solution
    source - N x N Matrix containing our source function evaluated on the grid
    solverinfo - Takes a solverinfo object to get info in/out of the solver
    """
    N = x.shape[0]
    nextx = np.zeros((N, N))
    h = 1/(N-1)
    # loop over the grid
    # Only iterate over interior points thus keeping the edges untouched and hence enforcing
    # the boundary condition that x = 0 at the edges.
    for i in range(1, N - 1):
        for j in range(1, N - 1):
            nextx[j, i] = 0.25 * (x[j+1, i] + x[j, i+1] + x[j-1, i] + x[j, i-1]) + 0.25*h**2*source[j, i]
    solverinfo.iterations += 1
    return nextx
```

The NAG routine that performs Anderson acceleration is `roots.sys_func_aa` and the [documentation for it](https://www.nag.com/numeric/py/nagdoc_latest/naginterfaces.library.roots.html) tells us that the function to be solved must have the form `fcn(x, data)` where `data` is any arbitrary Python object that can be used to get information in and out of `fcn`.  As such, we are going to need to change the `solverinfo` object so that it also holds the `source` term.

Two other constraints are going to change the way this function looks:

* The vector $x$ needs to be one-dimensional. As such, we'll need to change the shape to and from the two-dimensional representation we've been using until now.
* We need to return `jacobi(x) - x` instead of `jacobi(x)`

With all of this taken into account, the NAG-ified version of `jacobi` and the supporting `solverinfo` class is

```python
# This spec is required for Numba jit compilation
spec = [
    ('iterations', int64),
    ('w', float64),
    ('source', float64[:, :])
]

@jitclass(spec)
class NAG_solverinfo:
    """A class used to get information to/from a solver
    """
    def __init__(self, N=50, _w=1):
        self.iterations = 0
        self.w = 1 # Used for SOR and ignored for everything else
        self.source = np.zeros((N, N))
```

```python
@jit
def NAG_jacobi(x, solverinfo):
    """Performs one iteration of the jacobi method in the format required by NAG

    Arguments:
    x - N x N Matrix containing the previous iteration of the solution
    source - N x N Matrix containing our source function evaluated on the grid
    solverinfo - Takes a solverinfo object to get info in/out of the solver
    """
    N = int(np.sqrt(x.size))
    x.shape = (N, N) # Make x 2D because that's how I think
    nextx = np.zeros((N, N))
    h = 1/(N-1)
    # loop over the grid
    # Only iterate over interior points thus keeping the edges untouched and hence enforcing
    # the boundary condition that x = 0 at the edges.
    for i in range(1, N - 1):
        for j in range(1, N - 1):
            nextx[j, i] = 0.25 * (x[j+1, i] + x[j, i+1] + x[j-1, i] + x[j, i-1]) + 0.25*h**2*solverinfo.source[j, i]
    solverinfo.iterations += 1
    nextx -= x  # NAG requires this rather than nextx itself
    nextx.shape = N*N  # Make nextx 1D since that's what NAG needs
    return nextx
```

We can now call the NAG Anderson Acceleration routine.  If we set `m=0`, no acceleration is applied and the routine simply acts as a driver for our fixed point iterations.  We should get exactly the same behavior as our original code.  We can see that this is the case, right down to the 28,734 iterations required for the Jacobi method to solve.

This gives us confidence that our code changes haven't broken anything in the simulation.

```python
N = 100
x0 = init_problem(N)
x0.shape = N*N
NAG_jacobi_info = NAG_solverinfo()
nag_source = source(N)
NAG_jacobi_info.source = nag_source
tol = 1e-9
m = 0
NAG_jacobi_sol, fvec = roots.sys_func_aa(NAG_jacobi, x0, tol, eps, m, data=NAG_jacobi_info)
```

```python
print(f"NAG driven Jacobi on a {N} x {N} grid with no acceleration")
print(f"Solution found in {NAG_jacobi_info.iterations} iterations")
fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(10, 6))
NAG_jacobi_sol.shape = (N, N)
axes[0].imshow(NAG_jacobi_sol)
axes[0].set_title('Potential')
Ex, Ey = np.gradient(NAG_jacobi_sol)
E = np.sqrt(Ex**2+Ey**2) # Magnitude of Electric field
axes[1].imshow(E)
_ = axes[1].set_title('Electric Field')
```

Set `m=4` as suggested in the NAG documentation and convergence is much faster.  In fact, it's almost 3x faster than Gauss Seidel but we didn't need to change the mathematics in the Jacobi iteration at all.  As such, we can think of Anderson Acceleration as a Black-box accelerator for at least some types of fixed point iteration.

```python
N = 100
x0 = init_problem(N)
x0.shape = N*N
NAG_jacobi_info = NAG_solverinfo()
NAG_jacobi_info.source = source(N)
tol = 1e-9
m=4
NAG_jacobi_sol, fvec = roots.sys_func_aa(NAG_jacobi, x0, tol, eps, m, data=NAG_jacobi_info)
print(f"Anderson Accelerated Jacobi with m={m} on a {N} by {N} grid")
print(f"Solution found in {NAG_jacobi_info.iterations} iterations")
```

    Anderson Accelerated Jacobi with m=4 on a 100 by 100 grid
    Solution found in 5437 iterations

We will now investigate how the choice of `m` affects the number of iterations required for The Jacobi method.

```python
def find_best_m(m):
    N = 100
    x0 = init_problem(N)
    x0.shape = N*N
    NAG_jacobi_info = NAG_solverinfo()
    NAG_jacobi_info.source = source(N)
    tol = 1e-9
    _ = roots.sys_func_aa(NAG_jacobi, x0, tol, eps, m, data=NAG_jacobi_info)
    return NAG_jacobi_info.iterations
```

```python
mlist = np.arange(1, 200, 1)
iterations = [find_best_m(m) for m in mlist]
plt.plot(mlist, iterations)
plt.title('Anderson Accelerated Jacobi\nNumber of iterations vs m')
plt.xlabel('m')
_ = plt.ylabel('Iterations')
```

This is much nicer than the acceleration parameter for SOR.  We can just pick something high and be sure of good performance.  Let's go even higher than explored above and use `m=250` before moving on.

```python
N = 100
x0 = init_problem(N)
x0.shape = N*N
NAG_jacobi_info = NAG_solverinfo()
NAG_jacobi_info.source = source(N)
tol = 1e-9
m = 250
NAG_jacobi_sol, fvec = roots.sys_func_aa(NAG_jacobi, x0, tol, eps, m, data=NAG_jacobi_info)

print(f"Anderson Accelerated Jacobi with m={m} on a {N} by {N} grid")
print(f"Solution found in {NAG_jacobi_info.iterations} iterations")
fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(10, 6))
NAG_jacobi_sol.shape = (N, N)
axes[0].imshow(NAG_jacobi_sol)
axes[0].set_title('Potential')
Ex, Ey = np.gradient(NAG_jacobi_sol)
E = np.sqrt(Ex**2+Ey**2) # Magnitude of Electric field
axes[1].imshow(E)
_ = axes[1].set_title('Electric Field')
```

This is converging over 90x faster than the unaccelerated Jacobi method and even beats the optimal SOR method by a healthy margin.

### Anderson Acceleration and Gauss--Seidel

Next we combine Anderson Acceleration with the Gauss-Seidel method. Since Gauss-Seidel is significantly faster than Jacobi, we might hope that this new combination will be even better.

```python
@jit
def NAG_gauss_seidel(x, solverinfo):
    N = int(np.sqrt(x.size))
    x.shape = (N, N) # Make x 2D because that's how I think
    nextx = np.copy(x)
    h = 1/(N-1)
    for i in range(1, N - 1):
        for j in range(1, N - 1):
            nextx[j, i] = 0.25 * (nextx[j+1, i] + nextx[j, i+1] + nextx[j-1, i] + nextx[j, i-1]) + 0.25*h**2*solverinfo.source[j, i]
    solverinfo.iterations += 1
    nextx -= x  # NAG requires this rather than nextx itself
    nextx.shape = N*N  # Make nextx 1D since that's what NAG needs
    return nextx
```

```python
# Define and run the simulation using Gauss--Seidel and NAG Anderson Acceleration

N = 100
x0 = init_problem(N)
x0.shape = N*N
NAG_gs_info = NAG_solverinfo()
NAG_gs_info.source = source(N)
tol = 1e-9
m = 4
NAG_jacobi_sol, fvec = roots.sys_func_aa(NAG_gauss_seidel, x0, tol, eps, m, data=NAG_gs_info)
print(f"Anderson Accelerated Gauss-Seidel with m={m} on a {N} by {N} grid")
print(f"Solution found in {NAG_gs_info.iterations} iterations")
```

    Anderson Accelerated Gauss-Seidel with m=4 on a 100 by 100 grid
    Solution found in 577 iterations

With the standard choice of `m=4` we are doing almost as well as the best value of $\omega$ we could find for SOR.  Switching to `m=7` and we do even better than SOR.

```python
# Define and run the simulation using Gauss--Seidel and NAG Anderson Acceleration

N = 100
x0 = init_problem(N)
x0.shape = N*N
NAG_gs_info = NAG_solverinfo()
NAG_gs_info.source = source(N)
tol = 1e-9
m = 7
NAG_jacobi_sol, fvec = roots.sys_func_aa(NAG_gauss_seidel, x0, tol, eps, m, data=NAG_gs_info)
print(f"Anderson Accelerated Gauss--Seidel with m={m} on a {N} by {N} grid")
print(f"Solution found in {NAG_gs_info.iterations} iterations")
```

    Anderson Accelerated Gauss--Seidel with m=7 on a 100 by 100 grid
    Solution found in 445 iterations

Once again, we find ourselves needing to find the perfect acceleration parameter.  Let's see how the number of iterations varies with m.

```python
def find_best_m(m): # pylint: disable=function-redefined
    N = 100
    x0 = init_problem(N)
    x0.shape = N*N
    NAG_gs_info = NAG_solverinfo()
    NAG_gs_info.source = source(N)
    tol = 1e-9
    _ = roots.sys_func_aa(NAG_gauss_seidel, x0, tol, eps, m, data=NAG_gs_info)
    return NAG_gs_info.iterations
```

```python
mlist = np.arange(1, 200, 1)
iterations = [find_best_m(m) for m in mlist]
plt.plot(mlist, iterations)
plt.hlines(486, -1, 200, 'r')
plt.title('Anderson Accelerated Gauss--Seidel\nNumber of iterations vs m')
plt.xlabel('m')
plt.ylabel('Iterations')

label_x = 30
label_y = 1500
arrow_x = 20
arrow_y = 500

arrow_properties = dict(
    facecolor="black", width=0.5,
    headwidth=4, shrink=0.1)

_ = plt.annotate(
    "Best SOR result", xy=(arrow_x, arrow_y),
    xytext=(label_x, label_y),
    arrowprops=arrow_properties)
```

As with the accelerated Jacobi iterations, the situation with the Gauss--Seidel Anderson Acceleration parameter is rather better than the SOR parameter in that, for this simulation at least, we **always** do better than the best SOR result (486 iterations) for all `m>6` with the very best result being 285 iterations at `m=200`.  Furthermore, once you get past `m=10` for this simulation, the curve is pretty flat meaning that we are not being excessively punished for a bad choice of `m`.

```python
# Define and run the simulation using Gauss--Seidel and NAG Anderson Acceleration

N = 100
x0 = init_problem(N)
x0.shape = N*N
NAG_gs_info = NAG_solverinfo()
NAG_gs_info.source = source(N)
tol = 1e-9
m = 200
NAG_jacobi_sol, fvec = roots.sys_func_aa(NAG_gauss_seidel, x0, tol, eps, m, data=NAG_gs_info)
print(f"Anderson Accelerated Gauss--Seidel with m={m} on a {N} by {N} grid")
print(f"Solution found in {NAG_gs_info.iterations} iterations")
```

    Anderson Accelerated Gauss--Seidel with m=200 on a 100 by 100 grid
    Solution found in 285 iterations

## Combined SOR and Anderson Acceleration

The next natural question to ask is 'Could we combine both acceleration methods and do even better than anything seen so far'?

```python
@jit
def NAG_SOR(x, solverinfo):
    N = int(np.sqrt(x.size))
    x.shape = (N, N) # Make x 2D because that's how I think
    nextx = np.copy(x)
    h = 1/(N-1)
    w = solverinfo.w
    for i in range(1, N - 1):
        for j in range(1, N - 1):
            new = 0.25 * (nextx[j-1, i] + nextx[j+1, i] + nextx[j, i-1]+ nextx[j, i+1]) + 0.25*h**2*solverinfo.source[j, i]
            nextx[j, i] += w * (new - nextx[j, i])
    solverinfo.iterations += 1
    nextx -= x  # NAG requires this rather than nextx itself
    nextx.shape = N*N  # Make nextx 1D since that's what NAG needs
    return nextx
```

```python
# Define and run the simulation using SOR and NAG Anderson Acceleration
N = 100
x0 = init_problem(N)
x0.shape = N*N
NAG_SOR_info = NAG_solverinfo()
NAG_SOR_info.w = 1.94
NAG_SOR_info.source = source(N)
tol = 1e-9
m = 200
NAG_SOR_sol, fvec = roots.sys_func_aa(NAG_SOR, x0, tol, eps, m, data=NAG_SOR_info)
print(f"Anderson Accelerated SOR with m={m} and w={NAG_SOR_info.w} on a {N} by {N} grid")
print(f"Solution found in {NAG_SOR_info.iterations} iterations")
```

    Anderson Accelerated SOR with m=200 and w=1.94 on a 100 by 100 grid
    Solution found in 330 iterations

We see that using the best values of `m` and $\omega$  found so far gave us a result worse than using either of them independently.

Scanning over the range of `m` and $\omega$, we have the following results

```python
def find_best_mw(m, w):
    N = 100
    x0 = init_problem(N)
    x0.shape = N*N
    NAG_SOR_info = NAG_solverinfo()
    NAG_SOR_info.source = source(N)
    NAG_SOR_info.w = w
    tol = 1e-9
    _ = roots.sys_func_aa(NAG_SOR, x0, tol, eps, m, data=NAG_SOR_info)
    return NAG_SOR_info.iterations
```

```python
# This will take a LONG LONG time!
wlist = np.arange(1.01, 1.99, 0.01)
mlist = np.arange(5, 200, 1)
mw_list = [(m, w, find_best_mw(m, w)) for m in mlist for w in wlist]
```

```python
# Best 5
mw_list.sort(key=lambda x:x[2])
mw_list[0:10]
```

    [(199, 1.6700000000000006, 230),
     (198, 1.6600000000000006, 231),
     (198, 1.6700000000000006, 231),
     (198, 1.6800000000000006, 231),
     (198, 1.7300000000000006, 231),
     (199, 1.6400000000000006, 231),
     (199, 1.6500000000000006, 231),
     (199, 1.6600000000000006, 231),
     (199, 1.6800000000000006, 231),
     (199, 1.6900000000000006, 231)]

```python
# Worst 10
mw_list.sort(key=lambda x:x[2], reverse=True)
mw_list[0:10]
```

    [(5, 1.9800000000000009, 1543),
     (6, 1.9800000000000009, 1479),
     (8, 1.9800000000000009, 1442),
     (7, 1.9800000000000009, 1352),
     (9, 1.9800000000000009, 1326),
     (11, 1.9800000000000009, 1281),
     (10, 1.9800000000000009, 1192),
     (12, 1.9800000000000009, 1175),
     (15, 1.9800000000000009, 1122),
     (13, 1.9800000000000009, 1107)]

All of the best results have high values of the Anderson Acceleration factor, `m` so we could have restricted our search there to save time. For this simulation, more is better where `m` is concerned.  Finding the best value of $\omega$ for SOR, however is less obvious.

By using SOR and Anderson Acceleration we **can** do better than pure SOR or Anderson Accelerated Gauss--Seidel but it is more difficult to choose the best combination of $\omega$ and `m`.  Additionally, the returns are increasingly diminished.

# Conclusions

Anderson Acceleration is a general method for accelerating many fixed point iteration methods. We have shown that it can be used to reduce the number of iterations required in the traditional Jacobi, Gauss--Seidel and SOR methods for solving Laplace's and Poisson's equations in two dimensions.

Anderson Acceleration requires the user to select a parameter `m` which refers to the number of previous iterations used to produce the next one. For this simulation, high values of `m` give good results and getting the value of `m` wrong does not punish us too much.

| Method       | Unaccelerated | Accelerated |
|--------------|---------------|-------------|
| Jacobi       | 28734         | 313         |
| Gauss--Seidel| 14877         | 285         |
| SOR          | 486           | 230         |


Table of Contents

1  Logistic Regression

1.1  Logistic Function

1.2  Interpreting the Intercept

1.3  Defining The Cost Function

1.4  Gradient

1.5  Stochastic/Mini-batch Gradient

1.6  Implementation

1.7  Comparing Result and Convergence Behavior

1.8  Pros and Cons of Logistic Regression

2  Reference

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

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

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

```python
os.chdir(path)

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

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn import metrics
from sklearn.linear_model import LogisticRegression

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

    Ethen 2018-09-15 15:50:42

    CPython 3.6.4
    IPython 6.4.0

    numpy 1.14.1
    pandas 0.23.0
    matplotlib 2.2.2
    sklearn 0.19.1

# Logistic Regression

**Logistic regression** is an excellent tool to know for classification problems, which are problems where the output value that we wish to predict only takes on only a small number of discrete values. Here we'll focus on the binary classification problem, where the output can take on only two distinct classes. To make our examples more concrete, we will consider the Glass dataset.

```python
url = 'https://archive.ics.uci.edu/ml/machine-learning-databases/glass/glass.data'
col_names = ['id', 'ri', 'na', 'mg', 'al', 'si', 'k', 'ca', 'ba', 'fe', 'glass_type']
glass = pd.read_csv(url, names = col_names, index_col = 'id')
glass.sort_values('al', inplace = True)

# convert the glass type into binary outcome
# types 1, 2, 3 are window glass
# types 5, 6, 7 are household glass
glass['household'] = glass['glass_type'].map({1: 0, 2: 0, 3: 0, 5: 1, 6: 1, 7: 1})
glass.head()
```

      ri
      na
      mg
      al
      si
      k
      ca
      ba
      fe
      glass_type
      household

      id

      22
      1.51966
      14.77
      3.75
      0.29
      72.02
      0.03
      9.00
      0.0
      0.00
      1
      0

      185
      1.51115
      17.38
      0.00
      0.34
      75.41
      0.00
      6.65
      0.0
      0.00
      6
      1

      40
      1.52213
      14.21
      3.82
      0.47
      71.77
      0.11
      9.57
      0.0
      0.00
      1
      0

      39
      1.52213
      14.21
      3.82
      0.47
      71.77
      0.11
      9.57
      0.0
      0.00
      1
      0

      51
      1.52320
      13.72
      3.72
      0.51
      71.75
      0.09
      10.06
      0.0
      0.16
      1
      0

Our task is to predict the `household` column using the `al` column. Let's visualize the relationship between the input and output and also train the logsitic regression to see the outcome that it produces.

```python
logreg = LogisticRegression(C = 1e9)
X = glass['al'].reshape(-1, 1) # sklearn doesn't accept 1d-array, convert it to 2d
y = np.array(glass['household'])
logreg.fit(X, y)

# predict the probability that each observation belongs to class 1
# The first column indicates the predicted probability of class 0,
# and the second column indicates the predicted probability of class 1
glass['household_pred_prob'] = logreg.predict_proba(X)[:, 1]
```

```python
# plot the predicted probability (familiarize yourself with the S-shape)
# change default figure and font size
plt.rcParams['figure.figsize'] = 8, 6
plt.rcParams['font.size'] = 12

plt.scatter(glass['al'], glass['household'])
plt.plot(glass['al'], glass['household_pred_prob'])
plt.xlabel('al')
plt.ylabel('household')
plt.show()
```

As we can see, logistic regression can output the probabilities of observation belonging to a specific class and these probabilities can be converted into class predictions by choosing a cutoff value (e.g. probability higher than 0.5 is classified as class 1).

## Logistic Function

In **Logistic Regression**, the log-odds of a categorical response being "true" (1) is modeled as a linear combination of the features:

\begin{align*}
    \log \left({p\over 1-p}\right) &= w_0 + w_1x_1, ..., w_jx_j \nonumber \\
    &= w^Tx \nonumber
\end{align*}

Where:

- $w_{0}$ is the intercept term, and $w_1$ to $w_j$ represents the parameters for all the other features (a total of j features).
- By convention of we can assume that $x_0 = 1$, so that we can re-write the whole thing using the matrix notation $w^Tx$.

This is called the **logit function**. The equation can be re-arranged into the **logistic function**:

$$p = \frac{e^{w^Tx}} {1 + e^{w^Tx}}$$

Or in the more commonly seen form:

$$h_w(x) = \frac{1}{ 1 + e^{-w^Tx} }$$

Let's take a look at the plot of the function:

```python
x_values = np.linspace(-5, 5, 100)
y_values = [1 / (1 + np.exp(-x)) for x in x_values]
plt.plot(x_values, y_values)
plt.title('Logsitic Function')
plt.show()
```

The **logistic function** has some nice properties. The y-value represents the probability and it is always bounded between 0 and 1, which is want we wanted for probabilities. For an x value of 0 you get a 0.5 probability. Also as you get more positive x value you get a higher probability, on the other hand, a more negative x value results in a lower probability.

Toy sample code of how to predict the probability given the data and the weight is provided below.

```python
def predict_probability(data, weights):
    """probability predicted by the logistic regression"""
    score = np.dot(data, weights)
    predictions = 1 / (1 + np.exp(-score))
    return predictions
```

## Interpreting the Intercept

We can check logistic regression's coefficient does in fact generate the log-odds.

```python
# compute predicted log-odds for al = 2 using the equation
# convert log-odds to odds
# convert odds to probability
logodds = logreg.intercept_ + logreg.coef_[0] * 2
odds = np.exp(logodds)
prob = odds / (1 + odds)
print(prob)

logreg.predict_proba(2)[:, 1]
```

```python
# examine the coefficient for al
print('a1', logreg.coef_[0])
```

**Interpretation:** 1 unit increase in `al` is associated with a 4.18 unit increase in the log-odds of the observation being classified as `household 1`. We can confirm that again by doing the calculation ourselves.

```python
# increasing al by 1 (so that al now becomes 3)
# increases the log-odds by 4.18
logodds = logodds + logreg.coef_[0]
odds = np.exp(logodds)
prob = odds / (1 + odds)
print(prob)

logreg.predict_proba(3)[:, 1]
```

## Defining The Cost Function

When utilizing logistic regression, we are trying to learn the $w$ values in order to maximize the probability of correctly classifying our glasses. Let's say someone did give us some $w$ values of the logistic regression model, how would we determine if they were good values or not? What we would hope is that for the household of class 1, the probability values are close to 1 and for the household of class 0 the probability is close to 0.

But we don't care about getting the correct probability for just one observation, we want to correctly classify all our observations. If we assume our data are independent and identically distributed (think of it as all of them are treated equally), we can just take the product of all our individually calculated probabilities and that becomes the objective function we want to maximize. So in math:

$$\prod_{class1}h_w(x)\prod_{class0}1 - h_w(x)$$

The $\prod$ symbol means take the product of the $h_w(x)$ for the observations that are classified as that class. You will notice that for observations that are labeled as class 0, we are taking 1 minus the logistic function. That is because we are trying to find a value to maximize, and since observations that are labeled as class 0 should have a probability close to zero, 1 minus the probability should be close to 1. This procedure is also known as the **maximum likelihood estimation** and the following link contains a nice discussion of maximum likelihood using linear regression as an example. [Blog: The Principle of Maximum Likelihood](http://suriyadeepan.github.io/2017-01-22-mle-linear-regression/)

Next we will re-write the original cost function as:

$$\ell(w) = \sum_{i=1}^{N}y_{i}log(h_w(x_{i})) + (1-y_{i})log(1-h_w(x_{i}))$$

Where:

- We define $y_{i}$ to be 1 when the $i_{th}$ observation is labeled class 1 and 0 when labeled as class 0, then we only compute $h_w(x_{i})$ for observations that are labeled class 1 and $1 - h_w(x_{i})$ for observations that are labeled class 0, which is still the same idea as the original function.
- Next we'll transform the original $h_w(x_{i})$ by taking the log. As we'll later see this logarithm transformation will make our cost function more convenient to work with, and because the logarithm is a monotonically increasing function, the logarithm of a function achieves its maximum value at the same points as the function itself. When we take the log, our product across all data points, it becomes a sum. See [log rules](http://www.mathwords.com/l/logarithm_rules.htm) for more details (Hint: log(ab) = log(a) + log(b)).
- The $N$ simply represents the total number of the data.

Often times you'll also see the notation above be simplified in the form of a maximum likelihood estimator:

$$ \ell(w) = \sum_{i=1}^{N} log \big( P( y_i \mid x_i, w ) \big) $$

The equation above simply denotes the idea that , $\mathbf{w}$ represents the parameters we would like to estimate the parameters $w$ by maximizing conditional probability of $y_i$ given $x_i$.

Now by definition of probability in the logistic regression model: $h_w(x_{i}) = 1 \big/ 1 + e^{-w^T x_i}$ and $1- h_w(x_{i}) = e^{ -w^T x_i } \big/ ( 1 + e^{ -w^T x_i } )$. By substituting these expressions into our $\ell(w)$ equation and simplifying it further we can obtain  a simpler expression.

$$
\begin{align}
\ell(w)
&= \sum_{i=1}^{N}y_{i}log(h_w(x_{i})) + (1-y_{i})log(1-h_w(x_{i})) \nonumber \\
&= \sum_{i=1}^{N} y_{i} log( \frac{1}{ 1 + e^{ -w^T x_i } } ) + ( 1 - y_{i} )
log( \frac{ e^{ -w^T x_i } }{ 1 + e^{ -w^T x_i } } ) \nonumber \\
&= \sum_{i=1}^{N} -y_{i} log( 1 + e^{ -w^T x_i } ) + ( 1 - y_{i} )
( -w^T x_i - log( 1 + e^{ -w^T x_i } ) ) \nonumber \\
&= \sum_{i=1}^{N} ( y_{i} - 1 ) ( w^T x_i ) - log( 1 + e^{ -w^T x_i } ) \nonumber
\end{align}
$$

We'll use the formula above to compute the log likelihood for the entire dataset, which is used to assess the convergence of the algorithm. Toy code provided below.

```python
def compute_avg_log_likelihood(data, label, weights):
    """
    the function uses a simple check to prevent overflow problem,
    where numbers gets too large to represent and is converted to inf
    an example of overflow is provided below, when this problem occurs,
    simply use the original score (without taking the exponential)

    scores = np.array( [ -10000, 200, 300 ] )
    logexp = np.log( 1 + np.exp(-scores) )
    logexp
    """
    scores = np.dot(data, weights)
    logexp = np.log(1 + np.exp(-scores))

    # simple check to prevent overflow
    mask = np.isinf(logexp)
    logexp[mask] = -scores[mask]

    log_likelihood = np.sum((label - 1) * scores - logexp) / data.shape[0]
    return log_likelihood
```

**Note:** We made one tiny modification to the log likelihood function We added a ${1/N}$ term which averages the log likelihood across all data points. The ${1/N}$ term will make it easier for us to compare stochastic gradient ascent with batch gradient ascent later.

## Gradient

Now that we obtain the formula to assess our algorithm, we'll dive into the meat of the algorithm, which is to derive the gradient for the formula (the derivative of the formula with respect to each coefficient):

$$\ell(w) = \sum_{i=1}^{N} ( y_{i} - 1 ) ( w^T x_i ) - log( 1 + e^{ -w^T x_i } )$$

And it turns out the derivative of log likelihood with respect to to a single coefficient $w_j$ is as follows (the form is the same for all coefficients):

$$
\frac{\partial\ell(w)}{\partial w_j} = \sum_{i=1}^N (x_{ij})\left( y_i - \frac{1}{ 1 + e^{-w^Tx_i} } \right )
$$

To compute it, you simply need the following two terms:

- $\left( y_i - \frac{1}{ 1 + e^{-w^Tx_i} } \right )$ is the vector containing the difference between the predicted probability and the original label.
- $x_{ij}$ is the vector containing the $j_{th}$ feature's value.

For a step by step derivation, consider going through the following link. [Blog: Maximum likelihood and gradient descent demonstration](https://zlatankr.github.io/posts/2017/03/06/mle-gradient-descent), it uses a slightly different notation, but the walkthrough should still be pretty clear.

## Stochastic/Mini-batch Gradient

The problem with computing the gradient (or so called batched gradient) is the term $\sum_{i=1}^{N}$. This means that we must sum the contributions over all the data points to calculate the gradient, and this can be problematic if the dataset we're studying is extremely large. Thus, in stochastic gradient, we can use a single point as an approximation to the gradient:

$$
\frac{\partial\ell_i(w)}{\partial w_j} = (x_{ij})\left( y_i - \frac{1}{ 1 + e^{-w^Tx_i} } \right )
$$

**Note1:** Because the **Stochastic Gradient** algorithm uses each row of data in turn to update the gradient, if our data has some sort of implicit ordering, this will negatively affect the convergence of the algorithm. At an extreme, what if we had the data sorted so that all positive reviews came before negative reviews?  In that case, even if most reviews are negative, we might converge on an answer of +1 because we never get to see the other data. To avoid this, one practical trick is to shuffle the data before we begin so the rows are in random order.

**Note2:** Stochastic gradient compute the gradient using only 1 data point to update the the parameters, while batch gradient uses all $N$ data points. An alternative to these two extremes is a simple change that allows us to use a **mini-batch** of $B \leq N$ data points to calculate the gradient. This simple approach is faster than batch gradient but less noisy than stochastic gradient that uses only 1 data point. Given a mini-batch (or a set of data points) $\mathbf{x}_{i}, \mathbf{x}_{i+1} \ldots \mathbf{x}_{i+B}$, the gradient function for this mini-batch of data points is given by:

$$
\sum_{s = i}^{i+B} \frac{\partial\ell_s(w)}{\partial w_j} = \frac{1}{B} \sum_{s = i}^{i+B} (x_{sj})\left( y_i - \frac{1}{ 1 + e^{-w^Tx_i} } \right )
$$

Here, the $\frac{1}{B}$ means that we are normalizing the gradient update rule by the batch size $B$. In other words, we update the coefficients using the **average gradient over data points** (instead of using a pure summation). By using the average gradient, we ensure that the magnitude of the gradient is approximately the same for all batch sizes. This way, we can more easily compare various batch sizes and study the effect it has on the algorithm.

## Implementation

Recall our task is to find the optimal value for each individual weight to lower the cost. This requires taking the partial derivative of the cost/error function with respect to a single weight, and then running gradient descent for each individual weight to update them. Thus, for any individual weight $w_j$, we'll compute the following:

$$ w_j^{(t + 1)} = w_j^{(t)} + \alpha * \sum_{s = i}^{i+B} \frac{\partial\ell_s(w)}{\partial w_j}$$

Where:

- $\alpha$ denotes the the learning rate or so called step size, in other places you'll see it denoted as $\eta$.
- $w_j^{(t)}$ denotes the weight of the $j_{th}$ feature at iteration $t$.

And we'll do this iteratively for each weight, many times, until the whole network's cost function is minimized.

```python
# put the code together into one cell

def predict_probability(data, weights):
    """probability predicted by the logistic regression"""
    score = np.dot(data, weights)
    predictions = 1 / (1 + np.exp(-score))
    return predictions

def compute_avg_log_likelihood(data, label, weights):
    """
    the function uses a simple check to prevent overflow problem,
    where numbers gets too large to represent and is converted to inf
    an example of overflow is provided below, when this problem occurs,
    simply use the original score (without taking the exponential)

    scores = np.array([-10000, 200, 300])
    logexp = np.log(1 + np.exp(-scores))
    logexp
    """
    scores = np.dot(data, weights)
    logexp = np.log(1 + np.exp(-scores))

    # simple check to prevent overflow
    mask = np.isinf(logexp)
    logexp[mask] = -scores[mask]

    log_likelihood = np.sum((label - 1) * scores - logexp) / data.shape[0]
    return log_likelihood
```

```python
def logistic_regression(data, label, step_size, batch_size, max_iter):

    # weights of the model are initialized as zero
    data_num = data.shape[0]
    feature_num = data.shape[1]
    weights = np.zeros(data.shape[1])

    # `i` keeps track of the starting index of current batch
    # and shuffle the data before starting
    i = 0
    permutation = np.random.permutation(data_num)
    data, label = data[permutation], label[permutation]

    # do a linear scan over data, for each iteration update the weight using
    # batches of data, and store the log likelihood record to visualize convergence
    log_likelihood_record = []
    for _ in range(max_iter):

        # extract the batched data and label use it to compute
        # the predicted probability using the current weight and the errors
        batch = slice(i, i + batch_size)
        batch_data, batch_label = data[batch], label[batch]
        predictions = predict_probability(batch_data, weights)
        errors = batch_label - predictions

        # loop over each coefficient to compute the derivative and update the weight
        for j in range(feature_num):
            derivative = np.dot(errors, batch_data[:, j])
            weights[j] += step_size * derivative / batch_size

        # track whether log likelihood is increasing after
        # each weight update
        log_likelihood = compute_avg_log_likelihood(
            data = batch_data,
            label = batch_label,
            weights = weights
        )
        log_likelihood_record.append(log_likelihood)

        # update starting index of for the batches
        # and if we made a complete pass over data, shuffle again
        # and refresh the index that keeps track of the batch
        i += batch_size
        if i + batch_size > data_num:
            permutation = np.random.permutation(data_num)
            data, label = data[permutation], label[permutation]
            i = 0

    # We return the list of log likelihoods for plotting purposes.
    return weights, log_likelihood_record
```

## Comparing Result and Convergence Behavior

We'll use the logistic regression code that we've implemented and compare the predicted auc score with scikit-learn's implementation. This only serves to check that the predicted results are similar and that our toy code is correctly implemented. Then we'll also explore the convergence difference between batch gradient descent and stochastic gradient descent.

```python
# manually append the coefficient term,
# every good open-source library does not
# require this additional step from the user
data = np.c_[np.ones(X.shape[0]), X]

# using our logistic regression code
weights_batch, log_likelihood_batch = logistic_regression(
    data = data,
    label = np.array(y),
    step_size = 5e-1,
    batch_size = X.shape[0],  # batch gradient descent
    max_iter = 200
)
```

```python
# compare both logistic regression's auc score
logreg = LogisticRegression(C = 1e9)
logreg.fit(X, y)
pred_prob = logreg.predict_proba(X)[:, 1]

proba = predict_probability(data, weights_batch)

# check that the auc score is similar
auc1 = metrics.roc_auc_score(y, pred_prob)
auc2 = metrics.roc_auc_score(y, proba)
print('auc', auc1, auc2)
```

```python
weights_sgd, log_likelihood_sgd = logistic_regression(
    data = data,
    label = y,
    step_size = 5e-1,
    batch_size = 30,  # stochastic gradient descent
    max_iter = 200
)
```

```python
weights_minibatch, log_likelihood_minibatch = logistic_regression(
    data = data,
    label = y,
    step_size = 5e-1,
    batch_size = 100,  # mini-batch gradient descent
    max_iter = 200
)
```

```python
plt.figure(figsize = (10, 7))
plt.plot(log_likelihood_sgd, label = 'stochastic gradient descent')
plt.plot(log_likelihood_batch, label = 'batch gradient descent')
plt.plot(log_likelihood_minibatch, label = 'mini-batch gradient descent')
plt.legend(loc = 'best')
plt.xlabel('# of iterations')
plt.ylabel('Average log likelihood')
plt.title('Convergence Plot')
plt.show()
```

Based on the convergence plot above, we can see that the it's a good idea to use mini-batch gradient descent since it strikes a good balance between batch gradient, which convergences steadily but can be computationly too expensive when the dataset is too large, and stochastic gradient, which is faster to train, but the result can be too noisy.

## Pros and Cons of Logistic Regression

We'll end this notebook listing out some pros and cons of this method.

**Pros:**

- Highly interpretable (if you remember how).
- Model training and prediction are fast. Thus can be desirable in large-scale applications when we're dealing with millions of parameters.
- Almost no parameter tuning is required (excluding regularization).
- Outputs well-calibrated predicted probabilities.

**Cons:**

- Presumes a linear relationship between the features
- Performance is (generally) not competitive with the best supervised learning methods.
- Can't automatically learn feature interactions.

# Reference

- [Notebook: Logistic Regression](http://nbviewer.jupyter.org/github/justmarkham/DAT8/blob/master/notebooks/12_logistic_regression.ipynb)
- [Coursersa: Washington Classification](https://www.coursera.org/learn/ml-classification)
- [Blog: Maximum likelihood and gradient descent demonstration](https://zlatankr.github.io/posts/2017/03/06/mle-gradient-descent)


# Phase Shifts and Virtual Z gates

```python
import numpy as np
from pulser import Pulse, Sequence, Register
from pulser.devices import MockDevice
```

## Introduction

Under the right circumstances, phase shifts can be a great way of implementing so-called *virtual-Z gates*. Let's see how these arise and how we can use them to our advantage.
Consider an arbitrary 2x2 unitary matrix (up to a global phase) and a possible decompostion in terms of rotations around the $x$ and $z$ axes ($R_X$ and $R_Z$) of the Bloch sphere:

$$U(\gamma, \theta, \phi) =  R_Z(\gamma)R_X(\theta)R_Z(\phi)$$

Our goal is to be able to apply this transformation on our qubit states through our pulses. A pulse that is on-resonance with its target transition (i.e. has detuning $\delta = 0$) can be described as an arbitrary rotation around an axis contained in the Bloch sphere's equator, $\hat{n}= (\cos \phi, -\sin \phi, 0)$. From the general form of such a rotation, $R_{\hat{n}}(\theta) =     \exp\left(-i \frac{\theta}{2}  \hat{n}\cdot \vec{\sigma}\right)$, we arrive at:

$$
\begin{align}
 R_{\hat{n}(\phi)}(\theta) &= \exp\left(-i \frac{\theta}{2} \left[\cos(\phi)\sigma_x -\sin(\phi)\sigma_y\right]\right)\\
          &= e^{i\frac{\phi}{2}\sigma_z}e^{-i\frac{\theta}{2}\sigma_x}e^{-i\frac{\phi}{2}\sigma_z}\\
          &=  R_Z(-\phi)R_X(\theta)R_Z(\phi)
\end{align}
$$

Here, we have two free parameters: the angle of rotation $\theta$, which is determined by the integral of the amplitude waveform, and $\phi$, the pulse's phase. Thus, we can see that a pulse is a particular case of the arbitrary single-qubit gate $U$, where $\gamma = -\phi$, i.e.:

$$ R_{\hat{n}(\phi)}(\theta) = U(-\phi, \theta, \phi) $$

Thus, to reach the desired arbitrary single-qubit, we need an extra $R_Z$ gate, such that:

$$ U(\gamma, \theta, \phi) = R_Z(\gamma + \phi) R_Z(-\phi)R_X(\theta)R_Z(\phi) = R_Z(\gamma + \phi) U(-\phi, \theta, \phi)$$

Now, how do we implement such a gate? In fact, to physically change the phase of the qubit's state in this reference frame with a single pulse, we would have to be able to apply a detuned pulse of zero amplitude, which is no pulse at all! Instead, what we can do is change the frame of rotation such that the phase gate is applied *virtually*.

To understand how this can be done, we first have to realise that this last phase gate is irrelevant if it is the last one applied to the qubit before it is measured - a change of the phase between the $\sigma_z$ eigenstates will produce no change in the measurement outcome, since that occurs in the same basis. But what if it is not the last gate that we apply on this qubit? In that case, we can describe the situation as a new arbitrary gate being applied after the existent one, i.e.

$$ U(\alpha, \beta, \nu)~~U(\gamma, \theta, \phi) = R_Z(\alpha + \nu) R_Z(-\nu)R_X(\beta)R_Z(\nu) ~~ R_Z(k)~
U(-\phi, \theta, \phi) $$

where we define the *carry*, $k=\gamma + \phi$, as the phase of the unrealized phase gate. Now, we can restructure the previous expression such that:

$$
\begin{align}
& R_Z(\alpha + \nu) R_Z(-\nu)R_X(\beta)R_Z(\nu)~~ R_Z(k) ~U(-\phi, \theta, \phi) = \\
&= R_Z(\alpha + \nu + k)~~ R_Z(-\nu - k)R_X(\beta)R_Z(\nu + k) ~~U(-\phi, \theta, \phi) \\
&= R_Z(k') ~~U(-\nu - k, \beta, \nu + k)~~U(-\phi, \theta, \phi), ~~~k' = \alpha + \nu + k
\end{align}
$$

As shown, the previously existent phase gate of angle $k$ can be realized as a **shift** on the phase of the second pulse, $\nu \rightarrow \nu + k$. In this way, we go back to a situation where we have a phase gate at the end (with an updated carry $k'$). We can repeat this process until the moment we measure the qubit, at which point, as we've seen, the remaining phase gate is redundant.

This is the **virtual-Z gate**: the ability to perform phase gates through the adjustment of a qubit's phase reference.

## Phase shifts in Pulser

As shown above, implementing virtual-Z gates requires tracking the *carry* and adapting the phase of each provided pulse accordingly. Although this could certainly be done externally, Pulser offers a convenient way to automatically update a qubit's phase reference (i.e. the *carry*) through **phase shifts**.

A phase shift in Pulser is defined by three things: the *value* of the shift, the *target* qubits and the *basis*. A phase shift of *value* $\phi$ corresponds to a change in the *target*'s phase reference from $k \rightarrow k + \phi$.

It is important to realise that, unlike all other aspects so far, phase shifts are associated to the a qubit's transition, not a channel. Therefore, in principle, all qubit's can keep different phase references, as a result of being subjected to different phase shifts throughout a sequence.

### Programming a Hadamard gate

To exemplify the need for phase shifts, let's try to encode an Hadamard gate using only resonant pulses.
In our decomposition of a unitary matrix, the Hadamard is (you can check this for yourself):

$$ H = U\left(\frac{\pi}{2},\frac{\pi}{2}, \frac{\pi}{2}\right) = R_Z(\pi)~~U\left(-\frac{\pi}{2},\frac{\pi}{2}, \frac{\pi}{2}\right) $$

meaning that we have to apply a $\frac{\pi}{2}$-pulse with phase $\phi=\frac{\pi}{2}$, followed by a phase shift of $\pi$.

But first, let's create a simple setup.

```python
reg = Register({'q0': (0, 0)})
device = MockDevice
seq = Sequence(reg, device)
seq.available_channels
```

```python
seq.declare_channel('ch0', 'raman_local', initial_target = 'q0')
```

```python
# Defining the waveform for a pi/2 pulse
from pulser.waveforms import BlackmanWaveform
pi2_wf = BlackmanWaveform(1000, np.pi/2)   # Duration: 1us, Area: pi/2
pi2_wf.draw()
```

```python
# 2. Create the pi/2 pulse
pi_2 = Pulse.ConstantDetuning(pi2_wf, detuning=0, phase=np.pi/2)
pi_2.draw()
```

```python
#3. Applying the H gate

seq.add(pi_2, 'ch0')   # The first pi/2-pulse
# Now the phase shift of pi on 'q0', for the 'digital' basis, which is usually where phase shifts are useful
seq.phase_shift(np.pi, 'q0', basis='digital')

seq.draw(draw_phase_shifts=True)
```

This produced the desired effect: we have a $\frac{\pi}{2}$ pulse with $\phi=\frac{\pi}{2}$ followed by a phase shift of $\pi$. However, the need to specify the target qubit and basis is mildly incovenient. Moreover, it would be even better if the entire hadamard could be in the same pulse object. Fortunately, there's a way.

The `Pulse` object has an optional argument called `post_phase_shift`, with which the user can order a phase shift to be applied immediatly after the physical pulse. In this case, the target and basis are **implicitly defined** to be the channel's current target and basis.

Here's how we could define the Hadamard in a single pulse:

```python
h = Pulse.ConstantDetuning(pi2_wf, detuning=0, phase=np.pi/2, post_phase_shift=np.pi)

seq.add(h, 'ch0')
seq.draw(draw_phase_shifts=True)
```

Notice how the two pulse shapes are naturally identical, and are both followed by the adequate $\pi$ phase shift. However, we expect to see an adjustment on the phase of the second pulse, simply because it is following a phase shift. To inspect the phase of each pulse, we can print the sequence:

```python
print(seq)
```

Here, we see that the phase of the second pulse has the appropriate adjustment of $\pi$. What happens if we apply the Hadamard yet again?

```python
seq.add(h, 'ch0')
print(seq)
```

That's right, the phase of the third pulse is back to $\frac{\pi}{2}$ because, at the moment the third Hadamard was constructed, there had been two $\pi$ phase shifts, which means the reference is back to $\phi=0$.

### Phase shifts with multiple channels and different targets

Now, let's construct a more complex example where there are multiple channels and multiple qubits.

```python
reg = Register({'q0': (0, 0), 'q1': (5, 5)})
device = MockDevice
seq = Sequence(reg, device)
seq.declare_channel('raman', 'raman_local', initial_target = 'q0')
seq.declare_channel('ryd1', 'rydberg_local', initial_target = 'q0')
seq.declare_channel('ryd2', 'rydberg_local', initial_target = 'q0')
seq.declared_channels
```

We see that we have two qubits and three channels, all `Local`, with `raman` acting on the `digital` basis and the other two on the `ground-rydberg` basis. Let's use the Hadamard from before and add it to channels `raman` and `ryd1`, which are both targeting `q0` on different bases:

```python
seq.add(h, 'raman')
seq.add(h, 'ryd1')
seq.draw(draw_phase_shifts=True)
```

We can see that the pulse in channel `ryd1` waited for the pulse in `raman` to be finished (because they are acting on the same target). We also noticed that the phase shift in `raman` does not appear in the other channels, the reason for that being the fact that they act on **different bases**. We can check this by printing the sequence:

```python
print(seq)
```

Here, we confirm that the phase of the pulse in `ryd1` is $\frac{\pi}{2}$, which indicates a phase reference of $\phi=0$ as expected. What about when the phase shift targets the same qubit, the same basis, but the pulses are on different channels? In that case, we expect that the channel is irrelevant, and we can already see evidence of that in how the phase shift at the end of the Hadamard on `ryd1` also appears in `ryd2`. We can confirm this by adding another pulse to `ryd2` (e.g. the `pi_2` pulse we defined before) and then priting the sequence:

```python
seq.add(pi_2, 'ryd2')
print(seq)
```

Notice how the `pi_2` pulse has a phase of $\frac{3\pi}{2}$: it's phase of $\frac{\pi}{2}$ plus the shift of $\pi$ accrued by the Hadamard in channel `ryd1`.

Let's now shift our attention torwards what happens when the basis and the channel stay the same, but the target qubit changes. By now, you can already predict that qubit `q1` has a phase reference of $\phi=0$ on both basis, since all phase shifts so far were always targeting `q0`. We can see this if we change the target on some channels and apply pulses again:

```python
seq.target('q1', 'raman')
seq.add(h, 'raman')
seq.target('q1', 'ryd1')
seq.add(h, 'ryd1')
print(seq)
seq.draw(draw_phase_shifts=True)
```

We can see how the second pulse on the `raman` and `ryd1` channels are exactly identical to their first one, except for the target. Had the target not changed, they would have their phase shifted by $\pi$, just like we saw earlier. We also see how, this time, the phase shift in `ryd1` does not appear in `ryd2` like before because they have **different targets**. Notice what happens if we also make `ryd2` target `q1`:

```python
seq.target('q1', 'ryd2')
seq.draw(draw_phase_shifts=True)
```

As you can see, channel `ryd2` starts off with a phase of $\phi = \pi$, which it picked up from the phase shift induced by the Hadamard on `ryd1`. On the other hand, the Hadamard in the `raman` channel did not affect the phase of `ryd2`(we would have $\phi=0$ in that case) because, again, it acts on a different basis.

tutorials/composition/Phase Shifts and Virtual Z gates.ipynb


```python
from sympy import symbols, Eq, solve
import numpy as np
class Beta:

    def __init__(self, data):
        self.x = data
        self.calc_mean()
        self.calc_std()
        self._calc_paramters()

    def calc_mean(self):
        self.mean = np.mean(self.x)

    def calc_std(self):
        self.std = np.std(self.x)

    def _calc_paramters(self):
        self.alpha, self.beta = symbols("x y")
        eq1 = Eq(self.alpha/(self.alpha + self.beta), self.mean)
        eq2 = Eq((self.alpha * self.beta)/ (((self.alpha + self.beta)**2)*(self.alpha + self.beta + 1)), self.std**2)
        parameters = solve((eq1, eq2), (self.alpha, self.beta))
        self.alpha = float(parameters[self.alpha])
        self.beta = float(parameters[self.beta])

    def get_beta(self):
        return self.beta

    def get_alpha(self):
        return self.alpha

```


# License

https://raw.githubusercontent.com/computational-sediment-hyd/tridiagonal-algorithm-matrix-of-periodic-boundary-conditions/master/LICENSE

# tridiagonal algorithm matrix of periodic boundary conditions

## はじめに

- 周期境界の三重対角は代表的な問題．
- その解法は，Sherman-Morrison formulaを使う．
- 参考(一部間違いあり)
:https://www.cfd-online.com/Wiki/Tridiagonal_matrix_algorithm_-_TDMA_%28Thomas_algorithm%29

## Sherman-Morrison formulaとは
参考：https://mathtrain.jp/woodbury

- 逆行列の補助定理（Woodburyの恒等式）の特別な場合を示す．

逆行列の補助定理とは下式である．

$$
(A+BDC)^{-1}=A^{-1}-A^{-1}B(D^{-1}+CA^{-1}B)^{-1}CA^{-1}
$$

証明は省略する．

$
A:n\times n \,,\,
B:n\times k \,,\,
C:k\times n \,,\,
D:k\times k
$
が条件．

$k=1\,,\, D=1$のときは，

$$
(A+BC)^{-1}=A^{-1}-A^{-1}B(1+CA^{-1}B)^{-1}CA^{-1}
$$

であり，本式がシャーマンモリソンの公式と呼ばれる．

この式の良いところは，$A$の逆行列が計算できると$A+BC$の逆行列が計算できる点です．
今回の問題を例にとると，$A$にTDMAができれば，$A+BC$の逆行列が簡単に計算できるということ．

## 周期境界三重対角への適用

実際に計算してみます．

$$
A'x=d
$$

$$
A' =
\left(
\begin{array}{cccc}
b_1 & c_1 & 0 & a_1 \\
a_2 & b_2 & c_2 & 0 \\
0 & a_3 & b_3 & c_3 \\
c_4 & 0 & a_4 & b_4 \\
\end{array}
\right)
$$
とすると，

$$
A' = A + B \times C
$$

$$
A =
\left(
\begin{array}{cccc}
2 b_1 & c_1 & 0 & 0 \\
a_2 & b_2 & c_2 & 0 \\
0 & a_3 & b_3 & c_3 \\
0 & 0 & a_4 & b_4 + \dfrac{a_1 c_4}{b_1} \\
\end{array}
\right)
$$

$$
B =
\left(
\begin{array}{c}
-b_1  \\
0  \\
0  \\
c_4  \\
\end{array}
\right)
$$

$$
C =
\left(
\begin{array}{cccc}
1 & 0 & 0 & -\dfrac{a_1}{b_1}
\end{array}
\right)
$$

となるので，$A$にTDMAが適用できます

あとは，公式にあてはめて展開する.

$$
A y = d \,,\, A z = B \\
y = A^{-1}  d \,,\, z = A^{-1}  B
$$
とおくと

$$
x =  y - \dfrac{C \cdot y}{1 + C \cdot z}z
$$

となり，右辺第二項は$z$のスカラー倍のため，2回のTDMAで容易に計算可能である．

## 補足：TDMA（導出は省略）
 - 一般的にはメモリ節約版を使いますが（パタンカーの影響？），分かりにくいので通常版を使います．

三重対角行列
$$
a_i x_{i-1} + b_i x_i + c_i x_{i+1} = d_i
$$
について，

$$
\begin{align}
P_i &= -\frac{c_i}{a_i P_{i-1} + b_i} \\
Q_i &= \frac{-a_iQ_{i-1}+d_i}{a_i P_{i-1} + b_i} \\
P_1 &= -\frac{c_1}{b_1} \\
Q_1 &= \frac{d_1}{b_1}
\end{align}
$$

とすると，

$$
\begin{align}
x_n &= Q_n \\
x_i &= P_ix_{i+1} + Q_i
\end{align}
$$

## 実装

### Periodic boundary TDMA source code

```python
import numpy as np

def pTDMA(a,b,c,d):

    def TDMA(a,b,c,d):
        x, P, Q = np.empty_like(a), np.empty_like(a), np.empty_like(a)
        n = len(x)

        P[0], Q[0] = -c[0]/b[0], d[0]/b[0]
        for i in range(1, n):
            P[i] = -  c[i]                /(a[i]*P[i-1] + b[i])
            Q[i] =  (-a[i]*Q[i-1] + d[i])/(a[i]*P[i-1] + b[i])

        x[-1] = Q[-1]
        for i in range(n-2, -1, -1):
            x[i] = P[i]*x[i+1] + Q[i]

        return x

    a, b, c, d = np.array(a,dtype=float), np.array(b,dtype=float), np.array(c,dtype=float), np.array(d,dtype=float)
    a1, b1, c1, d1 = np.copy(a), np.copy(b), np.copy(c), np.copy(d) #deep copy
    a1[0], c1[-1] = 0.0, 0.0
    b1[0] = 2.0*b[0]
    b1[-1] = b[-1] + a[0]*c[-1]/b[0]
    d1[:] = 0.0
    d1[0] = -b[0]
    d1[-1] = c[-1]

    y = TDMA(a1,b1,c1,d)
    z = TDMA(a1,b1,c1,d1)

    w = ( y[0] - a[0]/b[0]*y[-1] ) / ( 1 + (z[0]- a[0]/b[0]*z[-1]) )
    x = y - w * z

    return x
```

### test

```python
n = 10
a = np.ones(n)
b = np.ones(n)
c = np.ones(n)
d = np.arange(1,n+1,1,dtype=float)

a *= -0.2
c *= 0.2

x = pTDMA(a,b,c,d)
print(x)
```

    [ 2.81456954  2.02649007  2.68211921  3.61589404  4.60264901  5.60264901
      6.58940397  7.65562914  8.31125828 11.09933775]

 - check print d

```python
print( b[0]*x[0] + c[0]*x[1] + a[0]*x[-1] )

for i in range(1, n-1):
    print( a[i]*x[i-1] + b[i]*x[i] + c[i]*x[i+1] )

print( a[-1]*x[-2] + b[-1]*x[-1] + c[-1]*x[0]  )
```

    1.0000000000000004
    2.0
    3.0
    3.9999999999999996
    5.0
    6.0
    7.0
    8.0
    9.0
    9.999999999999998

tridiagonal algorithm matrix of periodic boundary conditions.ipynb


```python
import time
import numpy as np
from scipy import sparse
from scipy import linalg as la
from scipy.integrate import solve_ivp
from matplotlib import animation, pyplot as plt

from IPython.display import Video
```

```python
π = np.pi
π2 = π**2
EMBED = False

plt.rcParams["figure.dpi"] = 300
plt.rcParams["figure.figsize"] = (10,2)
plt.rcParams["text.usetex"] = True
```

Model Reduction Example: Heat Equation

Shane McQuarrie

## Problem Setup

Consider the following initial/boundary value problem (IBVP) modeling diffusion of heat through a one-dimensional rod $\Omega = [0,1]$ with insulation at both ends:

$$
\begin{align*}
    x_t(t,\omega) &= \frac{1}{\pi^2}x_{\omega\omega}(t,\omega), & (t,\omega)&\in (0,T]\times\Omega,\\
    x_\omega(t,0) = x_\omega(t,1) &= 0, & t &> 0,\\
    x(0,\omega) = x_0(\omega) &= \cos(\pi \omega) + \varepsilon\cos(4\pi \omega) + 2, & \omega &\in\Omega.
\end{align*}
$$

The analytical solution to this equation is
$
x(t,\omega) = e^{-t}\cos(\pi \omega) + \varepsilon e^{-16t}\cos(4\pi \omega) + 2.
$

To solve the problem numerically, let $x_i:[0,T]\to\mathbb{R}$ denote the numerical approximation for $x(\cdot,\omega_i)$ for $i=1,\ldots,n$.
Discretizing the partial derivatives by second-order difference quotients (finite difference method) on a uniform grid, we arrive at the system

$$
\begin{align*}
    \dot{\mathbf{x}}(t) &= A\mathbf{x}(t),
    &
    \mathbf{x}(0) &= \mathbf{x}_0,
\end{align*}
$$

where $\mathbf{x}(t) = \begin{bmatrix}x_1(t) & \cdots & x_n(t)\end{bmatrix}^{\mathsf{T}}\in\mathbb{R}^n,\quad\mathbf{x}_0 = \begin{bmatrix}x_0(\omega_1) & \cdots & x_0(\omega_n)\end{bmatrix}^{\mathsf{T}}\in\mathbb{R}^n$, and

$$
A
= \frac{1}{\pi^2\Delta \omega^2}
\left[\begin{array}{ccccc}
-2 &      2 &        &        & 0  \\
 1 &     -2 &      1 &        &    \\
   & \ddots & \ddots & \ddots &    \\
   &        &      1 &     -2 &  1 \\
 0 &        &        &      2 & -2 \\
\end{array}\right] \in\mathbb{R}^{n\times n}.
$$

We can solve this with SciPy to get snapshots of the solution, from which we will construct the reduced order model.

```python
n = 500                         # Number of spatial points in the discretization.
k = 200                         # Number of points in time stepping.
L = 1                           # Length of spatial domain.
T = 4                           # Final simulation time.

# Define the spatial and temporal domains and true solution.
w = np.linspace(0, 1, n)
t = np.linspace(0, T, k)
δw = w[1] - w[0]
δt = t[1] - t[0]

ε = 1
analytic_solution = lambda t, w: np.exp(-t)*np.cos(π*w) + ε*np.exp(-16*t)*np.cos(4*π*w) + 2

Xtrue = analytic_solution(t, w[:,np.newaxis])
x0 = analytic_solution(0, w)
```

```python
# Construct the full-order model.
sub = np.concatenate((np.ones(n - 2), [2]))
A = sparse.diags([sub, -2, sub[::-1]],
                 offsets=[-1,0,1],
                 shape=(n,n)) / (δw**2 * π**2)

f = lambda t,x: A @ x
```

```python
%%time
Xcomp = solve_ivp(f,                        # Solve dx / dt = Ax(t)
                  [0,T],                    # over t = [0,T]
                  x0,                       # with initial condition x0
                  t_eval=t,                 # evaluated at the specified times
                  method="BDF").y           # using an implicit scheme.
```

    CPU times: user 243 ms, sys: 21.1 ms, total: 264 ms
    Wall time: 104 ms

## Constructing the Reduced-order Model

Next, we take a sampling of the computed solution `Xcomp` to use as the snapshot data $X$.
This is an easy problem, so we can get away with taking only every $10$th time step (including the initial condition) over the first half of the time domain.

```python
# Select snapshots from the computed trajectories: every 10th step for half of the time.
X = Xcomp[:,:Xcomp.shape[1]//2:10]
X.shape
```

    (500, 10)

Before computing a POD basis of the snapshot matrix $X$, we must choose $r$, the number of basis vectors to use.
The singular values inform this decision.

```python
snapshot_svdvals = la.svdvals(X)
plt.semilogy(np.arange(1,X.shape[1]+1), snapshot_svdvals, '.')
plt.xlabel(r"$j$")
plt.ylabel(r"$\sigma_{j}$")
plt.show()
```

For this example, it is clear that $r = 3$ components capture most of the information.
We therefore compute the POD basis of order $3$, which consists of the first $3$ left singular vectors of $X$.

```python
def POD_basis(X, r):
    """Compute the POD basis of order r corresponding to the data in X.

    Parameters
    ----------
    X : (n,k) ndarray
        Snapshot matrix, usually with n >> k.

    r : int
        Number of POD basis vectors to compute.

    Returns
    -------
    Vr : (n,r) ndarray
        The first r POD basis vectors of X.
    """
    Phi = la.svd(X, full_matrices=False)[0]
    return Phi[:,:r]
```

```python
# Compute the POD basis of rank r corresponding to the snapshots.
r = 3
Vr = POD_basis(X, r)
Vr.shape
```

    (500, 3)

To see if we're on the right track, we visualize each of the selected POD basis vectors, which live in the same space as the snapshots $\left(\mathbb{R}^n\right)$.
Note that the first three vectors capture the characteristics of the solution fairly well.

```python
rows = (r//3) + 1 if not r//3 else r//3
fig, axes = plt.subplots(rows, 3, sharex=True, sharey=True, figsize=(10,2*rows))
for j,ax in zip(range(r), axes.flatten()):
    ax.plot(w, Vr[:,j], f"C{j}", lw=1)
    ax.axhline(0, c='k', ls='--', lw=.5)
    ax.set_title(fr"Basis vector {j+1}: $\sigma_j\approx {snapshot_svdvals[j]:.4f}$", fontsize=8)
plt.show()
```

Now that we have $V_r$, we project the linear operator $A$ to the reduced operator $\hat{A} = V_r^\mathsf{T} A V_r$ and compute the reduced initial condition $\hat{\mathbf{x}}_0 = V_r^\mathsf{T}\mathbf{x}_0$.
This projection is called the _offline phase_, since it can be done a priori.

```python
%%time
# Offline phase: construct the reduced-order model.
x0_ = Vr.T @ x0              # Initial condition
A_ = Vr.T @ A @ Vr           # State matrix

f_ = lambda t,x_: A_ @ x_

print("Original system size:", A.shape)
print("Reduced system size:", A_.shape)
```

    Original system size: (500, 500)
    Reduced system size: (3, 3)
    CPU times: user 1.61 ms, sys: 1.82 ms, total: 3.43 ms
    Wall time: 4.06 ms

The _online phase_ is the actual simulation of the $r$-dimensional ROM.
Note that it is (obviously) much faster than the high-dimensional simulation.

```python
%%time
# Online phase: solve the reduced system.
X_rom_ = solve_ivp(f_,
                  [0,4],
                  x0_,
                  t_eval=t,
                  method="BDF").y
```

    CPU times: user 13.6 ms, sys: 1.4 ms, total: 15 ms
    Wall time: 13.9 ms

Finally, project the reduced solution in $\mathbb{R}^r$ back to the original space $\mathbb{R}^n$ as $\mathbf{x}_\text{ROM} = V_r \hat{\mathbf{x}}$.

```python
# Project the solution back up to the original space.
X_ROM = Vr @ X_rom_
assert Xcomp.shape == X_ROM.shape
```

## Results and Error Analysis

We'll do three things to visualize the results:
- Plot a few snapshots and the corresponding ROM outputs
- Compute relative error measures between the snapshots and the ROM outputs
- Animate both the snapshot set and the ROM outputs

```python
def plot_heat_data(Z, title, ax=None):
    if ax is None:
        fig, ax = plt.subplots(1, 1, figsize=(12,4))

    # Plot a few snapshots over the spatial domain.
    sample_columns = [0, 10, 20, 40, 80, 160]
    color = iter(plt.cm.viridis_r(np.linspace(0, .8, len(sample_columns))))
    for j in sample_columns:
        ax.plot(w, Z[:,j], color=next(color), label=rf"$x(\omega,t_{{{j}}})$")
    ax.set_xlim(0, 1)
    ax.set_xlabel(r"$\omega$")
    ax.set_ylabel(r"$x(\omega,t)$")
    ax.legend(loc="upper right")
    ax.set_title(title)
```

```python
fig, [ax1, ax2] = plt.subplots(1, 2, figsize=(12,4))
plot_heat_data(Xcomp, "Snapshots", ax1)
plot_heat_data(X_ROM, "ROM Output", ax2)
```

The relative $L^2$ error of the reduced order approximation $\mathbf{x}_\text{ROM}$ is given by

$$
\text{Err} %_\text{rel}
= \frac{\|\mathbf{x}(\cdot) - \mathbf{x}_\text{ROM}(\cdot)\|_{L^2([0,T])}}{\|\mathbf{x}(\cdot)\|_{L^2([0,T])}}.
$$

We approximate the $L^2([0,T])$ norms by numerical integration, for example with the Trapezoidal Rule:

$$
\begin{align}
    \|\mathbf{x}(\cdot)\|_{L^2([0,T])}^2
    &= \int_0^1 \|\mathbf{x}(t)\|_2^2\:dt\\
    &\approx \delta t\left(\frac{1}{2}\|\mathbf{x}(t_0)\|_2^2 + \sum_{j=1}^{k-1}\|\mathbf{x}(t_j)\|_2^2 + \frac{1}{2}\|\mathbf{x}(t_k)\|_2^2\right).
\end{align}
$$

```python
def relative_L2_error(x, y, t):
    """Compute the relative L2-norm error between x and y.

    Parameters
    ----------
        x : (n,k) ndarray
            True data.

        y : (n,k) ndarray
            Approximation to x.

        t : (k,) ndarray
            Time domain of data.
    """
    numerator = np.sqrt(np.trapz(np.sum((x - y)**2, axis=0), t))
    denominator = np.sqrt(np.trapz(np.sum(x**2, axis=0), t))
    return numerator / denominator

err = relative_L2_error(Xcomp, X_ROM, t)
print(f"Relative Error: {err*100:.6f}%")
```

    Relative Error: 0.026120%

Finally, we animate the solutions in the same plot and look for differences.

```python
def animate_1D_funcs(save_as, w, t, state_list, labels=None):
    """Animate one or more one-dimensional functions on R over a time domain.

    Parameters
    ----------
    save_as : str
        File to save the video as (e.g. "heat.mp4").

    w : (n,) ndarray
        Spatial domain of the function(s).

    t : (K,) ndarray
        Time domain of the function(s).

    state_list : list((n,K) ndarray)
        Computed function values. Each column gives the solution over the domain w at a given time.

    labels : list(str)
        Labels corresponding to the state list.
    """
    s = len(state_list)
    if labels is None:
        labels = [str(i) for i in range(s)]

    x0 = state_list[0][:,0]
    fig, ax = plt.subplots(1, 1, figsize=(10,4))
    lines = [ax.plot(w, x0, '-', lw=1, label=labels[i])[0] for i in range(s)]
    ax.axis([0, 1, min(X.min() for X in state_list),
                   max(X.max() for X in state_list)])
    ax.set_xlabel(r"$\omega$")
    ax.set_ylabel(r"$x(t,\omega)$")
    ax.legend(loc="lower left")

    def update_animation(index):
        for i in range(s):
            lines[i].set_ydata(state_list[i][:,index])

    dt = np.mean(np.diff(t))
    writer = animation.writers["ffmpeg"](fps=int(1/dt))
    with writer.saving(fig, save_as, 200):
        for index in range(len(t)):
            update_animation(index)
            writer.grab_frame()
    plt.close()
```

```python
video_name = "heat_linear_r=3.mp4"
animate_1D_funcs(video_name, w, t, [Xcomp, X_ROM], ["Snapshots", "ROM Output"])
Video(video_name, width=1200, height=500, embed=EMBED)
```

In this experiment, the training data and the ROM output are visibly identical.

## The Choice of Reduced Dimension Matters

With only $r = 3$ modes, we are able to accurately simulate this heat equation with high accuracy.
This is reasonable because the singular values of the snapshot matrix $X$ are essentially zero except for the first $3$.
In fact, we shouldn't expect much improvement at all after including the first 4 or 5 singular values.
On the other hand, using $r = 2$, we see some visible discrepancy between the ROM and the snapshot set.

```python
def linear_ROM(A, t, X, r, x0=None, method="BDF"):
    """Set up and solve the low-dimensional linear IVP

        dx_ / dt = A_ x_,   x_(t0) = x0_

    corresponding to the high-dimensional linear IVP

        dx / dt = A x,      x(t0) = x0.

    Parameters
    ----------
        A : (n,n) ndarray
            Linear evolution operator for the original, high-dimensional system.

        t : (K,) ndarray
            Times at which to evaluate the solution.

        X : (n,k) ndarray
            Snapshot matrix of sample solutions.

        r : int
            Dimension of desired reduced system.

        x0 : (n,) ndarrray
            High-dimensional initial condition. If not provided, use the first snapshot.

        method : str
            `method` keyword for scipy.integrate.solve_ivp().

    Returns
    -------
        X_ROM : (n,K) ndarray
            Approximate solution to the original, high-dimensional IVP.
    """
    if x0 is None:
        x0 = X[:,0]

    # Offline phase
    Vr = POD_basis(X, r)
    x0_ = Vr.T  @ x0
    A_ = Vr.T @ A @ Vr

    # Online phase
    X_ = solve_ivp(lambda t,x_: A_ @ x_,
                   (t[0], t[-1]),
                   x0_,
                   t_eval=t,
                   method=method).y

    return Vr @ X_
```

```python
X_rom2 = linear_ROM(A, t, X, 2, x0)
err = relative_L2_error(Xcomp, X_rom2, t)
print(f"Relative Error: {err*100:.6f}%")

video_name = "heat_linear_r=2.mp4"
animate_1D_funcs(video_name, w, t, [Xcomp, X_rom2], ["Snapshots", "ROM"])
Video(video_name, width=1200, height=500, embed=EMBED)
```

    Relative Error: 9.188259%

To see better how reduced dimension $r$ affects accuracy, we compute the relative errors for the ROM with $r=1,2,\ldots,12$.

```python
fig, [ax1, ax2] = plt.subplots(2, 1, sharex=True, figsize=(10,4))

errors, times = [], []
rs = np.arange(12) + 1
for r in rs:
    start_time = time.time()
    X_rom_r = linear_ROM(A, t, X, r, x0)
    times.append(time.time() - start_time)
    errors.append(relative_L2_error(Xcomp, X_rom_r, t))
ax1.semilogy(rs, errors, 'C3.-')
ax2.plot(rs, times, 'C4.-')

ax1.set_ylabel("Relative error")
ax2.set_xlabel(r"Reduced system dimension $r$")
ax2.set_ylabel("Computing time (seconds)")
ax2.set_xlim(0, rs.max()+1)
plt.show()
```


# Scaled up puzzles

Our implmentation of the 8-puzzle game is made to accept any 2D puzzle shape so let's see how it performs when we scale up the puzzles. Let's generate 3x3 puzzle, a 4x4 puzzle, a 5x5 puzzle, a 6x6 puzzle and try searching for a solution. We will not bother trying to perform uniform cost search on 5x5 as it could take more RAM than available on our local machines.

```python
import numpy as np
from tqdm import tqdm
import time
import matplotlib.pyplot as plt
import copy
import math

from board import Board
from node import Node
import search
import heuristics
```

### Let's create a function to run any puzzle on any algorithm with a specified timeout

we will choose manhattan distance as the heuristic because our previous [analysis](https://github.com/ribal-aladeeb/state-space-search/blob/main/analysis.ipynb) has demonstrated that it is on average our fasted heuristic.

```python
def search_solution(puzzle, algo, timeout=60, heuristic_func=heuristics.manhattan_distance):
    print(f'Puzzle {puzzle.shape[0]}x{puzzle.shape[1]}:')
    print(puzzle)
    print('Searching...')

    b = Board(puzzle)
    result: dict

    if algo == 'A*':
        result = search.a_star(b ,H=heuristic_func, timeout=timeout)

    elif algo == 'GBF':
        result = search.greedy_best_first(b ,H=heuristic_func, timeout=timeout)

    elif algo == 'UCS':
        result = search.uniform_cost(b, timeout=timeout)

    else:
        result = {'success': False, 'message': f'{algo} is not a valid argument position'}
        print(result['message'])

    if result['success']:
        print(f'Solution found in {round(result["runtime"],2)} seconds')
    else:
        print(result['message'])

    solution_path_length = 1
    node = result['current_node']

    while node.parent != None:
        solution_path_length += 1
        node = node.parent

    final_msg = f'Length of solution path: {solution_path_length}\n'
    final_msg += f'Cost: {result["current_node"].total_cost}\n'

    print(final_msg)

    return result

```

### Let's generate some random puzzles

```python
np.random.seed(4)
puzzle3 = np.random.permutation(3*3).reshape(3,3)
puzzle4 = np.random.permutation(4*4).reshape(4,4)
puzzle5 = np.random.permutation(5*5).reshape(5,5)
puzzle6 = np.random.permutation(6*6).reshape(6,6)

puzzles = [puzzle3, puzzle4, puzzle5, puzzle6]
print('Here are the random puzzles we generated:\n')
for p in puzzles:
    print(p)
    print()
```

    Here are the random puzzles we generated:

    [[3 4 6]
     [2 8 0]
     [1 5 7]]

    [[ 5  1  0  3]
     [ 6 11 10  2]
     [ 4  8  9 15]
     [14 13  7 12]]

    [[12 22 18 16 11]
     [ 4  5 20 13  1]
     [15  3 17  0  7]
     [10 19  8 14  2]
     [23  6  9 24 21]]

    [[30 21 15  8 14 26]
     [ 1 24 29  4 19  0]
     [33  5 34  2 32 20]
     [10 13  7  6 31 28]
     [11  3 22 18 23 12]
     [16 27  9 25 35 17]]

#### We expect GBF to converge relatively quickly even for very large search spaces (like 6x6 puzzle)

Therefore we will run GBF on all of these puzzles first to make sure that at the very least we can get to 6x6 using GBF. If all puzzles converge within a minute, we will try A*.

## Greedy Best First Search

```python
for p in puzzles:
    result = search_solution(p, 'GBF', timeout=60)

    if 'timeout' in result['message']:
        print('No point of searching through larger puzzles given the current timeout\n')
        break

    print()
```

    Puzzle 3x3:
    [[3 4 6]
     [2 8 0]
     [1 5 7]]
    Searching...
    Solution found in 0.02 seconds
    Length of solution path: 10
    Cost: 14

    Puzzle 4x4:
    [[ 5  1  0  3]
     [ 6 11 10  2]
     [ 4  8  9 15]
     [14 13  7 12]]
    Searching...
    Solution found in 2.26 seconds
    Length of solution path: 97
    Cost: 126

    Puzzle 5x5:
    [[12 22 18 16 11]
     [ 4  5 20 13  1]
     [15  3 17  0  7]
     [10 19  8 14  2]
     [23  6  9 24 21]]
    Searching...
    Solution found in 9.87 seconds
    Length of solution path: 250
    Cost: 297

    Puzzle 6x6:
    [[30 21 15  8 14 26]
     [ 1 24 29  4 19  0]
     [33  5 34  2 32 20]
     [10 13  7  6 31 28]
     [11  3 22 18 23 12]
     [16 27  9 25 35 17]]
    Searching...
    timeout after 60 seconds
    Length of solution path: 450
    Cost: 502

    No point of searching through larger puzzles given the current timeout

### Assumption proven wrong
#### It seems like GBF times out on this 6x6 puzzle after 1 minute, we can try setting the timeout to 5 min and see if we can find a solution

```python
_ = search_solution(puzzle6, 'GBF', timeout=5*60)
```

    Puzzle 6x6:
    [[30 21 15  8 14 26]
     [ 1 24 29  4 19  0]
     [33  5 34  2 32 20]
     [10 13  7  6 31 28]
     [11  3 22 18 23 12]
     [16 27  9 25 35 17]]
    Searching...
    Solution found in 200.24 seconds
    Length of solution path: 443
    Cost: 495

## Given what we know about GBF and its runtimes, Let's try A*
We will first try 1 minute timeouts to see what we can achieve. Our intuition is that A* will not be able to beat the runtimes of Greedy Best First and might possibly timeout on the 5x5 puzzle.

```python
for p in puzzles:
    result = search_solution(p, 'A*', timeout=60)

    if 'timeout' in result['message']:
        print('No point of searching through larger puzzles given the current timeout\n')
        break

    print()
```

    Puzzle 3x3:
    [[3 4 6]
     [2 8 0]
     [1 5 7]]
    Searching...
    Solution found in 0.05 seconds
    Length of solution path: 12
    Cost: 13

    Puzzle 4x4:
    [[ 5  1  0  3]
     [ 6 11 10  2]
     [ 4  8  9 15]
     [14 13  7 12]]
    Searching...
    timeout after 60 seconds
    Length of solution path: 12
    Cost: 15

    No point of searching through larger puzzles given the current timeout

#### A* timed-out on the 4x4 puzzle, which is sort of disappointing but kind of expected. let's see if we can do better with a 5 minute timeout

```python
_ = search_solution(puzzle4, 'A*', timeout=5*60)
```

    Puzzle 4x4:
    [[ 5  1  0  3]
     [ 6 11 10  2]
     [ 4  8  9 15]
     [14 13  7 12]]
    Searching...
    timeout after 300 seconds
    Length of solution path: 19
    Cost: 25

### Well that was disappointing but understandable:
We are using manhattan distance as our heuristic, which is admissible. This means the range of values that H(n) can take for any n is always between 0 and H*(n). With such a small domain, the function has a theoretical limit for how informative it can be. There's no free lunch, if you want to shave through the search space you need a more informative heuristic and that most probably means sacrificing admissibility (i.e you can no longer garantee that the search would yield the smallest cost path).

#### Let's try running A with an unadmissible heuristic (the algorithm is no longer A*)
We will use [permutation inversions](https://mathworld.wolfram.com/PermutationInversion.html). In our [heuristic analysis](https://github.com/ribal-aladeeb/state-space-search/blob/main/analysis.ipynb), we found that permutation inversions are the second best heuristic in terms of average runtime (after manhattan distance of course). Maybe it is the case that with scaled up puzzles, inversions will be able to find a solution fast enough.

```python
_ = search_solution(puzzle3, 'A*', timeout=5*60)
_ = search_solution(puzzle4, 'A*', timeout=5*60)
```

    Puzzle 3x3:
    [[3 4 6]
     [2 8 0]
     [1 5 7]]
    Searching...
    Solution found in 0.05 seconds
    Length of solution path: 12
    Cost: 13

    Puzzle 4x4:
    [[ 5  1  0  3]
     [ 6 11 10  2]
     [ 4  8  9 15]
     [14 13  7 12]]
    Searching...
    timeout after 300 seconds
    Length of solution path: 19
    Cost: 22

## Now Uniform Cost Search
We have low expectation for UCS, its runtimes are orders of magnitude larger than heuristic search so we expect it will timeout on the 4x4 puzzle.

```python
for p in puzzles:
    result = search_solution(p, 'UCS', timeout=60)

    if 'timeout' in result['message']:
        print('No point of searching through larger puzzles given the current timeout\n')
        break

    print()
```

    Puzzle 3x3:
    [[3 4 6]
     [2 8 0]
     [1 5 7]]
    Searching...
    Solution found in 12.09 seconds
    Length of solution path: 12
    Cost: 13

    Puzzle 4x4:
    [[ 5  1  0  3]
     [ 6 11 10  2]
     [ 4  8  9 15]
     [14 13  7 12]]
    Searching...
    timeout after 60 seconds
    Length of solution path: 12
    Cost: 14

    No point of searching through larger puzzles given the current timeout

### No surprise here
Well, as expected, UCS timed-out on the 4x4 puzzle (uninformed search sucks).

# Why is it so slow
You might be wondering why adding 1 row and 1 column increases the runtime by so much. Given a puzzle with R rows and C columns, the number of possible states configurations is:

 \begin{align}(R*C)!\end{align}

That is because there are RxC tiles and the first one can assume RxC possible values, the second can take RxC-1, and so on. So the number of of possible states is the number of permutations of a vector containing RxC unique values. This is what factorial growth looks like:

```python
def states(R,C): return math.factorial(R*C)
shapes = [
    (2,4),
    (3,3),
    (4,4),
    (5,5),
    (6,6),
    (7,7)
]

for s in shapes:
    print(f'For a {s[0]}x{s[1]} puzzle, the number of states is: {states(*s)}')
```

    For a 2x4 puzzle, the number of states is: 40320
    For a 3x3 puzzle, the number of states is: 362880
    For a 4x4 puzzle, the number of states is: 20922789888000
    For a 5x5 puzzle, the number of states is: 15511210043330985984000000
    For a 6x6 puzzle, the number of states is: 371993326789901217467999448150835200000000
    For a 7x7 puzzle, the number of states is: 608281864034267560872252163321295376887552831379210240000000000