CODEWITHSUNDEEP
pythonpython-3.xmatplotlibiterationnumerical-methods
mehek
mehek

Reputation: 27

How to turn this plot into an animation? I have a plot but would like to turn this into an animation somehow

I've been struggling for the last few nights how to turn my waves in the graph here into some sort of animation after each time step or after each x step. How can I modify and write the code so that my program animations each time step of the wave somehow. I'm very new to python and programming and never used the animation part of matplotlib before. 

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




T = 20.0 # t
X = 10.0 # x

n = 300
m = 300

#positions of time and space, xp= x position accross grid, tp = t position accross grid
tp = T / (n - 1)
xp = X / (m - 1)

C = 0.5

U = np.zeros((n, m))

# Set the initial values for each x position
for i in range(0, n):
    U[i, 0] = np.exp(- (i * xp - 2.5) ** 2)


for i in range(1, n): # across x
    for j in range(1, m): # across t
        U[i, j] = (xp * U[i, j - 1] + C * tp * U[i - 1, j]) / (xp + C * tp) # equation for new distribution value

fig = plt.figure(1)

#gives time position instead of time step
tn = np.zeros((m, 1))
for j in range(0, m):
    tn[j] = j * tp

#gives x position instead of x step
xn = np.zeros((n, 1))
for j in range(0, n):
    xn[j] = j * xp

for i in [0, 50, 100, 150, 200, 250, 299 ]: # selects which position of time 
    label = 't = ' + str(tn[i][0]) # lables legend
    subfig = fig.add_subplot(1, 1, 1)
    subfig.plot(xn, U[:, i], label=label)
    subfig.legend()
    subfig.grid(True)
    print(tn)

# Save Image
plt.xlabel('x: position')
plt.ylabel('u: u(x, t)')
plt.title(r'$\frac{\partial u}{\partial t} + C \frac{\partial u}{\partial x} = 0$')

plt.savefig('transport-equation')`

Upvotes: 0

Views: 129

Answers (1)

Bobby Ocean
Bobby Ocean

Reputation: 3324

plt is a tricky package to get used to. In general, graphics is not an easy thing to manage and plt tries to be as simple as possible, while also providing max flexibility. In general, when you use plt, there are a lot of global variables that are automatically generated, updated, cleaned, and handled for you. When you use things "plt.xlabel", you are really applying this to a particular axis in a particular figure, which are automatically determined for you. If you want more control in plt and/or you want to do something complicated like an animation, then you need to make your globals explicit. 

#Create xn and U. 

import matplotlib.pyplot as plt

figure = plt.figure()                     #This is the window that pops open. 
axis   = figure.add_subplot(1,1,1)        #This is a graph/grid.  
axis.grid(True)                           #Add a grid to the axis. 
label  = 't = ' + str(tn[i][0])
plots  = axis.plot(xn,U[:,0],label=label) #These are the plots of data with X and Y.

An X and Y arrays can generate more than one plot at a time, hence, plots is a list with one item in it. To get a sense of how this works, you can literally manipulate the data in real time and watch it change in the plt window. 

figure.show()
plots[0].set_ydata(U[:,10])
plots[0].set_ydata(U[:,-1])
# Close the window when done.

To make an animation we need to tell plt to apply an animation to a given figure. plt will then attempt to update the figure and everything that has been attached to it. If you already have the window open, the animation will still get applied and work, but you will also keep whatever was originally plotting in the figure (so we should close the window and re-code the animation). plt does NOT follow convention that executing one line at a time is the same as executing all the lines at once. plt behaves differently before and after a window is opened. 

#Create xn and U. 

import matplotlib.pyplot as plt

figure = plt.figure()
axis   = figure.add_subplot(1,1,1)
axis.grid(True)
label  = 't = ' + str(tn[i][0])
plots  = axis.plot(xn,U[:,0],label=label)

def animate_function(frame):       
    frame %= 300                   #frame is an integer that counts from 0. 
    plots[0].set_ydata(U[:,frame]) #Change which globals you want. 
    return plots                   #Return the changed items so plt knows. 

#Tell plt to apply this animation function to your figure. 
#Tell plt to wait approximately 10ms per frame. 
#Tell plt to only update pixels that actually change (called blit). 
#Save to a global variable so plt doesn't get upset at you. 
ani    = animation.FuncAnimation(figure,animate_function,interval=10,blit=True)

#Now open the window and show the figure. 
figure.show()

Upvotes: 1

Related Questions