Coverage for /usr/local/lib/python3.11/dist-packages/pyrocko/gui/drum/view.py: 0%

1# https://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

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

5 

6from __future__ import absolute_import, print_function, division 

7 

8import math 

9import time 

10import numpy as num 

11 

12from pyrocko.gui.qt_compat import qc, qg, qw 

13 

14from pyrocko.gui.talkie import TalkieConnectionOwner 

15from pyrocko.gui.drum.state import State, TextStyle 

16from pyrocko.gui_util import make_QPolygonF, PhaseMarker 

17from pyrocko import trace, util, pile 

18 

19 

20def lim(a, x, b): 

21 return min(max(a, x), b) 

22 

23 

24class Label(object): 

25 def __init__( 

26 self, x, y, label_str, 

27 anchor='BL', 

28 style=None, 

29 keep_inside=None, 

30 head=None): 

31 

32 if style is None: 

33 style = TextStyle() 

34 

35 text = qg.QTextDocument() 

36 font = style.qt_font 

37 if font: 

38 text.setDefaultFont(font) 

39 

40 color = style.color.qt_color 

41 

42 text.setDefaultStyleSheet('span { color: %s; }' % color.name()) 

43 text.setHtml('<span>%s</span>' % label_str) 

44 

45 self.position = x, y 

46 self.anchor = anchor 

47 self.text = text 

48 self.style = style 

49 self.keep_inside = keep_inside 

50 if head: 

51 self.head = head 

52 else: 

53 self.head = (0., 0.) 

54 

55 def draw(self, p): 

56 s = self.text.size() 

57 rect = qc.QRectF(0., 0., s.width(), s.height()) 

58 

59 tx, ty = x, y = self.position 

60 anchor = self.anchor 

61 

62 pxs = self.text.defaultFont().pointSize() 

63 

64 oy = self.head[0] * pxs 

65 ox = self.head[1]/2. * pxs 

66 

67 if 'L' in anchor: 

68 tx -= min(ox*2, rect.width()/2.) 

69 

70 elif 'R' in anchor: 

71 tx -= rect.width() - min(ox*2., rect.width()/2.) 

72 

73 elif 'C' in anchor: 

74 tx -= rect.width()/2. 

75 

76 if 'B' in anchor: 

77 ty -= rect.height() + oy 

78 

79 elif 'T' in anchor: 

80 ty += oy 

81 

82 elif 'M' in anchor: 

83 ty -= rect.height()/2. 

84 

85 rect.translate(tx, ty) 

86 

87 if self.keep_inside: 

88 keep_inside = self.keep_inside 

89 if rect.top() < keep_inside.top(): 

90 rect.moveTop(keep_inside.top()) 

91 

92 if rect.bottom() > keep_inside.bottom(): 

93 rect.moveBottom(keep_inside.bottom()) 

94 

95 if rect.left() < keep_inside.left(): 

96 rect.moveLeft(keep_inside.left()) 

97 

98 if rect.right() > keep_inside.right(): 

99 rect.moveRight(keep_inside.right()) 

100 

101 poly = None 

102 if self.head[0] != 0.: 

103 l, r, t, b = rect.left(), rect.right(), rect.top(), rect.bottom() 

104 

105 if 'T' in anchor: 

106 a, b = t, b 

107 elif 'B' in anchor: 

108 a, b = b, t 

109 elif 'M' in anchor: 

110 assert False, 'label cannot have head with M alignment' 

111 

112 c1, c2 = lim(l, x-ox, r), lim(l, x+ox, r) 

113 

114 px = (l, c1, x, c2, r, r, l) 

115 py = (a, a, y, a, a, b, b) 

116 

117 poly = make_QPolygonF(px, py) 

118 

119 tx = rect.left() 

120 ty = rect.top() 

121 

122 if self.style.outline or self.style.background_color: 

123 oldpen = p.pen() 

124 oldbrush = p.brush() 

125 if not self.style.outline: 

126 p.setPen(qg.QPen(qc.Qt.NoPen)) 

127 

128 p.setBrush(self.style.background_color.qt_color) 

129 if poly: 

130 p.drawPolygon(poly) 

