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, CodesNSLCE, CodesMatcher 

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(Filtering): 

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(Filtering): 

85 ''' 

86 Filter by codes pattern. 

87 ''' 

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

89 

90 def __init__(self, **kwargs): 

91 Filtering.__init__(self, **kwargs) 

92 if self.codes is not None: 

93 self._matcher = CodesMatcher(self.codes) 

94 else: 

95 self._matcher = None 

96 

97 def filter(self, it): 

98 if self._matcher is None: 

99 return list(it) 

100 else: 

101 return [ 

102 x for x in it if self._matcher.match(x)] 

103 

104 

105class Grouping(Object): 

106 ''' 

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

108 ''' 

109 

110 def key(self, codes): 

111 return codes 

112 

113 

114class RegexGrouping(Grouping): 

115 ''' 

116 Group by regex pattern. 

117 ''' 

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

119 

120 def __init__(self, **kwargs): 

121 Grouping.__init__(self, **kwargs) 

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

123 

124 def key(self, codes): 

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

126 

127 

128class NetworkGrouping(RegexGrouping): 

129 ''' 

130 Group by *network* code. 

131 ''' 

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

133 

134 

135class StationGrouping(RegexGrouping): 

136 ''' 

137 Group by *network.station* codes. 

138 ''' 

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

140 

141 

142class LocationGrouping(RegexGrouping): 

143 ''' 

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

145 ''' 

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

147 

148 

149class ChannelGrouping(RegexGrouping): 

150 ''' 

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

152 

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

154 the *extra* codes attribute. 

155 ''' 

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

157 

158 

159class SensorGrouping(RegexGrouping): 

160 ''' 

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

162 

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

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

165 or processings of a sensor. 

166 ''' 

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

168 

169 

170class Translation(Object): 

171 ''' 

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

173 ''' 

174 

175 def translate(self, codes): 

176 return codes 

177 

178 

179class AddSuffixTranslation(Translation): 

180 ''' 

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

182 ''' 

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

184 

185 def translate(self, codes): 

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

187 

188 

189class RegexTranslation(AddSuffixTranslation): 

190 ''' 

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

192 expression. 

193 ''' 

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

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

196 

197 def __init__(self, **kwargs): 

198 AddSuffixTranslation.__init__(self, **kwargs) 

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

200 

201 def translate(self, codes): 

202 return codes.__class__( 

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

204 

205 

206class ReplaceComponentTranslation(RegexTranslation): 

207 ''' 

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

209 component. 

210 ''' 

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

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

213 

214 

215def deregister(registry, group): 

216 for codes in group[2]: 

217 del registry[codes] 

218 

219 

220def register(registry, operator, group): 

221 for codes in group[2]: 

222 if codes in registry: 

223 logger.warning( 

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

225 registry[codes] = (operator, group) 

226 

227 

228class Operator(Object): 

229 

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

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

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

233 

234 def __init__(self, **kwargs): 

235 Object.__init__(self, **kwargs) 

236 self._output_names_cache = {} 

237 self._groups = {} 

238 self._available = [] 

239 

240 @property 

241 def name(self): 

242 return self.__class__.__name__ 

243 

244 def describe(self): 

245 return self.name 

246 

247 def iter_mappings(self): 

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

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

250 

251 def iter_in_codes(self): 

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

253 yield group[1] 

254 

255 def iter_out_codes(self): 

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

257 yield group[2] 

258 

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

260 available = list(available) 

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

262 

263 filt = self.filtering.filter 

264 gkey = self.grouping.key 

265 groups = self._groups 

266 

267 need_update = set() 

268 

269 for codes in filt(removed): 

270 k = gkey(codes) 

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

272 need_update.add(k) 

273 

274 for codes in filt(added): 

275 k = gkey(codes) 

276 if k not in groups: 

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

278 

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

280 need_update.add(k) 

281 

282 for k in need_update: 

283 group = groups[k] 

284 if registry is not None: 

285 deregister(registry, group) 

286 

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

288 if not group[1]: 

289 del groups[k] 

290 else: 

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

292 if registry is not None: 

293 register(registry, self, group) 

294 

295 self._available = available 

296 

297 def _out_codes(self, in_codes): 

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

299 

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

301 _, in_codes, out_codes = group 

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

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

304 channels_out = [] 

305 for channel in channels: 

306 channel_out = clone(channel) 

307 channel_out.codes = out_codes[0] 

308 channels_out.append(channel_out) 

309 

310 return channels_out 

311 

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

313 _, in_codes, out_codes = group 

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

315 

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

317 for tr in trs: 

318 tr.set_codes(*out_codes[0]) 

319 

320 return trs 

321 

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

323 # if codes is None: 

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

325 # for codes in 

326 

327 

328class Parameters(Object): 

329 pass 

330 

331 

332class RestitutionParameters(Parameters): 

333 frequency_min = Float.T() 

334 frequency_max = Float.T() 

335 frequency_taper_factor = Float.T(default=1.5) 

336 time_taper_factor = Float.T(default=2.0) 

337 

338 

339class Restitution(Operator): 

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

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

342 

343 @property 

344 def name(self): 

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

346 

347 def _out_codes(self, group): 

348 return [ 

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

350 quantity=self.quantity[0])) 

351 for codes in group] 

