1# http://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

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

5 

6import sys 

7import os 

8 

9import math 

10import logging 

11import threading 

12import queue 

13from collections import defaultdict 

14 

15from pyrocko.guts import Object, Int, List, Tuple, String, Timestamp, Dict 

16from pyrocko import util, trace 

17from pyrocko.progress import progress 

18from pyrocko.plot import nice_time_tick_inc_approx_secs 

19 

20from . import model, io, cache, dataset 

21 

22from .model import to_kind_id, WaveformOrder, to_kind, to_codes, \ 

23 STATION, CHANNEL, RESPONSE, EVENT, WAVEFORM, codes_patterns_list, \ 

24 codes_patterns_for_kind 

25from .client import fdsn, catalog 

26from .selection import Selection, filldocs 

27from .database import abspath 

28from .operators.base import Operator, CodesPatternFiltering 

29from . import client, environment, error 

30 

31logger = logging.getLogger('psq.base') 

32 

33guts_prefix = 'squirrel' 

34 

35 

36def make_task(*args): 

37 return progress.task(*args, logger=logger) 

38 

39 

40def lpick(condition, seq): 

41 ft = [], [] 

42 for ele in seq: 

43 ft[int(bool(condition(ele)))].append(ele) 

44 

45 return ft 

46 

47 

48def len_plural(obj): 

49 return len(obj), '' if len(obj) == 1 else 's' 

50 

51 

52def blocks(tmin, tmax, deltat, nsamples_block=100000): 

53 tblock = nice_time_tick_inc_approx_secs( 

54 util.to_time_float(deltat * nsamples_block)) 

55 iblock_min = int(math.floor(tmin / tblock)) 

56 iblock_max = int(math.ceil(tmax / tblock)) 

57 for iblock in range(iblock_min, iblock_max): 

58 yield iblock * tblock, (iblock+1) * tblock 

59 

60 

61def gaps(avail, tmin, tmax): 

62 assert tmin < tmax 

63 

64 data = [(tmax, 1), (tmin, -1)] 

65 for (tmin_a, tmax_a) in avail: 

66 assert tmin_a < tmax_a 

67 data.append((tmin_a, 1)) 

68 data.append((tmax_a, -1)) 

69 

70 data.sort() 

71 s = 1 

72 gaps = [] 

73 tmin_g = None 

74 for t, x in data: 

75 if s == 1 and x == -1: 

76 tmin_g = t 

77 elif s == 0 and x == 1 and tmin_g is not None: 

78 tmax_g = t 

79 if tmin_g != tmax_g: 

80 gaps.append((tmin_g, tmax_g)) 

81 

82 s += x 

83 

84 return gaps 

85 

86 

87def order_key(order): 

88 return (order.codes, order.tmin, order.tmax) 

89 

90 

91def _is_exact(pat): 

92 return not ('*' in pat or '?' in pat or ']' in pat or '[' in pat) 

93 

94 

95def prefix_tree(tups): 

96 if not tups: 

97 return [] 

98 

99 if len(tups[0]) == 1: 

100 return sorted((tup[0], []) for tup in tups) 

101 

102 d = defaultdict(list) 

103 for tup in tups: 

104 d[tup[0]].append(tup[1:]) 

105 

106 sub = [] 

107 for k in sorted(d.keys()): 

108 sub.append((k, prefix_tree(d[k]))) 

109 

110 return sub 

111 

112 

113def match_time_span(tmin, tmax, obj): 

114 return (obj.tmin is None or tmax is None or obj.tmin <= tmax) \ 

115 and (tmin is None or obj.tmax is None or tmin < obj.tmax) 

116 

117 

118class Batch(object): 

119 ''' 

120 Batch of waveforms from window-wise data extraction. 

121 

122 Encapsulates state and results yielded for each window in window-wise 

123 waveform extraction with the :py:meth:`Squirrel.chopper_waveforms` method. 

124 

125 *Attributes:* 

126 

127 .. py:attribute:: tmin 

128 

129 Start of this time window. 

130 

131 .. py:attribute:: tmax 

132 

133 End of this time window. 

134 

135 .. py:attribute:: i 

136 

137 Index of this time window in sequence. 

138 

139 .. py:attribute:: n 

140 

141 Total number of time windows in sequence. 

142 

143 .. py:attribute:: igroup 

144 

145 Index of this time window's sequence group. 

146 

147 .. py:attribute:: ngroups 

148 

149 Total number of sequence groups. 

150 

151 .. py:attribute:: traces 

152 

153 Extracted waveforms for this time window. 

154 ''' 

155 

156 def __init__(self, tmin, tmax, i, n, igroup, ngroups, traces): 

157 self.tmin = tmin 

158 self.tmax = tmax 

159 self.i = i 

160 self.n = n 

161 self.igroup = igroup 

162 self.ngroups = ngroups 

163 self.traces = traces 

164 

165 

166class Squirrel(Selection): 

167 ''' 

168 Prompt, lazy, indexing, caching, dynamic seismological dataset access. 

169 

170 :param env: 

171 Squirrel environment instance or directory path to use as starting 

172 point for its detection. By default, the current directory is used as 

173 starting point. When searching for a usable environment the directory 

174 ``'.squirrel'`` or ``'squirrel'`` in the current (or starting point) 

175 directory is used if it exists, otherwise the parent directories are 

176 search upwards for the existence of such a directory. If no such 

177 directory is found, the user's global Squirrel environment 

178 ``'$HOME/.pyrocko/squirrel'`` is used. 

179 :type env: 

180 :py:class:`~pyrocko.squirrel.environment.Environment` or 

181 :py:class:`str` 

182 

183 :param database: 

184 Database instance or path to database. By default the 

185 database found in the detected Squirrel environment is used. 

186 :type database: 

187 :py:class:`~pyrocko.squirrel.database.Database` or :py:class:`str` 

188 

189 :param cache_path: 

190 Directory path to use for data caching. By default, the ``'cache'`` 

191 directory in the detected Squirrel environment is used. 

192 :type cache_path: 

193 :py:class:`str` 

194 

195 :param persistent: 

196 If given a name, create a persistent selection. 

197 :type persistent: 

198 :py:class:`str` 

199 

200 This is the central class of the Squirrel framework. It provides a unified 

201 interface to query and access seismic waveforms, station meta-data and 

202 event information from local file collections and remote data sources. For 

203 prompt responses, a profound database setup is used under the hood. To 

204 speed up assemblage of ad-hoc data selections, files are indexed on first 

205 use and the extracted meta-data is remembered in the database for 

206 subsequent accesses. Bulk data is lazily loaded from disk and remote 

207 sources, just when requested. Once loaded, data is cached in memory to 

208 expedite typical access patterns. Files and data sources can be dynamically 

209 added to and removed from the Squirrel selection at runtime. 

210 

211 Queries are restricted to the contents of the files currently added to the 

212 Squirrel selection (usually a subset of the file meta-information 

213 collection in the database). This list of files is referred to here as the 

214 "selection". By default, temporary tables are created in the attached 

215 database to hold the names of the files in the selection as well as various 

216 indices and counters. These tables are only visible inside the application 

217 which created them and are deleted when the database connection is closed 

218 or the application exits. To create a selection which is not deleted at 

219 exit, supply a name to the ``persistent`` argument of the Squirrel 

220 constructor. Persistent selections are shared among applications using the 

221 same database. 

222 

223 **Method summary** 

224 

225 Some of the methods are implemented in :py:class:`Squirrel`'s base class 

226 :py:class:`~pyrocko.squirrel.selection.Selection`. 

227 

228 .. autosummary:: 

229 

230 ~Squirrel.add 

231 ~Squirrel.add_source 

232 ~Squirrel.add_fdsn 

233 ~Squirrel.add_catalog 

234 ~Squirrel.add_dataset 

235 ~Squirrel.add_virtual 

236 ~Squirrel.update 

237 ~Squirrel.update_waveform_promises 

238 ~Squirrel.advance_accessor 

239 ~Squirrel.clear_accessor 

240 ~Squirrel.reload 

241 ~pyrocko.squirrel.selection.Selection.iter_paths 

242 ~Squirrel.iter_nuts 

243 ~Squirrel.iter_kinds 

244 ~Squirrel.iter_deltats 

245 ~Squirrel.iter_codes 

246 ~pyrocko.squirrel.selection.Selection.get_paths 

247 ~Squirrel.get_nuts 

248 ~Squirrel.get_kinds 

249 ~Squirrel.get_deltats 

250 ~Squirrel.get_codes 

251 ~Squirrel.get_counts 

252 ~Squirrel.get_time_span 

253 ~Squirrel.get_deltat_span 

254 ~Squirrel.get_nfiles 

255 ~Squirrel.get_nnuts 

256 ~Squirrel.get_total_size 

257 ~Squirrel.get_stats 

258 ~Squirrel.get_content 

259 ~Squirrel.get_stations 

260 ~Squirrel.get_channels 

261 ~Squirrel.get_responses 

262 ~Squirrel.get_events 

263 ~Squirrel.get_waveform_nuts 

264 ~Squirrel.get_waveforms 

265 ~Squirrel.chopper_waveforms 

266 ~Squirrel.get_coverage 

267 ~Squirrel.pile 

268 ~Squirrel.snuffle 

269 ~Squirrel.glob_codes 

270 ~pyrocko.squirrel.selection.Selection.get_database 

271 ~Squirrel.print_tables 

272 ''' 

273 

274 def __init__( 

275 self, env=None, database=None, cache_path=None, persistent=None): 

276 

277 if not isinstance(env, environment.Environment): 

278 env = environment.get_environment(env) 

279 

280 if database is None: 

281 database = env.expand_path(env.database_path) 

282 

283 if cache_path is None: 

284 cache_path = env.expand_path(env.cache_path) 

285 

286 if persistent is None: 

287 persistent = env.persistent 

288 

289 Selection.__init__( 

290 self, database=database, persistent=persistent) 

291 

292 self.get_database().set_basepath(os.path.dirname(env.get_basepath())) 

293 

294 self._content_caches = { 

295 'waveform': cache.ContentCache(), 

296 'default': cache.ContentCache()} 

297 

298 self._cache_path = cache_path 

299 

300 self._sources = [] 

301 self._operators = [] 

302 self._operator_registry = {} 

303 

304 self._pending_orders = [] 

305 

306 self._pile = None 

307 self._n_choppers_active = 0 

308 

309 self._names.update({ 

310 'nuts': self.name + '_nuts', 

311 'kind_codes_count': self.name + '_kind_codes_count', 

312 'coverage': self.name + '_coverage'}) 

313 

314 with self.transaction('create tables') as cursor: 

315 self._create_tables_squirrel(cursor) 

316 

317 def _create_tables_squirrel(self, cursor): 

318 

319 cursor.execute(self._register_table(self._sql( 

320 ''' 

321 CREATE TABLE IF NOT EXISTS %(db)s.%(nuts)s ( 

322 nut_id integer PRIMARY KEY, 

323 file_id integer, 

324 file_segment integer, 

325 file_element integer, 

326 kind_id integer, 

327 kind_codes_id integer, 

328 tmin_seconds integer, 

329 tmin_offset integer, 

330 tmax_seconds integer, 

331 tmax_offset integer, 

332 kscale integer) 

333 '''))) 

334 

335 cursor.execute(self._register_table(self._sql( 

336 ''' 

337 CREATE TABLE IF NOT EXISTS %(db)s.%(kind_codes_count)s ( 

338 kind_codes_id integer PRIMARY KEY, 

339 count integer) 

340 '''))) 

341 

342 cursor.execute(self._sql( 

343 ''' 

344 CREATE UNIQUE INDEX IF NOT EXISTS %(db)s.%(nuts)s_file_element 

345 ON %(nuts)s (file_id, file_segment, file_element) 

346 ''')) 

347 

348 cursor.execute(self._sql( 

349 ''' 

350 CREATE INDEX IF NOT EXISTS %(db)s.%(nuts)s_index_file_id 

351 ON %(nuts)s (file_id) 

352 ''')) 

353 

354 cursor.execute(self._sql( 

355 ''' 

356 CREATE INDEX IF NOT EXISTS %(db)s.%(nuts)s_index_tmin_seconds 

357 ON %(nuts)s (kind_id, tmin_seconds) 

358 ''')) 

359 

360 cursor.execute(self._sql( 

361 ''' 

362 CREATE INDEX IF NOT EXISTS %(db)s.%(nuts)s_index_tmax_seconds 

363 ON %(nuts)s (kind_id, tmax_seconds) 

364 ''')) 

365 

366 cursor.execute(self._sql( 

367 ''' 

368 CREATE INDEX IF NOT EXISTS %(db)s.%(nuts)s_index_kscale 

369 ON %(nuts)s (kind_id, kscale, tmin_seconds) 

370 ''')) 

371 

372 cursor.execute(self._sql( 

373 ''' 

374 CREATE TRIGGER IF NOT EXISTS %(db)s.%(nuts)s_delete_nuts 

375 BEFORE DELETE ON main.files FOR EACH ROW 

376 BEGIN 

377 DELETE FROM %(nuts)s WHERE file_id == old.file_id; 

378 END 

379 ''')) 

380 

381 # trigger only on size to make silent update of mtime possible 

382 cursor.execute(self._sql( 

383 ''' 

384 CREATE TRIGGER IF NOT EXISTS %(db)s.%(nuts)s_delete_nuts2 

385 BEFORE UPDATE OF size ON main.files FOR EACH ROW 

386 BEGIN 

387 DELETE FROM %(nuts)s WHERE file_id == old.file_id; 

388 END 

389 ''')) 

390 

