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

1# https://pyrocko.org - GPLv3

3# The Pyrocko Developers, 21st Century

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

6import os

7import base64

9import numpy as num

11from pyrocko.plot import automap, mpl_get_cmap_names, mpl_get_cmap

12from pyrocko.guts import String, Float, StringChoice, Bool

13from pyrocko.plot import AutoScaler, AutoScaleMode

14from pyrocko.dataset import topo

16from pyrocko.gui.talkie import (TalkieRoot, TalkieConnectionOwner,

17 has_computed, computed)

19from pyrocko.gui.qt_compat import qc, qw

20from pyrocko.gui.vtk_util import cpt_to_vtk_lookuptable, ColorbarPipe

23from .. import common

24from ..state import \

25 state_bind_combobox, state_bind, state_bind_checkbox, state_bind_slider

28mpl_cmap_blacklist = [

29 "prism", "flag",

30 "Accent", "Dark2",

31 "Paired", "Pastel1", "Pastel2",

32 "Set1", "Set2", "Set3",

33 "tab10", "tab20", "tab20b", "tab20c"

34]

37def get_mpl_cmap_choices():

38 names = mpl_get_cmap_names()

39 for cmap_name in mpl_cmap_blacklist:

40 try:

41 names.remove(cmap_name)

42 names.remove("%s_r" % cmap_name)

43 except ValueError:

44 pass

46 return names

49def random_id():

50 return base64.urlsafe_b64encode(os.urandom(16)).decode('ascii')

53class ElementState(TalkieRoot):

55 element_id = String.T()

57 def __init__(self, **kwargs):

58 if 'element_id' not in kwargs:

59 kwargs['element_id'] = random_id()

61 TalkieRoot.__init__(self, **kwargs)

64class Element(TalkieConnectionOwner):

65 def __init__(self):

66 TalkieConnectionOwner.__init__(self)

67 self._parent = None

68 self._state = None

70 def remove(self):

71 if self._parent and self._state:

72 self._parent.state.elements.remove(self._state)

74 def set_parent(self, parent):

75 self._parent = parent

77 def unset_parent(self):

78 print(self)

79 raise NotImplementedError

81 def bind_state(self, state):

82 self._state = state

84 def unbind_state(self):

85 self.talkie_disconnect_all()

86 self._state = None

88 def update_visibility(self, visible):

89 self._state.visible = visible

91 def get_title_label(self):

92 title_label = common.MyDockWidgetTitleBarLabel(self.get_name())

94 def update_label(*args):

95 title_label.set_slug(self._state.element_id)

97 self.talkie_connect(

98 self._state, 'element_id', update_label)

100 update_label()

101 return title_label

102

103 def get_title_control_remove(self):

104 button = common.MyDockWidgetTitleBarButton('\u2716')

105 button.setStatusTip('Remove Element')

106 button.clicked.connect(self.remove)

107 return button

108

109 def get_title_control_visible(self):

110 assert hasattr(self._state, 'visible')

111

112 button = common.MyDockWidgetTitleBarButtonToggle('\u2b53', '\u2b54')

113 button.setStatusTip('Toggle Element Visibility')

114 button.toggled.connect(self.update_visibility)

115

116 def set_button_checked(*args):

117 button.blockSignals(True)

118 button.set_checked(self._state.visible)

119 button.blockSignals(False)

120

121 set_button_checked()

122

123 self.talkie_connect(

124 self._state, 'visible', set_button_checked)

125

126 return button

127

128

129class CPTChoice(StringChoice):

130

131 choices = ['slip_colors'] + get_mpl_cmap_choices()

132

133

134class ColorBarPositionChoice(StringChoice):

135 choices = ['bottom-left', 'bottom-right', 'top-left', 'top-right']

136

137

138@has_computed

139class CPTState(ElementState):

140 cpt_name = String.T(default=CPTChoice.choices[0])

141 cpt_mode = String.T(default=AutoScaleMode.choices[1])

142 cpt_scale_min = Float.T(optional=True)

143 cpt_scale_max = Float.T(optional=True)

144 cpt_revert = Bool.T(default=False)

145 cbar_show = Bool.T(default=True)

