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, depth=0): 

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, depth 

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, depth=-10*km) 

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( 

172 self._countries, depth=-50*km) 

173 if self._countries_pipe is None: 

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

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

176 self._countries_pipe = TextPipe( 

177 points, labels, camera=camera) 

178 

179 for actor in self._countries_pipe.actor: 

180 self._parent.add_actor(actor) 

181 

182 self._countries_pipe.set_label_size(state.label_size_countries) 

183 lightness = state.lightness_countries 

184 self._countries_pipe.set_colors( 

185 (lightness, lightness, lightness)) 

186 self._countries_pipe.set_orientation(COUNTRIES_ORIENTATION) 

187 else: 

188 self.remove_countries() 

189 

190 def update(self, *args): 

191 state = self._state 

192 

193 if state.visible: 

194 self.update_cities() 

195 self.update_countries() 

196 else: 

197 self.remove_cities() 

198 self.remove_countries() 

199 

200 self._parent.update_view() 

201 

202 def _get_controls(self): 

203 state = self._state 

204 if not self._controls: 

205 from ..state import state_bind_slider, state_bind_checkbox 

206 

207 frame = qw.QFrame() 

208 layout = qw.QGridLayout() 

209 frame.setLayout(layout) 

210 

211 # Cities 

212 il = 0 

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

214 

215 chb = qw.QCheckBox('show') 

216 layout.addWidget(chb, il, 1) 

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

218 

219 # Marker size 

220 il += 1 

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

222 

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

224 slider.setSizePolicy( 

225 qw.QSizePolicy( 

226 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed)) 

227 slider.setMinimum(0) 

228 slider.setMaximum(20) 

229 slider.setSingleStep(1) 

230 slider.setPageStep(1) 

231 layout.addWidget(slider, il, 2) 

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

233 

234 # Label size 

235 il += 1 

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

237 

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

239 slider.setSizePolicy( 

240 qw.QSizePolicy( 

241 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed)) 

242 slider.setMinimum(0.1) 

243 slider.setMaximum(3000) 

244 layout.addWidget(slider, il, 2) 

245 state_bind_slider( 

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

247 

248 # Lightness 

249 il += 1 

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

251 

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

253 slider.setSizePolicy( 

254 qw.QSizePolicy( 

255 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed)) 

256 slider.setMinimum(0) 

257 slider.setMaximum(1000) 

258 layout.addWidget(slider, il, 2) 

259 

260 state_bind_slider( 

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

262 

263 # Countries 

264 il += 1 

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

266 

267 chb = qw.QCheckBox('show') 

268 layout.addWidget(chb, il, 1) 

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

270 

271 # Label size 

272 il += 1 

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

274 

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

276 slider.setSizePolicy( 

277 qw.QSizePolicy( 

278 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed)) 

279 slider.setMinimum(0.1) 

280 slider.setMaximum(3000) 

281 layout.addWidget(slider, il, 2) 

282 state_bind_slider( 

283 self, state, 'label_size_countries', 

284 slider, factor=0.00001) 

285 

286 # Lightness 

287 il += 1 

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

289 

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

291 slider.setSizePolicy( 

292 qw.QSizePolicy( 

293 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed)) 

294 slider.setMinimum(0) 

295 slider.setMaximum(1000) 

296 layout.addWidget(slider, il, 2) 

297 

298 state_bind_slider( 

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

300 

301 self._controls = frame 

302 

303 return self._controls 

304 

305 

306__all__ = [ 

307 'GeonamesElement', 

308 'GeonamesState' 

309]