def main(lmax, resolution, steps):
scale = 0.5 * math.sqrt(4 * math.pi) / math.sqrt(2 * lmax + 1)
axis = dict(showbackground=False,
showticklabels=False,
showgrid=False,
zeroline=False,
title='',
nticks=3,
range=[-lmax / 2 - 0.5, lmax / 2 + 0.5])
layout = dict(width=resolution,
height=resolution,
scene=dict(
xaxis=axis,
yaxis=axis,
zaxis=axis,
aspectmode='manual',
aspectratio=dict(x=1, y=1, z=1),
camera=dict(
up=dict(x=0, y=0, z=1),
center=dict(x=0, y=0, z=0),
eye=dict(x=0, y=-1.3, z=0),
projection=dict(type='perspective'),
),
),
paper_bgcolor='rgba(0,0,0,0)',
plot_bgcolor='rgba(0,0,0,0)',
margin=dict(l=0, r=0, t=0, b=0))
if os.path.exists('sh'):
shutil.rmtree('sh')
os.makedirs('sh')
for i in tqdm.tqdm(range(steps)):
rot = 2 * math.pi * i / steps
a, b, c = 0, math.pi / 4, 0
abc = o3.compose(-c, -b, -a, *o3.compose(0, 0, rot, a, b, c))
surfaces = [
rsh_surface(l, m, scale, [
l + (m if m < 0 else 0) - lmax / 2, 0, lmax / 2 - l +
(m if m > 0 else 0)
], abc) for l in range(lmax + 1) for m in range(-l, l + 1)
]
fig = go.Figure(surfaces, layout=layout)
fig.write_image('sh/{:03d}.png'.format(i))
subprocess.check_output([
"convert", "-delay", "3", "-loop", "0", "-dispose", "2", "sh/*.png",
"output.gif"
])
def _test_is_representation(self, R):
"""
R(Z(a1) Y(b1) Z(c1) Z(a2) Y(b2) Z(c2)) = R(Z(a1) Y(b1) Z(c1)) R(Z(a2) Y(b2) Z(c2))
"""
with o3.torch_default_dtype(torch.float64):
g1 = o3.rand_angles()
g2 = o3.rand_angles()
g12 = o3.compose(*g1, *g2)
D12 = R(*g12)
D1D2 = R(*g1) @ R(*g2)
self.assertLess((D12 - D1D2).abs().max(), 1e-10 * D12.abs().max())
def _is_representation(D, eps):
I = D(0, 0, 0)
if not torch.allclose(I, I @ I):
return False
g1 = o3.rand_angles()
g2 = o3.rand_angles()
g12 = o3.compose(*g1, *g2)
D12 = D(*g12)
D1D2 = D(*g1) @ D(*g2)
return (D12 - D1D2).abs().max().item() < eps * D12.abs().max().item()
def _is_representation(D, eps, with_parity=False):
e = (0, 0, 0, 0) if with_parity else (0, 0, 0)
I = D(*e)
if not torch.allclose(I, I @ I):
return False
g1 = o3.rand_angles() + (0, ) if with_parity else o3.rand_angles()
g2 = o3.rand_angles() + (0, ) if with_parity else o3.rand_angles()
g12 = o3.compose_with_parity(*g1, *g2) if with_parity else o3.compose(
*g1, *g2)
D12 = D(*g12)
D1D2 = D(*g1) @ D(*g2)
return (D12 - D1D2).abs().max().item() < eps * D12.abs().max().item()
def _test_is_representation(self, R):
"""
R(Z(a1) Y(b1) Z(c1) Z(a2) Y(b2) Z(c2)) = R(Z(a1) Y(b1) Z(c1)) R(Z(a2) Y(b2) Z(c2))
"""
with o3.torch_default_dtype(torch.float64):
a1, b1, c1, a2, b2, c2 = torch.rand(6)
r1 = R(a1, b1, c1)
r2 = R(a2, b2, c2)
a, b, c = o3.compose(a1, b1, c1, a2, b2, c2)
r = R(a, b, c)
r_ = r1 @ r2
self.assertLess((r - r_).abs().max(), 1e-10 * r.abs().max())
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()
def test_spherical_harmonics(self):
"""
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 order 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 = o3.spherical_harmonics(order, ra, rb)
Y = o3.spherical_harmonics(order, a, b)
DrY = o3.irr_repr(order, alpha, beta, gamma) @ Y
self.assertLess((Yrx - DrY).abs().max(), 1e-10 * Y.abs().max())
