mariogeiger

# from sympy.functions.special.spherical_harmonics import Znm
# yl(x, y, z) = -(-1)**l * (-(-1)**m if m >=0 else 1) sqrt(4 pi) * Znm(l, m, acos(z), atan2(y, x))
# assuming x**2 + y**2 + z**2 = 1

import jax
import jax.numpy as jnp

@jax.jit
def y0(x, y, z):
    return jnp.stack([