#!/usr/bin/python3
# Crank-Nicolson method for the (1+1)-dimensional wave equation

import numpy as np
import matplotlib.pyplot as plt

L = .3       # string length
c = 100.     # phase velocity
N = 100      # number of discretization intervals
a = L/N      # lattice constant
tmax = .1    # calculate the solution for 0 <= t <= tmax
h = 1.E-4    # time increment
alpha = h * c**2 / 2 / a**2


# initial profile (some function which satisfies the BCs and is not an eigenmode)
u = np.zeros(2*N + 2)
for k in range(1, N+1):
    u[2*k] = np.sin(k / N * np.pi) * k / N / 100
tplot = [0, .01, .02, .03]  # intermediate times where the profile is plotted

# This is a very inefficient way of representing the matrices of the C-N formulas
# (in a realistic application, better use a sparse-matrix representation):
# Start with all-zero matrices
A = np.zeros([2*N + 2, 2*N + 2])
B = np.zeros([2*N + 2, 2*N + 2])
# The 2x2 blocks on the top left and on the bottom right are identity matrices
# (boundary conditions: both ends of the string are fixed and never move) 
A[0, 0] = A[1, 1] = A[2*N, 2*N] = A[2*N + 1, 2*N + 1] = 1.0
B[0, 0] = B[1, 1] = B[2*N, 2*N] = B[2*N + 1, 2*N + 1] = 1.0
# The other matrix elements:
for i in range(2, 2*N):
    A[i, i] = 1.0
    B[i, i] = 1.0
    if i % 2 == 0:   # even-index rows = phi_n
        A[i, i+1] = -h/2
        B[i, i+1] = h/2
    else:            # odd-index rows = psi_n
        A[i, i+1] = -alpha
        B[i, i+1] = alpha
        A[i, i-1] = 2*alpha
        B[i, i-1] = -2*alpha
        A[i, i-3] = -alpha
        B[i, i-3] = alpha

# main loop
t = 0.0
while t < tmax:
    # Instead of using a sparse matrix solver (which would be more efficient):
    # for demonstration purposes, use numpy.linalg.solve()
    u = np.linalg.solve(A, np.dot(B, u))
    
    for tp in tplot:
        if abs(t - tp) < 0.1 * h:
            plt.plot(u[::2], label = "t = " + str(tp) + " s")
    t += h

plt.xlabel("x")
plt.ylabel("phi(x)")   
plt.legend(loc = "upper left")  
plt.show()