146 cbar_position = ColorBarPositionChoice.T(default='bottom-right')

147 cbar_annotation_lightness = Float.T(default=1.0)

148 cbar_annotation_fontsize = Float.T(default=0.03)

149 cbar_height = Float.T(default=1.)

150 cbar_width = Float.T(default=1.)

151

152 @computed(['cpt_name', 'cpt_revert'])

153 def effective_cpt_name(self):

154 if self.cpt_revert:

155 return '%s_r' % self.cpt_name

156 else:

157 return self.cpt_name

158

159

160class CPTHandler(Element):

161

162 def __init__(self):

163

164 Element.__init__(self)

165 self._cpts = {}

166 self._autoscaler = None

167 self._lookuptable = None

168 self._cpt_combobox = None

169 self._values = None

170 self._state = None

171 self._cpt_scale_lineedit = None

172 self._cbar_pipe = None

173

174 def bind_state(self, cpt_state, update_function):

175 for state_attr in [

176 'effective_cpt_name', 'cpt_mode',

177 'cpt_scale_min', 'cpt_scale_max',

178 'cbar_show', 'cbar_position',

179 'cbar_annotation_lightness',

180 'cbar_annotation_fontsize',

181 'cbar_height', 'cbar_width']:

182

183 self.talkie_connect(

184 cpt_state, state_attr, update_function)

185

186 self._state = cpt_state

187

188 def unbind_state(self):

189 Element.unbind_state(self)

190 self._cpts = {}

191 self._lookuptable = None

192 self._values = None

193 self._autoscaler = None

194

195 def open_cpt_load_dialog(self):

196 caption = 'Select one *.cpt file to open'

197

