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