Coverage for /builds/ase/ase/ase/io/wannier90.py: 93.88%

1# fmt: off 

2 

3"""Read Wannier90 wout format.""" 

4from typing import IO, Any, Dict 

5 

6import numpy as np 

7 

8from ase import Atoms 

9 

10 

11def read_wout_all(fileobj: IO[str]) -> Dict[str, Any]: 

12 """Read atoms, wannier function centers and spreads.""" 

13 lines = fileobj.readlines() 

14 

15 for n, line in enumerate(lines): 

16 if line.strip().lower().startswith('lattice vectors (ang)'): 

17 break 

18 else: 

19 raise ValueError('Could not fine lattice vectors') 

20 

21 cell = [[float(x) for x in line.split()[-3:]] 

22 for line in lines[n + 1:n + 4]] 

23 

24 for n, line in enumerate(lines): 

25 if 'cartesian coordinate (ang)' in line.lower(): 

26 break 

27 else: 

28 raise ValueError('Could not find coordinates') 

29 

30 positions = [] 

31 symbols = [] 

32 n += 2 

33 while True: 

34 words = lines[n].split() 

35 if len(words) == 1: 

36 break 

37 positions.append([float(x) for x in words[-4:-1]]) 

38 symbols.append(words[1]) 

39 n += 1 

40 

41 atoms = Atoms(symbols, positions, cell=cell, pbc=True) 

42 

43 n = len(lines) - 1 

44 while n > 0: 

45 if lines[n].strip().lower().startswith('final state'): 

46 break 

47 n -= 1 

48 else: 

49 return {'atoms': atoms, 

50 'centers': np.zeros((0, 3)), 

51 'spreads': np.zeros((0,))} 

52 

53 n += 1 

54 centers = [] 

55 spreads = [] 

56 while True: 

57 line = lines[n].strip() 

58 if line.startswith('WF'): 

59 centers.append([float(x) 

60 for x in 

61 line.split('(')[1].split(')')[0].split(',')]) 

62 spreads.append(float(line.split()[-1])) 

63 n += 1 

64 else: 

65 break 

66 

67 return {'atoms': atoms, 

68 'centers': np.array(centers), 

69 'spreads': np.array(spreads)} 

70 

71 

72def read_wout(fileobj: IO[str], 

73 include_wannier_function_centers: bool = True) -> Atoms: 

74 """Read atoms and wannier function centers (as symbol X).""" 

75 dct = read_wout_all(fileobj) 

76 atoms = dct['atoms'] 

77 if include_wannier_function_centers: 

78 centers = dct['centers'] 

79 atoms += Atoms(f'X{len(centers)}', centers) 

80 return atoms