File: espresso.py

package info (click to toggle)
python-ase 3.12.0-2
  • links: PTS, VCS
  • area: main
  • in suites: stretch
  • size: 14,192 kB
  • ctags: 8,112
  • sloc: python: 93,375; sh: 99; makefile: 94
file content (221 lines) | stat: -rw-r--r-- 8,265 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
"""Reads quantum espresso files. Tested for output on PWSCF v.5.0.2, only
for typical output of input files made with ASE -- that is, ibrav=0."""

import numpy as np
from ase.atoms import Atoms, Atom
from ase import units
from ase.calculators.singlepoint import SinglePointCalculator


def read_espresso_out(fileobj, index):
    """Reads quantum espresso output text files."""
    if isinstance(fileobj, str):
        fileobj = open(fileobj, 'rU')
    lines = fileobj.readlines()
    images = []

    # Get unit cell info.
    bl_line = [line for line in lines if 'bravais-lattice index' in line]
    if len(bl_line) != 1:
        raise NotImplementedError('Unsupported: unit cell changing.')
    bl_line = bl_line[0].strip()
    brav_latt_index = bl_line.split('=')[1].strip()
    if brav_latt_index != '0':
        raise NotImplementedError('Supported only for Bravais-lattice '
                                  'index of 0 (free).')
    lp_line = [line for line in lines if 'lattice parameter (alat)' in
               line]
    if len(lp_line) != 1:
        raise NotImplementedError('Unsupported: unit cell changing.')
    lp_line = lp_line[0].strip().split('=')[1].strip().split()[0]
    lattice_parameter = float(lp_line) * units.Bohr
    ca_line_no = [number for (number, line) in enumerate(lines) if
                  'crystal axes: (cart. coord. in units of alat)' in line]
    if len(ca_line_no) != 1:
        raise NotImplementedError('Unsupported: unit cell changing.')
    ca_line_no = int(ca_line_no[0])
    cell = np.zeros((3, 3))
    for number, line in enumerate(lines[ca_line_no + 1: ca_line_no + 4]):
        line = line.split('=')[1].strip()[1:-1]
        values = [float(value) for value in line.split()]
        cell[number, 0] = values[0]
        cell[number, 1] = values[1]
        cell[number, 2] = values[2]
    cell *= lattice_parameter

    # Find atomic positions and add to images.
    for number, line in enumerate(lines):
        key = 'Begin final coordinates'  # these just reprint last posn.
        if key in line:
            break
        key = 'Cartesian axes'
        if key in line:
            atoms = make_atoms(number, lines, key, cell)
            images.append(atoms)
        key = 'ATOMIC_POSITIONS (crystal)'
        if key in line:
            atoms = make_atoms(number, lines, key, cell)
            images.append(atoms)
    return images[index]


def make_atoms(index, lines, key, cell):
    """Scan through lines to get the atomic positions."""
    atoms = Atoms()
    if key == 'Cartesian axes':
        for line in lines[index + 3:]:
            entries = line.split()
            if len(entries) == 0:
                break
            symbol = entries[1][:-1]
            x = float(entries[6])
            y = float(entries[7])
            z = float(entries[8])
            atoms.append(Atom(symbol, (x, y, z)))
        atoms.set_cell(cell)
    elif key == 'ATOMIC_POSITIONS (crystal)':
        for line in lines[index + 1:]:
            entries = line.split()
            if len(entries) == 0 or (entries[0] == 'End'):
                break
            symbol = entries[0][:-1]
            x = float(entries[1])
            y = float(entries[2])
            z = float(entries[3])
            atoms.append(Atom(symbol, (x, y, z)))
        atoms.set_cell(cell, scale_atoms=True)
    # Energy is located after positions.
    energylines = [number for number, line in enumerate(lines) if
                   ('!' in line and 'total energy' in line)]
    energyline = min([n for n in energylines if n > index])
    energy = float(lines[energyline].split()[-2]) * units.Ry
    # Forces are located after positions.
    forces = np.zeros((len(atoms), 3))
    forcelines = [number for number, line in enumerate(lines) if
                  'Forces acting on atoms (Ry/au):' in line]
    forceline = min([n for n in forcelines if n > index])
    for line in lines[forceline + 4:]:
        words = line.split()
        if len(words) == 0:
            break
        fx = float(words[-3])
        fy = float(words[-2])
        fz = float(words[-1])
        atom_number = int(words[1]) - 1
        forces[atom_number] = (fx, fy, fz)
    forces *= units.Ry / units.Bohr
    calc = SinglePointCalculator(atoms, energy=energy, forces=forces)
    atoms.set_calculator(calc)
    return atoms


