Mercurial > lcfOS
view python/other/bouncing_cube.py @ 296:9417caea2eb3
Directorized some backend files
author | Windel Bouwman |
---|---|
date | Sun, 01 Dec 2013 13:36:58 +0100 |
parents | 534b94b40aa8 |
children |
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 """ Test script that lets a dice bounce. Converted from 20-sim equations into python code. 20-sim website: http://www.20sim.com """ # OpenGL drawing madness: 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() # Math helper functions: 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): """ Create a H matrix from rotation and position """ H = mat(eye(4)) H[0:3, 0:3] = R H[0:3, 3] = p return H class Simulation: def simulate(self, endtime, dt): """ Evaluate the state of the system over a given time """ state = self.initial() times = np.arange(0.0, endtime, dt) y0 = state.T.A1 def fWrapper(y, t): y = mat(y).T dy = self.rateOfChange(y, t) return dy.T.A1 res = scipy.integrate.odeint(fWrapper, y0, times) states = [] res = res.T r,c = res.shape for ci in range(0,c): states.append(mat(res[:,ci]).T) self.result = states return states class BouncingCube(Simulation): """ Simulation of a bouncing cube """ def __init__(self): # Parameters: self.gravity = mat( array([0,0,-9.81]) ).T self.massI = mat(eye(6)) * 1.01 # Mass matrix def initial(self): """ Return initial pose """ 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) # The state vector! state = mat(zeros( (13,1) )) state[0:4, 0] = quatInitial state[4:7, 0] = posInitial state[7:13, 0] = PInitial return state def rateOfChange(self, states, thetime): """ Calculates rate of change given the current states. """ quat = states[0:4, 0] # Orientation (4) pos = states[4:7, 0] # Position (3) P = states[7:13, 0] # Momentum (6) massI, gravity = self.massI, self.gravity # 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 True: # Contact forces: Fx = 0 #np.exp(-5.0*vx) Fy = 0 # np.exp(-5.0*vy) Fz = np.exp(-z*10.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 draw(self, state): Dicequat = state[0:4, 0] Dicepos = state[4:7, 0] x, y, z = Dicepos.A1 R = quatToR(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) class RollingPendulum(Simulation): """ A 2D pendulum with a rolling contact. """ def __init__(self): self.R = 0.01 self.L = 1 self.M = 1 self.g = 9.81 def initial(self): i = mat(zeros((2,1))) i[0, 0] = return i def rateOfChange(self, states, thetime): phi, omega = states[0:2, 0] rates = mat(zeros((2,1))) return rates def draw(self, state): pass #drawCube(0.6) class WobbelingWheel(Simulation): """ A wheel pushed sideways over a surface """ def __init__(self): self.r1 = 0.5 self.r2 = 0.05 def initial(self): i = mat(zeros((5,1))) return i def rateOfChange(self, states, thetime): x, y = states[0:2, 0] rates = mat(zeros((5,1))) return rates def draw(self, state): pass #drawCube(0.6) class W(QGLWidget): time = 0.0 index = 0 def __init__(self, sim): super().__init__() self.firstRun = True self.savePNGS = False self.index = 0 self.sim = sim self.resize(500,500) self.UP() t = QTimer(self) t.timeout.connect(self.UP) t.start(40) def UP(self): if self.sim.result: if self.index + 1 < len(self.sim.result): self.index += 1 else: self.index = 0 # 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 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)) if self.sim.result: self.sim.draw(self.sim.result[self.index]) et = 1.5 dt = 0.04 # sim = BouncingCube() sim = WobbelingWheel() print('starting integration... endtime =', et, ' stepsize =', dt) t0 = time.time() states = sim.simulate(et, dt) t1 = time.time() print('That was heavy, it took me ', t1 -t0, ' seconds!') app = QApplication(sys.argv) w = W(sim) w.show() sys.exit(app.exec_())