Coverage for /builds/kinetik161/ase/ase/cli/run.py: 82.61%

161 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-12-10 11:04 +0000

1import sys 

2from typing import Any, Dict 

3 

4import numpy as np 

5 

6 

7class CLICommand: 

8 """Run calculation with one of ASE's calculators. 

9 

10 Four types of calculations can be done: 

11 

12 * single point 

13 * atomic relaxations 

14 * unit cell + atomic relaxations 

15 * equation-of-state 

16 

17 Examples of the four types of calculations: 

18 

19 ase run emt h2o.xyz 

20 ase run emt h2o.xyz -f 0.01 

21 ase run emt cu.traj -s 0.01 

22 ase run emt cu.traj -E 5,2.0 

23 """ 

24 

25 @staticmethod 

26 def add_arguments(parser): 

27 from ase.calculators.names import names 

28 parser.add_argument('calculator', 

29 help='Name of calculator to use. ' 

30 'Must be one of: {}.' 

31 .format(', '.join(names))) 

32 CLICommand.add_more_arguments(parser) 

33 

34 @staticmethod 

35 def add_more_arguments(parser): 

36 add = parser.add_argument 

37 add('name', nargs='?', default='-', 

38 help='Read atomic structure from this file.') 

39 add('-p', '--parameters', default='', 

40 metavar='key=value,...', 

41 help='Comma-separated key=value pairs of ' + 

42 'calculator specific parameters.') 

43 add('-t', '--tag', 

44 help='String tag added to filenames.') 

45 add('--properties', default='efsdMm', 

46 help='Default value is "efsdMm" meaning calculate energy, ' + 

47 'forces, stress, dipole moment, total magnetic moment and ' + 

48 'atomic magnetic moments.') 

49 add('-f', '--maximum-force', type=float, 

50 help='Relax internal coordinates.') 

51 add('--constrain-tags', 

52 metavar='T1,T2,...', 

53 help='Constrain atoms with tags T1, T2, ...') 

54 add('-s', '--maximum-stress', type=float, 

55 help='Relax unit-cell and internal coordinates.') 

56 add('-E', '--equation-of-state', 

57 help='Use "-E 5,2.0" for 5 lattice constants ranging from ' 

58 '-2.0 %% to +2.0 %%.') 

59 add('--eos-type', default='sjeos', help='Selects the type of eos.') 

60 add('-o', '--output', help='Write result to file (append mode).') 

61 add('--modify', metavar='...', 

62 help='Modify atoms with Python statement. ' + 

63 'Example: --modify="atoms.positions[-1,2]+=0.1".') 

64 add('--after', help='Perform operation after calculation. ' + 

65 'Example: --after="atoms.calc.write(...)"') 

66 

67 @staticmethod 

68 def run(args): 

69 runner = Runner() 

70 runner.parse(args) 

71 runner.run() 

72 

73 

74class Runner: 

75 def __init__(self): 

76 self.args = None 

77 self.calculator_name = None 

78 

79 def parse(self, args): 

80 self.calculator_name = args.calculator 

81 self.args = args 

82 

83 def run(self): 

84 args = self.args 

85 

86 atoms = self.build(args.name) 

87 if args.modify: 

88 exec(args.modify, {'atoms': atoms, 'np': np}) 

89 

90 if args.name == '-': 

91 args.name = 'stdin' 

92 

93 self.set_calculator(atoms, args.name) 

94 

95 self.calculate(atoms, args.name) 

96 

97 def calculate(self, atoms, name): 

98 from ase.io import write 

99 

100 args = self.args 

101 

102 if args.maximum_force or args.maximum_stress: 

103 self.optimize(atoms, name) 

104 if args.equation_of_state: 

105 self.eos(atoms, name) 

106 self.calculate_once(atoms) 

107 

108 if args.after: 

109 exec(args.after, {'atoms': atoms}) 

110 

111 if args.output: 

112 write(args.output, atoms, append=True) 

113 

114 def build(self, name): 

115 import ase.db as db 

116 from ase.io import read 

117 

118 if name == '-': 

119 con = db.connect(sys.stdin, 'json') 

120 return con.get_atoms(add_additional_information=True) 

121 else: 

122 atoms = read(name) 

123 if isinstance(atoms, list): 

124 assert len(atoms) == 1 

125 atoms = atoms[0] 

126 return atoms 

127 

128 def set_calculator(self, atoms, name): 

129 from ase.calculators.calculator import get_calculator_class 

130 

131 cls = get_calculator_class(self.calculator_name) 

132 parameters = str2dict(self.args.parameters) 

133 if getattr(cls, 'nolabel', False): 

134 atoms.calc = cls(**parameters) 

135 else: 

136 atoms.calc = cls(label=self.get_filename(name), **parameters) 

137 

138 def calculate_once(self, atoms): 

