Gauss-Newton update using Quaternions

gauss_newton.py

#!/usr/bin/env python3
#
# Rotation-only registration of point clouds using Quaternions and Gauss-Newton
#

import numpy as np
from sympy import (symbols, diff, Matrix, lambdify)

def rot_quat(v, a):
  """
  Initialize unit quaternion from rotation vector and angle
  """
  q = np.zeros(4)
  q[0:3] = v * np.sin(a / 2)
  q[3] = np.cos(a / 2)
  return q

def rot_vec(q,x):
  """
  Fast rotation, can be derived by working out q*x*conj(q) and
  using some identities, see
  "https://gamedev.stackexchange.com/questions/28395/rotating-vector3-by-a-quaternion"
  """
  w = q[3]
  v = q[0:3]
  return 2 * np.dot(v,x) * v + (2 * w**2 - 1.0) * x + 2.0 * w * np.cross(v,x)

def quat_to_rot(q):
  """
  Compute rotation matrix from quaternion.
  F. Landis Markey, NASA Goddard Space Flight Center
  doi: 10.2514/1.31730
  """
  a = q[0]
  b = q[1]
  c = q[2]
  d = q[3]
  
  R = np.matrix([
      [d * d + a * a - b * b - c * c, 2 * (a * b - c * d), 2 * (a * c + b * d)],
      [2 * (a * b + c * d), d * d + b * b - a * a - c * c, 2 * (b * c - a * d)],
      [2 * (a * c - b * d), 2 * (b * c + a * d), d * d + c * c - b * b - a * a]
  ])
  return R

def dR(q):
  """
  Partial derivatives of the rotation, R(q), q = [x,y,z,w]. Differentiation
  is carried out using SymPy.
  """
  if dR.func is None:
    x,y,z,w = symbols('x y z w')
    R = Matrix(quat_to_rot([x,y,z,w]))
    dR.func = lambdify((x,y,z,w), Matrix.vstack(
        diff(R,x),
        diff(R,y),
        diff(R,z),
        diff(R,w))) 
  return dR.func(q[0], q[1], q[2], q[3])
dR.func = None

def gn_estimate(source, target, normalize=True):
  nMaxIterations = 10;
  errs = []
  # Initial guess
  q = np.ones((1,4))
  for _ in range(nMaxIterations):
    if normalize:
      # Normalize
      # Question: Should we do this - it converges faster
      q = q / np.linalg.norm(q)
    # It wont properly flatten, unless it is a new copy!!!
    q0 = np.copy(q).flatten()
    # Gauss-newton update
    dq, err = gn_update(source, target, q0)
    
    print("err: %g" % (err))
    # q \in SO(3) - a unit quaternion
    # dq lies in tangent space at q, TqQ [not the tangent space at the identity, the Lie algebra]
    # qn is not a unit quaternion
    qn = q + dq.T
    # Questions:
    #  1. Can we take the logarithm of q, pull-back dq to the Lie algebra, update(+) and exponentiate
    #  2. Is it possible to make dq a valid rotation and update using q -> q * dq
    errs.append(err)
    q = qn
    if (err < 1e-10):
      break
    
  return np.copy(q).flatten(), errs

def gn_update(source, target, q):
  """
  Gauss-Newton update for rotation only registration
  """
  H = np.zeros((4,4))
  b = np.zeros((4,1))
  R = quat_to_rot(q)
  chi = 0.0
  for i in range(len(source)):
    s = source[i,:]
    t = target[i,:]
    J = np.dot(dR(q), s).reshape(4,3).T
    H = H + np.dot(J.T, J)
    r = (np.dot(R, s) - t)
    chi = chi + np.dot(r,r.T)
    b = b + np.dot(J.T, r.T)
  dq = -np.dot(np.linalg.pinv(H), b)

  return dq, chi[0,0]

if __name__ == "__main__":
  import sys
  noise = False
  q = rot_quat(np.r_[0, np.sqrt(0.5), np.sqrt(0.5)], np.pi/2.0)

  # Source data
  source = np.r_[np.c_[0, -1, 0],
                 np.c_[1, 0, 0]]
  # source = np.random.rand(50,3) # Random points

  # Target data
  target = np.r_[[rot_vec(q, source[i]) for i in range(len(source))]]

  # Add some noise
  if noise:
    target = target + 0.02 * np.random.randn(*target.shape)
  q1, errs = gn_estimate(source, target, normalize=True)
  print("q_est(%d): " % (len(errs)) + str(q1) +
        "\nq_true: " + str(q))

# Local variables: #
# tab-width: 2 #
# python-indent: 2 #
# indent-tabs-mode: nil #
# End: #