Coverage for /usr/local/lib/python3.11/dist-packages/grond/toy.py: 69%

1 

2import numpy as num 

3 

4from pyrocko.guts import Object, Float, Dict, List, String, Int 

5from pyrocko import gf 

6from grond import Problem, Parameter, MisfitTarget 

7from grond.optimisers.highscore.plot import HighScoreOptimiserPlot 

8 

9guts_prefix = 'grond.toy' 

10 

11 

12class ToyOptimiserPlot(HighScoreOptimiserPlot): 

13 

14 def set_source(self, source): 

15 self._source = source 

16 

17 def set_targets(self, targets): 

18 self._targets = targets 

19 

20 def set_contour_data(self, contour_data): 

21 self._contour_data = contour_data 

22 

23 def set_limits(self): 

24 self.axes.set_xlim(-10., 10.) 

25 self.axes.set_ylim(-10., 10.) 

26 

27 def start(self): 

28 HighScoreOptimiserPlot.start(self) 

29 x = [getattr(t, self.xpar_name) for t in self._targets] 

30 y = [getattr(t, self.ypar_name) for t in self._targets] 

31 self.axes.plot(x, y, '^', color='black') 

32 

33 for ibootstrap, (xc, yc, zc) in enumerate( 

34 self._contour_data['east', 'depth']): 

35 

36 zmin = num.min(zc) 

37 

38 if self.optimiser.nbootstrap < 5: 

39 alpha = 1.0 

40 else: 

41 alpha = 0.5 

42 

43 self.axes.contour( 

44 xc, yc, zc, [zmin + 0.01], 

45 colors=[self.bcolors[ibootstrap]], alpha=alpha) 

46 

47 self.axes.plot( 

48 getattr(self._source, self.xpar_name), 

49 getattr(self._source, self.ypar_name), 

50 '*', color='black') 

51 

52 

53class ToyTarget(MisfitTarget): 

54 north = Float.T() 

55 east = Float.T() 

56 depth = Float.T() 

57 obs_distance = Float.T() 

58 nmisfits = Int.T(default=1) 

59 

60 

61class ToySource(Object): 

62 north = Float.T() 

63 east = Float.T() 

64 depth = Float.T() 

65 

66 

67class ToyProblem(Problem): 

68 problem_parameters = [ 

69 Parameter('north', 'm', label='North'), 

70 Parameter('east', 'm', label='East'), 

71 Parameter('depth', 'm', label='Depth')] 

72 

73 ranges = Dict.T(String.T(), gf.Range.T()) 

74 

75 targets = List.T(ToyTarget.T()) 

76 base_source = ToySource.T() 

77 

78 def __init__(self, **kwargs): 

79 Problem.__init__(self, **kwargs) 

80 self._xtargets = None 

81 self._obs_distances = None 

82 

83 def pack(self, source): 

84 return num.array( 

85 [source.north, source.east, source.depth], dtype=float) 

86 

87 def _setup_modelling(self): 

88 if self._xtargets is None: 

89 self._xtargets = num.array( 

90 [(t.north, t.east, t.depth) for t in self.targets], 

91 dtype=float) 

92 

93 self._obs_distances = num.array( 

94 [t.obs_distance for t in self.targets], 

95 dtype=float) 

96 

97 def evaluate(self, x): 

98 raise NotImplementedError('Toy problem does not have evaluate()') 

99 

100 def misfits(self, x, mask=None): 

101 self._setup_modelling() 

102 distances = num.sqrt( 

103 num.sum((x[num.newaxis, :]-self._xtargets)**2, axis=1)) 

104 

105 misfits = num.zeros((self.ntargets, 2)) 

106 misfits[:, 0] = num.abs(distances - self._obs_distances) 

107 misfits[:, 1] = num.ones(self.ntargets) \ 

108 * num.mean(num.abs(self._obs_distances)) 

109 return misfits 

110 

111 def misfits_many(self, xs): 

112 self._setup_modelling() 

113 distances = num.sqrt( 

114 num.sum( 

115 (xs[:, num.newaxis, :]-self._xtargets[num.newaxis, :])**2, 

116 axis=2)) 

117 

118 misfits = num.zeros((xs.shape[0], self.ntargets, 2)) 

119 

120 misfits[:, :, 0] = num.abs( 

121 distances - self._obs_distances[num.newaxis, :]) 

122 

123 misfits[:, :, 1] = num.mean(num.abs(self._obs_distances)) 

124 

125 return misfits 

126 

127 def xref(self): 

128 base_source = self.base_source 

129 return num.array([ 

130 base_source.north, base_source.east, base_source.depth]) 

131 

132 def extract(self, xs, i): 

133 if xs.ndim == 1: 

134 return self.extract(xs[num.newaxis, :], i)[0] 

135 

136 if i < self.nparameters: 

137 return xs[:, i] 

138 else: 

139 return self.make_dependant( 

140 xs, self.dependants[i-self.nparameters].name) 

141 

142 

143def scenario(station_setup, noise_setup): 

144 

145 snorth = 0. 

146 seast = 0. 

147 sdepth = 5. 

148 

149 source = ToySource( 

150 north=snorth, 

151 east=seast, 

152 depth=sdepth) 

153 

154 n = 10 

155 num.random.seed(10) 

156 norths = num.random.uniform(-10., 10., n) 

157 

158 if station_setup == 'wellposed': 

159 easts = num.random.uniform(-10., 10., n) 

160 elif station_setup == 'illposed': 

161 easts = num.random.uniform(0, 10., n) 

162 else: 

163 assert False 

164 

165 depths = num.zeros(n) 

166 

167 distances = num.sqrt( 

168 (norths-snorth)**2 + 

169 (easts-seast)**2 + 

170 (depths-sdepth)**2) 

171 

172 if noise_setup == 'noisefree': 

173 measured_distances = distances 

174 elif noise_setup == 'lownoise': 

175 measured_distances = distances + num.random.normal(scale=0.4, size=n) 

176 elif noise_setup == 'highnoise': 

177 measured_distances = distances + num.random.normal(scale=0.8, size=n) 

178 

179 targets = [ 

180 ToyTarget( 

181 path='t%03i' % i, 

182 north=float(norths[i]), 

183 east=float(easts[i]), 

184 depth=float(depths[i]), 

185 obs_distance=float(measured_distances[i])) 

186 for i in range(n)] 

187 

188 return source, targets 

189 

190 

191__all__ = ''' 

192 ToyTarget 

193 ToySource 

194 ToyProblem 

195 scenario 

196'''.split()