Coverage for /usr/local/lib/python3.11/dist-packages/pyrocko/gui/sparrow/elements/stations.py: 27%

1# https://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

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

5 

6import numpy as num 

7 

8from pyrocko.guts import \ 

9 Object, Bool, Float, StringChoice, Timestamp, String, List 

10 

11from pyrocko import cake, table, model, geometry 

12from pyrocko.client import fdsn 

13from pyrocko.gui.qt_compat import qw, qc 

14 

15from pyrocko.gui.vtk_util import ScatterPipe 

16from .. import common 

17 

18from .base import Element, ElementState 

19 

20guts_prefix = 'sparrow' 

21 

22 

23def stations_to_points(stations): 

24 coords = num.zeros((len(stations), 3)) 

25 

26 for i, s in enumerate(stations): 

27 coords[i, :] = s.lat, s.lon, -s.elevation 

28 

29 station_table = table.Table() 

30 

31 station_table.add_col(('coords', '', ('lat', 'lon', 'depth')), coords) 

32 

33 return geometry.latlondepth2xyz( 

34 station_table.get_col('coords'), 

35 planetradius=cake.earthradius) 

36 

37 

38class FDSNSiteChoice(StringChoice): 

39 choices = [key.upper() for key in fdsn.g_site_abbr.keys()] 

40 

41 

42class StationSelection(Object): 

43 pass 

44 

45 

46class FDSNStationSelection(StationSelection): 

47 site = String.T() 

48 tmin = Timestamp.T() 

49 tmax = Timestamp.T() 

50 

51 def get_stations(self): 

52 return fdsn.station( 

53 site=self.site, 

54 format='text', 

55 level='channel', 

56 startbefore=self.tmin, 

57 endafter=self.tmax 

58 ).get_pyrocko_stations() 

59 

60 

61class FileStationSelection(StationSelection): 

62 paths = List.T(String.T()) 

63 

64 def get_stations(self): 

65 from pyrocko.io import stationxml 

66 

67 stations = [] 

68 for path in self.paths: 

69 if path.split('.')[-1].lower() in ['xml']: 

70 stxml = stationxml.load_xml(filename=path) 

71 stations.extend(stxml.get_pyrocko_stations()) 

72 

73 else: 

74 stations.extend(model.load_stations(path)) 

75 

76 return stations 

77 

78 

79class StationsState(ElementState): 

80 visible = Bool.T(default=True) 

81 size = Float.T(default=5.0) 

82 station_selection = StationSelection.T(optional=True) 

83 

84 @classmethod 

85 def get_name(cls): 

86 return 'Stations' 

87 

88 def create(self): 

89 element = StationsElement() 

90 return element 

91 

92 

93class StationsElement(Element): 

94 

95 def __init__(self): 

96 Element.__init__(self) 

97 self._parent = None 

98 self._pipe = None 

99 self._state = None 

100 self._controls = None 

101 self._points = num.array([]) 

102 

103 def bind_state(self, state): 

104 Element.bind_state(self, state) 

105 self.talkie_connect( 

106 state, 

107 ['visible', 'size', 'station_selection'], 

108 self.update) 

109 

110 self._current_selection = None 

111 

112 def get_name(self): 

113 return 'Stations' 

114 

115 def set_parent(self, parent): 

116 self._parent = parent 

117 self._parent.add_panel( 

118 self.get_title_label(), 

119 self._get_controls(), 

120 visible=True, 

121 title_controls=[ 

122 self.get_title_control_remove(), 

123 self.get_title_control_visible()]) 

124 self.update() 

125 

126 def unset_parent(self): 

127 self.unbind_state() 

128 if self._parent: 

129 if self._pipe: 

130 self._parent.remove_actor(self._pipe.actor) 

131 self._pipe = None 

132 

133 if self._controls: 

134 self._parent.remove_panel(self._controls) 

135 self._controls = None 

136 

137 self._parent.update_view() 

138 self._parent = None 

139 

140 def update(self, *args): 

141 state = self._state 

