1import numpy as np 

2 

3from ase.atoms import Atoms 

4from ase.units import Hartree 

5from ase.data import atomic_numbers 

6from ase.calculators.singlepoint import SinglePointCalculator 

7from ase.utils import writer, reader 

8 

9 

10@writer 

11def write_xsf(fileobj, images, data=None): 

12 is_anim = len(images) > 1 

13 

14 if is_anim: 

15 fileobj.write('ANIMSTEPS %d\n' % len(images)) 

16 

17 numbers = images[0].get_atomic_numbers() 

18 

19 pbc = images[0].get_pbc() 

20 npbc = sum(pbc) 

21 if pbc[2]: 

22 fileobj.write('CRYSTAL\n') 

23 assert npbc == 3 

24 elif pbc[1]: 

25 fileobj.write('SLAB\n') 

26 assert npbc == 2 

27 elif pbc[0]: 

28 fileobj.write('POLYMER\n') 

29 assert npbc == 1 

30 else: 

31 # (Header written as part of image loop) 

32 assert npbc == 0 

33 

34 cell_variable = False 

35 for image in images[1:]: 

36 if np.abs(images[0].cell - image.cell).max() > 1e-14: 

37 cell_variable = True 

38 break 

39 

40 for n, atoms in enumerate(images): 

41 anim_token = ' %d' % (n + 1) if is_anim else '' 

42 if pbc.any(): 

43 write_cell = (n == 0 or cell_variable) 

44 if write_cell: 

45 if cell_variable: 

46 fileobj.write('PRIMVEC%s\n' % anim_token) 

47 else: 

48 fileobj.write('PRIMVEC\n') 

49 cell = atoms.get_cell() 

50 for i in range(3): 

51 fileobj.write(' %.14f %.14f %.14f\n' % tuple(cell[i])) 

52 

53 fileobj.write('PRIMCOORD%s\n' % anim_token) 

54 else: 

55 fileobj.write('ATOMS%s\n' % anim_token) 

56 

57 # Get the forces if it's not too expensive: 

58 calc = atoms.calc 

59 if (calc is not None and 

60 (hasattr(calc, 'calculation_required') and 

61 not calc.calculation_required(atoms, ['forces']))): 

62 forces = atoms.get_forces() / Hartree 

63 else: 

64 forces = None 

65 

66 pos = atoms.get_positions() 

67 

68 if pbc.any(): 

69 fileobj.write(' %d 1\n' % len(pos)) 

70 for a in range(len(pos)): 

71 fileobj.write(' %2d' % numbers[a]) 

72 fileobj.write(' %20.14f %20.14f %20.14f' % tuple(pos[a])) 

73 if forces is None: 

74 fileobj.write('\n') 

75 else: 

76 fileobj.write(' %20.14f %20.14f %20.14f\n' % tuple(forces[a])) 

77 

78 if data is None: 

79 return 

80 

81 fileobj.write('BEGIN_BLOCK_DATAGRID_3D\n') 

82 fileobj.write(' data\n') 

83 fileobj.write(' BEGIN_DATAGRID_3Dgrid#1\n') 

84 

85 data = np.asarray(data) 

86 if data.dtype == complex: 

87 data = np.abs(data) 

88 

89 shape = data.shape 

90 fileobj.write(' %d %d %d\n' % shape) 

91 

92 cell = atoms.get_cell() 

93 origin = np.zeros(3) 

94 for i in range(3): 

95 if not pbc[i]: 

96 origin += cell[i] / shape[i] 

97 fileobj.write(' %f %f %f\n' % tuple(origin)) 

98 

99 for i in range(3): 

100 # XXXX is this not just supposed to be the cell? 

101 # What's with the strange division? 

102 # This disagrees with the output of Octopus. Investigate 

103 fileobj.write(' %f %f %f\n' % 

104 tuple(cell[i] * (shape[i] + 1) / shape[i])) 

105 

106 for k in range(shape[2]): 

107 for j in range(shape[1]): 

108 fileobj.write(' ') 

109 fileobj.write(' '.join(['%f' % d for d in data[:, j, k]])) 

110 fileobj.write('\n') 

111 fileobj.write('\n') 

112 

113 fileobj.write(' END_DATAGRID_3D\n') 

114 fileobj.write('END_BLOCK_DATAGRID_3D\n') 

115 

116 

117@reader 

118def iread_xsf(fileobj, read_data=False): 

