[FIXED] How can I add particles to an animated plot after x time?

Issue

Python noob here (week 2!) who has gotten a big headache from doing a lot of research and getting in over my head... so I appreciate your advice in advance.

I have adapted another creator's animated collision plot to drop 10 blue particles in a rectangle. This part works just fine!

"""
Animation of Elastic collisions with Gravity

author: Jake Vanderplas
email: [email protected]
website: http://jakevdp.github.com
license: BSD
Please feel free to use and modify this, but keep the above information. Thanks!
"""
import numpy as np
from scipy.spatial.distance import pdist, squareform

import matplotlib.pyplot as plt
import scipy.integrate as integrate
import matplotlib.animation as animation
from matplotlib.animation import FuncAnimation, PillowWriter  
import time

numballs = 10

class ParticleBox:
    """Orbits class
    
    init_state is an [N x 4] array, where N is the number of particles:
       [[x1, y1, vx1, vy1],
        [x2, y2, vx2, vy2],
        ...               ]

    bounds is the size of the box: [xmin, xmax, ymin, ymax]
    """
    def __init__(self,
                 init_state = [[1, 0, 0, -1],
                               [-0.5, 0.5, 0.5, 0.5],
                               [-0.5, -0.5, -0.5, 0.5]],
                 bounds = [-6, 6, -2.8, 2.8],
                 size = 0.04,
                 M = 0.05,#0.05
                 G = 9.8):
        self.init_state = np.asarray(init_state, dtype=float)
        self.M = M * np.ones(self.init_state.shape[0])
        self.size = size
        self.state = self.init_state.copy()
        self.time_elapsed = 0
        self.bounds = bounds
        self.G = G

    def step(self, dt):
        """step once by dt seconds"""
        self.time_elapsed += dt
        
        # update positions
        self.state[:, :2] += dt * self.state[:, 2:]

        # find pairs of particles undergoing a collision
        D = squareform(pdist(self.state[:, :2]))
        ind1, ind2 = np.where(D < 2 * self.size)
        unique = (ind1 < ind2)
        ind1 = ind1[unique]
        ind2 = ind2[unique]

        # update velocities of colliding pairs
        for i1, i2 in zip(ind1, ind2):
            # mass
            m1 = self.M[i1]
            m2 = self.M[i2]

            # location vector
            r1 = self.state[i1, :2]
            r2 = self.state[i2, :2]

            # velocity vector
            v1 = self.state[i1, 2:]
            v2 = self.state[i2, 2:]

            # relative location & velocity vectors
            r_rel = r1 - r2
            v_rel = v1 - v2

            # momentum vector of the center of mass
            v_cm = (m1 * v1 + m2 * v2) / (m1 + m2)

            # collisions of spheres reflect v_rel over r_rel
            rr_rel = np.dot(r_rel, r_rel)
            vr_rel = np.dot(v_rel, r_rel)
            v_rel = 2 * r_rel * vr_rel / rr_rel - v_rel

            # assign new velocities
            self.state[i1, 2:] = v_cm + v_rel * m2 / (m1 + m2)
            self.state[i2, 2:] = v_cm - v_rel * m1 / (m1 + m2) 

        # check for crossing boundary
        crossed_x1 = (self.state[:, 0] < self.bounds[0] + self.size)
        crossed_x2 = (self.state[:, 0] > self.bounds[1] - self.size)
        crossed_y1 = (self.state[:, 1] < self.bounds[2] + self.size)
        crossed_y2 = (self.state[:, 1] > self.bounds[3] - self.size)

        self.state[crossed_x1, 0] = self.bounds[0] + self.size
        self.state[crossed_x2, 0] = self.bounds[1] - self.size

        self.state[crossed_y1, 1] = self.bounds[2] + self.size
        self.state[crossed_y2, 1] = self.bounds[3] - self.size

        self.state[crossed_x1 | crossed_x2, 2] *= -1
        self.state[crossed_y1 | crossed_y2, 3] *= -1

        # add gravity
        self.state[:, 3] -= self.M * self.G * dt


#------------------------------------------------------------
# set up initial state
np.random.seed(0)
init_state = -0.5 + np.random.random((numballs, 4))
init_state[:, :2] *= 3.9

box = ParticleBox(init_state, size=0.01) #.004
dt = 1. / 30 # 30fps


