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()