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)
