from copy import copy, deepcopy
import sys
import numpy as np
from riigid.optimization_step import OptimizationStep
from riigid.optimizer.optimizer import Optimizer
[docs]
class GDWAS(Optimizer):
"""RIIGID optimizer: Gradient Descent with Adaptive Stepsize
Calculates force and torque on each fragment and moves them accordingly (like rigid bodies).
The goal of this optimizer is to converge to the nearest minimum. To achieve this, in case the
last update of the atomic positions lead to higher energies, it is dropped and the calculation
is continued from the previous structure. This way, climbing upwards in the potential energy
surface (PES) is prohibited and the probability for jumping into a different potential well is
lowered.
Additionally, a hard cap on the movement of individual fragments is enforced by max_trans and max_rot. This again
helps in reducing the probability of the structure leaving a local minimum.
The stepsize is adaptive and can change for two reasons:
1.) If the last update performed on the atomic positions lowered the total energy (i.e a larger
negative number), the stepsize is increased by multiplying with stepsize_factor_up.
On the other hand, if the energy increased, the stepsize is lowered by multiplying with
stepsize_factor_dn.
2.) As described above, a hard cap on the movement of individual fragments is enforced by max_trans and max_rot.
If a fragment is found to translate more than max_trans, the stepsize is lowered in such a way,
that the maximal displacement of the fragments is exactly max_trans.
On the other hand, if a fragment is rotated by more than max_rot, the stepsize is adapted, such that
the maximal rotation (i.e. the rotation of the fragment that rotates the most) is max_rot.
The user never needs to interact directly with the stepsize. Instead, via max_trans_0 and max_rot_0, the user
specifies how much the fragments (at most) shall move in the first optimization step. The stepsize
is then initialized accordingly and from there on adapted automatically, as described above.
Attributes
----------
optimization_history: list of riigid.Optimization_Step
The history of the optimization, which shall be checked for convergence.
iteration: int
Counts the number of finished optimization steps
stepsize: number
Timestep; [Da*Å**2/eV]
stepsize_factor_up: number > 1
Increase stepsize by this factor, if last optimization step lowered the total energy
stepsize_factor_dn: number < 1
Decrease stepsize by this factor, if last optimization step increased the total energy
max_trans: number
The maximum distance fragments are allowed to translate per optimization step; [Å]
max_rot: number
The maximum angle fragments are allowed to rotate per optimization step; [°]
max_trans_0, max_rot_0: number, number
In the first optimization step, the stepsize is chosen such that the fragment(s) translating/rotating
the most, translate/rotate by (one of) these value; [Å], [°]
start_with_random_step: bool
Shall the fragments forming the structure be randomly translated and rotated before the
first optimization step? This can be used to escape a saddle point starting-geometry.
The attributes with '_r0' at the end further specify this random step before the first
optimization step.
displacement_r0: number
How far shall the fragments be translated; [Å]
angle_r0: number
How much shall the fragments be rotated; [°]
respect_restrictions_r0: bool
If True, fragment.allowed_translation/rotation is respected.
If False, rotation and translation in arbitrary directions is allowed temporarily.
(After the random step, the restrictions are respected again.)
seed_r0: int
The random seed used to generate the translation directions and rotation axes
max_iter: int
The maximal number of optimization steps to be performed.
If the calculation does not converge within this limit, it is stopped.
start_structure: riigid.Structure
The structure to be optimized
calculator : ase.calculators.calculator.Calculator
The used ASE calculator object
convergence_criterion : riigid.convergence.Criterion
The used convergence criterion object
current_structure: riigid.Structure
The structure currently used by the optimizer
current_energy: number
The energy of current_structure; [eV]
current_forces: numpy.ndarray of shape (n_atoms_in_current_structure, 3)
The forces in current_structure; [eV/Å]
"""
def __init__(
self,
stepsize_factor_up=1.2,
stepsize_factor_dn=0.2,
max_trans=1.0,
max_rot=30.0,
max_trans_0=0.1,
max_rot_0=3.0,
start_with_random_step=True,
displacement_r0=0.01,
angle_r0=0.3,
respect_restrictions_r0=True,
seed_r0=1234,
max_iter=500,
):
"""Initialize the GDWAS optimizer.
Parameters
----------
stepsize_factor_up: number > 1, default: 1.2
Increase stepsize by this factor, if last optimization step lowered the total energy
stepsize_factor_dn: number < 1, default: 0.2
Decrease stepsize by this factor, if last optimization step increased the total energy
max_trans: number, default: 1.0
The maximum distance fragments are allowed to translate per optimization step; [Å]
max_rot: number, default: 30.0
The maximum angle fragments are allowed to rotate per optimization step; [°]
max_trans_0, max_rot_0: number, number, default: 0.1, 3.0
In the first optimization step, the stepsize is chosen such that the fragment(s) translating/rotating
the most, translate/rotate by (one of) these value; [Å], [°]
start_with_random_step: bool, default:True
Shall the fragments forming the structure be randomly translated and rotated before the
first optimization step? This can be used to escape a saddle point starting-geometry.
The parameters with '_r0' at the end further specify this random step before the first
optimization step.
displacement_r0: number, default:0.01
How far shall the fragments be translated; [Å]
angle_r0: number, default:0.3
How much shall the fragments be rotated; [°]
respect_restrictions_r0: bool, default:True
If True, fragment.allowed_translation/rotation is respected.
If False, rotation and translation in arbitrary directions is allowed temporarily.
(After the random step, the restrictions are respected again.)
seed_r0: int, default:1234
The random seed used to generate the translation directions and rotation axes
max_iter: int, default: 500
The maximal number of optimization steps to be performed.
If the calculation does not converge within this limit, it is stopped.
"""
super().__init__(max_iter=max_iter)
self.stepsize = 100 # initial stepsize, value doesn't matter, is later adapted to max_step_0
self.stepsize_factor_up = stepsize_factor_up
self.stepsize_factor_dn = stepsize_factor_dn
self.max_trans = max_trans
self.max_rot = max_rot
self.max_trans_0 = max_trans_0
self.max_rot_0 = max_rot_0
self.start_with_random_step = start_with_random_step
self.displacement_r0 = displacement_r0
self.angle_r0 = angle_r0
self.respect_restrictions_r0 = respect_restrictions_r0
self.seed_r0 = seed_r0
[docs]
def run(self, start_structure, calculator, convergence_criterion, callback=None):
"""Let the optimizer run its optimization on the structure.
Parameters
----------
start_structure: riigid.Structure
The structure to be optimized
calculator : ase.calculators.calculator.Calculator
The used ASE calculator object
convergence_criterion : riigid.convergence.criterion
The used convergence criterion object
callback : function, default:None
A callback function can be used to safe the optimization progress after each step.
"""
print("Starting GDWAS optimization...")
sys.stdout.flush() # Flush the output immediately
self.start_structure = start_structure
self.calculator = calculator
self.convergence_criterion = convergence_criterion
while not convergence_criterion.is_converged and self.iteration < self.max_iter:
print("Starting step " + str(self.iteration))
sys.stdout.flush()
# Get current structure (starting structure or updated structure from last step)
if self.iteration == 0:
self.current_structure = deepcopy(start_structure)
else:
self.current_structure = deepcopy(
self.optimization_history[-1].updated_structure
)
# Before first calculation, perform a random step, if start_with_random_step==True
if self.start_with_random_step and self.iteration == 0:
print("Doing random step before first calculation.")
sys.stdout.flush()
self.current_structure.move_random_step(
displacement=self.displacement_r0,
angle=self.angle_r0,
respect_restrictions=self.respect_restrictions_r0,
seed=self.seed_r0,
)
# Do Calculation to get energy and forces
(
self.current_energy,
self.current_forces,
) = self.current_structure.calculate_energy_and_forces(
calculator=calculator
)
# Adapt stepsize to energy change
self.adapt_stepsize_to_energy_change()
# undo last update if energy got larger
self.drop_last_step_if_energy_got_larger()
if (
self.iteration == 0
): # Initialize stepsize accordingly to max_rot/trans_0
self.initialize_stepsize_in_first_iteration()
else: # Adapt stepsize again, if necessary
self.adapt_stepsize_to_prevent_too_large_steps()
# Move atoms
updated_structure = deepcopy(self.current_structure)
updated_structure.move(forces=self.current_forces, stepsize=self.stepsize)
# Add Optimization step to history
new_step = OptimizationStep(
structure=self.current_structure,
forces_on_atoms=self.current_forces,
energy=self.current_energy,
updated_structure=updated_structure,
)
self.optimization_history.append(new_step)
# Check for convergence
self.convergence_criterion.check(
optimization_history=self.optimization_history
)
# Log some info (Note: We never know, if last step was successful or not.)
# Steps, that we know were not successful are not stored in optimization history
print(
"Created updated structure of step "
+ str(self.iteration)
+ " with stepsize: "
+ str(self.stepsize)
)
print(
"Successful/Finished Steps: "
+ str(len(self.optimization_history) - 1)
+ "/"
+ str(self.iteration + 1)
)
sys.stdout.flush()
# Prepare next iteration
self.iteration += 1
# If a callback function was provided, execute it. Useful to save data after every step
if callback is not None:
callback()
self.print_reason_for_end_of_optimization()
[docs]
def adapt_stepsize_to_energy_change(self):
"""Adapt the stepsize according to the last update step.
If the last update performed on the atomic positions lowered the total energy (i.e a larger
negative number), the stepsize is increased by multiplying with stepsize_factor_up.
On the other hand, if the energy increased, the stepsize is lowered by multiplying with
stepsize_factor_dn.
"""
if self.iteration != 0:
prev_energy = self.optimization_history[-1].energy
current_energy = self.current_energy
# Check if energy got smaller
if current_energy < prev_energy: # Yes -> make stepsize larger
self.stepsize *= self.stepsize_factor_up
else: # No -> make stepsize smaller
self.stepsize *= self.stepsize_factor_dn
[docs]
def drop_last_step_if_energy_got_larger(self):
"""Drop last update step if it increased the total energy."""
if self.iteration != 0:
prev_energy = self.optimization_history[-1].energy
current_energy = self.current_energy
if prev_energy < current_energy:
# Return to structure, forces and energy, previous to the last update
self.current_structure = deepcopy(
self.optimization_history[-1].structure
)
self.current_forces = deepcopy(
self.optimization_history[-1].forces_on_atoms
)
self.current_energy = deepcopy(self.optimization_history[-1].energy)
self.optimization_history.pop()
print("Previous step made energy larger and was undone!")
sys.stdout.flush()
[docs]
def adapt_stepsize_to_prevent_too_large_steps(self):
"""Prevent too large movement of the fragments.
If a fragment is found to translate more than self.max_trans, the stepsize is lowered in such a way,
that the maximal displacement of the fragments is exactly self.max_trans.
On the other hand, if a fragment is rotated by more than self.max_rot, the stepsize is adapted, such that
the maximal rotation (i.e. the rotation of the fragment that rotates the most) is self.max_rot.
(The weaker condition is fulfilled, i.e. the one with less movement/smaller stepsize.)
Note
----
Translation distances and rotation angles of fragments are directly proportional to the stepsize.
This is used here to turn down too large stepsizes.
"""
# Given the current stepsize, find the maximal rotation/translation of the fragments
test_structure = deepcopy(self.current_structure)
_, angles, shifts = test_structure.move(
forces=self.current_forces, stepsize=self.stepsize
)
max_found_translation_distance = np.max(
[np.linalg.norm(shift) for shift in shifts]
)
max_found_angle = np.max(np.abs(angles))
factor_rot = self.max_rot / max_found_angle
factor_trans = self.max_trans / max_found_translation_distance
factor = min([factor_rot, factor_trans])
factor = min([1.0, factor]) # this function never makes the stepsize larger
self.stepsize *= factor
[docs]
def initialize_stepsize_in_first_iteration(self):
"""Initialize the stepsize in the first iteration of the optimization.
During the first iteration of the optimizations, this function adapts the stepsize
such that the maximum displacement of the fragments is exactly self.max_trans_0 or the maximum
rotation is exactly self.max_rot_0.
(The weaker condition is fulfilled, i.e. the one with less movement/smaller stepsize.)
Note
----
Translation distances and rotation angles of fragments are directly proportional to the stepsize.
This is used here to turn down too large stepsizes.
"""
# Given the current stepsize, find the maximal rotation/translation of the fragments
test_structure = deepcopy(self.current_structure)
_, angles, shifts = test_structure.move(
forces=self.current_forces, stepsize=self.stepsize
)
max_found_translation_distance = np.max(
[np.linalg.norm(shift) for shift in shifts]
)
max_found_angle = np.max(np.abs(angles))
factor_rot = self.max_rot_0 / max_found_angle
factor_trans = self.max_trans_0 / max_found_translation_distance
factor = min([factor_rot, factor_trans])
self.stepsize *= factor