Coverage for /usr/local/lib/python3.11/dist-packages/pyrocko/squirrel/operators/base.py: 80%

1# http://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

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

5 

6import logging 

7import re 

8 

9from ..model import QuantityType, match_codes, CodesNSLCE 

10from .. import error 

11 

12from pyrocko.guts import Object, String, Duration, Float, clone, List 

13 

14guts_prefix = 'squirrel.ops' 

15 

16logger = logging.getLogger('psq.ops') 

17 

18 

19def odiff(a, b): 

20 ia = ib = 0 

21 only_a = [] 

22 only_b = [] 

23 while ia < len(a) or ib < len(b): 

24 # TODO remove when finished with implementation 

25 if ia > 0: 

26 assert a[ia] > a[ia-1] 

27 if ib > 0: 

28 assert b[ib] > b[ib-1] 

29 

30 if ib == len(b) or (ia < len(a) and a[ia] < b[ib]): 

31 only_a.append(a[ia]) 

32 ia += 1 

33 elif ia == len(a) or (ib < len(b) and a[ia] > b[ib]): 

34 only_b.append(b[ib]) 

35 ib += 1 

36 elif a[ia] == b[ib]: 

37 ia += 1 

38 ib += 1 

39 

40 return only_a, only_b 

41 

42 

43def _cglob_translate(creg): 

44 dd = [] 

45 for c in creg: 

46 if c == '*': 

47 d = r'[^.]*' 

48 elif c == '?': 

49 d = r'[^.]' 

50 elif c == '.': 

51 d = r'\.' 

52 else: 

53 d = c 

54 

55 dd.append(d) 

56 reg = ''.join(dd) 

57 return reg 

58 

59 

60class Filtering(Object): 

61 ''' 

62 Base class for :py:class:`pyrocko.squirrel.model.Nut` filters. 

63 ''' 

64 

65 def filter(self, it): 

66 return list(it) 

67 

68 

69class RegexFiltering(Object): 

70 ''' 

71 Filter by regex. 

72 ''' 

73 pattern = String.T(default=r'(.*)') 

74 

75 def __init__(self, **kwargs): 

76 Filtering.__init__(self, **kwargs) 

77 self._compiled_pattern = re.compile(self.pattern) 

78 

79 def filter(self, it): 

80 return [ 

81 x for x in it if self._compiled_pattern.fullmatch(x)] 

82 

83 

84class CodesPatternFiltering(Object): 

85 ''' 

86 Filter by codes pattern. 

87 ''' 

88 codes = List.T(CodesNSLCE.T(), optional=True) 

89 

90 def filter(self, it): 

91 if self.codes is None: 

92 return list(it) 

93 else: 

94 return [ 

95 x for x in it 

96 if any(match_codes(sc, x) for sc in self.codes)] 

97 

98 

99class Grouping(Object): 

100 ''' 

101 Base class for :py:class:`pyrocko.squirrel.model.Nut` grouping mechanisms. 

102 ''' 

103 

104 def key(self, codes): 

105 return codes 

106 

107 

108class RegexGrouping(Grouping): 

109 ''' 

110 Group by regex pattern. 

111 ''' 

112 pattern = String.T(default=r'(.*)') 

113 

114 def __init__(self, **kwargs): 

115 Grouping.__init__(self, **kwargs) 

116 self._compiled_pattern = re.compile(self.pattern) 

117 

118 def key(self, codes): 

119 return self._compiled_pattern.fullmatch(codes.safe_str).groups() 

120 

121 

122class NetworkGrouping(RegexGrouping): 

123 ''' 

124 Group by *network* code. 

125 ''' 

126 pattern = String.T(default=_cglob_translate('(*).*.*.*.*')) 

127 

128 

129class StationGrouping(RegexGrouping): 

130 ''' 

131 Group by *network.station* codes. 

132 ''' 

133 pattern = String.T(default=_cglob_translate('(*.*).*.*.*')) 

134 

135 

136class LocationGrouping(RegexGrouping): 

137 ''' 

138 Group by *network.station.location* codes. 

139 ''' 

140 pattern = String.T(default=_cglob_translate('(*.*.*).*.*')) 

141 

142 

143class ChannelGrouping(RegexGrouping): 

144 ''' 

145 Group by *network.station.location.channel* codes. 

146 

147 This effectively groups all processings of a channel, which may differ in 

148 the *extra* codes attribute. 

149 ''' 

150 pattern = String.T(default=_cglob_translate('(*.*.*.*).*')) 

151 

152 

153class SensorGrouping(RegexGrouping): 

154 ''' 

155 Group by *network.station.location.sensor* and *extra* codes. 

156 

157 For *sensor* all but the last character of the channel code (indicating the 

158 component) are used. This effectively groups all components of a sensor, 

159 or processings of a sensor. 

160 ''' 

