How to implement boundary value in Python using spectral methods

681 Views Asked by At

I want to solve the heat equation $$u_t=au_{xx},\quad 0\leq t, \ 0\leq x\leq\pi$$ in Python using spectral methods. I set $u(0,x)=\sin(x)$ as initial value for the time and choose $a=2$. My Code:

import numpy as np
import matplotlib.pyplot as plt

from mpl_toolkits.mplot3d import Axes3D
from scipy.integrate import odeint
from scipy.fftpack import diff as psdiff

def GenerateID(x):
    u = np.sin(x)
    return u

def hEq(u, t, L, a):
    uxx = psdiff(u, period=L,  order=2) 
    dudt = a*uxx
    return dudt

def hEq_solution(u0, t, L):
    sol = odeint(hEq, u0, t, args=(L,a,), mxstep=5000)
    return sol

if __name__ == "__main__":
    L = np.pi
    N = 1000
    x = np.linspace(0, L, N)
    a = 2
    u0 = GenerateID(x)
    Tfinal = 5.0
    t = np.linspace(0.0, Tfinal, 1000)
    sol = hEq_solution(u0, t, L)
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    sx, st = np.meshgrid(x, t)
    ax.plot_surface(sx, st, sol, cmap='jet')
    plt.xlabel('x')
    plt.ylabel('t')
    plt.show()

Which produces enter image description here This seems kind of right. Now my question: If I want to introduce the boundary condition $$u(t,0)=u(t,\pi)=0,$$ how would I do that? I know that considering all above conditions, the exact solution should be $$u(x,t)=\exp[-2t]\sin(x).$$