Marcel Rød marcelroed

2nd year PhD student at Stanford. Structured machine learning, generative models, compilers, and hardware accelerators. Advised by Jure Leskovec.

mariogeiger / spherical_harmonics.py

Created December 8, 2020 17:31

jax code for spherical harmonics

	# 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 x2 + y2 + z**2 = 1

	import jax
	import jax.numpy as jnp

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