Coverage for ase / io / gen.py: 93.33%

1# fmt: off 

2 

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

4 

5Refer to DFTB+ manual for GEN format description. 

6 

7Note: GEN format only supports single snapshot. 

8""" 

9import re 

10from typing import Dict, Sequence 

11 

12from ase.atoms import Atoms 

13from ase.utils import reader, writer 

14 

15 

16@reader 

17def read_gen(fileobj): 

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

19 Multiple snapshot are not allowed. """ 

20 image = Atoms() 

21 lines = fileobj.readlines() 

22 line = lines[0].split() 

23 natoms = int(line[0]) 

24 pb_flag = line[1] 

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

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

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

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

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

30 else: 

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

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

33 

34 # Read atomic symbols 

35 line = lines[1].split() 

36 symboldict = {symbolid: symb for symbolid, symb in enumerate(line, start=1)} 

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

38 del lines[:2] 

39 positions = [] 

40 symbols = [] 

41 tags = [] 

42 for line in lines[:natoms]: 

43 _dummy, symbolid, x, y, z = line.split()[:5] 

44 symbol = symboldict[int(symbolid)] 

45 parts = re.findall(r'\d+|[a-zA-Z]+', symbol) 

46 if len(parts) == 1: 

47 symbols.append(parts[0]) 

48 else: 

49 # If symbolid contains digits, it is a tag 

50 symbols.append(parts[0]) 

51 tags.append(int(parts[1])) 

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

53 if not tags: 

54 tags = None 

55 image = Atoms(symbols=symbols, positions=positions, tags=tags) 

56 del lines[:natoms] 

57 

58 # If Supercell, parse periodic vectors. 

59 # If Fraction, translate into Supercell. 

60 if pb_flag == 'C': 

61 return image 

62 else: 

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

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

65 del lines[:1] 

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

67 p = [] 

68 for i in range(3): 

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

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

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

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

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

74 if pb_flag == 'F': 

75 frac_positions = image.get_positions() 

76 image.set_scaled_positions(frac_positions) 

77 return image 

78 

79 

80@writer 

81def write_gen( 

82 fileobj, 

83 images: Atoms | Sequence[Atoms], 

84 fractional: bool = False, 

85 with_tags=False 

86): 

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

88 Multiple snapshots are not allowed. """ 

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

90 # GEN format doesn't support multiple snapshots 

91 if len(images) != 1: 

92 raise ValueError( 

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

94 'GEN format supports only single snapshot output.' 

95 ) 

96 atoms = images[0] 

97 else: 

98 atoms = images 

99 

100 symbols = atoms.get_chemical_symbols() 

101 tags = atoms.get_tags() 

102 

103 # Define a dictionary with symbols-id 

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

105 for sym, tag in zip(symbols, tags): 

106 sym = f'{sym}{tag}' if with_tags else sym 

107 if sym not in symboldict: 

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

109 # An ordered symbol list is needed as ordered dictionary 

110 # is just available in python 2.7 

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

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

113 orderedsymbols[num - 1] = sym 

114 

115 # Check whether the structure is periodic 

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

117 # a periodic structure is considered periodic in all the 

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

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

120 # vectors in the non-periodic directions 

121 if fractional: 

122 pb_flag = 'F' 

123 elif atoms.pbc.any(): 

124 pb_flag = 'S' 

125 else: 

126 pb_flag = 'C' 

127 

128 natoms = len(symbols) 

129 ind = 0 

130 

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

132 for sym in orderedsymbols: 

133 fileobj.write(f'{sym:<7s}') 

134 fileobj.write('\n') 

135 

136 if fractional: 

137 coords = atoms.get_scaled_positions(wrap=False) 

138 else: 

139 coords = atoms.get_positions(wrap=False) 

140 

141 for sym, tag, (x, y, z) in zip(symbols, tags, coords): 

142 ind += 1 

143 sym = f'{sym}{tag}' if with_tags else sym 

144 symbolid = symboldict[sym] 

145 fileobj.write( 

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

147 

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

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

150 cell = atoms.get_cell() 

151 for i in range(3): 

152 for j in range(3): 

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

154 fileobj.write('\n')