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_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_Ysize(ell_max):
# k = 0
for ell_max in range(ell_max + 1):
for ell_min in range(ell_max + 1):
a = sf.Ysize(ell_min, ell_max)
b = len(sf.Yrange(ell_min, ell_max))
assert a == b, ((ell_min, ell_max), a, b) #, k)
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_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_modes_grid_variants(ell_max, eps):
ell_max = max(3, ell_max)
s_max = 2
np.random.seed(1234)
ϵ = 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(-s_max, s_max + 1):
ell_min = abs(s)
a1 = np.random.rand(2, sf.Ysize(ell_min, ell_max) * 2).view(complex)
m1 = sf.Modes(a1, spin_weight=s, ell_min=ell_min, ell_max=ell_max)
fA = m1.grid(n_theta, n_phi, use_spinsfast=True)
fB = m1.grid(n_theta, n_phi, use_spinsfast=False)
assert np.allclose(fA.ndarray, fB.ndarray, rtol=ϵ, atol=ϵ), (
f"fA = np.array({fA.ndarray.tolist()})\n\n"
f"fB = np.array({fB.ndarray.tolist()})\n\n"
"\n"
f"max|fA-fB|={np.max(np.abs(fA.ndarray-fB.ndarray))} > ϵ={ϵ}; s={s}"
)