1import numpy as np 

2import xml.etree.ElementTree as ET 

3from xml.dom import minidom 

4 

5from ase.io.xsd import SetChild, _write_xsd_html 

6from ase import Atoms 

7 

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

9_image_footer = 'end\nend\n' 

10 

11 

12def _get_atom_str(an, xyz): 

13 s = '{:<5}'.format(an) 

14 s += '{:>15.9f}{:>15.9f}{:>15.9f}'.format(xyz[0], xyz[1], xyz[2]) 

15 s += ' XXXX 1 xx ' 

16 s += '{:<2}'.format(an) 

17 s += ' 0.000\n' 

18 return s 

19 

20 

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

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

23 atoms: Atoms object 

24 filename: path of the output file 

25 moviespeed: speed of animation. between 0 and 10 

26 

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

28 perodic system was inputted, full periodicity was assumed. 

29 """ 

30 if moviespeed < 0 or moviespeed > 10: 

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

32 

33 if hasattr(images, 'get_positions'): 

34 images = [images] 

35 

36 XSD, ATR = _write_xsd_html(images, connectivity) 

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

38 natoms = len(images[0]) 

39 

40 bonds = list() 

41 if connectivity is not None: 

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

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

44 if connectivity[i, j]: 

45 bonds.append([i, j]) 

46 

47 # non-periodic system 

48 s = '!BIOSYM archive 3\n' 

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

50 # Write trajectory 

51 SetChild(ATR, 'Trajectory', dict( 

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

53 Increment='-1', 

54 End=str(len(images)), 

55 Type='arc', 

56 Speed=str(moviespeed), 

57 FrameClassType='Atom', 

58 )) 

59 

60 # write frame information file 

61 s += 'PBC=OFF\n' 

62 for image in images: 

63 s += _image_header 

64 s += '\n' 

65 an = image.get_chemical_symbols() 

66 xyz = image.get_positions() 

67 for i in range(natoms): 

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

69 s += _image_footer 

70 

71 # periodic system 

72 else: 

73 SetChild(ATR, 'Trajectory', dict( 

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

75 Increment='-1', 

76 End=str(len(images)), 

77 Type='arc', 

78 Speed=str(moviespeed), 

79 FrameClassType='Atom', 

80 )) 

81 

82 # write frame information file 

83 s += 'PBC=ON\n' 

84 for image in images: 

85 s += _image_header 

86 s += 'PBC' 

87 vec = image.cell.lengths() 

88 s += '{:>10.4f}{:>10.4f}{:>10.4f}'.format(vec[0], vec[1], vec[2]) 

89 angles = image.cell.angles() 

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

91 s += '\n' 

92 an = image.get_chemical_symbols() 

93 

94 angrad = np.deg2rad(angles) 

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

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

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

98 * vec[1]) 

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

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

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

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

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

104 for i in range(natoms): 

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

106 s += _image_footer 

107 

108 # print arc file 

109 if isinstance(filename, str): 

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

111 else: 

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

113 

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

115 farc.write(s) 

116 

117 # check if file is an object or not. 

118 openandclose = False 

119 try: 

120 if isinstance(filename, str): 

121 fd = open(filename, 'w') 

122 openandclose = True 

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

124 fd = filename 

125 

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

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

128 reparsed = minidom.parseString(rough_string) 

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

130 

131 # write 

132 fd.write(Document) 

133 finally: 

134 if openandclose: 

135 fd.close() 

136 

137 

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

139 """Import xtd file (Materials Studio) 

140 

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

142 contains all the relevant information to make atoms 

143 so only Arc file needs to be read 

144 """ 

145 if isinstance(filename, str): 

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

147 else: 

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

149 

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

151 # common sense. 

152 with open(arcfilename, 'r') as fd: 

153 return read_arcfile(fd, index) 

154 

155 

156def read_arcfile(fd, index): 

157 images = [] 

158 

159 # the first line is comment 

160 fd.readline() 

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

162 L = fd.readline() 

163 while L != '': 

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

165 L = fd.readline() 

166 continue 

167 if pbc: 

168 L = fd.readline() 

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

170 else: 

171 fd.readline() 

172 symbols = [] 

173 coords = [] 

174 while True: 

175 line = fd.readline() 

176 L = line.split() 

177 if not line or 'end' in L: 

178 break 

179 symbols.append(L[0]) 

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

181 if pbc: 

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

183 else: 

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

185 images.append(image) 

186 L = fd.readline() 

187 

188 if not index: 

189 return images 

190 else: 

191 return images[index]