161 pattern = String.T(default=_cglob_translate('(*.*.*.*)?(.*)')) 

162 

163 

164class Translation(Object): 

165 ''' 

166 Base class for :py:class:`pyrocko.squirrel.model.Nut` translators. 

167 ''' 

168 

169 def translate(self, codes): 

170 return codes 

171 

172 

173class AddSuffixTranslation(Translation): 

174 ''' 

175 Add a suffix to :py:attr:`~pyrocko.squirrel.model.CodesNSLCEBase.extra`. 

176 ''' 

177 suffix = String.T(default='') 

178 

179 def translate(self, codes): 

180 return codes.replace(extra=codes.extra + self.suffix) 

181 

182 

183class RegexTranslation(AddSuffixTranslation): 

184 ''' 

185 Translate :py:class:`pyrocko.squirrel.model.Codes` using a regular 

186 expression. 

187 ''' 

188 pattern = String.T(default=r'(.*)') 

189 replacement = String.T(default=r'\1') 

190 

191 def __init__(self, **kwargs): 

192 AddSuffixTranslation.__init__(self, **kwargs) 

193 self._compiled_pattern = re.compile(self.pattern) 

194 

195 def translate(self, codes): 

196 return codes.__class__( 

197 self._compiled_pattern.sub(self.replacement, codes.safe_str)) 

198 

199 

200class ReplaceComponentTranslation(RegexTranslation): 

201 ''' 

202 Translate :py:class:`pyrocko.squirrel.model.Codes` by replacing a 

203 component. 

204 ''' 

205 pattern = String.T(default=_cglob_translate('(*.*.*.*)?(.*)')) 

206 replacement = String.T(default=r'\1{component}\2') 

207 

208 

209def deregister(registry, group): 

210 for codes in group[2]: 

211 del registry[codes] 

212 

213 

214def register(registry, operator, group): 

215 for codes in group[2]: 

216 if codes in registry: 

217 logger.warning( 

218 'duplicate operator output codes: %s' % codes) 

219 registry[codes] = (operator, group) 

220 

221 

222class Operator(Object): 

223 

224 filtering = Filtering.T(default=Filtering.D()) 

225 grouping = Grouping.T(default=Grouping.D()) 

226 translation = Translation.T(default=Translation.D()) 

227 

228 def __init__(self, **kwargs): 

229 Object.__init__(self, **kwargs) 

230 self._output_names_cache = {} 

231 self._groups = {} 

232 self._available = [] 

233 

234 @property 

235 def name(self): 

236 return self.__class__.__name__ 

237 

238 def describe(self): 

239 return self.name 

240 

241 def iter_mappings(self): 

242 for k, group in self._groups.items(): 

243 yield (group[1], group[2]) 

244 

245 def iter_in_codes(self): 

246 for k, group in self._groups.items(): 

247 yield group[1] 

248 

249 def iter_out_codes(self): 

250 for k, group in self._groups.items(): 

251 yield group[2] 

252 

253 def update_mappings(self, available, registry=None): 

254 available = list(available) 

255 removed, added = odiff(self._available, available) 

256 

257 filt = self.filtering.filter 

258 gkey = self.grouping.key 

259 groups = self._groups 

260 

261 need_update = set() 

262 

263 for codes in filt(removed): 

264 k = gkey(codes) 

265 groups[k][0].remove(codes) 

266 need_update.add(k) 

267 

268 for codes in filt(added): 

269 k = gkey(codes) 

270 if k not in groups: 

271 groups[k] = [set(), None, ()] 

272 

273 groups[k][0].add(codes) 

274 need_update.add(k) 

275 

276 for k in need_update: 

277 group = groups[k] 

278 if registry is not None: 

279 deregister(registry, group) 

280 

281 group[1] = tuple(sorted(group[0])) 

282 if not group[1]: 

283 del groups[k] 

284 else: 

285 group[2] = self._out_codes(group[1]) 

286 if registry is not None: 

287 register(registry, self, group) 

288 

289 self._available = available 

290 

291 def _out_codes(self, in_codes): 

292 return [self.translation.translate(codes) for codes in in_codes] 

293 

294 def get_channels(self, squirrel, group, *args, **kwargs): 

295 _, in_codes, out_codes = group 

296 assert len(in_codes) == 1 and len(out_codes) == 1 

297 channels = squirrel.get_channels(codes=in_codes[0], **kwargs) 

298 channels_out = [] 

299 for channel in channels: 

300 channel_out = clone(channel) 

301 channel_out.codes = out_codes[0] 

302 channels_out.append(channel_out) 

303 

304 return channels_out 

305 

306 def get_waveforms(self, squirrel, group, **kwargs): 

307 _, in_codes, out_codes = group 

308 assert len(in_codes) == 1 and len(out_codes) == 1 