#------------------------------------------------------------
# set up figure and animation
fig = plt.figure(facecolor = "none")
fig.subplots_adjust(left=0, right=1, bottom=0, top=1)
ax = fig.add_subplot(111, aspect='equal', autoscale_on=False, facecolor="none",
                     xlim=(-5.95, 5.95), ylim=(-2.8, 2.8))

# particles holds the locations of the particles
particles, = ax.plot([], [], 'bo', ms=6) #6

# rect is the box edge
rect = plt.Rectangle(box.bounds[::2],
                     box.bounds[1] - box.bounds[0],
                     box.bounds[3] - box.bounds[2],
                     ec='none', lw=2, fc='none')
ax.add_patch(rect)

def init():
    """initialize animation"""
    global box, rect
    particles.set_data([], [])
    rect.set_edgecolor('none')
    return particles, rect

def animate(i):
    """perform animation step"""
    global box, rect, dt, ax, fig
    box.step(dt)

    ms = int(fig.dpi * 2 * box.size * fig.get_figwidth()
             / np.diff(ax.get_xbound())[0])
    
    # update pieces of the animation
    rect.set_edgecolor('k')
    particles.set_data(box.state[:, 0], box.state[:, 1])
    particles.set_markersize(ms)
    return particles, rect


ani = animation.FuncAnimation(fig, animate, frames=60,
                              interval=10, blit=True, init_func=init) #frames = 600 ideally

plt.show()

I would like to add 20 red particles to drop in after 5 seconds (without the original 10 balls resetting to time 0). I have tried using the above code, then adding

time.sleep(5)
numballs = 30

and then pasting in all of the first block of code again, hoping (by a long-shot) that this would add 20 particles to the animation after 5 seconds. Not surprisingly, this did nothing.

How can I update this animation to show this change?

Also, since you're here :) I would love to know how to make the background of this plot transparent and how to make the balls drop initially from a different location (say, a circle in the top center instead of a big square). I don't need the exact code to those changes as they are not as important, but a hint as to where to look would really save me some advil! I've been able to change other features of the initial setup, but for some reason this evades me.

Thank you again, from a python rookie

Solution

Hey I know you are new so I will try to break down what I did as much as possible to do some of what you want.

Let's just try to add 20 balls after 5 seconds.

A nice place to start is to ask "how would it look if this was easy?". We can see in the code that this animation thing works by running an animate function very frequently.

So the easiest thing would be if we could call

box.add_balls(numballs=20, color='red')

inside this animate function. The problem with this is that we only want it to happen once. We don't want to add 20 balls on every animation frame. So we need the idea of some runtime_actions that happen only once, when a certain amount of time has elapsed.

You will want to pass these runtime_actions to the box when you instantiate it, like this:

box = ParticleBox(
    set_up_initial_state(),
    size=0.01,
    runtime_actions={5: ["add_balls", 20, "red"]}
)

I have chosen to model the runtime_actions as a dict, with time_threshold key and "method name with args" value. That might not make any sense but I will explain.

Inside the step method, you should look at the runtime_actions dictionary and see if any of the keys are now less than the total time elapsed (self.time_elapsed). If they are, then use their value. This value is a list of strings with first value the name of the method (add_balls) and the other values are arguments to the method.

This approach is more general than just adding balls. You could add any other methods to the ParticleBox class and call them in a similar way by just adding new key-value pairs to your runtime_actions dictionary.

So we just define this add_balls method on our ParticleBox class:

def add_balls(self, count=0, color="blue"):
        np.random.seed(0)
        append_state = -0.5 + np.random.random((count, 4))
        append_state[:, :2] *= 3.9
        append_array = np.asarray(append_state, dtype=float)
        self.state = (np.concatenate((self.state, append_array))).copy()
        self.M = self.base_M * np.ones(self.state.shape[0])

And there is one more subtle thing to note. The M property needs to be done in a slightly different way. We need to update self.M when the number of balls changes so we should store the initial M passed to the class. I have chosen to call this self.base_M.

Putting all this together, here is the working code that makes 20 balls appear after 5 seconds (the balls are still blue though):

"""
Animation of Elastic collisions with Gravity

author: Jake Vanderplas
email: [email protected]
website: http://jakevdp.github.com
license: BSD
Please feel free to use and modify this, but keep the above information. Thanks!
"""
import numpy as np
from scipy.spatial.distance import pdist, squareform

import matplotlib.pyplot as plt
import scipy.integrate as integrate
import matplotlib.animation as animation
from matplotlib.animation import FuncAnimation, PillowWriter
import time

