Add numerov solve

numerov_solve.py

@@ -0,0 +1,155 @@
+"""
+MO 26 April 2024
+modified from: plh 26 April 2024
+Template for Numerical integration using Numerov method.
+Ported from SHOsolve.py
+See https://physics.bu.edu/~py502/lectures4/schrod.pdf and
+https://mtdevans.com/2013/07/solving-the-schrodinger-equation-with-numerovs-algorithm/
+"""
+
+# 1-D Schrodinger wave equation
+#  (-hbar^2/2m) d^2/dx^2 psi(x) + V(x) psi(x) = E psi(x)
+#
+# Spherically symmetric  Psi(xvec)_L,Lz,n = R(r)_L,n Y(phi,theta)_L,Lz
+#  [(-hbar^2/2mu)(1/r)d^2/dr^2 + L(L+1)hbar^2/2mu-r^2 + V(r)] R(r)_L,n = E_L,n R(r)_L,n 
+# Define U(r) = r R(r)
+#  [(-hbar^2/2mu) d^2/dr^2 + L(L+1)hbar^2/2mu-r^2 + V(r)] U(r)_L,n = E_L,n U(r)_L,n 
+#
+
+from typing import Callable
+import numpy as np
+import pylab as pl
+
+# version 0.3
+# version history
+# 0.1 - ported from BohrSolve.py of 2018(ca), retains constants
+# 0.2 - simplified to calculating psi(0) for a list of energies
+# 0.3 - version used for 2024SP PS
+
+# constants
+hbar = 6.626070040e-34 / np.pi / 2  # J-s = Planck's constant
+light_speed = 299792458  # m s-1 = speed of light
+bohr_radius = 0.52917721067e-10  # m = Bohr radius
+mass_proton = 1.672621898e-27  # kg = mass of proton
+mass_electron = 9.10938356e-31  # kg  = mass of electron
+mass_electron_reduced = mass_electron / (1.0 + mass_electron / mass_proton)  # kg = reduced mass of electron in H atom
+electrical_constant = 9.0e9  # N-m^2/s^2 = electrical constant
+charge = 1.602176634e-19  # C = electronic charge
+potential_1 = -electrical_constant * charge * charge  # = potential at r=1m
+potential_0 = potential_1 / bohr_radius  # = potential at r=a0
+energy_1_ev = 13.6  # eV = ground state energy of H atom in eV
+energy_1_j = energy_1_ev * charge  # J = ground state energy of H atom in J = 21.88E-19 J
+
+# initial values
+intervals = 10000  # number of intervals
+radius_max = 200.0 * bohr_radius  # maximum radius = starting point
+radius_min = bohr_radius / 500000.0  # minimum radius = ending point
+# dr = a0/1000000.0   # radius step
+energy = -energy_1_j  # energy, turning point is where E=V(x)
+energy_min = -1 * energy_1_j  # min energy for scan
+energy_max = -0.005 * energy_1_j  # max energy for scan
+energy_intervals = 1000
+
+
+def bohr_potential(r):
+    """Models the Bohr atom electrostatic potential"""
+    return potential_1 / r
+
+
+def do_numerov(radial_vals: np.ndarray, energy: float, initial: tuple[float, float],
+               potential_func: Callable[[float], float], angular_momentum=0.0) -> np.ndarray[float]:
+    """
+    Solves the differential equation F''(r) + g(r) F(r) = 0 using the Numerov method.
+    In the context of the Schrödinger equation, g(r) = k^2 = (E - V(r))*2/m
+    Here, the argument 'potential' points to a function definition that calculates the potential V(r).
+    """
+    hbar_sqrd_over_2m = hbar * hbar / 2 / mass_electron_reduced
+    step_size = radial_vals[1] - radial_vals[0]
+    solution = np.zeros(len(radial_vals))
+    solution[0], solution[1] = initial
+
+    for i in range(2, len(radial_vals)):
+        g_factor = step_size * step_size * (((energy - potential_func(float(radial_vals[i]))) / hbar_sqrd_over_2m)
+                                            - (angular_momentum * (angular_momentum + 1) / radial_vals[i] ** 2)) / 12
+        solution[i] = ((solution[i - 1] * (2.0 - 10.0 * g_factor) - solution[i - 2] * (1.0 + g_factor))
+                       / (1.0 + g_factor))
+    # Normalize the solution so that the total probability equals 1 in the center
+    normalization_factor = solution.max()
+    solution = solution / normalization_factor
+    return solution
+
+
+def opposite_signs(a, b):
+    return a * b < 0
+
+
+def main():
+    # Main program begins here
+    x_vals = np.linspace(radius_max, radius_min, intervals)  # generate x values
+    dx = abs(x_vals[1] - x_vals[0])
+    energies = np.linspace(energy_min, energy_max, energy_intervals).tolist()
+    results = []  # these will be wave function(at r near 0) for each of the energies
+    zeros = []
+    print(f'\nTesting {len(energies)} energies from {energy_min} to {energy_max} for solutions...')
+    for i, energy in enumerate(energies):
+        if i % 100 == 0:
+            print(f' - Testing energy: {i}/{len(energies)} = {energy}...')
+        initial = (0.0, 1.00)  # initial value of function and slope, at large r
+        numerov = do_numerov(x_vals, energy, initial, bohr_potential)  # does integration, returns wave function
+        if len(results) != 0:
+            if opposite_signs(results[-1], numerov[-1]):
+                zero = (energy + energies[i - 1]) / 2
+                print(f'\tFound solution at {zero}')
+                zeros.append(zero)
+        results.append(numerov[-1])
+    print(f'Done testing energies\n')
+    # guess at best solution E
+    # first plot the results as a function of energy
+    pl.figure(4)
+    pl.xlabel('Energy')
+    pl.ylabel('Wave function near r=0')
+    pl.title('Trial results vs Energy')
+    for zero in zeros:
+        pl.axvline(x=zero, color='red')
+    pl.plot(energies, results)
+
+    for n, energy_solution in enumerate(zeros):
+        print(f'Trying E = {energy_solution:.3e} J, {energy_solution / charge:.3} eV')
+        initial = (0.0, 1.0)
+        func_reduced = do_numerov(x_vals, energy_solution, initial, bohr_potential)  # U(r) = reduced wave function
+        # Figure 1 = reduced radial wave function
+        pl.figure(1)
+        pl.xlabel('r')
+        pl.ylabel('U(r)')
+        pl.title('Reduced U(r) = r R(r)')
+        pl.plot(x_vals, func_reduced, label=f'n={n:d}')
+        pl.legend()
+
+        prob = dx * func_reduced.dot(func_reduced)  # calculate total probability, should we x 4pi?
+        func_reduced = func_reduced / np.sqrt(prob)  # total probability should = 1.0!
+        func_radial = func_reduced / x_vals  # get the radial wave function R(r)
+        mean_x = dx * (x_vals * func_reduced).dot(func_reduced)
+        print(f'Expectation value <r> = {mean_x:.3e} m')
+
+        # Figure 2 = actual radial wave function R(r)
+        pl.figure(2)
+        pl.title('Radial wave function R(r) = U(r)/r')
+        pl.plot(x_vals, func_radial, label=f'n={n:d}')
+        pl.xlabel('r')
+        pl.ylabel('R(r)')
+        # pl.ylim(-1,1.1*E)
+        pl.legend()
+
+        # Figure 3 = probability integrand r^2 R(r)^2 = U^2
+        pl.figure(3)
+        pl.title('Probability Integrand [r^2 R(r)^2]')
+        pl.plot(x_vals, func_reduced * func_reduced, label=f'n={n:d}')
+        pl.xlabel('r')
+        pl.ylabel('P(r)')
+        # pl.ylim(-1,1.1*E)
+        pl.legend()
+
+    pl.show()  # show graph window until closed
+
+
+main()