Coverage for /builds/kinetik161/ase/ase/gui/graphs.py: 43.86%

1import pickle 

2import sys 

3 

4import numpy as np 

5 

6import ase.gui.ui as ui 

7from ase.gui.i18n import _ 

8 

9graph_help_text = _("""\ 

10Symbols: 

11<c>e</c>: total energy 

12<c>epot</c>: potential energy 

13<c>ekin</c>: kinetic energy 

14<c>fmax</c>: maximum force 

15<c>fave</c>: average force 

16<c>R[n,0-2]</c>: position of atom number <c>n</c> 

17<c>d(n₁,n₂)</c>: distance between two atoms \ 

18<c>n₁</c> and <c>n₂</c> 

19<c>i</c>: current image number 

20<c>E[i]</c>: energy of image number <c>i</c> 

21<c>F[n,0-2]</c>: force on atom number <c>n</c> 

22<c>V[n,0-2]</c>: velocity of atom number <c>n</c> 

23<c>M[n]</c>: magnetic moment of atom number <c>n</c> 

24<c>A[0-2,0-2]</c>: unit-cell basis vectors 

25<c>s</c>: path length 

26<c>a(n1,n2,n3)</c>: angle between atoms <c>n₁</c>, \ 

27<c>n₂</c> and <c>n₃</c>, centered on <c>n₂</c> 

28<c>dih(n1,n2,n3,n4)</c>: dihedral angle between <c>n₁</c>, \ 

29<c>n₂</c>, <c>n₃</c> and <c>n₄</c> 

30<c>T</c>: temperature (K)\ 

31""") 

32 

33 

34class Graphs: 

35 def __init__(self, gui): 

36 win = ui.Window('Graphs', wmtype='utility') 

37 self.expr = ui.Entry('', 50, self.plot) 

38 win.add([self.expr, ui.helpbutton(graph_help_text)]) 

39 

40 win.add([ui.Button(_('Plot'), self.plot, 'xy'), 

41 ' x, y1, y2, ...'], 'w') 

42 win.add([ui.Button(_('Plot'), self.plot, 'y'), 

43 ' y1, y2, ...'], 'w') 

44 win.add([ui.Button(_('Save'), self.save)], 'w') 

45 

46 self.gui = gui 

47 

48 def plot(self, type=None, expr=None, ignore_if_nan=False): 

49 if expr is None: 

50 expr = self.expr.value 

51 else: 

52 self.expr.value = expr 

53 

54 try: 

55 data = self.gui.images.graph(expr) 

56 except Exception as ex: 

57 ui.error(ex) 

58 return 

59 

60 if ignore_if_nan and len(data) == 2 and np.isnan(data[1]).all(): 

61 return 

62 pickledata = (data, self.gui.frame, expr, type) 

63 self.gui.pipe('graph', pickledata) 

64 

65 def save(self): 

66 dialog = ui.SaveFileDialog(self.gui.window.win, 

67 _('Save data to file ... ')) 

68 # fix tkinter not automatically setting dialog type 

69 # remove from Python3.8+ 

70 # see https://github.com/python/cpython/pull/25187 

71 # and https://bugs.python.org/issue43655 

72 # and https://github.com/python/cpython/pull/25592 

73 ui.set_windowtype(dialog.top, 'dialog') 

74 filename = dialog.go() 

75 if filename: 

76 expr = self.expr.value 

77 data = self.gui.images.graph(expr) 

78 np.savetxt(filename, data.T, header=expr) 

79 

80 

81def make_plot(data, i, expr, type, show=True): 

82 import matplotlib.pyplot as plt 

83 basesize = 4 

84 plt.figure(figsize=(basesize * 2.5**0.5, basesize)) 

85 m = len(data) 

86 

87 if type is None: 

88 if m == 1: 

89 type = 'y' 

90 else: 

91 type = 'xy' 

92 

93 if type == 'y': 

94 for j in range(m): 

95 plt.plot(data[j]) 

96 plt.plot([i], [data[j, i]], 'o') 

97 else: 

98 for j in range(1, m): 

99 plt.plot(data[0], data[j]) 

100 plt.plot([data[0, i]], [data[j, i]], 'o') 

101 plt.title(expr) 

102 if show: 

103 plt.show() 

104 

105 

106if __name__ == '__main__': 

107 make_plot(*pickle.load(sys.stdin.buffer))