75 lines
2.5 KiB
Python
Executable File
75 lines
2.5 KiB
Python
Executable File
#!/usr/bin/python
|
|
# -*- coding:utf-8 -*-
|
|
"""
|
|
Implementation of the plotting and visualization functions.
|
|
"""
|
|
import numpy as np
|
|
import time
|
|
import matplotlib.pyplot as plt
|
|
|
|
class DynamicUpdate():
|
|
#Suppose we know the x range
|
|
min_x = -10
|
|
max_x = 10
|
|
|
|
plt.ion()
|
|
|
|
def __init__(self, dyn_syst):
|
|
self.dyn_syst = dyn_syst
|
|
|
|
def set_lims(self, factor=1.5):
|
|
self.ax.set_xlim(factor*self.min_x, factor*self.max_x)
|
|
self.ax.set_ylim(factor*self.min_x, factor*self.max_x)
|
|
self.ax.set_zlim(factor*self.min_x, factor*self.max_x)
|
|
|
|
def on_launch(self):
|
|
#Set up plot
|
|
self.fig = plt.figure(figsize=(10,10))
|
|
self.ax = self.fig.add_subplot(projection='3d')
|
|
self.lines = []
|
|
for i,body in enumerate(self.dyn_syst.bodylist):
|
|
x, y, z = body.q
|
|
lines, = self.ax.plot([x],[y],[z],'o',color="C{0:d}".format(i),label="{0:s}".format(str(body)))
|
|
self.lines.append(lines)
|
|
self.lines = np.array(self.lines)
|
|
#Autoscale on unknown axis and known lims on the other
|
|
self.ax.set_autoscaley_on(True)
|
|
self.set_lims()
|
|
#Other stuff
|
|
self.ax.grid()
|
|
self.ax.legend()
|
|
|
|
def on_running(self, dyn_syst, step=None, label=None):
|
|
xdata, ydata, zdata = dyn_syst.get_positions()
|
|
values = np.sqrt(np.sum((np.array((xdata,ydata,zdata))**2).T,axis=1))
|
|
self.min_x, self.max_x = -np.max([np.abs(values).max(),self.max_x]), np.max([np.abs(values).max(),self.max_x])
|
|
self.set_lims()
|
|
#Update data (with the new _and_ the old points)
|
|
for i,body in enumerate(dyn_syst.bodylist):
|
|
x, y, z = body.q
|
|
self.lines[i].set_data_3d([x], [y], [z])
|
|
if not label is None:
|
|
self.ax.set_title(label)
|
|
#Need both of these in order to rescale
|
|
self.ax.relim()
|
|
self.ax.autoscale_view()
|
|
#We need to draw *and* flush
|
|
self.fig.canvas.draw()
|
|
self.fig.canvas.flush_events()
|
|
if not step is None and step%1000==0:
|
|
self.fig.savefig("tmp/{0:06d}.png".format(step),bbox_inches="tight")
|
|
|
|
#Example
|
|
def __call__(self):
|
|
import numpy as np
|
|
import time
|
|
self.on_launch()
|
|
xdata = []
|
|
ydata = []
|
|
for x in np.arange(0,10,0.5):
|
|
xdata.append(x)
|
|
ydata.append(np.exp(-x**2)+10*np.exp(-(x-7)**2))
|
|
self.on_running(xdata, ydata)
|
|
time.sleep(1)
|
|
return xdata, ydata
|