Coverage for ase / io / xtd.py: 69.91%

1# fmt: off 

2 

3import xml.etree.ElementTree as ET 

4from xml.dom import minidom 

5 

6import numpy as np 

7 

8from ase import Atoms 

9from ase.io.utils import connectivity2bonds 

10from ase.io.xsd import SetChild, _write_xsd_html 

11 

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

13_image_footer = 'end\nend\n' 

14 

15 

16def _get_atom_str(an, xyz): 

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

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

19 s += ' XXXX 1 xx ' 

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

21 s += ' 0.000\n' 

22 return s 

23 

24 

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

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

27 atoms: Atoms object 

28 filename: path of the output file 

29 moviespeed: speed of animation. between 0 and 10 

30 

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

32 perodic system was inputted, full periodicity was assumed. 

33 """ 

34 if moviespeed < 0 or moviespeed > 10: 

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

36 

37 if hasattr(images, 'get_positions'): 

38 images = [images] 

39 

40 XSD, ATR = _write_xsd_html(images, connectivity) 

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

42 natoms = len(images[0]) 

43 

44 if connectivity is not None: 

45 bonds = connectivity2bonds(connectivity) 

46 else: 

47 bonds = [] 

48 

49 # non-periodic system 

50 s = '!BIOSYM archive 3\n' 

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

52 # Write trajectory 

53 SetChild(ATR, 'Trajectory', dict( 

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

55 Increment='-1', 

56 End=str(len(images)), 

57 Type='arc', 

58 Speed=str(moviespeed), 

59 FrameClassType='Atom', 

60 )) 

61 

62 # write frame information file 

63 s += 'PBC=OFF\n' 

64 for image in images: 

65 s += _image_header 

66 s += '\n' 

67 an = image.get_chemical_symbols() 

68 xyz = image.get_positions() 

69 for i in range(natoms): 

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

71 s += _image_footer 

72 

73 # periodic system 

74 else: 

75 SetChild(ATR, 'Trajectory', dict( 

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

77 Increment='-1', 

78 End=str(len(images)), 

79 Type='arc', 

80 Speed=str(moviespeed), 

81 FrameClassType='Atom', 

82 )) 

83 

84 # write frame information file 

85 s += 'PBC=ON\n' 

86 for image in images: 

87 s += _image_header 

88 s += 'PBC' 

89 vec = image.cell.lengths() 

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

91 angles = image.cell.angles() 

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

93 s += '\n' 

94 an = image.get_chemical_symbols() 

95 

96 angrad = np.deg2rad(angles) 

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

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

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

100 * vec[1]) 

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

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

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

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

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

106 for i in range(natoms): 

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

108 s += _image_footer 

109 

110 # print arc file 

111 if isinstance(filename, str): 

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

113 else: 

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

115 

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

117 farc.write(s) 

118 

119 # check if file is an object or not. 

120 openandclose = False 

121 try: 

122 if isinstance(filename, str): 

123 fd = open(filename, 'w') 

124 openandclose = True 

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

126 fd = filename 

127 

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

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

130 reparsed = minidom.parseString(rough_string) 

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

132 

133 # write 

134 fd.write(Document) 

135 finally: 

136 if openandclose: 

137 fd.close() 

138 

139 

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

141 """Import xtd file (Materials Studio) 

142 

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

144 contains all the relevant information to make atoms 

145 so only Arc file needs to be read 

146 """ 

147 if isinstance(filename, str): 

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

149 else: 

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

151 

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

153 # common sense. 

154 with open(arcfilename) as fd: 

155 return read_arcfile(fd, index) 

156 

157 

158def read_arcfile(fd, index): 

159 images = [] 

160 

161 # the first line is comment 

162 fd.readline() 

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

164 L = fd.readline() 

165 while L != '': 

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

167 L = fd.readline() 

168 continue 

169 if pbc: 

170 L = fd.readline() 

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

172 else: 

173 fd.readline() 

174 symbols = [] 

175 coords = [] 

176 while True: 

177 line = fd.readline() 

178 L = line.split() 

179 if not line or 'end' in L: 

180 break 

181 symbols.append(L[0]) 

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

183 if pbc: 

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

185 else: 

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

187 images.append(image) 

188 L = fd.readline() 

189 

190 if not index: 

191 return images 

192 else: 

193 return images[index]