391 cursor.execute(self._sql( 

392 ''' 

393 CREATE TRIGGER IF NOT EXISTS 

394 %(db)s.%(file_states)s_delete_files 

395 BEFORE DELETE ON %(db)s.%(file_states)s FOR EACH ROW 

396 BEGIN 

397 DELETE FROM %(nuts)s WHERE file_id == old.file_id; 

398 END 

399 ''')) 

400 

401 cursor.execute(self._sql( 

402 ''' 

403 CREATE TRIGGER IF NOT EXISTS %(db)s.%(nuts)s_inc_kind_codes 

404 BEFORE INSERT ON %(nuts)s FOR EACH ROW 

405 BEGIN 

406 INSERT OR IGNORE INTO %(kind_codes_count)s VALUES 

407 (new.kind_codes_id, 0); 

408 UPDATE %(kind_codes_count)s 

409 SET count = count + 1 

410 WHERE new.kind_codes_id 

411 == %(kind_codes_count)s.kind_codes_id; 

412 END 

413 ''')) 

414 

415 cursor.execute(self._sql( 

416 ''' 

417 CREATE TRIGGER IF NOT EXISTS %(db)s.%(nuts)s_dec_kind_codes 

418 BEFORE DELETE ON %(nuts)s FOR EACH ROW 

419 BEGIN 

420 UPDATE %(kind_codes_count)s 

421 SET count = count - 1 

422 WHERE old.kind_codes_id 

423 == %(kind_codes_count)s.kind_codes_id; 

424 END 

425 ''')) 

426 

427 cursor.execute(self._register_table(self._sql( 

428 ''' 

429 CREATE TABLE IF NOT EXISTS %(db)s.%(coverage)s ( 

430 kind_codes_id integer, 

431 time_seconds integer, 

432 time_offset integer, 

433 step integer) 

434 '''))) 

435 

436 cursor.execute(self._sql( 

437 ''' 

438 CREATE UNIQUE INDEX IF NOT EXISTS %(db)s.%(coverage)s_time 

439 ON %(coverage)s (kind_codes_id, time_seconds, time_offset) 

440 ''')) 

441 

442 cursor.execute(self._sql( 

443 ''' 

444 CREATE TRIGGER IF NOT EXISTS %(db)s.%(nuts)s_add_coverage 

445 AFTER INSERT ON %(nuts)s FOR EACH ROW 

446 BEGIN 

447 INSERT OR IGNORE INTO %(coverage)s VALUES 

448 (new.kind_codes_id, new.tmin_seconds, new.tmin_offset, 0) 

449 ; 

450 UPDATE %(coverage)s 

451 SET step = step + 1 

452 WHERE new.kind_codes_id == %(coverage)s.kind_codes_id 

453 AND new.tmin_seconds == %(coverage)s.time_seconds 

454 AND new.tmin_offset == %(coverage)s.time_offset 

455 ; 

456 INSERT OR IGNORE INTO %(coverage)s VALUES 

457 (new.kind_codes_id, new.tmax_seconds, new.tmax_offset, 0) 

458 ; 

459 UPDATE %(coverage)s 

460 SET step = step - 1 

461 WHERE new.kind_codes_id == %(coverage)s.kind_codes_id 

462 AND new.tmax_seconds == %(coverage)s.time_seconds 

463 AND new.tmax_offset == %(coverage)s.time_offset 

464 ; 

465 DELETE FROM %(coverage)s 

466 WHERE new.kind_codes_id == %(coverage)s.kind_codes_id 

467 AND new.tmin_seconds == %(coverage)s.time_seconds 

468 AND new.tmin_offset == %(coverage)s.time_offset 

469 AND step == 0 

470 ; 

471 DELETE FROM %(coverage)s 

472 WHERE new.kind_codes_id == %(coverage)s.kind_codes_id 

473 AND new.tmax_seconds == %(coverage)s.time_seconds 

474 AND new.tmax_offset == %(coverage)s.time_offset 

475 AND step == 0 

476 ; 

477 END 

478 ''')) 

479 

480 cursor.execute(self._sql( 

481 ''' 

482 CREATE TRIGGER IF NOT EXISTS %(db)s.%(nuts)s_remove_coverage 

483 BEFORE DELETE ON %(nuts)s FOR EACH ROW 

484 BEGIN 

485 INSERT OR IGNORE INTO %(coverage)s VALUES 

486 (old.kind_codes_id, old.tmin_seconds, old.tmin_offset, 0) 

487 ; 

488 UPDATE %(coverage)s 

489 SET step = step - 1 

490 WHERE old.kind_codes_id == %(coverage)s.kind_codes_id 

491 AND old.tmin_seconds == %(coverage)s.time_seconds 

492 AND old.tmin_offset == %(coverage)s.time_offset 

493 ; 

494 INSERT OR IGNORE INTO %(coverage)s VALUES 

495 (old.kind_codes_id, old.tmax_seconds, old.tmax_offset, 0) 

496 ; 

497 UPDATE %(coverage)s 

498 SET step = step + 1 

499 WHERE old.kind_codes_id == %(coverage)s.kind_codes_id 

500 AND old.tmax_seconds == %(coverage)s.time_seconds 

501 AND old.tmax_offset == %(coverage)s.time_offset 

502 ; 

503 DELETE FROM %(coverage)s 

504 WHERE old.kind_codes_id == %(coverage)s.kind_codes_id 

505 AND old.tmin_seconds == %(coverage)s.time_seconds 

506 AND old.tmin_offset == %(coverage)s.time_offset 

507 AND step == 0 

508 ; 

509 DELETE FROM %(coverage)s 

510 WHERE old.kind_codes_id == %(coverage)s.kind_codes_id 

511 AND old.tmax_seconds == %(coverage)s.time_seconds 

512 AND old.tmax_offset == %(coverage)s.time_offset 

513 AND step == 0 

514 ; 

515 END 

516 ''')) 

517 

518 def _delete(self): 

519 '''Delete database tables associated with this Squirrel.''' 

520 

521 with self.transaction('delete tables') as cursor: 

522 for s in ''' 

523 DROP TRIGGER %(db)s.%(nuts)s_delete_nuts; 

524 DROP TRIGGER %(db)s.%(nuts)s_delete_nuts2; 

525 DROP TRIGGER %(db)s.%(file_states)s_delete_files; 

526 DROP TRIGGER %(db)s.%(nuts)s_inc_kind_codes; 

527 DROP TRIGGER %(db)s.%(nuts)s_dec_kind_codes; 

528 DROP TABLE %(db)s.%(nuts)s; 

529 DROP TABLE %(db)s.%(kind_codes_count)s; 

530 DROP TRIGGER IF EXISTS %(db)s.%(nuts)s_add_coverage; 

531 DROP TRIGGER IF EXISTS %(db)s.%(nuts)s_remove_coverage; 

532 DROP TABLE IF EXISTS %(db)s.%(coverage)s; 

533 '''.strip().splitlines(): 

534 

535 cursor.execute(self._sql(s)) 

536 

537 Selection._delete(self) 

538 

539 @filldocs 

540 def add(self, 

541 paths, 

542 kinds=None, 

543 format='detect', 

544 include=None, 

545 exclude=None, 

546 check=True): 

547 

548 ''' 

549 Add files to the selection. 

550 

551 :param paths: 

552 Iterator yielding paths to files or directories to be added to the 

553 selection. Recurses into directories. If given a ``str``, it 

554 is treated as a single path to be added. 

555 :type paths: 

556 :py:class:`list` of :py:class:`str` 

557 

558 :param kinds: 

559 Content types to be made available through the Squirrel selection. 

560 By default, all known content types are accepted. 

561 :type kinds: 

562 :py:class:`list` of :py:class:`str` 

563 

564 :param format: 

565 File format identifier or ``'detect'`` to enable auto-detection 

566 (available: %(file_formats)s). 

567 :type format: 

568 str 

569 

570 :param include: 

571 If not ``None``, files are only included if their paths match the 

572 given regular expression pattern. 

573 :type format: 

574 str 

575 

576 :param exclude: 

577 If not ``None``, files are only included if their paths do not 

578 match the given regular expression pattern. 

579 :type format: 

580 str 

581 

582 :param check: 

583 If ``True``, all file modification times are checked to see if 

584 cached information has to be updated (slow). If ``False``, only 

585 previously unknown files are indexed and cached information is used 

586 for known files, regardless of file state (fast, corrresponds to 

587 Squirrel's ``--optimistic`` mode). File deletions will go 

588 undetected in the latter case. 

589 :type check: 

590 bool 

591 

592 :Complexity: 

593 O(log N) 

594 ''' 

595 

596 if isinstance(kinds, str): 

597 kinds = (kinds,) 

598 

599 if isinstance(paths, str): 

600 paths = [paths] 

601 

602 kind_mask = model.to_kind_mask(kinds) 

603 

604 with progress.view(): 

605 Selection.add( 

606 self, util.iter_select_files( 

607 paths, 

608 show_progress=False, 

609 include=include, 

610 exclude=exclude, 

611 pass_through=lambda path: path.startswith('virtual:') 

612 ), kind_mask, format) 

613 

614 self._load(check) 

615 self._update_nuts() 

616 

617 def reload(self): 

618 ''' 

619 Check for modifications and reindex modified files. 

620 

621 Based on file modification times. 

622 ''' 

623 

624 self._set_file_states_force_check() 

625 self._load(check=True) 

626 self._update_nuts() 

627 

628 def add_virtual(self, nuts, virtual_paths=None): 

629 ''' 

630 Add content which is not backed by files. 

631 

632 :param nuts: 

633 Content pieces to be added. 

634 :type nuts: 

635 iterator yielding :py:class:`~pyrocko.squirrel.model.Nut` objects 

636 

637 :param virtual_paths: 

638 List of virtual paths to prevent creating a temporary list of the 

639 nuts while aggregating the file paths for the selection. 

640 :type virtual_paths: 

641 :py:class:`list` of :py:class:`str` 

642 

643 Stores to the main database and the selection. 

644 ''' 

645 

646 if isinstance(virtual_paths, str): 

647 virtual_paths = [virtual_paths] 

648 

649 if virtual_paths is None: 

650 if not isinstance(nuts, list): 

651 nuts = list(nuts) 

652 virtual_paths = set(nut.file_path for nut in nuts) 

653 

654 Selection.add(self, virtual_paths) 

655 self.get_database().dig(nuts) 

656 self._update_nuts() 

657 

658 def add_volatile(self, nuts): 

659 if not isinstance(nuts, list): 

660 nuts = list(nuts) 

661 

662 paths = list(set(nut.file_path for nut in nuts)) 

663 io.backends.virtual.add_nuts(nuts) 

664 self.add_virtual(nuts, paths) 

665 self._volatile_paths.extend(paths) 

666 

667 def add_volatile_waveforms(self, traces): 

668 ''' 

669 Add in-memory waveforms which will be removed when the app closes. 

670 ''' 

671 

672 name = model.random_name() 

673 

674 path = 'virtual:volatile:%s' % name 

675 

676 nuts = [] 

677 for itr, tr in enumerate(traces): 

678 assert tr.tmin <= tr.tmax 

679 tmin_seconds, tmin_offset = model.tsplit(tr.tmin) 

680 tmax_seconds, tmax_offset = model.tsplit( 

681 tr.tmin + tr.data_len()*tr.deltat) 

682 

683 nuts.append(model.Nut( 

684 file_path=path, 

685 file_format='virtual', 

686 file_segment=itr, 

687 file_element=0, 

688 file_mtime=0, 

689 codes=tr.codes, 

690 tmin_seconds=tmin_seconds, 

691 tmin_offset=tmin_offset, 

692 tmax_seconds=tmax_seconds, 

693 tmax_offset=tmax_offset, 

694 deltat=tr.deltat, 

695 kind_id=to_kind_id('waveform'), 

696 content=tr)) 

697 

698 self.add_volatile(nuts) 

699 return path 

700 

701 def _load(self, check): 

702 for _ in io.iload( 

703 self, 

704 content=[], 

705 skip_unchanged=True, 

706 check=check): 

707 pass 

708 

709 def _update_nuts(self, transaction=None): 

710 transaction = transaction or self.transaction('update nuts') 

711 with make_task('Aggregating selection') as task, \ 

712 transaction as cursor: 

713 

714 self._conn.set_progress_handler(task.update, 100000) 

715 nrows = cursor.execute(self._sql( 

716 ''' 

717 INSERT INTO %(db)s.%(nuts)s 

718 SELECT NULL, 

719 nuts.file_id, nuts.file_segment, nuts.file_element, 

720 nuts.kind_id, nuts.kind_codes_id, 

721 nuts.tmin_seconds, nuts.tmin_offset, 

722 nuts.tmax_seconds, nuts.tmax_offset, 

723 nuts.kscale 

724 FROM %(db)s.%(file_states)s 

725 INNER JOIN nuts 

726 ON %(db)s.%(file_states)s.file_id == nuts.file_id 

727 INNER JOIN kind_codes 

728 ON nuts.kind_codes_id == 

729 kind_codes.kind_codes_id 

730 WHERE %(db)s.%(file_states)s.file_state != 2 

731 AND (((1 << kind_codes.kind_id) 

732 & %(db)s.%(file_states)s.kind_mask) != 0) 

733 ''')).rowcount 

734 

735 task.update(nrows) 

736 self._set_file_states_known(transaction) 

737 self._conn.set_progress_handler(None, 0) 

738 

739 def add_source(self, source, check=True): 

740 ''' 

741 Add remote resource. 

742 

743 :param source: 

744 Remote data access client instance. 

745 :type source: 

746 subclass of :py:class:`~pyrocko.squirrel.client.base.Source` 

747 ''' 

748 

749 self._sources.append(source) 

750 source.setup(self, check=check) 

751 

752 def add_fdsn(self, *args, **kwargs): 

753 ''' 

754 Add FDSN site for transparent remote data access. 

755 

756 Arguments are passed to 

757 :py:class:`~pyrocko.squirrel.client.fdsn.FDSNSource`. 

758 ''' 

759 

760 self.add_source(fdsn.FDSNSource(*args, **kwargs)) 

761 

762 def add_catalog(self, *args, **kwargs): 

763 ''' 

764 Add online catalog for transparent event data access. 

765 

766 Arguments are passed to 

767 :py:class:`~pyrocko.squirrel.client.catalog.CatalogSource`. 

768 ''' 

769 

770 self.add_source(catalog.CatalogSource(*args, **kwargs)) 

771 

772 def add_dataset(self, ds, check=True): 

773 ''' 

774 Read dataset description from file and add its contents. 

775 

776 :param ds: 

777 Path to dataset description file or dataset description object 

778 . See :py:mod:`~pyrocko.squirrel.dataset`. 

779 :type ds: 

780 :py:class:`str` or :py:class:`~pyrocko.squirrel.dataset.Dataset` 

781 

782 :param check: 

783 If ``True``, all file modification times are checked to see if 

784 cached information has to be updated (slow). If ``False``, only 

785 previously unknown files are indexed and cached information is used 

786 for known files, regardless of file state (fast, corrresponds to 

787 Squirrel's ``--optimistic`` mode). File deletions will go 

788 undetected in the latter case. 

789 :type check: 

790 bool 

791 ''' 

792 if isinstance(ds, str): 

793 ds = dataset.read_dataset(ds) 

794 

795 ds.setup(self, check=check) 

796 

797 def _get_selection_args( 

798 self, kind_id, 

799 obj=None, tmin=None, tmax=None, time=None, codes=None): 

800 

801 if codes is not None: 

802 codes = codes_patterns_for_kind(kind_id, codes) 

803 

804 if time is not None: 

805 tmin = time 

806 tmax = time 

807 

808 if obj is not None: 

809 tmin = tmin if tmin is not None else obj.tmin 

810 tmax = tmax if tmax is not None else obj.tmax 

811 codes = codes if codes is not None else codes_patterns_for_kind( 

812 kind_id, obj.codes) 

813 

814 return tmin, tmax, codes 

815 

816 def _get_selection_args_str(self, *args, **kwargs): 

817 

818 tmin, tmax, codes = self._get_selection_args(*args, **kwargs) 

819 return 'tmin: %s, tmax: %s, codes: %s' % ( 

820 util.time_to_str(tmin) if tmin is not None else 'none', 

821 util.time_to_str(tmax) if tmax is not None else 'none', 

822 ','.join(str(entry) for entry in codes)) 

823 

824 def _selection_args_to_kwargs( 

825 self, obj=None, tmin=None, tmax=None, time=None, codes=None): 

826 

827 return dict(obj=obj, tmin=tmin, tmax=tmax, time=time, codes=codes) 

828 

829 def _timerange_sql(self, tmin, tmax, kind, cond, args, naiv): 

830 

831 tmin_seconds, tmin_offset = model.tsplit(tmin) 

832 tmax_seconds, tmax_offset = model.tsplit(tmax) 

833 if naiv: 

834 cond.append('%(db)s.%(nuts)s.tmin_seconds <= ?') 

835 args.append(tmax_seconds) 

836 else: 

837 tscale_edges = model.tscale_edges 

838 tmin_cond = [] 

839 for kscale in range(tscale_edges.size + 1): 

840 if kscale != tscale_edges.size: 

841 tscale = int(tscale_edges[kscale]) 

842 tmin_cond.append(''' 

843 (%(db)s.%(nuts)s.kind_id = ? 

844 AND %(db)s.%(nuts)s.kscale == ? 

845 AND %(db)s.%(nuts)s.tmin_seconds BETWEEN ? AND ?) 

846 ''') 

847 args.extend( 

848 (to_kind_id(kind), kscale, 

849 tmin_seconds - tscale - 1, tmax_seconds + 1)) 

850 

851 else: 

852 tmin_cond.append(''' 

853 (%(db)s.%(nuts)s.kind_id == ? 

854 AND %(db)s.%(nuts)s.kscale == ? 

855 AND %(db)s.%(nuts)s.tmin_seconds <= ?) 

856 ''') 

857 

858 args.extend( 

859 (to_kind_id(kind), kscale, tmax_seconds + 1)) 

860 if tmin_cond: 

861 cond.append(' ( ' + ' OR '.join(tmin_cond) + ' ) ') 

862 

863 cond.append('%(db)s.%(nuts)s.tmax_seconds >= ?') 

864 args.append(tmin_seconds) 

865 

866 def _codes_match_sql(self, kind_id, codes, cond, args): 

867 pats = codes_patterns_for_kind(kind_id, codes) 

868 if pats is None: 

869 return 

870 

871 pats_exact = [] 

872 pats_nonexact = [] 

873 for pat in pats: 

874 spat = pat.safe_str 

875 (pats_exact if _is_exact(spat) else pats_nonexact).append(spat) 

876 

877 cond_exact = None 

878 if pats_exact: 

879 cond_exact = ' ( kind_codes.codes IN ( %s ) ) ' % ', '.join( 

880 '?'*len(pats_exact)) 

881 

882 args.extend(pats_exact) 

883 

884 cond_nonexact = None 

885 if pats_nonexact: 

886 cond_nonexact = ' ( %s ) ' % ' OR '.join( 

887 ('kind_codes.codes GLOB ?',) * len(pats_nonexact)) 

888 

889 args.extend(pats_nonexact) 

890 

891 if cond_exact and cond_nonexact: 

892 cond.append(' ( %s OR %s ) ' % (cond_exact, cond_nonexact)) 

893 

894 elif cond_exact: 

895 cond.append(cond_exact) 

896 

897 elif cond_nonexact: 

898 cond.append(cond_nonexact) 

899 

900 def iter_nuts( 

901 self, kind=None, tmin=None, tmax=None, codes=None, naiv=False, 

902 kind_codes_ids=None, path=None, limit=None): 

903 

904 ''' 

905 Iterate over content entities matching given constraints. 

906 

907 :param kind: 

908 Content kind (or kinds) to extract. 

909 :type kind: 

910 :py:class:`str`, :py:class:`list` of :py:class:`str` 

911 

912 :param tmin: 

913 Start time of query interval. 

914 :type tmin: 

915 timestamp 

916 

917 :param tmax: 

918 End time of query interval. 

919 :type tmax: 

920 timestamp 

921 

922 :param codes: 

923 List of code patterns to query. 

924 :type codes: 

925 :py:class:`list` of :py:class:`~pyrocko.squirrel.model.Codes` 

926 objects appropriate for the queried content type, or anything which 

927 can be converted to such objects. 

928 

929 :param naiv: 

930 Bypass time span lookup through indices (slow, for testing). 

931 :type naiv: 

932 :py:class:`bool` 

933 

934 :param kind_codes_ids: 

935 Kind-codes IDs of contents to be retrieved (internal use). 

936 :type kind_codes_ids: 

937 :py:class:`list` of :py:class:`int` 

938 

939 :yields: 

940 :py:class:`~pyrocko.squirrel.model.Nut` objects representing the 

941 intersecting content. 

942 

943 :complexity: 

944 O(log N) for the time selection part due to heavy use of database 

945 indices. 

946 

947 Query time span is treated as a half-open interval ``[tmin, tmax)``. 

948 However, if ``tmin`` equals ``tmax``, the edge logics are modified to 

949 closed-interval so that content intersecting with the time instant ``t 

950 = tmin = tmax`` is returned (otherwise nothing would be returned as 

951 ``[t, t)`` never matches anything). 

952 

953 Time spans of content entities to be matched are also treated as half 

954 open intervals, e.g. content span ``[0, 1)`` is matched by query span 

955 ``[0, 1)`` but not by ``[-1, 0)`` or ``[1, 2)``. Also here, logics are 

956 modified to closed-interval when the content time span is an empty 

957 interval, i.e. to indicate a time instant. E.g. time instant 0 is 

958 matched by ``[0, 1)`` but not by ``[-1, 0)`` or ``[1, 2)``. 

959 ''' 

960 

961 if not isinstance(kind, str): 

962 if kind is None: 

963 kind = model.g_content_kinds 

964 for kind_ in kind: 

965 for nut in self.iter_nuts(kind_, tmin, tmax, codes): 

966 yield nut 

967 

968 return 

969 

970 kind_id = to_kind_id(kind) 

971 

972 cond = [] 

973 args = [] 

974 if tmin is not None or tmax is not None: 

975 assert kind is not None 

976 if tmin is None: 

977 tmin = self.get_time_span()[0] 

978 if tmax is None: 

979 tmax = self.get_time_span()[1] + 1.0 

980 

981 self._timerange_sql(tmin, tmax, kind, cond, args, naiv) 

982 

983 cond.append('kind_codes.kind_id == ?') 

984 args.append(kind_id) 

985 

986 if codes is not None: 

987 self._codes_match_sql(kind_id, codes, cond, args) 

988 

989 if kind_codes_ids is not None: 

990 cond.append( 

991 ' ( kind_codes.kind_codes_id IN ( %s ) ) ' % ', '.join( 

992 '?'*len(kind_codes_ids))) 

993 

994 args.extend(kind_codes_ids) 

995 

996 db = self.get_database() 

997 if path is not None: 

998 cond.append('files.path == ?') 

999 args.append(db.relpath(abspath(path))) 

1000 

1001 sql = (''' 

1002 SELECT 

1003 files.path, 

1004 files.format, 

1005 files.mtime, 

1006 files.size, 

1007 %(db)s.%(nuts)s.file_segment, 

1008 %(db)s.%(nuts)s.file_element, 

1009 kind_codes.kind_id, 

1010 kind_codes.codes, 

1011 %(db)s.%(nuts)s.tmin_seconds, 

1012 %(db)s.%(nuts)s.tmin_offset, 

1013 %(db)s.%(nuts)s.tmax_seconds, 

1014 %(db)s.%(nuts)s.tmax_offset, 

1015 kind_codes.deltat 

1016 FROM files 

1017 INNER JOIN %(db)s.%(nuts)s 

1018 ON files.file_id == %(db)s.%(nuts)s.file_id 

1019 INNER JOIN kind_codes 

1020 ON %(db)s.%(nuts)s.kind_codes_id == kind_codes.kind_codes_id 

1021 ''') 

1022 

1023 if cond: 

1024 sql += ''' WHERE ''' + ' AND '.join(cond) 

1025 

1026 if limit is not None: 

1027 sql += ''' LIMIT %i''' % limit 

1028 

1029 sql = self._sql(sql) 

1030 if tmin is None and tmax is None: 

1031 for row in self._conn.execute(sql, args): 

1032 row = (db.abspath(row[0]),) + row[1:] 

1033 nut = model.Nut(values_nocheck=row) 

1034 yield nut 

1035 else: 

1036 assert tmin is not None and tmax is not None 

1037 if tmin == tmax: 

1038 for row in self._conn.execute(sql, args): 

1039 row = (db.abspath(row[0]),) + row[1:] 

1040 nut = model.Nut(values_nocheck=row) 

1041 if (nut.tmin <= tmin < nut.tmax) \ 

1042 or (nut.tmin == nut.tmax and tmin == nut.tmin): 

1043 

1044 yield nut 

1045 else: 

1046 for row in self._conn.execute(sql, args): 

1047 row = (db.abspath(row[0]),) + row[1:] 

1048 nut = model.Nut(values_nocheck=row) 

1049 if (tmin < nut.tmax and nut.tmin < tmax) \ 

1050 or (nut.tmin == nut.tmax 

1051 and tmin <= nut.tmin < tmax): 

1052 

1053 yield nut 

1054 

1055 def get_nuts(self, *args, **kwargs): 

1056 ''' 

1057 Get content entities matching given constraints. 

1058 

1059 Like :py:meth:`iter_nuts` but returns results as a list. 

1060 ''' 

1061 

1062 return list(self.iter_nuts(*args, **kwargs)) 

1063 

1064 def _split_nuts( 

1065 self, kind, tmin=None, tmax=None, codes=None, path=None): 

1066 

1067 kind_id = to_kind_id(kind) 

1068 tmin_seconds, tmin_offset = model.tsplit(tmin) 

1069 tmax_seconds, tmax_offset = model.tsplit(tmax) 

1070 

1071 names_main_nuts = dict(self._names) 

1072 names_main_nuts.update(db='main', nuts='nuts') 

1073 

1074 db = self.get_database() 

1075 

1076 def main_nuts(s): 

1077 return s % names_main_nuts 

1078 

1079 with self.transaction('split nuts') as cursor: 

1080 # modify selection and main 

1081 for sql_subst in [ 

1082 self._sql, main_nuts]: 

1083 

1084 cond = [] 

1085 args = [] 

1086 

1087 self._timerange_sql(tmin, tmax, kind, cond, args, False) 

1088 

1089 if codes is not None: 

1090 self._codes_match_sql(kind_id, codes, cond, args) 

1091 

1092 if path is not None: 

1093 cond.append('files.path == ?') 

1094 args.append(db.relpath(abspath(path))) 

1095 

1096 sql = sql_subst(''' 

1097 SELECT 

1098 %(db)s.%(nuts)s.nut_id, 

1099 %(db)s.%(nuts)s.tmin_seconds, 

1100 %(db)s.%(nuts)s.tmin_offset, 

1101 %(db)s.%(nuts)s.tmax_seconds, 

1102 %(db)s.%(nuts)s.tmax_offset, 

1103 kind_codes.deltat 

1104 FROM files 

1105 INNER JOIN %(db)s.%(nuts)s 

1106 ON files.file_id == %(db)s.%(nuts)s.file_id 

1107 INNER JOIN kind_codes 

1108 ON %(db)s.%(nuts)s.kind_codes_id == kind_codes.kind_codes_id 

1109 WHERE ''' + ' AND '.join(cond)) # noqa 

