Skip to content

Commit 22694d2

Browse files
authored
Jupyter notebook version added
1 parent 5b87f20 commit 22694d2

7 files changed

+606
-0
lines changed

ch09/Program_9.2_qmshoot.ipynb

+98
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "code",
5+
"execution_count": null,
6+
"metadata": {},
7+
"outputs": [],
8+
"source": [
9+
"#\n",
10+
"# Program 9.2: Shooting method for double well (qmshoot.ipynb)\n",
11+
"# J Wang, Computational modeling and visualization with Python\n",
12+
"#\n",
13+
"\n",
14+
"import numpy as np, rootfinder as rtf, ode\n",
15+
"import matplotlib.pyplot as plt\n",
16+
"%matplotlib notebook\n",
17+
"\n",
18+
"def V(x): # potential\n",
19+
" return 0 if (abs(x) > a/2. or abs(x) < b/2.) else -V0\n",
20+
" \n",
21+
"def sch(psi, x): # Schrodinger eqn\n",
22+
" return [psi[1], 2*(V(x)-E)*psi[0]]\n",
23+
" \n",
24+
"def intsch(psi, n, x, h): # integrate Sch eqn for n steps\n",
25+
" psix = [psi[0]]\n",
26+
" for i in range(n):\n",
27+
" psi = ode.RK45n(sch, psi, x, h)\n",
28+
" x += h\n",
29+
" psix.append(psi[0])\n",
30+
" return psix, psi # return WF and last point\n",
31+
" \n",
32+
"def shoot(En): # calc diff of derivatives at given E\n",
33+
" global E # global E, needed in sch()\n",
34+
" E = En\n",
35+
" wfup, psiup = intsch([0., .1], N, xL, h) # march upward \n",
36+
" wfdn, psidn = intsch([0., .1], N, xR, -h) # march downward\n",
37+
" return psidn[0]*psiup[1] - psiup[0]*psidn[1]\n",
38+
"\n",
39+
"a, b, V0 = 4.0, 1.0, 6. # double well widths, depth\n",
40+
"xL, xR, N = -4*a, 4*a, 500 # limits, intervals\n",
41+
"xa = np.linspace(xL, xR, 2*N+1) # grid\n",
42+
"h, M = xa[1]-xa[0], 2 # step size, M=matching point\n",
43+
"E1, dE, j = -V0, 0.01, 1\n",
44+
"\n",
45+
"plt.figure()\n",
46+
"while (E1 < 0): # find E, calc and plot wave function\n",
47+
" if (shoot(E1) * shoot(E1 + dE) < 0): # bracket E\n",
48+
" E = rtf.bisect(shoot, E1, E1 + dE, 1.e-8)\n",
49+
" print ('Energy found: %.3f' %(E))\n",
50+
" wfup, psiup = intsch([0., .1], N+M, xL, h) # compute WF\n",
51+
" wfdn, psidn = intsch([0., .1], N-M, xR, -h)\n",
52+
" psix = np.concatenate((wfup[:-1], wfdn[::-1])) # combine WF\n",
53+
" scale = psiup[0]/psidn[0]\n",
54+
" psix[N+M:] *= scale # match WF\n",
55+
" ax = plt.subplot(2,2,j)\n",
56+
" ax.plot(xa, psix/max(psix)) # plot WF\n",
57+
" ax.plot(xa, np.vectorize(V)(xa)/(2*V0)) # overlay V\n",
58+
" ax.set_xlim(-a,a)\n",
59+
" ax.text(2.2,0.7, '%.3f' %(E))\n",
60+
" if (j == 1 or j == 3): ax.set_ylabel(r'$\\psi$')\n",
61+
" if (j == 3 or j == 4): ax.set_xlabel('$x$')\n",
62+
" if (j<4): j += 1 # 4 plots max\n",
63+
" E1 += dE\n",
64+
"plt.show()"
65+
]
66+
},
67+
{
68+
"cell_type": "code",
69+
"execution_count": null,
70+
"metadata": {
71+
"collapsed": true
72+
},
73+
"outputs": [],
74+
"source": []
75+
}
76+
],
77+
"metadata": {
78+
"kernelspec": {
79+
"display_name": "Python 3",
80+
"language": "python",
81+
"name": "python3"
82+
},
83+
"language_info": {
84+
"codemirror_mode": {
85+
"name": "ipython",
86+
"version": 3
87+
},
88+
"file_extension": ".py",
89+
"mimetype": "text/x-python",
90+
"name": "python",
91+
"nbconvert_exporter": "python",
92+
"pygments_lexer": "ipython3",
93+
"version": "3.6.1"
94+
}
95+
},
96+
"nbformat": 4,
97+
"nbformat_minor": 2
98+
}

