Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1""" 

2Implementation of the Precon abstract base class and subclasses 

3""" 

4 

5import sys 

6import time 

7import copy 

8import warnings 

9from abc import ABC, abstractmethod 

10 

11import numpy as np 

12from scipy import sparse 

13from scipy.sparse.linalg import spsolve 

14from scipy.interpolate import CubicSpline 

15 

16 

17from ase.constraints import Filter, FixAtoms 

18from ase.utils import longsum 

19from ase.geometry import find_mic 

20import ase.utils.ff as ff 

21import ase.units as units 

22from ase.optimize.precon.neighbors import (get_neighbours, 

23 estimate_nearest_neighbour_distance) 

24from ase.neighborlist import neighbor_list 

25from ase.utils import tokenize_version 

26 

27try: 

28 import pyamg 

29except ImportError: 

30 have_pyamg = False 

31else: 

32 have_pyamg = True 

33 

34 

35def create_pyamg_solver(P, max_levels=15): 

36 if not have_pyamg: 

37 raise RuntimeError( 

38 'pyamg not available: install with `pip install pyamg`') 

39 filter_key = 'filter' 

40 if tokenize_version(pyamg.__version__) >= tokenize_version('4.2.1'): 

41 filter_key = 'filter_entries' 

42 return pyamg.smoothed_aggregation_solver( 

43 P, B=None, 

44 strength=('symmetric', {'theta': 0.0}), 

45 smooth=( 

46 'jacobi', {filter_key: True, 'weighting': 'local'}), 

47 improve_candidates=([('block_gauss_seidel', 

48 {'sweep': 'symmetric', 'iterations': 4})] + 

49 [None] * (max_levels - 1)), 

50 aggregate='standard', 

51 presmoother=('block_gauss_seidel', 

52 {'sweep': 'symmetric', 'iterations': 1}), 

53 postsmoother=('block_gauss_seidel', 

54 {'sweep': 'symmetric', 'iterations': 1}), 

55 max_levels=max_levels, 

56 max_coarse=300, 

57 coarse_solver='pinv') 

58 

59 

60THz = 1e12 * 1. / units.s 

61 

62 

63class Precon(ABC): 

64 

65 @abstractmethod 

66 def make_precon(self, atoms, reinitialize=None): 

67 """ 

68 Create a preconditioner matrix based on the passed set of atoms. 

69 

70 Creates a general-purpose preconditioner for use with optimization 

71 algorithms, based on examining distances between pairs of atoms in the 

72 lattice. The matrix will be stored in the attribute self.P and 

73 returned. 

74 

75 Args: 

76 atoms: the Atoms object used to create the preconditioner. 

77 

78 reinitialize: if True, parameters of the preconditioner 

79 will be recalculated before the preconditioner matrix is 

80 created. If False, they will be calculated only when they 

81 do not currently have a value (ie, the first time this 

82 function is called). 

83 

84 Returns: 

85 P: A sparse scipy csr_matrix. BE AWARE that using 

86 numpy.dot() with sparse matrices will result in 

87 errors/incorrect results - use the .dot method directly 

88 on the matrix instead. 

89 """ 

90 ... 

91 

92 @abstractmethod 

93 def Pdot(self, x): 

94 """ 

95 Return the result of applying P to a vector x 

96 """ 

97 ... 

98 

99 def dot(self, x, y): 

100 """ 

101 Return the preconditioned dot product <P x, y> 

102 

103 Uses 128-bit floating point math for vector dot products 

104 """ 

105 return longsum(self.Pdot(x) * y) 

106 

107 def norm(self, x): 

108 """ 

109 Return the P-norm of x, where |x|_P = sqrt(<Px, x>) 

110 """ 

111 return np.sqrt(self.dot(x, x)) 

112 

113 def vdot(self, x, y): 

114 return self.dot(x.reshape(-1), 

115 y.reshape(-1)) 

116 

117 @abstractmethod 

118 def solve(self, x): 

119 """ 

120 Solve the (sparse) linear system P x = y and return y 

121 """ 

122 ... 

123 

124 def apply(self, forces, atoms): 

125 """ 

126 Convenience wrapper that combines make_precon() and solve() 

127 

128 Parameters 

129 ---------- 

130 forces: array 

131 (len(atoms)*3) array of input forces 

132 atoms: ase.atoms.Atoms 

133 

134 Returns 

135 ------- 

136 precon_forces: array 

137 (len(atoms), 3) array of preconditioned forces 

138 residual: float 

139 inf-norm of original forces, i.e. maximum absolute force 

140 """ 

141 self.make_precon(atoms) 

142 residual = np.linalg.norm(forces, np.inf) 

143 precon_forces = self.solve(forces) 

144 return precon_forces, residual 

145 

146 @abstractmethod 

147 def copy(self): 

148 ... 

149 

150 @abstractmethod 

151 def asarray(self): 

152 """ 

153 Array representation of preconditioner, as a dense matrix 

154 """ 

155 ... 

156 

157 

158class Logfile: 

159 def __init__(self, logfile=None): 

160 if isinstance(logfile, str): 

161 if logfile == "-": 

162 logfile = sys.stdout 

163 else: 

164 logfile = open(logfile, "a") 

165 self.logfile = logfile 

166 

167 def write(self, *args): 

168 if self.logfile is None: 

169 return 

170 self.logfile.write(*args) 

171 

172 

173class SparsePrecon(Precon): 

174 def __init__(self, r_cut=None, r_NN=None, 

175 mu=None, mu_c=None, 

176 dim=3, c_stab=0.1, force_stab=False, 

177 reinitialize=False, array_convention='C', 

178 solver="auto", solve_tol=1e-8, 

179 apply_positions=True, apply_cell=True, 

180 estimate_mu_eigmode=False, logfile=None, rng=None, 

181 neighbour_list=neighbor_list): 

