1from math import sqrt 

2 

3import numpy as np 

4 

5from ase import Atoms 

6from ase.calculators.singlepoint import SinglePointCalculator 

7from ase.constraints import FixAtoms 

8from ase.data import covalent_radii 

9from ase.gui.defaults import read_defaults 

10from ase.io import read, write, string2index 

11from ase.gui.i18n import _ 

12from ase.geometry import find_mic 

13 

14import warnings 

15 

16 

17class Images: 

18 def __init__(self, images=None): 

19 self.covalent_radii = covalent_radii.copy() 

20 self.config = read_defaults() 

21 self.atom_scale = self.config['radii_scale'] 

22 if images is None: 

23 images = [Atoms()] 

24 self.initialize(images) 

25 

26 def __len__(self): 

27 return len(self._images) 

28 

29 def __getitem__(self, index): 

30 return self._images[index] 

31 

32 def __iter__(self): 

33 return iter(self._images) 

34 

35 # XXXXXXX hack 

36 # compatibility hacks while allowing variable number of atoms 

37 def get_dynamic(self, atoms): 

38 dynamic = np.ones(len(atoms), bool) 

39 for constraint in atoms.constraints: 

40 if isinstance(constraint, FixAtoms): 

41 dynamic[constraint.index] = False 

42 return dynamic 

43 

44 def set_dynamic(self, mask, value): 

45 # Does not make much sense if different images have different 

46 # atom counts. Attempts to apply mask to all images, 

47 # to the extent possible. 

48 for atoms in self: 

49 dynamic = self.get_dynamic(atoms) 

50 dynamic[mask[:len(atoms)]] = value 

51 atoms.constraints = [c for c in atoms.constraints 

52 if not isinstance(c, FixAtoms)] 

53 atoms.constraints.append(FixAtoms(mask=~dynamic)) 

54 

55 def scale_radii(self, scaling_factor): 

56 self.covalent_radii *= scaling_factor 

57 

58 def get_energy(self, atoms): 

59 try: 

60 e = atoms.get_potential_energy() * self.repeat.prod() 

61 except RuntimeError: 

62 e = np.nan 

63 return e 

64 

65 def get_forces(self, atoms): 

66 try: 

67 F = atoms.get_forces(apply_constraint=False) 

68 except RuntimeError: 

69 return None 

70 else: 

71 return F 

72 

73 def initialize(self, images, filenames=None): 

74 nimages = len(images) 

75 if filenames is None: 

76 filenames = [None] * nimages 

77 self.filenames = filenames 

78 

79 warning = False 

80 

81 self._images = [] 

82 

83 # Whether length or chemical composition changes: 

84 self.have_varying_species = False 

85 for i, atoms in enumerate(images): 

86 # copy atoms or not? Not copying allows back-editing, 

87 # but copying actually forgets things like the attached 

88 # calculator (might have forces/energies 

89 self._images.append(atoms) 

90 self.have_varying_species |= not np.array_equal(self[0].numbers, 

91 atoms.numbers) 

92 if hasattr(self, 'Q'): 

93 assert False # XXX askhl fix quaternions 

94 self.Q[i] = atoms.get_quaternions() 

95 if (atoms.pbc != self[0].pbc).any(): 

96 warning = True 

97 

98 if warning: 

99 import warnings 

100 warnings.warn('Not all images have the same boundary conditions!') 

101 

102 self.maxnatoms = max(len(atoms) for atoms in self) 

103 self.selected = np.zeros(self.maxnatoms, bool) 

104 self.selected_ordered = [] 

105 self.visible = np.ones(self.maxnatoms, bool) 

106 self.repeat = np.ones(3, int) 

107 

108 def get_radii(self, atoms): 

109 radii = np.array([self.covalent_radii[z] for z in atoms.numbers]) 

110 radii *= self.atom_scale 

111 return radii 

112 

113 def read(self, filenames, default_index=':', filetype=None): 

114 if isinstance(default_index, str): 

115 default_index = string2index(default_index) 

116 

117 images = [] 

118 names = [] 

119 for filename in filenames: 

120 from ase.io.formats import parse_filename 

121 

122 if '@' in filename and 'postgres' not in filename or \ 

123 'postgres' in filename and filename.count('@') == 2: 

124 actual_filename, index = parse_filename(filename, None) 

125 else: 

126 actual_filename, index = parse_filename(filename, 

127 default_index) 

128 

129 # Read from stdin: 

130 if filename == '-': 

131 import sys 

132 from io import BytesIO 

133 buf = BytesIO(sys.stdin.buffer.read()) 

134 buf.seek(0) 

135 filename = buf 

136 filetype = 'traj' 

137 

138 imgs = read(filename, index, filetype) 

139 if hasattr(imgs, 'iterimages'): 

140 imgs = list(imgs.iterimages()) 

141 

142 images.extend(imgs) 

143 

