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
|
# fmt: off
import time
import numpy as np
from ase.filters import UnitCellFilter
from ase.optimize.optimize import Optimizer
class PreconFIRE(Optimizer):
def __init__(self, atoms, restart=None, logfile='-', trajectory=None,
dt=0.1, maxmove=0.2, dtmax=1.0, Nmin=5, finc=1.1, fdec=0.5,
astart=0.1, fa=0.99, a=0.1, theta=0.1,
precon=None, use_armijo=True, variable_cell=False, **kwargs):
"""
Preconditioned version of the FIRE optimizer
In time this implementation is expected to replace
:class:`~ase.optimize.fire.FIRE`.
Parameters
----------
atoms: :class:`~ase.Atoms`
The Atoms object to relax.
restart: string
JSON file used to store hessian matrix. If set, file with
such a name will be searched and hessian matrix stored will
be used, if the file exists.
trajectory: string
Trajectory file used to store optimisation path.
logfile: file object or str
If *logfile* is a string, a file with that name will be opened.
Use '-' for stdout.
variable_cell: bool
If True, wrap atoms in UnitCellFilter to relax cell and positions.
kwargs : dict, optional
Extra arguments passed to
:class:`~ase.optimize.optimize.Optimizer`.
"""
if variable_cell:
atoms = UnitCellFilter(atoms)
Optimizer.__init__(self, atoms, restart, logfile, trajectory, **kwargs)
self._actual_atoms = atoms
self.dt = dt
self.Nsteps = 0
self.maxmove = maxmove
self.dtmax = dtmax
self.Nmin = Nmin
self.finc = finc
self.fdec = fdec
self.astart = astart
self.fa = fa
self.a = a
self.theta = theta
self.precon = precon
self.use_armijo = use_armijo
def initialize(self):
self.v = None
self.skip_flag = False
self.e1 = None
def read(self):
self.v, self.dt = self.load()
def step(self, f=None):
atoms = self._actual_atoms
if f is None:
f = atoms.get_forces()
r = atoms.get_positions()
if self.precon is not None:
# Can this be moved out of the step method?
self.precon.make_precon(atoms)
invP_f = self.precon.solve(f.reshape(-1)).reshape(len(atoms), -1)
if self.v is None:
self.v = np.zeros((len(self._actual_atoms), 3))
else:
if self.use_armijo:
if self.precon is None:
v_test = self.v + self.dt * f
else:
v_test = self.v + self.dt * invP_f
r_test = r + self.dt * v_test
self.skip_flag = False
func_val = self.func(r_test)
self.e1 = func_val
if (func_val > self.func(r) -
self.theta * self.dt * np.vdot(v_test, f)):
self.v[:] *= 0.0
self.a = self.astart
self.dt *= self.fdec
self.Nsteps = 0
self.skip_flag = True
if not self.skip_flag:
v_f = np.vdot(self.v, f)
if v_f > 0.0:
if self.precon is None:
self.v = (1.0 - self.a) * self.v + self.a * f / \
np.sqrt(np.vdot(f, f)) * \
np.sqrt(np.vdot(self.v, self.v))
else:
self.v = (
(1.0 - self.a) * self.v +
self.a *
(np.sqrt(self.precon.dot(self.v.reshape(-1),
self.v.reshape(-1))) /
np.sqrt(np.dot(f.reshape(-1),
invP_f.reshape(-1))) * invP_f))
if self.Nsteps > self.Nmin:
self.dt = min(self.dt * self.finc, self.dtmax)
self.a *= self.fa
self.Nsteps += 1
else:
self.v[:] *= 0.0
self.a = self.astart
self.dt *= self.fdec
self.Nsteps = 0
if self.precon is None:
self.v += self.dt * f
else:
self.v += self.dt * invP_f
dr = self.dt * self.v
normdr = np.sqrt(np.vdot(dr, dr))
if normdr > self.maxmove:
dr = self.maxmove * dr / normdr
atoms.set_positions(r + dr)
self.dump((self.v, self.dt))
def func(self, x):
"""Objective function for use of the optimizers"""
self._actual_atoms.set_positions(x.reshape(-1, 3))
potl = self._actual_atoms.get_potential_energy()
return potl
def run(self, fmax=0.05, steps=100000000, smax=None):
if smax is None:
smax = fmax
self.smax = smax
return Optimizer.run(self, fmax, steps)
def converged(self, gradient):
"""Did the optimization converge?"""
# XXX ignoring gradient
forces = self._actual_atoms.get_forces()
if isinstance(self._actual_atoms, UnitCellFilter):
natoms = len(self._actual_atoms.atoms)
forces, stress = forces[:natoms], self._actual_atoms.stress
fmax_sq = (forces**2).sum(axis=1).max()
smax_sq = (stress**2).max()
return (fmax_sq < self.fmax**2 and smax_sq < self.smax**2)
else:
fmax_sq = (forces**2).sum(axis=1).max()
return fmax_sq < self.fmax**2
def log(self, gradient):
forces = self._actual_atoms.get_forces()
if isinstance(self._actual_atoms, UnitCellFilter):
natoms = len(self._actual_atoms.atoms)
forces, stress = forces[:natoms], self._actual_atoms.stress
fmax = np.sqrt((forces**2).sum(axis=1).max())
smax = np.sqrt((stress**2).max())
else:
fmax = np.sqrt((forces**2).sum(axis=1).max())
if self.e1 is not None:
# reuse energy at end of line search to avoid extra call
e = self.e1
else:
e = self._actual_atoms.get_potential_energy()
T = time.localtime()
if self.logfile is not None:
name = self.__class__.__name__
if isinstance(self._actual_atoms, UnitCellFilter):
self.logfile.write(
'%s: %3d %02d:%02d:%02d %15.6f %12.4f %12.4f\n' %
(name, self.nsteps, T[3], T[4], T[5], e, fmax, smax))
else:
self.logfile.write(
'%s: %3d %02d:%02d:%02d %15.6f %12.4f\n' %
(name, self.nsteps, T[3], T[4], T[5], e, fmax))
self.logfile.flush()
|
