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_())
+