Coverage for ase / gui / graphs.py: 44.44%

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

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

71 if filename: 

72 expr = self.expr.value 

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

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

75 

76 

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

78 import matplotlib.pyplot as plt 

79 basesize = 4 

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

81 m = len(data) 

82 

83 if type is None: 

84 if m == 1: 

85 type = 'y' 

86 else: 

87 type = 'xy' 

88 

89 if type == 'y': 

90 for j in range(m): 

91 plt.plot(data[j]) 

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

93 else: 

94 for j in range(1, m): 

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

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

97 plt.title(expr) 

98 if show: 

99 plt.show() 

100 

101 

102if __name__ == '__main__': 

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