Coverage for /builds/ase/ase/ase/io/xsf.py: 96.32%

1# fmt: off 

2 

3import numpy as np 

4 

5from ase.atoms import Atoms 

6from ase.calculators.singlepoint import SinglePointCalculator 

7from ase.data import atomic_numbers 

8from ase.units import Hartree 

9from ase.utils import reader, writer 

10 

11 

12@writer 

13def write_xsf(fileobj, images, data=None, origin=None, span_vectors=None): 

14 is_anim = len(images) > 1 

15 

16 if is_anim: 

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

18 

19 numbers = images[0].get_atomic_numbers() 

20 

21 pbc = images[0].get_pbc() 

22 npbc = sum(pbc) 

23 if pbc[2]: 

24 fileobj.write('CRYSTAL\n') 

25 assert npbc == 3 

26 elif pbc[1]: 

27 fileobj.write('SLAB\n') 

28 assert npbc == 2 

29 elif pbc[0]: 

30 fileobj.write('POLYMER\n') 

31 assert npbc == 1 

32 else: 

33 # (Header written as part of image loop) 

34 assert npbc == 0 

35 

36 cell_variable = False 

37 for image in images[1:]: 

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

39 cell_variable = True 

40 break 

41 

42 for n, atoms in enumerate(images): 

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

44 if pbc.any(): 

45 write_cell = (n == 0 or cell_variable) 

46 if write_cell: 

47 if cell_variable: 

48 fileobj.write(f'PRIMVEC{anim_token}\n') 

49 else: 

50 fileobj.write('PRIMVEC\n') 

51 cell = atoms.get_cell() 

52 for i in range(3): 

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

54 

55 fileobj.write(f'PRIMCOORD{anim_token}\n') 

56 else: 

57 fileobj.write(f'ATOMS{anim_token}\n') 

58 

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

60 calc = atoms.calc 

61 if (calc is not None and 

62 (hasattr(calc, 'calculation_required') and 

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

64 forces = atoms.get_forces() / Hartree 

65 else: 

66 forces = None 

67 

68 pos = atoms.get_positions() 

69 

70 if pbc.any(): 

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

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

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

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

75 if forces is None: 

76 fileobj.write('\n') 

77 else: 

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

79 

80 if data is None: 

81 return 

82 

83 fileobj.write('BEGIN_BLOCK_DATAGRID_3D\n') 

84 fileobj.write(' data\n') 

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

86 

87 data = np.asarray(data) 

88 if data.dtype == complex: 

89 data = np.abs(data) 

90 

91 shape = data.shape 

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

93 

94 cell = atoms.get_cell() 

95 if origin is None: 

96 origin = np.zeros(3) 

97 for i in range(3): 

98 if not pbc[i]: 

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

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

101 

102 for i in range(3): 

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

104 # What's with the strange division? 

105 # This disagrees with the output of Octopus. Investigate 

106 if span_vectors is None: 

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

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

109 else: 

110 fileobj.write(' %f %f %f\n' % tuple(span_vectors[i])) 

111 

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

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

114 fileobj.write(' ') 

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

116 fileobj.write('\n') 

117 fileobj.write('\n') 

118 

119 fileobj.write(' END_DATAGRID_3D\n') 

120 fileobj.write('END_BLOCK_DATAGRID_3D\n') 

121 

122 

123@reader 

124def iread_xsf(fileobj, read_data=False): 

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

126 

127 Yields image1, image2, ..., imageN[, data, origin, 

128 span_vectors]. 

129 

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

131 

132 It also returns the origin of the simulation box 

133 as a numpy array and its spanning vectors as a 

134 list of numpy arrays, if data is returned. 

135 

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

137 def _line_generator_func(): 

138 for line in fileobj: 

139 line = line.strip() 

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

141 continue # Discard comments and empty lines 

142 yield line 

143 

144 _line_generator = _line_generator_func() 

145 

146 def readline(): 

147 return next(_line_generator) 

148 

149 line = readline() 

150 

151 if line.startswith('ANIMSTEPS'): 

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

153 line = readline() 

154 else: 

155 nimages = 1 

156 

157 if line == 'CRYSTAL': 

158 pbc = (True, True, True) 

159 elif line == 'SLAB': 

160 pbc = (True, True, False) 

161 elif line == 'POLYMER': 

162 pbc = (True, False, False) 

163 else: 

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

165 pbc = (False, False, False) 

166 

167 cell = None 

168 for n in range(nimages): 

169 if any(pbc): 

170 line = readline() 

171 if line.startswith('PRIMCOORD'): 

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

173 else: 

174 assert line.startswith('PRIMVEC') 

175 cell = [] 

176 for i in range(3): 

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

178 

179 line = readline() 

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

181 for i in range(3): 

182 readline() 

183 line = readline() 

184 

185 assert line.startswith('PRIMCOORD') 

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

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

188 else: 

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

190 line = readline() 

191 lines = [] 

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

193 lines.append(line) 

194 try: 

195 line = readline() 

196 except StopIteration: 

197 break 

198 if line.startswith('BEGIN'): 

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

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

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

202 # We will remember the line until later then. 

203 data_header_line = line 

204 

205 numbers = [] 

206 positions = [] 

207 for positionline in lines: 

208 tokens = positionline.split() 

209 symbol = tokens[0] 

210 if symbol.isdigit(): 

211 numbers.append(int(symbol)) 

212 else: 

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

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

215 

216 positions = np.array(positions) 

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

218 forces = None 

219 else: 

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

221 positions = positions[:, :3] 

222 

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

224 

225 if forces is not None: 

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

227 yield image 

228 

229 if read_data: 

230 if any(pbc): 

231 line = readline() 

232 else: 

233 line = data_header_line 

234 assert line.startswith('BEGIN_BLOCK_DATAGRID_3D') 

235 readline() # name 

236 line = readline() 

237 assert line.startswith('BEGIN_DATAGRID_3D') 

238 

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

240 assert len(shape) == 3 

241 origin = [float(x) for x in readline().split()] 

242 origin = np.array(origin) 

243 

244 span_vectors = [] 

245 for i in range(3): 

246 span_vector = [float(x) for x in readline().split()] 

247 span_vector = np.array(span_vector) 

248 span_vectors.append(span_vector) 

249 span_vectors = np.array(span_vectors) 

250 assert len(span_vectors) == len(shape) 

251 

252 npoints = np.prod(shape) 

253 

254 data = [] 

255 line = readline() # First line of data 

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

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

258 line = readline() 

259 assert len(data) == npoints 

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

261 # Note that data array is Fortran-ordered 

262 yield data, origin, span_vectors 

263 

264 

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

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

267 if read_data: 

268 array, origin, span_vectors = images[-1] 

269 images = images[:-1] 

270 return array, origin, span_vectors, images[index] 

271 return images[index]