exam square wave

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import math
from scipy.sparse import lil_matrix
from scipy.sparse import csr_matrix
from scipy.sparse.linalg import spsolve
import time
import sys


def solver(n, manufactured=False):

    x0 = 2
    x1 = 50
    t0 = 0
    t1 = 30
    imax = int(50*n)
    jmax = int(50*n)
    plot = True


    def non_homogeneous(x, t):
        return math.cos(x)*math.sin(t)-math.sin(x)*math.cos(t)-x*math.sin(x)*math.sin(t)

    def assumed(x, t):
        return math.sin(x)*math.sin(t)

    def assumeddx(x, t):
        return math.cos(x)*math.sin(t)

    dx = (x1-x0)/imax
    dt = (t1-t0)/jmax

    x = np.linspace(x0, x1, imax)
    t = np.linspace(t0, t1, jmax)

    xx, tt = np.meshgrid(x, t)

    #initial condition
    start = time.time()

    forced = np.zeros((imax*jmax))
    mat = lil_matrix((imax*jmax, imax*jmax))

    def vec_index(i, j):
        return i+j*imax

    count = 0

    for j in range(0, jmax):
        for i in range(0, imax):
            if manufactured:
                forced[vec_index(i,j)] = non_homogeneous(x[i], t[j])
            #Drichilet at t = 0
            if j == 0:
                forced[vec_index(i,j)] = 1 if -1<x[i]<1 else 0
                if manufactured:
                    forced[vec_index(i,j)] = assumed(x[i], t[j])
                mat[vec_index(i,j),vec_index(i,j)] = 1
                continue
            #Drichilet at left boundary u = 0
            if i == 0 and False:
                forced[vec_index(i,j)] = 0
                if manufactured:
                    forced[vec_index(i,j)] = assumed(x[i], t[j])
                mat[vec_index(i,j),vec_index(i,j)] = 1
                continue
            #Drichilet at right boundary u = 0
            if i == imax-1:
                forced[vec_index(i,j)] = 0
                if manufactured:
                    forced[vec_index(i,j)] = assumed(x[i], t[j])
                mat[vec_index(i,j),vec_index(i,j)] = 1
                continue

            # du/dx
            # forawrd diff
            if i == 0:
                #i: -3/2/dx, i+1: 2/dx, i+2: -1/2/dx
                mat[vec_index(i,j), vec_index(i,j)] += -3/2/dx
                mat[vec_index(i,j), vec_index(i+1,j)] += 2/dx
                mat[vec_index(i,j), vec_index(i+2,j)] += -1/2/dx
            # backward diff
            elif i == imax-1:
                #i: 3/2/dx, i-1: -2/dx, i-2: 1/2/dx
                mat[vec_index(i,j), vec_index(i,j)] += 3/2/dx
                mat[vec_index(i,j), vec_index(i-1,j)] += -2/dx
                mat[vec_index(i,j), vec_index(i-2,j)] += 1/2/dx
            # central diff
            else:
                #i+1: 1/2/dx, i-1: -1/2/dx
                mat[vec_index(i,j), vec_index(i+1,j)] += 1/2/dx
                mat[vec_index(i,j), vec_index(i-1,j)] += -1/2/dx

            # -du/dt
            # forawrd diff
            if j == 0:
                #i: -3/2/dx, i+1: 2/dx, i+2: -1/2/dx
                mat[vec_index(i,j), vec_index(i,j)] -= -3/2/dx
                mat[vec_index(i,j), vec_index(i,j+1)] -= 2/dx
                mat[vec_index(i,j), vec_index(i,j+2)] -= -1/2/dx
            # backward diff
            elif j == jmax-1:
                #i: 3/2/dx, i-1: -2/dx, i-2: 1/2/dx
                mat[vec_index(i,j), vec_index(i,j)] -= 3/2/dx
                mat[vec_index(i,j), vec_index(i,j-1)] -= -2/dx
                mat[vec_index(i,j), vec_index(i,j-2)] -= 1/2/dx
            # central diff
            else:
                #i+1: 1/2/dx, i-1: -1/2/dx
                mat[vec_index(i,j), vec_index(i,j+1)] -= 1/2/dx
                mat[vec_index(i,j), vec_index(i,j-1)] -= -1/2/dx

            #x*d2u/dx2
            #forawrd diff
            if i == 0:
                #i: 2/dx/dx, i+1: -5/dx/dx, i+2: 4/dx/dx, i+3: -1/dx/dx
                mat[vec_index(i,j), vec_index(i,j)] += 2/dx/dx*x[i]
                mat[vec_index(i,j), vec_index(i+1,j)] += -5/dx/dx*x[i]
                mat[vec_index(i,j), vec_index(i+2,j)] += 4/dx/dx*x[i]
                mat[vec_index(i,j), vec_index(i+3,j)] += -1/dx/dx*x[i]
            # backward diff
            elif i == imax-1:
                #i: 2/dx/dx, i+1: -5/dx/dx, i+2: 4/dx/dx, i+3: -1/dx/dx
                mat[vec_index(i,j), vec_index(i,j)] += 2/dx/dx*x[i]
                mat[vec_index(i,j), vec_index(i-1,j)] += -5/dx/dx*x[i]
                mat[vec_index(i,j), vec_index(i-2,j)] += 4/dx/dx*x[i]
                mat[vec_index(i,j), vec_index(i-3,j)] += -1/dx/dx*x[i]

            # central diff
            else:
                #i-1: 1/dx/dx, i: -1/dx/dx, i+1: 1/dx/dx,
                mat[vec_index(i,j), vec_index(i-1,j)] += 1/dx/dx*x[i]
                mat[vec_index(i,j), vec_index(i,j)] += -2/dx/dx*x[i]
                mat[vec_index(i,j), vec_index(i+1,j)] += 1/dx/dx*x[i]

    R = spsolve(csr_matrix(mat), forced).reshape((jmax,imax))
    if manufactured:
        real = np.zeros((jmax,imax))
        for j in range(jmax):
            for i in range(imax):
                real[j,i] = assumed(x[i], t[j])
        return (dx, np.linalg.norm((R-real).reshape((imax*jmax))))
    print(time.time()-start)

    if plot and not manufactured:
        fig = plt.figure(figsize=(18,10))

        ax1 = plt.subplot2grid((2,2), (0,0), colspan=2, rowspan=2, projection='3d')
        ax1.set_xlabel('x')
        ax1.set_ylabel('t')
        ax1.set_zlabel('u')
        ax1.plot_surface(xx, tt, R, shade=True, rstride=1,cstride=1,
        cmap=plt.cm.CMRmap, linewidth=0, antialiased=True);

        ax1.view_init(30, -120)
        plt.show()


if len(sys.argv) > 1 and sys.argv[1] == "TEST":
    dxs = []
    error = []
    for n in np.logspace(-1,1,25):
        a, b = solver(n, manufactured=True)
        dxs.append(a)
        error.append(b)
    print(dxs)
    plt.plot(dxs,error, 'x-')
    plt.xlabel('log(\u0394x)')
    plt.ylabel('log(\u03B5)')
    plt.yscale('log')
    plt.xscale('log')
    plt.show()
else:
    solver(1)