Coverage for /builds/kinetik161/ase/ase/optimize/mdmin.py: 96.88%

1from typing import IO, Optional, Union 

2 

3import numpy as np 

4 

5from ase import Atoms 

6from ase.optimize.optimize import Optimizer 

7 

8 

9class MDMin(Optimizer): 

10 # default parameters 

11 defaults = {**Optimizer.defaults, 'dt': 0.2} 

12 

13 def __init__( 

14 self, 

15 atoms: Atoms, 

16 restart: Optional[str] = None, 

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

18 trajectory: Optional[str] = None, 

19 dt: Optional[float] = None, 

20 maxstep: Optional[float] = None, 

21 master: Optional[bool] = None, 

22 ): 

23 """Parameters: 

24 

25 atoms: Atoms object 

26 The Atoms object to relax. 

27 

28 restart: string 

29 Pickle file used to store hessian matrix. If set, file with 

30 such a name will be searched and hessian matrix stored will 

31 be used, if the file exists. 

32 

33 trajectory: string 

34 Pickle file used to store trajectory of atomic movement. 

35 

36 logfile: string 

37 Text file used to write summary information. 

38 

39 dt: float 

40 Time step for integrating the equation of motion. 

41 

42 maxstep: float 

43 Spatial step limit in Angstrom. This allows larger values of dt 

44 while being more robust to instabilities in the optimization. 

45 

46 master: boolean 

47 Defaults to None, which causes only rank 0 to save files. If 

48 set to true, this rank will save files. 

49 """ 

50 Optimizer.__init__(self, atoms, restart, logfile, trajectory, master) 

51 

52 self.dt = dt or self.defaults['dt'] 

53 self.maxstep = maxstep or self.defaults['maxstep'] 

54 

55 def initialize(self): 

56 self.v = None 

57 

58 def read(self): 

59 self.v, self.dt = self.load() 

60 

61 def step(self, forces=None): 

62 optimizable = self.optimizable 

63 

64 if forces is None: 

65 forces = optimizable.get_forces() 

66 

67 if self.v is None: 

68 self.v = np.zeros((len(optimizable), 3)) 

69 else: 

70 self.v += 0.5 * self.dt * forces 

71 # Correct velocities: 

72 vf = np.vdot(self.v, forces) 

73 if vf < 0.0: 

74 self.v[:] = 0.0 

75 else: 

76 self.v[:] = forces * vf / np.vdot(forces, forces) 

77 

78 self.v += 0.5 * self.dt * forces 

79 pos = optimizable.get_positions() 

80 dpos = self.dt * self.v 

81 

82 # For any dpos magnitude larger than maxstep, scaling 

83 # is <1. We add a small float to prevent overflows/zero-div errors. 

84 # All displacement vectors (rows) of dpos which have a norm larger 

85 # than self.maxstep are scaled to it. 

86 scaling = self.maxstep / (1e-6 + np.max(np.linalg.norm(dpos, axis=1))) 

87 dpos *= np.clip(scaling, 0.0, 1.0) 

88 optimizable.set_positions(pos + dpos) 

89 self.dump((self.v, self.dt))