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

1# fmt: off 

2 

3import pickle 

4import sys 

5 

6import numpy as np 

7 

8import ase.gui.ui as ui 

9from ase.gui.i18n import _ 

10 

11graph_help_text = _("""\ 

12Symbols: 

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

33""") 

34 

35 

36class Graphs: 

37 def __init__(self, gui): 

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

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

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

41 

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

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

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

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

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

47 

48 self.gui = gui 

49 

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

51 if expr is None: 

52 expr = self.expr.value 

53 else: 

54 self.expr.value = expr 

55 

56 try: 

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

58 except Exception as ex: 

59 ui.error(ex) 

60 return 

61 

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

63 return 

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

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

66 

67 def save(self): 

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

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

70 # fix tkinter not automatically setting dialog type 

71 # remove from Python3.8+ 

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

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

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

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

76 filename = dialog.go() 

77 if filename: 

78 expr = self.expr.value 

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

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

81 

82 

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

84 import matplotlib.pyplot as plt 

85 basesize = 4 

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

87 m = len(data) 

88 

89 if type is None: 

90 if m == 1: 

91 type = 'y' 

92 else: 

93 type = 'xy' 

94 

95 if type == 'y': 

96 for j in range(m): 

97 plt.plot(data[j]) 

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

99 else: 

100 for j in range(1, m): 

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

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

103 plt.title(expr) 

104 if show: 

105 plt.show() 

106 

107 

108if __name__ == '__main__': 

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