%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
from matplotlib import animation
from IPython.display import HTML

"""
Explanation: Content provided under a Creative Commons Attribution license, CC-BY 4.0; code under MIT License. (c)2015 David I. Ketcheson
Version 0.3 - May 2022
Pseudospectral methods for wave equations in Python
Welcome to PseudoSpectralPython, a short course that will teach you how to solve wave equations using pseudospectral collocation methods.  This notebook is the first lesson, on solving linear problems.  Pseudospectral methods are great for wave problems where:

The solution is smooth (no shocks)
The domain is simple; e.g., cartesian or spherical domains.  In the course, we'll focus on periodic Cartesian domains.

Table of contents:

Advection-diffusion
Approximate solution
General linear PDEs

Advection-diffusion
Let's get started!  Run the code cell below to import a bunch of Python libraries we'll use.
End of explanation
"""


# Spatial grid
m=64                            # Number of grid points in space
L = 2 * np.pi                   # Width of spatial domain
x = np.arange(-m/2,m/2)*(L/m)   # Grid points
dx = x[1]-x[0]                  # Grid spacing

# Temporal grid
tmax = 4.0     # Final time
N = 25       # number grid points in time
k = tmax/N   # interval between output times

x

"""
Explanation: To begin, let's consider the partial differential equation (PDE)
$$u_t + a u_x = \epsilon u_{xx}$$
referred to as the advection-diffusion equation, for reasons we'll soon discover.  Here we wish to find $u(x,t)$, which might be the density or concentration of some substance.  The subscripts denote partial differentiation; e.g. $u_t$ is the partial derivative of $u$ with respect to $t$.  The coefficients $a$ and $\epsilon$ are constants that determine the strength of the advective and diffusive effects.
Exact solution by Fourier analysis
Let's solve this equation on a periodic domain $[-\pi,\pi]$, with some initial data
$$u(x,0) = u_0(x).$$
If we suppose for a moment that our solution is composed of a single Fourier mode with wavenumber $\xi$ and time-dependent amplitude $\hat{u}$:
$$u(x,t; \xi) = \hat{u}(t) e^{i\xi x},$$
Then we obtain a simple ordinary differential equation (ODE) for $\hat{u}$:
$$\hat{u}'(t; \xi) + i\xi a \hat{u} = -\xi^2 \epsilon \hat{u}$$
We can solve this scalar ODE exactly:
$$\hat{u}(t; \xi) = e^{(-i \xi a - \epsilon \xi^2)t} \hat{u}(0).$$
We've transformed the original PDE into a simple ODE, but you may wonder whether this is useful, since we assumed a very simple form for the solution.  The marvelous fact is that every solution of our advection-diffusion equation can be written as a linear combination (a superposition) of simple solutions of the form above, with different wavenumbers $\xi$.  We can construct the general solution as follows.
First, we take a Fourier transform of the initial data:
$$\hat{u}(t=0;\xi) = \frac{1}{\sqrt{2\pi}} \int_{-\infty}^\infty u_0(x) e^{-i\xi x}dx.$$
Then each mode evolves according to the solution of the ODE above:
$$\hat{u}'(t; \xi) = e^{(-i \xi a - \epsilon \xi^2)t} \hat{u}(0;\xi).$$
Finally, we construct the solution again by taking the inverse Fourier transform.  This just means summing up all the Fourier modes:
$$u(x,t) = \frac{1}{\sqrt{2\pi}} \int_{-\infty}^\infty \hat{u}(t; \xi) e^{i\xi x}d\xi.$$
If you haven't seen Fourier analysis at all before, now is a good time to go read up a bit and then come back.
Approximate solution by discrete Fourier transforms
We can't evaluate the integrals above exactly on the computer (at least, not for arbitrary initial data $u_0$).  Instead, we need to discretize.  To do so, we introduce a grid with a finite set of points in space and time:
End of explanation
"""


xi = np.fft.fftfreq(m)*m*2*np.pi/L  # Wavenumber "grid"
# (this is the order in which numpy's FFT gives the frequencies)

"""
Explanation: and a corresponding set of discrete wavenumber values $\xi$:
End of explanation
"""


# Initial data
u = np.sin(2*x)**2 * (x<-L/4)
uhat0 = np.fft.fft(u)

"""
Explanation: The functions $u, \hat{u}$ discussed above are replaced by finite-dimensional vectors.  These vectors are related through the discrete version of the Fourier transform, aptly called the discrete Fourier transform (DFT).  We'll look at the DFT in more detail in the next lesson.  For now, let's set the initial condition to
$$u_0(x) = \begin{cases} \sin^2(2x) & -\pi \le x < -\pi/2 \ 0 & x>-\pi/2 \end{cases}$$
and compute its DFT:
End of explanation
"""


epsilon=0.01  # Diffusion coefficient
a = 1.0       # Advection coefficient

# Store solutions in a list for plotting later
frames = [u.copy()]

# Now we solve the problem
for n in range(1,N+1):
    t = n*k
    uhat = np.exp(-(1.j*xi*a + epsilon*xi**2)*t) * uhat0
    u = np.real(np.fft.ifft(uhat))
    frames.append(u.copy())

"""
Explanation: Next, we set a value for epsilon and  compute the solution:
End of explanation
"""


# Set up plotting
fig = plt.figure(figsize=(9,4)); axes = fig.add_subplot(111)
line, = axes.plot([],[],lw=3)
axes.set_xlim((x[0],x[-1])); axes.set_ylim((0.,1.))
plt.close()

def plot_frame(i):
    #fig = plt.figure()
    #plt.plot(x,frames[i])
    line.set_data(x,frames[i])
    axes.set_title('t='+str(i*k))
    fig.canvas.draw()
    return fig

# Animate the solution
anim = matplotlib.animation.FuncAnimation(fig, plot_frame,
                                   frames=len(frames),
                                   interval=200,
                                   repeat=False)

HTML(anim.to_jshtml())

"""
Explanation: We have computed and stored the solution.  The code below plots it as an animation.
End of explanation
"""


