3d_analytical.py (1183B)
1 import numpy as np 2 from scipy.special import factorial, sph_harm, genlaguerre 3 import matplotlib.pyplot as plt 4 from dolfin import * 5 import pandas as pd 6 from mayavi import mlab 7 from skimage import measure 8 9 # build wave-function: 10 def wave_function(n, l, m, x, y, z, a0): 11 12 r = np.sqrt(x**2+y**2+z**2) 13 theta = np.arccos(z/r) 14 phi = np.arctan(y/x) 15 16 # replace possible NaNs wirh 0 17 theta[np.argwhere(np.isnan(theta))] = 0 18 phi[np.argwhere(np.isnan(phi))] = 0 19 20 # Radial part 21 R = np.sqrt((2/n/a0)**3*factorial(n-l-1)/(2*n*(factorial(n+l))**3)) * (2*r/n/a0)**l * np.exp(-r/n/a0) * genlaguerre(n-l-1,2*l+1)(2*r/n/a0) 22 # spherical harmonics 23 Y = sph_harm(m, l, phi, theta) 24 # complete wavefunction 25 wf = R * Y 26 return wf 27 28 29 def main(): 30 31 d=0.1 32 min=-10 33 max=10 34 X = np.arange(min,max,d) 35 Y = np.arange(min,max,d) 36 Z = np.arange(min,max,d) 37 x, y, z = np.meshgrid(X, Y, Z) 38 a0 = 1 39 40 # mlab.figure() 41 42 wf = np.abs(wave_function(2, 1, 0, x, y, z, a0))**2 43 mlab.contour3d(wf, transparent=True) 44 45 mlab.savefig('s-orbital.png') 46 mlab.colorbar() 47 mlab.outline() 48 mlab.show() 49 50 51 if __name__ == "__main__": 52 main()