def main(number_of_time_intervals, time_step, want_to_plot, want_to_animate):
#
# ================ Make Directory for Plots if it's not there already =============
#
# detect the current working directory and print it
path = os.getcwd()
print("The current working directory is %s" % path)
# define the name of the directory to be created
Plot_Output = path + "/Plot_Output"
if want_to_plot is True:
if os.path.exists(Plot_Output):
print("Directory for wave function plots already exists", Plot_Output)
else:
print(
"Attempting to create directory for wave function plots ", Plot_Output,
)
try:
os.mkdir(Plot_Output)
except OSError:
print("Creation of the directory %s failed" % Plot_Output)
else:
print("Successfully created the directory %s " % Plot_Output)
# =====================================FEM_DVR===================================
# Set up the FEM DVR grid given only the Finite Element boundaries and order
# of Gauss Lobatto quadrature, and compute the Kinetic Energy matrix for
# this FEM DVR for Mass = mass set in call (atomic units).
#
# Here is where the reduced mass is set
# H_Mass = 1.007825032 #H atom atomic mass
# H_Mass = 1.00727647 #proton atomic mass
# standard atomic weight natural abundance, seems to be what Turner, McCurdy used
H_Mass = 1.0078
Daltons_to_eMass = 1822.89
mu = (H_Mass / 2.0) * Daltons_to_eMass
print("reduced mass mu = ", mu)
bohr_to_Angstrom = 0.529177
Hartree_to_eV = 27.211386245988 # NIST ref
eV_to_wavennumber = 8065.54393734921 # NIST ref on constants + conversions
# value from NIST ref on constants + conversions
Hartree_to_wavenumber = 2.1947463136320e5
HartreeToKelvin = 315773
atu_to_fs = 24.18884326509 / 1000
# Set up the FEM-DVR grid
n_order = 35
FEM_boundaries = [
0.4,
1.0,
2.0,
3.0,
4.0,
5.0,
6.0,
7.0,
8.0,
10.0,
11.0,
12.0,
15.0,
20.0,
25.0,
]
# parameters for Fig 2 in Turner, McCurdy paper
# Julia Turner and C. William McCurdy, Chemical Physics 71(1982) 127-133
scale_factor = np.exp(1j * 20.0 * np.pi / 180.0)
R0 = 10.0 # potential must be analytic for R >= R0
fem_dvr = FEM_DVR(
n_order, FEM_boundaries, Mass=mu, Complex_scale=scale_factor, R0_scale=R0,
)
print("\nFEM-DVR basis of ", fem_dvr.nbas, " functions")
pertubation = Potential()
# ==================================================================================
# Plot potential on the DVR grid points on which the wavefunction is defined
print("\n Plot potential ")
time = 0.0
x_Plot = []
pot_Plot = []
for j in range(0, fem_dvr.nbas):
x_Plot.append(np.real(fem_dvr.x_pts[j + 1]))
pot_Plot.append(np.real(pertubation.V_Bernstein(fem_dvr.x_pts[j + 1], time)))
if want_to_plot is True:
plt.suptitle("V(x) at DVR basis function nodes", fontsize=14, fontweight="bold")
string = "V"
plt.plot(x_Plot, pot_Plot, "ro", label=string)
plt.plot(x_Plot, pot_Plot, "-b")
plt.legend(loc="best")
plt.xlabel(" x ", fontsize=14)
plt.ylabel("V", fontsize=14)
print(
"\n Running from terminal, close figure window to proceed and make .pdf file of figure"
)
# Insert limits if necessary
# Generally comment this logic. Here I am matching the Turner McCurdy Figure 2
# CWM: need to use float() to get plt.xlim to work to set x limits
ymax = float(0.05)
plt.ylim([-0.18, ymax])
# save plot to .pdf file
plt.savefig("Plot_Output/" + "Plot_potential" + ".pdf", transparent=False)
plt.show()
#
# =============Build Hamiltonian (at t=0 if time-dependent)=================================
# Pass name of potential function explicitly here
time = 0.0
H_mat = fem_dvr.Hamiltonian(pertubation.vectorized_V_Bernstein, time)
print("\n Completed construction of Hamiltonian at t = 0")
# ====================================================================================
#
# Find all the eigenvalues of the Hamiltonian so we can compare with known bound state energies
# or make a plot of the spectrum -- For a time-independent Hamiltonian example here
#
# EigenVals = LA.eigvals(H_mat) # eigenvalues only for general matrix. LA.eigvalsh for Hermitian
print(
"Calculating ",
fem_dvr.nbas,
" eigenvalues and eigenvectors for plotting eigenfunctions",
)
EigenVals, EigenVecs = LA.eig(H_mat, right=True, homogeneous_eigvals=True)
print("after LA.eig()")
#
n_energy = fem_dvr.nbas
file_opened = open("Spectrum_ECS.dat", "w")
print("EigenVals shape ", EigenVals.shape)
for i in range(0, n_energy):
print("E( ", i, ") = ", EigenVals[0, i], " hartrees")
print(
np.real(EigenVals[0, i]), " ", np.imag(EigenVals[0, i]), file=file_opened,
)
# ====================================================================================
#
# Extract the n_Plot'th eigenfunction for plotting
#
# pick one of the bound states of Morse Potential to plot
# numbering can depend on numpy and python installation that determines
# behavior of the linear algebra routines.
n_Plot = n_energy - 1 # This is generally the highest energy continuum eigenvalue
n_Plot = 426
wfcnPlot = []
for j in range(0, fem_dvr.nbas):
wfcnPlot.append(EigenVecs[j, n_Plot])
#
# Normalize wave function from diagonalization
# using integration of square on the contour
# following the original Rescigno McCurdy idea for partial widths
# from the paper
# "Normalization of resonance wave functions and the calculation of resonance widths"
# Rescigno, McCurdy, Phys Rev A 34,1882 (1986)
norm_squared = 0.0
for j in range(0, fem_dvr.nbas):
norm_squared = norm_squared + (wfcnPlot[j]) ** 2
wfcnPlot = wfcnPlot / np.sqrt(norm_squared)
norm_squared = 0.0
gamma_residue = 0.0
# background momentum defined with Re(Eres)
k_momentum = np.sqrt(2.0 * mu * np.real(EigenVals[0, n_Plot]))
# k_momentum = np.real(np.sqrt(2.0*mu*EigenVals[0,n_Plot]))
for j in range(0, fem_dvr.nbas):
norm_squared = norm_squared + (wfcnPlot[j]) ** 2
if fem_dvr.x_pts[j + 1] < 5.8:
free_wave = (2.0 * np.sqrt(mu / k_momentum)) * np.sin(
k_momentum * fem_dvr.x_pts[j + 1]
)
gamma_residue = gamma_residue + wfcnPlot[j] * pertubation.V_Bernstein(
fem_dvr.x_pts[j + 1], time
) * free_wave * np.sqrt(fem_dvr.w_pts[j + 1])
print("Complex symmetric inner product (psi|psi) is being used")
if want_to_plot is True:
print(
"Norm of wave function from int psi^2 on contour being plotted is ",
np.sqrt(norm_squared),
)
print(" For this state the asymptotic value of k = ", k_momentum)
print(
"gamma from int = ", gamma_residue, " |gamma|^2 = ", np.abs(gamma_residue) ** 2,
)
# Plot wave function -- It must be type np.complex
Cinitial = np.zeros((fem_dvr.nbas), dtype=np.complex)
wfcnInitialPlot = np.zeros((fem_dvr.nbas), dtype=np.complex)
for j in range(0, fem_dvr.nbas):
Cinitial[j] = wfcnPlot[j]
if want_to_plot is True:
#
# plot n_Plot'th eigenfunction
#
print(
"\n Plot Hamiltonian eigenfunction number ",
n_Plot,
" with energy ",
EigenVals[0, n_Plot],
)
tau = atu_to_fs / (-2.0 * np.imag(EigenVals[0, n_Plot]))
print(" Lifetime tau = 1/Gamma = ", tau, " fs")
number_string = str(n_Plot)
title = "Wavefunction number = " + number_string
wfn_plot_points = 2000
x_Plot_array, Psi_plot_array = fem_dvr.Plot_Psi(
Cinitial,
plot_title_string=title,
N_plot_points=wfn_plot_points,
make_plot=True,
)
#
# Make data file for n_Plot'th eigenfunction
# Psi and the integrand of the residue factors, gamma, of the free-free matrix element
# of the full Green's function, as per
#
filename = "wavefunction" + number_string + ".dat"
file_opened = open(filename, "w")
print_points = len(x_Plot_array)
print("x_Plot_array shape ", print_points)
for i in range(print_points):
free_wave = (2.0 * np.sqrt(mu / k_momentum)) * np.sin(
k_momentum * x_Plot_array[i]
)
# for partial width gamma
integrand = (
Psi_plot_array[i]
* pertubation.V_Bernstein(x_Plot_array[i], time)
* free_wave
)
print(
np.real(x_Plot_array[i]),
" ",
np.imag(x_Plot_array[i]),
" ",
np.real(Psi_plot_array[i]),
" ",
np.imag(Psi_plot_array[i]),
" ",
np.real(integrand),
np.imag(integrand),
file=file_opened,
)
#
# ================# Initialize wave function at t = 0 ================================
# It must be type np.complex
Cinitial = np.zeros((fem_dvr.nbas), dtype=np.complex)
wfcnInitialPlot = np.zeros((fem_dvr.nbas), dtype=np.complex)
for j in range(0, fem_dvr.nbas):
#
# Displaced Gaussian packet in well. Displace to larger R because displacing
# to R < Req produces significant amplitude in the continuum
#
alpha = 40.0
shift = 4.0
Cinitial[j] = (
np.sqrt(np.sqrt(2.0 * alpha / np.pi))
* np.exp(-alpha * (fem_dvr.x_pts[j + 1] - shift) ** 2)
* np.sqrt(fem_dvr.w_pts[j + 1])
)
wfcnInitialPlot[j] = Cinitial[j]
#
# =====================Propagation from t = tinitial to tfinal ============================
# initialize arrays for storing the information for all the frames in the animation of
# the propagation of the wave function
times_array = []
x_Plot_time_array = []
Psi_plot_time_array = []
tinitial = 0.0
# specify final time and specify number of intervals for plotting
# commented info for Morse oscillator revivals, in atu
# a = 1.0277
# d = 0.1746
# omega = np.sqrt((2.e0/mu)*d*a**2)
# T_revival = 4*np.pi*mu/a**2 # for the morse oscillator
# number_of_time_intervals = 25*np.int((tfinal-tinitial)*omega/(2.0*np.pi))
# print("T_revival = ",T_revival," atomic time units ",T_revival*24.1888/1000.0," femtoseconds")
tfinal = 5000 # specified in atomic time units
print(
"\nOverall propagation will be from ",
tinitial,
" to ",
tfinal,
" atu in ",
number_of_time_intervals,
" intervals",
)
if want_to_plot is True:
#
# plot initial wave packet
#
print("\n Plot initial wave function at t = ", tinitial)
number_string = str(0.0)
title = "Wavefunction at t = " + number_string
x_Plot_array, Psi_plot_array = fem_dvr.Plot_Psi(
Cinitial, plot_title_string=title, N_plot_points=1250, make_plot=True,
)
times_array.append(tinitial)
x_Plot_time_array.append(x_Plot_array)
Psi_plot_time_array.append(Psi_plot_array)
#
# Loop over number_of_time_intervals intervals that make up t = 0 to tfinal
#
t_interval = (tfinal - tinitial) / number_of_time_intervals
# for H2 reduced mass, Deltat = 0.05 in Morse oscillator seems to work for 10s of fs
# for D2 reduced mass, Deltat = 0.1 in Morse oscillator seems to work for at least 2 ps
# based on comparison with time steps 5 to 10 times smaller
time_step = 0.1 # atomic time units
# time_step = 0.05 # atomic time units
# time_step = 0.015 # atomic time units
for i_time_interval in range(0, number_of_time_intervals):
t_start = tinitial + i_time_interval * t_interval
t_finish = tinitial + (i_time_interval + 1) * t_interval
N_time_steps = np.int(t_interval / time_step)
Delta_t = t_interval / np.float(N_time_steps)
#
# Check norm of initial wavefunction on real part of ECS Contour
#
# find FEM DVR function at the last real DVR point
first_ECS_element = fem_dvr.i_elem_scale
last_real_dvr_fcn = 1
for i in range(0, fem_dvr.nbas):
if fem_dvr.x_pts[i + 1] < FEM_boundaries[first_ECS_element]:
last_real_dvr_fcn = last_real_dvr_fcn + 1
# the point last_real_dvr_fcn still has complex bridging function weight.
# small error for integral (0,R0) may result.
norm_initial = 0
for j in range(0, last_real_dvr_fcn):
norm_initial = norm_initial + np.abs(Cinitial[j]) ** 2
print(
"\nNorm of starting wave function on real part of ECS contour",
norm_initial,
)
#
# use the propagator from DVR()
#
clock_start = timeclock.time()
Ctfinal = fem_dvr.Crank_Nicolson(
t_start,
t_finish,
N_time_steps,
Cinitial,
pertubation.vectorized_V_Bernstein,
Is_H_time_dependent=False,
)
clock_finish = timeclock.time()
print(
"Time for ",
N_time_steps,
" Crank-Nicolson steps was ",
clock_finish - clock_start,
" secs.",
)
#
# Check norms of initial and final wavefunctions
#
norm_final = 0
for j in range(0, last_real_dvr_fcn):
# should be change to real part of contour.
norm_final = norm_final + np.abs(Ctfinal[j]) ** 2
print(
"Norm of final wave function on real part of ECS contour", norm_final,
)
if want_to_plot is True:
#
# Plot of packet at end of each interval using Plot_Psi from DVR()
#
print(
"Plot function propagated to t = ",
t_finish,
" atomic time units ",
t_finish * atu_to_fs,
" fs",
)
number_string = str(t_finish)
title = "Wavefunction at t = " + number_string
x_Plot_array, Psi_plot_array = fem_dvr.Plot_Psi(
Ctfinal, plot_title_string=title, N_plot_points=1250, make_plot=False,
)
times_array.append(t_finish)
x_Plot_time_array.append(x_Plot_array)
Psi_plot_time_array.append(Psi_plot_array)
#
# reset initial packet as final packet for this interval
for i in range(0, fem_dvr.nbas):
Cinitial[i] = Ctfinal[i]
#
if want_to_plot is True:
x_Plot_array, Psi_plot_array = fem_dvr.Plot_Psi(
Ctfinal, plot_title_string=title, N_plot_points=1250, make_plot=True
)
print(
"Final frame at tfinal = ",
tfinal,
"atomic time units is showing -- ready to prepare animation ",
)
print("\n **** close the plot window to proceed ****")
# ==============================================================================
# # initialization function: plot the background of each frame
# ==============================================================================
def init():
line.set_data([], [])
time_text.set_text("")
ax.set_xlabel(" x (bohr) ", fontsize=16, fontweight="bold")
ax.set_ylabel(" Psi(t): Re, Im & Abs ", fontsize=16, fontweight="bold")
fig.suptitle("Wave Packet in H2 Potential", fontsize=16, fontweight="bold")
# put in a line at the value Phi = 0
ax.plot([x_Plot[0], x_Plot[len(x_Plot) - 1]], [0, 0], "k")
return line, time_text
nframes = number_of_time_intervals
def animate(i):
# ==============================================================================
#
# commands to specify changing elements of the animation
# We precomputed all the lines at all the times so that this
# wouldn't recompute everything at each frame, which is very slow...
#
# ==============================================================================
time_string = str(times_array[i] * 24.189 / 1000.0) # string for a plot label
re_array = np.real(Psi_plot_time_array[i])
im_array = np.imag(Psi_plot_time_array[i])
abs_array = np.abs(Psi_plot_time_array[i])
line1.set_data(x_Plot_array, re_array)
line2.set_data(x_Plot_array, im_array)
line3.set_data(x_Plot_array, abs_array)
time_text.set_text("time = " + time_string + " fs")
return (line1, line2, line3, time_text)
if want_to_animate is True:
# ==============================================================================
# Now plot the animation
# ==============================================================================
# reinitialize the figure for the next plot which is the animation
fig = plt.figure()
ymax = 1.25
xmin = x_Plot[0]
xmax = 20.0
ax = fig.add_subplot(
111, autoscale_on=False, xlim=(xmin, xmax), ylim=(-ymax, +ymax)
)
(line,) = ax.plot([], [], "-r", lw=2)
(line1,) = ax.plot([], [], "-r", lw=2)
(line2,) = ax.plot([], [], "-b", lw=2)
(line3,) = ax.plot([], [], "k", lw=2)
# define the object that will be the time printed on each frame
time_text = ax.text(0.02, 0.95, "", transform=ax.transAxes)
# ==============================================================================
# call the animator. blit=True means only re-draw the parts that have changed.
# ==============================================================================
anim = animation.FuncAnimation(
fig,
animate,
init_func=init,
frames=nframes + 1,
interval=200,
blit=True,
repeat=True,
)
# ==============================================================================
# show the animation
# ==============================================================================
anim.save("Plot_Output/H2_wavepacket.mp4")
plt.show()
print("done")
if want_to_plot is False:
print(
"\n\nSet the command line option want_to_plot=True to see figures and create plotting directory.\n\n"
)
if want_to_animate is False:
print("Set the command line option want_to_animate=True to see animation.\n\n")
def main(want_to_plot):
#
# ============== Make Directory for Plots if it's not there already ==========
#
# detect the current working directory and print it
path = os.getcwd()
print("The current working directory is %s" % path)
# define the name of the directory to be created
Plot_Output = path + "/Plot_Output"
if want_to_plot is True:
if os.path.exists(Plot_Output):
print("Directory for wave function plots already exists",
Plot_Output)
else:
print(
"Attempting to create directory for wave function plots ",
Plot_Output,
)
try:
os.mkdir(Plot_Output)
except OSError:
print("Creation of the directory %s failed" % Plot_Output)
else:
print("Successfully created the directory %s " % Plot_Output)
#
# === Make Directory for for the .dat output if it's not there already ======
#
# detect the current working directory and print it
path = os.getcwd()
print("The current working directory is %s" % path)
# define the name of the directory to be created
Data_Output = path + "/Data_Output"
if os.path.exists(Data_Output):
print("Directory for output .dat files already exists", Plot_Output)
else:
print(
"Attempting to create directory for .dat output files ",
Plot_Output,
)
try:
os.mkdir(Data_Output)
except OSError:
print("Creation of the directory %s failed" % Data_Output)
else:
print("Successfully created the directory %s " % Data_Output)
# ===================================================================
if want_to_plot is True:
wfcn_plotfile = open("Data_Output/wavefunctions.dat",
"w") # data file for saving wavefunctions
#
# =============Constants and conversion factors ==============
Daltons_to_eMass = 1822.89
bohr_to_Angstrom = 0.529177
Hartree_to_eV = 27.211386245988 # NIST ref
eV_to_wavennumber = 8065.54393734921 # NIST ref on constants + conversions
Hartree_to_wavenumber = (
2.1947463136320e5 # value from NIST ref on constants + conversions
)
atu_to_fs = 24.18884326509 / 1000
HartreeToKelvin = 315773
#
# =====================================FEM_DVR===================================
# Set up the FEM DVR grid given only the Finite Element boundaries and order
# of Gauss Lobatto quadrature, and compute the Kinetic Energy matrix for
# this FEM DVR for Mass = mass set in call (atomic units).
#
# Here is where the reduced mass (or mass for any 1D problem) is set
O_Mass = 15.99491461957 # O 16 mass from NIST tables
C_Mass = 12.0 # C 12 mass
mu = (O_Mass * C_Mass / (O_Mass + C_Mass)) * Daltons_to_eMass
n_order = 30
FEM_boundaries = [
0.801,
1.25,
1.5,
2.0,
2.5,
3.0,
3.5,
4.0,
4.5,
5.0,
5.5,
6.0,
6.5,
7.0,
7.5,
8.0,
8.5,
8.99,
]
fem_dvr = FEM_DVR(n_order, FEM_boundaries, Mass=mu)
print("\nFEM-DVR basis of ", fem_dvr.nbas, " functions")
#
# Function to define potential at x and t (if potential is time-dependent)
# DVRHelper class library expects function for V(x,t) in general
#
# =================================Potential=====================================
# Read in files with points for potential curve in hartrees
# and load in arrays for interpolation
path = Path(__file__).parent.absolute()
perturbation = Potential(path / "potcurve_CISDT_CO_ccpvDZ.dat")
n_vals_pot = perturbation.r_data.shape[0]
#
# ===========================================================================================
# Plot potential on the DVR grid points on which the wavefunction is defined
# and ALSO the interpolation to check we are using the potential that we mean to.
#
if want_to_plot is True:
print("\n Plot potential ")
x_Plot = []
pot_Plot = []
for j in range(0, n_vals_pot):
x_Plot.append(np.real(perturbation.r_data[j]))
pot_Plot.append(np.real(perturbation.V_data[j]))
plt.suptitle("V(r) interpolation", fontsize=14, fontweight="bold")
string = "V input points"
plt.plot(x_Plot, pot_Plot, "ro", label=string)
#
x_Plot = []
pot_Plot = []
Number_plot_points = 731
dx = (np.real(fem_dvr.x_pts[fem_dvr.nbas - 1]) -
np.real(fem_dvr.x_pts[0])) / float(Number_plot_points - 1)
time = 0.0 # dummy time in general call to potential function
for j in range(0, Number_plot_points):
x = np.real(fem_dvr.x_pts[0]) + j * dx
try:
x >= perturbation.r_data[0] and x <= perturbation.r_data[
n_vals_pot - 1]
except IndexError:
print(
"Number of plot points is out of range of pertubation data"
)
x_Plot.append(x)
pot_Plot.append(perturbation.V_Interpolated(x, time))
plt.plot(x_Plot,
pot_Plot,
"-b",
label="Interpolation on DVR grid range")
plt.legend(loc="best")
plt.xlabel(" x ", fontsize=14)
plt.ylabel("V", fontsize=14)
print(
"\n Running from terminal, close figure window to proceed and make .pdf file of figure"
)
# Insert limits if necessary
# xmax = float(rmax) # CWM: need to use float() to get plt.xlim to work to set x limits
# plt.xlim([0,xmax])
# number_string = str(a)
plt.savefig("Plot_Output/" + "Plot_potential" + ".pdf",
transparent=False)
plt.show()
#
#
# =============Build Hamiltonian (using general routine with dummy time t=0)=========
# Pass name of potential function explicitly here
time = 0.0
H_mat = fem_dvr.Hamiltonian(perturbation.vectorized_V_Interpolated, time)
print("\n Completed construction of Hamiltonian ")
# ====================================================================================
#
# Find all the eigenvalues of the Hamiltonian so we can compare with known bound state energies
# or make a plot of the spectrum -- For a time-independent Hamiltonian example here
#
EigenVals = LA.eigvalsh(H_mat)
#
n_energy = 20
for i in range(0, n_energy):
print(
"E( ",
i,
") = ",
EigenVals[i],
" hartrees, excitation energy = ",
(EigenVals[i] - EigenVals[0]) * Hartree_to_wavenumber,
" cm^-1",
)
if want_to_plot is True:
# ====================================================================================
#
# Extract the n_Plot'th eigenfunction for plotting and use as initial wave function
# to test propagation
#
number_of_eigenvectors = 20
#
# Here n_Plot picks which eigenfunction to plot
n_Plot = 10 # pick a state of this potential to plot < number_of_eigenvectors -1
#
print(
"Calculating ",
number_of_eigenvectors,
" eigenvectors for plotting eigenfunctions",
)
EigenVals, EigenVecs = LA.eigh(H_mat,
eigvals=(0, number_of_eigenvectors))
wfcnPlot = []
for j in range(0, fem_dvr.nbas):
wfcnPlot.append(EigenVecs[j, n_Plot])
#
# normalize wave function from diagonalization
#
norm_squared = 0.0
for j in range(0, fem_dvr.nbas):
norm_squared = norm_squared + np.abs(wfcnPlot[j])**2
wfcnPlot = wfcnPlot / np.sqrt(norm_squared)
norm_squared = 0.0
for j in range(0, fem_dvr.nbas):
norm_squared = norm_squared + np.abs(wfcnPlot[j])**2
print("Norm of wave function being plotted is ", np.sqrt(norm_squared))
#
# ================# Plot the wave function specified by n_Plot above======================
# It must be type np.complex for this general wave function plotting logic
#
Cinitial = np.zeros((fem_dvr.nbas), dtype=np.complex)
wfcnInitialPlot = np.zeros((fem_dvr.nbas), dtype=np.complex)
for j in range(0, fem_dvr.nbas):
Cinitial[j] = wfcnPlot[j]
#
print(
"\n Plot wave function ",
n_Plot,
" (numbered in order of increasing energy)",
)
title = "Wavefunction" + str(n_Plot)
# note that the dvr.Plot_Psi function makes a .pdf file in the Plot_Output directory
# That's what make_plot=True controls.
x_Plot_array, Psi_plot_array = fem_dvr.Plot_Psi(
Cinitial,
plot_title_string=title,
N_plot_points=Number_plot_points,
make_plot=True,
)
# write the data in file also
for j in range(0, Number_plot_points):
print(
x_Plot_array[j],
" ",
np.real(Psi_plot_array[j]),
" ",
np.imag(Psi_plot_array[j]),
file=wfcn_plotfile,
)
print("& \n ", file=wfcn_plotfile)
#
# plot square of wave function (radial probability distribution)
#
Csquared = np.zeros(fem_dvr.nbas, dtype=np.complex)
raverage = 0.0
for i in range(fem_dvr.nbas):
Csquared[i] = (Cinitial[i]**2) / np.sqrt(fem_dvr.w_pts[i + 1])
# compute <r> for this wave function
# note that Cinitial[i] contains the necessary weight, sqrt(dvr.w_pts[i+1])
raverage = raverage + Cinitial[i]**2 * fem_dvr.x_pts[i + 1]
print("\n Average value of r using DVR for the integral, <r> = ",
raverage)
title = "Wavefunction" + str(n_Plot) + "^2"
# note that the dvr.Plot_Psi function makes a .pdf file in the Plot_Output directory
# That's what make_plot=True controls.
x_Plot_array, Psi_plot_array = fem_dvr.Plot_Psi(
Csquared,
plot_title_string=title,
N_plot_points=Number_plot_points,
make_plot=True,
)
# write the data in file also
for j in range(0, Number_plot_points):
print(
x_Plot_array[j],
" ",
np.real(Psi_plot_array[j]),
" ",
np.imag(Psi_plot_array[j]),
file=wfcn_plotfile,
)
print("& \n ", file=wfcn_plotfile)
def main(want_to_plot):
# Making directories for plots and data
if want_to_plot is True:
path = os.getcwd()
Plot_Output = path + "/Plot_Output_2D"
if os.path.exists(Plot_Output):
print("Directory for wave function plots already exists",
Plot_Output)
else:
print(
"Attempting to create directory for wave function plots ",
Plot_Output,
)
try:
os.mkdir(Plot_Output)
except OSError:
print("Creation of the directory %s failed" % Plot_Output)
else:
print("Successfully created the directory %s " % Plot_Output)
#
Data_Output = path + "/Data_Output_2D"
if os.path.exists(Data_Output):
print("Directory for output .dat files already exists",
Data_Output)
else:
print(
"Attempting to create directory for .dat output files ",
Data_Output,
)
try:
os.mkdir(Data_Output)
except OSError:
print("Creation of the directory %s failed" % Data_Output)
else:
print("Successfully created the directory %s " % Data_Output)
#%%
# ===============================FEM_DVR=========================================
# Set up the FEM DVR grid given the Finite Element boundaries and order
# of Gauss Lobatto quadrature, and compute the Kinetic Energy matrix for
# this FEM DVR for Mass = mass set in call (atomic units).
#
# Here is where the reduced mass in kinetic energy is set
Elec_Mass = 1.0
mu = Elec_Mass
print("reduced mass mu = ", mu)
bohr_to_Angstrom = 0.529177
Hartree_to_eV = 27.211386245988 # NIST ref
eV_to_wavennumber = 8065.54393734921 # NIST ref on constants + conversions
Hartree_to_wavenumber = (
2.1947463136320e5 # value from NIST ref on constants + conversions
)
HartreeToKelvin = 315773
atu_to_fs = 24.18884326509 / 1000
#
# Specify the FEM-DVR grid
#
n_order = (
15 # grid setup here is for calculating bound states of He or H- target
)
FEM_boundaries = [0.0, 1.0, 5.0, 10.0, 20.0, 30.0]
theta_scale = 20
scale_factor = np.exp(1j * theta_scale * np.pi /
180.0) # scale_factor = 1 (Integer) for no scaling
R0 = 19.0 # Complex scaling begins at greatest FEM boundary < R0 specified here
print("ECS scaling by ", theta_scale, " degrees at least FEM boundary > ",
R0)
fem_dvr = FEM_DVR(
n_order,
FEM_boundaries,
Mass=mu,
Complex_scale=scale_factor,
R0_scale=R0,
)
print("\nFEM-DVR basis of ", fem_dvr.nbas, " functions")
# Initialize perturbations
pertubation = Potential()
# =============Build one-electron Hamiltonian ===============================
# Pass name of one-body potential function explicitly here
time = 0.0
H_mat = fem_dvr.Hamiltonian(pertubation.vectorized_V_Coulomb, time)
print("\n Completed construction of 1D Hamiltonian at t = 0")
#
# ========= Calculate two-electron integrals in FEM-DVR basis ===============
# These are diagonal in the FEM-DVR basis satisfying
# <ij|kl> = delta_ik delta_jl v_ij See
# C W McCurdy et al 2004 J. Phys. B: At. Mol. Opt. Phys. 37 R137
#
# They therefore have the same form as a local potential V(r1,r2)
# that would be represented by V(r1_i,r2_j) = v_ij at points on the grid
# We pass the matrix v_ij to fem_dvr.Hamiltonian_2D() to build the Hamiltonian
# in both cases
#
Temkin_Poet = (
True # Temkin-Poet model is the radial limit model 1/r12 = 1/r>
)
T_mat = 2.0 * fem_dvr.KE_mat
Two_electron_ints = LA.inv(T_mat)
R0_ECS = fem_dvr.FEM_boundaries[fem_dvr.i_elem_scale]
R_max = fem_dvr.FEM_boundaries[len(fem_dvr.FEM_boundaries) - 1]
if Temkin_Poet:
print("\n Evaluating two-electron integrals for Temkin-Poet model")
print(
" Grid has R0_ECS = ",
R0_ECS,
" R_max ",
R_max,
" shape of inverse of KE ",
Two_electron_ints.shape,
)
for i in range(fem_dvr.nbas):
for j in range(fem_dvr.nbas):
if Temkin_Poet:
Two_electron_ints[i, j] = (
Two_electron_ints[i, j] /
(fem_dvr.x_pts[i + 1] * fem_dvr.x_pts[j + 1] *
np.sqrt(fem_dvr.w_pts[i + 1] * fem_dvr.w_pts[j + 1])) +
1.0 / R_max)
else:
Two_electron_ints[i, j] = pertubation.V_colinear_model(
fem_dvr.x_pts[i + 1],
fem_dvr.x_pts[j + 1]) # colinear model
#
# ============ Build two-electron Hamiltonian ================================
# pass matrix of two-electron integrals in FEM-DVR or matrix of FEM-DVR
# representation of local two-electron interaction
#
H_mat_2D = fem_dvr.Hamiltonian_2D(H_mat, Two_electron_ints, time)
print("\n Completed construction of 2D Hamiltonian")
# =============================================================================
#
# ====== Find all the eigenvalues and eigenvectorsof the 2D Hamiltonian =====
#
# Here we use LA.eig, which solves the complex generalized eigenvalue & eigenvector problem
# (use LA.eigvalsh() for Hermitian matrices)
# Called here with arguments right=True for right hand eigenvectors
#
print(
"Calculating ",
fem_dvr.nbas**2,
" eigenvalues and eigenvectors of 2D Hamiltonian ",
)
EigenVals, EigenVecs = LA.eig(H_mat_2D,
right=True,
homogeneous_eigvals=False)
print("after LA.eig()")
#
# Sort eigenvalues and eigenvectors using np.argsort() function
#
print(" shape of EigenVals on return from LA.eig ", EigenVals.shape)
print(" shape of EigenVecs on return from LA.eig ", EigenVecs.shape)
size_eigen_system = EigenVals.shape[
0] # note shape is an array two long = (2,dimension)
print("size_eigen_system = ", size_eigen_system)
real_parts_eigenvals = np.empty(size_eigen_system)
for i in range(0, size_eigen_system):
real_parts_eigenvals[i] = np.real(EigenVals[i])
idx = np.argsort(
real_parts_eigenvals) # argsort() returns indices to sort array
EigenVals = EigenVals[
idx] # mysterious python syntax to indirectly address the array
EigenVecs = EigenVecs[:,
idx] # mysterious python syntax to indirectly address 2D array
#
# Make a file with sorted complex eigenvalues
#
n_energy = fem_dvr.nbas * fem_dvr.nbas
n_lowest = 10
print("\n Lowest ", n_lowest, " eigenvalues of 2D Hamiltonian")
for i in range(0, min(n_lowest, n_energy)):
print(
"E( ",
i,
") = ",
np.real(EigenVals[i]),
" ",
np.imag(EigenVals[i]),
" i",
" hartrees",
)
#
if want_to_plot is True:
file_opened = open(Data_Output + "/Spectrum_ECS.dat", "w")
for i in range(0, n_energy):
print(
np.real(EigenVals[i]),
" ",
np.imag(EigenVals[i]),
file=file_opened,
)
print(" Complete spectrum, sorted on Re(energy) written to file ")
#
# ============= Plot of eigenfunctions of 2D Hamiltonian ==================
#
n_state_plot = 31 # n_state_plot = 0 is lowest energy (by real part) state
print(
"Plotting eigenstate ",
n_state_plot,
" with energy ",
EigenVals[n_state_plot],
)
Coef_vector = []
for i in range(0, n_energy):
Coef_vector.append(EigenVecs[i, n_state_plot])
# normalize vector
sum = 0.0
for i in range(0, n_energy):
sum = sum + np.abs(Coef_vector[i])**2
Norm_const = 1.0 / np.sqrt(sum)
for i in range(0, n_energy):
Coef_vector[i] = Norm_const * Coef_vector[i]
#
# ======== Optional plot at FEM-DVR points only without interpolation ============
#
Plot_at_DVR_only = False
if want_to_plot is True:
if Plot_at_DVR_only:
print(" making plot data file at DVR grid points")
plot_file_opened = open(Data_Output + "/wave_function_2D.dat", "w")
for i in range(1, fem_dvr.nbas + 1):
for j in range(1, fem_dvr.nbas + 1):
print(
np.real(fem_dvr.x_pts[i]),
" ",
np.real(fem_dvr.x_pts[j]),
" ",
np.real(Coef_vector[j + (i - 1) * fem_dvr.nbas - 1] /
np.sqrt(fem_dvr.w_pts[i] * fem_dvr.w_pts[j])),
file=plot_file_opened,
)
print(" ", file=plot_file_opened)
exit()
#
# =========== Plot of wave function on a grid evenly spaced within FEMs ==============
#
if want_to_plot is True:
plot_pts = 50
x_Plot_array, y_Plot_array, Psi_plot_array = fem_dvr.Plot_Psi_2D(
Coef_vector,
plot_title_string="Real_Psi_on_2D_ECS_grid",
N_plot_points=plot_pts,
make_plot=True,
)
plot_file_opened = open(Data_Output + "/wave_function_2D.dat", "w")
print(" Input number of plotting points in each dimension ", plot_pts)
print(
" Actual number of plotting ponts ",
len(x_Plot_array),
" evenly spaced in elements on ECS contour",
)
# write data file for external plotting
ipt = 0
plot_pts_on_contour = len(
x_Plot_array
) # can differ from plot_pts, PLot_Psi_2D puts an integer number in each element
for i in range(plot_pts_on_contour):
for j in range(plot_pts_on_contour):
print(
np.real(x_Plot_array[i]),
" ",
np.real(y_Plot_array[j]),
" ",
np.real(Psi_plot_array[ipt]),
file=plot_file_opened,
)
ipt = ipt + 1
print(" ", file=plot_file_opened)
print(
" File of points for plot written: ",
Data_Output + "/wave_function_2D.dat",
)
def main(number_of_time_intervals, time_step, want_to_plot_animate):
# ================ Make Directory for Plots if it's not there already =============
#
# detect the current working directory and print it
path = os.getcwd()
print("The current working directory is %s" % path)
# define the name of the directory to be created
Plot_Output = path + "/Plot_Output"
if want_to_plot_animate is True:
if os.path.exists(Plot_Output):
print("Directory for wave function plots already exists",
Plot_Output)
else:
print(
"Attempting to create directory for wave function plots ",
Plot_Output,
)
try:
os.mkdir(Plot_Output)
except OSError:
print("Creation of the directory %s failed" % Plot_Output)
else:
print("Successfully created the directory %s " % Plot_Output)
# =====================================FEM_DVR===================================
# Set up the FEM DVR grid given only the Finite Element boundaries and order
# of Gauss Lobatto quadrature, and compute the Kinetic Energy matrix for
# this FEM DVR for Mass = mass set in call (atomic units).
#
massH = 1836.0
mu = massH / 2.0
n_order = 30 # 40th order with finite elements of 1 or 2 bohr OK for light masses here
FEM_boundaries = [0.0, 1.0, 2.0, 3.0, 4.0, 6.0, 8.0]
#
# constants and conversions
#
Daltons_to_eMass = 1822.89
bohr_to_Angstrom = 0.529177
Hartree_to_wavenumber = 8065.54 * 27.211
Hartree_to_eV = 27.211
atu_to_fs = 24.18884 / 1000.0
c_light = 137.035999084 # speed of light in atomic units
a0_to_m = bohr_to_Angstrom * 10**(-10)
I_0 = (
3.5095 * 10**16
) # conversion for intensity in W/cm^2 I =I_0*(A_0*omega/c_light)**2
#
print("Mass = ", mu)
fem_dvr = FEM_DVR(n_order, FEM_boundaries, Mass=mu)
print("\nFEM-DVR basis of ", fem_dvr.nbas, " functions")
#
# pulse parameters for the dipole coupling between the two
# electronic states
#
T_pulse = 3.0 # set in femtoseconds
omega = 0.2 # central frequency 0.113903 hartrees is 400 nm
A_0 = 10.0 # strength of A field 6.422 atomic units is 10^12 W/cm^2 for 400 nm
print("\n\n*** Radiation Pulse ***")
print(" Duration = ", T_pulse, " fs, ")
print(
" Intensity of pulse = ",
"%15.3e" % (I_0 * (A_0 * omega / c_light)**2),
" W/cm^2",
)
print(
" Central energy = ",
omega * Hartree_to_eV,
" eV, lambda = ",
(10**9) * a0_to_m * 2 * np.pi * c_light / omega,
" nm",
)
def E_field(time):
#
# Definition of E in - mu E(t) radiative coupling between electronic states
# One pulse here, followed by A = 0, durations and other parameters set in main program
#
if 0 <= time * atu_to_fs <= T_pulse:
T = T_pulse / atu_to_fs # pulse duration converted from fs
field = ((A_0 * omega / c_light) * (np.sin(np.pi * time / T)**2) *
np.cos(omega * time - omega * T / 2))
else:
field = 0.0
return field
def Vcoupling(x, time):
#
# mu_dipole * E(t)
#
field = E_field(time)
mu_dipole = 1.0 # assumed value of dipole matrix element
coupling = mu_dipole * field
return coupling
def Vzero(x, time):
Vcoupzero = 0.0
return Vcoupzero
pertubation = Potential()
# =============Build Hamiltonian without coupling to pick an initial state=========
# Pass name of potential function explicitly here
time = 0.0
H_mat = fem_dvr.Hamiltonian_Two_States(pertubation.V_morse_1,
pertubation.V_morse_2, Vzero, time)
print("\n Completed construction of Hamiltonian at t = 0")
# ====================================================================================
#
# Find all the eigenvalues of the Hamiltonian so we can compare with known bound state energies
# or make a plot of the spectrum -- For a time-independent Hamiltonian example here
#
EigenVals_diabatic = LA.eigvalsh(H_mat)
# ====================================================================================
#
# Print the lowest few eigenvalues as a check, and help pick an initial state in one diabatic well
print(
"\nEigenvalues of the uncoupled Hamiltonians for the two potential curves"
)
print("n E_uncoupled")
for i in range(0, 50):
print(i, " ", EigenVals_diabatic[i])
#
# ===========================================================================================
# Plot potential on the DVR grid points on which the wavefunction is defined
time = (1 / 2) * T_pulse / atu_to_fs # Maximum of pulse
#
print("\n Plot potential ")
x_Plot = []
pot1_Plot = []
pot2_Plot = []
coup_Plot = []
for j in range(0, fem_dvr.nbas):
x_Plot.append(fem_dvr.x_pts[j + 1])
pot1_Plot.append(pertubation.V_morse_1(fem_dvr.x_pts[j + 1], time))
pot2_Plot.append(pertubation.V_morse_2(fem_dvr.x_pts[j + 1], time))
coup_Plot.append(10.0 * Vcoupling(fem_dvr.x_pts[j + 1], time))
if want_to_plot_animate is True:
plt.suptitle(
"Uncoupled potentials and max value of coupling",
fontsize=14,
fontweight="bold",
)
string = "V1"
# plt.plot(x_Plot,pot1_Plot,'ro',label=string) # plot points only
plt.plot(x_Plot, pot1_Plot, "-r", label=string)
string = "V2"
plt.plot(x_Plot, pot2_Plot, "-b", label=string)
string = "Max Vcoup x 10"
plt.plot(x_Plot, coup_Plot, "-g", label=string)
plt.legend(loc="best")
plt.xlabel(" x (bohr)", fontsize=14)
plt.ylabel("V (hartrees)", fontsize=14)
print(
"\n Running from terminal, close figure window to proceed and make .pdf file of figure"
)
# Insert limits if necessary
# xmax = float(rmax) # CWM: need to use float() to get plt.xlim to work to set x limits
# plt.xlim([0,xmax])
plt.ylim([-0.2, 0.2])
# number_string = str(a)
plt.savefig(
"Plot_Output/" + "Plot_two_state_potentials" + ".pdf",
transparent=False,
)
plt.show()
#
# ====================================================================================
#
# Extract the n_Plot'th eigenfunction of uncoupled Hamiltonians for plotting
#
#
np1 = 20
number_of_eigenvectors = np1
# Initial wave function chosen here -- change n_Plot to change it
n_Plot = 0 # pick one of the bound states of Potential to plot and use as initial state below
print("Calculating ", np1, " eigenvectors for plotting eigenfunctions")
EigenVals = []
EigenVals, EigenVecs = LA.eigh(H_mat, eigvals=(0, number_of_eigenvectors))
wfcnPlot = []
for j in range(0, 2 * fem_dvr.nbas):
wfcnPlot.append(EigenVecs[j, n_Plot])
#
# normalize 2 component wave function from diagonalization
# so that it has the right overall normalization
#
norm_squared = 0.0
for j in range(0, 2 * fem_dvr.nbas):
norm_squared = norm_squared + np.abs(wfcnPlot[j])**2
wfcnPlot = wfcnPlot / np.sqrt(norm_squared)
norm_squared = 0.0
for j in range(0, 2 * fem_dvr.nbas):
norm_squared = norm_squared + np.abs(wfcnPlot[j])**2
print("Norm of wave function being plotted is ", np.sqrt(norm_squared))
# now extract the two components
wfcnPlot1 = []
wfcnPlot2 = []
for j in range(0, fem_dvr.nbas):
wfcnPlot1.append(wfcnPlot[j, ])
wfcnPlot2.append(wfcnPlot[j + fem_dvr.nbas])
#
# Uncomment to both components of the eigenfunction. For uncoupled Hamiltonians
# only one component will be nonzero for each eigenfunction.
#
# x_Plot_array, Psi1_plot_array, Psi2_plot_array = fem_dvr.Plot_Psi_Two_States(wfcnPlot1, wfcnPlot2,plot_title_string="psi_1 and psi_2 with n ="+str(n_Plot),N_plot_points=750,make_plot=True)
#
#
# =============Build Hamiltonian at t=0==============================================
# Pass name of potential function explicitly here
time = 0.0
H_mat = fem_dvr.Hamiltonian_Two_States(pertubation.V_morse_1,
pertubation.V_morse_2, Vcoupling,
time)
print("\n Completed construction of Hamiltonian at t = 0")
#
# ================# Initialize wave function at t = 0 ================================
#
# Bound state of one of the two uncoupled Hamiltonians chosen as initial state
#
Cinitial = np.zeros(2 * fem_dvr.nbas, dtype=np.complex)
for i in range(0, 2 * fem_dvr.nbas):
Cinitial[i] = (
wfcnPlot[i] + 0.0 * 1j
) # this loop was necessary to make Cinitial an array of complex numbers
wfcnInitialPlot1 = np.zeros(fem_dvr.nbas, dtype=np.complex)
wfcnInitialPlot2 = np.zeros(fem_dvr.nbas, dtype=np.complex)
norm_1 = 0.0
norm_2 = 0.0
for j in range(0, fem_dvr.nbas):
wfcnInitialPlot1[j] = Cinitial[j]
wfcnInitialPlot2[j] = Cinitial[j + fem_dvr.nbas]
norm_1 = norm_1 + np.abs(wfcnInitialPlot1[j])**2
norm_2 = norm_2 + np.abs(wfcnInitialPlot2[j])**2
#
# plot initial wave packet
#
tinitial = 0.0
x_Plot_time_array = []
Psi1_plot_time_array = []
Psi2_plot_time_array = []
if want_to_plot_animate is True:
print("\n Plot initial wave function at t = ", tinitial)
number_string = str(0.0)
title = "Wavefunction at t = " + number_string
(
x_Plot_array,
Psi1_plot_array,
Psi2_plot_array,
) = fem_dvr.Plot_Psi_Two_States(
wfcnInitialPlot1,
wfcnInitialPlot2,
plot_title_string=title,
N_plot_points=750,
make_plot=True,
)
x_Plot_time_array.append(x_Plot_array)
Psi1_plot_time_array.append(Psi1_plot_array)
Psi2_plot_time_array.append(Psi2_plot_array)
#
# initialize arrays for storing the information for all the frames in the animation of
# the propagation of the wave function
times_array = []
norm_1_array = []
norm_2_array = []
times_array.append(tinitial)
norm_1_array.append(norm_1)
norm_2_array.append(norm_2)
#
# =====================Propagation from t = tinitial to tfinal ============================
# specify final time and specify number of intervals for plotting
#
t_final_propagation = 15.0 # femtoseconds
tfinal = t_final_propagation / atu_to_fs
print(
"\nOverall propagation will be from ",
tinitial,
" to ",
tfinal,
" = ",
tfinal * 24.189 / 1000.0,
"fs in ",
number_of_time_intervals,
" intervals",
)
#
# Loop over n_intervals intervals that make up t = 0 to tfinal
#
t_interval = (tfinal - tinitial) / number_of_time_intervals
time_step = (
0.1 # atomic time units -- denser grids require smaller time steps
)
for i_time_interval in range(0, number_of_time_intervals):
t_start = tinitial + i_time_interval * t_interval
t_finish = tinitial + (i_time_interval + 1) * t_interval
N_time_steps = np.int(t_interval / time_step)
Delta_t = t_interval / np.float(N_time_steps)
#
# Check norm of initial wavefunction
#
norm_initial = 0.0
norm_1 = 0.0
for j in range(0, 2 * fem_dvr.nbas):
norm_initial = norm_initial + np.abs(Cinitial[j])**2
if j < fem_dvr.nbas:
norm_1 = norm_1 + np.abs(Cinitial[j])**2
norm_2 = norm_initial - norm_1
print(
"\nNorm of starting wave function",
norm_initial,
" norm on state 1 ",
norm_1,
" norm on state 2 ",
norm_2,
)
#
# use the propagator from DVR()
#
clock_start = timeclock.time()
if t_start * atu_to_fs < T_pulse:
Ctfinal = fem_dvr.Crank_Nicolson_Two_States(
t_start,
t_finish,
N_time_steps,
Cinitial,
pertubation.V_morse_1,
pertubation.V_morse_2,
Vcoupling,
Is_H_time_dependent=True,
)
else:
Ctfinal = fem_dvr.Crank_Nicolson_Two_States(
t_start,
t_finish,
N_time_steps,
Cinitial,
pertubation.V_morse_1,
pertubation.V_morse_2,
Vzero,
Is_H_time_dependent=False,
)
clock_finish = timeclock.time()
print(
"Time for ",
N_time_steps,
" Crank-Nicolson steps was ",
clock_finish - clock_start,
" secs.",
)
#
# Check norms of initial and final wavefunctions
#
norm_final = 0
norm_1 = 0.0
for j in range(0, 2 * fem_dvr.nbas):
norm_final = norm_final + np.abs(Ctfinal[j])**2
if j < fem_dvr.nbas:
norm_1 = norm_1 + np.abs(Ctfinal[j])**2
norm_2 = norm_final - norm_1
print(
"Norm of final wave function",
norm_final,
" norm on state 1 ",
norm_1,
" norm on state 2 ",
norm_2,
)
#
# Plot of packet at end of each interval using Plot_Psi from DVR()
# "False" returns the values but doesn't make a plot, so the values
# can be used in the animation.
#
print(
"Plot function propagated to t = ",
t_finish,
" atomic time units = ",
t_finish * 24.189 / 1000.0,
" fs",
)
number_string = str(t_finish)
title = "Wavefunction at t = " + number_string
wfcn_at_t_Plot1 = np.zeros(fem_dvr.nbas, dtype=np.complex)
wfcn_at_t_Plot2 = np.zeros(fem_dvr.nbas, dtype=np.complex)
for j in range(0, fem_dvr.nbas):
wfcn_at_t_Plot1[j] = Ctfinal[j]
wfcn_at_t_Plot2[j] = Ctfinal[j + fem_dvr.nbas]
# plot using the Plot_Psi_Two_States function in DVRHelper, returns two
# components of wave function separately
if want_to_plot_animate is True:
(
x_Plot_array,
Psi1_plot_array,
Psi2_plot_array,
) = fem_dvr.Plot_Psi_Two_States(
wfcn_at_t_Plot1,
wfcn_at_t_Plot2,
plot_title_string=title,
N_plot_points=750,
make_plot=False,
)
times_array.append(t_finish)
norm_1_array.append(norm_1)
norm_2_array.append(norm_2)
x_Plot_time_array.append(x_Plot_array)
Psi1_plot_time_array.append(Psi1_plot_array)
Psi2_plot_time_array.append(Psi2_plot_array)
#
# reset initial packet as final packet for this interval
for i in range(0, 2 * fem_dvr.nbas):
Cinitial[i] = Ctfinal[i]
#
if want_to_plot_animate is True:
(
x_Plot_array,
Psi1_plot_array,
Psi2_plot_array,
) = fem_dvr.Plot_Psi_Two_States(
wfcn_at_t_Plot1,
wfcn_at_t_Plot2,
plot_title_string=title,
N_plot_points=750,
make_plot=True,
)
print(
"Final frame at tfinal = ",
tfinal,
"atomic time units is showing -- ready to prepare animation ",
)
print("\n **** close the plot window to proceed ****")
#
# print a file with the norms of the components of the wave function as a function of time
#
norms_file = open("./Plot_Output/Wavepacket_norms.dat", "w")
for k in range(0, number_of_time_intervals):
norms_file.write(
"%12.8f %12.8f %12.8f\n" %
(times_array[k], norm_1_array[k], norm_2_array[k]))
norms_file.close()
# ==============================================================================
# # initialization function: plot the background of each frame
# ==============================================================================
def init():
line.set_data([], [])
time_text.set_text("")
ax.set_xlabel(" x (bohr) ", fontsize=16, fontweight="bold")
ax.set_ylabel(" |Psi_1(t)| and |Psi_2(t)| ",
fontsize=16,
fontweight="bold")
fig.suptitle("Both Components of Wave Packet",
fontsize=16,
fontweight="bold")
ax.plot([x_Plot[0], x_Plot[len(x_Plot) - 1]], [0, 0],
"k") # put in a line at the value Phi = 0
return line, time_text
nframes = number_of_time_intervals
def animate(i):
# ==============================================================================
#
# commands to specify changing elements of the animation
# We precomputed all the lines at all the times so that this
# wouldn't recompute everything at each frame, which is very slow...
#
# ==============================================================================
time_string = str(times_array[i] * 24.189 /
1000.0) # string for a plot label
# re_array = np.real(Psi1_plot_time_array[i])
# im_array = np.imag(Psi1_plot_time_array[i])
abs_array1 = np.abs(Psi1_plot_time_array[i])
abs_array2 = np.abs(Psi2_plot_time_array[i])
line1.set_data(x_Plot_array, abs_array1)
line2.set_data(x_Plot_array, abs_array2)
# line3.set_data(x_Plot_array,abs_array)
time_text.set_text("time = " + time_string + " fs")
# return (line1,line2,line3,time_text)
return (line1, line2, time_text)
if want_to_plot_animate is True:
# ==============================================================================
# Now plot the animation
# ==============================================================================
# reinitialize the figure for the next plot which is the animation
fig = plt.figure()
ymax = 2.0
xmin = x_Plot[0]
xmax = x_Plot[len(x_Plot) - 1]
ax = fig.add_subplot(111,
autoscale_on=False,
xlim=(xmin, xmax),
ylim=(0, +ymax))
(line, ) = ax.plot([], [], "-r", lw=2)
(line1, ) = ax.plot([], [], "-r", lw=2)
(line2, ) = ax.plot([], [], "-b", lw=2)
# line3, = ax.plot([], [], 'k',lw=2)
# define the object that will be the time printed on each frame
time_text = ax.text(0.02, 0.95, "", transform=ax.transAxes)
# ==============================================================================
# call the animator. blit=True means only re-draw the parts that have changed.
# ==============================================================================
anim = animation.FuncAnimation(
fig,
animate,
init_func=init,
frames=nframes + 1,
interval=200,
blit=True,
repeat=True,
)
# ==============================================================================
# show the animation
# ==============================================================================
anim.save("Non_adiabatic_coupled_Morse_pots.mp4")
plt.show()
print("done")
if want_to_plot_animate is False:
print(
"\n\nSet the command line option want_to_plot_animate=True to see figures, "
"animation, and create plotting directory.\n\n")
def main(want_to_plot):
#
# ================ Make Directory for Plots if it's not there already =============
#
# detect the current working directory and print it
path = os.getcwd()
print("The current working directory is %s" % path)
# define the name of the directory to be created
Plot_Output = path + "/Plot_Output"
if want_to_plot == True:
if os.path.exists(Plot_Output):
print(
"Directory for wave function plots already exists", Plot_Output
)
else:
print(
"Attempting to create directory for wave function plots ",
Plot_Output,
)
try:
os.mkdir(Plot_Output)
except OSError:
print("Creation of the directory %s failed" % Plot_Output)
else:
print("Successfully created the directory %s " % Plot_Output)
# =====================================FEM_DVR===================================
# Set up the FEM DVR grid given only the Finite Element boundaries and order
# of Gauss Lobatto quadrature, and compute the Kinetic Energy matrix for
# this FEM DVR for Mass = mass set in call (atomic units).
#
# Here is where the reduced mass is set
# H_Mass = 1.007825032 #H atom atomic mass
# H_Mass = 1.00727647 #proton atomic mass
# standard atomic weight natural abundance, seems to be what Turner, McCurdy used
H_Mass = 1.0078
Daltons_to_eMass = 1822.89
mu = (H_Mass / 2.0) * Daltons_to_eMass
print("reduced mass mu = ", mu)
bohr_to_Angstrom = 0.529177
Hartree_to_eV = 27.211386245988 # NIST ref
eV_to_wavennumber = 8065.54393734921 # NIST ref on constants + conversions
# value from NIST ref on constants + conversions
Hartree_to_wavenumber = 2.1947463136320e5
HartreeToKelvin = 315773
atu_to_fs = 24.18884326509 / 1000
# Set up the FEM-DVR grid
n_order = 25
FEM_boundaries = [
0.4,
1.0,
2.0,
3.0,
4.0,
5.0,
6.0,
7.0,
8.0,
10.0,
15.0,
20.0,
25.0,
30.0,
50,
]
# parameters for Fig 2 in Turner, McCurdy paper
# Julia Turner and C. William McCurdy, Chemical Physics 71(1982) 127-133
scale_factor = np.exp(1j * 37.0 * np.pi / 180.0)
R0 = 22.75
fem_dvr = FEM_DVR(
n_order,
FEM_boundaries,
Mass=mu,
Complex_scale=scale_factor,
R0_scale=R0,
)
print("\nFEM-DVR basis of ", fem_dvr.nbas, " functions")
#
pertubation = Potential()
# ==================================================================================
# Plot potential on the DVR grid points on which the wavefunction is defined
if want_to_plot is True:
print("\n Plot potential ")
print("Test V", pertubation.V_Bernstein(5.0 + 0.0 * 1j, 0.0))
print("Test V", pertubation.V_Bernstein(10.0 + 0.5 * 1j, 0.0))
time = 0.0
x_Plot = []
pot_Plot = []
for j in range(0, fem_dvr.nbas):
x_Plot.append(np.real(fem_dvr.x_pts[j + 1]))
pot_Plot.append(
np.real(pertubation.V_Bernstein(fem_dvr.x_pts[j + 1], time))
)
if want_to_plot is True:
plt.suptitle(
"V(x) at DVR basis function nodes", fontsize=14, fontweight="bold"
)
string = "V"
plt.plot(x_Plot, pot_Plot, "ro", label=string)
plt.plot(x_Plot, pot_Plot, "-b")
plt.legend(loc="best")
plt.xlabel(" x ", fontsize=14)
plt.ylabel("V", fontsize=14)
print(
"\n Running from terminal, close figure window to proceed and make .pdf file of figure"
)
# Insert limits if necessary
# Generally comment this logic. Here I am matching the Turner McCurdy Figure 2
# CWM: need to use float() to get plt.xlim to work to set x limits
ymax = float(0.05)
plt.ylim([-0.18, ymax])
# save plot to .pdf file
plt.savefig(
"Plot_Output/" + "Plot_potential" + ".pdf", transparent=False
)
plt.show()
#
# =============Build Hamiltonian (at t=0 if time-dependent)=================================
# Pass name of potential function explicitly here
time = 0.0
H_mat = fem_dvr.Hamiltonian(pertubation.vectorized_V_Bernstein, time)
print("\n Completed construction of Hamiltonian at t = 0")
# ====================================================================================
#
# Find all the eigenvalues of the Hamiltonian so we can compare with known bound state energies
# or make a plot of the spectrum -- For a time-independent Hamiltonian example here
#
# EigenVals = LA.eigvals(H_mat) # eigenvalues only for general matrix. LA.eigvalsh for Hermitian
print(
"Calculating ",
fem_dvr.nbas,
" eigenvalues and eigenvectors for plotting eigenfunctions",
)
EigenVals, EigenVecs = LA.eig(H_mat, right=True, homogeneous_eigvals=True)
print("after LA.eig()")
#
n_energy = fem_dvr.nbas
file_opened = open("Spectrum_ECS.dat", "w")
print("EigenVals shape ", EigenVals.shape)
for i in range(0, n_energy):
print("E( ", i, ") = ", EigenVals[0, i], " hartrees")
print(
np.real(EigenVals[0, i]),
" ",
np.imag(EigenVals[0, i]),
file=file_opened,
)
# ====================================================================================
#
# Extract the n_Plot'th eigenfunction for plotting
#
# pick one of the bound states of Morse Potential to plot
# numbering can depend on numpy and python installation that determines
# behavior of the linear algebra routines.
n_Plot = (
n_energy - 1
) # This is generally the highest energy continuum eigenvalue
n_Plot = 292
wfcnPlot = []
for j in range(0, fem_dvr.nbas):
wfcnPlot.append(EigenVecs[j, n_Plot])
#
# Normalize wave function from diagonalization
# using integration of square on the contour
# following the original Rescigno McCurdy idea for partial widths
# from the paper
# "Normalization of resonance wave functions and the calculation of resonance widths"
# Rescigno, McCurdy, Phys Rev A 34,1882 (1986)
norm_squared = 0.0
for j in range(0, fem_dvr.nbas):
norm_squared = norm_squared + (wfcnPlot[j]) ** 2
wfcnPlot = wfcnPlot / np.sqrt(norm_squared)
norm_squared = 0.0
gamma_residue = 0.0
# background momentum defined with Re(Eres)
k_momentum = np.sqrt(2.0 * mu * np.real(EigenVals[0, n_Plot]))
# k_momentum = np.real(np.sqrt(2.0*mu*EigenVals[0,n_Plot]))
for j in range(0, fem_dvr.nbas):
norm_squared = norm_squared + (wfcnPlot[j]) ** 2
if fem_dvr.x_pts[j + 1] < 5.8:
free_wave = (2.0 * np.sqrt(mu / k_momentum)) * np.sin(
k_momentum * fem_dvr.x_pts[j + 1]
)
gamma_residue = gamma_residue + wfcnPlot[
j
] * pertubation.V_Bernstein(
fem_dvr.x_pts[j + 1], time
) * free_wave * np.sqrt(
fem_dvr.w_pts[j + 1]
)
print("Complex symmetric inner product (psi|psi) is being used")
if want_to_plot is True:
print(
"Norm of wave function from int psi^2 on contour being plotted is ",
np.sqrt(norm_squared),
)
print(" For this state the asymptotic value of k = ", k_momentum)
print(
"gamma from int = ",
gamma_residue,
" |gamma|^2 = ",
np.abs(gamma_residue) ** 2,
)
# Plot wave function -- It must be type np.complex
Cinitial = np.zeros((fem_dvr.nbas), dtype=np.complex)
wfcnInitialPlot = np.zeros((fem_dvr.nbas), dtype=np.complex)
for j in range(0, fem_dvr.nbas):
Cinitial[j] = wfcnPlot[j]
if want_to_plot is True:
#
# plot n_Plot'th eigenfunction
#
print(
"\n Plot Hamiltonian eigenfunction number ",
n_Plot,
" with energy ",
EigenVals[0, n_Plot],
)
tau = atu_to_fs / (-2.0 * np.imag(EigenVals[0, n_Plot]))
print(" Lifetime tau = 1/Gamma = ", tau, " fs")
number_string = str(n_Plot)
title = "Wavefunction number = " + number_string
wfn_plot_points = 2000
x_Plot_array, Psi_plot_array = fem_dvr.Plot_Psi(
Cinitial,
plot_title_string=title,
N_plot_points=wfn_plot_points,
make_plot=True,
)
#
# Make data file for n_Plot'th eigenfunction
# Psi and the integrand of the residue factors, gamma, of the free-free matrix element
# of the full Green's function, as per
#
filename = "wavefunction" + number_string + ".dat"
file_opened = open(filename, "w")
print_points = len(x_Plot_array)
print("x_Plot_array shape ", print_points)
for i in range(print_points):
free_wave = (2.0 * np.sqrt(mu / k_momentum)) * np.sin(
k_momentum * x_Plot_array[i]
)
# for partial width gamma
integrand = (
Psi_plot_array[i]
* pertubation.V_Bernstein(x_Plot_array[i], time)
* free_wave
)
print(
np.real(x_Plot_array[i]),
" ",
np.imag(x_Plot_array[i]),
" ",
np.real(Psi_plot_array[i]),
" ",
np.imag(Psi_plot_array[i]),
" ",
np.real(integrand),
np.imag(integrand),
file=file_opened,
)
#
# exit()
# ====================================================================================
#
# Extract the n_Plot'th eigenfunction for plotting
#
# pick one of the eigenstates of Potential to plot
# numbering can depend on numpy and python installation that determines
# behavior of the linear algebra routines.
n_Plot = 292
if want_to_plot is True:
print(
"Calculating ",
fem_dvr.nbas,
" eigenvectors for plotting eigenfunctions",
)
EigenVals2, EigenVecs = LA.eig(H_mat, right=True, homogeneous_eigvals=True)
wfcnPlot = []
for j in range(0, fem_dvr.nbas):
wfcnPlot.append(EigenVecs[j, n_Plot])
#
# normalize wave function from diagonalization
# using integration of square on the contour
# following the old Rescigno McCurdy idea for partial widths
norm_squared = 0.0
for j in range(0, fem_dvr.nbas):
norm_squared = norm_squared + (wfcnPlot[j]) ** 2
wfcnPlot = wfcnPlot / np.sqrt(norm_squared)
norm_squared = 0.0
for j in range(0, fem_dvr.nbas):
norm_squared = norm_squared + (wfcnPlot[j]) ** 2
if want_to_plot is True:
print(
"Norm of wave function from int psi^2 on contour being plotted is ",
np.sqrt(norm_squared),
)
# Plot wave function -- It must be type np.complex
Cinitial = np.zeros((fem_dvr.nbas), dtype=np.complex)
wfcnInitialPlot = np.zeros((fem_dvr.nbas), dtype=np.complex)
for j in range(0, fem_dvr.nbas):
Cinitial[j] = wfcnPlot[j]
if want_to_plot is True:
#
# plot n_Plot'th eigenfunction
#
print(
"\n Plot Hamiltonian eigenfunction number ",
n_Plot,
" with energy ",
EigenVals[0, n_Plot],
)
number_string = str(n_Plot)
title = "Wavefunction number = " + number_string
wfn_plot_points = 1000
x_Plot_array, Psi_plot_array = fem_dvr.Plot_Psi(
Cinitial,
plot_title_string=title,
N_plot_points=wfn_plot_points,
make_plot=True,
)
#
filename = "wavefunction" + number_string + ".dat"
file_opened = open(filename, "w")
print_points = len(x_Plot_array)
# print("x_Plot_array shape ",print_points)
for i in range(print_points):
print(
np.real(x_Plot_array[i]),
" ",
np.imag(x_Plot_array[i]),
" ",
np.real(Psi_plot_array[i]),
" ",
np.imag(Psi_plot_array[i]),
file=file_opened,
)
if want_to_plot is False:
print(
"\n\n Set the command line option want_to_plot=True to see figures and create plotting directory.\n\n"
)