Coverage for /builds/ase/ase/ase/gui/gui.py: 64.94%

1# fmt: off

3import pickle

4import subprocess

5import sys

6from functools import partial

7from time import time

9import numpy as np

11import ase.gui.ui as ui

12from ase import Atoms, __version__

13from ase.gui.defaults import read_defaults

14from ase.gui.i18n import _

15from ase.gui.images import Images

16from ase.gui.nanoparticle import SetupNanoparticle

17from ase.gui.nanotube import SetupNanotube

18from ase.gui.observer import Observers

19from ase.gui.save import save_dialog

20from ase.gui.settings import Settings

21from ase.gui.status import Status

22from ase.gui.surfaceslab import SetupSurfaceSlab

23from ase.gui.view import View

26class GUIObservers:

27 def __init__(self):

28 self.new_atoms = Observers()

29 self.set_atoms = Observers()

30 self.change_atoms = Observers()

33class GUI(View):

34 ARROWKEY_SCAN = 0

35 ARROWKEY_MOVE = 1

36 ARROWKEY_ROTATE = 2

38 def __init__(self, images=None,

39 rotations='',

40 show_bonds=False, expr=None):

42 if not isinstance(images, Images):

43 images = Images(images)

45 self.images = images

47 # Ordinary observers seem unused now, delete?

48 self.observers = []

49 self.obs = GUIObservers()

51 self.config = read_defaults()

52 if show_bonds:

53 self.config['show_bonds'] = True

55 menu = self.get_menu_data()

57 self.window = ui.ASEGUIWindow(close=self.exit, menu=menu,

58 config=self.config, scroll=self.scroll,

59 scroll_event=self.scroll_event,

60 press=self.press, move=self.move,

61 release=self.release,

62 resize=self.resize)

64 super().__init__(rotations)

65 self.status = Status(self)

67 self.subprocesses = [] # list of external processes

68 self.movie_window = None

69 self.simulation = {} # Used by modules on Calculate menu.

70 self.module_state = {} # Used by modules to store their state.

72 self.arrowkey_mode = self.ARROWKEY_SCAN

73 self.move_atoms_mask = None

75 self.set_frame(len(self.images) - 1, focus=True)

77 # Used to move the structure with the mouse

78 self.prev_pos = None

79 self.last_scroll_time = time()

80 self.orig_scale = self.scale

82 if len(self.images) > 1:

83 self.movie()

85 if expr is None:

86 expr = self.config['gui_graphs_string']

88 if expr is not None and expr != '' and len(self.images) > 1:

89 self.plot_graphs(expr=expr, ignore_if_nan=True)

91 @property

92 def moving(self):

93 return self.arrowkey_mode != self.ARROWKEY_SCAN

95 def run(self):

96 self.window.run()

98 def toggle_move_mode(self, key=None):

99 self.toggle_arrowkey_mode(self.ARROWKEY_MOVE)

100

101 def toggle_rotate_mode(self, key=None):

102 self.toggle_arrowkey_mode(self.ARROWKEY_ROTATE)

103

104 def toggle_arrowkey_mode(self, mode):

105 # If not currently in given mode, activate it.

106 # Else, deactivate it (go back to SCAN mode)

107 assert mode != self.ARROWKEY_SCAN

108

109 if self.arrowkey_mode == mode:

110 self.arrowkey_mode = self.ARROWKEY_SCAN

111 self.move_atoms_mask = None

112 else:

113 self.arrowkey_mode = mode

114 self.move_atoms_mask = self.images.selected.copy()

115

116 self.draw()

117

118 def step(self, key):

119 d = {'Home': -10000000,

120 'Page-Up': -1,

121 'Page-Down': 1,

122 'End': 10000000}[key]

123 i = max(0, min(len(self.images) - 1, self.frame + d))

124 self.set_frame(i)

125 if self.movie_window is not None:

126 self.movie_window.frame_number.value = i

127

128 def copy_image(self, key=None):

129 self.images._images.append(self.atoms.copy())

130 self.images.filenames.append(None)

131

132 if self.movie_window is not None:

133 self.movie_window.frame_number.scale.configure(to=len(self.images))

134 self.step('End')

