1# http://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

4# ---|P------/S----------~Lg---------- 

5from __future__ import absolute_import 

6from ..snuffling import Snuffling, Param, Switch 

7import numpy as num 

8from matplotlib import cm 

9 

10 

11def window(freqs, fc, b): 

12 if fc == 0.: 

13 w = num.zeros(len(freqs)) 

14 w[freqs == 0] = 1. 

15 return w 

16 T = num.log10(freqs/fc)*b 

17 w = (num.sin(T)/T)**4 

18 w[freqs == fc] = 1. 

19 w[freqs == 0.] = 0. 

20 w /= num.sum(w) 

21 return w 

22 

23 

24class AmpSpec(Snuffling): 

25 ''' 

26 <html> 

27 <head> 

28 <style type="text/css"> 

29 body { margin-left:10px }; 

30 </style> 

31 <body> 

32 <h1 align='center'>Plot Amplitude Spectrum</h1> 

33 <p> 

34 When smoothing is activated, a smoothing algorithm is applied as proposed 

35 by Konno and Ohmachi, (1998). </p> 

36 <p style='font-family:courier'> 

37 Konno, K. and Omachi, T. (1998). Ground-motion characteristics 

38 estimated from spectral ratio between horizontal and vertical 

39 components of microtremor, Bull. Seism. Soc. Am., 88, 228-241 

40 </p> 

41 </body> 

42 </html> 

43 ''' 

44 def setup(self): 

45 '''Customization of the snuffling.''' 

46 

47 self.set_name('Plot Amplitude Spectrum') 

48 self.add_parameter(Switch('Smoothing', 'want_smoothing', False)) 

49 self.add_parameter(Param('Smoothing band width', 'b', 40., 1., 100.)) 

50 self.set_live_update(False) 

51 self._wins = {} 

52 

53 def call(self): 

54 '''Main work routine of the snuffling.''' 

55 

56 all = [] 

57 for traces in self.chopper_selected_traces(fallback=True): 

58 for trace in traces: 

59 all.append(trace) 

60 

61 extrema = [] 

62 colors = iter(cm.Accent(num.linspace(0., 1., len(all)))) 

63 if self.want_smoothing: 

64 alpha = 0.2 

65 additional = 'and smoothing' 

66 else: 

67 alpha = 1. 

68 additional = '' 

69 

70 minf = 0. 

71 maxf = 0. 

72 pblabel = 'Calculating amplitude spectra %s' % additional 

73 pb = self.get_viewer().parent().get_progressbars() 

74 pb.set_status(pblabel, 0) 

75 num_traces = len(all) 

76 maxval = float(num_traces) 

77 plot_data = [] 

78 plot_data_supplement = [] 

79 for i_tr, tr in enumerate(all): 

80 val = i_tr/maxval*100. 

81 pb.set_status(pblabel, val) 

82 

83 tr.ydata = tr.ydata.astype(float) 

84 tr.ydata -= tr.ydata.mean() 

85 f, a = tr.spectrum() 

86 minf = min([f.min(), minf]) 

87 maxf = max([f.max(), maxf]) 

88 absa = num.abs(a) 

89 labsa = num.log(absa) 

90 stdabsa = num.std(labsa) 

91 meanabsa = num.mean(labsa) 

92 mi, ma = meanabsa - 3*stdabsa, meanabsa + 3*stdabsa 

93 extrema.append(mi) 

94 extrema.append(ma) 

95 c = next(colors) 

96 plot_data.append((f, num.abs(a))) 

97 plot_data_supplement.append((c, alpha, '.'.join(tr.nslc_id))) 

98 if self.want_smoothing: 

99 smoothed = self.konnoohmachi(num.abs(a), f, self.b) 

100 plot_data.append((f, num.abs(smoothed))) 

101 plot_data_supplement.append((c, 1., '.'.join(tr.nslc_id))) 

102 self.get_viewer().update() 

103 

104 pb.set_status(pblabel, 100.) 

105 

106 fig = self.figure(name='Amplitude Spectra') 

107 p = fig.add_subplot(111) 

108 args = ('c', 'alpha', 'label') 

109 for d, s in zip(plot_data, plot_data_supplement): 

110 p.plot(*d, **dict(zip(args, s))) 

111 mi, ma = min(extrema), max(extrema) 

112 p.set_xscale('log') 

113 p.set_yscale('log') 

114 p.set_ylim(num.exp(mi), num.exp(ma)) 

115 p.set_xlim(minf, maxf) 

116 p.set_xlabel('Frequency [Hz]') 

117 p.set_ylabel('Counts') 

118 handles, labels = p.get_legend_handles_labels() 

119 leg_dict = dict(zip(labels, handles)) 

120 if num_traces > 1: 

121 p.legend(list(leg_dict.values()), list(leg_dict.keys()), 

122 loc=2, 

123 borderaxespad=0., 

124 bbox_to_anchor=((1.05, 1.))) 

125 fig.subplots_adjust(right=0.8, 

126 left=0.1) 

127 else: 

128 p.set_title(list(leg_dict.keys())[0], fontsize=16) 

129 fig.subplots_adjust(right=0.9, 

130 left=0.1) 

131 fig.canvas.draw() 

132 

133 def konnoohmachi(self, amps, freqs, b=20): 

134 smooth = num.zeros(len(freqs), dtype=freqs.dtype) 

135 amps = num.array(amps) 

136 for i, fc in enumerate(freqs): 

137 fkey = tuple((self.b, fc, freqs[0], freqs[1], freqs[-1])) 

138 if fkey in self._wins.keys(): 

139 win = self._wins[fkey] 

140 else: 

141 win = window(freqs, fc, b) 

142 self._wins[fkey] = win 

143 smooth[i] = num.sum(win*amps) 

144 

145 return smooth 

146 

147 

148def __snufflings__(): 

149 ''' 

150 Returns a list of snufflings to be exported by this module. 

151 ''' 

152 

153 return [AmpSpec()]