Mercurial > lcfOS
view python/bouncing_cube.py @ 150:4ae0e02599de
Added type check start and analyze phase
author | Windel Bouwman |
---|---|
date | Fri, 01 Mar 2013 16:53:22 +0100 |
parents | 4a37d6992bd3 |
children | ef683881c64e |
line wrap: on
line source
from PyQt4.QtGui import * from PyQt4.QtCore import * from PyQt4.QtOpenGL import QGLWidget from OpenGL.GL import * from OpenGL.GLU import gluPerspective import sys from random import random from math import pi, cos, sin, fabs, sqrt from numpy import mat, array, ones, zeros, eye from numpy.linalg import norm import numpy as np import time import scipy.integrate #import pyopencl """ Test script that lets a dice bounce. Converted from 20-sim equations into python code. 20-sim website: http://www.20sim.com """ def drawCube(w): glBegin(GL_QUADS) # Start Drawing The Cube glColor3f(0.0,1.0,0.0) # Set The Color To Blue glVertex3f( w, w,-w) # Top Right Of The Quad (Top) glVertex3f(-w, w,-w) # Top Left Of The Quad (Top) glVertex3f(-w, w, w) # Bottom Left Of The Quad (Top) glVertex3f( w, w, w) # Bottom Right Of The Quad (Top) glColor3f(1.0,0.5,0.0) # Set The Color To Orange glVertex3f( w,-w, w) # Top Right Of The Quad (Bottom) glVertex3f(-w,-w, w) # Top Left Of The Quad (Bottom) glVertex3f(-w,-w,-w) # Bottom Left Of The Quad (Bottom) glVertex3f( w,-w,-w) # Bottom Right Of The Quad (Bottom) glColor3f(1.0,0.0,0.0) # Set The Color To Red glVertex3f( w, w, w) # Top Right Of The Quad (Front) glVertex3f(-w, w, w) # Top Left Of The Quad (Front) glVertex3f(-w,-w, w) # Bottom Left Of The Quad (Front) glVertex3f( w,-w, w) # Bottom Right Of The Quad (Front) glColor3f(1.0,1.0,0.0) # Set The Color To Yellow glVertex3f( w,-w,-w) # Bottom Left Of The Quad (Back) glVertex3f(-w,-w,-w) # Bottom Right Of The Quad (Back) glVertex3f(-w, w,-w) # Top Right Of The Quad (Back) glVertex3f( w, w,-w) # Top Left Of The Quad (Back) glColor3f(0.0,0.0,1.0) # Set The Color To Blue glVertex3f(-w, w, w) # Top Right Of The Quad (Left) glVertex3f(-w, w,-w) # Top Left Of The Quad (Left) glVertex3f(-w,-w,-w) # Bottom Left Of The Quad (Left) glVertex3f(-w,-w, w) # Bottom Right Of The Quad (Left) glColor3f(1.0,0.0,1.0) # Set The Color To Violet glVertex3f( w, w,-w) # Top Right Of The Quad (Right) glVertex3f( w, w, w) # Top Left Of The Quad (Right) glVertex3f( w,-w, w) # Bottom Left Of The Quad (Right) glVertex3f( w,-w,-w) # Bottom Right Of The Quad (Right) glEnd() # Done Drawing The Quad def drawFloor(w, h): glBegin(GL_QUADS) # Start Drawing The Cube glColor3f(1.0,0.5,0.0) # Set The Color To Orange glVertex3f( w,-w,h)# Top Right Of The Quad (Bottom) glVertex3f(-w,-w,h)# Top Left Of The Quad (Bottom) glVertex3f(-w,w,h)# Bottom Left Of The Quad (Bottom) glVertex3f( w,w,h)# Bottom Right Of The Quad (Bottom) glEnd() # Done Drawing The Quad def drawAxis(): glLineWidth(0.5) glBegin(GL_LINES) glColor3f(1.0, 0.0, 0.0) glVertex3f(0,0,0) glVertex3f(1,0,0) glColor3f(0.0, 1.0, 0.0) glVertex3f(0,0,0) glVertex3f(0,1,0) glColor3f(0.0, 0.0, 1.0) glVertex3f(0,0,0) glVertex3f(0,0,1) glEnd() def cross(A, B): a = A.A1 b = B.A1 return mat(np.cross(a, b)).T def skew(X): Y = mat(zeros( (3, 3) )) a,b,c = X.A1 Y[0,1] = -c Y[0,2] = b Y[1,0] = c Y[1,2] = -a Y[2,0] = -b Y[2,1] = a return Y def adjoint(T): W = T[0:3, 0] V = T[3:6, 0] a = mat(zeros( (6,6) ) ) a[0:3, 0:3] = skew(W) a[3:6, 0:3] = skew(V) a[3:6, 3:6] = skew(W) return a def Adjoint(H): R = H[0:3, 0:3] P = H[0:3, 3] a = mat(zeros( (6,6) ) ) a[0:3, 0:3] = R a[3:6, 3:6] = R a[3:6, 0:3] = skew(P) * R return a def quatToR(q): x, y, z, w = q.A1 r = mat(eye(3)) r[0,0] = 1 - (2*y**2+2*z**2) r[0,1] = 2*x*y+2*z*w r[0,2] = 2*x*z - 2*y*w r[1,0] = 2*x*y-2*z*w r[1,1] = 1 - (2*x**2 + 2*z**2) r[1,2] = 2*y*z + 2*x*w r[2,0] = 2*x*z+2*y*w r[2,1] = 2*y*z - 2*x*w r[2,2] = 1 - (2*x**2+2*y**2) return r def rotateAbout(axis, angle): ax, ay, az = (axis/norm(axis)).A1 qx = ax*sin(angle/2.0) qy = ay*sin(angle/2.0) qz = az*sin(angle/2.0) qw = cos(angle/2.0) q = mat(array([qx,qy,qz,qw])).T return q def normalizeQuaternion(quat): x,y,z,w = quat.A1 magnitude = sqrt(x*x + y*y + z*z + w*w) x = x / magnitude y = y / magnitude z = z / magnitude w = w / magnitude quat[0, 0] = x quat[1, 0] = y quat[2, 0] = z quat[3, 0] = w return quat def VTo4x4(V): v1, v2, v3 = V.A1 return mat(array( \ [[0.0, -v3, v2, -v1], \ [ v3, 0.0, -v1, -v2], \ [-v2, v1, 0.0, -v3], \ [v1, v2, v3, 0.0] ])) def homogeneous(R,p): H = mat(eye(4)) H[0:3, 0:3] = R H[0:3, 3] = p return H def rateOfChange(states, thetime, parameters): quat = states[0:4, 0] # Orientation (4) pos = states[4:7, 0] # Position (3) P = states[7:13, 0] # Momentum (6) massI, gravity = parameters # Rigid body parts: # Forward Kinematic chain: H = homogeneous(quatToR(quat), pos) # Forward kinematics AdjX2 = mat(eye(6)) # The connectionpoint in the real world adjAdjX2_1 = adjoint(AdjX2[0:6,0]) adjAdjX2_2 = adjoint(AdjX2[0:6,1]) adjAdjX2_3 = adjoint(AdjX2[0:6,2]) adjAdjX2_4 = adjoint(AdjX2[0:6,3]) adjAdjX2_5 = adjoint(AdjX2[0:6,4]) adjAdjX2_6 = adjoint(AdjX2[0:6,5]) AdjInv2 = Adjoint(H.I) M2 = AdjInv2.T * (massI * AdjInv2) # Transfor mass to base frame MassMatrix = M2 wrenchGrav2 = mat( zeros((1,6)) ) wrenchGrav2[0, 0:3] = -cross(gravity, pos).T wrenchGrav2[0, 3:6] = gravity.T Bk = mat( zeros( (6,6) )) Bk[0:3, 0:3] = skew(P[0:3, 0]) Bk[3:6, 0:3] = skew(P[3:6, 0]) Bk[0:3, 3:6] = skew(P[3:6, 0]) # TODO: do this a cholesky: v = np.linalg.solve(MassMatrix, P) # Matrix inverse like thingy ! T2_00 = v # Calculate the relative twist! TM2 = T2_00.T * M2 twistExternal = T2_00 # Twist van het blokje TMSum2 = TM2 PDotBodies = mat( zeros( (6,1)) ) PDotBodies[0,0] = TMSum2 * (adjAdjX2_1 * T2_00) PDotBodies[1,0] = TMSum2 * (adjAdjX2_2 * T2_00) PDotBodies[2,0] = TMSum2 * (adjAdjX2_3 * T2_00) PDotBodies[3,0] = TMSum2 * (adjAdjX2_4 * T2_00) PDotBodies[4,0] = TMSum2 * (adjAdjX2_5 * T2_00) PDotBodies[5,0] = TMSum2 * (adjAdjX2_6 * T2_00) PDot = -PDotBodies - Bk * v PDot += wrenchGrav2.T ##### Contact wrench part: HB_W = H # Is H-matrix van het blokje WrenchB = mat(zeros( (1,6) )) for px in [-0.5, 0.5]: for py in [-0.5, 0.5]: for pz in [-0.5, 0.5]: HB1_B = homogeneous(mat(eye(3)), mat([px,py,pz]).T) HB1_W = HB_W * HB1_B HW1_W = homogeneous(mat(eye(3)), HB1_W[0:3,3]) HW_W1 = HW1_W.I HB_W1 = HW_W1 * HB_W AdjHB_W1 = Adjoint(HB_W1) TB_W1_W1 = AdjHB_W1 * twistExternal z = HB1_W[2, 3] vx, vy, vz = TB_W1_W1[3:6, 0].A1 if z < 0: # Contact forces: Fx = -50.0*vx Fy = -50.0*vy Fz = -z*50000.0 else: Fx = 0.0 Fy = 0.0 Fz = 0.0 # TODO: reflect impulse WrenchW1 = mat([0,0,0,0,0,Fz]) # Transform it back: WrenchB += (AdjHB_W1.T * WrenchW1.T).T ##### End of contact wrench PDot += (WrenchB * AdjInv2).T # Position and orientation rates: QOmega = VTo4x4(v[0:3, 0]) quatDot = 0.5 * QOmega * quat vel = v[3:6, 0] posDot = skew(v[0:3]) * pos + vel # The rate vector: rates = mat(zeros( (13,1) )) rates[0:4, 0] = quatDot rates[4:7, 0] = posDot rates[7:13, 0] = PDot return rates def fWrapper(y, t, parameters): y = mat(y).T dy = rateOfChange(y, t, parameters) return dy.T.A1 def SCIPY(endtime, dt, state, parameters): times = np.arange(0.0, endtime, dt) y0 = state.T.A1 res = scipy.integrate.odeint(fWrapper, y0, times, args=(parameters,)) states = [] res = res.T r,c = res.shape for ci in range(0,c): states.append(mat(res[:,ci]).T) return states def RK4(endtime, dt, state, parameters): t = 0.0 states = [] while t < endtime: newstate = mat (zeros( state.shape )) # Create a new object #### Runge Kutta integration: k1 = rateOfChange(state, t, parameters) k2 = rateOfChange(state + 0.5*dt*k1, t, parameters) k3 = rateOfChange(state + 0.5*dt*k1, t, parameters) k4 = rateOfChange(state + dt*k3, t, parameters) newstate = state + (dt/6.0)*(k1+2*k2+2*k3+k4) # Normalize quat: newstate[0:4, 0] = normalizeQuaternion(newstate[0:4, 0]) states.append(newstate) state = newstate t += dt print state[6,0], t, ' ', (t/endtime)*100.0, '%' return states def simulate(endtime, dt): PInitial = mat( zeros((6,1)) ) posInitial = mat(array([-1.2, -1.3, 2.8])).T quatInitial = rotateAbout(mat(array([0.2, 1.0, 0.4])).T, 0.5) # Parameters: gravity = mat( array([0,0,-9.81]) ).T massI = mat(eye(6)) * 1.01 # Mass matrix parameters = (massI, gravity) # The state vector! state = mat(zeros( (13,1) )) state[0:4, 0] = quatInitial state[4:7, 0] = posInitial state[7:13, 0] = PInitial return SCIPY(endtime, dt, state, parameters) class W(QGLWidget): time = 0.0 index = 0 def __init__(self, states, dt, parent=None): super(W, self).__init__(parent) self.firstRun = True self.savePNGS = False self.dt = dt self.resize(500,500) self.states = states self.UP() t = QTimer(self) t.timeout.connect(self.UP) t.start(self.dt*1000.0) def UP(self): self.time += self.dt if self.states: state = self.states[self.index] self.Dicequat = state[0:4, 0] self.Dicepos = state[4:7, 0] self.index += 1 if self.index == len(self.states): self.index = 0 self.firstRun = False # Paint: self.update() if self.firstRun: # Create png images for the movie: if self.savePNGS: pm = self.renderPixmap() pm.save('image'+str(self.index)+'.png') def initializeGL(self): glClearColor(0.0, 0.5, 0.0, 1.0) glEnable(GL_DEPTH_TEST) glDepthFunc(GL_LESS) glShadeModel(GL_SMOOTH) def resizeGL(self,w,h): glViewport(0, 0, w, h) glMatrixMode(GL_PROJECTION) glLoadIdentity() gluPerspective(45.0, float(w)/float(h), 0.1, 100.0) glMatrixMode(GL_MODELVIEW) def paintGL(self): glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT) # Clear buffers glLoadIdentity() # Reset The View glLoadIdentity() glTranslatef(0.0,-2.0,-10.0) # Move Left And Into The Screen glRotatef(-90.0, 1.0, 0.0, 0.0) drawFloor(2.0, 0.0) drawAxis() self.renderText(1.0, 0.0, 0.0, 'X') self.renderText(0.0, 1.0, 0.0, 'Y') self.renderText(0.0, 0.0, 1.0, 'Z') self.renderText(0.0,0.0,1.2,str(self.time)) x,y,z = self.Dicepos.A1 R = quatToR(self.Dicequat) glTranslatef(x, y, z) # Trick to rotate the openGL matrix: r = R.A1 rotR = (r[0], r[3], r[6], 0.0, r[1], r[4], r[7], 0.0, r[2], r[5], r[8], 0.0, 0.0, 0.0, 0.0, 1.0) glMultMatrixd(rotR) drawCube(0.6) et = 20.0 dt = 0.04 print 'starting integration... endtime =', et, ' stepsize =', dt t0 = time.time() states = simulate(et, dt) t1 = time.time() print 'That was heavy, it took me ', t1-t0, ' seconds!' app = QApplication(sys.argv) w = W(states, dt) w.show() sys.exit(app.exec_())