def setUp(self):
self.maxl = 3
self.sphs = SphericalHarmonics(maxl=self.maxl, sh_norm='qm')
sphs_conj = SphericalHarmonics(maxl=self.maxl, conj=True, sh_norm='qm')
# Generate some reference point(s)
phi_refs = np.array([
np.pi / 2,
-np.pi / 2,
])
theta_refs = np.pi / 2 * np.ones_like(phi_refs)
theta_phi_refs = np.stack([theta_refs, phi_refs], axis=-1)
xyz_refs = spherical_to_cartesian(theta_phi_refs)
y_lms_conj = sphs_conj.forward(
torch.tensor(xyz_refs, dtype=torch.float))
self.a_lms_1 = estimate_alms(y_lms_conj)
# Another set of a_lms
phi_refs = np.array([np.pi / 3])
theta_refs = np.pi / 3 * np.ones_like(phi_refs)
theta_phi_refs = np.stack([theta_refs, phi_refs], axis=-1)
xyz_refs = spherical_to_cartesian(theta_phi_refs)
y_lms_conj = sphs_conj.forward(
torch.tensor(xyz_refs, dtype=torch.float))
self.a_lms_2 = estimate_alms(y_lms_conj)
def test_invariant(self):
max_ell = 4
sphs_conj = SphericalHarmonics(maxl=max_ell, conj=True, sh_norm='unit')
atomic_scalars = AtomicScalars(maxl=max_ell)
theta_phi = np.array([[np.pi / 3, np.pi / 4],
[2 * np.pi / 3, np.pi / 2]])
xyz_refs = spherical_to_cartesian(theta_phi)
y_lms_conj = sphs_conj.forward(
torch.tensor(xyz_refs, dtype=torch.float))
a_lms = estimate_alms(y_lms_conj)
invariant = atomic_scalars(a_lms)
self.assertTrue(invariant.shape[-1],
atomic_scalars.get_output_dim(channels=1))
random_rotation = SO3WignerD.euler(maxl=max_ell, dtype=torch.float)
a_lms_rotated = rotate_rep(random_rotation, a_lms)
self.assertFalse(
np.allclose(to_numpy(a_lms[1]), to_numpy(a_lms_rotated[1])))
invariant_rotated = atomic_scalars(a_lms_rotated)
self.assertTrue(np.allclose(invariant, invariant_rotated))
def test_concat(self):
theta_phi = np.array([[np.pi / 2, np.pi / 2]])
xyz_refs = spherical_to_cartesian(theta_phi)
y_lms_conj = self.sphs_conj.forward(
torch.tensor(xyz_refs, dtype=torch.float))
a_lms = estimate_alms(y_lms_conj)
a_lms = concat_so3vecs([a_lms] * 3)
self.assertTrue(all(a_lm.shape[0] == 3 for a_lm in a_lms))
def test_normalization(self):
theta_phi = np.array([[np.pi / 2, np.pi / 2]])
xyz_refs = spherical_to_cartesian(theta_phi)
y_lms_conj = self.sphs_conj.forward(
torch.tensor(xyz_refs, dtype=torch.float))
a_lms = estimate_alms(y_lms_conj)
k1 = get_normalization_constant(a_lms)
self.assertTrue(k1.shape, (1, ))
# If sh_norm='unit', sum over m = 1.
self.assertTrue(k1.item(), self.maxl + 1)
normalized_a_lms = normalize_alms(a_lms)
k2 = get_normalization_constant(normalized_a_lms)
self.assertAlmostEqual(k2.item(), 1.0)