182 """Initialise a preconditioner object based on passed parameters. 

183 

184 Parameters: 

185 r_cut: float. This is a cut-off radius. The preconditioner matrix 

186 will be created by considering pairs of atoms that are within a 

187 distance r_cut of each other. For a regular lattice, this is 

188 usually taken somewhere between the first- and second-nearest 

189 neighbour distance. If r_cut is not provided, default is 

190 2 * r_NN (see below) 

191 r_NN: nearest neighbour distance. If not provided, this is 

192 calculated 

193 from input structure. 

194 mu: float 

195 energy scale for position degreees of freedom. If `None`, mu 

196 is precomputed using finite difference derivatives. 

197 mu_c: float 

198 energy scale for cell degreees of freedom. Also precomputed 

199 if None. 

200 estimate_mu_eigmode: 

201 If True, estimates mu based on the lowest eigenmodes of 

202 unstabilised preconditioner. If False it uses the sine based 

203 approach. 

204 dim: int; dimensions of the problem 

205 c_stab: float. The diagonal of the preconditioner matrix will have 

206 a stabilisation constant added, which will be the value of 

207 c_stab times mu. 

208 force_stab: 

209 If True, always add the stabilisation to diagnonal, regardless 

210 of the presence of fixed atoms. 

211 reinitialize: if True, the value of mu will be recalculated when 

212 self.make_precon is called. This can be overridden in specific 

213 cases with reinitialize argument in self.make_precon. If it 

214 is set to True here, the value passed for mu will be 

215 irrelevant unless reinitialize is set to False the first time 

216 make_precon is called. 

217 array_convention: Either 'C' or 'F' for Fortran; this will change 

218 the preconditioner to reflect the ordering of the indices in 

219 the vector it will operate on. The C convention assumes the 

220 vector will be arranged atom-by-atom (ie [x1, y1, z1, x2, ...]) 

221 while the F convention assumes it will be arranged component 

222 by component (ie [x1, x2, ..., y1, y2, ...]). 

223 solver: One of "auto", "direct" or "pyamg", specifying whether to 

224 use a direct sparse solver or PyAMG to solve P x = y. 

225 Default is "auto" which uses PyAMG if available, falling 

226 back to sparse solver if not. solve_tol: tolerance used for 

227 PyAMG sparse linear solver, if available. 

228 apply_positions: bool 

229 if True, apply preconditioner to position DoF 

230 apply_cell: bool 

231 if True, apply preconditioner to cell DoF 

232 logfile: file object or str 

233 If *logfile* is a string, a file with that name will be opened. 

234 Use '-' for stdout. 

235 rng: None or np.random.RandomState instance 

236 Random number generator to use for initialising pyamg solver 

237 neighbor_list: function (optional). Optionally replace the built-in 

238 ASE neighbour list with an alternative with the same call 

239 signature, e.g. `matscipy.neighbours.neighbour_list`. 

240 

241 Raises: 

242 ValueError for problem with arguments 

243 

244 """ 

245 self.r_NN = r_NN 

246 self.r_cut = r_cut 

247 self.mu = mu 

248 self.mu_c = mu_c 

249 self.estimate_mu_eigmode = estimate_mu_eigmode 

250 self.c_stab = c_stab 

251 self.force_stab = force_stab 

252 self.array_convention = array_convention 

253 self.reinitialize = reinitialize 

254 self.P = None 

255 self.old_positions = None 

256 

257 use_pyamg = False 

258 if solver == "auto": 

259 use_pyamg = have_pyamg 

260 elif solver == "direct": 

261 use_pyamg = False 

262 elif solver == "pyamg": 

263 if not have_pyamg: 

264 raise RuntimeError("solver='pyamg', PyAMG can't be imported!") 

265 use_pyamg = True 

266 else: 

267 raise ValueError('unknown solver - ' 

268 'should be "auto", "direct" or "pyamg"') 

269 

270 self.use_pyamg = use_pyamg 

271 self.solve_tol = solve_tol 

272 self.apply_positions = apply_positions 

273 self.apply_cell = apply_cell 

274 

275 if dim < 1: 

276 raise ValueError('Dimension must be at least 1') 

277 self.dim = dim 

278 self.logfile = Logfile(logfile) 

279 

280 if rng is None: 

281 rng = np.random.RandomState() 

282 self.rng = rng 

283 

284 self.neighbor_list = neighbor_list 

285 

286 def copy(self): 

287 return copy.deepcopy(self) 

288 

289 def Pdot(self, x): 

290 return self.P.dot(x) 

291 

292 def solve(self, x): 

293 start_time = time.time() 

294 if self.use_pyamg and have_pyamg: 

295 y = self.ml.solve(x, x0=self.rng.random(self.P.shape[0]), 

296 tol=self.solve_tol, 

297 accel='cg', 

298 maxiter=300, 

299 cycle='W') 

300 else: 

301 y = spsolve(self.P, x) 

302 self.logfile.write('--- Precon applied in %s seconds ---\n' % 

303 (time.time() - start_time)) 

304 return y 

305 

306 def estimate_mu(self, atoms, H=None): 

307 r"""Estimate optimal preconditioner coefficient \mu 

308 

309 \mu is estimated from a numerical solution of 

310 

311 [dE(p+v) - dE(p)] \cdot v = \mu < P1 v, v > 

312 

313 with perturbation 

314 

315 v(x,y,z) = H P_lowest_nonzero_eigvec(x, y, z) 

316 

317 or 

318 

319 v(x,y,z) = H (sin(x / Lx), sin(y / Ly), sin(z / Lz)) 

320 

321 After the optimal \mu is found, self.mu will be set to its value. 

322 

323 If `atoms` is an instance of Filter an additional \mu_c 

324 will be computed for the cell degrees of freedom . 

325 

326 Args: 

327 atoms: Atoms object for initial system 

328 

329 H: 3x3 array or None 

330 Magnitude of deformation to apply. 

331 Default is 1e-2*rNN*np.eye(3) 

332 

333 Returns: 

334 mu : float 

335 mu_c : float or None 

336 """ 

337 logfile = self.logfile 

338 

339 if self.dim != 3: 

340 raise ValueError('Automatic calculation of mu only possible for ' 

341 'three-dimensional preconditioners. Try setting ' 

342 'mu manually instead.') 

343 

344 if self.r_NN is None: 

345 self.r_NN = estimate_nearest_neighbour_distance(atoms, 

346 self.neighbor_list) 

347 

348 # deformation matrix, default is diagonal 

349 if H is None: 

350 H = 1e-2 * self.r_NN * np.eye(3) 

351 

352 # compute perturbation 

353 p = atoms.get_positions() 

354 

355 if self.estimate_mu_eigmode: 

356 self.mu = 1.0 

357 self.mu_c = 1.0 

358 c_stab = self.c_stab 

359 self.c_stab = 0.0 

360 

361 if isinstance(atoms, Filter): 

362 n = len(atoms.atoms) 

363 else: 

364 n = len(atoms) 

365 self._make_sparse_precon(atoms, initial_assembly=True) 

366 self.P = self.P[:3 * n, :3 * n] 

367 eigvals, eigvecs = sparse.linalg.eigsh(self.P, k=4, which='SM') 

368 

369 logfile.write('estimate_mu(): lowest 4 eigvals = %f %f %f %f\n' % 

370 (eigvals[0], eigvals[1], eigvals[2], eigvals[3])) 

371 # check eigenvalues 

372 if any(eigvals[0:3] > 1e-6): 

373 raise ValueError('First 3 eigenvalues of preconditioner matrix' 

374 'do not correspond to translational modes.') 

375 elif eigvals[3] < 1e-6: 

376 raise ValueError('Fourth smallest eigenvalue of ' 

377 'preconditioner matrix ' 

378 'is too small, increase r_cut.') 

379 

380 x = np.zeros(n) 

381 for i in range(n): 

382 x[i] = eigvecs[:, 3][3 * i] 

383 x = x / np.linalg.norm(x) 

384 if x[0] < 0: 

385 x = -x 

386 

387 v = np.zeros(3 * len(atoms)) 

388 for i in range(n): 

389 v[3 * i] = x[i] 

390 v[3 * i + 1] = x[i] 

391 v[3 * i + 2] = x[i] 

