def get_step_func(self, eigvals, gradient, grad_rms_thresh=1e-2):
positive_definite = (eigvals < 0).sum() == 0
gradient_small = rms(gradient) < grad_rms_thresh
if self.adapt_step_func and gradient_small and positive_definite:
return self.get_newton_step_on_trust, self.quadratic_model
# RFO fallback
else:
return self.get_rs_step, self.rfo_model
def dimer_method(geom0,
N,
R,
calc_getter,
max_cycles=50,
f_thresh=1e-3,
max_step=0.3,
rot_kwargs=None,
trans_kwargs=None):
if rot_kwargs is None:
rot_kwargs = {}
if trans_kwargs is None:
trans_kwargs = {}
geom1, geom2 = get_dimer_ends(geom0, N, R, calc_getter)
coords = [
(geom1.coords, geom0.coords, geom2.coords),
]
restrict_step = partial(restrict_step_length, max_length=max_step)
trans_optimizer = lbfgs_closure(lambda _, *args: get_f_trans(*args),
restrict_step=restrict_step)
for i in range(max_cycles):
f0 = geom0.forces
f0_rms = rms(f0)
print(f"{i:0d} rms(f0)={f0_rms:.6f}")
if f0_rms < f_thresh:
print("Converged!")
break
# Rotation
rot_result = rotate_dimer(geom0, geom1, geom2, N, R, **rot_kwargs)
geom1, geom2, N, C, _ = rot_result
# Translation
step, f_trans = trans_optimizer(geom0.coords, geom0, N, C)
new_coords0 = geom0.coords + step
geom0.coords = new_coords0
# Steepet descent
# f_trans = get_f_trans(geom0, N, C)
# alpha = 0.5
# step = alpha * f_trans
# new_coords0 = geom0.coords + alpha*step
# geom0.coords = new_coords0
update_dimer_ends(geom0, geom1, geom2, N, R)
coords.append((geom1.coords, geom0.coords, geom2.coords))
# TODO: return dimerresults
return coords
def do_dimer_rotations(self, rotation_thresh=None):
self.log("Doing dimer rotations")
if rotation_thresh is None:
rotation_thresh = self.rotation_thresh
self.log(f"\tThreshold norm(rot_force)={rotation_thresh:.6f}")
lbfgs = small_lbfgs_closure()
try:
N_first = self.N
prev_step = None
for i in range(
self.rotation_max_cycles): # lgtm [py/redundant-else]
N_cur = self.N
rot_force = self.rot_force
rms_rot_force = rms(rot_force)
if self.should_bias_f1:
C_real = self.C
C_bias = -self.bias_rotation_a * (self.N.dot(
self.N_init))**2
C = C_real + C_bias
C_str = f"C={C: .6f}, C_real={C_real: .6f}, C_bias={C_bias: .6f}"
else:
C_str = f"C={self.C: .6f}"
self.log(
f"\t{i:02d}: rms(rot_force)={rms_rot_force:.6f} {C_str}")
if rms_rot_force <= rotation_thresh:
self.log("\trms(rot_force) is below threshold!")
raise RotationConverged
coords1_old = self.coords1
self.rotation_method(lbfgs, prev_step)
actual_step = self.coords1 - coords1_old
prev_step = actual_step
rot_deg = np.rad2deg(np.arccos(N_cur.dot(self.N)))
self.log(f"\t\tRotated by {rot_deg:.1f}°")
else:
msg = "\tDimer rotation did not converge in " \
f"{self.rotation_max_cycles}"
except RotationConverged:
msg = f"\tDimer rotation converged in {i+1} cycle(s)."
self.log(msg)
self.log("\tN after rotation:\n\t" + str(self.N))
self.log()
# Restrict to interval [-1,1] where arccos is defined
rot_deg = np.rad2deg(
np.arccos(max(min(N_first.dot(self.N), 1.0), -1.0)))
self.log(f"\tRotated by {rot_deg:.1f}° w.r.t. the orientation "
"before rotation(s).")
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 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 optimize(self):
energy, gradient, H, big_eigvals, big_eigvecs, resetted = self.housekeeping()
# Reference RFO step, used for judging the proposed GDIIS step
ref_gradient = gradient.copy()
ref_rfo_step = self.get_rs_step(big_eigvals, big_eigvecs, gradient, name="RS-RFO")
# Right everything is in place to check for convergence. If all values are below
# the thresholds, there is no need to do additional inter/extrapolations.
if self.check_convergence(ref_rfo_step):
self.log("Convergence achieved! Skipping inter/extrapolation.")
return ref_rfo_step
# Try to interpolate an intermediate geometry, either from GDIIS or line search.
#
# Set some defaults
ip_gradient = None
ip_step = None
diis_result = None
# Check if we can do GDIIS or GEDIIS. If we (can) do a line search is decided
# after trying GDIIS.
rms_forces = rms(gradient)
rms_step = rms(ref_rfo_step)
can_diis = (rms_step <= self.gdiis_thresh) and (not resetted)
can_gediis = (rms_forces <= self.gediis_thresh) and (not resetted)
# GDIIS / GEDIIS, prefer GDIIS over GEDIIS
if self.gdiis and can_diis:
# Gradients as error vectors
err_vecs = -np.array(self.forces)
diis_result = gdiis(err_vecs, self.coords, self.forces, ref_rfo_step)
# Don't try GEDIIS if GDIIS failed. If GEDIIS should be tried after GDIIS failed
# comment the line below and uncomment the line following it.
elif self.gediis and can_gediis:
# if self.gediis and can_gediis and (diis_result == None):
diis_result = gediis(self.coords, self.energies, self.forces, hessian=H)
try:
ip_coords = diis_result.coords
ip_step = ip_coords - self.geometry.coords
ip_gradient = -diis_result.forces
# When diis_result is None
except AttributeError:
self.log("GDIIS didn't succeed.")
# Try line search if GDIIS failed or not requested
if self.line_search and (diis_result is None) and (not resetted):
ip_energy, ip_gradient, ip_step = self.poly_line_search()
# Use the interpolated gradient for the RFO step if interpolation succeeded
if (ip_gradient is not None) and (ip_step is not None):
gradient = ip_gradient
# Keep the original gradient when the interpolation failed, but recreate
# ip_step, as it will be returned as None from poly_line_search().
else:
ip_step = np.zeros_like(gradient)
# RFO step (from intermediate geometry) with (interpolated) gradient
rfo_step = self.get_rs_step(big_eigvals, big_eigvecs, gradient, name="RS-RFO")
# Form full step. If we did not interpolate or it failed ip_step will be zero.
step = rfo_step + ip_step
# Use the original, actually calculated, gradient
quadratic_prediction = step @ ref_gradient + 0.5 * step @ H @ step
rfo_prediction = quadratic_prediction / (1 + step @ step)
self.predicted_energy_changes.append(rfo_prediction)
return step
def step(self):
##################
# PREDICTOR STEP #
##################
mw_grad = self.mw_gradient
energy = self.energy
if self.cur_cycle > 0:
if self.hessian_recalc and (self.cur_cycle % self.hessian_recalc
== 0):
self.mw_hessian = self.geometry.mw_hessian
h5_fn = f"hess_calc_irc_{self.direction}_cyc{self.cur_cycle}.h5"
save_hessian(h5_fn, self.geometry)
self.log("Calculated excact hessian")
else:
dx = self.mw_coords - self.irc_mw_coords[-2]
dg = mw_grad - self.irc_mw_gradients[-2]
dH, key = self.hessian_update_func(self.mw_hessian, dx, dg)
self.mw_hessian += dH
self.log(f"Did {key} hessian update before predictor step.")
self.dwi.update(self.mw_coords.copy(), energy, mw_grad,
self.mw_hessian.copy())
# Create a copy of the inital coordinates for the determination
# of the actual step size in the predictor Euler integration.
init_mw_coords = self.mw_coords.copy()
get_integration_length = self.get_integration_length_func(
init_mw_coords)
# Calculate predictor Euler-integration step length. See get_conv_fact
# method definition for a comment on this.
conv_fact = self.get_conv_fact(mw_grad)
euler_step_length = self.step_length / (self.max_pred_steps /
conv_fact)
def taylor_gradient(step):
"""Return gradient from Taylor expansion of energy to 2nd order."""
return mw_grad + self.mw_hessian @ step
# These variables will hold the coordinates and gradients along
# the Euler integration and will be updated frequently.
euler_mw_coords = self.mw_coords.copy()
euler_mw_grad = mw_grad.copy()
self.log(
f"Predictor-Euler-integration with Δs={euler_step_length:.6f} "
f"for up to {self.max_pred_steps} steps")
prev_cur_length = 0.
for i in range(self.max_pred_steps):
# Calculate step length in non-mass-weighted coordinates
cur_length = get_integration_length(euler_mw_coords)
if i % 50 == 0:
diff = cur_length - prev_cur_length
self.log(f"\t{i:03d}: {cur_length:.4f} Δ={diff:.4f}")
prev_cur_length = cur_length
# Check if we achieved the desired step length.
if cur_length >= self.step_length:
self.log(
"Predictor-Euler integration converged with "
f"Δs={cur_length:.4f} (desired Δs={self.step_length:.4f}) "
f"after {i+1} steps!")
break
step_ = euler_step_length * -euler_mw_grad / np.linalg.norm(
euler_mw_grad)
euler_mw_coords += step_
# Determine actual step by comparing the current and the initial coordinates
euler_step = euler_mw_coords - init_mw_coords
euler_mw_grad = taylor_gradient(euler_step)
else:
self.log(f"Predictor-Euler integration did not converge in {i+1} "
f"steps. Δs={cur_length:.4f}.")
# Check if we are already sufficiently converged. If so signal
# convergence.
self.mw_coords = euler_mw_coords
# Use rms of gradient from taylor expansion for convergence check.
euler_grad = self.unweight_vec(euler_mw_grad)
rms_grad = rms(euler_grad)
# Or check true gradient? But this would need an additional calculation,
# so I disabled it for now.
# rms_grad = rms(self.gradient)
# if rms_grad <= 5*self.rms_grad_thresh:
if rms_grad <= self.rms_grad_thresh:
self.log("Sufficient convergence achieved on rms(grad)")
self.converged = True
return
self.log("")
# Calculate energy and gradient at new predicted geometry. Update the
# hessian accordingly. These results will be added to the DWI for use
# in the corrector step.
self.mw_coords = euler_mw_coords
self.log("Calculating energy and gradient at predictor step geometry.")
mw_grad = self.mw_gradient
energy = self.energy
# Hessian update
dx = self.mw_coords - self.irc_mw_coords[-1]
dg = mw_grad - self.irc_mw_gradients[-1]
dH, key = self.hessian_update_func(self.mw_hessian, dx, dg)
self.mw_hessian += dH
self.log(f"Did {key} hessian update after predictor step.\n")
self.dwi.update(self.mw_coords.copy(), energy, mw_grad,
self.mw_hessian.copy())
if self.dump_dwi:
self.dwi.dump(
f"dwi_{self.cur_direction}_{self.cur_cycle:0{self.cycle_places}d}.h5"
)
corrected_mw_coords = self.corr_func(init_mw_coords, self.step_length,
self.dwi)
self.mw_coords = corrected_mw_coords
corr_step_length = get_integration_length(self.mw_coords)
self.log(f"Corrected unweighted step length: {corr_step_length:.6f}")