ch09/Program_9.3_qmfem.ipynb

+88
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "code",
5+
"execution_count": null,
6+
"metadata": {},
7+
"outputs": [],
8+
"source": [
9+
"#\n",
10+
"# Program 9.3: Eigenenergies by FEM (qmfem.ipynb)\n",
11+
"# J Wang, Computational modeling and visualization with Python\n",
12+
"#\n",
13+
"\n",
14+
"import numpy as np\n",
15+
"from scipy.sparse.linalg import eigsh\n",
16+
"\n",
17+
"def V(x): # potential\n",
18+
" return 0. if (abs(x) >= a/2.) else -V0\n",
19+
"\n",
20+
"def TB_mat(n): # fill in T and and B matrices\n",
21+
" Tm, B = np.diag([2.]*n), np.diag([4.]*n)\n",
22+
" Tm += np.diag([-1.]*(n-1),1) + np.diag([-1.]*(n-1),-1) # off diag\n",
23+
" B += np.diag([1.]*(n-1),1) + np.diag([1.]*(n-1),-1) \n",
24+
" return Tm/(2.*h), B*h/6.\n",
25+
" \n",
26+
"def Vij(i, j): # pot. matrix Vij over $[x_i,x_{i+1}]$ by order-4 Gaussian\n",
27+
" x=np.array([0.3399810435848564, 0.8611363115940525]) # abscissa\n",
28+
" w=np.array([0.6521451548625463, 0.3478548451374545]) # weight\n",
29+
" phi = lambda i, x: 1.0 - abs(x-xa[i])/h # tent function\n",
30+
" vV, hh = np.vectorize(V), h/2.0 # vectorize V\n",
31+
" x1, x2 = xa[i] + hh - hh*x, xa[i] + hh + hh*x\n",
32+
" return hh*np.sum(w * (vV(x1) * phi(i, x1) * phi(j, x1) +\n",
33+
" vV(x2) * phi(i, x2) * phi(j, x2)) )\n",
34+
" \n",
35+
"def V_mat(): # fill potential matrix\n",
36+
" Vm = np.zeros((N-1,N-1))\n",
37+
" for i in range(N): # for each element # contribution to:\n",
38+
" if (i>0): Vm[i-1,i-1] += Vij(i, i) # left node\n",
39+
" if (i < N-1): Vm[i,i] += Vij(i+1, i+1) # right node\n",
40+
" if (i>0 and i< N-1):\n",
41+
" Vm[i-1,i] += Vij(i, i+1) # off diagonals\n",
42+
" Vm[i,i-1] = Vm[i-1,i] # symmetry\n",
43+
" return Vm\n",
44+
"\n",
45+
"a, V0 = 4.0, 4.0 # well width, depth\n",
46+
"xL, xR, N = -4*a, 4*a, 600 # boundaries, num. of elements\n",
47+
"xa = np.linspace(xL, xR, N+1) # nodes\n",
48+
"h = xa[1]-xa[0] # size of element\n",
49+
"\n",
50+
"Tm, B = TB_mat(N-1) # obtain T, B, V matrices\n",
51+
"Vm = V_mat()\n",
52+
" \n",
53+
"E, u = eigsh(Tm + Vm, 10, B, which='SA') # get lowest 10 states \n",
54+
"print (E)"
55+
]
56+
},
57+
{
58+
"cell_type": "code",
59+
"execution_count": null,
60+
"metadata": {
61+
"collapsed": true
62+
},
63+
"outputs": [],
64+
"source": []
65+
}
66+
],
67+
"metadata": {
68+
"kernelspec": {
69+
"display_name": "Python 3",
70+
"language": "python",
71+
"name": "python3"
72+
},
73+
"language_info": {
74+
"codemirror_mode": {
75+
"name": "ipython",
76+
"version": 3
77+
},
78+
"file_extension": ".py",
79+
"mimetype": "text/x-python",
80+
"name": "python",
81+
"nbconvert_exporter": "python",
82+
"pygments_lexer": "ipython3",
83+
"version": "3.6.1"
84+
}
85+
},
86+
"nbformat": 4,
87+
"nbformat_minor": 2
88+
}

