Coverage for /builds/debichem-team/python-ase/ase/io/wannier90.py : 91.49%
Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""Read Wannier90 wout format."""
2from typing import IO, Dict, Any
4import numpy as np
6from ase import Atoms
9def read_wout_all(fileobj: IO[str]) -> Dict[str, Any]:
10 """Read atoms, wannier function centers and spreads."""
11 lines = fileobj.readlines()
13 for n, line in enumerate(lines):
14 if line.strip().lower().startswith('lattice vectors (ang)'):
15 break
16 else:
17 raise ValueError('Could not fine lattice vectors')
19 cell = [[float(x) for x in line.split()[-3:]]
20 for line in lines[n + 1:n + 4]]
22 for n, line in enumerate(lines):
23 if 'cartesian coordinate (ang)' in line.lower():
24 break
25 else:
26 raise ValueError('Could not find coordinates')
28 positions = []
29 symbols = []
30 n += 2
31 while True:
32 words = lines[n].split()
33 if len(words) == 1:
34 break
35 positions.append([float(x) for x in words[-4:-1]])
36 symbols.append(words[1])
37 n += 1
39 atoms = Atoms(symbols, positions, cell=cell, pbc=True)
41 n = len(lines) - 1
42 while n > 0:
43 if lines[n].strip().lower().startswith('final state'):
44 break
45 n -= 1
46 else:
47 return {'atoms': atoms,
48 'centers': np.zeros((0, 3)),
49 'spreads': np.zeros((0,))}
51 n += 1
52 centers = []
53 spreads = []
54 while True:
55 line = lines[n].strip()
56 if line.startswith('WF'):
57 centers.append([float(x)
58 for x in
59 line.split('(')[1].split(')')[0].split(',')])
60 spreads.append(float(line.split()[-1]))
61 n += 1
62 else:
63 break
65 return {'atoms': atoms,
66 'centers': np.array(centers),
67 'spreads': np.array(spreads)}
70def read_wout(fileobj: IO[str],
71 include_wannier_function_centers: bool = True) -> Atoms:
72 """Read atoms and wannier function centers (as symbol X)."""
73 dct = read_wout_all(fileobj)
74 atoms = dct['atoms']
75 if include_wannier_function_centers:
76 centers = dct['centers']
77 atoms += Atoms(f'X{len(centers)}', centers)
78 return atoms