def __init__(self,
geometry,
step_length=0.1,
max_cycles=75,
downhill=False,
forward=True,
backward=True,
mode=0,
hessian_init=None,
displ="energy",
displ_energy=5e-4,
displ_length=0.1,
rms_grad_thresh=3e-3,
energy_thresh=1e-6,
force_inflection=True,
dump_fn="irc_data.h5",
dump_every=5):
assert (step_length > 0), "step_length must be positive"
assert (max_cycles > 0), "max_cycles must be positive"
self.logger = logging.getLogger("irc")
self.geometry = geometry
assert self.geometry.coord_type == "cart"
self.step_length = step_length
self.max_cycles = max_cycles
self.downhill = downhill
# Disable forward/backward when downhill is set
self.forward = not self.downhill and forward
self.backward = not self.downhill and backward
self.mode = mode
if hessian_init is None:
hessian_init = "calc" if not self.downhill else "unit"
self.hessian_init = hessian_init
self.displ = displ
assert self.displ in ("energy", "length"), \
"displ must be either 'energy' or 'length'"
self.displ_energy = float(displ_energy)
self.displ_length = float(displ_length)
self.rms_grad_thresh = float(rms_grad_thresh)
self.energy_thresh = float(energy_thresh)
self.force_inflection = force_inflection
self.dump_fn = dump_fn
self.dump_every = int(dump_every)
self._m_sqrt = np.sqrt(self.geometry.masses_rep)
self.all_energies = list()
self.all_coords = list()
self.all_gradients = list()
self.all_mw_coords = list()
self.all_mw_gradients = list()
# step length dE max(|grad|) rms(grad)
col_fmts = "int float float float float".split()
header = ("Step", "IRC length", "dE / au", "max(|grad|)", "rms(grad)")
self.table = TablePrinter(header, col_fmts)
self.cycle_places = ceil(log(self.max_cycles, 10))
def do_analysis(geom, nu_thresh=-5.):
# Do calculations
hess = geom.hessian
gradient = geom.gradient
energy = geom.energy
# Calculate hessian
hessian = geom.hessian
mw_hess = geom.mw_hessian
proj_hess = geom.eckart_projection(mw_hess)
w, v = np.linalg.eigh(proj_hess)
nus = eigval_to_wavenumber(w)
neg_inds = nus <= nu_thresh
grad_norm = np.linalg.norm(gradient)
print(f"norm(grad)={grad_norm:.6f}")
print()
# Overlaps between gradient and imaginary modes. We use the absolute
# values, as the sign of the eigenvectors is ambiguous.
overlaps = np.einsum("ij,i->j", v[:, neg_inds], gradient)
print(f"Found {neg_inds.sum()} imaginary frequencies < {nu_thresh} cm⁻¹.")
tp = TablePrinter("# nu/cm⁻¹ <grad|qi>".split(), "int float float".split())
tp.print_header()
for i, (nu, ovlp) in enumerate(zip(nus[neg_inds], overlaps)):
tp.print_row((i, nu, ovlp))
res = NMResult(
energy=energy,
gradient=gradient,
hessian=hess,
w=w,
v=v,
nus=nus,
neg_inds=neg_inds,
neg_nus=nus[neg_inds],
)
return res
class IRC:
def __init__(self,
geometry,
step_length=0.1,
max_cycles=75,
downhill=False,
forward=True,
backward=True,
mode=0,
hessian_init=None,
displ="energy",
displ_energy=5e-4,
displ_length=0.1,
rms_grad_thresh=3e-3,
energy_thresh=1e-6,
force_inflection=True,
dump_fn="irc_data.h5",
dump_every=5):
assert (step_length > 0), "step_length must be positive"
assert (max_cycles > 0), "max_cycles must be positive"
self.logger = logging.getLogger("irc")
self.geometry = geometry
assert self.geometry.coord_type == "cart"
self.step_length = step_length
self.max_cycles = max_cycles
self.downhill = downhill
# Disable forward/backward when downhill is set
self.forward = not self.downhill and forward
self.backward = not self.downhill and backward
self.mode = mode
if hessian_init is None:
hessian_init = "calc" if not self.downhill else "unit"
self.hessian_init = hessian_init
self.displ = displ
assert self.displ in ("energy", "length"), \
"displ must be either 'energy' or 'length'"
self.displ_energy = float(displ_energy)
self.displ_length = float(displ_length)
self.rms_grad_thresh = float(rms_grad_thresh)
self.energy_thresh = float(energy_thresh)
self.force_inflection = force_inflection
self.dump_fn = dump_fn
self.dump_every = int(dump_every)
self._m_sqrt = np.sqrt(self.geometry.masses_rep)
self.all_energies = list()
self.all_coords = list()
self.all_gradients = list()
self.all_mw_coords = list()
self.all_mw_gradients = list()
# step length dE max(|grad|) rms(grad)
col_fmts = "int float float float float".split()
header = ("Step", "IRC length", "dE / au", "max(|grad|)", "rms(grad)")
self.table = TablePrinter(header, col_fmts)
self.cur_cycle = 0
self.converged = False
self.cur_direction = None
self.cycle_places = ceil(log(self.max_cycles, 10))
@property
def coords(self):
return self.geometry.coords
@coords.setter
def coords(self, coords):
self.geometry.coords = coords
@property
def mw_coords(self):
return self.geometry.mw_coords
@mw_coords.setter
def mw_coords(self, mw_coords):
self.geometry.mw_coords = mw_coords
@property
def energy(self):
return self.geometry.energy
@property
def gradient(self):
return self.geometry.gradient
@property
def mw_gradient(self):
return self.geometry.mw_gradient
# @property
# def mw_hessian(self):
# # TODO: This can be removed when the mw_hessian property is updated
# # in Geometry.py.
# return self.geometry.mw_hessian
def log(self, msg):
# self.logger.debug(f"step {self.cur_cycle:03d}, {msg}")
self.logger.debug(msg)
@property
def m_sqrt(self):
return self._m_sqrt
def unweight_vec(self, vec):
return self.m_sqrt * vec
def mass_weigh_hessian(self, hessian):
return self.geometry.mass_weigh_hessian(hessian)
def prepare(self, direction):
self.cur_cycle = 0
self.direction = direction
self.converged = False
self.past_inflection = not self.force_inflection
self.irc_energies = list()
# Not mass-weighted
self.irc_coords = list()
self.irc_gradients = list()
# Mass-weighted
self.irc_mw_coords = list()
self.irc_mw_gradients = list()
# Over the course of the IRC the hessian may get updated.
# Copying the initial hessian here ensures a clean start in combined
# forward and backward runs. Otherwise we may accidentally use
# the updated hessian from the end of the first run for the second
# run.
self.mw_hessian = self.mass_weigh_hessian(self.init_hessian)
# We don't need an initial displacement when going downhill
if self.downhill:
return
# Do inital displacement from the TS
init_factor = 1 if (direction == "forward") else -1
initial_step = init_factor * self.init_displ
self.coords = self.ts_coords + initial_step
initial_step_length = np.linalg.norm(initial_step)
self.logger.info(
f"Did inital step of length {initial_step_length:.4f} "
"from the TS.")
def initial_displacement(self):
"""Returns a non-mass-weighted step in angstrom for an initial
displacement from the TS along the transition vector.
See https://aip.scitation.org/doi/pdf/10.1063/1.454172?class=pdf
"""
mm_sqr_inv = self.geometry.mm_sqrt_inv
mw_hessian = self.mass_weigh_hessian(self.init_hessian)
try:
if not self.geometry.calculator.analytical_2d:
mw_hessian = self.geometry.eckart_projection(mw_hessian)
except AttributeError:
pass
eigvals, eigvecs = np.linalg.eigh(mw_hessian)
neg_inds = eigvals < -1e-8
assert sum(neg_inds
) > 0, "The hessian does not have any negative eigenvalues!"
min_eigval = eigvals[self.mode]
self.log(f"Transition vector is mode {self.mode} with wavenumber "
f"{eigval_to_wavenumber(min_eigval):.2f} cm⁻¹.")
mw_trans_vec = eigvecs[:, self.mode]
self.mw_transition_vector = mw_trans_vec
# Un-mass-weight the transition vector
trans_vec = mm_sqr_inv.dot(mw_trans_vec)
self.transition_vector = trans_vec / np.linalg.norm(trans_vec)
if self.downhill:
step = np.zeros_like(self.transition_vector)
elif self.displ == "length":
self.log("Using length-based initial displacement from the TS.")
step = self.displ_length * self.transition_vector
else:
# Calculate the length of the initial step away from the TS to initiate
# the IRC/MEP. We assume a quadratic potential and calculate the
# displacement for a given energy lowering.
# dE = (k*dq**2)/2 (dE = energy lowering, k = eigenvalue corresponding
# to the transition vector/imaginary mode, dq = step length)
# dq = sqrt(dE*2/k)
# See 10.1021/ja00295a002 and 10.1063/1.462674
# 10.1002/jcc.540080808 proposes 3 kcal/mol as initial energy lowering
self.log("Using energy-based initial displacement from the TS.")
step_length = np.sqrt(self.displ_energy * 2 / np.abs(min_eigval))
# This calculation is derived from the mass-weighted hessian, so we
# probably have to multiply this step length with the mass-weighted
# mode and un-weigh it.
mw_step = step_length * mw_trans_vec
step = mw_step / self.m_sqrt
print(f"Norm of initial displacement step: {np.linalg.norm(step):.4f}")
self.log("")
return step
def get_conv_fact(self, mw_grad, min_fact=2.):
# Numerical integration of differential equations requires a step length and/or
# we have to terminate the integration at some point, e.g. when the desired
# step length is reached. IRCs are integrated in mass-weighted coordinates,
# but self.step_length is given in unweighted coordinates. Unweighting a step
# in mass-weighted coordinates will reduce its norm as we divide by sqrt(m).
#
# If we want to do an Euler-integration we have to decide on a step size
# when a desired integration length is to be reached in a given number of steps.
# [3] proposes using Δs/250 with a maximum of 500 steps, so something like
# Δs/(max_steps / 2). It seems we can't use this because (at
# least for the systems I tested) this will lead to a step length that is too
# small, so the predictor Euler-integration will fail to converge in the
# prescribed number of cycles. It fails because simply dividing the desired
# step length in unweighted coordinates does not take into account the mass
# dependence. Such a step size is appropriate for integrations in unweighted
# coordinates, but not when using mass-weighted coordinates.
#
# We determine a conversion factor from comparing the magnitudes (norms) of
# the mass-weighted and un-mass-weighted gradients. This takes into account
# which atoms are actually moving, so it should be a good guess.
norm_mw_grad = np.linalg.norm(mw_grad)
norm_grad = np.linalg.norm(self.unweight_vec(mw_grad))
conv_fact = norm_grad / norm_mw_grad
conv_fact = max(min_fact, conv_fact)
self.log(
f"Un-weighted / mass-weighted conversion factor {conv_fact:.4f}")
return conv_fact
def irc(self, direction):
self.log(highlight_text(f"IRC {direction}", level=1))
self.cur_direction = direction
self.prepare(direction)
# Calculate gradient
self.gradient
self.irc_energies.append(self.energy)
# Non mass-weighted
self.irc_coords.append(self.coords)
self.irc_gradients.append(self.gradient)
# Mass-weighted
self.irc_mw_coords.append(self.mw_coords)
self.irc_mw_gradients.append(self.mw_gradient)
self.table.print_header()
while True:
self.log(highlight_text(f"IRC step {self.cur_cycle:03d}") + "\n")
if self.cur_cycle == self.max_cycles:
print("IRC steps exceeded. Stopping.")
print()
break
# Do macroiteration/IRC step to update the geometry
self.step()
# Calculate gradient and energy on the new geometry
# Non mass-weighted
self.log("Calculating energy and gradient at new geometry.")
self.irc_coords.append(self.coords)
self.irc_gradients.append(self.gradient)
self.irc_energies.append(self.energy)
# Mass-weighted
self.irc_mw_coords.append(self.mw_coords)
self.irc_mw_gradients.append(self.mw_gradient)
rms_grad = rms(self.gradient)
# Only update once
if not self.past_inflection:
self.past_inflection = rms_grad >= self.rms_grad_thresh
_ = "" if self.past_inflection else "not yet"
self.log(f"(rms(grad) > threshold) {_} fullfilled!")
irc_length = np.linalg.norm(self.irc_mw_coords[0] -
self.irc_mw_coords[-1])
dE = self.irc_energies[-1] - self.irc_energies[-2]
max_grad = np.abs(self.gradient).max()
row_args = (self.cur_cycle, irc_length, dE, max_grad, rms_grad)
self.table.print_row(row_args)
try:
# The derived IRC classes may want to do some printing
add_info = self.get_additional_print()
self.table.print(add_info)
except AttributeError:
pass
last_energy = self.irc_energies[-2]
this_energy = self.irc_energies[-1]
break_msg = ""
if self.converged:
break_msg = "Integrator indicated convergence!"
elif self.past_inflection and (rms_grad <= self.rms_grad_thresh):
break_msg = "rms(grad) converged!"
self.converged = True
# TODO: Allow some threshold?
elif this_energy > last_energy:
break_msg = "Energy increased!"
elif abs(last_energy - this_energy) <= self.energy_thresh:
break_msg = "Energy converged!"
self.converged = True
dumped = (self.cur_cycle % self.dump_every) == 0
if dumped:
dump_fn = f"{direction}_{self.dump_fn}"
self.dump_data(dump_fn)
if break_msg:
self.table.print(break_msg)
break
self.cur_cycle += 1
if check_for_stop_sign():
break
self.log("")
sys.stdout.flush()
if direction == "forward":
self.irc_energies.reverse()
self.irc_coords.reverse()
self.irc_gradients.reverse()
self.irc_mw_coords.reverse()
self.irc_mw_gradients.reverse()
if not dumped:
self.dump_data(dump_fn)
self.cur_direction = None
def set_data(self, prefix):
energies_name = f"{prefix}_energies"
coords_name = f"{prefix}_coords"
grad_name = f"{prefix}_gradients"
mw_coords_name = f"{prefix}_mw_coords"
mw_grad_name = f"{prefix}_mw_gradients"
setattr(self, coords_name, self.irc_coords)
setattr(self, grad_name, self.irc_gradients)
setattr(self, mw_coords_name, self.irc_mw_coords)
setattr(self, mw_grad_name, self.irc_mw_gradients)
setattr(self, energies_name, self.irc_energies)
self.all_energies.extend(getattr(self, energies_name))
self.all_coords.extend(getattr(self, coords_name))
self.all_gradients.extend(getattr(self, grad_name))
self.all_mw_coords.extend(getattr(self, mw_coords_name))
self.all_mw_gradients.extend(getattr(self, mw_grad_name))
setattr(self, f"{prefix}_is_converged", self.converged)
setattr(self, f"{prefix}_cycle", self.cur_cycle)
self.write_trj(".", prefix, getattr(self, mw_coords_name))
def run(self):
# Calculate data at TS and create backup
self.ts_coords = self.coords.copy()
self.ts_mw_coords = self.mw_coords.copy()
print("Calculating energy and gradient at the TS")
self.ts_gradient = self.gradient.copy()
self.ts_mw_gradient = self.mw_gradient.copy()
self.ts_energy = self.energy
ts_grad_norm = np.linalg.norm(self.ts_gradient)
ts_grad_max = np.abs(self.ts_gradient).max()
ts_grad_rms = rms(self.ts_gradient)
self.log("Transition state (TS):\n"
f"\tnorm(grad)={ts_grad_norm:.6f}\n"
f"\t max(grad)={ts_grad_max:.6f}\n"
f"\t rms(grad)={ts_grad_rms:.6f}")
print("IRC length in mw. coords, max(|grad|) and rms(grad) in "
"unweighted coordinates.")
self.init_hessian, hess_str = get_guess_hessian(
self.geometry,
self.hessian_init,
cart_gradient=self.ts_gradient,
h5_fn=f"hess_init_irc.h5",
)
self.log(f"Initial hessian: {hess_str}")
# For forward/backward runs from a TS we need an intial displacement,
# calculated from the transition vector (imaginary mode) of the TS
# hessian. If we need/want a Hessian for a downhill run from a
# non-stationary point (with non-vanishing gradient) depends on the
# actual IRC integrator (e.g. EulerPC and LQA need a Hessian).
if not self.downhill:
self.init_displ = self.initial_displacement()
if self.forward:
print("\n" + highlight_text("IRC - Forward") + "\n")
self.irc("forward")
self.set_data("forward")
# Add TS/starting data
self.all_energies.append(self.ts_energy)
self.all_coords.append(self.ts_coords)
self.all_gradients.append(self.ts_gradient)
self.all_mw_coords.append(self.ts_mw_coords)
self.all_mw_gradients.append(self.ts_mw_gradient)
self.ts_index = len(self.all_energies) - 1
if self.backward:
print("\n" + highlight_text("IRC - Backward") + "\n")
self.irc("backward")
self.set_data("backward")
if self.downhill:
print("\n" + highlight_text("IRC - Downhill") + "\n")
self.irc("downhill")
self.set_data("downhill")
self.all_mw_coords = np.array(self.all_mw_coords)
self.all_energies = np.array(self.all_energies)
self.postprocess()
if not self.downhill:
self.write_trj(".", "finished")
# Dump the whole IRC to HDF5
dump_fn = "finished_" + self.dump_fn
self.dump_data(dump_fn, full=True)
# Convert to arrays
[
setattr(self, name, np.array(getattr(self, name)))
for name in "all_energies all_coords all_gradients "
"all_mw_coords all_mw_gradients".split()
]
# Right now self.all_mw_coords is still in mass-weighted coordinates.
# Convert them to un-mass-weighted coordinates.
self.all_mw_coords_umw = self.all_mw_coords / self.m_sqrt
def postprocess(self):
pass
def write_trj(self, path, prefix, coords=None):
path = pathlib.Path(path)
atoms = self.geometry.atoms
if coords is None:
coords = self.all_mw_coords
coords = coords.copy()
coords /= self.m_sqrt
coords = coords.reshape(-1, len(atoms), 3) * BOHR2ANG
# all_mw_coords = self.all_mw_coords.flatten()
trj_string = make_trj_str(atoms, coords, comments=self.all_energies)
trj_fn = f"{prefix}_irc.trj"
with open(path / trj_fn, "w") as handle:
handle.write(trj_string)
first_coords = coords[0]
first_fn = f"{prefix}_first.xyz"
with open(path / first_fn, "w") as handle:
handle.write(make_xyz_str(atoms, first_coords))
last_coords = coords[-1]
first_fn = f"{prefix}_last.xyz"
with open(path / first_fn, "w") as handle:
handle.write(make_xyz_str(atoms, last_coords))
def get_irc_data(self):
data_dict = {
"energies": np.array(self.irc_energies, dtype=float),
"coords": np.array(self.irc_coords, dtype=float),
"gradients": np.array(self.irc_gradients, dtype=float),
"mw_coords": np.array(self.irc_mw_coords, dtype=float),
"mw_gradients": np.array(self.irc_mw_gradients, dtype=float),
}
return data_dict
def get_full_irc_data(self):
data_dict = {
"energies": np.array(self.all_energies, dtype=float),
"coords": np.array(self.all_coords, dtype=float),
"gradients": np.array(self.all_gradients, dtype=float),
"mw_coords": np.array(self.all_mw_coords, dtype=float),
"mw_gradients": np.array(self.all_mw_gradients, dtype=float),
"ts_index": np.array(self.ts_index, dtype=int),
}
return data_dict
def dump_data(self, dump_fn=None, full=False):
get_data = self.get_full_irc_data if full else self.get_irc_data
data_dict = get_data()
data_dict.update({
"atoms": np.array(self.geometry.atoms, dtype="S"),
"rms_grad_thresh": np.array(self.rms_grad_thresh),
})
if dump_fn is None:
dump_fn = self.dump_fn
with h5py.File(dump_fn, "w") as handle:
for key, val in data_dict.items():
handle.create_dataset(name=key, dtype=val.dtype, data=val)
def dimer_method(geoms, calc_getter, N_init=None,
max_step=0.1, max_cycles=50,
max_rots=10, interpolate=True,
rot_type="fourier",
rot_opt="lbfgs", trans_opt="lbfgs",
trans_memory=5,
trial_angle=5, angle_tol=0.5, dR_base=0.01,
restrict_step="scale", ana_2dpot=False,
f_thresh=1e-3, rot_f_thresh=2e-3,
zero_weights=[], dimer_pickle=None,
f_tran_mod=True,
multiple_translations=False, max_translations=10):
"""Dimer method using steepest descent for rotation and translation.
See
# Original paper
[1] https://doi.org/10.1063/1.480097
# Improved dimer
[2] https://doi.org/10.1063/1.2104507
# Several trial rotations
[3] https://doi.org/10.1063/1.1809574
# Superlinear dimer
[4] https://doi.org/10.1063/1.2815812
# Modified Broyden
[5] https://doi.org/10.1021/ct9005147
To add:
Comparing curvatures and adding π/2 if appropriate.
Default parameters from [1]
max_step = 0.1 bohr
dR_base = = 0.01 bohr
"""
# Parameters
rad_tol = np.deg2rad(angle_tol)
rms_f_thresh = float(f_thresh)
max_f_thresh = 1.5 * rms_f_thresh
max_translations = max_translations if multiple_translations else 1
opt_name_dict = {
"cg": "conjugate gradient",
"lbfgs": "L-BFGS",
"mb": "modified Broyden",
"sd": "steepst descent",
}
assert rot_opt in opt_name_dict.keys()
assert rot_type in "fourier direct".split()
# Translation using L-BFGS as described in [4]
# Modified broyden as proposed in [5] 10.1021/ct9005147
trans_closures = {
"lbfgs": closures.lbfgs_closure_,
"mb": closures.modified_broyden_closure,
}
print(f"Using {opt_name_dict[rot_opt]} for dimer rotation.")
print(f"Using {opt_name_dict[trans_opt]} for dimer translation.")
print(f"Keeping information of last {trans_memory} cycles for "
"translation optimization.")
f_tran_dict = {
True: get_f_tran_mod,
False: get_f_tran_org
}
get_f_tran = f_tran_dict[f_tran_mod]
rot_header = "Rot._cycle Curvature rms(f_rot)".split()
rot_fmts = "int float float".split()
rot_table = TablePrinter(rot_header, rot_fmts, width=16, shift_left=2)
trans_header = "Trans._cycle Curvature max(|f0|) rms(f0) rms(step)".split()
trans_fmts = "int float float float float".split()
trans_table = TablePrinter(trans_header, trans_fmts, width=16)
geom_getter = get_geom_getter(geoms[0], calc_getter)
assert len(geoms) in (1, 2), "geoms argument must be of length 1 or 2!"
# if dimer_pickle:
# with open(dimer_pickle, "rb") as handle:
# dimer_tuple = cloudpickle.load(handle)
if len(geoms) == 2:
geom1, geom2 = geoms
dR = np.linalg.norm(geom1.coords - geom2.coords) / 2
# N points from geom2 to geom1
N = make_unit_vec(geom1.coords, geom2.coords)
coords0 = geom2.coords + dR*N
geom0 = geom_getter(coords0)
# This handles cases where only one geometry is supplied. We use the
# given geometry as midpoint, and use the given N_init. If N_init is
# none, we select a random direction and derive geom1 and geom2 from it.
else:
geom0 = geoms[0]
coords0 = geom0.coords
# Assign random unit vector and use dR_base to for dR
coord_size = geom0.coords.size
if N_init is None:
N_init = np.random.rand(coord_size)
N = np.array(N_init)
if ana_2dpot:
N[2] = 0
N /= np.linalg.norm(N)
coords1 = geom0.coords + dR_base*N
geom1 = geom_getter(coords1)
coords2 = geom0.coords - dR_base*N
geom2 = geom_getter(coords2)
dR = dR_base
geom0.calculator.base_name += "_image0"
geom1.calculator.base_name += "_image1"
geom2.calculator.base_name += "_image2"
dimer_pickle = DimerPickle(coords0, N, dR)
with open("dimer_pickle", "wb") as handle:
cloudpickle.dump(dimer_pickle, handle)
dimer_cycles = list()
print("Using N:", N)
def f_tran_getter(coords, N, C):
# The force for the given coord should be already set.
np.testing.assert_allclose(geom0.coords, coords)
# Only return f_tran and drop the parallel and perpendicular component.
return get_f_tran(geom0.forces, N, C)[0]
def rot_force_getter(N, f1, f2):
return get_f_perp(f1, f2, N)
def restrict_max_step_comp(x, step, max_step=max_step):
step_max = np.abs(step).max()
if step_max > max_step:
factor = max_step / step_max
step *= factor
return step
def restrict_step_length(x, step, max_step_length=max_step):
step_norm = np.linalg.norm(step)
if step_norm > max_step_length:
step_direction = step / step_norm
step = step_direction * max_step_length
return step
rstr_dict = {
"max": restrict_max_step_comp,
"scale": restrict_step_length,
}
rstr_func = rstr_dict[restrict_step]
tot_rot_force_evals = 0
# Create the translation optimizer in the first cycle of the loop.
trans_opt_kwargs = {
"restrict_step": rstr_func,
"M": trans_memory,
# "beta": 0.01,
}
if trans_opt == "mb":
trans_opt_kwargs["beta"] = 0.01
trans_optimizer = trans_closures[trans_opt](f_tran_getter, **trans_opt_kwargs)
def cbm_rot_force_getter(coords1, N, f1, f2):
return get_f_perp(f1, f2, N)
converged = False
# In the first dimer cycle f0 and f1 aren't set and have to be calculated.
# For cycles i > 0 f0 and f1 will be already set here.
add_force_evals = 2
for i in range(max_cycles):
logger.debug(f"Dimer macro cycle {i:03d}")
print(highlight_text(f"Dimer Cycle {i:03d}"))
f0 = geom0.forces
f1 = geom1.forces
f2 = 2*f0 - f1
coords0 = geom0.coords
coords1 = geom1.coords
coords2 = geom2.coords
C = get_curvature(f1, f2, N, dR)
f0_rms = get_rms(f0)
f0_max = np.abs(f0).max()
trans_table.print(f"@ {i:03d}: C={C:.6f}; max(|f0|)={f0_max:.6f}; rms(f0)={f0_rms:.6f}")
print()
converged = C < 0 and f0_rms <= rms_f_thresh and f0_max <= max_f_thresh
if converged:
rot_table.print("@ Converged!")
break
rot_force_evals = 0
rot_force0 = get_f_perp(f1, f2, N)
# Initialize some data structure for the rotation optimizers
if rot_opt == "cg":
# Lists for conjugate gradient
G_perps = [rot_force0, ]
F_rots = [rot_force0, ]
# LBFGS optimizer
elif rot_opt == "lbfgs" :
rot_lbfgs = lbfgs_closure_(rot_force_getter)
# Modified broyden
elif rot_opt == "mb":
rot_mb = closures.modified_broyden_closure(rot_force_getter)
def cbm_restrict(coords1, step, coords0=coords0, dR=dR):
"""Constrain of R1 back on hypersphere."""
coords1_rot = coords1 + step
coords1_dir = coords1_rot - coords0
coords1_dir /= np.linalg.norm(coords1_dir)
coords1_rot = coords0 + coords1_dir*dR
return coords1_rot - coords1
rot_cbm_kwargs = {
"restrict_step": cbm_restrict,
"beta": 0.01,
"M": 3,
}
rot_cbm = closures.modified_broyden_closure(cbm_rot_force_getter, **rot_cbm_kwargs)
for j in range(max_rots):
logger.debug(f"Rotation cycle {j:02d}")
if rot_type == "fourier":
C = get_curvature(f1, f2, N, dR)
logger.debug(f"C={C:.6f}")
# Theta from L-BFGS as implemented in DL-FIND dlf_dimer.f90
if rot_opt == "lbfgs":
theta_dir, rot_force = rot_lbfgs(N, f1, f2)
elif rot_opt == "mb":
theta_dir, rot_force = rot_mb(N, f1, f2)
# Theta from conjugate gradient
elif j > 0 and rot_opt == "cg":
# f_perp = weights.dot(get_f_perp(f1, f2, N))
f_perp = get_f_perp(f1, f2, N)
rot_force = f_perp
gamma = (f_perp - F_rots[-1]).dot(f_perp) / f_perp.dot(f_perp)
G_last = G_perps[-1]
# theta will be present with j > 0.
G_perp = f_perp + gamma * np.linalg.norm(G_last)*theta # noqa: F821
theta_dir = G_perp
F_rots.append(f_perp)
G_perps.append(G_perp)
# Theta from plain steepest descent F_rot/|F_rot|
else:
# theta_dir = weights.dot(get_f_perp(f1, f2, N))
theta_dir = get_f_perp(f1, f2, N)
# Remove component that is parallel to N
theta_dir = theta_dir - theta_dir.dot(N)*N
theta = theta_dir / np.linalg.norm(theta_dir)
# Get rotated endpoint geometries. The rotation takes place in a plane
# spanned by N and theta. Theta is a unit vector perpendicular to N that
# can be formed from the perpendicular components of the forces at the
# endpoints.
# Derivative of the curvature, Eq. (29) in [2]
# (f2 - f1) or -(f1 - f2)
dC = 2*(f0-f1).dot(theta)/dR
rad_trial = -0.5*np.arctan2(dC, 2*abs(C))
logger.debug(f"rad_trial={rad_trial:.2f}")
# print(f"rad_trial={rad_trial:.2f} rad_tol={rad_tol:.2f}")
if np.abs(rad_trial) < rad_tol:
logger.debug(f"rad_trial={rad_trial:.2f} below threshold. Breaking.")
break
# Trial rotation for finite difference calculation of rotational force
# and rotational curvature.
coords1_trial = rotate_R1(coords0, rad_trial, N, theta, dR)
f1_trial = geom1.get_energy_and_forces_at(coords1_trial)["forces"]
rot_force_evals += 1
f2_trial = 2*f0 - f1_trial
N_trial = make_unit_vec(coords1_trial, coords0)
C_trial = get_curvature(f1_trial, f2_trial, N_trial, dR)
b1 = 0.5 * dC
a1 = (C - C_trial + b1*np.sin(2*rad_trial)) / (1-np.cos(2*rad_trial))
a0 = 2 * (C - a1)
rad_min = 0.5 * np.arctan(b1/a1)
logger.debug(f"rad_min={rad_min:.2f}")
def get_C(theta_rad):
return a0/2 + a1*np.cos(2*theta_rad) + b1*np.sin(2*theta_rad)
C_min = get_C(rad_min) # lgtm [py/multiple-definition]
if C_min > C:
rad_min += np.deg2rad(90)
C_min_new = get_C(rad_min)
logger.debug( "Predicted theta_min lead us to a curvature maximum "
f"(C(theta)={C_min:.6f}). Adding pi/2 to theta_min. "
f"(C(theta+pi/2)={C_min_new:.6f})"
)
C_min = C_min_new
# TODO: handle cases where the curvature is still positive, but
# the angle is small, so the rotation is skipped.
# Don't do rotation for small angles
if np.abs(rad_min) < rad_tol:
logger.debug(f"rad_min={rad_min:.2f} below threshold. Breaking.")
break
coords1_rot = rotate_R1(coords0, rad_min, N, theta, dR)
# Interpolate force at coords1_rot; see Eq. (12) in [4]
if interpolate:
f1 = (np.sin(rad_trial-rad_min)/np.sin(rad_trial)*f1
+ np.sin(rad_min)/np.sin(rad_trial)*f1_trial
+ (1 - np.cos(rad_min) - np.sin(rad_min)
* np.tan(rad_trial/2))*f0
)
else:
f1 = geom1.get_energy_and_forces_at(coords1_rot)["forces"]
rot_force_evals += 1
elif rot_type == "direct":
rot_force = get_f_perp(f1, f2, N)
rot_force_rms = get_rms(rot_force)
# print(f"rot_force_rms={rot_force_rms:.4e}, thresh={rot_f_thresh:.4e}")
if rot_force_rms <= rot_f_thresh:
break
rot_step, rot_f = rot_cbm(coords1, N, f1, f2)
coords1_rot = coords1 + rot_step
geom1.coords = coords1_rot
coords1 = coords1_rot
f1 = geom1.forces
rot_force_evals += 1
else:
raise NotImplementedError("Invalid 'rot_type'!")
N = make_unit_vec(coords1_rot, coords0)
coords2_rot = coords0 - N*dR
f2 = 2*f0 - f1
C = get_curvature(f1, f2, N, dR)
rot_force = get_f_perp(f1, f2, N)
rot_force_rms = get_rms(rot_force)
logger.debug("")
if j == 0:
rot_table.print_header()
rot_table.print_row((j, C, rot_force_rms))
tot_rot_force_evals += rot_force_evals
rot_str = (f"Did {rot_force_evals} force evaluation(s) and {j} "
"dimer rotation(s).")
rot_table.print(rot_str)
logger.debug(rot_str)
print()
# If multiple_translations == False then max_translations is 1
# and we will do only one iteration.
for trans_i in range(max_translations):
_, f_parallel, f_perp = get_f_tran(f0, N, C)
prev_f_par_rms = get_rms(f_parallel)
# prev_f_perp_rms = get_rms(f_perp)
prev_C = C
f0_rms = get_rms(f0)
f0_max = max(np.abs(f0))
converged = C < 0 and f0_rms <= rms_f_thresh and f0_max <= max_f_thresh
if converged:
break
step, f_tran = trans_optimizer(coords0, N, C)
step_rms = get_rms(step)
coords0_trans = coords0 + step
coords1_trans = coords0_trans + dR*N
coords2_trans = coords0_trans - dR*N
geom0.coords = coords0_trans
geom1.coords = coords1_trans
geom2.coords = coords2_trans
coords0 = geom0.coords
# Calculate new forces for translated dimer
f0 = geom0.forces
f1 = geom1.forces
add_force_evals += 2
f2 = 2*f0 - f1
C = get_curvature(f1, f2, N, dR)
f0_rms = get_rms(f0)
f0_max = max(np.abs(f0))
if multiple_translations:
if trans_i == 0:
trans_table.print_header()
trans_args = (trans_i, C, f0_max, f0_rms, step_rms)
trans_table.print_row(trans_args)
else:
trans_table.print("Did dimer translation.")
_, f_parallel, f_perp = get_f_tran(f0, N, C)
f_par_rms = get_rms(f_parallel)
# f_perp_rms = get_rms(f_perp)
# Check for sign change of curvature
if (prev_C < 0) and (np.sign(C/prev_C) < 0):
trans_table.print("Curvature became positive!")
break
if (C < 0) and (f_par_rms > prev_f_par_rms):
break
# elif ((C > 0) and (f_par_rms < prev_f_par_rms)
# or (f_perp_rms > prev_f_perp_rms)):
elif (C > 0) and (f_par_rms < prev_f_par_rms):
break
prev_f_par_rms = f_par_rms # lgtm [py/multiple-definition]
# prev_f_perp_rms = f_perp_rms
# Save cycle information
org_coords = np.array((coords1, coords0, coords2))
try:
rot_coords = np.array((coords1_rot, coords0, coords2_rot))
except NameError:
rot_coords = np.array((coords1, coords0, coords2))
trans_coords = np.array((coords1_trans, coords0_trans, coords2_trans))
dc = DimerCycle(org_coords, rot_coords, trans_coords, f0, f_tran)
dimer_cycles.append(dc)
write_progress(geom0)
if check_for_stop_sign():
break
logger.debug("")
print()
sys.stdout.flush()
print(f"@Did {tot_rot_force_evals} force evaluations in the rotation steps "
f"using the {opt_name_dict[rot_opt]} optimizer.")
tot_force_evals = tot_rot_force_evals + add_force_evals
print(f"@Used {add_force_evals} additional force evaluations for a total of "
f"{tot_force_evals} force evaluations.")
dimer_results = DimerResult(dimer_cycles=dimer_cycles,
force_evals=tot_force_evals,
geom0=geom0,
converged=converged)
return dimer_results
def block_davidson(
cart_coords,
masses,
forces_getter,
guess_modes,
lowest=None,
trial_step_size=0.01,
hessian_precon=None,
max_cycles=25,
res_rms_thresh=1e-4,
start_precon=5,
remove_trans_rot=True,
print_level=1,
):
num = len(guess_modes)
B_full = np.zeros((len(guess_modes[0]), num * max_cycles))
S_full = np.zeros_like(B_full)
I = np.eye(cart_coords.size)
masses_rep = np.repeat(masses, 3)
msqrt = np.sqrt(masses_rep)
# Projector to remove translation and rotation
P = get_trans_rot_projector(cart_coords, masses)
guess_modes = [
NormalMode(P.dot(mode.l_mw) / msqrt, masses_rep)
for mode in guess_modes
]
col_fmts = "int int int float float str".split()
header = ("#", "subspace size", "mode", "ṽ / cm⁻¹", "rms(r)", "Conv")
table = TablePrinter(header, col_fmts, width=11)
if print_level == 1:
table.print_header()
b_prev = np.array([mode.l_mw for mode in guess_modes]).T
for i in range(max_cycles):
# Add new basis vectors to B matrix
b = np.array([mode.l_mw for mode in guess_modes]).T
from_ = i * num
to_ = (i + 1) * num
B_full[:, from_:to_] = b
# Estimate action of Hessian by finite differences.
for j in range(num):
mode = guess_modes[j]
# Get a step size in mass-weighted coordinates that results
# in the desired 'trial_step_size' in not-mass-weighted coordinates.
mw_step_size = mode.mw_norm_for_norm(trial_step_size)
# Actual step in non-mass-weighted coordinates
step = trial_step_size * mode.l
S_full[:, from_ + j] = (
# Convert to mass-weighted coordinates
forces_fin_diff(forces_getter, cart_coords, step, mw_step_size)
/ msqrt)
# Views on columns that are actually set
B = B_full[:, :to_]
S = S_full[:, :to_]
# Calculate and symmetrize approximate hessian
Hm = B.T.dot(S)
Hm = (Hm + Hm.T) / 2
# Diagonalize small Hessian
w, v = np.linalg.eigh(Hm)
# Approximations to exact eigenvectors in current cycle
approx_modes = (v * B[:, :, None]).sum(axis=1).T
# Calculate overlaps between previous root and the new approximate
# normal modes for root following.
if lowest is None:
# 2D overlap array. approx_modes in row, b_prev in columns.
overlaps = np.einsum("ij,jk->ik", approx_modes, b_prev)
mode_inds = np.abs(overlaps).argmax(axis=0)
else:
mode_inds = np.arange(lowest)
b_prev = approx_modes[mode_inds].T
# Eq. (7) in [1]
residues = (v * (S[:, :, None] - w * B[:, :, None])).sum(axis=1)
# Determine preconditioner matrix
#
# Use supplied matrix
if hessian_precon is not None:
precon_mat = hessian_precon
# Reconstruct Hessian, but only start after some cycles
elif i >= start_precon:
precon_mat = B.dot(Hm).dot(B.T)
# No preconditioning if no matrix was supplied or we are in an early cycle.
else:
precon_mat = None
# Construct new basis vector from residuum of selected mode
b = np.zeros_like(b_prev)
for j, mode_ind in enumerate(mode_inds):
r = residues[:, mode_ind]
if precon_mat is not None:
# Construct actual preconditioner X
X = np.linalg.pinv(precon_mat - w[mode_ind] * I, rcond=1e-8)
b[:, j] = X.dot(r)
else:
b[:, j] = r
# Project out translation and rotation from new mode guess
if remove_trans_rot:
b = P.dot(b)
# Orthogonalize new vectors against preset vectors
b = np.linalg.qr(np.concatenate((B, b), axis=1))[0][:, -num:]
# New NormalMode from non-mass-weighted displacements
guess_modes = [NormalMode(b_ / msqrt, masses_rep) for b_ in b.T]
# Calculate wavenumbers
nus = eigval_to_wavenumber(w)
# Check convergence criteria
max_res = np.abs(residues).max(axis=0)
res_rms = np.sqrt(np.mean(residues**2, axis=0))
converged = res_rms < res_rms_thresh
# Print progress if requested
if print_level == 2:
print(f"Cycle {i:02d}")
print("\t # | ṽ / cm⁻¹| rms(r) | max(|r|) ")
print("\t------------------------------------------")
for j, (nu, rms, mr) in enumerate(zip(nus, res_rms, max_res)):
sel_str = "*" if (j in mode_inds) else " "
conv_str = "✓" if converged[j] else ""
print(
f"\t{j:02d}{sel_str} | {nu:> 10.2f} | {rms:.8f} | {mr:.8f} {conv_str}"
)
print()
elif print_level == 1:
for j in mode_inds:
conv_str = "✓" if converged[j] else "✗"
table.print_row(
(i, B.shape[1], j, nus[j], res_rms[j], conv_str))
# Convergence is signalled using only the roots we are actually interested in
modes_converged = all(converged[mode_inds])
if modes_converged:
if print_level > 0:
print(f"\tDavidson procedure converged in {i+1} cycles!")
if lowest is not None:
nus_str = np.array2string(nus[mode_inds], precision=2)
print(f"\tLowest {lowest} wavenumbers: {nus_str} cm⁻¹")
neg_nus = sum(nus[mode_inds] < 0)
type_ = "minimum" if (
neg_nus == 0) else f"saddle point of index {neg_nus}"
print(f"\tThis geometry seems to be a {type_} on the PES.")
break
result = DavidsonResult(
cur_cycle=i,
converged=modes_converged,
final_modes=guess_modes,
qs=approx_modes,
nus=nus,
mode_inds=mode_inds,
res_rms=res_rms,
)
return result