wave/wave.py

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

E=7.8e10  # Young's modulus (Pa)
rho=2700   # density (kg/m^3)
c=np.sqrt(E/rho) # wave velocity
L=1.0      # length (m)
dx=0.05
step=10    # spatial step (m)
T=1   # total simulation time (s)
dt=T/step   # time step (s) for stability
n_x=int(L/dx) # space steps
n_t=int(T/dt) # time steps
x=np.linspace(0,L,n_x) # space array
t=np.linspace(0,T,n_t) # time array
C=c*dt/dx # courant number

# initial conditions
u=np.zeros((n_x,n_t))
u[:,0]=np.cos(2*np.pi*x/L)
print(u.shape)
# step one
for i in range(1,n_x-1):
    u[i,1]=u[i,0]-.5*C**2*(u[i+1,0]-2*u[i,0]+u[i-1,0])
# boundary conditions
u[[0,-1],:]=0 # Dirichlet

# for loop for time steps
for j in range(1,n_t-1):
    for i in range(1,n_x-1):
        u[i,j+1]=u[i,j-1]+2*u[i,j]+C**2*(u[i+1,j]-2*u[i,j]+u[i-1,j])
    u[[0,-1],:]=0 # Dirichlet
 
# plot

# #
fig, ax = plt.subplots()
line, = ax.plot([], [], lw=2)
ax.set_xlim(0, L)
# ax.set_ylim(min(u[:,0]), max(u[:,0])) # Adjust y-limits based on expected amplitude
ax.set_xlabel("Position (x)")
ax.set_ylabel("Displacement (u)")
ax.set_title("Wave Propagation")

def init():
    line.set_data([], [])
    return line,

def update(frame):
    line.set_data(x, u[:,frame])
    return line,

# Animate every 10th frame to reduce file size and processing time
ani = animation.FuncAnimation(
    fig, update, frames=range(0, n_t, 10),
    init_func=init, blit=True, interval=100
)

# Save the animation to an MP4 file
ani.save('wave_propagation.gif', writer='pillow', fps=100)
np.savetxt("wave_state.txt", u)
plt.show()