1110 

1111 insert = [] 

1112 delete = [] 

1113 for row in cursor.execute(sql, args): 

1114 nut_id, nut_tmin_seconds, nut_tmin_offset, \ 

1115 nut_tmax_seconds, nut_tmax_offset, nut_deltat = row 

1116 

1117 nut_tmin = model.tjoin( 

1118 nut_tmin_seconds, nut_tmin_offset) 

1119 nut_tmax = model.tjoin( 

1120 nut_tmax_seconds, nut_tmax_offset) 

1121 

1122 if nut_tmin < tmax and tmin < nut_tmax: 

1123 if nut_tmin < tmin: 

1124 insert.append(( 

1125 nut_tmin_seconds, nut_tmin_offset, 

1126 tmin_seconds, tmin_offset, 

1127 model.tscale_to_kscale( 

1128 tmin_seconds - nut_tmin_seconds), 

1129 nut_id)) 

1130 

1131 if tmax < nut_tmax: 

1132 insert.append(( 

1133 tmax_seconds, tmax_offset, 

1134 nut_tmax_seconds, nut_tmax_offset, 

1135 model.tscale_to_kscale( 

1136 nut_tmax_seconds - tmax_seconds), 

1137 nut_id)) 

1138 

1139 delete.append((nut_id,)) 

1140 

1141 sql_add = ''' 

1142 INSERT INTO %(db)s.%(nuts)s ( 

1143 file_id, file_segment, file_element, kind_id, 

1144 kind_codes_id, tmin_seconds, tmin_offset, 

1145 tmax_seconds, tmax_offset, kscale ) 

1146 SELECT 

1147 file_id, file_segment, file_element, 

1148 kind_id, kind_codes_id, ?, ?, ?, ?, ? 

1149 FROM %(db)s.%(nuts)s 

1150 WHERE nut_id == ? 

1151 ''' 

1152 cursor.executemany(sql_subst(sql_add), insert) 

1153 

1154 sql_delete = ''' 

1155 DELETE FROM %(db)s.%(nuts)s WHERE nut_id == ? 

1156 ''' 

1157 cursor.executemany(sql_subst(sql_delete), delete) 

1158 

1159 def get_time_span(self, kinds=None): 

1160 ''' 

1161 Get time interval over all content in selection. 

1162 

1163 :param kinds: 

1164 If not ``None``, restrict query to given content kinds. 

1165 :type kind: 

1166 list of str 

1167 

1168 :complexity: 

1169 O(1), independent of the number of nuts. 

1170 

1171 :returns: 

1172 ``(tmin, tmax)``, combined time interval of queried content kinds. 

1173 ''' 

1174 

1175 sql_min = self._sql(''' 

1176 SELECT MIN(tmin_seconds), MIN(tmin_offset) 

1177 FROM %(db)s.%(nuts)s 

1178 WHERE kind_id == ? 

1179 AND tmin_seconds == ( 

1180 SELECT MIN(tmin_seconds) 

1181 FROM %(db)s.%(nuts)s 

1182 WHERE kind_id == ?) 

1183 ''') 

1184 

1185 sql_max = self._sql(''' 

1186 SELECT MAX(tmax_seconds), MAX(tmax_offset) 

1187 FROM %(db)s.%(nuts)s 

1188 WHERE kind_id == ? 

1189 AND tmax_seconds == ( 

1190 SELECT MAX(tmax_seconds) 

1191 FROM %(db)s.%(nuts)s 

1192 WHERE kind_id == ?) 

1193 ''') 

1194 

1195 gtmin = None 

1196 gtmax = None 

1197 

1198 if isinstance(kinds, str): 

1199 kinds = [kinds] 

1200 

1201 if kinds is None: 

1202 kind_ids = model.g_content_kind_ids 

1203 else: 

1204 kind_ids = model.to_kind_ids(kinds) 

1205 

1206 for kind_id in kind_ids: 

1207 for tmin_seconds, tmin_offset in self._conn.execute( 

1208 sql_min, (kind_id, kind_id)): 

1209 tmin = model.tjoin(tmin_seconds, tmin_offset) 

1210 if tmin is not None and (gtmin is None or tmin < gtmin): 

1211 gtmin = tmin 

1212 

1213 for (tmax_seconds, tmax_offset) in self._conn.execute( 

1214 sql_max, (kind_id, kind_id)): 

1215 tmax = model.tjoin(tmax_seconds, tmax_offset) 

1216 if tmax is not None and (gtmax is None or tmax > gtmax): 

1217 gtmax = tmax 

1218 

1219 return gtmin, gtmax 

1220 

1221 def has(self, kinds): 

1222 ''' 

1223 Check availability of given content kinds. 

1224 

1225 :param kinds: 

1226 Content kinds to query. 

1227 :type kind: 

1228 list of str 

1229 

1230 :returns: 

1231 ``True`` if any of the queried content kinds is available 

1232 in the selection. 

1233 ''' 

1234 self_tmin, self_tmax = self.get_time_span(kinds) 

1235 

1236 return None not in (self_tmin, self_tmax) 

1237 

1238 def get_deltat_span(self, kind): 

1239 ''' 

1240 Get min and max sampling interval of all content of given kind. 

1241 

1242 :param kind: 

1243 Content kind 

1244 :type kind: 

1245 str 

1246 

1247 :returns: ``(deltat_min, deltat_max)`` 

1248 ''' 

1249 

1250 deltats = [ 

1251 deltat for deltat in self.get_deltats(kind) 

1252 if deltat is not None] 

1253 

1254 if deltats: 

1255 return min(deltats), max(deltats) 

1256 else: 

1257 return None, None 

1258 

1259 def iter_kinds(self, codes=None): 

1260 ''' 

1261 Iterate over content types available in selection. 

1262 

1263 :param codes: 

1264 If given, get kinds only for selected codes identifier. 

1265 Only a single identifier may be given here and no pattern matching 

1266 is done, currently. 

1267 :type codes: 

1268 :py:class:`~pyrocko.squirrel.model.Codes` 

1269 

1270 :yields: 

1271 Available content kinds as :py:class:`str`. 

1272 

1273 :complexity: 

1274 O(1), independent of number of nuts. 

1275 ''' 

1276 

1277 return self._database._iter_kinds( 

1278 codes=codes, 

1279 kind_codes_count='%(db)s.%(kind_codes_count)s' % self._names) 

1280 

1281 def iter_deltats(self, kind=None): 

1282 ''' 

1283 Iterate over sampling intervals available in selection. 

1284 

1285 :param kind: 

1286 If given, get sampling intervals only for a given content type. 

1287 :type kind: 

1288 str 

1289 

1290 :yields: 

1291 :py:class:`float` values. 

1292 

1293 :complexity: 

1294 O(1), independent of number of nuts. 

1295 ''' 

1296 return self._database._iter_deltats( 

1297 kind=kind, 

1298 kind_codes_count='%(db)s.%(kind_codes_count)s' % self._names) 

1299 

1300 def iter_codes(self, kind=None): 

1301 ''' 

1302 Iterate over content identifier code sequences available in selection. 

1303 

1304 :param kind: 

1305 If given, get codes only for a given content type. 

1306 :type kind: 

1307 str 

1308 

1309 :yields: 

1310 :py:class:`tuple` of :py:class:`str` 

1311 

1312 :complexity: 

1313 O(1), independent of number of nuts. 

1314 ''' 

1315 return self._database._iter_codes( 

1316 kind=kind, 

1317 kind_codes_count='%(db)s.%(kind_codes_count)s' % self._names) 

1318 

1319 def _iter_codes_info(self, kind=None, codes=None): 

1320 ''' 

1321 Iterate over number of occurrences of any (kind, codes) combination. 

1322 

1323 :param kind: 

1324 If given, get counts only for selected content type. 

1325 :type kind: 

1326 str 

1327 

1328 :yields: 

1329 Tuples of the form ``(kind, codes, deltat, kind_codes_id, count)``. 

1330 

1331 :complexity: 

1332 O(1), independent of number of nuts. 

1333 ''' 

1334 return self._database._iter_codes_info( 

1335 kind=kind, 

1336 codes=codes, 

1337 kind_codes_count='%(db)s.%(kind_codes_count)s' % self._names) 

1338 

1339 def get_kinds(self, codes=None): 

1340 ''' 

1341 Get content types available in selection. 

1342 

1343 :param codes: 

1344 If given, get kinds only for selected codes identifier. 

1345 Only a single identifier may be given here and no pattern matching 

1346 is done, currently. 

1347 :type codes: 

1348 :py:class:`~pyrocko.squirrel.model.Codes` 

1349 

1350 :returns: 

1351 Sorted list of available content types. 

1352 :rtype: 

1353 py:class:`list` of :py:class:`str` 

1354 

1355 :complexity: 

1356 O(1), independent of number of nuts. 

1357 

1358 ''' 

1359 return sorted(list(self.iter_kinds(codes=codes))) 

1360 

1361 def get_deltats(self, kind=None): 

1362 ''' 

1363 Get sampling intervals available in selection. 

1364 

1365 :param kind: 

1366 If given, get sampling intervals only for selected content type. 

1367 :type kind: 

1368 str 

1369 

1370 :complexity: 

1371 O(1), independent of number of nuts. 

1372 

1373 :returns: Sorted list of available sampling intervals. 

1374 ''' 

1375 return sorted(list(self.iter_deltats(kind=kind))) 

1376 

1377 def get_codes(self, kind=None): 

1378 ''' 

1379 Get identifier code sequences available in selection. 

1380 

1381 :param kind: 

1382 If given, get codes only for selected content type. 

1383 :type kind: 

1384 str 

1385 

1386 :complexity: 

1387 O(1), independent of number of nuts. 

1388 

1389 :returns: Sorted list of available codes as tuples of strings. 

1390 ''' 

1391 return sorted(list(self.iter_codes(kind=kind))) 

1392 

1393 def get_counts(self, kind=None): 

1394 ''' 

1395 Get number of occurrences of any (kind, codes) combination. 

1396 

1397 :param kind: 

1398 If given, get codes only for selected content type. 

1399 :type kind: 

1400 str 

1401 

1402 :complexity: 

1403 O(1), independent of number of nuts. 

1404 

1405 :returns: ``dict`` with ``counts[kind][codes]`` or ``counts[codes]`` 

1406 if kind is not ``None`` 

1407 ''' 

1408 d = {} 

1409 for kind_id, codes, _, _, count in self._iter_codes_info(kind=kind): 

1410 if kind_id not in d: 

1411 v = d[kind_id] = {} 

1412 else: 

1413 v = d[kind_id] 

1414 

1415 if codes not in v: 

1416 v[codes] = 0 

1417 

1418 v[codes] += count 

1419 

1420 if kind is not None: 

1421 return d[to_kind_id(kind)] 

1422 else: 

1423 return dict((to_kind(kind_id), v) for (kind_id, v) in d.items()) 

1424 

1425 def glob_codes(self, kind, codes): 

1426 ''' 

1427 Find codes matching given patterns. 

1428 

1429 :param kind: 

1430 Content kind to be queried. 

1431 :type kind: 

1432 str 

1433 

1434 :param codes: 

1435 List of code patterns to query. 

1436 :type codes: 

1437 :py:class:`list` of :py:class:`~pyrocko.squirrel.model.Codes` 

1438 objects appropriate for the queried content type, or anything which 

1439 can be converted to such objects. 

1440 

1441 :returns: 

1442 List of matches of the form ``[kind_codes_id, codes, deltat]``. 

1443 ''' 

1444 

1445 kind_id = to_kind_id(kind) 

1446 args = [kind_id] 

1447 pats = codes_patterns_for_kind(kind_id, codes) 

1448 

1449 if pats: 

1450 codes_cond = 'AND ( %s ) ' % ' OR '.join( 

1451 ('kind_codes.codes GLOB ?',) * len(pats)) 

1452 

1453 args.extend(pat.safe_str for pat in pats) 

1454 else: 

1455 codes_cond = '' 

1456 

1457 sql = self._sql(''' 

1458 SELECT kind_codes_id, codes, deltat FROM kind_codes 

1459 WHERE 

1460 kind_id == ? ''' + codes_cond) 

1461 

1462 return list(map(list, self._conn.execute(sql, args))) 

1463 

1464 def update(self, constraint=None, **kwargs): 

1465 ''' 

1466 Update or partially update channel and event inventories. 

1467 

1468 :param constraint: 

1469 Selection of times or areas to be brought up to date. 

1470 :type constraint: 

1471 :py:class:`~pyrocko.squirrel.client.base.Constraint` 

1472 

1473 :param \\*\\*kwargs: 

1474 Shortcut for setting ``constraint=Constraint(**kwargs)``. 

1475 

1476 This function triggers all attached remote sources, to check for 

1477 updates in the meta-data. The sources will only submit queries when 

1478 their expiration date has passed, or if the selection spans into 

1479 previously unseen times or areas. 

1480 ''' 

1481 

1482 if constraint is None: 

1483 constraint = client.Constraint(**kwargs) 

1484 

1485 for source in self._sources: 

1486 source.update_channel_inventory(self, constraint) 

1487 source.update_event_inventory(self, constraint) 

1488 

1489 def update_waveform_promises(self, constraint=None, **kwargs): 

