Coverage for /builds/debichem-team/python-ase/ase/gui/view.py: 63.41%

1from math import cos, sin, sqrt 

2from os.path import basename 

3 

4import numpy as np 

5 

6from ase.calculators.calculator import PropertyNotImplementedError 

7from ase.data import atomic_numbers 

8from ase.data.colors import jmol_colors 

9from ase.geometry import complete_cell 

10from ase.gui.colors import ColorWindow 

11from ase.gui.i18n import ngettext 

12from ase.gui.render import Render 

13from ase.gui.repeat import Repeat 

14from ase.gui.rotate import Rotate 

15from ase.gui.utils import get_magmoms 

16from ase.utils import rotate 

17 

18GREEN = '#74DF00' 

19PURPLE = '#AC58FA' 

20BLACKISH = '#151515' 

21 

22 

23def get_cell_coordinates(cell, shifted=False): 

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

25 nn = [] 

26 for c in range(3): 

27 v = cell[c] 

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

29 if d < 1e-12: 

30 n = 0 

31 else: 

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

33 nn.append(n) 

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

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

36 n1 = 0 

37 for c, n in enumerate(nn): 

38 n2 = n1 + n 

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

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

41 

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

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

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

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

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

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

48 n1 = n2 

49 B1.shape = (-1, 3) 

50 B2.shape = (-1, 3) 

51 if shifted: 

52 B1 -= 0.5 

53 B2 -= 0.5 

54 return B1, B2 

55 

56 

57def get_bonds(atoms, covalent_radii): 

58 from ase.neighborlist import PrimitiveNeighborList 

59 

60 nl = PrimitiveNeighborList( 

61 covalent_radii * 1.5, 

62 skin=0.0, 

63 self_interaction=False, 

64 bothways=False, 

65 ) 

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

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

