I'm trying to write an N-body simulation, and I've run into a bit of a roadblock with Matplotlib's animation capabilities. I'm not entirely sure how to update the data inside the scatter plot, so I think that might be where the problem is.
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.animation as animation
system = []
xlist = []
ylist = []
zlist = []
poslist = []
nbodies = 3
scale = 100
mass = 0
timestep = 0.1
step = 0.01
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
i = 1
plt.ion()
scat = ax.scatter(xlist, ylist, zlist, c='r', marker='o', animated=True)
class Body:
def __init__(self, posvel, mass):
self.state = np.zeros((3, 3))
self.state[0:2] = posvel * scale
self.state[2] = np.array([0.0, 0.0, 0.0])
self.mass = mass
def setPos(self, pos):
self.state[0] = pos
def setVel(self, vel):
self.state[1] = vel
def setAcc(self, acc):
self.state[2] = acc
def getPos(self):
return self.state[0]
def getVel(self):
return self.state[1]
def getAcc(self):
return self.state[2]
def startSystem():
for i in range(0,nbodies):
system.append(Body(np.random.rand(2,3), np.random.random_sample()))
curpos = system[i].getPos()
xlist.append(curpos[0])
ylist.append(curpos[1])
zlist.append(curpos[2])
def stepSystem():
xlist = []
ylist = []
zlist = []
poslist = []
for i in range(0,nbodies):
curpos = system[i].getPos()
system[i].setPos(curpos + 1*step)
poslist.append(curpos + 1*step)
xlist.append(curpos[0] + 1*step)
ylist.append(curpos[1] + 1*step)
zlist.append(curpos[2] + 1*step)
yield xlist, ylist, zlist
def simPoints(stepSystem):
xlist, ylist,zlist = stepSystem[0], stepSystem[1], stepSystem[2]
scat.set_offsets(poslist)
return scat
def Plot():
ani = animation.FuncAnimation(fig, simPoints, stepSystem, blit=False,\
interval=100, repeat=True)
plt.show()
def main():
startSystem()
Plot()