1490 ''' 

1491 Permit downloading of remote waveforms. 

1492 

1493 :param constraint: 

1494 Remote waveforms compatible with the given constraint are enabled 

1495 for download. 

1496 :type constraint: 

1497 :py:class:`~pyrocko.squirrel.client.base.Constraint` 

1498 

1499 :param \\*\\*kwargs: 

1500 Shortcut for setting ``constraint=Constraint(**kwargs)``. 

1501 

1502 Calling this method permits Squirrel to download waveforms from remote 

1503 sources when processing subsequent waveform requests. This works by 

1504 inserting so called waveform promises into the database. It will look 

1505 into the available channels for each remote source and create a promise 

1506 for each channel compatible with the given constraint. If the promise 

1507 then matches in a waveform request, Squirrel tries to download the 

1508 waveform. If the download is successful, the downloaded waveform is 

1509 added to the Squirrel and the promise is deleted. If the download 

1510 fails, the promise is kept if the reason of failure looks like being 

1511 temporary, e.g. because of a network failure. If the cause of failure 

1512 however seems to be permanent, the promise is deleted so that no 

1513 further attempts are made to download a waveform which might not be 

1514 available from that server at all. To force re-scheduling after a 

1515 permanent failure, call :py:meth:`update_waveform_promises` 

1516 yet another time. 

1517 ''' 

1518 

1519 if constraint is None: 

1520 constraint = client.Constraint(**kwargs) 

1521 

1522 for source in self._sources: 

1523 source.update_waveform_promises(self, constraint) 

1524 

1525 def remove_waveform_promises(self, from_database='selection'): 

1526 ''' 

1527 Remove waveform promises from live selection or global database. 

1528 

1529 Calling this function removes all waveform promises provided by the 

1530 attached sources. 

1531 

1532 :param from_database: 

1533 Remove from live selection ``'selection'`` or global database 

1534 ``'global'``. 

1535 ''' 

1536 for source in self._sources: 

1537 source.remove_waveform_promises(self, from_database=from_database) 

1538 

1539 def update_responses(self, constraint=None, **kwargs): 

1540 if constraint is None: 

1541 constraint = client.Constraint(**kwargs) 

1542 

1543 for source in self._sources: 

1544 source.update_response_inventory(self, constraint) 

1545 

1546 def get_nfiles(self): 

1547 ''' 

1548 Get number of files in selection. 

1549 ''' 

1550 

1551 sql = self._sql('''SELECT COUNT(*) FROM %(db)s.%(file_states)s''') 

1552 for row in self._conn.execute(sql): 

1553 return row[0] 

1554 

1555 def get_nnuts(self): 

1556 ''' 

1557 Get number of nuts in selection. 

1558 ''' 

1559 

1560 sql = self._sql('''SELECT COUNT(*) FROM %(db)s.%(nuts)s''') 

1561 for row in self._conn.execute(sql): 

1562 return row[0] 

1563 

1564 def get_total_size(self): 

1565 ''' 

1566 Get aggregated file size available in selection. 

1567 ''' 

1568 

1569 sql = self._sql(''' 

1570 SELECT SUM(files.size) FROM %(db)s.%(file_states)s 

1571 INNER JOIN files 

1572 ON %(db)s.%(file_states)s.file_id = files.file_id 

1573 ''') 

1574 

1575 for row in self._conn.execute(sql): 

1576 return row[0] or 0 

1577 

1578 def get_stats(self): 

1579 ''' 

1580 Get statistics on contents available through this selection. 

1581 ''' 

1582 

1583 kinds = self.get_kinds() 

1584 time_spans = {} 

1585 for kind in kinds: 

1586 time_spans[kind] = self.get_time_span([kind]) 

1587 

1588 return SquirrelStats( 

1589 nfiles=self.get_nfiles(), 

1590 nnuts=self.get_nnuts(), 

1591 kinds=kinds, 

1592 codes=self.get_codes(), 

1593 total_size=self.get_total_size(), 

1594 counts=self.get_counts(), 

1595 time_spans=time_spans, 

1596 sources=[s.describe() for s in self._sources], 

1597 operators=[op.describe() for op in self._operators]) 

1598 

1599 @filldocs 

1600 def check( 

1601 self, obj=None, tmin=None, tmax=None, time=None, codes=None, 

1602 ignore=[]): 

1603 ''' 

1604 Check for common data/metadata problems. 

1605 

1606 %(query_args)s 

1607 

1608 :param ignore: 

1609 Problem types to be ignored. 

1610 :type ignore: 

1611 :class:`list` of :class:`str` 

1612 (:py:class:`~pyrocko.squirrel.check.SquirrelCheckProblemType`) 

1613 

1614 :returns: 

1615 :py:class:`~pyrocko.squirrel.check.SquirrelCheck` object 

1616 containing the results of the check. 

1617 

1618 See :py:func:`~pyrocko.squirrel.check.do_check`. 

1619 ''' 

1620 

1621 from .check import do_check 

1622 tmin, tmax, codes = self._get_selection_args( 

1623 CHANNEL, obj, tmin, tmax, time, codes) 

1624 

1625 return do_check(self, tmin=tmin, tmax=tmax, codes=codes, ignore=ignore) 

1626 

1627 def get_content( 

1628 self, 

1629 nut, 

1630 cache_id='default', 

1631 accessor_id='default', 

1632 show_progress=False, 

1633 model='squirrel'): 

1634 

1635 ''' 

1636 Get and possibly load full content for a given index entry from file. 

1637 

1638 Loads the actual content objects (channel, station, waveform, ...) from 

1639 file. For efficiency, sibling content (all stuff in the same file 

1640 segment) will also be loaded as a side effect. The loaded contents are 

1641 cached in the Squirrel object. 

1642 ''' 

1643 

1644 content_cache = self._content_caches[cache_id] 

1645 if not content_cache.has(nut): 

1646 

1647 for nut_loaded in io.iload( 

1648 nut.file_path, 

1649 segment=nut.file_segment, 

1650 format=nut.file_format, 

1651 database=self._database, 

1652 update_selection=self, 

1653 show_progress=show_progress): 

1654 

1655 content_cache.put(nut_loaded) 

1656 

1657 try: 

1658 return content_cache.get(nut, accessor_id, model) 

1659 

1660 except KeyError: 

1661 raise error.NotAvailable( 

1662 'Unable to retrieve content: %s, %s, %s, %s' % nut.key) 

1663 

1664 def advance_accessor(self, accessor_id='default', cache_id=None): 

1665 ''' 

1666 Notify memory caches about consumer moving to a new data batch. 

1667 

1668 :param accessor_id: 

1669 Name of accessing consumer to be advanced. 

1670 :type accessor_id: 

1671 str 

1672 

1673 :param cache_id: 

1674 Name of cache to for which the accessor should be advanced. By 

1675 default the named accessor is advanced in all registered caches. 

1676 By default, two caches named ``'default'`` and ``'waveform'`` are 

1677 available. 

1678 :type cache_id: 

1679 str 

1680 

1681 See :py:class:`~pyrocko.squirrel.cache.ContentCache` for details on how 

1682 Squirrel's memory caching works and can be tuned. Default behaviour is 

1683 to release data when it has not been used in the latest data 

1684 window/batch. If the accessor is never advanced, data is cached 

1685 indefinitely - which is often desired e.g. for station meta-data. 

1686 Methods for consecutive data traversal, like 

1687 :py:meth:`chopper_waveforms` automatically advance and clear 

1688 their accessor. 

1689 ''' 

1690 for cache_ in ( 

1691 self._content_caches.keys() 

1692 if cache_id is None 

1693 else [cache_id]): 

1694 

1695 self._content_caches[cache_].advance_accessor(accessor_id) 

1696 

1697 def clear_accessor(self, accessor_id, cache_id=None): 

1698 ''' 

1699 Notify memory caches about a consumer having finished. 

1700 

1701 :param accessor_id: 

1702 Name of accessor to be cleared. 

1703 :type accessor_id: 

1704 str 

1705 

1706 :param cache_id: 

1707 Name of cache for which the accessor should be cleared. By default 

1708 the named accessor is cleared from all registered caches. By 

1709 default, two caches named ``'default'`` and ``'waveform'`` are 

1710 available. 

1711 :type cache_id: 

1712 str 

1713 

1714 Calling this method clears all references to cache entries held by the 

1715 named accessor. Cache entries are then freed if not referenced by any 

1716 other accessor. 

1717 ''' 

1718 

1719 for cache_ in ( 

1720 self._content_caches.keys() 

1721 if cache_id is None 

1722 else [cache_id]): 

1723 

1724 self._content_caches[cache_].clear_accessor(accessor_id) 

1725 

1726 def get_cache_stats(self, cache_id): 

1727 return self._content_caches[cache_id].get_stats() 

1728 

1729 @filldocs 

1730 def get_stations( 

1731 self, obj=None, tmin=None, tmax=None, time=None, codes=None, 

1732 model='squirrel'): 

1733 

1734 ''' 

1735 Get stations matching given constraints. 

1736 

1737 %(query_args)s 

1738 

1739 :param model: 

1740 Select object model for returned values: ``'squirrel'`` to get 

1741 Squirrel station objects or ``'pyrocko'`` to get Pyrocko station 

1742 objects with channel information attached. 

1743 :type model: 

1744 str 

1745 

1746 :returns: 

1747 List of :py:class:`pyrocko.squirrel.Station 

1748 <pyrocko.squirrel.model.Station>` objects by default or list of 

1749 :py:class:`pyrocko.model.Station <pyrocko.model.station.Station>` 

1750 objects if ``model='pyrocko'`` is requested. 

1751 

1752 See :py:meth:`iter_nuts` for details on time span matching. 

1753 ''' 

1754 

1755 if model == 'pyrocko': 

1756 return self._get_pyrocko_stations(obj, tmin, tmax, time, codes) 

1757 elif model in ('squirrel', 'stationxml', 'stationxml+'): 

1758 args = self._get_selection_args( 

1759 STATION, obj, tmin, tmax, time, codes) 

1760 

1761 nuts = sorted( 

1762 self.iter_nuts('station', *args), key=lambda nut: nut.dkey) 

1763 

1764 return [self.get_content(nut, model=model) for nut in nuts] 

1765 else: 

1766 raise ValueError('Invalid station model: %s' % model) 

1767 

1768 @filldocs 

1769 def get_channels( 

1770 self, obj=None, tmin=None, tmax=None, time=None, codes=None, 

1771 model='squirrel'): 

1772 

1773 ''' 

1774 Get channels matching given constraints. 

1775 

1776 %(query_args)s 

1777 

1778 :returns: 

1779 List of :py:class:`~pyrocko.squirrel.model.Channel` objects. 

1780 

1781 See :py:meth:`iter_nuts` for details on time span matching. 

1782 ''' 

1783 

1784 args = self._get_selection_args( 

1785 CHANNEL, obj, tmin, tmax, time, codes) 

1786 

1787 nuts = sorted( 

1788 self.iter_nuts('channel', *args), key=lambda nut: nut.dkey) 

1789 

1790 return [self.get_content(nut, model=model) for nut in nuts] 

1791 

1792 @filldocs 

1793 def get_sensors( 

1794 self, obj=None, tmin=None, tmax=None, time=None, codes=None): 

1795 

1796 ''' 

1797 Get sensors matching given constraints. 

1798 

1799 %(query_args)s 

1800 

1801 :returns: 

1802 List of :py:class:`~pyrocko.squirrel.model.Sensor` objects. 

1803 

1804 See :py:meth:`iter_nuts` for details on time span matching. 

1805 ''' 

1806 

1807 tmin, tmax, codes = self._get_selection_args( 

1808 CHANNEL, obj, tmin, tmax, time, codes) 

1809 

1810 if codes is not None: 

1811 codes = codes_patterns_list( 

1812 (entry.replace(channel=entry.channel[:-1] + '?') 

1813 if entry.channel != '*' else entry) 

1814 for entry in codes) 

1815 

1816 nuts = sorted( 

1817 self.iter_nuts( 

1818 'channel', tmin, tmax, codes), key=lambda nut: nut.dkey) 

1819 

1820 return [ 

1821 sensor for sensor in model.Sensor.from_channels( 

1822 self.get_content(nut) for nut in nuts) 

1823 if match_time_span(tmin, tmax, sensor)] 

1824 

1825 @filldocs 

1826 def get_responses( 

1827 self, obj=None, tmin=None, tmax=None, time=None, codes=None, 

1828 model='squirrel'): 

1829 

1830 ''' 

1831 Get instrument responses matching given constraints. 

1832 

1833 %(query_args)s 

1834 

1835 :returns: 

1836 List of :py:class:`~pyrocko.squirrel.model.Response` objects. 

1837 

1838 See :py:meth:`iter_nuts` for details on time span matching. 

1839 ''' 

1840 

1841 args = self._get_selection_args( 

1842 RESPONSE, obj, tmin, tmax, time, codes) 

1843 

1844 nuts = sorted( 

1845 self.iter_nuts('response', *args), key=lambda nut: nut.dkey) 

1846 

1847 return [self.get_content(nut, model=model) for nut in nuts] 

1848 

1849 @filldocs 

1850 def get_response( 

1851 self, obj=None, tmin=None, tmax=None, time=None, codes=None, 

1852 model='squirrel'): 

1853 

1854 ''' 

1855 Get instrument response matching given constraints. 

1856 

1857 %(query_args)s 

1858 

1859 :returns: 

1860 :py:class:`~pyrocko.squirrel.model.Response` object. 

1861 

1862 Same as :py:meth:`get_responses` but returning exactly one response. 

1863 Raises :py:exc:`~pyrocko.squirrel.error.NotAvailable` if zero or more 

1864 than one is available. 

1865 

1866 See :py:meth:`iter_nuts` for details on time span matching. 

1867 ''' 

1868 

1869 if model == 'stationxml': 

1870 model_ = 'stationxml+' 

1871 else: 

1872 model_ = model 

1873 

1874 responses = self.get_responses( 

1875 obj, tmin, tmax, time, codes, model=model_) 

