import logging
from pathlib import Path
import h5py
import numpy as np
from pysisyphus.calculators.Calculator import Calculator
from pysisyphus.helpers import get_tangent_trj_str
from pysisyphus.helpers_pure import rms
from pysisyphus.intcoords.helpers import get_weighted_bond_mode
from pysisyphus.linalg import perp_comp, make_unit_vec
from pysisyphus.optimizers.closures import small_lbfgs_closure
from pysisyphus.optimizers.restrict_step import get_scale_max
[docs]
class RotationConverged(Exception):
pass
[docs]
class Gaussian:
def __init__(self, height, center, std, N):
self.center = np.array(center)
self.height = float(height)
self.std = float(std)
self.N = np.array(N)
self.s0 = self.center.dot(self.N)
[docs]
def energy(self, R, height=None):
if height is None:
height = self.height
return height * np.exp(-((R.dot(self.N) - self.s0) ** 2) / (2 * self.std**2))
[docs]
def forces(self, R, height=None):
if height is None:
height = self.height
s_diff = R.dot(self.N) - self.s0
return (
height
* np.exp(-(s_diff**2) / (2 * self.std**2))
* s_diff
/ self.std**2
* self.N
)
def __str__(self):
return (
f"Gaussian(height={self.height:.4f}, center={self.center}, "
f"s0={self.s0:.4f}, std={self.std:.4f}"
)
# [1] https://aip.scitation.org/doi/abs/10.1063/1.480097
# Original Dimer paper
# Henkelmann, 1999
# [2] https://doi.org/10.1063/1.1809574
# Comparison of TS search methods
# Olsen, 2004
# [3] https://doi.org/10.1063/1.2104507
# Comparison of Dimer and P-RFO
# Heyden, 2005
# [4] https://aip.scitation.org/doi/abs/10.1063/1.2815812
# Superlinear Dimer method
# Kästner, 2008
[docs]
class Dimer(Calculator):
def __init__(
self,
calculator,
*args,
N_raw=None,
length=0.0189,
rotation_max_cycles=15,
rotation_method="fourier",
rotation_thresh=1e-4,
rotation_tol=1,
rotation_max_element=0.001,
rotation_interpolate=True,
rotation_disable=False,
rotation_disable_pos_curv=True,
rotation_remove_trans=True,
trans_force_f_perp=True,
bonds=None,
N_hessian=None,
bias_rotation=False,
bias_translation=False,
bias_gaussian_dot=0.1,
seed=None,
write_orientations=True,
forward_hessian=True,
**kwargs,
):
super().__init__(*args, **kwargs)
self.logger = logging.getLogger("dimer")
self.calculator = calculator
self.length = float(length)
# Rotation parameters
self.rotation_max_cycles = int(rotation_max_cycles)
rotation_methods = {
"direct": self.direct_rotation,
"fourier": self.fourier_rotation,
}
try:
self.rotation_method = rotation_methods[rotation_method]
except KeyError as err:
print(
f"Invalid rotation_method={rotation_method}! Valid types are: "
f"{tuple(self.rotation_methods.keys())}"
)
raise err
self.rotation_thresh = float(rotation_thresh)
self.rotation_tol = np.deg2rad(rotation_tol)
self.rotation_max_element = float(rotation_max_element)
self.rotation_interpolate = bool(rotation_interpolate)
self.rotation_disable = bool(rotation_disable)
self.rotation_disable_pos_curv = bool(rotation_disable_pos_curv)
self.rotation_remove_trans = bool(rotation_remove_trans)
self.trans_force_f_perp = trans_force_f_perp
self.forward_hessian = forward_hessian
# Regarding generation of initial orientation
self.bonds = bonds
self.N_hessian = N_hessian
# Bias
self.bias_rotation = bool(bias_rotation)
self.bias_rotation_a = 0.0
self.bias_translation = bool(bias_translation)
self.bias_gaussian_dot = float(bias_gaussian_dot)
self.write_orientations = write_orientations
restrict_steps = {
"direct": get_scale_max(self.rotation_max_element),
"fourier": None,
}
self.restrict_step = restrict_steps[rotation_method]
self._N = None
self._coords0 = None
self._energy0 = None
self._f0 = None
self._f1 = None
self.force_evals = 0
self.gaussians = list()
# Set dimer direction if given
self.N_raw = N_raw
if self.N_raw is not None:
msg = "Setting initial orientation from given 'N_raw'"
if isinstance(self.N_raw, str) and Path(self.N_raw).exists():
fn = self.N_raw
self.N_raw = np.loadtxt(fn)
msg = f"Read initial orientation from file '{fn}'"
N_raw = self.N_raw.copy()
self.N = N_raw
self.log(msg)
self.N_init = None
if seed is None:
# 2**32 - 1
seed = np.random.randint(4294967295)
np.random.seed(seed)
msg = f"Using seed {seed} to initialize the random number generator.\n"
print(msg)
self.log(msg)
@property
def N(self):
return self._N
@N.setter
def N(self, N_new):
N_new = np.array(N_new, dtype=float).flatten()
if self.rotation_remove_trans:
N_new = self.remove_translation(N_new)
N_new /= np.linalg.norm(N_new)
self._N = N_new
self._f1 = None
@property
def coords0(self):
return self._coords0
@coords0.setter
def coords0(self, coords0_new):
self._coords0 = coords0_new
self._energy0 = None
self._f0 = None
self._f1 = None
@property
def coords1(self):
return self.coords0 + self.length * self.N
@coords1.setter
def coords1(self, coords1_new):
N_new = coords1_new - self.coords0
self.N = N_new
self._f1 = None
@property
def energy0(self):
if self._energy0 is None:
results = self.calculator.get_energy(self.atoms, self.coords0)["energy"]
self._energy0 = results["energy"]
return self._energy0
@property
def f0(self):
if self._f0 is None:
results = self.calculator.get_forces(self.atoms, self.coords0)
self.force_evals += 1
self._f0 = results["forces"]
self._energy0 = results["energy"]
return self._f0
@property
def f1(self):
if self._f1 is None:
results = self.calculator.get_forces(self.atoms, self.coords1)
self.force_evals += 1
self._f1 = results["forces"]
return self._f1
@f1.setter
def f1(self, f1_new):
self._f1 = f1_new
@property
def f2(self):
"""Never calculated explicitly, but estimated from f0 and f1."""
return 2 * self.f0 - self.f1
@property
def can_bias_f1(self):
return (
self.bias_rotation
and (self.N_init is not None)
and (self.bias_rotation_a is not None)
and (self.bias_rotation_a > 0.0)
)
@property
def should_bias_f1(self):
"""May lead to calculation of f0 and/or f1 if present!"""
return self.can_bias_f1 and self.C > 0.0
@property
def can_bias_f0(self):
return self.bias_translation and (self.N is not None)
@property
def should_bias_f0(self):
"""May lead to calculation of f0 and/or f1 if present!"""
return self.can_bias_f0 and self.C > 0.0
@property
def f1_bias(self):
# Apply bias force to f1 if desired. Dont apply bias if N_init is not (yet)
# set. When N_raw was converged to a reasonable N_init we can add
# the bias.
assert self.bias_rotation_a >= 0.0, "This should not be negative!"
fN = self.bias_rotation_a * self.length * self.N.dot(self.N_init) * self.N_init
return fN
@property
def rot_force(self):
f1_perp = perp_comp(self.f1, self.N)
f2_perp = perp_comp(self.f2, self.N)
f_perp = f1_perp - f2_perp
# Don't bias rotational force if curvature is already negative
if self.should_bias_f1:
f1_bias_perp = perp_comp(self.f1_bias, self.N)
f_perp += f1_bias_perp
return f_perp
[docs]
def curvature(self, f1, f2, N):
"""Curvature of the mode represented by the dimer."""
return (f2 - f1).dot(N) / (2 * self.length)
@property
def C(self):
"""Shortcut for the curvature."""
return self.curvature(self.f1, self.f2, self.N)
[docs]
def get_gaussian_energies(self, coords, sum_=True):
energies = [gauss.energy(coords) for gauss in self.gaussians]
if sum_:
energies = sum(energies)
return energies
[docs]
def get_gaussian_forces(self, coords, sum_=True):
forces = [gauss.forces(coords) for gauss in self.gaussians]
if sum_:
forces = np.sum(forces, axis=0)
return forces
[docs]
def add_gaussian(
self, atoms, center, N, height=0.1, std=0.0529, max_cycles=50, dot_ref=None
):
# Create new gaussian object with default height that will be
# refined later.
gaussian = Gaussian(height=height, center=center, std=std, N=N)
if dot_ref is None:
dot_ref = self.bias_gaussian_dot
# Calculate real forces at inflection point of new gaussian
infl_coords = center + gaussian.std * N
infl_results = self.calculator.get_forces(atoms, infl_coords)
self.force_evals += 1
infl_forces = infl_results["forces"]
assert (
infl_forces.dot(N) < 0
), "We probably overstepped the TS. See Section 2.3 in the paper."
forces = self.get_gaussian_forces(infl_coords) + infl_forces
def get_dot(height):
"""Dot product of forces for a given height and orientation N."""
tmp_forces = forces.copy()
tmp_forces += gaussian.forces(infl_coords, height=height)
dot = tmp_forces.dot(N)
return dot
def can_break(dot):
"""Convergence indicator."""
return abs(dot - dot_ref) <= 1e-3
def bisect(
min_,
max_,
):
for i in range(max_cycles):
if abs(min_ - max_) <= 1e-10:
raise Exception("min_ and max_ became too similar!")
# Determie value at half of the internval
height = min_ + (max_ - min_) / 2
dot = get_dot(height)
if can_break(dot):
break
if dot > dot_ref:
max_ = height
elif dot < dot_ref:
min_ = height
return height
# Determine appropriate height
grow = 2
min_height = 0
assert get_dot(0) < dot_ref
for i in range(max_cycles):
dot = get_dot(height)
if can_break(dot):
break
if 0 < dot < dot_ref:
min_height = height
height *= grow
elif dot < dot_ref:
height *= grow
else:
height = bisect(min_height, height)
break
dot = get_dot(height)
gaussian.height = height
self.gaussians.append(gaussian)
self.log(f"Added gaussian with height={height:.6f}")
self.log(f"There are now {len(self.gaussians)} gaussians.")
return gaussian
[docs]
def get_N_raw_from_hessian(self, h5_fn, root=0):
with h5py.File(h5_fn, "r") as handle:
hessian = handle["hessian"][:]
w, v = np.linalg.eigh(hessian)
assert w[root] < -1e-3
N_raw = v[:, root]
return N_raw
[docs]
def set_N_raw(self, coords):
self.log("No initial orientation given. Generating one.")
if self.bonds is not None:
N_raw = get_weighted_bond_mode(self.bonds, coords.reshape(-1, 3))
msg = "weighted bond mode"
elif self.N_hessian is not None:
N_raw = self.get_N_raw_from_hessian(self.N_hessian)
msg = "first imaginary mode of HDF5 Hessian"
else:
msg = "random guess"
N_raw = np.random.rand(coords.size)
self.log(f"Obtained initial orientation from {msg}.")
# Make N_raw translationally invariant and normalize
self.N = N_raw
# Now we keep the normalized dimer orientation
self.N_raw = self.N
[docs]
def remove_translation(self, displacement):
# Average vector components over cartesian directions (x,y,z)
# Sum over each direction should equal zero if translationally invariant
average = displacement.reshape(-1, 3).mean(axis=0)
if max(abs(average)) > 1e-8:
self.log(
f"N-vector not translationally invariant. Removing average before normalization."
)
else:
return displacement
# Subtract the average component along each direction to make sum zero
invariant_displacement = (
displacement.reshape(-1, 3) - average[None, :]
).flatten()
return invariant_displacement
[docs]
def rotate_coords1(self, rad, theta):
"""Rotate dimer and produce new coords1."""
return self.coords0 + (self.N * np.cos(rad) + theta * np.sin(rad)) * self.length
[docs]
def direct_rotation(self, optimizer, prev_step):
rot_step = optimizer(self.rot_force, prev_step)
rot_step = self.restrict_step(rot_step)
# Strictly speaking rot_step should be constrained to conserve the desired
# dimer length (coords1 - coords0)*2. This step is unconstrained.
# Later on we calculate the actual step between the old coords1 and the new
# coords1 that have been reconstrained.
self.coords1 = self.coords1 + rot_step
[docs]
def fourier_rotation(self, optimizer, prev_step):
theta_dir = optimizer(self.rot_force, prev_step)
# Remove component that is parallel to N
theta_dir = theta_dir - theta_dir.dot(self.N) * self.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.
C = self.C
# Derivative of the curvature, Eq. (29) in [2]
# (f2 - f1) or -(f1 - f2)
dC = 2 * (self.f0 - self.f1).dot(theta) / self.length
rad_trial = -0.5 * np.arctan2(dC, 2 * abs(C))
# self.log(f"rad_trial={rad_trial:.2f}")
if np.abs(rad_trial) < self.rotation_tol:
self.log(f"rad_trial={rad_trial:.2f} below threshold. Breaking.")
raise RotationConverged
# Trial rotation for finite difference calculation of rotational force
# and rotational curvature.
coords1_trial = self.rotate_coords1(rad_trial, theta)
f1_trial = self.calculator.get_forces(self.atoms, coords1_trial)["forces"]
self.force_evals += 1
f2_trial = 2 * self.f0 - f1_trial
N_trial = make_unit_vec(coords1_trial, self.coords0)
C_trial = self.curvature(f1_trial, f2_trial, N_trial)
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)
# self.log(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)
# self.log("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})"
# )
# 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) < self.rotation_tol:
# self.log(f"rad_min={rad_min:.2f} below threshold. Breaking.")
raise RotationConverged
f1 = None
# Interpolate force at coords1_rot; see Eq. (12) in [4]
if self.rotation_interpolate:
f1 = (
np.sin(rad_trial - rad_min) / np.sin(rad_trial) * self.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))
* self.f0
)
self.coords1 = self.rotate_coords1(rad_min, theta)
self.f1 = f1
[docs]
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(gamma_mult=True)
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)."
)
[docs]
def update_orientation(self, coords):
# Generate random guess for the dimer orientation if not yet set
if self.N is None:
self.set_N_raw(coords)
# Refine N_raw to N_init if not yet done
if self.bias_rotation and self.N_init is None and self.C > 0.0:
# Run initial sweep with a much softer convergence threshold
self.log("Initial sweep to refine N_raw to N_init.")
self.do_dimer_rotations(10 * self.rotation_thresh)
if self.C > 0:
self.N_init = self.N
rot_rad = np.arccos(self.N_raw.dot(self.N_init))
rot_deg = np.rad2deg(rot_rad)
self.log(f"N_raw:\n\t{self.N_raw}")
self.log(f"Rotated N_raw by {rot_deg:.1f}° to N_init")
self.log(f"N_init:\n\t{self.N_init}")
C = self.C
self.log(f"Curvature after intial sweep is C={C:.6f}")
self.log("Determining proper scaling factor for bias potential.")
# Determine proper scaling factor for the quadratic bias potential
# from the current curvature.
scale_fact = 1.5
self.bias_rotation_a = scale_fact * self.C
self.log(f"Using a={scale_fact}*C={self.bias_rotation_a:.6f}")
else:
self.log(
f"Curvature after initial sweep C={self.C:.6f} is "
"already negative. Not setting N_init and bias_rotation_a!"
)
self.do_dimer_rotations()
[docs]
def get_forces(self, atoms, coords):
self.atoms = atoms
self.coords0 = coords
try:
N_backup = self.N.copy()
except AttributeError:
N_backup = None
if not self.rotation_disable:
self.update_orientation(coords)
if (N_backup is not None) and self.rotation_disable_pos_curv and self.C > 0:
self.log(
"Rotation did not yield a negative curvature. "
"Restoring previous unrotated N."
)
self.N = N_backup
# Now we (have an updated self.N and) can do the force projections
N = self.N
# self.log(f"Orientation N:\n\t{N}")
# Save orientation N in human-readable format, aka .trj
# file in Angstrom.
if self.write_orientations:
trj_str = get_tangent_trj_str(atoms, coords, N)
trj_fn = self.make_fn("N.trj")
with open(trj_fn, "w") as handle:
handle.write(trj_str)
self.log(f"Wrote current orientation animation to '{trj_fn}'")
# Always save orientation in Bohr
N_fn = self.make_fn("N")
np.savetxt(N_fn, N)
energy = self.energy0
self.log(f"\tenergy={self.energy0:.8f} au")
f0 = self.f0
if self.should_bias_f0:
self.log("Biasing translation forces")
self.log(f"There are currently {len(self.gaussians)} gaussians present.")
bias_energy = self.get_gaussian_energies(coords)
energy += bias_energy
bias_forces = self.get_gaussian_forces(coords, sum_=False)
try:
bias_norms = np.linalg.norm(bias_forces, axis=1)
bias_norms_str = np.array2string(bias_norms, precision=4)
self.log(f"\tnorm(bias_forces)={bias_norms_str}")
except np.AxisError:
self.log("Skipping calculation of norm(bias_forces)")
f0 += np.sum(bias_forces, axis=0)
norm_f0 = np.linalg.norm(f0)
self.log(f"\tnorm(forces)={norm_f0:.6f}")
f_parallel = f0.dot(N) * N
norm_parallel = np.linalg.norm(f_parallel)
self.log(f"\tnorm(forces_parallel)={norm_parallel:.6f}")
# self.log(f"\tforce_parallel:\n\t{f_parallel}")
f_perp = f0 - f_parallel
norm_perp = np.linalg.norm(f_perp)
self.log(f"\tnorm(forces_perp)={norm_perp:.6f}")
# self.log(f"\tforce_perp:\n\t{f_perp}")
# Only return perpendicular component when curvature is negative
f_tran = -f_parallel
curv_str = "positive" if self.C > 0 else "negative"
force_str = "reversed parallel component of"
if (self.C < 0) or self.trans_force_f_perp:
f_tran += f_perp
force_str = "full"
self.log(f"Curvature is {curv_str}. Returning {force_str} f_tran.")
# self.log(f"\tf_tran:\n\t{f_tran}")
self.log()
self.calc_counter += 1
results = {"energy": energy, "forces": f_tran}
return results
[docs]
def get_hessian(self, atoms, coords):
if not self.forward_hessian:
raise Exception("Actual Hessian method not forwarded by Dimer calculator!")
results = self.calculator.get_hessian(atoms, coords)
return results