392 v = v / np.linalg.norm(v) 

393 v = v.reshape((-1, 3)) 

394 

395 self.c_stab = c_stab 

396 else: 

397 Lx, Ly, Lz = [p[:, i].max() - p[:, i].min() for i in range(3)] 

398 logfile.write('estimate_mu(): Lx=%.1f Ly=%.1f Lz=%.1f\n' % 

399 (Lx, Ly, Lz)) 

400 

401 x, y, z = p.T 

402 # sine_vr = [np.sin(x/Lx), np.sin(y/Ly), np.sin(z/Lz)], but we need 

403 # to take into account the possibility that one of Lx/Ly/Lz is 

404 # zero. 

405 sine_vr = [x, y, z] 

406 

407 for i, L in enumerate([Lx, Ly, Lz]): 

408 if L == 0: 

409 warnings.warn( 

410 'Cell length L[%d] == 0. Setting H[%d,%d] = 0.' % 

411 (i, i, i)) 

412 H[i, i] = 0.0 

413 else: 

414 sine_vr[i] = np.sin(sine_vr[i] / L) 

415 

416 v = np.dot(H, sine_vr).T 

417 

418 natoms = len(atoms) 

419 if isinstance(atoms, Filter): 

420 natoms = len(atoms.atoms) 

421 eps = H / self.r_NN 

422 v[natoms:, :] = eps 

423 

424 v1 = v.reshape(-1) 

425 

426 # compute LHS 

427 dE_p = -atoms.get_forces().reshape(-1) 

428 atoms_v = atoms.copy() 

429 atoms_v.calc = atoms.calc 

430 if isinstance(atoms, Filter): 

431 atoms_v = atoms.__class__(atoms_v) 

432 if hasattr(atoms, 'constant_volume'): 

433 atoms_v.constant_volume = atoms.constant_volume 

434 atoms_v.set_positions(p + v) 

435 dE_p_plus_v = -atoms_v.get_forces().reshape(-1) 

436 

437 # compute left hand side 

438 LHS = (dE_p_plus_v - dE_p) * v1 

439 

440 # assemble P with \mu = 1 

441 self.mu = 1.0 

442 self.mu_c = 1.0 

443 

444 self._make_sparse_precon(atoms, initial_assembly=True) 

445 

446 # compute right hand side 

447 RHS = self.P.dot(v1) * v1 

448 

449 # use partial sums to compute separate mu for positions and cell DoFs 

450 self.mu = longsum(LHS[:3 * natoms]) / longsum(RHS[:3 * natoms]) 

451 if self.mu < 1.0: 

452 logfile.write('estimate_mu(): mu (%.3f) < 1.0, ' 

453 'capping at mu=1.0' % self.mu) 

454 self.mu = 1.0 

455 

456 if isinstance(atoms, Filter): 

457 self.mu_c = longsum(LHS[3 * natoms:]) / longsum(RHS[3 * natoms:]) 

458 if self.mu_c < 1.0: 

459 logfile.write('estimate_mu(): mu_c (%.3f) < 1.0, ' 

460 'capping at mu_c=1.0\n' % self.mu_c) 

461 self.mu_c = 1.0 

462 

463 logfile.write('estimate_mu(): mu=%r, mu_c=%r\n' % 

464 (self.mu, self.mu_c)) 

465 

466 self.P = None # force a rebuild with new mu (there may be fixed atoms) 

467 return (self.mu, self.mu_c) 

468 

469 def asarray(self): 

470 return np.array(self.P.todense()) 

471 

472 def one_dim_to_ndim(self, csc_P, N): 

473 """ 

474 Expand an N x N precon matrix to self.dim*N x self.dim * N 

475 

476 Args: 

477 csc_P (sparse matrix): N x N sparse matrix, in CSC format 

478 """ 

479 start_time = time.time() 

480 if self.dim == 1: 

481 P = csc_P 

482 elif self.array_convention == 'F': 

483 csc_P = csc_P.tocsr() 

484 P = csc_P 

485 for i in range(self.dim - 1): 

486 P = sparse.block_diag((P, csc_P)).tocsr() 

487 else: 

488 # convert back to triplet and read the arrays 

489 csc_P = csc_P.tocoo() 

490 i = csc_P.row * self.dim 

491 j = csc_P.col * self.dim 

492 z = csc_P.data 

493 

494 # N-dimensionalise, interlaced coordinates 

495 I = np.hstack([i + d for d in range(self.dim)]) 

496 J = np.hstack([j + d for d in range(self.dim)]) 

497 Z = np.hstack([z for d in range(self.dim)]) 

498 P = sparse.csc_matrix((Z, (I, J)), 

499 shape=(self.dim * N, self.dim * N)) 

500 P = P.tocsr() 

501 self.logfile.write('--- N-dim precon created in %s s ---\n' % 

502 (time.time() - start_time)) 

503 return P 

504 

505 def create_solver(self): 

506 if self.use_pyamg and have_pyamg: 

507 start_time = time.time() 

508 self.ml = create_pyamg_solver(self.P) 

509 self.logfile.write('--- multi grid solver created in %s ---\n' % 

510 (time.time() - start_time)) 

511 

512 

513class SparseCoeffPrecon(SparsePrecon): 

514 def _make_sparse_precon(self, atoms, initial_assembly=False, 

515 force_stab=False): 

516 """Create a sparse preconditioner matrix based on the passed atoms. 

517 

518 Creates a general-purpose preconditioner for use with optimization 

519 algorithms, based on examining distances between pairs of atoms in the 

520 lattice. The matrix will be stored in the attribute self.P and 

521 returned. Note that this function will use self.mu, whatever it is. 

522 

523 Args: 

524 atoms: the Atoms object used to create the preconditioner. 

525 

526 Returns: 

527 A scipy.sparse.csr_matrix object, representing a d*N by d*N matrix 

528 (where N is the number of atoms, and d is the value of self.dim). 

529 BE AWARE that using numpy.dot() with this object will result in 

530 errors/incorrect results - use the .dot method directly on the 

531 sparse matrix instead. 

532 

533 """ 

534 logfile = self.logfile 

535 logfile.write('creating sparse precon: initial_assembly=%r, ' 

536 'force_stab=%r, apply_positions=%r, apply_cell=%r\n' % 

537 (initial_assembly, force_stab, self.apply_positions, 

538 self.apply_cell)) 

539 

540 N = len(atoms) 

541 diag_i = np.arange(N, dtype=int) 

542 start_time = time.time() 

543 if self.apply_positions: 

544 # compute neighbour list 

545 i, j, rij, fixed_atoms = get_neighbours( 

546 atoms, self.r_cut, 

547 neighbor_list=self.neighbor_list) 

548 logfile.write('--- neighbour list created in %s s --- \n' % 

549 (time.time() - start_time)) 

550 

551 # compute entries in triplet format: without the constraints 

552 start_time = time.time() 

553 coeff = self.get_coeff(rij) 

554 diag_coeff = np.bincount(i, -coeff, minlength=N).astype(np.float64) 