135

136 def _do_zoom(self, x):

137 """Utility method for zooming"""

138 self.scale *= x

139 self.draw()

140

141 def zoom(self, key):

142 """Zoom in/out on keypress or clicking menu item"""

143 x = {'+': 1.2, '-': 1 / 1.2}[key]

144 self._do_zoom(x)

145

146 def scroll_event(self, event):

147 """Zoom in/out when using mouse wheel"""

148 SHIFT = event.modifier == 'shift'

149 x = 1.0

150 if event.button == 4 or event.delta > 0:

151 x = 1.0 + (1 - SHIFT) * 0.2 + SHIFT * 0.01

152 elif event.button == 5 or event.delta < 0:

153 x = 1.0 / (1.0 + (1 - SHIFT) * 0.2 + SHIFT * 0.01)

154 self._do_zoom(x)

155

156 def settings(self):

157 return Settings(self)

158

159 def scroll(self, event):

160 shift = 0x1

161 ctrl = 0x4

162 alt_l = 0x8 # Also Mac Command Key

163 mac_option_key = 0x10

164

165 use_small_step = bool(event.state & shift)

166 rotate_into_plane = bool(event.state & (ctrl | alt_l | mac_option_key))

167

168 dxdydz = {'up': (0, 1 - rotate_into_plane, rotate_into_plane),

169 'down': (0, -1 + rotate_into_plane, -rotate_into_plane),

170 'right': (1, 0, 0),

171 'left': (-1, 0, 0)}.get(event.key, None)

172

173 # Get scroll direction using shift + right mouse button

174 # event.type == '6' is mouse motion, see:

175 # http://infohost.nmt.edu/tcc/help/pubs/tkinter/web/event-types.html

176 if event.type == '6':

177 cur_pos = np.array([event.x, -event.y])

178 # Continue scroll if button has not been released

179 if self.prev_pos is None or time() - self.last_scroll_time > .5:

180 self.prev_pos = cur_pos

181 self.last_scroll_time = time()

182 else:

183 dxdy = cur_pos - self.prev_pos

184 dxdydz = np.append(dxdy, [0])

185 self.prev_pos = cur_pos

186 self.last_scroll_time = time()

187

188 if dxdydz is None:

189 return

190

191 vec = 0.1 * np.dot(self.axes, dxdydz)

192 if use_small_step:

193 vec *= 0.1

194

195 if self.arrowkey_mode == self.ARROWKEY_MOVE:

196 self.atoms.positions[self.move_atoms_mask[:len(self.atoms)]] += vec

197 self.set_frame()

198 elif self.arrowkey_mode == self.ARROWKEY_ROTATE:

199 # For now we use atoms.rotate having the simplest interface.

200 # (Better to use something more minimalistic, obviously.)

201 mask = self.move_atoms_mask[:len(self.atoms)]

202 center = self.atoms.positions[mask].mean(axis=0)

203 tmp_atoms = self.atoms[mask]

204 tmp_atoms.positions -= center

205 tmp_atoms.rotate(50 * np.linalg.norm(vec), vec)

206 self.atoms.positions[mask] = tmp_atoms.positions + center

207 self.set_frame()

208 else:

209 # The displacement vector is scaled

210 # so that the cursor follows the structure

211 # Scale by a third works for some reason

212 scale = self.orig_scale / (3 * self.scale)

213 self.center -= vec * scale

214

215 # dx * 0.1 * self.axes[:, 0] - dy * 0.1 * self.axes[:, 1])

216

217 self.draw()

218

219 def delete_selected_atoms(self, widget=None, data=None):

220 import ase.gui.ui as ui

221 nselected = sum(self.images.selected)

222 if nselected and ui.ask_question(_('Delete atoms'),

223 _('Delete selected atoms?')):

224 self.really_delete_selected_atoms()

225

226 def really_delete_selected_atoms(self):

227 mask = self.images.selected[:len(self.atoms)]

228 del self.atoms[mask]

229

230 # Will remove selection in other images, too

231 self.images.selected[:] = False

232 self.set_frame()

233 self.draw()

234

235 def constraints_window(self):

236 from ase.gui.constraints import Constraints

237 return Constraints(self)

238

239 def set_selected_atoms(self, selected):

240 newmask = np.zeros(len(self.images.selected), bool)

241 newmask[selected] = True

242

243 if np.array_equal(newmask, self.images.selected):

244 return

245

246 # (By creating newmask, we can avoid resetting the selection in

247 # case the selected indices are invalid)

248 self.images.selected[:] = newmask

249 self.draw()

250

251 def select_all(self, key=None):

252 self.images.selected[:] = True

253 self.draw()

254

255 def invert_selection(self, key=None):

256 self.images.selected[:] = ~self.images.selected

257 self.draw()

258

259 def select_constrained_atoms(self, key=None):

260 self.images.selected[:] = ~self.images.get_dynamic(self.atoms)

261 self.draw()

262

263 def select_immobile_atoms(self, key=None):

264 if len(self.images) > 1:

265 R0 = self.images[0].positions

266 for atoms in self.images[1:]:

267 R = atoms.positions

268 self.images.selected[:] = ~(np.abs(R - R0) > 1.0e-10).any(1)

269 self.draw()

270

271 def movie(self):

272 from ase.gui.movie import Movie

273 self.movie_window = Movie(self)

274

275 def plot_graphs(self, key=None, expr=None, ignore_if_nan=False):

276 from ase.gui.graphs import Graphs

277 g = Graphs(self)

278 if expr is not None:

279 g.plot(expr=expr, ignore_if_nan=ignore_if_nan)

280

281 def pipe(self, task, data):

282 process = subprocess.Popen([sys.executable, '-m', 'ase.gui.pipe'],

283 stdout=subprocess.PIPE,

284 stdin=subprocess.PIPE)

285 pickle.dump((task, data), process.stdin)

286 process.stdin.close()

287 # Either process writes a line, or it crashes and line becomes ''

288 line = process.stdout.readline().decode('utf8').strip()

289

290 if line != 'GUI:OK':

291 if line == '': # Subprocess probably crashed

292 line = _('Failure in subprocess')

293 self.bad_plot(line)

294 else:

295 self.subprocesses.append(process)

296 return process

297

298 def bad_plot(self, err, msg=''):

299 ui.error(_('Plotting failed'), '\n'.join([str(err), msg]).strip())

300

301 def neb(self):

302 from ase.utils.forcecurve import fit_images

303 try:

304 forcefit = fit_images(self.images)

305 except Exception as err:

306 self.bad_plot(err, _('Images must have energies and forces, '

307 'and atoms must not be stationary.'))

308 else:

309 self.pipe('neb', forcefit)

310

311 def bulk_modulus(self):

312 try:

313 v = [abs(np.linalg.det(atoms.cell)) for atoms in self.images]

314 e = [self.images.get_energy(a) for a in self.images]

315 from ase.eos import EquationOfState

316 eos = EquationOfState(v, e)

317 plotdata = eos.getplotdata()

318 except Exception as err:

319 self.bad_plot(err, _('Images must have energies '

320 'and varying cell.'))

321 else:

322 self.pipe('eos', plotdata)

323

324 def reciprocal(self):

325 if self.atoms.cell.rank != 3:

326 self.bad_plot(_('Requires 3D cell.'))

327 return None

328

329 cell = self.atoms.cell.uncomplete(self.atoms.pbc)

330 bandpath = cell.bandpath(npoints=0)

331 return self.pipe('reciprocal', bandpath)

332

333 def open(self, button=None, filename=None):

334 chooser = ui.ASEFileChooser(self.window.win)

335

336 filename = filename or chooser.go()

337 format = chooser.format

338 if filename:

339 try:

340 self.images.read([filename], slice(None), format)

341 except Exception as err:

342 ui.show_io_error(filename, err)

343 return # Hmm. Is self.images in a consistent state?

344 self.set_frame(len(self.images) - 1, focus=True)

345

346 def modify_atoms(self, key=None):

347 from ase.gui.modify import ModifyAtoms

348 return ModifyAtoms(self)

349

350 def add_atoms(self, key=None):

351 from ase.gui.add import AddAtoms

352 return AddAtoms(self)

353

354 def cell_editor(self, key=None):

355 from ase.gui.celleditor import CellEditor

356 return CellEditor(self)

357

358 def atoms_editor(self, key=None):

359 from ase.gui.atomseditor import AtomsEditor

360 return AtomsEditor(self)

361

362 def quick_info_window(self, key=None):

363 from ase.gui.quickinfo import info

364 info_win = ui.Window(_('Quick Info'), wmtype='utility')

365 info_win.add(info(self))

366

367 # Update quickinfo window when we change frame

368 def update(window):

369 exists = window.exists

370 if exists:

371 # Only update if we exist

372 window.things[0].text = info(self)

373 return exists

374 self.attach(update, info_win)

375 return info_win

376

377 def surface_window(self):

378 return SetupSurfaceSlab(self)

379

380 def nanoparticle_window(self):

381 return SetupNanoparticle(self)

382

383 def nanotube_window(self):

384 return SetupNanotube(self)

385

386 def new_atoms(self, atoms):

387 "Set a new atoms object."

388 rpt = getattr(self.images, 'repeat', None)

389 self.images.repeat_images(np.ones(3, int))

390 self.images.initialize([atoms])

391 self.frame = 0 # Prevent crashes

392 self.images.repeat_images(rpt)

393 self.set_frame(frame=0, focus=True)

394 self.obs.new_atoms.notify()

395

396 def exit(self, event=None):

397 for process in self.subprocesses:

398 process.terminate()

399 self.window.close()

400

401 def new(self, key=None):

402 subprocess.Popen([sys.executable, '-m', 'ase', 'gui'])

403

404 def save(self, key=None):

405 return save_dialog(self)

406

407 def external_viewer(self, name):

408 from ase.visualize import view

409 return view(list(self.images), viewer=name)

410

411 def selected_atoms(self):

412 selection_mask = self.images.selected[:len(self.atoms)]

413 return self.atoms[selection_mask]

414

415 def wrap_atoms(self, key=None):

416 """Wrap atoms around the unit cell."""

417 for atoms in self.images:

418 atoms.wrap()

419 self.set_frame()

420

421 @property

422 def clipboard(self):

423 from ase.gui.clipboard import AtomsClipboard

424 return AtomsClipboard(self.window.win)

425

426 def cut_atoms_to_clipboard(self, event=None):

427 self.copy_atoms_to_clipboard(event)

428 self.really_delete_selected_atoms()

429

430 def copy_atoms_to_clipboard(self, event=None):

431 atoms = self.selected_atoms()

432 self.clipboard.set_atoms(atoms)

433

434 def paste_atoms_from_clipboard(self, event=None):

435 try:

436 atoms = self.clipboard.get_atoms()

437 except Exception as err:

438 ui.error(

439 'Cannot paste atoms',

440 'Pasting currently works only with the ASE JSON format.\n\n'

441 f'Original error:\n\n{err}')

442 return

443

444 if self.atoms == Atoms():

445 self.atoms.cell = atoms.cell

446 self.atoms.pbc = atoms.pbc

447 self.paste_atoms_onto_existing(atoms)

448

449 def paste_atoms_onto_existing(self, atoms):

450 selection = self.selected_atoms()

451 if len(selection):

452 paste_center = selection.positions.sum(axis=0) / len(selection)

453 # atoms.center() is a no-op in directions without a cell vector.

454 # But we actually want the thing centered nevertheless!

455 # Therefore we have to set the cell.

456 atoms = atoms.copy()

457 atoms.cell = (1, 1, 1) # arrrgh.

458 atoms.center(about=paste_center)

459

460 self.add_atoms_and_select(atoms)

461 self.move_atoms_mask = self.images.selected.copy()

462 self.arrowkey_mode = self.ARROWKEY_MOVE

463 self.draw()

464

465 def add_atoms_and_select(self, new_atoms):

466 atoms = self.atoms

467 atoms += new_atoms

468

469 if len(atoms) > self.images.maxnatoms:

470 self.images.initialize(list(self.images),

471 self.images.filenames)

472

473 selected = self.images.selected

474 selected[:] = False

475 # 'selected' array may be longer than current atoms

476 selected[len(atoms) - len(new_atoms):len(atoms)] = True

477

478 self.set_frame()

479 self.draw()

480

481 def get_menu_data(self):

482 M = ui.MenuItem

483 return [

484 (_('_File'),

485 [M(_('_Open'), self.open, 'Ctrl+O'),

486 M(_('_New'), self.new, 'Ctrl+N'),

487 M(_('_Save'), self.save, 'Ctrl+S'),

488 M('---'),

489 M(_('_Quit'), self.exit, 'Ctrl+Q')]),

490

491 (_('_Edit'),

492 [M(_('Select _all'), self.select_all),

493 M(_('_Invert selection'), self.invert_selection),

494 M(_('Select _constrained atoms'), self.select_constrained_atoms),

495 M(_('Select _immobile atoms'), self.select_immobile_atoms),

496 # M('---'),

497 M(_('_Cut'), self.cut_atoms_to_clipboard, 'Ctrl+X'),

498 M(_('_Copy'), self.copy_atoms_to_clipboard, 'Ctrl+C'),

499 M(_('_Paste'), self.paste_atoms_from_clipboard, 'Ctrl+V'),

500 M('---'),

501 M(_('Hide selected atoms'), self.hide_selected),

502 M(_('Show selected atoms'), self.show_selected),

503 M('---'),

504 M(_('_Modify'), self.modify_atoms, 'Ctrl+Y'),

505 M(_('_Add atoms'), self.add_atoms, 'Ctrl+A'),

506 M(_('_Delete selected atoms'), self.delete_selected_atoms,

507 'Backspace'),

508 M(_('Edit _cell …'), self.cell_editor, 'Ctrl+E'),

509 M(_('Edit _atoms …'), self.atoms_editor, 'A'),

510 M('---'),

511 M(_('_First image'), self.step, 'Home'),

512 M(_('_Previous image'), self.step, 'Page-Up'),

513 M(_('_Next image'), self.step, 'Page-Down'),

514 M(_('_Last image'), self.step, 'End'),

515 M(_('Append image copy'), self.copy_image)]),

516

517 (_('_View'),

518 [M(_('Show _unit cell'), self.toggle_show_unit_cell, 'Ctrl+U',

519 value=self.config['show_unit_cell']),

520 M(_('Show _axes'), self.toggle_show_axes,

521 value=self.config['show_axes']),

522 M(_('Show _bonds'), self.toggle_show_bonds, 'Ctrl+B',

523 value=self.config['show_bonds']),

524 M(_('Show _velocities'), self.toggle_show_velocities, 'Ctrl+G',

525 value=False),

526 M(_('Show _forces'), self.toggle_show_forces, 'Ctrl+F',

527 value=False),

528 M(_('Show _magmoms'), self.toggle_show_magmoms,

529 value=False),

530 M(_('Show _Labels'), self.show_labels,

531 choices=[_('_None'),

532 _('Atom _Index'),

533 _('_Magnetic Moments'), # XXX check if exist

534 _('_Element Symbol'),

535 _('_Initial Charges'), # XXX check if exist

536 ]),

537 M('---'),

538 M(_('Quick Info ...'), self.quick_info_window, 'Ctrl+I'),

539 M(_('Repeat ...'), self.repeat_window, 'R'),

540 M(_('Rotate ...'), self.rotate_window),

541 M(_('Colors ...'), self.colors_window, 'C'),

542 # TRANSLATORS: verb

543 M(_('Focus'), self.focus, 'F'),

544 M(_('Zoom in'), self.zoom, '+'),

545 M(_('Zoom out'), self.zoom, '-'),

546 M(_('Change View'),

547 submenu=[

548 M(_('Reset View'), self.reset_view, '='),

549 M(_('xy-plane'), self.set_view, 'Z'),

550 M(_('yz-plane'), self.set_view, 'X'),

551 M(_('zx-plane'), self.set_view, 'Y'),

552 M(_('yx-plane'), self.set_view, 'Shift+Z'),

553 M(_('zy-plane'), self.set_view, 'Shift+X'),

554 M(_('xz-plane'), self.set_view, 'Shift+Y'),

555 M(_('a2,a3-plane'), self.set_view, 'I'),

556 M(_('a3,a1-plane'), self.set_view, 'J'),

557 M(_('a1,a2-plane'), self.set_view, 'K'),

558 M(_('a3,a2-plane'), self.set_view, 'Shift+I'),

559 M(_('a1,a3-plane'), self.set_view, 'Shift+J'),

560 M(_('a2,a1-plane'), self.set_view, 'Shift+K')]),

561 M(_('Settings ...'), self.settings),

562 M('---'),

563 M(_('VMD'), partial(self.external_viewer, 'vmd')),

564 M(_('RasMol'), partial(self.external_viewer, 'rasmol')),

565 M(_('xmakemol'), partial(self.external_viewer, 'xmakemol')),

566 M(_('avogadro'), partial(self.external_viewer, 'avogadro'))]),

567

568 (_('_Tools'),

569 [M(_('Graphs ...'), self.plot_graphs),

570 M(_('Movie ...'), self.movie),

571 M(_('Constraints ...'), self.constraints_window),

572 M(_('Render scene ...'), self.render_window),

573 M(_('_Move selected atoms'), self.toggle_move_mode, 'Ctrl+M'),

574 M(_('_Rotate selected atoms'), self.toggle_rotate_mode,

575 'Ctrl+R'),

576 M(_('NE_B plot'), self.neb),

577 M(_('B_ulk Modulus'), self.bulk_modulus),

578 M(_('Reciprocal space ...'), self.reciprocal),

579 M(_('Wrap atoms'), self.wrap_atoms, 'Ctrl+W')]),

580

581 # TRANSLATORS: Set up (i.e. build) surfaces, nanoparticles, ...

582 (_('_Setup'),

583 [M(_('_Surface slab'), self.surface_window, disabled=False),

584 M(_('_Nanoparticle'),

585 self.nanoparticle_window),

586 M(_('Nano_tube'), self.nanotube_window)]),

587

588 # (_('_Calculate'),

589 # [M(_('Set _Calculator'), self.calculator_window, disabled=True),

590 # M(_('_Energy and Forces'), self.energy_window, disabled=True),

591 # M(_('Energy Minimization'), self.energy_minimize_window,

592 # disabled=True)]),

593

594 (_('_Help'),

595 [M(_('_About'), partial(ui.about, 'ASE-GUI',

596 version=__version__,

597 webpage='https://wiki.fysik.dtu.dk/'

598 'ase/ase/gui/gui.html')),

599 M(_('Webpage ...'), webpage)])]

600

601 def attach(self, function, *args, **kwargs):

602 self.observers.append((function, args, kwargs))

603

604 def call_observers(self):

605 # Use function return value to determine if we keep observer

606 self.observers = [(function, args, kwargs) for (function, args, kwargs)

607 in self.observers if function(*args, **kwargs)]

608

609 def repeat_poll(self, callback, ms, ensure_update=True):

610 """Invoke callback(gui=self) every ms milliseconds.

611

612 This is useful for polling a resource for updates to load them

613 into the GUI. The GUI display will be hence be updated after

614 each call; pass ensure_update=False to circumvent this.

615

616 Polling stops if the callback function raises StopIteration.

617

618 Example to run a movie manually, then quit::

619

620 from ase.collections import g2

621 from ase.gui.gui import GUI

622

623 names = iter(g2.names)

624

625 def main(gui):

626 try:

627 name = next(names)

628 except StopIteration:

629 gui.window.win.quit()

630 else:

631 atoms = g2[name]

632 gui.images.initialize([atoms])

633

634 gui = GUI()

635 gui.repeat_poll(main, 30)

636 gui.run()"""

637

638 def callbackwrapper():

639 try:

640 callback(gui=self)

641 except StopIteration:

642 pass

643 finally:

644 # Reinsert self so we get called again:

645 self.window.win.after(ms, callbackwrapper)

646

647 if ensure_update:

648 self.set_frame()

649 self.draw()

650

651 self.window.win.after(ms, callbackwrapper)

652

653

654def webpage():

655 import webbrowser

656 webbrowser.open('https://wiki.fysik.dtu.dk/ase/ase/gui/gui.html')