def propagate_trajectory(traj, dt):
"""Propagates a single trajectory with VV."""
# position update
propagate_position(traj, dt)
# update electronic structure
surface.update_pes_traj(traj)
# momentum/phase update
propagate_momentum(traj, dt)
def propagate_trajectory(traj, dt):
"""Propagates a single trajectory with RKF45."""
global h_traj
ncrd = traj.dim
kx = np.zeros((rk_ordr, ncrd))
kp = np.zeros((rk_ordr, ncrd))
kg = np.zeros((rk_ordr))
t = 0.
if h_traj is None:
h_traj = dt
while abs(t) < abs(dt):
hstep = np.sign(dt) * min(abs(h_traj), abs(dt - t))
for rk in range(rk_ordr):
tmptraj = traj.copy()
propagate_rk(tmptraj, hstep, rk, kx, kp, kg)
# update the PES to evaluate new gradients
if rk < rk_ordr - 1:
surface.update_pes_traj(tmptraj)
# calculate the 4th and 5th order changes and the error
dx_lo = np.sum(wgt_lo[:, np.newaxis] * kx, axis=0)
dx_hi = np.sum(wgt_hi[:, np.newaxis] * kx, axis=0)
dp_lo = np.sum(wgt_lo[:, np.newaxis] * kp, axis=0)
dp_hi = np.sum(wgt_hi[:, np.newaxis] * kp, axis=0)
if propphase:
dg_lo = np.sum(wgt_lo * kg)
dg_hi = np.sum(wgt_hi * kg)
# err = np.max((np.abs(dx_hi-dx_lo), np.abs(dp_hi-dp_lo),
# np.abs(dg_hi-dg_lo)))
err = 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_traj = hstep * max(safety * (tol / err)**0.25, 0.1)
else:
# scale the time step and update the position
t += h_traj
err = max(err, tol * 1e-5)
h_traj *= min(safety * (tol / err)**0.2, 5.)
traj.update_x(traj.x() + dx_lo)
traj.update_p(traj.p() + dp_lo)
if propphase:
traj.update_phase(traj.phase() + dg_lo)
surface.update_pes_traj(traj)
def propagate_trajectory(traj, dt):
"""Propagates a single trajectory with RK4."""
ncrd = traj.dim
kx = np.zeros((rk_ordr, ncrd))
kp = np.zeros((rk_ordr, ncrd))
kg = np.zeros((rk_ordr, ncrd))
for rk in range(rk_ordr):
tmptraj = traj.copy()
propagate_rk(tmptraj, dt, rk, kx, kp, kg)
# update the PES to evaluate new gradients
if rk < rk_ordr - 1:
surface.update_pes_traj(tmptraj)
# update to the final position
traj.update_x(traj.x() + np.sum(wgt[:,np.newaxis]*kx, axis=0))
traj.update_p(traj.p() + np.sum(wgt[:,np.newaxis]*kp, axis=0))
if propphase:
traj.update_phase(traj.phase() + np.sum(wgt[:,np.newaxis]*kg, axis=0))
surface.update_pes_traj(traj)
def set_initial_state(master):
"""Sets the initial state of the trajectories in the bundle."""
# initialize to the state with largest transition dipole moment
if glbl.sampling['init_brightest']:
# set all states to the ground state
for i in range(master.n_traj()):
master.traj[i].state = 0
# compute transition dipoles
surface.update_pes_traj(master.traj[i])
# set the initial state to the one with largest t. dip.
for i in range(master.n_traj()):
if 'tr_dipole' not in master.traj[i].pes_data.data_keys:
raise KeyError('ERROR, trajectory ' + str(i) +
': Cannot set state by transition moments - ' +
'tr_dipole not in pes_data.data_keys')
tdip = np.array([
np.linalg.norm(master.traj[i].pes_data.dipoles[:, 0, j])
for j in range(1, glbl.propagate['n_states'])
])
fileio.print_fms_logfile('general',
['Initializing trajectory '+str(i)+
' to state '+str(np.argmax(tdip)+1)+
' | tr. dipople array='+np.array2string(tdip, \
formatter={'float_kind':lambda x: "%.4f" % x})])
master.traj[i].state = np.argmax(tdip) + 1
# use "init_state" to set the initial state
elif len(glbl.sampling['init_states']) == master.n_traj():
for i in range(master.n_traj()):
master.traj[i].state = glbl.sampling['init_states'][i]
else:
raise ValueError('Ambiguous initial state assignment.')
return
def make_origin_traj():
"""Construct a trajectory basis function at the origin
specified in the input files"""
ndim = len(glbl.nuclear_basis['geometries'][0])
m_vec = np.array(glbl.nuclear_basis['masses'])
w_vec = np.array(glbl.nuclear_basis['widths'])
x_vec = np.array(glbl.nuclear_basis['geometries'][0])
p_vec = np.array(glbl.nuclear_basis['momenta'][0])
origin = trajectory.Trajectory(glbl.propagate['n_states'],
ndim,
width=w_vec,
mass=m_vec,
parent=0)
origin.update_x(x_vec)
origin.update_p(p_vec)
origin.state = 0
# if we need pes data to evaluate overlaps, determine that now
if glbl.integrals.overlap_requires_pes:
surface.update_pes_traj(origin)
return origin
def propagate_trajectory(traj, dt):
"""Propagates a single trajectory with BS."""
global h_traj
ncrd = traj.dim
t = 0.
if h_traj is None:
h_traj = dt
while abs(t) < abs(dt):
hstep = np.sign(dt) * min(abs(h_traj), abs(dt - t))
reduced = False
tsav = np.zeros(kmax_traj)
err = np.zeros(kmax_traj)
Tx = np.zeros((kmax_traj, ncrd))
Tp = np.zeros((kmax_traj, ncrd))
Tg = np.zeros((kmax_traj, ncrd))
for k in range(kmax):
tmptraj = traj.copy()
x0 = np.zeros(ncrd)
p0 = np.zeros(ncrd)
g0 = np.zeros(ncrd)
x1 = np.zeros(ncrd)
p1 = np.zeros(ncrd)
g1 = np.zeros(ncrd)
# step through n modified midpoint steps
for n in range(nstep[k]):
mm_step(tmptraj, hstep/nstep[k], x0, x1, p0, p1, g0, g1, n)
surface.update_pes_traj(tmptraj)
# compute the modified midpoint estimate
x1 = 0.5*(x1 + x0 + hstep/nstep[k]*tmptraj.velocity())
p1 = 0.5*(p1 + p0 + hstep/nstep[k]*tmptraj.force())
if propphase:
g1 = 0.5*(g1 + g0 + hstep/nstep[k]*tmptraj.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(red, 0.7))
h_traj = sfac * hstep
break
else:
# scale the time step if possible
t += h_traj
h_traj = increase_tstep(k, err, reduced) * hstep
# update to the final position
xnew = np.sum(Tx, axis=0)
pnew = np.sum(Tp, axis=0)
traj.update_x(xnew)
traj.update_p(pnew)
if propphase:
gnew = np.sum(Tg, axis=0)
traj.update_phase(gnew)
surface.update_pes_traj(traj)
break
def spawn(master, dt):
"""Propagates to the point of maximum coupling, spawns a new
basis function, then propagates the function to the current time."""
global coup_hist
basis_grown = False
current_time = master.time
# list of added trajectories
# we want to know the history of the coupling for each trajectory
# in order to assess spawning criteria -- make sure coup_hist has a slot
# for every trajectory
if len(coup_hist) < master.n_traj():
n_add = master.n_traj() - len(coup_hist)
for i in range(n_add):
coup_hist.append(np.zeros((master.nstates, 3)))
#--------------- iterate over all trajectories in bundle ---------------------
for i in range(master.n_traj()):
# only live trajectories can spawn
if not master.traj[i].alive:
continue
for st in range(master.nstates):
# can only spawn to different electronic states
if master.traj[i].state == st:
continue
# compute magnitude of coupling to state j
coup = master.traj[i].eff_coup(st)
coup_hist[i][st, :] = np.roll(coup_hist[i][st, :], 1)
coup_hist[i][st, 0] = coup
# if we satisfy spawning conditions, begin spawn process
if spawn_trajectory(master, i, st, coup_hist[i][st, :],
current_time):
# we're going to messing with this trajectory -- mess with a copy
parent = master.traj[i].copy()
# propagate the parent forward in time until coupling maximized
[success, child, parent_spawn, spawn_time,
exit_time] = spawn_forward(parent, st, current_time, dt)
# set the spawn attempt in master, even if spawn failed (avoid repeated fails)
master.traj[i].last_spawn[st] = spawn_time
master.traj[i].exit_time[st] = exit_time
if success:
# at this point, child is at the spawn point. Propagate
# backwards in time until we reach the current time
child_spawn = child.copy()
# need electronic structure at current geometry -- on correct state
surface.update_pes_traj(child)
spawn_backward(child, spawn_time, current_time, -dt)
bundle_overlap = utils.overlap_with_bundle(child, master)
if not bundle_overlap:
basis_grown = True
master.add_trajectory(child)
child_spawn.label = master.traj[
-1].label # a little hacky...
utils.write_spawn_log(current_time, spawn_time,
exit_time, parent_spawn,
child_spawn)
else:
fileio.print_fms_logfile(
'spawn_bad_step',
['overlap with bundle too large'])
# let caller known if the basis has been changed
return basis_grown