def run():
ovlp_types = "wf tden nto_org nto".split()
# ovlp_types = ("nto", )
ovlp_withs = "adapt first previous".split()
for i, (ovlp_type,
ovlp_with) in enumerate(it.product(ovlp_types, ovlp_withs)):
# ovlp_type = "wf"
# ovlp_with = "adapt"
print(
highlight_text(
f"i={i:02d}, ovlp_type={ovlp_type}, ovlp_with={ovlp_with}"))
geom = geom_from_library("cytosin.xyz", coord_type="redund")
calc = get_calc(ovlp_type, ovlp_with)
geom.set_calculator(calc)
opt = RFOptimizer(geom)
opt.run()
assert calc.root_flips[2] # == True
assert all([
flipped == False for i, flipped in enumerate(calc.root_flips)
if i != 2
])
assert calc.root == 2
assert opt.cur_cycle == 4
print()
def run(self):
if not self.restarted:
prep_start_time = time.time()
self.prepare_opt()
prep_end_time = time.time()
prep_time = prep_end_time - prep_start_time
print(f"Spent {prep_time:.1f} s preparing the first cycle.")
self.print_header()
self.stopped = False
# Actual optimization loop
for self.cur_cycle in range(self.last_cycle, self.max_cycles):
start_time = time.time()
self.log(highlight_text(f"Cycle {self.cur_cycle:03d}"))
if self.is_cos and self.check_coord_diffs:
image_coords = [
image.cart_coords for image in self.geometry.images
]
align = len(image_coords[0]) > 3
cds = get_coords_diffs(image_coords, align=align)
# Differences of coordinate differences ;)
cds_diffs = np.diff(cds)
min_ind = cds_diffs.argmin()
if cds_diffs[min_ind] < self.coord_diff_thresh:
similar_inds = min_ind, min_ind + 1
msg = (
f"Cartesian coordinates of images {similar_inds} are "
"too similar. Stopping optimization!")
# I should improve my logging :)
print(msg)
self.log(msg)
break
# Check if something considerably changed in the optimization,
# e.g. new images were added/interpolated. Then the optimizer
# should be reset.
reset_flag = False
if self.cur_cycle > 0 and self.is_cos:
reset_flag = self.geometry.prepare_opt_cycle(
self.coords[-1], self.energies[-1], self.forces[-1])
# Reset when number of coordinates changed
elif self.cur_cycle > 0:
reset_flag = reset_flag or (self.geometry.coords.size !=
self.coords[-1].size)
if reset_flag:
self.reset()
self.coords.append(self.geometry.coords.copy())
self.cart_coords.append(self.geometry.cart_coords.copy())
# Determine and store number of currenctly actively optimized images
try:
image_inds = self.geometry.image_inds
image_num = len(image_inds)
except AttributeError:
image_inds = [
0,
]
image_num = 1
self.image_inds.append(image_inds)
self.image_nums.append(image_num)
step = self.optimize()
if step is None:
# Remove the previously added coords
self.coords.pop(-1)
self.cart_coords.pop(-1)
continue
if self.is_cos:
self.tangents.append(self.geometry.get_tangents().flatten())
self.steps.append(step)
# Convergence check
self.is_converged = self.check_convergence()
end_time = time.time()
elapsed_seconds = end_time - start_time
self.cycle_times.append(elapsed_seconds)
if self.dump:
self.write_cycle_to_file()
with open(self.current_fn, "w") as handle:
handle.write(self.geometry.as_xyz())
if (self.dump and self.dump_restart
and (self.cur_cycle % self.dump_restart) == 0):
self.dump_restart_info()
self.print_opt_progress()
if self.is_converged:
print("Converged!")
print()
break
# Update coordinates
new_coords = self.geometry.coords.copy() + step
try:
self.geometry.coords = new_coords
# Use the actual step. It may differ from the proposed step
# when internal coordinates are used, as the internal-Cartesian
# transformation is done iteratively.
self.steps[-1] = self.geometry.coords - self.coords[-1]
except RebuiltInternalsException as exception:
print("Rebuilt internal coordinates")
with open("rebuilt_primitives.xyz", "w") as handle:
handle.write(self.geometry.as_xyz())
if self.is_cos:
for image in self.geometry.images:
image.reset_coords(exception.typed_prims)
self.reset()
if hasattr(self.geometry, "reparametrize"):
reparametrized = self.geometry.reparametrize()
cur_coords = self.geometry.coords
prev_coords = self.coords[-1]
if reparametrized and (cur_coords.size == prev_coords.size):
self.log("Did reparametrization")
rms = np.sqrt(np.mean((prev_coords - cur_coords)**2))
self.log(
f"rms of coordinates after reparametrization={rms:.6f}"
)
self.is_converged = rms < self.reparam_thresh
if self.is_converged:
print("Insignificant coordinate change after "
"reparametrization. Signalling convergence!")
print()
break
sys.stdout.flush()
sign = check_for_end_sign()
if sign == "stop":
self.stopped = True
break
elif sign == "converged":
self.converged = True
print("Operator indicated convergence!")
break
self.log("")
else:
print("Number of cycles exceeded!")
# Outside loop
if self.dump:
self.out_trj_handle.close()
if (not self.is_cos) and (not self.stopped):
print(self.final_summary())
# Remove 'current_geometry.xyz' file
try:
os.remove(self.current_fn)
except FileNotFoundError:
self.log(
f"Tried to delete '{self.current_fn}'. Couldn't find it.")
with open(self.final_fn, "w") as handle:
handle.write(self.geometry.as_xyz())
print(
f"Wrote final, hopefully optimized, geometry to '{self.final_fn.name}'"
)
sys.stdout.flush()
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 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 run(self):
for self.cur_cycle in range(self.max_cycles):
cycle_start = time.time()
self.log(highlight_text(f"Cycle {self.cur_cycle}"))
input_geoms = [self.get_input_geom(self.initial_geom)
for _ in range(self.cycle_size)]
# Write input geometries to disk
self.write_geoms_to_trj(input_geoms, f"cycle_{self.cur_cycle:03d}_input.trj")
# Run optimizations on input geometries
calc_start = time.time()
opt_geoms = list()
for i, geom in enumerate(input_geoms, 1):
print(f"Optimizing geometry {i:03d}/{self.cycle_size:03d}", end="\r")
opt_geoms.append(self.run_geom_opt(geom))
print()
calc_end = time.time()
calc_duration = calc_end - calc_start
self.log(f"Optimizations took {calc_duration:.0f} s.")
kept_geoms = list()
for geom in opt_geoms:
# Do all the filtering and reject all invalid geometries
if not self.geom_is_valid(geom):
continue
energy = geom.energy
i = bisect.bisect_left(self.new_energies, energy)
self.new_energies.insert(i, energy)
self.new_geoms.insert(i, geom)
kept_geoms.append(geom)
if i == 0 and len(self.new_energies) > 1:
last_minimum = self.new_energies[1]
diff = abs(energy - last_minimum)
self.log(f"It is a new global minimum at {energy:.4f} au! "
f"Last one was {diff:.4f} au higher.")
kept_num = len(kept_geoms)
trj_filtered_fn = f"cycle_{self.cur_cycle:03d}.trj"
# Sort by energy
kept_geoms = sorted(kept_geoms, key=lambda g: g.energy)
if kept_geoms:
self.write_geoms_to_trj(kept_geoms, trj_filtered_fn)
self.log(f"Kicks in cycle {self.cur_cycle} produced "
f"{kept_num} new geometries.")
self.break_in = self.break_after
elif self.break_in == 0:
self.log("Didn't find any new geometries in the last "
f"{self.break_after} cycles. Exiting!")
break
else:
self.log(f"Cycle {self.cur_cycle} produced no new geometries.")
self.break_in -= 1
cycle_end = time.time()
cycle_duration = cycle_end - cycle_start
self.log(f"Cycle {i} took {cycle_duration:.0f} s.")
self.log("")
if check_for_stop_sign():
break
self.log(f"Run produced {len(self.new_energies)} geometries!")
# Return empty list of nothing was found
if not self.new_energies:
return []
fn = "final.trj"
self.write_geoms_to_trj(self.new_geoms, fn)
# self.new_energies = np.array(new_energies)
np.savetxt("energies.dat", self.new_energies)
first_geom = self.new_geoms[0]
first_geom.standard_orientation()
first_geom.energy = self.new_energies[0]
if self.is_analytical2d:
return self.new_geoms
matched_geoms = [first_geom, ]
for geom, energy in zip(self.new_geoms[1:], self.new_energies):
rmsd, (_, matched_geom) = matched_rmsd(first_geom, geom)
matched_geom.energy = energy
matched_geoms.append(matched_geom)
fn_matched = "final_matched.trj"
self.write_geoms_to_trj(matched_geoms, fn_matched)
return matched_geoms
def run_calculations(geom,
charge,
mult,
calc_getter_gas,
calc_getter_solv,
opt=False):
def get_name(base):
return f"{base}_{charge}_{mult}"
print(highlight_text(f"Charge={charge}, Mult={mult}"))
gas_name = get_name("gas")
calc_kwargs = {
"charge": charge,
"mult": mult,
"base_name": gas_name,
}
opt_kwargs = {
"thresh": "gau",
"h5_group_name": f"opt_{charge}_{mult}",
"dump": True,
"prefix": f"{gas_name}_",
}
# Gas phase calculation, optimization and frequency
gas_calc = calc_getter_gas(calc_kwargs)
geom.set_calculator(gas_calc)
if opt:
opt = RFOptimizer(geom, **opt_kwargs)
opt.run()
assert opt.is_converged
print(highlight_text("Gas phase optimization finished!", level=1))
Hg = geom.cart_hessian
energyg = geom.energy
save_hessian(f"{get_name('gas')}.h5",
geom,
cart_hessian=Hg,
energy=energyg,
mult=mult)
print(highlight_text("Gas phase hessian finished!", level=1))
# Solvent calculation, frequency
solv_name = get_name("solv")
solv_kwargs = calc_kwargs.copy()
solv_kwargs["base_name"] = solv_name
solv_calc = calc_getter_solv(solv_kwargs)
solv_geom = geom.copy()
solv_geom.set_calculator(solv_calc)
energys = solv_geom.energy
with open(f"{solv_name}.energy", "w") as handle:
handle.write(str(energys))
print(highlight_text("Solvated energy finished!", level=1))
dG_solv = energys - energyg
res = RedoxResult(
energy_gas=energyg,
hessian_gas=Hg,
energy_solv=energys,
dG_solv=dG_solv,
geom_gas=geom,
charge=charge,
mult=mult,
)
return res
def mdp(
geom,
steps,
dt,
term_funcs=None,
steps_init=None,
E_excess=0.0,
displ_length=0.1,
epsilon=5e-4,
ascent_alpha=0.05,
max_ascent_steps=25,
max_init_trajs=10,
dump=True,
seed=None,
external_md=False,
):
# Sanity checks and forcing some types
dt = float(dt)
assert dt > 0.0
steps = int(steps)
t = dt * steps
# assert t > dt
if steps_init is None:
steps_init = steps // 10
print(f"No 'steps_init' provided! Using {steps_init}")
E_excess = float(E_excess)
assert E_excess >= 0.0
displ_length = float(displ_length)
assert displ_length >= 0.0
if term_funcs is None:
term_funcs = {}
for k, v in term_funcs.items():
if callable(v):
continue
elif isinstance(v, str):
term_funcs[k] = parse_raw_term_func(v)
else:
raise Exception(f"Invalid term function '{k}: {v}' encountered!")
print(highlight_text("Minimum dynamic path calculation"))
if seed is None:
# 2**32 - 1
seed = np.random.randint(4294967295)
np.random.seed(seed)
print(f"Using seed {seed} to initialize the random number generator.\n")
E_TS = geom.energy
E_tot = E_TS + E_excess
# Distribute E_excess evenly on E_pot and E_kin
E_pot_diff = 0.5 * E_excess
E_pot_desired = E_TS + E_pot_diff
print(f"E_TS={E_TS:.6f} au")
# Determine transition vector
w, v = np.linalg.eigh(geom.hessian)
assert w[0] < -1e-8
trans_vec = v[:, 0]
# Disable removal of translation/rotation for analytical potentials
remove_com_v = remove_rot_v = geom.cart_coords.size > 3
if E_excess == 0.0:
print("MDP without excess energy.")
# Without excess energy we have to do an initial displacement along
# the transition vector to get a non-vanishing gradient.
initial_displacement = displ_length * trans_vec
x0_plus = geom.coords + initial_displacement
x0_minus = geom.coords - initial_displacement
v0_zero = np.zeros_like(geom.coords)
md_kwargs = {
"v0": v0_zero.copy(),
"t": t,
"dt": dt,
"term_funcs": term_funcs,
"external": external_md,
}
geom.coords = x0_plus
md_fin_plus = run_md(geom, **md_kwargs)
geom.coords = x0_minus
md_fin_minus = run_md(geom, **md_kwargs)
if dump:
dump_coords(geom.atoms, md_fin_plus.coords, "mdp_plus.trj")
dump_coords(geom.atoms, md_fin_minus.coords, "mdp_minus.trj")
mdp_result = MDPResult(
ascent_xs=None,
md_init_plus=None,
md_init_minus=None,
md_fin_plus=md_fin_plus,
md_fin_minus=md_fin_minus,
)
return mdp_result
print(f"E_excess={E_excess:.6f} au, ({E_excess*AU2KJPERMOL:.1f} kJ/mol)")
print(f"E_pot,desired=E_TS + {E_pot_diff*AU2KJPERMOL:.1f} kJ/mol")
print()
# Generate random vector perpendicular to transition vector
perp_vec = np.random.rand(*trans_vec.shape)
# Zero last element if we have an analytical surface
if perp_vec.size == 3:
perp_vec[2] = 0
# Orthogonalize vector
perp_vec = perp_vec - (perp_vec @ trans_vec) * trans_vec
perp_vec /= np.linalg.norm(perp_vec)
# Initial displacement from x_TS to x, generating a point with
# non-vanishing gradient.
x = geom.coords + epsilon * perp_vec
geom.coords = x
# Do steepest ascent until E_tot is reached
E_pot = geom.energy
ascent_xs = list()
for i in range(max_ascent_steps):
ascent_xs.append(geom.coords.copy())
ascent_converged = E_pot >= E_pot_desired
if ascent_converged:
break
gradient = geom.gradient
E_pot = geom.energy
direction = gradient / np.linalg.norm(gradient)
step = ascent_alpha * direction
new_coords = geom.coords + step
geom.coords = new_coords
# calc = geom.calculator
# class Opt:
# pass
# _opt = Opt()
# _opt.coords = np.array(ascent_xs)
# calc.plot_opt(_opt, show=True)
assert ascent_converged, "Steepest ascent didn't converge!"
assert (E_tot - E_pot) > 0.0, (
"Potential energy after steepst ascent is greater than the desired "
f"total energy ({E_pot:.6f} > {E_tot:.6f}). Maybe try a smaller epsilon? "
f"The current value Ɛ={epsilon:.6f} may be too big!")
ascent_xs = np.array(ascent_xs)
if dump:
dump_coords(geom.atoms, ascent_xs, "mdp_ee_ascent.trj")
x0 = geom.coords.copy()
print(highlight_text("Runninig initialization trajectories", level=1))
for i in range(max_init_trajs):
# Determine random momentum vector for the given kinetic energy
E_kin = E_tot - E_pot
T = temperature_for_kinetic_energy(len(geom.atoms), E_kin)
v0 = get_mb_velocities_for_geom(geom,
T,
remove_com_v=remove_com_v,
remove_rot_v=remove_rot_v).flatten()
# Zero last element if we have an analytical surface
if v0.size == 3:
v0[2] = 0
# Run initial MD to check if both trajectories run towards different
# basins of attraction.
# First MD with positive v0
md_init_kwargs = {
"v0": v0.copy(),
"steps": steps_init,
"dt": dt,
"external": external_md,
}
geom.coords = x0.copy()
md_init_plus = run_md(geom, **md_init_kwargs)
# Second MD with negative v0
geom.coords = x0.copy()
md_init_kwargs["v0"] = -v0.copy()
md_init_minus = run_md(geom, **md_init_kwargs)
dump_coords(geom.atoms, md_init_plus.coords,
f"mdp_ee_init_plus_{i:02d}.trj")
dump_coords(geom.atoms, md_init_minus.coords,
f"mdp_ee_init_minus_{i:02d}.trj")
# Check if both MDs run into different basins of attraction.
# We (try to) do this by calculating the overlap between the
# transition vector and the normalized vector defined by the
# difference between x0 and the endpoint of the respective
# test trajectory. Both overlaps should have different sings.
end_plus = md_init_plus.coords[-1]
pls = end_plus - x0
pls /= np.linalg.norm(pls)
end_minus = md_init_minus.coords[-1]
minus = end_minus - x0
minus /= np.linalg.norm(minus)
p = trans_vec @ pls
m = trans_vec @ minus
init_trajs_converged = np.sign(p) != np.sign(m)
if init_trajs_converged:
print("Trajectories ran into different basins. Breaking.")
break
if dump:
dump_coords(geom.atoms, md_init_plus.coords, "mdp_ee_init_plus.trj")
dump_coords(geom.atoms, md_init_minus.coords, "mdp_ee_init_minus.trj")
assert init_trajs_converged
print(f"Ran 2*{i+1} initialization trajectories.")
print()
# Run actual trajectories, using the supplied termination functions if possible.
print(highlight_text("Running actual full trajectories.", level=1))
def print_status(terminated, step):
if terminated:
msg = f"\tTerminated by '{terminated}' in step {step}."
else:
msg = "\tMax time steps reached!"
print(msg)
# "Production"/Final MDs
md_fin_kwargs = {
"v0": v0.copy(),
"steps": steps,
"dt": dt,
"term_funcs": term_funcs,
"external": external_md,
}
# MD with positive v0.
geom.coords = x0.copy()
md_fin_plus = run_md(geom, **md_fin_kwargs)
print_status(md_fin_plus.terminated, md_fin_plus.step)
# MD with negative v0.
geom.coords = x0.copy()
md_fin_kwargs["v0"] = -v0
md_fin_minus = run_md(geom, **md_fin_kwargs)
print_status(md_fin_minus.terminated, md_fin_minus.step)
md_fin_plus_term = md_fin_plus.terminated
md_fin_minus_term = md_fin_minus.terminated
if dump:
dump_coords(geom.atoms, md_fin_plus.coords, "mdp_ee_fin_plus.trj")
dump_coords(geom.atoms, md_fin_minus.coords, "mdp_ee_fin_minus.trj")
mdp_result = MDPResult(
ascent_xs=ascent_xs,
md_init_plus=md_init_plus,
md_init_minus=md_init_minus,
md_fin_plus=md_fin_plus,
md_fin_minus=md_fin_minus,
md_fin_plus_term=md_fin_plus_term,
md_fin_minus_term=md_fin_minus_term,
)
return mdp_result