68 number_of_pbc_neighbors = sum( 

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

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

71 nbonds = number_of_neighbors + number_of_pbc_neighbors 

72 

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

74 if nbonds == 0: 

75 return bonds 

76 

77 n1 = 0 

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

79 indices, offsets = nl.get_neighbors(a) 

80 n2 = n1 + len(indices) 

81 bonds[n1:n2, 0] = a 

82 bonds[n1:n2, 1] = indices 

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

84 n1 = n2 

85 

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

87 pbcbonds = bonds[:n2][i] 

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

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

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

91 return bonds 

92 

93 

94class View: 

95 def __init__(self, rotations): 

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

97 self.labels = None 

98 self.axes = rotate(rotations) 

99 self.configured = False 

100 self.frame = None 

101 

102 # XXX 

103 self.colormode = 'jmol' 

104 self.colors = { 

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

106 for i, rgb in enumerate(jmol_colors) 

107 } 

108 # scaling factors for vectors 

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

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

111 

112 # buttons 

113 self.b1 = 1 # left 

114 self.b3 = 3 # right 

115 if self.config['swap_mouse']: 

116 self.b1 = 3 

117 self.b3 = 1 

118 

119 @property 

120 def atoms(self): 

121 return self.images[self.frame] 

122 

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

124 if frame is None: 

125 frame = self.frame 

126 assert frame < len(self.images) 

127 self.frame = frame 

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

129 

130 fname = self.images.filenames[frame] 

131 if fname is None: 

132 header = 'ase.gui' 

133 else: 

134 # fname is actually not necessarily the filename but may 

135 # contain indexing like filename@0 

136 header = basename(fname) 

137 

138 images_loaded_text = ngettext( 

139 'one image loaded', 

140 '{} images loaded', 

141 len(self.images) 

142 ).format(len(self.images)) 

143 

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

145 

146 if focus: 

147 self.focus() 

148 else: 

149 self.draw() 

150 

151 def get_bonds(self, atoms): 

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

153 # so that it can be overridden by external libraries 

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

155 

156 def set_atoms(self, atoms): 

157 natoms = len(atoms) 

158 

159 if self.showing_cell(): 

160 B1, B2 = get_cell_coordinates(atoms.cell, 

161 self.config['shift_cell']) 

162 else: 

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

164 

165 if self.showing_bonds(): 

166 atomscopy = atoms.copy() 

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

168 bonds = self.get_bonds(atomscopy) 

169 else: 

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

171 

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

173 # like bonds and cell segments. 

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

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

176 

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

178 

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

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

181 self.X_pos[:] = atoms.positions 

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

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

184 

185 cell = atoms.cell 

186 ncellparts = len(B1) 

187 nbonds = len(bonds) 

188 

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

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

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

192 

193 if nbonds > 0: 

194 P = atoms.positions 

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

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

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

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

199 r = 0.65 * self.get_covalent_radii() 

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

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

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

203 b *= 1.0 - x0 - x1 

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

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

206 

207 def showing_bonds(self): 

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

209 

210 def showing_cell(self): 

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

212 

213 def toggle_show_unit_cell(self, key=None): 

214 self.set_frame() 

215 

216 def update_labels(self): 

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

218 if index == 0: 

219 self.labels = None 

220 elif index == 1: 

221 self.labels = list(range(len(self.atoms))) 

222 elif index == 2: 

223 self.labels = list(get_magmoms(self.atoms)) 

224 elif index == 4: 

225 Q = self.atoms.get_initial_charges() 

226 self.labels = [f'{q:.4g}' for q in Q] 

227 else: 

228 self.labels = self.atoms.get_chemical_symbols() 

229 

230 def show_labels(self): 

231 self.update_labels() 

232 self.draw() 

233 

234 def toggle_show_axes(self, key=None): 

235 self.draw() 

236 

237 def toggle_show_bonds(self, key=None): 

238 self.set_frame() 

239 

240 def toggle_show_velocities(self, key=None): 

241 self.draw() 

242 

243 def get_forces(self): 

244 if self.atoms.calc is not None: 

245 try: 

246 return self.atoms.get_forces() 

247 except PropertyNotImplementedError: 

248 pass 

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

250 

251 def toggle_show_forces(self, key=None): 

252 self.draw() 

253 

254 def hide_selected(self): 

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

256 self.draw() 

257 

258 def show_selected(self): 

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

260 self.draw() 

261 

262 def repeat_window(self, key=None): 

263 return Repeat(self) 

264 

265 def rotate_window(self): 

266 return Rotate(self) 

267 

268 def colors_window(self, key=None): 

269 win = ColorWindow(self) 

270 self.register_vulnerable(win) 

271 return win 

272 

273 def focus(self, x=None): 

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

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

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

277 self.scale = 20.0 

278 self.center = np.zeros(3) 

279 self.draw() 

280 return 

281 

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

283 # including the covalent_radii used for drawing the atoms 

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

285 n = len(self.atoms) 

286 covalent_radii = self.get_covalent_radii() 

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

288 P1 = P.min(0) 

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

290 P2 = P.max(0) 

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

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

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

294 S = 1.3 * (P2 - P1) 

295 w, h = self.window.size 

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

297 self.scale = h / S[1] 

298 elif S[0] > 0.0001: 

299 self.scale = w / S[0] 

300 else: 

301 self.scale = 1.0 

302 self.draw() 

303 

304 def reset_view(self, menuitem): 

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

306 self.set_frame() 

307 self.focus(self) 

308 

309 def set_view(self, key): 

310 if key == 'Z': 

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

312 elif key == 'X': 

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

314 elif key == 'Y': 

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

316 elif key == 'Alt+Z': 

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

318 elif key == 'Alt+X': 

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

320 elif key == 'Alt+Y': 

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

322 else: 

323 if key == '3': 

324 i, j = 0, 1 

325 elif key == '1': 

326 i, j = 1, 2 

327 elif key == '2': 

328 i, j = 2, 0 

329 elif key == 'Alt+3': 

330 i, j = 1, 0 

331 elif key == 'Alt+1': 

332 i, j = 2, 1 

333 elif key == 'Alt+2': 

334 i, j = 0, 2 

335 

336 A = complete_cell(self.atoms.cell) 

337 x1 = A[i] 

338 x2 = A[j] 

339 

340 norm = np.linalg.norm 

341 

342 x1 = x1 / norm(x1) 

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

344 x2 /= norm(x2) 

345 x3 = np.cross(x1, x2) 

346 

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

348 

349 self.set_frame() 

350 

351 def get_colors(self, rgb=False): 

352 if rgb: 

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

354 for _rgb in self.get_colors()] 

355 

356 if self.colormode == 'jmol': 

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

358 

359 if self.colormode == 'neighbors': 

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

361 for Z in self.get_color_scalars()] 

362 

363 colorscale, cmin, cmax = self.colormode_data 

364 N = len(colorscale) 

365 colorswhite = colorscale + ['#ffffff'] 

366 if cmin == cmax: 

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

368 else: 

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

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

371 0.5).astype(int), 

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

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

374 

375 def get_color_scalars(self, frame=None): 

376 if self.colormode == 'tag': 

377 return self.atoms.get_tags() 

378 if self.colormode == 'force': 

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

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

381 elif self.colormode == 'velocity': 

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

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

384 return self.atoms.get_initial_charges() 

385 elif self.colormode == 'magmom': 

386 return get_magmoms(self.atoms) 

387 elif self.colormode == 'neighbors': 

388 from ase.neighborlist import NeighborList 

389 n = len(self.atoms) 

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

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

392 nl.update(self.atoms) 

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

394 else: 

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

396 dtype=float) 

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

398 

399 def get_covalent_radii(self, atoms=None): 

400 if atoms is None: 

401 atoms = self.atoms 

402 return self.images.get_radii(atoms) 

403 

404 def draw(self, status=True): 

405 self.window.clear() 

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

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

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

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

410 n = len(self.atoms) 

411 

412 # The indices enumerate drawable objects in z order: 

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

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

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

416 r *= 0.65 

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

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

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

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

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

422 axes)).round().astype(int) 

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

424 

425 vector_arrays = [] 

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

427 # Scale ugly? 

