Coverage for ase / gui / view.py: 69.26%

1# fmt: off 

2 

3from math import cos, sin, sqrt 

4from os.path import basename 

5 

6import numpy as np 

7 

8from ase.calculators.calculator import PropertyNotImplementedError 

9from ase.data import atomic_numbers 

10from ase.data.colors import jmol_colors 

11from ase.geometry import complete_cell 

12from ase.gui.colors import ColorWindow 

13from ase.gui.i18n import _, ngettext 

14from ase.gui.render import Render 

15from ase.gui.repeat import Repeat 

16from ase.gui.rotate import Rotate 

17from ase.gui.utils import get_magmoms 

18from ase.utils import rotate 

19 

20GREEN = '#74DF00' 

21PURPLE = '#AC58FA' 

22BLACKISH = '#151515' 

23 

24 

25def get_cell_coordinates(cell, shifted=False): 

26 """Get start and end points of lines segments used to draw cell.""" 

27 nn = [] 

28 for c in range(3): 

29 v = cell[c] 

30 d = sqrt(np.dot(v, v)) 

31 if d < 1e-12: 

32 n = 0 

33 else: 

34 n = max(2, int(d / 0.3)) 

35 nn.append(n) 

36 B1 = np.zeros((2, 2, sum(nn), 3)) 

37 B2 = np.zeros((2, 2, sum(nn), 3)) 

38 n1 = 0 

39 for c, n in enumerate(nn): 

40 n2 = n1 + n 

41 h = 1.0 / (2 * n - 1) 

42 R = np.arange(n) * (2 * h) 

43 

44 for i, j in [(0, 0), (0, 1), (1, 0), (1, 1)]: 

45 B1[i, j, n1:n2, c] = R 

46 B1[i, j, n1:n2, (c + 1) % 3] = i 

47 B1[i, j, n1:n2, (c + 2) % 3] = j 

48 B2[:, :, n1:n2] = B1[:, :, n1:n2] 

49 B2[:, :, n1:n2, c] += h 

50 n1 = n2 

51 B1.shape = (-1, 3) 

52 B2.shape = (-1, 3) 

53 if shifted: 

54 B1 -= 0.5 

55 B2 -= 0.5 

56 return B1, B2 

57 

58 

59def get_bonds(atoms, covalent_radii): 

60 from ase.neighborlist import PrimitiveNeighborList 

61 

62 nl = PrimitiveNeighborList( 

63 covalent_radii * 1.5, 

64 skin=0.0, 

65 self_interaction=False, 

66 bothways=False, 

67 ) 

68 nl.update(atoms.pbc, atoms.get_cell(complete=True), atoms.positions) 

69 number_of_neighbors = sum(indices.size for indices in nl.neighbors) 

70 number_of_pbc_neighbors = sum( 

71 offsets.any(axis=1).sum() for offsets in nl.displacements 

72 ) # sum up all neighbors that have non-zero supercell offsets 

73 nbonds = number_of_neighbors + number_of_pbc_neighbors 

74 

75 bonds = np.empty((nbonds, 5), int) 

76 if nbonds == 0: 

77 return bonds 

78 

79 n1 = 0 

80 for a in range(len(atoms)): 

81 indices, offsets = nl.get_neighbors(a) 

82 n2 = n1 + len(indices) 

83 bonds[n1:n2, 0] = a 

84 bonds[n1:n2, 1] = indices 

85 bonds[n1:n2, 2:] = offsets 

86 n1 = n2 

87 

88 i = bonds[:n2, 2:].any(1) 

89 pbcbonds = bonds[:n2][i] 

90 bonds[n2:, 0] = pbcbonds[:, 1] 

91 bonds[n2:, 1] = pbcbonds[:, 0] 

92 bonds[n2:, 2:] = -pbcbonds[:, 2:] 

93 return bonds 

94 

95 

96class View: 

97 def __init__(self, rotations): 

98 self.colormode = 'jmol' # The default colors 

99 self.axes = rotate(rotations) 

100 self.configured = False 

101 self.frame = None 

102 

103 # XXX 

104 self.colormode = 'jmol' 

105 self.colors = { 

106 i: ('#{:02X}{:02X}{:02X}'.format(*(int(x * 255) for x in rgb))) 

107 for i, rgb in enumerate(jmol_colors) 

108 } 

