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

1# fmt: off 

2 

3import sys 

4from typing import Any, Dict 

5 

6import numpy as np 

7 

8 

9class CLICommand: 

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

11 

12 Four types of calculations can be done: 

13 

14 * single point 

15 * atomic relaxations 

16 * unit cell + atomic relaxations 

17 * equation-of-state 

18 

19 Examples of the four types of calculations: 

20 

21 ase run emt h2o.xyz 

22 ase run emt h2o.xyz -f 0.01 

23 ase run emt cu.traj -s 0.01 

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

25 """ 

26 

27 @staticmethod 

28 def add_arguments(parser): 

29 from ase.calculators.names import names 

30 parser.add_argument('calculator', 

31 help='Name of calculator to use. ' 

32 'Must be one of: {}.' 

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

34 CLICommand.add_more_arguments(parser) 

35 

36 @staticmethod 

37 def add_more_arguments(parser): 

38 add = parser.add_argument 

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

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

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

42 metavar='key=value,...', 

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

44 'calculator specific parameters.') 

45 add('-t', '--tag', 

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

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

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

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

50 'atomic magnetic moments.') 

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

52 help='Relax internal coordinates.') 

53 add('--constrain-tags', 

54 metavar='T1,T2,...', 

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

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

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

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

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

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

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

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

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

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

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

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

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

68 

69 @staticmethod 

70 def run(args): 

71 runner = Runner() 

72 runner.parse(args) 

73 runner.run() 

74 

75 

76class Runner: 

77 def __init__(self): 

78 self.args = None 

79 self.calculator_name = None 

80 

81 def parse(self, args): 

82 self.calculator_name = args.calculator 

83 self.args = args 

84 

85 def run(self): 

86 args = self.args 

87 

88 atoms = self.build(args.name) 

89 if args.modify: 

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

91 

92 if args.name == '-': 

93 args.name = 'stdin' 

94 

95 self.set_calculator(atoms, args.name) 

96 

97 self.calculate(atoms, args.name) 

98 

99 def calculate(self, atoms, name): 

100 from ase.io import write 

101 

102 args = self.args 

103 

104 if args.maximum_force or args.maximum_stress: 

105 self.optimize(atoms, name) 

106 if args.equation_of_state: 

107 self.eos(atoms, name) 

108 self.calculate_once(atoms) 

109 

110 if args.after: 

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

112 

113 if args.output: 

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

115 

116 def build(self, name): 

117 import ase.db as db 

118 from ase.io import read 

119 

120 if name == '-': 

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

122 return con.get_atoms(add_additional_information=True) 

123 else: 

124 atoms = read(name) 

125 if isinstance(atoms, list): 

126 assert len(atoms) == 1 

127 atoms = atoms[0] 

128 return atoms 

129 

130 def set_calculator(self, atoms, name): 

131 from ase.calculators.calculator import get_calculator_class 

132 

133 cls = get_calculator_class(self.calculator_name) 

134 parameters = str2dict(self.args.parameters) 

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

136 atoms.calc = cls(**parameters) 

137 else: 

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

139 

140 def calculate_once(self, atoms): 

141 from ase.calculators.calculator import PropertyNotImplementedError 

142 

143 args = self.args 

144 

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

146 _property, method = {'e': ('energy', 'get_potential_energy'), 

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

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

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

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

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

152 try: 

153 getattr(atoms, method)() 

154 except PropertyNotImplementedError: 

155 pass 

156 

157 def optimize(self, atoms, name): 

158 from ase.constraints import FixAtoms, UnitCellFilter 

159 from ase.io import Trajectory 

160 from ase.optimize import LBFGS 

161 

162 args = self.args 

163 if args.constrain_tags: 

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

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

166 atoms.constraints = FixAtoms(mask=mask) 

167 

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

169 if args.maximum_stress: 

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

171 fmax = args.maximum_stress 

172 else: 

173 optimizer = LBFGS(atoms, logfile=logfile) 

174 fmax = args.maximum_force 

175 

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

177 optimizer.attach(trajectory) 

178 optimizer.run(fmax=fmax) 

179 

180 def eos(self, atoms, name): 

181 from ase.eos import EquationOfState 

182 from ase.io import Trajectory 

183 

184 args = self.args 

185 

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

187 

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

189 N = int(N) 

190 eps = float(eps) / 100 

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

192 v1 = atoms.get_volume() 

193 volumes = strains**3 * v1 

194 energies = [] 

195 cell1 = atoms.cell.copy() 

196 for s in strains: 

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

198 energies.append(atoms.get_potential_energy()) 

199 traj.write(atoms) 

200 traj.close() 

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

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

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

204 from ase.parallel import parprint as p 

205 p('volumes:', volumes) 

206 p('energies:', energies) 

207 p('fitted energy:', e0) 

208 p('fitted volume:', v0) 

209 p('bulk modulus:', B) 

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

211 

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

213 if '.' in name: 

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

215 if self.args.tag is not None: 

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

217 if ext: 

218 name += '.' + ext 

219 return name 

220 

221 

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

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

224 

225 Examples: 

226 

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

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

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

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

231 """ 

232 

233 def myeval(value): 

234 try: 

235 value = eval(value, namespace) 

236 except (NameError, SyntaxError): 

237 pass 

238 return value 

239 

240 dct = {} 

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

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

243 key = strings[i] 

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

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

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

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

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

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

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

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

252 else: 

253 value = myeval(value) 

254 dct[key] = value 

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

256 return dct