1from functools import partial 

2 

3from ase.gui.i18n import _ 

4 

5import ase.gui.ui as ui 

6from ase.gui.widgets import Element 

7from ase.gui.utils import get_magmoms 

8 

9 

10class ModifyAtoms: 

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

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

13 """ 

14 

15 def __init__(self, gui): 

16 self.gui = gui 

17 selected = self.selection() 

18 if not selected.any(): 

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

20 return 

21 

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

23 element = Element(callback=self.set_element) 

24 win.add(element) 

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

26 partial(self.set_element, element))) 

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

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

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

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

31 

32 atoms = self.gui.atoms 

33 sym = atoms.symbols[selected] 

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

35 element.symbol = sym[0] 

36 

37 tags = atoms.get_tags()[selected] 

38 if tags.ptp() == 0: 

39 self.tag.value = tags[0] 

40 

41 magmoms = get_magmoms(atoms)[selected] 

42 if magmoms.round(2).ptp() == 0.0: 

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

44 

45 def selection(self): 

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

47 

48 def set_element(self, element): 

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

50 self.gui.draw() 

51 

52 def set_tag(self): 

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

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

55 self.gui.atoms.set_tags(tags) 

56 self.gui.draw() 

57 

58 def set_magmom(self): 

59 magmoms = get_magmoms(self.gui.atoms) 

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

61 self.gui.atoms.set_initial_magnetic_moments(magmoms) 

62 self.gui.draw()