#!/usr/bin/env python # Copyright (c) 2019, Michael Boyle # See LICENSE file for details: <https://github.com/moble/spherical_functions/blob/master/LICENSE> from __future__ import print_function, division, absolute_import # Try to keep imports to a minimum; from the standard library as much # as possible. We have to conda install all dependencies, and it's # not right to make Travis do too much work. import math import cmath import numpy as np import quaternion import spherical_functions as sf import pytest slow = pytest.mark.slow precision_Wigner_D_element = 4.e-14 def test_Wigner_D_linear_indices(ell_max): for l_min in range(ell_max): for l_max in range(l_min + 1, ell_max + 1): LMpM = sf.LMpM_range(l_min, l_max) assert len(LMpM) == sf._total_size_D_matrices(l_min, l_max) for ell in range(l_min, l_max + 1): assert list(LMpM[sf._linear_matrix_offset(ell, l_min)]) == [ell, -ell, -ell] for ell in range(l_min, l_max + 1): for mp in range(-ell, ell + 1): i = sf._linear_matrix_diagonal_index(ell, mp) o = sf._linear_matrix_offset(ell, l_min) assert list(LMpM[i + o]) == [ell, mp, mp] for m in range(-ell, ell + 1): i = sf._linear_matrix_index(ell, mp, m) o = sf._linear_matrix_offset(ell, l_min) assert list(LMpM[i + o]) == [ell, mp, m] def test_Wigner_D_matrices_negative_argument(Rs, ell_max): # For integer ell, D(R)=D(-R) LMpM = sf.LMpM_range(0, ell_max) a = np.empty((LMpM.shape[0],), dtype=complex) b = np.empty((LMpM.shape[0],), dtype=complex) for R in Rs: sf._Wigner_D_matrices(R.a, R.b, 0, ell_max, a) sf._Wigner_D_matrices(-R.a, -R.b, 0, ell_max, b) assert np.allclose(a, b, rtol=ell_max * precision_Wigner_D_element) @slow def test_Wigner_D_elements_representation_property(Rs, ell_max): # Test the representation property for special and random angles # Try half-integers, too ell_max = min(8, ell_max) twoLMpM = np.round(2*sf.LMpM_range_half_integer(0, ell_max)).astype(int) print("") D1 = np.empty((twoLMpM.shape[0],), dtype=complex) D2 = np.empty((twoLMpM.shape[0],), dtype=complex) D12 = np.empty((twoLMpM.shape[0],), dtype=complex) for i, R1 in enumerate(Rs): print("\t{0} of {1}: R1 = {2}".format(i+1, len(Rs), R1)) for j, R2 in enumerate(Rs): # print("\t\t{0} of {1}: R2 = {2}".format(j+1, len(Rs), R2)) R12 = R1 * R2 sf._Wigner_D_element(R1.a, R1.b, twoLMpM, D1) sf._Wigner_D_element(R2.a, R2.b, twoLMpM, D2) sf._Wigner_D_element(R12.a, R12.b, twoLMpM, D12) M12 = np.array([np.sum([D1[sf._Wigner_index(twoell, twomp, twompp)] * D2[sf._Wigner_index(twoell, twompp, twom)] for twompp in range(-twoell, twoell + 1, 2)]) for twoell in range(2*ell_max + 1) for twomp in range(-twoell, twoell + 1, 2) for twom in range(-twoell, twoell + 1, 2)]) # if not np.allclose(M12, D12, atol=ell_max * precision_Wigner_D_element): # for k in range(min(100, M12.size)): # print(twoLMpM[k], "\t", abs(D12[k]-M12[k]), "\t\t", D12[k], "\t", M12[k]) # print(D12.shape, M12.shape) assert np.allclose(M12, D12, atol=ell_max * precision_Wigner_D_element) @slow def test_Wigner_D_matrices_representation_property(Rs, ell_max): # Test the representation property for special and random angles # Can't try half-integers because Wigner_D_matrices doesn't accept them ell_max = min(8, ell_max) LMpM = sf.LMpM_range(0, ell_max) print("") D1 = np.empty((LMpM.shape[0],), dtype=complex) D2 = np.empty((LMpM.shape[0],), dtype=complex) D12 = np.empty((LMpM.shape[0],), dtype=complex) for i, R1 in enumerate(Rs): print("\t{0} of {1}: R1 = {2}".format(i+1, len(Rs), R1)) for j, R2 in enumerate(Rs[i:]): # print("\t\t{0} of {1}: R2 = {2}".format(j+1, len(Rs), R2)) R12 = R1 * R2 sf._Wigner_D_matrices(R1.a, R1.b, 0, ell_max, D1) sf._Wigner_D_matrices(R2.a, R2.b, 0, ell_max, D2) sf._Wigner_D_matrices(R12.a, R12.b, 0, ell_max, D12) M12 = np.array([np.sum([D1[sf.LMpM_index(ell, mp, mpp, 0)] * D2[sf.LMpM_index(ell, mpp, m, 0)] for mpp in range(-ell, ell + 1)]) for ell in range(ell_max + 1) for mp in range(-ell, ell + 1) for m in range(-ell, ell + 1)]) assert np.allclose(M12, D12, atol=ell_max * precision_Wigner_D_element) @slow def test_Wigner_D_matrix_inverse(Rs, ell_max): # Ensure that the matrix of the inverse rotation is the inverse of # the matrix of the rotation print() for i, R in enumerate(Rs): print("\t{0} of {1}: {2}".format(i+1, len(Rs), R)) for twoell in range(2*ell_max + 1): LMpM = np.array([[twoell/2, twomp/2, twom/2] for twomp in range(-twoell, twoell+1, 2) for twom in range(-twoell, twoell+1, 2)]) D1 = sf.Wigner_D_element(R.a, R.b, LMpM) D2 = sf.Wigner_D_element(R.a.conjugate(), -R.b, LMpM) D1 = D1.reshape((twoell + 1, twoell + 1)) D2 = D2.reshape((twoell + 1, twoell + 1)) assert np.allclose(D1.dot(D2), np.identity(twoell + 1), atol=ell_max ** 4 * precision_Wigner_D_element, rtol=ell_max ** 4 * precision_Wigner_D_element) @slow def test_Wigner_D_element_symmetries(Rs, ell_max): LMpM = sf.LMpM_range(0, ell_max) # D_{mp,m}(R) = (-1)^{mp+m} \bar{D}_{-mp,-m}(R) MpPM = np.array([mp + m for ell in range(ell_max + 1) for mp in range(-ell, ell + 1) for m in range(-ell, ell + 1)]) LmMpmM = np.array( [[ell, -mp, -m] for ell in range(ell_max + 1) for mp in range(-ell, ell + 1) for m in range(-ell, ell + 1)]) print() for R in Rs: print("\t", R) assert np.allclose(sf.Wigner_D_element(R, LMpM), (-1.) ** MpPM * np.conjugate(sf.Wigner_D_element(R, LmMpmM)), atol=ell_max ** 2 * precision_Wigner_D_element, rtol=ell_max ** 2 * precision_Wigner_D_element) # D is a unitary matrix, so its conjugate transpose is its # inverse. D(R) should equal the matrix inverse of D(R^{-1}). # So: D_{mp,m}(R) = \bar{D}_{m,mp}(\bar{R}) LMMp = np.array( [[ell, m, mp] for ell in range(ell_max + 1) for mp in range(-ell, ell + 1) for m in range(-ell, ell + 1)]) for R in Rs: print("\t", R) assert np.allclose(sf.Wigner_D_element(R, LMpM), np.conjugate(sf.Wigner_D_element(R.inverse(), LMMp)), atol=ell_max ** 4 * precision_Wigner_D_element, rtol=ell_max ** 4 * precision_Wigner_D_element) def test_Wigner_D_element_roundoff(Rs, ell_max): LMpM = sf.LMpM_range(0, ell_max) # Test rotations with |Ra|<1e-15 expected = [((-1.) ** ell if mp == -m else 0.0) for ell in range(ell_max + 1) for mp in range(-ell, ell + 1) for m in range(-ell, ell + 1)] assert np.allclose(sf.Wigner_D_element(quaternion.x, LMpM), expected, atol=ell_max * precision_Wigner_D_element, rtol=ell_max * precision_Wigner_D_element) expected = [((-1.) ** (ell + m) if mp == -m else 0.0) for ell in range(ell_max + 1) for mp in range(-ell, ell + 1) for m in range(-ell, ell + 1)] assert np.allclose(sf.Wigner_D_element(quaternion.y, LMpM), expected, atol=ell_max * precision_Wigner_D_element, rtol=ell_max * precision_Wigner_D_element) for theta in np.linspace(0, 2 * np.pi): expected = [((-1.) ** (ell + m) * (np.cos(theta) + 1j * np.sin(theta)) ** (2 * m) if mp == -m else 0.0) for ell in range(ell_max + 1) for mp in range(-ell, ell + 1) for m in range(-ell, ell + 1)] assert np.allclose(sf.Wigner_D_element(np.cos(theta) * quaternion.y + np.sin(theta) * quaternion.x, LMpM), expected, atol=ell_max * precision_Wigner_D_element, rtol=ell_max * precision_Wigner_D_element) # Test rotations with |Rb|<1e-15 expected = [(1.0 if mp == m else 0.0) for ell in range(ell_max + 1) for mp in range(-ell, ell + 1) for m in range(-ell, ell + 1)] assert np.allclose(sf.Wigner_D_element(quaternion.one, LMpM), expected, atol=ell_max * precision_Wigner_D_element, rtol=ell_max * precision_Wigner_D_element) expected = [((-1.) ** m if mp == m else 0.0) for ell in range(ell_max + 1) for mp in range(-ell, ell + 1) for m in range(-ell, ell + 1)] assert np.allclose(sf.Wigner_D_element(quaternion.z, LMpM), expected, atol=ell_max * precision_Wigner_D_element, rtol=ell_max * precision_Wigner_D_element) for theta in np.linspace(0, 2 * np.pi): expected = [((np.cos(theta) + 1j * np.sin(theta)) ** (2 * m) if mp == m else 0.0) for ell in range(ell_max + 1) for mp in range(-ell, ell + 1) for m in range(-ell, ell + 1)] assert np.allclose(sf.Wigner_D_element(np.cos(theta) * quaternion.one + np.sin(theta) * quaternion.z, LMpM), expected, atol=ell_max * precision_Wigner_D_element, rtol=ell_max * precision_Wigner_D_element) def test_Wigner_D_element_underflow(Rs, ell_max): # NOTE: This is a delicate test, which depends on the result underflowing exactly when expected. # In particular, it should underflow to 0.0 when |mp+m|>32, but should never undeflow to 0.0 # when |mp+m|<32. So it's not the end of the world if this test fails, but it does mean that # the underflow properties have changed, so it might be worth a look. assert sf.ell_max >= 15 # Test can't work if this has been set lower eps = 1.e-10 ell_max = max(sf.ell_max, ell_max) # Test |Ra|=1e-10 LMpM = np.array( [[ell, mp, m] for ell in range(ell_max + 1) for mp in range(-ell, ell + 1) for m in range(-ell, ell + 1) if abs(mp + m) > 32]) R = np.quaternion(eps, 1, 0, 0).normalized() assert np.all(sf.Wigner_D_element(R, LMpM) == 0j) LMpM = np.array( [[ell, mp, m] for ell in range(ell_max + 1) for mp in range(-ell, ell + 1) for m in range(-ell, ell + 1) if abs(mp + m) < 32]) R = np.quaternion(eps, 1, 0, 0).normalized() assert np.all(sf.Wigner_D_element(R, LMpM) != 0j) # Test |Rb|=1e-10 LMpM = np.array( [[ell, mp, m] for ell in range(ell_max + 1) for mp in range(-ell, ell + 1) for m in range(-ell, ell + 1) if abs(m - mp) > 32]) R = np.quaternion(1, eps, 0, 0).normalized() assert np.all(sf.Wigner_D_element(R, LMpM) == 0j) LMpM = np.array( [[ell, mp, m] for ell in range(ell_max + 1) for mp in range(-ell, ell + 1) for m in range(-ell, ell + 1) if abs(m - mp) < 32]) R = np.quaternion(1, eps, 0, 0).normalized() assert np.all(sf.Wigner_D_element(R, LMpM) != 0j) def test_Wigner_D_element_overflow(Rs, ell_max): assert sf.ell_max >= 15 # Test can't work if this has been set lower ell_max = max(sf.ell_max, ell_max) LMpM = sf.LMpM_range(0, ell_max) # Test |Ra|=1e-10 R = np.quaternion(1.e-10, 1, 0, 0).normalized() assert np.all(np.isfinite(sf.Wigner_D_element(R, LMpM))) # Test |Rb|=1e-10 R = np.quaternion(1, 1.e-10, 0, 0).normalized() assert np.all(np.isfinite(sf.Wigner_D_element(R, LMpM))) def slow_Wignerd(beta, ell, mp, m): # https://en.wikipedia.org/wiki/Wigner_D-matrix#Wigner_.28small.29_d-matrix Prefactor = math.sqrt( math.factorial(ell + mp) * math.factorial(ell - mp) * math.factorial(ell + m) * math.factorial(ell - m)) s_min = int(round(max(0, round(m - mp)))) s_max = int(round(min(round(ell + m), round(ell - mp)))) assert isinstance(s_max, int), type(s_max) assert isinstance(s_min, int), type(s_min) return Prefactor * sum([((-1.) ** (mp - m + s) * math.cos(beta / 2.) ** (2 * ell + m - mp - 2 * s) * math.sin(beta / 2.) ** (mp - m + 2 * s) / float(math.factorial(ell + m - s) * math.factorial(s) * math.factorial(mp - m + s) * math.factorial(ell - mp - s))) for s in range(s_min, s_max + 1)]) def slow_Wigner_D_element(alpha, beta, gamma, ell, mp, m): # https://en.wikipedia.org/wiki/Wigner_D-matrix#Definition_of_the_Wigner_D-matrix return cmath.exp(-1j * mp * alpha) * slow_Wignerd(beta, ell, mp, m) * cmath.exp(-1j * m * gamma) @slow def test_Wigner_D_element_values(special_angles, ell_max): LMpM = sf.LMpM_range_half_integer(0, ell_max // 2) # Compare with more explicit forms given in Euler angles print("") for alpha in special_angles: print("\talpha={0}".format(alpha)) # Need to show some progress to Travis for beta in special_angles: print("\t\tbeta={0}".format(beta)) for gamma in special_angles: a = np.conjugate(np.array([slow_Wigner_D_element(alpha, beta, gamma, ell, mp, m) for ell,mp,m in LMpM])) b = sf.Wigner_D_element(quaternion.from_euler_angles(alpha, beta, gamma), LMpM) # if not np.allclose(a, b, # atol=ell_max ** 6 * precision_Wigner_D_element, # rtol=ell_max ** 6 * precision_Wigner_D_element): # for i in range(min(a.shape[0], 100)): # print(LMpM[i], "\t", abs(a[i]-b[i]), "\t\t", a[i], "\t", b[i]) assert np.allclose(a, b, atol=ell_max ** 6 * precision_Wigner_D_element, rtol=ell_max ** 6 * precision_Wigner_D_element) @slow def test_Wigner_D_matrix(Rs, ell_max): for l_min in [0, 1, 2, ell_max // 2, ell_max - 1]: print("") for l_max in range(l_min + 1, ell_max + 1): print("\tWorking on (l_min,l_max)=({0},{1})".format(l_min, l_max)) LMpM = sf.LMpM_range(l_min, l_max) for R in Rs: elements = sf.Wigner_D_element(R, LMpM) matrix = np.empty(LMpM.shape[0], dtype=complex) sf._Wigner_D_matrices(R.a, R.b, l_min, l_max, matrix) assert np.allclose(elements, matrix, atol=1e3 * l_max * ell_max * precision_Wigner_D_element, rtol=1e3 * l_max * ell_max * precision_Wigner_D_element) @slow def test_Wigner_D_input_types(Rs, special_angles, ell_max): LMpM = sf.LMpM_range(0, ell_max // 2) # Compare with more explicit forms given in Euler angles print("") for alpha in special_angles: print("\talpha={0}".format(alpha)) # Need to show some progress to Travis for beta in special_angles: for gamma in special_angles: a = sf.Wigner_D_element(alpha, beta, gamma, LMpM) b = sf.Wigner_D_element(quaternion.from_euler_angles(alpha, beta, gamma), LMpM) assert np.allclose(a, b, atol=ell_max ** 6 * precision_Wigner_D_element, rtol=ell_max ** 6 * precision_Wigner_D_element) for R in Rs: a = sf.Wigner_D_element(R, LMpM) b = sf.Wigner_D_element(R.a, R.b, LMpM) assert np.allclose(a, b, atol=ell_max ** 6 * precision_Wigner_D_element, rtol=ell_max ** 6 * precision_Wigner_D_element) def test_Wigner_D_signatures(Rs): """There are two ways to call the WignerD function: with an array of Rs, or with an array of (ell,mp,m) values. This test ensures that the results are the same in both cases.""" # from spherical_functions.WignerD import _Wigner_D_elements ell_max = 6 ell_mp_m = sf.LMpM_range(0, ell_max) Ds1 = np.zeros((Rs.size, ell_mp_m.shape[0]), dtype=np.complex) Ds2 = np.zeros_like(Ds1) for i, R in enumerate(Rs): Ds1[i, :] = sf.Wigner_D_element(R, ell_mp_m) for i, (ell, mp, m) in enumerate(ell_mp_m): Ds2[:, i] = sf.Wigner_D_element(Rs, ell, mp, m) assert np.allclose(Ds1, Ds2, rtol=3e-15, atol=3e-15)