Coverage for /usr/local/lib/python3.11/dist-packages/grond/targets/plot.py: 94%

1import numpy as num 

2import logging 

3import math 

4 

5from pyrocko import plot 

6from pyrocko.guts import Tuple, Float 

7 

8from matplotlib import pyplot as plt, colors as mcolors, cm as mcm 

9from matplotlib import lines 

10from matplotlib.ticker import FuncFormatter 

11 

12from grond.plot.config import PlotConfig 

13 

14guts_prefix = 'grond' 

15km = 1e3 

16d2r = math.pi / 180. 

17 

18logger = logging.getLogger('targets.plot') 

19 

20 

21def ndig_inc(vinc): 

22 assert vinc != 0.0 

23 ndig = -math.floor(math.log10(abs(vinc))) 

24 if ndig <= 0: 

25 return 0 

26 else: 

27 return ndig + len(('%5.3f' % (vinc * 10**ndig)).rstrip('0')) - 2 

28 

29 

30def make_scale(vmin, vmax, approx_ticks=3, mode='0-max', snap=True): 

31 cscale = plot.AutoScaler(approx_ticks=approx_ticks, mode=mode, snap=snap) 

32 vmin, vmax, vinc = cscale.make_scale((vmin, vmax)) 

33 imin = math.ceil(vmin/vinc) 

34 imax = math.floor(vmax/vinc) 

35 vs = num.arange(imin, imax+1) * vinc 

36 vexp = cscale.make_exp(vinc) 

37 vexp_factor = 10**vexp 

38 vinc_base = vinc / vexp_factor 

39 ndig = ndig_inc(vinc_base) 

40 return vs, vexp, '%%.%if' % ndig 

41 

42 

43def darken(f, c): 

44 c = num.asarray(c) 

45 return f*c 

46 

47 

48class StationDistributionPlot(PlotConfig): 

49 ''' Plot showing all waveform fits for the ensemble of solutions''' 

50 

51 name = 'seismic_stations_base' 

52 size_cm = Tuple.T( 

53 2, Float.T(), 

54 default=(16., 13.), 

55 help='width and length of the figure in cm') 

56 font_size = Float.T( 

57 default=10, 

58 help='font size of all text, except station labels') 

59 font_size_labels = Float.T( 

60 default=6, 

61 help='font size of station labels') 

62 

63 def plot_station_distribution( 

64 self, azimuths, distances, weights, labels=None, 

65 colors=None, cmap=None, cnorm=None, clabel=None, 

66 scatter_kwargs=dict(), annotate_kwargs=dict(), maxsize=10**2, 

67 legend_title=''): 

68 

69 invalid_color = plot.mpl_color('aluminium3') 

70 

71 scatter_default = { 

72 'alpha': .5, 

73 'zorder': 10, 

74 'c': plot.mpl_color('skyblue2'), 

75 } 

76 

77 annotate_default = { 

78 'alpha': .8, 

79 'color': 'k', 

80 'fontsize': self.font_size_labels, 

81 'ha': 'right', 

82 'va': 'top', 

83 'xytext': (-5, -5), 

84 'textcoords': 'offset points' 

85 } 

86 

87 scatter_default.update(scatter_kwargs) 

88 annotate_default.update(annotate_kwargs) 

89 

90 fig = plt.figure(figsize=self.size_inch) 

91 

92 plot.mpl_margins( 

93 fig, nw=1, nh=1, left=3., right=10., top=3., bottom=3., 

94 units=self.font_size) 

95 

96 ax = fig.add_subplot(111, projection='polar') 

97 

98 valid = num.isfinite(weights) 

99 valid[valid] = num.logical_and(valid[valid], weights[valid] > 0.0) 

100 

101 weights = weights.copy() 

102 if num.sum(valid) == 0: 

103 weights[:] = 1.0 

104 weights_ref = 1.0 

105 else: 

106 weights[~valid] = weights[valid].min() 

107 weights_ref = plot.nice_value(weights[valid].max()) 

108 

109 if weights_ref == 0.: 

110 weights_ref = 1.0 

