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

1# https://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

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

5 

6import numpy as num 

7import vtk 

8 

9from pyrocko.guts import Bool, StringChoice, Float 

10from pyrocko.gui.qt_compat import qw, qc 

11from pyrocko.color import Color 

12 

13 

14from pyrocko.gui import vtk_util 

15from .. import common 

16from .base import Element, ElementState 

17from pyrocko.dataset.gshhg import Coastlines, Rivers, Borders 

18 

19 

20guts_prefix = 'sparrow' 

21 

22gshhg_dataset_mapping = { 

23 'coastlines': Coastlines, 

24 'rivers': Rivers, 

25 'borders': Borders, 

26} 

27 

28 

29class GSHHGDatasetChoice(StringChoice): 

30 choices = ['coastlines', 'borders', 'rivers'] 

31 

32 

33class GSHHGResolutionChoice(StringChoice): 

34 choices = [ 

35 'crude', 

36 'low', 

37 'intermediate', 

38 'high', 

39 'full'] 

40 

41 

42class GSHHGPipe(object): 

43 def __init__(self, dataset, resolution='low', levels=None): 

44 

45 self.mapper = vtk.vtkDataSetMapper() 

46 self.plane = vtk.vtkPlane() 

47 self.plane.SetOrigin(0.0, 0.0, 0.0) 

48 coll = vtk.vtkPlaneCollection() 

49 coll.AddItem(self.plane) 

50 self.mapper.SetClippingPlanes(coll) 

51 

52 self._polyline_grid = {} 

53 self._opacity = 1.0 

54 self._line_width = 1.0 

55 self._color = Color('white') 

56 self.set_resolution(dataset, resolution, levels) 

57 

58 actor = vtk.vtkActor() 

59 actor.SetMapper(self.mapper) 

60 

61 prop = actor.GetProperty() 

62 prop.SetDiffuseColor(1, 1, 1) 

63 

64 self.prop = prop 

65 self.actor = actor 

66 

67 def set_resolution(self, dataset, resolution, levels): 

68 assert resolution in GSHHGResolutionChoice.choices 

69 assert dataset in GSHHGDatasetChoice.choices 

70 assert levels is None or isinstance(levels, tuple) 

71 

72 if (resolution, levels) not in self._polyline_grid: 

73 pb = common.get_viewer().progressbars 

74 if pb: 

75 mess = 'Loading %s resolution %s' % (resolution, dataset) 

76 pb.set_status(mess, 0, can_abort=False) 

77 

78 dataset = gshhg_dataset_mapping[dataset] 

79 

80 g = getattr(dataset, resolution)() 

81 g.load_all() 

82 

83 lines = [] 

84 npoly = len(g.polygons) 

85 for ipoly, poly in enumerate(g.polygons): 

86 if pb: 

87 pb.set_status( 

88 mess, float(ipoly) / npoly * 100., can_abort=False) 

89 

90 if levels is None or poly.level_no in levels: 

91 lines.append(poly.points) 

92 

93 self._polyline_grid[resolution, levels] \ 

94 = vtk_util.make_multi_polyline(lines_latlon=lines, depth=-200.) 

95 

96 if pb: 

97 pb.set_status(mess, 100, can_abort=False) 

98 

99 vtk_util.vtk_set_input( 

100 self.mapper, self._polyline_grid[resolution, levels]) 

101 

102 def set_opacity(self, opacity): 

103 opacity = float(opacity) 

104 if self._opacity != opacity: 

105 self.prop.SetOpacity(opacity) 

106 self._opacity = opacity 

107 

108 def set_color(self, color): 

109 if self._color != color: 

110 self.prop.SetDiffuseColor(color.rgb) 

111 self._color = color 

112 

113 def set_line_width(self, width): 

114 width = float(width) 

115 if self._line_width != width: 

116 self.prop.SetLineWidth(width) 

117 self._line_width = width 

118 

119 def set_clipping_plane(self, origin, normal): 

120 self.plane.SetOrigin(*origin) 

121 self.plane.SetNormal(*normal) 

122 

123 

124class GSHHGState(ElementState): 

125 visible = Bool.T(default=True) 

126 dataset = GSHHGDatasetChoice.T(default='coastlines') 

127 resolution = GSHHGResolutionChoice.T(default='low') 

128 opacity = Float.T(default=0.4) 

129 color = Color.T(default=Color.D('white')) 

130 line_width = Float.T(default=1.0) 

131 

132 def create(self): 

133 element = GSHHGElement() 

