Coverage for /builds/debichem-team/python-ase/ase/optimize/gpmin/gpmin.py: 69.92%
123 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-03-06 04:00 +0000
1import warnings
3import numpy as np
4from scipy.optimize import minimize
6from ase.io.jsonio import write_json
7from ase.optimize.gpmin.gp import GaussianProcess
8from ase.optimize.gpmin.kernel import SquaredExponential
9from ase.optimize.gpmin.prior import ConstantPrior
10from ase.optimize.optimize import Optimizer
11from ase.parallel import world
14class GPMin(Optimizer, GaussianProcess):
15 def __init__(self, atoms, restart=None, logfile='-', trajectory=None,
16 prior=None, kernel=None, noise=None, weight=None, scale=None,
17 batch_size=None, bounds=None, update_prior_strategy="maximum",
18 update_hyperparams=False, **kwargs):
19 """Optimize atomic positions using GPMin algorithm, which uses both
20 potential energies and forces information to build a PES via Gaussian
21 Process (GP) regression and then minimizes it.
23 Default behaviour:
24 --------------------
25 The default values of the scale, noise, weight, batch_size and bounds
26 parameters depend on the value of update_hyperparams. In order to get
27 the default value of any of them, they should be set up to None.
28 Default values are:
30 update_hyperparams = True
31 scale : 0.3
32 noise : 0.004
33 weight: 2.
34 bounds: 0.1
35 batch_size: 1
37 update_hyperparams = False
38 scale : 0.4
39 noise : 0.005
40 weight: 1.
41 bounds: irrelevant
42 batch_size: irrelevant
44 Parameters
45 ----------
47 atoms: :class:`~ase.Atoms`
48 The Atoms object to relax.
50 restart: str
51 JSON file used to store the training set. If set, file with
52 such a name will be searched and the data in the file incorporated
53 to the new training set, if the file exists.
55 logfile: file object or str
56 If *logfile* is a string, a file with that name will be opened.
57 Use '-' for stdout
59 trajectory: str
60 File used to store trajectory of atomic movement.
62 prior: Prior object or None
63 Prior for the GP regression of the PES surface
64 See ase.optimize.gpmin.prior
65 If *prior* is None, then it is set as the
66 ConstantPrior with the constant being updated
67 using the update_prior_strategy specified as a parameter
69 kernel: Kernel object or None
70 Kernel for the GP regression of the PES surface
71 See ase.optimize.gpmin.kernel
72 If *kernel* is None the SquaredExponential kernel is used.
73 Note: It needs to be a kernel with derivatives!!!!!
75 noise: float
76 Regularization parameter for the Gaussian Process Regression.
78 weight: float
79 Prefactor of the Squared Exponential kernel.
80 If *update_hyperparams* is False, changing this parameter
81 has no effect on the dynamics of the algorithm.
83 update_prior_strategy: str
84 Strategy to update the constant from the ConstantPrior
85 when more data is collected. It does only work when
86 Prior = None
88 options:
89 'maximum': update the prior to the maximum sampled energy
90 'init' : fix the prior to the initial energy
91 'average': use the average of sampled energies as prior
93 scale: float
94 scale of the Squared Exponential Kernel
96 update_hyperparams: bool
97 Update the scale of the Squared exponential kernel
98 every batch_size-th iteration by maximizing the
99 marginal likelihood.
101 batch_size: int
102 Number of new points in the sample before updating
103 the hyperparameters.
104 Only relevant if the optimizer is executed in update_hyperparams
105 mode: (update_hyperparams = True)
107 bounds: float, 0<bounds<1
108 Set bounds to the optimization of the hyperparameters.
109 Let t be a hyperparameter. Then it is optimized under the
110 constraint (1-bound)*t_0 <= t <= (1+bound)*t_0
111 where t_0 is the value of the hyperparameter in the previous
112 step.
113 If bounds is False, no constraints are set in the optimization of
114 the hyperparameters.
116 kwargs : dict, optional
117 Extra arguments passed to
118 :class:`~ase.optimize.optimize.Optimizer`.
120 .. warning:: The memory of the optimizer scales as O(n²N²) where
121 N is the number of atoms and n the number of steps.
122 If the number of atoms is sufficiently high, this
123 may cause a memory issue.
124 This class prints a warning if the user tries to
125 run GPMin with more than 100 atoms in the unit cell.
126 """
127 # Warn the user if the number of atoms is very large
128 if len(atoms) > 100:
129 warning = ('Possible Memory Issue. There are more than '
130 '100 atoms in the unit cell. The memory '
131 'of the process will increase with the number '
132 'of steps, potentially causing a memory issue. '
133 'Consider using a different optimizer.')
135 warnings.warn(warning)
137 # Give it default hyperparameters
138 if update_hyperparams: # Updated GPMin
139 if scale is None:
140 scale = 0.3
141 if noise is None:
142 noise = 0.004
143 if weight is None:
144 weight = 2.
145 if bounds is None:
146 self.eps = 0.1
147 elif bounds is False:
148 self.eps = None
149 else:
150 self.eps = bounds
151 if batch_size is None:
152 self.nbatch = 1
153 else:
154 self.nbatch = batch_size
155 else: # GPMin without updates
156 if scale is None:
157 scale = 0.4
158 if noise is None:
159 noise = 0.001
160 if weight is None:
161 weight = 1.
162 if bounds is not None:
163 warning = ('The parameter bounds is of no use '
164 'if update_hyperparams is False. '
165 'The value provided by the user '
166 'is being ignored.')
167 warnings.warn(warning, UserWarning)
168 if batch_size is not None:
169 warning = ('The parameter batch_size is of no use '
170 'if update_hyperparams is False. '
171 'The value provided by the user '
172 'is being ignored.')
173 warnings.warn(warning, UserWarning)
175 # Set the variables to something anyways
176 self.eps = False
177 self.nbatch = None
179 self.strategy = update_prior_strategy
180 self.update_hp = update_hyperparams
181 self.function_calls = 1
182 self.force_calls = 0
183 self.x_list = [] # Training set features
184 self.y_list = [] # Training set targets
186 Optimizer.__init__(self, atoms, restart=restart, logfile=logfile,
187 trajectory=trajectory, **kwargs)
188 if prior is None:
189 self.update_prior = True
190 prior = ConstantPrior(constant=None)
191 else:
192 self.update_prior = False
194 if kernel is None:
195 kernel = SquaredExponential()
196 GaussianProcess.__init__(self, prior, kernel)
197 self.set_hyperparams(np.array([weight, scale, noise]))
199 def acquisition(self, r):
200 e = self.predict(r)
201 return e[0], e[1:]
203 def update(self, r, e, f):
204 """Update the PES
206 Update the training set, the prior and the hyperparameters.
207 Finally, train the model
208 """
209 # update the training set
210 self.x_list.append(r)
211 f = f.reshape(-1)
212 y = np.append(np.array(e).reshape(-1), -f)
213 self.y_list.append(y)
215 # Set/update the constant for the prior
216 if self.update_prior:
217 if self.strategy == 'average':
218 av_e = np.mean(np.array(self.y_list)[:, 0])
219 self.prior.set_constant(av_e)
220 elif self.strategy == 'maximum':
221 max_e = np.max(np.array(self.y_list)[:, 0])
222 self.prior.set_constant(max_e)
223 elif self.strategy == 'init':
224 self.prior.set_constant(e)
225 self.update_prior = False
227 # update hyperparams
228 if (self.update_hp and self.function_calls % self.nbatch == 0 and
229 self.function_calls != 0):
230 self.fit_to_batch()
232 # build the model
233 self.train(np.array(self.x_list), np.array(self.y_list))
235 def relax_model(self, r0):
236 result = minimize(self.acquisition, r0, method='L-BFGS-B', jac=True)
237 if result.success:
238 return result.x
239 else:
240 self.dump()
241 raise RuntimeError("The minimization of the acquisition function "
242 "has not converged")
244 def fit_to_batch(self):
245 """Fit hyperparameters keeping the ratio noise/weight fixed"""
246 ratio = self.noise / self.kernel.weight
247 self.fit_hyperparameters(np.array(self.x_list),
248 np.array(self.y_list), eps=self.eps)
249 self.noise = ratio * self.kernel.weight
251 def step(self, f=None):
252 optimizable = self.optimizable
253 if f is None:
254 f = optimizable.get_forces()
256 r0 = optimizable.get_positions().reshape(-1)
257 e0 = optimizable.get_potential_energy()
258 self.update(r0, e0, f)
260 r1 = self.relax_model(r0)
261 optimizable.set_positions(r1.reshape(-1, 3))
262 e1 = optimizable.get_potential_energy()
263 f1 = optimizable.get_forces()
264 self.function_calls += 1
265 self.force_calls += 1
266 count = 0
267 while e1 >= e0:
268 self.update(r1, e1, f1)
269 r1 = self.relax_model(r0)
271 optimizable.set_positions(r1.reshape(-1, 3))
272 e1 = optimizable.get_potential_energy()
273 f1 = optimizable.get_forces()
274 self.function_calls += 1
275 self.force_calls += 1
276 if self.converged(f1):
277 break
279 count += 1
280 if count == 30:
281 raise RuntimeError("A descent model could not be built")
282 self.dump()
284 def dump(self):
285 """Save the training set"""
286 if world.rank == 0 and self.restart is not None:
287 with open(self.restart, 'w') as fd:
288 write_json(fd, (self.x_list, self.y_list))
290 def read(self):
291 self.x_list, self.y_list = self.load()