hemmelig
hemmelig

Reputation: 350

Python - Animating a numerical solution to the wave equation

I am trying to animate a solution to the wave equation - I am plotting the stress and the displacement against x, but I want it to evolve with time.

I think a solution would be to treat every position x as a single 'particle' and then have that particle obey the function that defines y, and animate that in t? However, the only way I'm seeing to implement this seems a bit too brute force, leading to bulky code for what should really be short and easy.

I intend to extend this code to allow for mu and rho to be variables dependent upon position as opposed to constants, but for now I just want to get the animation working.

Code:

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

# Definition of parameters
dt = 0.04
t = np.arange(0.0,40,dt)
x = np.linspace(0, 15, 1000) # x position
Y0 = np.array([0,10])       # initial conditions
mu = 1.5
rho = 1.2
omega = 1
def dY_dx(Y, t=0):
    """ Return the gradient of y1 and y2"""
    return np.array([Y[1] / mu, - (omega ** 2) * rho * Y[0]])

from scipy import integrate

Y = integrate.odeint(dY_dx, Y0, x)

y1, y2 = Y.T

# set up the figure, the axis, and the plot element we want to animate
fig = plt.figure()
ax = fig.add_subplot(111, autoscale_on=False, xlim=(0,15), ylim=(-10,10))
ax.grid()

line1, = ax.plot([], [], 'o', ms=2)
line2, = ax.plot([], [], '-', lw=2)
time_template = 'time = %.lfs'
time_text = ax.text(0.05, 0.9, '', transform=ax.transAxes)

# initialisation function: plot the background of each frame
def init():
    line1.set_data([], [])
    line2.set_data([], [])
    time_text.set_text('')
    return line1, line2, time_text

# animation function: this is called sequentially
def animate(i):

    line1.set_data(x, y1)
    line2.set_data(x, y2)

    time_text.set_text(time_template%(i*dt))
    return line1, line2, time_text

# call the animator. blit=True means only re-draw the parts that have changed
ani = animation.FuncAnimation(fig, animate, np.arange(1, len(Y)), interval=25, blit=True, init_func=init)

#ani.save('waveEquation.mp4', fps=15)
plt.show()

I have tried:

# animation function: this is called sequentially
def animate(i):

    line1.set_data(x[i], y1[i])

and

# animation function: this is called sequentially
def animate(i):

    line1.set_data(x, y1[i])

but neither give the result I desire.

Could the issue be that I am integrating my solutions over t before they are plotted, and then not including t as a variable in my animation?

I know I can use code (found on a similar question here):

def animate(i):
    thisx = x
    thisy = np.sin(2 * np.pi * (x - 0.01 * i))

    line1.set_data(thisx, thisy)

to animate a sin wave, but I don't want to be using an analytic solution to calculate y, I want to do it numerically (for later problems).

Upvotes: 4

Views: 2254

Answers (1)

rjonnal
rjonnal

Reputation: 1207

ConorB,

You asked:

Could the issue be that I am integrating my solutions over t before they are plotted, and then not including t as a variable in my animation?

Yes, that is precisely your issue. The animate definition should include updating of t, and then--presumably--should require re-integration and recalculation of the lines to plot, y1 and y2. It's hard to know exactly what you want your plot to show, without additional knowledge of the general problem domain, but your animate function should look something like this:

def animate(i):
    t = np.arange(0.0,40,dt) + i*dt
    Y = integrate.odeint(dY_dx,Y0,t)
    y1,y2 = Y.T

    line1.set_data(x, y1)
    line2.set_data(x, y2)

    time_text.set_text(time_template%t)
    return line1, line2, time_text

Also, your code wouldn't run, due to invocations of array instead of np.array and some errant whitespace. I cleaned it up, pending peer-review.

Upvotes: 4

Related Questions