---
execution:
  timeout: 60
jupytext:
  cell_metadata_filter: -all
  formats: md:myst,ipynb
  main_language: python
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.13
    jupytext_version: 1.16.7
kernelspec:
  display_name: Python 3 (ipykernel)
  language: python
  name: python3
---

```{code-cell}
:tags: [hide-cell]

import mmf_setup; mmf_setup.nbinit()
import os, sys
from pathlib import Path
FIG_DIR = Path(mmf_setup.ROOT) / '../Docs/_build/figures/'
os.makedirs(FIG_DIR, exist_ok=True)
import logging; logging.getLogger("matplotlib").setLevel(logging.CRITICAL)
%matplotlib
import numpy as np, matplotlib.pyplot as plt
try: from myst_nb import glue
except: glue = None

from matplotlib.animation import FuncAnimation
from gpe.contexts import FPS
```

# Spectral Methods: Example 1

In §1.2, {cite}`Boyd:1989` solves the following equation using a low-order polynomial
expansion:
\begin{gather*}
  u_{,xx} - (x^6 + 3x^2)u = 0, \qquad
  u(-1) = u(1) = 1,
\end{gather*}
which has the analytic solution
\begin{gather*}
  u(x) = e^{(x^4 - 1)/4}.
\end{gather*}

```{code-cell}
:tags: [margin, hide-input]
import sympy
x = sympy.symbols('x', real=True)
u_exact = sympy.exp((x**4-1)/4)

def deq(u):
    return (u.diff(x, x) - (x**6 + 3*x**2)*u)

assert deq(u_exact).simplify() == 0
assert u_exact.subs([(x, -1)]) == 1
assert u_exact.subs([(x, 1)]) == 1
```

Read the solution there, then repeat that analysis using a Fourier basis using
everything you learn, including appropriately chosen collocation points, and considering
the discussion about $a_1 = 0$.

## Solution

Here we use a basis constructed from the Chebyshev polynomials:
\begin{gather*}
  T_n(x) = \cos(n\cos^{-1}x).
\end{gather*}

As Boyd points out, the solutions must be even (prove this!), hence we need only
consider the even terms in this series.  We use $T_0(x) = 1$ to satisfy the boundary
conditions, and then use the following basis, which 
\begin{gather*}
  \phi_n(x) = T_{2n}(x) - 1, \qquad n \in \{1, 2, \cdots, N\},\\
  u_{\vect{a}}(x) = 1 + \sum_{n=1} a_n \phi_n(x) 
                  = 1 + \sum_{n=1} a_n \bigl(\cos(2n\cos^{-1}x) -1\bigr).
\end{gather*}

The goal is to determine the parameters $a_n$ in order to minimize some appropriate
**residual** $R(\vect{a})$.

### Spectral vs Pseudo-Spectral

Consider directly fitting the solution $u_*(x)$.  The idea of a **spectral** method
would be to use some sort of residual like the least-squares overlap:
\begin{gather*}
  R(\vect{a}) = \int_{-1}^{1}\d{x} \bigl(u(x) - u_{\vect{a}}(x)\bigr)^2.
\end{gather*}
Given that the Chebyshev polynomials are orthogonal with a weight, a better option might
be 
\begin{gather*}
  R(\vect{a}) = \int_{-1}^{1}\frac{\d{x}}{\sqrt{1-x^2}} \bigl(u(x) - u_{\vect{a}}(x)\bigr)^2.
\end{gather*}
A **pseudo-spectral** method replaces these integrals by evaluating the residual at a
carefully chosen set of **collocation points** $x_i$ and **quadrature weights** $w_i$:
\begin{gather*}
  R(\vect{a}) = \sum_{i=1}^{N} w_i \bigl(u(x_i) - u_{\vect{a}}(x_i)\bigr)^2.
\end{gather*}
As {cite}`Boyd:1989` discusses, if these are optimally chosen, then this method is as
good as the spectral approach.

### Best Spectral Fit to Solution

Since we know the solution, we can compute the best fit solution, in the least-squared

