Code Solutions
This section provides the solutions to all of the coding exercises provided in the text. The solutions are written in Python and use the NumPy library for numerical operations. The code snippets are self-contained and can be run in any Python environment. The solutions are organized by the exercise they correspond to and are presented in the same order as in the text. For convenience, the full code solution files resmet.py, neqdet.py and shcdet.py can be saved by clicking the links.
Restricted Electronic Structure Methods
#!/usr/bin/env python
import argparse as ap, itertools as it, numpy as np
np.set_printoptions(edgeitems=30, linewidth=100000, formatter=dict(float=lambda x: "%20.14f" % x))
ATOM = {
"H" : 1, "He": 2,
"Li": 3, "Be": 4, "B" : 5, "C" : 6, "N" : 7, "O" : 8, "F" : 9, "Ne": 10,
"Na": 11, "Mg": 12, "Al": 13, "Si": 14, "P" : 15, "S" : 16, "Cl": 17, "Ar": 18,
"K" : 19, "Ca": 20, "Ga": 31, "Ge": 32, "As": 33, "Se": 34, "Br": 35, "Kr": 36,
"Rb": 37, "Sr": 38, "In": 49, "Sn": 50, "Sb": 51, "Te": 52, "I" : 53, "Xe": 54,
}
# SECTION FOR PARSING COMMAND LINE ARGUMENTS =====================================================================================
# create the parser
parser = ap.ArgumentParser(
prog="resmet.py", description="Restricted Electronic Structure Methods Educational Toolkit",
formatter_class=lambda prog: ap.HelpFormatter(prog, max_help_position=128),
add_help=False, allow_abbrev=False
)
# add the arguments
parser.add_argument("-m", "--molecule", help="Molecule file in the .xyz format. (default: %(default)s)", default="molecule.xyz")
parser.add_argument("-t", "--threshold", help="Convergence threshold. (default: %(default)s)", type=float, default=1e-12)
parser.add_argument("-h", "--help", help="This help message.", action=ap.BooleanOptionalAction)
# add integral options
parser.add_argument("--int", help="Integral files. (default: %(default)s)", default=["T_AO.mat", "V_AO.mat", "S_AO.mat", "J_AO.mat"])
# optional flags
parser.add_argument("--diis", help="Flag ofor the DIIS convergence accelerator.", action=ap.BooleanOptionalAction)
# method switches
parser.add_argument("--fci", help="Perform the full configuration interaction calculation.", action=ap.BooleanOptionalAction)
parser.add_argument("--ccd", help="Perform the doubles coupled clusters calculation.", action=ap.BooleanOptionalAction)
parser.add_argument("--ccsd", help="Perform the singles/doubles coupled clusters calculation.", action=ap.BooleanOptionalAction)
parser.add_argument("--mp2", help="Perform the second order Moller-Plesset calculation.", action=ap.BooleanOptionalAction)
parser.add_argument("--mp3", help="Perform the third order Moller-Plesset calculation.", action=ap.BooleanOptionalAction)
# parse the arguments
args = parser.parse_args()
# print the help message if the flag is set
if args.help: parser.print_help(); exit()
# forward declaration of integrals and energies in MS basis (just to avoid NameError in linting)
Hms, Fms, Jms, Jmsa, Emss, Emsd = np.zeros([0]), np.zeros([0]), np.zeros([0]), np.zeros([0]), np.zeros([0]), np.zeros([0])
# OBTAIN THE MOLECULE AND ATOMIC INTEGRALS =======================================================================================
# get the atomic numbers and coordinates of all atoms
atoms = np.array([ATOM[line.split()[0]] for line in open(args.molecule).readlines()[2:]], dtype=int )
coords = np.array([line.split()[1:] for line in open(args.molecule).readlines()[2:]], dtype=float)
# convert coordinates to bohrs and forward declare orbital slices and atom count
coords *= 1.8897261254578281; o, v = slice(0, 0), slice(0, 0); natoms = len(atoms)
# load the one-electron integrals from the files
H, S = np.loadtxt(args.int[0], skiprows=1) + np.loadtxt(args.int[1], skiprows=1), np.loadtxt(args.int[2], skiprows=1);
# load the two-electron integrals from files
J = np.loadtxt(args.int[3], skiprows=1).reshape(4 * [S.shape[1]])
# HARTREE-FOCK METHOD ============================================================================================================
# energies, iteration counter, number of basis functions and number of occupied/virtual orbitals
E_HF, E_HF_P, VNN, iter, nbf, nocc = 0, 1, 0, 0, S.shape[0], sum(atoms) // 2; nvirt = nbf - nocc
# exchange integrals and the guess density matrix
K, D = J.transpose(0, 3, 2, 1), np.zeros((nbf, nbf))
# Fock matrix, coefficient matrix and orbital energies initialized to zero
F, C, eps = np.zeros((nbf, nbf)), np.zeros((nbf, nbf)), np.zeros((nbf))
# DIIS containers
DIIS_F, DIIS_E = [], []
# the X matrix which is the inverse of the square root of the overlap matrix
SEP = np.linalg.eigh(S); X = SEP[1] @ np.diag(1 / np.sqrt(SEP[0])) @ SEP[1].T
# the scf loop
while abs(E_HF - E_HF_P) > args.threshold:
# build the Fock matrix and increment the iteration counter
F = H + np.einsum("ijkl,ij->kl", J - 0.5 * K, D, optimize=True); iter += 1
# DIIS extrapolation
if args.diis and iter > 1:
# append the DIIS matrices
DIIS_F.append(F); DIIS_E.append(S @ D @ F - F @ D @ S);
# truncate the DIIS subspace
if len(DIIS_F) > 5: DIIS_F.pop(0), DIIS_E.pop(0)
# build the DIIS system
A = np.ones ((len(DIIS_F) + 1, len(DIIS_F) + 1)); A[-1, -1] = 0
b = np.zeros((len(DIIS_F) + 1 )); b[-1] = 1
# fill the DIIS matrix
for i, j in it.product(range(len(DIIS_F)), range(len(DIIS_F))):
A[i, j] = A[j, i] = np.einsum("ij,ij", DIIS_E[i], DIIS_E[j])
# solve the DIIS equations and extrapolate the Fock matrix
c = np.linalg.solve(A, b); F = np.einsum("i,ijk->jk", c[:-1], DIIS_F)
# solve the Fock equations
eps, C = np.linalg.eigh(X @ F @ X); C = X @ C
# build the density from coefficients
D = 2 * np.einsum("ij,kj->ik", C[:, :nocc], C[:, :nocc])
# save the previous energy and calculate the current electron energy
E_HF_P, E_HF = E_HF, 0.5 * np.einsum("ij,ij", D, H + F, optimize=True)
# calculate nuclear-nuclear repulsion
for i, j in ((i, j) for i, j in it.product(range(natoms), range(natoms)) if i != j):
VNN += 0.5 * atoms[i] * atoms[j] / np.linalg.norm(coords[i, :] - coords[j, :])
# print the results
print(" RHF ENERGY: {:.8f} ({} ITERATIONS)".format(E_HF + VNN, iter))
# INTEGRAL TRANSFORMS FOR POST-HARTREE-FOCK METHODS ==============================================================================
if args.mp2 or args.mp3 or args.ccd or args.ccsd or args.fci:
# define the occ and virt spinorbital slices shorthand
o, v = slice(0, 2 * nocc), slice(2 * nocc, 2 * nbf)
# define the tiling matrix for the Jmsa coefficients and energy placeholders
P = np.array([np.eye(nbf)[:, i // 2] for i in range(2 * nbf)]).T
# define the spin masks
M = np.repeat([1 - np.arange(2 * nbf, dtype=int) % 2], nbf, axis=0)
N = np.repeat([ np.arange(2 * nbf, dtype=int) % 2], nbf, axis=0)
# tile the coefficient matrix, apply the spin mask and tile the orbital energies
Cms, epsms = np.block([[C @ P], [C @ P]]) * np.block([[M], [N]]), np.repeat(eps, 2)
# transform the core Hamiltonian and Fock matrix to the molecular spinorbital basis
Hms = np.einsum("ip,ij,jq->pq", Cms, np.kron(np.eye(2), H), Cms, optimize=True)
Fms = np.einsum("ip,ij,jq->pq", Cms, np.kron(np.eye(2), F), Cms, optimize=True)
# transform the coulomb integrals to the MS basis in chemists' notation
Jms = np.einsum("ip,jq,ijkl,kr,ls->pqrs",
Cms, Cms, np.kron(np.eye(2), np.kron(np.eye(2), J).T), Cms, Cms, optimize=True
);
# antisymmetrized two-electron integrals in physicists' notation
Jmsa = (Jms - Jms.swapaxes(1, 3)).transpose(0, 2, 1, 3)
# tensor epsilon_i^a
Emss = epsms[o] - epsms[v, None]
# tensor epsilon_ij^ab
Emsd = epsms[o] + epsms[o, None] - epsms[v, None, None] - epsms[v, None, None, None]
# MOLLER-PLESSET PERTURBATION THEORY =============================================================================================
if args.mp2 or args.mp3:
# energy containers
E_MP2, E_MP3 = 0, 0
# calculate the MP2 correlation energy
if args.mp2 or args.mp3:
E_MP2 += 0.25 * np.einsum("abij,ijab,abij",
Jmsa[v, v, o, o], Jmsa[o, o, v, v], Emsd**-1, optimize=True
)
print(" MP2 ENERGY: {:.8f}".format(E_HF + E_MP2 + VNN))
# calculate the MP3 correlation energy
if args.mp3:
E_MP3 += 0.125 * np.einsum("abij,cdab,ijcd,abij,cdij",
Jmsa[v, v, o, o], Jmsa[v, v, v, v], Jmsa[o, o, v, v], Emsd**-1, Emsd**-1, optimize=True
)
E_MP3 += 0.125 * np.einsum("abij,ijkl,klab,abij,abkl",
Jmsa[v, v, o, o], Jmsa[o, o, o, o], Jmsa[o, o, v, v], Emsd**-1, Emsd**-1, optimize=True
)
E_MP3 += 1.000 * np.einsum("abij,cjkb,ikac,abij,acik",
Jmsa[v, v, o, o], Jmsa[v, o, o, v], Jmsa[o, o, v, v], Emsd**-1, Emsd**-1, optimize=True
)
print(" MP3 ENERGY: {:.8f}".format(E_HF + E_MP2 + E_MP3 + VNN))
# COUPLED CLUSTER METHOD =========================================================================================================
if args.ccd or args.ccsd:
# energy containers for all the CC methods
E_CCD, E_CCD_P, E_CCSD, E_CCSD_P = 0, 1, 0, 1
# initialize the first guess for the t-amplitudes as zeros
t1, t2 = np.zeros((2 * nvirt, 2 * nocc)), np.zeros(2 * [2 * nvirt] + 2 * [2 * nocc])
# CCD loop
if args.ccd:
while abs(E_CCD - E_CCD_P) > args.threshold:
# collect all the distinct LCCD terms
lccd1 = 0.5 * np.einsum("abcd,cdij->abij", Jmsa[v, v, v, v], t2, optimize=True)
lccd2 = 0.5 * np.einsum("klij,abkl->abij", Jmsa[o, o, o, o], t2, optimize=True)
lccd3 = 1.0 * np.einsum("akic,bcjk->abij", Jmsa[v, o, o, v], t2, optimize=True)
# apply the permuation operator and add it to the corresponding LCCD terms
lccd3 += lccd3.transpose(1, 0, 3, 2) - lccd3.swapaxes(0, 1) - lccd3.swapaxes(2, 3)
# collect all the remaining CCD terms
ccd1 = -0.50 * np.einsum("klcd,acij,bdkl->abij", Jmsa[o, o, v, v], t2, t2, optimize=True)
ccd2 = -0.50 * np.einsum("klcd,abik,cdjl->abij", Jmsa[o, o, v, v], t2, t2, optimize=True)
ccd3 = +0.25 * np.einsum("klcd,cdij,abkl->abij", Jmsa[o, o, v, v], t2, t2, optimize=True)
ccd4 = +1.00 * np.einsum("klcd,acik,bdjl->abij", Jmsa[o, o, v, v], t2, t2, optimize=True)
# permutation operators
ccd1 -= ccd1.swapaxes(0, 1)
ccd2 -= ccd2.swapaxes(2, 3)
ccd4 -= ccd4.swapaxes(2, 3)
# update the t-amplitudes
t2 = (Jmsa[v, v, o, o] + lccd1 + lccd2 + lccd3 + ccd1 + ccd2 + ccd3 + ccd4) / Emsd
# evaluate the energy
E_CCD_P, E_CCD = E_CCD, 0.25 * np.einsum("ijab,abij", Jmsa[o, o, v, v], t2)
# print the CCD energy
print(" CCD ENERGY: {:.8f}".format(E_HF + E_CCD + VNN))
# CCSD loop
if args.ccsd:
while abs(E_CCSD - E_CCSD_P) > args.threshold:
# define taus
tau, ttau = t2.copy(), t2.copy()
# add contributions to the tilde tau
ttau += 0.5 * np.einsum("ai,bj->abij", t1, t1, optimize=True).swapaxes(0, 0)
ttau -= 0.5 * np.einsum("ai,bj->abij", t1, t1, optimize=True).swapaxes(2, 3)
# add the contributions to tau
tau += np.einsum("ai,bj->abij", t1, t1, optimize=True).swapaxes(0, 0)
tau -= np.einsum("ai,bj->abij", t1, t1, optimize=True).swapaxes(2, 3)
# define the deltas for Fae and Fmi
dae, dmi = np.eye(2 * nvirt), np.eye(2 * nocc)
# define Fae, Fmi and Fme
Fae, Fmi, Fme = (1 - dae) * Fms[v, v], (1 - dmi) * Fms[o, o], Fms[o, v].copy()
# add the contributions to Fae
Fae -= 0.5 * np.einsum("me,am->ae", Fms[o, v], t1, optimize=True)
Fae += 1.0 * np.einsum("mafe,fm->ae", Jmsa[o, v, v, v], t1, optimize=True)
Fae -= 0.5 * np.einsum("mnef,afmn->ae", Jmsa[o, o, v, v], ttau, optimize=True)
# add the contributions to Fmi
Fmi += 0.5 * np.einsum("me,ei->mi", Fms[o, v], t1, optimize=True)
Fmi += 1.0 * np.einsum("mnie,en->mi", Jmsa[o, o, o, v], t1, optimize=True)
Fmi += 0.5 * np.einsum("mnef,efin->mi", Jmsa[o, o, v, v], ttau, optimize=True)
# add the contributions to Fme
Fme += np.einsum("mnef,fn->me", Jmsa[o, o, v, v], t1, optimize=True)
# define Wmnij, Wabef and Wmbej
Wmnij = Jmsa[o, o, o, o].copy()
Wabef = Jmsa[v, v, v, v].copy()
Wmbej = Jmsa[o, v, v, o].copy()
# define some complementary variables used in the Wmbej intermediate
t12 = 0.5 * t2 + np.einsum("fj,bn->fbjn", t1, t1, optimize=True)
# add contributions to Wmnij
Wmnij += 0.25 * np.einsum("efij,mnef->mnij", tau, Jmsa[o, o, v, v], optimize=True)
Wabef += 0.25 * np.einsum("abmn,mnef->abef", tau, Jmsa[o, o, v, v], optimize=True)
Wmbej += 1.00 * np.einsum("fj,mbef->mbej", t1, Jmsa[o, v, v, v], optimize=True)
Wmbej -= 1.00 * np.einsum("bn,mnej->mbej", t1, Jmsa[o, o, v, o], optimize=True)
Wmbej -= 1.00 * np.einsum("fbjn,mnef->mbej", t12, Jmsa[o, o, v, v], optimize=True)
# define the permutation arguments for Wmnij and Wabef and add them
PWmnij = np.einsum("ej,mnie->mnij", t1, Jmsa[o, o, o, v], optimize=True)
PWabef = np.einsum("bm,amef->abef", t1, Jmsa[v, o, v, v], optimize=True)
# add the permutations to Wmnij and Wabef
Wmnij += PWmnij - PWmnij.swapaxes(2, 3)
Wabef += PWabef.swapaxes(0, 1) - PWabef
# define the right hand side of the T1 and T2 amplitude equations
rhs_T1, rhs_T2 = Fms[v, o].copy(), Jmsa[v, v, o, o].copy()
# calculate the right hand side of the CCSD equation for T1
rhs_T1 += 1.0 * np.einsum("ei,ae->ai", t1, Fae, optimize=True)
rhs_T1 -= 1.0 * np.einsum("am,mi->ai", t1, Fmi, optimize=True)
rhs_T1 += 1.0 * np.einsum("aeim,me->ai", t2, Fme, optimize=True)
rhs_T1 -= 1.0 * np.einsum("fn,naif->ai", t1, Jmsa[o, v, o, v], optimize=True)
rhs_T1 -= 0.5 * np.einsum("efim,maef->ai", t2, Jmsa[o, v, v, v], optimize=True)
rhs_T1 -= 0.5 * np.einsum("aemn,nmei->ai", t2, Jmsa[o, o, v, o], optimize=True)
# contracted F matrices that used in the T2 equations
Faet = np.einsum("bm,me->be", t1, Fme, optimize=True)
Fmet = np.einsum("ej,me->mj", t1, Fme, optimize=True)
# define the permutation arguments for all terms in the equation for T2
P1 = np.einsum("aeij,be->abij", t2, Fae - 0.5 * Faet, optimize=True)
P2 = np.einsum("abim,mj->abij", t2, Fmi + 0.5 * Fmet, optimize=True)
P3 = np.einsum("aeim,mbej->abij", t2, Wmbej, optimize=True)
P3 -= np.einsum("ei,am,mbej->abij", t1, t1, Jmsa[o, v, v, o], optimize=True)
P4 = np.einsum("ei,abej->abij", t1, Jmsa[v, v, v, o], optimize=True)
P5 = np.einsum("am,mbij->abij", t1, Jmsa[o, v, o, o], optimize=True)
# calculate the right hand side of the CCSD equation for T2
rhs_T2 += 0.5 * np.einsum("abmn,mnij->abij", tau, Wmnij, optimize=True)
rhs_T2 += 0.5 * np.einsum("efij,abef->abij", tau, Wabef, optimize=True)
rhs_T2 += P1.transpose(0, 1, 2, 3) - P1.transpose(1, 0, 2, 3)
rhs_T2 -= P2.transpose(0, 1, 2, 3) - P2.transpose(0, 1, 3, 2)
rhs_T2 += P3.transpose(0, 1, 2, 3) - P3.transpose(0, 1, 3, 2)
rhs_T2 -= P3.transpose(1, 0, 2, 3) - P3.transpose(1, 0, 3, 2)
rhs_T2 += P4.transpose(0, 1, 2, 3) - P4.transpose(0, 1, 3, 2)
rhs_T2 -= P5.transpose(0, 1, 2, 3) - P5.transpose(1, 0, 2, 3)
# Update T1 and T2 amplitudes and save the previous iteration
t1, t2 = rhs_T1 / Emss, rhs_T2 / Emsd; E_CCSD_P = E_CCSD
# evaluate the energy
E_CCSD = 1.00 * np.einsum("ia,ai", Fms[o, v], t1 )
E_CCSD += 0.25 * np.einsum("ijab,abij", Jmsa[o, o, v, v], t2 )
E_CCSD += 0.50 * np.einsum("ijab,ai,bj", Jmsa[o, o, v, v], t1, t1)
# print the CCSD energy
print(" CCSD ENERGY: {:.8f}".format(E_HF + E_CCSD + VNN))
# CONFIGURATION INTERACTION ======================================================================================================
if args.fci:
# generate the determiants
dets = [np.array(det) for det in it.combinations(range(2 * nbf), 2 * nocc)]
# define the CI Hamiltonian
Hci = np.zeros([len(dets), len(dets)])
# define the Slater-Condon rules, "so" is an array of unique and common spinorbitals
slater0 = lambda so: (
sum(np.diag(Hms)[so]) + sum([0.5 * Jmsa[m, n, m, n] for m, n in it.product(so, so)])
)
slater1 = lambda so: (
Hms[so[0], so[1]] + sum([Jmsa[so[0], m, so[1], m] for m in so[2:]])
)
slater2 = lambda so: (
Jmsa[so[0], so[1], so[2], so[3]]
)
# filling of the CI Hamiltonian
for i in range(0, Hci.shape[0]):
for j in range(i, Hci.shape[1]):
# aligned determinant and the sign
aligned, sign = dets[j].copy(), 1
# align the determinant "j" to "i" and calculate the sign
for k in (k for k in range(len(aligned)) if aligned[k] != dets[i][k]):
while len(l := np.where(dets[i] == aligned[k])[0]) and l[0] != k:
aligned[[k, l[0]]] = aligned[[l[0], k]]; sign *= -1
# find the unique and common spinorbitals
so = np.block(list(map(lambda l: np.array(l), [
[aligned[k] for k in range(len(aligned)) if aligned[k] not in dets[i]],
[dets[i][k] for k in range(len(dets[j])) if dets[i][k] not in aligned],
[aligned[k] for k in range(len(aligned)) if aligned[k] in dets[i]]
]))).astype(int)
# apply the Slater-Condon rules and multiply by the sign
if ((aligned - dets[i]) != 0).sum() == 0: Hci[i, j] = Hci[j, i] = slater0(so) * sign
if ((aligned - dets[i]) != 0).sum() == 1: Hci[i, j] = Hci[j, i] = slater1(so) * sign
if ((aligned - dets[i]) != 0).sum() == 2: Hci[i, j] = Hci[j, i] = slater2(so) * sign
# solve the eigensystem and assign energy
eci, Cci = np.linalg.eigh(Hci); E_FCI = eci[0] - E_HF
# print the results
print(" FCI ENERGY: {:.8f}".format(E_HF + E_FCI + VNN))
Nonadiabatic Exact Quantum Dynamics
#!/usr/bin/env python
import argparse as ap, itertools as it, matplotlib.animation as anm, matplotlib.pyplot as plt, numpy as np, scipy as sp, scipy.linalg
np.set_printoptions(edgeitems=30, linewidth=100000, formatter=dict(float=lambda x: "%20.14f" % x)); np.random.seed(0)
# SECTION FOR PARSING COMMAND LINE ARGUMENTS =====================================================================================
# create the parser
parser = ap.ArgumentParser(
prog="resmet.py", description="Numerically Exact Quantum Dynamics Educational Toolkit",
formatter_class=lambda prog: ap.HelpFormatter(prog, max_help_position=128),
add_help=False, allow_abbrev=False
)
# add the arguments
parser.add_argument("-a", "--align", help="Align the wavefunction plot to the potential.", action=ap.BooleanOptionalAction)
parser.add_argument("-c", "--imaginary", help="Perform imaginary-time propagation for n states.", type=int, default=0)
parser.add_argument("-d", "--damp", help="Gaussian damping parameter. (default: %(default)s)", type=float, default=0.003)
parser.add_argument("-f", "--factor", help="Factor to scale the wavefunction. (default: %(default)s)", type=float, default=10)
parser.add_argument("-g", "--guess", help="Initial guess. (default: %(default)s)", type=str, default="[np.exp(-(r1-1)**2)]")
parser.add_argument("-h", "--help", help="Print this help message.", action=ap.BooleanOptionalAction)
parser.add_argument("-i", "--iterations", help="Number of iterations. (default: %(default)s)", type=int, default=500)
parser.add_argument("-l", "--limit", help="Grid limits. (default: %(default)s)", type=float, default=8)
parser.add_argument("-m", "--mass", help="Mass of the particle. (default: %(default)s)", type=float, default=1)
parser.add_argument("-n", "--ntraj", help="Number of Bohmian trajectories. (default: %(default)s)", type=int, default=100)
parser.add_argument("-p", "--points", help="Number of grid points. (default: %(default)s)", type=int, default=128)
parser.add_argument("-t", "--timestep", help="Time step of the simulation. (default: %(default)s)", type=float, default=0.1)
parser.add_argument("-u", "--adiabatic", help="Adiabatic transform. (default: %(default)s)", action=ap.BooleanOptionalAction)
parser.add_argument("-v", "--potential", help="Potential matrix. (default: %(default)s)", type=str, default="[[0.5*r1**2]]")
# parse the arguments
args = parser.parse_args()
# replace the variables in the string expressions
for i in range(1, 10): args.guess = args.guess.replace(f"r{i}", f"r[:, {i - 1}]")
for i in range(1, 10): args.potential = args.potential.replace(f"r{i}", f"r[:, {i - 1}]")
# print the help message if the flag is set
if args.help: parser.print_help(); exit()
# PREPARE VARIABLES FOR THE DYNAMICS =============================================================================================
# extract the number of dimensions without using regex
ndim = max([int(e[0]) for e in args.potential.split("r[:, ")[1:]]) + 1
# define the function to evaluate potential
potf = lambda: np.array(eval(args.potential)).transpose(2, 0, 1);
# define the function to evaluate the guess wavefunction on a specified grid
psif = lambda: np.array(list(map(lambda x: x * np.ones(args.points**ndim), eval(args.guess))), dtype=complex).T
# REAL AND IMAGINARY QUANTUM DYNAMICS ============================================================================================
# calculate the space step and the grid limits in each dimension
dr, grid = 2 * args.limit / (args.points - 1), [np.linspace(-args.limit, args.limit, args.points)] * ndim
# create the grid in real and fourier space
r = np.stack(np.meshgrid(*[np.linspace(-args.limit, args.limit, args.points)] * ndim, indexing="ij"), axis=-1).reshape(-1, ndim)
k = np.stack(np.meshgrid(*[2 * np.pi * np.fft.fftfreq(args.points, dr) ] * ndim, indexing="ij"), axis=-1).reshape(-1, ndim)
# create the potential, wavefunction and extract wavefunction dimensions
V, psi = potf(), psif(); shape = ndim * [args.points] + [psi.shape[1]]
# normalize the guess wavefunction
psi /= np.sqrt(dr * np.einsum("ij,ij->", psi.conj(), psi))
# create the containers for the observables and wavefunctions
position, momentum, ekin, epot = [], [], [], []; wfnopt, wfn = [], [psi]
# create the function to apply an operator in the momentum space to the n dimensional wavefunction
fapply = lambda op, psik: np.fft.ifftn((op * psik).reshape(shape), axes=range(ndim)).reshape(psi.shape)
# iterate over the propagations
for i in range(args.imaginary if args.imaginary else 1):
# print the propagation header
print() if i else None; print("PROPAGATION OF STATE %d " % (i))
# get the random coordinate indices
rind = np.random.choice(psi.shape[0], size=args.ntraj, p=(np.abs(psi)**2 * dr).sum(axis=1))
# create the initial points for Bohmian trajectories
trajs = np.concatenate((r[rind ][:, None, :], np.zeros((args.ntraj, args.iterations, ndim))), axis=1)
# initialize the initial wavefunction from the wfn container and clear all containers
psi = wfn[0]; wfn.clear(), position.clear(), momentum.clear(), ekin.clear(), epot.clear()
# get the propagator exponent unit
unit = -0.5 * (1 if args.imaginary else 1j) * args.timestep
# calculate the propagators for each point in the grid
K = np.array([sp.linalg.expm(unit * np.sum(k[i, :]**2) / args.mass * np.eye(psi.shape[1])) for i in range(r.shape[0])])
R = np.array([sp.linalg.expm(unit * V[i] ) for i in range(r.shape[0])])
# print the propagation header
print("%6s %12s %12s %12s" % ("ITER", "EKIN", "EPOT", "ETOT", ))
# propagate the wavefunction
for j in range(args.iterations + 1):
# propagate in real space
if (j): psi = np.einsum("ijk,ik->ij", R, psi)
# fourier transform the wavefunction
if (j): psi = np.fft.fftn(psi.reshape(shape), axes=range(ndim)).reshape(psi.shape)
# propagate in momentum space
if (j): psi = np.einsum("ijk,ik->ij", K, psi)
# inverse fourier transform the wavefunction
if (j): psi = np.fft.ifftn(psi.reshape(shape), axes=range(ndim)).reshape(psi.shape)
# propagate in real space
if (j): psi = np.einsum("ijk,ik->ij", R, psi)
# orthogonalize the wavefunction
for k in range(len(wfnopt)): psi -= np.sum(wfnopt[k].conj() * psi) * wfnopt[k] * dr
# normalize the wavefunction
if args.imaginary: psi /= np.sqrt(dr * np.einsum("ij,ij->", psi.conj(), psi))
# append the potential energy and the wavefunction
epot.append(np.einsum("ij,ijk,ik->", psi.conj(), V, psi).real * dr); wfn.append(psi.copy())
# create a n dimensional copy of the wavefunction and its fourier transform
psid, psik = psi.reshape(shape), np.fft.fftn(psi.reshape(shape), axes=range(ndim)).reshape(psi.shape)
# append the kinetic energy
ekin.append((psi.conj() * fapply((0.5 * np.sum(k**2, axis=1) / args.mass)[:, None], psik)).real.sum() * dr)
# append the position
position.append(np.sum(r * np.sum(np.abs(psi)**2, axis=1, keepdims=True), axis=0) * dr)
# append the momentum
momentum.append([(psi.conj() * fapply(1j * k[:, l:l + 1], psik)).imag.sum() * dr for l in range(ndim)])
# Bohmian velocity container
v = np.zeros(shape[:-1] + [ndim])
# calculate probability current in each dimension
if (j): currents = [(np.conjugate(psid) * np.gradient(psid, dr, axis=l)).sum(axis=-1).imag for l in range(ndim)]
# calculate the velocity of the Bohmian trajectories
if (j): v[..., :] = np.array([currents[l] / ((np.abs(psid)**2).sum(axis=-1) + 1e-14) / args.mass for l in range(ndim)]).T
# propagate the Bohmian trajectories
if (j): trajs[:, j] = (lambda r: r + sp.interpolate.interpn(points=grid, values=v, xi=r) * args.timestep)(trajs[:, j - 1])
# print the iteration info
if j % 100 == 0: print("%6d %12.6f %12.6f %12.6f" % (j, ekin[-1], epot[-1], ekin[-1] + epot[-1]))
# append the optimized wavefunction to the container
if args.imaginary: wfnopt.append(psi.copy())
# ADIABATIC TRANSFORM AND SPECTRUM ===============================================================================================
# calculate the adiabatic eigenstates
A = [np.linalg.eigh(V[i])[0] for i in range(V.shape[0])]
U = [np.linalg.eigh(V[i])[1] for i in range(V.shape[0])]
# adiabatize the potential and wavefunctions
V = np.array ([np.diag(a) for a in A]) if args.adiabatic else np.array(V )
wfn = np.einsum("ikl,jik->jil", U, wfn ) if args.adiabatic else np.array(wfn)
# calculate the density matrices and acf
density = np.einsum("jia,jib->jab", wfn, wfn.conj()).real * dr
acf = np.einsum("ij,tij->t", wfn[0].conj(), wfn ) * dr
# symmetrize the acf
acf = np.concatenate((np.flip(acf)[:-1], np.array(acf).conj()))
# apply the damping to the acf
acf *= np.exp(-args.damp * (np.arange(-args.iterations, args.iterations + 1) * args.timestep)**2)
# create the padded acf
acfpad = np.pad(acf, 2 * [10 * len(acf)], mode="constant")
# calculate the spectrum of the zero-padded acf and the corresponding energies
spectrum = np.abs(np.fft.fft(acfpad))**2; omega = 2 * np.pi * np.fft.fftfreq(len(spectrum), args.timestep)
# PRINT AND PLOT THE RESULTS =====================================================================================================
# print the final population
print(); print("FINAL POPULATION: %s" % np.diag(density[-1]))
# scale the wavefunction and add the potential to the wavefunction
wfn = args.factor * np.array(wfn); wfn += (1 + 1j) * np.einsum("ijj,k->kij", V, np.ones(args.iterations + 1)) if args.align else 0
# extract the momenta from the Bohmian trajectories
ptrajs = np.gradient(trajs, args.timestep, axis=1) * args.mass
# create the stationary subplots
fig, axs = plt.subplots(2, 3, figsize=(12, 6)); time = np.arange(args.iterations + 1) * args.timestep
# plot the energies
axs[0, 0].plot(time, ekin, label="Kinetic Energy" )
axs[0, 0].plot(time, epot, label="Potential Energy")
# plot the population
axs[0, 1].plot(time, [np.diag(rho) for rho in density], label=[f"S$_{i}$" for i in range(density.shape[1])])
# print the acf
axs[1, 0].plot(np.arange(-args.iterations, args.iterations + 1) * args.timestep, np.array(acf).real, label="Re(ACF)")
axs[1, 0].plot(np.arange(-args.iterations, args.iterations + 1) * args.timestep, np.array(acf).imag, label="Im(ACF)")
# plot the spectrum
axs[1, 1].plot(omega[np.argsort(omega)], spectrum[np.argsort(omega)] / np.max(spectrum))
# plot the position and momentum of the Bohmian trajectories
[axs[0, 2].plot(time, trajs[i, :, 0], alpha=0.05, color="tab:blue") for i in range(args.ntraj)]
[axs[1, 2].plot(time, ptrajs[i, :, 0], alpha=0.05, color="tab:blue") for i in range(args.ntraj)]
# plot the numerically exact expectation values of position and momentum
axs[0, 2].plot(time, np.array(position)[:, 0], color="tab:orange", label="$<\Psi|\hat{r_x}|\Psi>$")
axs[1, 2].plot(time, np.array(momentum)[:, 0], color="tab:orange", label="$<\Psi|\hat{p_x}|\Psi>$")
# plot the mean of the Bohmian trajectories
axs[0, 2].plot(time, np.mean( trajs[:, :, 0], axis=0), "--", color="black", label="$<r_x>$")
axs[1, 2].plot(time, np.mean(ptrajs[:, :, 0], axis=0), "--", color="black", label="$<p_x>$")
# set the labels for the stationary plot
axs[0, 0].set_xlabel("Time (a.u.)" ); axs[0, 0].set_ylabel("Energy (a.u.)" )
axs[0, 1].set_xlabel("Time (a.u.)" ); axs[0, 1].set_ylabel("Population" )
axs[0, 2].set_xlabel("Time (a.u.)" ); axs[0, 2].set_ylabel("Position (a.u.)" )
axs[1, 0].set_xlabel("Time (a.u.)" ); axs[1, 0].set_ylabel("Autocorrelation Function")
axs[1, 1].set_xlabel("Energy (a.u.)"); axs[1, 1].set_ylabel("Normalized Intensity" )
axs[1, 2].set_xlabel("Time (a.u.)" ); axs[1, 2].set_ylabel("Momentum (a.u.)" )
# set the domain for the spectrum plot, the end will be as last element from the end less than some value
axs[1, 1].set_xlim(0, omega[np.argsort(omega)][np.where(spectrum[np.argsort(omega)] / np.max(spectrum) > 1e-6)][-1])
# enable legends
axs[0, 0].legend(); axs[0, 1].legend(); axs[1, 0].legend(); axs[0, 2].legend(); axs[1, 2].legend()
# only for 1D
if ndim == 1:
# creace the wavefunction plot
wfig, wax = plt.subplots(1, 1)
# plot the wavefunction
wfnplot = np.array([
[wax.plot(r, wfn[0][:, i].real, label="Re($\Psi$)")[0], wax.plot(r, wfn[0][:, i].imag, label="Im($\Psi$)")[0]]
for i in range(psi.shape[1])]).flatten()
# set the labels for the wavefunction plot and enable legend
wax.set_xlabel("Position (a.u.)"); wax.set_ylabel("Wavefunction"); wax.legend()
# extract the wfn min and max
minwfn, maxwfn = min(np.array(wfn).real.min(), np.array(wfn).imag.min()), max(np.array(wfn).real.max(), np.array(wfn).imag.max())
# set the limits for the wavefunction animation
wax.set_ylim(minwfn - 0.1 * (maxwfn - minwfn), maxwfn + 0.1 * (maxwfn - minwfn))
# define the update function for the wavefunction animation
update = lambda i: [
wfnplot[j].set_ydata(wfn[i][:, j // 2].real if j % 2 == 0 else wfn[i][:, j // 2].imag) for j in range(2 * psi.shape[1])
]
# make the wavefunction plot animation and set the layout
ani = anm.FuncAnimation(wfig, update, frames=range(len(wfn)), repeat=True, interval=30); wfig.tight_layout()
# show the plot
fig.tight_layout(); plt.show()
Surface Hopping Classical Dynamics
#!/usr/bin/env python
import argparse as ap, itertools as it, matplotlib.animation as anm, matplotlib.pyplot as plt, numpy as np, scipy as sp, scipy.linalg
np.set_printoptions(edgeitems=30, linewidth=100000, formatter=dict(float=lambda x: "%20.14f" % x)); np.random.seed(0)
# SECTION FOR PARSING COMMAND LINE ARGUMENTS ===========================================================================
# create the parser
parser = ap.ArgumentParser(
prog="resmet.py", description="Surface Hopping Classical Dynamics Educational Toolkit",
formatter_class=lambda prog: ap.HelpFormatter(prog, max_help_position=128),
add_help=False, allow_abbrev=False
)
# add the arguments
parser.add_argument("-h", "--help", help="Print this help message.", action=ap.BooleanOptionalAction)
parser.add_argument("-i", "--iterations", help="Number of iterations. (default: %(default)s)", type=int, default=5000)
parser.add_argument("-m", "--mass", help="Mass of the particle. (default: %(default)s)", type=float, default=1)
parser.add_argument("-n", "--trajectories", help="Number of classical trajectories. (default: %(default)s)", type=int, default=1000)
parser.add_argument("-s", "--state", help="Initial state of the trajectories. (default: %(default)s)", type=int, default=0)
parser.add_argument("-t", "--timestep", help="Time step of the simulation. (default: %(default)s)", type=float, default=0.1)
parser.add_argument("-p", "--momentum", help="Momentum distribution. (default: %(default)s)", type=float, nargs= "+", default=[0, 1])
parser.add_argument("-r", "--position", help="Coords distribution. (default: %(default)s)", type=float, nargs= "+", default=[1, 0.5])
parser.add_argument("-u", "--adiabatic", help="Adiabatic transform. (default: %(default)s)", action=ap.BooleanOptionalAction)
parser.add_argument("-v", "--potential", help="Potential matrix. (default: %(default)s)", type=str, default="[[0.5*r1**2]]")
# surface hopping arguments
parser.add_argument("--lzsh", help="Enable LZSH algorithm. (default: %(default)s)", action=ap.BooleanOptionalAction)
parser.add_argument("--fssh", help="Enable FSSH algorithm. (default: %(default)s)", action=ap.BooleanOptionalAction)
parser.add_argument("--mash", help="Enable MASH algorithm. (default: %(default)s)", action=ap.BooleanOptionalAction)
# parse the arguments
args = parser.parse_args()
# replace the variables in the string expressions
for i in range(1, 10): args.potential = args.potential.replace(f"r{i}", f"r[:, {i - 1}]")
# print the help message if the flag is set
if args.help: parser.print_help(); exit()
# SURFACE HOPPING ALGORITHMS =====================================================================================================
# define the state vector and containers for the potential and transformation matrices
s = np.zeros((args.trajectories, args.iterations + 1), dtype=int) + args.state; Vs, Us = [], []
# define the Fewest Switches Surface Hopping algorithm
def fssh(i, substeps=10):
# fill the coefficients on the first iteration
if i == 1:
for j in range(args.trajectories): C[j, s[j, 0]] = 1
# compute the eigenvector overlap and the trajectory indices
O = Us[-1].swapaxes(1, 2) @ Us[-2]; traji = np.arange(args.trajectories)
# calculate the time derivative coupling
TDC = (np.transpose(O, (0, 2, 1)) - O) / (2 * args.timestep)
# vectorized derivative function
dC = lambda C: -1j * np.einsum("ijj->ij", Vs[-1]) * C - np.einsum("ijk,ik->ij", TDC, C)
# loop over the substeps
for _ in range(substeps):
# calculate the Runge-Kutta coefficients
k1 = dC(C )
k2 = dC(C + 0.5 * args.timestep * k1 / substeps)
k3 = dC(C + 0.5 * args.timestep * k2 / substeps)
k4 = dC(C + 1.0 * args.timestep * k3 / substeps)
# update the electronic coefficients and get trajectory indices
C[:] = C + args.timestep * (k1 + 2 * k2 + 2 * k3 + k4) / substeps / 6.0;
# calculate the denominator for the hopping probability
denominator = 0.5 * (np.abs(C[traji, s[:, i]])**2 + 1e-14) * substeps / args.timestep
# calculate the hopping probability
p = TDC[traji, s[:, i], :] * (C * np.conjugate(C[traji, s[:, i]][:, None])).real / denominator[:, None]
# calculate the hopping mask
hopmask = (np.random.rand(args.trajectories)[:, None] < p)
# perform the hop using the mask to the state that has maximum hopping probability
if hopmask.any(): s[hopmask.any(axis=1), i:] = np.argmax(p, axis=1)[hopmask.any(axis=1), None]
# define the Landau-Zener Surface Hopping algorithm
def lzsh(i):
# loop over the trajectories
for j in range(args.trajectories):
# generate random number
rn = np.random.rand()
# loop over the states
for k in (l for l in range(V.shape[1]) if l != s[j, i]):
# skip current state
if k == s[j, i]: continue
# calculate the energy differences
Z0 = abs(Vs[-1][j, k, k] - Vs[-1][j, s[j, i], s[j, i]])
Z1 = abs(Vs[-2][j, k, k] - Vs[-2][j, s[j, i], s[j, i]])
Z2 = abs(Vs[-3][j, k, k] - Vs[-3][j, s[j, i], s[j, i]])
# calculate first derivatives of the energy differences
dZ0, dZ1 = (Z0 - Z1) / args.timestep, (Z1 - Z2) / args.timestep
# calculate the second derivative
ddZ0 = (Z0 - 2 * Z1 + Z2) / args.timestep**2
# check if the trajectory is in the place for a hop
if (dZ0 * dZ1 > 0 or (dZ0 * dZ1 < 0 and ddZ0 < 0)): continue
# calculate the hopping probability
p = np.exp(-0.5 * np.pi * np.sqrt(Z0**3 / ddZ0)) if (Z0**3 * ddZ0 > 0) else 0
# hop to another state
if rn < p:
s[j, i:] = k; break
# define the Mapping Approach to Surface Hopping algorithm
def mash(i):
# sample the bloch vectors if i == 1
for j in (k for k in range(args.trajectories) if i == 1):
# define the angles
p, ct = 2 * np.pi * np.random.rand(), np.sqrt(np.random.rand()); st = np.sqrt(1 - ct**2)
# fill the row
for l in range(len(SPHR)): SPHR[l, j] = np.array([st * np.cos(p), st * np.sin(p), ct])
# compute the eigenvector overlap and store spheres
O = Us[-1].swapaxes(1, 2) @ Us[-2]; SPHRP = SPHR.copy()
# calculate the time derivative coupling
TDC = (np.transpose(O, (0, 2, 1)) - O) / (2 * args.timestep)
# loop over the trajectories
for j in range(args.trajectories):
# loop over all spheres
for k in range(len(SPHR)):
# get upper and lower index
su, sd = max(SIND[j, k]), min(SIND[j, k])
# calculate the propagator
omega = sp.linalg.expm(np.array([
[0, Vs[-1][j, sd, sd] - Vs[-1][j, su, su], 2 * TDC[j, sd, su]],
[Vs[-1][j, su, su] - Vs[-1][j, sd, sd], 0, 0 ],
[-2 * TDC[j, sd, su], 0, 0 ]
]) * args.timestep)
# propagate the Bloch vector
SPHR[k, j] = omega @ SPHR[k, j]
# loop over spheres
for k in range(len(SPHR)):
# perform the jump if the sphere crossed the equator
if SPHRP[k, j, 2] * SPHR[k, j, 2] < 0: s[j, i:] = SIND[j, k][0] if SIND[j, k][1] == s[j, i] else SIND[j, k][1]
# loop over sphere indices if the jump occured
for k in (l for l in range(len(SIND[j])) if s[j, i] != s[j, i - 1]):
# update the indices of the spheres
if SIND[j, k][0] == s[j, i - 1] and SIND[j, k][1] != s[j, i]: SIND[j, k][0] = s[j, i]
if SIND[j, k][1] == s[j, i - 1] and SIND[j, k][0] != s[j, i]: SIND[j, k][1] = s[j, i]
# PERFORM THE CLASSICAL DYNAMICS =================================================================================================
# create the initial conditions for the trajectories
r = np.column_stack([np.random.normal(mu, sigma, len(s)) for mu, sigma in zip(args.position[0::2], args.position[1::2])])
v = np.column_stack([np.random.normal(mu, sigma, len(s)) for mu, sigma in zip(args.momentum[0::2], args.momentum[1::2])])
# extract the number of states
nstate = np.array(eval(args.potential)).shape[0]
# create the electronic coefficients and one Bloch sphere
C = np.zeros((args.trajectories, nstate), dtype=complex)
S = np.zeros((args.trajectories, 3 ), dtype=float )
# create the array of spheres and sphere indices for the MASH algorithm
SPHR = np.array([ S.copy() for i in range(nstate - 1) ] )
SIND = np.array([[[s[0, 0], i] for i in range(nstate - 0) if i != s[0, 0]] for j in range(args.trajectories)])
# divide momentum by mass and define acceleration
v /= args.mass; a = np.zeros_like(r); diff = 0.0001;
# print the propagation header
print("%6s %12s %12s %12s" % ("ITER", "EKIN", "EPOT", "ETOT"))
# loop over the iterations
for i in range(args.iterations + 1):
# save the previous values
rp = r.copy(); vp = v.copy(); ap = a.copy()
# loop over the dimensions and propagate
for j in (j for j in range(r.shape[1]) if i):
# calculate offsets
rplus = r.copy(); rplus[:, j] += diff
rmins = r.copy(); rmins[:, j] -= diff
# calculate the potential energy at the offsetted positions
Vplus = np.array(eval(args.potential.replace("r", "rplus"))).transpose(2, 0, 1)
Vmins = np.array(eval(args.potential.replace("r", "rmins"))).transpose(2, 0, 1)
# transform the potential energy to the adiabatic basis
if args.adiabatic: Vplus = np.einsum("ij,jk->ijk", np.linalg.eigvalsh(Vplus), np.eye(V.shape[1]))
if args.adiabatic: Vmins = np.einsum("ij,jk->ijk", np.linalg.eigvalsh(Vmins), np.eye(V.shape[1]))
Vplusi = np.einsum("ijj->ij", Vplus)[np.arange(args.trajectories), s[:, i]]
Vminsi = np.einsum("ijj->ij", Vmins)[np.arange(args.trajectories), s[:, i]]
# calculate the acceleration
a[:, j] = 0.5 * (Vminsi - Vplusi) / (diff * args.mass)
# update the positions and velocities
v[:, j] += 0.5 * (a[:, j] + ap[:, j]) * args.timestep; r[:, j] += (v[:, j] + 0.5 * a[:, j] * args.timestep) * args.timestep
# calculate the diabatic potential
V = np.array(eval(args.potential)).transpose(2, 0, 1)
# adiabatize the potential
E, U = np.linalg.eigh(V); V = np.einsum("ij,jk->ijk", E, np.eye(V.shape[1])); Vs.append(V); Us.append(U)
# maximize overlap between the current and previous eigenvectors
if i > 1: Us[-1] = Us[-1] * np.where(np.sum(Us[-1] * Us[-2], axis=1) < 0, -1, 1)[:, None, :]
# surface hopping algorithm
if (args.fssh and i > 0): fssh(i)
if (args.lzsh and i > 1): lzsh(i)
if (args.mash and i > 0): mash(i)
# calculate the potential and kinetic energy for each trajectory
Epot = V[np.arange(args.trajectories), s[:, i], s[:, i]]; Ekin = 0.5 * args.mass * np.sum(v**2, axis=1)
# get the mask for trajectories that hopped to another state
mask = (s[:, i] != s[:, i - bool(i)]) & (Ekin > (Epot - V[np.arange(args.trajectories), s[:, i - bool(i)], s[:, i - bool(i)]]))
# rescale the velocities of the trajectories that hopped to another state
if mask.any(): v[mask] *= np.sqrt((Ekin - Epot + V[np.arange(args.trajectories), s[:, i - 1], s[:, i - 1]]) / Ekin)[mask, None]
# recalculate the kinetic energy if the velocities were rescaled
if mask.any(): Ekin = 0.5 * args.mass * np.sum(v**2, axis=1)
# print the iteration
if i % 1000 == 0: print("%6d %12.6f %12.6f %12.6f" % (i, np.mean(Ekin), np.mean(Epot), np.mean(Ekin + Epot)))
# PRINT AND PLOT THE RESULTS =====================================================================================================
# print the final populations
print(); print("FINAL POPULATION: %s" % (np.array([(s[:, -1] == j).sum() / s.shape[0] for j in range(V.shape[1])])))
# create the subplots
fig, axs = plt.subplots(1, 1, figsize=(8, 6));
# plot the population
axs.plot(np.arange(args.iterations + 1) * args.timestep, [[(s[:, i] == j).sum() / s.shape[0] for j in range(V.shape[1])] for i in range(s.shape[1])])
# set the labels
axs.set_xlabel("Time (a.u.)"); axs.set_ylabel("Population")
# show the plot
plt.tight_layout(); plt.show()