Skip to content

Instantly share code, notes, and snippets.

@admalledd

admalledd/mechanum_test.py

Created January 11, 2013 02:03
Show Gist options
  • Save admalledd/4507354 to your computer and use it in GitHub Desktop.
Save admalledd/4507354 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
mechanum_test.py
import math
import pygame
from pygame.locals import *
FONT = None
class robot_(object):
def __init__(self,joystick):
self.joystick=joystick
self.robot_box = pygame.Rect((0,0),(200,200))
self.robot_box.center = 400,300
#debug displays
#wheel motor values, not in self.displays due to use in other parts than self.draw()
self.wheel_a = var_display('A',(50,50))
self.wheel_b = var_display('B',(50,75))
self.wheel_c = var_display('C',(50,100))
self.wheel_d = var_display('D',(50,125))
self.displays = {
'angle':var_display('AGL',(175,50)),
'magnitude':var_display('MAG',(175,75)),
'vx':var_display('VX',(175,100)),
'vy':var_display('VY',(175,125)),
'z':var_display('ROT',(175,150))
}
def calculate_move_vector(self):
#get joystick movement and output to self.m_vector angle in radians and magnitude.
self.x = self.joystick.get_axis(0)
self.y = self.joystick.get_axis(1)
self.z = self.joystick.get_axis(3)
vx,vy = self.x,self.y #incase we need to math on these, use intermediate values
distance = math.hypot(vx,vy)
if distance >1:
distance = 1.0
elif distance <-1:
distance = -1.0
try:
angle = math.atan2(vy,vx)
angle = angle * -1.0
if angle < 0:
angle = angle + (2*math.pi)
except Exception as e:
angle = math.pi/2
if distance < 0.01:
angle = 0
self.displays['angle'].set(math.degrees(angle))
self.displays['magnitude'].set(distance)
self.displays['vx'].set(vx)
self.displays['vy'].set(vy)
self.m_vector=(angle,distance,(vx*100,vy*100))
def calculate_wheel_values(self):
self.displays['z'].set(self.z)
x,y,z = self.x,self.y,self.z #save my sanity...
if x == 0 and y == 0:
y = 0.00001 #set to make sure that we fall into ONE of these calculations for rotation.
#Upper Left Quadrant
if x < 0 and y < 0:
if z> 0:
a = (-(x*x)+(y*y)+(z*z))/(-x-y+z)
b = (-(x*x)-(y*y)+(z*z))/(-x-y+z)
c = ((x*x)+(y*y)+(z*z))/(-x-y+z)
d = ((x*x)-(y*y)+(z*z))/(-x-y+z)
elif z < 0:
a = (-(x*x)+(y*y)-(z*z))/(-x-y-z)
b = (-(x*x)-(y*y)-(z*z))/(-x-y-z)
c = ((x*x)+(y*y)-(z*z))/(-x-y-z)
d = ((x*x)-(y*y)-(z*z))/(-x-y-z)
else:
a = (-(x*x)+(y*y)+(z*z))/(-x-y+z)
b = (-(x*x)-(y*y)+(z*z))/(-x-y+z)
c = ((x*x)+(y*y)+(z*z))/(-x-y+z)
d = ((x*x)-(y*y)+(z*z))/(-x-y+z)
#Upper Right Quadrant
elif x> 0 and y < 0:
if z> 0:
a = ((x*x)+(y*y)+(z*z)) / (x-y+z)
b = ((x*x)-(y*y)+(z*z)) / (x-y+z)
c = (-(x*x)+(y*y)+(z*z)) / (x-y+z)
d = (-(x*x)-(y*y)+(z*z)) / (x-y+z)
elif z < 0:
a = ((x*x)+(y*y)-(z*z)) / (x-y-z)
b = ((x*x)-(y*y)-(z*z)) / (x-y-z)
c = (-(x*x)+(y*y)-(z*z)) / (x-y-z)
d = (-(x*x)-(y*y)-(z*z)) / (x-y-z)
else:
a = ((x*x)+(y*y)-(z*z)) / (x-y+z)
b = ((x*x)-(y*y)-(z*z)) / (x-y+z)
c = (-(x*x)+(y*y)-(z*z)) / (x-y+z)
d = (-(x*x)-(y*y)-(z*z)) / (x-y+z)
#Back Right Quadrant
elif x> 0 and y> 0:
if z> 0:
a = ((x*x)-(y*y)+(z*z))/(x+y+z)
b = ((x*x)+(y*y)+(z*z))/(x+y+z)
c = (-(x*x)-(y*y)+(z*z))/(x+y+z)
d = (-(x*x)+(y*y)+(z*z))/(x+y+z)
elif z < 0:
a = ((x*x)-(y*y)-(z*z))/(x+y-z)
b = ((x*x)+(y*y)-(z*z))/(x+y-z)
c = (-(x*x)-(y*y)-(z*z))/(x+y-z)
d = (-(x*x)+(y*y)-(z*z))/(x+y-z)
else:
a = ((x*x)-(y*y)+(z*z))/(x+y+z)
b = ((x*x)+(y*y)+(z*z))/(x+y+z)
c = (-(x*x)-(y*y)+(z*z))/(x+y+z)
d = (-(x*x)+(y*y)+(z*z))/(x+y+z)
#Back Left Quadrant
elif x < 0 and y> 0:
if z> 0:
a = (-(x*x)-(y*y)+(z*z))/(-x+y+z)
b = (-(x*x)+(y*y)+(z*z))/(-x+y+z)
c = ((x*x)-(y*y)+(z*z))/(-x+y+z)
d = ((x*x)+(y*y)+(z*z))/(-x+y+z)
elif z < 0:
a = (-(x*x)-(y*y)-(z*z))/(-x+y-z)
b = (-(x*x)+(y*y)-(z*z))/(-x+y-z)
c = ((x*x)-(y*y)-(z*z))/(-x+y-z)
d = ((x*x)+(y*y)-(z*z))/(-x+y-z)
else:
a = (-(x*x)-(y*y)+(z*z))/(-x+y+z)
b = (-(x*x)+(y*y)+(z*z))/(-x+y+z)
c = ((x*x)-(y*y)+(z*z))/(-x+y+z)
d = ((x*x)+(y*y)+(z*z))/(-x+y+z)
#Strafe Right
elif x> 0 and y == 0:
if z> 0:
a = ((x*x)+(z*z))/(x+z)
b = ((x*x)+(z*z))/(x+z)
c = (-(x*x)+(z*z))/(x+z)
d = (-(x*x)+(z*z))/(x+z)
elif z < 0:
a = ((x*x)-(z*z))/(x-z)
b = ((x*x)-(z*z))/(x-z)
c = (-(x*x)-(z*z))/(x-z)
d = (-(x*x)-(z*z))/(x-z)
else:
a = ((x*x)+(z*z))/(x+z)
b = ((x*x)+(z*z))/(x+z)
c = (-(x*x)+(z*z))/(x+z)
d = (-(x*x)+(z*z))/(x+z)
#Strafe Left
elif x < 0 and y == 0:
if z> 0:
a = (-(x*x)+(z*z))/(-x+z)
b = (-(x*x)+(z*z))/(-x+z)
c = ((x*x)+(z*z))/(-x+z)
d = ((x*x)+(z*z))/(-x+z)
elif z < 0:
a = (-(x*x)-(z*z))/(-x-z)
b = (-(x*x)-(z*z))/(-x-z)
c = ((x*x)-(z*z))/(-x-z)
d = ((x*x)-(z*z))/(-x-z)
else:
a = (-(x*x)+(z*z))/(-x+z)
b = (-(x*x)+(z*z))/(-x+z)
c = ((x*x)+(z*z))/(-x+z)
d = ((x*x)+(z*z))/(-x+z)
#Backward
elif y> 0 and x == 0:
if z> 0:
a = (-(y*y)+(z*z))/(y+z)
b = ((y*y)+(z*z))/(y+z)
c = (-(y*y)+(z*z))/(y+z)
d = ((y*y)+(z*z))/(y+z)
elif z < 0:
a = (-(y*y)-(z*z))/(y-z)
b = ((y*y)-(z*z))/(y-z)
c = (-(y*y)-(z*z))/(y-z)
d = ((y*y)-(z*z))/(y-z)
else:
a = (-(y*y)+(z*z))/(y+z)
b = ((y*y)+(z*z))/(y+z)
c = (-(y*y)+(z*z))/(y+z)
d = ((y*y)+(z*z))/(y+z)
#Forward
elif y < 0 and x == 0:
if z> 0:
a = ((y*y)+(z*z))/(-y+z)
b = (-(y*y)+(z*z))/(-y+z)
c = ((y*y)+(z*z))/(-y+z)
d = (-(y*y)+(z*z))/(-y+z)
elif z < 0:
a = ((y*y)-(z*z))/(-y-z)
b = (-(y*y)-(z*z))/(-y-z)
c = ((y*y)-(z*z))/(-y-z)
d = (-(y*y)-(z*z))/(-y-z)
else:
a = ((y*y)-(z*z))/(-y+z)
b = (-(y*y)-(z*z))/(-y+z)
c = ((y*y)-(z*z))/(-y+z)
d = (-(y*y)-(z*z))/(-y+z)
self.wheel_a.set(a)
self.wheel_b.set(b)
self.wheel_c.set(c)
self.wheel_d.set(d)
def draw(self,screen):
#draw robot "base" first, layer other things over it
pygame.draw.rect(screen,(0,255,75),self.robot_box,2)
#get new movement vectors to draw
self.calculate_move_vector()
self.calculate_wheel_values()
pygame.draw.arc(screen, (255,255,255), self.robot_box, 0, self.m_vector[0], 1)
vector_end = (self.m_vector[2][0]+self.robot_box.center[0]),(self.m_vector[2][1]+self.robot_box.center[1])
pygame.draw.line(screen, (255,0,0),self.robot_box.center,vector_end,1)
#rot line
rot_line_start = self.robot_box.midtop[0], self.robot_box.midtop[1]-10 #10 pixels above
rot_line_end = rot_line_start[0]+(self.z*100),rot_line_start[1]
pygame.draw.line(screen,(255,0,255),rot_line_start,rot_line_end,4)
self.wheel_a.draw(screen)
self.wheel_b.draw(screen)
self.wheel_c.draw(screen)
self.wheel_d.draw(screen)
for key,disp in self.displays.items():
disp.draw(screen)
class var_display(object):
def __init__(self,name,cpos):
self.center = cpos
self.val = 0
self.name = name
def set(self,val):
self.val = val
def draw(self,screen):
surf = FONT.render('%s:%0.3f'%(self.name,self.val),True,(255,255,255))
rect = surf.get_rect()
rect.center = self.center
screen.blit(surf,rect)
def main():
pygame.init()
pygame.joystick.init()
global FONT
FONT = pygame.font.SysFont("Courier", 20)
screen = pygame.display.set_mode((800,600),pygame.SRCALPHA)
clock=pygame.time.Clock()
j = pygame.joystick.Joystick(0)
j.init()
robot = robot_(j)
while True:
timedelta = clock.tick(30)
#print "fps: %0.3f"%clock.get_fps()
events = pygame.event.get()
for event in events:
if event.type == QUIT:
pygame.quit()
return
elif event.type == MOUSEBUTTONDOWN and event.button==1:
print event.pos
elif event.type == KEYDOWN:
if event.key == K_ESCAPE:
pygame.quit()
return
elif event.key == K_SPACE:
#clear screen
screen.fill((0, 0, 0))
elif event.type == JOYBUTTONDOWN:
if event.button == 9:
#clear screen
screen.fill((0, 0, 0))
elif event.button == 10:
#left stick
joy.left_stick.pos=(SIZE[0]/2,SIZE[1]/2)
joy.left_stick.x_spd=0.0
joy.left_stick.y_spd=0.0
elif event.button == 11:
#right stick
joy.right_stick.pos=(SIZE[0]/2,SIZE[1]/2)
joy.right_stick.x_spd=0.0
joy.right_stick.y_spd=0.0
elif event.button == 6:
#left color
joy.left_stick.new_color()
elif event.button == 7:
#right color
joy.right_stick.new_color()
screen.fill((0,0,0))
robot.draw(screen)
pygame.display.flip()
if __name__ == '__main__':
main()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment