Coverage for /builds/kinetik161/ase/ase/gui/gui.py: 63.77%

1import pickle

2import subprocess

3import sys

4import weakref

5from functools import partial

6from time import time

8import numpy as np

10import ase.gui.ui as ui

11from ase import Atoms, __version__

12from ase.gui.defaults import read_defaults

13from ase.gui.i18n import _

14from ase.gui.images import Images

15from ase.gui.nanoparticle import SetupNanoparticle

16from ase.gui.nanotube import SetupNanotube

17from ase.gui.save import save_dialog

18from ase.gui.settings import Settings

19from ase.gui.status import Status

20from ase.gui.surfaceslab import SetupSurfaceSlab

21from ase.gui.view import View

24class GUI(View, Status):

25 ARROWKEY_SCAN = 0

26 ARROWKEY_MOVE = 1

27 ARROWKEY_ROTATE = 2

29 def __init__(self, images=None,

30 rotations='',

31 show_bonds=False, expr=None):

33 if not isinstance(images, Images):

34 images = Images(images)

36 self.images = images

37 self.observers = []

39 self.config = read_defaults()

40 if show_bonds:

41 self.config['show_bonds'] = True

43 menu = self.get_menu_data()

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

46 config=self.config, scroll=self.scroll,

47 scroll_event=self.scroll_event,

48 press=self.press, move=self.move,

49 release=self.release,

50 resize=self.resize)

52 View.__init__(self, rotations)

53 Status.__init__(self)

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

56 self.movie_window = None

57 self.vulnerable_windows = []

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

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

61 self.arrowkey_mode = self.ARROWKEY_SCAN

62 self.move_atoms_mask = None

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

66 # Used to move the structure with the mouse

67 self.prev_pos = None

68 self.last_scroll_time = time()

69 self.orig_scale = self.scale

71 if len(self.images) > 1:

72 self.movie()

74 if expr is None:

75 expr = self.config['gui_graphs_string']

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

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

80 @property

81 def moving(self):

82 return self.arrowkey_mode != self.ARROWKEY_SCAN

84 def run(self):

85 self.window.run()

87 def toggle_move_mode(self, key=None):

88 self.toggle_arrowkey_mode(self.ARROWKEY_MOVE)

90 def toggle_rotate_mode(self, key=None):

91 self.toggle_arrowkey_mode(self.ARROWKEY_ROTATE)

93 def toggle_arrowkey_mode(self, mode):

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

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

96 assert mode != self.ARROWKEY_SCAN

98 if self.arrowkey_mode == mode:

99 self.arrowkey_mode = self.ARROWKEY_SCAN

100 self.move_atoms_mask = None

101 else:

102 self.arrowkey_mode = mode

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

104

105 self.draw()

106

107 def step(self, key):

108 d = {'Home': -10000000,

109 'Page-Up': -1,

110 'Page-Down': 1,

111 'End': 10000000}[key]

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

113 self.set_frame(i)

114 if self.movie_window is not None:

115 self.movie_window.frame_number.value = i + 1

116

117 def copy_image(self, key=None):

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

119 self.images.filenames.append(None)

120

121 if self.movie_window is not None:

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

123 self.step('End')

124

125 def _do_zoom(self, x):

126 """Utility method for zooming"""

127 self.scale *= x

128 self.draw()

129

130 def zoom(self, key):

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

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

133 self._do_zoom(x)

134

135 def scroll_event(self, event):

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

137 SHIFT = event.modifier == 'shift'

138 x = 1.0

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

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

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

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

143 self._do_zoom(x)

144

145 def settings(self):

146 return Settings(self)

147

148 def scroll(self, event):

149 CTRL = event.modifier == 'ctrl'

150

151 # Bug: Simultaneous CTRL + shift is the same as just CTRL.

152 # Therefore movement in Z direction does not support the

153 # shift modifier.

154 dxdydz = {'up': (0, 1 - CTRL, CTRL),

155 'down': (0, -1 + CTRL, -CTRL),

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

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

158

159 # Get scroll direction using shift + right mouse button

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

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

162 if event.type == '6':

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

164 # Continue scroll if button has not been released

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

166 self.prev_pos = cur_pos

167 self.last_scroll_time = time()

168 else:

169 dxdy = cur_pos - self.prev_pos

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

171 self.prev_pos = cur_pos

172 self.last_scroll_time = time()

173

174 if dxdydz is None:

175 return

176

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

178 if event.modifier == 'shift':

179 vec *= 0.1

180

181 if self.arrowkey_mode == self.ARROWKEY_MOVE:

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

183 self.set_frame()

184 elif self.arrowkey_mode == self.ARROWKEY_ROTATE:

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

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

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

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

189 tmp_atoms = self.atoms[mask]

190 tmp_atoms.positions -= center

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

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

193 self.set_frame()

194 else:

195 # The displacement vector is scaled

196 # so that the cursor follows the structure

197 # Scale by a third works for some reason

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

199 self.center -= vec * scale

200

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

202

203 self.draw()

204

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

206 import ase.gui.ui as ui

207 nselected = sum(self.images.selected)

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

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

210 self.really_delete_selected_atoms()

211

212 def really_delete_selected_atoms(self):

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

214 del self.atoms[mask]

215

216 # Will remove selection in other images, too

217 self.images.selected[:] = False

218 self.set_frame()

219 self.draw()

220

221 def constraints_window(self):

222 from ase.gui.constraints import Constraints

223 return Constraints(self)

224

225 def select_all(self, key=None):

226 self.images.selected[:] = True

227 self.draw()

228

229 def invert_selection(self, key=None):

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

231 self.draw()

232

233 def select_constrained_atoms(self, key=None):

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

235 self.draw()

236

237 def select_immobile_atoms(self, key=None):

238 if len(self.images) > 1:

239 R0 = self.images[0].positions

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

241 R = atoms.positions

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

243 self.draw()

244

245 def movie(self):

246 from ase.gui.movie import Movie

247 self.movie_window = Movie(self)

248

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

250 from ase.gui.graphs import Graphs

251 g = Graphs(self)

252 if expr is not None:

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

254

255 def pipe(self, task, data):

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

257 stdout=subprocess.PIPE,

258 stdin=subprocess.PIPE)

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

260 process.stdin.close()

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

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

263

264 if line != 'GUI:OK':

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

266 line = _('Failure in subprocess')

267 self.bad_plot(line)

268 else:

269 self.subprocesses.append(process)

270 return process

271

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

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

274

275 def neb(self):

276 from ase.utils.forcecurve import fit_images

277 try:

278 forcefit = fit_images(self.images)

279 except Exception as err:

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

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

282 else:

283 self.pipe('neb', forcefit)

284

285 def bulk_modulus(self):

286 try:

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

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

289 from ase.eos import EquationOfState

290 eos = EquationOfState(v, e)

291 plotdata = eos.getplotdata()

292 except Exception as err:

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

294 'and varying cell.'))

295 else:

296 self.pipe('eos', plotdata)

297

298 def reciprocal(self):

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

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

301 return

302

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

304 bandpath = cell.bandpath(npoints=0)

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

306

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

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

309

310 filename = filename or chooser.go()

311 format = chooser.format

312 if filename:

313 try:

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

315 except Exception as err:

316 ui.show_io_error(filename, err)

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

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

319

320 def modify_atoms(self, key=None):

321 from ase.gui.modify import ModifyAtoms

322 return ModifyAtoms(self)

323

324 def add_atoms(self, key=None):

325 from ase.gui.add import AddAtoms

326 return AddAtoms(self)

327

328 def cell_editor(self, key=None):

329 from ase.gui.celleditor import CellEditor

330 return CellEditor(self)

331

332 def quick_info_window(self, key=None):

333 from ase.gui.quickinfo import info

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

335 info_win.add(info(self))

336

337 # Update quickinfo window when we change frame

338 def update(window):

339 exists = window.exists

340 if exists:

341 # Only update if we exist

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

343 return exists

344 self.attach(update, info_win)

345 return info_win

346

347 def surface_window(self):

348 return SetupSurfaceSlab(self)

349

350 def nanoparticle_window(self):

351 return SetupNanoparticle(self)

352

353 def nanotube_window(self):

354 return SetupNanotube(self)

355

356 def new_atoms(self, atoms):

357 "Set a new atoms object."

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

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

360 self.images.initialize([atoms])

361 self.frame = 0 # Prevent crashes

362 self.images.repeat_images(rpt)

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

364 self.notify_vulnerable()

365

366 def notify_vulnerable(self):

367 """Notify windows that would break when new_atoms is called.

368

369 The notified windows may adapt to the new atoms. If that is not

370 possible, they should delete themselves.

371 """

372 new_vul = [] # Keep weakrefs to objects that still exist.

373 for wref in self.vulnerable_windows:

374 ref = wref()

375 if ref is not None:

376 new_vul.append(wref)

377 ref.notify_atoms_changed()

378 self.vulnerable_windows = new_vul

379

380 def register_vulnerable(self, obj):

381 """Register windows that are vulnerable to changing the images.

382

383 Some windows will break if the atoms (and in particular the

384 number of images) are changed. They can register themselves

385 and be closed when that happens.

386 """

387 self.vulnerable_windows.append(weakref.ref(obj))

388

389 def exit(self, event=None):

390 for process in self.subprocesses:

391 process.terminate()

392 self.window.close()

393

394 def new(self, key=None):

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

396

397 def save(self, key=None):

398 return save_dialog(self)

399

400 def external_viewer(self, name):

401 from ase.visualize import view

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

403

404 def selected_atoms(self):

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

406 return self.atoms[selection_mask]

407

408 @property

409 def clipboard(self):

410 from ase.gui.clipboard import AtomsClipboard

411 return AtomsClipboard(self.window.win)

412

413 def cut_atoms_to_clipboard(self, event=None):

414 self.copy_atoms_to_clipboard(event)

415 self.really_delete_selected_atoms()

416

417 def copy_atoms_to_clipboard(self, event=None):

418 atoms = self.selected_atoms()

419 self.clipboard.set_atoms(atoms)

420

421 def paste_atoms_from_clipboard(self, event=None):

422 try:

423 atoms = self.clipboard.get_atoms()

424 except Exception as err:

425 ui.error(

426 'Cannot paste atoms',

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

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

429 return

430

431 if self.atoms == Atoms():

432 self.atoms.cell = atoms.cell

433 self.atoms.pbc = atoms.pbc

434 self.paste_atoms_onto_existing(atoms)

435

436 def paste_atoms_onto_existing(self, atoms):

437 selection = self.selected_atoms()

438 if len(selection):

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

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

441 # But we actually want the thing centered nevertheless!

442 # Therefore we have to set the cell.

443 atoms = atoms.copy()

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

445 atoms.center(about=paste_center)

446

447 self.add_atoms_and_select(atoms)

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

449 self.arrowkey_mode = self.ARROWKEY_MOVE

450 self.draw()

451

452 def add_atoms_and_select(self, new_atoms):

453 atoms = self.atoms

454 atoms += new_atoms

455

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

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

458 self.images.filenames)

459

460 selected = self.images.selected

461 selected[:] = False

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

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

464

465 self.set_frame()

466 self.draw()

467

468 def get_menu_data(self):

469 M = ui.MenuItem

470 return [

471 (_('_File'),

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

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

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

475 M('---'),

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

477

478 (_('_Edit'),

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

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

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

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

483 # M('---'),

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

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

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

487 M('---'),

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

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

490 M('---'),

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

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

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

494 'Backspace'),

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

496 M('---'),

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

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

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

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

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

502

503 (_('_View'),

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

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

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

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

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

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

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

511 value=False),

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

513 value=False),

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

515 choices=[_('_None'),

516 _('Atom _Index'),

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

518 _('_Element Symbol'),

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

520 ]),

521 M('---'),

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

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

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

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

526 # TRANSLATORS: verb

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

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

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

530 M(_('Change View'),

531 submenu=[

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

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

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

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

536 M(_('yx-plane'), self.set_view, 'Alt+Z'),

537 M(_('zy-plane'), self.set_view, 'Alt+X'),

538 M(_('xz-plane'), self.set_view, 'Alt+Y'),

539 M(_('a2,a3-plane'), self.set_view, '1'),

540 M(_('a3,a1-plane'), self.set_view, '2'),

541 M(_('a1,a2-plane'), self.set_view, '3'),

542 M(_('a3,a2-plane'), self.set_view, 'Alt+1'),

543 M(_('a1,a3-plane'), self.set_view, 'Alt+2'),

544 M(_('a2,a1-plane'), self.set_view, 'Alt+3')]),

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

546 M('---'),

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

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

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

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

551

552 (_('_Tools'),

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

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

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

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

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

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

559 'Ctrl+R'),

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

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

562 M(_('Reciprocal space ...'), self.reciprocal)]),

563

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

565 (_('_Setup'),

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

567 M(_('_Nanoparticle'),

568 self.nanoparticle_window),

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

570

571 # (_('_Calculate'),

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

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

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

575 # disabled=True)]),

576

577 (_('_Help'),

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

579 version=__version__,

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

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

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

583

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

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

586

587 def call_observers(self):

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

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

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

591

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

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

594

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

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

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

598

599 Polling stops if the callback function raises StopIteration.

600

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

602

603 from ase.collections import g2

604 from ase.gui.gui import GUI

605

606 names = iter(g2.names)

607

608 def main(gui):

609 try:

610 name = next(names)

611 except StopIteration:

612 gui.window.win.quit()

613 else:

614 atoms = g2[name]

615 gui.images.initialize([atoms])

616

617 gui = GUI()

618 gui.repeat_poll(main, 30)

619 gui.run()"""

620

621 def callbackwrapper():

622 try:

623 callback(gui=self)

624 except StopIteration:

625 pass

626 finally:

627 # Reinsert self so we get called again:

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

629

630 if ensure_update:

631 self.set_frame()

632 self.draw()

633

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

635

636

637def webpage():

638 import webbrowser

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