numballs = 10


class ParticleBox:
    """Orbits class

    init_state is an [N x 4] array, where N is the number of particles:
       [[x1, y1, vx1, vy1],
        [x2, y2, vx2, vy2],
        ...               ]

    bounds is the size of the box: [xmin, xmax, ymin, ymax]
    """

    def __init__(
        self,
        init_state=[[1, 0, 0, -1], [-0.5, 0.5, 0.5, 0.5], [-0.5, -0.5, -0.5, 0.5]],
        bounds=[-6, 6, -2.8, 2.8],
        size=0.04,
        M=0.05,  # 0.05
        G=9.8,
        runtime_actions=None,
    ):
        self.runtime_actions = runtime_actions
        self.base_M=M
        self.init_state = np.asarray(init_state, dtype=float)
        self.M = self.base_M * np.ones(self.init_state.shape[0])
        self.size = size
        self.state = self.init_state.copy()
        self.time_elapsed = 0
        self.bounds = bounds
        self.G = G

    def add_balls(self, count=0, color="blue"):
        np.random.seed(0)
        append_state = -0.5 + np.random.random((count, 4))
        append_state[:, :2] *= 3.9
        append_array = np.asarray(append_state, dtype=float)
        self.state = (np.concatenate((self.state, append_array))).copy()
        self.M = self.base_M * np.ones(self.state.shape[0])

    def step(self, dt):
        """step once by dt seconds"""
        self.time_elapsed += dt

        print(self.time_elapsed)
        if self.runtime_actions:
            time_thresholds_to_delete = []
            for time_threshold, action in self.runtime_actions.items():
                if time_threshold < self.time_elapsed:
                    # this next line does `self.add_balls(numballs=20, color='red')`
                    getattr(self, action[0])(*action[1:]).
                    time_thresholds_to_delete.append(time_threshold)
            for time_threshold_to_delete in time_thresholds_to_delete:
                del self.runtime_actions[time_threshold_to_delete]


        # update positions
        self.state[:, :2] += dt * self.state[:, 2:]

        # find pairs of particles undergoing a collision
        D = squareform(pdist(self.state[:, :2]))
        ind1, ind2 = np.where(D < 2 * self.size)
        unique = ind1 < ind2
        ind1 = ind1[unique]
        ind2 = ind2[unique]

        # update velocities of colliding pairs
        for i1, i2 in zip(ind1, ind2):
            # mass
            m1 = self.M[i1]
            m2 = self.M[i2]

            # location vector
            r1 = self.state[i1, :2]
            r2 = self.state[i2, :2]

            # velocity vector
            v1 = self.state[i1, 2:]
            v2 = self.state[i2, 2:]

            # relative location & velocity vectors
            r_rel = r1 - r2
            v_rel = v1 - v2

            # momentum vector of the center of mass
            v_cm = (m1 * v1 + m2 * v2) / (m1 + m2)

            # collisions of spheres reflect v_rel over r_rel
            rr_rel = np.dot(r_rel, r_rel)
            vr_rel = np.dot(v_rel, r_rel)
            v_rel = 2 * r_rel * vr_rel / rr_rel - v_rel

            # assign new velocities
            self.state[i1, 2:] = v_cm + v_rel * m2 / (m1 + m2)
            self.state[i2, 2:] = v_cm - v_rel * m1 / (m1 + m2)

        # check for crossing boundary
        crossed_x1 = self.state[:, 0] < self.bounds[0] + self.size
        crossed_x2 = self.state[:, 0] > self.bounds[1] - self.size
        crossed_y1 = self.state[:, 1] < self.bounds[2] + self.size
        crossed_y2 = self.state[:, 1] > self.bounds[3] - self.size

        self.state[crossed_x1, 0] = self.bounds[0] + self.size
        self.state[crossed_x2, 0] = self.bounds[1] - self.size

        self.state[crossed_y1, 1] = self.bounds[2] + self.size
        self.state[crossed_y2, 1] = self.bounds[3] - self.size

        self.state[crossed_x1 | crossed_x2, 2] *= -1
        self.state[crossed_y1 | crossed_y2, 3] *= -1

        # add gravity
        self.state[:, 3] -= self.M * self.G * dt


# ------------------------------------------------------------
# set up initial state
np.random.seed(0)
init_state = -0.5 + np.random.random((numballs, 4))
init_state[:, :2] *= 3.9

