NSSC/Exercise_02/plot.py

38 lines
865 B
Python

#!/usr/bin/env python
import sys
from state import State
from matplotlib import pyplot as plt
def main(argv):
path = argv[1]
with open(path) as f:
lines = f.readlines()
count = int(lines[0])
size = float(lines[2])
split = count + 3
states = [lines[i:i+split] for i in range(0, len(lines), split)]
states = [State.deserialize(state) for state in states]
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
current = 0
while True:
ax.cla()
plt.xlim([0,size])
plt.ylim([0,size])
plt.title(f'Step {current}')
ax.set_zlim([0,size])
ax.scatter(
states[current].molecules[:,0],
states[current].molecules[:,1],
states[current].molecules[:,2],
c=range(count)
)
plt.pause(0.02)
current += 1
current %= len(states)
if __name__ == '__main__':
main(sys.argv)