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

1import numpy as np 

2 

3from ase.atoms import Atoms 

4from ase.calculators.singlepoint import SinglePointCalculator 

5from ase.data import atomic_numbers 

6from ase.units import Hartree 

7from ase.utils import reader, writer 

8 

9 

10@writer 

11def write_xsf(fileobj, images, data=None, origin=None, span_vectors=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(f'PRIMVEC{anim_token}\n') 

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(f'PRIMCOORD{anim_token}\n') 

54 else: 

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

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 if origin is None: 

94 origin = np.zeros(3) 

95 for i in range(3): 

96 if not pbc[i]: 

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

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

99 

100 for i in range(3): 

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

102 # What's with the strange division? 

103 # This disagrees with the output of Octopus. Investigate 

104 if span_vectors is None: 

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

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

107 else: 

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

109 

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

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

112 fileobj.write(' ') 

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

114 fileobj.write('\n') 

115 fileobj.write('\n') 

116 

117 fileobj.write(' END_DATAGRID_3D\n') 

118 fileobj.write('END_BLOCK_DATAGRID_3D\n') 

119 

120 

121@reader 

122def iread_xsf(fileobj, read_data=False): 

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

124 

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

126 span_vectors]. 

127 

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

129 

130 It also returns the origin of the simulation box 

131 as a numpy array and its spanning vectors as a 

132 list of numpy arrays, if data is returned. 

133 

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

135 def _line_generator_func(): 

136 for line in fileobj: 

137 line = line.strip() 

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

139 continue # Discard comments and empty lines 

140 yield line 

141 

142 _line_generator = _line_generator_func() 

143 

144 def readline(): 

145 return next(_line_generator) 

146 

147 line = readline() 

148 

149 if line.startswith('ANIMSTEPS'): 

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

151 line = readline() 

152 else: 

153 nimages = 1 

154 

155 if line == 'CRYSTAL': 

156 pbc = (True, True, True) 

157 elif line == 'SLAB': 

158 pbc = (True, True, False) 

159 elif line == 'POLYMER': 

160 pbc = (True, False, False) 

161 else: 

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

163 pbc = (False, False, False) 

164 

165 cell = None 

166 for n in range(nimages): 

167 if any(pbc): 

168 line = readline() 

169 if line.startswith('PRIMCOORD'): 

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

171 else: 

172 assert line.startswith('PRIMVEC') 

173 cell = [] 

174 for i in range(3): 

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

176 

177 line = readline() 

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

179 for i in range(3): 

180 readline() 

181 line = readline() 

182 

183 assert line.startswith('PRIMCOORD') 

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

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

186 else: 

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

188 line = readline() 

189 lines = [] 

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

191 lines.append(line) 

192 try: 

193 line = readline() 

194 except StopIteration: 

195 break 

196 if line.startswith('BEGIN'): 

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

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

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

200 # We will remember the line until later then. 

201 data_header_line = line 

202 

203 numbers = [] 

204 positions = [] 

205 for positionline in lines: 

206 tokens = positionline.split() 

207 symbol = tokens[0] 

208 if symbol.isdigit(): 

209 numbers.append(int(symbol)) 

210 else: 

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

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

213 

214 positions = np.array(positions) 

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

216 forces = None 

217 else: 

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

219 positions = positions[:, :3] 

220 

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

222 

223 if forces is not None: 

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

225 yield image 

226 

227 if read_data: 

228 if any(pbc): 

229 line = readline() 

230 else: 

231 line = data_header_line 

232 assert line.startswith('BEGIN_BLOCK_DATAGRID_3D') 

233 readline() # name 

234 line = readline() 

235 assert line.startswith('BEGIN_DATAGRID_3D') 

236 

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

238 assert len(shape) == 3 

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

240 origin = np.array(origin) 

241 

242 span_vectors = [] 

243 for i in range(3): 

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

245 span_vector = np.array(span_vector) 

246 span_vectors.append(span_vector) 

247 span_vectors = np.array(span_vectors) 

248 assert len(span_vectors) == len(shape) 

249 

250 npoints = np.prod(shape) 

251 

252 data = [] 

253 line = readline() # First line of data 

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

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

256 line = readline() 

257 assert len(data) == npoints 

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

259 # Note that data array is Fortran-ordered 

260 yield data, origin, span_vectors 

261 

262 

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

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

265 if read_data: 

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

267 images = images[:-1] 

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

269 return images[index]