def test_exp_Im_alpha_second_rank():
n = 100
cos_alpha = (np.random.rand(n) - 0.5) * 2.0
exp_im_alpha = clib.get_exp_Im_alpha(n, cos_alpha, False)
arr = np.zeros((4, n), dtype=np.complex128)
alpha = np.arccos(cos_alpha)
arr[-1] = np.exp(1j * alpha)
arr[-2] = np.exp(2j * alpha)
start = 2 * n
assert np.allclose(exp_im_alpha[start:], arr[2:].ravel())
def test_exp_Im_alpha_fourth_rank():
n = 100
cos_alpha = (np.random.rand(n) - 0.5) * 2.0
exp_im_alpha = clib.get_exp_Im_alpha(n, cos_alpha, True)
arr = np.zeros((4, n), dtype=np.complex128)
alpha = np.arccos(cos_alpha)
arr[-1] = np.exp(1j * alpha)
arr[-2] = np.exp(2j * alpha)
arr[-3] = np.exp(3j * alpha)
arr[-4] = np.exp(4j * alpha)
assert np.allclose(exp_im_alpha, arr.ravel())
def test__batch_wigner_rotation():
n = 40
n_octants = 1
alpha = np.random.rand(n) * np.pi / 2.0
beta = np.random.rand(n) * np.pi / 2.0
cos_beta = np.cos(beta)
sin_beta = np.sqrt(1.0 - cos_beta ** 2)
exp_I_beta = cos_beta + 1j * sin_beta
wigner_2j_matrices = clib.wigner_d_matrices_from_exp_I_beta(
2, True, exp_I_beta
).ravel()
wigner_4j_matrices = clib.wigner_d_matrices_from_exp_I_beta(
4, True, exp_I_beta
).ravel()
R4 = np.asarray(
[0 - 0.2j, 0, 0 + 0.5j, 0, 0 + 0.1j, 0, 0 - 0.5j, 0, 0 + 0.2j],
dtype=np.complex128,
)
R2 = np.asarray([0 + 0.5j, 0, 0 + 0.1j, 0, 0 - 0.5j], dtype=np.complex128)
cos_alpha = np.cos(alpha)
exp_im_alpha = clib.get_exp_Im_alpha(n, cos_alpha, True).ravel()
exp_im_alpha_in = np.empty(4 * n, dtype=np.complex128)
exp_im_alpha_in[:] = exp_im_alpha.copy()
w2, w4 = clib.__batch_wigner_rotation(
n, n_octants, wigner_2j_matrices, R2, wigner_4j_matrices, R4, exp_im_alpha
)
# assert np.allclose(exp_im_alpha_in, exp_im_alpha, atol=1e-15)
alpha_octants = []
for i in range(n_octants):
alpha_octants.append(alpha + i * np.pi / 2.0)
alpha_octants = np.asarray(alpha_octants).ravel()
cos_alpha_octants = np.cos(alpha_octants)
beta_octants = []
for i in range(n_octants):
beta_octants.append(beta)
beta_octants = np.asarray(beta_octants).ravel()
cos_beta_octants = np.cos(beta_octants)
w2_1 = clib.__wigner_rotation_2(2, cos_alpha_octants, cos_beta_octants, R2).ravel()
w4_1 = clib.__wigner_rotation_2(4, cos_alpha_octants, cos_beta_octants, R4).ravel()
assert np.allclose(w2, w2_1, atol=1e-12)
np.testing.assert_almost_equal(w4, w4_1, decimal=8)