198 fns, _ = qw.QFileDialog.getOpenFileNames(

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

200

201 if fns:

202 self.load_cpt_file(fns[0])

203

204 def load_cpt_file(self, path):

205 cpt_name = 'USR' + os.path.basename(path).split('.')[0]

206 self._cpts.update([(cpt_name, automap.read_cpt(path))])

207

208 self._state.cpt_name = cpt_name

209

210 self._update_cpt_combobox()

211 self.update_cpt()

212

213 def _update_cpt_combobox(self):

214 from pyrocko import config

215 conf = config.config()

216

217 if self._cpt_combobox is None:

218 raise ValueError('CPT combobox needs init before updating!')

219

220 cb = self._cpt_combobox

221

222 if cb is not None:

223 cb.clear()

224

225 for s in CPTChoice.choices:

226 if s not in self._cpts:

227 try:

228 cpt = automap.read_cpt(topo.cpt(s))

229 except Exception:

230 cmap = mpl_get_cmap(s)

231 cpt = automap.CPT.from_numpy(cmap(range(256))[:, :-1])

232

233 self._cpts.update([(s, cpt)])

234

235 cpt_dir = conf.colortables_dir

236 if os.path.isdir(cpt_dir):

237 for f in [

238 f for f in os.listdir(cpt_dir)

239 if f.lower().endswith('.cpt')]:

240

241 s = 'USR' + os.path.basename(f).split('.')[0]

242 self._cpts.update(

243 [(s, automap.read_cpt(os.path.join(cpt_dir, f)))])

244

245 for i, (s, cpt) in enumerate(self._cpts.items()):

246 if s[-2::] != "_r":

247 cb.insertItem(i, s, qc.QVariant(self._cpts[s]))

248 cb.setItemData(i, qc.QVariant(s), qc.Qt.ToolTipRole)

249

250 cb.setCurrentIndex(cb.findText(self._state.effective_cpt_name))

251

252 def _update_cptscale_lineedit(self):

253 le = self._cpt_scale_lineedit

254 if le is not None:

255 le.clear()

256

257 self._cptscale_to_lineedit(self._state, le)

258

259 def _cptscale_to_lineedit(self, state, widget):

260 # sel = widget.selectedText() == widget.text()

261

262 crange = (None, None)

263 if self._lookuptable is not None:

264 crange = self._lookuptable.GetRange()

265

266 if state.cpt_scale_min is not None and state.cpt_scale_max is not None:

267 crange = state.cpt_scale_min, state.cpt_scale_max

268

269 fmt = ', '.join(['%s' if item is None else '%g' for item in crange])

270

271 widget.setText(fmt % crange)

272

273 # if sel:

274 # widget.selectAll()

275

276 def update_cpt(self, mask_zeros=False):

277 state = self._state

278

279 if self._autoscaler is None:

280 self._autoscaler = AutoScaler()

281

282 if self._cpt_scale_lineedit:

283 if state.cpt_mode == 'off':

284 self._cpt_scale_lineedit.setEnabled(True)

285 else:

286 self._cpt_scale_lineedit.setEnabled(False)

287

288 if state.cpt_scale_min is not None:

289 state.cpt_scale_min = None

290

291 if state.cpt_scale_max is not None:

292 state.cpt_scale_max = None

293

294 if state.effective_cpt_name is not None and self._values is not None:

295 if self._values.size == 0:

296 vscale = (0., 1.)

297 else:

298 vscale = (num.nanmin(self._values), num.nanmax(self._values))

299

300 vmin, vmax = None, None

301 if None not in (state.cpt_scale_min, state.cpt_scale_max):

302 vmin, vmax = state.cpt_scale_min, state.cpt_scale_max

303 else:

304 vmin, vmax, _ = self._autoscaler.make_scale(

305 vscale, override_mode=state.cpt_mode)

306

307 self._cpts[state.effective_cpt_name].scale(vmin, vmax)

308 cpt = self._cpts[state.effective_cpt_name]

309 vtk_lut = cpt_to_vtk_lookuptable(cpt, mask_zeros=mask_zeros)

310 vtk_lut.SetNanColor(0.0, 0.0, 0.0, 0.0)

311

312 self._lookuptable = vtk_lut

313 self._update_cptscale_lineedit()

314

315 elif state.effective_cpt_name and self._values is None:

316 raise ValueError('No values passed to colormapper!')

317

318 def update_cbar(self, display_parameter):

319

320 state = self._state

321 lut = self._lookuptable

322

323 if state.cbar_show and lut:

324 sx, sy = 1, 1

325 off = 0.08 * sy

326 pos = {

327 'top-left': (off, sy/2 + off, 0, 2),

328 'top-right': (sx - off, sy/2 + off, 2, 2),

329 'bottom-left': (off, off, 0, 0),

330 'bottom-right': (sx - off, off, 2, 0)}

331 x, y, _, _ = pos[state.cbar_position]

332

333 if not isinstance(self._cbar_pipe, ColorbarPipe):

334 self._cbar_pipe = ColorbarPipe(

335 parent_pipe=self._parent,

336 lut=lut,

337 cbar_title=display_parameter,

338 position=(x, y))

339 self._parent.add_actor(self._cbar_pipe.actor)

340 else:

341 self._cbar_pipe.set_lookuptable(lut)

342 self._cbar_pipe.set_title(display_parameter)

343 self._cbar_pipe._set_position(x, y)

344

345 sx, sy = self._parent.gui_state.size

346 fontsize = round(state.cbar_annotation_fontsize*sy)

347 lightness = 0.9 * state.cbar_annotation_lightness

348 self._cbar_pipe._format_text(

349 lightness=lightness, fontsize=fontsize)

350

351 height_px = int(round(sy / 3 * state.cbar_height))

352 width_px = int(round(50 * state.cbar_width))

353 self._cbar_pipe._format_size(height_px, width_px)

354

355 else:

356 self.remove_cbar_pipe()

357

358 def remove_cbar_pipe(self):

359 if self._cbar_pipe is not None:

360 self._parent.remove_actor(self._cbar_pipe.actor)

361

362 self._cbar_pipe = None

363

364 def cpt_controls(self, parent, state, layout):

365 self._parent = parent

366

367 iy = layout.rowCount() + 1

368

369 layout.addWidget(qw.QLabel('Color Map'), iy, 0)

370

371 cb = common.CPTComboBox()

372 layout.addWidget(cb, iy, 1)

373 state_bind_combobox(

374 self, state, 'cpt_name', cb)

375

376 self._cpt_combobox = cb

377

378 pb = qw.QPushButton('Load CPT')

379 layout.addWidget(pb, iy, 2)

380 pb.clicked.connect(self.open_cpt_load_dialog)

381

382 iy += 1

383 layout.addWidget(qw.QLabel('Color Scaling'), iy, 0)

384

385 cb = common.string_choices_to_combobox(AutoScaleMode)

386 layout.addWidget(cb, iy, 1)

387 state_bind_combobox(

388 self, state, 'cpt_mode', cb)

389

390 le = qw.QLineEdit()

391 le.setEnabled(False)

392 layout.addWidget(le, iy, 2)

393 state_bind(

394 self, state,

395 ['cpt_scale_min', 'cpt_scale_max'], _lineedit_to_cptscale,

396 le, [le.editingFinished, le.returnPressed],

397 self._cptscale_to_lineedit)

398

399 self._cpt_scale_lineedit = le

400

401 iy += 1

402 cb = qw.QCheckBox('Revert')

403 layout.addWidget(cb, iy, 1)

404 state_bind_checkbox(self, state, 'cpt_revert', cb)

405

406 # color bar

407 iy += 1

408 layout.addWidget(qw.QLabel('Color Bar'), iy, 0)

409

410 chb = qw.QCheckBox('show')

411 layout.addWidget(chb, iy, 1)

412 state_bind_checkbox(self, state, 'cbar_show', chb)

413

414 cb = common.string_choices_to_combobox(

415 ColorBarPositionChoice)

416 layout.addWidget(cb, iy, 2)

417 state_bind_combobox(

418 self, self._state, 'cbar_position', cb)

419

420 # cbar text

421 iy += 1

422 layout.addWidget(qw.QLabel('Lightness'), iy, 1)

423

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

425 slider.setSizePolicy(

426 qw.QSizePolicy(

427 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))

428 slider.setMinimum(0)

429 slider.setMaximum(1000)

430 layout.addWidget(slider, iy, 2)

431

432 state_bind_slider(

433 self,

434 self._state,

435 'cbar_annotation_lightness',

436 slider,

437 factor=0.001)

438

439 iy += 1

440 layout.addWidget(qw.QLabel('Fontsize'), iy, 1)

441

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

443 slider.setSizePolicy(

444 qw.QSizePolicy(

445 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))

446 slider.setMinimum(0)

447 slider.setMaximum(100)

448 layout.addWidget(slider, iy, 2)

449

450 state_bind_slider(

451 self,

452 self._state,

453 'cbar_annotation_fontsize',

454 slider,

455 factor=0.001)

456

457 # cbar size

458 iy += 1

459 layout.addWidget(qw.QLabel('Height'), iy, 1)

460

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

462 slider.setSizePolicy(

463 qw.QSizePolicy(

464 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))

465 slider.setMinimum(1)

466 slider.setMaximum(200)

467 layout.addWidget(slider, iy, 2)

468

469 state_bind_slider(

470 self,

471 self._state,

472 'cbar_height',

473 slider,

474 factor=0.01)

475

476 iy += 1

477 layout.addWidget(qw.QLabel('Width'), iy, 1)

478

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

480 slider.setSizePolicy(

481 qw.QSizePolicy(

482 qw.QSizePolicy.Expanding, qw.QSizePolicy.Fixed))

483 slider.setMinimum(1)

484 slider.setMaximum(200)

485 layout.addWidget(slider, iy, 2)

486

487 state_bind_slider(

488 self,

489 self._state,

490 'cbar_width',

491 slider,

492 factor=0.01)

493

494

495def _lineedit_to_cptscale(widget, cpt_state):

496 s = str(widget.text())

497 s = s.replace(',', ' ')

498

499 crange = tuple((float(i) for i in s.split()))

500 crange = tuple((

501 crange[0],

502 crange[0]+0.01 if crange[0] >= crange[1] else crange[1]))

503

504 try:

505 cpt_state.cpt_scale_min, cpt_state.cpt_scale_max = crange

506 except Exception:

507 raise ValueError(

508 'need two numerical values: <vmin>, <vmax>')

509

510

511__all__ = [

512 'Element',

513 'ElementState',

514 'random_id',

515]