```{code-cell}
import numpy as np
from scipy.optimize import curve_fit

def model(x, *a):
    n_ = 1 + np.arange(len(a), dtype=float)[None, :]
    a_ = np.array(a)[None, :]
    x_ = np.asarray(x)[:, None]
    return 1 + (a_ * (np.cos(2*n_*np.arccos(x_))-1)).sum(axis=-1)

exact = sympy.lambdify(x, u_exact, "numpy")

fig, (ax, ax1) = plt.subplots(1, 2, figsize=(8, 3), constrained_layout=True)
xdata = np.linspace(-1, 1, 500)
ydata = exact(xdata)
ax.plot(xdata, ydata, '-k')
max_errs_spec = []
Ns = np.arange(2, 10)
for N in Ns:
    N = int(N)
    p, pcov = curve_fit(model, xdata, ydata, p0=np.ones(N))
    ax.plot(xdata, model(xdata, *p), '--', label=f"{N=}")
    max_errs_spec.append((abs(model(xdata, *p) - exact(xdata))).max())
ax.legend(loc='upper center')
ax.set(xlabel="$x$", ylabel="$u(x)$")

ax1.loglog(Ns, max_errs_spec, '-x')
ax1.set(xlabel="$N$", ylabel="max err", xticks=Ns, xticklabels=Ns)
ax1.yaxis.tick_right()
ax1.yaxis.set_ticks_position('both')
p2, pcov = curve_fit(model, xdata, ydata, p0=np.ones(2))
print(p2)
```


```{code-cell}
:tags: [margin, hide-input]
th = np.linspace(0, 2*np.pi, 1000)[:-1]
dth = np.diff(th).mean()

@np.vectorize
def get_f(k):
    return np.sum(np.exp(-1j*k*th + np.cos(th)**4))*dth

k = 2*np.arange(1, 10)
fk = abs(get_f(k))
fk_ = 1/k**(k/4+0.5)
fk_ *= fk[0]/fk_[0]

fig, ax = plt.subplots(figsize=(4, 3))
ax.semilogy(k, fk, label=r"$|\tilde{f}_{k}|$")
ax.semilogy(k, fk_, '--', label=r"$\propto 1/k^{k/4+1/2}$")
ax.set(xlabel="$k=2N$", ylabel=r"$\tilde{f}_{k}$", 
       title=r"Saddle-point approximation of $\tilde{f}_{k}$")
ax.legend();
```

```{code-cell}
:tags: [margin, hide-input]
fig, ax = plt.subplots(figsize=(4, 3))
max_errs_spec = []
Ns = np.arange(2, 10)
for N in Ns:
    p, pcov = curve_fit(model, xdata, ydata, p0=np.ones(N), 
                        xtol=1e-12, maxfev=10000)
    max_errs_spec.append((abs(model(xdata, *p) - exact(xdata))).max())
ax.loglog(Ns, max_errs_spec, '-x', label="Error")
ax.loglog(Ns, max_errs_spec[0]*np.exp(-Ns**(1.4) + Ns[0]**(1.4)), '--', 
          label=r"$\propto e^{-N^{1.4}}$")
k = 2*(Ns+1)
err = 1/k**(k/4 + 0.5)
ax.loglog(Ns, max_errs_spec[0]*err/err[0], '-', label=r"$\propto 1/k_{\max}^{k_{\max}/4+1/2}$")
ax.set(xlabel="$N$", ylabel="max err", xticks=Ns, xticklabels=Ns);
ax.legend();
```
:::{margin}
Absolute error as a function of $N$ and saddle-point approximation $\abs{\tilde{f}_{k_{\max}}}
\propto 1/k_{\max}^{k_{\max}/4+1/2}$ with $k_{\max} \approx 2(N+1)$.  If you can explain
the deviation, please let me know.
:::
:::::{admonition} Questions

The right plot shows the absolute errors on a log-log plot.

1. Why do the errors bottom out at $\sim 10^{-9}$?
2. How do the errors scale as a function of the number of basis functions?

:::{solution}
   
1. This is because {func}`scipy.optimize.curve_fit` calls
   {func}`scipy.optimize.least_squares`, which has a default tolerance of $10^{-8}$.
   Increasing the tolerance gives the plot in the margin.
