def test_sh_dirac(self): with o3.torch_default_dtype(torch.float64): for l in range(5): angles = torch.tensor(1.2), torch.tensor(2.1) a = sphten.spherical_harmonics_dirac(torch.stack(o3.angles_to_xyz(*angles), dim=-1), l) v = sphten.SphericalTensor(a, 1, l).value(*angles) self.assertAlmostEqual(v.item(), 1)
def plotly_sphere(fun, n=240, radius=False, center=None, relu=False): """ surface = plotly_sphere(partial(spherical_harmonic_signal, x)) fig = go.Figure(data=[surface]) fig.show() """ import plotly.graph_objs as go a = torch.linspace(0, 2 * math.pi, n, dtype=torch.float64) b = torch.linspace(0, math.pi, n, dtype=torch.float64) a, b = torch.meshgrid(a, b) f = fun(a, b) f = f.relu() if relu else f x, y, z = o3.angles_to_xyz(a, b) if radius: r = f.abs() x *= r y *= r z *= r if center is not None: x += center[0] y += center[1] z += center[2] return go.Surface(x=x.numpy(), y=y.numpy(), z=z.numpy(), surfacecolor=f.numpy())
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 precompute(self, R): a = torch.linspace(0, 2 * math.pi, 2 * self.n) b = torch.linspace(0, math.pi, self.n)[2:-2] a, b = torch.meshgrid(a, b) xyz = torch.stack(o3.angles_to_xyz(a, b), dim=-1) @ R.t() a, b = o3.xyz_to_angles(xyz) proj = SphericalHarmonicsProject(a, b, self.lmax) return xyz, proj
def spherical_surface(n): beta = torch.linspace(1e-16, math.pi - 1e-16, 2 * n) alpha = torch.linspace(0, 2 * math.pi, 2 * n) beta_, alpha_ = torch.meshgrid(beta, alpha) x, y, z = o3.angles_to_xyz(beta_, alpha_) beta = 0.5 * (beta[1:] + beta[:-1]) alpha = 0.5 * (alpha[1:] + alpha[:-1]) beta, alpha = torch.meshgrid(beta, alpha) return x, y, z, alpha, beta
def test_xyz(float_tolerance): R = o3.rand_matrix(10) assert (R @ R.transpose(-1, -2) - torch.eye(3)).abs().max() < float_tolerance a, b, c = o3.matrix_to_angles(R) pos1 = o3.angles_to_xyz(a, b) pos2 = R @ torch.tensor([0, 1.0, 0]) assert torch.allclose(pos1, pos2, atol=float_tolerance) a2, b2 = o3.xyz_to_angles(pos2) assert (a - a2).abs().max() < float_tolerance assert (b - b2).abs().max() < float_tolerance
def grid(self): """ grid of positions """ beta, alpha = torch.meshgrid(self.betas, self.alphas) return o3.angles_to_xyz(alpha, beta)
def grid(self): beta, alpha = torch.meshgrid(self.betas, self.alphas) return o3.angles_to_xyz(alpha, beta)
def signal_on_sphere(self, n=100): grid = ToS2Grid(self.lmax, res=n) beta, alpha = torch.meshgrid(grid.betas, grid.alphas) # [beta, alpha] x, y, z = o3.angles_to_xyz(alpha, beta) r = torch.stack([x, y, z], dim=-1) # [beta, alpha, 3] return r, grid(self.signal)