Skip to content

Commit 1994537

Browse files
mkphuthimkphuthi
committed
Merge branch 'mkphuthi/develop' of https://github.com/BattModels/asimtools into mkphuthi/develop
2 parents 97470a1 + 2ff2827 commit 1994537

File tree

4 files changed

+245
-31
lines changed

4 files changed

+245
-31
lines changed

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ ASE source for this as ASIMTools can't go into the calculator code.
2424
- asimtools.utils.write_yaml now stops sorting keys to help with readability of
2525
written yamls
2626
- write_atoms now more universally used and recommended in asimmodules
27+
- VASP calculation now fails if the max number of iterations is reached for ibrion=1,2,3
2728

2829
### Fixed
2930
- Minor bugs in geometry optimizations

‎asimtools/asimmodules/vasp/vasp.py‎

Lines changed: 55 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -13,15 +13,48 @@
1313
from numpy.random import randint
1414
import subprocess
1515
import logging
16+
import shutil
1617
from ase.io import read
1718
from ase import Atoms
1819
from pymatgen.io.ase import AseAtomsAdaptor
1920
from pymatgen.io.vasp import Poscar, Incar, Potcar, Kpoints, VaspInput
2021
import pymatgen.io.vasp.sets
2122
from asimtools.utils import (
2223
get_atoms,
24+
get_str_btn,
2325
)
2426