2. Playing a bit, you can find that the absolute error behaves like
   \begin{gather*}
     \epsilon \sim e^{-N^{1.4}}.
   \end{gather*}
   I doubt this has an significance, but it does capture the behavior over the shown
   range.
   
   To gain a better understanding, note that the Chebyshev basis is really a Fourier basis in
   disguise.  Thus, the convergence is determined by the asymptotic behaviour of the
   Fourier transform of
   \begin{gather*}
     f(\theta) \sim u(x) \sim e^{\cos^4\theta}, \qquad 
     x = \cos\theta.
   \end{gather*}
   This can be estimated using the [saddle-point approximation][].  Let $x_0 = \cos\theta_0$:
   \begin{gather*}
     \tilde{f}_{k} = \int_{0}^{2\pi} e^{-\I k \theta + \cos^4\theta}\d{\theta}
                   \approx \pm \I e^{-\I k \theta_0 + x_0^4}
                   \sqrt{\frac{2\pi}{4x_0^2(3 - 4x_0^2)}}
   \end{gather*}
   where $\theta_0$ is a saddle-point of the integrand:
   \begin{gather*}
     0 = \diff{(-\I k \theta_0 + \cos^4\theta_0)}{\theta_0} =
     -\I k - 4\cos^3\theta_0\sin\theta_0 \\ 
     = -\I k - 4x_0^3\sqrt{1-x_0^2}.
   \end{gather*}
   There is a saddle-point on the negative imaginary axis $\theta_0 = -\I b$ where
   \begin{gather*}
     k = 4\cosh^3{b} \sinh{b} \approx \frac{e^{4b}}{4}, \qquad
     b \approx \frac{\ln 4k}{4},\\
     -\I k\theta_0 \approx \frac{- k\ln 4k}{4}, \\
     x_0 = \cos\theta_0 \approx \cosh b \approx \frac{e^{b}}{2}
         = \frac{(4k)^{1/4}}{2}, \qquad
     x_0^4 \approx \frac{k}{4},\\
     -\I k \theta_0 + x_0^4 \approx \frac{k(1 - \ln 4k)}{4},\\
     \tilde{f}_{k} \sim \frac{1}{k^{k/4+1/2}}.
   \end{gather*}
   The largest $k$ in our basis corresponds to $\cos(2N \theta)$ so $k_{\max} = 2N$ and
   the truncation error should be proportional to the missing terms with $k = 2(N+1)$.
   This is approximately correct. (See the plots to the right, which verify that the
   saddle-point approximation is good, but there is some deviation in this scaling --
   the basis is slightly better than expected.)
:::
:::::

### Best Pseudo-Spectral Fit

For the Chebyshev polynomials, the collocation points are the roots of the $T_{N}(x)$,
in other words, where $\cos(N \cos^{-1} x) = 0$, or $N\cos^{-1}x = (i + \tfrac{1}{2})\pi$:
\begin{gather*}
  x_{i} = \cos\frac{(i + \tfrac{1}{2}) \pi}{N}.
\end{gather*}
In our cases, an $N$ term solution is built upon $T_{2n}$ up to $T_{2N}(x)$.  Recall
that we skip the odd orders because of the parity of our solution: likewise, we need
only keep the positive collocation points:
\begin{gather*}
  x_{i} = \left.\cos\frac{(i + \tfrac{1}{2}) \pi}{2N}\right|_{i=0}^{N-1}.
\end{gather*}
With these, we can exactly solve for the coefficients $a_n$ using
{func}`numpy.linalg.solve`, minimizing the pseudo-spectral residual to exactly zero:

```{code-cell}
def get_xi(N):
    """Return the N collocation points."""
    i = np.arange(N)
    xi = np.cos((i + 0.5) * np.pi / 2 / N)
    return xi
    
def get_a(N, exact=exact):
    x_ = get_xi(N)[:, None]
    n_ = 1 + np.arange(N, dtype=float)[None, :]
    k_ = 2*n_
    A = np.cos(k_*np.arccos(x_)) - 1
    a = np.linalg.solve(A, exact(x_.ravel()) - 1)
    return a
```
```{code-cell}
:tags: [hide-input]

fig, (ax, ax1) = plt.subplots(1, 2, figsize=(8, 3), constrained_layout=True)
xdata = np.linspace(-1, 1, 500)
ydata = exact(xdata)
ax.plot(xdata, ydata, '-k')
max_errs_pspec = []
Ns = np.arange(2, 16)
for N in Ns:
    N = int(N)
    a = get_a(N)
    max_errs_pspec.append((abs(model(xdata, *a) - exact(xdata))).max())
    if N > 6:
        continue
    ax.plot(xdata, model(xdata, *a), '--', label=f"{N=}")
ax.legend(loc='upper center')
ax.set(xlabel="$x$", ylabel="$u(x)$")

ax1.loglog(Ns, max_errs_pspec, '-x', label="Pseudo-spectral")
ax1.loglog(Ns[:len(max_errs_spec)], max_errs_spec, ':+', label="Spectral")

k = 2*(Ns+1)
err = 1/k**(k/4 + 0.5)
ax1.loglog(Ns, max_errs_pspec[0]*err/err[0], '--', label=r"$\propto 1/k_{\max}^{k_{\max}/4+1/2}$")
ax1.set(xlabel="$N$", ylabel="max err", xticks=Ns, xticklabels=Ns)
ax1.yaxis.tick_right()
ax1.yaxis.set_ticks_position('both')
ax1.legend(loc="lower left")
a2 = get_a(2)
print(a2)
```
Although the errors are slightly larger than with the spectral method (which directly
minimizes the residual), they still decrease exponentially, and are very comparable, and
code is much simpler (no need to call {func}`~scipy.optimize.curve_fit`).