1876 if len(responses) == 0: 

1877 raise error.NotAvailable( 

1878 'No instrument response available (%s).' 

1879 % self._get_selection_args_str( 

1880 RESPONSE, obj, tmin, tmax, time, codes)) 

1881 

1882 elif len(responses) > 1: 

1883 if model_ == 'squirrel': 

1884 resps_sq = responses 

1885 elif model_ == 'stationxml+': 

1886 resps_sq = [resp[0] for resp in responses] 

1887 else: 

1888 raise ValueError('Invalid response model: %s' % model) 

1889 

1890 rinfo = ':\n' + '\n'.join( 

1891 ' ' + resp.summary for resp in resps_sq) 

1892 

1893 raise error.NotAvailable( 

1894 'Multiple instrument responses matching given constraints ' 

1895 '(%s)%s' % ( 

1896 self._get_selection_args_str( 

1897 RESPONSE, obj, tmin, tmax, time, codes), rinfo)) 

1898 

1899 if model == 'stationxml': 

1900 return responses[0][1] 

1901 else: 

1902 return responses[0] 

1903 

1904 @filldocs 

1905 def get_events( 

1906 self, obj=None, tmin=None, tmax=None, time=None, codes=None): 

1907 

1908 ''' 

1909 Get events matching given constraints. 

1910 

1911 %(query_args)s 

1912 

1913 :returns: 

1914 List of :py:class:`~pyrocko.model.event.Event` objects. 

1915 

1916 See :py:meth:`iter_nuts` for details on time span matching. 

1917 ''' 

1918 

1919 args = self._get_selection_args(EVENT, obj, tmin, tmax, time, codes) 

1920 nuts = sorted( 

1921 self.iter_nuts('event', *args), key=lambda nut: nut.dkey) 

1922 

1923 return [self.get_content(nut) for nut in nuts] 

1924 

1925 def _redeem_promises(self, *args, order_only=False): 

1926 

1927 def split_promise(order): 

1928 self._split_nuts( 

1929 'waveform_promise', 

1930 order.tmin, order.tmax, 

1931 codes=order.codes, 

1932 path=order.source_id) 

1933 

1934 tmin, tmax, _ = args 

1935 

1936 waveforms = list(self.iter_nuts('waveform', *args)) 

1937 promises = list(self.iter_nuts('waveform_promise', *args)) 

1938 

1939 codes_to_avail = defaultdict(list) 

1940 for nut in waveforms: 

1941 codes_to_avail[nut.codes].append((nut.tmin, nut.tmax)) 

1942 

1943 def tts(x): 

1944 if isinstance(x, tuple): 

1945 return tuple(tts(e) for e in x) 

1946 elif isinstance(x, list): 

1947 return list(tts(e) for e in x) 

1948 else: 

1949 return util.time_to_str(x) 

1950 

1951 orders = [] 

1952 for promise in promises: 

1953 waveforms_avail = codes_to_avail[promise.codes] 

1954 for block_tmin, block_tmax in blocks( 

1955 max(tmin, promise.tmin), 

1956 min(tmax, promise.tmax), 

1957 promise.deltat): 

1958 

1959 orders.append( 

1960 WaveformOrder( 

1961 source_id=promise.file_path, 

1962 codes=promise.codes, 

1963 tmin=block_tmin, 

1964 tmax=block_tmax, 

1965 deltat=promise.deltat, 

1966 gaps=gaps(waveforms_avail, block_tmin, block_tmax))) 

1967 

1968 orders_noop, orders = lpick(lambda order: order.gaps, orders) 

1969 

1970 order_keys_noop = set(order_key(order) for order in orders_noop) 

1971 if len(order_keys_noop) != 0 or len(orders_noop) != 0: 

1972 logger.info( 

1973 'Waveform orders already satisified with cached/local data: ' 

1974 '%i (%i)' % (len(order_keys_noop), len(orders_noop))) 

1975 

1976 for order in orders_noop: 

1977 split_promise(order) 

1978 

1979 if order_only: 

1980 if orders: 

1981 self._pending_orders.extend(orders) 

1982 logger.info( 

1983 'Enqueuing %i waveform order%s.' 

1984 % len_plural(orders)) 

1985 return 

1986 else: 

1987 if self._pending_orders: 

1988 orders.extend(self._pending_orders) 

1989 logger.info( 

1990 'Adding %i previously enqueued order%s.' 

1991 % len_plural(self._pending_orders)) 

1992 

1993 self._pending_orders = [] 

1994 

1995 source_ids = [] 

1996 sources = {} 

1997 for source in self._sources: 

1998 if isinstance(source, fdsn.FDSNSource): 

1999 source_ids.append(source._source_id) 

2000 sources[source._source_id] = source 

2001 

2002 source_priority = dict( 

2003 (source_id, i) for (i, source_id) in enumerate(source_ids)) 

2004 

2005 order_groups = defaultdict(list) 

2006 for order in orders: 

2007 order_groups[order_key(order)].append(order) 

2008 

2009 for k, order_group in order_groups.items(): 

2010 order_group.sort( 

2011 key=lambda order: source_priority[order.source_id]) 

2012 

2013 n_order_groups = len(order_groups) 

2014 

2015 if len(order_groups) != 0 or len(orders) != 0: 

2016 logger.info( 

2017 'Waveform orders standing for download: %i (%i)' 

2018 % (len(order_groups), len(orders))) 

2019 

2020 task = make_task('Waveform orders processed', n_order_groups) 

2021 else: 

2022 task = None 

2023 

2024 def release_order_group(order): 

2025 okey = order_key(order) 

2026 for followup in order_groups[okey]: 

2027 split_promise(followup) 

2028 

2029 del order_groups[okey] 

2030 

2031 if task: 

2032 task.update(n_order_groups - len(order_groups)) 

2033 

2034 def noop(order): 

2035 pass 

2036 

2037 def success(order): 

2038 release_order_group(order) 

2039 split_promise(order) 

2040 

2041 def batch_add(paths): 

2042 self.add(paths) 

2043 

2044 calls = queue.Queue() 

2045 

2046 def enqueue(f): 

2047 def wrapper(*args): 

2048 calls.put((f, args)) 

2049 

2050 return wrapper 

2051 

2052 while order_groups: 

2053 

2054 orders_now = [] 

2055 empty = [] 

2056 for k, order_group in order_groups.items(): 

2057 try: 

2058 orders_now.append(order_group.pop(0)) 

2059 except IndexError: 

2060 empty.append(k) 

2061 

2062 for k in empty: 

2063 del order_groups[k] 

2064 

2065 by_source_id = defaultdict(list) 

2066 for order in orders_now: 

2067 by_source_id[order.source_id].append(order) 

2068 

2069 threads = [] 

2070 for source_id in by_source_id: 

2071 def download(): 

2072 try: 

2073 sources[source_id].download_waveforms( 

2074 by_source_id[source_id], 

2075 success=enqueue(success), 

2076 error_permanent=enqueue(split_promise), 

2077 error_temporary=noop, 

2078 batch_add=enqueue(batch_add)) 

2079 

2080 finally: 

2081 calls.put(None) 

2082 

2083 thread = threading.Thread(target=download) 

2084 thread.start() 

2085 threads.append(thread) 

2086 

2087 ndone = 0 

2088 while ndone < len(threads): 

2089 ret = calls.get() 

2090 if ret is None: 

2091 ndone += 1 

2092 else: 

2093 ret[0](*ret[1]) 

2094 

2095 for thread in threads: 

2096 thread.join() 

2097 

2098 if task: 

2099 task.update(n_order_groups - len(order_groups)) 

2100 

2101 if task: 

2102 task.done() 

2103 

2104 @filldocs 

2105 def get_waveform_nuts( 

2106 self, obj=None, tmin=None, tmax=None, time=None, codes=None, 

2107 order_only=False): 

2108 

2109 ''' 

2110 Get waveform content entities matching given constraints. 

2111 

2112 %(query_args)s 

2113 

2114 Like :py:meth:`get_nuts` with ``kind='waveform'`` but additionally 

2115 resolves matching waveform promises (downloads waveforms from remote 

2116 sources). 

2117 

2118 See :py:meth:`iter_nuts` for details on time span matching. 

2119 ''' 

2120 

2121 args = self._get_selection_args(WAVEFORM, obj, tmin, tmax, time, codes) 

2122 self._redeem_promises(*args, order_only=order_only) 

2123 return sorted( 

2124 self.iter_nuts('waveform', *args), key=lambda nut: nut.dkey) 

2125 

2126 @filldocs 

2127 def have_waveforms( 

2128 self, obj=None, tmin=None, tmax=None, time=None, codes=None): 

2129 

2130 ''' 

2131 Check if any waveforms or waveform promises are available for given 

2132 constraints. 

2133 

2134 %(query_args)s 

2135 ''' 

2136 

2137 args = self._get_selection_args(WAVEFORM, obj, tmin, tmax, time, codes) 

2138 return bool(list( 

2139 self.iter_nuts('waveform', *args, limit=1))) \ 

2140 or bool(list( 

2141 self.iter_nuts('waveform_promise', *args, limit=1))) 

2142 

2143 @filldocs 

2144 def get_waveforms( 

2145 self, obj=None, tmin=None, tmax=None, time=None, codes=None, 

2146 uncut=False, want_incomplete=True, degap=True, maxgap=5, 

2147 maxlap=None, snap=None, include_last=False, load_data=True, 

2148 accessor_id='default', operator_params=None, order_only=False): 

2149 

2150 ''' 

2151 Get waveforms matching given constraints. 

2152 

2153 %(query_args)s 

2154 

2155 :param uncut: 

2156 Set to ``True``, to disable cutting traces to [``tmin``, ``tmax``] 

2157 and to disable degapping/deoverlapping. Returns untouched traces as 

2158 they are read from file segment. File segments are always read in 

2159 their entirety. 

2160 :type uncut: 

2161 bool 

2162 

2163 :param want_incomplete: 

2164 If ``True``, gappy/incomplete traces are included in the result. 

2165 :type want_incomplete: 

2166 bool 

2167 

2168 :param degap: 

2169 If ``True``, connect traces and remove gaps and overlaps. 

2170 :type degap: 

2171 bool 

2172 

2173 :param maxgap: 

2174 Maximum gap size in samples which is filled with interpolated 

2175 samples when ``degap`` is ``True``. 

2176 :type maxgap: 

2177 int 

2178 

2179 :param maxlap: 

2180 Maximum overlap size in samples which is removed when ``degap`` is 

2181 ``True``. 

2182 :type maxlap: 

2183 int 

2184 

2185 :param snap: 

2186 Rounding functions used when computing sample index from time 

2187 instance, for trace start and trace end, respectively. By default, 

2188 ``(round, round)`` is used. 

2189 :type snap: 

2190 tuple of 2 callables 

2191 

2192 :param include_last: 

2193 If ``True``, add one more sample to the returned traces (the sample 

2194 which would be the first sample of a query with ``tmin`` set to the 

2195 current value of ``tmax``). 

2196 :type include_last: 

2197 bool 

2198 

2199 :param load_data: 

2200 If ``True``, waveform data samples are read from files (or cache). 

2201 If ``False``, meta-information-only traces are returned (dummy 

2202 traces with no data samples). 

2203 :type load_data: 

2204 bool 

2205 

2206 :param accessor_id: 

2207 Name of consumer on who's behalf data is accessed. Used in cache 

2208 management (see :py:mod:`~pyrocko.squirrel.cache`). Used as a key 

2209 to distinguish different points of extraction for the decision of 

2210 when to release cached waveform data. Should be used when data is 

2211 alternately extracted from more than one region / selection. 

2212 :type accessor_id: 

2213 str 

2214 

2215 See :py:meth:`iter_nuts` for details on time span matching. 

2216 

2217 Loaded data is kept in memory (at least) until 

2218 :py:meth:`clear_accessor` has been called or 

2219 :py:meth:`advance_accessor` has been called two consecutive times 

2220 without data being accessed between the two calls (by this accessor). 

2221 Data may still be further kept in the memory cache if held alive by 

2222 consumers with a different ``accessor_id``. 

2223 ''' 

2224 

2225 tmin, tmax, codes = self._get_selection_args( 

2226 WAVEFORM, obj, tmin, tmax, time, codes) 

2227 

2228 self_tmin, self_tmax = self.get_time_span( 

2229 ['waveform', 'waveform_promise']) 

2230 

2231 if None in (self_tmin, self_tmax): 

2232 logger.warning( 

2233 'No waveforms available.') 

2234 return [] 

2235 

2236 tmin = tmin if tmin is not None else self_tmin 

2237 tmax = tmax if tmax is not None else self_tmax 

2238 

2239 if codes is not None and len(codes) == 1: 

2240 # TODO: fix for multiple / mixed codes 

2241 operator = self.get_operator(codes[0]) 

2242 if operator is not None: 

2243 return operator.get_waveforms( 

2244 self, codes[0], 

2245 tmin=tmin, tmax=tmax, 

2246 uncut=uncut, want_incomplete=want_incomplete, degap=degap, 

2247 maxgap=maxgap, maxlap=maxlap, snap=snap, 

2248 include_last=include_last, load_data=load_data, 

2249 accessor_id=accessor_id, params=operator_params) 

2250 

2251 nuts = self.get_waveform_nuts( 

2252 obj, tmin, tmax, time, codes, order_only=order_only) 

2253 

2254 if order_only: 

2255 return [] 

2256 

2257 if load_data: 

2258 traces = [ 

2259 self.get_content(nut, 'waveform', accessor_id) for nut in nuts] 

2260 

2261 else: 

2262 traces = [ 

2263 trace.Trace(**nut.trace_kwargs) for nut in nuts] 