309 

310 trs = squirrel.get_waveforms(codes=in_codes[0], **kwargs) 

311 for tr in trs: 

312 tr.set_codes(*out_codes[0]) 

313 

314 return trs 

315 

316 # def update_waveforms(self, squirrel, tmin, tmax, codes): 

317 # if codes is None: 

318 # for _, in_codes, out_codes in self._groups.values(): 

319 # for codes in 

320 

321 

322class Parameters(Object): 

323 pass 

324 

325 

326class RestitutionParameters(Parameters): 

327 frequency_min = Float.T() 

328 frequency_max = Float.T() 

329 frequency_taper_factor = Float.T(default=1.5) 

330 time_taper_factor = Float.T(default=2.0) 

331 

332 

333class Restitution(Operator): 

334 translation = AddSuffixTranslation(suffix='R{quantity}') 

335 quantity = QuantityType.T(default='displacement') 

336 

337 @property 

338 def name(self): 

339 return 'Restitution(%s)' % self.quantity[0] 

340 

341 def _out_codes(self, group): 

342 return [ 

343 codes.__class__(self.translation.translate(codes).safe_str.format( 

344 quantity=self.quantity[0])) 

345 for codes in group] 

346 

347 def get_tpad(self, params): 

348 return params.time_taper_factor / params.frequency_min 

349 

350 def get_waveforms( 

351 self, squirrel, codes, params, tmin, tmax, **kwargs): 

352 

353 self_, (_, in_codes, out_codes) = squirrel.get_operator_group(codes) 

354 assert self is self_ 

355 assert len(in_codes) == 1 and len(out_codes) == 1 

356 

357 tpad = self.get_tpad(params) 

358 

359 tmin_raw = tmin - tpad 

360 tmax_raw = tmax + tpad 

361 

362 trs = squirrel.get_waveforms( 

363 codes=in_codes[0], tmin=tmin_raw, tmax=tmax_raw, **kwargs) 

364 

365 try: 

366 resp = squirrel.get_response( 

367 codes=in_codes[0], 

368 tmin=tmin_raw, 

369 tmax=tmax_raw).get_effective(self.quantity) 

370 

371 except error.NotAvailable: 

372 return [] 

373 

374 freqlimits = ( 

375 params.frequency_min / params.frequency_taper_factor, 

376 params.frequency_min, 

377 params.frequency_max, 

378 params.frequency_max * params.frequency_taper_factor) 

379 

380 trs_rest = [] 

381 for tr in trs: 

382 tr_rest = tr.transfer( 

383 tfade=tpad, 

384 freqlimits=freqlimits, 

385 transfer_function=resp, 

386 invert=True) 

387 

388 tr_rest.set_codes(*out_codes[0]) 

389 

390 trs_rest.append(tr_rest) 

391 

392 return trs_rest 

393 

394 

395class Shift(Operator): 

396 translation = AddSuffixTranslation(suffix='S') 

397 delay = Duration.T() 

398 

399 

400class Transform(Operator): 

401 grouping = Grouping.T(default=SensorGrouping.D()) 

402 translation = ReplaceComponentTranslation(suffix='T{system}') 

403 

404 def _out_codes(self, group): 

405 return [ 

406 self.translation.translate(group[0]).format( 

407 component=c, system=self.components.lower()) 

408 for c in self.components] 

409 

410 

411class ToENZ(Transform): 

412 components = 'ENZ' 

413 

414 def get_waveforms( 

415 self, squirrel, in_codes, out_codes, params, tmin, tmax, **kwargs): 

416 

417 trs = squirrel.get_waveforms( 

418 codes=in_codes, tmin=tmin, tmax=tmax, **kwargs) 

419 

420 for tr in trs: 

421 print(tr) 

422 

423 

424class ToRTZ(Transform): 

425 components = 'RTZ' 

426 backazimuth = Float.T() 

427 

428 

429class ToLTQ(Transform): 

430 components = 'LTQ' 

431 

432 

433class Composition(Operator): 

434 g = Operator.T() 

435 f = Operator.T() 

436 

437 def __init__(self, g, f, **kwargs): 

438 Operator.__init__(self, g=g, f=f, **kwargs) 

439 

440 @property 

441 def name(self): 

442 return '(%s ○ %s)' % (self.g.name, self.f.name) 

443 

444 

445__all__ = [ 

446 'Grouping', 

447 'RegexGrouping', 

448 'NetworkGrouping', 

449 'StationGrouping', 

450 'LocationGrouping', 

451 'SensorGrouping', 

452 'ChannelGrouping', 

453 'Operator', 

454 'RestitutionParameters', 

455 'Restitution', 

456 'Shift', 

457 'ToENZ', 

458 'ToRTZ', 

459 'ToLTQ', 

460 'Composition']