def drawWheel(ax, shape, body):
r = shape.radius
pos = body.transform * shape.pos
drawCircle(ax, pos, r)
for a in [0, 2 * np.pi / 3, 4 * np.pi / 3]:
v = vrotate((1, 0), body.angle + a)
x = [pos[0], pos[0] + 0.99 * r * v[0]]
y = [pos[1], pos[1] + 0.99 * r * v[1]]
ax.plot(x, y, linestyle='-', color='b', lw=1)
def drawWheel(ax,shape,body):
r = shape.radius
pos = body.transform*shape.pos
drawCircle(ax,pos,r)
for a in [0,2*np.pi/3,4*np.pi/3]:
v = vrotate((1,0),body.angle+a)
x = [pos[0],pos[0]+0.99*r*v[0]]
y = [pos[1],pos[1]+0.99*r*v[1]]
ax.plot(x,y, linestyle='-', color='b', lw=1)
def createCircle(position, r=0.3, dynamic=True, bMatplotlib = True):
global world, fig, ax
bodyDef = Box2D.b2BodyDef()
fixtureDef = Box2D.b2FixtureDef()
if dynamic:
bodyDef.type = Box2D.b2_dynamicBody
fixtureDef.density = 1
else:
bodyDef.type = Box2D.b2_staticBody
fixtureDef.density = 0
bodyDef.position = position
if(abs(world.gravity[1]) > 1):
bodyDef.linearDamping = 0.1
bodyDef.angularDamping = 0.1
else:
bodyDef.linearDamping = 70
bodyDef.angularDamping = 30
body = world.CreateBody(bodyDef)
fixture = body.CreateFixture(shape=Box2D.b2CircleShape(radius=r), density=1.0, friction=0.3)
if(bMatplotlib):
createGlobalFigure()
fixture.userData = drawCircle(ax,position,r)
return body
def drawEpuck(ax, shape, body):
r = shape.radius
pos = body.position
drawCircle(ax, pos, r)
irangles = body.userData["IRAngles"]
nir = len(irangles)
for a in [-np.pi / 4, np.pi / 4]:
v = vrotate((1, 0), body.angle + a)
x = [pos[0], pos[0] + 0.99 * r * v[0]]
y = [pos[1], pos[1] + 0.99 * r * v[1]]
ax.plot(x, y, linestyle='-', color='b', lw=1)
for k in range(nir):
v = vrotate((1, 0), body.angle + irangles[k])
c = pos + [0.97 * r * v[0], 0.97 * r * v[1]]
value = body.userData["IRValues"][k]
if (value > 1): value = 1
drawCircle(ax, c, r=0.07, color=[1 - value, value, 0])
def plotWorld(ax, alpha=0.3, nao=None, obj=None, bDrawGround=False, color='b', centers=[], specials=[], cradius=0.1, ccolor='r', label='_'):
global world
if(nao and obj):
pobj = [obj.position[0], obj.position[1]]
pnao = nao.ini_pos
#ax.plot([pobj[0],pnao[0]], [pobj[1],pnao[1]], linestyle='--', color='g', lw=2)
for body in world.bodies:
for fixture in body.fixtures:
shape = fixture.shape
if(body.active):
if(isinstance(shape,Box2D.b2PolygonShape)):
if(fixture.density == 0.0): drawBox2D(ax,body,fixture,color=color,alpha=0.25)
else: drawBox2D(ax,body,fixture,color=color,alpha=alpha)
if(isinstance(shape,Box2D.b2CircleShape)):
drawWheel(ax,shape,body)
else:
if(isinstance(shape,Box2D.b2PolygonShape)): drawBox2D(ax,body,fixture,color=color,alpha=0.25,fill=False,linestyle='dashed')
if(isinstance(shape,Box2D.b2CircleShape)): drawCircle(ax,body.position,0.3,fill=False,linestyle='dashed')
if(len(centers)>0): plotVectors(ax,centers, cradius=cradius, specials=specials, ccolor=ccolor, label=label)
plt.grid()
def plotVectors(ax,
centers,
specials=[],
dirs=[],
label='_',
cradius=0.1,
ccolor='r'):
for i, p in enumerate(centers):
dx, dy = -2 * cradius, cradius
if (p[0] >= 0): dx, dy = cradius, cradius / 2.0
if (label != '_'):
plt.annotate(r'$%s_%d$' % (label, i),
xy=p,
xytext=(p[0] + dx, p[1] + dy))
if (len(specials) > 0):
if (i in specials): drawCircle(ax, p, cradius, color=ccolor)
else: drawCircle(ax, p, cradius - 0.05, color='b')
else: drawCircle(ax, p, cradius, color=ccolor)
if (len(dirs) > 0):
d = dirs[i]
vel = (3 * d[0], 3 * d[1])
if (VectorFigUtils.vnorm(vel)):
ax.arrow(p[0],
p[1],
vel[0],
vel[1],
head_width=0.07,
head_length=0.11,
fc='k',
ec='k')
def plotVectors(ax, centers, specials=[], dirs = [], label='_', cradius=0.1, ccolor='r'):
for i,p in enumerate(centers):
dx,dy = -2*cradius,cradius
if(p[0] >= 0): dx,dy = cradius,cradius/2.0
if(label != '_'):
plt.annotate(r'$%s_%d$'%(label,i), xy=p, xytext=(p[0]+dx,p[1]+dy))
if(len(specials)>0):
if(i in specials): drawCircle(ax,p,cradius,color=ccolor)
else: drawCircle(ax,p,cradius-0.05,color='b')
else: drawCircle(ax,p,cradius,color=ccolor)
if(len(dirs)>0):
d = dirs[i]
vel = (3*d[0],3*d[1])
if(VectorFigUtils.vnorm(vel)):
ax.arrow(p[0], p[1], vel[0], vel[1], head_width=0.07, head_length=0.11, fc='k', ec='k')
def plotWorld(ax, alpha=0.3, nao=None, obj=None, bDrawGround=False, color='b', centers=[], specials=[], cradius=0.1, ccolor='r', label='_'):
global world
# ax.plot([pobj[0],pnao[0]], [pobj[1],pnao[1]], linestyle='--', color='g', lw=2)
for body in world.bodies:
name = body.userData["name"]
for fixture in body.fixtures:
shape = fixture.shape
if(body.active):
if(isinstance(shape, Box2D.b2PolygonShape)):
if(name == "boxA"):
drawBox2D(ax, body, fixture, color='g', alpha=1)
elif(name == "boxB"):
drawBox2D(ax, body, fixture, color='r', alpha=1)
elif(name == "bar"):
drawBox2D(ax, body, fixture, color=color, alpha=0.1)
else:
if(fixture.density == 0.0):
drawBox2D(ax, body, fixture, color=color, alpha=0.25)
else:
drawBox2D(ax, body, fixture, color=color, alpha=alpha)
if(isinstance(shape, Box2D.b2CircleShape)):
if(name == "epuck"):
drawEpuck(ax, shape, body)
elif(name == "reward"):
drawCircle(ax, body.position, shape.radius, color=[0.5, 0.7, 0.3])
elif(name == "toy"):
drawCircle(ax, body.position, shape.radius)
else:
drawWheel(ax, shape, body)
else:
if(isinstance(shape, Box2D.b2PolygonShape)):
drawBox2D(ax, body, fixture, color=color, alpha=0.25, fill=False, linestyle='dashed')
if(isinstance(shape, Box2D.b2CircleShape)):
drawCircle(ax, body.position, 0.3, fill=False, linestyle='dashed')
if(len(centers) > 0):
plotVectors(ax, centers, cradius=cradius, specials=specials, ccolor=ccolor, label=label)
plt.grid()
