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 # The below seems to be about "quaternions" 

80 if 0: # XXXXXXXXXXXXXXXXXXXX hasattr(images[0], 'get_shapes'): 

81 self.Q = np.empty((nimages, self.natoms, 4)) 

82 self.shapes = images[0].get_shapes() 

83 import os as os 

84 if os.path.exists('shapes'): 

85 shapesfile = open('shapes') 

86 lines = shapesfile.readlines() 

87 shapesfile.close() 

88 if '#{type:(shape_x,shape_y,shape_z), .....,}' in lines[0]: 

89 shape = eval(lines[1]) 

90 shapes = [] 

91 for an in images[0].get_atomic_numbers(): 

92 shapes.append(shape[an]) 

93 self.shapes = np.array(shapes) 

94 else: 

95 print('shape file has wrong format') 

96 else: 

97 print('no shapesfile found: default shapes were used!') 

98 

99 else: 

100 self.shapes = None 

101 

102 warning = False 

103 

104 self._images = [] 

105 

106 # Whether length or chemical composition changes: 

107 self.have_varying_species = False 

108 for i, atoms in enumerate(images): 

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

110 # but copying actually forgets things like the attached 

111 # calculator (might have forces/energies 

112 self._images.append(atoms) 

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

114 atoms.numbers) 

115 if hasattr(self, 'Q'): 

116 assert False # XXX askhl fix quaternions 

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

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

119 warning = True 

120 

121 if warning: 

122 import warnings 

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

124 

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

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

127 self.selected_ordered = [] 

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

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

130 

131 def get_radii(self, atoms): 

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

133 radii *= self.atom_scale 

134 return radii 

135 

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

137 if isinstance(default_index, str): 

138 default_index = string2index(default_index) 

139 

140 images = [] 

141 names = [] 

142 for filename in filenames: 

143 from ase.io.formats import parse_filename 

144 

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

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

147 actual_filename, index = parse_filename(filename, None) 

148 else: 

149 actual_filename, index = parse_filename(filename, 

150 default_index) 

151 

152 # Read from stdin: 

153 if filename == '-': 

154 import sys 

155 from io import BytesIO 

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

157 buf.seek(0) 

158 filename = buf 

159 filetype = 'traj' 

160 

161 imgs = read(filename, index, filetype) 

162 if hasattr(imgs, 'iterimages'): 

163 imgs = list(imgs.iterimages()) 

164 

165 images.extend(imgs) 

166 

167 # Name each file as filename@index: 

168 if isinstance(index, slice): 

169 start = index.start or 0 

170 step = index.step or 1 

171 else: 

172 start = index 

173 step = 1 

174 for i, img in enumerate(imgs): 

175 if isinstance(start, int): 

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

177 actual_filename, start + i * step)) 

178 else: 

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

180 

181 self.initialize(images, names) 

182 

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

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

185 to the repeated amount of atoms. 

