Coverage for /usr/local/lib/python3.11/dist-packages/pyrocko/multitrace.py: 52%

1# http://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

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

5 

6''' 

7Multi-component waveform data model. 

8''' 

9 

10 

11import numpy as num 

12import numpy.ma as ma 

13 

14from . import trace 

15from .trace import Trace 

16from .guts import Object, Float, Timestamp, List 

17from .guts_array import Array 

18from .squirrel.model import CodesNSLCE 

19 

20 

21class MultiTrace(Object): 

22 ''' 

23 Container for multi-component waveforms with common time span and sampling. 

24 

25 Instances of this class can be used to efficiently represent 

26 multi-component waveforms of a single sensor or of a sensor array. The data 

27 samples are stored in a single 2D array where the first index runs over 

28 components and the second index over time. Metadata contains sampling rate, 

29 start-time and :py:class:`~pyrocko.squirrel.model.CodesNSLCE` identifiers 

30 for the contained traces. 

31 

32 :param traces: 

33 If given, construct multi-trace from given single-component waveforms 

34 (see :py:func:`~pyrocko.trace.get_traces_data_as_array`) and ignore 

35 any other arguments. 

36 :type traces: 

37 :py:class:`list` of :py:class:`~pyrocko.trace.Trace` 

38 ''' 

39 

40 codes = List.T( 

41 CodesNSLCE.T(), 

42 help='List of codes identifying the components.') 

43 data = Array.T( 

44 shape=(None, None), 

45 help='Array containing the data samples indexed as ' 

46 '``(icomponent, isample)``.') 

47 tmin = Timestamp.T( 

48 default=Timestamp.D('1970-01-01 00:00:00'), 

49 help='Start time.') 

50 deltat = Float.T( 

51 default=1.0, 

52 help='Sampling interval [s]') 

53 

54 def __init__( 

55 self, 

56 traces=None, 

57 assemble='concatenate', 

58 data=None, 

59 codes=None, 

60 tmin=None, 

61 deltat=None): 

62 

63 if traces is not None: 

64 if len(traces) == 0: 

65 data = ma.zeros((0, 0)) 

66 else: 

67 if assemble == 'merge': 

68 data, codes, tmin, deltat \ 

69 = trace.merge_traces_data_as_array(traces) 

70 

71 elif assemble == 'concatenate': 

72 data = ma.array(trace.get_traces_data_as_array(traces)) 

73 codes = [tr.codes for tr in traces] 

74 tmin = traces[0].tmin 

75 deltat = traces[0].deltat 

76 

77 self.ntraces, self.nsamples = data.shape 

78 

79 if codes is None: 

80 codes = [CodesNSLCE()] * self.ntraces 

81 

82 if len(codes) != self.ntraces: 

83 raise ValueError( 

84 'MultiTrace construction: mismatch between number of traces ' 

85 'and number of codes given.') 

86 

87 if deltat is None: 

88 deltat = self.T.deltat.default() 

89 

90 if tmin is None: 

91 tmin = self.T.tmin.default() 

92 

93 Object.__init__(self, codes=codes, data=data, tmin=tmin, deltat=deltat) 

94 

95 def __len__(self): 

96 ''' 

97 Get number of components. 

98 ''' 

99 return self.ntraces 

100 

101 def __getitem__(self, i): 

102 ''' 

103 Get single component waveform (shared data). 

104 

105 :param i: 

106 Component index. 

107 :type i: 

108 int 

109 ''' 

110 return self.get_trace(i) 

111 

112 def copy(self, data='copy'): 

113 

114 if isinstance(data, str): 

115 assert data in ('copy', 'reference') 

116 data = self.data.copy() if data == 'copy' else self.data 

117 else: 

118 assert isinstance(data, ma.MaskedArray) 

119 

120 return MultiTrace( 

121 data=data, 

122 codes=list(self.codes), 

123 tmin=self.tmin, 

124 deltat=self.deltat) 

125 

126 @property 

127 def tmax(self): 

128 ''' 

129 End time (time of last sample, read-only). 

130 ''' 

131 return self.tmin + (self.nsamples - 1) * self.deltat 

132 

133 def get_trace(self, i, span=slice(None)): 

134 ''' 

135 Get single component waveform (shared data). 

136 

137 :param i: 

138 Component index. 

139 :type i: 

140 int 

141 ''' 

142 

143 network, station, location, channel, extra = self.codes[i] 

144 return Trace( 

145 network=network, 

146 station=station, 

147 location=location, 

148 channel=channel, 

149 extra=extra, 

150 tmin=self.tmin + (span.start or 0) * self.deltat, 

151 deltat=self.deltat, 

152 ydata=self.data.data[i, span]) 

153 

154 def iter_valid_traces(self): 

155 if self.data.mask is ma.nomask: 

156 return iter(self) 

157 

158 for irow, row in enumerate(ma.notmasked_contiguous(self.data, axis=1)): 

159 for slice in row: 

160 yield self.get_trace(irow, slice) 

161 

162 def get_traces(self): 

163 return list(self) 

164 

165 def get_valid_traces(self): 

166 return list(self.iter_valid_traces()) 

167 

168 def snuffle(self): 

169 ''' 

170 Show in Snuffler. 

171 ''' 

172 trace.snuffle(list(self)) 

173 

174 def snuffle_valid(self): 

175 trace.snuffle(self.get_valid_traces()) 

176 

177 def bleed_mask(self, t): 

178 if self.data.mask is ma.nomask: 

179 return 

180 

181 nt = int(num.round(abs(t)/self.deltat)) 

182 for irow, row in enumerate(ma.notmasked_contiguous(self.data, axis=1)): 

183 for span in row: 

184 self.data.mask[irow, span.start:span.start+nt] = True 

185 self.data.mask[irow, max(0, span.stop-nt):span.stop] = True 

186 

187 

188def correlate(a, b, mode='valid', normalization=None, use_fft=False): 

189 

190 if isinstance(a, Trace) and isinstance(b, Trace): 

191 return trace.correlate( 

192 a, b, mode=mode, normalization=normalization, use_fft=use_fft) 

193 

194 elif isinstance(a, Trace) and isinstance(b, MultiTrace): 

195 return MultiTrace([ 

196 trace.correlate( 

197 a, b_, 

198 mode=mode, normalization=normalization, use_fft=use_fft) 

199 for b_ in b]) 

200 

201 elif isinstance(a, MultiTrace) and isinstance(b, Trace): 

202 return MultiTrace([ 

203 trace.correlate( 

204 a_, b, 

205 mode=mode, normalization=normalization, use_fft=use_fft) 

206 for a_ in a]) 

207 

208 elif isinstance(a, MultiTrace) and isinstance(b, MultiTrace): 

209 return MultiTrace([ 

210 trace.correlate( 

211 a_, b_, 

212 mode=mode, normalization=normalization, use_fft=use_fft) 

213 

214 for a_ in a for b_ in b])