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)
def test_l_1(self):
theta_phi = np.array([np.pi / 2, 0.0])
pos = spherical_to_cartesian(theta_phi)
pos_tensor = torch.tensor(pos, dtype=torch.float32)
# To match the definition Mathematica uses sh_norm='qm' is required
sph = SphericalHarmonics(maxl=1, normalize=True, sh_norm='qm')
output = sph.forward(pos_tensor)
# Mathematica output:
expected = np.array([
[0.345494, 0],
[0, 0],
[-0.345494, 0],
],
dtype=np.float32)
self.assertTrue(np.allclose(output[1].cpu().detach().numpy(),
expected))
def test_l_2(self):
theta_phi = np.array([np.pi / 3, np.pi / 4])
pos = spherical_to_cartesian(theta_phi)
pos_tensor = torch.tensor(pos, dtype=torch.float32)
# To match the definition Mathematica uses sh_norm='qm' is required
sph = SphericalHarmonics(maxl=2, normalize=False, sh_norm='qm')
output = sph.forward(pos_tensor)
# Mathematica output:
expected = np.array([
[0, -0.289706],
[0.236544, -0.236544],
[-0.0788479, 0],
[-0.236544, -0.236544],
[0, 0.289706],
],
dtype=np.float32)
self.assertTrue(np.allclose(output[2].cpu().detach().numpy(),
expected))
def test_conversion(self):
theta_phi = np.array([np.pi / 3, np.pi / 4])
pos = spherical_to_cartesian(theta_phi)
expected = np.array([0.612372, 0.612372, 0.5])
self.assertTrue(np.allclose(pos, expected))
def test_cycle_2(self):
theta_phi = np.array([[0.3, -1.2]])
xyz = spherical_to_cartesian(theta_phi)
theta_phi_2 = cartesian_to_spherical(xyz)
self.assertTrue(np.all(np.isclose(theta_phi, theta_phi_2)))
def test_cycle(self):
xyz = np.array([[0.0, -1.0, 0.0]])
theta_phi = cartesian_to_spherical(xyz)
xyz_new = spherical_to_cartesian(theta_phi)
self.assertTrue(np.all(np.isclose(xyz, xyz_new)))
def test_spherical_to_cartesian_2(self):
theta_phi = np.array([[np.pi / 2, 3 / 2 * np.pi]])
xyz = spherical_to_cartesian(theta_phi)
self.assertTrue(np.all(np.isclose(xyz, np.array([[0.0, -1.0, 0.0]]))))