186 """ 

187 def getresult(name, get_quantity): 

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

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

190 # so that is an alternative option. 

191 try: 

192 if (not atoms.calc or 

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

194 quantity = None 

195 else: 

196 quantity = get_quantity() 

197 except Exception as err: 

198 quantity = None 

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

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

201 warnings.warn(errmsg) 

202 return quantity 

203 

204 if repeat is None: 

205 repeat = self.repeat.prod() 

206 if oldprod is None: 

207 oldprod = self.repeat.prod() 

208 

209 results = {} 

210 

211 original_length = len(atoms) // oldprod 

212 newprod = repeat.prod() 

213 

214 # Read the old properties 

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

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

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

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

219 

220 # Update old properties to the repeated image 

221 if magmoms is not None: 

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

223 results['magmoms'] = magmoms 

224 

225 if magmom is not None: 

226 magmom = magmom * newprod / oldprod 

227 results['magmom'] = magmom 

228 

229 if forces is not None: 

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

231 results['forces'] = forces 

232 

233 if energy is not None: 

234 energy = energy * newprod / oldprod 

235 results['energy'] = energy 

236 

237 return results 

238 

239 def repeat_unit_cell(self): 

240 for atoms in self: 

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

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

243 oldprod=self.repeat.prod()) 

244 

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

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

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

248 

249 def repeat_images(self, repeat): 

250 from ase.constraints import FixAtoms 

251 repeat = np.array(repeat) 

252 oldprod = self.repeat.prod() 

253 images = [] 

254 constraints_removed = False 

255 

256 for i, atoms in enumerate(self): 

257 refcell = atoms.get_cell() 

258 fa = [] 

259 for c in atoms._constraints: 

260 if isinstance(c, FixAtoms): 

261 fa.append(c) 

262 else: 

263 constraints_removed = True 

264 atoms.set_constraint(fa) 

265 

266 # Update results dictionary to repeated atoms 

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

268 

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

270 

271 atoms *= repeat 

272 atoms.cell = refcell 

273 

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

275 

276 images.append(atoms) 

277 

278 if constraints_removed: 

279 from ase.gui.ui import tk, showwarning 

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

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

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

283 tmpwindow = tk.Tk() 

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

285 showwarning(_('Constraints discarded'), 

286 _('Constraints other than FixAtoms ' 

287 'have been discarded.')) 

288 tmpwindow.destroy() 

289 

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

291 self.repeat = repeat 

292 

293 def center(self): 

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

295 cell constant.""" 

296 for atoms in self: 

297 atoms.center() 

298 

299 def graph(self, expr): 

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

301 string expr.""" 

302 import ase.units as units 

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

304 

305 nimages = len(self) 

306 

307 def d(n1, n2): 

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

309 

310 def a(n1, n2, n3): 

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

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

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

314 if arg > 1.0: 

315 arg = 1.0 

316 if arg < -1.0: 

317 arg = -1.0 

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

319 

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

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

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

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

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

325 bxa = np.cross(b, a) 

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

327 cxb = np.cross(c, b) 

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

329 angle = np.vdot(bxa, cxb) 

330 # check for numerical trouble due to finite precision: 

331 if angle < -1: 

332 angle = -1 

333 if angle > 1: 

334 angle = 1 

335 angle = np.arccos(angle) 

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

337 angle = 2 * np.pi - angle 

338 return angle * 180.0 / np.pi 

339 

340 # get number of mobile atoms for temperature calculation 

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

342 

343 s = 0.0 

344 

345 # Namespace for eval: 

346 ns = {'E': E, 

347 'd': d, 'a': a, 'dih': dih} 

348 

349 data = [] 

350 for i in range(nimages): 

351 ns['i'] = i 

352 ns['s'] = s 

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

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

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

356 if F is not None: 

357 ns['F'] = F 

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

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

360 # XXX askhl verify: 

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

362 if F is not None: 

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

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

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

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

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

368 ns['e'] = epot + ekin 

369 ndynamic = dynamic.sum() 

370 if ndynamic > 0: 

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

372 data = eval(code, ns) 

373 if i == 0: 

374 nvariables = len(data) 

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

376 xy[:, i] = data 

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

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

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

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

381 return xy 

382 

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

384 **kwargs): 

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

386 indices = range(len(self)) 

387 p = filename.rfind('@') 

388 if p != -1: 

389 try: 

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

391 except ValueError: 

392 pass 

393 else: 

394 indices = indices[slice] 

395 filename = filename[:p] 

396 if isinstance(indices, int): 

397 indices = [indices] 

398 

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

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

401 write(filename, images, 

402 rotation=rotations, 

403 bbox=bbox, **kwargs) 

404 else: 

405 write(filename, images, **kwargs) 

406 

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

408 atoms = self[frame] 

409 try: 

410 E = atoms.get_potential_energy() 

411 except RuntimeError: 

412 E = None 

413 try: 

414 F = atoms.get_forces() 

415 except RuntimeError: 

416 F = None 

417 

418 # Remove hidden atoms if applicable 

419 if remove_hidden: 

420 atoms = atoms[self.visible] 

421 if F is not None: 

422 F = F[self.visible] 

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

424 return atoms 

425 

426 def delete(self, i): 

427 self.images.pop(i) 

428 self.filenames.pop(i) 

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