109 # scaling factors for vectors 

110 self.force_vector_scale = self.config['force_vector_scale'] 

111 self.velocity_vector_scale = self.config['velocity_vector_scale'] 

112 self.magmom_vector_scale = self.config['magmom_vector_scale'] 

113 

114 # buttons 

115 self.b1 = 1 # left 

116 self.b3 = 3 # right 

117 if self.config['swap_mouse']: 

118 self.b1 = 3 

119 self.b3 = 1 

120 

121 @property 

122 def atoms(self): 

123 return self.images[self.frame] 

124 

125 def set_frame(self, frame=None, focus=False): 

126 if frame is None: 

127 frame = self.frame 

128 assert frame < len(self.images) 

129 self.frame = frame 

130 self.set_atoms(self.images[frame]) 

131 

132 fname = self.images.filenames[frame] 

133 if fname is None: 

134 header = 'ase.gui' 

135 else: 

136 # fname is actually not necessarily the filename but may 

137 # contain indexing like filename@0 

138 header = basename(fname) 

139 

140 images_loaded_text = ngettext( 

141 'one image loaded', 

142 '{} images loaded', 

143 len(self.images) 

144 ).format(len(self.images)) 

145 

146 self.window.title = f'{header} — {images_loaded_text}' 

147 

148 if self.movie_window is not None: 

149 self.movie_window.frame_number.value = frame 

150 

151 if focus: 

152 self.focus() 

153 else: 

154 self.draw() 

155 

156 def get_bonds(self, atoms): 

157 # this method exists rather than just using the standalone function 

158 # so that it can be overridden by external libraries 

159 return get_bonds(atoms, self.get_covalent_radii(atoms)) 

160 

161 def set_atoms(self, atoms): 

162 natoms = len(atoms) 

163 

164 if self.showing_cell(): 

165 B1, B2 = get_cell_coordinates(atoms.cell, 

166 self.config['shift_cell']) 

167 else: 

168 B1 = B2 = np.zeros((0, 3)) 

169 

170 if self.showing_bonds(): 

171 atomscopy = atoms.copy() 

172 atomscopy.cell *= self.images.repeat[:, np.newaxis] 

173 bonds = self.get_bonds(atomscopy) 

174 else: 

175 bonds = np.empty((0, 5), int) 

176 

177 # X is all atomic coordinates, and starting points of vectors 

178 # like bonds and cell segments. 

179 # The reason to have them all in one big list is that we like to 

180 # eventually rotate/sort it by Z-order when rendering. 

181 

182 # Also B are the end points of line segments. 

183 

184 self.X = np.empty((natoms + len(B1) + len(bonds), 3)) 

185 self.X_pos = self.X[:natoms] 

186 self.X_pos[:] = atoms.positions 

187 self.X_cell = self.X[natoms:natoms + len(B1)] 

188 self.X_bonds = self.X[natoms + len(B1):] 

189 

190 cell = atoms.cell 

191 ncellparts = len(B1) 

192 nbonds = len(bonds) 

193 

194 self.X_cell[:] = np.dot(B1, cell) 

195 self.B = np.empty((ncellparts + nbonds, 3)) 

196 self.B[:ncellparts] = np.dot(B2, cell) 

197 

198 if nbonds > 0: 

199 P = atoms.positions 

200 Af = self.images.repeat[:, np.newaxis] * cell 

201 a = P[bonds[:, 0]] 

202 b = P[bonds[:, 1]] + np.dot(bonds[:, 2:], Af) - a 

203 d = (b**2).sum(1)**0.5 

204 r = 0.65 * self.get_covalent_radii() 

205 x0 = (r[bonds[:, 0]] / d).reshape((-1, 1)) 

206 x1 = (r[bonds[:, 1]] / d).reshape((-1, 1)) 

207 self.X_bonds[:] = a + b * x0 

208 b *= 1.0 - x0 - x1 

209 b[bonds[:, 2:].any(1)] *= 0.5 

210 self.B[ncellparts:] = self.X_bonds + b 

211 

212 self.obs.set_atoms.notify() 

213 

214 def showing_bonds(self): 

215 return self.window['toggle-show-bonds'] 

216 

217 def showing_cell(self): 

218 return self.window['toggle-show-unit-cell'] 

219 