144 # Name each file as filename@index: 

145 if isinstance(index, slice): 

146 start = index.start or 0 

147 step = index.step or 1 

148 else: 

149 start = index 

150 step = 1 

151 for i, img in enumerate(imgs): 

152 if isinstance(start, int): 

153 names.append('{}@{}'.format( 

154 actual_filename, start + i * step)) 

155 else: 

156 names.append('{}@{}'.format(actual_filename, start)) 

157 

158 self.initialize(images, names) 

159 

160 def repeat_results(self, atoms, repeat=None, oldprod=None): 

161 """Return a dictionary which updates the magmoms, energy and forces 

162 to the repeated amount of atoms. 

163 """ 

164 def getresult(name, get_quantity): 

165 # ase/io/trajectory.py line 170 does this by using 

166 # the get_property(prop, atoms, allow_calculation=False) 

167 # so that is an alternative option. 

168 try: 

169 if (not atoms.calc or 

170 atoms.calc.calculation_required(atoms, [name])): 

171 quantity = None 

172 else: 

173 quantity = get_quantity() 

174 except Exception as err: 

175 quantity = None 

176 errmsg = ('An error occurred while retrieving {} ' 

177 'from the calculator: {}'.format(name, err)) 

178 warnings.warn(errmsg) 

179 return quantity 

180 

181 if repeat is None: 

182 repeat = self.repeat.prod() 

183 if oldprod is None: 

184 oldprod = self.repeat.prod() 

185 

186 results = {} 

187 

188 original_length = len(atoms) // oldprod 

189 newprod = repeat.prod() 

190 

191 # Read the old properties 

192 magmoms = getresult('magmoms', atoms.get_magnetic_moments) 

193 magmom = getresult('magmom', atoms.get_magnetic_moment) 

194 energy = getresult('energy', atoms.get_potential_energy) 

195 forces = getresult('forces', atoms.get_forces) 

196 

197 # Update old properties to the repeated image 

198 if magmoms is not None: 

199 magmoms = np.tile(magmoms[:original_length], newprod) 

200 results['magmoms'] = magmoms 

201 

202 if magmom is not None: 

203 magmom = magmom * newprod / oldprod 

204 results['magmom'] = magmom 

205 

206 if forces is not None: 

207 forces = np.tile(forces[:original_length].T, newprod).T 

208 results['forces'] = forces 

209 

210 if energy is not None: 

211 energy = energy * newprod / oldprod 

212 results['energy'] = energy 

213 

214 return results 

215 

216 def repeat_unit_cell(self): 

217 for atoms in self: 

218 # Get quantities taking into account current repeat():' 

219 results = self.repeat_results(atoms, self.repeat.prod(), 

220 oldprod=self.repeat.prod()) 

221 

222 atoms.cell *= self.repeat.reshape((3, 1)) 

223 atoms.calc = SinglePointCalculator(atoms, **results) 

224 self.repeat = np.ones(3, int) 

225 

226 def repeat_images(self, repeat): 

227 from ase.constraints import FixAtoms 

228 repeat = np.array(repeat) 

229 oldprod = self.repeat.prod() 

230 images = [] 

231 constraints_removed = False 

232 

233 for i, atoms in enumerate(self): 

234 refcell = atoms.get_cell() 

235 fa = [] 

236 for c in atoms._constraints: 

237 if isinstance(c, FixAtoms): 

238 fa.append(c) 

239 else: 

240 constraints_removed = True 

241 atoms.set_constraint(fa) 

242 

243 # Update results dictionary to repeated atoms 

244 results = self.repeat_results(atoms, repeat, oldprod) 

245 

