Skip to main content
`Ctrl`+`K`

#

# Source code for ase.md.contour_exploration

    # fmt: off

    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)

    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, gradient):
            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.normal(size=(len(self._actual_atoms), 3))
            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