352 

353 def get_tpad(self, params): 

354 return params.time_taper_factor / params.frequency_min 

355 

356 def get_waveforms( 

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

358 

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

360 assert self is self_ 

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

362 

363 tpad = self.get_tpad(params) 

364 

365 tmin_raw = tmin - tpad 

366 tmax_raw = tmax + tpad 

367 

368 trs = squirrel.get_waveforms( 

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

370 

371 try: 

372 resp = squirrel.get_response( 

373 codes=in_codes[0], 

374 tmin=tmin_raw, 

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

376 

377 except error.NotAvailable: 

378 return [] 

379 

380 freqlimits = ( 

381 params.frequency_min / params.frequency_taper_factor, 

382 params.frequency_min, 

383 params.frequency_max, 

384 params.frequency_max * params.frequency_taper_factor) 

385 

386 trs_rest = [] 

387 for tr in trs: 

388 tr_rest = tr.transfer( 

389 tfade=tpad, 

390 freqlimits=freqlimits, 

391 transfer_function=resp, 

392 invert=True) 

393 

394 tr_rest.set_codes(*out_codes[0]) 

395 

396 trs_rest.append(tr_rest) 

397 

398 return trs_rest 

399 

400 

401class Shift(Operator): 

402 translation = AddSuffixTranslation(suffix='S') 

403 delay = Duration.T() 

404 

405 

406class Transform(Operator): 

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

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

409 

410 def _out_codes(self, group): 

411 return [ 

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

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

414 for c in self.components] 

415 

416 

417class ToENZ(Transform): 

418 components = 'ENZ' 

419 

420 def get_waveforms( 

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

422 

423 trs = squirrel.get_waveforms( 

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

425 

426 for tr in trs: 

427 print(tr) 

428 

429 

430class ToRTZ(Transform): 

431 components = 'RTZ' 

432 backazimuth = Float.T() 

433 

434 

435class ToLTQ(Transform): 

436 components = 'LTQ' 

437 

438 

439class Composition(Operator): 

440 g = Operator.T() 

441 f = Operator.T() 

442 

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

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

445 

446 @property 

447 def name(self): 

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

449 

450 

451__all__ = [ 

452 'Grouping', 

453 'RegexGrouping', 

454 'NetworkGrouping', 

455 'StationGrouping', 

456 'LocationGrouping', 

457 'SensorGrouping', 

458 'ChannelGrouping', 

459 'Operator', 

460 'RestitutionParameters', 

461 'Restitution', 

462 'Shift', 

463 'ToENZ', 

464 'ToRTZ', 

465 'ToLTQ', 

466 'Composition']