428 v = self.atoms.get_velocities() 

429 if v is not None: 

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

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

432 f = self.get_forces() 

433 vector_arrays.append(f * self.force_vector_scale) 

434 

435 for array in vector_arrays: 

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

437 

438 colors = self.get_colors() 

439 circle = self.window.circle 

440 arc = self.window.arc 

441 line = self.window.line 

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

443 

444 selected = self.images.selected 

445 visible = self.images.visible 

446 ncell = len(self.X_cell) 

447 bond_linewidth = self.scale * 0.15 

448 

449 self.update_labels() 

450 

451 if self.arrowkey_mode == self.ARROWKEY_MOVE: 

452 movecolor = GREEN 

453 elif self.arrowkey_mode == self.ARROWKEY_ROTATE: 

454 movecolor = PURPLE 

455 

456 for a in self.indices: 

457 if a < n: 

458 ra = d[a] 

459 if visible[a]: 

460 try: 

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

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

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

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

465 fill = '#ffffff' 

466 circle(fill, selected[a], 

467 A[a, 0], A[a, 1], 

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

469 start = 0 

470 # start with the dominant species 

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

472 key=lambda x: x[1], 

473 reverse=True): 

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

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

476 A[a, 0], A[a, 1], 

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

478 else: 

479 # jmol colors for the moment 

480 extent = 360. * occ 

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

482 selected[a], 

483 start, extent, 

484 A[a, 0], A[a, 1], 

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

486 start += extent 

487 except KeyError: 

488 # legacy behavior 

489 # Draw the atoms 

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

491 and self.move_atoms_mask[a]): 

492 circle(movecolor, False, 

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

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

495 

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

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

498 

499 # Draw labels on the atoms 

500 if self.labels is not None: 

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

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

503 str(self.labels[a])) 

504 

505 # Draw cross on constrained atoms 

506 if constrained[a]: 

507 R1 = int(0.14644 * ra) 

508 R2 = int(0.85355 * ra) 

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

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

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

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

513 

514 # Draw velocities and/or forces 

515 for v in vector_arrays: 

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

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

518 width=2) 

519 else: 

520 # Draw unit cell and/or bonds: 

521 a -= n 

522 if a < ncell: 

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

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

525 else: 

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

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

528 width=bond_linewidth) 

529 

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

531 self.draw_axes() 

532 

533 if len(self.images) > 1: 

534 self.draw_frame_number() 

535 

536 self.window.update() 

537 

538 if status: 

539 self.status(self.atoms) 

540 

541 def arrow(self, coords, width): 

542 line = self.window.line 

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

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

545 line(coords, width) 

546 

547 vec = end - begin 

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

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

550 

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

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

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

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

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

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

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

558 

559 def draw_axes(self): 

560 axes_length = 15 

561 

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

563 

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

565 a = 20 

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

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

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

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

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

571 

572 def draw_frame_number(self): 

573 x, y = self.window.size 

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

575 anchor='SE') 

576 

577 def release(self, event): 

578 if event.button in [4, 5]: 

579 self.scroll_event(event) 

580 return 

581 

582 if event.button != self.b1: 

583 return 

584 

585 selected = self.images.selected 

586 selected_ordered = self.images.selected_ordered 

587 

588 if event.time < self.t0 + 200: # 200 ms 

589 d = self.P - self.xy 

590 r = self.get_covalent_radii() 

591 hit = np.less((d**2).sum(1), (self.scale * r)**2) 

592 for a in self.indices[::-1]: 

593 if a < len(self.atoms) and hit[a]: 

594 if event.modifier == 'ctrl': 

595 selected[a] = not selected[a] 

596 if selected[a]: 

597 selected_ordered += [a] 

598 elif len(selected_ordered) > 0: 

599 if selected_ordered[-1] == a: 

600 selected_ordered = selected_ordered[:-1] 

601 else: 

602 selected_ordered = [] 

603 else: 

604 selected[:] = False 

605 selected[a] = True 

606 selected_ordered = [a] 

607 break 

608 else: 

609 selected[:] = False 

610 selected_ordered = [] 

611 self.draw() 

612 else: 

613 A = (event.x, event.y) 

614 C1 = np.minimum(A, self.xy) 

615 C2 = np.maximum(A, self.xy) 

616 hit = np.logical_and(self.P > C1, self.P < C2) 

617 indices = np.compress(hit.prod(1), np.arange(len(hit))) 

618 if event.modifier != 'ctrl': 

619 selected[:] = False 

620 selected[indices] = True 

621 if (len(indices) == 1 and 

622 indices[0] not in self.images.selected_ordered): 

623 selected_ordered += [indices[0]] 

624 elif len(indices) > 1: 

625 selected_ordered = [] 

626 self.draw() 

627 

628 # XXX check bounds 

629 natoms = len(self.atoms) 

630 indices = np.arange(natoms)[self.images.selected[:natoms]] 

631 if len(indices) != len(selected_ordered): 

632 selected_ordered = [] 

633 self.images.selected_ordered = selected_ordered