220 def toggle_show_unit_cell(self, key=None): 

221 self.set_frame() 

222 

223 def get_labels(self): 

224 index = self.window['show-labels'] 

225 if index == 0: 

226 return None 

227 

228 if index == 1: 

229 return list(range(len(self.atoms))) 

230 

231 if index == 2: 

232 return list(get_magmoms(self.atoms)) 

233 

234 if index == 4: 

235 Q = self.atoms.get_initial_charges() 

236 return [f'{q:.4g}' for q in Q] 

237 

238 return self.atoms.symbols 

239 

240 def show_labels(self): 

241 self.draw() 

242 

243 def toggle_show_axes(self, key=None): 

244 self.draw() 

245 

246 def toggle_show_bonds(self, key=None): 

247 self.set_frame() 

248 

249 def toggle_show_velocities(self, key=None): 

250 self.draw() 

251 

252 def get_forces(self): 

253 if self.atoms.calc is not None: 

254 try: 

255 return self.atoms.get_forces() 

256 except PropertyNotImplementedError: 

257 pass 

258 return np.zeros((len(self.atoms), 3)) 

259 

260 def toggle_show_forces(self, key=None): 

261 self.draw() 

262 

263 def toggle_show_magmoms(self, key=None): 

264 self.draw() 

265 

266 def hide_selected(self): 

267 self.images.visible[self.images.selected] = False 

268 self.draw() 

269 

270 def show_selected(self): 

271 self.images.visible[self.images.selected] = True 

272 self.draw() 

273 

274 def repeat_window(self, key=None): 

275 return Repeat(self) 

276 

277 def rotate_window(self): 

278 return Rotate(self) 

279 

280 def colors_window(self, key=None): 

281 win = ColorWindow(self) 

282 self.obs.new_atoms.register(win.notify_atoms_changed) 

283 return win 

284 

285 def focus(self, x=None): 

286 cell = (self.window['toggle-show-unit-cell'] and 

287 self.images[0].cell.any()) 

288 if (len(self.atoms) == 0 and not cell): 

289 self.scale = 20.0 

290 self.center = np.zeros(3) 

291 self.draw() 

292 return 

293 

294 # Get the min and max point of the projected atom positions 

295 # including the covalent_radii used for drawing the atoms 

296 P = np.dot(self.X, self.axes) 

297 n = len(self.atoms) 

298 covalent_radii = self.get_covalent_radii() 

299 P[:n] -= covalent_radii[:, None] 

300 P1 = P.min(0) 

301 P[:n] += 2 * covalent_radii[:, None] 

302 P2 = P.max(0) 

303 self.center = np.dot(self.axes, (P1 + P2) / 2) 

304 self.center += self.atoms.get_celldisp().reshape((3,)) / 2 

305 # Add 30% of whitespace on each side of the atoms 

306 S = 1.3 * (P2 - P1) 

307 w, h = self.window.size 

308 if S[0] * h < S[1] * w: 

309 self.scale = h / S[1] 

310 elif S[0] > 0.0001: 

311 self.scale = w / S[0] 

312 else: 

313 self.scale = 1.0 

314 self.draw() 

315 

316 def reset_view(self, menuitem): 

317 self.axes = rotate('0.0x,0.0y,0.0z') 

318 self.set_frame() 

319 self.focus(self) 

320 

321 def set_view(self, key): 

322 if key == 'Z': 

323 self.axes = rotate('0.0x,0.0y,0.0z') 

324 elif key == 'X': 

325 self.axes = rotate('-90.0x,-90.0y,0.0z') 

326 elif key == 'Y': 

327 self.axes = rotate('90.0x,0.0y,90.0z') 

328 elif key == 'Shift+Z': 

329 self.axes = rotate('180.0x,0.0y,90.0z') 

330 elif key == 'Shift+X': 

331 self.axes = rotate('0.0x,90.0y,0.0z') 

332 elif key == 'Shift+Y': 

333 self.axes = rotate('-90.0x,0.0y,0.0z') 

334 else: 

335 if key == 'I': 

336 i, j = 1, 2 

337 elif key == 'J': 

338 i, j = 2, 0 

339 elif key == 'K': 

340 i, j = 0, 1 

341 elif key == 'Shift+I': 

342 i, j = 2, 1 

343 elif key == 'Shift+J': 

344 i, j = 0, 2 

