dimanche 12 juillet 2020

How to close matplotlib animation window after a given time?

This is an animation I made with matplotlib, where dots are randomly moving in a circle.
I want the animation to stop (and the window to close) after 4 seconds.
I tried to set a timer and put the plotting in a while loop, but it doesn't do the work, it keeps running.

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
from scipy.spatial import distance
import random
import datetime, time

fig, ax = plt.subplots()
ax.set_xlim(0, 20)
ax.set_ylim(0, 20)
then = datetime.datetime.now() + datetime.timedelta(seconds=10)
circle_r=2
def get_initial_coordinates():
    x_coord =[random.uniform(3, 7) for i in range(n_particles)]
    y_coord = [random.uniform(3, 7) for i in range(n_particles)]
    return x_coord, y_coord

def get_initial_velocities():
    x_vel = [3 * (np.random.random() - 0.5) * box_width for i in range(n_particles)]
    y_vel = [3 * (np.random.random() - 0.5) * box_width for i in range(n_particles)]
    return x_vel, y_vel

def take_step(x_coord, y_coord, x_vel, y_vel):
    for i in range(n_particles):
     x_coord[i] += x_vel[i]*dt
     y_coord[i] += y_vel[i]*dt
     if distance.euclidean([5, 5], [x_coord[i],y_coord[i]]) >= 3:
         x_vel[i] = -x_vel[i]
         x_coord[i] += x_vel[i] * dt
         y_vel[i] = -y_vel[i]
         y_coord[i] += y_vel[i] * dt

    return x_coord, y_coord, x_vel, y_vel

n_particles = 40
box_width = 10
n_steps = 5000
dt = 0.001

x_coord, y_coord = get_initial_coordinates()
x_vel, y_vel = get_initial_velocities()
for i in range(n_steps):
    x_coord,y_coord,x_vel,y_vel= take_step(x_coord,y_coord,x_vel,y_vel)
#---------------------
fig = plt.figure()
ax = plt.axes(xlim=(0, 10), ylim=(0, 10))
ax.set_aspect('equal')
d, = ax.plot([x_coord[i] for i in range(n_particles)],
             [y_coord[i] for i in range(n_particles)], 'ro')
x_coord, y_coord = get_initial_coordinates()
x_vel, y_vel = get_initial_velocities()
circle = plt.Circle((5, 5), 3, color='blue', fill=False)
ax.add_artist(circle)
def animate(i):
    take_step(x_coord, y_coord, x_vel, y_vel)
    d.set_data([x_coord[i]for i in range(n_particles)],
               [y_coord[i]for i in range(n_particles)])
    return d,

anim = animation.FuncAnimation(fig, animate, frames=200, interval=20)
while then > datetime.datetime.now():
    plt.show()
    time.sleep(1)
else:
    plt.close()



Aucun commentaire:

Enregistrer un commentaire