1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

import logging 

import hashlib 

import time 

import numpy as num 

from scipy import signal 

 

from matplotlib.cm import get_cmap 

from matplotlib.colors import Normalize 

 

from pyrocko.trace import t2ind 

from .qt_compat import qg 

 

logger = logging.getLogger(__name__) 

 

 

DEFAULT_CMAP = 'viridis' 

 

 

class TraceWaterfall: 

 

def __init__(self): 

self.tmin = 0. 

self.tmax = 0. 

self.traces = [] 

 

self._current_cmap = None 

self.cmap = None 

self.norm = Normalize() 

 

self._data_cache = None 

 

self._show_absolute = False 

self._integrate = False 

self._clip_min = 0. 

self._clip_max = 1. 

self._common_scale = True 

 

self.set_cmap(DEFAULT_CMAP) 

 

def set_traces(self, traces): 

self.traces = traces 

 

def set_time_range(self, tmin, tmax): 

self.tmin = tmin 

self.tmax = tmax 

 

def set_clip(self, clip_min, clip_max): 

assert 0. <= clip_min < clip_max <= 1. 

self._clip_min = clip_min 

self._clip_max = clip_max 

 

def set_integrate(self, integrate): 

self._integrate = integrate 

 

def show_absolute_values(self, show_absolute): 

self._show_absolute = show_absolute 

 

def set_cmap(self, cmap): 

if cmap == self._current_cmap: 

return 

logger.debug('setting colormap to %s', cmap) 

self.cmap = get_cmap(cmap) 

self._current_cmap = cmap 

 

def set_common_scale(self, _common_scale): 

self._common_scale = _common_scale 

 

def get_state_hash(self): 

sha1 = hashlib.sha1() 

sha1.update(self.tmin.hex().encode()) 

sha1.update(self.tmax.hex().encode()) 

sha1.update(self._clip_min.hex().encode()) 

sha1.update(self._clip_max.hex().encode()) 

sha1.update(self.cmap.name.encode()) 

sha1.update(bytes(self._show_absolute)) 

sha1.update(bytes(self._integrate)) 

sha1.update(bytes(len(self.traces))) 

for tr in self.traces: 

sha1.update(tr.hash(unsafe=True).encode()) 

 

return sha1 

 

def get_image(self, px_x, px_y): 

hash = self.get_state_hash() 

hash.update(bytes(px_x)) 

hash.update(bytes(px_y)) 

 

data_hash = hash.hexdigest() 

 

if self._data_cache and self._data_cache[-1] == data_hash: 

logger.debug('using cached image') 

return self._data_cache 

 

# Undersample in space 

traces_step = int(len(self.traces) // px_y) + 1 

traces = self.traces[::traces_step] 

img_rows = len(traces) 

 

# Undersample in time 

raw_deltat = min(tr.deltat for tr in traces) 

raw_nsamples = int(round((self.tmax - self.tmin) / raw_deltat)) + 1 

 

img_undersample = max(1, int(raw_nsamples // (2*px_x))) 

img_deltat = raw_deltat * img_undersample 

img_nsamples = int(round((self.tmax - self.tmin) / img_deltat)) + 1 

 

dtypes = set(tr.ydata.dtype for tr in traces) 

dtype = num.float64 if num.float64 in dtypes else num.float32 

 

data = num.zeros((img_rows, img_nsamples), dtype=dtype) 

empty_data = num.ones_like(data, dtype=num.bool) 

 

deltats = num.zeros(img_rows) if self._integrate else None 

 

logger.debug( 

'image render: using [::%d] traces at %d time undersampling' 

' - rect (%d, %d), data: (%d, %d)', 

traces_step, img_undersample, px_y, px_x, *data.shape) 

 

for itr, tr in enumerate(traces): 

tr_data = tr.ydata 

 

if tr.deltat != img_deltat: 

time_vec = tr.tmin \ 

+ num.arange((tr.tmax - tr.tmin) // img_deltat) \ 

* img_deltat 

tr_data = num.interp(time_vec, tr.get_xdata(), tr.ydata) 

 

ibeg = max(0, t2ind(self.tmin - tr.tmin, img_deltat, round)) 

iend = min( 

tr_data.size, 

t2ind(self.tmax - tr.tmin, img_deltat, round)) 

tr_tmin = tr.tmin + ibeg * img_deltat 

 

img_ibeg = max(0, t2ind(tr_tmin - self.tmin, img_deltat, round)) 

img_iend = img_ibeg + (iend - ibeg) 

 

data[itr, img_ibeg:img_iend] = tr_data[ibeg:iend] 

empty_data[itr, img_ibeg:img_iend] = False 

 

if self._integrate: 

deltats[itr] = tr.deltat 

 

if self._integrate: 

data = num.cumsum(data, axis=1) * deltats[:, num.newaxis] 

# data -= data.mean(axis=1)[:, num.newaxis] 

 

if self._common_scale: 

data /= num.abs(data).max(axis=1)[:, num.newaxis] 

 

if self._show_absolute: 

data = num.abs(signal.hilbert(data, axis=1)) 

vmax = data.max() 

vmin = data.min() 

else: 

vmax = num.abs(data).max() 

vmin = -vmax 

 

vrange = vmax - vmin 

 

self.norm.vmin = vmin + self._clip_min*vrange 

self.norm.vmax = vmax - (1. - self._clip_max)*vrange 

 

tstart = time.time() 

img_data = self.norm(data) 

t_norm = time.time() - tstart 

tstart = time.time() 

img_data = self.cmap(img_data, alpha=None, bytes=True) 

t_cmap = time.time() - tstart 

logger.debug('normalizing: %.3f cmap: %.3f', t_norm, t_cmap) 

 

# Mask out empty data 

img_data[empty_data, 3] = 0 

 

px_x, px_y = data.shape 

img = qg.QImage( 

img_data, 

px_y, px_x, qg.QImage.Format_RGBA8888) 

 

self._data_cache = (data, img, data_hash) 

return self._data_cache 

 

def draw_waterfall(self, p, rect=None): 

if not self.traces: 

raise AttributeError('No traces to paint.') 

 

rect = rect or p.window() 

trace_data, img, *_ = self.get_image( 

int(rect.width()), int(rect.height())) 

p.drawImage(rect, img)