555 if force_stab or len(fixed_atoms) == 0: 

556 logfile.write('adding stabilisation to precon') 

557 diag_coeff += self.mu * self.c_stab 

558 else: 

559 diag_coeff = np.ones(N) 

560 

561 # precon is mu_c * identity for cell DoF 

562 if isinstance(atoms, Filter): 

563 if self.apply_cell: 

564 diag_coeff[-3:] = self.mu_c 

565 else: 

566 diag_coeff[-3:] = 1.0 

567 logfile.write('--- computed triplet format in %s s ---\n' % 

568 (time.time() - start_time)) 

569 

570 if self.apply_positions and not initial_assembly: 

571 # apply the constraints 

572 start_time = time.time() 

573 mask = np.ones(N) 

574 mask[fixed_atoms] = 0.0 

575 coeff *= mask[i] * mask[j] 

576 diag_coeff[fixed_atoms] = 1.0 

577 logfile.write('--- applied fixed_atoms in %s s ---\n' % 

578 (time.time() - start_time)) 

579 

580 if self.apply_positions: 

581 # remove zeros 

582 start_time = time.time() 

583 inz = np.nonzero(coeff) 

584 i = np.hstack((i[inz], diag_i)) 

585 j = np.hstack((j[inz], diag_i)) 

586 coeff = np.hstack((coeff[inz], diag_coeff)) 

587 logfile.write('--- remove zeros in %s s ---\n' % 

588 (time.time() - start_time)) 

589 else: 

590 i = diag_i 

591 j = diag_i 

592 coeff = diag_coeff 

593 

594 # create an N x N precon matrix in compressed sparse column (CSC) format 

595 start_time = time.time() 

596 csc_P = sparse.csc_matrix((coeff, (i, j)), shape=(N, N)) 

597 logfile.write('--- created CSC matrix in %s s ---\n' % 

598 (time.time() - start_time)) 

599 

600 self.P = self.one_dim_to_ndim(csc_P, N) 

601 self.create_solver() 

602 

603 def make_precon(self, atoms, reinitialize=None): 

604 if self.r_NN is None: 

605 self.r_NN = estimate_nearest_neighbour_distance(atoms, 

606 self.neighbor_list) 

607 

608 if self.r_cut is None: 

609 # This is the first time this function has been called, and no 

610 # cutoff radius has been specified, so calculate it automatically. 

611 self.r_cut = 2.0 * self.r_NN 

612 elif self.r_cut < self.r_NN: 

613 warning = ('WARNING: r_cut (%.2f) < r_NN (%.2f), ' 

614 'increasing to 1.1*r_NN = %.2f' % (self.r_cut, 

615 self.r_NN, 

616 1.1 * self.r_NN)) 

617 warnings.warn(warning) 

618 self.r_cut = 1.1 * self.r_NN 

619 

620 if reinitialize is None: 

621 # The caller has not specified whether or not to recalculate mu, 

622 # so the Precon's setting is used. 

623 reinitialize = self.reinitialize 

624 

625 if self.mu is None: 

626 # Regardless of what the caller has specified, if we don't 

627 # currently have a value of mu, then we need one. 

628 reinitialize = True 

629 

630 if reinitialize: 

631 self.estimate_mu(atoms) 

632 

633 if self.P is not None: 

634 real_atoms = atoms 

635 if isinstance(atoms, Filter): 

636 real_atoms = atoms.atoms 

637 if self.old_positions is None: 

638 self.old_positions = real_atoms.positions 

639 displacement, _ = find_mic(real_atoms.positions - 

640 self.old_positions, 

641 real_atoms.cell, real_atoms.pbc) 

642 self.old_positions = real_atoms.get_positions() 

643 max_abs_displacement = abs(displacement).max() 

644 self.logfile.write('max(abs(displacements)) = %.2f A (%.2f r_NN)' % 

645 (max_abs_displacement, 

646 max_abs_displacement / self.r_NN)) 

647 if max_abs_displacement < 0.5 * self.r_NN: 

648 return 

649 

650 start_time = time.time() 

651 self._make_sparse_precon(atoms, force_stab=self.force_stab) 

652 self.logfile.write('--- Precon created in %s seconds --- \n' % 

653 (time.time() - start_time)) 

654 

655 @abstractmethod 

656 def get_coeff(self, r): 

657 ... 

658 

659 

660class Pfrommer(Precon): 

661 """ 

662 Use initial guess for inverse Hessian from Pfrommer et al. as a 

663 simple preconditioner 

664 

665 J. Comput. Phys. vol 131 p233-240 (1997) 

666 """ 

667 

668 def __init__(self, bulk_modulus=500 * units.GPa, phonon_frequency=50 * THz, 

669 apply_positions=True, apply_cell=True): 

670 """ 

671 Default bulk modulus is 500 GPa and default phonon frequency is 50 THz 

672 """ 

673 self.bulk_modulus = bulk_modulus 

674 self.phonon_frequency = phonon_frequency 

675 self.apply_positions = apply_positions 

676 self.apply_cell = apply_cell 

677 self.H0 = None 

678 

679 def make_precon(self, atoms, reinitialize=None): 

680 if self.H0 is not None: 

681 # only build H0 on first call 

682 return 

683 

684 variable_cell = False 

685 if isinstance(atoms, Filter): 

686 variable_cell = True 

687 atoms = atoms.atoms 

688 

689 # position DoF 

690 omega = self.phonon_frequency 

691 mass = atoms.get_masses().mean() 

692 block = np.eye(3) / (mass * omega**2) 

693 blocks = [block] * len(atoms) 

694 

695 # cell DoF 

696 if variable_cell: 

697 coeff = 1.0 

698 if self.apply_cell: 

699 coeff = 1.0 / (3 * self.bulk_modulus) 

700 blocks.append(np.diag([coeff] * 9)) 

701 

702 self.H0 = sparse.block_diag(blocks, format='csr') 

703 return 

704 

705 def Pdot(self, x): 

706 return self.H0.solve(x) 

707 

708 def solve(self, x): 

709 y = self.H0.dot(x) 

710 return y 

711 

712 def copy(self): 

713 return Pfrommer(self.bulk_modulus, 

714 self.phonon_frequency, 

715 self.apply_positions, 

716 self.apply_cell) 

717 

718 def asarray(self): 

719 return np.array(self.H0.todense()) 

720 

721 

722class IdentityPrecon(Precon): 

723 """ 

724 Dummy preconditioner which does not modify forces 

725 """ 

726 

727 def make_precon(self, atoms, reinitialize=None): 

728 self.atoms = atoms 

729 

730 def Pdot(self, x): 

731 return x 

732 

733 def solve(self, x): 

734 return x 

735 

736 def copy(self): 

737 return IdentityPrecon() 

738 

739 def asarray(self): 

740 return np.eye(3 * len(self.atoms)) 

741 

742 

743class C1(SparseCoeffPrecon): 

744 """Creates matrix by inserting a constant whenever r_ij is less than r_cut. 

745 """ 

746 

