Coverage for /builds/alexhroom/ase/ase/optimize/cellawarebfgs.py: 98.59%

1import time 

2from typing import IO, Optional, Union 

3 

4import numpy as np 

5 

6from ase import Atoms 

7from ase.parallel import world 

8from ase.geometry import cell_to_cellpar 

9from ase.optimize import BFGS 

10from ase.optimize.optimize import Dynamics 

11from ase.units import GPa 

12 

13 

14def calculate_isotropic_elasticity_tensor(bulk_modulus, poisson_ratio, 

15 suppress_rotation=0): 

16 """ 

17 Parameters: 

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

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

20 

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

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

23 

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

25 default 6 non-zero eigenvalues, because energy is 

26 invariant to orthonormal rotations of the cell 

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

28 zero eigenvalues. Suppress rotation sets a value for 

29 those zero eigenvalues. 

30 

31 Returns C_ijkl 

32 """ 

33 

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

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

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

37 

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

39 g_ij = np.eye(3) 

40 

41 # Construct 4th rank Elasticity tensor for isotropic systems 

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

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

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

45 

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

47 # now 0 eigenvalues. 

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

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

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

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

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

53 * suppress_rotation / 2) 

54 

55 return C_ijkl 

56 

57 

58class CellAwareBFGS(BFGS): 

59 def __init__( 

60 self, 

61 atoms: Atoms, 

62 restart: Optional[str] = None, 

63 logfile: Union[IO, str] = '-', 

64 trajectory: Optional[str] = None, 

65 append_trajectory: bool = False, 

66 maxstep: Optional[float] = None, 

67 master: Optional[bool] = None, 

68 bulk_modulus: Optional[float] = 145 * GPa, 

69 poisson_ratio: Optional[float] = 0.3, 

70 alpha: Optional[float] = None, 

71 long_output: Optional[bool] = False, 

72 comm=world, 

73 ): 

74 self.bulk_modulus = bulk_modulus 

75 self.poisson_ratio = poisson_ratio 

76 self.long_output = long_output 

77 BFGS.__init__(self, atoms=atoms, restart=restart, logfile=logfile, 

78 trajectory=trajectory, maxstep=maxstep, master=master, 

79 alpha=alpha, append_trajectory=append_trajectory, 

80 comm=comm) 

81 assert not isinstance(atoms, Atoms) 

82 if hasattr(atoms, 'exp_cell_factor'): 

83 assert atoms.exp_cell_factor == 1.0 

84 

85 def initialize(self): 

86 BFGS.initialize(self) 

87 C_ijkl = calculate_isotropic_elasticity_tensor( 

88 self.bulk_modulus, 

89 self.poisson_ratio, 

90 suppress_rotation=self.alpha) 

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

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

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

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

95 

96 def converged(self, forces=None): 

97 if forces is None: 

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

99 stress = self.atoms.atoms.get_stress() 

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

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

102 

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

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

105 self.fmax = fmax 

106 self.smax = smax 

107 if steps is not None: 

108 self.max_steps = steps 

109 return Dynamics.run(self) 

110 

111 def log(self, forces=None): 

112 if forces is None: 

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

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

115 e = self.optimizable.get_potential_energy() 

116 T = time.localtime() 

117 smax = abs(self.atoms.atoms.get_stress()).max() 

118 volume = self.atoms.atoms.cell.volume 

119 if self.logfile is not None: 

120 name = self.__class__.__name__ 

121 if self.nsteps == 0: 

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

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

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

125 if self.long_output: 

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

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

128 msg += '\n' 

129 self.logfile.write(msg) 

130 

131 ast = '' 

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

133 volume) 

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

135 args) 

136 if self.long_output: 

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

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

139 msg += '\n' 

140 self.logfile.write(msg) 

141 

142 self.logfile.flush()