def read_espresso_in(fileobj):
    """Reads espresso input files."""
    if isinstance(fileobj, str):
        fileobj = open(fileobj, 'rU')
    data, extralines = read_fortran_namelist(fileobj)
    positions, method = get_atomic_positions(extralines,
                                             n_atoms=data['system']['nat'])
    cell = get_cell_parameters(extralines)
    if data['system']['ibrav'] == 0:
        atoms = build_atoms(positions, method, cell,
                            data['system']['celldm(1)'])
    else:
        raise NotImplementedError('ibrav=%i not implemented.' %
                                  data['system']['ibrav'])
    return atoms


def build_atoms(positions, method, cell, alat):
    """Creates the atoms for a quantum espresso in file."""
    if method != 'crystal':
        raise NotImplementedError('Only supported for crystal method of '
                                  'ATOMIC_POSITIONS, not %s.' % method)
    atoms = Atoms()
    for el, (x, y, z) in positions:
        atoms.append(Atom(el, (x, y, z)))
    cell *= alat * units.Bohr
    atoms.set_cell(cell, scale_atoms=True)
    return atoms


def get_atomic_positions(lines, n_atoms):
    """Returns the atomic positions of the atoms as an (ordered) list from
    the lines of text of the espresso input file."""
    atomic_positions = []
    line = [n for (n, l) in enumerate(lines) if 'ATOMIC_POSITIONS' in l]
    if len(line) == 0:
        return None
    if len(line) > 1:
        raise RuntimeError('More than one ATOMIC_POSITIONS section?')
    line_no = line[0]
    for line in lines[line_no + 1:line_no + n_atoms + 1]:
        el, x, y, z = line.split()
        atomic_positions.append([el, (f2f(x), f2f(y), f2f(z))])
    line = lines[line_no]
    if '{' in line:
        method = line[line.find('{') + 1:line.find('}')]
    elif '(' in line:
        method = line[line.find('(') + 1:line.find(')')]
    else:
        method = None
    return atomic_positions, method


def get_cell_parameters(lines):
    """Returns the cell parameters as a matrix."""
    cell_parameters = np.zeros((3, 3))
    line = [n for (n, l) in enumerate(lines) if 'CELL_PARAMETERS' in l]
    if len(line) == 0:
        return None
    if len(line) > 1:
        raise RuntimeError('More than one CELL_PARAMETERS section?')
    line_no = line[0]
    for vector, line in enumerate(lines[line_no + 1:line_no + 4]):
        x, y, z = line.split()
        cell_parameters[vector] = (f2f(x), f2f(y), f2f(z))
    return cell_parameters


def str2value(string):
    """Convert string into int, float, or bool, if possible, else return it."""
    for datatype in [int, float]:
        try:
            return datatype(string)
        except ValueError:
            pass
    return {'.true.': True, '.false.': False}.get(string, string)


def read_fortran_namelist(fileobj):
    """Takes a fortran-namelist formatted file and returns appropriate
    dictionaries, followed by lines of text that do not fit this pattern.
    """
    data = {}
    extralines = []
    indict = False
    fileobj.seek(0)
    for line in fileobj.readlines():
        if indict and line.strip().startswith('/'):
            indict = False
        elif line.strip().startswith('&'):
            indict = True
            dictname = line.strip()[1:].lower()
            data[dictname] = {}
        elif (not indict) and (len(line.strip()) > 0):
            extralines.append(line)
        elif indict:
            key, value = line.strip().split('=')
            if value.endswith(','):
                value = value[:-1]
            value = str2value(value.strip())
            data[dictname][key.strip()] = value
    return data, extralines


def f2f(value):
    """Converts a fortran-formatted double precision number (e.g., 2.323d2)
    to a python float. value should be a string."""
    value = value.replace('d', 'e')
    value = value.replace('D', 'e')
    return float(value)