747 def __init__(self, r_cut=None, mu=None, mu_c=None, dim=3, c_stab=0.1, 

748 force_stab=False, 

749 reinitialize=False, array_convention='C', 

750 solver="auto", solve_tol=1e-9, 

751 apply_positions=True, apply_cell=True, logfile=None): 

752 super().__init__(r_cut=r_cut, mu=mu, mu_c=mu_c, 

753 dim=dim, c_stab=c_stab, 

754 force_stab=force_stab, 

755 reinitialize=reinitialize, 

756 array_convention=array_convention, 

757 solver=solver, solve_tol=solve_tol, 

758 apply_positions=apply_positions, 

759 apply_cell=apply_cell, 

760 logfile=logfile) 

761 

762 def get_coeff(self, r): 

763 return -self.mu * np.ones_like(r) 

764 

765 

766class Exp(SparseCoeffPrecon): 

767 """Creates matrix with values decreasing exponentially with distance. 

768 """ 

769 

770 def __init__(self, A=3.0, r_cut=None, r_NN=None, mu=None, mu_c=None, 

771 dim=3, c_stab=0.1, 

772 force_stab=False, reinitialize=False, array_convention='C', 

773 solver="auto", solve_tol=1e-9, 

774 apply_positions=True, apply_cell=True, 

775 estimate_mu_eigmode=False, logfile=None): 

776 """ 

777 Initialise an Exp preconditioner with given parameters. 

778 

779 Args: 

780 r_cut, mu, c_stab, dim, sparse, reinitialize, array_convention: see 

781 precon.__init__() 

782 A: coefficient in exp(-A*r/r_NN). Default is A=3.0. 

783 """ 

784 super().__init__(r_cut=r_cut, r_NN=r_NN, 

785 mu=mu, mu_c=mu_c, dim=dim, c_stab=c_stab, 

786 force_stab=force_stab, 

787 reinitialize=reinitialize, 

788 array_convention=array_convention, 

789 solver=solver, 

790 solve_tol=solve_tol, 

791 apply_positions=apply_positions, 

792 apply_cell=apply_cell, 

793 estimate_mu_eigmode=estimate_mu_eigmode, 

794 logfile=logfile) 

795 

796 self.A = A 

797 

798 def get_coeff(self, r): 

799 return -self.mu * np.exp(-self.A * (r / self.r_NN - 1)) 

800 

801 

802def ij_to_x(i, j): 

803 x = [3 * i, 3 * i + 1, 3 * i + 2, 

804 3 * j, 3 * j + 1, 3 * j + 2] 

805 return x 

806 

807 

808def ijk_to_x(i, j, k): 

809 x = [3 * i, 3 * i + 1, 3 * i + 2, 

810 3 * j, 3 * j + 1, 3 * j + 2, 

811 3 * k, 3 * k + 1, 3 * k + 2] 

812 return x 

813 

814 

815def ijkl_to_x(i, j, k, l): 

816 x = [3 * i, 3 * i + 1, 3 * i + 2, 

817 3 * j, 3 * j + 1, 3 * j + 2, 

818 3 * k, 3 * k + 1, 3 * k + 2, 

819 3 * l, 3 * l + 1, 3 * l + 2] 

820 return x 

821 

822 

823def apply_fixed(atoms, P): 

824 fixed_atoms = [] 

825 for constraint in atoms.constraints: 

826 if isinstance(constraint, FixAtoms): 

827 fixed_atoms.extend(list(constraint.index)) 

828 else: 

829 raise TypeError( 

830 'only FixAtoms constraints are supported by Precon class') 

831 if len(fixed_atoms) != 0: 

832 P = P.tolil() 

833 for i in fixed_atoms: 

834 P[i, :] = 0.0 

835 P[:, i] = 0.0 

836 P[i, i] = 1.0 

837 return P 

838 

839 

840class FF(SparsePrecon): 

841 """Creates matrix using morse/bond/angle/dihedral force field parameters. 

842 """ 

843 

844 def __init__(self, dim=3, c_stab=0.1, force_stab=False, 

845 array_convention='C', solver="auto", solve_tol=1e-9, 

846 apply_positions=True, apply_cell=True, 

847 hessian='spectral', morses=None, bonds=None, angles=None, 

848 dihedrals=None, logfile=None): 

849 """Initialise an FF preconditioner with given parameters. 

850 

851 Args: 

852 dim, c_stab, force_stab, array_convention, use_pyamg, solve_tol: 

853 see SparsePrecon.__init__() 

854 morses: Morse instance 

855 bonds: Bond instance 

856 angles: Angle instance 

857 dihedrals: Dihedral instance 

858 """ 

859 

860 if (morses is None and bonds is None and angles is None and 

861 dihedrals is None): 

862 raise ImportError( 

863 'At least one of morses, bonds, angles or dihedrals must be ' 

864 'defined!') 

865 

866 super().__init__(dim=dim, c_stab=c_stab, 

867 force_stab=force_stab, 

868 array_convention=array_convention, 

869 solver=solver, 

870 solve_tol=solve_tol, 

871 apply_positions=apply_positions, 

872 apply_cell=apply_cell, logfile=logfile) 

873 

874 self.hessian = hessian 

875 self.morses = morses 

876 self.bonds = bonds 

877 self.angles = angles 

878 self.dihedrals = dihedrals 

879 

880 def make_precon(self, atoms, reinitialize=None): 

881 start_time = time.time() 

882 self._make_sparse_precon(atoms, force_stab=self.force_stab) 

883 self.logfile.write('--- Precon created in %s seconds ---\n' 

884 % (time.time() - start_time)) 

885 

886 def add_morse(self, morse, atoms, row, col, data, conn=None): 

887 if self.hessian == 'reduced': 

888 i, j, Hx = ff.get_morse_potential_reduced_hessian( 

889 atoms, morse) 

890 elif self.hessian == 'spectral': 

891 i, j, Hx = ff.get_morse_potential_hessian( 

892 atoms, morse, spectral=True) 

893 else: 

894 raise NotImplementedError('Not implemented hessian') 

895 x = ij_to_x(i, j) 

896 row.extend(np.repeat(x, 6)) 

897 col.extend(np.tile(x, 6)) 

898 data.extend(Hx.flatten()) 

899 if conn is not None: 

900 conn[i, j] = True 

901 conn[j, i] = True 

902 

903 def add_bond(self, bond, atoms, row, col, data, conn=None): 

904 if self.hessian == 'reduced': 

905 i, j, Hx = ff.get_bond_potential_reduced_hessian( 

906 atoms, bond, self.morses) 

907 elif self.hessian == 'spectral': 

908 i, j, Hx = ff.get_bond_potential_hessian( 

909 atoms, bond, self.morses, spectral=True) 

910 else: 

911 raise NotImplementedError('Not implemented hessian') 

912 x = ij_to_x(i, j) 

913 row.extend(np.repeat(x, 6)) 

914 col.extend(np.tile(x, 6)) 

915 data.extend(Hx.flatten()) 

916 if conn is not None: 

917 conn[i, j] = True 

918 conn[j, i] = True 

919 

