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

1# http://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

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

5 

6''' 

7Reader for `SAC 

8<http://ds.iris.edu/ds/nodes/dmc/software/downloads/sac/>`_ pole-zero files. 

9 

10See also :py:mod:`~pyrocko.io.enhanced_sacpz`. 

11''' 

12 

13import math 

14import numpy as num 

15try: 

16 from StringIO import StringIO as BytesIO 

17except ImportError: 

18 from io import BytesIO 

19 

20from pyrocko import trace 

21 

22d2r = math.pi/180. 

23 

24 

25class SacPoleZeroError(Exception): 

26 pass 

27 

28 

29def read_sac_zpk(filename=None, file=None, string=None, get_comments=False): 

30 ''' 

31 Read SAC pole-zero file. 

32 

33 :returns: ``(zeros, poles, constant)`` or 

34 ``(zeros, poles, constant, comments)`` if ``get_comments`` is True. 

35 ''' 

36 

37 if filename is not None: 

38 f = open(filename, 'rb') 

39 

40 elif file is not None: 

41 f = file 

42 

43 elif string is not None: 

44 f = BytesIO(string) 

45 

46 sects = ('ZEROS', 'POLES', 'CONSTANT') 

47 sectdata = {'ZEROS': [], 'POLES': []} 

48 npoles = 0 

49 nzeros = 0 

50 constant = 1.0 

51 atsect = None 

52 comments = [] 

53 for iline, line in enumerate(f): 

54 line = str(line.decode('ascii')) 

55 toks = line.split() 

56 if len(toks) == 0: 

57 continue 

58 

59 if toks[0][0] in '*#': 

60 comments.append(line) 

61 continue 

62 

63 if len(toks) != 2: 

64 f.close() 

65 raise SacPoleZeroError( 

66 'Expected 2 tokens in line %i of file %s' 

67 % (iline+1, filename)) 

68 

69 if toks[0].startswith('*'): 

70 continue 

71 

72 lsect = toks[0].upper() 

73 if lsect in sects: 

74 atsect = lsect 

75 sectdata[atsect] = [] 

76 if lsect.upper() == 'ZEROS': 

77 nzeros = int(toks[1]) 

78 elif toks[0].upper() == 'POLES': 

79 npoles = int(toks[1]) 

80 elif toks[0].upper() == 'CONSTANT': 

81 constant = float(toks[1]) 

82 else: 

83 if atsect: 

84 sectdata[atsect].append( 

85 complex(float(toks[0]), float(toks[1]))) 

86 

87 if f != file: 

88 f.close() 

89 

90 poles = sectdata['POLES'] 

91 zeros = sectdata['ZEROS'] 

92 npoles_ = len(poles) 

93 nzeros_ = len(zeros) 

94 if npoles_ > npoles: 

95 raise SacPoleZeroError( 

96 'Expected %i poles but found %i in pole-zero file "%s"' 

97 % (npoles, npoles_, filename)) 

98 if nzeros_ > nzeros: 

99 raise SacPoleZeroError( 

100 'Expected %i zeros but found %i in pole-zero file "%s"' 

101 % (nzeros, nzeros_, filename)) 

102 

103 if npoles_ < npoles: 

104 poles.extend([complex(0.)]*(npoles-npoles_)) 

105 

106 if nzeros_ < npoles: 

107 zeros.extend([complex(0.)]*(nzeros-nzeros_)) 

108 

109 if len(poles) == 0 and len(zeros) == 0: 

110 raise SacPoleZeroError( 

111 'No poles and zeros found in file "%s"' % (filename)) 

112 

113 if not num.all(num.isfinite(poles)): 

114 raise SacPoleZeroError( 

115 'Not finite pole(s) found in pole-zero file "%s"' 

116 % filename) 

117 

118 if not num.all(num.isfinite(zeros)): 

119 raise SacPoleZeroError( 

120 'Not finite zero(s) found in pole-zero file "%s"' 

121 % filename) 

122 

123 if not num.isfinite(constant): 

124 raise SacPoleZeroError( 

125 'Ivalid constant (%g) found in pole-zero file "%s"' 

126 % (constant, filename)) 

127 

128 if get_comments: 

129 return zeros, poles, constant, comments 

130 else: 

131 return zeros, poles, constant 

132 

133 

134def write_sac_zpk(zeros, poles, constant, filename): 

135 if hasattr(filename, 'write'): 

