from typing import IO, Optional, Union
import numpy as np
from ase import Atoms
from ase.optimize.optimize import Dynamics
def subtract_projection(a, b):
'''returns new vector that removes vector a's projection vector b. Is
also equivalent to the vector rejection.'''
aout = a - np.vdot(a, b) / np.vdot(b, b) * b
return aout
def normalize(a):
'''Makes a unit vector out of a vector'''
return a / np.linalg.norm(a)
[docs]class ContourExploration(Dynamics):
def __init__(
self,
atoms: Atoms,
maxstep: float = 0.5,
parallel_drift: float = 0.1,
energy_target: Optional[float] = None,
angle_limit: Optional[float] = 20.0,
potentiostat_step_scale: Optional[float] = None,
remove_translation: bool = False,
use_frenet_serret: bool = True,
initialization_step_scale: float = 1e-2,
use_target_shift: bool = True,
target_shift_previous_steps: int = 10,
use_tangent_curvature: bool = False,
rng=np.random,
force_consistent: Optional[bool] = None,
trajectory: Optional[str] = None,
logfile: Optional[Union[IO, str]] = None,
append_trajectory: bool = False,
loginterval: int = 1,
):
"""Contour Exploration object.
Parameters:
atoms: Atoms object
The Atoms object to operate on. Atomic velocities are required for
the method. If the atoms object does not contain velocities,
random ones will be applied.
maxstep: float
Used to set the maximum distance an atom can move per
iteration (default value is 0.5 Å).
parallel_drift: float
The fraction of the update step that is parallel to the contour but
in a random direction. Used to break symmetries.
energy_target: float
The total system potential energy for that the potentiostat attepts
to maintain. (defaults the initial potential energy)
angle_limit: float or None
Limits the stepsize to a maximum change of direction angle using the
curvature. Gives a scale-free means of tuning the stepsize on the
fly. Typically less than 30 degrees gives reasonable results but
lower angle limits result in higher potentiostatic accuracy. Units
of degrees. (default 20°)
potentiostat_step_scale: float or None
Scales the size of the potentiostat step. The potentiostat step is
determined by linear extrapolation from the current potential energy
to the target_energy with the current forces. A
potentiostat_step_scale > 1.0 overcorrects and < 1.0
undercorrects. By default, a simple heuristic is used to selected
the valued based on the parallel_drift. (default None)
remove_translation: boolean
When True, the net momentum is removed at each step. Improves
potentiostatic accuracy slightly for bulk systems but should not be
used with constraints. (default False)
use_frenet_serret: Bool
Controls whether or not the Taylor expansion of the Frenet-Serret
formulas for curved path extrapolation are used. Required for using
angle_limit based step scalling. (default True)
initialization_step_scale: float
Controls the scale of the initial step as a multiple of maxstep.
(default 1e-2)
use_target_shift: boolean
Enables shifting of the potentiostat target to compensate for
systematic undercorrection or overcorrection by the potentiostat.
Uses the average of the *target_shift_previous_steps* to prevent
coupled occilations. (default True)
target_shift_previous_steps: int
The number of pevious steps to average when using use_target_shift.
(default 10)
use_tangent_curvature: boolean
Use the velocity unit tangent rather than the contour normals from
forces to compute the curvature. Usually not as accurate.
(default False)
rng: a random number generator
Lets users control the random number generator for the
parallel_drift vector. (default numpy.random)
force_consistent: boolean
(default None)
trajectory: Trajectory object or str (optional)
Attach trajectory object. If *trajectory* is a string a
Trajectory will be constructed. Default: None.
logfile: file object or str (optional)
If *logfile* is a string, a file with that name will be opened.
Use '-' for stdout. Default: None.
loginterval: int (optional)
Only write a log line for every *loginterval* time steps.
Default: 1
append_trajectory: boolean
Defaults to False, which causes the trajectory file to be
overwriten each time the dynamics is restarted from scratch.
If True, the new structures are appended to the trajectory
file instead.
"""
if potentiostat_step_scale is None:
# a heuristic guess since most systems will overshoot when there is
# drift
self.potentiostat_step_scale = 1.1 + 0.6 * parallel_drift
else:
self.potentiostat_step_scale = potentiostat_step_scale
self.rng = rng
self.remove_translation = remove_translation
self.use_frenet_serret = use_frenet_serret
self.use_tangent_curvature = use_tangent_curvature
self.initialization_step_scale = initialization_step_scale
self.maxstep = maxstep
self.angle_limit = angle_limit
self.parallel_drift = parallel_drift
self.use_target_shift = use_target_shift
# These will be populated once self.step() is called, but could be set
# after instantiating with ce = ContourExploration(...) like so:
# ce.Nold = Nold
# ce.r_old = atoms_old.get_positions()
# ce.Told = Told
# to resume a previous contour trajectory.
self.T = None
self.Told = None
self.N = None
self.Nold = None
self.r_old = None
self.r = None
if energy_target is None:
self.energy_target = atoms.get_potential_energy(
force_consistent=True)
else:
self.energy_target = energy_target
# Initizing the previous steps at the target energy slows
# target_shifting untill the system has had
# 'target_shift_previous_steps' steps to equilibrate and should prevent
# occilations. These need to be initialized before the initialize_old
# step to prevent a crash
self.previous_energies = np.full(target_shift_previous_steps,
self.energy_target)
# these first two are purely for logging,
# auto scaling will still occur
# and curvature will still be found if use_frenet_serret == True
self.step_size = 0.0
self.curvature = 0
# loginterval exists for the MolecularDynamics class but not for
# the more general Dynamics class
Dynamics.__init__(self, atoms,
logfile, trajectory, # loginterval,
append_trajectory=append_trajectory,
)
self._actual_atoms = atoms
# we need velocities or NaNs will be produced,
# if none are provided we make random ones
velocities = self._actual_atoms.get_velocities()
if np.linalg.norm(velocities) < 1e-6:
# we have to pass dimension since atoms are not yet stored
atoms.set_velocities(self.rand_vect())
# Required stuff for Dynamics
def todict(self):
return {'type': 'contour-exploration',
'dyn-type': self.__class__.__name__,
'stepsize': self.step_size}
def run(self, steps=50):
""" Call Dynamics.run and adjust max_steps """
return Dynamics.run(self, steps=steps)
def log(self):
if self.logfile is not None:
# name = self.__class__.__name__
if self.nsteps == 0:
args = (
"Step",
"Energy_Target",
"Energy",
"Curvature",
"Step_Size",
"Energy_Deviation_per_atom")
msg = "# %4s %15s %15s %12s %12s %15s\n" % args
self.logfile.write(msg)
e = self._actual_atoms.get_potential_energy(force_consistent=True)
dev_per_atom = (e - self.energy_target) / len(self._actual_atoms)
args = (
self.nsteps,
self.energy_target,
e,
self.curvature,
self.step_size,
dev_per_atom)
msg = "%6d %15.6f %15.6f %12.6f %12.6f %24.9f\n" % args
self.logfile.write(msg)
self.logfile.flush()
def rand_vect(self):
'''Returns a random (Natoms,3) vector'''
vect = self.rng.random((len(self._actual_atoms), 3)) - 0.5
return vect
def create_drift_unit_vector(self, N, T):
'''Creates a random drift unit vector with no projection on N or T and
with out a net translation so systems don't wander'''
drift = self.rand_vect()
drift = subtract_projection(drift, N)
drift = subtract_projection(drift, T)
# removes net translation, so systems don't wander
drift = drift - drift.sum(axis=0) / len(self._actual_atoms)
D = normalize(drift)
return D
def compute_step_contributions(self, potentiostat_step_size):
'''Computes the orthogonal component sizes of the step so that the net
step obeys the smaller of step_size or maxstep.'''
if abs(potentiostat_step_size) < self.step_size:
delta_s_perpendicular = potentiostat_step_size
contour_step_size = np.sqrt(
self.step_size**2 - potentiostat_step_size**2)
delta_s_parallel = np.sqrt(
1 - self.parallel_drift**2) * contour_step_size
delta_s_drift = contour_step_size * self.parallel_drift
else:
# in this case all priority goes to potentiostat terms
delta_s_parallel = 0.0
delta_s_drift = 0.0
delta_s_perpendicular = np.sign(
potentiostat_step_size) * self.step_size
return delta_s_perpendicular, delta_s_parallel, delta_s_drift
def _compute_update_without_fs(self, potentiostat_step_size, scale=1.0):
'''Only uses the forces to compute an orthogonal update vector'''
# Without the use of curvature there is no way to estimate the
# limiting step size
self.step_size = self.maxstep * scale
delta_s_perpendicular, delta_s_parallel, delta_s_drift = \
self.compute_step_contributions(
potentiostat_step_size)
dr_perpendicular = self.N * delta_s_perpendicular
dr_parallel = delta_s_parallel * self.T
D = self.create_drift_unit_vector(self.N, self.T)
dr_drift = D * delta_s_drift
dr = dr_parallel + dr_drift + dr_perpendicular
dr = self.step_size * normalize(dr)
return dr
def _compute_update_with_fs(self, potentiostat_step_size):
'''Uses the Frenet–Serret formulas to perform curvature based
extrapolation to compute the update vector'''
# this should keep the dr clear of the constraints
# by using the actual change, not a velocity vector
delta_r = self.r - self.rold
delta_s = np.linalg.norm(delta_r)
# approximation of delta_s we use this incase an adaptive step_size
# algo get used
delta_T = self.T - self.Told
delta_N = self.N - self.Nold
dTds = delta_T / delta_s
dNds = delta_N / delta_s
if self.use_tangent_curvature:
curvature = np.linalg.norm(dTds)
# on a perfect trajectory, the normal can be computed this way,
# But the normal should always be tied to forces
# N = dTds / curvature
else:
# normals are better since they are fixed to the reality of
# forces. I see smaller forces and energy errors in bulk systems
# using the normals for curvature
curvature = np.linalg.norm(dNds)
self.curvature = curvature
if self.angle_limit is not None:
phi = np.pi / 180 * self.angle_limit
self.step_size = np.sqrt(2 - 2 * np.cos(phi)) / curvature
self.step_size = min(self.step_size, self.maxstep)
# now we can compute a safe step
delta_s_perpendicular, delta_s_parallel, delta_s_drift = \
self.compute_step_contributions(
potentiostat_step_size)
N_guess = self.N + dNds * delta_s_parallel
T_guess = self.T + dTds * delta_s_parallel
# the extrapolation is good at keeping N_guess and T_guess
# orthogonal but not normalized:
N_guess = normalize(N_guess)
T_guess = normalize(T_guess)
dr_perpendicular = delta_s_perpendicular * (N_guess)
dr_parallel = delta_s_parallel * self.T * \
(1 - (delta_s_parallel * curvature)**2 / 6.0) \
+ self.N * (curvature / 2.0) * delta_s_parallel**2
D = self.create_drift_unit_vector(N_guess, T_guess)
dr_drift = D * delta_s_drift
# combine the components
dr = dr_perpendicular + dr_parallel + dr_drift
dr = self.step_size * normalize(dr)
# because we guess our orthonormalization directions,
# we renormalize to ensure a correct step size
return dr
def update_previous_energies(self, energy):
'''Updates the energy history in self.previous_energies to include the
current energy.'''
# np.roll shifts the values to keep nice sequential ordering.
self.previous_energies = np.roll(self.previous_energies, 1)
self.previous_energies[0] = energy
def compute_potentiostat_step_size(self, forces, energy):
'''Computes the potentiostat step size by linear extrapolation of the
potential energy using the forces. The step size can be positive or
negative depending on whether or not the energy is too high or too low.
'''
if self.use_target_shift:
target_shift = self.energy_target - np.mean(self.previous_energies)
else:
target_shift = 0.0
# deltaU is the potential error that will be corrected for
deltaU = energy - (self.energy_target + target_shift)
f_norm = np.linalg.norm(forces)
# can be positive or negative
potentiostat_step_size = (deltaU / f_norm) * \
self.potentiostat_step_scale
return potentiostat_step_size
def step(self, f=None):
atoms = self._actual_atoms
if f is None:
f = atoms.get_forces()
# get the velocity vector and old kinetic energy for momentum rescaling
velocities = atoms.get_velocities()
KEold = atoms.get_kinetic_energy()
energy = atoms.get_potential_energy(force_consistent=True)
self.update_previous_energies(energy)
potentiostat_step_size = self.compute_potentiostat_step_size(f, energy)
self.N = normalize(f)
self.r = atoms.get_positions()
# remove velocity projection on forces
v_parallel = subtract_projection(velocities, self.N)
self.T = normalize(v_parallel)
if self.use_frenet_serret:
if self.Nold is not None and self.Told is not None:
dr = self._compute_update_with_fs(potentiostat_step_size)
else:
# we must have the old positions and vectors for an FS step
# if we don't, we can only do a small step
dr = self._compute_update_without_fs(
potentiostat_step_size,
scale=self.initialization_step_scale)
else: # of course we can run less accuratly without FS.
dr = self._compute_update_without_fs(potentiostat_step_size)
# now that dr is done, we check if there is translation
if self.remove_translation:
net_motion = dr.sum(axis=0) / len(atoms)
# print(net_motion)
dr = dr - net_motion
dr_unit = dr / np.linalg.norm(dr)
dr = dr_unit * self.step_size
# save old positions before update
self.Nold = self.N
self.rold = self.r
self.Told = self.T
# if we have constraints then this will do the first part of the
# RATTLE algorithm:
# If we can avoid using momenta, this will be simpler.
masses = atoms.get_masses()[:, np.newaxis]
atoms.set_positions(self.r + dr)
new_momenta = (atoms.get_positions() - self.r) * masses # / self.dt
# We need to store the momenta on the atoms before calculating
# the forces, as in a parallel Asap calculation atoms may
# migrate during force calculations, and the momenta need to
# migrate along with the atoms.
atoms.set_momenta(new_momenta, apply_constraint=False)
# Now we get the new forces!
f = atoms.get_forces(md=True)
# I don't really know if removing md=True from above will break
# compatibility with RATTLE, leaving it alone for now.
f_constrained = atoms.get_forces()
# but this projection needs the forces to be consistent with the
# constraints. We have to set the new velocities perpendicular so they
# get logged properly in the trajectory files.
vnew = subtract_projection(atoms.get_velocities(), f_constrained)
# using the md = True forces like this:
# vnew = subtract_projection(atoms.get_velocities(), f)
# will not work with constraints
atoms.set_velocities(vnew)
# rescaling momentum to maintain constant kinetic energy.
KEnew = atoms.get_kinetic_energy()
Ms = np.sqrt(KEold / KEnew) # Ms = Momentum_scale
atoms.set_momenta(Ms * atoms.get_momenta())
# Normally this would be the second part of RATTLE
# will be done here like this:
# atoms.set_momenta(atoms.get_momenta() + 0.5 * self.dt * f)
return f