Untitled

import numpy as np


def hat(v):
    """
    Hat function, creating for a vector its skew symmetric matrix.

    Input:
        v (np.array<float>(..., 3, 1)): The input vector.

    Output:
        (np.array<float>(..., 3, 3)): The output skew symmetric matrix.

    """
    #   Lie algebra generators
    E1 = np.array([[0., 0., 0.], [0., 0., -1.], [0., 1., 0.]])
    E2 = np.array([[0., 0., 1.], [0., 0., 0.], [-1., 0., 0.]])
    E3 = np.array([[0., -1., 0.], [1., 0., 0.], [0., 0., 0.]])
    
    return v[..., 0:1, :] * E1 + v[..., 1:2, :] * E2 + v[..., 2:3, :] * E3


def exp(v, der=False):
    """
    Exponential map. Can return derivatives

    Input:
        v (np.array<float>(..., 3, 1)): The input axis-angle vector.
        der (bool, optional) Indicates if derivative is needed or not.

    Returns:
        R (np.array<float>(..., 3, 3)): The corresponding rotation matrix.
        [dR (np.array<float>(3, ..., 3, 3)): The derivative of each rotation matrix.
                                            The matrix dR[i, ..., :, :] corresponds to
                                            the derivative d R[..., :, :] / d v[..., i, :],
                                            so the derivative of the rotation R gained 
                                            through the axis-angle vector v with respect
                                            to v_i. Note that this is not a Jacobian of
                                            any form but a vectorized version of derivatives.]

    """
    n = np.linalg.norm(v, axis=-2, keepdims=True)
    H = hat(v)
    
    with np.errstate(all='ignore'):
        R = np.identity(3) + (np.sin(n) / n) * H + ((1 - np.cos(n)) / n**2) * (H @ H)
    R = np.where(n == 0, np.identity(3), R)
    
    if der:
        sh = (3,) + tuple(1 for _ in range(v.ndim - 2)) + (3, 1)
        dR = np.swapaxes(np.expand_dims(v, axis=0), 0, -2) * H
        dR = dR + hat(np.cross(v, ((np.identity(3) - R) @ np.identity(3).reshape(sh)), axis=-2))
        dR = dR @ R
        
        n = n**2  
        with np.errstate(all='ignore'):
            dR = dR / n
        dR = np.where(n == 0, hat(np.identity(3).reshape(sh)), dR)
            
        return R, dR
    
    else:
        return R
 
    
# generate two sets of points which differ by a rotation  
np.random.seed(1001)
# Number of points.
n = 100  
#   Point set
p_1 = np.random.randn(n, 3, 1)
#   Actual axis-angle rotation
v = np.array([0.3, -0.2, 0.1]).reshape(3, 1)  # the axis-angle vector
#   Transforming the point set and adding noise.
p_2 = exp(v) @ p_1 + np.random.randn(n, 3, 1) * 1e-2


# estimate v with least sqaures, so the objective function  becomes:
# minimize v over f(v) = sum_[1<=i<=n] (||p_1_i - exp(v)p_2_i||^2)
# Due to the way least_squres is implemented we have to pass the
# individual residuals ||p_1_i - exp(v)p_2_i||^2 as ||p_1_i - exp(v)p_2_i||.
from scipy.optimize import least_squares


def loss_fn(x):
    """
        Loss is \sum_{i=1}^{K}||p2-R*p1||_2^2.
    
        Input:
            x: rotation in axis-angle format.
        Output:
            y: error value.    
    """
    R = exp(x.reshape(1, 3, 1))
    y = p_2 - R @ p_1
    y = np.linalg.norm(y, axis=-2).squeeze(-1)
    
    return y


def loss_derivative(x):
    """
        Derivative wrt the loss.
    """
    R, d_R = exp(x.reshape(1, 3, 1), der=True)  #   Rotation and derivative of rotation.
    y = p_2 - R @ p_1
    d_y = -d_R @ p_1
        
    d_y = np.sum(y * d_y, axis=-2) / np.linalg.norm(y, axis=-2)
    d_y = d_y.squeeze(-1).T
    
    return d_y


x0 = np.zeros((3))      #   
res = least_squares(loss_fn, x0, loss_derivative)

print('Original axis-angle vector: {}'.format(v.reshape(-1)))
print('Estimated axis-angle vector: {}'.format(res.x))