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