1D Crank-Nicholson Implicit

##########################################################
#                                                        #
#              PARABOLIC EQ LINEAR CODE                  #
#          Implicit Crank-Nicholson Method               #
#      COPYRIGHT CHRISTOPHER GOMEZ, 2020, 2023, 2024     #
#                                                        #
##########################################################

import numpy as np
import matplotlib.pyplot as plt
from scipy.sparse import diags
from scipy.sparse.linalg import spsolve

# Space parameters
L = 1.0  # Length of domain
nx = 100  # Number of spatial points
dx = L / (nx - 1)  # Spatial step size
x = np.linspace(0, L, nx)

# Time parameters
T = 0.5  # Total time
nt = 1000  # Number of time steps
dt = T / nt  # Time step size

# Diffusion coefficient
alpha = 0.1

# Compute stability parameter
r = alpha * dt / (dx**2)

# Initialize solution array
u = np.zeros((nt + 1, nx))

# Set initial condition (sine wave)
u[0, :] = np.sin(2 * np.pi * x)

# Set boundary conditions (u = 0 at both ends)
u[:, 0] = 0
u[:, -1] = 0

# Create implicit matrix (I - r/2 * A)
diagonal = 1 + r
off_diagonal = -r/2
diagonals_implicit = [diagonal * np.ones(nx), 
                     off_diagonal * np.ones(nx-1), 
                     off_diagonal * np.ones(nx-1)]
A_implicit = diags(diagonals_implicit, [0, -1, 1]).tocsr()

# Modify matrices for boundary conditions
A_implicit[0, :] = 0
A_implicit[0, 0] = 1
A_implicit[-1, :] = 0
A_implicit[-1, -1] = 1

# Create explicit matrix (I + r/2 * A)
diagonal = 1 - r
off_diagonal = r/2
diagonals_explicit = [diagonal * np.ones(nx), 
                     off_diagonal * np.ones(nx-1), 
                     off_diagonal * np.ones(nx-1)]
A_explicit = diags(diagonals_explicit, [0, -1, 1]).tocsr()

# Modify matrices for boundary conditions
A_explicit[0, :] = 0
A_explicit[0, 0] = 1
A_explicit[-1, :] = 0
A_explicit[-1, -1] = 1

# Time stepping
for n in range(nt):
    # Compute RHS (explicit part)
    b = A_explicit.dot(u[n, :])
    
    # Enforce boundary conditions
    b[0] = 0
    b[-1] = 0
    
    # Solve system (implicit part)
    u[n + 1, :] = spsolve(A_implicit, b)

# Plot solution at different time points
plt.figure(figsize=(10, 6))
time_points = [0, nt//4, nt//2, 3*nt//4, -1]

for t in time_points:
    plt.plot(x, u[t, :], label=f't = {t*dt:.2f}')

plt.xlabel('x')
plt.ylabel('u(x,t)')
plt.title('Crank-Nicolson Solution of Heat Equation')
plt.legend()
plt.grid(True)
plt.show()

# Create animation
from IPython.display import HTML
import matplotlib.animation as animation

fig, ax = plt.subplots(figsize=(10, 6))
line, = ax.plot(x, u[0, :])

ax.set_xlim(x.min(), x.max())
ax.set_ylim(u.min(), u.max())
ax.set_xlabel('x')
ax.set_ylabel('u(x,t)')
ax.grid(True)

num_frames = 50

def animate(frame):
    idx = frame * (len(u) // num_frames)
    line.set_ydata(u[idx, :])
    ax.set_title(f'Time: {idx*dt:.3f}')
    return line,

anim = animation.FuncAnimation(fig, animate, frames=num_frames, 
                             interval=100, blit=True)
plt.close()

# Display animation in notebook
HTML(anim.to_jshtml())