Data generation

data.py

def get_R(angles):
    ''' Get rotation matrix from three rotation angles (radians). right-handed.
    Args:
        angles: [3,]. x, y, z angles
    Returns:
        R: [3, 3]. rotation matrix.
    '''
    x, y, z = angles.astype(np.float32)
    # x
    Rx = np.array([[1, 0, 0],
                   [0, np.cos(x), -np.sin(x)],
                   [0, np.sin(x), np.cos(x)]])
    # y
    Ry = np.array([[np.cos(y), 0, np.sin(y)],
                   [0, 1, 0],
                   [-np.sin(y), 0, np.cos(y)]])
    # z
    Rz = np.array([[np.cos(z), -np.sin(z), 0],
                   [np.sin(z), np.cos(z), 0],
                   [0, 0, 1]])

    R = Rz.dot(Ry.dot(Rx))
    return R

INIT_AXES = np.array([[1,0,0], [0,1,0], [0,0,1]]).astype(np.float32)

def get_data(min_angle_rad=-np.pi, max_angle_rad=np.pi, nb_examples=30000, split=0.1):
    data = defaultdict(lambda: [])
    for i in tqdm(range(nb_examples)):
        angles = np.random.uniform(min_angle_rad, max_angle_rad, size=3)
        R = get_R(angles)
        q = Quaternion(matrix=R, rtol=1e-05, atol=1e-05).elements.astype(np.float32)
        data['R'].append(R)
        data['angles'].append(angles)
        # full quaternion
        data['q'].append(q)
        # quaternion constrained to one hemisphere
        data['qh'].append(-q if q[0] < 0 else q)
        data['rotated_axes'].append(R.dot(INIT_AXES.T).T)
    for key in data.keys():
        data[key] = np.array(data[key])
    return data