Coverage for ase / optimize / cellawarebfgs.py: 100.00%

1# fmt: off 

2 

3import time 

4from typing import IO 

5 

6import numpy as np 

7 

8from ase import Atoms 

9from ase.geometry import cell_to_cellpar 

10from ase.optimize import BFGS 

11from ase.optimize.optimize import Dynamics 

12from ase.units import GPa 

13 

14 

15def calculate_isotropic_elasticity_tensor(bulk_modulus, poisson_ratio, 

16 suppress_rotation=0): 

17 """ 

18 Parameters 

19 ---------- 

20 bulk_modulus Bulk Modulus of the isotropic system used to set up the 

21 Hessian (in ASE units (eV/Å^3)). 

22 

23 poisson_ratio Poisson ratio of the isotropic system used to set up the 

24 initial Hessian (unitless, between -1 and 0.5). 

25 

26 suppress_rotation The rank-2 matrix C_ijkl.reshape((9,9)) has by 

27 default 6 non-zero eigenvalues, because energy is 

28 invariant to orthonormal rotations of the cell 

29 vector. This serves as a bad initial Hessian due to 3 

30 zero eigenvalues. Suppress rotation sets a value for 

31 those zero eigenvalues. 

32 

33 Returns C_ijkl 

34 """ 

35 

36 # https://scienceworld.wolfram.com/physics/LameConstants.html 

37 _lambda = 3 * bulk_modulus * poisson_ratio / (1 + 1 * poisson_ratio) 

38 _mu = _lambda * (1 - 2 * poisson_ratio) / (2 * poisson_ratio) 

39 

40 # https://en.wikipedia.org/wiki/Elasticity_tensor 

41 g_ij = np.eye(3) 

42 

43 # Construct 4th rank Elasticity tensor for isotropic systems 

44 C_ijkl = _lambda * np.einsum('ij,kl->ijkl', g_ij, g_ij) 

45 C_ijkl += _mu * (np.einsum('ik,jl->ijkl', g_ij, g_ij) + 

46 np.einsum('il,kj->ijkl', g_ij, g_ij)) 

47 

48 # Supplement the tensor with suppression of pure rotations that are right 

49 # now 0 eigenvalues. 

50 # Loop over all basis vectors of skew symmetric real matrix 

51 for i, j in ((0, 1), (0, 2), (1, 2)): 

52 Q = np.zeros((3, 3)) 

53 Q[i, j], Q[j, i] = 1, -1 

54 C_ijkl += (np.einsum('ij,kl->ijkl', Q, Q) 

55 * suppress_rotation / 2) 

56 

57 return C_ijkl 

58 

59 

60class CellAwareBFGS(BFGS): 

61 def __init__( 

62 self, 

63 atoms: Atoms, 

64 restart: str | None = None, 

65 logfile: IO | str = '-', 

66 trajectory: str | None = None, 

67 append_trajectory: bool = False, 

68 maxstep: float | None = None, 

69 bulk_modulus: float | None = 145 * GPa, 

70 poisson_ratio: float | None = 0.3, 

71 alpha: float | None = None, 

72 long_output: bool | None = False, 

73 **kwargs, 

74 ): 

75 self.bulk_modulus = bulk_modulus 

76 self.poisson_ratio = poisson_ratio 

77 self.long_output = long_output 

78 super().__init__( 

79 atoms=atoms, restart=restart, logfile=logfile, 

80 trajectory=trajectory, maxstep=maxstep, 

81 alpha=alpha, append_trajectory=append_trajectory, 

82 **kwargs) 

83 assert not isinstance(atoms, Atoms) 

84 if hasattr(atoms, 'exp_cell_factor'): 

85 assert atoms.exp_cell_factor == 1.0 

86 

87 def initialize(self): 

88 super().initialize() 

89 C_ijkl = calculate_isotropic_elasticity_tensor( 

90 self.bulk_modulus, 

91 self.poisson_ratio, 

92 suppress_rotation=self.alpha) 

93 cell_H = self.H0[-9:, -9:] 

94 ind = np.where(self.atoms.mask.ravel() != 0)[0] 

95 cell_H[np.ix_(ind, ind)] = C_ijkl.reshape((9, 9))[ 

96 np.ix_(ind, ind)] * self.atoms.atoms.cell.volume 

97 

98 def gradient_converged(self, gradient): 

99 # XXX currently ignoring gradient 

100 forces = self.atoms.atoms.get_forces() 

101 stress = self.atoms.atoms.get_stress(voigt=False) * self.atoms.mask 

102 return np.max(np.sum(forces**2, axis=1))**0.5 < self.fmax and \ 

103 np.max(np.abs(stress)) < self.smax 

104 

105 def run(self, fmax=0.05, smax=0.005, steps=None): 

106 """ call Dynamics.run and keep track of fmax""" 

107 self.fmax = fmax 

108 self.smax = smax 

109 if steps is not None: 

110 return Dynamics.run(self, steps=steps) 

111 return Dynamics.run(self) 

112 

113 def log(self, gradient): 

114 # XXX ignoring gradient 

115 forces = self.atoms.atoms.get_forces() 

116 fmax = (forces ** 2).sum(axis=1).max() ** 0.5 

117 e = self.optimizable.get_value() 

118 T = time.localtime() 

119 smax = abs(self.atoms.atoms.get_stress(voigt=False) * 

120 self.atoms.mask).max() 

121 volume = self.atoms.atoms.cell.volume 

122 if self.logfile is not None: 

123 name = self.__class__.__name__ 

124 if self.nsteps == 0: 

125 args = (" " * len(name), 

126 "Step", "Time", "Energy", "fmax", "smax", "volume") 

127 msg = "\n%s %4s %8s %15s %15s %15s %15s" % args 

128 if self.long_output: 

129 msg += ("%8s %8s %8s %8s %8s %8s" % 

130 ('A', 'B', 'C', 'α', 'β', 'γ')) 

131 msg += '\n' 

132 self.logfile.write(msg) 

133 

134 ast = '' 

135 args = (name, self.nsteps, T[3], T[4], T[5], e, ast, fmax, smax, 

136 volume) 

137 msg = ("%s: %3d %02d:%02d:%02d %15.6f%1s %15.6f %15.6f %15.6f" % 

138 args) 

139 if self.long_output: 

140 msg += ("%8.3f %8.3f %8.3f %8.3f %8.3f %8.3f" % 

141 tuple(cell_to_cellpar(self.atoms.atoms.cell))) 

142 msg += '\n' 

143 self.logfile.write(msg)