File: climbfixinternals.py

package info (click to toggle)
python-ase 3.26.0-3
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 15,484 kB
  • sloc: python: 148,112; xml: 2,728; makefile: 110; javascript: 47
file content (213 lines) | stat: -rw-r--r-- 8,979 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
# fmt: off

from typing import IO, Any, Dict, List, Optional, Type, Union

from numpy.linalg import norm

from ase import Atoms
from ase.constraints import FixInternals
from ase.optimize.bfgs import BFGS
from ase.optimize.optimize import Optimizer


class BFGSClimbFixInternals(BFGS):
    """Class for transition state search and optimization

    Climbs the 1D reaction coordinate defined as constrained internal coordinate
    via the :class:`~ase.constraints.FixInternals` class while minimizing all
    remaining degrees of freedom.

    Details: Two optimizers, 'A' and 'B', are applied orthogonal to each other.
    Optimizer 'A' climbs the constrained coordinate while optimizer 'B'
    optimizes the remaining degrees of freedom after each climbing step.
    Optimizer 'A' uses the BFGS algorithm to climb along the projected force of
    the selected constraint. Optimizer 'B' can be any ASE optimizer
    (default: BFGS).

    In combination with other constraints, the order of constraints matters.
    Generally, the FixInternals constraint should come first in the list of
    constraints.
    This has been tested with the :class:`~ase.constraints.FixAtoms` constraint.

    Inspired by concepts described by P. N. Plessow. [1]_

    .. [1] Plessow, P. N., Efficient Transition State Optimization of Periodic
           Structures through Automated Relaxed Potential Energy Surface Scans.
           J. Chem. Theory Comput. 2018, 14 (2), 981–990.
           https://doi.org/10.1021/acs.jctc.7b01070.

    .. note::
       Convergence is based on 'fmax' of the total forces which is the sum of
       the projected forces and the forces of the remaining degrees of freedom.
       This value is logged in the 'logfile'. Optimizer 'B' logs 'fmax' of the
       remaining degrees of freedom without the projected forces. The projected
       forces can be inspected using the :meth:`get_projected_forces` method:

       >>> for _ in dyn.irun():
       ...     projected_forces = dyn.get_projected_forces()

    Example
    -------
    .. literalinclude:: ../../ase/test/optimize/test_climb_fix_internals.py
       :end-before: # end example for documentation
    """

    def __init__(
        self,
        atoms: Atoms,
        restart: Optional[str] = None,
        logfile: Union[IO, str] = '-',
        trajectory: Optional[str] = None,
        maxstep: Optional[float] = None,
        alpha: Optional[float] = None,
        climb_coordinate: Optional[List[FixInternals]] = None,
        optB: Type[Optimizer] = BFGS,
        optB_kwargs: Optional[Dict[str, Any]] = None,
        optB_fmax: float = 0.05,
        optB_fmax_scaling: float = 0.0,
        **kwargs,
    ):
        """Allowed parameters are similar to the parent class
        :class:`~ase.optimize.bfgs.BFGS` with the following additions:

        Parameters
        ----------
        climb_coordinate: list
            Specifies which subconstraint of the
            :class:`~ase.constraints.FixInternals` constraint is to be climbed.
            Provide the corresponding nested list of indices
            (including coefficients in the case of Combo constraints).
            See the example above.

        optB: any ASE optimizer, optional
            Optimizer 'B' for optimization of the remaining degrees of freedom
            after each climbing step.

        optB_kwargs: dict, optional
            Specifies keyword arguments to be passed to optimizer 'B' at its
            initialization. By default, optimizer 'B' writes a logfile and
            trajectory (optB_{...}.log, optB_{...}.traj) where {...} is the
            current value of the ``climb_coordinate``. Set ``logfile`` to '-'
            for console output. Set ``trajectory`` to 'None' to suppress
            writing of the trajectory file.

        optB_fmax: float, optional
            Specifies the convergence criterion 'fmax' of optimizer 'B'.

        optB_fmax_scaling: float, optional
            Scaling factor to dynamically tighten 'fmax' of optimizer 'B' to
            the value of ``optB_fmax`` when close to convergence.
            Can speed up the climbing process. The scaling formula is

            'fmax' = ``optB_fmax`` + ``optB_fmax_scaling``
            :math:`\\cdot` norm_of_projected_forces

            The final optimization with optimizer 'B' is
            performed with ``optB_fmax`` independent of ``optB_fmax_scaling``.
        """
        self.targetvalue = None  # may be assigned during restart in self.read()
        super().__init__(atoms, restart=restart, logfile=logfile,
                         trajectory=trajectory, maxstep=maxstep,
                         alpha=alpha, **kwargs)

        self.constr2climb = get_constr2climb(
            self.optimizable.atoms, climb_coordinate)
        self.targetvalue = self.targetvalue or self.constr2climb.targetvalue

        self.optB = optB
        self.optB_kwargs = optB_kwargs or {}
        self.optB_fmax = optB_fmax
        self.scaling = optB_fmax_scaling
        # log optimizer 'B' in logfiles named after current value of constraint
        self.autolog = 'logfile' not in self.optB_kwargs
        self.autotraj = 'trajectory' not in self.optB_kwargs

    def Nx3(self, array):
        return array.reshape(-1, 3)

    def read(self):
        (self.H, self.pos0, self.forces0, self.maxstep,
         self.targetvalue) = self.load()

    def step(self):
        self.relax_remaining_dof()  # optimization with optimizer 'B'

        pos, dpos = self.pretend2climb()  # with optimizer 'A'
        self.update_positions_and_targetvalue(pos, dpos)  # obey other constr.

        self.dump((self.H, self.pos0, self.forces0, self.maxstep,
                   self.targetvalue))

    def pretend2climb(self):
        """Get directions for climbing and climb with optimizer 'A'."""
        proj_forces = self.get_projected_forces()
        pos = self.optimizable.get_x()
        dpos, steplengths = self.prepare_step(pos, proj_forces)
        dpos = self.determine_step(dpos, steplengths)
        return pos, dpos

    def update_positions_and_targetvalue(self, pos, dpos):
        """Adjust constrained targetvalue of constraint and update positions."""
        self.constr2climb.adjust_positions(
            self.Nx3(pos), self.Nx3(pos + dpos))  # update sigma
        self.targetvalue += self.constr2climb.sigma          # climb constraint
        self.constr2climb.targetvalue = self.targetvalue     # adjust positions
        # XXX very magical ...
        self.optimizable.set_x(self.optimizable.get_x())   # to targetvalue

    def relax_remaining_dof(self):
        """Optimize remaining degrees of freedom with optimizer 'B'."""
        if self.autolog:
            self.optB_kwargs['logfile'] = f'optB_{self.targetvalue}.log'
        if self.autotraj:
            self.optB_kwargs['trajectory'] = f'optB_{self.targetvalue}.traj'
        fmax = self.get_scaled_fmax()
        with self.optB(self.optimizable.atoms, **self.optB_kwargs) as opt:
            opt.run(fmax)  # optimize with scaled fmax
            grad = self.optimizable.get_gradient()
            if self.converged(grad) and fmax > self.optB_fmax:
                # (final) optimization with desired fmax
                opt.run(self.optB_fmax)

    def get_scaled_fmax(self):
        """Return the adaptive 'fmax' based on the estimated distance to the
        transition state."""
        return (self.optB_fmax +
                self.scaling * norm(self.constr2climb.projected_forces))

    def get_projected_forces(self):
        """Return the projected forces along the constrained coordinate in
        uphill direction (negative sign)."""
        forces = self.constr2climb.projected_forces
        # XXX simplify me once optimizable shape shenanigans have converged
        forces = -forces.ravel()
        return forces

    def get_total_forces(self):
        """Return forces obeying all constraints plus projected forces."""
        forces = self.optimizable.get_gradient()
        return forces + self.get_projected_forces()

    def converged(self, gradient):
        """Did the optimization converge based on the total forces?"""
        # XXX ignoring gradient
        gradient = self.get_total_forces().ravel()
        return super().converged(gradient=gradient)

    def log(self, gradient):
        forces = self.get_total_forces()
        super().log(gradient=forces.ravel())


def get_fixinternals(atoms):
    """Get pointer to the FixInternals constraint on the atoms object."""
    all_constr_types = list(map(type, atoms.constraints))
    index = all_constr_types.index(FixInternals)  # locate constraint
    return atoms.constraints[index]


def get_constr2climb(atoms, climb_coordinate):
    """Get pointer to the subconstraint that is to be climbed.
    Identification by its definition via indices (and coefficients)."""
    constr = get_fixinternals(atoms)
    return constr.get_subconstraint(atoms, climb_coordinate)