Python Animation of projectile giving straight line

Question

The goal is to animate/simulate a sprinkler. The basic idea is to create multiple droplets at an instance, each exiting the sprinkler at a different angle. However, instead of having a point that moves along, I'm getting a straight, static line. I've tried the same code with only one or 2 drops/points, but I still get the same straight line. At first I thought it was because I included drag into the animation (in previous trials I got strange/unexpected results when I included the drag, so I was hoping it was the same problem)

Below is the code I wrote.

EDIT: dragx and dragy are no longer set to the same values. The initial velocities are calculated outside the update function.

Drag forces are recalculated at each index increase, since the drag depends on the velocity.

The imports are also now included.

I have also now tried to use a scatter plot instead of a line, which gave me a single dot at the origin. I reverted back to the line since then

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from mpl_toolkits.mplot3d import Axes3D 
import matplotlib.animation as animation
from math import sin, cos
import matplotlib.axes as axess

#pressure in system
def pressure(pump):
    return 10.197*pump*10**5/998.2

#Initial velocity based on the pressure of the system
def vinit(init_head):
    return (2*9.81*init_head)**0.5
# Number of droplets per iteration
n = 2
v0 = vinit(pressure(5.8))
cd = 0.5                                             #drag coefficient of a sphere;   
rho = 1.225                                          #density of air;                 kg/m^3

fps = 20
runtime = 1*20
numframes = fps*runtime
time = np.linspace(0,runtime,numframes)

#Initialize droplets

droplets = np.zeros(numframes, dtype = [('position', float, 2),
                                ('velocity', float, 2),
                                ('force', float, 2),
                                ('angle',float, 1),
                                ('radius', float, 1)])

droplets['radius'] = np.random.normal(0.0045, 0.001, numframes)
A = 4*np.pi*droplets['radius']**2
mass = ((4*np.pi*droplets['radius']**3)/3)*1000     #mass of waterdroplet
drag = np.zeros(numframes,dtype = [('dragx', float, 1),('dragy', float, 1)])

rads = np.radians(np.random.normal(37,1,numframes))
droplets['angle'] = rads
droplets['force'][:,1] = -9.8*mass

#Initialize velocity
for i in range(0, numframes):
    droplets['velocity'][i,0] = v0*cos(droplets['angle'][i])
    droplets['velocity'][i,1] = v0*sin(droplets['angle'][i])    

#Initialize the drag
drag['dragx'] = -0.5*rho*cd*A*droplets['velocity'][:,0]
drag['dragy'] = -0.5*rho*cd*A*droplets['velocity'][:,1]


droplets['position'][:,0] = 0


#Initialize the figure
fig,ax = plt.subplots()
axess.Axes.set_frame_on(ax,True)
ax.autoscale(True)
ax.set_xlim(0)
ax.set_ylim(0)

line = ax.plot(droplets['position'][0, 0], droplets['position'][0, 1],'b.')[0]
line = ax.plot(droplets['position'][:, 0], droplets['position'][:, 1],'b.')[0]
xdata = [droplets['position'][0, 0]]
ydata = [droplets['position'][0, 1]]
ln = plt.plot([],[],'b')[0]

def initfunc():
    droplets['position'][0,:] = 0
    ln.set_data(droplets['position'][0, 0], droplets['position'][0, 1])
    return ln

def update(framenum):
    index = framenum
    cd = 0.5                                          #drag coefficient of a sphere;   
    rho = 1.225                                       #density of air;                 kg/m^3
    A = 4*np.pi*droplets['radius']**2                 #surface area of droplet;        m^2
    mass = ((4*np.pi*droplets['radius']**3)/3)*1000   #mass of waterdroplet

    #Update the drag force on the droplet
    drag['dragx'] = -0.5*rho*cd*A*droplets['velocity'][:,0]
    drag['dragy'] = -0.5*rho*cd*A*droplets['velocity'][:,1]

    droplets['force'][index,0] += drag['dragx'][index]
    droplets['force'][index,1] += drag['dragy'][index]
    #droplets['position'][0] = [0,0]
    droplets['velocity'][index,0] = droplets['velocity'][index,0] + (droplets['force'][index,0]/mass[index])*index
    droplets['velocity'][index,1] = droplets['velocity'][index,1] + (droplets['force'][index,1]/mass[index])*index
    droplets['position'][index,0] = droplets['position'][index,0] + (droplets['velocity'][index,0])*index  
    droplets['position'][index,1] = droplets['position'][index,1] + (droplets['velocity'][index,1])*index  

    xdata.append(droplets['position'][index,0])
    ydata.append(droplets['position'][index,1])
    ln.set_data(xdata,ydata)

    line.set_data(droplets['position'][:,0], droplets['position'][:,1])

    return ln



sprink = animation.FuncAnimation(fig, update,init_func = initfunc,interval= 200,  frames = numframes)
plt.show()

sprink.save('Sprinkler.mp4', fps = fps)

Answer 1