345 elif key == 'Shift+K': 

346 i, j = 1, 0 

347 

348 A = complete_cell(self.atoms.cell) 

349 x1 = A[i] 

350 x2 = A[j] 

351 

352 norm = np.linalg.norm 

353 

354 x1 = x1 / norm(x1) 

355 x2 = x2 - x1 * np.dot(x1, x2) 

356 x2 /= norm(x2) 

357 x3 = np.cross(x1, x2) 

358 

359 self.axes = np.array([x1, x2, x3]).T 

360 

361 self.set_frame() 

362 

363 def get_colors(self, rgb=False): 

364 if rgb: 

365 return [tuple(int(_rgb[i:i + 2], 16) / 255 for i in range(1, 7, 2)) 

366 for _rgb in self.get_colors()] 

367 

368 if self.colormode == 'jmol': 

369 return [self.colors.get(Z, BLACKISH) for Z in self.atoms.numbers] 

370 

371 if self.colormode == 'neighbors': 

372 return [self.colors.get(Z, BLACKISH) 

373 for Z in self.get_color_scalars()] 

374 

375 colorscale, cmin, cmax = self.colormode_data 

376 N = len(colorscale) 

377 colorswhite = colorscale + ['#ffffff'] 

378 if cmin == cmax: 

379 indices = [N // 2] * len(self.atoms) 

380 else: 

381 scalars = np.ma.array(self.get_color_scalars()) 

382 indices = np.clip(((scalars - cmin) / (cmax - cmin) * N + 

383 0.5).astype(int), 

384 0, N - 1).filled(N) 

385 return [colorswhite[i] for i in indices] 

386 

387 def get_color_scalars(self, frame=None): 

388 if self.colormode == 'tag': 

389 return self.atoms.get_tags() 

390 if self.colormode == 'force': 

391 f = (self.get_forces()**2).sum(1)**0.5 

392 return f * self.images.get_dynamic(self.atoms) 

393 elif self.colormode == 'velocity': 

394 return (self.atoms.get_velocities()**2).sum(1)**0.5 

395 elif self.colormode == 'initial charge': 

396 return self.atoms.get_initial_charges() 

397 elif self.colormode == 'magmom': 

398 return get_magmoms(self.atoms) 

399 elif self.colormode == 'neighbors': 

400 from ase.neighborlist import NeighborList 

401 n = len(self.atoms) 

402 nl = NeighborList(self.get_covalent_radii(self.atoms) * 1.5, 

403 skin=0, self_interaction=False, bothways=True) 

404 nl.update(self.atoms) 

405 return [len(nl.get_neighbors(i)[0]) for i in range(n)] 

406 else: 

407 scalars = np.array(self.atoms.get_array(self.colormode), 

408 dtype=float) 

409 return np.ma.array(scalars, mask=np.isnan(scalars)) 

410 

411 def get_covalent_radii(self, atoms=None): 

412 if atoms is None: 

413 atoms = self.atoms 

414 return self.images.get_radii(atoms) 

415 

416 def draw(self, status=True): 

417 self.window.clear() 

418 axes = self.scale * self.axes * (1, -1, 1) 

419 offset = np.dot(self.center, axes) 

420 offset[:2] -= 0.5 * self.window.size 

421 X = np.dot(self.X, axes) - offset 

422 n = len(self.atoms) 

423 

424 # The indices enumerate drawable objects in z order: 

425 self.indices = X[:, 2].argsort() 

426 r = self.get_covalent_radii() * self.scale 

427 if self.window['toggle-show-bonds']: 

428 r *= 0.65 

429 P = self.P = X[:n, :2] 

430 A = (P - r[:, None]).round().astype(int) 

431 X1 = X[n:, :2].round().astype(int) 

432 X2 = (np.dot(self.B, axes) - offset).round().astype(int) 

433 disp = (np.dot(self.atoms.get_celldisp().reshape((3,)), 

434 axes)).round().astype(int) 

435 d = (2 * r).round().astype(int) 

436 

437 vector_arrays = [] 

438 if self.window['toggle-show-velocities']: 

439 # Scale ugly? 

440 v = self.atoms.get_velocities() 

441 if v is not None: 

442 vector_arrays.append(v * 10.0 * self.velocity_vector_scale) 

443 if self.window['toggle-show-forces']: 

444 f = self.get_forces() 

445 vector_arrays.append(f * self.force_vector_scale) 

446 

447 if self.window['toggle-show-magmoms']: 

448 magmom = get_magmoms(self.atoms) 

449 # Turn this into a 3D vector if it is a scalar 

450 magmom_vecs = [] 

451 for i in range(len(magmom)): 

452 if isinstance(magmom[i], (int, float)): 

453 magmom_vecs.append(np.array([0, 0, magmom[i]])) 

454 elif isinstance(magmom[i], np.ndarray) and len(magmom[i]) == 3: 

455 magmom_vecs.append(magmom[i]) 

456 else: 

457 raise TypeError('Magmom is not a 3-component vector ' 

458 'or a scalar') 

459 magmom_vecs = np.array(magmom_vecs) 

460 vector_arrays.append(magmom_vecs * 0.5 * self.magmom_vector_scale) 

461 

462 for array in vector_arrays: 

463 array[:] = np.dot(array, axes) + X[:n] 

464 

465 colors = self.get_colors() 

466 circle = self.window.circle 

467 arc = self.window.arc 

468 line = self.window.line 

469 constrained = ~self.images.get_dynamic(self.atoms) 

470 

471 selected = self.images.selected 

472 visible = self.images.visible 

473 ncell = len(self.X_cell) 

474 bond_linewidth = self.scale * 0.15 

475 

476 labels = self.get_labels() 

477 

478 if self.arrowkey_mode == self.ARROWKEY_MOVE: 

479 movecolor = GREEN 

480 elif self.arrowkey_mode == self.ARROWKEY_ROTATE: 

481 movecolor = PURPLE 

482 

483 for a in self.indices: 

484 if a < n: 

485 ra = d[a] 

486 if visible[a]: 

487 try: 

488 kinds = self.atoms.arrays['spacegroup_kinds'] 

489 site_occ = self.atoms.info['occupancy'][str(kinds[a])] 

490 # first an empty circle if a site is not fully occupied 

491 if (np.sum([v for v in site_occ.values()])) < 1.0: 

492 fill = '#ffffff' 

493 circle(fill, selected[a], 

494 A[a, 0], A[a, 1], 

495 A[a, 0] + ra, A[a, 1] + ra) 

496 start = 0 

497 # start with the dominant species 

498 for sym, occ in sorted(site_occ.items(), 

499 key=lambda x: x[1], 

500 reverse=True): 

501 if np.round(occ, decimals=4) == 1.0: 

502 circle(colors[a], selected[a], 

503 A[a, 0], A[a, 1], 

504 A[a, 0] + ra, A[a, 1] + ra) 

505 else: 

506 # jmol colors for the moment 

507 extent = 360. * occ 

508 arc(self.colors[atomic_numbers[sym]], 

509 selected[a], 

510 start, extent, 

511 A[a, 0], A[a, 1], 

512 A[a, 0] + ra, A[a, 1] + ra) 

513 start += extent 

514 except KeyError: 

515 # legacy behavior 

516 # Draw the atoms 

517 if (self.moving and a < len(self.move_atoms_mask) 

518 and self.move_atoms_mask[a]): 

519 circle(movecolor, False, 

520 A[a, 0] - 4, A[a, 1] - 4, 

521 A[a, 0] + ra + 4, A[a, 1] + ra + 4) 

522 

523 circle(colors[a], selected[a], 

524 A[a, 0], A[a, 1], A[a, 0] + ra, A[a, 1] + ra) 

525 

526 # Draw labels on the atoms 

527 if labels is not None: 

528 self.window.text(A[a, 0] + ra / 2, 

529 A[a, 1] + ra / 2, 

530 str(labels[a])) 

531 

532 # Draw cross on constrained atoms 

533 if constrained[a]: 

534 R1 = int(0.14644 * ra) 

535 R2 = int(0.85355 * ra) 

536 line((A[a, 0] + R1, A[a, 1] + R1, 

537 A[a, 0] + R2, A[a, 1] + R2)) 

538 line((A[a, 0] + R2, A[a, 1] + R1, 

539 A[a, 0] + R1, A[a, 1] + R2)) 

540 

541 # Draw velocities and/or forces 

542 for v in vector_arrays: 

543 assert not np.isnan(v).any() 

544 self.arrow((X[a, 0], X[a, 1], v[a, 0], v[a, 1]), 

545 width=2) 

546 else: 

547 # Draw unit cell and/or bonds: 

548 a -= n 

549 if a < ncell: 

550 line((X1[a, 0] + disp[0], X1[a, 1] + disp[1], 

551 X2[a, 0] + disp[0], X2[a, 1] + disp[1])) 

552 else: 

553 line((X1[a, 0], X1[a, 1], 

554 X2[a, 0], X2[a, 1]), 

555 width=bond_linewidth) 

556 

557 if self.window['toggle-show-axes']: 

558 self.draw_axes() 

559 

560 if self.arrowkey_mode != self.ARROWKEY_SCAN: 

561 self.draw_arrowkey_hint() 

562 else: 

563 self.hide_arrowkey_hint() 

564 if len(self.images) > 1: 

565 self.draw_frame_number() 

566 

567 self.window.update() 

568 

569 if status: 

570 self.status.status(self.atoms) 

571 

572 # Currently we change the atoms all over the place willy-nilly 

573 # and then call draw(). For which reason we abuse draw() to notify 

574 # the observers about general changes. 

575 # 

576 # We should refactor so change_atoms is only emitted 

577 # when when atoms actually change, and maybe have a separate signal 

578 # to listen to e.g. changes of view. 

579 self.obs.change_atoms.notify() 

580 

581 def arrow(self, coords, width): 

582 line = self.window.line 

583 begin = np.array((coords[0], coords[1])) 

584 end = np.array((coords[2], coords[3])) 

585 line(coords, width) 

586 

587 vec = end - begin 

588 length = np.sqrt((vec[:2]**2).sum()) 

589 length = min(length, 0.3 * self.scale) 

590 

591 angle = np.arctan2(end[1] - begin[1], end[0] - begin[0]) + np.pi 

592 x1 = (end[0] + length * np.cos(angle - 0.3)).round().astype(int) 

593 y1 = (end[1] + length * np.sin(angle - 0.3)).round().astype(int) 

594 x2 = (end[0] + length * np.cos(angle + 0.3)).round().astype(int) 

595 y2 = (end[1] + length * np.sin(angle + 0.3)).round().astype(int) 

596 line((x1, y1, end[0], end[1]), width) 

597 line((x2, y2, end[0], end[1]), width) 

598 

599 def draw_axes(self): 

600 axes_length = 15 

601 

602 rgb = ['red', 'green', 'blue'] 

603 

604 for i in self.axes[:, 2].argsort(): 

605 a = 20 

606 b = self.window.size[1] - 20 

607 c = int(self.axes[i][0] * axes_length + a) 

608 d = int(-self.axes[i][1] * axes_length + b) 

609 self.window.line((a, b, c, d)) 

610 self.window.text(c, d, 'XYZ'[i], color=rgb[i]) 

611 

612 def draw_arrowkey_hint(self): 

613 if self.arrowkey_mode == self.ARROWKEY_ROTATE: 

614 hint = _('ROTATING') 

615 bg = PURPLE 

616 tip_text = _( 

617 'Ctrl + Up/Down: rotate around the view axis\n' 

618 'Shift + Arrow keys: rotate in smaller increments' 

619 ) 

620 else: 

621 hint = _('MOVING') 

622 bg = GREEN 

623 tip_text = _( 

624 'Ctrl + Up/Down: move along the view axis\n' 

625 'Shift + Arrow keys: move in smaller increments' 

626 ) 

627 self.arrowkey_hint.label.configure( 

628 text=hint, 

629 padx=3, 

630 bg=bg, 

631 ) 

632 self.arrowkey_hint.tooltip.configure( 

633 text=tip_text, 

634 justify='right', 

635 font='TkTooltipFont', 

636 ) 

637 self.arrowkey_hint.place(relx=1.0, rely=1.0, anchor='se') 

638 self.arrowkey_hint.exists = True 

639 

640 def hide_arrowkey_hint(self): 

641 if self.arrowkey_hint.exists: 

642 self.arrowkey_hint.place_forget() 

643 self.arrowkey_hint.exists = False 

644 

645 def draw_frame_number(self): 

646 x, y = self.window.size 

647 self.window.text(x, y, '{}'.format(self.frame), 

648 anchor='SE') 

649