Mercurial > lcfOS

comparison python/bouncing_cube.py @ 95:4a37d6992bd3

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
movage
author windel
date Mon, 24 Dec 2012 13:24:59 +0100
parents python/apps/bouncing_cube.py@32078200cdd6
children ef683881c64e
comparison
equal deleted inserted replaced
94:1be00bcfaabb 95:4a37d6992bd3
1 from PyQt4.QtGui import *
2 from PyQt4.QtCore import *
3 from PyQt4.QtOpenGL import QGLWidget
4 from OpenGL.GL import *
5 from OpenGL.GLU import gluPerspective
6 import sys
7 from random import random
8 from math import pi, cos, sin, fabs, sqrt
9 from numpy import mat, array, ones, zeros, eye
10 from numpy.linalg import norm
11 import numpy as np
12 import time
13 import scipy.integrate
14 #import pyopencl
15
16 """
17 Test script that lets a dice bounce.
18 Converted from 20-sim equations into python code.
19
20 20-sim website:
21 http://www.20sim.com
22
23 """
24 def drawCube(w):
25 glBegin(GL_QUADS) # Start Drawing The Cube
26 glColor3f(0.0,1.0,0.0) # Set The Color To Blue
27 glVertex3f( w, w,-w) # Top Right Of The Quad (Top)
28 glVertex3f(-w, w,-w) # Top Left Of The Quad (Top)
29 glVertex3f(-w, w, w) # Bottom Left Of The Quad (Top)
30 glVertex3f( w, w, w) # Bottom Right Of The Quad (Top)
31
32 glColor3f(1.0,0.5,0.0) # Set The Color To Orange
33 glVertex3f( w,-w, w) # Top Right Of The Quad (Bottom)
34 glVertex3f(-w,-w, w) # Top Left Of The Quad (Bottom)
35 glVertex3f(-w,-w,-w) # Bottom Left Of The Quad (Bottom)
36 glVertex3f( w,-w,-w) # Bottom Right Of The Quad (Bottom)
37
38 glColor3f(1.0,0.0,0.0) # Set The Color To Red
39 glVertex3f( w, w, w) # Top Right Of The Quad (Front)
40 glVertex3f(-w, w, w) # Top Left Of The Quad (Front)
41 glVertex3f(-w,-w, w) # Bottom Left Of The Quad (Front)
42 glVertex3f( w,-w, w) # Bottom Right Of The Quad (Front)
43
44 glColor3f(1.0,1.0,0.0) # Set The Color To Yellow
45 glVertex3f( w,-w,-w) # Bottom Left Of The Quad (Back)
46 glVertex3f(-w,-w,-w) # Bottom Right Of The Quad (Back)
47 glVertex3f(-w, w,-w) # Top Right Of The Quad (Back)
48 glVertex3f( w, w,-w) # Top Left Of The Quad (Back)
49
50 glColor3f(0.0,0.0,1.0) # Set The Color To Blue
51 glVertex3f(-w, w, w) # Top Right Of The Quad (Left)
52 glVertex3f(-w, w,-w) # Top Left Of The Quad (Left)
53 glVertex3f(-w,-w,-w) # Bottom Left Of The Quad (Left)
54 glVertex3f(-w,-w, w) # Bottom Right Of The Quad (Left)
55
56 glColor3f(1.0,0.0,1.0) # Set The Color To Violet
57 glVertex3f( w, w,-w) # Top Right Of The Quad (Right)
58 glVertex3f( w, w, w) # Top Left Of The Quad (Right)
59 glVertex3f( w,-w, w) # Bottom Left Of The Quad (Right)
60 glVertex3f( w,-w,-w) # Bottom Right Of The Quad (Right)
61 glEnd() # Done Drawing The Quad
62
63 def drawFloor(w, h):
64 glBegin(GL_QUADS) # Start Drawing The Cube
65
66 glColor3f(1.0,0.5,0.0) # Set The Color To Orange
67 glVertex3f( w,-w,h)# Top Right Of The Quad (Bottom)
68 glVertex3f(-w,-w,h)# Top Left Of The Quad (Bottom)
69 glVertex3f(-w,w,h)# Bottom Left Of The Quad (Bottom)
70 glVertex3f( w,w,h)# Bottom Right Of The Quad (Bottom)
71 glEnd() # Done Drawing The Quad
72
73 def drawAxis():
74 glLineWidth(0.5)
75 glBegin(GL_LINES)
76 glColor3f(1.0, 0.0, 0.0)
77 glVertex3f(0,0,0)
78 glVertex3f(1,0,0)
79 glColor3f(0.0, 1.0, 0.0)
80 glVertex3f(0,0,0)
81 glVertex3f(0,1,0)
82 glColor3f(0.0, 0.0, 1.0)
83 glVertex3f(0,0,0)
84 glVertex3f(0,0,1)
85 glEnd()
86
87
88 def cross(A, B):
89 a = A.A1
90 b = B.A1
91 return mat(np.cross(a, b)).T
92
93 def skew(X):
94 Y = mat(zeros( (3, 3) ))
95 a,b,c = X.A1
96 Y[0,1] = -c
97 Y[0,2] = b
98 Y[1,0] = c
99 Y[1,2] = -a
100 Y[2,0] = -b
101 Y[2,1] = a
102 return Y
103
104 def adjoint(T):
105 W = T[0:3, 0]
106 V = T[3:6, 0]
107 a = mat(zeros( (6,6) ) )
108 a[0:3, 0:3] = skew(W)
109 a[3:6, 0:3] = skew(V)
110 a[3:6, 3:6] = skew(W)
111 return a
112
113 def Adjoint(H):
114 R = H[0:3, 0:3]
115 P = H[0:3, 3]
116 a = mat(zeros( (6,6) ) )
117 a[0:3, 0:3] = R
118 a[3:6, 3:6] = R
119 a[3:6, 0:3] = skew(P) * R
120 return a
121
122 def quatToR(q):
123 x, y, z, w = q.A1
124 r = mat(eye(3))
125 r[0,0] = 1 - (2*y**2+2*z**2)
126 r[0,1] = 2*x*y+2*z*w
127 r[0,2] = 2*x*z - 2*y*w
128 r[1,0] = 2*x*y-2*z*w
129 r[1,1] = 1 - (2*x**2 + 2*z**2)
130 r[1,2] = 2*y*z + 2*x*w
131 r[2,0] = 2*x*z+2*y*w
132 r[2,1] = 2*y*z - 2*x*w
133 r[2,2] = 1 - (2*x**2+2*y**2)
134 return r
135
136 def rotateAbout(axis, angle):
137 ax, ay, az = (axis/norm(axis)).A1
138 qx = ax*sin(angle/2.0)
139 qy = ay*sin(angle/2.0)
140 qz = az*sin(angle/2.0)
141 qw = cos(angle/2.0)
142 q = mat(array([qx,qy,qz,qw])).T
143 return q
144
145 def normalizeQuaternion(quat):
146 x,y,z,w = quat.A1
147 magnitude = sqrt(x*x + y*y + z*z + w*w)
148 x = x / magnitude
149 y = y / magnitude
150 z = z / magnitude
151 w = w / magnitude
152 quat[0, 0] = x
153 quat[1, 0] = y
154 quat[2, 0] = z
155 quat[3, 0] = w
156 return quat
157
158 def VTo4x4(V):
159 v1, v2, v3 = V.A1
160 return mat(array( \
161 [[0.0, -v3, v2, -v1], \
162 [ v3, 0.0, -v1, -v2], \
163 [-v2, v1, 0.0, -v3], \
164 [v1, v2, v3, 0.0] ]))
165
166 def homogeneous(R,p):
167 H = mat(eye(4))
168 H[0:3, 0:3] = R
169 H[0:3, 3] = p
170 return H
171
172 def rateOfChange(states, thetime, parameters):
173 quat = states[0:4, 0] # Orientation (4)
174 pos = states[4:7, 0] # Position (3)
175 P = states[7:13, 0] # Momentum (6)
176 massI, gravity = parameters
177 # Rigid body parts:
178 # Forward Kinematic chain:
179 H = homogeneous(quatToR(quat), pos) # Forward kinematics
180
181 AdjX2 = mat(eye(6)) # The connectionpoint in the real world
182 adjAdjX2_1 = adjoint(AdjX2[0:6,0])
183 adjAdjX2_2 = adjoint(AdjX2[0:6,1])
184 adjAdjX2_3 = adjoint(AdjX2[0:6,2])
185 adjAdjX2_4 = adjoint(AdjX2[0:6,3])
186 adjAdjX2_5 = adjoint(AdjX2[0:6,4])
187 adjAdjX2_6 = adjoint(AdjX2[0:6,5])
188 AdjInv2 = Adjoint(H.I)
189 M2 = AdjInv2.T * (massI * AdjInv2) # Transfor mass to base frame
190 MassMatrix = M2
191
192 wrenchGrav2 = mat( zeros((1,6)) )
193 wrenchGrav2[0, 0:3] = -cross(gravity, pos).T
194 wrenchGrav2[0, 3:6] = gravity.T
195
196 Bk = mat( zeros( (6,6) ))
197 Bk[0:3, 0:3] = skew(P[0:3, 0])
198 Bk[3:6, 0:3] = skew(P[3:6, 0])
199 Bk[0:3, 3:6] = skew(P[3:6, 0])
200
201 # TODO: do this a cholesky:
202 v = np.linalg.solve(MassMatrix, P) # Matrix inverse like thingy !
203
204 T2_00 = v # Calculate the relative twist!
205 TM2 = T2_00.T * M2
206 twistExternal = T2_00 # Twist van het blokje
207 TMSum2 = TM2
208
209 PDotBodies = mat( zeros( (6,1)) )
210 PDotBodies[0,0] = TMSum2 * (adjAdjX2_1 * T2_00)
211 PDotBodies[1,0] = TMSum2 * (adjAdjX2_2 * T2_00)
212 PDotBodies[2,0] = TMSum2 * (adjAdjX2_3 * T2_00)
213 PDotBodies[3,0] = TMSum2 * (adjAdjX2_4 * T2_00)
214 PDotBodies[4,0] = TMSum2 * (adjAdjX2_5 * T2_00)
215 PDotBodies[5,0] = TMSum2 * (adjAdjX2_6 * T2_00)
216
217 PDot = -PDotBodies - Bk * v
218 PDot += wrenchGrav2.T
219
220 ##### Contact wrench part:
221 HB_W = H # Is H-matrix van het blokje
222 WrenchB = mat(zeros( (1,6) ))
223 for px in [-0.5, 0.5]:
224 for py in [-0.5, 0.5]:
225 for pz in [-0.5, 0.5]:
226 HB1_B = homogeneous(mat(eye(3)), mat([px,py,pz]).T)
227 HB1_W = HB_W * HB1_B
228 HW1_W = homogeneous(mat(eye(3)), HB1_W[0:3,3])
229 HW_W1 = HW1_W.I
230 HB_W1 = HW_W1 * HB_W
231
232 AdjHB_W1 = Adjoint(HB_W1)
233 TB_W1_W1 = AdjHB_W1 * twistExternal
234 z = HB1_W[2, 3]
235 vx, vy, vz = TB_W1_W1[3:6, 0].A1
236 if z < 0:
237 # Contact forces:
238 Fx = -50.0*vx
239 Fy = -50.0*vy
240 Fz = -z*50000.0
241 else:
242 Fx = 0.0
243 Fy = 0.0
244 Fz = 0.0
245 # TODO: reflect impulse
246 WrenchW1 = mat([0,0,0,0,0,Fz])
247 # Transform it back:
248 WrenchB += (AdjHB_W1.T * WrenchW1.T).T
249 ##### End of contact wrench
250
251 PDot += (WrenchB * AdjInv2).T
252
253 # Position and orientation rates:
254 QOmega = VTo4x4(v[0:3, 0])
255 quatDot = 0.5 * QOmega * quat
256 vel = v[3:6, 0]
257 posDot = skew(v[0:3]) * pos + vel
258 # The rate vector:
259 rates = mat(zeros( (13,1) ))
260 rates[0:4, 0] = quatDot
261 rates[4:7, 0] = posDot
262 rates[7:13, 0] = PDot
263 return rates
264
265 def fWrapper(y, t, parameters):
266 y = mat(y).T
267 dy = rateOfChange(y, t, parameters)
268 return dy.T.A1
269
270 def SCIPY(endtime, dt, state, parameters):
271 times = np.arange(0.0, endtime, dt)
272 y0 = state.T.A1
273 res = scipy.integrate.odeint(fWrapper, y0, times, args=(parameters,))
274 states = []
275 res = res.T
276 r,c = res.shape
277 for ci in range(0,c):
278 states.append(mat(res[:,ci]).T)
279 return states
280
281 def RK4(endtime, dt, state, parameters):
282 t = 0.0
283 states = []
284 while t < endtime:
285 newstate = mat (zeros( state.shape )) # Create a new object
286
287 #### Runge Kutta integration:
288 k1 = rateOfChange(state, t, parameters)
289 k2 = rateOfChange(state + 0.5*dt*k1, t, parameters)
290 k3 = rateOfChange(state + 0.5*dt*k1, t, parameters)
291 k4 = rateOfChange(state + dt*k3, t, parameters)
292 newstate = state + (dt/6.0)*(k1+2*k2+2*k3+k4)
293
294 # Normalize quat:
295 newstate[0:4, 0] = normalizeQuaternion(newstate[0:4, 0])
296 states.append(newstate)
297
298 state = newstate
299
300 t += dt
301 print state[6,0], t, ' ', (t/endtime)*100.0, '%'
302 return states
303
304
305 def simulate(endtime, dt):
306 PInitial = mat( zeros((6,1)) )
307 posInitial = mat(array([-1.2, -1.3, 2.8])).T
308 quatInitial = rotateAbout(mat(array([0.2, 1.0, 0.4])).T, 0.5)
309 # Parameters:
310 gravity = mat( array([0,0,-9.81]) ).T
311 massI = mat(eye(6)) * 1.01 # Mass matrix
312 parameters = (massI, gravity)
313
314 # The state vector!
315 state = mat(zeros( (13,1) ))
316 state[0:4, 0] = quatInitial
317 state[4:7, 0] = posInitial
318 state[7:13, 0] = PInitial
319
320 return SCIPY(endtime, dt, state, parameters)
321
322 class W(QGLWidget):
323 time = 0.0
324 index = 0
325 def __init__(self, states, dt, parent=None):
326 super(W, self).__init__(parent)
327 self.firstRun = True
328 self.savePNGS = False
329 self.dt = dt
330 self.resize(500,500)
331 self.states = states
332 self.UP()
333 t = QTimer(self)
334 t.timeout.connect(self.UP)
335 t.start(self.dt*1000.0)
336
337 def UP(self):
338 self.time += self.dt
339 if self.states:
340 state = self.states[self.index]
341 self.Dicequat = state[0:4, 0]
342 self.Dicepos = state[4:7, 0]
343 self.index += 1
344 if self.index == len(self.states):
345 self.index = 0
346 self.firstRun = False
347 # Paint:
348 self.update()
349 if self.firstRun:
350 # Create png images for the movie:
351 if self.savePNGS:
352 pm = self.renderPixmap()
353 pm.save('image'+str(self.index)+'.png')
354
355 def initializeGL(self):
356 glClearColor(0.0, 0.5, 0.0, 1.0)
357 glEnable(GL_DEPTH_TEST)
358 glDepthFunc(GL_LESS)
359 glShadeModel(GL_SMOOTH)
360 def resizeGL(self,w,h):
361 glViewport(0, 0, w, h)
362 glMatrixMode(GL_PROJECTION)
363 glLoadIdentity()
364 gluPerspective(45.0, float(w)/float(h), 0.1, 100.0)
365 glMatrixMode(GL_MODELVIEW)
366 def paintGL(self):
367 glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT) # Clear buffers
368 glLoadIdentity() # Reset The View
369
370 glLoadIdentity()
371 glTranslatef(0.0,-2.0,-10.0) # Move Left And Into The Screen
372 glRotatef(-90.0, 1.0, 0.0, 0.0)
373 drawFloor(2.0, 0.0)
374 drawAxis()
375 self.renderText(1.0, 0.0, 0.0, 'X')
376 self.renderText(0.0, 1.0, 0.0, 'Y')
377 self.renderText(0.0, 0.0, 1.0, 'Z')
378
379 self.renderText(0.0,0.0,1.2,str(self.time))
380
381 x,y,z = self.Dicepos.A1
382 R = quatToR(self.Dicequat)
383
384 glTranslatef(x, y, z)
385 # Trick to rotate the openGL matrix:
386 r = R.A1
387 rotR = (r[0], r[3], r[6], 0.0,
388 r[1], r[4], r[7], 0.0,
389 r[2], r[5], r[8], 0.0,
390 0.0, 0.0, 0.0, 1.0)
391 glMultMatrixd(rotR)
392
393 drawCube(0.6)
394
395 et = 20.0
396 dt = 0.04
397 print 'starting integration... endtime =', et, ' stepsize =', dt
398 t0 = time.time()
399 states = simulate(et, dt)
400 t1 = time.time()
401 print 'That was heavy, it took me ', t1-t0, ' seconds!'
402 app = QApplication(sys.argv)
403 w = W(states, dt)
404 w.show()
405 sys.exit(app.exec_())
406