920 def add_angle(self, angle, atoms, row, col, data, conn=None): 

921 if self.hessian == 'reduced': 

922 i, j, k, Hx = ff.get_angle_potential_reduced_hessian( 

923 atoms, angle, self.morses) 

924 elif self.hessian == 'spectral': 

925 i, j, k, Hx = ff.get_angle_potential_hessian( 

926 atoms, angle, self.morses, spectral=True) 

927 else: 

928 raise NotImplementedError('Not implemented hessian') 

929 x = ijk_to_x(i, j, k) 

930 row.extend(np.repeat(x, 9)) 

931 col.extend(np.tile(x, 9)) 

932 data.extend(Hx.flatten()) 

933 if conn is not None: 

934 conn[i, j] = conn[i, k] = conn[j, k] = True 

935 conn[j, i] = conn[k, i] = conn[k, j] = True 

936 

937 def add_dihedral(self, dihedral, atoms, row, col, data, conn=None): 

938 if self.hessian == 'reduced': 

939 i, j, k, l, Hx = \ 

940 ff.get_dihedral_potential_reduced_hessian( 

941 atoms, dihedral, self.morses) 

942 elif self.hessian == 'spectral': 

943 i, j, k, l, Hx = ff.get_dihedral_potential_hessian( 

944 atoms, dihedral, self.morses, spectral=True) 

945 else: 

946 raise NotImplementedError('Not implemented hessian') 

947 x = ijkl_to_x(i, j, k, l) 

948 row.extend(np.repeat(x, 12)) 

949 col.extend(np.tile(x, 12)) 

950 data.extend(Hx.flatten()) 

951 if conn is not None: 

952 conn[i, j] = conn[i, k] = conn[i, l] = conn[ 

953 j, k] = conn[j, l] = conn[k, l] = True 

954 conn[j, i] = conn[k, i] = conn[l, i] = conn[ 

955 k, j] = conn[l, j] = conn[l, k] = True 

956 

957 def _make_sparse_precon(self, atoms, initial_assembly=False, 

958 force_stab=False): 

959 N = len(atoms) 

960 

961 row = [] 

962 col = [] 

963 data = [] 

964 

965 if self.morses is not None: 

966 for morse in self.morses: 

967 self.add_morse(morse, atoms, row, col, data) 

968 

969 if self.bonds is not None: 

970 for bond in self.bonds: 

971 self.add_bond(bond, atoms, row, col, data) 

972 

973 if self.angles is not None: 

974 for angle in self.angles: 

975 self.add_angle(angle, atoms, row, col, data) 

976 

977 if self.dihedrals is not None: 

978 for dihedral in self.dihedrals: 

979 self.add_dihedral(dihedral, atoms, row, col, data) 

980 

981 # add the diagonal 

982 row.extend(range(self.dim * N)) 

983 col.extend(range(self.dim * N)) 

984 data.extend([self.c_stab] * self.dim * N) 

985 

986 # create the matrix 

987 start_time = time.time() 

988 self.P = sparse.csc_matrix( 

989 (data, (row, col)), shape=(self.dim * N, self.dim * N)) 

990 self.logfile.write('--- created CSC matrix in %s s ---\n' % 

991 (time.time() - start_time)) 

992 

993 self.P = apply_fixed(atoms, self.P) 

994 self.P = self.P.tocsr() 

995 self.logfile.write('--- N-dim precon created in %s s ---\n' % 

996 (time.time() - start_time)) 

997 self.create_solver() 

998 

999 

1000class Exp_FF(Exp, FF): 

1001 """Creates matrix with values decreasing exponentially with distance. 

1002 """ 

1003 

1004 def __init__(self, A=3.0, r_cut=None, r_NN=None, mu=None, mu_c=None, 

1005 dim=3, c_stab=0.1, 

1006 force_stab=False, reinitialize=False, array_convention='C', 

1007 solver="auto", solve_tol=1e-9, 

1008 apply_positions=True, apply_cell=True, 

1009 estimate_mu_eigmode=False, 

1010 hessian='spectral', morses=None, bonds=None, angles=None, 

1011 dihedrals=None, logfile=None): 

1012 """Initialise an Exp+FF preconditioner with given parameters. 

1013 

1014 Args: 

1015 r_cut, mu, c_stab, dim, reinitialize, array_convention: see 

1016 precon.__init__() 

1017 A: coefficient in exp(-A*r/r_NN). Default is A=3.0. 

1018 """ 

1019 if (morses is None and bonds is None and angles is None and 

1020 dihedrals is None): 

1021 raise ImportError( 

1022 'At least one of morses, bonds, angles or dihedrals must ' 

1023 'be defined!') 

1024 

1025 SparsePrecon.__init__(self, r_cut=r_cut, r_NN=r_NN, 

1026 mu=mu, mu_c=mu_c, dim=dim, c_stab=c_stab, 

1027 force_stab=force_stab, 

1028 reinitialize=reinitialize, 

1029 array_convention=array_convention, 

1030 solver=solver, 

1031 solve_tol=solve_tol, 

1032 apply_positions=apply_positions, 

1033 apply_cell=apply_cell, 

1034 estimate_mu_eigmode=estimate_mu_eigmode, 

1035 logfile=logfile) 

1036 

1037 self.A = A 

1038 self.hessian = hessian 

1039 self.morses = morses 

1040 self.bonds = bonds 

1041 self.angles = angles 

1042 self.dihedrals = dihedrals 

1043 

1044 def make_precon(self, atoms, reinitialize=None): 

1045 if self.r_NN is None: 

1046 self.r_NN = estimate_nearest_neighbour_distance(atoms, 

1047 self.neighbor_list) 

1048 

1049 if self.r_cut is None: 

1050 # This is the first time this function has been called, and no 

1051 # cutoff radius has been specified, so calculate it automatically. 

1052 self.r_cut = 2.0 * self.r_NN 

1053 elif self.r_cut < self.r_NN: 

1054 warning = ('WARNING: r_cut (%.2f) < r_NN (%.2f), ' 

1055 'increasing to 1.1*r_NN = %.2f' % (self.r_cut, 

1056 self.r_NN, 

1057 1.1 * self.r_NN)) 

1058 warnings.warn(warning) 

1059 self.r_cut = 1.1 * self.r_NN 

1060 

1061 if reinitialize is None: 

1062 # The caller has not specified whether or not to recalculate mu, 

1063 # so the Precon's setting is used. 

1064 reinitialize = self.reinitialize 

1065 

1066 if self.mu is None: 

1067 # Regardless of what the caller has specified, if we don't 

1068 # currently have a value of mu, then we need one. 

1069 reinitialize = True 

1070 

1071 if reinitialize: 

1072 self.estimate_mu(atoms) 

1073 

1074 if self.P is not None: 

1075 real_atoms = atoms 

1076 if isinstance(atoms, Filter): 

1077 real_atoms = atoms.atoms 

1078 if self.old_positions is None: 

1079 self.old_positions = real_atoms.positions 

