|
1 | | -# ----------------------------------------------------------------------------- |
2 | | -# Rain simulation |
3 | | -# Author: Nicolas P. Rougier |
4 | | -# ----------------------------------------------------------------------------- |
5 | | -import matplotlib |
| 1 | +""" |
| 2 | +Rain simulation |
| 3 | +
|
| 4 | +Simulates rain drops on a surface by animating the scale and opacity |
| 5 | +of 50 scatter points. |
| 6 | +
|
| 7 | +Author: Nicolas P. Rougier |
| 8 | +""" |
6 | 9 | import numpy as np |
7 | 10 | import matplotlib.pyplot as plt |
8 | 11 | from matplotlib.animation import FuncAnimation |
9 | 12 |
|
10 | 13 |
|
11 | | -# Create new figure |
| 14 | +# Create new Figure and an Axes which fills it. |
12 | 15 | fig = plt.figure(figsize=(7,7)) |
13 | | -ax = fig.add_axes([0,0,1,1], frameon=False, aspect=1) |
14 | | - |
15 | | -# Create rain data |
16 | | -P = np.zeros(50, dtype=[('position', float, 2), |
17 | | - ('size', float, 1), |
18 | | - ('growth', float, 1), |
19 | | - ('color', float, 4)]) |
20 | | - |
21 | | -# Scatter plot is used to animate rain drops |
22 | | -scat = ax.scatter(P['position'][:,0], P['position'][:,1], P['size'], |
23 | | - lw=0.5, edgecolors = P['color'], facecolors='none') |
| 16 | +ax = fig.add_axes([0, 0, 1, 1], frameon=False) |
24 | 17 | ax.set_xlim(0,1), ax.set_xticks([]) |
25 | 18 | ax.set_ylim(0,1), ax.set_yticks([]) |
26 | 19 |
|
| 20 | +# Create rain data |
| 21 | +n_drops = 50 |
| 22 | +rain_drops = np.zeros(n_drops, dtype=[('position', float, 2), |
| 23 | + ('size', float, 1), |
| 24 | + ('growth', float, 1), |
| 25 | + ('color', float, 4)]) |
27 | 26 |
|
28 | | -def update(frame): |
29 | | - i = frame % len(P) |
30 | | - |
31 | | - # Make all colors more transparent |
32 | | - P['color'][:,3] -= 1.0/len(P) |
33 | | - P['color'][:,3] = np.clip(P['color'][:,3], 0, 1) |
| 27 | +# Initialize the raindrops in random positions and with |
| 28 | +# random growth rates. |
| 29 | +rain_drops['position'] = np.random.uniform(0, 1, (n_drops, 2)) |
| 30 | +rain_drops['growth'] = np.random.uniform(50, 200, n_drops) |
34 | 31 |
|
35 | | - # Make all circles bigger |
36 | | - P['size'] += P['growth'] |
| 32 | +# Construct the scatter which we will update during animation |
| 33 | +# as the raindrops develop. |
| 34 | +scat = ax.scatter(rain_drops['position'][:,0], rain_drops['position'][:,1], |
| 35 | + s=rain_drops['size'], lw=0.5, edgecolors=rain_drops['color'], |
| 36 | + facecolors='none') |
37 | 37 |
|
38 | | - # Pick a new position for oldest rain drop |
39 | | - P['position'][i] = np.random.uniform(0, 1, 2) |
40 | 38 |
|
41 | | - # Reset size |
42 | | - P['size'][i] = 5 |
| 39 | +def update(frame_number): |
| 40 | + # Get an index which we can use to re-spawn the oldest raindrop. |
| 41 | + current_index = frame_number % n_drops |
43 | 42 |
|
44 | | - # Reset color |
45 | | - P['color'][i] = (0, 0, 0, 1) |
| 43 | + # Make all colors more transparent as time progresses. |
| 44 | + rain_drops['color'][:, 3] -= 1.0/len(rain_drops) |
| 45 | + rain_drops['color'][:,3] = np.clip(rain_drops['color'][:,3], 0, 1) |
46 | 46 |
|
47 | | - # Choose a random growth factor |
48 | | - P['growth'][i] = np.random.uniform(50, 200) |
| 47 | + # Make all circles bigger. |
| 48 | + rain_drops['size'] += rain_drops['growth'] |
49 | 49 |
|
50 | | - # Update scatter plot |
51 | | - scat.set_edgecolors(P['color']) |
52 | | - scat.set_sizes(P['size']) |
53 | | - scat.set_offsets(P['position']) |
| 50 | + # Pick a new position for oldest rain drop, resetting its size, |
| 51 | + # color and growth factor. |
| 52 | + rain_drops['position'][current_index] = np.random.uniform(0, 1, 2) |
| 53 | + rain_drops['size'][current_index] = 5<
829B
/div> |
| 54 | + rain_drops['color'][current_index] = (0, 0, 0, 1) |
| 55 | + rain_drops['growth'][current_index] = np.random.uniform(50, 200) |
54 | 56 |
|
55 | | - return scat, |
| 57 | + # Update the scatter collection, with the new colors, sizes and positions. |
| 58 | + scat.set_edgecolors(rain_drops['color']) |
| 59 | + scat.set_sizes(rain_drops['size']) |
| 60 | + scat.set_offsets(rain_drops['position']) |
| 61 | + |
56 | 62 |
|
| 63 | +# Construct the animation, using the update function as the animation |
| 64 | +# director. |
57 | 65 | animation = FuncAnimation(fig, update, interval=10) |
58 | 66 | plt.show() |
0 commit comments