Mercurial > lcfOS

comparison python/bouncing_cube_opencl.py @ 97:5a965e9664f2

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:32:54 +0100
parents
children
comparison
equal deleted inserted replaced
96:a350055d6119 97:5a965e9664f2
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