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
« prev ^ index » next coverage.py v7.2.7, created at 2023-12-10 11:04 +0000
1import sys
2from typing import Any, Dict
4import numpy as np
7class CLICommand:
8 """Run calculation with one of ASE's calculators.
10 Four types of calculations can be done:
12 * single point
13 * atomic relaxations
14 * unit cell + atomic relaxations
15 * equation-of-state
17 Examples of the four types of calculations:
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 """
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)
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(...)"')
67 @staticmethod
68 def run(args):
69 runner = Runner()
70 runner.parse(args)
71 runner.run()
74class Runner:
75 def __init__(self):
76 self.args = None
77 self.calculator_name = None
79 def parse(self, args):
80 self.calculator_name = args.calculator
81 self.args = args
83 def run(self):
84 args = self.args
86 atoms = self.build(args.name)
87 if args.modify:
88 exec(args.modify, {'atoms': atoms, 'np': np})
90 if args.name == '-':
91 args.name = 'stdin'
93 self.set_calculator(atoms, args.name)
95 self.calculate(atoms, args.name)
97 def calculate(self, atoms, name):
98 from ase.io import write
100 args = self.args
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)
108 if args.after:
109 exec(args.after, {'atoms': atoms})
111 if args.output:
112 write(args.output, atoms, append=True)
114 def build(self, name):
115 import ase.db as db
116 from ase.io import read
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
128 def set_calculator(self, atoms, name):
129 from ase.calculators.calculator import get_calculator_class
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)
138 def calculate_once(self, atoms):
139 from ase.calculators.calculator import PropertyNotImplementedError
141 args = self.args
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
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
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)
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
174 trajectory = Trajectory(self.get_filename(name, 'traj'), 'w', atoms)
175 optimizer.attach(trajectory)
176 optimizer.run(fmax=fmax)
178 def eos(self, atoms, name):
179 from ase.eos import EquationOfState
180 from ase.io import Trajectory
182 args = self.args
184 traj = Trajectory(self.get_filename(name, 'traj'), 'w', atoms)
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)
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
220def str2dict(s: str, namespace={}, sep: str = '=') -> Dict[str, Any]:
221 """Convert comma-separated key=value string to dictionary.
223 Examples:
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 """
231 def myeval(value):
232 try:
233 value = eval(value, namespace)
234 except (NameError, SyntaxError):
235 pass
236 return value
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