ch09/Program_9.4_bem.ipynb

+75
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "code",
5+
"execution_count": null,
6+
"metadata": {},
7+
"outputs": [],
8+
"source": [
9+
"#\n",
10+
"# Program 9.4: Eigenenergies by BEM (bem.ipynb)\n",
11+
"# J Wang, Computational modeling and visualization with Python\n",
12+
"#\n",
13+
"\n",
14+
"import numpy as np, integral as itg\n",
15+
"from scipy.sparse.linalg import eigsh\n",
16+
"from scipy.special import eval_hermite, gamma, ai_zeros\n",
17+
"\n",
18+
"def V(x): # potential\n",
19+
" return beta*abs(x)\n",
20+
"\n",
21+
"def uVu(x): # integrand $u_m V u_n$\n",
22+
" y, c = x/a0, a0*np.sqrt(np.pi*gamma(m+1)*gamma(n+1)*2**(m+n))\n",
23+
" return V(x)*eval_hermite(m,y)*eval_hermite(n,y)*np.exp(-y*y)/c\n",
24+
"\n",
25+
"beta, omega, N = 1., 1., 20 # pot slope, nat freq, num. basis states\n",
26+
"a0, x0 = 1/np.sqrt(omega), 1/(2*beta)**(1./3) # length scale, x0\n",
27+
"\n",
28+
"m, Vm = np.arange(N-2), np.zeros((N,N))\n",
29+
"mm = np.sqrt((m+1)*(m+2)) # calc T matrix\n",
30+
"Tm = np.diag(np.arange(1, N+N, 2,float)) # diagonal\n",
31+
"Tm -= np.diag(mm, 2) + np.diag(mm, -2) # off diagonal by +/-2\n",
32+
"for m in range(N): # calc V matrix\n",
33+
" for n in range(m, N, 2): # m+n is even every 2 steps\n",
34+
" Vm[m, n] = 2*itg.gauss(uVu, 0., 10*a0) # use symm.\n",
35+
" if (m != n): Vm[n,m] = Vm[m,n]\n",
36+
"Tm = Tm*omega/4.\n",
37+
"E, u = eigsh(Tm + Vm, 6, which='SA') # get lowest 6 states\n",
38+
"\n",
39+
"zn, zpn, zbn, zbpn = ai_zeros(3) # find exact values\n",
40+
"Ex = - beta*x0*np.insert(zpn, range(1, 4), zn) # combine even/odd \n",
41+
"print (E, E/Ex)"
42+
]
43+
},
44+
{
45+
"cell_type": "code",
46+
"execution_count": null,
47+
"metadata": {
48+
"collapsed": true
49+
},
50+
"outputs": [],
51+
"source": []
52+
}
53+
],
54+
"metadata": {
55+
"kernelspec": {
56+
"display_name": "Python 3",
57+
"language": "python",
58+
"name": "python3"
59+
},
60+
"language_info": {
61+
"codemirror_mode": {
62+
"name": "ipython",
63+
"version": 3
64+
},
65+
"file_extension": ".py",
66+
"mimetype": "text/x-python",
67+
"name": "python",
68+
"nbconvert_exporter": "python",
69+
"pygments_lexer": "ipython3",
70+
"version": "3.6.1"
71+
}
72+
},
73+
"nbformat": 4,
74+
"nbformat_minor": 2
75+
}

ch09/Program_9.5_hydrogen.ipynb

