Coverage for ase / optimize / basin.py: 97.65%

1# fmt: off 

2 

3from typing import IO, Type, Union 

4 

5import numpy as np 

6 

7from ase import Atoms, units 

8from ase.io.trajectory import Trajectory 

9from ase.optimize.fire import FIRE 

10from ase.optimize.optimize import Dynamics, Optimizer 

11from ase.parallel import world 

12 

13 

14class BasinHopping(Dynamics): 

15 """Basin hopping algorithm. 

16 

17 After Wales and Doye, J. Phys. Chem. A, vol 101 (1997) 5111-5116 

18 

19 and 

20 

21 David J. Wales and Harold A. Scheraga, Science, Vol. 285, 1368 (1999) 

22 """ 

23 

24 def __init__( 

25 self, 

26 atoms: Atoms, 

27 temperature: float = 100 * units.kB, 

28 optimizer: Type[Optimizer] = FIRE, 

29 fmax: float = 0.1, 

30 dr: float = 0.1, 

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

32 trajectory: str = 'lowest.traj', 

33 optimizer_logfile: str = '-', 

34 local_minima_trajectory: str = 'local_minima.traj', 

35 adjust_cm: bool = True, 

36 ): 

37 """Parameters: 

38 

39 atoms: Atoms object 

40 The Atoms object to operate on. 

41 

42 trajectory: string 

43 Trajectory file used to store optimisation path. 

44 

45 logfile: file object or str 

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

47 Use '-' for stdout. 

48 """ 

49 self.kT = temperature 

50 self.optimizer = optimizer 

51 self.fmax = fmax 

52 self.dr = dr 

53 if adjust_cm: 

54 self.cm = atoms.get_center_of_mass() 

55 else: 

56 self.cm = None 

57 

58 self.optimizer_logfile = optimizer_logfile 

59 self.lm_trajectory = local_minima_trajectory 

60 if isinstance(local_minima_trajectory, str): 

61 self.lm_trajectory = self.closelater( 

62 Trajectory(local_minima_trajectory, 'w', atoms)) 

63 

64 Dynamics.__init__(self, atoms, logfile, trajectory) 

65 self.initialize() 

66 

67 def todict(self): 

68 d = {'type': 'optimization', 

69 'optimizer': self.__class__.__name__, 

70 'local-minima-optimizer': self.optimizer.__name__, 

71 'temperature': self.kT, 

72 'max-force': self.fmax, 

73 'maximal-step-width': self.dr} 

74 return d 

75 

76 def initialize(self): 

77 positions = self.optimizable.get_x().reshape(-1, 3) 

78 self.positions = np.zeros_like(positions) 

79 self.Emin = self.get_energy(positions) or 1.e32 

80 self.rmin = self.optimizable.get_x().reshape(-1, 3) 

81 self.positions = self.optimizable.get_x().reshape(-1, 3) 

82 self.call_observers() 

83 self.log(-1, self.Emin, self.Emin) 

84 

85 def run(self, steps): 

86 """Hop the basins for defined number of steps.""" 

87 

88 ro = self.positions 

89 Eo = self.get_energy(ro) 

90 

91 for step in range(steps): 

92 En = None 

93 while En is None: 

94 rn = self.move(ro) 

95 En = self.get_energy(rn) 

96 

97 if En < self.Emin: 

98 # new minimum found 

99 self.Emin = En 

100 self.rmin = self.optimizable.get_x().reshape(-1, 3) 

101 self.call_observers() 

102 self.log(step, En, self.Emin) 

103 

104 accept = np.exp((Eo - En) / self.kT) > np.random.uniform() 

105 if accept: 

106 ro = rn.copy() 

107 Eo = En 

108 

109 def log(self, step, En, Emin): 

110 if self.logfile is None: 

111 return 

112 name = self.__class__.__name__ 

113 self.logfile.write('%s: step %d, energy %15.6f, emin %15.6f\n' 

114 % (name, step, En, Emin)) 

115 

116 def _atoms(self): 

117 from ase.optimize.optimize import OptimizableAtoms 

118 assert isinstance(self.optimizable, OptimizableAtoms) 

119 # Some parts of the basin code cannot work on Filter objects. 

120 # They evidently need an actual Atoms object - at least until 

121 # someone changes the code so it doesn't need that. 

122 return self.optimizable.atoms 

123 

124 def move(self, ro): 

125 """Move atoms by a random step.""" 

126 atoms = self._atoms() 

127 # displace coordinates 

128 disp = np.random.uniform(-1., 1., (len(atoms), 3)) 

129 rn = ro + self.dr * disp 

130 atoms.set_positions(rn) 

131 if self.cm is not None: 

132 cm = atoms.get_center_of_mass() 

133 atoms.translate(self.cm - cm) 

134 rn = atoms.get_positions() 

135 world.broadcast(rn, 0) 

136 atoms.set_positions(rn) 

137 return atoms.get_positions() 

138 

139 def get_minimum(self): 

140 """Return minimal energy and configuration.""" 

141 atoms = self._atoms().copy() 

142 atoms.set_positions(self.rmin) 

143 return self.Emin, atoms 

144 

145 def get_energy(self, positions): 

146 """Return the energy of the nearest local minimum.""" 

147 if np.any(self.positions != positions): 

148 self.positions = positions 

149 self.optimizable.set_x(positions.ravel()) 

150 

151 with self.optimizer(self.optimizable, 

152 logfile=self.optimizer_logfile) as opt: 

153 opt.run(fmax=self.fmax) 

154 if self.lm_trajectory is not None: 

155 self.lm_trajectory.write(self.optimizable) 

156 

157 self.energy = self.optimizable.get_value() 

158 

159 return self.energy