1# https://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

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

5 

6import numpy as num 

7 

8from pyrocko import table, geometry, cake 

9from pyrocko.guts import Bool, Float 

10from pyrocko.gui.qt_compat import qw, qc 

11 

12from pyrocko.dataset.geonames import \ 

13 get_countries_region, get_cities_region 

14from pyrocko.gui.vtk_util import ScatterPipe, TextPipe 

15 

16from .base import Element, ElementState 

17 

18guts_prefix = 'sparrow' 

19 

20km = 1e3 

21 

22 

23CITIES_ORIENTATION = (0, -45, 45) 

24COUNTRIES_ORIENTATION = (0, 0, 0) 

25 

26 

27def locations_to_points(locations): 

28 locations = [loc for loc in locations if loc.lat is not None] 

29 coords = num.zeros((len(locations), 3)) 

30 

31 for i, v in enumerate(locations): 

32 coords[i, :] = v.lat, v.lon, 0 - 10*km 

33 

34 loc_table = table.Table() 

35 loc_table.add_col(('coords', '', ('lat', 'lon', 'depth')), coords) 

36 

37 points = geometry.latlondepth2xyz( 

38 loc_table.get_col('coords'), 

39 planetradius=cake.earthradius) 

40 return points, locations 

41 

42 

43def cities_to_color(cities, lightness): 

44 return num.full((len(cities), 3), lightness) 

45 

46 

47class GeonamesState(ElementState): 

48 visible = Bool.T(default=True) 

49 show_cities = Bool.T(default=False) 

50 show_country_names = Bool.T(default=False) 

51 marker_size_cities = Float.T(default=3.0) 

52 label_size_cities = Float.T(default=0.01) 

53 label_size_countries = Float.T(default=0.01) 

54 lightness_cities = Float.T(default=0.9) 

55 lightness_countries = Float.T(default=0.9) 

56 

57 def create(self): 

58 element = GeonamesElement() 

59 return element 

60 

61 

62class GeonamesElement(Element): 

63 

64 def __init__(self): 

65 Element.__init__(self) 

66 self._cities_pipe = None 

67 self._cities_label_pipe = None 

68 self._countries_pipe = None 

69 self._controls = None 

70 self._countries = None 

71 self._cities = None 

72 

73 def bind_state(self, state): 

74 Element.bind_state(self, state) 

75 self.talkie_connect( 

76 state, 

77 ['visible', 'show_cities', 

78 'show_country_names', 'marker_size_cities', 

79 'label_size_countries', 'label_size_cities', 

80 'lightness_countries', 'lightness_cities'], 

81 self.update) 

82 

83 def get_name(self): 

84 return 'Geonames' 

85 

86 def set_parent(self, parent): 

87 self._parent = parent 

88 if not self._countries: 

89 self._countries = get_countries_region( 

90 minpop=1e6, minarea=8000) 

91 

92 if not self._cities: 

93 self._cities = get_cities_region(minpop=1e6) 

94 

95 self._parent.add_panel( 

96 self.get_title_label(), 

97 self._get_controls(), 

98 visible=True, 

99 title_controls=[ 

100 self.get_title_control_remove(), 

101 self.get_title_control_visible()]) 

102 

103 self.update() 

104 

105 def remove_cities(self): 

106 if self._cities_pipe is not None: 

107 self._parent.remove_actor(self._cities_pipe.actor) 

108 self._cities_pipe = None 

109 

110 if self._cities_label_pipe is not None: 

111 for act in self._cities_label_pipe.actor: 

112 self._parent.remove_actor(act) 

113 self._cities_label_pipe = None 

114 

115 def remove_countries(self): 

116 if self._countries_pipe is not None: 

117 for act in self._countries_pipe.actor: 

118 self._parent.remove_actor(act) 

119 self._countries_pipe = None 

120 

121 def unset_parent(self): 

122 self.unbind_state() 

123 if not self._parent: 

124 return 

125 

126 self.remove_cities() 

127 self.remove_countries() 

128 

129 self._parent.remove_panel(self._controls) 

130 self._controls = None 

131 

132 self._parent.update_view() 

133 self._parent = None 

134 

135 def update_cities(self): 

136 state = self._state 

137 

138 if state.show_cities: 

139 points, _ = locations_to_points(self._cities) 

140 colors = cities_to_color(self._cities, state.lightness_cities) 

141 if self._cities_pipe is None: 

142 self._cities_pipe = ScatterPipe(points) 

143 self._cities_pipe.set_colors(colors) 

144 self._parent.add_actor(self._cities_pipe.actor) 

145 

146 self._cities_pipe.set_size(state.marker_size_cities) 

147 self._cities_pipe.set_symbol('sphere') 

148 self._cities_pipe.set_colors(colors) 

149 

150 if self._cities_label_pipe is None: 

151 camera = self._parent.ren.GetActiveCamera() 

152 labels = [c.asciiname for c in self._cities] 

