test_poisson_spherical.py

import numpy as np
import scipy as sp
from matplotlib import pyplot as plt
from poisson_solvers import PoissonSpherical

r_max = 30
nr = 5000
r = np.linspace(0, r_max, nr + 2)

poisson = PoissonSpherical(r, l_max=3)

psi_1s = np.exp(-r)  # Groundstate of the hydrogen atom
psi_1s /= np.sqrt(sp.integrate.simps(psi_1s**2 * r**2, r))  # Normalize

psi_HO = np.exp(
    -(r**2) / 2
)  # Groundstate of the harmonic oscillator with omega = 1
psi_HO /= np.sqrt(sp.integrate.simps(psi_HO**2 * r**2, r))  # Normalize


rho_HO = np.abs(psi_HO[1:-1]) ** 2 / np.sqrt(
    4 * np.pi
)  # The factor 1/(4pi) is due to the integraton over the spherical harmonics
rho_1s = np.abs(psi_1s[1:-1]) ** 2 / np.sqrt(4 * np.pi)

b_HO = -4 * np.pi * r[1:-1] * rho_HO
b_1s = -4 * np.pi * r[1:-1] * rho_1s
W_rmax = np.sqrt(4 * np.pi) * sp.integrate.simps(
    psi_HO**2 * r**2, r
)  # Boundary condition at r_max when l=m=0.


W00_HO = poisson.solve(l=0, f=b_HO, w_rmax=W_rmax)
W00_1s = poisson.solve(l=0, f=b_1s, w_rmax=W_rmax)

W00_exact_HO = np.sqrt(4 * np.pi) * sp.special.erf(r)
W00_exact_1s = np.sqrt(4 * np.pi) * (1 - (r + 1) * np.exp(-2 * r))


plt.figure()
plt.subplot(211)
plt.title(
    "Electrostatic potential generated by the groundstate of the Harmonic oscillator"
)
plt.plot(
    r[1:-1], W00_HO, color="red", linestyle="dashed", label=r"$W^{FDM}(r)$"
)
plt.plot(r[1:-1], W00_exact_HO[1:-1], color="black", label=r"$W^{exact}(r)$")
plt.legend()
plt.subplot(212)
plt.plot(
    r[1:-1],
    np.abs(W00_HO - W00_exact_HO[1:-1]),
    color="black",
    label=r"$|W^{FDM}(r)-W^{exact}(r)|$",
)
plt.legend()

plt.figure()
plt.subplot(211)
plt.title(
    "Electrostatic potential generated by the groundstate of the hydrogen atom"
)
plt.plot(
    r[1:-1], W00_1s, color="red", linestyle="dashed", label=r"$W^{FDM}(r)$"
)
plt.plot(r[1:-1], W00_exact_1s[1:-1], color="black", label=r"$W^{exact}(r)$")
plt.legend()
plt.subplot(212)
plt.plot(
    r[1:-1],
    np.abs(W00_1s - W00_exact_1s[1:-1]),
    color="black",
    label=r"$|W^{FDM}(r)-W^{exact}(r)|$",
)

plt.show()