2264 

2265 if uncut: 

2266 return traces 

2267 

2268 if snap is None: 

2269 snap = (round, round) 

2270 

2271 chopped = [] 

2272 for tr in traces: 

2273 if not load_data and tr.ydata is not None: 

2274 tr = tr.copy(data=False) 

2275 tr.ydata = None 

2276 

2277 try: 

2278 chopped.append(tr.chop( 

2279 tmin, tmax, 

2280 inplace=False, 

2281 snap=snap, 

2282 include_last=include_last)) 

2283 

2284 except trace.NoData: 

2285 pass 

2286 

2287 processed = self._process_chopped( 

2288 chopped, degap, maxgap, maxlap, want_incomplete, tmin, tmax) 

2289 

2290 return processed 

2291 

2292 @filldocs 

2293 def chopper_waveforms( 

2294 self, obj=None, tmin=None, tmax=None, time=None, codes=None, 

2295 tinc=None, tpad=0., 

2296 want_incomplete=True, snap_window=False, 

2297 degap=True, maxgap=5, maxlap=None, 

2298 snap=None, include_last=False, load_data=True, 

2299 accessor_id=None, clear_accessor=True, operator_params=None, 

2300 grouping=None): 

2301 

2302 ''' 

2303 Iterate window-wise over waveform archive. 

2304 

2305 %(query_args)s 

2306 

2307 :param tinc: 

2308 Time increment (window shift time) (default uses ``tmax-tmin``). 

2309 :type tinc: 

2310 timestamp 

2311 

2312 :param tpad: 

2313 Padding time appended on either side of the data window (window 

2314 overlap is ``2*tpad``). 

2315 :type tpad: 

2316 timestamp 

2317 

2318 :param want_incomplete: 

2319 If ``True``, gappy/incomplete traces are included in the result. 

2320 :type want_incomplete: 

2321 bool 

2322 

2323 :param snap_window: 

2324 If ``True``, start time windows at multiples of tinc with respect 

2325 to system time zero. 

2326 :type snap_window: 

2327 bool 

2328 

2329 :param degap: 

2330 If ``True``, connect traces and remove gaps and overlaps. 

2331 :type degap: 

2332 bool 

2333 

2334 :param maxgap: 

2335 Maximum gap size in samples which is filled with interpolated 

2336 samples when ``degap`` is ``True``. 

2337 :type maxgap: 

2338 int 

2339 

2340 :param maxlap: 

2341 Maximum overlap size in samples which is removed when ``degap`` is 

2342 ``True``. 

2343 :type maxlap: 

2344 int 

2345 

2346 :param snap: 

2347 Rounding functions used when computing sample index from time 

2348 instance, for trace start and trace end, respectively. By default, 

2349 ``(round, round)`` is used. 

2350 :type snap: 

2351 tuple of 2 callables 

2352 

2353 :param include_last: 

2354 If ``True``, add one more sample to the returned traces (the sample 

2355 which would be the first sample of a query with ``tmin`` set to the 

2356 current value of ``tmax``). 

2357 :type include_last: 

2358 bool 

2359 

2360 :param load_data: 

2361 If ``True``, waveform data samples are read from files (or cache). 

2362 If ``False``, meta-information-only traces are returned (dummy 

2363 traces with no data samples). 

2364 :type load_data: 

2365 bool 

2366 

2367 :param accessor_id: 

2368 Name of consumer on who's behalf data is accessed. Used in cache 

2369 management (see :py:mod:`~pyrocko.squirrel.cache`). Used as a key 

2370 to distinguish different points of extraction for the decision of 

2371 when to release cached waveform data. Should be used when data is 

2372 alternately extracted from more than one region / selection. 

2373 :type accessor_id: 

2374 str 

2375 

2376 :param clear_accessor: 

2377 If ``True`` (default), :py:meth:`clear_accessor` is called when the 

2378 chopper finishes. Set to ``False`` to keep loaded waveforms in 

2379 memory when the generator returns. 

2380 :type clear_accessor: 

2381 bool 

2382 

2383 :param grouping: 

2384 By default, traversal over the data is over time and all matching 

2385 traces of a time window are yielded. Using this option, it is 

2386 possible to traverse the data first by group (e.g. station or 

2387 network) and second by time. This can reduce the number of traces 

2388 in each batch and thus reduce the memory footprint of the process. 

2389 :type grouping: 

2390 :py:class:`~pyrocko.squirrel.operator.Grouping` 

2391 

2392 :yields: 

2393 A list of :py:class:`~pyrocko.trace.Trace` objects for every 

2394 extracted time window. 

2395 

2396 See :py:meth:`iter_nuts` for details on time span matching. 

2397 ''' 

2398 

2399 tmin, tmax, codes = self._get_selection_args( 

2400 WAVEFORM, obj, tmin, tmax, time, codes) 

2401 

2402 self_tmin, self_tmax = self.get_time_span( 

2403 ['waveform', 'waveform_promise']) 

2404 

2405 if None in (self_tmin, self_tmax): 

2406 logger.warning( 

2407 'Content has undefined time span. No waveforms and no ' 

2408 'waveform promises?') 

2409 return 

2410 

2411 if snap_window and tinc is not None: 

2412 tmin = tmin if tmin is not None else self_tmin 

2413 tmax = tmax if tmax is not None else self_tmax 

2414 tmin = math.floor(tmin / tinc) * tinc 

2415 tmax = math.ceil(tmax / tinc) * tinc 

2416 else: 

2417 tmin = tmin if tmin is not None else self_tmin + tpad 

2418 tmax = tmax if tmax is not None else self_tmax - tpad 

2419 

2420 tinc = tinc if tinc is not None else tmax - tmin 

2421 

2422 try: 

2423 if accessor_id is None: 

2424 accessor_id = 'chopper%i' % self._n_choppers_active 

2425 

2426 self._n_choppers_active += 1 

2427 

2428 eps = tinc * 1e-6 

2429 if tinc != 0.0: 

2430 nwin = int(((tmax - eps) - tmin) / tinc) + 1 

2431 else: 

2432 nwin = 1 

2433 

2434 if grouping is None: 

2435 codes_list = [codes] 

2436 else: 

2437 operator = Operator( 

2438 filtering=CodesPatternFiltering(codes=codes), 

2439 grouping=grouping) 

2440 

2441 available = set(self.get_codes(kind='waveform')) 

2442 available.update(self.get_codes(kind='waveform_promise')) 

2443 operator.update_mappings(sorted(available)) 

2444 

2445 codes_list = [ 

2446 codes_patterns_list(scl) 

2447 for scl in operator.iter_in_codes()] 

2448 

2449 ngroups = len(codes_list) 

2450 for igroup, scl in enumerate(codes_list): 

2451 for iwin in range(nwin): 

2452 wmin, wmax = tmin+iwin*tinc, min(tmin+(iwin+1)*tinc, tmax) 

2453 

2454 chopped = self.get_waveforms( 

2455 tmin=wmin-tpad, 

2456 tmax=wmax+tpad, 

2457 codes=scl, 

2458 snap=snap, 

2459 include_last=include_last, 

2460 load_data=load_data, 

2461 want_incomplete=want_incomplete, 

2462 degap=degap, 

2463 maxgap=maxgap, 

2464 maxlap=maxlap, 

2465 accessor_id=accessor_id, 

2466 operator_params=operator_params) 

2467 

2468 self.advance_accessor(accessor_id) 

2469 

2470 yield Batch( 

2471 tmin=wmin, 

2472 tmax=wmax, 

2473 i=iwin, 

2474 n=nwin, 

2475 igroup=igroup, 

2476 ngroups=ngroups, 

2477 traces=chopped) 

2478 

2479 finally: 

2480 self._n_choppers_active -= 1 

2481 if clear_accessor: 

2482 self.clear_accessor(accessor_id, 'waveform') 

2483 

2484 def _process_chopped( 

2485 self, chopped, degap, maxgap, maxlap, want_incomplete, tmin, tmax): 

2486 

2487 chopped.sort(key=lambda a: a.full_id) 

2488 if degap: 

2489 chopped = trace.degapper(chopped, maxgap=maxgap, maxlap=maxlap) 

2490 

2491 if not want_incomplete: 

2492 chopped_weeded = [] 

2493 for tr in chopped: 

2494 emin = tr.tmin - tmin 

2495 emax = tr.tmax + tr.deltat - tmax 

2496 if (abs(emin) <= 0.5*tr.deltat and abs(emax) <= 0.5*tr.deltat): 

2497 chopped_weeded.append(tr) 

2498 

2499 elif degap: 

2500 if (0. < emin <= 5. * tr.deltat 

2501 and -5. * tr.deltat <= emax < 0.): 

2502 

2503 tr.extend(tmin, tmax-tr.deltat, fillmethod='repeat') 

2504 chopped_weeded.append(tr) 

2505 

2506 chopped = chopped_weeded 

2507 

2508 return chopped 

2509 

2510 def _get_pyrocko_stations( 

2511 self, obj=None, tmin=None, tmax=None, time=None, codes=None): 

2512 

2513 from pyrocko import model as pmodel 

2514 

2515 if codes is not None: 

2516 codes = codes_patterns_for_kind(STATION, codes) 

2517 

2518 by_nsl = defaultdict(lambda: (list(), list())) 

2519 for station in self.get_stations(obj, tmin, tmax, time, codes): 

2520 sargs = station._get_pyrocko_station_args() 

2521 by_nsl[station.codes.nsl][0].append(sargs) 

2522 

2523 if codes is not None: 

2524 codes = [model.CodesNSLCE(c) for c in codes] 

2525 

2526 for channel in self.get_channels(obj, tmin, tmax, time, codes): 

2527 sargs = channel._get_pyrocko_station_args() 

2528 sargs_list, channels_list = by_nsl[channel.codes.nsl] 

2529 sargs_list.append(sargs) 

2530 channels_list.append(channel) 

2531 

2532 pstations = [] 

2533 nsls = list(by_nsl.keys()) 

2534 nsls.sort() 

2535 for nsl in nsls: 

2536 sargs_list, channels_list = by_nsl[nsl] 

2537 sargs = util.consistency_merge( 

2538 [('',) + x for x in sargs_list]) 

2539 

2540 by_c = defaultdict(list) 

2541 for ch in channels_list: 

2542 by_c[ch.codes.channel].append(ch._get_pyrocko_channel_args()) 

2543 

2544 chas = list(by_c.keys()) 

2545 chas.sort() 

2546 pchannels = [] 

2547 for cha in chas: 

2548 list_of_cargs = by_c[cha] 

2549 cargs = util.consistency_merge( 

2550 [('',) + x for x in list_of_cargs]) 

2551 pchannels.append(pmodel.Channel(*cargs)) 

2552 

2553 pstations.append( 

2554 pmodel.Station(*sargs, channels=pchannels)) 

2555 

2556 return pstations 

2557 

2558 @property 

2559 def pile(self): 

2560 

2561 ''' 

2562 Emulates the older :py:class:`pyrocko.pile.Pile` interface. 

2563 

2564 This property exposes a :py:class:`pyrocko.squirrel.pile.Pile` object, 

2565 which emulates most of the older :py:class:`pyrocko.pile.Pile` methods 

2566 but uses the fluffy power of the Squirrel under the hood. 

2567 

2568 This interface can be used as a drop-in replacement for piles which are 

2569 used in existing scripts and programs for efficient waveform data 

2570 access. The Squirrel-based pile scales better for large datasets. Newer 

2571 scripts should use Squirrel's native methods to avoid the emulation 

2572 overhead. 

2573 ''' 

2574 from . import pile 

2575 

2576 if self._pile is None: 

2577 self._pile = pile.Pile(self) 

2578 

2579 return self._pile 

2580 

2581 def snuffle(self): 

2582 ''' 

2583 Look at dataset in Snuffler. 

2584 ''' 

2585 self.pile.snuffle() 

2586 

2587 def _gather_codes_keys(self, kind, gather, selector): 

2588 return set( 

2589 gather(codes) 

2590 for codes in self.iter_codes(kind) 

2591 if selector is None or selector(codes)) 

2592 

2593 def __str__(self): 

2594 return str(self.get_stats()) 

2595 

2596 def get_coverage( 

2597 self, kind, tmin=None, tmax=None, codes=None, limit=None): 

2598 

2599 ''' 

2600 Get coverage information. 

2601 

2602 Get information about strips of gapless data coverage. 

2603 

2604 :param kind: 

2605 Content kind to be queried. 

2606 :type kind: 

2607 str 

2608 

2609 :param tmin: 

2610 Start time of query interval. 

2611 :type tmin: 

2612 timestamp 

2613 

2614 :param tmax: 

2615 End time of query interval. 

2616 :type tmax: 

2617 timestamp 

2618 

2619 :param codes: 

2620 If given, restrict query to given content codes patterns. 

2621 :type codes: 

2622 :py:class:`list` of :py:class:`~pyrocko.squirrel.model.Codes` 

2623 objects appropriate for the queried content type, or anything which 

2624 can be converted to such objects. 

2625 

2626 :param limit: 

2627 Limit query to return only up to a given maximum number of entries 

2628 per matching time series (without setting this option, very gappy 

2629 data could cause the query to execute for a very long time). 

2630 :type limit: 

2631 int 

2632 

2633 :returns: 

2634 Information about time spans covered by the requested time series 

2635 data. 

2636 :rtype: 

2637 :py:class:`list` of :py:class:`Coverage` objects 

2638 ''' 

2639 

2640 tmin_seconds, tmin_offset = model.tsplit(tmin) 

2641 tmax_seconds, tmax_offset = model.tsplit(tmax) 

2642 kind_id = to_kind_id(kind) 

2643 

2644 codes_info = list(self._iter_codes_info(kind=kind)) 