134 return element 

135 

136 

137class GSHHGElement(Element): 

138 

139 def __init__(self, levels=None): 

140 Element.__init__(self) 

141 self._parent = None 

142 self._controls = None 

143 self._lines = None 

144 self._levels = levels 

145 

146 def bind_state(self, state): 

147 Element.bind_state(self, state) 

148 self.talkie_connect( 

149 state, 

150 ['visible', 'resolution', 'opacity', 'color', 'line_width'], 

151 self.update) 

152 

153 def set_parent(self, parent): 

154 self._parent = parent 

155 self._parent.add_panel( 

156 self.get_title_label(), 

157 self._get_controls(), 

158 visible=True, 

159 title_controls=[ 

160 self.get_title_control_remove(), 

161 self.get_title_control_visible()]) 

162 

163 self.talkie_connect( 

164 self._parent.state, 

165 ['lat', 'lon', 'depth', 'distance', 'azimuth', 'dip'], 

166 self.update_clipping) 

167 

168 self.update() 

169 self.update_clipping() 

170 

171 def unset_parent(self): 

172 self.unbind_state() 

173 if self._parent: 

174 if self._lines: 

175 self._parent.remove_actor(self._lines.actor) 

176 self._lines = None 

177 

178 if self._controls: 

179 self._parent.remove_panel(self._controls) 

180 self._controls = None 

181 

182 self._parent.update_view() 

183 self._parent = None 

184 

185 def update(self, *args): 

186 state = self._state 

187 if not state.visible and self._lines: 

188 self._parent.remove_actor(self._lines.actor) 

189 

190 if state.visible: 

191 if not self._lines: 

192 self._lines = GSHHGPipe( 

193 dataset=state.dataset, 

194 resolution=state.resolution, 

195 levels=self._levels) 

196 

197 self._parent.add_actor(self._lines.actor) 

198 self._lines.set_resolution( 

199 state.dataset, state.resolution, self._levels) 

200 self._lines.set_opacity(state.opacity) 

201 self._lines.set_color(state.color) 

202 self._lines.set_line_width(state.line_width) 

203 

204 self._parent.update_view() 

205 

206 def update_clipping(self, *args): 

207 if self._state.visible and self._lines: 

208 cam = self._parent.camera_params[0] 

209 origin = cam / num.linalg.norm(cam)**2 

210 self._lines.set_clipping_plane(origin, cam) 

211 

212 def _get_controls(self): 

213 if not self._controls: 

214 from ..state import state_bind_combobox, \ 

215 state_bind_slider, state_bind_combobox_color 

216 

217 frame = qw.QFrame() 

218 layout = qw.QGridLayout() 

219 frame.setLayout(layout) 

220 

221 layout.addWidget(qw.QLabel('Resolution'), 0, 0) 

222 

223 cb = common.string_choices_to_combobox(GSHHGResolutionChoice) 

224 layout.addWidget(cb, 0, 1) 

225 state_bind_combobox(self, self._state, 'resolution', cb) 

226 

227 # opacity 

228 

229 layout.addWidget(qw.QLabel('Opacity'), 1, 0) 

230 

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

232 slider.setSizePolicy( 

233 qw.QSizePolicy( 

234 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed)) 

235 slider.setMinimum(0) 

236 slider.setMaximum(1000) 

237 layout.addWidget(slider, 1, 1) 

238 

239 state_bind_slider( 

240 self, self._state, 'opacity', slider, factor=0.001) 

241 

242 # color 

243 

244 layout.addWidget(qw.QLabel('Color'), 2, 0) 

245 

246 cb = common.strings_to_combobox( 

247 ['black', 'white', 'blue', 'red']) 

248 

249 layout.addWidget(cb, 2, 1) 

250 state_bind_combobox_color( 

251 self, self._state, 'color', cb) 

252 

253 layout.addWidget(qw.QLabel('Line width'), 3, 0) 

254 

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

256 slider.setSizePolicy( 

257 qw.QSizePolicy( 

258 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed)) 

259 slider.setMinimum(0) 

260 slider.setMaximum(100) 

261 layout.addWidget(slider, 3, 1) 

262 state_bind_slider( 

263 self, self._state, 'line_width', slider, factor=0.1) 

264 

265 layout.addWidget(qw.QFrame(), 5, 0, 1, 2) 

266 

267 self._controls = frame 

268 

269 return self._controls 

270 

271 

272__all__ = [ 

273 'GSHHGElement', 

274 'GSHHGState']