Coverage for ase / _4 / symopt / relax.py: 91.99%
337 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-05-21 15:52 +0000
1from dataclasses import dataclass
3import numpy as np
4import scipy
6from ase import Atoms
7from ase._4.symopt.relax_print import (
8 pprint_atoms,
9 pretty,
10 pretty_atomic_dofs,
11 pretty_dofs,
12 pretty_header,
13 pretty_subheader,
14)
15from ase.cell import Cell
16from ase.utils.abc import Optimizable
19def green(text: str) -> str:
20 return f'\x1b[32m{text}\x1b[0m'
23def chol_derivative(A, dA):
24 """
25 Compute the derivative of the Cholesky factorization.
27 A = L L^T
29 where L is lower-triangular. For a small symmetric perturbation `dA`,
30 this function returns an approximation of the differential `dL` of
31 the Cholesky factor:
33 L + dL ≈ chol(A + dA).
35 Either L or A must be given as an input.
36 """
37 A = (A + A.T) / 2
38 dA = (dA + dA.T) / 2
40 L = np.linalg.cholesky(A)
41 Linv = np.linalg.inv(L)
42 S = Linv @ dA @ Linv.T
43 X = np.tril(S)
44 X[np.diag_indices_from(X)] *= 0.5
45 return L @ X
48def symmetrize_atoms(
49 S_ac: np.ndarray, U_scc: np.ndarray, f_sc: np.ndarray, atommap_sa, tol=1e-12
50):
51 """
52 Symmetrize fractional atomic coordinates under a space-group.
54 Given atomic scaled positions `S_ac` and a set of space-group operations
55 (U_scc, f_sc), this function projects the positions onto the symmetry-
56 invariant subspace by averaging over all symmetry-related images.
58 Parameters
59 ----------
60 S_ac : ndarray, shape (na, 3)
61 Scaled atomic coordinates.
63 U_scc : ndarray, shape (ns, 3, 3)
64 Rotation matrices.
66 f_sc : ndarray, shape (ns, 3)
67 Translation vectors.
69 atommap_sa : ndarray, shape (ns, na)
70 Mapping such that atommap_sa[s, a] gives the index of the atom
71 to which atom `a` is mapped by symmetry operation `s`.
73 tol : float, optional
74 Tolerance for snapping values close to 0 or 1 back to 0.
76 Returns
77 -------
78 Ssym_ac : ndarray, shape (na, 3)
79 Symmetrized fractional atomic coordinates in [0, 1).
81 Notes
82 -----
83 The symmetrization is done in the complex phase representation
84 exp(2πi x) to correctly average periodic fractional coordinates.
85 """
86 ns, na = atommap_sa.shape
87 tmp_Ssym_ac = np.zeros_like(S_ac, dtype=np.complex128)
88 for a in range(na):
89 for s in range(ns):
90 new = U_scc[s].T @ S_ac[a] - f_sc[s]
91 tmp_Ssym_ac[atommap_sa[s, a]] += np.exp(2j * np.pi * new)
92 Ssym_ac = (np.angle(tmp_Ssym_ac) / (2 * np.pi)) % 1.0 % 1.0
93 Ssym_ac[np.abs(Ssym_ac) < tol] = 0.0
94 Ssym_ac[np.abs(Ssym_ac - 1.0) < tol] = 0.0
95 return Ssym_ac
98@dataclass
99class AtomsSymmetries:
100 """
101 Dataclass to contain symmetry information from atoms.
103 This is to set up an interface for spglib/GPAW or whatever
104 source of symmetry operations.
105 """
107 rotation_scc: np.ndarray
108 atommap_sa: np.ndarray
109 translation_sc: np.ndarray
110 symmorphic: bool
111 symprec: float
113 @classmethod
114 def from_GPAW(cls, atoms, log=None, *, tolerance, symmorphic):
115 from gpaw.new.symmetry import create_symmetries_object
117 gpaw_symmetries = create_symmetries_object(
118 atoms, tolerance=tolerance, symmorphic=symmorphic
119 )
120 log(gpaw_symmetries)
121 sym = AtomsSymmetries(
122 gpaw_symmetries.rotation_scc,
123 gpaw_symmetries.atommap_sa,
124 gpaw_symmetries.translation_sc,
125 symmorphic,
126 tolerance,
127 )
128 return sym
131@dataclass
132class SymmeryAdaptedCellCoordinates:
133 """Class for defining symmetry adapted cell coordinates
135 Note: This is not symmetry adapted cell, it just provides the set of
136 generalized coordinates for the symmetry adapted cell.
137 To get the cell, call C_cv = get_cell(cell_z).
139 sacc = SymmeryAdaptedCellCoordinates(...)
140 sacc.get_cell(cell_z), where cell_z is 1D array of the cell coordinates.
142 Thus, the M_cc and C_cv here is just the origin of the coordinate system.
143 """
145 # Symmetrized cell
146 M_cc: np.ndarray # Rename to M0_cc
147 C_cv: np.ndarray # Rename to C0_cv
148 dM_zcc: np.ndarray
149 dM_zvv: np.ndarray
150 rot_vv: np.ndarray
152 def get_cell(self, cell_z):
153 """
154 Construct the real-space unit cell from symmetry-adapted coordinates.
156 Given generalized cell coordinates `cell_z`, this method reconstructs
157 the metric tensor (see get_M_cc) and then computes a corresponding
158 cell matrix C_cv via a Cholesky factorization of M_cc,
159 followed by a fixed rotation `rot_vv`:
161 C_cv = chol(M_cc) @ rot_vv.T
163 Parameters
164 ----------
165 cell_z : ndarray, shape (nz,)
166 Symmetry-adapted cell coordinates.
168 Returns
169 -------
170 cell : ase.geometry.Cell
171 The reconstructed unit cell.
172 """
174 M_cc = self.get_M_cc(cell_z)
175 try:
176 C_cv = np.linalg.cholesky(M_cc) @ self.rot_vv.T
177 except np.linalg.LinAlgError:
178 raise RuntimeError('Failed to create cell from metric', M_cc)
180 return Cell(C_cv)
182 def get_M_cc(self, cell_z):
183 """
184 Reconstruct the metric tensor from symmetry-adapted coordinates.
186 Computes the metric tensor as a linear expansion around a reference
187 metric M0_cc in the symmetry-allowed tangent directions:
189 M_cc = M0_cc + sum_z cell_z[z] * dM_zcc[z]
191 Parameters
192 ----------
193 cell_z : ndarray, shape (nz,)
194 Symmetry-adapted cell coordinates.
196 Returns
197 -------
198 M_cc : ndarray, shape (3, 3)
199 Symmetrized metric tensor corresponding to `cell_z`.
200 """
201 return self.M_cc + np.einsum('z,zcd->cd', cell_z, self.dM_zcc)
203 @classmethod
204 def build(cls, cell, pbc_c, rotation_scc: np.ndarray, *, log):
205 return cls(*cls.unit_cell_symmetry(cell, rotation_scc, pbc_c, log=log))
207 @classmethod
208 def symmetrize_cell(cls, C_cv, rotation_scc):
209 """Symmetrize the cell
211 Calculates the cell metric, and applies the rotation operations to it.
212 New cell lower diagonal cell is calculated via Cholesky decomposition.
213 By doing polar decomposition to the deformation gradient, the rotation
214 back to the original cell rotation is obtained.
216 Returns osymC_cV, symC_cv, M_cc, rot_vv
218 Where osymC_cv is the symmetrized original like cell
219 symC_cv is the lower diagonal symmetrized cell
220 M_cc is the symmetrized cell metric
221 rot_vv is the rotation matrix between osymC_Cv and symC_cv
222 such that osymC_cv = symC_cv @ rot_vv.T
224 """
225 # Calculate the cell metric
226 M_cc = C_cv @ C_cv.T
227 # Symmetrize the cell metric
228 M_cc = np.einsum(
229 'scd,de,sfe->cf', rotation_scc, M_cc, rotation_scc, optimize=True
230 ) / len(rotation_scc)
232 symC_cv = np.linalg.cholesky(M_cc)
234 # Deformation gradient
235 F_vv = np.linalg.inv(C_cv) @ symC_cv
237 # Sanity check
238 assert np.allclose(C_cv @ F_vv, symC_cv)
240 # Do a polar decomposition to rotate the symmetrized cell back
241 rot_vv, P_vv = scipy.linalg.polar(F_vv)
242 osymC_cv = symC_cv @ rot_vv.T
244 return osymC_cv, symC_cv, M_cc, rot_vv
246 @classmethod
247 def unit_cell_symmetry(
248 cls, C_cv, rotation_scc, pbc_c, units='Å^2', log=None
249 ):
250 pretty(
251 C_cv @ C_cv.T, "Cell metric (M_cc' = C_cv C_c'v)", units, log=log
252 )
253 osymC_cv, symC_cv, M_cc, rot_vv = cls.symmetrize_cell(
254 C_cv, rotation_scc
255 )
256 pretty(
257 M_cc, "Symmetrized cell metric (M_cc' = C_cv C_c'v)", units, log=log
258 )
260 # Now we can construct exact Cartesian rotation matrices
261 iosymC_cv = np.linalg.inv(osymC_cv)
262 U_svv = [osymC_cv.T @ U_cc.T @ iosymC_cv.T for U_cc in rotation_scc]
263 U_svv = np.array(U_svv)
265 # Build unit vector in symmetric matrix space
266 def e(i, j):
267 eps_ij = np.zeros((3, 3))
268 eps_ij[i, j] = 1.0
269 eps_ij[j, i] = 1.0
270 return eps_ij
272 eps_ijk = np.zeros((3, 3, 6))
273 k = 0
274 for i in range(3):
275 for j in range(i, 3):
276 if i == j:
277 s = 1.0
278 else:
279 s = 2 ** (-0.5)
280 eps_ijk[i, j, k] = s
281 eps_ijk[j, i, k] = s
282 k += 1
284 A_blocks = []
285 for U_vv in U_svv:
286 rows = []
287 for k in range(6):
288 row = U_vv @ eps_ijk[:, :, k] @ U_vv.T - eps_ijk[:, :, k]
289 rows.append(row.reshape((9,)))
290 A_blocks.append(np.vstack(rows))
291 for c in range(3):
292 if not pbc_c[c]:
293 A_blocks.append(e(c, c).reshape((9,)))
294 A = np.vstack(A_blocks)
295 A = A @ eps_ijk.reshape((9, 6))
296 # Compute null space via SVD
297 U, S, Vh = np.linalg.svd(A)
298 tol = 1e-6
299 null_mask = S < tol
300 nullspace = Vh[null_mask]
301 dM_zcc = []
302 dM_zvv = []
303 for B in nullspace:
304 B = B @ eps_ijk.reshape((9, 6)).T
305 dM_vv = B.reshape((3, 3))
306 dof = osymC_cv @ dM_vv @ osymC_cv.T
307 dM_zcc.append(dof)
308 dM_zvv.append(rot_vv @ dM_vv @ rot_vv.T)
309 dM_zcc = np.array(dM_zcc).reshape((-1, 3, 3))
310 dM_zvv = np.array(dM_zvv).reshape((-1, 3, 3))
312 # Do a QR decomposition, try to get more zeros to coordinates
313 basis = np.array(dM_zcc).reshape((-1, 9))
314 Q, R = np.linalg.qr(basis)
315 dM_zcc = (Q.T @ basis).reshape((-1, 3, 3))
317 # Normalize tangent space vectors
318 Cinv = np.linalg.inv(C_cv)
319 for z in range(len(dM_zcc)):
320 dC = chol_derivative(M_cc, dM_zcc[z]) @ rot_vv.T
321 eps = 0.5 * (Cinv @ dC + dC.T @ Cinv.T)
323 dM_zcc[z] /= np.sum(np.abs(eps)) * np.linalg.det(C_cv)
324 dM_zcc[z] *= 40
325 # Define the direction of the tangent vector such that
326 # it increases the volume. Sign cannot be used because of shear
327 if np.trace(np.linalg.inv(C_cv) @ dC) < 0:
328 dM_zcc[z] *= -1
330 pretty_dofs(dM_zcc, M_cc, rot_vv, osymC_cv, log=log)
332 # TODO: Move U_svv
333 return M_cc, osymC_cv, dM_zcc, dM_zvv, rot_vv
336@dataclass
337class SymmetryAdaptedScaledCoordinates:
338 dof_zac: np.ndarray
339 s0_ac: np.ndarray
341 def get_scaled_coordinates(self, atoms_z: np.ndarray):
342 return self.s0_ac + np.einsum('zac,z->ac', self.dof_zac, atoms_z)
344 @classmethod
345 def build(
346 cls,
347 s_ac,
348 rotation_scc,
349 translation_sc,
350 atommap_sa,
351 symprec,
352 C_cv,
353 *,
354 log,
355 ):
356 ns, na = atommap_sa.shape
357 B_ascac = np.zeros((na, ns, 3, na, 3), int)
358 for s, U_cc in enumerate(rotation_scc):
359 for a in range(na):
360 a2 = atommap_sa[s, a]
361 B_ascac[a, s, :, a] = U_cc.T
362 B_ascac[a, s, :, a2] -= np.eye(3, dtype=int)
363 B_EA = B_ascac.reshape((na * ns * 3, na * 3))
364 # Extra translational gauge degrees of freedom
365 B_A = np.zeros((na * 3, 3))
366 for a in range(na):
367 B_A[(a * 3) : (a * 3 + 3), :] = np.eye(3)
368 B_EA = np.vstack([B_EA, B_A.T])
369 # import sympy as sp
371 # nullspace = np.array(sp.Matrix(np.array(B_EA,
372 # dtype=int)).nullspace(), dtype=float)
374 # Make sure the old svd code reproduces the same result
375 U, S, Vh = np.linalg.svd(B_EA, False)
376 tol = 1e-6
377 null_mask = S < tol
378 nullspace = Vh[null_mask]
380 # def same_rowspace(N, M, tol=1e-10):
381 # A = np.vstack([N, M])
382 # rA = np.linalg.matrix_rank(A, tol)
383 # rN = np.linalg.matrix_rank(N, tol)
384 # rM = np.linalg.matrix_rank(M, tol)
385 # return rA == rN == rM
386 # assert same_rowspace(nullspace, nullspace2)
388 # Just make the printing prettyer for now
389 nullspace = np.where(np.abs(nullspace) < 1e-10, 0, nullspace)
391 s0_ac = symmetrize_atoms(
392 s_ac,
393 rotation_scc,
394 translation_sc,
395 atommap_sa,
396 )
398 log(f'Atomic degrees of freedom: {len(nullspace)}')
400 dof_zac = nullspace.reshape((-1, na, 3))
402 if len(dof_zac):
403 dof_zav = np.einsum('zac,cv->zav', dof_zac, C_cv)
404 # Normalize such that the distance in Cartesian real space
405 # is reflected on the generalized coordinate
406 dof_zac /= np.max(np.linalg.norm(dof_zav, axis=2), axis=1)[
407 :, None, None
408 ]
410 sasc = SymmetryAdaptedScaledCoordinates(dof_zac, s0_ac)
412 return sasc
415class SymmetryAdaptedAtoms(Optimizable):
416 """Implementation of symmetry adapted atoms
418 Symmetry adapted atoms WILL symmetrize the actual_atoms given to init.
420 SymmetryAdaptedAtoms does not behave like Atoms object, but will expose the
421 __ase_optimizable__ protocol, so it can be optimized with ASE.
422 """
424 def __init__(
425 self, actual_atoms: Atoms, symmetries: AtomsSymmetries, log=print
426 ):
427 self.actual_atoms = actual_atoms
428 self.symmetries = symmetries
429 self.symmetry_force_violation = np.inf
430 self.fmax = 0.01
432 pretty_subheader('Symmetry adapted cell coordinates', log)
433 self.cell_coordinates: SymmeryAdaptedCellCoordinates = (
434 SymmeryAdaptedCellCoordinates.build(
435 self.actual_atoms.cell,
436 self.actual_atoms.pbc,
437 self.symmetries.rotation_scc,
438 log=log,
439 )
440 )
442 pretty_subheader('Symmetry adapted atomic coordinates', log)
443 self.atom_coordinates = SymmetryAdaptedScaledCoordinates.build(
444 self.actual_atoms.get_scaled_positions(),
445 self.symmetries.rotation_scc,
446 self.symmetries.translation_sc,
447 self.symmetries.atommap_sa,
448 self.symmetries.symprec,
449 self.cell_coordinates.C_cv,
450 log=log,
451 )
452 pretty_atomic_dofs(actual_atoms, self.atom_coordinates.dof_zac, log=log)
454 assert isinstance(
455 self.atom_coordinates, SymmetryAdaptedScaledCoordinates
456 )
457 # s_ac = dof_zac s_z -> ds_ac/d_sz = dof_zac
458 # dR_av / dsz = dR_av / d_sac ds_ac / ds_z
459 # R_av = s_ac C_cv
460 #
461 # self.actual_atoms.set_cell(self.cell_coordinates.C_cv,
462 # scale_atoms=True)
463 #
464 # self.actual_atoms.wrap()
465 # self.actual_atoms.set_scaled_positions(self.S_ac)
466 if 1:
467 log('Skipping sanity checks for now')
468 else:
469 pass
470 # new_positions = atoms.get_positions()
471 # dR_av = new_positions - old_positions
472 # s_ac = np.linalg.solve(self.C_cv, dR_av.T)
473 # assert (
474 # np.max(np.abs(new_positions.flatten() -
475 # old_positions.flatten()))
476 # < symprec
477 # )
479 self._ndofs_cell = len(self.cell_coordinates.dM_zcc)
480 self._ndofs_atoms = len(self.atom_coordinates.dof_zac)
481 self._ndofs = self._ndofs_cell + self._ndofs_atoms
483 self.value_z = np.zeros((self._ndofs))
484 # !!! This actually symmetrizes actual atoms
485 self.set_x(self.value_z)
486 self.actual_atoms.wrap()
488 @classmethod
489 def from_atoms(cls, atoms, log=print, *, symprec, symmorphic):
490 symmetries = AtomsSymmetries.from_GPAW(
491 atoms,
492 tolerance=symprec,
493 symmorphic=symmorphic,
494 log=log,
495 )
496 return cls(atoms, symmetries, log=log)
498 @property
499 def stress_conv(self):
500 S_vv = self.actual_atoms.get_stress(voigt=False)
501 C_cv = self.actual_atoms.cell
502 S_cc = C_cv @ S_vv @ np.linalg.inv(C_cv)
503 for c, periodic in enumerate(self.actual_atoms.pbc):
504 if periodic:
505 continue
506 S_cc[c, :] = 0.0
507 S_cc[:, c] = 0.0
508 S_vv = np.linalg.inv(C_cv) @ S_cc @ C_cv
509 return np.max(np.max(np.linalg.norm(S_vv)))
511 # Properties for internal degrees of freedom
512 @property
513 def cell_z(self):
514 return self.get_x()[: self._ndofs_cell]
516 # From here on out, these are the __ase_optimizable__ interface
517 def ndofs(self):
518 return self._ndofs
520 def get_x(self):
521 return self.value_z.copy()
523 def set_x(self, x):
524 self.value_z[:] = x
525 self.actual_atoms.set_cell(self.cell_coordinates.get_cell(self.cell_z))
527 self.actual_atoms.set_scaled_positions(
528 self.atom_coordinates.get_scaled_coordinates(self.atoms_z)
529 )
531 def get_gradient(self):
532 grad_z = np.zeros(self._ndofs_cell)
533 S_vv = self.actual_atoms.get_stress(voigt=False)
534 C_cv = self.cell_coordinates.get_cell(self.cell_z)
535 V = np.linalg.det(C_cv)
536 Cinv = np.linalg.inv(C_cv)
538 M_cc = self.cell_coordinates.get_M_cc(self.cell_z)
540 # TODO: Move to SymmetryAdaptedCellCoordinates
541 # dE/deps_vv deps_vv/dC_cv dC_cv/dz
543 ncellz = len(self.cell_coordinates.dM_zcc)
544 for z in range(ncellz):
545 dC_cv = (
546 chol_derivative(M_cc, self.cell_coordinates.dM_zcc[z])
547 @ self.cell_coordinates.rot_vv.T
548 )
549 grad_z[z] = V * np.sum(S_vv * (Cinv @ dC_cv + dC_cv.T @ Cinv.T) / 2)
551 F_av = self.actual_atoms.get_forces()
552 # dE/ds_z = dE/dR_av dR_av/ds_ac ds_ac/ds_z
553 # R_av = ds_ac C_cv
554 # ds_ac = self.dof_zac S_z
555 atoms_grad_z = -np.einsum(
556 'av,cv,zac->z', F_av, C_cv, self.atom_coordinates.dof_zac
557 )
559 natomz = len(self.atom_coordinates.dof_zac)
560 # For sanity check, we want to project the atomic gradient back
561 # minimizing the Cartesian metrix.
562 if natomz > 0:
563 dof_zX = np.einsum(
564 'cv,zac->zav', C_cv, self.atom_coordinates.dof_zac
565 ).reshape((natomz, -1))
566 back_Fav = -(
567 dof_zX.T @ np.linalg.inv(dof_zX @ dof_zX.T) @ atoms_grad_z
568 ).reshape(F_av.shape)
569 else:
570 # Even if there is degrees of freedom, it is possible to
571 # get symmetry violation
572 back_Fav = np.zeros_like(F_av)
574 dF_av = F_av - back_Fav
575 dF = np.max(np.linalg.norm(dF_av, axis=1))
576 self.symmetry_force_violation = dF
577 self.back_Fav = back_Fav
579 if dF > self.fmax / 20:
580 # Should probably be logged somehow instead of being a warning
581 # as such. This may happen if the code's forces are noisy.
582 import warnings
584 warning_chunks = [
585 'Warning!!! Back projection of symmetry adapted'
586 f' forces to Cartesian space failed by {dF:7.13f}\n'
587 'atom Obtained force Back projected force'
588 ]
590 for a, (F_v, F2_v) in enumerate(zip(F_av, back_Fav)):
591 warning_chunks.append(
592 f'{a:5d} {F_v[0]:7.4f} {F_v[1]:7.4f} {F_v[2]:7.4f}'
593 f' {F2_v[0]:7.4f} {F2_v[1]:7.4f} {F2_v[2]:7.4f}'
594 )
596 warnings.warn('\n'.join(warning_chunks))
598 return np.hstack([grad_z, atoms_grad_z])
600 def gradient_norm(self, grad_z):
601 # Go actually to cell metric
602 return np.max(np.abs(grad_z))
604 def get_value(self):
605 return self.actual_atoms.get_potential_energy()
607 def iterimages(self):
608 return [self.actual_atoms]
610 def converged(self, gradient, fmax):
611 # Convergence needs to be from the back projected forces.
612 # The symmetry violating forces will never converge.
613 Fconv = np.max(np.linalg.norm(self.back_Fav, axis=1))
614 return Fconv < self.fmax and self.stress_conv < self.smax
616 @property
617 def atoms_z(self):
618 return self.get_x()[self._ndofs_cell :]
621class Relax:
622 """General utility class to log and perform symmetry adapted relax"""
624 def __init__(
625 self,
626 symmorphic=False,
627 logfile=None,
628 teelog=True,
629 *,
630 atoms: Atoms,
631 calc,
632 optimizer_factory,
633 symprec,
634 comm,
635 ):
636 self.comm = comm
637 self.logfile = logfile
638 self.logf = None
639 if self.logfile:
640 if self.comm.rank == 0:
641 self.logf = open(self.logfile, 'w')
642 self.teelog = teelog
644 if atoms.calc is not None:
645 raise ValueError('Do not attach a calculator to Atoms yet.')
647 self.symprec = symprec
649 pretty_header('Symmetry adapted Cell and Atomic Relaxation', self.log)
650 pretty_subheader('Original atoms', self.log)
651 self.original_atoms = atoms.copy()
652 pprint_atoms(self.original_atoms, self.log)
654 self.atoms = atoms
655 self.symmetry_adapted_atoms = SymmetryAdaptedAtoms.from_atoms(
656 self.atoms, log=self.log, symmorphic=False, symprec=symprec
657 )
659 # Now, with cell and atoms symmetrized,
660 # it is safe to assign the calculator
661 # TODO: Implement Setter or something
662 self.calc = calc
663 self.symmetry_adapted_atoms.actual_atoms.calc = calc()
665 pretty_subheader('Symmetrized atoms', self.log)
666 pprint_atoms(self.symmetry_adapted_atoms.actual_atoms, self.log)
668 self.optimizer_factory = optimizer_factory
669 self.optimizer = self.optimizer_factory(self.symmetry_adapted_atoms)
671 def log(self, *args, **kwargs):
672 if self.comm.rank == 0:
673 if self.logf:
674 print(*args, **kwargs, flush=True, file=self.logf)
675 if self.teelog:
676 print(*args, **kwargs)
678 def run(self, *, fmax=0.01, smax=0.0001, steps=20):
679 # Why would symmetry adapted atoms care about smax and fmax
680 # But it needs to be ase optimizable, so it needs to do that
681 self.symmetry_adapted_atoms.smax = smax
682 self.symmetry_adapted_atoms.fmax = fmax
684 self.smax = smax
685 self.fmax = fmax
686 self.maxiter = steps
687 i = 0
688 dtitles = ' '.join(
689 [
690 f'q{i:02d}'
691 for i in range(len(self.symmetry_adapted_atoms.atoms_z))
692 ]
693 )
694 self.log(
695 f'iter time E maxF maxS maxG a1 a2'
696 f' a3 L1 L2 L3 {dtitles} log_10 viol.'
697 )
698 for _ in self.optimizer.irun(fmax=fmax):
699 import time
701 T = time.localtime()
702 tstr = '%02d:%02d:%02d' % (T[3], T[4], T[5])
703 E = self.symmetry_adapted_atoms.actual_atoms.get_potential_energy()
704 F = self.symmetry_adapted_atoms.back_Fav # get_forces()
705 g = self.symmetry_adapted_atoms.get_gradient()
706 Fmax = np.max(np.linalg.norm(F, axis=1))
707 sFmax = f'{Fmax:7.3f}'
708 if Fmax < self.fmax:
709 sFmax = green(sFmax)
711 Smax = self.symmetry_adapted_atoms.stress_conv
712 sSmax = f'{Smax:7.4f}'
713 if Smax < self.smax:
714 sSmax = green(sSmax)
716 gmax = np.max(np.abs(g))
718 cell = self.symmetry_adapted_atoms.actual_atoms.cell
719 a = cell.angles()
720 l = cell.lengths()
721 cell = f'{a[0]:5.1f} {a[1]:5.1f} {a[2]:5.1f} '
722 cell += f'{l[0]:7.3f} {l[1]:7.3f} {l[2]:7.3f}'
724 dofs = ''
725 for Z in self.symmetry_adapted_atoms.atoms_z:
726 dofs += f' {Z:6.3f}'
727 syviol = np.log10(
728 self.symmetry_adapted_atoms.symmetry_force_violation
729 )
730 symviol = f'{syviol:4.1f}'
731 if syviol < self.fmax:
732 symviol = green(symviol)
733 self.log(
734 f'{i:5d} {tstr} {E:9.5f} {sFmax} {sSmax} {gmax:7.3f}'
735 f' {cell}{dofs} {symviol}'
736 )
737 i += 1
738 if i > self.maxiter or i > 40:
739 self.log(f'Not converged in {self.maxiter} or 40 steps.')
740 return False
742 return True
744 def visualize_modes(self):
745 from ase.io.trajectory import Trajectory
747 traj = Trajectory('modes.traj', 'w')
748 for z in range(self.ndofs()):
749 x = np.zeros((self.ndofs(),))
750 for i in np.arange(0, 6 * np.pi, 0.1):
751 x[z] = np.sin(i) * 0.004
752 self.set_x(x)
753 traj.write(self.atoms.copy())