+104
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "code",
5+
"execution_count": null,
6+
"metadata": {},
7+
"outputs": [],
8+
"source": [
9+
"#\n",
10+
"# Program 9.5: Hydrogen atom by Numerov's method (hydrogen.ipynb)\n",
11+
"# J Wang, Computational modeling and visualization with Python\n",
12+
"#\n",
13+
"\n",
14+
"import matplotlib.pyplot as plt\n",
15+
"import numpy as np, rootfinder as rtf\n",
16+
"%matplotlib notebook\n",
17+
"\n",
18+
"def Veff(r): # effective potential\n",
19+
" return (L*(L+1)/(2*mass*r)-1)/r\n",
20+
" \n",
21+
"def f(r): # Sch eqn in Numerov form\n",
22+
" return 2*mass*(E-Veff(r))\n",
23+
" \n",
24+
"def numerov(f, u, n, x, h): # Numerov integrator for $u''+f(x)u=0$\n",
25+
" nodes, c = 0, h*h/12. # given $[u_0,u_1]$, return $[u_0,u_1,...,u_{n+1}]$\n",
26+
" f0, f1 = 0., f(x+h)\n",
27+
" for i in range(n):\n",
28+
" x += h\n",
29+
" f2 = f(x+h) # Numerov method below, \n",
30+
" u.append((2*(1-5*c*f1)*u[i+1] - (1+c*f0)*u[i])/(1+c*f2)) \n",
31+
" f0, f1 = f1, f2\n",
32+
" if (u[-1]*u[-2] < 0.0): nodes += 1\n",
33+
" return u, nodes # return u, nodes\n",
34+
" \n",
35+
"def shoot(En):\n",
36+
" global E # E needed in f(r)\n",
37+
" E, c, xm = En, (h*h)/6., xL + M*h\n",
38+
" wfup, nup = numerov(f, [0,.1], M, xL, h)\n",
39+
" wfdn, ndn = numerov(f, [0,.1], N-M, xR, -h) # $f'$ from \n",
40+
" dup = ((1+c*f(xm+h))*wfup[-1] - (1+c*f(xm-h))*wfup[-3])/(h+h)\n",
41+
" ddn = ((1+c*f(xm+h))*wfdn[-3] - (1+c*f(xm-h))*wfdn[-1])/(h+h)\n",
42+
" return dup*wfdn[-2] - wfup[-2]*ddn\n",
43+
"\n",
44+
"xL, xR, N = 0., 120., 2200 # limits, intervals\n",
45+
"h, mass = (xR-xL)/N, 1.0 # step size, mass\n",
46+
"Lmax, EL, M = 4, [], 100 # M = matching point\n",
47+
"\n",
48+
"Estart, dE = -.5/np.arange(1, Lmax+1)**2-.1, 0.001 # $\\sim -1/2n^2$\n",
49+
"for L in range(Lmax):\n",
50+
" n, E1, Ea = L+1, Estart[L], []\n",
51+
" while (E1 < -4*dE): # sweep E range for each L\n",
52+
" E1 += dE\n",
53+
" if (shoot(E1)*shoot(E1 + dE) > 0): continue\n",
54+
" E = rtf.bisect(shoot, E1, E1 + dE, 1.e-8)\n",
55+
" Ea.append(E)\n",
56+
" wfup, nup = numerov(f, [0,.1], M-1, xL, h) # calc wf\n",
57+
" wfdn, ndn = numerov(f, [0,.1], N-M-1, xR, -h)\n",
58+
" psix = np.concatenate((wfup[:-1], wfdn[::-1]))\n",
59+
" psix[M:] *= wfup[-1]/wfdn[-1] # match\n",
60+
" print ('nodes, n,l,E=', nup+ndn, n, L, E)\n",
61+
" n += 1\n",
62+
" EL.append(Ea)\n",
63+
" \n",
64+
"plt.figure() # plot energy levels\n",
65+
"for L in range(Lmax):\n",
66+
" for i in range(len(EL[L])):\n",
67+
" plt.plot([L-.3, L+.3], [EL[L][i]]*2, 'k-')\n",
68+
" plt.xlabel('$l$'), plt.ylabel('$E$')\n",
69+
" plt.ylim(-.51, 0), plt.xticks(range(Lmax))\n",
70+
"plt.show()"
71+
]
72+
},
73+
{
74+
"cell_type": "code",
75+
"execution_count": null,
76+
"metadata": {
77+
"collapsed": true
78+
},
79+
"outputs": [],
80+
"source": []
81+
}
82+
],
83+
"metadata": {
84+
"kernelspec": {
85+
"display_name": "Python 3",
86+
"language": "python",
87+
"name": "python3"
88+
},
89+
"language_info": {
90+
"codemirror_mode": {
91+
"name": "ipython",
92+
"version": 3
93+
},
94+
"file_extension": ".py",
95+
"mimetype": "text/x-python",
96+
"name": "python",
97+
"nbconvert_exporter": "python",
98+
"pygments_lexer": "ipython3",
99+
"version": "3.6.1"
100+
}
101+
},
102+
"nbformat": 4,
103+
"nbformat_minor": 2
104+
}

0 commit comments

Comments
 (0)