139 from ase.calculators.calculator import PropertyNotImplementedError 

140 

141 args = self.args 

142 

143 for p in args.properties or 'efsdMm': 

144 property, method = {'e': ('energy', 'get_potential_energy'), 

145 'f': ('forces', 'get_forces'), 

146 's': ('stress', 'get_stress'), 

147 'd': ('dipole', 'get_dipole_moment'), 

148 'M': ('magmom', 'get_magnetic_moment'), 

149 'm': ('magmoms', 'get_magnetic_moments')}[p] 

150 try: 

151 getattr(atoms, method)() 

152 except PropertyNotImplementedError: 

153 pass 

154 

155 def optimize(self, atoms, name): 

156 from ase.constraints import FixAtoms, UnitCellFilter 

157 from ase.io import Trajectory 

158 from ase.optimize import LBFGS 

159 

160 args = self.args 

161 if args.constrain_tags: 

162 tags = [int(t) for t in args.constrain_tags.split(',')] 

163 mask = [t in tags for t in atoms.get_tags()] 

164 atoms.constraints = FixAtoms(mask=mask) 

165 

166 logfile = self.get_filename(name, 'log') 

167 if args.maximum_stress: 

168 optimizer = LBFGS(UnitCellFilter(atoms), logfile=logfile) 

169 fmax = args.maximum_stress 

170 else: 

171 optimizer = LBFGS(atoms, logfile=logfile) 

172 fmax = args.maximum_force 

173 

174 trajectory = Trajectory(self.get_filename(name, 'traj'), 'w', atoms) 

175 optimizer.attach(trajectory) 

176 optimizer.run(fmax=fmax) 

177 

178 def eos(self, atoms, name): 

179 from ase.eos import EquationOfState 

180 from ase.io import Trajectory 

181 

182 args = self.args 

183 

184 traj = Trajectory(self.get_filename(name, 'traj'), 'w', atoms) 

185 

186 N, eps = args.equation_of_state.split(',') 

187 N = int(N) 

188 eps = float(eps) / 100 

189 strains = np.linspace(1 - eps, 1 + eps, N) 

190 v1 = atoms.get_volume() 

191 volumes = strains**3 * v1 

192 energies = [] 

193 cell1 = atoms.cell.copy() 

194 for s in strains: 

195 atoms.set_cell(cell1 * s, scale_atoms=True) 

196 energies.append(atoms.get_potential_energy()) 

197 traj.write(atoms) 

198 traj.close() 

199 eos = EquationOfState(volumes, energies, args.eos_type) 

200 v0, e0, B = eos.fit() 

201 atoms.set_cell(cell1 * (v0 / v1)**(1 / 3), scale_atoms=True) 

202 from ase.parallel import parprint as p 

203 p('volumes:', volumes) 

204 p('energies:', energies) 

205 p('fitted energy:', e0) 

206 p('fitted volume:', v0) 

207 p('bulk modulus:', B) 

208 p('eos type:', args.eos_type) 

209 

210 def get_filename(self, name: str, ext: str = '') -> str: 

211 if '.' in name: 

212 name = name.rsplit('.', 1)[0] 

213 if self.args.tag is not None: 

214 name += '-' + self.args.tag 

215 if ext: 

216 name += '.' + ext 

217 return name 

218 

219 

220def str2dict(s: str, namespace={}, sep: str = '=') -> Dict[str, Any]: 

221 """Convert comma-separated key=value string to dictionary. 

222 

223 Examples: 

224 

225 >>> str2dict('xc=PBE,nbands=200,parallel={band:4}') 

226 {'xc': 'PBE', 'nbands': 200, 'parallel': {'band': 4}} 

227 >>> str2dict('a=1.2,b=True,c=ab,d=1,2,3,e={f:42,g:cd}') 

228 {'a': 1.2, 'c': 'ab', 'b': True, 'e': {'g': 'cd', 'f': 42}, 'd': (1, 2, 3)} 

229 """ 

230 

231 def myeval(value): 

232 try: 

233 value = eval(value, namespace) 

234 except (NameError, SyntaxError): 

235 pass 

236 return value 

237 

238 dct = {} 

239 strings = (s + ',').split(sep) 

240 for i in range(len(strings) - 1): 

241 key = strings[i] 

242 m = strings[i + 1].rfind(',') 

243 value: Any = strings[i + 1][:m] 

244 if value[0] == '{': 

245 assert value[-1] == '}' 

246 value = str2dict(value[1:-1], namespace, ':') 

247 elif value[0] == '(': 

248 assert value[-1] == ')' 

249 value = [myeval(t) for t in value[1:-1].split(',')] 

250 else: 

251 value = myeval(value) 

252 dct[key] = value 

253 strings[i + 1] = strings[i + 1][m + 1:] 

254 return dct