box = ParticleBox(
    init_state,
    size=0.01,
    runtime_actions={5: ["add_balls", 20, "red"]}
)
dt = 1.0 / 30  # 30fps


# ------------------------------------------------------------
# set up figure and animation
fig = plt.figure(facecolor="none")
fig.subplots_adjust(left=0, right=1, bottom=0, top=1)
ax = fig.add_subplot(
    111,
    aspect="equal",
    autoscale_on=False,
    facecolor="none",
    xlim=(-5.95, 5.95),
    ylim=(-2.8, 2.8),
)

# particles holds the locations of the particles
(particles,) = ax.plot([], [], "bo", ms=6)  # 6

# rect is the box edge
rect = plt.Rectangle(
    box.bounds[::2],
    box.bounds[1] - box.bounds[0],
    box.bounds[3] - box.bounds[2],
    ec="none",
    lw=2,
    fc="none",
)
ax.add_patch(rect)


def init():
    """initialize animation"""
    global box, rect
    particles.set_data([], [])
    rect.set_edgecolor("none")
    return particles, rect


def animate(i):
    """perform animation step"""
    global box, rect, dt, ax, fig
    box.step(dt)

    ms = int(fig.dpi * 2 * box.size * fig.get_figwidth() / np.diff(ax.get_xbound())[0])

    # update pieces of the animation
    rect.set_edgecolor("k")
    particles.set_data(box.state[:, 0], box.state[:, 1])
    particles.set_markersize(ms)
    return particles, rect


ani = animation.FuncAnimation(
    fig, animate, frames=60, interval=10, blit=True, init_func=init
)  # frames = 600 ideally

plt.show()

Notice that I am not yet using the color argument passed to add_balls. This might require rewriting more bits of the code since we can do

particles.set_color("red")

in the animate function but that will affect all the particles.

Here is a "hacky" way of splitting the box data into 2 separate particles lines (I call the second one particles2). It's not good because it's so specific but it shows what needs to be done (i.e. using separate line objects for the different groups of particles).

"""
Animation of Elastic collisions with Gravity

author: Jake Vanderplas
email: [email protected]
website: http://jakevdp.github.com
license: BSD
Please feel free to use and modify this, but keep the above information. Thanks!
"""
import numpy as np
from scipy.spatial.distance import pdist, squareform

import matplotlib.pyplot as plt
import scipy.integrate as integrate
import matplotlib.animation as animation
from matplotlib.animation import FuncAnimation, PillowWriter
import time

numballs = 10