In addition to the mistakes pointed out by @JohanC, you are using the surface area (4 * pi * r^2) instead of the cross-sectional area (pi * r^2) of the water droplets (approximated as spheres). Furthermore, you appear to be trying to using equations of motion which assume a vacuum, those equations break down when air resistance or other nonconstant forces are added.

I would recommend you read this answer for a more complete explanation of calculating ballistic projectile motion with air resistance, but here is a simplified version.

Calculating the position of a particle moving through a medium is essentially solving a first order ODE, there are many ways to solve such an equation numerically. The simplest (albeit least robust) way is to break the problem into small segments and approximate the change from the current point to the next point. For this problem such a method works quite well, assuming you use a small enough time step.

In your case we could do something like this

import matplotlib
import matplotlib.pyplot as plt
from matplotlib import animation
import numpy as np
from math import sin, cos

g = 9.81
rho = 1.225
v0 = 50.0    # This can be initialized as desired, I chose 50 m/s for demonstration
C = 0.5

nframes = 100
ndrops = 10
tend = 8.0
time = np.linspace(0.1, tend, nframes)   # Runtime of 4s chosen arbitrarily
dt = time[1] - time[0]                    # change in time between steps
maxs = [0.0, 0.0]
r = [0.001, 0.0045]                       # This stores the range of radii to be used

# I have implemented a class to store the drops for brevity
class drop:

    def __init__(self, pos, vel, r):
        self.pos = pos
        self.vel = vel
        self.r = r

class sprinkler:
# I implemented this as a class to simplify updating the scatter plot
    def __init__(self):
        self.fig, self.ax = plt.subplots()
        self.drops = [None]*ndrops
        self.maxt = 0.0
        ## This step is simply to figure out how far the water drops can travel
        ## in order to set the bounds of the plot accordingly
        angles = np.linspace(0, np.pi/2, 90)    # We only need to sample from 0 to pi/2
        for angle in angles:
            m = [drop([0.0, 0.1], [v0*cos(angle), v0*sin(angle)], 0.001),
                 drop([0.0, 0.1], [v0*cos(angle), v0*sin(angle)], 0.0045)]
            for d in m:
                t = 0.0
                coef = - 0.5*C*np.pi*d.r**2*rho
                mass = 4/3*np.pi*d.r**3*1000
                while d.pos[1] > 0:
                    a = np.power(d.vel, 2) * coef * np.sign(d.vel)/mass
                    a[1] -= g
                    d.pos += (d.vel + a * dt) * dt
                    d.vel += a * dt
                    t += dt
                    if d.pos[1] > maxs[1]:
                        maxs[1] = d.pos[1]
                    if d.pos[0] > maxs[0]:
                        maxs[0] = d.pos[0]
                    if d.pos[1] < 0.0:
                        if t > self.maxt:
                            self.maxt = t
                        break
        for ii in range(ndrops):    # Make some randomly distributed water drops
            self.drops[ii] = drop([0.0, 0.0], [cos(np.random.random()*np.pi)*v0,
                             sin(np.random.random()*np.pi)*v0],
                             np.random.random()*(r[1]-r[0]) + r[0])

        anim = animation.FuncAnimation(self.fig, self.update, init_func=self.setup,
                                       interval=200, frames=nframes)
        anim.save('Sprinkler.gif', fps = 20, writer='imagemagick')

    def setup(self):
        self.scat = self.ax.scatter([d.pos[0] for d in self.drops],
                               [d.pos[1] for d in self.drops], marker='.', color='k')
        self.ax.set_xlim([-maxs[0]*1.15, maxs[0]*1.15])
        self.ax.set_ylim([0, maxs[1]*1.15])

        return self.scat,

    def update(self,frame):  # Use set_offsets to move the water drops
        self.step()          # Advance to the next 'step'
        self.scat.set_offsets(np.stack([[d.pos[0] for d in self.drops],
                                       [d.pos[1] for d in self.drops]], 1))

        return self.scat,

    def step(self):
        for ii in range(ndrops):
            coef = - 0.5*C*np.pi*self.drops[ii].r**2*rho    # Aggregated coefficient
            mass = 4/3*np.pi*self.drops[ii].r**3*1000
            a = np.power(self.drops[ii].vel, 2)*coef * np.sign(self.drops[ii].vel)/mass
            a[1] -= g
            # Approximate how much the position and velocity would change if we assume
            # a(t) does not change between t and t+dt
            self.drops[ii].pos += np.array(self.drops[ii].vel) * dt + 0.5 * a * dt**2
            self.drops[ii].vel += a * dt
            if self.drops[ii].pos[1] < 0.0:     # Check if the drop has hit the ground
                self.drops[ii].pos[1] = 0.0
                self.drops[ii].vel = [0.0, 0.0]

sprinkler()

This may not produce exactly what you are after but it correctly models the motion of the water droplets.

