1import pickle 

2import sys 

3 

4import numpy as np 

5 

6from ase.gui.i18n import _ 

7import ase.gui.ui as ui 

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') 

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 filename = dialog.go() 

69 if filename: 

70 expr = self.expr.value 

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

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

73 

74 

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

76 import matplotlib.pyplot as plt 

77 basesize = 4 

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

79 m = len(data) 

80 

81 if type is None: 

82 if m == 1: 

83 type = 'y' 

84 else: 

85 type = 'xy' 

86 

87 if type == 'y': 

88 for j in range(m): 

89 plt.plot(data[j]) 

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

91 else: 

92 for j in range(1, m): 

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

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

95 plt.title(expr) 

96 if show: 

97 plt.show() 

98 

99 

100if __name__ == '__main__': 

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