class ParticleBox:
    """Orbits class

    init_state is an [N x 4] array, where N is the number of particles:
       [[x1, y1, vx1, vy1],
        [x2, y2, vx2, vy2],
        ...               ]

    bounds is the size of the box: [xmin, xmax, ymin, ymax]
    """

    def __init__(
        self,
        init_state=[[1, 0, 0, -1], [-0.5, 0.5, 0.5, 0.5], [-0.5, -0.5, -0.5, 0.5]],
        bounds=[-6, 6, -2.8, 2.8],
        size=0.04,
        M=0.05,  # 0.05
        G=9.8,
        runtime_actions=None,
    ):
        self.runtime_actions = runtime_actions
        self.base_M=M
        self.init_state = np.asarray(init_state, dtype=float)
        self.M = self.base_M * np.ones(self.init_state.shape[0])
        self.size = size
        self.state = self.init_state.copy()
        self.time_elapsed = 0
        self.bounds = bounds
        self.G = G

    def add_balls(self, count=0, color="blue"):
        np.random.seed(0)
        append_state = -0.5 + np.random.random((count, 4))
        append_state[:, :2] *= 3.9
        append_array = np.asarray(append_state, dtype=float)
        self.state = (np.concatenate((self.state, append_array))).copy()
        self.M = self.base_M * np.ones(self.state.shape[0])

    def step(self, dt):
        """step once by dt seconds"""
        self.time_elapsed += dt

        print(self.time_elapsed)
        if self.runtime_actions:
            time_thresholds_to_delete = []
            for time_threshold, action in self.runtime_actions.items():
                if time_threshold < self.time_elapsed:
                    # this next line does `self.add_balls(numballs=20, color='red')`
                    getattr(self, action[0])(*action[1:])
                    time_thresholds_to_delete.append(time_threshold)
            for time_threshold_to_delete in time_thresholds_to_delete:
                del self.runtime_actions[time_threshold_to_delete]

        # update positions
        self.state[:, :2] += dt * self.state[:, 2:]

        # find pairs of particles undergoing a collision
        D = squareform(pdist(self.state[:, :2]))
        ind1, ind2 = np.where(D < 2 * self.size)
        unique = ind1 < ind2
        ind1 = ind1[unique]
        ind2 = ind2[unique]

        # update velocities of colliding pairs
        for i1, i2 in zip(ind1, ind2):
            # mass
            m1 = self.M[i1]
            m2 = self.M[i2]

            # location vector
            r1 = self.state[i1, :2]
            r2 = self.state[i2, :2]

            # velocity vector
            v1 = self.state[i1, 2:]
            v2 = self.state[i2, 2:]

            # relative location & velocity vectors
            r_rel = r1 - r2
            v_rel = v1 - v2

            # momentum vector of the center of mass
            v_cm = (m1 * v1 + m2 * v2) / (m1 + m2)

            # collisions of spheres reflect v_rel over r_rel
            rr_rel = np.dot(r_rel, r_rel)
            vr_rel = np.dot(v_rel, r_rel)
            v_rel = 2 * r_rel * vr_rel / rr_rel - v_rel

            # assign new velocities
            self.state[i1, 2:] = v_cm + v_rel * m2 / (m1 + m2)
            self.state[i2, 2:] = v_cm - v_rel * m1 / (m1 + m2)

        # check for crossing boundary
        crossed_x1 = self.state[:, 0] < self.bounds[0] + self.size
        crossed_x2 = self.state[:, 0] > self.bounds[1] - self.size
        crossed_y1 = self.state[:, 1] < self.bounds[2] + self.size
        crossed_y2 = self.state[:, 1] > self.bounds[3] - self.size

        self.state[crossed_x1, 0] = self.bounds[0] + self.size
        self.state[crossed_x2, 0] = self.bounds[1] - self.size

        self.state[crossed_y1, 1] = self.bounds[2] + self.size
        self.state[crossed_y2, 1] = self.bounds[3] - self.size

        self.state[crossed_x1 | crossed_x2, 2] *= -1
        self.state[crossed_y1 | crossed_y2, 3] *= -1

        # add gravity
        self.state[:, 3] -= self.M * self.G * dt


# ------------------------------------------------------------
# set up initial state
np.random.seed(0)
init_state = -0.5 + np.random.random((numballs, 4))
init_state[:, :2] *= 3.9

box = ParticleBox(
    init_state,
    size=0.01,
    runtime_actions={5: ["add_balls", 20, "red"]}
)
dt = 1.0 / 30  # 30fps


# ------------------------------------------------------------
# set up figure and animation
fig = plt.figure(facecolor="none")
fig.subplots_adjust(left=0, right=1, bottom=0, top=1)
ax = fig.add_subplot(
    111,
    aspect="equal",
    autoscale_on=False,
    facecolor="none",
    xlim=(-5.95, 5.95),
    ylim=(-2.8, 2.8),
)

# particles holds the locations of the particles
(particles,) = ax.plot([], [], "bo", ms=6)  # 6
(particles2,) = ax.plot([], [], "bo", ms=6)  # 6

# rect is the box edge
rect = plt.Rectangle(
    box.bounds[::2],
    box.bounds[1] - box.bounds[0],
    box.bounds[3] - box.bounds[2],
    ec="none",
    lw=2,
    fc="none",
)
ax.add_patch(rect)


def init():
    """initialize animation"""
    global box, rect
    particles.set_data([], [])
    rect.set_edgecolor("none")
    return particles, rect


def animate(i):
    """perform animation step"""
    global box, rect, dt, ax, fig
    box.step(dt)

    ms = int(fig.dpi * 2 * box.size * fig.get_figwidth() / np.diff(ax.get_xbound())[0])

    # update pieces of the animation
    rect.set_edgecolor("k")
    rect.set_color("none")
    particles.set_data(box.state[:numballs, 0], box.state[:numballs, 1])
    particles2.set_data(box.state[numballs:, 0], box.state[numballs:, 1])
    particles.set_markersize(ms)
    particles2.set_markersize(ms)
    particles2.set_color("red")
    return particles, particles2, rect


ani = animation.FuncAnimation(
    fig, animate, frames=60, interval=10, blit=True, init_func=init
)  # frames = 600 ideally

plt.show()

Answered By - Peaceful James