Created
April 5, 2020 23:38
-
-
Save AndrewILWilliams/975221ba3d235252fabecd6342fb2660 to your computer and use it in GitHub Desktop.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| from matplotlib.animation import FuncAnimation | |
| # 22/04/2019: Script to plot animation of | |
| # balls falling and bouncing under gravity (in 1-D, for now) | |
| # Set up the plotting area | |
| fig, ax = plt.subplots() | |
| ax.set(xlim=(0, 10), ylim=(0, 10)) | |
| # set up | |
| n = 20 # number of balls | |
| g = -10 # constant acc. due to gravity | |
| dt = 0.005 # time step per frame | |
| pos = (5 * np.random.sample(n*2)+2).reshape(n, 2) | |
| vel = np.zeros(n) # have balls falling from rest for now | |
| sizes = 100 * np.random.sample(n) + 100 | |
| # colors -- each row is (red, green, blue, alpha). Each can go | |
| # from 0 to 1, alpha is the transparency. | |
| colors = np.random.sample([n, 4]) | |
| # Plot the first line | |
| circles = plt.scatter(pos[:,0], pos[:,1], marker='o', s=sizes, c=colors) | |
| # Create a function to update the line | |
| def animate(i): | |
| global pos, vel | |
| vel += g * dt | |
| pos[:, 1] += vel * dt | |
| bounce = pos[:, 1] <= 0 # Find balls that are at floor | |
| vel[bounce] = - vel[bounce] # Bounce if hit floor | |
| pos[:, 1][bounce] = 0 # Stop the balls sinking below floor... | |
| circles.set_offsets(pos) # Change the positions | |
| # Call FuncAnimation and show | |
| anim = FuncAnimation(fig, animate, frames=10**3, interval = 10) | |
| #plt.draw() | |
| #plt.show() | |
| anim.save('filename.mp4') |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment