def spawn_backward(child, current_time, end_time, dt):
"""Propagates the child backwards in time until the current time
is reached."""
nstep = int(round( np.abs((current_time-end_time) / dt) ))
back_time = current_time
for i in range(nstep):
step.fms_step_trajectory(child, back_time, dt)
back_time = back_time + dt
fileio.print_fms_logfile('spawn_back', [back_time])
def init_bundle(master):
"""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
glbl.pes.init_interface()
# now load the initial trajectories into the bundle
if glbl.sampling['restart']:
init_restart(master)
else:
# first generate the initial nuclear coordinates and momenta
# and add the resulting trajectories to the bundle
glbl.distrib.set_initial_coords(master)
# set the initial state of the trajectories in bundle. This may
# require evaluation of electronic structure
set_initial_state(master)
# set the initial amplitudes of the basis functions
set_initial_amplitudes(master)
# add virtual basis functions, if desired (i.e. virtual basis = true)
if glbl.sampling['virtual_basis'] and not glbl.sampling['restart']:
virtual_basis(master)
# update all pes info for all trajectories and centroids (where necessary)
surface.update_pes(master)
# compute the hamiltonian matrix...
master.update_matrices()
# so that we may appropriately renormalize to unity
master.renormalize()
# this is the bundle at time t=0. Save in order to compute auto
# correlation function
set_initial_bundle(master)
# write to the log files
if glbl.mpi['rank'] == 0:
master.update_logs()
fileio.print_fms_logfile(
't_step',
[master.time, glbl.propagate['default_time_step'], master.nalive])
return master.time
def set_initial_amplitudes(master):
"""Sets the initial amplitudes."""
# if init_amp_overlap is set, overwrite 'amplitudes' that was
# set in nomad.input
if glbl.nuclear_basis['init_amp_overlap']:
origin = make_origin_traj()
# update all pes info for all trajectories and centroids (where necessary)
if glbl.integrals.overlap_requires_pes:
surface.update_pes(master)
# Calculate the initial expansion coefficients via projection onto
# the initial wavefunction that we are sampling
ovec = np.zeros(master.n_traj(), dtype=complex)
for i in range(master.n_traj()):
ovec[i] = glbl.integrals.traj_overlap(master.traj[i],
origin,
nuc_only=True)
smat = np.zeros((master.n_traj(), master.n_traj()), dtype=complex)
for i in range(master.n_traj()):
for j in range(i + 1):
smat[i,
j] = glbl.integrals.traj_overlap(master.traj[i],
master.traj[j])
if i != j:
smat[j, i] = smat[i, j].conjugate()
sinv = sp_linalg.pinvh(smat)
glbl.nuclear_basis['amplitudes'] = np.dot(sinv, ovec)
# if we don't have a sufficient number of amplitudes, append
# amplitudes with "zeros" as necesary
if len(glbl.nuclear_basis['amplitudes']) < master.n_traj():
dif = master.n_traj() - len(glbl.nuclear_basis['amplitudes'])
fileio.print_fms_logfile(
'warning',
['appending ' + str(dif) + ' values of 0+0j to amplitudes'])
glbl.nuclear_basis['amplitudes'].extend([0 + 0j for i in range(dif)])
# finally -- update amplitudes in the bundle
for i in range(master.n_traj()):
master.traj[i].update_amplitude(glbl.nuclear_basis['amplitudes'][i])
return
def fms_step_trajectory(traj, init_time, dt):
"""Propagates a single trajectory.
Used to backward/forward propagate a trajectory during spawning.
NOTE: fms_step_bundle and fms_step_trajectory could/should probably
be integrated somehow...
"""
current_time = init_time
end_time = init_time + dt
time_step = dt
min_time_step = abs(dt / 2.**5)
while not step_complete(current_time, end_time, time_step):
# save the bundle from previous step in case step rejected
traj0 = traj.copy()
# propagate single trajectory
glbl.integrator.propagate_trajectory(traj, time_step)
# update current time
proposed_time = current_time + time_step
# check time_step is fine, energy/amplitude conserved
accept = check_step_trajectory(traj0, traj)
# if everything is ok..
if accept:
current_time = proposed_time
else:
# redo time step
# recall -- this time trying to propagate
# to the failed step
time_step *= 0.5
if abs(time_step) < min_time_step:
fileio.print_fms_logfile(
'general', ['minimum time step exceeded -- STOPPING.'])
raise ValueError('Trajectory minimum step exceeded.')
# reset the beginning of the time step and go to beginning of loop
traj = traj0.copy()
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 set_initial_coords(master):
"""Samples a v=0 Wigner distribution
"""
# Set the coordinate type: Cartesian or normal mode coordinates
if glbl.interface['interface'] == 'vibronic':
coordtype = 'normal'
ham = glbl.pes.ham
else:
coordtype = 'cart'
# if multiple geometries in geometry.dat -- just take the first one
x_ref = np.array(glbl.nuclear_basis['geometries'][0], dtype=float)
p_ref = np.array(glbl.nuclear_basis['momenta'][0], dtype=float)
w_vec = np.array(glbl.nuclear_basis['widths'], dtype=float)
m_vec = np.array(glbl.nuclear_basis['masses'], dtype=float)
ndim = len(x_ref)
# create template trajectory basis function
template = trajectory.Trajectory(glbl.propagate['n_states'],
ndim,
width=w_vec,
mass=m_vec,
parent=0)
template.update_x(x_ref)
template.update_p(p_ref)
# If Cartesian coordinates are being used, then set up the
# mass-weighted Hessian and diagonalise to obtain the normal modes
# and frequencies
if coordtype == 'cart':
hessian = np.array(glbl.nuclear_basis['hessian'], dtype=float)
invmass = np.asarray([
1. / np.sqrt(m_vec[i]) if m_vec[i] != 0. else 0
for i in range(len(m_vec))
],
dtype=float)
mw_hess = invmass * hessian * invmass[:, np.newaxis]
evals, evecs = sp_linalg.eigh(mw_hess)
f_cutoff = 0.0001
freq_list = []
mode_list = []
for i in range(len(evals)):
if evals[i] >= 0 and np.sqrt(evals[i]) >= f_cutoff:
freq_list.append(np.sqrt(evals[i]))
mode_list.append(evecs[:, i].tolist())
n_modes = len(freq_list)
freqs = np.asarray(freq_list)
modes = np.asarray(mode_list).transpose()
# confirm that modes * tr(modes) = 1
m_chk = np.dot(modes, np.transpose(modes))
fileio.print_fms_logfile('string',
['\n -- frequencies from hessian.dat --\n'])
# If normal modes are being used, set the no. modes
# equal to the total number of modes of the model
# Hamiltonian and load only the active frequencies
if coordtype == 'normal':
n_modes = len(w_vec)
# we multiply by 0.5 below -- multiply by 2 here (i.e. default width for
# vibronic hamiltonans is 1/2, assuming frequency weighted coords
freqs = 2. * w_vec
fileio.print_fms_logfile(
'string', ['\n -- widths employed in coordinate sampling --\n'])
# write out frequencies
fstr = '\n'.join([
'{0:.5f} au {1:10.1f} cm^-1'.format(
freqs[j], freqs[j] * glbl.constants['au2cm'])
for j in range(n_modes)
])
fileio.print_fms_logfile('string', [fstr + '\n'])
# loop over the number of initial trajectories
max_try = 1000
ntraj = glbl.sampling['n_init_traj']
for i in range(ntraj):
delta_x = np.zeros(n_modes)
delta_p = np.zeros(n_modes)
x_sample = np.zeros(n_modes)
p_sample = np.zeros(n_modes)
for j in range(n_modes):
alpha = 0.5 * freqs[j]
if alpha > glbl.constants['fpzero']:
sigma_x = (glbl.sampling['distrib_compression'] *
np.sqrt(0.25 / alpha))
sigma_p = (glbl.sampling['distrib_compression'] *
np.sqrt(alpha))
itry = 0
while itry <= max_try:
dx = random.gauss(0., sigma_x)
dp = random.gauss(0., sigma_p)
itry += 1
if mode_overlap(alpha, dx,
dp) > glbl.sampling['init_mode_min_olap']:
break
if mode_overlap(alpha, dx,
dp) < glbl.sampling['init_mode_min_olap']:
print('Cannot get mode overlap > ' +
str(glbl.sampling['init_mode_min_olap']) +
' within ' + str(max_try) + ' attempts. Exiting...')
delta_x[j] = dx
delta_p[j] = dp
# If Cartesian coordinates are being used, displace along each
# normal mode to generate the final geometry...
if coordtype == 'cart':
disp_x = np.dot(modes, delta_x) / np.sqrt(m_vec)
disp_p = np.dot(modes, delta_p) / np.sqrt(m_vec)
# ... else if mass- and frequency-scaled normal modes are
# being used, then take the frequency-scaled normal mode
# displacements and momenta as the inital point in phase
# space
elif coordtype == 'normal':
disp_x = delta_x * np.sqrt(freqs)
disp_p = delta_p * np.sqrt(freqs)
x_sample = x_ref + disp_x
p_sample = p_ref + disp_p
# add new trajectory to the bundle
new_traj = template.copy()
new_traj.update_x(x_sample)
new_traj.update_p(p_sample)
# Add the trajectory to the bundle
master.add_trajectory(new_traj)
def fms_step_bundle(master, dt):
"""Propagates the wave packet using a run-time selected propagator."""
# save the bundle from previous step in case step rejected
end_time = master.time + dt
time_step = dt
min_time_step = dt / 2.**5
while not step_complete(master.time, end_time, dt):
# save the bundle from previous step in case step rejected
#try:
# del master0
#except NameError:
# pass
master0 = master.copy()
# propagate each trajectory in the bundle
time_step = min(time_step, end_time - master.time)
# propagate amplitudes for 1/2 time step using x0
master.update_amplitudes(0.5 * dt, update_ham=False)
# the propagators update the potential energy surface as need be.
glbl.integrator.propagate_bundle(master, time_step)
# propagate amplitudes for 1/2 time step using x1
master.update_amplitudes(0.5 * dt)
# Renormalization
if glbl.propagate['renorm'] == 1:
master.renormalize()
# check time_step is fine, energy/amplitude conserved
accept, error_msg = check_step_bundle(master0, master, time_step)
# if everything is ok..
if accept:
# update the bundle time
master.time += time_step
# spawn new basis functions if necessary
basis_grown = glbl.spawn.spawn(master, time_step)
# kill the dead trajectories
basis_pruned = master.prune()
# 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.integrals.require_centroids:
surface.update_pes(master)
# update the Hamiltonian and associated matrices
if basis_grown or basis_pruned:
master.update_matrices()
# re-expression of the basis using the matching pursuit
# algorithm
if glbl.propagate['matching_pursuit'] == 1:
mp.reexpress_basis(master)
# update the running log
fileio.print_fms_logfile('t_step',
[master.time, time_step, master.nalive])
else:
# recall -- this time trying to propagate to the failed step
time_step *= 0.5
fileio.print_fms_logfile('new_step', [error_msg, time_step])
if time_step < min_time_step:
fileio.print_fms_logfile(
'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
master = master0.copy()
return master
def init_interface():
"""Reads geometry.dat, freq.dat and the operator file.
Note that the order of modes is determined by geometry.dat and
the active modes are determined from the freq.dat file.
As such, we must read the labels in geometry.dat followed by
freq.dat file BEFORE reading the operator file.
"""
global kecoeff, ham
# Read in geometry labels, frequency and operator files
ham = VibHam()
# if 'geometry.dat' present, read info from there, else, from inputfile
if glbl.nuclear_basis['geomfile'] != '':
ham.rdgeomfile(fileio.home_path + '/geometry.dat')
else:
ham.nmode_total = len(glbl.nuclear_basis['geometries'][0])
ham.mlbl_total = glbl.nuclear_basis['labels']
# I propose discontinuing 'freq.dat' file. This can be entered in
# input file. Need way to differentiate between active/inactive modes
# I presume
#ham.rdfreqfile(fileio.home_path + '/freq.dat')
ham.nmode_active = len(glbl.nuclear_basis['freqs'])
ham.mlbl_active = ham.mlbl_total
ham.freq = np.array(glbl.nuclear_basis['freqs'])
for i in range(len(ham.freq)):
ham.freqmap[ham.mlbl_active[i]] = ham.freq[i]
# operator file will always be a separate file
ham.rdoperfile(fileio.home_path + '/' + glbl.interface['opfile'])
# KE operator coefficients, mass- and frequency-scaled normal mode
# coordinates, a_i = 0.5*omega_i
kecoeff = np.zeros(ham.nmode_total)
kecoeff[ham.mrange] = 0.5 * ham.freq
# Ouput some information about the Hamiltonian
fileio.print_fms_logfile('string', ['*' * 72])
fileio.print_fms_logfile('string',
['* Vibronic Coupling Hamiltonian Information'])
fileio.print_fms_logfile('string', ['*' * 72])
fileio.print_fms_logfile('string',
['Operator file: ' + glbl.interface['opfile']])
fileio.print_fms_logfile(
'string', ['Number of Hamiltonian terms: ' + str(ham.nterms)])
string = 'Total no. modes: ' + str(ham.nmode_total)
fileio.print_fms_logfile('string', [string])
string = 'No. active modes: ' + str(ham.nmode_active)
fileio.print_fms_logfile('string', [string])
fileio.print_fms_logfile('string', ['Active mode labels:'])
for i in range(ham.nmode_active):
string = str(i + 1) + ' ' + ham.mlbl_active[i]
fileio.print_fms_logfile('string', [string])
fileio.print_fms_logfile('string', ['Active mode frequencies (a.u.):'])
for i in range(ham.nmode_active):
string = str(i + 1) + ' ' + str(ham.freq[i])
fileio.print_fms_logfile('string', [string])
def spawn(master, dt):
"""Spawns a basis function if the minimum overlap drops below a given
threshold."""
basis_grown = False
current_time = master.time
for i in range(master.n_traj()):
if not master.traj[i].alive:
continue
parent = master.traj[i]
for st in range(master.nstates):
# don't check overlap with basis functions on same state
if st == parent.state:
continue
s_array = [
abs(
glbl.integrals.traj_overlap(
parent, master.traj[j], nuc_only=True))
if master.traj[j].state == st and master.traj[j].alive else 0.
for j in range(master.n_traj())
]
if max(s_array, key=abs) < glbl.spawning['continuous_min_overlap']:
child = parent.copy()
child.amplitude = 0j
child.state = st
child.parent = parent.label
success = utilities.adjust_child(
parent, child, parent.nact(parent.state, child.state))
sij = glbl.integrals.traj_overlap(parent, child, nuc_only=True)
# try to set up the child
if not success:
pass
# fileio.print_fms_logfile('spawn_bad_step',
# ['cannot adjust kinetic energy of child'])
elif abs(sij) < glbl.spawning['spawn_olap_thresh']:
pass
# fileio.print_fms_logfile('spawn_bad_step',
# ['child-parent overlap too small'])
else:
child_created = True
spawn_time = current_time
parent.last_spawn[child.state] = spawn_time
child.last_spawn[parent.state] = spawn_time
bundle_overlap = utilities.overlap_with_bundle(
child, master)
if not bundle_overlap:
basis_grown = True
master.add_trajectory(child)
fileio.print_fms_logfile(
'spawn_success', [current_time, parent.label, st])
utilities.write_spawn_log(current_time, current_time,
current_time, parent,
master.traj[-1])
else:
err_msg = ('Traj ' + str(parent.label) +
' from state ' + str(parent.state) +
' to state ' + str(st) + ': ' +
'overlap with bundle too large,' +
' s_max=' +
str(glbl.propagate['sij_thresh']))
fileio.print_fms_logfile('spawn_bad_step', [err_msg])
return basis_grown
def spawn_forward(parent, child_state, initial_time, dt):
"""Propagates the parent forward (into the future) until the coupling
decreases."""
parent_state = parent.state
current_time = initial_time
spawn_time = initial_time
exit_time = initial_time
child_created = False
parent_at_spawn = None
child_at_spawn = None
coup = np.zeros(3)
fileio.print_fms_logfile('spawn_start',
[parent.label, parent_state, child_state])
while True:
coup = np.roll(coup,1)
coup[0] = abs(parent.eff_coup(child_state))
child_attempt = parent.copy()
child_attempt.state = child_state
adjust_success = utils.adjust_child(parent, child_attempt,
parent.derivative(parent_state,
child_state))
sij = abs(glbl.integrals.traj_overlap(parent, child_attempt, nuc_only=True))
# if the coupling has already peaked, either we exit with a successful
# spawn from previous step, or we exit with a fail
if np.all(coup[0] < coup[1:]):
sp_str = 'no [decreasing coupling]'
fileio.print_fms_logfile('spawn_step',
[current_time, coup[0], sij, sp_str])
if child_created:
fileio.print_fms_logfile('spawn_success', [spawn_time])
child_at_spawn.exit_time[parent_state] = current_time
else:
fileio.print_fms_logfile('spawn_failure', [current_time])
# exit, we're done trying to spawn
exit_time = current_time
break
# coupling still increasing
else:
# try to set up the child
if not adjust_success:
sp_str = 'no [momentum adjust fail]'
elif sij < glbl.spawning['spawn_olap_thresh']:
sp_str = 'no [overlap too small]'
elif not np.all(coup[0] > coup[1:]):
sp_str = 'no [decreasing coupling]'
else:
spawn_time = current_time
child_created = True
parent_at_spawn = parent.copy()
child_at_spawn = child_attempt.copy()
child_at_spawn.last_spawn[parent_state] = spawn_time
child_at_spawn.amplitude = 0j
child_at_spawn.parent = parent.label
child_at_spawn.label = parent.label
sp_str = 'yes'
fileio.print_fms_logfile('spawn_step',
[current_time, coup[0], sij, sp_str])
step.fms_step_trajectory(parent, current_time, dt)
current_time = current_time + dt
return child_created, child_at_spawn, parent_at_spawn, spawn_time, exit_time
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