Test the performance of scipy vs the rotation functions in litebird_sim

scipy_rotations_test.py

#!/usr/bin/env python3
# -*- encoding: utf-8 -*-

import numpy as np
from litebird_sim import all_compute_pointing_and_polangle
from scipy.spatial.transform import Rotation as R


def polangle_scipy(theta, phi, poldir):
    cos_psi = np.clip(-np.sin(phi) * poldir[:, 0] + np.cos(phi) * poldir[:, 1], -1, 1)
    sin_psi = np.clip(
        (-np.cos(theta) * np.cos(phi) * poldir[:, 0])
        + (-np.cos(theta) * np.sin(phi) * poldir[:, 1])
        + (np.sin(theta) * poldir[:, 2]),
        -1,
        1,
    )

    return np.arctan2(sin_psi, cos_psi)


def test_scipy(quat_matrix):
    N = quat_matrix.shape[0]
    result_matrix = np.empty((N, 3))

    r = R.from_quat(quat_matrix)

    result = r.apply(np.repeat(np.array([0.0, 0.0, 1.0]).reshape(1, -1), N, axis=0))
    theta = np.arccos(result[:, 2])
    phi = np.arctan2(result[:, 1], result[:, 0])

    poldir = r.apply(np.repeat(np.array([1.0, 0.0, 0.0]).reshape(1, -1), N, axis=0))
    psi = polangle_scipy(theta, phi, poldir)

    return np.array([theta, phi, psi]).transpose()


def test_lbs(quat_matrix):
    result_matrix = np.empty((quat_matrix.shape[0], 3))
    all_compute_pointing_and_polangle(result_matrix, quat_matrix)
    return result_matrix


quat_matrix = np.random.rand(10000, 4)
# Normalize the quaternions
quat_matrix /= np.sqrt(np.sum(quat_matrix * quat_matrix, axis=1))[:, None]

print(
    "Consistency check:", np.allclose(test_scipy(quat_matrix), test_lbs(quat_matrix),),
)

# poetry run ipython
# In [1]: %run scipy_rotations_test.py
# Consistency test: True
# 
# In [2]: %timeit test_lbs(quat_matrix)
# 1.55 ms ± 12.2 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
#
# In [3]: %timeit test_scipy(quat_matrix)
# 3.48 ms ± 18.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)