119 """Yield images and optionally data from xsf file. 

120 

121 Yields image1, image2, ..., imageN[, data]. 

122 

123 Images are Atoms objects and data is a numpy array. 

124 

125 Presently supports only a single 3D datagrid.""" 

126 def _line_generator_func(): 

127 for line in fileobj: 

128 line = line.strip() 

129 if not line or line.startswith('#'): 

130 continue # Discard comments and empty lines 

131 yield line 

132 

133 _line_generator = _line_generator_func() 

134 

135 def readline(): 

136 return next(_line_generator) 

137 

138 line = readline() 

139 

140 if line.startswith('ANIMSTEPS'): 

141 nimages = int(line.split()[1]) 

142 line = readline() 

143 else: 

144 nimages = 1 

145 

146 if line == 'CRYSTAL': 

147 pbc = (True, True, True) 

148 elif line == 'SLAB': 

149 pbc = (True, True, False) 

150 elif line == 'POLYMER': 

151 pbc = (True, False, False) 

152 else: 

153 assert line.startswith('ATOMS'), line # can also be ATOMS 1 

154 pbc = (False, False, False) 

155 

156 cell = None 

157 for n in range(nimages): 

158 if any(pbc): 

159 line = readline() 

160 if line.startswith('PRIMCOORD'): 

161 assert cell is not None # cell read from previous image 

162 else: 

163 assert line.startswith('PRIMVEC') 

164 cell = [] 

165 for i in range(3): 

166 cell.append([float(x) for x in readline().split()]) 

167 

168 line = readline() 

169 if line.startswith('CONVVEC'): # ignored; 

170 for i in range(3): 

171 readline() 

172 line = readline() 

173 

174 assert line.startswith('PRIMCOORD') 

175 natoms = int(readline().split()[0]) 

176 lines = [readline() for _ in range(natoms)] 

177 else: 

178 assert line.startswith('ATOMS'), line 

179 line = readline() 

180 lines = [] 

181 while not (line.startswith('ATOMS') or line.startswith('BEGIN')): 

182 lines.append(line) 

183 try: 

184 line = readline() 

185 except StopIteration: 

186 break 

187 if line.startswith('BEGIN'): 

188 # We read "too far" and accidentally got the header 

189 # of the data section. This happens only when parsing 

190 # ATOMS blocks, because one cannot infer their length. 

191 # We will remember the line until later then. 

192 data_header_line = line 

193 

194 numbers = [] 

195 positions = [] 

196 for positionline in lines: 

197 tokens = positionline.split() 

198 symbol = tokens[0] 

199 if symbol.isdigit(): 

200 numbers.append(int(symbol)) 

201 else: 

202 numbers.append(atomic_numbers[symbol.capitalize()]) 

203 positions.append([float(x) for x in tokens[1:]]) 

204 

205 positions = np.array(positions) 

206 if len(positions[0]) == 3: 

207 forces = None 

208 else: 

209 forces = positions[:, 3:] * Hartree 

210 positions = positions[:, :3] 

211 

212 image = Atoms(numbers, positions, cell=cell, pbc=pbc) 

213 

214 if forces is not None: 

215 image.calc = SinglePointCalculator(image, forces=forces) 

216 yield image 

217 

218 if read_data: 

219 if any(pbc): 

220 line = readline() 

221 else: 

222 line = data_header_line 

223 assert line.startswith('BEGIN_BLOCK_DATAGRID_3D') 

224 readline() # name 

225 line = readline() 

226 assert line.startswith('BEGIN_DATAGRID_3D') 

227 

228 shape = [int(x) for x in readline().split()] 

229 assert len(shape) == 3 

230 readline() # start 

231 # XXX what to do about these? 

232 

233 for i in range(3): 

234 readline() # Skip 3x3 matrix for some reason 

235 

236 npoints = np.prod(shape) 

237 

238 data = [] 

239 line = readline() # First line of data 

240 while not line.startswith('END_DATAGRID_3D'): 

241 data.extend([float(x) for x in line.split()]) 

242 line = readline() 

243 assert len(data) == npoints 

244 data = np.array(data, float).reshape(shape[::-1]).T 

245 # Note that data array is Fortran-ordered 

246 yield data 

247 

248 

249def read_xsf(fileobj, index=-1, read_data=False): 

250 images = list(iread_xsf(fileobj, read_data=read_data)) 

251 if read_data: 

252 array = images[-1] 

253 images = images[:-1] 

254 return array, images[index] 

255 return images[index]