wave/wave_equation.py

# Import necessary libraries
import numpy as np  # For numerical computations and arrays
import matplotlib.pyplot as plt  # For plotting
from matplotlib.animation import FuncAnimation  # For creating animations

# Simulation parameters
c = 1.0001  # Wave propagation speed
L = 1.0  # Length of the spatial domain
T = 1.0  # Total simulation time
dx = 0.01  # Spatial discretization step
dt = 0.01  # Time discretization step
r = c * dt / dx  # Courant number (must be <= 1 for stability)

# Spatial grid setup
N = int(L / dx) + 1  # Number of grid points
x = np.linspace(0, L, N)  # Spatial coordinates array

# Initial conditions for the wave equation
u0 = np.sin(2 * np.pi * x / L)  # Initial displacement at t=0
g = np.cos(2 * np.pi * x / L)  # Initial velocity function at t=0
u1 = u0 + dt * g  # Displacement at first time step (t=dt), incorporating initial velocity

# External force function applied to the wave equation
def force(x, t):
    return 1 * np.sin(2 * np.pi * x / L) * np.sin(2 * np.pi * t / T)  # Sinusoidal force in space and time

# Initialize arrays for time-stepping (finite difference method)
u_prev = u0.copy()  # Displacement at previous time step (t=0)
u_curr = u1.copy()  # Displacement at current time step (t=dt)

# Time-stepping loop to solve the wave equation using finite differences
frames = []  # List to store displacement at each time step for animation
for n in range(int(T / dt)):  # Loop over time steps
    t = n * dt  # Current time
    u_next = np.zeros(N)  # Array for next time step displacement
    f = force(x[1:-1], t)  # Compute external force at interior points
    # Vectorized update for interior points: ∂²u/∂t² = c² ∂²u/∂x² + f
    u_next[1:-1] = 2 * u_curr[1:-1] - u_prev[1:-1] + r**2 * (u_curr[2:] - 2*u_curr[1:-1] + u_curr[:-2]) + dt**2 * f
    # Apply Dirichlet boundary conditions (fixed ends)
    u_next[0] = 0  # Left boundary
    u_next[-1] = 0  # Right boundary
    # Shift arrays for next iteration
    u_prev = u_curr.copy()
    u_curr = u_next.copy()
    frames.append(u_curr.copy())  # Store frame for animation

# Create animation of the wave propagation
fig, ax = plt.subplots()  # Create figure and axis
line, = ax.plot(x, u0)  # Initial plot line
ax.set_ylim(-1, 1)  # Set y-axis limits
ax.set_xlim(0, L)  # Set x-axis limits

def animate(frame):  # Animation function
    line.set_ydata(frame)  # Update line data
    return line,

ani = FuncAnimation(fig, animate, frames=frames, interval=50)  # Create animation
ani.save("wave_equation.gif", writer='pillow')  # Save as GIF
plt.show()  # Display the plot (may not work in headless environment)
np.savetxt("final_wave_state.txt", frames)  # Save final frames to text file