dte.py

#
# imports
import numpy as np

# import typing


def skew_symmetric_matrix(vector: np.typing.NDArray) -> np.typing.NDArray:
    """
    Create sklew symmetric matrix from input vector (i.e. a square matrix whose transpose equals its negative or a_ij = -a_ji)
    :param vector: input vector
    :return: skew symmetrix matrix
    """
    x, y, z = vector
    return np.array([
        [0, -z, y],
        [z, 0, -x],
        [-y, x, 0]
    ])


def exponential_map(vector: np.typing.NDArray) -> np.typing.NDArray:
    """
    Convert rotation vector into corresponding rotation matrix
    :param vector: rotation vector
    :type vector: np.typing.NDArray
    :return: rotation matrix
    :rtype: NDArray[Any]
    """
    angle = np.linalg.norm(vector)
    # For very small angles, use Taylor expansion with the rotation vector ω = vector
    if angle < 1e-12:
        ssk = skew_symmetric_matrix(vector)
        return np.eye(3) + ssk + 0.5 * (ssk @ ssk)
    axis = vector / angle
    ss_matrix = skew_symmetric_matrix(axis)
    ss_squared = ss_matrix @ ss_matrix
    # rotation matrix (Rodrigues)
    return (np.eye(3) + ss_matrix * np.sin(angle) +
            ss_squared * (1 - np.cos(angle)))


def log_map(R: np.typing.NDArray) -> np.typing.NDArray:
    """
    Convert rotation matrix into corresopnding rotation vector
    
    :param R: Rotation matrix
    :type R: np.typing.NDArray
    :return: rotation vector
    :rtype: NDArray[Any]
    """
    
    # calculate rotation angle
    angle_cos = (np.trace(R) - 1) / 2
    angle_cos = np.clip(angle_cos, -1, 1)
    theta = np.arccos(angle_cos)

    # return no rotation when angle is approximately zero
    if theta < 1e-10:
        return np.zeros((3,))

    sin_theta = np.sin(theta)
    # handle numerical instability when theta is near pi
    if np.abs(sin_theta) < 1e-6:
        # theta ~ pi, extract axis robustly from diagonal
        rx = np.sqrt(max(0.0, (R[0, 0] + 1) / 2))
        ry = np.sqrt(max(0.0, (R[1, 1] + 1) / 2))
        rz = np.sqrt(max(0.0, (R[2, 2] + 1) / 2))
        # determine signs from off-diagonals
        if (R[2, 1] - R[1, 2]) < 0:
            rx = -rx
        if (R[0, 2] - R[2, 0]) < 0:
            ry = -ry
        if (R[1, 0] - R[0, 1]) < 0:
            rz = -rz
        omega = np.array([rx, ry, rz])
    else:
        # standard case
        omega = (1.0 / (2.0 * sin_theta)) * np.array([
            (R[2, 1] - R[1, 2]),
            (R[0, 2] - R[2, 0]),
            (R[1, 0] - R[0, 1])
        ])

    # calculate and return rotation vector
    v = omega * theta
    return v


def random_rotation(angle_deg: float, rng: np.random.Generator | None = None) -> np.typing.NDArray:
    """
    :param angle_deg: rotation angle in degrees
    :return: rotation matrix
    """
    if rng is None:
        r = np.random.rand(3)
    else:
        r = rng.random(3)
    axis = r - 0.5
    axis = axis / np.linalg.norm(axis)
    angle = angle_deg/180*np.pi
    # convert to rotation vector
    rotation_vector = angle * axis
    return exponential_map(rotation_vector)


def weiszfeld_algorithm(X: np.typing.NDArray, n_iterations: int, threshold: float) -> np.typing.NDArray: 
    """
    :param X: input matrix (3xn) containing n xyz coordinate points
    :param n_iterations: iterations to compute
    :param threshold: threshold to stop estimation at
    :return: computed minimum point
    """
    
    X = np.asarray(X, dtype=float)
    (a, _) = X.shape
    assert (a == 3)
    # initial guess
    y = np.mean(X, axis=1, keepdims=True)  # (3,1)
    epsilon = np.finfo(float).eps
    
    for _ in range(n_iterations):
        # calculate norm
        distance = np.linalg.norm(X - y, axis=0)  # (1, n)
        
        # (special case) check if estimate in points
        if (np.any(distance < epsilon)): # special case, break and return y
            return y
        
        # calculate weights
        weights = np.sum(1 / distance)  # (scalar)
        num = np.sum(X / distance, axis=1, keepdims=True)  # (3, 1)
        
        # take sum
        estimate = num / weights
        
        # check if within threshold and return if so 
        if (np.linalg.norm(y - estimate) < threshold):
            y = estimate
            break
            
        # update Y
        y = estimate
        
    # return estimate minimum
    return y