153 self._cities_label_pipe = TextPipe( 

154 points, labels, camera=camera) 

155 

156 for actor in self._cities_label_pipe.actor: 

157 self._parent.add_actor(actor) 

158 

159 self._cities_label_pipe.set_label_size(state.label_size_cities) 

160 lightness = state.lightness_cities 

161 self._cities_label_pipe.set_colors( 

162 (lightness, lightness, lightness)) 

163 self._cities_label_pipe.set_orientation(CITIES_ORIENTATION) 

164 else: 

165 self.remove_cities() 

166 

167 def update_countries(self): 

168 state = self._state 

169 

170 if state.show_country_names: 

171 points, countries = locations_to_points(self._countries) 

172 if self._countries_pipe is None: 

173 camera = self._parent.ren.GetActiveCamera() 

174 labels = [c.name for c in countries] 

175 self._countries_pipe = TextPipe( 

176 points, labels, camera=camera) 

177 

178 for actor in self._countries_pipe.actor: 

179 self._parent.add_actor(actor) 

180 

181 self._countries_pipe.set_label_size(state.label_size_countries) 

182 lightness = state.lightness_countries 

183 self._countries_pipe.set_colors( 

184 (lightness, lightness, lightness)) 

185 self._countries_pipe.set_orientation(COUNTRIES_ORIENTATION) 

186 else: 

187 self.remove_countries() 

188 

189 def update(self, *args): 

190 state = self._state 

191 

192 if state.visible: 

193 self.update_cities() 

194 self.update_countries() 

195 else: 

196 self.remove_cities() 

197 self.remove_countries() 

198 

199 self._parent.update_view() 

200 

201 def _get_controls(self): 

202 state = self._state 

203 if not self._controls: 

204 from ..state import state_bind_slider, state_bind_checkbox 

205 

206 frame = qw.QFrame() 

207 layout = qw.QGridLayout() 

208 frame.setLayout(layout) 

209 

210 # Cities 

211 il = 0 

212 layout.addWidget(qw.QLabel('Cities'), il, 0) 

213 

214 chb = qw.QCheckBox('show') 

215 layout.addWidget(chb, il, 1) 

216 state_bind_checkbox(self, state, 'show_cities', chb) 

217 

218 # Marker size 

219 il += 1 

220 layout.addWidget(qw.QLabel('Marker size'), il, 1) 

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(20) 

228 slider.setSingleStep(1) 

229 slider.setPageStep(1) 

230 layout.addWidget(slider, il, 2) 

231 state_bind_slider(self, state, 'marker_size_cities', slider) 

232 

233 # Label size 

234 il += 1 

235 layout.addWidget(qw.QLabel('Label size'), il, 1) 

236 

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

238 slider.setSizePolicy( 

239 qw.QSizePolicy( 

240 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed)) 

241 slider.setMinimum(0.1) 

242 slider.setMaximum(3000) 

243 layout.addWidget(slider, il, 2) 

244 state_bind_slider( 

245 self, state, 'label_size_cities', slider, factor=0.00001) 

246 

247 # Lightness 

248 il += 1 

249 layout.addWidget(qw.QLabel('Lightness'), il, 1) 

250 

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

252 slider.setSizePolicy( 

253 qw.QSizePolicy( 

254 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed)) 

255 slider.setMinimum(0) 

256 slider.setMaximum(1000) 

257 layout.addWidget(slider, il, 2) 

258 

259 state_bind_slider( 

260 self, state, 'lightness_cities', slider, factor=0.001) 

261 

262 # Countries 

263 il += 1 

264 layout.addWidget(qw.QLabel('Countries'), il, 0) 

265 

266 chb = qw.QCheckBox('show') 

267 layout.addWidget(chb, il, 1) 

268 state_bind_checkbox(self, state, 'show_country_names', chb) 

269 

270 # Label size 

271 il += 1 

272 layout.addWidget(qw.QLabel('Label size'), il, 1) 

273 

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

275 slider.setSizePolicy( 

276 qw.QSizePolicy( 

277 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed)) 

278 slider.setMinimum(0.1) 

279 slider.setMaximum(3000) 

280 layout.addWidget(slider, il, 2) 

281 state_bind_slider( 

282 self, state, 'label_size_countries', 

283 slider, factor=0.00001) 

284 

285 # Lightness 

286 il += 1 

287 layout.addWidget(qw.QLabel('Lightness'), il, 1) 

288 

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

290 slider.setSizePolicy( 

291 qw.QSizePolicy( 

292 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed)) 

293 slider.setMinimum(0) 

294 slider.setMaximum(1000) 

295 layout.addWidget(slider, il, 2) 

296 

297 state_bind_slider( 

298 self, state, 'lightness_countries', slider, factor=0.001) 

299 

300 self._controls = frame 

301 

302 return self._controls 

303 

304 

305__all__ = [ 

306 'GeonamesElement', 

307 'GeonamesState' 

308]