Coverage for /builds/kinetik161/ase/ase/io/gen.py: 92.68%

1"""Extension to ASE: read and write structures in GEN format 

2 

3Refer to DFTB+ manual for GEN format description. 

4 

5Note: GEN format only supports single snapshot. 

6""" 

7from typing import Dict, Sequence, Union 

8 

9from ase.atoms import Atoms 

10from ase.utils import reader, writer 

11 

12 

13@reader 

14def read_gen(fileobj): 

15 """Read structure in GEN format (refer to DFTB+ manual). 

16 Multiple snapshot are not allowed. """ 

17 image = Atoms() 

18 lines = fileobj.readlines() 

19 line = lines[0].split() 

20 natoms = int(line[0]) 

21 pb_flag = line[1] 

22 if line[1] not in ['C', 'F', 'S']: 

23 if line[1] == 'H': 

24 raise OSError('Error in line #1: H (Helical) is valid but not ' 

25 'supported. Only C (Cluster), S (Supercell) ' 

26 'or F (Fraction) are supported options') 

27 else: 

28 raise OSError('Error in line #1: only C (Cluster), S (Supercell) ' 

29 'or F (Fraction) are supported options') 

30 

31 # Read atomic symbols 

32 line = lines[1].split() 

33 # Define a dictionary with symbols-id 

34 symboldict = {} 

35 symbolid = 1 

36 for symb in line: 

37 symboldict[symbolid] = symb 

38 symbolid += 1 

39 

40 # Read atoms (GEN format supports only single snapshot) 

41 del lines[:2] 

42 positions = [] 

43 symbols = [] 

44 for line in lines[:natoms]: 

45 dummy, symbolid, x, y, z = line.split()[:5] 

46 symbols.append(symboldict[int(symbolid)]) 

47 positions.append([float(x), float(y), float(z)]) 

48 image = Atoms(symbols=symbols, positions=positions) 

49 del lines[:natoms] 

50 

51 # If Supercell, parse periodic vectors. 

52 # If Fraction, translate into Supercell. 

53 if pb_flag == 'C': 

54 return image 

55 else: 

56 # Dummy line: line after atom positions is not uniquely defined 

57 # in gen implementations, and not necessary in DFTB package 

58 del lines[:1] 

59 image.set_pbc([True, True, True]) 

60 p = [] 

61 for i in range(3): 

62 x, y, z = lines[i].split()[:3] 

63 p.append([float(x), float(y), float(z)]) 

64 image.set_cell([(p[0][0], p[0][1], p[0][2]), 

65 (p[1][0], p[1][1], p[1][2]), 

66 (p[2][0], p[2][1], p[2][2])]) 

67 if pb_flag == 'F': 

68 frac_positions = image.get_positions() 

69 image.set_scaled_positions(frac_positions) 

70 return image 

71 

72 

73@writer 

74def write_gen( 

75 fileobj, 

76 images: Union[Atoms, Sequence[Atoms]], 

77 fractional: bool = False, 

78): 

79 """Write structure in GEN format (refer to DFTB+ manual). 

80 Multiple snapshots are not allowed. """ 

81 if isinstance(images, (list, tuple)): 

82 # GEN format doesn't support multiple snapshots 

83 if len(images) != 1: 

84 raise ValueError( 

85 '"images" contains more than one structure. ' 

86 'GEN format supports only single snapshot output.' 

87 ) 

88 atoms = images[0] 

89 else: 

90 atoms = images 

91 

92 symbols = atoms.get_chemical_symbols() 

93 

94 # Define a dictionary with symbols-id 

95 symboldict: Dict[str, int] = {} 

96 for sym in symbols: 

97 if sym not in symboldict: 

98 symboldict[sym] = len(symboldict) + 1 

99 # An ordered symbol list is needed as ordered dictionary 

100 # is just available in python 2.7 

101 orderedsymbols = list(['null'] * len(symboldict.keys())) 

102 for sym, num in symboldict.items(): 

103 orderedsymbols[num - 1] = sym 

104 

105 # Check whether the structure is periodic 

106 # GEN cannot describe periodicity in one or two direction, 

107 # a periodic structure is considered periodic in all the 

108 # directions. If your structure is not periodical in all 

109 # the directions, be sure you have set big periodicity 

110 # vectors in the non-periodic directions 

111 if fractional: 

112 pb_flag = 'F' 

113 elif atoms.pbc.any(): 

114 pb_flag = 'S' 

115 else: 

116 pb_flag = 'C' 

117 

118 natoms = len(symbols) 

119 ind = 0 

120 

121 fileobj.write(f'{natoms:d} {pb_flag:<5s}\n') 

122 for sym in orderedsymbols: 

123 fileobj.write(f'{sym:<5s}') 

124 fileobj.write('\n') 

125 

126 if fractional: 

127 coords = atoms.get_scaled_positions(wrap=False) 

128 else: 

129 coords = atoms.get_positions(wrap=False) 

130 

131 for sym, (x, y, z) in zip(symbols, coords): 

132 ind += 1 

133 symbolid = symboldict[sym] 

134 fileobj.write( 

135 f'{ind:-6d} {symbolid:d} {x:22.15f} {y:22.15f} {z:22.15f}\n') 

136 

137 if atoms.pbc.any() or fractional: 

138 fileobj.write(f'{0.0:22.15f} {0.0:22.15f} {0.0:22.15f} \n') 

139 cell = atoms.get_cell() 

140 for i in range(3): 

141 for j in range(3): 

142 fileobj.write(f'{cell[i, j]:22.15f} ') 

143 fileobj.write('\n')