131 else: 

132 p.drawRect(rect) 

133 

134 if self.style.background_color: 

135 p.fillRect(rect, self.style.background_color.qt_color) 

136 

137 p.setPen(oldpen) 

138 p.setBrush(oldbrush) 

139 

140 p.translate(tx, ty) 

141 self.text.drawContents(p) 

142 p.translate(-tx, -ty) 

143 

144 

145class ProjectXU(object): 

146 def __init__(self, xr=(0., 1.), ur=(0., 1.)): 

147 (self.xmin, self.xmax) = xr 

148 (self.umin, self.umax) = ur 

149 self.cxu = (self.umax - self.umin) / (self.xmax - self.xmin) 

150 

151 def u(self, x): 

152 return self.umin + (x-self.xmin) * self.cxu 

153 

154 def x(self, u): 

155 return self.xmin + (u-self.umin) / self.cxu 

156 

157 

158class ProjectYV(object): 

159 def __init__(self, yr=(0., 1.), vr=(0., 1.)): 

160 (self.ymin, self.ymax) = yr 

161 (self.vmin, self.vmax) = vr 

162 self.cyv = (self.vmax - self.vmin) / (self.ymax - self.ymin) 

163 

164 def v(self, y): 

165 return self.vmin + (y-self.ymin) * self.cyv 

166 

167 def y(self, v): 

168 return self.ymin + (v-self.vmin) / self.cyv 

169 

170 

171class ProjectXYUV(object): 

172 def __init__(self, xyr=((0., 1.), (0., 1.)), uvr=((0., 1.), (0., 1.))): 

173 (self.xmin, self.xmax), (self.ymin, self.ymax) = xyr 

174 (self.umin, self.umax), (self.vmin, self.vmax) = uvr 

175 self.cxu = (self.umax - self.umin) / (self.xmax - self.xmin) 

176 self.cyv = (self.vmax - self.vmin) / (self.ymax - self.ymin) 

177 

178 def u(self, x): 

179 return self.umin + (x-self.xmin) * self.cxu 

180 

181 def v(self, y): 

182 return self.vmin + (y-self.ymin) * self.cyv 

183 

184 def x(self, u): 

185 return self.xmin + (u-self.umin) / self.cxu 

186 

187 def y(self, v): 

188 return self.ymin + (v-self.vmin) / self.cyv 

189 

190 

191class PlotEnv(qg.QPainter): 

192 def __init__(self, *args): 

193 qg.QPainter.__init__(self, *args) 

194 self.umin = 0. 

195 self.umax = 1. 

196 self.vmin = 0. 

197 self.vmax = 1. 

198 

199 def draw_vline(self, project, x, y0, y1): 

200 u = project.u(x) 

201 line = qc.QLineF(u, project.v(y0), u, project.v(y1)) 

202 self.drawLine(line) 

203 

204 def projector(self, xyr): 

205 return ProjectXYUV(xyr, self.uvrange) 

206 

207 @property 

208 def uvrange(self): 

209 return (self.umin, self.umax), (self.vmin, self.vmax) 

210 

211 

212def time_fmt_drumline(t): 

213 ti = int(t) 

214 if ti % 60 == 0: 

215 fmt = '%H:%M' 

216 else: 

217 fmt = '%H:%M:%S' 

218 

219 if ti % (3600*24) == 0: 

220 fmt = fmt + ' %Y-%m-%d' 

221 

222 return fmt 

223 

224 

225class Empty(Exception): 

226 pass 

227 

228 

229class DrumLine(qc.QObject, TalkieConnectionOwner): 

230 def __init__(self, iline, tmin, tmax, traces, state): 

231 qc.QObject.__init__(self) 

232 self.traces = traces 

233 self.tmin = tmin 

234 self.tmax = tmax 

235 self.ymin = 0. 

236 self.ymax = 1. 

237 self.ymin_data = 0. 

238 self.ymax_data = 1. 

239 self.iline = iline 

240 self._last_mode = None 

241 self._ydata_cache = {} 

242 self._time_per_pixel = None 

243 self.access_counter = 0 

244 self.talkie_connect(state, 'style.trace_resolution', self._empty_cache) 

