spherical_harmonics.py

import torch
from scipy.special import sph_harm


# Converted to numpy arrays in sph_harm
def real_sph_harm(m, l, phi, theta):
    z = sph_harm(abs(m), l, phi, theta)

    if m < 0:
        z = 2**0.5 * (-1)**m * z.imag
    elif m > 0:
        z = 2**0.5 * (-1)**m * z.real
    z = torch.abs(z)

    return z


# Ignores the first, constant, harmonic
# This is with θ in [0, π] and φ in [0, 2π]
def real_sph_harm_up_to_l(num_l, phi, theta):
    return torch.stack(
        [real_sph_harm(m, l, phi, theta) for l in range(1, num_l) for m in range(-l, l + 1)],
        dim=-1
    )