def test_Wigner_D_representation_property(Rs, ell_max_slow, eps): # Test the representation property for special and random angles # For each l, 𝔇ˡₘₚ,ₘ(R1 * R2) = Σₘₚₚ 𝔇ˡₘₚ,ₘₚₚ(R1) * 𝔇ˡₘₚₚ,ₘ(R2) import time print("") t1 = time.perf_counter() D1 = np.zeros(sf.WignerDsize(0, ell_max_slow), dtype=complex) D2 = np.zeros(sf.WignerDsize(0, ell_max_slow), dtype=complex) D12 = np.zeros(sf.WignerDsize(0, ell_max_slow), dtype=complex) wigner = sf.Wigner(ell_max_slow) for i, R1 in enumerate(Rs): print(f"\t{i+1} of {len(Rs)}: R1 = {R1}") for j, R2 in enumerate(Rs): # print(f"\t\t{j+1} of {len(Rs)}: R2 = {R2}") R12 = R1 * R2 wigner.D(R1, out=D1) wigner.D(R2, out=D2) wigner.D(R12, out=D12) for ell in range(ell_max_slow + 1): ϵ = (2 * ell + 1)**2 * eps i1 = sf.WignerDindex(ell, -ell, -ell) i2 = sf.WignerDindex(ell, ell, ell) shape = (2 * ell + 1, 2 * ell + 1) Dˡ1 = D1[i1:i2 + 1].reshape(shape) Dˡ2 = D2[i1:i2 + 1].reshape(shape) Dˡ12 = D12[i1:i2 + 1].reshape(shape) assert np.allclose(Dˡ1 @ Dˡ2, Dˡ12, rtol=ϵ, atol=ϵ), ell t2 = time.perf_counter() print(f"\tFinished in {t2-t1:.4f} seconds.")
def test_modes_grid(ell_max, eps): ell_max = max(3, ell_max) np.random.seed(1234) wigner = sf.Wigner(ell_max) ϵ = 10 * (2 * ell_max + 1) * eps n_theta = n_phi = 2 * ell_max + 1 rotors = quaternionic.array.from_spherical_coordinates( sf.theta_phi(n_theta, n_phi)) for s in range(-2, 2 + 1): ell_min = abs(s) a1 = np.random.rand(11, sf.Ysize(ell_min, ell_max) * 2).view(complex) m1 = sf.Modes(a1, spin_weight=s, ell_min=ell_min, ell_max=ell_max) f1 = m1.grid(n_theta, n_phi) assert f1.shape == m1.shape[:-1] + rotors.shape[:-1] sYlm = np.zeros((sf.Ysize(0, ell_max), ) + rotors.shape[:-1], dtype=complex) for i, Rs in enumerate(rotors): for j, R in enumerate(Rs): wigner.sYlm(s, R, out=sYlm[:, i, j]) f2 = np.tensordot(m1.view(np.ndarray), sYlm, axes=([-1], [0])) assert f2.shape == m1.shape[:-1] + rotors.shape[:-1] assert np.allclose( f1.ndarray, f2, rtol=ϵ, atol=ϵ), f"max|f1-f2|={np.max(np.abs(f1.ndarray-f2))} > ϵ={ϵ}"
def test_Wigner_D_inverse_property(Rs, ell_max, eps): # Test the inverse property for special and random angles # For each l, 𝔇ˡₘₚ,ₘ(R⁻¹) should be the inverse matrix of 𝔇ˡₘₚ,ₘ(R) D1 = np.zeros(sf.WignerDsize(0, ell_max), dtype=complex) D2 = np.zeros(sf.WignerDsize(0, ell_max), dtype=complex) wigner = sf.Wigner(ell_max) for i, R in enumerate(Rs): # print(f"\t{i+1} of {len(Rs)}: R = {R}") wigner.D(R, out=D1) wigner.D(R.inverse, out=D2) for ell in range(ell_max + 1): ϵ = (2 * ell + 1)**2 * eps i1 = sf.WignerDindex(ell, -ell, -ell) i2 = sf.WignerDindex(ell, ell, ell) shape = (2 * ell + 1, 2 * ell + 1) Dˡ1 = D1[i1:i2 + 1].reshape(shape) Dˡ2 = D2[i1:i2 + 1].reshape(shape) assert np.allclose(Dˡ1 @ Dˡ2, np.identity(2 * ell + 1), rtol=ϵ, atol=ϵ), ell # print(f"\t{i+1} of {len(Rs)}: R = {R}") D1 = wigner.D(R) D2 = wigner.D(R.inverse) for ell in range(ell_max + 1): ϵ = (2 * ell + 1)**2 * eps i1 = sf.WignerDindex(ell, -ell, -ell) i2 = sf.WignerDindex(ell, ell, ell) shape = (-1, 2 * ell + 1, 2 * ell + 1) Dˡ1 = D1[..., i1:i2 + 1].reshape(shape) Dˡ2 = D2[..., i1:i2 + 1].reshape(shape) assert np.allclose(Dˡ1 @ Dˡ2, np.identity(2 * ell + 1), rtol=ϵ, atol=ϵ), ell
def test_Wigner_D_underflow(Rs, ell_max, eps): # 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 underflow 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. epsilon = 1.e-10 ϵ = 5 * ell_max * eps D = np.zeros(sf.WignerDsize(0, ell_max), dtype=complex) wigner = sf.Wigner(ell_max) ell_mp_m = sf.WignerDrange(0, ell_max) # Test |Ra|=1e-10 R = quaternionic.array(epsilon, 1, 0, 0).normalized wigner.D(R, out=D) # print(R.to_euler_angles.tolist()) # print(D.tolist()) non_underflowing_indices = np.abs(ell_mp_m[:, 1] + ell_mp_m[:, 2]) < 32 assert np.all(D[non_underflowing_indices] != 0j) underflowing_indices = np.abs(ell_mp_m[:, 1] + ell_mp_m[:, 2]) > 32 assert np.all(D[underflowing_indices] == 0j) # Test |Rb|=1e-10 R = quaternionic.array(1, epsilon, 0, 0).normalized wigner.D(R, out=D) non_underflowing_indices = np.abs(ell_mp_m[:, 1] - ell_mp_m[:, 2]) < 32 assert np.all(D[non_underflowing_indices] != 0j) underflowing_indices = np.abs(ell_mp_m[:, 1] - ell_mp_m[:, 2]) > 32 assert np.all(D[underflowing_indices] == 0j)
def test_wigner_evaluate(horner, ell_max_slow, eps): import time ell_max = max(3, ell_max_slow) np.random.seed(1234) ϵ = 10 * (2 * ell_max + 1) * eps n_theta = n_phi = 2 * ell_max + 1 max_s = 2 wigner = sf.Wigner(ell_max, mp_max=max_s) max_error = 0.0 total_time = 0.0 for rotors in [ quaternionic.array.from_spherical_coordinates( sf.theta_phi(n_theta, n_phi)), quaternionic.array(np.random.rand(n_theta, n_phi, 4)).normalized ]: for s in range(-max_s, max_s + 1): ell_min = abs(s) a1 = np.random.rand(7, sf.Ysize(ell_min, ell_max) * 2).view(complex) a1[:, sf.Yindex(ell_min, -ell_min, ell_min):sf. Yindex(abs(s), -abs(s), ell_min)] = 0.0 m1 = sf.Modes(a1, spin_weight=s, ell_min=ell_min, ell_max=ell_max) t1 = time.perf_counter() f1 = wigner.evaluate(m1, rotors, horner=horner) t2 = time.perf_counter() assert f1.shape == m1.shape[:-1] + rotors.shape[:-1] # print(f"Evaluation for s={s} took {t2-t1:.4f} seconds") # print(f1.shape) sYlm = np.zeros((sf.Ysize(0, ell_max), ) + rotors.shape[:-1], dtype=complex) for i, Rs in enumerate(rotors): for j, R in enumerate(Rs): wigner.sYlm(s, R, out=sYlm[:, i, j]) f2 = np.tensordot(m1.view(np.ndarray), sYlm, axes=([-1], [0])) assert f2.shape == m1.shape[:-1] + rotors.shape[:-1] assert np.allclose(f1, f2, rtol=ϵ, atol=ϵ), ( f"max|f1-f2|={np.max(np.abs(f1-f2))} > ϵ={ϵ}\n\n" f"s = {s}\n\nrotors = {rotors.tolist()}\n\n" # f"f1 = {f1.tolist()}\n\nf2 = {f2.tolist()}" ) max_error = max(np.max(np.abs(f1 - f2)), max_error) total_time += t2 - t1 print() print(f"\tmax_error[{horner}] = {max_error}") print(f"\ttotal_time[{horner}] = {total_time}")
def test_Wigner_D_non_overflow(ell_max): D = np.zeros(sf.WignerDsize(0, ell_max), dtype=complex) wigner = sf.Wigner(ell_max) # Test |Ra|=1e-10 R = quaternionic.array(1.e-10, 1, 0, 0).normalized assert np.all(np.isfinite(wigner.D(R, out=D))) # Test |Rb|=1e-10 R = quaternionic.array(1, 1.e-10, 0, 0).normalized assert np.all(np.isfinite(wigner.D(R, out=D)))
def test_sYlm_conjugation(special_angles, ell_max_slow, eps): # {s}Y{l,m}.conjugate() = (-1.)**(s+m) {-s}Y{l,-m} ϵ = 2 * ell_max_slow * eps wigner = sf.Wigner(ell_max_slow) m = sf.Yrange(0, ell_max_slow)[:, 1] flipped_indices = np.array([sf.Yindex(ell, -m, 0) for ell, m in sf.Yrange(0, ell_max_slow)]) for iota in special_angles: for phi in special_angles: R = quaternionic.array.from_spherical_coordinates(iota, phi) for s in range(-ell_max_slow, ell_max_slow + 1): Y1 = wigner.sYlm(s, R).conjugate() Y2 = (-1.0)**(s+m) * wigner.sYlm(-s, R)[flipped_indices] assert np.allclose(Y1, Y2, atol=ϵ, rtol=ϵ)
def test_constant_as_ell_0_mode(special_angles): ell_max = 1 wigner = sf.Wigner(ell_max) np.random.seed(123) for imaginary_part in [0.0, 1.0j]: # Test both real and imaginary constants for _ in range(1000): constant = np.random.uniform(-1, 1) + imaginary_part * np.random.uniform(-1, 1) const_ell_m = sf.constant_as_ell_0_mode(constant) assert abs(constant - sf.constant_from_ell_0_mode(const_ell_m)) < 1e-15 for theta in special_angles: for phi in special_angles: Y = wigner.sYlm(0, quaternionic.array.from_spherical_coordinates(theta, phi)) dot = np.dot(const_ell_m, Y[0]) assert abs(constant - dot) < 1e-15, imaginary_part
def test_sYlm_vs_NINJA(special_angles, ell_max_slow, eps): ϵ = 5 * ell_max_slow**6 * eps # This is mostly due to the expressions above being inaccurate wigner = sf.Wigner(ell_max_slow) for iota in special_angles: for phi in special_angles: R = quaternionic.array.from_euler_angles(phi, iota, 0) Y1 = np.array([ [ sYlm_NINJA(s, ell, m, iota, phi) for ell in range(ell_max_slow + 1) for m in range(-ell, ell + 1) ] for s in range(-ell_max_slow, ell_max_slow + 1) ]) Y2 = np.array([wigner.sYlm(s, R) for s in range(-ell_max_slow, ell_max_slow + 1)]) assert np.allclose(Y1, Y2, rtol=ϵ, atol=ϵ)
def test_sYlm_spin_behavior(Rs, special_angles, ell_max_slow, eps): # We expect that the SWSHs behave according to # sYlm( R * exp(gamma*z/2) ) = sYlm(R) * exp(-1j*s*gamma) # See http://moble.github.io/spherical/SWSHs.html#fn:2 # for a more detailed explanation # print("") ϵ = 2 * ell_max_slow * eps wigner = sf.Wigner(ell_max_slow) for i, R in enumerate(Rs): # print("\t{0} of {1}: R = {2}".format(i, len(Rs), R)) for gamma in special_angles: Rgamma = R * quaternionic.array(math.cos(gamma / 2.), 0, 0, math.sin(gamma / 2.)) for s in range(-ell_max_slow, ell_max_slow + 1): Y1 = wigner.sYlm(s, Rgamma) Y2 = wigner.sYlm(s, R) * cmath.exp(-1j * s * gamma) assert np.allclose(Y1, Y2, atol=ϵ, rtol=ϵ)
def test_wigner_rotate_composition(horner, Rs, ell_max_slow, eps): import time ell_min = 0 ell_max = max(3, ell_max_slow) np.random.seed(1234) ϵ = (10 * (2 * ell_max + 1))**2 * eps wigner = sf.Wigner(ell_max) skipping = 5 print() max_error = 0.0 total_time = 0.0 Rs = Rs[::skipping] for i, R1 in enumerate(Rs): # print(f"\tR1[{i+1}] of {len(Rs)}") for j, R2 in enumerate(Rs): for spin_weight in range(-2, 2 + 1): a1 = np.random.rand(7, sf.Ysize(ell_min, ell_max) * 2).view(complex) a1[:, sf.Yindex(ell_min, -ell_min, ell_min):sf. Yindex(abs(spin_weight), -abs(spin_weight), ell_min)] = 0.0 m1 = sf.Modes(a1, spin_weight=spin_weight, ell_min=ell_min, ell_max=ell_max) t1 = time.perf_counter() fA = wigner.rotate(wigner.rotate(m1, R1, horner=horner), R2, horner=horner) fB = wigner.rotate(m1, R1 * R2, horner=horner) t2 = time.perf_counter() max_error = max(np.max(np.abs(fA - fB)), max_error) total_time += t2 - t1 # import warnings # warnings.warn("Eliminating assert for debugging") assert np.allclose( fA, fB, rtol=ϵ, atol=ϵ ), f"{np.max(np.abs(fA-fB))} > {ϵ} for R1={R1} R2={R2}" print(f"\tmax_error[{horner}] = {max_error}") print(f"\ttotal_time[{horner}] = {total_time}")
def test_Wigner_D_negative_argument(Rs, ell_max, eps): # For integer ell, D(R)=D(-R) # # This test passes (specifically, using these tolerances) for at least # ell_max=100, but takes a few minutes at that level. a = np.zeros(sf.WignerDsize(0, ell_max), dtype=complex) b = np.zeros(sf.WignerDsize(0, ell_max), dtype=complex) wigner = sf.Wigner(ell_max) for R in Rs: wigner.D(R, out=a) wigner.D(-R, out=b) # sf.wigner_D(R, 0, ell_max, out=a) # sf.wigner_D(-R, 0, ell_max, out=b) assert np.allclose(a, b, rtol=ell_max * eps, atol=2 * ell_max * eps) assert np.allclose(wigner.D(R), wigner.D(-R), rtol=ell_max * eps, atol=2 * ell_max * eps)
def test_sYlm_WignerD_expression(special_angles, ell_max_slow, eps): # ₛYₗₘ(R) = (-1)ˢ √((2ℓ+1)/(4π)) 𝔇ˡₘ₋ₛ(R) # = (-1)ˢ √((2ℓ+1)/(4π)) 𝔇̄ˡ₋ₛₘ(R̄) ϵ = 2 * ell_max_slow * eps wigner = sf.Wigner(ell_max_slow) for iota in special_angles: for phi in special_angles: R = quaternionic.array.from_euler_angles(phi, iota, 0) D_R̄ = wigner.D(R.conjugate()) for s in range(-ell_max_slow, ell_max_slow + 1): Y = wigner.sYlm(s, R) for ell in range(abs(s), ell_max_slow + 1): Y_ℓ = Y[sf.Yindex(ell, -ell):sf.Yindex(ell, ell)+1] Y_D = ( (-1.) ** (s) * math.sqrt((2 * ell + 1) / (4 * np.pi)) * D_R̄[sf.WignerDindex(ell, -s, -ell):sf.WignerDindex(ell, -s, ell)+1].conjugate() ) assert np.allclose(Y_ℓ, Y_D, atol=ϵ, rtol=ϵ)
def test_Wigner_D_symmetries(Rs, ell_max, eps): # We have two obvious symmetries to test. First, # # D_{mp,m}(R) = (-1)^{mp+m} \bar{D}_{-mp,-m}(R) # # Second, since D is a unitary matrix, its conjugate transpose is its # inverse; because of the representation property, D(R) should equal the # matrix inverse of D(R⁻¹). Thus, # # D_{mp,m}(R) = \bar{D}_{m,mp}(\bar{R}) ϵ = 5 * ell_max * eps D1 = np.zeros(sf.WignerDsize(0, ell_max), dtype=complex) D2 = np.zeros(sf.WignerDsize(0, ell_max), dtype=complex) wigner = sf.Wigner(ell_max) ell_mp_m = sf.WignerDrange(0, ell_max) flipped_indices = np.array( [sf.WignerDindex(ell, -mp, -m) for ell, mp, m in ell_mp_m]) swapped_indices = np.array( [sf.WignerDindex(ell, m, mp) for ell, mp, m in ell_mp_m]) signs = (-1)**np.abs(np.sum(ell_mp_m[:, 1:], axis=1)) for R in Rs: wigner.D(R, out=D1) wigner.D(R.inverse, out=D2) # D_{mp,m}(R) = (-1)^{mp+m} \bar{D}_{-mp,-m}(R) a = D1 b = signs * D1[flipped_indices].conjugate() assert np.allclose(a, b, rtol=ϵ, atol=ϵ) # D_{mp,m}(R) = \bar{D}_{m,mp}(\bar{R}) b = D2[swapped_indices].conjugate() assert np.allclose(a, b, rtol=ϵ, atol=ϵ) D1 = wigner.D(Rs) D2 = wigner.D(Rs.inverse) # D_{mp,m}(R) = (-1)^{mp+m} \bar{D}_{-mp,-m}(R) a = D1 b = signs * D1[..., flipped_indices].conjugate() assert np.allclose(a, b, rtol=ϵ, atol=ϵ) # D_{mp,m}(R) = \bar{D}_{m,mp}(\bar{R}) b = D2[..., swapped_indices].conjugate() assert np.allclose(a, b, rtol=ϵ, atol=ϵ)
def test_Wigner_D_vs_sympy(special_angles, ell_max_slow, eps): from sympy import S, N from sympy.physics.quantum.spin import WignerD as Wigner_D_sympy # Note that this does not fully respect ell_max_slow because # this test is extraordinarily slow ell_max = min(4, ell_max_slow) ϵ = 2 * ell_max * eps wigner = sf.Wigner(ell_max) max_error = 0.0 j = 0 k = 0 print() a = special_angles[::4] b = special_angles[len(special_angles) // 2::4] c = special_angles[::2] for α in a: for β in b: for γ in c: R = quaternionic.array.from_euler_angles(α, β, γ) 𝔇 = wigner.D(R) k += 1 print(f"\tAngle iteration {k} of {a.size*b.size*c.size}") for ell in range(wigner.ell_max + 1): for mp in range(-ell, ell + 1): for m in range(-ell, ell + 1): sympyD = N(Wigner_D_sympy(ell, mp, m, α, β, γ).doit(), n=24).conjugate() sphericalD = 𝔇[wigner.Dindex(ell, mp, m)] error = float(abs(sympyD - sphericalD)) assert error < ϵ, ( f"Testing Wigner d recursion: ell={ell}, m'={mp}, m={m}, " f"sympy:{sympyD}, spherical:{sphericalD}, error={error}" ) max_error = max(error, max_error) print( f"\tmax_error={max_error} after checking {(len(special_angles)**3)*wigner.Dsize} values" )
def test_vector_as_ell_1_modes(special_angles): ell_min = 1 ell_max = 1 wigner = sf.Wigner(ell_max, ell_min=ell_min) def nhat(theta, phi): return np.array([math.sin(theta) * math.cos(phi), math.sin(theta) * math.sin(phi), math.cos(theta)]) np.random.seed(123) for _ in range(1000): vector = np.random.uniform(-1, 1, size=(3,)) vec_ell_m = sf.vector_as_ell_1_modes(vector) assert np.allclose(vector, sf.vector_from_ell_1_modes(vec_ell_m), atol=1.0e-16, rtol=1.0e-15) for theta in special_angles: for phi in special_angles: dot1 = np.dot(vector, nhat(theta, phi)) dot2 = np.dot(vec_ell_m, wigner.sYlm(0, quaternionic.array.from_spherical_coordinates(theta, phi))).real assert abs(dot1 - dot2) < 1e-15
def test_wigner_evaluate_vs_spinsfast(horner, ell_max, eps): import time ell_max = max(3, ell_max) np.random.seed(1234) ϵ = ell_max * (2 * ell_max + 1) * eps n_theta = n_phi = 2 * ell_max + 1 max_s = 2 wigner = sf.Wigner(ell_max, mp_max=max_s) rotors = quaternionic.array.from_spherical_coordinates( sf.theta_phi(n_theta, n_phi)) max_error = 0.0 total_time = 0.0 for s in range(-max_s, max_s + 1): ell_min = abs(s) a1 = np.random.rand(7, sf.Ysize(ell_min, ell_max) * 2).view(complex) m1 = sf.Modes(a1, spin_weight=s, ell_min=ell_min, ell_max=ell_max) t1 = time.perf_counter() f1 = wigner.evaluate(m1, rotors, horner=horner) t2 = time.perf_counter() assert f1.shape == m1.shape[:-1] + rotors.shape[:-1] f2 = m1.grid(n_theta, n_phi, use_spinsfast=True) assert f2.shape == m1.shape[:-1] + rotors.shape[:-1] assert np.allclose(f1, f2.ndarray, rtol=ϵ, atol=ϵ), ( f"max|f1-f2|={np.max(np.abs(f1-f2.ndarray))} > ϵ={ϵ}\n\n" f"s = {s}\n\nrotors = {rotors.tolist()}\n\n" # f"f1 = {f1.tolist()}\n\nf2 = {f2.tolist()}" ) max_error = max(np.max(np.abs(f1 - f2.ndarray)), max_error) total_time += t2 - t1 print() print(f"\tmax_error[{horner}] = {max_error}") print(f"\ttotal_time[{horner}] = {total_time}")
def test_wigner_rotate_vector(horner, special_angles, Rs, eps): """Rotating a vector == rotating the mode-representation of that vector Note that the wigner.rotate function rotates the *basis* in which the modes are represented, so we rotate the modes by the inverse of the rotation we apply to the vector. """ ell_min = 1 ell_max = 1 wigner = sf.Wigner(ell_max, ell_min=ell_min) def nhat(theta, phi): return quaternionic.array.from_vector_part([ math.sin(theta) * math.cos(phi), math.sin(theta) * math.sin(phi), math.cos(theta) ]) for theta in special_angles[special_angles >= 0]: for phi in special_angles: v = nhat(theta, phi) vₗₘ = sf.Modes(sf.vector_as_ell_1_modes(v.vector), ell_min=ell_min, ell_max=ell_max, spin_weight=0) for R in Rs: vprm1 = (R * v * R.conjugate()).vector vₗₙ = wigner.rotate( vₗₘ, R.conjugate(), horner=horner).ndarray[1:] # See note above vprm2 = sf.vector_from_ell_1_modes(vₗₙ).real assert np.allclose(vprm1, vprm2, atol=5 * eps, rtol=0), (f"\ntheta: {theta}\n" f"phi: {phi}\n" f"R: {R}\n" f"v: {v}\n" f"vprm1: {vprm1}\n" f"vprm2: {vprm2}\n" f"vprm1-vprm2: {vprm1-vprm2}\n")
def test_Wigner_D_vs_Wikipedia(special_angles, ell_max_slow, eps): ell_max = ell_max_slow ϵ = 5 * ell_max**6 * eps D = np.zeros(sf.WignerDsize(0, ell_max), dtype=complex) wigner = sf.Wigner(ell_max) ell_mp_m = sf.WignerDrange(0, ell_max) print("") for alpha in special_angles: print("\talpha={0}".format(alpha)) # Need to show some progress for CI for beta in special_angles[len(special_angles) // 2:]: # Skip beta < 0 print("\t\tbeta={0}".format(beta)) for gamma in special_angles: a = np.conjugate( np.array([ Wigner_D_Wikipedia(alpha, beta, gamma, ell, mp, m) for ell, mp, m in ell_mp_m ])) b = wigner.D(quaternionic.array.from_euler_angles( alpha, beta, gamma), out=D) assert np.allclose(a, b, rtol=ϵ, atol=ϵ)
def test_H_vs_sympy(eps): """Eq. (29) of arxiv:1403.7698: d^{m',m}_{n}(β) = ϵ(m') ϵ(-m) H^{m',m}_{n}(β)""" from sympy.physics.quantum.spin import WignerD as Wigner_D_sympy def ϵ(m): m = np.asarray(m) eps = np.ones_like(m) eps[m >= 0] = (-1)**m[m >= 0] return eps ell_max = 4 alpha, beta, gamma = 0.0, 0.1, 0.0 max_error = 0.0 print() for mp_max in range(ell_max + 1): print(f"Checking mp_max={mp_max} (going up to {ell_max})") w = sf.Wigner(ell_max, mp_max=mp_max) workspace = w.new_workspace() Hwedge, Hv, Hextra, _, _, _ = w._split_workspace(workspace) Hnmpm = w.H(np.exp(1j * beta), Hwedge, Hv, Hextra) for n in range(w.ell_max + 1): for mp in range(-min(n, mp_max), min(n, mp_max) + 1): for m in range(-n, n + 1): sympyd = sympy.re( sympy.N( Wigner_D_sympy(n, mp, m, alpha, beta, gamma).doit())) sphericald = ϵ(mp) * ϵ(-m) * Hnmpm[sf.WignerHindex( n, mp, m, mp_max)] error = float(abs(sympyd - sphericald)) assert error < 1.1 * ell_max * eps, ( f"Testing Wigner d recursion with n={n}, m'={mp}, m={m}, mp_max={mp_max}, " f"sympyd={sympyd}, sphericald={sphericald}, error={error}" )
def test_sYlm_vs_scipy(special_angles, ell_max, eps): from scipy.special import sph_harm ϵ = 2 * ell_max * eps wigner = sf.Wigner(ell_max) # Test one quaternion at a time # print() for theta in np.arange(0, 1 * np.pi + 0.1, np.pi / 8.): for phi in np.arange(0, 2 * np.pi + 0.1, np.pi / 8.): R = quaternionic.array.from_spherical_coordinates(theta, phi) # Note that sph_harm names its arguments (theta, phi) in that order, but swaps the # meaning of theta and phi, so we have to swap the order here. Y1 = np.array([ sph_harm(m, ell, phi, theta) for ell in range(ell_max+1) for m in range(-ell, ell+1) ]) Y2 = wigner.sYlm(0, R) # print(Y1.shape, Y2.shape, theta, phi) assert np.allclose(Y1, Y2, rtol=ϵ, atol=ϵ) # Test N quaternions at a time theta_phi = [ [theta, phi] for theta in np.arange(0, 1 * np.pi + 0.1, np.pi / 8.) for phi in np.arange(0, 2 * np.pi + 0.1, np.pi / 8.) ] R = quaternionic.array.from_spherical_coordinates(theta_phi) Y1 = np.array([ [ sph_harm(m, ell, phi, theta) for ell in range(ell_max+1) for m in range(-ell, ell+1) ] for theta, phi in theta_phi ]) Y2 = wigner.sYlm(0, R) # print() # print(Y1.shape, Y2.shape) assert np.allclose(Y1, Y2, rtol=ϵ, atol=ϵ) # Test NxM quaternions at a time theta_phi = [ [ [theta, phi] for theta in np.arange(0, 1 * np.pi + 0.1, np.pi / 8.) ] for phi in np.arange(0, 2 * np.pi + 0.1, np.pi / 8.) ] R = quaternionic.array.from_spherical_coordinates(theta_phi) Y1 = np.array([ [ [ sph_harm(m, ell, phi, theta) for ell in range(ell_max+1) for m in range(-ell, ell+1) ] for theta, phi in tp ] for tp in theta_phi ]) Y2 = wigner.sYlm(0, R) # print() # print(Y1.shape, Y2.shape) assert np.allclose(Y1, Y2, rtol=ϵ, atol=ϵ)
def calc_rotdens(grid_3d, WDMATS, params): #calc rotdens at a point wavepacket_file = params['rot_wf_file'] coef_file = params['rot_coeffs_file'] Jmax = params['Jmax'] itime = int(params['rv_wavepacket_time'] / params['rv_wavepacket_dt']) states = read_coefficients(coef_file, coef_thresh=1.0e-16) time, coefs, quanta_all = read_wavepacket(wavepacket_file, coef_thresh=1.0e-16) quanta = quanta_all[itime] npoints_3d = grid_3d.shape[0] # mapping between wavepacket and rovibrational states print("shape of grid_3d " + str(grid_3d.shape)) ind_state = [] print(np.shape(quanta)) print(quanta) for q in quanta: j = q[1] #q[0] = M id = q[2] ideg = q[3] istate = [(state["j"], state["id"], state["ideg"]) for state in states].index((j, id, ideg)) """ find in which position in the states list we have quanta j, id, ideg - from the current wavepacket """ ind_state.append( istate ) #at each time we append an array of indices which locate the current wavepacket in the states dictionary # lists of J and m quantum numbers jlist = list(set(j for j in quanta[:, 1])) mlist = [] for j in jlist: mlist.append( list(set(m for m, jj in zip(quanta[:, 0], quanta[:, 1]) if jj == j))) print("List of J-quanta:", jlist) print("List of m-quanta:", mlist) # precompute symmetric-top functions on a 3D grid of Euler angles for given J, m=J, and k=-J..J wigner = spherical.Wigner(Jmax) R = quaternionic.array.from_euler_angles(grid_3d) print("\nPrecompute symmetric-top functions...") symtop = [] for J, ml, ij in zip(jlist, mlist, range(len(jlist))): print("J = ", J) print("m = ", ml) print("ij = ", ij) Jfac = np.sqrt((2 * J + 1) / (8 * np.pi**2)) symtop.append([]) #D[wigner.Dindex(J, m, k)] #wigner.wiglib.DJ_m_k(int(J), int(m), grid_3d[:,:]) #grid_3d= (3,npoints_3d). Returns wig = array (npoints, 2*J+1) for each k, #wig Contains values of D-functions on grid, #D_{m,k}^{(J)} = wig[ipoint,k+J], so that the range for the second argument is 0,...,2J symtop[ij].append(np.conj(WDMATS[int(J)]) * Jfac) print("...done") #print(np.shape(symtop)) #print(np.shape(symtop[0])) # compute rotational density vmax = max([max([v for v in state["v"]]) for state in states]) func = np.zeros((npoints_3d, vmax + 1), dtype=np.complex128) tot_func = np.zeros((npoints_3d, vmax + 1), dtype=np.complex128) dens = np.zeros(npoints_3d, dtype=np.complex128) for q, cc, istate in zip( quanta, coefs, ind_state): #loop over coefficients in the wavepacket m = q[0] j = q[1] state = states[istate] ind_j = jlist.index(j) ind_m = mlist[ind_j].index(m) print(ind_j) print(ind_m) print(m + int(j)) # primitive rovibrational function on Euler grid func[:, :] = 0 print(np.shape(symtop[ind_j])) print(np.shape(WDMATS)) print("shape of grid", np.shape(grid_3d)) for v, k, c in zip( state["v"], state["k"], state["coef"] ): #loop over coefficients of primitive symmetric top functions comprising individual components of the wavepacket func[:, v] += c * np.conj(WDMATS[int(j)][int(m) + int(j), k + int(j), :]) * Jfac # total function tot_func[:, :] += func[:, :] * cc # reduced rotational density on Euler grid dens = np.einsum('ij,ji->i', tot_func, np.conj(tot_func.T)) * np.sin( grid_3d[:, 1]) #tensor contraction: element-wise multuplication of tot_func and transpose of np.conj(tot_func.T)) * np.sin(grid_3d[1,ipoint0:ipoint1] and we take diagonal elements of the output. # This is to remove the vibrational index. return grid_3d, dens
def cached_swsh_grid( size, num_points, spin_weight, ell_max, clip_y_normal, clip_z_normal, cache_dir=None, ): logger = logging.getLogger(__name__) X = np.linspace(-size, size, num_points) Y = np.linspace(-size, 0, num_points // 2) if clip_y_normal else X Z = np.linspace(-size, 0, num_points // 2) if clip_z_normal else X x, y, z = map(lambda arr: arr.flatten(order="F"), np.meshgrid(X, Y, Z, indexing="ij")) r = np.sqrt(x**2 + y**2 + z**2) swsh_grid = None if cache_dir: swsh_grid_id = ( round(float(size), 3), int(num_points), int(spin_weight), int(ell_max), bool(clip_y_normal), bool(clip_z_normal), ) # Create a somewhat unique filename swsh_grid_hash = (int( hashlib.md5(repr(swsh_grid_id).encode("utf-8")).hexdigest(), 16) % 10**8) swsh_grid_cache_file = os.path.join( cache_dir, f"swsh_grid_D{int(size)}_N{int(num_points)}_{str(swsh_grid_hash)}.npy", ) if os.path.exists(swsh_grid_cache_file): logger.debug( f"Loading SWSH grid from file '{swsh_grid_cache_file}'...") swsh_grid = np.load(swsh_grid_cache_file) else: logger.debug(f"No SWSH grid file '{swsh_grid_cache_file}' found.") if swsh_grid is None: logger.info("No cached SWSH grid found, computing now...") logger.info("Loading 'spherical' module...") import quaternionic import spherical logger.info("'spherical' module loaded.") start_time = time.time() th = np.arccos(z / r) phi = np.arctan2(y, x) angles = quaternionic.array.from_spherical_coordinates(th, phi) swsh_grid = spherical.Wigner(ell_max).sYlm(s=spin_weight, R=angles) logger.info(f"SWSH grid computed in {time.time() - start_time:.3f}s.") if cache_dir: if not os.path.exists(cache_dir): os.makedirs(cache_dir) if not os.path.exists(swsh_grid_cache_file): np.save(swsh_grid_cache_file, swsh_grid) logger.debug( f"SWSH grid cache saved to file '{swsh_grid_cache_file}'.") return swsh_grid, r
def test_Wigner_D_roundoff(Rs, ell_max, eps): # Testing rotations in special regions with simple expressions for 𝔇 ϵ = 5 * ell_max * eps D = np.zeros(sf.WignerDsize(0, ell_max), dtype=complex) wigner = sf.Wigner(ell_max) # Test rotations with |Ra|<1e-15 wigner.D(quaternionic.x, out=D) actual = D expected = np.array([((-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(actual, expected, rtol=ϵ, atol=ϵ) wigner.D(quaternionic.y, out=D) actual = D expected = np.array([((-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(actual, expected, rtol=ϵ, atol=ϵ) for theta in np.linspace(0, 2 * np.pi): wigner.D(np.cos(theta) * quaternionic.y + np.sin(theta) * quaternionic.x, out=D) actual = D expected = np.array([ ((-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(actual, expected, rtol=ϵ, atol=ϵ) # Test rotations with |Rb|<1e-15 wigner.D(quaternionic.one, out=D) actual = D expected = np.array([(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(actual, expected, rtol=ϵ, atol=ϵ) wigner.D(quaternionic.z, out=D) actual = D expected = np.array([((-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(actual, expected, rtol=ϵ, atol=ϵ) for theta in np.linspace(0, 2 * np.pi): wigner.D(np.cos(theta) * quaternionic.one + np.sin(theta) * quaternionic.z, out=D) actual = D expected = np.array([ ((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(actual, expected, rtol=ϵ, atol=ϵ)