def test_scipy_spherical_harmonics():
with o3.torch_default_dtype(torch.float64):
ls = [0, 1, 2, 3, 4, 5]
beta = torch.linspace(1e-3, math.pi - 1e-3, 100, requires_grad=True).reshape(1, -1)
alpha = torch.linspace(0, 2 * math.pi, 100, requires_grad=True).reshape(-1, 1)
Y1 = rsh.spherical_harmonics_alpha_beta(ls, alpha, beta)
Y2 = rsh.spherical_harmonics_alpha_beta(ls, alpha.detach(), beta.detach())
assert (Y1 - Y2).abs().max() < 1e-10
def test_sh_is_in_irrep():
with o3.torch_default_dtype(torch.float64):
for l in range(4 + 1):
a, b = 3.14 * torch.rand(2) # works only for beta in [0, pi]
Y = rsh.spherical_harmonics_alpha_beta([l], a, b) * math.sqrt(4 * math.pi) / math.sqrt(2 * l + 1) * (-1) ** l
D = o3.irr_repr(l, a, b, 0)
assert (Y - D[:, l]).norm() < 1e-10
def rsh_surface(l, m, scale, tr, rot):
n = 50
a = torch.linspace(0, 2 * math.pi, 2 * n)
b = torch.linspace(0, math.pi, n)
a, b = torch.meshgrid(a, b)
f = rsh.spherical_harmonics_alpha_beta([l], a, b)
f = torch.einsum('ij,...j->...i', o3.irr_repr(l, *rot), f)
f = f[..., l + m]
x, y, z = o3.angles_to_xyz(a, b)
r = f.abs()
x = scale * r * x + tr[0]
y = scale * r * y + tr[1]
z = scale * r * z + tr[2]
max_value = 0.5
return go.Surface(
x=x.numpy(),
y=y.numpy(),
z=z.numpy(),
surfacecolor=f.numpy(),
showscale=False,
cmin=-max_value,
cmax=max_value,
colorscale=[[0, 'rgb(0,50,255)'], [0.5, 'rgb(200,200,200)'],
[1, 'rgb(255,50,0)']],
)
def signal_alpha_beta(self, alpha, beta):
"""
Evaluate the signal on the sphere
"""
sh = rsh.spherical_harmonics_alpha_beta(list(range(self.lmax + 1)),
alpha, beta)
dim = (self.lmax + 1)**2
output = torch.einsum('ai,zi->za', sh.reshape(-1, dim),
self.signal.reshape(-1, dim))
return output.reshape((*self.signal.shape[:-1], *alpha.shape))
def test_sh_equivariance():
"""
This test tests that
- irr_repr
- compose
- spherical_harmonics
are compatible
Y(Z(alpha) Y(beta) Z(gamma) x) = D(alpha, beta, gamma) Y(x)
with x = Z(a) Y(b) eta
"""
for l in range(7):
with o3.torch_default_dtype(torch.float64):
a, b = torch.rand(2)
alpha, beta, gamma = torch.rand(3)
ra, rb, _ = o3.compose(alpha, beta, gamma, a, b, 0)
Yrx = rsh.spherical_harmonics_alpha_beta([l], ra, rb)
Y = rsh.spherical_harmonics_alpha_beta([l], a, b)
DrY = o3.irr_repr(l, alpha, beta, gamma) @ Y
assert (Yrx - DrY).abs().max() < 1e-10 * Y.abs().max()