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

1"""Read Wannier90 wout format.""" 

2from typing import IO, Any, Dict 

3 

4import numpy as np 

5 

6from ase import Atoms 

7 

8 

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

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

11 lines = fileobj.readlines() 

12 

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') 

18 

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

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

21 

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') 

27 

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 

38 

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

40 

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,))} 

50 

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 

64 

65 return {'atoms': atoms, 

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

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

68 

69 

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