Untitled

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.ticker import LinearLocator
import scipy.sparse as ssp

def matrix(N):
    dx = 1/(N+1)
    matrix = (1/dx**2)*(np.diag(np.ones(N-1), -1) + np.diag(-2*np.ones(N)) + np.diag(np.ones(N-1), 1))
    return matrix

def eulerstep(Tdx, uold, dt):
    unew = uold + dt*np.dot(Tdx,uold)
    return unew

def trapztep(Tdx, uold, dt):
    unew = np.linalg.solve(np.eye(len(Tdx)) - np.multiply(Tdx, dt/2), uold + dt/2*np.dot(Tdx,uold))
    return unew

def step(choice, Tdx, uold, dt):
    if choice == 0:
        return eulerstep(Tdx, uold, dt)
    if choice == 1:
        return trapztep(Tdx, uold, dt)

def plot1(m, T, X):
    fig = plt.figure()
    ax = plt.axes(projection = '3d')
    wireframe = ax.plot_surface(T,X, np.transpose(m), cmap = cm.inferno)
    ax.zaxis.set_major_locator(LinearLocator(10))
    ax.zaxis.set_major_formatter('{x:.02f}')
    fig.colorbar(wireframe, shrink = 0.5, aspect = 5)
    ax.set_xlabel('T')
    ax.set_ylabel('X')
    ax.set_zlabel('u(t,x)')
    ax.set_title('Solution to Diffusion Equation Custom $u(0,x) = g(x)$')
    plt.show()

def task1(N , M, g, t0, tf, choice):
    # Next line should compute a 2-dimensional vector u(m∆t,n∆x)
    # u(t,x) has u(t,0) = u(t,1) = 0
    # u(0,x) = g(x)
    dt = (tf-t0)/M
    xx = np.linspace(0,1,N+2)
    tt = np.linspace(t0,tf,M+1)
    T,X = np.meshgrid(tt, xx) 
    uold = [g(i, i) for i in range(1, N+1)]
    m = []
    m.append(np.array([0, *uold, 0]))
    for _ in range(1, np.size(tt)):
        uold = step(choice, Tdx = matrix(N), uold = uold, dt = dt )
        m.append(np.array([0, *uold, 0]))
    m = np.array(m)
    return m, T, X
    

if __name__ == '__main__':
    c = 1e-6
    N = 100
    tstart = 0
    tend= 0.6
    t = tend - tstart
    M_CFL =  2*t*(N+1)**2 # CFL-condition dt/dx**2 <= 1/2 \iff t/M * (N+1)^2
    chc = 1
    if chc == 0:
        Mnbr = int(np.ceil(M_CFL))
    if chc == 1:
        Mnbr = 10*N
    # m, T, X = task1(N = 100, M = Mnbr, g = lambda x,t: c*(x)**3 + c/((t)**3), t0= tstart, tf=tend, choice = chc )
    # m, T, X = task1(N = N, M = Mnbr, g = lambda x,t: c*(x)**3 + c/((t)**3), t0= tstart, tf=tend, choice = chc ) 
    # m, T, X = task1(N = N, M = Mnbr-10, g = lambda x,t: c*(x)**3 + c/((t)**3), t0= tstart, tf=tend, choice = chc) #BREAK IT!
    m, T, X = task1(N = N, M = Mnbr, g = lambda x,t: np.exp(x /t), t0= tstart, tf = tend, choice = chc )
    # m, T, X = task1(N = N, M = Mnbr-16, g = lambda x,t: np.exp(x /t), t0= tstart, tf = tend, choice = chc) #BREAK IT AND GET PANTS!
    plot1(m = m, T = T, X = X)