def _sample_sh_cube(size, J, version=3): # pylint: disable=W0613
'''
Sample spherical harmonics in a cube.
No bandlimiting considered, aliased regions need to be cut by windowing!
:param size: side length of the kernel
:param J: order of the spherical harmonics
'''
with torch_default_dtype(torch.float64):
rng = torch.linspace(-((size - 1) / 2), (size - 1) / 2, steps=size)
Y_J = torch.zeros(2 * J + 1, size, size, size, dtype=torch.float64)
for idx_x, x in enumerate(rng):
for idx_y, y in enumerate(rng):
for idx_z, z in enumerate(rng):
if x == y == z == 0: # angles at origin are nan, special treatment
if J == 0: # Y^0 is angularly independent, choose any angle
Y_J[:, idx_x, idx_y,
idx_z] = spherical_harmonics(0, 123,
321) # [m]
else: # insert zeros for Y^J with J!=0
Y_J[:, idx_x, idx_y, idx_z] = 0
else: # not at the origin, sample spherical harmonic
alpha, beta = x_to_alpha_beta([x, y, z])
Y_J[:, idx_x, idx_y,
idx_z] = spherical_harmonics(J, alpha, beta) # [m]
assert Y_J.dtype == torch.float64
return Y_J # [m, x, y, z]
def rotation_from_orientations(t1, f1, t2, f2):
from se3cnn.SO3 import x_to_alpha_beta, rot
zero = t1.new_tensor(0)
r_e_t1 = rot(*x_to_alpha_beta(t1), zero)
r_e_t2 = rot(*x_to_alpha_beta(t2), zero)
f1_e = r_e_t1.t() @ f1
f2_e = r_e_t2.t() @ f2
c = torch.atan2(f2_e[1], f2_e[0]) - torch.atan2(f1_e[1], f1_e[0])
r_f1_f2 = rot(zero, zero, c)
r = r_e_t2 @ r_f1_f2 @ r_e_t1.t()
# t2 = r @ t1
# f2 ~ r @ f1
return r