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=ϵ)