find rotation matrix to align two normalised vectors in 3d avoiding trig calls

align_two_3d_vectors.py

import numpy as np

def align(a, b):
    """
    Find r such that r @ a = b when a and b are normalised vectors
    :param a: normalised vector of length 3
    :param b: normalised vector of length 3
    :return: rotation matrix
    """
    # cross product to find axis about which rotation should occur
    axis = np.cross(a, b)
    # dot product equals cosine of angle between normalised vectors
    cos_angle = np.dot(a, b)
    # k is a constant which appears as a factor in the rotation matrix
    k = 1 / (1 + cos_angle)

    # construct rotation matrix
    r = np.empty((3, 3))
    r[0, 0] = (axis[0] * axis[0] * k) + cos_angle
    r[0, 1] = (axis[1] * axis[0] * k) - axis[2]
    r[0, 2] = (axis[2] * axis[0] * k) + axis[1]
    r[1, 0] = (axis[0] * axis[1] * k) + axis[2]
    r[1, 1] = (axis[1] * axis[1] * k) + cos_angle
    r[1, 2] = (axis[2] * axis[1] * k) - axis[0]
    r[2, 0] = (axis[0] * axis[2] * k) - axis[1]
    r[2, 1] = (axis[1] * axis[2] * k) + axis[0]
    r[2, 2] = (axis[2] * axis[2] * k) + cos_angle
    
    return r