136 f = filename 

137 else: 

138 f = open('w', filename) 

139 

140 def write_complex(x): 

141 f.write('%12.8g %12.8g\n' % (complex(x).real, complex(x).imag)) 

142 

143 f.write('POLES %i\n' % len(poles)) 

144 for p in poles: 

145 if p != 0.0: 

146 write_complex(p) 

147 

148 f.write('ZEROS %i\n' % len(zeros)) 

149 for z in zeros: 

150 if z != 0.0: 

151 write_complex(z) 

152 

153 f.write('CONSTANT %12.8g\n' % constant) 

154 if not hasattr(filename, 'write'): 

155 f.close() 

156 

157 

158def evaluate(zeros, poles, constant, fmin=0.001, fmax=100., nf=100): 

159 

160 logfmin = math.log(fmin) 

161 logfmax = math.log(fmax) 

162 logf = num.linspace(logfmin, logfmax, nf) 

163 f = num.exp(logf) 

164 trans = trace.PoleZeroResponse(zeros, poles, constant) 

165 return f, trans.evaluate(f) 

166 

167 

168def evaluate_at(zeros, poles, constant, f): 

169 jomeg = 1.0j * 2. * math.pi * f 

170 

171 a = constant 

172 for z in zeros: 

173 a *= jomeg-z 

174 for p in poles: 

175 a /= jomeg-p 

176 

177 return a 

178 

179 

180def read_to_pyrocko_response(filename=None, file=None, string=None): 

181 ''' 

182 Read SAC pole-zero file into Pyrocko response object. 

183 

184 :returns: Response as a :py:class:~pyrocko.trace.PoleZeroResponse` object. 

185 ''' 

186 

187 from pyrocko import trace 

188 

189 zeros, poles, constant = read_sac_zpk( 

190 filename=filename, file=file, string=string) 

191 return trace.PoleZeroResponse(zeros, poles, constant) 

192 

193 

194def read_to_stationxml_response( 

195 input_unit, output_unit, normalization_frequency=1.0, 

196 filename=None, file=None, string=None): 

197 ''' 

198 Read SAC pole-zero file into StationXML response object. 

199 

200 :param input_unit: Input unit to be reported in the StationXML response. 

201 :type input_unit: str 

202 :param output_unit: Output unit to be reported in the StationXML response. 

203 :type output_unit: str 

204 :param normalization_frequency: Frequency where the normalization factor 

205 for the StationXML response should be computed. 

206 :type normalization_frequency: float 

207 

208 :returns: Response as a :py:class:~pyrocko.io.stationxml.Response` object 

209 with a single pole-zero response stage. 

210 ''' 

211 

212 from pyrocko.io import stationxml 

213 

214 presponse = read_to_pyrocko_response( 

215 filename=filename, file=file, string=string) 

216 

217 return stationxml.Response.from_pyrocko_pz_response( 

218 presponse, input_unit, output_unit, normalization_frequency) 

219 

220 

221def plot_amplitudes_zpk( 

222 zpks, filename_pdf, 

223 fmin=0.001, 

224 fmax=100., 

225 nf=100, 

226 fnorm=None): 

227 

228 from pyrocko.plot import gmtpy 

229 

230 p = gmtpy.LogLogPlot(width=30*gmtpy.cm, yexp=0) 

231 for i, (zeros, poles, constant) in enumerate(zpks): 

232 f, h = evaluate(zeros, poles, constant, fmin, fmax, nf) 

233 if fnorm is not None: 

234 h /= evaluate_at(zeros, poles, constant, fnorm) 

235 

236 amp = num.abs(h) 

237 p.plot((f, amp), '-W2p,%s' % gmtpy.color(i)) 

238 

239 p.save(filename_pdf) 

240 

241 

242def plot_phases_zpk(zpks, filename_pdf, fmin=0.001, fmax=100., nf=100): 

243 

244 from pyrocko.plot import gmtpy 

245 

246 p = gmtpy.LogLinPlot(width=30*gmtpy.cm) 

247 for i, (zeros, poles, constant) in enumerate(zpks): 

248 f, h = evaluate(zeros, poles, constant, fmin, fmax, nf) 

249 phase = num.unwrap(num.angle(h)) / d2r 

250 p.plot((f, phase), '-W1p,%s' % gmtpy.color(i)) 

251 

252 p.save(filename_pdf)