def compute_optimal_rotation(R_gt: np.typing.NDArray, R_est: np.typing.NDArray):
    """
    :resources: https://trumpf.id.au/pubs/Hartley_Aftab_Trumpf_CVPR2011.pdf
    :param R_gt: rotation ground truth
    :param R_est: rotation estimate
    :return:  R_geo, errors, mean_error, rms_error
    """
    # init
    (_, _, n) = R_gt.shape
    errors = np.zeros((n))
    R_transform = np.zeros((3,3,n))

    # compute relative rotations
    for i in range(n):
        R_transform[:,:,i] = R_gt[:,:, i]@R_est[:,:,i].T
        
    # initial guess
    vectors = R_transform.reshape((9,n))
    
    vectors_median = np.median(vectors, axis=1) # (9,)
    M_median = vectors_median.reshape((3,3))
    
    (U, S, Vt) = np.linalg.svd(M_median)
    R_median = U @ Vt
    
    if (np.linalg.det(R_median) < 0):
        U[:, -1] *= -1
        R_median = U @ Vt
    
    R_geo = R_median # (3,3)
    
    # iterative L1-type refinement
    for _ in range(10):
        step_num = np.zeros((3,))
        step_den = 0
        
        for i in range(n):
            # map relative rotation vector to tagent vector at identity
            R_rel = R_transform[:,:,i] @ R_geo.T
            v = log_map(R_rel) # map to R3
            
            if np.linalg.norm(v) < 1e-6:
                v = v + np.random.randn(3)*1e-4

            v_norm = np.linalg.norm(v)

            step_num = step_num + (v) / v_norm
            step_den += (1 / v_norm)
            
        # weighted average direction 
        delta = step_num / step_den # (3,)
        
        # scalar angle
        delta_angle = np.linalg.norm(delta)
        
        # break if slight angle
        if delta_angle < 1e-12:
            break
        
        # unit vector
        delta_axis = delta / delta_angle
        
        ss_delta = skew_symmetric_matrix(delta_axis)
        
        # Rodrigues'formula
        R_delta = np.eye(3) + ss_delta*np.sin(delta_angle) + (ss_delta@ss_delta) * (1 - np.cos(delta_angle))
        
        R_geo = R_delta @ R_geo # (3,3)
    
    # compute angular errors
    for i in range(n):
        R_aligned = R_geo @ R_est[:,:,i] # (3,3)
        R_err = R_gt[:,:,i] @ R_aligned.T # (3,3)
        
        tr = np.trace(R_err)
        cos_angle = (tr - 1) / 2
        cos_angle = np.clip(cos_angle, -1, 1)
        angle_rad = np.arccos(cos_angle)
        errors[i] = np.abs(angle_rad)
        
    mean_error = np.mean(errors)
    rms_error = np.sqrt((np.mean(errors*errors)))
    
    return (R_geo, errors, mean_error, rms_error)
    

def compute_relative_scale(t_gtge: np.typing.NDArray, t_gt: np.typing.NDArray,t_estge: np.typing.NDArray,t_est: np.typing.NDArray) -> float:
    """
    Compute relative scale given 
    
    :param t_gtge: trajectory groundtruth geometric median
    :type t_gtge: np.typing.NDArray
    :param t_gt: trajectory groundtruth 
    :type t_gt: np.typing.NDArray
    :param t_estge: trajectory estimated geometric median
    :type t_estge: np.typing.NDArray
    :param t_est: trajectory estimated 
    :type t_est: np.typing.NDArray
    """
    
    est_med = np.median(np.linalg.norm(t_est - t_estge, axis=0)) # (n)
    gt_med = np.median(np.linalg.norm(t_gt - t_gtge, axis=0)) # (n)
    
    if (est_med < 1e-12):
        est_med = np.finfo(float).eps
    scale = gt_med/est_med # scalar
    
    return scale


def compute_translation(t_gc: np.typing.NDArray, scale: float, R: np.typing.NDArray, t_estge: np.typing.NDArray):
    """
    Computer translation 
    
    :param t_gc: ground truth reference point
    :type t_gc: np.typing.NDArray
    :param scale: scale factor that is applied to geometry
    :type scale: float
    :param R: rotation matrix aligning estimated point to ground truth
    :type R: np.typing.NDArray
    :param t_estge: eference point from estimated trajectory (e.g., geometric median).
    :type t_estge: np.typing.NDArray
    """
    
    return t_gc - scale*R@t_estge

    
