Skip to content

Commit e2a4eb6

Browse files
dev-zerodev-zero
committed
refactor structure loading
1 parent 14d8e1b commit e2a4eb6

File tree

3 files changed

+105
-94
lines changed

3 files changed

+105
-94
lines changed

‎aiida_cp2k/cli/data/structure.py‎

Lines changed: 5 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,13 @@
11
# -*- coding: utf-8 -*-
22
"""Command line utilities to create and inspect `StructureData` nodes from CP2K input files."""
33

4-
import re
5-
64
import click
7-
import numpy as np
85

96
from aiida.cmdline.params import options
107
from aiida.cmdline.utils import decorators, echo
118

129
from . import cmd_data
10+
from ..utils.structure import structure_from_cp2k_inp
1311

1412

1513
@cmd_data.group('structure')
@@ -24,79 +22,10 @@ def cmd_structure():
2422
def cmd_import(filename, dry_run):
2523
"""Import a `StructureData` from a CP2K input file."""
2624

27-
# pylint: disable=import-outside-toplevel,invalid-name,too-many-locals,too-many-statements,too-many-branches
28-
29-
from cp2k_input_tools.parser import CP2KInputParser
30-
from aiida.orm.nodes.data.structure import StructureData, Kind, Site, symop_fract_from_ortho
31-
from ase.geometry.cell import cell_to_cellpar, cellpar_to_cell
32-
33-
# the following was taken from aiida-quantumespresso
34-
VALID_ELEMENTS_REGEX = re.compile(
35-
r"""
36-
^(
37-
H | He |
38-
Li | Be | B | C | N | O | F | Ne |
39-
Na | Mg | Al | Si | P | S | Cl | Ar |
40-
K | Ca | Sc | Ti | V | Cr | Mn | Fe | Co | Ni | Cu | Zn | Ga | Ge | As | Se | Br | Kr |
41-
Rb | Sr | Y | Zr | Nb | Mo | Tc | Ru | Rh | Pd | Ag | Cd | In | Sn | Sb | Te | I | Xe |
42-
Cs | Ba | Hf | Ta | W | Re | Os | Ir | Pt | Au | Hg | Tl | Pb | Bi | Po | At | Rn |
43-
Fr | Ra | Rf | Db | Sg | Bh | Hs | Mt |
44-
La | Ce | Pr | Nd | Pm | Sm | Eu | Gd | Tb | Dy | Ho | Er | Tm | Yb | Lu | # Lanthanides
45-
Ac | Th | Pa | U | Np | Pu | Am | Cm | Bk | Cf | Es | Fm | Md | No | Lr | # Actinides
46-
)
47-
""", re.VERBOSE | re.IGNORECASE)
48-
49-
parser = CP2KInputParser()
50-
tree = parser.parse(filename)
51-
force_eval_no = -1
52-
53-
for force_eval_no, force_eval in enumerate(tree['+force_eval']):
54-
try:
55-
cell = force_eval['+subsys']['+cell']
56-
kinds = force_eval['+subsys']['+kind']
57-
coord = force_eval['+subsys']['+coord']
58-
break # for now grab the first &COORD found
59-
except KeyError:
60-
continue
61-
else:
62-
echo.echo_critical('no STRUCTURE, CELL, or KIND found in the given input file')
63-
64-
# CP2K can get its cell information in two ways:
65-
# - A, B, C: cell vectors
66-
# - ABC: scaling of cell vectors, ALPHA_BETA_GAMMA: angles between the cell vectors (optional)
67-
68-
if 'a' in cell:
69-
unit_cell = np.array([cell['a'], cell['b'], cell['c']]) # unit vectors given
70-
cellpar = cell_to_cellpar(unit_cell)
71-
elif 'abc' in cell:
72-
cellpar = cell['abc'] + cell.get('alpha_beta_gamma', [90., 90., 90.])
73-
unit_cell = cellpar_to_cell(cellpar)
74-
else:
75-
echo.echo_critical('incomplete &CELL section')
76-
77-
pbc = [c in cell.get('periodic', 'XYZ') for c in 'XYZ']
78-
79-
structure = StructureData(cell=unit_cell, pbc=pbc)
80-
81-
if coord.get('scaled', False):
82-
tmat = symop_fract_from_ortho(cellpar)
83-
else:
84-
tmat = np.eye(3)
85-
86-
for kind in kinds:
87-
name = kind['_']
88-
89-
try:
90-
# prefer the ELEMENT keyword, fallback to extracting from name
91-
element = kind.get('element', VALID_ELEMENTS_REGEX.search(name)[0])
92-
except TypeError:
93-
echo.echo_critical('ELEMENT not set and unable to extract from {name}'.format(name=name))
94-
95-
structure.append_kind(Kind(name=name, symbols=element))
96-
97-
for name, position, _ in parser.coords(force_eval_no):
98-
# positions can be scaled, apply transformation matrix
99-
structure.append_site(Site(kind_name=name, position=tmat @ np.array(position)))
25+
try:
26+
structure = structure_from_cp2k_inp(filename)
27+
except ValueError as exc:
28+
echo.echo_critical(str(exc))
10029

10130
formula = structure.get_formula()
10231

