851b618c26a3ab02076e7f40a54d7542e404013b,se3cnn/SO3.py,,spherical_harmonics_xyz,#,109

Before Change

// xyz tensor [,,, 3]
def spherical_harmonics_xyz(x, y, z, J):
    with torch_default_dtype(torch.float64):
        if x == y == z == 0:  // angles at origin are nan, special treatment
            if J == 0:  // Y^0 is angularly independent, choose any angle
                return spherical_harmonics(0, 123, 321)  // [m]
            else:  // insert zeros for Y^J with J!=0
                return 0
        else:  // not at the origin, sample spherical harmonic
            alpha, beta = x_to_alpha_beta([x, y, z])
            return spherical_harmonics(J, alpha, beta)  // [m]


def compose(a1, b1, c1, a2, b2, c2):
    
    (a, b, c) = (a1, b1, c1) composed with (a2, b2, c2)
    

After Change

    :param xyz: tensor of shape [A, 3]
    :return: tensor of shape [m, A]
    
    if not isinstance(order, list):
        order = [order]

    with torch_default_dtype(torch.float64):
        alpha, beta = x_to_alpha_beta(xyz)  // two tensors of shape [A]
        out = spherical_harmonics(order, alpha, beta)  // [m, A]

Instances