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 def __init__(self, gui): 

15 self.gui = gui 

16 selected = self.selection() 

17 if not selected.any(): 

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

19 return 

20 

21 win = ui.Window(_('Modify')) 

22 element = Element(callback=self.set_element) 

23 win.add(element) 

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

25 partial(self.set_element, element))) 

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

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

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

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

30 

31 atoms = self.gui.atoms 

32 sym = atoms.symbols[selected] 

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

34 element.symbol = sym[0] 

35 

36 tags = atoms.get_tags()[selected] 

37 if tags.ptp() == 0: 

38 self.tag.value = tags[0] 

39 

40 magmoms = get_magmoms(atoms)[selected] 

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

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

43 

44 def selection(self): 

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

46 

47 def set_element(self, element): 

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

49 self.gui.draw() 

50 

51 def set_tag(self): 

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

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

54 self.gui.atoms.set_tags(tags) 

55 self.gui.draw() 

56 

57 def set_magmom(self): 

58 magmoms = get_magmoms(self.gui.atoms) 

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

60 self.gui.atoms.set_initial_magnetic_moments(magmoms) 

61 self.gui.draw()