1import sys 

2from typing import Dict, Any 

3import numpy as np 

4 

5 

6class CLICommand: 

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

8 

9 Four types of calculations can be done: 

10 

11 * single point 

12 * atomic relaxations 

13 * unit cell + atomic relaxations 

14 * equation-of-state 

15 

16 Examples of the four types of calculations: 

17 

18 ase run emt h2o.xyz 

19 ase run emt h2o.xyz -f 0.01 

20 ase run emt cu.traj -s 0.01 

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

22 """ 

23 

24 @staticmethod 

25 def add_arguments(parser): 

26 from ase.calculators.names import names 

27 parser.add_argument('calculator', 

28 help='Name of calculator to use. ' 

29 'Must be one of: {}.' 

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

31 CLICommand.add_more_arguments(parser) 

32 

33 @staticmethod 

34 def add_more_arguments(parser): 

35 add = parser.add_argument 

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

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

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

39 metavar='key=value,...', 

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

41 'calculator specific parameters.') 

42 add('-t', '--tag', 

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

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

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

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

47 'atomic magnetic moments.') 

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

49 help='Relax internal coordinates.') 

50 add('--constrain-tags', 

51 metavar='T1,T2,...', 

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

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

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

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

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

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

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

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

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

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

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

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

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

65 

66 @staticmethod 

67 def run(args): 

68 runner = Runner() 

69 runner.parse(args) 

70 runner.run() 

71 

72 

73class Runner: 

74 def __init__(self): 

75 self.args = None 

76 self.calculator_name = None 

77 

78 def parse(self, args): 

79 self.calculator_name = args.calculator 

80 self.args = args 

81 

82 def run(self): 

83 args = self.args 

84 

85 atoms = self.build(args.name) 

86 if args.modify: 

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

88 

89 if args.name == '-': 

90 args.name = 'stdin' 

91 

92 self.set_calculator(atoms, args.name) 

93 

94 self.calculate(atoms, args.name) 

95 

96 def calculate(self, atoms, name): 

97 from ase.io import write 

98 

99 args = self.args 

100 

101 if args.maximum_force or args.maximum_stress: 

102 self.optimize(atoms, name) 

103 if args.equation_of_state: 

104 self.eos(atoms, name) 

105 self.calculate_once(atoms) 

106 

107 if args.after: 

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

109 

110 if args.output: 

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

112 

113 def build(self, name): 

114 import ase.db as db 

115 from ase.io import read 

116 

117 if name == '-': 

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

119 return con.get_atoms(add_additional_information=True) 

120 else: 

121 atoms = read(name) 

122 if isinstance(atoms, list): 

123 assert len(atoms) == 1 

124 atoms = atoms[0] 

125 return atoms 

126 

127 def set_calculator(self, atoms, name): 

128 from ase.calculators.calculator import get_calculator_class 

129 

130 cls = get_calculator_class(self.calculator_name) 

131 parameters = str2dict(self.args.parameters) 

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

133 atoms.calc = cls(**parameters) 

134 else: 

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

136 

137 def calculate_once(self, atoms): 

138 from ase.calculators.calculator import PropertyNotImplementedError 

139 

140 args = self.args 

141 

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

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

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

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

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

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

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

149 try: 

150 getattr(atoms, method)() 

151 except PropertyNotImplementedError: 

152 pass 

153 

154 def optimize(self, atoms, name): 

155 from ase.constraints import FixAtoms, UnitCellFilter 

156 from ase.io import Trajectory 

157 from ase.optimize import LBFGS 

158 

159 args = self.args 

160 if args.constrain_tags: 

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

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

163 atoms.constraints = FixAtoms(mask=mask) 

164 

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

166 if args.maximum_stress: 

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

168 fmax = args.maximum_stress 

169 else: 

170 optimizer = LBFGS(atoms, logfile=logfile) 

171 fmax = args.maximum_force 

172 

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

174 optimizer.attach(trajectory) 

175 optimizer.run(fmax=fmax) 

176 

177 def eos(self, atoms, name): 

178 from ase.eos import EquationOfState 

179 from ase.io import Trajectory 

180 

181 args = self.args 

182 

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

184 

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

186 N = int(N) 

187 eps = float(eps) / 100 

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

189 v1 = atoms.get_volume() 

190 volumes = strains**3 * v1 

191 energies = [] 

192 cell1 = atoms.cell.copy() 

193 for s in strains: 

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

195 energies.append(atoms.get_potential_energy()) 

196 traj.write(atoms) 

197 traj.close() 

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

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

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

201 from ase.parallel import parprint as p 

202 p('volumes:', volumes) 

203 p('energies:', energies) 

204 p('fitted energy:', e0) 

205 p('fitted volume:', v0) 

206 p('bulk modulus:', B) 

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

208 

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

210 if '.' in name: 

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

212 if self.args.tag is not None: 

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

214 if ext: 

215 name += '.' + ext 

216 return name 

217 

218 

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

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

221 

222 Examples: 

223 

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

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

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

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

228 """ 

229 

230 def myeval(value): 

231 try: 

232 value = eval(value, namespace) 

233 except (NameError, SyntaxError): 

234 pass 

235 return value 

236 

237 dct = {} 

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

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

240 key = strings[i] 

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

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

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

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

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

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

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

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

249 else: 

250 value = myeval(value) 

251 dct[key] = value 

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

253 return dct