1080 displacement, _ = find_mic(real_atoms.positions - 

1081 self.old_positions, 

1082 real_atoms.cell, real_atoms.pbc) 

1083 self.old_positions = real_atoms.get_positions() 

1084 max_abs_displacement = abs(displacement).max() 

1085 self.logfile.write('max(abs(displacements)) = %.2f A (%.2f r_NN)' % 

1086 (max_abs_displacement, 

1087 max_abs_displacement / self.r_NN)) 

1088 if max_abs_displacement < 0.5 * self.r_NN: 

1089 return 

1090 

1091 # Create the preconditioner: 

1092 start_time = time.time() 

1093 self._make_sparse_precon(atoms, force_stab=self.force_stab) 

1094 self.logfile.write('--- Precon created in %s seconds ---\n' % 

1095 (time.time() - start_time)) 

1096 

1097 def _make_sparse_precon(self, atoms, initial_assembly=False, 

1098 force_stab=False): 

1099 """Create a sparse preconditioner matrix based on the passed atoms. 

1100 

1101 Args: 

1102 atoms: the Atoms object used to create the preconditioner. 

1103 

1104 Returns: 

1105 A scipy.sparse.csr_matrix object, representing a d*N by d*N matrix 

1106 (where N is the number of atoms, and d is the value of self.dim). 

1107 BE AWARE that using numpy.dot() with this object will result in 

1108 errors/incorrect results - use the .dot method directly on the 

1109 sparse matrix instead. 

1110 

1111 """ 

1112 self.logfile.write('creating sparse precon: initial_assembly=%r, ' 

1113 'force_stab=%r, apply_positions=%r, ' 

1114 'apply_cell=%r\n' % 

1115 (initial_assembly, force_stab, 

1116 self.apply_positions, self.apply_cell)) 

1117 

1118 N = len(atoms) 

1119 start_time = time.time() 

1120 if self.apply_positions: 

1121 # compute neighbour list 

1122 i_list, j_list, rij_list, fixed_atoms = get_neighbours( 

1123 atoms, self.r_cut, self.neighbor_list) 

1124 self.logfile.write('--- neighbour list created in %s s ---\n' % 

1125 (time.time() - start_time)) 

1126 

1127 row = [] 

1128 col = [] 

1129 data = [] 

1130 

1131 # precon is mu_c*identity for cell DoF 

1132 start_time = time.time() 

1133 if isinstance(atoms, Filter): 

1134 i = N - 3 

1135 j = N - 2 

1136 k = N - 1 

1137 x = ijk_to_x(i, j, k) 

1138 row.extend(x) 

1139 col.extend(x) 

1140 if self.apply_cell: 

1141 data.extend(np.repeat(self.mu_c, 9)) 

1142 else: 

1143 data.extend(np.repeat(self.mu_c, 9)) 

1144 self.logfile.write('--- computed triplet format in %s s ---\n' % 

1145 (time.time() - start_time)) 

1146 

1147 conn = sparse.lil_matrix((N, N), dtype=bool) 

1148 

1149 if self.apply_positions and not initial_assembly: 

1150 if self.morses is not None: 

1151 for morse in self.morses: 

1152 self.add_morse(morse, atoms, row, col, data, conn) 

1153 

1154 if self.bonds is not None: 

1155 for bond in self.bonds: 

1156 self.add_bond(bond, atoms, row, col, data, conn) 

1157 

1158 if self.angles is not None: 

1159 for angle in self.angles: 

1160 self.add_angle(angle, atoms, row, col, data, conn) 

1161 

1162 if self.dihedrals is not None: 

1163 for dihedral in self.dihedrals: 

1164 self.add_dihedral(dihedral, atoms, row, col, data, conn) 

1165 

1166 if self.apply_positions: 

1167 for i, j, rij in zip(i_list, j_list, rij_list): 

1168 if not conn[i, j]: 

1169 coeff = self.get_coeff(rij) 

1170 x = ij_to_x(i, j) 

1171 row.extend(x) 

1172 col.extend(x) 

1173 data.extend(3 * [-coeff] + 3 * [coeff]) 

1174 

1175 row.extend(range(self.dim * N)) 

1176 col.extend(range(self.dim * N)) 

1177 if initial_assembly: 

1178 data.extend([self.mu * self.c_stab] * self.dim * N) 

1179 else: 

1180 data.extend([self.c_stab] * self.dim * N) 

1181 

1182 # create the matrix 

1183 start_time = time.time() 

1184 self.P = sparse.csc_matrix( 

1185 (data, (row, col)), shape=(self.dim * N, self.dim * N)) 

1186 self.logfile.write('--- created CSC matrix in %s s ---\n' % 

1187 (time.time() - start_time)) 

1188 

1189 if not initial_assembly: 

1190 self.P = apply_fixed(atoms, self.P) 

1191 

1192 self.P = self.P.tocsr() 

1193 self.create_solver() 

1194 

1195 

1196def make_precon(precon, atoms=None, **kwargs): 

1197 """ 

1198 Construct preconditioner from a string and optionally build for Atoms 

1199 

1200 Parameters 

1201 ---------- 

1202 precon - one of 'C1', 'Exp', 'Pfrommer', 'FF', 'Exp_FF', 'ID', None 

1203 or an instance of a subclass of `ase.optimize.precon.Precon` 

1204 

1205 atoms - ase.atoms.Atoms instance, optional 

1206 If present, build apreconditioner for this Atoms object 

1207 

1208 **kwargs - additional keyword arguments to pass to Precon constructor 

1209 

1210 Returns 

1211 ------- 

1212 precon - instance of relevant subclass of `ase.optimize.precon.Precon` 

1213 """ 

1214 lookup = { 

1215 'C1': C1, 

1216 'Exp': Exp, 

1217 'Pfrommer': Pfrommer, 

1218 'FF': FF, 

1219 'Exp_FF': Exp_FF, 

1220 'ID': IdentityPrecon, 

1221 'IdentityPrecon': IdentityPrecon, 

1222 None: IdentityPrecon 

1223 } 

1224 if isinstance(precon, str) or precon is None: 

1225 cls = lookup[precon] 

1226 precon = cls(**kwargs) 

1227 if atoms is not None: 

1228 precon.make_precon(atoms) 

1229 return precon 

1230 

1231 

1232class SplineFit: 

1233 """ 

1234 Fit a cubic spline fit to images 

1235 """ 

1236 

1237 def __init__(self, s, x): 

1238 self._s = s 

1239 self._x_data = x 

1240 self._x = CubicSpline(self._s, x, bc_type='not-a-knot') 

1241 self._dx_ds = self._x.derivative() 

1242 self._d2x_ds2 = self._x.derivative(2) 

1243 

1244 @property 

1245 def s(self): 

1246 return self._s 

1247 

1248 @property 

1249 def x_data(self): 

1250 return self._x_data 

1251 

1252 @property 

1253 def x(self): 

1254 return self._x 

1255 

1256 @property 

1257 def dx_ds(self): 

1258 return self._dx_ds 

1259 

1260 @property 

1261 def d2x_ds2(self): 

1262 return self._d2x_ds2 