142 if self._pipe and \ 

143 self._current_selection is not state.station_selection: 

144 

145 self._parent.remove_actor(self._pipe.actor) 

146 self._pipe = None 

147 

148 if self._pipe and not state.visible: 

149 self._parent.remove_actor(self._pipe.actor) 

150 self._pipe.set_size(state.size) 

151 

152 if state.visible: 

153 if self._current_selection is not state.station_selection: 

154 stations = state.station_selection.get_stations() 

155 points = stations_to_points(stations) 

156 self._pipe = ScatterPipe(points) 

157 self._parent.add_actor(self._pipe.actor) 

158 elif self._pipe: 

159 self._parent.add_actor(self._pipe.actor) 

160 

161 if self._pipe: 

162 self._pipe.set_size(state.size) 

163 

164 self._parent.update_view() 

165 self._current_selection = state.station_selection 

166 

167 def open_file_load_dialog(self): 

168 caption = 'Select one or more files to open' 

169 

170 fns, _ = qw.QFileDialog.getOpenFileNames( 

171 self._parent, caption, options=common.qfiledialog_options) 

172 

173 self._state.station_selection = FileStationSelection( 

174 paths=[str(fn) for fn in fns]) 

175 

176 def open_fdsn_load_dialog(self): 

177 dialog = qw.QDialog(self._parent) 

178 dialog.setWindowTitle('Get stations from FDSN web service') 

179 

180 layout = qw.QHBoxLayout(dialog) 

181 

182 layout.addWidget(qw.QLabel('Site')) 

183 

184 sites = [key.upper() for key in fdsn.g_site_abbr.keys()] 

185 

186 cb = qw.QComboBox() 

187 for i, s in enumerate(sites): 

188 cb.insertItem(i, s) 

189 

190 layout.addWidget(cb) 

191 

192 pb = qw.QPushButton('Cancel') 

193 pb.clicked.connect(dialog.reject) 

194 layout.addWidget(pb) 

195 

196 pb = qw.QPushButton('Ok') 

197 pb.clicked.connect(dialog.accept) 

198 layout.addWidget(pb) 

199 

200 dialog.exec_() 

201 

202 site = str(cb.currentText()).lower() 

203 

204 vstate = self._parent.state 

205 

206 if dialog.result() == qw.QDialog.Accepted: 

207 self._state.station_selection = FDSNStationSelection( 

208 site=site, 

209 tmin=vstate.tmin, 

210 tmax=vstate.tmax) 

211 

212 def _get_controls(self): 

213 if not self._controls: 

214 from ..state import state_bind_slider 

215 

216 frame = qw.QFrame() 

217 layout = qw.QGridLayout() 

218 frame.setLayout(layout) 

219 

220 layout.addWidget(qw.QLabel('Size'), 0, 0) 

221 

222 slider = qw.QSlider(qc.Qt.Horizontal) 

223 slider.setSizePolicy( 

224 qw.QSizePolicy( 

225 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed)) 

226 slider.setMinimum(0) 

227 slider.setMaximum(10) 

228 slider.setSingleStep(1) 

229 slider.setPageStep(1) 

230 layout.addWidget(slider, 0, 1) 

231 state_bind_slider(self, self._state, 'size', slider) 

232 

233 lab = qw.QLabel('Load from:') 

234 pb_file = qw.QPushButton('File') 

235 pb_fdsn = qw.QPushButton('FDSN') 

236 

237 layout.addWidget(lab, 1, 0) 

238 layout.addWidget(pb_file, 1, 1) 

239 layout.addWidget(pb_fdsn, 1, 2) 

240 

241 pb_file.clicked.connect(self.open_file_load_dialog) 

242 pb_fdsn.clicked.connect(self.open_fdsn_load_dialog) 

243 

244 layout.addWidget(qw.QFrame(), 2, 0, 1, 3) 

245 

246 self._controls = frame 

247 

248 return self._controls 

249 

250 

251__all__ = [ 

252 'StationsElement', 

253 'StationsState', 

254]