### Solving the BVP

This previous analysis shows how well we might do with (pseudo-)spectral methods, but
cheats by using the exact solution.  We now apply the method to directly solve the ODE.
The only change is that we now use the ODE as our residual:
\begin{gather*}
  \left.u_{,xx}(x) - (x^6 + 3x^2)u(x) = 0\right|_{x\in \{x_{i}\}}, \qquad
  x_i = \left.\cos \frac{i + \tfrac{1}{2}}{2N}\pi\right|_{i=0}^{N-1}.
\end{gather*}
Note that these equations are inhomogeneous because of our incorporation of the boundary
conditions into $u_{\vect{a}}(x)$:
\begin{gather*}
  \sum_{n}[\phi_{n}''(x_{i}) - (x_{i}^6 + 3x_{i}^2)\phi_{n}(x_{i})]a_{n}= x_{i}^6 + 3x_{i}^2.
\end{gather*}
The only thing we need is the second derivative of the basis functions, which can be
computed using the chain rule:
:::{margin}
These can be simplified and expressed in terms of other basis functions, or computed
using the DCT, but we can just use them for now.  See {ref}`sec:ChebyshevBasis` for
details.
:::
\begin{gather*}
  x = \cos \theta, \qquad T_{n}(x) = \cos n \theta,\\
  \diff{x}{\theta} = -\sin \theta, \qquad
  T'_{n}(x) = -n\sin n \theta \diff{\theta}{x}
            = \frac{n\sin n \theta}{\sin \theta},\\
  T''_{n}(x) = -\frac{n^2\cos n \theta \sin \theta - n\sin n \theta \cos \theta}{\sin^3 \theta}.
\end{gather*}

```{code-cell}
def phi(x, n, d=0):
    """Return the dth derivative of phi_n(x)"""
    th = np.arccos(x)
    cn = np.cos(2*n*th)

    if d == 0:
        return cn - 1
    
    sn = np.sin(2*n*th)
    s = np.sin(th)

    if d == 1:
        return (2*n) * sn / s

    c = np.cos(th)
    if d == 2:
        return -((2*n)**2 * cn * s - (2*n) * sn * c) / s**3


# Check that these are correct... Don't use x0 = 0.5!
x0 = 0.3
h = 1e-8
n = 3
assert np.allclose((phi(x0+h, n) - phi(x0-h, n))/2/h, phi(x0, n, d=1))
assert np.allclose((phi(x0+h, n, d=1) - phi(x0-h, n, d=1))/2/h, phi(x0, n, d=2))

def get_a(N):
    i = np.arange(N)
    xi = np.cos((i + 0.5) * np.pi / 2 / N)
    phis = np.array([phi(xi, n=n) for n in 1+np.arange(N)])
    d2phis = np.array([phi(xi, n=n, d=2) for n in 1+np.arange(N)])
    rhs = xi**6 + 3*xi**2
    A = np.transpose(d2phis - rhs[None, :]*phis)
    a = np.linalg.solve(A, rhs)
    return a

a2 = get_a(2)
print(a2)
```
```{code-cell}

fig, (ax, ax1) = plt.subplots(1, 2, figsize=(8, 3), constrained_layout=True)
xdata = np.linspace(-1, 1, 500)
ydata = exact(xdata)
ax.plot(xdata, ydata, '-k')
max_errs_ode = []
Ns = np.arange(2, 16)
for N in Ns:
    N = int(N)
    a = get_a(N)
    max_errs_ode.append((abs(model(xdata, *a) - exact(xdata))).max())
    if N > 6:
        continue
    ax.plot(xdata, model(xdata, *a), '--', label=f"{N=}")
ax.legend(loc='upper center')
ax.set(xlabel="$x$", ylabel="$u(x)$")

ax1.loglog(Ns, max_errs_ode, '-x', label="Pseudo-spectral ODE errors")
ax1.loglog(Ns[:len(max_errs_spec)], max_errs_spec, ':+', label="Spectral best fit")

k = 2*(Ns+1)
err = 1/k**(k/4 + 0.5)
ax1.loglog(Ns, max_errs_pspec[0]*err/err[0], '--', label=r"$\propto 1/k_{\max}^{k_{\max}/4+1/2}$")
ax1.set(xlabel="$N$", ylabel="max err", xticks=Ns, xticklabels=Ns)
ax1.yaxis.tick_right()
ax1.yaxis.set_ticks_position('both')
ax1.legend(loc="lower left")
a2 = get_a(2)
print(a2)
```

[saddle-point approximation]: <https://physics-571-math-methods.readthedocs.io/en/latest/11_ComplexAnalysis.html#saddle-point-approximation>