1263 

1264 

1265class PreconImages: 

1266 def __init__(self, precon, images, **kwargs): 

1267 """ 

1268 Wrapper for a list of Precon objects and associated images 

1269 

1270 This is used when preconditioning a NEB object. Equation references 

1271 refer to Paper IV in the :class:`ase.neb.NEB` documentation, i.e. 

1272 

1273 S. Makri, C. Ortner and J. R. Kermode, J. Chem. Phys. 

1274 150, 094109 (2019) 

1275 https://dx.doi.org/10.1063/1.5064465 

1276 

1277 Args: 

1278 precon (str or list): preconditioner(s) to use 

1279 images (list of Atoms): Atoms objects that define the state 

1280 

1281 """ 

1282 self.images = images 

1283 if isinstance(precon, list): 

1284 if len(precon) != len(images): 

1285 raise ValueError(f'length mismatch: len(precon)={len(precon)} ' 

1286 f'!= len(images)={len(images)}') 

1287 self.precon = precon 

1288 return 

1289 P0 = make_precon(precon, images[0], **kwargs) 

1290 self.precon = [P0] 

1291 for image in images[1:]: 

1292 P = P0.copy() 

1293 P.make_precon(image) 

1294 self.precon.append(P) 

1295 self._spline = None 

1296 

1297 def __len__(self): 

1298 return len(self.precon) 

1299 

1300 def __iter__(self): 

1301 return iter(self.precon) 

1302 

1303 def __getitem__(self, index): 

1304 return self.precon[index] 

1305 

1306 def apply(self, all_forces, index=None): 

1307 """Apply preconditioners to stored images 

1308 

1309 Args: 

1310 all_forces (array): forces on images, shape (nimages, natoms, 3) 

1311 index (slice, optional): Which images to include. Defaults to all. 

1312 

1313 Returns: 

1314 precon_forces: array of preconditioned forces 

1315 """ 

1316 if index is None: 

1317 index = slice(None) 

1318 precon_forces = [] 

1319 for precon, image, forces in zip(self.precon[index], 

1320 self.images[index], 

1321 all_forces): 

1322 f_vec = forces.reshape(-1) 

1323 pf_vec, _ = precon.apply(f_vec, image) 

1324 precon_forces.append(pf_vec.reshape(-1, 3)) 

1325 

1326 return np.array(precon_forces) 

1327 

1328 def average_norm(self, i, j, dx): 

1329 """Average norm between images i and j 

1330 

1331 Args: 

1332 i (int): left image 

1333 j (int): right image 

1334 dx (array): vector 

1335 

1336 Returns: 

1337 norm: norm of vector wrt average of precons at i and j 

1338 """ 

1339 return np.sqrt(0.5 * (self.precon[i].dot(dx, dx) + 

1340 self.precon[j].dot(dx, dx))) 

1341 

1342 def get_tangent(self, i): 

1343 """Normalised tangent vector at image i 

1344 

1345 Args: 

1346 i (int): image of interest 

1347 

1348 Returns: 

1349 tangent: tangent vector, normalised with appropriate precon norm 

1350 """ 

1351 tangent = self.spline.dx_ds(self.spline.s[i]) 

1352 tangent /= self.precon[i].norm(tangent) 

1353 return tangent.reshape(-1, 3) 

1354 

1355 def get_residual(self, i, imgforce): 

1356 # residuals computed according to eq. 11 in the paper 

1357 P_dot_imgforce = self.precon[i].Pdot(imgforce.reshape(-1)) 

1358 return np.linalg.norm(P_dot_imgforce, np.inf) 

1359 

1360 def get_spring_force(self, i, k1, k2, tangent): 

1361 """Spring force on image 

1362 

1363 Args: 

1364 i (int): image of interest 

1365 k1 (float): spring constant for left spring 

1366 k2 (float): spring constant for right spring 

1367 tangent (array): tangent vector, shape (natoms, 3) 

1368 

1369 Returns: 

1370 eta: NEB spring forces, shape (natoms, 3) 

1371 """ 

1372 # Definition following Eq. 9 in Paper IV 

1373 nimages = len(self.images) 

1374 k = 0.5 * (k1 + k2) / (nimages ** 2) 

1375 curvature = self.spline.d2x_ds2(self.spline.s[i]).reshape(-1, 3) 

1376 # complete Eq. 9 by including the spring force 

1377 eta = k * self.precon[i].vdot(curvature, tangent) * tangent 

1378 return eta 

1379 

1380 def get_coordinates(self, positions=None): 

1381 """Compute displacements wrt appropriate precon metric for each image 

1382 

1383 Args: 

1384 positions (list or array, optional) - images positions. 

1385 Shape either (nimages * natoms, 3) or ((nimages-2)*natoms, 3) 

1386 

1387 Returns: 

1388 s : array shape (nimages,), reaction coordinates, in range [0, 1] 

1389 x : array shape (nimages, 3 * natoms), flat displacement vectors 

1390 """ 

1391 nimages = len(self.precon) 

1392 natoms = len(self.images[0]) 

1393 d_P = np.zeros(nimages) 

1394 x = np.zeros((nimages, 3 * natoms)) # flattened positions 

1395 if positions is None: 

1396 positions = [image.positions for image in self.images] 

1397 elif isinstance(positions, np.ndarray) and len(positions.shape) == 2: 

1398 positions = positions.reshape(-1, natoms, 3) 

1399 positions = [positions[i, :, :] for i in range(len(positions))] 

1400 if len(positions) == len(self.images) - 2: 

1401 # prepend and append the non-moving images 

1402 positions = ([self.images[0].positions] + positions + 

1403 [self.images[-1].positions]) 

1404 assert len(positions) == len(self.images) 

1405 

1406 x[0, :] = positions[0].reshape(-1) 

1407 for i in range(1, nimages): 

1408 x[i, :] = positions[i].reshape(-1) 

1409 dx, _ = find_mic(positions[i] - positions[i - 1], 

1410 self.images[i - 1].cell, 

1411 self.images[i - 1].pbc) 

1412 dx = dx.reshape(-1) 

1413 d_P[i] = self.average_norm(i, i - 1, dx) 

1414 

1415 s = d_P.cumsum() / d_P.sum() # Eq. A1 in paper IV 

1416 return s, x 

1417 

1418 def spline_fit(self, positions=None): 

1419 """Fit 3 * natoms cubic splines as a function of reaction coordinate 

1420 

1421 Returns: 

1422 fit : :class:`ase.optimize.precon.SplineFit` object 

1423 """ 

1424 s, x = self.get_coordinates(positions) 

1425 return SplineFit(s, x) 

1426 

1427 @property 

1428 def spline(self): 

1429 s, x = self.get_coordinates() 

1430 if self._spline and (np.abs(s - self._old_s).max() < 1e-6 and 

1431 np.abs(x - self._old_x).max() < 1e-6): 

1432 return self._spline 

1433 

1434 self._spline = self.spline_fit() 

1435 self._old_s = s 

1436 self._old_x = x 

1437 return self._spline