---
execution:
  timeout: 300
jupytext:
  formats: ipynb,md:myst
  notebook_metadata_filter: all
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.13
    jupytext_version: 1.13.0
kernelspec:
  display_name: Python 3 (phys-581)
  language: python
  name: phys-581
language_info:
  codemirror_mode:
    name: ipython
    version: 3
  file_extension: .py
  mimetype: text/x-python
  name: python
  nbconvert_exporter: python
  pygments_lexer: ipython3
  version: 3.9.7
toc:
  base_numbering: 1
  nav_menu: {}
  number_sections: true
  sideBar: true
  skip_h1_title: false
  title_cell: Table of Contents
  title_sidebar: Contents
  toc_cell: false
  toc_position: {}
  toc_section_display: true
  toc_window_display: false
varInspector:
  cols:
    lenName: 16
    lenType: 16
    lenVar: 40
  kernels_config:
    python:
      delete_cmd_postfix: ''
      delete_cmd_prefix: 'del '
      library: var_list.py
      varRefreshCmd: print(var_dic_list())
    r:
      delete_cmd_postfix: ') '
      delete_cmd_prefix: rm(
      library: var_list.r
      varRefreshCmd: 'cat(var_dic_list()) '
  types_to_exclude: [module, function, builtin_function_or_method, instance, _Feature]
  window_display: false
---

```{code-cell} ipython3
:cell_style: center
:hide_input: false

import mmf_setup;mmf_setup.nbinit()
import logging;logging.getLogger('matplotlib').setLevel(logging.CRITICAL)
%matplotlib inline
import numpy as np, matplotlib.pyplot as plt
```

# Chebyshev Bases

+++

For testing, we will use a Gaussian.

```{code-cell} ipython3
def get_f(x, d=0, x0=0.0, sigma=1.2):
    """Return a gaussian or its derivatives."""
    f = np.exp(-((x-x0)/sigma)**2/2)
    if d == 0:
        return f
    elif d == 1:
        return (x0-x)/sigma**2 * f
    elif d == 2:
        return ((x0-x)**2 - sigma**2)/sigma**4 * f
    else:
        raise NotImplementedError
```

(fourier-bases)=
## Fourier Basis

+++

:::{margin}
We are using the "numerical" normalization of the basis functions here so that under the
dot product, the basis vector are orthogonal.  The physical normalization does not have
the factor $N/L$ in the measure, and the corresponding normalization of the states is

$$
  f_{m}(x) = \frac{1}{\sqrt{L}} e^{\I k_m (x + L/2)}.
$$
:::

We start by reviewing how derivatives are computed using the Fourier basis, where we can
represent periodic functions on $N$ equally spaced points in a box of length $L$ with
spacing $h = L/N$ (`dx` in the code):

$$
  x_n = \left.hn - \frac{L}{2}\right|_{n=0}^{N-1}, \qquad
  f_m(x) = \frac{1}{\sqrt{N}}e^{\I k_m (x+L/2)}, \qquad
  k_m = \left.\frac{2\pi m}{L}\right|_{m=0}^{N-1}.
$$

*(The wavevectors $k_m$ should be interpreted modulo $2\pi$ and wrapped so that they extend from $-\pi/L$.  See the code below.  On the grid, these values are correct because of aliasing, but should not be used for derivatives.  The ordering of $m$ goes from $0$ to $N/2-1$, then from $-N/2$ back up to $-1$.  See the
code below for a precise specification.  (Practically, use {func}`numpy.fft.fftfreq`.))* 

These are used to represent a function $f(x)$ in the box $-L/2 \leq x < L/2$, and are
orthonormal under the following inner product:

$$
  \braket{f|g} = \frac{N}{L}\int_{-L/2}^{L/2} \!\!\!\d{x}\; f^*(x) g(x),\\
  \braket{x|k_m} = f_{m}(x) = \frac{1}{\sqrt{N}}e^{\I k_m (x + L/2)},\\
  \frac{N}{L}\int_{-L/2}^{L/2} \!\!\!\d{x}\; \ket{x}\bra{x} =
  \sum_{m} \ket{k_m}\bra{k_m} = \mat{1},\\
  \braket{k_m|k_n} 
  = \frac{N}{L}\int_{-L/2}^{L/2} \!\!\!\d{x}\; \frac{e^{2\pi \I (x + L/2)(n-m)/L}}{N}
  = \delta_{mn}.
$$

```{code-cell} ipython3
N = 32
L = 20.0
dx = L/N
x = np.arange(N) * dx - L / 2
k = 2*np.pi * np.fft.fftfreq(N, dx)
k_ = 2*np.pi * ((N//2 + np.arange(N)) % N - N//2) / L
assert np.allclose(k, k_)

# Basis vectors: Note, the np.fft assumes that the
# abscissa go from 0 to L, so we need to include the -L/2
# shift as a phase
U = np.exp(1j*(x[:, None] + L / 2) * k[None, :])/np.sqrt(N)
assert np.allclose(U.T.conj() @ U, np.eye(N))

# Derivative operators
def get_D(d=1):
    """Return the matrix implementing the d'th derivative."""
    return (U * ((1j*k)**d)[None, :]) @ U.T.conj()


def D(f, d=1):
    """Return the `d`'th derivative of `f`."""
    coeff = (1j*k)**d
    if d % 2 == 1:
        coeff[N//2] = 0
    return np.fft.ifft(coeff*np.fft.fft(f))

f_x = get_f(x)

# Check that the basis does the FFT.  In the FFT, the normalization 
# is only in the inverse transform.
assert np.allclose(np.fft.fft(f_x), U.T.conj() @ f_x * np.sqrt(N))

fig, ax = plt.subplots()
for d in [0, 1, 2]:
    l, = ax.plot(x, get_f(x, d=d), '-', label=f'$f^{{({d})}}(x)$')
    l, = ax.plot(x, D(f_x, d=d).real, ':', c=l.get_c())    
    assert np.allclose(D(f_x, d=d), get_f(x, d=d), atol=2e-8) 
ax.set(xlabel='$x$')
ax.legend()
```

## Dirichlet Boundary Conditions

+++

If you want to impose Dirichlet boundary conditions, use the DST-II transform with endpoints half a grid from the boundary.

$$
  x_n = \left.h(n + \tfrac{1}{2}) - \frac{L}{2}\right|_{n=0}^{N-1}, \qquad
  f_m(x) = \frac{1}{\sqrt{N}}\sin\bigl(k_m (x + L/2)\bigr), \qquad
  k_m = \left.\frac{2\pi (m+1)}{L}\right|_{m=0}^{N-1}.
$$

These are used to represent a function $f(x)$ in the box $-L/2 \leq x < L/2$, and are
orthonormal under the following inner product:

$$
  \braket{f|g} = \frac{N}{L}\int_{-L/2}^{L/2} \!\!\!\d{x}\; f^*(x) g(x),\\
  \braket{x|k_m} = f_{m}(x) = \frac{1}{\sqrt{N}}\sin\bigl(k_m (x + \tfrac{L}{2})\bigr),\\
  \frac{N}{L}\int_{-L/2}^{L/2} \!\!\!\d{x}\; \ket{x}\bra{x} =
  \sum_{m} \ket{k_m}\bra{k_m} = \mat{1},\\
  \braket{k_m|k_n} 
  = \frac{N}{L}\int_{-L/2}^{L/2} \!\!\!\d{x}\; \frac{\sin\bigl(2\pi m(x+L/2)/L\bigr)\sin\bigl(2\pi n(x+L/2)/L\bigr)}{N}
  = \delta_{mn}.
$$

+++

### Odd Derivatives

+++

Odd derivatives are a little tricky because the the momenta differ.  In particular, the DCT allows $k_{0} = 0$.  Thus, to take the derivative, we must shift the coefficients before calling the DCT (dropping the highest frequency component from the DST).

```{code-cell} ipython3
# Find the fastest way of shifting and dropping.
x = np.random.random(64)

def p1(x):
    res = np.roll(x, 1)
    res[0] = 0
    return res

def p2(x):
    res = np.zeros_like(x)
    res[1:] = x[:-1]
    return res

def p3(x):
    res = np.concatenate([[0], x[:-1]])
    return res


assert np.allclose(p1(x), p2(x))
assert np.allclose(p1(x), p3(x))
%timeit p1(x)
%timeit p2(x)
%timeit p3(x)
```

```{code-cell} ipython3
N = 32
L = 20.0
dx = L/N
x = np.arange(N) * dx - L/2 + dx / 2
k = np.pi * (np.arange(N) + 1) / L

def D(f, d=1):
    """Return the `d`'th derivative of `f`."""
    f_k = sp.fft.dst(f, type=2)
    coeff = (-1)**(d // 2) * k ** d
    c_f_k = coeff * f_k
    if d % 2 == 0:
        return sp.fft.idst(c_f_k, type=2)
    else:
        _c = np.concatenate([[0], c_f_k[:-1]])
        return sp.fft.idct(_c, type=2)
    

f_x = get_f(x)

fig, ax = plt.subplots()
for d in [0, 1, 2]:
    l, = ax.plot(x, get_f(x, d=d), label=f'$f^{{({d})}}(x)$')
    l, = ax.plot(x, D(f_x, d=d), ':', c=l.get_c())
    #print(abs(D(f_x, d=d) -  get_f(x, d=d)).max())
    assert np.allclose(D(f_x, d=d), get_f(x, d=d), atol=1e-8) 
ax.set(xlabel='$x$')
ax.legend()    
```

## Chebyshev Basis (Incomplete)

In particular, we can define this basis on angles $\theta \in [0, \pi)$ by choosing $L =
\pi$ and shifting by $\d{\theta}/2$ so that we avoid the endpoints.

The idea of Chebyshev basis is to use the Fourier basis, but on the transformed abscissa

$$
  x = -\frac{L}{2}\cos \theta, \qquad
  \theta = \cos^{-1}\left(\frac{-2x}{L}\right), \\
  \diff{x}{\theta} = \frac{L}{2}\sin\theta 
                   = \frac{L}{2}\sqrt{1 - \frac{4x^2}{L^2}}
                   = \sqrt{\frac{L^2}{4} - x^2}.
$$

Hence, to compute the derivative of $f(x)$, we use the FFT to compute the derivative of
$f(\theta)$ as above, then transform back with the chain rule:

$$
  \diff{f(x)}{x} = \diff{f(\theta)}{\theta}\diff{\theta}{x} = 
  f'(\theta)\frac{2}{L \sin \theta}\\
  \diff[2]{f(x)}{x} 
  f''(\theta)\frac{4}{L^2 \sin^2 \theta} -
  f'(\theta)\frac{2\cos\theta}{L \sin^2 \theta}.
$$


$$
  \braket{f|g} = \frac{N}{\pi}\int_{0}^{\theta}\d{\theta}\; f^*(\theta)g(\theta)
  %= \frac{2N}{\pi L}\int_{-L/2}^{L/2}\!\!\!\d{x}\; \frac{f^*(\theta)g(\theta)}{\sin\theta}\\
  = \frac{N}{\pi}\int_{-L/2}^{L/2}\!\!\!\d{x}\; \frac{f^*(x)g(x)}{\sqrt{\frac{L^2}{4} - x^2}}
$$


The basis vectors are thus

$$
  f_m(x) = \frac{1}{\sqrt{N}}e^{\I k_m \theta(x)}
  
$$

+++



+++

Although this approach works, a problem is that the matrix corresponding to the derivative operator is not anti-Hermitian.  To fix this, we look for an alternative proceedure 
$$
  \diff{f(x)}{x} = \diff{f(\theta)}{\theta}\diff{\theta}{x} = 
  f'(\theta)\frac{2}{L \sin \theta}\\
$$

```{code-cell} ipython3
N = 3
L = 20.0

L_theta = np.pi
d_theta = L_theta / N 
theta = np.arange(N) * d_theta + d_theta / 2
x = -L/2 * np.cos(theta)
k_theta = 2*np.pi * np.fft.fftfreq(N, d_theta)
U_theta = np.exp(1j*(theta[:, None] - d_theta / 2) * k_theta[None, :])/np.sqrt(N)

def D(f, d=1):
    assert d == 1
    coeff = (1j*k_theta)**d
    df_dtheta = np.fft.ifft(coeff*np.fft.fft(f))
    df_dx = 2/L * df_dtheta / np.sin(theta)
    return df_dx


# Derivative operators
def get_D(d=1):
    """Return the matrix implementing the d'th derivative."""
    assert d == 1
    a = np.sqrt(2 / (L * np.sin(theta)))
    b = np.cos(theta) / (L * np.sin(theta))**2
    U = a[:, None] * U_theta
    D_k = ((1j*k_theta)**d)
    return ((a**2)[:, None] * U_theta * D_k[None, :]) @ U_theta.T.conj()
    return (U * D_k[None, :]) @ U.T.conj() - np.diag(b)


#assert np.allclose(D(f_x, d=1), get_D(d=1) @ f_x)
f_x = get_f(x, d=0)
plt.plot(x, f_x, '-+')
plt.plot(x, get_f(x, d=1), '-')
plt.plot(x, D(f_x, d=1), ':')



print(abs(D(f_x, d=1) - get_D(d=1) @ f_x).max())
```

```{code-cell} ipython3

```

```{code-cell} ipython3
D1 = get_D(d=1)
abs(D1 + D1.T.conj()).max()
```

```{code-cell} ipython3
N = 32
L = 20.0

L_theta = np.pi
d_theta = L_theta / N 
theta = np.arange(N) * d_theta
x = -L/2 * np.cos(theta)
k_theta = 2*np.pi * np.fft.fftfreq(N, d_theta)

U_theta = np.exp(1j*(theta[:, None]) * k_theta[None, :])/np.sqrt(N)
assert np.allclose(U.T.conj() @ U, np.eye(N))

# Derivative operators
def get_D(d=1):
    """Return the matrix implementing the d'th derivative."""
    return (U * ((1j*k)**d)[:, None]) @ U.T.conj()


def get_f(x, d=0, x0=0.0, sigma=1.2):
    """Return a gaussian or its derivatives."""
    f = np.exp(-((x-x0)/sigma)**2/2)
    if d == 0:
        return f
    elif d == 1:
        return (x0-x)/sigma**2 * f
    elif d == 2:
        return ((x0-x)**2 - sigma**2)/sigma**4 * f
    else:
       raise NotImplementedError


def D(f, d=1):
    """Return the `d`'th derivative of `f`."""
    coeff = (1j*k)**d
    return np.fft.ifft(coeff*np.fft.fft(f))

f_x = get_f(x)

# Check that the basis does the FFT.  In the FFT, the normalization 
# is only in the inverse transform.
assert np.allclose(np.fft.fft(f_x), U.T.conj() @ f_x * np.sqrt(N))

fig, ax = plt.subplots()
for d in [0, 1, 2]:
    l, = ax.plot(x, get_f(x, d=d), label=f'$f^{{({d})}}(x)$')
    assert np.allclose(D(f_x, d=d), get_f(x, d=d), atol=2e-8) 
ax.set(xlabel='$x$')
ax.legend()
```

## Trefethan

+++

Here is an implementation of the formula provided by Trefethen:

```{code-cell} ipython3

def get_x_D(N):
    """Return `(x, D)`, the points and derivative matrix."""
    if N == 1:
        x = np.array([1])
        D = np.array([[0]])
        return x, D
    theta = np.pi * np.arange(N)/(N-1)
    x = np.cos(theta)
    c = np.ones(N)
    c[0] = c[-1] = 2.0
    c = c[:, np.newaxis] * (-1) ** np.arange(N)
    dX = x[:, None] - x[None, :]
    D = c/c.T / (dX + np.eye(N))
    D -= np.diag(D.sum(axis=-1))
    return x, D

x, D = get_x_D(21)
f_x = np.exp(x)*np.sin(5*x)
df_x = np.exp(x)*(np.sin(5*x) + 5*np.cos(5*x))
plt.plot(x, D@f_x - df_x, '-o')
```

```{code-cell} ipython3

```