def compute_distance(k: int, t_gc: np.typing.NDArray, t_gcge: np.typing.NDArray, s_ge: float, R_ge:np.typing.NDArray, t_ec:np.typing.NDArray, t_ge: np.typing.NDArray):
    """
    Docstring for compute_distance
    
    :param k: upper bound parameter
    :type k: int
    :param t_gc: ground truth points (3, 3, n)
    :type t_gc: np.typing.NDArray
    :param t_gcge: ground truth point's geometric median (3, 3)
    :type t_gcge: np.typing.NDArray
    :param s_ge: scale factor that is applied to geometry
    :type s_ge: float
    :param R_ge: rotation matrix that aligns estimated point to ground truth
    :type R_ge: np.typing.NDArray
    :param t_ec: estimated point
    :type t_ec: np.typing.NDArray
    :param t_ge: translation from estimated point to ground truth
    :type t_ge: np.typing.NDArray
    """
    
    t_p = s_ge*R_ge@t_ec + t_ge
    d = np.linalg.norm(t_gc - t_p, axis = 0) # (n)
    

    u = k * np.median(np.linalg.norm(t_gc - t_gcge, axis=0))

    epsilon = np.minimum(d, u) / u
    
    return epsilon
    

def generate_synthetic_data(n_total: int, n_outliers: int,
                            sigma_xyz: float, sigma_R: float):
    """
    Generate synthetic data
    :param n_total: total number of points
    :param n_outliers: number of outlier points
    :param sigma_xyz: ?
    :sigma_R: ?
    """
    # init RNG
    rng = np.random.default_rng()

    # 1. generate ground truth points and rotations
    xyz_gt = np.empty((3, n_total))
    R_gt = np.empty((3, 3, n_total))
    for i in range(n_total):
        # in range [-0.5, 0.5]
        xyz_gt[:, i] = rng.random(3) - 0.5
        # in range [0, 360]
        R_gt[:, :, i] = random_rotation(rng.random() * 360.0, rng)

    # 2. generate outlier points and rotations
    xyz_outliers = np.empty((3, n_total))
    R_outliers = np.empty((3, 3, n_total))
    for i in range(n_total):
        # in range [-5, 5]
        xyz_outliers[:, i] = (rng.random(3) - 0.5) * 10.0
        # in range [0, 360]
        R_outliers[:, :, i] = random_rotation(rng.random() * 360.0, rng)

    # 3. build input points with noise (concatenate GT and outliers)
    xyz_input = np.concatenate((xyz_gt[:, : n_total - n_outliers],
                                xyz_outliers[:, :n_outliers]), axis=1)
    xyz_input = xyz_input + rng.normal(0.0, sigma_xyz, xyz_input.shape)
    R_input = np.empty(R_gt.shape)

    # random global transform
    R_transform = random_rotation(rng.random() * 360.0, rng)  # gt to est
    scale_transform = rng.random() * 10.0
    translation_transform = rng.random(3) * 100.0

    # 4. Apply global similarity transform to all points/rotations
    for i in range(n_total):
        # rotation noise then apply GT
        R_input[:, :, i] = random_rotation(abs(rng.normal(0.0, sigma_R)), rng) @ R_gt[:, :, i]
        # multiply by rotation matrix (matrix multiplication)
        R_input[:, :, i] = R_transform @ R_input[:, :, i]
        # multiply by scale and translate
        xyz_input[:, i] = (scale_transform * xyz_input[:, i] + translation_transform)

    # 5. Replace last rotations with outlier rotations
    start = n_total - n_outliers
    if start < n_total:
        R_input[:, :, start:] = R_outliers[:, :, :n_outliers]

    return (xyz_gt, R_gt, xyz_input, R_input)



def compute_dte_dtr(k: int, xyz_gt: np.typing.NDArray,
                    R_gt: np.typing.NDArray,
                    xyz_est: np.typing.NDArray,
                    R_est: np.typing.NDArray):
    # 1. compute geometric medians
    t_gt = weiszfeld_algorithm(xyz_gt, 50, 1e-5)
    t_ec = weiszfeld_algorithm(xyz_est, 50, 1e-5)
    
    # 2. find optimal rotation
    (R_geo, errors, mean_error, rms_error) = compute_optimal_rotation(R_gt, R_est)
    
    # 3. compute relative scale
    scale = compute_relative_scale(t_gt, xyz_gt, t_ec, xyz_est)
    
    # 4. compute translation
    translation = compute_translation(t_gt, scale, R_geo, t_ec)
    
    # 5. compute and winorize distances
    epsilon = compute_distance(k, xyz_gt, t_gt, scale, R_geo, xyz_est, translation)
    n = epsilon.size
    
    # 6. define DTE
    suml1 = np.sum(epsilon / n) 
    suml2 = np.sqrt(np.sum(epsilon*epsilon / n))
    
    dte = 1/2 * (suml1 + suml2)
    # print(f"sum L1: {suml1}\nsum L2: {suml2}")
    
    dtr = 1/2*(mean_error + rms_error)
    
    return (dte, dtr)
    

# main

if __name__ == "__main__":
    (xyz_gt, R_gt, xyz_input, R_input) = generate_synthetic_data(1000, 100,
                                                                 0.005, 10)
    
    k = 5
    dte, dtr = compute_dte_dtr(k, xyz_gt, R_gt, xyz_input, R_input)
    
    print(f"DTE: {dte}")
    print(f"DTR: {dtr}")