‎aiida_cp2k/cli/utils/structure.py‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
# -*- coding: utf-8 -*-
2+
"""Helper functions for building CLI functions dealing with structures"""
3+
4+
import re
5+
6+
import numpy as np
7+
8+
9+
def structure_from_cp2k_inp(filename):
10+
"""Create an AiiDA StructureData from a structure inside a CP2K input file"""
11+
# pylint: disable=import-outside-toplevel,invalid-name,too-many-locals,too-many-statements,too-many-branches
12+
13+
from cp2k_input_tools.parser import CP2KInputParser
14+
from aiida.orm.nodes.data.structure import StructureData, Kind, Site, symop_fract_from_ortho
15+
from ase.geometry.cell import cell_to_cellpar, cellpar_to_cell
16+
import ase.io
17+
18+
# the following was taken from aiida-quantumespresso
19+
VALID_ELEMENTS_REGEX = re.compile(
20+
r"""
21+
^(
22+
H | He |
23+
Li | Be | B | C | N | O | F | Ne |
24+
Na | Mg | Al | Si | P | S | Cl | Ar |
25+
K | Ca | Sc | Ti | V | Cr | Mn | Fe | Co | Ni | Cu | Zn | Ga | Ge | As | Se | Br | Kr |
26+
Rb | Sr | Y | Zr | Nb | Mo | Tc | Ru | Rh | Pd | Ag | Cd | In | Sn | Sb | Te | I | Xe |
27+
Cs | Ba | Hf | Ta | W | Re | Os | Ir | Pt | Au | Hg | Tl | Pb | Bi | Po | At | Rn |
28+
Fr | Ra | Rf | Db | Sg | Bh | Hs | Mt |
29+
La | Ce | Pr | Nd | Pm | Sm | Eu | Gd | Tb | Dy | Ho | Er | Tm | Yb | Lu | # Lanthanides
30+
Ac | Th | Pa | U | Np | Pu | Am | Cm | Bk | Cf | Es | Fm | Md | No | Lr | # Actinides
31+
)
32+
""", re.VERBOSE | re.IGNORECASE)
33+
34+
parser = CP2KInputParser()
35+
tree = parser.parse(filename)
36+
force_eval_no = -1
37+
force_eval = None
38+
39+
for force_eval_no, force_eval in enumerate(tree['+force_eval']):
40+
try:
41+
cell = force_eval['+subsys']['+cell']
42+
kinds = force_eval['+subsys']['+kind']
43+
break # for now grab the first &COORD found
44+
except KeyError:
45+
continue
46+
else:
47+
raise ValueError('no CELL, or KIND found in the given input file')
48+
49+
# CP2K can get its cell information in two ways:
50+
# - A, B, C: cell vectors
51+
# - ABC: scaling of cell vectors, ALPHA_BETA_GAMMA: angles between the cell vectors (optional)
52+
53+
if 'a' in cell:
54+
unit_cell = np.array([cell['a'], cell['b'], cell['c']]) # unit vectors given
55+
cellpar = cell_to_cellpar(unit_cell)
56+
elif 'abc' in cell:
57+
cellpar = cell['abc'] + cell.get('alpha_beta_gamma', [90., 90., 90.])
58+
unit_cell = cellpar_to_cell(cellpar)
59+
else:
60+
raise ValueError('incomplete &CELL section')
61+
62+
pbc = [c in cell.get('periodic', 'XYZ') for c in 'XYZ']
63+
64+
structure = StructureData(cell=unit_cell, pbc=pbc)
65+
66+
if force_eval['+subsys'].get('+coord', {}).get('scaled', False):
67+
tmat = symop_fract_from_ortho(cellpar)
68+
else:
69+
tmat = np.eye(3)
70+
71+
for kind in kinds:
72+
name = kind['_']
73+
74+
try:
75+
# prefer the ELEMENT keyword, fallback to extracting from name
76+
element = kind.get('element', VALID_ELEMENTS_REGEX.search(name)[0])
77+
except TypeError:
78+
raise ValueError('ELEMENT not set and unable to extract from {name}'.format(name=name))
79+
80+
structure.append_kind(Kind(name=name, symbols=element))
81+
82+
try:
83+
structfn = force_eval["+subsys"]["+topology"]["coord_file_name"]
84+
atoms = ase.io.read(structfn)
85+
86+
for name, position in zip(atoms.symbols, atoms.positions):
87+
structure.append_site(Site(kind_name=name, position=tmat @ np.array(position)))
88+
89+
except KeyError:
90+
for name, position, _ in parser.coords(force_eval_no):
91+
# positions can be scaled, apply transformation matrix
92+
structure.append_site(Site(kind_name=name, position=tmat @ np.array(position)))
93+
94+
return structure

‎aiida_cp2k/cli/workflows/base.py‎

Lines changed: 6 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -9,19 +9,7 @@
99

1010
from . import cmd_launch
1111
from ..utils import cp2k_options, launch_process
12-
13-
14-
def _import_struct(structfn):
15-
"""Import a data structure from a filename using ASE"""
16-
17-
from aiida.orm import StructureData
18-
import ase.io
19-
20-
asecell = ase.io.read(structfn)
21-
structure = StructureData(ase=asecell)
22-
structure.store()
23-
24-
return structure
12+
from ..utils.structure import structure_from_cp2k_inp
2513

2614

2715
@cmd_launch.command('base')
@@ -60,11 +48,11 @@ def cmd_launch_workflow(code, daemon, cp2k_input_file, label, description, struc
6048

6149
else:
6250
try:
63-
structfn = tree["FORCE_EVAL"]["SUBSYS"]["TOPOLOGY"].pop("COORD_FILE_NAME")
64-
builder.cp2k.structure = _import_struct(structfn)
65-
echo.echo("Created StructureData<{}> from file {}]\n".format(builder.cp2k.structure.pk, structfn))
66-
except KeyError:
67-
pass
51+
with open(cp2k_input_file.name, "r") as fhandle:
52+
builder.cp2k.structure = structure_from_cp2k_inp(fhandle)
53+
builder.cp2k.structure.store()
54+
echo.echo("Created StructureData<{}> from file {}]\n".format(builder.cp2k.structure.pk,
55+
cp2k_input_file.name))
6856
except ValueError as err:
6957
echo.echo_critical(str(err))
7058

0 commit comments

Comments
 (0)