245 

246 def data_range(self, mode='min-max'): 

247 if not self.traces: 

248 raise Empty() 

249 

250 modemap = { 

251 'min-max': 'minmax', 

252 'mean-plusminus-1-sigma': 1., 

253 'mean-plusminus-2-sigma': 2., 

254 'mean-plusminus-4-sigma': 4., 

255 } 

256 

257 if self._last_mode != mode: 

258 mi, ma = trace.minmax( 

259 self.traces, key=lambda tr: None, mode=modemap[mode])[None] 

260 self.ymin_data = mi 

261 self.ymax_data = ma 

262 self._last_mode = mode 

263 

264 return self.ymin_data, self.ymax_data 

265 

266 def set_yrange(self, ymin, ymax): 

267 if ymax == ymin: 

268 ymax = 0.5*(ymin + ymax) + 1.0 

269 ymin = 0.5*(ymin + ymax) - 1.0 

270 

271 self.ymin = ymin 

272 self.ymax = ymax 

273 

274 @property 

275 def tyrange(self): 

276 return (self.tmin, self.tmax), (self.ymin, self.ymax) 

277 

278 def draw(self, plotenv, markers): 

279 self._draw_traces(plotenv) 

280 self._draw_time_label(plotenv) 

281 self._draw_markers(plotenv, markers) 

282 

283 def _draw_time_label(self, plotenv): 

284 text = util.time_to_str(self.tmin, format=time_fmt_drumline(self.tmin)) 

285 font = plotenv.style.label_textstyle.qt_font 

286 lab = Label( 

287 font.pointSize(), plotenv.vmin, text, 

288 anchor='ML', style=plotenv.style.label_textstyle) 

289 

290 lab.draw(plotenv) 

291 

292 def _draw_markers(self, plotenv, markers): 

293 project = plotenv.projector(((self.tmin, self.tmax), (-1., 1.))) 

294 plotenv.setPen(plotenv.style.marker_color.qt_color) 

295 

296 rect = qc.QRectF( 

297 plotenv.umin, 0., 

298 plotenv.umax-plotenv.umin, plotenv.widget.height()) 

299 

300 for marker in markers: 

301 plotenv.draw_vline(project, marker.tmin, -1., 1.) 

302 s = marker.get_label() 

303 if s: 

304 if isinstance(marker, PhaseMarker): 

305 anchor = 'BC' 

306 v = project.v(-1.) 

307 else: 

308 anchor = 'BR' 

309 v = project.v(-1.) 

310 

311 lab = Label( 

312 project.u(marker.tmin), v, s, 

313 anchor=anchor, 

314 style=plotenv.style.marker_textstyle, 

315 keep_inside=rect, 

316 head=(1.0, 3.0)) 

317 

318 lab.draw(plotenv) 

319 

320 def _draw_traces(self, plotenv): 

321 project = plotenv.projector(self.tyrange) 

322 tpp = (project.xmax - project.xmin) / (project.umax - project.umin) 

323 if self._time_per_pixel != tpp: 

324 self._empty_cache() 

325 self._time_per_pixel = tpp 

326 

327 for tr in self.traces: 

328 udata, vdata = self._projected_trace_data( 

329 tr, 

330 project, 

331 plotenv.style.trace_resolution) 

332 

333 qpoints = make_QPolygonF(udata, vdata) 

334 plotenv.setPen(plotenv.style.trace_color.qt_color) 

335 plotenv.drawPolyline(qpoints) 

336 

337 def _projected_trace_data(self, tr, project, trace_resolution): 

338 n = tr.data_len() 

339 if trace_resolution > 0 \ 

340 and n > 2 \ 

341 and tr.deltat < 0.5 / trace_resolution*self._time_per_pixel: 

342 

343 spp = int(self._time_per_pixel / tr.deltat / trace_resolution) 

344 if tr not in self._ydata_cache: 

