def propagate_wfn(wfn, dt):
"""Propagates the Bundle object with RK4."""
ncrd = wfn.traj[0].dim
ntraj = wfn.n_traj()
kx = np.zeros((ntraj, rk_ordr, ncrd))
kp = np.zeros((ntraj, rk_ordr, ncrd))
kg = np.zeros((ntraj, rk_ordr))
for rk in range(rk_ordr):
tmp_wfn = wfn.copy()
for i in range(ntraj):
if tmp_wfn.traj[i].active:
propagate_rk(tmp_wfn.traj[i], dt, rk, kx[i], kp[i], kg[i])
# update the PES to evaluate new gradients
if rk < rk_ordr - 1:
evaluate.update_pes(tmp_wfn, update_integrals=False)
# update to the final position
for i in range(ntraj):
if wfn.traj[i].active:
wfn.traj[i].update_x(wfn.traj[i].x() +
np.sum(wgt[:, np.newaxis] * kx[i], axis=0))
wfn.traj[i].update_p(wfn.traj[i].p() +
np.sum(wgt[:, np.newaxis] * kp[i], axis=0))
if propphase:
wfn.traj[i].update_phase(wfn.traj[i].phase() +
np.sum(wgt * kg[i, :]))
evaluate.update_pes(wfn)
def propagate_wfn(wfn, dt):
"""Propagates the Bundle object with VV."""
# update position
for i in range(wfn.n_traj()):
if wfn.traj[i].active:
propagate_position(wfn.traj[i], dt)
# update electronic structure for all trajectories
# and centroids (where necessary)
evaluate.update_pes(wfn)
# finish update of momentum and phase
for i in range(wfn.n_traj()):
if wfn.traj[i].active:
propagate_momentum(wfn.traj[i], dt)
def propagate_wfn(wfn, dt):
"""Propagates the Bundle object with BS."""
global h
ncrd = wfn.traj[0].dim
ntraj = wfn.n_traj()
t = 0.
if h is None:
h = dt
while abs(t) < abs(dt):
hstep = np.sign(dt) * min(abs(h), abs(dt - t))
reduced = False
tsav = np.zeros(kmax)
err = np.zeros(kmax)
Tx = np.zeros((kmax, ntraj, ncrd))
Tp = np.zeros((kmax, ntraj, ncrd))
Tg = np.zeros((kmax, ntraj))
for k in range(kmax):
tmp_wfn = wfn.copy()
x0 = np.zeros((ntraj, ncrd))
p0 = np.zeros((ntraj, ncrd))
g0 = np.zeros((ntraj))
x1 = np.zeros((ntraj, ncrd))
p1 = np.zeros((ntraj, ncrd))
g1 = np.zeros((ntraj))
# step through n modified midpoint steps
for n in range(nstep[k]):
for i in range(ntraj):
if tmp_wfn.traj[i].active:
mm_step(tmp_wfn.traj[i], hstep / nstep[k], x0[i],
x1[i], p0[i], p1[i], g0[i], g1[i], n)
evaluate.update_pes(tmp_wfn, update_integrals=False)
# compute the modified midpoint estimate
for i in range(ntraj):
if wfn.traj[i].active:
x1[i] = 0.5 * (x1[i] + x0[i] + hstep / nstep[k] *
tmp_wfn.traj[i].velocity())
p1[i] = 0.5 * (p1[i] + p0[i] +
hstep / nstep[k] * tmp_wfn.traj[i].force())
if propphase:
g1[i] = 0.5 * (g1[i] + g0[i] + hstep / nstep[k] *
tmp_wfn.traj[i].phase_dot())
# extrapolate from modified midpoint results
poly_extrapolate(k, (hstep / nstep[k])**2, tsav, x1, p1, g1, Tx,
Tp, Tg)
if k > 0:
errmax = np.amax((abs(Tx[k]), abs(Tp[k]), abs(Tg[k]))) / tol
err[k - 1] = (errmax / 0.25)**(1. / (2. * k + 1.))
if k >= kopt - 2:
if errmax > 1:
# scale the time step and try again
sfac, reduced = reduce_tstep(k, err)
if reduced:
red = max(1e-5, min(reduced, 0.7))
h = sfac * hstep
break
else:
# scale the time step if possible
t += h
h = increase_tstep(k, err, reduced) * hstep
# update to the final position
xnew = np.sum(Tx, axis=0)
pnew = np.sum(Tp, axis=0)
gnew = np.sum(Tg, axis=0)
for i in range(ntraj):
if wfn.traj[i].active:
wfn.traj[i].update_x(xnew[i])
wfn.traj[i].update_p(pnew[i])
if propphase:
wfn.traj[i].update_phase(gnew[i])
evaluate.update_pes(
wfn, update_integrals=(abs(t) >= abs(dt)))
break
def step_wavefunction(dt):
"""Propagates the wave packet using a run-time selected propagator."""
# save the wavefunction from previous step in case step rejected
end_time = min(glbl.modules['wfn'].time + dt,
glbl.properties['simulation_time'])
time_step = min(
dt, glbl.properties['simulation_time'] - glbl.modules['wfn'].time)
min_time_step = dt / 2.**5
while not step_complete(glbl.modules['wfn'].time, end_time, dt):
# save the wavefunction from previous step in case step rejected
wfn_start = glbl.modules['wfn'].copy()
# propagate each trajectory in the wavefunction
time_step = min(time_step, end_time - glbl.modules['wfn'].time)
# propagate amplitudes for 1/2 time step using x0
glbl.modules['wfn'].update_amplitudes(0.5 * dt)
# the propagators update the potential energy surface as need be.
glbl.modules['propagator'].propagate_wfn(glbl.modules['wfn'],
time_step)
# update the couplings for all the trajectories
for i in range(glbl.modules['wfn'].n_traj()):
glbl.modules['interface'].evaluate_coupling(
glbl.modules['wfn'].traj[i])
# propagate amplitudes for 1/2 time step using x1
glbl.modules['matrices'].build(glbl.modules['wfn'],
glbl.modules['integrals'])
glbl.modules['wfn'].update_matrices(glbl.modules['matrices'])
glbl.modules['wfn'].update_amplitudes(0.5 * dt)
# Renormalization
if glbl.properties['renorm'] == 1:
glbl.modules['wfn'].renormalize()
# check time_step is fine, energy/amplitude conserved
accept, error_msg = check_step_wfn(wfn_start, glbl.modules['wfn'],
time_step)
# if everything is ok..
if accept:
# update the wavefunction time
glbl.modules['wfn'].time += time_step
# spawn new basis functions if necessary
basis_grown = glbl.modules['adapt'].spawn(glbl.modules['wfn'],
time_step)
# kill the dead trajectories
basis_pruned = utilities.prune(glbl.modules['wfn'])
# if a trajectory has been added, then call update_pes
# to get the electronic structure information at the associated
# centroids. This is necessary in order to propagate the amplitudes
# at the start of the next time step.
if basis_grown and glbl.modules['integrals'].require_centroids:
evaluate.update_pes(glbl.modules['wfn'])
# update the Hamiltonian and associated matrices
if basis_grown or basis_pruned:
glbl.modules['matrices'].build(glbl.modules['wfn'],
glbl.modules['integrals'])
glbl.modules['wfn'].update_matrices(glbl.modules['matrices'])
for i in range(glbl.modules['wfn'].n_traj()):
glbl.modules['interface'].evaluate_coupling(
glbl.modules['wfn'].traj[i])
# re-expression of the basis using the matching pursuit
# algorithm
#if glbl.properties['matching_pursuit'] == 1:
# mp.reexpress_basis(master)
# update the running log
log.print_message('t_step', [
glbl.modules['wfn'].time, time_step, glbl.modules['wfn'].nalive
])
#del wfn_start
else:
# recall -- this time trying to propagate to the failed step
time_step *= 0.5
log.print_message('new_step', [error_msg, time_step])
if time_step < min_time_step:
log.print_message('general',
['minimum time step exceeded -- STOPPING.'])
raise ValueError('Bundle minimum step exceeded.')
# reset the beginning of the time step and go to beginning of loop
#del master
glbl.modules['wfn'] = wfn_start.copy()
def init_wavefunction():
"""Initializes the trajectories."""
# initialize the interface we'll be using the determine the
# the PES. There are some details here that trajectories
# will want to know about
log.print_message(
'string',
['\n **************\n' + ' initialization\n' + ' **************\n'])
# creat the wave function instance
glbl.modules['wfn'] = wavefunction.Wavefunction()
glbl.modules['interface'] = __import__('nomad.interfaces.' +
glbl.methods['interface'],
fromlist=['NA'])
glbl.modules['interface'].init_interface()
# create glbl modules
glbl.modules['matrices'] = matrices.Matrices()
glbl.modules['init_conds'] = __import__('nomad.initconds.' +
glbl.methods['init_conds'],
fromlist=['NA'])
glbl.modules['adapt'] = __import__('nomad.adapt.' +
glbl.methods['adapt_basis'],
fromlist=['a'])
glbl.modules['propagator'] = __import__('nomad.propagators.' +
glbl.methods['propagator'],
fromlist=['a'])
# WARNING: this should be done more eloquently
# need to determine form of the coefficients on the kinetic energy
# operator, i.e. cartesian vs. normal mode coordinate basis
# this is messy -- should be re-worked mvoing forward...
if glbl.methods['interface'] == 'vibronic':
# if normal modes, read frequencies from vibronic interface
kecoef = 0.5 * glbl.modules['interface'].ham.freq
else:
kecoef = 0.5 / glbl.properties['crd_masses']
glbl.modules['integrals'] = integral.Integral(
kecoef, glbl.methods['ansatz'], glbl.methods['integral_eval'])
# now load the initial trajectories into the bundle
if glbl.properties['restart']:
# print to logfile that we're restarting simulation
log.print_message('string', [
' restarting simulation from checkpoint file: ' +
str(glbl.paths['chkpt_file']) + '\n'
])
# retrieve current wave function, no arguments defaults to most recent simulation
wfn0 = None
ints0 = None
if glbl.mpi['rank'] == 0:
[glbl.modules['wfn'],
glbl.modules['integrals']] = checkpoint.retrieve_simulation(
time=glbl.properties['restart_time'])
[wfn0, ints0] = checkpoint.retrieve_simulation(time=0.)
# only root reads checkpoint file -- then broadcasts contents to other proceses
if glbl.mpi['parallel']:
# synchronize tasks
glbl.mpi['comm'].barrier()
glbl.modules['wfn'] = glbl.mpi['comm'].bcast(glbl.modules['wfn'],
root=0)
wfn0 = glbl.mpi['comm'].bcast(wfn0, root=0)
if glbl.modules['integrals'].require_centroids:
glbl.modules['integrals'].centroid_required = glbl.mpi[
'comm'].bcast(glbl.modules['integrals'].centroid_required,
root=0)
glbl.modules['integrals'].centroids = glbl.mpi['comm'].bcast(
glbl.modules['integrals'].centroids, root=0)
# save copy of t=0 wfn for autocorrelation function
save_initial_wavefunction(wfn0)
# check that we have all the data we need to propagate the first step -- otherwise, we
# need to update the potential
update_surface = False
for i in range(glbl.modules['wfn'].n_traj()):
if glbl.modules['wfn'].traj[i].alive and glbl.modules['wfn'].traj[
i].active:
pes_data = glbl.modules['wfn'].traj[i].pes
if ('potential' not in pes_data.avail_data()
or 'derivative' not in pes_data.avail_data()):
update_suface = True
if glbl.modules['integrals'].require_centroids:
for j in range(i):
pes_data = glbl.modules['integrals'].centroids[i][j].pes
if glbl.modules['integrals'].centroid_required[i][j]:
if ('potential' not in pes_data.avail_data()
and glbl.modules['wfn'].traj[i].state
== glbl.modules['wfn'].traj[j].state):
update_surface = True
if ('derivative' not in pes_data.avail_data()
and glbl.modules['wfn'].traj[i].state !=
glbl.modules['wfn'].traj[j].state):
update_surface = True
if update_surface:
evaluate.update_pes(glbl.modules['wfn'])
# update the couplings for all the trajectories
for i in range(glbl.modules['wfn'].n_traj()):
glbl.modules['interface'].evaluate_coupling(
glbl.modules['wfn'].traj[i])
# build the necessary matrices
glbl.modules['matrices'].build(glbl.modules['wfn'],
glbl.modules['integrals'])
glbl.modules['wfn'].update_matrices(glbl.modules['matrices'])
else:
# first generate the initial nuclear coordinates and momenta
# and add the resulting trajectories to the bundle
glbl.modules['init_conds'].set_initial_coords(glbl.modules['wfn'])
# set the initial state of the trajectories in bundle. This may
# require evaluation of electronic structure
set_initial_state(glbl.modules['wfn'])
# set the initial amplitudes of the basis functions
set_initial_amplitudes(glbl.modules['wfn'])
# add virtual basis functions, if desired
if glbl.properties['virtual_basis']:
virtual_basis(glbl.modules['wfn'])
# update the integrals
glbl.modules['integrals'].update(glbl.modules['wfn'])
# once phase space position and states of the basis functions
# are set, update the potential information
evaluate.update_pes(glbl.modules['wfn'])
# update the couplings for all the trajectories
for i in range(glbl.modules['wfn'].n_traj()):
glbl.modules['interface'].evaluate_coupling(
glbl.modules['wfn'].traj[i])
# compute the hamiltonian matrix...
glbl.modules['matrices'].build(glbl.modules['wfn'],
glbl.modules['integrals'])
glbl.modules['wfn'].update_matrices(glbl.modules['matrices'])
# so that we may appropriately renormalize to unity
glbl.modules['wfn'].renormalize()
# this is the bundle at time t=0. Save in order to compute auto
# correlation function
save_initial_wavefunction(glbl.modules['wfn'])
# write the wavefunction to the archive
if glbl.mpi['rank'] == 0:
checkpoint.archive_simulation(glbl.modules['wfn'],
glbl.modules['integrals'])
log.print_message(
'string', ['\n ***********\n' + ' propagation\n' + ' ***********\n\n'])
log.print_message('t_step', [
glbl.modules['wfn'].time, glbl.properties['default_time_step'],
glbl.modules['wfn'].nalive
])
return glbl.modules['wfn'].time
def propagate_wfn(wfn, dt):
"""Propagates the Bundle object with RKF45."""
global h
ncrd = wfn.traj[0].dim
ntraj = wfn.n_traj()
kx = np.zeros((ntraj, rk_ordr, ncrd))
kp = np.zeros((ntraj, rk_ordr, ncrd))
kg = np.zeros((ntraj, rk_ordr))
t = 0.
if h is None:
h = dt
while abs(t) < abs(dt):
hstep = np.sign(dt) * min(abs(h), abs(dt - t))
for rk in range(rk_ordr):
tmp_wfn = wfn.copy()
for i in range(ntraj):
if tmp_wfn.traj[i].active:
propagate_rk(tmp_wfn.traj[i], hstep, rk, kx[i], kp[i],
kg[i])
# update the PES to evaluate new gradients
if rk < rk_ordr - 1:
evaluate.update_pes(tmp_wfn, update_integrals=False)
# calculate the 4th and 5th order changes and the error
dx_lo = np.zeros((wfn.nalive, ncrd))
dx_hi = np.zeros((wfn.nalive, ncrd))
dp_lo = np.zeros((wfn.nalive, ncrd))
dp_hi = np.zeros((wfn.nalive, ncrd))
dg_lo = np.zeros((wfn.nalive))
dg_hi = np.zeros((wfn.nalive))
for i in range(ntraj):
if wfn.traj[i].active:
dx_lo[i] = np.sum(wgt_lo[:, np.newaxis] * kx[i], axis=0)
dx_hi[i] = np.sum(wgt_hi[:, np.newaxis] * kx[i], axis=0)
dp_lo[i] = np.sum(wgt_lo[:, np.newaxis] * kp[i], axis=0)
dp_hi[i] = np.sum(wgt_hi[:, np.newaxis] * kp[i], axis=0)
if propphase:
for i in range(ntraj):
if wfn.traj[i].active:
dg_lo[i] = np.sum(wgt_lo * kg[i])
dg_hi[i] = np.sum(wgt_hi * kg[i])
#print("a="+str(np.abs(dx_hi-dx_lo).flatten()))
#print("b="+str(np.abs(dp_hi-dp_lo).flatten()))
#print("c="+str(np.abs(dg_hi-dg_lo)))
err = np.max(
np.max(np.abs(dg_hi - dg_lo)),
np.max((np.abs(dx_hi - dx_lo).flatten(),
np.abs(dp_hi - dp_lo).flatten())))
else:
err = np.max((np.abs(dx_hi - dx_lo), np.abs(dp_hi - dp_lo)))
if err > tol:
# scale the time step and try again
h = hstep * max(safety * (tol / err)**0.25, 0.1)
else:
# scale the time step and update the position
t += h
err = max(err, tol * 1e-5)
h *= min(safety * (tol / err)**0.2, 5.)
for i in range(ntraj):
if wfn.traj[i].active:
wfn.traj[i].update_x(wfn.traj[i].x() + dx_lo[i])
wfn.traj[i].update_p(wfn.traj[i].p() + dp_lo[i])
if propphase:
wfn.traj[i].update_phase(wfn.traj[i].phase() +
dg_lo[i])
evaluate.update_pes(wfn, update_integrals=(abs(t) >= abs(dt)))