Mercurial > lcfOS
diff python/bouncing_cube.py @ 95:4a37d6992bd3
movage
author | windel |
---|---|
date | Mon, 24 Dec 2012 13:24:59 +0100 |
parents | python/apps/bouncing_cube.py@32078200cdd6 |
children | ef683881c64e |
line wrap: on
line diff
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/python/bouncing_cube.py Mon Dec 24 13:24:59 2012 +0100 @@ -0,0 +1,406 @@ +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_()) +