Coverage for /builds/debichem-team/python-ase/ase/optimize/precon/precon.py : 84.90%
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"""
5import sys
6import time
7import copy
8import warnings
9from abc import ABC, abstractmethod
11import numpy as np
12from scipy import sparse
13from scipy.sparse.linalg import spsolve
14from scipy.interpolate import CubicSpline
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
26try:
27 from pyamg import smoothed_aggregation_solver
28 have_pyamg = True
30 def create_pyamg_solver(P, max_levels=15):
31 return smoothed_aggregation_solver(
32 P, B=None,
33 strength=('symmetric', {'theta': 0.0}),
34 smooth=(
35 'jacobi', {'filter': True, 'weighting': 'local'}),
36 improve_candidates=([('block_gauss_seidel',
37 {'sweep': 'symmetric', 'iterations': 4})] +
38 [None] * (max_levels - 1)),
39 aggregate='standard',
40 presmoother=('block_gauss_seidel',
41 {'sweep': 'symmetric', 'iterations': 1}),
42 postsmoother=('block_gauss_seidel',
43 {'sweep': 'symmetric', 'iterations': 1}),
44 max_levels=max_levels,
45 max_coarse=300,
46 coarse_solver='pinv')
48except ImportError:
49 have_pyamg = False
51THz = 1e12 * 1. / units.s
54class Precon(ABC):
56 @abstractmethod
57 def make_precon(self, atoms, reinitialize=None):
58 """
59 Create a preconditioner matrix based on the passed set of atoms.
61 Creates a general-purpose preconditioner for use with optimization
62 algorithms, based on examining distances between pairs of atoms in the
63 lattice. The matrix will be stored in the attribute self.P and
64 returned.
66 Args:
67 atoms: the Atoms object used to create the preconditioner.
69 reinitialize: if True, parameters of the preconditioner
70 will be recalculated before the preconditioner matrix is
71 created. If False, they will be calculated only when they
72 do not currently have a value (ie, the first time this
73 function is called).
75 Returns:
76 P: A sparse scipy csr_matrix. BE AWARE that using
77 numpy.dot() with sparse matrices will result in
78 errors/incorrect results - use the .dot method directly
79 on the matrix instead.
80 """
81 ...
83 @abstractmethod
84 def Pdot(self, x):
85 """
86 Return the result of applying P to a vector x
87 """
88 ...
90 def dot(self, x, y):
91 """
92 Return the preconditioned dot product <P x, y>
94 Uses 128-bit floating point math for vector dot products
95 """
96 return longsum(self.Pdot(x) * y)
98 def norm(self, x):
99 """
100 Return the P-norm of x, where |x|_P = sqrt(<Px, x>)
101 """
102 return np.sqrt(self.dot(x, x))
104 def vdot(self, x, y):
105 return self.dot(x.reshape(-1),
106 y.reshape(-1))
108 @abstractmethod
109 def solve(self, x):
110 """
111 Solve the (sparse) linear system P x = y and return y
112 """
113 ...
115 def apply(self, forces, atoms):
116 """
117 Convenience wrapper that combines make_precon() and solve()
119 Parameters
120 ----------
121 forces: array
122 (len(atoms)*3) array of input forces
123 atoms: ase.atoms.Atoms
125 Returns
126 -------
127 precon_forces: array
128 (len(atoms), 3) array of preconditioned forces
129 residual: float
130 inf-norm of original forces, i.e. maximum absolute force
131 """
132 self.make_precon(atoms)
133 residual = np.linalg.norm(forces, np.inf)
134 precon_forces = self.solve(forces)
135 return precon_forces, residual
137 @abstractmethod
138 def copy(self):
139 ...
141 @abstractmethod
142 def asarray(self):
143 """
144 Array representation of preconditioner, as a dense matrix
145 """
146 ...
149class Logfile:
150 def __init__(self, logfile=None):
151 if isinstance(logfile, str):
152 if logfile == "-":
153 logfile = sys.stdout
154 else:
155 logfile = open(logfile, "a")
156 self.logfile = logfile
158 def write(self, *args):
159 if self.logfile is None:
160 return
161 self.logfile.write(*args)
164class SparsePrecon(Precon):
165 def __init__(self, r_cut=None, r_NN=None,
166 mu=None, mu_c=None,
167 dim=3, c_stab=0.1, force_stab=False,
168 reinitialize=False, array_convention='C',
169 solver="auto", solve_tol=1e-8,
170 apply_positions=True, apply_cell=True,
171 estimate_mu_eigmode=False, logfile=None, rng=None,
172 neighbour_list=neighbor_list):
173 """Initialise a preconditioner object based on passed parameters.
175 Parameters:
176 r_cut: float. This is a cut-off radius. The preconditioner matrix
177 will be created by considering pairs of atoms that are within a
178 distance r_cut of each other. For a regular lattice, this is
179 usually taken somewhere between the first- and second-nearest
180 neighbour distance. If r_cut is not provided, default is
181 2 * r_NN (see below)
182 r_NN: nearest neighbour distance. If not provided, this is
183 calculated
184 from input structure.
185 mu: float
186 energy scale for position degreees of freedom. If `None`, mu
187 is precomputed using finite difference derivatives.
188 mu_c: float
189 energy scale for cell degreees of freedom. Also precomputed
190 if None.
191 estimate_mu_eigmode:
192 If True, estimates mu based on the lowest eigenmodes of
193 unstabilised preconditioner. If False it uses the sine based
194 approach.
195 dim: int; dimensions of the problem
196 c_stab: float. The diagonal of the preconditioner matrix will have
197 a stabilisation constant added, which will be the value of
198 c_stab times mu.
199 force_stab:
200 If True, always add the stabilisation to diagnonal, regardless
201 of the presence of fixed atoms.
202 reinitialize: if True, the value of mu will be recalculated when
203 self.make_precon is called. This can be overridden in specific
204 cases with reinitialize argument in self.make_precon. If it
205 is set to True here, the value passed for mu will be
206 irrelevant unless reinitialize is set to False the first time
207 make_precon is called.
208 array_convention: Either 'C' or 'F' for Fortran; this will change
209 the preconditioner to reflect the ordering of the indices in
210 the vector it will operate on. The C convention assumes the
211 vector will be arranged atom-by-atom (ie [x1, y1, z1, x2, ...])
212 while the F convention assumes it will be arranged component
213 by component (ie [x1, x2, ..., y1, y2, ...]).
214 solver: One of "auto", "direct" or "pyamg", specifying whether to
215 use a direct sparse solver or PyAMG to solve P x = y.
216 Default is "auto" which uses PyAMG if available, falling
217 back to sparse solver if not. solve_tol: tolerance used for
218 PyAMG sparse linear solver, if available.
219 apply_positions: bool
220 if True, apply preconditioner to position DoF
221 apply_cell: bool
222 if True, apply preconditioner to cell DoF
223 logfile: file object or str
224 If *logfile* is a string, a file with that name will be opened.
225 Use '-' for stdout.
226 rng: None or np.random.RandomState instance
227 Random number generator to use for initialising pyamg solver
228 neighbor_list: function (optional). Optionally replace the built-in
229 ASE neighbour list with an alternative with the same call
230 signature, e.g. `matscipy.neighbours.neighbour_list`.
232 Raises:
233 ValueError for problem with arguments
235 """
236 self.r_NN = r_NN
237 self.r_cut = r_cut
238 self.mu = mu
239 self.mu_c = mu_c
240 self.estimate_mu_eigmode = estimate_mu_eigmode
241 self.c_stab = c_stab
242 self.force_stab = force_stab
243 self.array_convention = array_convention
244 self.reinitialize = reinitialize
245 self.P = None
246 self.old_positions = None
248 use_pyamg = False
249 if solver == "auto":
250 use_pyamg = have_pyamg
251 elif solver == "direct":
252 use_pyamg = False
253 elif solver == "pyamg":
254 if not have_pyamg:
255 raise RuntimeError("solver='pyamg', PyAMG can't be imported!")
256 use_pyamg = True
257 else:
258 raise ValueError('unknown solver - '
259 'should be "auto", "direct" or "pyamg"')
261 self.use_pyamg = use_pyamg
262 self.solve_tol = solve_tol
263 self.apply_positions = apply_positions
264 self.apply_cell = apply_cell
266 if dim < 1:
267 raise ValueError('Dimension must be at least 1')
268 self.dim = dim
269 self.logfile = Logfile(logfile)
271 if rng is None:
272 rng = np.random.RandomState()
273 self.rng = rng
275 self.neighbor_list = neighbor_list
277 def copy(self):
278 return copy.deepcopy(self)
280 def Pdot(self, x):
281 return self.P.dot(x)
283 def solve(self, x):
284 start_time = time.time()
285 if self.use_pyamg and have_pyamg:
286 y = self.ml.solve(x, x0=self.rng.rand(self.P.shape[0]),
287 tol=self.solve_tol,
288 accel='cg',
289 maxiter=300,
290 cycle='W')
291 else:
292 y = spsolve(self.P, x)
293 self.logfile.write('--- Precon applied in %s seconds ---\n' %
294 (time.time() - start_time))
295 return y
297 def estimate_mu(self, atoms, H=None):
298 r"""Estimate optimal preconditioner coefficient \mu
300 \mu is estimated from a numerical solution of
302 [dE(p+v) - dE(p)] \cdot v = \mu < P1 v, v >
304 with perturbation
306 v(x,y,z) = H P_lowest_nonzero_eigvec(x, y, z)
308 or
310 v(x,y,z) = H (sin(x / Lx), sin(y / Ly), sin(z / Lz))
312 After the optimal \mu is found, self.mu will be set to its value.
314 If `atoms` is an instance of Filter an additional \mu_c
315 will be computed for the cell degrees of freedom .
317 Args:
318 atoms: Atoms object for initial system
320 H: 3x3 array or None
321 Magnitude of deformation to apply.
322 Default is 1e-2*rNN*np.eye(3)
324 Returns:
325 mu : float
326 mu_c : float or None
327 """
328 logfile = self.logfile
330 if self.dim != 3:
331 raise ValueError('Automatic calculation of mu only possible for '
332 'three-dimensional preconditioners. Try setting '
333 'mu manually instead.')
335 if self.r_NN is None:
336 self.r_NN = estimate_nearest_neighbour_distance(atoms,
337 self.neighbor_list)
339 # deformation matrix, default is diagonal
340 if H is None:
341 H = 1e-2 * self.r_NN * np.eye(3)
343 # compute perturbation
344 p = atoms.get_positions()
346 if self.estimate_mu_eigmode:
347 self.mu = 1.0
348 self.mu_c = 1.0
349 c_stab = self.c_stab
350 self.c_stab = 0.0
352 if isinstance(atoms, Filter):
353 n = len(atoms.atoms)
354 else:
355 n = len(atoms)
356 self._make_sparse_precon(atoms, initial_assembly=True)
357 self.P = self.P[:3 * n, :3 * n]
358 eigvals, eigvecs = sparse.linalg.eigsh(self.P, k=4, which='SM')
360 logfile.write('estimate_mu(): lowest 4 eigvals = %f %f %f %f\n' %
361 (eigvals[0], eigvals[1], eigvals[2], eigvals[3]))
362 # check eigenvalues
363 if any(eigvals[0:3] > 1e-6):
364 raise ValueError('First 3 eigenvalues of preconditioner matrix'
365 'do not correspond to translational modes.')
366 elif eigvals[3] < 1e-6:
367 raise ValueError('Fourth smallest eigenvalue of '
368 'preconditioner matrix '
369 'is too small, increase r_cut.')
371 x = np.zeros(n)
372 for i in range(n):
373 x[i] = eigvecs[:, 3][3 * i]
374 x = x / np.linalg.norm(x)
375 if x[0] < 0:
376 x = -x
378 v = np.zeros(3 * len(atoms))
379 for i in range(n):
380 v[3 * i] = x[i]
381 v[3 * i + 1] = x[i]
382 v[3 * i + 2] = x[i]
383 v = v / np.linalg.norm(v)
384 v = v.reshape((-1, 3))
386 self.c_stab = c_stab
387 else:
388 Lx, Ly, Lz = [p[:, i].max() - p[:, i].min() for i in range(3)]
389 logfile.write('estimate_mu(): Lx=%.1f Ly=%.1f Lz=%.1f\n' %
390 (Lx, Ly, Lz))
392 x, y, z = p.T
393 # sine_vr = [np.sin(x/Lx), np.sin(y/Ly), np.sin(z/Lz)], but we need
394 # to take into account the possibility that one of Lx/Ly/Lz is
395 # zero.
396 sine_vr = [x, y, z]
398 for i, L in enumerate([Lx, Ly, Lz]):
399 if L == 0:
400 warnings.warn(
401 'Cell length L[%d] == 0. Setting H[%d,%d] = 0.' %
402 (i, i, i))
403 H[i, i] = 0.0
404 else:
405 sine_vr[i] = np.sin(sine_vr[i] / L)
407 v = np.dot(H, sine_vr).T
409 natoms = len(atoms)
410 if isinstance(atoms, Filter):
411 natoms = len(atoms.atoms)
412 eps = H / self.r_NN
413 v[natoms:, :] = eps
415 v1 = v.reshape(-1)
417 # compute LHS
418 dE_p = -atoms.get_forces().reshape(-1)
419 atoms_v = atoms.copy()
420 atoms_v.calc = atoms.calc
421 if isinstance(atoms, Filter):
422 atoms_v = atoms.__class__(atoms_v)
423 if hasattr(atoms, 'constant_volume'):
424 atoms_v.constant_volume = atoms.constant_volume
425 atoms_v.set_positions(p + v)
426 dE_p_plus_v = -atoms_v.get_forces().reshape(-1)
428 # compute left hand side
429 LHS = (dE_p_plus_v - dE_p) * v1
431 # assemble P with \mu = 1
432 self.mu = 1.0
433 self.mu_c = 1.0
435 self._make_sparse_precon(atoms, initial_assembly=True)
437 # compute right hand side
438 RHS = self.P.dot(v1) * v1
440 # use partial sums to compute separate mu for positions and cell DoFs
441 self.mu = longsum(LHS[:3 * natoms]) / longsum(RHS[:3 * natoms])
442 if self.mu < 1.0:
443 logfile.write('estimate_mu(): mu (%.3f) < 1.0, '
444 'capping at mu=1.0' % self.mu)
445 self.mu = 1.0
447 if isinstance(atoms, Filter):
448 self.mu_c = longsum(LHS[3 * natoms:]) / longsum(RHS[3 * natoms:])
449 if self.mu_c < 1.0:
450 logfile.write('estimate_mu(): mu_c (%.3f) < 1.0, '
451 'capping at mu_c=1.0\n' % self.mu_c)
452 self.mu_c = 1.0
454 logfile.write('estimate_mu(): mu=%r, mu_c=%r\n' %
455 (self.mu, self.mu_c))
457 self.P = None # force a rebuild with new mu (there may be fixed atoms)
458 return (self.mu, self.mu_c)
460 def asarray(self):
461 return np.array(self.P.todense())
463 def one_dim_to_ndim(self, csc_P, N):
464 """
465 Expand an N x N precon matrix to self.dim*N x self.dim * N
467 Args:
468 csc_P (sparse matrix): N x N sparse matrix, in CSC format
469 """
470 start_time = time.time()
471 if self.dim == 1:
472 P = csc_P
473 elif self.array_convention == 'F':
474 csc_P = csc_P.tocsr()
475 P = csc_P
476 for i in range(self.dim - 1):
477 P = sparse.block_diag((P, csc_P)).tocsr()
478 else:
479 # convert back to triplet and read the arrays
480 csc_P = csc_P.tocoo()
481 i = csc_P.row * self.dim
482 j = csc_P.col * self.dim
483 z = csc_P.data
485 # N-dimensionalise, interlaced coordinates
486 I = np.hstack([i + d for d in range(self.dim)])
487 J = np.hstack([j + d for d in range(self.dim)])
488 Z = np.hstack([z for d in range(self.dim)])
489 P = sparse.csc_matrix((Z, (I, J)),
490 shape=(self.dim * N, self.dim * N))
491 P = P.tocsr()
492 self.logfile.write('--- N-dim precon created in %s s ---\n' %
493 (time.time() - start_time))
494 return P
496 def create_solver(self):
497 if self.use_pyamg and have_pyamg:
498 start_time = time.time()
499 self.ml = create_pyamg_solver(self.P)
500 self.logfile.write('--- multi grid solver created in %s ---\n' %
501 (time.time() - start_time))
504class SparseCoeffPrecon(SparsePrecon):
505 def _make_sparse_precon(self, atoms, initial_assembly=False,
506 force_stab=False):
507 """Create a sparse preconditioner matrix based on the passed atoms.
509 Creates a general-purpose preconditioner for use with optimization
510 algorithms, based on examining distances between pairs of atoms in the
511 lattice. The matrix will be stored in the attribute self.P and
512 returned. Note that this function will use self.mu, whatever it is.
514 Args:
515 atoms: the Atoms object used to create the preconditioner.
517 Returns:
518 A scipy.sparse.csr_matrix object, representing a d*N by d*N matrix
519 (where N is the number of atoms, and d is the value of self.dim).
520 BE AWARE that using numpy.dot() with this object will result in
521 errors/incorrect results - use the .dot method directly on the
522 sparse matrix instead.
524 """
525 logfile = self.logfile
526 logfile.write('creating sparse precon: initial_assembly=%r, '
527 'force_stab=%r, apply_positions=%r, apply_cell=%r\n' %
528 (initial_assembly, force_stab, self.apply_positions,
529 self.apply_cell))
531 N = len(atoms)
532 diag_i = np.arange(N, dtype=int)
533 start_time = time.time()
534 if self.apply_positions:
535 # compute neighbour list
536 i, j, rij, fixed_atoms = get_neighbours(atoms, self.r_cut,
537 neighbor_list=self.neighbor_list)
538 logfile.write('--- neighbour list created in %s s --- \n' %
539 (time.time() - start_time))
541 # compute entries in triplet format: without the constraints
542 start_time = time.time()
543 coeff = self.get_coeff(rij)
544 diag_coeff = np.bincount(i, -coeff, minlength=N).astype(np.float64)
545 if force_stab or len(fixed_atoms) == 0:
546 logfile.write('adding stabilisation to precon')
547 diag_coeff += self.mu * self.c_stab
548 else:
549 diag_coeff = np.ones(N)
551 # precon is mu_c * identity for cell DoF
552 if isinstance(atoms, Filter):
553 if self.apply_cell:
554 diag_coeff[-3:] = self.mu_c
555 else:
556 diag_coeff[-3:] = 1.0
557 logfile.write('--- computed triplet format in %s s ---\n' %
558 (time.time() - start_time))
560 if self.apply_positions and not initial_assembly:
561 # apply the constraints
562 start_time = time.time()
563 mask = np.ones(N)
564 mask[fixed_atoms] = 0.0
565 coeff *= mask[i] * mask[j]
566 diag_coeff[fixed_atoms] = 1.0
567 logfile.write('--- applied fixed_atoms in %s s ---\n' %
568 (time.time() - start_time))
570 if self.apply_positions:
571 # remove zeros
572 start_time = time.time()
573 inz = np.nonzero(coeff)
574 i = np.hstack((i[inz], diag_i))
575 j = np.hstack((j[inz], diag_i))
576 coeff = np.hstack((coeff[inz], diag_coeff))
577 logfile.write('--- remove zeros in %s s ---\n' %
578 (time.time() - start_time))
579 else:
580 i = diag_i
581 j = diag_i
582 coeff = diag_coeff
584 # create an N x N precon matrix in compressed sparse column (CSC) format
585 start_time = time.time()
586 csc_P = sparse.csc_matrix((coeff, (i, j)), shape=(N, N))
587 logfile.write('--- created CSC matrix in %s s ---\n' %
588 (time.time() - start_time))
590 self.P = self.one_dim_to_ndim(csc_P, N)
591 self.create_solver()
593 def make_precon(self, atoms, reinitialize=None):
594 if self.r_NN is None:
595 self.r_NN = estimate_nearest_neighbour_distance(atoms,
596 self.neighbor_list)
598 if self.r_cut is None:
599 # This is the first time this function has been called, and no
600 # cutoff radius has been specified, so calculate it automatically.
601 self.r_cut = 2.0 * self.r_NN
602 elif self.r_cut < self.r_NN:
603 warning = ('WARNING: r_cut (%.2f) < r_NN (%.2f), '
604 'increasing to 1.1*r_NN = %.2f' % (self.r_cut,
605 self.r_NN,
606 1.1 * self.r_NN))
607 warnings.warn(warning)
608 self.r_cut = 1.1 * self.r_NN
610 if reinitialize is None:
611 # The caller has not specified whether or not to recalculate mu,
612 # so the Precon's setting is used.
613 reinitialize = self.reinitialize
615 if self.mu is None:
616 # Regardless of what the caller has specified, if we don't
617 # currently have a value of mu, then we need one.
618 reinitialize = True
620 if reinitialize:
621 self.estimate_mu(atoms)
623 if self.P is not None:
624 real_atoms = atoms
625 if isinstance(atoms, Filter):
626 real_atoms = atoms.atoms
627 if self.old_positions is None:
628 self.old_positions = real_atoms.positions
629 displacement, _ = find_mic(real_atoms.positions -
630 self.old_positions,
631 real_atoms.cell, real_atoms.pbc)
632 self.old_positions = real_atoms.get_positions()
633 max_abs_displacement = abs(displacement).max()
634 self.logfile.write('max(abs(displacements)) = %.2f A (%.2f r_NN)' %
635 (max_abs_displacement,
636 max_abs_displacement / self.r_NN))
637 if max_abs_displacement < 0.5 * self.r_NN:
638 return
640 start_time = time.time()
641 self._make_sparse_precon(atoms, force_stab=self.force_stab)
642 self.logfile.write('--- Precon created in %s seconds --- \n' %
643 (time.time() - start_time))
645 @abstractmethod
646 def get_coeff(self, r):
647 ...
650class Pfrommer(Precon):
651 """
652 Use initial guess for inverse Hessian from Pfrommer et al. as a
653 simple preconditioner
655 J. Comput. Phys. vol 131 p233-240 (1997)
656 """
658 def __init__(self, bulk_modulus=500 * units.GPa, phonon_frequency=50 * THz,
659 apply_positions=True, apply_cell=True):
660 """
661 Default bulk modulus is 500 GPa and default phonon frequency is 50 THz
662 """
663 self.bulk_modulus = bulk_modulus
664 self.phonon_frequency = phonon_frequency
665 self.apply_positions = apply_positions
666 self.apply_cell = apply_cell
667 self.H0 = None
669 def make_precon(self, atoms, reinitialize=None):
670 if self.H0 is not None:
671 # only build H0 on first call
672 return
674 variable_cell = False
675 if isinstance(atoms, Filter):
676 variable_cell = True
677 atoms = atoms.atoms
679 # position DoF
680 omega = self.phonon_frequency
681 mass = atoms.get_masses().mean()
682 block = np.eye(3) / (mass * omega**2)
683 blocks = [block] * len(atoms)
685 # cell DoF
686 if variable_cell:
687 coeff = 1.0
688 if self.apply_cell:
689 coeff = 1.0 / (3 * self.bulk_modulus)
690 blocks.append(np.diag([coeff] * 9))
692 self.H0 = sparse.block_diag(blocks, format='csr')
693 return
695 def Pdot(self, x):
696 return self.H0.solve(x)
698 def solve(self, x):
699 y = self.H0.dot(x)
700 return y
702 def copy(self):
703 return Pfrommer(self.bulk_modulus,
704 self.phonon_frequency,
705 self.apply_positions,
706 self.apply_cell)
708 def asarray(self):
709 return np.array(self.H0.todense())
712class IdentityPrecon(Precon):
713 """
714 Dummy preconditioner which does not modify forces
715 """
717 def make_precon(self, atoms, reinitialize=None):
718 self.atoms = atoms
720 def Pdot(self, x):
721 return x
723 def solve(self, x):
724 return x
726 def copy(self):
727 return IdentityPrecon()
729 def asarray(self):
730 return np.eye(3 * len(self.atoms))
733class C1(SparseCoeffPrecon):
734 """Creates matrix by inserting a constant whenever r_ij is less than r_cut.
735 """
737 def __init__(self, r_cut=None, mu=None, mu_c=None, dim=3, c_stab=0.1,
738 force_stab=False,
739 reinitialize=False, array_convention='C',
740 solver="auto", solve_tol=1e-9,
741 apply_positions=True, apply_cell=True, logfile=None):
742 super().__init__(r_cut=r_cut, mu=mu, mu_c=mu_c,
743 dim=dim, c_stab=c_stab,
744 force_stab=force_stab,
745 reinitialize=reinitialize,
746 array_convention=array_convention,
747 solver=solver, solve_tol=solve_tol,
748 apply_positions=apply_positions,
749 apply_cell=apply_cell,
750 logfile=logfile)
752 def get_coeff(self, r):
753 return -self.mu * np.ones_like(r)
756class Exp(SparseCoeffPrecon):
757 """Creates matrix with values decreasing exponentially with distance.
758 """
760 def __init__(self, A=3.0, r_cut=None, r_NN=None, mu=None, mu_c=None,
761 dim=3, c_stab=0.1,
762 force_stab=False, reinitialize=False, array_convention='C',
763 solver="auto", solve_tol=1e-9,
764 apply_positions=True, apply_cell=True,
765 estimate_mu_eigmode=False, logfile=None):
766 """
767 Initialise an Exp preconditioner with given parameters.
769 Args:
770 r_cut, mu, c_stab, dim, sparse, reinitialize, array_convention: see
771 precon.__init__()
772 A: coefficient in exp(-A*r/r_NN). Default is A=3.0.
773 """
774 super().__init__(r_cut=r_cut, r_NN=r_NN,
775 mu=mu, mu_c=mu_c, dim=dim, c_stab=c_stab,
776 force_stab=force_stab,
777 reinitialize=reinitialize,
778 array_convention=array_convention,
779 solver=solver,
780 solve_tol=solve_tol,
781 apply_positions=apply_positions,
782 apply_cell=apply_cell,
783 estimate_mu_eigmode=estimate_mu_eigmode,
784 logfile=logfile)
786 self.A = A
788 def get_coeff(self, r):
789 return -self.mu * np.exp(-self.A * (r / self.r_NN - 1))
792def ij_to_x(i, j):
793 x = [3 * i, 3 * i + 1, 3 * i + 2,
794 3 * j, 3 * j + 1, 3 * j + 2]
795 return x
798def ijk_to_x(i, j, k):
799 x = [3 * i, 3 * i + 1, 3 * i + 2,
800 3 * j, 3 * j + 1, 3 * j + 2,
801 3 * k, 3 * k + 1, 3 * k + 2]
802 return x
805def ijkl_to_x(i, j, k, l):
806 x = [3 * i, 3 * i + 1, 3 * i + 2,
807 3 * j, 3 * j + 1, 3 * j + 2,
808 3 * k, 3 * k + 1, 3 * k + 2,
809 3 * l, 3 * l + 1, 3 * l + 2]
810 return x
813def apply_fixed(atoms, P):
814 fixed_atoms = []
815 for constraint in atoms.constraints:
816 if isinstance(constraint, FixAtoms):
817 fixed_atoms.extend(list(constraint.index))
818 else:
819 raise TypeError(
820 'only FixAtoms constraints are supported by Precon class')
821 if len(fixed_atoms) != 0:
822 P = P.tolil()
823 for i in fixed_atoms:
824 P[i, :] = 0.0
825 P[:, i] = 0.0
826 P[i, i] = 1.0
827 return P
830class FF(SparsePrecon):
831 """Creates matrix using morse/bond/angle/dihedral force field parameters.
832 """
834 def __init__(self, dim=3, c_stab=0.1, force_stab=False,
835 array_convention='C', solver="auto", solve_tol=1e-9,
836 apply_positions=True, apply_cell=True,
837 hessian='spectral', morses=None, bonds=None, angles=None,
838 dihedrals=None, logfile=None):
839 """Initialise an FF preconditioner with given parameters.
841 Args:
842 dim, c_stab, force_stab, array_convention, use_pyamg, solve_tol:
843 see SparsePrecon.__init__()
844 morses: Morse instance
845 bonds: Bond instance
846 angles: Angle instance
847 dihedrals: Dihedral instance
848 """
850 if (morses is None and bonds is None and angles is None and
851 dihedrals is None):
852 raise ImportError(
853 'At least one of morses, bonds, angles or dihedrals must be '
854 'defined!')
856 super().__init__(dim=dim, c_stab=c_stab,
857 force_stab=force_stab,
858 array_convention=array_convention,
859 solver=solver,
860 solve_tol=solve_tol,
861 apply_positions=apply_positions,
862 apply_cell=apply_cell, logfile=logfile)
864 self.hessian = hessian
865 self.morses = morses
866 self.bonds = bonds
867 self.angles = angles
868 self.dihedrals = dihedrals
870 def make_precon(self, atoms, reinitialize=None):
871 start_time = time.time()
872 self._make_sparse_precon(atoms, force_stab=self.force_stab)
873 self.logfile.write('--- Precon created in %s seconds ---\n'
874 % (time.time() - start_time))
876 def add_morse(self, morse, atoms, row, col, data, conn=None):
877 if self.hessian == 'reduced':
878 i, j, Hx = ff.get_morse_potential_reduced_hessian(
879 atoms, morse)
880 elif self.hessian == 'spectral':
881 i, j, Hx = ff.get_morse_potential_hessian(
882 atoms, morse, spectral=True)
883 else:
884 raise NotImplementedError('Not implemented hessian')
885 x = ij_to_x(i, j)
886 row.extend(np.repeat(x, 6))
887 col.extend(np.tile(x, 6))
888 data.extend(Hx.flatten())
889 if conn is not None:
890 conn[i, j] = True
891 conn[j, i] = True
893 def add_bond(self, bond, atoms, row, col, data, conn=None):
894 if self.hessian == 'reduced':
895 i, j, Hx = ff.get_bond_potential_reduced_hessian(
896 atoms, bond, self.morses)
897 elif self.hessian == 'spectral':
898 i, j, Hx = ff.get_bond_potential_hessian(
899 atoms, bond, self.morses, spectral=True)
900 else:
901 raise NotImplementedError('Not implemented hessian')
902 x = ij_to_x(i, j)
903 row.extend(np.repeat(x, 6))
904 col.extend(np.tile(x, 6))
905 data.extend(Hx.flatten())
906 if conn is not None:
907 conn[i, j] = True
908 conn[j, i] = True
910 def add_angle(self, angle, atoms, row, col, data, conn=None):
911 if self.hessian == 'reduced':
912 i, j, k, Hx = ff.get_angle_potential_reduced_hessian(
913 atoms, angle, self.morses)
914 elif self.hessian == 'spectral':
915 i, j, k, Hx = ff.get_angle_potential_hessian(
916 atoms, angle, self.morses, spectral=True)
917 else:
918 raise NotImplementedError('Not implemented hessian')
919 x = ijk_to_x(i, j, k)
920 row.extend(np.repeat(x, 9))
921 col.extend(np.tile(x, 9))
922 data.extend(Hx.flatten())
923 if conn is not None:
924 conn[i, j] = conn[i, k] = conn[j, k] = True
925 conn[j, i] = conn[k, i] = conn[k, j] = True
927 def add_dihedral(self, dihedral, atoms, row, col, data, conn=None):
928 if self.hessian == 'reduced':
929 i, j, k, l, Hx = \
930 ff.get_dihedral_potential_reduced_hessian(
931 atoms, dihedral, self.morses)
932 elif self.hessian == 'spectral':
933 i, j, k, l, Hx = ff.get_dihedral_potential_hessian(
934 atoms, dihedral, self.morses, spectral=True)
935 else:
936 raise NotImplementedError('Not implemented hessian')
937 x = ijkl_to_x(i, j, k, l)
938 row.extend(np.repeat(x, 12))
939 col.extend(np.tile(x, 12))
940 data.extend(Hx.flatten())
941 if conn is not None:
942 conn[i, j] = conn[i, k] = conn[i, l] = conn[
943 j, k] = conn[j, l] = conn[k, l] = True
944 conn[j, i] = conn[k, i] = conn[l, i] = conn[
945 k, j] = conn[l, j] = conn[l, k] = True
947 def _make_sparse_precon(self, atoms, initial_assembly=False,
948 force_stab=False):
949 N = len(atoms)
951 row = []
952 col = []
953 data = []
955 if self.morses is not None:
956 for morse in self.morses:
957 self.add_morse(morse, atoms, row, col, data)
959 if self.bonds is not None:
960 for bond in self.bonds:
961 self.add_bond(bond, atoms, row, col, data)
963 if self.angles is not None:
964 for angle in self.angles:
965 self.add_angle(angle, atoms, row, col, data)
967 if self.dihedrals is not None:
968 for dihedral in self.dihedrals:
969 self.add_dihedral(dihedral, atoms, row, col, data)
971 # add the diagonal
972 row.extend(range(self.dim * N))
973 col.extend(range(self.dim * N))
974 data.extend([self.c_stab] * self.dim * N)
976 # create the matrix
977 start_time = time.time()
978 self.P = sparse.csc_matrix(
979 (data, (row, col)), shape=(self.dim * N, self.dim * N))
980 self.logfile.write('--- created CSC matrix in %s s ---\n' %
981 (time.time() - start_time))
983 self.P = apply_fixed(atoms, self.P)
984 self.P = self.P.tocsr()
985 self.logfile.write('--- N-dim precon created in %s s ---\n' %
986 (time.time() - start_time))
987 self.create_solver()
990class Exp_FF(Exp, FF):
991 """Creates matrix with values decreasing exponentially with distance.
992 """
994 def __init__(self, A=3.0, r_cut=None, r_NN=None, mu=None, mu_c=None,
995 dim=3, c_stab=0.1,
996 force_stab=False, reinitialize=False, array_convention='C',
997 solver="auto", solve_tol=1e-9,
998 apply_positions=True, apply_cell=True,
999 estimate_mu_eigmode=False,
1000 hessian='spectral', morses=None, bonds=None, angles=None,
1001 dihedrals=None, logfile=None):
1002 """Initialise an Exp+FF preconditioner with given parameters.
1004 Args:
1005 r_cut, mu, c_stab, dim, reinitialize, array_convention: see
1006 precon.__init__()
1007 A: coefficient in exp(-A*r/r_NN). Default is A=3.0.
1008 """
1009 if (morses is None and bonds is None and angles is None and
1010 dihedrals is None):
1011 raise ImportError(
1012 'At least one of morses, bonds, angles or dihedrals must '
1013 'be defined!')
1015 SparsePrecon.__init__(self, r_cut=r_cut, r_NN=r_NN,
1016 mu=mu, mu_c=mu_c, dim=dim, c_stab=c_stab,
1017 force_stab=force_stab,
1018 reinitialize=reinitialize,
1019 array_convention=array_convention,
1020 solver=solver,
1021 solve_tol=solve_tol,
1022 apply_positions=apply_positions,
1023 apply_cell=apply_cell,
1024 estimate_mu_eigmode=estimate_mu_eigmode,
1025 logfile=logfile)
1027 self.A = A
1028 self.hessian = hessian
1029 self.morses = morses
1030 self.bonds = bonds
1031 self.angles = angles
1032 self.dihedrals = dihedrals
1034 def make_precon(self, atoms, reinitialize=None):
1035 if self.r_NN is None:
1036 self.r_NN = estimate_nearest_neighbour_distance(atoms,
1037 self.neighbor_list)
1039 if self.r_cut is None:
1040 # This is the first time this function has been called, and no
1041 # cutoff radius has been specified, so calculate it automatically.
1042 self.r_cut = 2.0 * self.r_NN
1043 elif self.r_cut < self.r_NN:
1044 warning = ('WARNING: r_cut (%.2f) < r_NN (%.2f), '
1045 'increasing to 1.1*r_NN = %.2f' % (self.r_cut,
1046 self.r_NN,
1047 1.1 * self.r_NN))
1048 warnings.warn(warning)
1049 self.r_cut = 1.1 * self.r_NN
1051 if reinitialize is None:
1052 # The caller has not specified whether or not to recalculate mu,
1053 # so the Precon's setting is used.
1054 reinitialize = self.reinitialize
1056 if self.mu is None:
1057 # Regardless of what the caller has specified, if we don't
1058 # currently have a value of mu, then we need one.
1059 reinitialize = True
1061 if reinitialize:
1062 self.estimate_mu(atoms)
1064 if self.P is not None:
1065 real_atoms = atoms
1066 if isinstance(atoms, Filter):
1067 real_atoms = atoms.atoms
1068 if self.old_positions is None:
1069 self.old_positions = real_atoms.positions
1070 displacement, _ = find_mic(real_atoms.positions -
1071 self.old_positions,
1072 real_atoms.cell, real_atoms.pbc)
1073 self.old_positions = real_atoms.get_positions()
1074 max_abs_displacement = abs(displacement).max()
1075 self.logfile.write('max(abs(displacements)) = %.2f A (%.2f r_NN)' %
1076 (max_abs_displacement,
1077 max_abs_displacement / self.r_NN))
1078 if max_abs_displacement < 0.5 * self.r_NN:
1079 return
1081 # Create the preconditioner:
1082 start_time = time.time()
1083 self._make_sparse_precon(atoms, force_stab=self.force_stab)
1084 self.logfile.write('--- Precon created in %s seconds ---\n' %
1085 (time.time() - start_time))
1087 def _make_sparse_precon(self, atoms, initial_assembly=False,
1088 force_stab=False):
1089 """Create a sparse preconditioner matrix based on the passed atoms.
1091 Args:
1092 atoms: the Atoms object used to create the preconditioner.
1094 Returns:
1095 A scipy.sparse.csr_matrix object, representing a d*N by d*N matrix
1096 (where N is the number of atoms, and d is the value of self.dim).
1097 BE AWARE that using numpy.dot() with this object will result in
1098 errors/incorrect results - use the .dot method directly on the
1099 sparse matrix instead.
1101 """
1102 self.logfile.write('creating sparse precon: initial_assembly=%r, '
1103 'force_stab=%r, apply_positions=%r, '
1104 'apply_cell=%r\n' %
1105 (initial_assembly, force_stab,
1106 self.apply_positions, self.apply_cell))
1108 N = len(atoms)
1109 start_time = time.time()
1110 if self.apply_positions:
1111 # compute neighbour list
1112 i_list, j_list, rij_list, fixed_atoms = get_neighbours(
1113 atoms, self.r_cut, self.neighbor_list)
1114 self.logfile.write('--- neighbour list created in %s s ---\n' %
1115 (time.time() - start_time))
1117 row = []
1118 col = []
1119 data = []
1121 # precon is mu_c*identity for cell DoF
1122 start_time = time.time()
1123 if isinstance(atoms, Filter):
1124 i = N - 3
1125 j = N - 2
1126 k = N - 1
1127 x = ijk_to_x(i, j, k)
1128 row.extend(x)
1129 col.extend(x)
1130 if self.apply_cell:
1131 data.extend(np.repeat(self.mu_c, 9))
1132 else:
1133 data.extend(np.repeat(self.mu_c, 9))
1134 self.logfile.write('--- computed triplet format in %s s ---\n' %
1135 (time.time() - start_time))
1137 conn = sparse.lil_matrix((N, N), dtype=bool)
1139 if self.apply_positions and not initial_assembly:
1140 if self.morses is not None:
1141 for morse in self.morses:
1142 self.add_morse(morse, atoms, row, col, data, conn)
1144 if self.bonds is not None:
1145 for bond in self.bonds:
1146 self.add_bond(bond, atoms, row, col, data, conn)
1148 if self.angles is not None:
1149 for angle in self.angles:
1150 self.add_angle(angle, atoms, row, col, data, conn)
1152 if self.dihedrals is not None:
1153 for dihedral in self.dihedrals:
1154 self.add_dihedral(dihedral, atoms, row, col, data, conn)
1156 if self.apply_positions:
1157 for i, j, rij in zip(i_list, j_list, rij_list):
1158 if not conn[i, j]:
1159 coeff = self.get_coeff(rij)
1160 x = ij_to_x(i, j)
1161 row.extend(x)
1162 col.extend(x)
1163 data.extend(3 * [-coeff] + 3 * [coeff])
1165 row.extend(range(self.dim * N))
1166 col.extend(range(self.dim * N))
1167 if initial_assembly:
1168 data.extend([self.mu * self.c_stab] * self.dim * N)
1169 else:
1170 data.extend([self.c_stab] * self.dim * N)
1172 # create the matrix
1173 start_time = time.time()
1174 self.P = sparse.csc_matrix(
1175 (data, (row, col)), shape=(self.dim * N, self.dim * N))
1176 self.logfile.write('--- created CSC matrix in %s s ---\n' %
1177 (time.time() - start_time))
1179 if not initial_assembly:
1180 self.P = apply_fixed(atoms, self.P)
1182 self.P = self.P.tocsr()
1183 self.create_solver()
1186def make_precon(precon, atoms=None, **kwargs):
1187 """
1188 Construct preconditioner from a string and optionally build for Atoms
1190 Parameters
1191 ----------
1192 precon - one of 'C1', 'Exp', 'Pfrommer', 'FF', 'Exp_FF', 'ID', None
1193 or an instance of a subclass of `ase.optimize.precon.Precon`
1195 atoms - ase.atoms.Atoms instance, optional
1196 If present, build apreconditioner for this Atoms object
1198 **kwargs - additional keyword arguments to pass to Precon constructor
1200 Returns
1201 -------
1202 precon - instance of relevant subclass of `ase.optimize.precon.Precon`
1203 """
1204 lookup = {
1205 'C1': C1,
1206 'Exp': Exp,
1207 'Pfrommer': Pfrommer,
1208 'FF': FF,
1209 'Exp_FF': Exp_FF,
1210 'ID': IdentityPrecon,
1211 'IdentityPrecon': IdentityPrecon,
1212 None: IdentityPrecon
1213 }
1214 if isinstance(precon, str) or precon is None:
1215 cls = lookup[precon]
1216 precon = cls(**kwargs)
1217 if atoms is not None:
1218 precon.make_precon(atoms)
1219 return precon
1222class SplineFit:
1223 """
1224 Fit a cubic spline fit to images
1225 """
1226 def __init__(self, s, x):
1227 self._s = s
1228 self._x_data = x
1229 self._x = CubicSpline(self._s, x, bc_type='not-a-knot')
1230 self._dx_ds = self._x.derivative()
1231 self._d2x_ds2 = self._x.derivative(2)
1233 @property
1234 def s(self):
1235 return self._s
1237 @property
1238 def x_data(self):
1239 return self._x_data
1241 @property
1242 def x(self):
1243 return self._x
1245 @property
1246 def dx_ds(self):
1247 return self._dx_ds
1249 @property
1250 def d2x_ds2(self):
1251 return self._d2x_ds2
1254class PreconImages:
1255 def __init__(self, precon, images, **kwargs):
1256 """
1257 Wrapper for a list of Precon objects and associated images
1259 This is used when preconditioning a NEB object. Equation references
1260 refer to Paper IV in the :class:`ase.neb.NEB` documentation, i.e.
1262 S. Makri, C. Ortner and J. R. Kermode, J. Chem. Phys.
1263 150, 094109 (2019)
1264 https://dx.doi.org/10.1063/1.5064465
1266 Args:
1267 precon (str or list): preconditioner(s) to use
1268 images (list of Atoms): Atoms objects that define the state
1270 """
1271 self.images = images
1272 if isinstance(precon, list):
1273 if len(precon) != len(images):
1274 raise ValueError(f'length mismatch: len(precon)={len(precon)} '
1275 f'!= len(images)={len(images)}')
1276 self.precon = precon
1277 return
1278 P0 = make_precon(precon, images[0], **kwargs)
1279 self.precon = [P0]
1280 for image in images[1:]:
1281 P = P0.copy()
1282 P.make_precon(image)
1283 self.precon.append(P)
1284 self._spline = None
1286 def __len__(self):
1287 return len(self.precon)
1289 def __iter__(self):
1290 return iter(self.precon)
1292 def __getitem__(self, index):
1293 return self.precon[index]
1295 def apply(self, all_forces, index=None):
1296 """Apply preconditioners to stored images
1298 Args:
1299 all_forces (array): forces on images, shape (nimages, natoms, 3)
1300 index (slice, optional): Which images to include. Defaults to all.
1302 Returns:
1303 precon_forces: array of preconditioned forces
1304 """
1305 if index is None:
1306 index = slice(None)
1307 precon_forces = []
1308 for precon, image, forces in zip(self.precon[index],
1309 self.images[index],
1310 all_forces):
1311 f_vec = forces.reshape(-1)
1312 pf_vec, _ = precon.apply(f_vec, image)
1313 precon_forces.append(pf_vec.reshape(-1, 3))
1315 return np.array(precon_forces)
1317 def average_norm(self, i, j, dx):
1318 """Average norm between images i and j
1320 Args:
1321 i (int): left image
1322 j (int): right image
1323 dx (array): vector
1325 Returns:
1326 norm: norm of vector wrt average of precons at i and j
1327 """
1328 return np.sqrt(0.5 * (self.precon[i].dot(dx, dx) +
1329 self.precon[j].dot(dx, dx)))
1331 def get_tangent(self, i):
1332 """Normalised tangent vector at image i
1334 Args:
1335 i (int): image of interest
1337 Returns:
1338 tangent: tangent vector, normalised with appropriate precon norm
1339 """
1340 tangent = self.spline.dx_ds(self.spline.s[i])
1341 tangent /= self.precon[i].norm(tangent)
1342 return tangent.reshape(-1, 3)
1344 def get_residual(self, i, imgforce):
1345 # residuals computed according to eq. 11 in the paper
1346 P_dot_imgforce = self.precon[i].Pdot(imgforce.reshape(-1))
1347 return np.linalg.norm(P_dot_imgforce, np.inf)
1349 def get_spring_force(self, i, k1, k2, tangent):
1350 """Spring force on image
1352 Args:
1353 i (int): image of interest
1354 k1 (float): spring constant for left spring
1355 k2 (float): spring constant for right spring
1356 tangent (array): tangent vector, shape (natoms, 3)
1358 Returns:
1359 eta: NEB spring forces, shape (natoms, 3)
1360 """
1361 # Definition following Eq. 9 in Paper IV
1362 nimages = len(self.images)
1363 k = 0.5 * (k1 + k2) / (nimages ** 2)
1364 curvature = self.spline.d2x_ds2(self.spline.s[i]).reshape(-1, 3)
1365 # complete Eq. 9 by including the spring force
1366 eta = k * self.precon[i].vdot(curvature, tangent) * tangent
1367 return eta
1369 def get_coordinates(self, positions=None):
1370 """Compute displacements wrt appropriate precon metric for each image
1372 Args:
1373 positions (list or array, optional) - images positions.
1374 Shape either (nimages * natoms, 3) or ((nimages-2)*natoms, 3)
1376 Returns:
1377 s : array shape (nimages,), reaction coordinates, in range [0, 1]
1378 x : array shape (nimages, 3 * natoms), flat displacement vectors
1379 """
1380 nimages = len(self.precon)
1381 natoms = len(self.images[0])
1382 d_P = np.zeros(nimages)
1383 x = np.zeros((nimages, 3 * natoms)) # flattened positions
1384 if positions is None:
1385 positions = [image.positions for image in self.images]
1386 elif isinstance(positions, np.ndarray) and len(positions.shape) == 2:
1387 positions = positions.reshape(-1, natoms, 3)
1388 positions = [positions[i, :, :] for i in range(len(positions))]
1389 if len(positions) == len(self.images) - 2:
1390 # prepend and append the non-moving images
1391 positions = ([self.images[0].positions] + positions +
1392 [self.images[-1].positions])
1393 assert len(positions) == len(self.images)
1395 x[0, :] = positions[0].reshape(-1)
1396 for i in range(1, nimages):
1397 x[i, :] = positions[i].reshape(-1)
1398 dx, _ = find_mic(positions[i] - positions[i - 1],
1399 self.images[i - 1].cell,
1400 self.images[i - 1].pbc)
1401 dx = dx.reshape(-1)
1402 d_P[i] = self.average_norm(i, i - 1, dx)
1404 s = d_P.cumsum() / d_P.sum() # Eq. A1 in paper IV
1405 return s, x
1407 def spline_fit(self, positions=None):
1408 """Fit 3 * natoms cubic splines as a function of reaction coordinate
1410 Returns:
1411 fit : :class:`ase.optimize.precon.SplineFit` object
1412 """
1413 s, x = self.get_coordinates(positions)
1414 return SplineFit(s, x)
1416 @property
1417 def spline(self):
1418 s, x = self.get_coordinates()
1419 if self._spline and (np.abs(s - self._old_s).max() < 1e-6 and
1420 np.abs(x - self._old_x).max() < 1e-6):
1421 return self._spline
1423 self._spline = self.spline_fit()
1424 self._old_s = s
1425 self._old_x = x
1426 return self._spline