In [1]:
import k3d
import numpy as np
import scipy.special
import scipy.misc

r = lambda x,y,z: np.sqrt(x**2+y**2+z**2)
theta = lambda x,y,z: np.arccos(z/r(x,y,z))
phi = lambda x,y,z: np.arctan(y/x)

a0 = 1.
R = lambda r,n,l: (2*r/n/a0)**l * np.exp(-r/n/a0) * scipy.special.genlaguerre(n-l-1,2*l+1)(2*r/n/a0)
WF = lambda r,theta,phi,n,l,m: R(r,n,l) * scipy.special.sph_harm(m,l,phi,theta)
absWF = lambda r,theta,phi,n,l,m: abs(WF(r,theta,phi,n,l,m))**2
N=100j
a = 200.0
x,y,z = np.ogrid[-a:a:N,-a:a:N,-a:a:N]
x = x.astype(np.float32)
y = y.astype(np.float32)
z = z.astype(np.float32)

In [2]:
orbital = absWF(r(x,y,z),theta(x,y,z),phi(x,y,z),1,0,0) # 1s

In [3]:
plot = k3d.plot()
plot.display()

In [4]:
plot.grid_auto_fit = False

In [ ]:
E = 10
volume_animation = {}
label_animation = {}
i = 0

for l in range(E):
print(l, '/', E-1, end='\r')
for m in range(-l,l+1):
psi2 = absWF(r(x, y, z), theta(x, y, z), phi(x, y, z), E, l, m)

volume_animation[str(i)] =  (psi2/np.max(psi2))

i += 0.1

9 / 9

In [ ]:
plot += k3d.text2d(label_animation, (0.,0.))

In [ ]:
plot += k3d.volume(volume_animation, color_map=k3d.colormaps.basic_color_maps.CoolWarm,
color_range=(0.0,0.1))

In [ ]:
np.array(volume_animation).tolist()['0'].dtype

In [ ]:
plot.colorbar_object_id = 0

In [ ]:
plot.start_auto_play()

In [ ]:
plot.stop_auto_play()

In [ ]:
np.float32

In [ ]: