Coverage for /builds/debichem-team/python-ase/ase/gui/modify.py: 95.65%

1from functools import partial 

2 

3import numpy as np 

4 

5import ase.gui.ui as ui 

6from ase.gui.i18n import _ 

7from ase.gui.utils import get_magmoms 

8from ase.gui.widgets import Element 

9 

10 

11class ModifyAtoms: 

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

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

14 """ 

15 

16 def __init__(self, gui): 

17 self.gui = gui 

18 selected = self.selection() 

19 if not selected.any(): 

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

21 return 

22 

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

24 element = Element(callback=self.set_element) 

25 win.add(element) 

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

27 partial(self.set_element, element))) 

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

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

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

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

32 

33 atoms = self.gui.atoms 

34 sym = atoms.symbols[selected] 

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

36 element.symbol = sym[0] 

37 

38 tags = atoms.get_tags()[selected] 

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

40 self.tag.value = tags[0] 

41 

42 magmoms = get_magmoms(atoms)[selected] 

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

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

45 

46 def selection(self): 

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

48 

49 def set_element(self, element): 

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

51 self.gui.draw() 

52 

53 def set_tag(self): 

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

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

56 self.gui.atoms.set_tags(tags) 

57 self.gui.draw() 

58 

59 def set_magmom(self): 

60 magmoms = get_magmoms(self.gui.atoms) 

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

62 self.gui.atoms.set_initial_magnetic_moments(magmoms) 

63 self.gui.draw()