Coverage for /builds/kinetik161/ase/ase/io/xtd.py: 67.83%

1import xml.etree.ElementTree as ET 

2from xml.dom import minidom 

3 

4import numpy as np 

5 

6from ase import Atoms 

7from ase.io.xsd import SetChild, _write_xsd_html 

8 

9_image_header = ' ' * 74 + '0.0000\n!DATE Jan 01 00:00:00 2000\n' 

10_image_footer = 'end\nend\n' 

11 

12 

13def _get_atom_str(an, xyz): 

14 s = f'{an:<5}' 

15 s += f'{xyz[0]:>15.9f}{xyz[1]:>15.9f}{xyz[2]:>15.9f}' 

16 s += ' XXXX 1 xx ' 

17 s += f'{an:<2}' 

18 s += ' 0.000\n' 

19 return s 

20 

21 

22def write_xtd(filename, images, connectivity=None, moviespeed=10): 

23 """Takes Atoms object, and write materials studio file 

24 atoms: Atoms object 

25 filename: path of the output file 

26 moviespeed: speed of animation. between 0 and 10 

27 

28 note: material studio file cannot use a partial periodic system. If partial 

29 perodic system was inputted, full periodicity was assumed. 

30 """ 

31 if moviespeed < 0 or moviespeed > 10: 

32 raise ValueError('moviespeed only between 0 and 10 allowed') 

33 

34 if hasattr(images, 'get_positions'): 

35 images = [images] 

36 

37 XSD, ATR = _write_xsd_html(images, connectivity) 

38 ATR.attrib['NumChildren'] = '2' 

39 natoms = len(images[0]) 

40 

41 bonds = [] 

42 if connectivity is not None: 

43 for i in range(connectivity.shape[0]): 

44 for j in range(i + 1, connectivity.shape[0]): 

45 if connectivity[i, j]: 

46 bonds.append([i, j]) 

47 

48 # non-periodic system 

49 s = '!BIOSYM archive 3\n' 

50 if not images[0].pbc.all(): 

51 # Write trajectory 

52 SetChild(ATR, 'Trajectory', dict( 

53 ID=str(natoms + 3 + len(bonds)), 

54 Increment='-1', 

55 End=str(len(images)), 

56 Type='arc', 

57 Speed=str(moviespeed), 

58 FrameClassType='Atom', 

59 )) 

60 

61 # write frame information file 

62 s += 'PBC=OFF\n' 

63 for image in images: 

64 s += _image_header 

65 s += '\n' 

66 an = image.get_chemical_symbols() 

67 xyz = image.get_positions() 

68 for i in range(natoms): 

69 s += _get_atom_str(an[i], xyz[i, :]) 

70 s += _image_footer 

71 

72 # periodic system 

73 else: 

74 SetChild(ATR, 'Trajectory', dict( 

75 ID=str(natoms + 9 + len(bonds)), 

76 Increment='-1', 

77 End=str(len(images)), 

78 Type='arc', 

79 Speed=str(moviespeed), 

80 FrameClassType='Atom', 

81 )) 

82 

83 # write frame information file 

84 s += 'PBC=ON\n' 

85 for image in images: 

86 s += _image_header 

87 s += 'PBC' 

88 vec = image.cell.lengths() 

89 s += f'{vec[0]:>10.4f}{vec[1]:>10.4f}{vec[2]:>10.4f}' 

90 angles = image.cell.angles() 

91 s += '{:>10.4f}{:>10.4f}{:>10.4f}'.format(*angles) 

92 s += '\n' 

93 an = image.get_chemical_symbols() 

94 

95 angrad = np.deg2rad(angles) 

96 cell = np.zeros((3, 3)) 

97 cell[0, :] = [vec[0], 0, 0] 

98 cell[1, :] = (np.array([np.cos(angrad[2]), np.sin(angrad[2]), 0]) 

99 * vec[1]) 

100 cell[2, 0] = vec[2] * np.cos(angrad[1]) 

101 cell[2, 1] = ((vec[1] * vec[2] * np.cos(angrad[0]) 

102 - cell[1, 0] * cell[2, 0]) / cell[1, 1]) 

103 cell[2, 2] = np.sqrt(vec[2]**2 - cell[2, 0]**2 - cell[2, 1]**2) 

104 xyz = np.dot(image.get_scaled_positions(), cell) 

105 for i in range(natoms): 

106 s += _get_atom_str(an[i], xyz[i, :]) 

107 s += _image_footer 

108 

109 # print arc file 

110 if isinstance(filename, str): 

111 farcname = filename[:-3] + 'arc' 

112 else: 

113 farcname = filename.name[:-3] + 'arc' 

114 

115 with open(farcname, 'w') as farc: 

116 farc.write(s) 

117 

118 # check if file is an object or not. 

119 openandclose = False 

120 try: 

121 if isinstance(filename, str): 

122 fd = open(filename, 'w') 

123 openandclose = True 

124 else: # Assume it's a 'file-like object' 

125 fd = filename 

126 

127 # Return a pretty-printed XML string for the Element. 

128 rough_string = ET.tostring(XSD, 'utf-8') 

129 reparsed = minidom.parseString(rough_string) 

130 Document = reparsed.toprettyxml(indent='\t') 

131 

132 # write 

133 fd.write(Document) 

134 finally: 

135 if openandclose: 

136 fd.close() 

137 

138 

139def read_xtd(filename, index=-1): 

140 """Import xtd file (Materials Studio) 

141 

142 Xtd files always come with arc file, and arc file 

143 contains all the relevant information to make atoms 

144 so only Arc file needs to be read 

145 """ 

146 if isinstance(filename, str): 

147 arcfilename = filename[:-3] + 'arc' 

148 else: 

149 arcfilename = filename.name[:-3] + 'arc' 

150 

151 # This trick with opening a totally different file is a gross violation of 

152 # common sense. 

153 with open(arcfilename) as fd: 

154 return read_arcfile(fd, index) 

155 

156 

157def read_arcfile(fd, index): 

158 images = [] 

159 

160 # the first line is comment 

161 fd.readline() 

162 pbc = 'ON' in fd.readline() 

163 L = fd.readline() 

164 while L != '': 

165 if '!' not in L: # flag for the start of an image 

166 L = fd.readline() 

167 continue 

168 if pbc: 

169 L = fd.readline() 

170 cell = [float(d) for d in L.split()[1:]] 

171 else: 

172 fd.readline() 

173 symbols = [] 

174 coords = [] 

175 while True: 

176 line = fd.readline() 

177 L = line.split() 

178 if not line or 'end' in L: 

179 break 

180 symbols.append(L[0]) 

181 coords.append([float(x) for x in L[1:4]]) 

182 if pbc: 

183 image = Atoms(symbols, positions=coords, cell=cell, pbc=pbc) 

184 else: 

185 image = Atoms(symbols, positions=coords, pbc=pbc) 

186 images.append(image) 

187 L = fd.readline() 

188 

189 if not index: 

190 return images 

191 else: 

192 return images[index]