Coverage for /usr/local/lib/python3.13/dist-packages/pyrocko/squirrel/tool/commands/summon.py: 15%

1# http://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

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

5 

6''' 

7Implementation of :app:`squirrel summon`. 

8''' 

9 

10import math 

11import logging 

12 

13from pyrocko.squirrel.error import SquirrelError 

14from pyrocko import progress, guts, util 

15 

16logger = logging.getLogger('psq.cli.summon') 

17 

18headline = 'Fill local cache.' 

19 

20 

21def make_subparser(subparsers): 

22 return subparsers.add_parser( 

23 'summon', 

24 help=headline, 

25 description=headline) 

26 

27 

28def setup(parser): 

29 parser.add_squirrel_selection_arguments() 

30 parser.add_squirrel_query_arguments() 

31 parser.add_argument( 

32 '--channel-priorities', 

33 dest='channel_priorities', 

34 metavar='CHA', 

35 help=''' 

36List of 2-character band/instrument code combinations to try. For example, 

37giving ```HH,BH``` would first try to get ```HH?``` channels and then fallback 

38to ```BH?``` if these are not available. The first matching waveforms are 

39returned. Use in combination with ``--sample-rate-min`` and 

40``--sample-rate-max`` to constrain the sample rate. 

41'''.strip()) 

42 

43 parser.add_argument( 

44 '--sample-rate-min', 

45 dest='sample_rate_min', 

46 metavar='FLOAT', 

47 type=float, 

48 help='Minimum sample rate [Hz] to consider.') 

49 

50 parser.add_argument( 

51 '--sample-rate-max', 

52 dest='sample_rate_max', 

53 metavar='FLOAT', 

54 type=float, 

55 help='Maximum sample rate [Hz] to consider.') 

56 

57 parser.add_argument( 

58 '--get-events', 

59 dest='get_events', 

60 action='store_true', 

61 help='Get only time spans with events.') 

62 

63 parser.add_argument( 

64 '--event-time-span', 

65 dest='event_time_span', 

66 metavar='DURATION,DURATION', 

67 help='Set time span to be completed for events, e.g. `-30m,2h`.') 

68 

69 parser.add_argument( 

70 '--tinc', 

71 dest='tinc', 

72 type=float, 

73 default=3600., 

74 metavar='SECONDS', 

75 help='Download loop time window length [s].') 

76 

77 

78def run(parser, args): 

79 d = args.squirrel_query 

80 squirrel = args.make_squirrel() 

81 

82 if not squirrel.has(['waveform_promise']): 

83 logger.warning( 

84 'No waveform promises available. It may be necessary to run ' 

85 '`squirrel update --promises ...`.') 

86 

87 if args.get_events: 

88 events = squirrel.get_events( 

89 tmin=d.get('tmin', None), 

90 tmax=d.get('tmax', None)) 

91 

92 if not args.event_time_span: 

93 raise SquirrelError( 

94 'Time span setting need. Currently --event-time-span') 

95 

96 try: 

97 span_tmin, span_tmax = [ 

98 guts.parse_duration(s) 

99 for s in args.event_time_span.split(',')] 

100 

101 except Exception: 

102 raise SquirrelError( 

103 'Invalid argument to --event-time-span: %s' 

104 % args.event_time_span) 

105 

106 groups = [ 

107 (event.time + span_tmin, event.time + span_tmax, event.name) 

108 for event in events] 

109 else: 

110 p_tmin, p_tmax = squirrel.get_time_span( 

111 kinds=['waveform_promise'], 

112 dummy_limits=False) 

113 

114 if d.get('tmin') is None and p_tmin is not None: 

115 d['tmin'] = p_tmin 

116 

117 if d.get('tmax') is None and p_tmax is not None: 

118 d['tmax'] = p_tmax 

119 

120 if 'tmin' not in d or 'tmax' not in d: 

121 raise SquirrelError('Time span required.') 

122 

123 groups = [(d['tmin'], d['tmax'], '')] 

124 

125 d.pop('tmin', None) 

126 d.pop('tmax', None) 

127 

128 tinc = args.tinc 

129 

130 channel_priorities = None 

131 if args.channel_priorities: 

132 channel_priorities = [ 

133 cha.strip() for cha in args.channel_priorities.split(',')] 

134 

135 with progress.view(): 

136 task_group = progress.task('Group', len(groups)) 

137 

138 for igroup, (group_tmin, group_tmax, group_name) in enumerate(groups): 

139 group_label = '%s - %s%s' % ( 

140 util.time_to_str(group_tmin, '%Y-%m-%d %H:%M:%S'), 

141 util.time_to_str(group_tmax, '%Y-%m-%d %H:%M:%S'), 

142 '(%s)' % group_name if group_name else '') 

143 

144 task_group.update(igroup, group_label) 

145 

146 if group_name: 

147 logger.info('Summoning group: %s' % group_label) 

148 

149 tmin = math.floor(group_tmin / tinc) * tinc 

150 tmax = math.ceil(group_tmax / tinc) * tinc 

151 nwindows = int(round((tmax - tmin) / tinc)) 

152 task = progress.task('Summoning', nwindows) 

153 iwindow = 0 

154 try: 

155 for batch in squirrel.chopper_waveforms( 

156 tinc=tinc, 

157 load_data=False, 

158 channel_priorities=channel_priorities, 

159 sample_rate_min=args.sample_rate_min, 

160 sample_rate_max=args.sample_rate_max, 

161 tmin=tmin, 

162 tmax=tmax, 

163 **d): 

164 

165 iwindow += 1 

166 task.update(iwindow, '%s - %s' % ( 

167 util.time_to_str(batch.tmin, '%Y-%m-%d %H:%M:%S'), 

168 util.time_to_str(batch.tmax, '%Y-%m-%d %H:%M:%S'))) 

169 

170 finally: 

171 task.done() 

172 

173 stats = str(squirrel) 

174 stats = '\n'.join(' ' + s for s in stats.splitlines()) 

175 

176 logger.info('Squirrel stats:\n%s' % stats)