def DrawGraph(self, lastFlag, randNode=None):
"""
Updates the Plot with uncertainty ellipse and trajectory at each time step
Input Parameters:
lastFlag: Flag to denote if its the last iteration
randNode: Node data representing the randomly sampled point
"""
xValues = []
yValues = []
widthValues = []
heightValues = []
angleValues = []
lineObjects = []
for ellipseNode in self.nodeList:
if ellipseNode is not None and ellipseNode.parent is not None:
ellNodeShape = ellipseNode.means.shape
xPlotValues = []
yPlotValues = []
# Prepare the trajectory x and y vectors and plot them
for k in range(ellNodeShape[0]):
xPlotValues.append(ellipseNode.means[k,0,0])
yPlotValues.append(ellipseNode.means[k,1,0])
# Plotting the risk bounded trajectories
lx, = plt.plot(xPlotValues, yPlotValues, "#636D97", linewidth=2.0)
lineObjects.append(lx)
# Plot only the last ellipse in the trajectory
alfa = math.atan2(ellipseNode.means[-1,1,0], ellipseNode.means[-1,0,0])
elcovar = np.asarray(ellipseNode.covar[-1,:,:])
elE, elV = LA.eig(elcovar[0:2,0:2])
xValues.append(ellipseNode.means[-1,0,0])
yValues.append(ellipseNode.means[-1,1,0])
widthValues.append(math.sqrt(elE[0]))
heightValues.append(math.sqrt(elE[1]))
angleValues.append(alfa*360)
# Plot the randomly sampled point
rx, = plt.plot(randNode.means[-1,0,:], randNode.means[-1,1,:], "^k")
# Plot the Safe Ellipses
XY = np.column_stack((xValues, yValues))
ec = EllipseCollection(widthValues,
heightValues,
angleValues,
units='x',
offsets=XY,
facecolors="#C59434",
transOffset=plt.axes().transData)
ec.set_alpha(0.5)
plt.axes().add_collection(ec)
plt.pause(0.0001)
if not lastFlag:
ec.remove()
rx.remove()
for lx in lineObjects:
lx.remove()
class Animator(object):
def __init__(self, positions, diameter):
self.circ_count = positions.shape[0]
plt.ion()
fig = plt.figure(figsize=(10, 10))
ax = fig.gca()
self.ax = ax
diameters = np.ones(self.circ_count) * diameter
circ_colors = [(0.0, 0.0, 1.0) for _ in range(self.circ_count)]
#add circles
self.circles = EllipseCollection(widths=diameters,
heights=diameters,
angles=np.zeros_like(diameters),
units='xy',
offsets=positions,
transOffset=ax.transData,
edgecolor='face',
facecolor=circ_colors)
ax.add_collection(self.circles)
#add polygons
self.poly_count = 3
verts = [[(0, 1), (1, 0), (2, 2)], [(6, 5), (3, 7), (7, 6)],
[(0, -1), (-1, 0), (4, 5)]]
poly_colors = [(0.0, 0.0, 1.0) for _ in range(self.poly_count)]
self.polygons = PolyCollection(verts, facecolors=poly_colors)
ax.add_collection(self.polygons)
ax.axis([-20, 20, -20, 20])
ax.get_xaxis().set_visible(True)
ax.get_yaxis().set_visible(True)
#ax.set_axis_bgcolor('black')
plt.draw()
plt.pause(0.1)
def update(self, positions):
#update number of balls
self.circles.set_offsets(positions)
colors = [(1.0, 0.0, 0.0) for _ in range(self.circ_count)]
self.circles.set_facecolors(colors)
#redefine circles
diameter = 2
diameters = np.ones(self.circ_count + 100) * diameter
circ_colors = [(0.0, 0.0, 1.0) for _ in range(self.circ_count)]
#add circles
self.circles.remove()
self.circles = EllipseCollection(widths=diameters,
heights=diameters,
angles=np.zeros_like(diameters),
units='xy',
offsets=positions,
transOffset=self.ax.transData,
edgecolor='face',
facecolor=circ_colors)
self.ax.add_collection(self.circles)
#label
text_labels = [None] * positions.shape[0]
for i in range(positions.shape[0]):
text_labels[i] = plt.text(positions[i, 0],
positions[i, 1],
str(i),
color='white')
#remove polygon
#self.polygons.remove()
poly_color = np.random.uniform(0.0, 0.89, (3, )) + 0.1
verts = [[(10, 1), (9, 0), (8, 2), (10, 2)], [(8, 8), (9, 9), (6, 7)]]
self.polygons.set_verts(verts)
plt.draw()
plt.pause(0.1)
#remove labels
for i in range(positions.shape[0]):
text_labels[i].remove()
class Animator(object):
def __init__(self, template_w, template_h):
#set up figure
plt.ion()
#fig = plt.figure(figsize = (10,10))
fig = plt.figure(figsize=(template_w, template_h))
#print template_w, template_h
#raw_input()
ax = fig.gca()
#set axis dimension
ax.axis([
-2 * template_w / 2.0, 2 * template_w / 2.0, -2 * template_h / 2.0,
2 * template_h / 2.0
])
#ax.axis([-template_w/2.0, template_w/2.0,-template_h/2.0,template_h/2.0])
ax.get_xaxis().set_visible(True)
ax.get_yaxis().set_visible(True)
#draw template
self.plot_template(template_w, template_h)
#draw to screen
plt.draw()
plt.pause(0.1)
#save axes to object
self.ax = ax
#save animator canvas area
self.width = template_w
self.height = template_h
self.title = plt.text(0, 20, 'start title', color='k')
#plot template
def plot_template(self, l=20.0, w=20.0):
x_pos = [-l / 2, l / 2, l / 2, -l / 2, -l / 2]
y_pos = [-w / 2, -w / 2, w / 2, w / 2, -w / 2]
plt.plot(x_pos, y_pos, 'k-', linewidth=3.0)
plt.xlabel('x')
plt.ylabel('y')
def generate_color_array(self, collision):
#define colors (red = collision, blue = no collision)
colors = [None] * len(collision)
for i in range(len(collision)):
if collision[i] == True:
colors[i] = (1.0, 0.0, 0.0)
else:
colors[i] = (0.0, 0.0, 1.0)
return colors
#add circular objects
def add_circular_objects(self,
diameters,
positions,
collision,
pause_time=0.1):
circ_colors = self.generate_color_array(collision)
self.circles = EllipseCollection(widths=diameters,
heights=diameters,
angles=np.zeros_like(diameters),
units='xy',
offsets=positions,
transOffset=self.ax.transData,
edgecolor='k',
facecolor=circ_colors)
self.ax.add_collection(self.circles)
#add text label
text_labels = [None] * len(collision)
for i in range(len(collision)):
text_labels[i] = plt.text(positions[i, 0],
positions[i, 1],
str(i),
color='k')
self.circle_labels = text_labels
#draw to screen
#plt.draw()
#plt.pause(pause_time)
#remove text labels
#for i in range(len(collision)):
# text_labels[i].remove()
#add polygon objects
def add_polygon_objects(self, verts, collision, pause_time=0.1):
poly_colors = self.generate_color_array(collision)
self.polygons = PolyCollection(verts, facecolors=poly_colors)
self.ax.add_collection(self.polygons)
#add text label
num_circles = self.circles.get_offsets().shape[0]
#print 'number of circles =', num_circles
text_labels = [None] * len(collision)
for i in range(len(collision)):
temp = np.array(verts[i])
x = np.mean(temp[:, 0])
y = np.mean(temp[:, 1])
text_labels[i] = plt.text(x, y, str(i + num_circles), color='k')
self.polygon_labels = text_labels
plt.draw()
plt.pause(pause_time)
#remove text labels
#for i in range(len(collision)):
# text_labels[i].remove()
#remove circular objects:
def remove_circles(self):
self.circles.remove()
for label in self.circle_labels:
label.remove()
#remove polygon objects:
def remove_polygons(self):
self.polygons.remove()
for label in self.polygon_labels:
label.remove()
#update circular objects
def update_circular_objects(self, positions, collision, pause_time=0.1):
#set circle colors
circ_colors = self.generate_color_array(collision)
self.circles.set_facecolors(circ_colors)
#set circle positions
self.circles.set_offsets(positions)
#remove labels
for label in self.circle_labels:
label.remove()
#add labels
text_labels = [None] * len(collision)
for i in range(len(collision)):
text_labels[i] = plt.text(positions[i, 0],
positions[i, 1],
str(i),
color='k')
self.circle_labels = text_labels
plt.draw()
plt.pause(pause_time)
#update polygon objects
def update_polygon_objects(self, positions, collision, pause_time=0.1):
#set polygon colors
poly_colors = self.generate_color_array(collision)
self.polygons.set_facecolors(poly_colors)
#set polygon positions
#print 'new verts positions=' , positions
self.polygons.set_verts(positions)
#remove labels
for label in self.polygon_labels:
label.remove()
#self.title.remove()
#add new labels
num_circles = self.circles.get_offsets().shape[0]
#print self.polygons.get_offsets()
#assert(0)
text_labels = [None] * len(collision)
for i in range(len(collision)):
temp = np.array(positions[i])
x = np.mean(temp[:, 0])
y = np.mean(temp[:, 1])
text_labels[i] = plt.text(x, y, str(i + num_circles), color='k')
self.polygon_labels = text_labels
plt.draw()
plt.pause(pause_time)
#display the total area covered
def show_title(self, area, pause_time=0.1):
title = plt.text(-self.width / 4,
self.height,
'total covered area = ' + str(area),
color='k',
fontsize=20)
self.title.remove()
self.title = title
class Animator(object):
def __init__(self, template_w, template_h):
#set up figure
plt.ion()
fig = plt.figure(figsize = (10,10))
ax = fig.gca()
ax.axis([-template_w/2.0, template_w/2.0,-template_h/2.0,template_h/2.0])
ax.get_xaxis().set_visible(True)
ax.get_yaxis().set_visible(True)
#draw template
self.plot_template(template_w,template_h)
#draw to screen
plt.draw()
plt.pause(0.1)
#save axes to object
self.ax = ax
#plot template
def plot_template(self,l = 20.0, w = 20.0):
x_pos = [-l/2, l/2, l/2, -l/2, -l/2 ]
y_pos = [-w/2, -w/2, w/2, w/2, -w/2]
plt.plot(x_pos,y_pos,'k-',linewidth = 3.0)
plt.xlabel('x')
plt.ylabel('y')
def generate_color_array(self,collision):
#define colors (red = collision, blue = no collision)
colors = [None] * len(collision)
for i in range(len(collision)):
if collision[i] == True:
colors[i] = (1.0,0.0,0.0)
else:
colors[i] = (0.0,0.0,1.0)
return colors
#add circular objects
def add_circular_objects(self,diameters,positions,collision):
circ_colors = self.generate_color_array(collision)
self.circles = EllipseCollection(widths=diameters,
heights=diameters,
angles=np.zeros_like(diameters),
units='xy',
offsets=positions,
transOffset=self.ax.transData,
edgecolor='k', facecolor=circ_colors)
self.ax.add_collection(self.circles)
#add text label
text_labels = [None] * len(collision)
for i in range(len(collision)):
text_labels[i]= plt.text(positions[i,0],positions[i,1],str(i),color = 'w')
self.circle_labels = text_labels
#draw to screen
plt.draw()
plt.pause(0.1)
#remove text labels
#for i in range(len(collision)):
# text_labels[i].remove()
#add polygon objects
def add_polygon_objects(self,verts,collision):
poly_colors = self.generate_color_array(collision)
self.polygons = PolyCollection(verts, facecolors=poly_colors)
self.ax.add_collection(self.polygons)
#add text label
num_circles = self.circles.get_offsets().shape[1]
text_labels = [None] * len(collision)
for i in range(len(collision)):
temp = np.array(verts[i])
x = np.mean(temp[:,0])
y = np.mean(temp[:,1])
text_labels[i]= plt.text(x,y,str(i+num_circles),color = 'w')
self.polygon_labels = text_labels
plt.draw()
plt.pause(0.1)
#remove text labels
#for i in range(len(collision)):
# text_labels[i].remove()
#remove circular objects:
def remove_circles(self):
self.circles.remove()
for label in self.circle_labels:
label.remove()
#remove polygon objects:
def remove_polygons(self):
self.polygons.remove()
for label in self.polygon_labels:
label.remove()
#update circular objects
def update_circular_objects(self,positions,collision):
#set circle colors
circ_colors = self.generate_color_array(collision)
self.circles.set_facecolors(circ_colors)
#set circle positions
self.circles.set_offsets(positions)
#remove labels
for label in self.circle_labels:
label.remove()
#add labels
text_labels = [None] * len(collision)
for i in range(len(collision)):
text_labels[i]= plt.text(positions[i,0],positions[i,1],str(i),color = 'w')
self.circle_labels = text_labels
plt.draw()
plt.pause(0.1)
#update polygon objects
def update_polygon_objects(self,positions,collision):
#set polygon colors
poly_colors = self.generate_color_array(collision)
self.polygons.set_facecolors(poly_colors)
#set polygon positions
#print 'new verts positions=' , positions
self.polygons.set_verts(positions)
#remove labels
for label in self.polygon_labels:
label.remove()
#add new labels
num_circles = self.circles.get_offsets().shape[1]
#print self.polygons.get_offsets()
#assert(0)
text_labels = [None] * len(collision)
for i in range(len(collision)):
temp = np.array(positions[i])
x = np.mean(temp[:,0])
y = np.mean(temp[:,1])
text_labels[i]= plt.text(x,y,str(i+num_circles),color = 'w')
self.polygon_labels = text_labels
plt.draw()
plt.pause(0.1)
class Quiver(_Quiver):
def __init__(self, ax, *args, **kwargs):
if 'sigma' in kwargs:
scale_units = kwargs.get('scale_units', 'xy')
kwargs['scale_units'] = scale_units
if kwargs['scale_units'] != 'xy':
raise ValueError(
'scale units must be "xy" when sigma is given')
angles = kwargs.get('angles', 'xy')
kwargs['angles'] = angles
if kwargs['angles'] != 'xy':
raise ValueError('angles must be "xy" when sigma is given')
sigma = kwargs.pop('sigma', None)
ellipse_kwargs = kwargs.pop('ellipse_kwargs', {})
if 'offsets' in ellipse_kwargs:
raise ValueError('cannot specify ellipse offsets')
if 'units' in ellipse_kwargs:
raise ValueError('cannot specify ellipse units')
self.ellipse_kwargs = {
'edgecolors': 'k',
'facecolors': 'none',
'linewidths': 1.0
}
self.ellipse_kwargs.update(ellipse_kwargs)
self.ellipsoids = None
_Quiver.__init__(self, ax, *args, **kwargs)
if sigma is not None:
if self.scale is None:
self.scale = _estimate_scale(self.X, self.Y, self.U, self.V)
su, sv, rho = sigma[0], sigma[1], sigma[2]
self._update_ellipsoids(su, sv, rho)
def _update_ellipsoids(self, su, sv, rho):
self.scale_units = 'xy'
self.angles = 'xy'
tips_x = self.X + self.U / self.scale
tips_y = self.Y + self.V / self.scale
tips = np.array([tips_x, tips_y]).transpose()
a, b, angle = compute_abphi(su, sv, rho)
width = 2.0 * a / self.scale
height = 2.0 * b / self.scale
if self.ellipsoids is not None:
self.ellipsoids.remove()
# do not draw ellipses which are too small
too_small = 0.001
length = np.sqrt((self.U / self.scale)**2 + (self.V / self.scale)**2)
with warnings.catch_warnings():
# do not print out zero division warning
warnings.simplefilter("ignore")
is_not_too_small = ((np.nan_to_num(width / length) > too_small) |
(np.nan_to_num(height / length) > too_small))
width = width[is_not_too_small]
height = height[is_not_too_small]
angle = angle[is_not_too_small]
tips = tips[is_not_too_small]
# dont add ellipses if there are no ellipses to add
if any(is_not_too_small):
self.ellipsoids = EllipseCollection(width,
height,
angle,
units=self.scale_units,
offsets=tips,
transOffset=self.ax.transData,
**self.ellipse_kwargs)
self.ax.add_collection(self.ellipsoids)
else:
self.ellipsoids = None
def set_UVC(self, u, v, C=None, sigma=None):
if C is None:
_Quiver.set_UVC(self, u, v)
else:
_Quiver.set_UVC(self, u, v, C)
if sigma is not None:
su, sv, rho = sigma[0], sigma[1], sigma[2]
self._update_ellipsoids(su, sv, rho)
def remove(self):
# remove the quiver and ellipsoid collection
_Quiver.remove(self)
if self.ellipsoids is not None:
self.ellipsoids.remove()
class Quiver(_Quiver):
def __init__(self,ax,*args,**kwargs):
if 'sigma' in kwargs:
scale_units = kwargs.get('scale_units','xy')
kwargs['scale_units'] = scale_units
if kwargs['scale_units'] != 'xy':
raise ValueError('scale units must be "xy" when sigma is given')
angles = kwargs.get('angles','xy')
kwargs['angles'] = angles
if kwargs['angles'] != 'xy':
raise ValueError('angles must be "xy" when sigma is given')
sigma = kwargs.pop('sigma',None)
ellipse_kwargs = kwargs.pop('ellipse_kwargs',{})
if 'offsets' in ellipse_kwargs:
raise ValueError('cannot specify ellipse offsets')
if 'units' in ellipse_kwargs:
raise ValueError('cannot specify ellipse units')
self.ellipse_kwargs = {'edgecolors':'k',
'facecolors':'none',
'linewidths':1.0}
self.ellipse_kwargs.update(ellipse_kwargs)
self.ellipsoids = None
_Quiver.__init__(self,ax,*args,**kwargs)
if sigma is not None:
if self.scale is None:
self.scale = _estimate_scale(self.X,self.Y,self.U,self.V)
su,sv,rho = sigma[0],sigma[1],sigma[2]
self._update_ellipsoids(su,sv,rho)
def _update_ellipsoids(self,su,sv,rho):
self.scale_units = 'xy'
self.angles = 'xy'
tips_x = self.X + self.U/self.scale
tips_y = self.Y + self.V/self.scale
tips = np.array([tips_x,tips_y]).transpose()
a,b,angle = compute_abphi(su,sv,rho)
width = 2.0*a/self.scale
height = 2.0*b/self.scale
if self.ellipsoids is not None:
self.ellipsoids.remove()
# do not draw ellipses which are too small
too_small = 0.001
length = np.sqrt((self.U/self.scale)**2 + (self.V/self.scale)**2)
with warnings.catch_warnings():
# do not print out zero division warning
warnings.simplefilter("ignore")
is_not_too_small = ((np.nan_to_num(width/length) > too_small) |
(np.nan_to_num(height/length) > too_small))
width = width[is_not_too_small]
height = height[is_not_too_small]
angle = angle[is_not_too_small]
tips = tips[is_not_too_small]
# dont add ellipses if there are no ellipses to add
if any(is_not_too_small):
self.ellipsoids = EllipseCollection(width,height,angle,
units=self.scale_units,
offsets = tips,
transOffset=self.ax.transData,
**self.ellipse_kwargs)
self.ax.add_collection(self.ellipsoids)
else:
self.ellipsoids = None
def set_UVC(self,u,v,C=None,sigma=None):
if C is None:
_Quiver.set_UVC(self,u,v)
else:
_Quiver.set_UVC(self,u,v,C)
if sigma is not None:
su,sv,rho = sigma[0],sigma[1],sigma[2]
self._update_ellipsoids(su,sv,rho)
def remove(self):
# remove the quiver and ellipsoid collection
_Quiver.remove(self)
if self.ellipsoids is not None:
self.ellipsoids.remove()
