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

7 

8from ase.atoms import Atoms 

9from ase.utils import reader, writer 

10 

11 

12@reader 

13def read_gen(fileobj): 

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

15 Multiple snapshot are not allowed. """ 

16 image = Atoms() 

17 lines = fileobj.readlines() 

18 line = lines[0].split() 

19 natoms = int(line[0]) 

20 pb_flag = line[1] 

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

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

23 'or F (Fraction) are valid options') 

24 

25 # Read atomic symbols 

26 line = lines[1].split() 

27 # Define a dictionary with symbols-id 

28 symboldict = dict() 

29 symbolid = 1 

30 for symb in line: 

31 symboldict[symbolid] = symb 

32 symbolid += 1 

33 

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

35 del lines[:2] 

36 positions = [] 

37 symbols = [] 

38 for line in lines[:natoms]: 

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

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

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

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

43 del lines[:natoms] 

44 

45 # If Supercell, parse periodic vectors. 

46 # If Fraction, translate into Supercell. 

47 if pb_flag == 'C': 

48 return image 

49 else: 

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

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

52 del lines[:1] 

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

54 p = [] 

55 for i in range(3): 

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

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

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

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

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

61 if pb_flag == 'F': 

62 frac_positions = image.get_positions() 

63 image.set_scaled_positions(frac_positions) 

64 return image 

65 

66 

67@writer 

68def write_gen(fileobj, images): 

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

70 Multiple snapshots are not allowed. """ 

71 if not isinstance(images, (list, tuple)): 

72 images = [images] 

73 

74 # Images is kept in a list but a size > 0 is not allowed 

75 # as GEN format doesn't support multiple snapshots. 

76 # Images is used as a list for consistency with the other 

77 # output modules 

78 if len(images) != 1: 

79 raise ValueError('images contains more than one structure\n' + 

80 'GEN format supports only single snapshot output') 

81 

82 symbols = images[0].get_chemical_symbols() 

83 

84 # Define a dictionary with symbols-id 

85 symboldict = dict() 

86 for sym in symbols: 

87 if not (sym in symboldict): 

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

89 # An ordered symbol list is needed as ordered dictionary 

90 # is just available in python 2.7 

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

92 for sym in symboldict.keys(): 

93 orderedsymbols[symboldict[sym] - 1] = sym 

94 

95 # Check whether the structure is periodic 

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

97 # a periodic structure is considered periodic in all the 

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

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

100 # vectors in the non-periodic directions 

101 if images[0].pbc.any(): 

102 pb_flag = 'S' 

103 else: 

104 pb_flag = 'C' 

105 

106 natoms = len(symbols) 

107 ind = 0 

108 for atoms in images: 

109 fileobj.write('%d %-5s\n' % (natoms, pb_flag)) 

110 for s in orderedsymbols: 

111 fileobj.write('%-5s' % s) 

112 fileobj.write('\n') 

113 for sym, (x, y, z) in zip(symbols, atoms.get_positions()): 

114 ind += 1 

115 symbolid = symboldict[sym] 

116 fileobj.write('%-6d %d %22.15f %22.15f %22.15f\n' % (ind, 

117 symbolid, x, y, z)) 

118 if images[0].pbc.any(): 

119 fileobj.write('%22.15f %22.15f %22.15f \n' % (0.0, 0.0, 0.0)) 

120 fileobj.write('%22.15f %22.15f %22.15f \n' % 

121 (images[0].get_cell()[0][0], 

122 images[0].get_cell()[0][1], 

123 images[0].get_cell()[0][2])) 

124 fileobj.write('%22.15f %22.15f %22.15f \n' % 

125 (images[0].get_cell()[1][0], 

126 images[0].get_cell()[1][1], 

127 images[0].get_cell()[1][2])) 

128 fileobj.write('%22.15f %22.15f %22.15f \n' % 

129 (images[0].get_cell()[2][0], 

130 images[0].get_cell()[2][1], 

131 images[0].get_cell()[2][2]))