1import sys 

2from typing import Dict, Any 

3import numpy as np 

4 

5from ase.calculators.calculator import (get_calculator_class, 

6 names as calcnames, 

7 PropertyNotImplementedError) 

8from ase.constraints import FixAtoms, UnitCellFilter 

9from ase.eos import EquationOfState 

10from ase.io import read, write, Trajectory 

11from ase.optimize import LBFGS 

12import ase.db as db 

13 

14 

15class CLICommand: 

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

17 

18 Four types of calculations can be done: 

19 

20 * single point 

21 * atomic relaxations 

22 * unit cell + atomic relaxations 

23 * equation-of-state 

24 

25 Examples of the four types of calculations: 

26 

27 ase run emt h2o.xyz 

28 ase run emt h2o.xyz -f 0.01 

29 ase run emt cu.traj -s 0.01 

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

31 """ 

32 

33 @staticmethod 

34 def add_arguments(parser): 

35 parser.add_argument('calculator', 

36 help='Name of calculator to use. ' 

37 'Must be one of: {}.' 

38 .format(', '.join(calcnames))) 

39 CLICommand.add_more_arguments(parser) 

40 

41 @staticmethod 

42 def add_more_arguments(parser): 

43 add = parser.add_argument 

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

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

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

47 metavar='key=value,...', 

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

49 'calculator specific parameters.') 

50 add('-t', '--tag', 

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

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

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

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

55 'atomic magnetic moments.') 

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

57 help='Relax internal coordinates.') 

58 add('--constrain-tags', 

59 metavar='T1,T2,...', 

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

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

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

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

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

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

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

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

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

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

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

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

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

73 

74 @staticmethod 

75 def run(args): 

76 runner = Runner() 

77 runner.parse(args) 

78 runner.run() 

79 

80 

81class Runner: 

82 def __init__(self): 

83 self.args = None 

84 self.calculator_name = None 

85 

86 def parse(self, args): 

87 self.calculator_name = args.calculator 

88 self.args = args 

89 

90 def run(self): 

91 args = self.args 

92 

93 atoms = self.build(args.name) 

94 if args.modify: 

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

96 

97 if args.name == '-': 

98 args.name = 'stdin' 

99 

100 self.set_calculator(atoms, args.name) 

101 

102 self.calculate(atoms, args.name) 

103 

104 def calculate(self, atoms, name): 

105 args = self.args 

106 

107 if args.maximum_force or args.maximum_stress: 

108 self.optimize(atoms, name) 

109 if args.equation_of_state: 

110 self.eos(atoms, name) 

111 self.calculate_once(atoms) 

112 

113 if args.after: 

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

115 

116 if args.output: 

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

118 

119 def build(self, name): 

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 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 args = self.args 

140 

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

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

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

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

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

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

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

148 try: 

149 getattr(atoms, method)() 

150 except PropertyNotImplementedError: 

151 pass 

152 

153 def optimize(self, atoms, name): 

154 args = self.args 

155 if args.constrain_tags: 

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

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

158 atoms.constraints = FixAtoms(mask=mask) 

159 

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

161 if args.maximum_stress: 

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

163 fmax = args.maximum_stress 

164 else: 

165 optimizer = LBFGS(atoms, logfile=logfile) 

166 fmax = args.maximum_force 

167 

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

169 optimizer.attach(trajectory) 

170 optimizer.run(fmax=fmax) 

171 

172 def eos(self, atoms, name): 

173 args = self.args 

174 

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

176 

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

178 N = int(N) 

179 eps = float(eps) / 100 

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

181 v1 = atoms.get_volume() 

182 volumes = strains**3 * v1 

183 energies = [] 

184 cell1 = atoms.cell 

185 for s in strains: 

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

187 energies.append(atoms.get_potential_energy()) 

188 traj.write(atoms) 

189 traj.close() 

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

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

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

193 from ase.parallel import parprint as p 

194 p('volumes:', volumes) 

195 p('energies:', energies) 

196 p('fitted energy:', e0) 

197 p('fitted volume:', v0) 

198 p('bulk modulus:', B) 

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

200 

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

202 if '.' in name: 

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

204 if self.args.tag is not None: 

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

206 if ext: 

207 name += '.' + ext 

208 return name 

209 

210 

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

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

213 

214 Examples: 

215 

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

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

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

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

220 """ 

221 

222 def myeval(value): 

223 try: 

224 value = eval(value, namespace) 

225 except (NameError, SyntaxError): 

226 pass 

227 return value 

228 

229 dct = {} 

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

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

232 key = strings[i] 

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

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

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

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

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

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

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

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

241 else: 

242 value = myeval(value) 

243 dct[key] = value 

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

245 return dct