def test_swsft_forward_backward(self, width, spin):
r"""Given bandlimited spherical function, SWSFT -> ISWFT must return it."""
sphere_gt = np.linspace(-1, 1, width**2).reshape((width, width))
# This effectively bandlimits the spherical function.
sphere_gt = np_spin_spherical_harmonics.swsft_backward_naive(
np_spin_spherical_harmonics.swsft_forward_naive(sphere_gt, spin),
spin)
coeffs = np_spin_spherical_harmonics.swsft_forward_naive(
sphere_gt, spin)
sphere = np_spin_spherical_harmonics.swsft_backward_naive(coeffs, spin)
self.assertAllClose(sphere, sphere_gt)
def test_swsft_backward_matches_np(self, num_coefficients, spin):
transformer = _get_transformer()
coeffs_np = (jnp.linspace(-1, 1, num_coefficients) +
1j * jnp.linspace(0, 1, num_coefficients))
coeffs_jax = spin_spherical_harmonics.coefficients_to_matrix(coeffs_np)
sphere_np = np_spin_spherical_harmonics.swsft_backward_naive(
coeffs_np, spin)
sphere_jax = transformer.swsft_backward(coeffs_jax, spin)
self.assertAllClose(sphere_np, sphere_jax)
def test_swsft_backward_forward(self, n_coeffs, spin):
r"""Given coefficients, ISWSFT -> SWFT must return them."""
coeffs_gt = np.linspace(-1, 1,
n_coeffs) + 1j * np.linspace(0, 1, n_coeffs)
# Coefficients for ell < abs(spin) are always zero.
coeffs_gt[:spin**2] = 0
sphere = np_spin_spherical_harmonics.swsft_backward_naive(
coeffs_gt, spin)
coeffs = np_spin_spherical_harmonics.swsft_forward_naive(sphere, spin)
self.assertAllClose(coeffs, coeffs_gt)
def test_spherical_harmonics_backward(self, width, ell, m):
r"""ISWSFT is Y_m^\ell when the only nonzero coeffs is 1. at (ell, m)."""
n_coeffs = (width // 2)**2
coeffs = np.zeros(n_coeffs, dtype='complex128')
coeffs[np_spin_spherical_harmonics._get_swsft_coeff_index(ell, m)] = 1
sphere = np_spin_spherical_harmonics.swsft_backward_naive(coeffs, 0)
phi_g, theta_g = sphere_utils.make_equiangular_grid(width)
sphere_gt = scipy.special.sph_harm(m, ell, phi_g, theta_g)
self.assertAllClose(sphere, sphere_gt)