27+
def execute_vasp_run_command(
28+
command: str,
29+
) -> None:
30+
command = command.split(' ')
31+
completed_process = subprocess.run(
32+
command, check=False, capture_output=True, text=True,
33+
)
34+
35+
with open('vasp_stdout.txt', 'a+', encoding='utf-8') as f:
36+
f.write(completed_process.stdout)
37+
38+
if completed_process.returncode != 0:
39+
err_txt = f'VASP failed with error code: {completed_process.returncode}'
40+
err_txt += '\nSee vasp_stderr.txt for details.'
41+
logging.error(err_txt)
42+
with open('vasp_stderr.txt', 'a+', encoding='utf-8') as f:
43+
f.write(completed_process.stderr)
44+
completed_process.check_returncode()
45+
if "ZBRENT: fatal error in bracketing" in completed_process.stderr:
46+
filenum = 0
47+
while os.path.exists(f'OUTCAR.{filenum}'):
48+
filenum += 1
49+
shutil.move('OUTCAR', f'OUTCAR.{filenum}')
50+
shutil.move('POSCAR', f'POSCAR.{filenum}')
51+
shutil.move('CONTCAR', f'POSCAR')
52+
execute_vasp_run_command(command)
53+
else:
54+
logging.info(
55+
f'VASP run completed with code {completed_process.returncode}.'
56+
)
57+
2558
def vasp(
2659
image: Optional[Dict],
2760
user_incar_settings: Optional[Dict] = None,
@@ -58,10 +91,6 @@ def vasp(
5891

5992
struct = get_atoms(**image, return_type='pymatgen')
6093
if mpset is not None:
61-
# if not ((incar is None) or (potcar is None) or (kpoints is None)):
62-
# raise ValueError(
63-
# 'Provide either mpset or all of incar and kpoints'
64-
# )
6594
try:
6695
set_ = getattr(pymatgen.io.vasp.sets, mpset)
6796
except:
@@ -107,32 +136,31 @@ def vasp(
107136
**vaspinput_kwargs
108137
)
109138

110-
111139
vasp_input.write_input("./")
112-
# if incar_kwargs is not None:
113-
# with open('INCAR', 'a+') as fp:
114-
# for k, v in incar_kwargs.items():
115-
# fp.write(f'\n{k} = {v}')
116140

117141
if run_vasp:
118-
command = command.split(' ')
119-
completed_process = subprocess.run(
120-
command, check=False, capture_output=True, text=True,
121-
)
122-
123-
with open('vasp_stdout.txt', 'a+', encoding='utf-8') as f:
124-
f.write(completed_process.stdout)
125-
126-
if completed_process.returncode != 0:
127-
err_txt = f'Failed to run VASP\n'
128-
err_txt += 'See vasp_stderr.txt for details.'
129-
logging.error(err_txt)
130-
with open('vasp_stderr.txt', 'a+', encoding='utf-8') as f:
131-
f.write(completed_process.stderr)
132-
completed_process.check_returncode()
133-
return {}
134-
135-
if write_image_output:
142+
optimization_failed = False
143+
execute_vasp_run_command(command)
144+
incar = vasp_input.incar
145+
ibrion = incar.get('IBRION', -1)
146+
nsw = incar.get('NSW', 0)
147+
# Don't write result if running a relaxation that didn't finish
148+
if ibrion in (1,2,3):
149+
if nsw > 0:
150+
with open('OUTCAR', 'r') as f:
151+
lines = f.readlines()
152+
lines = lines[::-1]
153+
for line in lines:
154+
if 'Ionic step' in line:
155+
last_step = int(get_str_btn(line, 'Ionic step', '--'))
156+
if last_step == nsw:
157+
optimization_failed = True
158+
raise RuntimeError(
159+
'VASP relaxation did not complete. '
160+
'Check OUTCAR for details.'
161+
)
162+
163+
if write_image_output and not optimization_failed:
136164
atoms_output = read('OUTCAR')
137165
atoms_output.write(
138166
'image_output.xyz',

‎asimtools/asimmodules/vasp/vasp.py.bak‎

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
#!/usr/bin/env python
2+
'''
3+
Runs VASP based on input files and optionally MP settings.
4+
Heavily uses pymatgen for IO and MP settings.
5+
VASP must be installed
6+
7+
Author: mkphuthi@github.com
8+
'''
9+
from typing import Dict, Optional, Sequence
10+
import os
11+
import sys
12+
from pathlib import Path
13+
from numpy.random import randint
14+
import subprocess
15+
import logging
16+
from ase.io import read
17+
from ase import Atoms
18+
from pymatgen.io.ase import AseAtomsAdaptor
19+
from pymatgen.io.vasp import Poscar, Incar, Potcar, Kpoints, VaspInput
20+
import pymatgen.io.vasp.sets
21+
from asimtools.utils import (
22+
get_atoms,
23+
)
24+
25+
def vasp(
26+
image: Optional[Dict],
27+
user_incar_settings: Optional[Dict] = None,
28+
user_kpoints_settings: Optional[Dict] = None,
29+
user_potcar_functional: str = 'PBE_64',
30+
potcar: Optional[Dict] = None,
31+
vaspinput_kwargs: Optional[Dict] = None,
32+
command: str = 'srun vasp_std',
33+
mpset: Optional[str] = None,
34+
prev_calc: Optional[os.PathLike] = None,
35+
write_image_output: bool = True,
36+
run_vasp: bool = True,
37+
) -> Dict:
38+
"""Run VASP with given input files and specified image
39+
40+
:param image: Initial image for VASP calculation. Image specification,
41+
see :func:`asimtools.utils.get_atoms`
42+
:type image: Dict
43+
:param vaspinput_args: Dictionary of pymatgen's VaspInput arguments.
44+
See :class:`pymatgen.io.vasp.inputs.VaspInput`
45+
:type vaspinput_args: Dict
46+
:param command: Command with which to run VASP, defaults to 'vasp_std'
47+
:type command: str, optional
48+
:param mpset: Materials Project VASP set to use see
49+
:mod:`pymatgen.io.vasp.sets`, defaults to None
50+
:type mpset: str, optional
51+
:param write_image_output: Whether to write output image in standard
52+
asimtools format to file, defaults to False
53+
:type write_image_output: bool, optional
54+
"""
55+
56+
if vaspinput_kwargs is None:
57+
vaspinput_kwargs = {}
58+
59+
struct = get_atoms(**image, return_type='pymatgen')
60+
if mpset is not None:
61+
# if not ((incar is None) or (potcar is None) or (kpoints is None)):
62+
# raise ValueError(
63+
# 'Provide either mpset or all of incar and kpoints'
64+
# )
65+
try:
66+
set_ = getattr(pymatgen.io.vasp.sets, mpset)
67+
except:
68+
raise ImportError(
69+
f'Unknown mpset: {mpset}. See available sets in pymatgen.')
70+
71+
if prev_calc is not None:
72+
vasp_input = set_.from_prev_calc(
73+
prev_calc,
74+
user_incar_settings=user_incar_settings,
75+
user_kpoints_settings=user_kpoints_settings,
76+
user_potcar_functional=user_potcar_functional,
77+
**vaspinput_kwargs
78+
)
79+
else:
80+
vasp_input = set_(
81+
struct,
82+
user_incar_settings=user_incar_settings,
83+
user_kpoints_settings=user_kpoints_settings,
84+
user_potcar_functional=user_potcar_functional,
85+
**vaspinput_kwargs
86+
)
87+
88+
else:
89+
90+
incar = Incar(user_incar_settings)
91+
incar.check_params()
92+
if potcar is not None:
93+
potcar = Potcar(potcar)
94+
if user_kpoints_settings is not None:
95+
kpoints = Kpoints(user_kpoints_settings)
96+
else:
97+
kpoints=None
98+
99+
if vaspinput_args is None:
100+
vaspinput_args = {}
101+
102+
vasp_input = VaspInput(
103+
incar=incar,
104+
kpoints=kpoints,
105+
poscar=Poscar(struct),
106+
potcar=potcar,
107+
**vaspinput_kwargs
108+
)
109+
110+
111+
vasp_input.write_input("./")
112+
# if incar_kwargs is not None:
113+
# with open('INCAR', 'a+') as fp:
114+
# for k, v in incar_kwargs.items():
115+
# fp.write(f'\n{k} = {v}')
116+
117+
if run_vasp:
118+
command = command.split(' ')
119+
completed_process = subprocess.run(
120+
command, check=False, capture_output=True, text=True,
121+
)
122+
123+
with open('vasp_stdout.txt', 'a+', encoding='utf-8') as f:
124+
f.write(completed_process.stdout)
125+
126+
if completed_process.returncode != 0:
127+
err_txt = f'Failed to run VASP\n'
128+
err_txt += 'See vasp_stderr.txt for details.'
129+
logging.error(err_txt)
130+
with open('vasp_stderr.txt', 'a+', encoding='utf-8') as f:
131+
f.write(completed_process.stderr)
132+
completed_process.check_returncode()
133+
return {}
134+
135+
if write_image_output:
136+
atoms_output = read('OUTCAR')
137+
atoms_output.write(
138+
'image_output.xyz',
139+
format='extxyz',
140+
)
141+
142+
return {}

‎asimtools/calculators.py‎

Lines changed: 47 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -306,14 +306,14 @@ def load_matgl(calc_params):
306306

307307
return calc
308308

309-
def load_omat24(calc_params):
310-
"""Load any OMAT24 calculator
309+
def load_fairchemV1(calc_params):
310+
"""Load any fairchemV1 calculator
311311
312312
:param calc_params: parameters to be passed to fairchem.core.OCPCalculator.
313313
Must include a key "model" that points to the model files used to
314314
instantiate the potential
315315
:type calc_params: Dict
316-
:return: OMAT24 calculator
316+
:return: fairchem calculator
317317
:rtype: :class:`fairchem.core.OCPCalculator`
318318
"""
319319
from fairchem.core import OCPCalculator
@@ -328,6 +328,46 @@ def load_omat24(calc_params):
328328

329329
return calc
330330

331+
def load_fairchemV2(calc_params):
332+
"""Load any fairchemV1 calculator
333+
334+
:param calc_params: parameters to be passed to fairchem.core.FAIRChemCalculator.
335+
Must include a key "model" that points to the model files used to
336+
instantiate the potential
337+
:type calc_params: Dict
338+
:return: fairchem calculator
339+
:rtype: :class:`fairchem.core.FAIRChemCalculator`
340+
341+
Examples
342+
--------
343+
>>> from asimtools.calculators import load_calc
344+
>>> calc_params = {
345+
... 'name': 'fairchem',
346+
... 'args': {
347+
... 'model_name': 'uma-s-1',
348+
... 'device': 'cuda',
349+
... 'task_name': 'oc20' # Also 'omol','omat','oc20','odac' or 'omc'
350+
... }
351+
... }
352+
>>> calc = load_calc(calc_params=calc_params)
353+
354+
"""
355+
from fairchem.core import pretrained_mlip, FAIRChemCalculator
356+
og_calc_params = deepcopy(calc_params)
357+
task_name = calc_params['args'].pop('task_name', None)
358+
predictor = pretrained_mlip.get_predict_unit(**calc_params['args'])
359+
360+
try:
361+
calc = FAIRChemCalculator(predictor, task_name=task_name)
362+
except Exception:
363+
logging.error(
364+
"Failed to load FAIRChemCalculator with parameters:\n %s", \
365+
og_calc_params
366+
)
367+
raise
368+
369+
return calc
370+
331371
def load_ase_dftd3(calc_params):
332372
"""Load any calculator with DFTD3 correction as implemented in ASE
333373
@@ -370,6 +410,9 @@ def load_ase_dftd3(calc_params):
370410
'EspressoProfile': load_espresso_profile,
371411
'M3GNet': load_m3gnet,
372412
'MatGL': load_matgl,
373-
'OMAT24': load_omat24,
413+
'OMAT24': load_fairchemV1,
414+
'fairchemV1': load_fairchemV1,
415+
'fairchemV2': load_fairchemV2,
416+
'fairchem': load_fairchemV2,
374417
'ASEDFTD3': load_ase_dftd3,
375418
}

0 commit comments

Comments
 (0)