Coverage for /builds/debichem-team/python-ase/ase/io/gen.py: 92.31%

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 symboldict = {symbolid: symb for symbolid, symb in enumerate(line, start=1)} 

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( 

69 fileobj, 

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

71 fractional: bool = False, 

72): 

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

74 Multiple snapshots are not allowed. """ 

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

76 # GEN format doesn't support multiple snapshots 

77 if len(images) != 1: 

78 raise ValueError( 

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

80 'GEN format supports only single snapshot output.' 

81 ) 

82 atoms = images[0] 

83 else: 

84 atoms = images 

85 

86 symbols = atoms.get_chemical_symbols() 

87 

88 # Define a dictionary with symbols-id 

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

90 for sym in symbols: 

91 if sym not in symboldict: 

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

93 # An ordered symbol list is needed as ordered dictionary 

94 # is just available in python 2.7 

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

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

97 orderedsymbols[num - 1] = sym 

98 

99 # Check whether the structure is periodic 

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

101 # a periodic structure is considered periodic in all the 

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

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

104 # vectors in the non-periodic directions 

105 if fractional: 

106 pb_flag = 'F' 

107 elif atoms.pbc.any(): 

108 pb_flag = 'S' 

109 else: 

110 pb_flag = 'C' 

111 

112 natoms = len(symbols) 

113 ind = 0 

114 

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

116 for sym in orderedsymbols: 

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

118 fileobj.write('\n') 

119 

120 if fractional: 

121 coords = atoms.get_scaled_positions(wrap=False) 

122 else: 

123 coords = atoms.get_positions(wrap=False) 

124 

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

126 ind += 1 

127 symbolid = symboldict[sym] 

128 fileobj.write( 

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

130 

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

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

133 cell = atoms.get_cell() 

134 for i in range(3): 

135 for j in range(3): 

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

137 fileobj.write('\n')