246 del atoms[len(atoms) // oldprod:] # Original atoms 

247 

248 atoms *= repeat 

249 atoms.cell = refcell 

250 

251 atoms.calc = SinglePointCalculator(atoms, **results) 

252 

253 images.append(atoms) 

254 

255 if constraints_removed: 

256 from ase.gui.ui import tk, showwarning 

257 # We must be able to show warning before the main GUI 

258 # has been created. So we create a new window, 

259 # then show the warning, then destroy the window. 

260 tmpwindow = tk.Tk() 

261 tmpwindow.withdraw() # Host window will never be shown 

262 showwarning(_('Constraints discarded'), 

263 _('Constraints other than FixAtoms ' 

264 'have been discarded.')) 

265 tmpwindow.destroy() 

266 

267 self.initialize(images, filenames=self.filenames) 

268 self.repeat = repeat 

269 

270 def center(self): 

271 """Center each image in the existing unit cell, keeping the 

272 cell constant.""" 

273 for atoms in self: 

274 atoms.center() 

275 

276 def graph(self, expr): 

277 """Routine to create the data in graphs, defined by the 

278 string expr.""" 

279 import ase.units as units 

280 code = compile(expr + ',', '<input>', 'eval') 

281 

282 nimages = len(self) 

283 

284 def d(n1, n2): 

285 return sqrt(((R[n1] - R[n2])**2).sum()) 

286 

287 def a(n1, n2, n3): 

288 v1 = R[n1] - R[n2] 

289 v2 = R[n3] - R[n2] 

290 arg = np.vdot(v1, v2) / (sqrt((v1**2).sum() * (v2**2).sum())) 

291 if arg > 1.0: 

292 arg = 1.0 

293 if arg < -1.0: 

294 arg = -1.0 

295 return 180.0 * np.arccos(arg) / np.pi 

296 

297 def dih(n1, n2, n3, n4): 

298 # vector 0->1, 1->2, 2->3 and their normalized cross products: 

299 a = R[n2] - R[n1] 

300 b = R[n3] - R[n2] 

301 c = R[n4] - R[n3] 

302 bxa = np.cross(b, a) 

303 bxa /= np.sqrt(np.vdot(bxa, bxa)) 

304 cxb = np.cross(c, b) 

305 cxb /= np.sqrt(np.vdot(cxb, cxb)) 

306 angle = np.vdot(bxa, cxb) 

307 # check for numerical trouble due to finite precision: 

308 if angle < -1: 

309 angle = -1 

310 if angle > 1: 

311 angle = 1 

312 angle = np.arccos(angle) 

313 if np.vdot(bxa, c) > 0: 

314 angle = 2 * np.pi - angle 

315 return angle * 180.0 / np.pi 

316 

317 # get number of mobile atoms for temperature calculation 

318 E = np.array([self.get_energy(atoms) for atoms in self]) 

319 

320 s = 0.0 

321 

322 # Namespace for eval: 

323 ns = {'E': E, 

324 'd': d, 'a': a, 'dih': dih} 

325 

326 data = [] 

327 for i in range(nimages): 

328 ns['i'] = i 

329 ns['s'] = s 

330 ns['R'] = R = self[i].get_positions() 

331 ns['V'] = self[i].get_velocities() 

332 F = self.get_forces(self[i]) 

333 if F is not None: 

334 ns['F'] = F 

335 ns['A'] = self[i].get_cell() 

336 ns['M'] = self[i].get_masses() 

337 # XXX askhl verify: 

338 dynamic = self.get_dynamic(self[i]) 

339 if F is not None: 

340 ns['f'] = f = ((F * dynamic[:, None])**2).sum(1)**.5 

341 ns['fmax'] = max(f) 

342 ns['fave'] = f.mean() 

343 ns['epot'] = epot = E[i] 

344 ns['ekin'] = ekin = self[i].get_kinetic_energy() 

345 ns['e'] = epot + ekin 

346 ndynamic = dynamic.sum() 

347 if ndynamic > 0: 

348 ns['T'] = 2.0 * ekin / (3.0 * ndynamic * units.kB) 

349 data = eval(code, ns) 

350 if i == 0: 

351 nvariables = len(data) 

352 xy = np.empty((nvariables, nimages)) 

353 xy[:, i] = data 

354 if i + 1 < nimages and not self.have_varying_species: 

355 dR = find_mic(self[i + 1].positions - R, self[i].get_cell(), 

356 self[i].get_pbc())[0] 

357 s += sqrt((dR**2).sum()) 

358 return xy 

359 

360 def write(self, filename, rotations='', bbox=None, 

361 **kwargs): 

362 # XXX We should show the unit cell whenever there is one 

363 indices = range(len(self)) 

364 p = filename.rfind('@') 

365 if p != -1: 

366 try: 

367 slice = string2index(filename[p + 1:]) 

368 except ValueError: 

369 pass 

370 else: 

371 indices = indices[slice] 

372 filename = filename[:p] 

373 if isinstance(indices, int): 

374 indices = [indices] 

375 

376 images = [self.get_atoms(i) for i in indices] 

377 if len(filename) > 4 and filename[-4:] in ['.eps', '.png', '.pov']: 

378 write(filename, images, 

379 rotation=rotations, 

380 bbox=bbox, **kwargs) 

381 else: 

382 write(filename, images, **kwargs) 

383 

384 def get_atoms(self, frame, remove_hidden=False): 

385 atoms = self[frame] 

386 try: 

387 E = atoms.get_potential_energy() 

388 except RuntimeError: 

389 E = None 

390 try: 

391 F = atoms.get_forces() 

392 except RuntimeError: 

393 F = None 

394 

395 # Remove hidden atoms if applicable 

396 if remove_hidden: 

397 atoms = atoms[self.visible] 

398 if F is not None: 

399 F = F[self.visible] 

400 atoms.calc = SinglePointCalculator(atoms, energy=E, forces=F) 

401 return atoms 

402 

403 def delete(self, i): 

404 self.images.pop(i) 

405 self.filenames.pop(i) 

406 self.initialize(self.images, self.filenames)