111 

112 if colors is None: 

113 scolors = [ 

114 scatter_default['c'] if s else invalid_color for s in valid] 

115 else: 

116 if cnorm is None: 

117 cnorm = mcolors.Normalize() 

118 elif isinstance(cnorm, tuple): 

119 cnorm = mcolors.Normalize(vmin=cnorm[0], vmax=cnorm[1]) 

120 

121 if cmap is None: 

122 cmap = plt.get_cmap('viridis') 

123 elif isinstance(cmap, str): 

124 cmap = plt.get_cmap(cmap) 

125 

126 sm = mcm.ScalarMappable(norm=cnorm, cmap=cmap) 

127 sm.set_array(colors) 

128 

129 scolors = [ 

130 sm.to_rgba(value) for value in colors] 

131 

132 scolors = num.array(scolors) 

133 

134 scatter_default.pop('c') 

135 

136 weights_scaled = (weights / weights_ref) * maxsize 

137 

138 ws, exp, fmt = make_scale(0., weights_ref) 

139 ws = ws[1:] 

140 weight_clip_min = ws[0] 

141 weight_clip_min_scaled = (weight_clip_min / weights_ref) * maxsize 

142 weights_scaled = num.maximum(weight_clip_min_scaled, weights_scaled) 

143 

144 stations = ax.scatter( 

145 azimuths*d2r, distances, s=weights_scaled, c=scolors, 

146 edgecolors=darken(0.5, scolors), 

147 linewidths=1.0, 

148 **scatter_default) 

149 

150 if len(labels) < 30: # TODO: remove after impl. of collision detection 

151 if labels is not None: 

152 for ilbl, label in enumerate(labels): 

153 ax.annotate( 

154 label, (azimuths[ilbl]*d2r, distances[ilbl]), 

155 **annotate_default) 

156 

157 ax.set_theta_zero_location('N') 

158 ax.set_theta_direction(-1) 

159 ax.tick_params('y', labelsize=self.font_size, labelcolor='gray') 

160 ax.grid(alpha=.2) 

161 ax.set_ylim(0, distances.max()*1.1) 

162 ax.yaxis.set_major_locator(plt.MaxNLocator(4)) 

163 ax.yaxis.set_major_formatter( 

164 FuncFormatter(lambda x, pos: '%d km' % (x/km) if x != 0.0 else '')) 

165 

166 # Legend 

167 valid_marker = num.argmax(valid) 

168 ecl = stations.get_edgecolor() 

169 fc = tuple(stations.get_facecolor()[valid_marker]) 

170 ec = tuple(ecl[min(valid_marker, len(ecl)-1)]) 

171 

172 legend_data = [] 

173 for w in ws[::-1]: 

174 legend_data.append(( 

175 lines.Line2D( 

176 [0], [0], 

177 ls='none', 

178 marker='o', 

179 ms=num.sqrt(w/weights_ref*maxsize), 

180 mfc=fc, 

181 mec=ec), 

182 fmt % (w/10**exp) + ( 

183 '$\\times 10^{%i}$' % exp if exp != 0 else ''))) 

184 

185 if not num.all(valid): 

186 legend_data.append(( 

187 lines.Line2D( 

188 [0], [0], 

189 marker='o', 

190 ms=num.sqrt(maxsize), 

191 mfc=invalid_color, 

192 mec=darken(0.5, invalid_color), 

193 mew=1.0), 

194 'Excluded')) 

195 

196 legend = fig.legend( 

197 *zip(*legend_data), 

198 fontsize=self.font_size, 

199 loc='upper left', bbox_to_anchor=(0.77, 0.4), 

200 markerscale=1, numpoints=1, 

201 frameon=False) 

202 

203 legend.set_title( 

204 legend_title, 

205 prop=dict(size=self.font_size)) 

206 

207 cb_axes = fig.add_axes([0.8, 0.6, 0.02, 0.3]) 

208 

209 if clabel is not None: 

210 fig.colorbar(sm, cax=cb_axes, label=clabel, extend='both') 

211 

212 return fig, ax, legend