Coverage for /builds/ase/ase/ase/gui/modify.py: 95.65%

1# fmt: off 

2 

3from functools import partial 

4 

5import numpy as np 

6 

7import ase.gui.ui as ui 

8from ase.gui.i18n import _ 

9from ase.gui.utils import get_magmoms 

10from ase.gui.widgets import Element 

11 

12 

13class ModifyAtoms: 

14 """Presents a dialog box where the user is able to change the 

15 atomic type, the magnetic moment and tags of the selected atoms. 

16 """ 

17 

18 def __init__(self, gui): 

19 self.gui = gui 

20 selected = self.selection() 

21 if not selected.any(): 

22 ui.error(_('No atoms selected!')) 

23 return 

24 

25 win = ui.Window(_('Modify'), wmtype='utility') 

26 element = Element(callback=self.set_element) 

27 win.add(element) 

28 win.add(ui.Button(_('Change element'), 

29 partial(self.set_element, element))) 

30 self.tag = ui.SpinBox(0, -1000, 1000, 1, self.set_tag) 

31 win.add([_('Tag'), self.tag]) 

32 self.magmom = ui.SpinBox(0.0, -10, 10, 0.1, self.set_magmom) 

33 win.add([_('Moment'), self.magmom]) 

34 

35 atoms = self.gui.atoms 

36 sym = atoms.symbols[selected] 

37 if len(sym.species()) == 1: 

38 element.symbol = sym[0] 

39 

40 tags = atoms.get_tags()[selected] 

41 if np.ptp(tags) == 0: 

42 self.tag.value = tags[0] 

43 

44 magmoms = get_magmoms(atoms)[selected] 

45 if np.ptp(magmoms.round(2)) == 0.0: 

46 self.magmom.value = round(magmoms[0], 2) 

47 

48 def selection(self): 

49 return self.gui.images.selected[:len(self.gui.atoms)] 

50 

51 def set_element(self, element): 

52 self.gui.atoms.numbers[self.selection()] = element.Z 

53 self.gui.draw() 

54 

55 def set_tag(self): 

56 tags = self.gui.atoms.get_tags() 

57 tags[self.selection()] = self.tag.value 

58 self.gui.atoms.set_tags(tags) 

59 self.gui.draw() 

60 

61 def set_magmom(self): 

62 magmoms = get_magmoms(self.gui.atoms) 

63 magmoms[self.selection()] = self.magmom.value 

64 self.gui.atoms.set_initial_magnetic_moments(magmoms) 

65 self.gui.draw()