Coverage for /builds/ase/ase/ase/io/gen.py: 92.31%

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""" 

9from typing import Dict, Sequence, Union 

10 

11from ase.atoms import Atoms 

12from ase.utils import reader, writer 

13 

14 

15@reader 

16def read_gen(fileobj): 

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

18 Multiple snapshot are not allowed. """ 

19 image = Atoms() 

20 lines = fileobj.readlines() 

21 line = lines[0].split() 

22 natoms = int(line[0]) 

23 pb_flag = line[1] 

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

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

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

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

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

29 else: 

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

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

32 

33 # Read atomic symbols 

34 line = lines[1].split() 

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

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

37 del lines[:2] 

38 positions = [] 

39 symbols = [] 

40 for line in lines[:natoms]: 

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

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

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

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

45 del lines[:natoms] 

46 

47 # If Supercell, parse periodic vectors. 

48 # If Fraction, translate into Supercell. 

49 if pb_flag == 'C': 

50 return image 

51 else: 

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

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

54 del lines[:1] 

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

56 p = [] 

57 for i in range(3): 

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

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

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

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

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

63 if pb_flag == 'F': 

64 frac_positions = image.get_positions() 

65 image.set_scaled_positions(frac_positions) 

66 return image 

67 

68 

69@writer 

70def write_gen( 

71 fileobj, 

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

73 fractional: bool = False, 

74): 

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

76 Multiple snapshots are not allowed. """ 

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

78 # GEN format doesn't support multiple snapshots 

79 if len(images) != 1: 

80 raise ValueError( 

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

82 'GEN format supports only single snapshot output.' 

83 ) 

84 atoms = images[0] 

85 else: 

86 atoms = images 

87 

88 symbols = atoms.get_chemical_symbols() 

89 

90 # Define a dictionary with symbols-id 

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

92 for sym in symbols: 

93 if sym not in symboldict: 

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

95 # An ordered symbol list is needed as ordered dictionary 

96 # is just available in python 2.7 

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

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

99 orderedsymbols[num - 1] = sym 

100 

101 # Check whether the structure is periodic 

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

103 # a periodic structure is considered periodic in all the 

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

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

106 # vectors in the non-periodic directions 

107 if fractional: 

108 pb_flag = 'F' 

109 elif atoms.pbc.any(): 

110 pb_flag = 'S' 

111 else: 

112 pb_flag = 'C' 

113 

114 natoms = len(symbols) 

115 ind = 0 

116 

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

118 for sym in orderedsymbols: 

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

120 fileobj.write('\n') 

121 

122 if fractional: 

123 coords = atoms.get_scaled_positions(wrap=False) 

124 else: 

125 coords = atoms.get_positions(wrap=False) 

126 

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

128 ind += 1 

129 symbolid = symboldict[sym] 

130 fileobj.write( 

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

132 

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

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

135 cell = atoms.get_cell() 

136 for i in range(3): 

137 for j in range(3): 

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

139 fileobj.write('\n')