3D animation of sprinkler

[Marlisevt]

Hi,
I am trying to animate/simulate a sprinkler. I tried it in 2D, and got the output I wanted (moving points, iterative creation of points/droplets, hitting the ground and stopping).
I thought changing it to 3D would not be an issue, however I am only getting a single stationary point as my output on the 3D graph. The point is located at the starting position of (0.0, 0.0, 0.0). I think the problem may by with the update function, where set_offsets may only be for a 2D graph? I looked through the matplotlib documentation, and couldn't find anything similar for a 3D plot.
https://matplotlib.org/api/_as_gen/mpl_t...s3d.Axes3D

Going from 2D to 3D, I changed the position, velocity, and acceleration from 1x2 arrays ro 1x3 arrays. I also changed the max positions from maxs[1] to maxs[2], etc.

Is there an error in my code/logic or is the problem with the update function?

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156

import matplotlib import matplotlib.pyplot as plt from matplotlib import animation import numpy as np from math import sin, cos from mpl_toolkits.mplot3d import Axes3D import random   #Parameters rho = 1.225 c = 0.5 v0 = 50 g = 9.81   #Timing tmax = 10 nframes = 100 time = np.linspace(0,tmax, nframes) dt = time[1]-time[0]   #Waterdroplets ndrops = 100   #Positioning maxs = [0.0, 0.0, 0.0] rmax = [0.0008, 0.0065]   #Droplet sizing theta = np.radians(np.random.normal(37, 8, 80)) phi = np.radians(np.random.normal(3, 1.07, 80)) radii = np.random.normal(0.004, 0.001, 20)   class drop:           def __init__(self,pos,vel,r):         self.pos = pos         self.vel = vel         self.r = r   class sprinkler:           def __init__(self):         self.fig = plt.figure()         self.ax = self.fig.add_subplot(111, projection = '3d')         self.drops = [None]*ndrops    # creates empty list of length ndrops         self.maxt = 0.0                   #Find the maximum flight time for each droplet         #and the maximum distance the droplets will travel         theta = np.radians(np.random.normal(37, 8, 80))         phi = np.radians(np.random.normal(3, 1.07, 80))         radii = np.random.normal(0.004, 0.001, 80)                      for i in range(len(phi)):             m = [drop([0.0, 0.1, 0.1], [v0*cos(theta[i])*cos(phi[i]),                                         v0*cos(theta[i])*sin(theta[i]),                                         v0*sin(theta[i])],0.0008),                  drop([0.0, 0.1, 0.1], [v0*cos(theta[i])*cos(phi[i]),                                        v0*cos(theta[i])*sin(theta[i]),                                        v0*sin(theta[i])],0.0065)]             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[2] > 0:                     a = np.power(d.vel, 2) * coef * np.sign(d.vel)/mass                     a[2] -= g                     d.pos += (d.vel + a * dt) * dt                     d.vel += a * dt                     t += dt                     if d.pos[2] > maxs[2]:                         maxs[2] = d.pos[2]                                        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[2] < 0.0:                         if t > self.maxt:                             self.maxt = t                         break         print('Max time is:',self.maxt)         print('Max positions are:', maxs)                   #create initial droplets                   for ii in range(ndrops):             phiang = random.randint(0,len(phi)-1)             thetang = random.randint(0,len(theta)-1)             rad = random.randint(0,len(radii)-1)                           self.drops[ii] = drop([0.0, 0.0, 0.1],                                  [v0*cos(theta[thetang])*cos(phi[phiang]),                                   v0*cos(theta[thetang])*sin(phi[phiang]),                                   v0*sin(theta[thetang])],                                   radii[random.randint(0,len(radii)-1)])         #initiate animation         ani = animation.FuncAnimation(self.fig, self.update, init_func = self.setup,                                           interval = 200, frames = nframes)         ani.save('Sprinkler.mp4', fps = 20, extra_args=['-vcodec', 'libx264'])         plt.show()           #Setup is used to initiate the 3D axis and plot     def setup(self):         self.scat = self.ax.scatter([d.pos[0] for d in self.drops],                                     [d.pos[1] for d in self.drops],                                     [d.pos[2] for d in self.drops],                                     'b.')                       self.ax.set_xlim(-1, 100)         self.ax.set_ylim(-1, 100)         self.ax.set_zlim(0, 50)         self.ax.set_xlabel('X Distance')         self.ax.set_ylabel('Y Distance')         self.ax.set_zlabel('Height')                       return self.scat           #Update function updates the animation each frame.           def update(self, frame):         if time[frame] <(tmax-self.maxt*1.1):             self.create(ndrops)         self.step()         #offset moves each point along accoring to the step function         self.scat.set_offsets([[d.pos[0] for d in self.drops],                               [d.pos[1] for d in self.drops],                               [d.pos[2] for d in self.drops]])         return self.scat,            #New droplets are created at each frame to have iterative droplets coming from the sprinkler     def create(self, i):         for l in range(i):             phiang = random.randint(0,len(phi)-1)             thetang = random.randint(0,len(theta)-1)             rad = random.randint(0,len(radii)-1)             self.drops.append(drop([0.0, 0.0, 0.0],                                    [v0*cos(theta[thetang])*cos(phi[phiang]),                                     v0*cos(theta[thetang])*sin(phi[phiang]),                                     v0*sin(theta[thetang])],                                    radii[rad]))           #step function moves each droplet along at each interval dt             def step(self):                   for x in range(len(self.drops)):             coef = -0.5*c*np.pi*self.drops[x].r**2*rho  #drag force coefficient             mass = 4/3*np.pi*self.drops[x].r**3*1000    #mass of droplet             a = np.power(self.drops[x].vel,2) * coef * np.sign(self.drops[x].vel)/mass  #acceleration             a[2] = a[2]-g                               self.drops[x].pos += np.array(self.drops[x].vel)*dt +0.5*a*dt**2             self.drops[x].vel += a*dt             #When droplet hits ground, it stops moving             if self.drops[x].pos[2] < 0.0:                 self.drops[x].pos[2] = 0.0                 self.drops[x].vel = [0.0, 0.0, 0.0]