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

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.xsd import SetChild, _write_xsd_html 

10 

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

12_image_footer = 'end\nend\n' 

13 

14 

15def _get_atom_str(an, xyz): 

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

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

18 s += ' XXXX 1 xx ' 

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

20 s += ' 0.000\n' 

21 return s 

22 

23 

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

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

26 atoms: Atoms object 

27 filename: path of the output file 

28 moviespeed: speed of animation. between 0 and 10 

29 

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

31 perodic system was inputted, full periodicity was assumed. 

32 """ 

33 if moviespeed < 0 or moviespeed > 10: 

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

35 

36 if hasattr(images, 'get_positions'): 

37 images = [images] 

38 

39 XSD, ATR = _write_xsd_html(images, connectivity) 

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

41 natoms = len(images[0]) 

42 

43 bonds = [] 

44 if connectivity is not None: 

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

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

47 if connectivity[i, j]: 

48 bonds.append([i, j]) 

49 

50 # non-periodic system 

51 s = '!BIOSYM archive 3\n' 

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

53 # Write trajectory 

54 SetChild(ATR, 'Trajectory', dict( 

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

56 Increment='-1', 

57 End=str(len(images)), 

58 Type='arc', 

59 Speed=str(moviespeed), 

60 FrameClassType='Atom', 

61 )) 

62 

63 # write frame information file 

64 s += 'PBC=OFF\n' 

65 for image in images: 

66 s += _image_header 

67 s += '\n' 

68 an = image.get_chemical_symbols() 

69 xyz = image.get_positions() 

70 for i in range(natoms): 

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

72 s += _image_footer 

73 

74 # periodic system 

75 else: 

76 SetChild(ATR, 'Trajectory', dict( 

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

78 Increment='-1', 

79 End=str(len(images)), 

80 Type='arc', 

81 Speed=str(moviespeed), 

82 FrameClassType='Atom', 

83 )) 

84 

85 # write frame information file 

86 s += 'PBC=ON\n' 

87 for image in images: 

88 s += _image_header 

89 s += 'PBC' 

90 vec = image.cell.lengths() 

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

92 angles = image.cell.angles() 

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

94 s += '\n' 

95 an = image.get_chemical_symbols() 

96 

97 angrad = np.deg2rad(angles) 

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

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

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

101 * vec[1]) 

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

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

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

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

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

107 for i in range(natoms): 

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

109 s += _image_footer 

110 

111 # print arc file 

112 if isinstance(filename, str): 

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

114 else: 

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

116 

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

118 farc.write(s) 

119 

120 # check if file is an object or not. 

121 openandclose = False 

122 try: 

123 if isinstance(filename, str): 

124 fd = open(filename, 'w') 

125 openandclose = True 

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

127 fd = filename 

128 

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

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

131 reparsed = minidom.parseString(rough_string) 

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

133 

134 # write 

135 fd.write(Document) 

136 finally: 

137 if openandclose: 

138 fd.close() 

139 

140 

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

142 """Import xtd file (Materials Studio) 

143 

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

145 contains all the relevant information to make atoms 

146 so only Arc file needs to be read 

147 """ 

148 if isinstance(filename, str): 

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

150 else: 

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

152 

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

154 # common sense. 

155 with open(arcfilename) as fd: 

156 return read_arcfile(fd, index) 

157 

158 

159def read_arcfile(fd, index): 

160 images = [] 

161 

162 # the first line is comment 

163 fd.readline() 

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

165 L = fd.readline() 

166 while L != '': 

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

168 L = fd.readline() 

169 continue 

170 if pbc: 

171 L = fd.readline() 

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

173 else: 

174 fd.readline() 

175 symbols = [] 

176 coords = [] 

177 while True: 

178 line = fd.readline() 

179 L = line.split() 

180 if not line or 'end' in L: 

181 break 

182 symbols.append(L[0]) 

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

184 if pbc: 

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

186 else: 

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

188 images.append(image) 

189 L = fd.readline() 

190 

191 if not index: 

192 return images 

193 else: 

194 return images[index]