2645 

2646 kdata_all = [] 

2647 if codes is None: 

2648 for _, codes_entry, deltat, kind_codes_id, _ in codes_info: 

2649 kdata_all.append( 

2650 (codes_entry, kind_codes_id, codes_entry, deltat)) 

2651 

2652 else: 

2653 for codes_entry in codes: 

2654 pattern = to_codes(kind_id, codes_entry) 

2655 for _, codes_entry, deltat, kind_codes_id, _ in codes_info: 

2656 if model.match_codes(pattern, codes_entry): 

2657 kdata_all.append( 

2658 (pattern, kind_codes_id, codes_entry, deltat)) 

2659 

2660 kind_codes_ids = [x[1] for x in kdata_all] 

2661 

2662 counts_at_tmin = {} 

2663 if tmin is not None: 

2664 for nut in self.iter_nuts( 

2665 kind, tmin, tmin, kind_codes_ids=kind_codes_ids): 

2666 

2667 k = nut.codes, nut.deltat 

2668 if k not in counts_at_tmin: 

2669 counts_at_tmin[k] = 0 

2670 

2671 counts_at_tmin[k] += 1 

2672 

2673 coverages = [] 

2674 for pattern, kind_codes_id, codes_entry, deltat in kdata_all: 

2675 entry = [pattern, codes_entry, deltat, None, None, []] 

2676 for i, order in [(0, 'ASC'), (1, 'DESC')]: 

2677 sql = self._sql(''' 

2678 SELECT 

2679 time_seconds, 

2680 time_offset 

2681 FROM %(db)s.%(coverage)s 

2682 WHERE 

2683 kind_codes_id == ? 

2684 ORDER BY 

2685 kind_codes_id ''' + order + ''', 

2686 time_seconds ''' + order + ''', 

2687 time_offset ''' + order + ''' 

2688 LIMIT 1 

2689 ''') 

2690 

2691 for row in self._conn.execute(sql, [kind_codes_id]): 

2692 entry[3+i] = model.tjoin(row[0], row[1]) 

2693 

2694 if None in entry[3:5]: 

2695 continue 

2696 

2697 args = [kind_codes_id] 

2698 

2699 sql_time = '' 

2700 if tmin is not None: 

2701 # intentionally < because (== tmin) is queried from nuts 

2702 sql_time += ' AND ( ? < time_seconds ' \ 

2703 'OR ( ? == time_seconds AND ? < time_offset ) ) ' 

2704 args.extend([tmin_seconds, tmin_seconds, tmin_offset]) 

2705 

2706 if tmax is not None: 

2707 sql_time += ' AND ( time_seconds < ? ' \ 

2708 'OR ( ? == time_seconds AND time_offset <= ? ) ) ' 

2709 args.extend([tmax_seconds, tmax_seconds, tmax_offset]) 

2710 

2711 sql_limit = '' 

2712 if limit is not None: 

2713 sql_limit = ' LIMIT ?' 

2714 args.append(limit) 

2715 

2716 sql = self._sql(''' 

2717 SELECT 

2718 time_seconds, 

2719 time_offset, 

2720 step 

2721 FROM %(db)s.%(coverage)s 

2722 WHERE 

2723 kind_codes_id == ? 

2724 ''' + sql_time + ''' 

2725 ORDER BY 

2726 kind_codes_id, 

2727 time_seconds, 

2728 time_offset 

2729 ''' + sql_limit) 

2730 

2731 rows = list(self._conn.execute(sql, args)) 

2732 

2733 if limit is not None and len(rows) == limit: 

2734 entry[-1] = None 

2735 else: 

2736 counts = counts_at_tmin.get((codes_entry, deltat), 0) 

2737 tlast = None 

2738 if tmin is not None: 

2739 entry[-1].append((tmin, counts)) 

2740 tlast = tmin 

2741 

2742 for row in rows: 

2743 t = model.tjoin(row[0], row[1]) 

2744 counts += row[2] 

2745 entry[-1].append((t, counts)) 

2746 tlast = t 

2747 

2748 if tmax is not None and (tlast is None or tlast != tmax): 

2749 entry[-1].append((tmax, counts)) 

2750 

2751 coverages.append(model.Coverage.from_values(entry + [kind_id])) 

2752 

2753 return coverages 

2754 

2755 def get_stationxml( 

2756 self, obj=None, tmin=None, tmax=None, time=None, codes=None, 

2757 level='response'): 

2758 

2759 ''' 

2760 Get station/channel/response metadata in StationXML representation. 

2761 

2762 %(query_args)s 

2763 

2764 :returns: 

2765 :py:class:`~pyrocko.io.stationxml.FDSNStationXML` object. 

2766 ''' 

2767 

2768 if level not in ('network', 'station', 'channel', 'response'): 

2769 raise ValueError('Invalid level: %s' % level) 

2770 

2771 tmin, tmax, codes = self._get_selection_args( 

2772 CHANNEL, obj, tmin, tmax, time, codes) 

2773 

2774 filtering = CodesPatternFiltering(codes=codes) 

2775 

2776 nslcs = list(set( 

2777 codes.nslc for codes in 

2778 filtering.filter(self.get_codes(kind='channel')))) 

2779 

2780 from pyrocko.io import stationxml as sx 

2781 

2782 networks = [] 

2783 for net, stas in prefix_tree(nslcs): 

2784 network = sx.Network(code=net) 

2785 networks.append(network) 

2786 

2787 if level not in ('station', 'channel', 'response'): 

2788 continue 

2789 

2790 for sta, locs in stas: 

2791 stations = self.get_stations( 

2792 tmin=tmin, 

2793 tmax=tmax, 

2794 codes=(net, sta, '*'), 

2795 model='stationxml') 

2796 

2797 errors = sx.check_overlaps( 

2798 'Station', (net, sta), stations) 

2799 

2800 if errors: 

2801 raise sx.Inconsistencies( 

2802 'Inconsistencies found:\n %s' 

2803 % '\n '.join(errors)) 

2804 

2805 network.station_list.extend(stations) 

2806 

2807 if level not in ('channel', 'response'): 

2808 continue 

2809 

2810 for loc, chas in locs: 

2811 for cha, _ in chas: 

2812 channels = self.get_channels( 

2813 tmin=tmin, 

2814 tmax=tmax, 

2815 codes=(net, sta, loc, cha), 

2816 model='stationxml') 

2817 

2818 errors = sx.check_overlaps( 

2819 'Channel', (net, sta, loc, cha), channels) 

2820 

2821 if errors: 

2822 raise sx.Inconsistencies( 

2823 'Inconsistencies found:\n %s' 

2824 % '\n '.join(errors)) 

2825 

2826 for channel in channels: 

2827 station = sx.find_containing(stations, channel) 

2828 if station is not None: 

2829 station.channel_list.append(channel) 

2830 else: 

2831 raise sx.Inconsistencies( 

2832 'No station or station epoch found for ' 

2833 'channel: %s' % '.'.join( 

2834 (net, sta, loc, cha))) 

2835 

2836 if level != 'response': 

2837 continue 

2838 

2839 response_sq, response_sx = self.get_response( 

2840 codes=(net, sta, loc, cha), 

2841 tmin=channel.start_date, 

2842 tmax=channel.end_date, 

2843 model='stationxml+') 

2844 

2845 if not ( 

2846 sx.eq_open( 

2847 channel.start_date, response_sq.tmin) 

2848 and sx.eq_open( 

2849 channel.end_date, response_sq.tmax)): 

2850 

2851 raise sx.Inconsistencies( 

2852 'Response time span does not match ' 

2853 'channel time span: %s' % '.'.join( 

2854 (net, sta, loc, cha))) 

2855 

2856 channel.response = response_sx 

2857 

2858 return sx.FDSNStationXML( 

2859 source='Generated by Pyrocko Squirrel.', 

2860 network_list=networks) 

2861 

2862 def add_operator(self, op): 

2863 self._operators.append(op) 

2864 

2865 def update_operator_mappings(self): 

2866 available = self.get_codes(kind=('channel')) 

2867 

2868 for operator in self._operators: 

2869 operator.update_mappings(available, self._operator_registry) 

2870 

2871 def iter_operator_mappings(self): 

2872 for operator in self._operators: 

2873 for in_codes, out_codes in operator.iter_mappings(): 

2874 yield operator, in_codes, out_codes 

2875 

2876 def get_operator_mappings(self): 

2877 return list(self.iter_operator_mappings()) 

2878 

2879 def get_operator(self, codes): 

2880 try: 

2881 return self._operator_registry[codes][0] 

2882 except KeyError: 

2883 return None 

2884 

2885 def get_operator_group(self, codes): 

2886 try: 

2887 return self._operator_registry[codes] 

2888 except KeyError: 

2889 return None, (None, None, None) 

2890 

2891 def iter_operator_codes(self): 

2892 for _, _, out_codes in self.iter_operator_mappings(): 

2893 for codes in out_codes: 

2894 yield codes 

2895 

2896 def get_operator_codes(self): 

2897 return list(self.iter_operator_codes()) 

2898 

2899 def print_tables(self, table_names=None, stream=None): 

2900 ''' 

2901 Dump raw database tables in textual form (for debugging purposes). 

2902 

2903 :param table_names: 

2904 Names of tables to be dumped or ``None`` to dump all. 

2905 :type table_names: 

2906 :py:class:`list` of :py:class:`str` 

2907 

2908 :param stream: 

2909 Open file or ``None`` to dump to standard output. 

2910 ''' 

2911 

2912 if stream is None: 

2913 stream = sys.stdout 

2914 

2915 if isinstance(table_names, str): 

2916 table_names = [table_names] 

2917 

2918 if table_names is None: 

2919 table_names = [ 

2920 'selection_file_states', 

2921 'selection_nuts', 

2922 'selection_kind_codes_count', 

2923 'files', 'nuts', 'kind_codes', 'kind_codes_count'] 

2924 

2925 m = { 

2926 'selection_file_states': '%(db)s.%(file_states)s', 

2927 'selection_nuts': '%(db)s.%(nuts)s', 

2928 'selection_kind_codes_count': '%(db)s.%(kind_codes_count)s', 

2929 'files': 'files', 

2930 'nuts': 'nuts', 

2931 'kind_codes': 'kind_codes', 

2932 'kind_codes_count': 'kind_codes_count'} 

2933 

2934 for table_name in table_names: 

2935 self._database.print_table( 

2936 m[table_name] % self._names, stream=stream) 

2937 

2938 

2939class SquirrelStats(Object): 

2940 ''' 

2941 Container to hold statistics about contents available from a Squirrel. 

2942 

2943 See also :py:meth:`Squirrel.get_stats`. 

2944 ''' 

2945 

2946 nfiles = Int.T( 

2947 help='Number of files in selection.') 

2948 nnuts = Int.T( 

2949 help='Number of index nuts in selection.') 

2950 codes = List.T( 

2951 Tuple.T(content_t=String.T()), 

2952 help='Available code sequences in selection, e.g. ' 

2953 '(agency, network, station, location) for stations nuts.') 

2954 kinds = List.T( 

2955 String.T(), 

2956 help='Available content types in selection.') 

2957 total_size = Int.T( 

2958 help='Aggregated file size of files is selection.') 

2959 counts = Dict.T( 

2960 String.T(), Dict.T(Tuple.T(content_t=String.T()), Int.T()), 

2961 help='Breakdown of how many nuts of any content type and code ' 

2962 'sequence are available in selection, ``counts[kind][codes]``.') 

2963 time_spans = Dict.T( 

2964 String.T(), Tuple.T(content_t=Timestamp.T()), 

2965 help='Time spans by content type.') 

2966 sources = List.T( 

2967 String.T(), 

2968 help='Descriptions of attached sources.') 

2969 operators = List.T( 

2970 String.T(), 

2971 help='Descriptions of attached operators.') 

2972 

2973 def __str__(self): 

2974 kind_counts = dict( 

2975 (kind, sum(self.counts[kind].values())) for kind in self.kinds) 

2976 

2977 scodes = model.codes_to_str_abbreviated(self.codes) 

2978 

2979 ssources = '<none>' if not self.sources else '\n' + '\n'.join( 

2980 ' ' + s for s in self.sources) 

2981 

2982 soperators = '<none>' if not self.operators else '\n' + '\n'.join( 

2983 ' ' + s for s in self.operators) 

2984 

2985 def stime(t): 

2986 return util.tts(t) if t is not None and t not in ( 

2987 model.g_tmin, model.g_tmax) else '<none>' 

2988 

2989 def stable(rows): 

2990 ns = [max(len(w) for w in col) for col in zip(*rows)] 

2991 return '\n'.join( 

2992 ' '.join(w.ljust(n) for n, w in zip(ns, row)) 

2993 for row in rows) 

2994 

2995 def indent(s): 

2996 return '\n'.join(' '+line for line in s.splitlines()) 

2997 

2998 stspans = '<none>' if not self.kinds else '\n' + indent(stable([( 

2999 kind + ':', 

3000 str(kind_counts[kind]), 

3001 stime(self.time_spans[kind][0]), 

3002 '-', 

3003 stime(self.time_spans[kind][1])) for kind in sorted(self.kinds)])) 

3004 

3005 s = ''' 

3006Number of files: %i 

3007Total size of known files: %s 

3008Number of index nuts: %i 

3009Available content kinds: %s 

3010Available codes: %s 

3011Sources: %s 

3012Operators: %s''' % ( 

3013 self.nfiles, 

3014 util.human_bytesize(self.total_size), 

3015 self.nnuts, 

3016 stspans, scodes, ssources, soperators) 

3017 

3018 return s.lstrip() 

3019 

3020 

3021__all__ = [ 

3022 'Squirrel', 

3023 'SquirrelStats', 

3024]