Implementing the drops as a class and storing them in a list as I have makes it easy to create more droplets to simulate many 'iterations' of the above by simply modifying the update function (create method added for simplicity):

    def update(self,frame):  # Use set_offsets to move the water drops
        # Create between 0 and ndrops new drops, but only until tend-maxt
        if time[frame] < tend - self.maxt*1.1:
            self.create(np.random.randint(0, ndrops))
        self.step()          # Advance to the next 'step'
        self.scat.set_offsets(np.stack([[d.pos[0] for d in self.drops],
                                        [d.pos[1] for d in self.drops]], 1))

        return self.scat,

    def create(self, i):
        for ii in range(i):
            self.drops.append(drop([0.0, 0.0], [cos(np.random.random()*np.pi)*v0,
                             sin(np.random.random()*np.pi)*v0],
                             np.random.random()*(r[1]-r[0]) + r[0]))

Edit - The entire code to produce the second animation

import matplotlib
import matplotlib.pyplot as plt
from matplotlib import animation
import numpy as np
from math import sin, cos

g = 9.81
rho = 1.225
v0 = 50.0    # This can be initialized as desired, I chose 50 m/s for demonstration
C = 0.5

nframes = 100
ndrops = 10
tend = 8.0
time = np.linspace(0.1, tend, nframes)   # Runtime of 4s chosen arbitrarily
dt = time[1] - time[0]                    # change in time between steps
maxs = [0.0, 0.0]
r = [0.001, 0.0045]                       # This stores the range of radii to be used

# I have implemented a class to store the drops for brevity
class drop:

    def __init__(self, pos, vel, r):
        self.pos = pos
        self.vel = vel
        self.r = r

class sprinkler:
# I implemented this as a class to simplify updating the scatter plot
    def __init__(self):
        self.fig, self.ax = plt.subplots()
        self.drops = [None]*ndrops
        self.maxt = 0.0
        ## This step is simply to figure out how far the water drops can travel
        ## in order to set the bounds of the plot accordingly
        angles = np.linspace(0, np.pi/2, 90)    # We only need to sample from 0 to pi/2
        for angle in angles:
            m = [drop([0.0, 0.1], [v0*cos(angle), v0*sin(angle)], 0.001),
                 drop([0.0, 0.1], [v0*cos(angle), v0*sin(angle)], 0.0045)]
            for d in m:
                t = 0.0
                coef = - 0.5*C*np.pi*d.r**2*rho
                mass = 4/3*np.pi*d.r**3*1000
                while d.pos[1] > 0:
                    a = np.power(d.vel, 2) * coef * np.sign(d.vel)/mass
                    a[1] -= g
                    d.pos += (d.vel + a * dt) * dt
                    d.vel += a * dt
                    t += dt
                    if d.pos[1] > maxs[1]:
                        maxs[1] = d.pos[1]
                    if d.pos[0] > maxs[0]:
                        maxs[0] = d.pos[0]
                    if d.pos[1] < 0.0:
                        if t > self.maxt:
                            self.maxt = t
                        break
        for ii in range(ndrops):    # Make some randomly distributed water drops
            self.drops[ii] = drop([0.0, 0.0], [cos(np.random.random()*np.pi)*v0,
                             sin(np.random.random()*np.pi)*v0],
                             np.random.random()*(r[1]-r[0]) + r[0])

        anim = animation.FuncAnimation(self.fig, self.update, init_func=self.setup,
                                       interval=200, frames=nframes)
        anim.save('Sprinkler.gif', fps = 20, writer='imagemagick')

    def setup(self):
        self.scat = self.ax.scatter([d.pos[0] for d in self.drops],
                                    [d.pos[1] for d in self.drops], marker='.', color='k')
        self.ax.set_xlim([-maxs[0]*1.15, maxs[0]*1.15])
        self.ax.set_ylim([0, maxs[1]*1.15])

        return self.scat,

    def update(self,frame):  # Use set_offsets to move the water drops
        # Create between 0 and ndrops new drops, but only until tend-maxt
        if time[frame] < tend - self.maxt*1.1:
            self.create(np.random.randint(0, ndrops))
        self.step()          # Advance to the next 'step'
        self.scat.set_offsets(np.stack([[d.pos[0] for d in self.drops],
                                        [d.pos[1] for d in self.drops]], 1))

        return self.scat,

    def create(self, i):
        for ii in range(i):
            self.drops.append(drop([0.0, 0.0], [cos(np.random.random()*np.pi)*v0,
                             sin(np.random.random()*np.pi)*v0],
                             np.random.random()*(r[1]-r[0]) + r[0]))
    def step(self):
        for ii in range(len(self.drops)):
            coef = - 0.5*C*np.pi*self.drops[ii].r**2*rho    # Aggregated coefficient
            mass = 4/3*np.pi*self.drops[ii].r**3*1000
            a = np.power(self.drops[ii].vel, 2) * coef * np.sign(self.drops[ii].vel)/mass
            a[1] -= g
            # Approximate how much the position and velocity would change if we assume
            # a(t) does not change between t and t+dt
            self.drops[ii].pos += np.array(self.drops[ii].vel) * dt + 0.5 * a * dt**2
            self.drops[ii].vel += a * dt
            if self.drops[ii].pos[1] < 0.0:     # Check if the drop has hit the ground
                self.drops[ii].pos[1] = 0.0
                self.drops[ii].vel = [0.0, 0.0]

sprinkler()