345 nok = (tr.data_len() // spp) * spp 

346 ydata_rs = tr.ydata[:nok].reshape((-1, spp)) 

347 ydata = num.empty((nok // spp)*2) 

348 ydata[::2] = num.min(ydata_rs, axis=1) 

349 ydata[1::2] = num.max(ydata_rs, axis=1) 

350 self._ydata_cache[tr] = ydata 

351 else: 

352 ydata = self._ydata_cache[tr] 

353 

354 udata_min = float( 

355 project.u(tr.tmin)) 

356 udata_max = float( 

357 project.u(tr.tmin+0.5*tr.deltat*spp*(ydata.size-1))) 

358 else: 

359 ydata = tr.ydata 

360 udata_min = float(project.u(tr.tmin)) 

361 udata_max = float(project.u(tr.tmin+tr.deltat*(n-1))) 

362 

363 vdata = project.v(ydata) 

364 udata = num.linspace(udata_min, udata_max, vdata.size) 

365 return udata, vdata 

366 

367 def _empty_cache(self): 

368 self._ydata_cache = {} 

369 

370 

371def tlen(x): 

372 return x.tmax-x.tmin 

373 

374 

375def is_relevant(x, tmin, tmax): 

376 return tmax >= x.tmin and x.tmax >= tmin 

377 

378 

379class MarkerStore(object): 

380 def __init__(self): 

381 self.empty() 

382 

383 def empty(self): 

384 self._by_tmin = pile.Sorted([], 'tmin') 

385 self._by_tmax = pile.Sorted([], 'tmax') 

386 self._by_tlen = pile.Sorted([], tlen) 

387 self._adjust_minmax() 

388 self._listeners = [] 

389 

390 def relevant(self, tmin, tmax, selector=None): 

391 if not self._by_tmin or not is_relevant(self, tmin, tmax): 

392 return [] 

393 

394 if selector is None: 

395 return [ 

396 x for x in self._by_tmin.with_key_in(tmin-self.tlenmax, tmax) 

397 if is_relevant(x, tmin, tmax)] 

398 else: 

399 return [ 

400 x for x in self._by_tmin.with_key_in(tmin-self.tlenmax, tmax) 

401 if is_relevant(x, tmin, tmax) and selector(x)] 

402 

403 def insert(self, x): 

404 self._by_tmin.insert(x) 

405 self._by_tmax.insert(x) 

406 self._by_tlen.insert(x) 

407 

408 self._adjust_minmax() 

409 self._notify_listeners('insert', x) 

410 

411 def remove(self, x): 

412 self._by_tmin.remove(x) 

413 self._by_tmax.remove(x) 

414 self._by_tlen.remove(x) 

415 

416 self._adjust_minmax() 

417 self._notify_listeners('remove', x) 

418 

419 def insert_many(self, x): 

420 self._by_tmin.insert_many(x) 

421 self._by_tmax.insert_many(x) 

422 self._by_tlen.insert_many(x) 

423 

424 self._adjust_minmax() 

425 self._notify_listeners('insert_many', x) 

426 

427 def remove_many(self, x): 

428 self._by_tmin.remove_many(x) 

429 self._by_tmax.remove_many(x) 

430 self._by_tlen.remove_many(x) 

431 

432 self._adjust_minmax() 

433 self._notify_listeners('remove_many', x) 

434 

435 def add_listener(self, obj): 

436 self._listeners.append(util.smart_weakref(obj)) 

437 

438 def _adjust_minmax(self): 

439 if self._by_tmin: 

440 self.tmin = self._by_tmin.min().tmin 

441 self.tmax = self._by_tmax.max().tmax 

442 self.tlenmax = tlen(self._by_tlen.max()) 

443 else: 

444 self.tmin = None 

445 self.tmax = None 

446 self.tlenmax = None 

447 

448 def _notify_listeners(self, what, x): 

449 for ref in self._listeners: 

450 obj = ref() 

451 if obj: 

452 obj(what, x) 

453 

454 def __iter__(self): 

455 return iter(self._by_tmin) 

456 

457 

458class DrumViewMain(qw.QWidget, TalkieConnectionOwner): 

459 

460 def __init__(self, pile, *args): 

461 qw.QWidget.__init__(self, *args) 

462 

463 self.setAttribute(qc.Qt.WA_AcceptTouchEvents, True) 

464 

465 st = self.state = State() 

466 self.markers = MarkerStore() 

467 self.markers.add_listener(self._markers_changed) 

468 

469 self.pile = pile 

470 self.pile.add_listener(self._pile_changed) 

471 

472 self._drumlines = {} 

473 self._wheel_pos = 0 

474 self._iline_float = None 

475 self._project_iline_to_screen = ProjectYV( 

476 (st.iline-0.5, st.iline+st.nlines+0.5), (0., self.height())) 

477 self._access_counter = 0 

478 self._waiting_for_first_data = True 

479 

480 self._init_touch() 

481 self._init_following() 

482 

483 self.talkie_connect( 

484 self.state, 

485 '', 

486 self._state_changed, 

487 drop_args=True) 

488 

489 self.talkie_connect( 

490 self.state, 

491 ['filters', 'nslc', 'tline'], 

492 self._drop_cached_drumlines, 

493 drop_args=True) 

494 

495 self.talkie_connect( 

496 self.state, 

497 'style.background_color', 

498 self._adjust_background_color, 

499 drop_args=True) 

500 

501 self.talkie_connect( 

502 self.state, 

503 'follow', 

504 self._adjust_follow, 

505 drop_args=True) 

506 

507 self._adjust_background_color() 

508 self._adjust_follow() 

509 

510 def goto_data_begin(self): 

511 if self.pile.tmin: 

512 self.state.tmin = self.pile.tmin 

513 

514 def goto_data_end(self): 

515 if self.pile.tmax: 

516 self.state.tmax = self.pile.tmax 

517 

518 def next_nslc(self): 

519 nslc_ids = sorted(self.pile.nslc_ids.keys()) 

520 if nslc_ids: 

521 try: 

522 i = nslc_ids.index(self.state.nslc) 

523 except ValueError: 

524 i = -1 

525 

526 self.state.nslc = nslc_ids[(i+1) % len(nslc_ids)] 

527 

528 def title(self): 

529 return ' '.join(x for x in self.state.nslc if x) 

530 

531 def _state_changed(self): 

532 self.update() 

533 

534 def _markers_changed(self, *_): 

535 self.update() 

536 

537 def _pile_changed(self, what, content): 

538 self._adjust_first_data() 

539 

540 delete = [] 

541 tline = self.state.tline 

542 for iline in self._drumlines.keys(): 

543 for c in content: 

544 if c.overlaps(iline*tline, (iline+1)*tline): 

545 delete.append(iline) 

546 

547 for iline in delete: 

548 del self._drumlines[iline] 

549 

550 self.update() 

551 

552 def _init_following(self): 

553 self._follow_timer = None 

554 self._following_interrupted = False 

555 self._following_interrupted_tstart = None 

556 

557 def interrupt_following(self): 

558 self._following_interrupted = True 

559 self._following_interrupted_tstart = time.time() 

560 

561 def _follow_update(self): 

562 if self._following_interrupted: 

563 now = time.time() 

564 if self._following_interrupted_tstart < now - 20.: 

565 self._following_interrupted = False 

566 else: 

567 return 

568 

569 now = time.time() 

570 iline = int(math.ceil(now / self.state.tline))-self.state.nlines 

571 if iline != self.state.iline: 

572 self.state.iline = iline 

573 

574 def _adjust_follow(self): 

575 follow = self.state.follow 

576 if follow and not self._follow_timer: 

577 self._follow_timer = qc.QTimer(self) 

578 self._follow_timer.timeout.connect(self._follow_update) 

579 self._follow_timer.setInterval(1000) 

580 self._follow_update() 

581 self._follow_timer.start() 

582 

583 elif not follow and self._follow_timer: 

584 self._follow_timer.stop() 

585 

586 def _draw(self, plotenv): 

587 self._adjust_first_data() 

588 plotenv.umin = 0. 

589 plotenv.umax = self.width() 

590 self._draw_title(plotenv) 

591 self._draw_time_axis(plotenv) 

592 self._draw_lines(plotenv) 

593 

594 def _draw_title(self, plotenv): 

595 font = plotenv.style.title_textstyle.qt_font 

596 lab = Label( 

597 0.5*(plotenv.umin + plotenv.umax), 

598 font.pointSize(), 

599 self.title(), 

600 anchor='TC', 

601 style=plotenv.style.title_textstyle) 

602 

603 lab.draw(plotenv) 

604 

605 def _draw_time_axis(self, plotenv): 

606 pass 

607 

608 def _draw_lines(self, plotenv): 

609 st = self.state 

610 drumlines_seen = [] 

611 for iline in range(st.iline, st.iline+st.nlines): 

612 self._update_line(iline) 

613 drumline = self._drumlines.get(iline, None) 

614 if drumline: 

615 drumlines_seen.append(drumline) 

616 

617 self._autoscale(drumlines_seen) 

618 

619 top_margin = 50. 

620 bottom_margin = 50. 

621 

622 self._project_iline_to_screen = ProjectYV( 

623 (st.iline-0.5, st.iline+st.nlines-0.5), 

624 (top_margin, self.height()-bottom_margin)) 

625 

626 for drumline in drumlines_seen: 

627 plotenv.vmin = self._project_iline_to_screen.v(drumline.iline-0.5) 

628 plotenv.vmax = self._project_iline_to_screen.v(drumline.iline+0.5) 

629 markers = self._relevant_markers(drumline.tmin, drumline.tmax) 

630 drumline.draw(plotenv, markers) 

631 drumline.access_counter = self._access_counter 

632 self._access_counter += 1 

633 

634 drumlines_by_access = sorted( 

635 self._drumlines.values(), key=lambda dl: dl.access_counter) 

636 

637 for drumline in drumlines_by_access[:-st.npages_cache*st.nlines]: 

638 del self._drumlines[drumline.iline] 

639 

640 def _relevant_markers(self, tmin, tmax): 

641 return self.markers.relevant( 

642 tmin, tmax, 

643 lambda m: not m.nslc_ids or m.match_nslc(self.state.nslc)) 

644 

645 def _autoscale(self, drumlines): 

646 if not drumlines: 

647 return 

648 

649 st = self.state 

650 

651 data = [] 

652 for drumline in drumlines: 

653 try: 

654 data.append(drumline.data_range(st.scaling.base)) 

655 except Empty: 

656 pass 

657 

658 if not data: 

659 data = [[0, 0]] 

660 

661 mi, ma = num.array(data, dtype=float).T 

662 gain = st.scaling.gain 

663 if st.scaling.mode == 'same': 

664 ymin, ymax = mi.min(), ma.max() 

665 for drumline in drumlines: 

666 drumline.set_yrange(ymin/gain, ymax/gain) 

667 elif st.scaling.mode == 'individual': 

668 for drumline, ymin, ymax in zip(drumlines, mi, ma): 

669 drumline.set_yrange(ymin/gain, ymax/gain) 

670 elif st.scaling.mode == 'fixed': 

671 for drumline in drumlines: 

672 drumline.set_yrange(st.scaling.min/gain, st.scaling.max/gain) 

673 

674 def _update_line(self, iline): 

675 

676 if iline not in self._drumlines: 

677 st = self.state 

678 tmin = iline*st.tline 

679 tmax = (iline+1)*st.tline 

680 if st.filters: 

681 tpad = max(x.tpad() for x in st.filters) 

682 else: 

683 tpad = 0.0 

684 

685 traces = self.pile.all( 

686 tmin=iline*st.tline, 

687 tmax=(iline+1)*st.tline, 

688 tpad=tpad, 

689 trace_selector=lambda tr: tr.nslc_id == st.nslc, 

690 keep_current_files_open=True, 

691 accessor_id=id(self)) 

692 

693 for tr in traces: 

694 for filter in st.filters: 

695 filter.apply(tr) 

696 

697 self._drumlines[iline] = DrumLine( 

698 iline, tmin, tmax, traces, self.state) 

699 

700 def _drop_cached_drumlines(self): 

701 self._drumlines = {} 

702 

703 def _adjust_background_color(self): 

704 color = self.state.style.background_color.qt_color 

705 

706 p = qg.QPalette() 

707 p.setColor(qg.QPalette.Background, color) 

708 self.setAutoFillBackground(True) 

709 self.setPalette(p)