Coverage for /usr/local/lib/python3.11/dist-packages/pyrocko/squirrel/dataset.py: 36%

1# http://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

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

5 

6''' 

7Portable dataset description. 

8 

9The :py:class:`Dataset` class defines sets of local and remote data-sources to 

10be used in combination in Squirrel-based programs. By convention, 

11Squirrel-based programs accept the ``--dataset`` option to read such dataset 

12descriptions from file. To add a dataset programmatically, to a 

13:py:class:`~pyrocko.squirrel.base.Squirrel` instance, use 

14:py:meth:`~pyrocko.squirrel.base.Squirrel.add_dataset`. 

15''' 

16 

17import os.path as op 

18import logging 

19 

20from pyrocko.guts import List, load, StringPattern, String 

21 

22from ..has_paths import HasPaths 

23from .client.base import Source 

24from .client.catalog import CatalogSource 

25from .client.fdsn import FDSNSource 

26from .error import SquirrelError 

27from .selection import re_persistent_name 

28from .operators.base import Operator 

29 

30guts_prefix = 'squirrel' 

31 

32logger = logging.getLogger('psq.dataset') 

33 

34km = 1000. 

35 

36 

37class PersistentID(StringPattern): 

38 pattern = re_persistent_name 

39 

40 

41def make_builtin_datasets(): 

42 datasets = {} 

43 for site in ['isc', 'geofon', 'gcmt']: 

44 for magnitude_min in [4.0, 5.0, 6.0, 7.0]: 

45 name = 'events-%s-m%g' % (site, magnitude_min) 

46 datasets[name] = Dataset( 

47 sources=[ 

48 CatalogSource( 

49 catalog=site, 

50 query_args=dict(magmin=magnitude_min))], 

51 comment='Event catalog: %s, minimum magnitude: %g' % ( 

52 site, magnitude_min)) 

53 

54 for site in ['geofon', 'gcmt']: 

55 depth_min = 100*km 

56 for magnitude_min in [6.0, 7.0]: 

57 name = 'events-%s-m%g-d%g' % (site, magnitude_min, depth_min/km) 

58 datasets[':' + name] = Dataset( 

59 sources=[ 

60 CatalogSource( 

61 catalog=site, 

62 query_args=dict( 

63 magmin=magnitude_min, 

64 depthmin=depth_min))], 

65 comment='Event catalog: %s, minimum magnitude: %g, ' 

66 'minimum depth: %g km' % ( 

67 site, 

68 magnitude_min, 

69 depth_min/km)) 

70 

71 for site, network, cha in [ 

72 ('bgr', 'gr', 'lh'), 

73 ('up', None, None)]: 

74 name = 'fdsn-' + '-'.join( 

75 x for x in (site, network, cha) if x is not None) 

76 

77 query_args = {} 

78 comments = ['FDSN: %s' % site] 

79 

80 if network is not None: 

81 query_args['network'] = network.upper() 

82 comments.append('network: %s' % query_args['network']) 

83 

84 if cha is not None: 

85 query_args['channel'] = cha.upper() + '?' 

86 comments.append('channels: %s' % query_args['channel']) 

87 

88 datasets[':' + name] = Dataset( 

89 sources=[FDSNSource(site=site, query_args=query_args)], 

90 comment=', '.join(comments)) 

91 

92 for site, net, sta in [ 

93 ('bgr', 'gr', 'bfo')]: 

94 name = 'fdsn-%s-%s-%s' % (site, net, sta) 

95 sta = sta.upper() 

96 net = net.upper() 

97 datasets[name] = Dataset( 

98 sources=[ 

99 FDSNSource( 

100 site=site, 

101 query_args=dict(network=net, station=sta))], 

102 comment='FDSN: %s, network: %s, ' 

103 'station: %s' % (site, net, sta)) 

104 

105 return datasets 

106 

107 

108g_builtin_datasets = None 

109 

110 

111def get_builtin_datasets(): 

112 global g_builtin_datasets 

113 g_builtin_datasets = make_builtin_datasets() 

114 return g_builtin_datasets 

115 

116 

117class Dataset(HasPaths): 

118 ''' 

119 Dataset description. 

120 ''' 

121 sources = List.T(Source.T()) 

122 operators = List.T(Operator.T()) 

123 comment = String.T(optional=True) 

124 

125 def setup(self, squirrel, check=True, upgrade=False): 

126 for source in self.sources: 

127 squirrel.add_source( 

128 source, check=check, upgrade=upgrade) 

129 

130 for operator in self.operators: 

131 squirrel.add_operator(operator) 

132 

133 squirrel.update_operator_mappings() 

134 

135 

136def read_dataset(path): 

137 ''' 

138 Read dataset description file. 

139 ''' 

140 

141 if path.startswith(':'): 

142 name = path[1:] 

143 datasets = get_builtin_datasets() 

144 try: 

145 return datasets[name] 

146 except KeyError: 

147 raise SquirrelError( 

148 ('No dataset name given. ' 

149 if not name else 'Named dataset not found: %s' % name) + 

150 '\n Use `squirrel dataset` to get information about ' 

151 'available datasets. Available:\n' 

152 ' %s' % '\n '.join( 

153 sorted(datasets.keys()))) 

154 

155 try: 

156 dataset = load(filename=path) 

157 except OSError: 

158 raise SquirrelError( 

159 'Cannot read dataset file: %s' % path) 

160 

161 if not isinstance(dataset, Dataset): 

162 raise SquirrelError('Invalid dataset file "%s".' % path) 

163 

164 dataset.set_basepath(op.dirname(path) or '.') 

165 return dataset 

166 

167 

168__all__ = [ 

169 'PersistentID', 

170 'Dataset', 

171 'read_dataset', 

172]