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

1import logging 

2import os 

3import shutil 

4import os.path as op 

5 

6from pyrocko import gf, scenario, util, guts 

7from pyrocko.gui import marker as pmarker 

8import math 

9 

10import grond 

11 

12 

13DEFAULT_STATIC_STORE = 'ak135_static' 

14DEFAULT_WAVEFORM_STORE = 'global2s' 

15 

16logger = logging.getLogger('grond.scenario') 

17km = 1e3 

18 

19 

20class GrondScenario(object): 

21 

22 def __init__(self, project_dir, 

23 center_lat=23., center_lon=52., radius=230*km, 

24 problem=None, observations=None): 

25 

26 if observations is None: 

27 observations = [] 

28 

29 self.project_dir = project_dir 

30 self.data_dir = op.join('data', 'scenario') 

31 

32 self.center_lat = center_lat 

33 self.center_lon = center_lon 

34 self.radius = radius 

35 

36 self.problem = problem 

37 self.observations = observations 

38 

39 self.rebuild = False 

40 self._scenario = None 

41 

42 def get_gf_stores_dir(self): 

43 return op.join(self.project_dir, 'gf_stores') 

44 

45 def create_project_dir(self, force=False): 

46 prj_dir = self.project_dir 

47 

48 if op.exists(prj_dir) and not force: 

49 raise grond.GrondError( 

50 'Directory "%s" already exists! Use --force to overwrite' 

51 % prj_dir) 

52 elif op.exists(prj_dir) and force: 

53 logger.info('Overwriting directory %s.', prj_dir) 

54 shutil.rmtree(prj_dir) 

55 

56 util.ensuredir(prj_dir) 

57 

58 @property 

59 def stores_wanted(self): 

60 target_generators = self.get_scenario().target_generators 

61 return set([t.store_id for t in target_generators]) 

62 

63 def symlink_gfstores(self, engine): 

64 logger.info('Symlinking Green\'s function stores...') 

65 

66 for store_id in self.stores_wanted: 

67 store = engine.get_store(store_id) 

68 dtarget = op.join(self.get_gf_stores_dir(), store_id) 

69 if not op.exists(dtarget): 

70 os.symlink(store.store_dir, dtarget) 

71 

72 def add_observation(self, observation): 

73 logger.info('Adding %s.' % observation.__class__.__name__) 

74 self.observations.append(observation) 

75 

76 def set_problem(self, problem): 

77 self.problem = problem 

78 

79 def get_dataset_config(self): 

80 events_path = op.join(self.data_dir, 'events.txt') 

81 

82 dataset_config = grond.DatasetConfig( 

83 events_path=events_path) 

84 

85 for obs in self.observations: 

86 obs.update_dataset_config(dataset_config, self.data_dir) 

87 return dataset_config 

88 

89 def set_config(self, filename): 

90 self.config_fn = filename 

91 

92 def get_scenario(self): 

93 if not self._scenario: 

94 if self.rebuild: 

95 scenario_file = op.join(self.project_dir, 'data', 'scenario', 

96 'scenario.yml') 

97 self._scenario = guts.load(filename=scenario_file) 

98 

99 else: 

100 if not self.observations: 

101 raise AttributeError('No observations set,' 

102 ' use .add_observation(Observation)') 

103 if not self.problem: 

104 raise AttributeError('No Source Problem set,' 

105 ' use .set_problem(Problem).') 

106 

107 self._scenario = scenario.ScenarioGenerator( 

108 center_lat=self.center_lat, 

109 center_lon=self.center_lon, 

110 radius=self.radius, 

111 target_generators=[obs.get_scenario_target_generator() 

112 for obs in self.observations], 

113 source_generator= self.problem.get_scenario_source_generator()) # noqa 

114 

115 return self._scenario 

116 

117 def create_scenario( 

118 self, 

119 force=False, 

120 interactive=True, 

121 gf_store_superdirs=None, 

122 make_map=True): 

123 

124 logger.info('Creating scenario...') 

125 

126 scenario = self.get_scenario() 

127 self.create_project_dir(force) 

128 util.ensuredir(self.get_gf_stores_dir()) 

129 

130 data_dir = op.join(self.project_dir, self.data_dir) 

131 util.ensuredir(data_dir) 

132 

133 scenario.dump(filename=op.join(data_dir, 'scenario.yml')) 

134 

135 if gf_store_superdirs is None: 

136 engine1 = gf.LocalEngine( 

137 use_config=True, 

138 store_superdirs=[self.get_gf_stores_dir()]) 

139 else: 

140 engine1 = gf.LocalEngine( 

141 use_config=False, 

142 store_superdirs=gf_store_superdirs) 

143 

144 scenario.init_modelling(engine=engine1) 

145 

146 scenario.ensure_gfstores( 

147 interactive=interactive) 

148 self.symlink_gfstores(engine1) 

149 

150 engine2 = gf.LocalEngine( 

151 use_config=False, 

152 store_superdirs=[self.get_gf_stores_dir()]) 

153 

154 scenario.init_modelling(engine=engine2) 

155 

156 scenario.dump_data(path=data_dir) 

157 if make_map: 

158 scenario.make_map(op.join(self.project_dir, 'scenario_map.pdf')) 

159 

160 shutil.move(op.join(data_dir, 'sources.yml'), 

161 op.join(data_dir, 'scenario_sources.yml')) 

162 

163 markers = scenario.get_onsets() 

164 marker_path = op.join(data_dir, 'picks', 'picks.markers') 

165 if markers: 

166 util.ensuredirs(marker_path) 

167 pmarker.save_markers(markers, marker_path) 

168 

169 def get_grond_config(self): 

170 engine_config = grond.EngineConfig( 

171 gf_stores_from_pyrocko_config=False, 

172 gf_store_superdirs=['gf_stores']) 

173 

174 optimiser_config = grond.HighScoreOptimiserConfig() 

175 

176 config = grond.Config( 

177 rundir_template=op.join('runs', '${problem_name}.grun'), 

178 dataset_config=self.get_dataset_config(), 

179 target_groups=[obs.get_grond_target_group() 

180 for obs in self.observations], 

181 problem_config=self.problem.get_grond_problem_config(), 

182 optimiser_config=optimiser_config, 

183 engine_config=engine_config) 

184 

185 config.set_basepath(self.project_dir) 

186 return config 

187 

188 def get_grond_config_path(self): 

189 return op.join('config', 'scenario.gronf') 

190 

191 def create_grond_files(self): 

192 logger.info('Creating Grond configuration for %s.' 

193 % ' and '.join([obs.name for obs in self.observations])) 

194 

195 config_path = op.join(self.project_dir, self.get_grond_config_path()) 

196 util.ensuredirs(config_path) 

197 grond.write_config(self.get_grond_config(), config_path) 

198 

199 def build( 

200 self, 

201 force=False, 

202 interactive=False, 

203 gf_store_superdirs=None, 

204 make_map=True): 

205 

206 logger.info('Building scenario...') 

207 

208 self.create_scenario( 

209 force=force, 

210 interactive=interactive, 

211 gf_store_superdirs=gf_store_superdirs, 

212 make_map=make_map) 

213 

214 self.create_grond_files() 

215 

216 

217class Observation(object): 

218 

219 name = 'Observation Prototype' 

220 

221 def __init__(self, store_id, *args, **kwargs): 

222 self.store_id = store_id 

223 

224 def update_dataset_config(self, dataset_config, data_dir): 

225 return dataset_config 

226 

227 def get_scenario_target_generator(self): 

228 pass 

229 

230 def get_grond_target_group(self): 

231 pass 

232 

233 

234class WaveformObservation(Observation): 

235 

236 name = 'seismic waveforms' 

237 

238 def __init__(self, store_id=DEFAULT_WAVEFORM_STORE, nstations=25, 

239 stationxml_paths=None, stations_paths=None): 

240 self.nstations = nstations 

241 self.store_id = store_id 

242 self.stationxml_paths = stationxml_paths 

243 self.stations_paths = stations_paths 

244 

245 def update_dataset_config(self, dataset_config, data_dir): 

246 ds = dataset_config 

247 ds.waveform_paths = [op.join(data_dir, 'waveforms')] 

248 ds.stations_stationxml_paths = [ 

249 op.join(data_dir, 'meta', 'waveform_response.xml')] 

250 ds.responses_stationxml_paths = [ 

251 op.join(data_dir, 'meta', 'waveform_response.xml')] 

252 return ds 

253 

254 def get_scenario_target_generator(self): 

255 if self.stationxml_paths or self.stations_paths: 

256 station_generator = scenario.station.ImportStationGenerator( 

257 stations_paths=self.stations_paths, 

258 stations_stationxml_paths=self.stationxml_paths) 

259 else: 

260 station_generator = scenario.station.RandomStationGenerator( 

261 nstations=self.nstations) 

262 return scenario.targets.WaveformGenerator( 

263 station_generators=[station_generator], 

264 store_id=self.store_id, 

265 tabulated_phases_from_store=True, 

266 tabulated_phases_noise_scale=0.3, 

267 seismogram_quantity='displacement') 

268 

269 def get_grond_target_group(self): 

270 return grond.WaveformTargetGroup( 

271 normalisation_family='time_domain', 

272 path='all', 

273 distance_min=10*km, 

274 distance_max=1000*km, 

275 channels=['Z', 'R', 'T'], 

276 interpolation='multilinear', 

277 store_id=self.store_id, 

278 misfit_config=grond.WaveformMisfitConfig( 

279 tmin='vel_surface:5.5', 

280 tmax='vel_surface:3.0', 

281 fmin=0.01, 

282 fmax=0.1)) 

283 

284 

285class InSARObservation(Observation): 

286 

287 name = 'InSAR displacement' 

288 

289 def __init__(self, store_id=DEFAULT_STATIC_STORE): 

290 self.store_id = store_id 

291 

292 def update_dataset_config(self, dataset_config, data_dir): 

293 dataset_config.kite_scene_paths = [op.join(data_dir, 'insar')] 

294 return dataset_config 

295 

296 def get_scenario_target_generator(self): 

297 return scenario.targets.InSARGenerator( 

298 mask_water=False, 

299 store_id=self.store_id) 

300 

301 def get_grond_target_group(self): 

302 logger.info('''Hint: 

303 

304Use kite's `spool` to configure the InSAR scene's quadtree and covariance! 

305''') 

306 return grond.SatelliteTargetGroup( 

307 normalisation_family='insar_target', 

308 path='all', 

309 interpolation='multilinear', 

310 store_id=self.store_id, 

311 kite_scenes=['*all'], 

312 misfit_config=grond.SatelliteMisfitConfig( 

313 optimise_orbital_ramp=True, 

314 ranges={ 

315 'offset': '-0.5 .. 0.5', 

316 'ramp_north': '-1e-4 .. 1e-4', 

317 'ramp_east': '-1e-4 .. 1e-4' 

318 }) 

319 ) 

320 

321 

322class GNSSCampaignObservation(Observation): 

323 

324 name = 'GNSS campaign' 

325 

326 def __init__(self, store_id=DEFAULT_STATIC_STORE, nstations=25): 

327 self.store_id = store_id 

328 self.nstations = nstations 

329 

330 def update_dataset_config(self, dataset_config, data_dir): 

331 dataset_config.gnss_campaign_paths = [op.join(data_dir, 'gnss')] 

332 return dataset_config 

333 

334 def get_scenario_target_generator(self): 

335 return scenario.targets.GNSSCampaignGenerator( 

336 station_generators=[scenario.targets.RandomStationGenerator( 

337 nstations=self.nstations)], 

338 store_id=self.store_id) 

339 

340 def get_grond_target_group(self): 

341 return grond.GNSSCampaignTargetGroup( 

342 gnss_campaigns=['*all'], 

343 normalisation_family='gnss_target', 

344 path='all', 

345 interpolation='multilinear', 

346 store_id=self.store_id, 

347 misfit_config=grond.GNSSCampaignMisfitConfig() 

348 ) 

349 

350 

351class SourceProblem(object): 

352 

353 def __init__(self, magnitude_min=6., magnitude_max=7., radius=10*km, 

354 nevents=1): 

355 self.magnitude_min = magnitude_min 

356 self.magnitude_max = magnitude_max 

357 self.radius = radius 

358 self.nevents = nevents 

359 

360 def get_scenario_source_generator(self): 

361 return 

362 

363 def get_grond_problem_config(self): 

364 return 

365 

366 

367class DCSourceProblem(SourceProblem): 

368 

369 def get_scenario_source_generator(self): 

370 return scenario.sources.DCSourceGenerator( 

371 magnitude_min=self.magnitude_min, 

372 magnitude_max=self.magnitude_max, 

373 radius=self.radius, 

374 depth_min=5*km, 

375 depth_max=20*km, 

376 nevents=self.nevents) 

377 

378 def get_grond_problem_config(self): 

379 s2 = math.sqrt(2) 

380 return grond.CMTProblemConfig( 

381 name_template='cmt_${event_name}', 

382 distance_min=2.*km, 

383 mt_type='full', 

384 ranges=dict( 

385 time=gf.Range(-5.0, 5.0, relative='add'), 

386 north_shift=gf.Range(-15*km, 15*km), 

387 east_shift=gf.Range(-15*km, 15*km), 

388 depth=gf.Range(5*km, 20*km), 

389 magnitude=gf.Range(self.magnitude_min, self.magnitude_max), 

390 rmnn=gf.Range(-1.0, 1.0), 

391 rmee=gf.Range(-1.0, 1.0), 

392 rmdd=gf.Range(-1.0, 1.0), 

393 rmne=gf.Range(-s2, s2), 

394 rmnd=gf.Range(-s2, s2), 

395 rmed=gf.Range(-s2, s2), 

396 duration=gf.Range(1.0, 15.0)) 

397 ) 

398 

399 

400class RectangularSourceProblem(SourceProblem): 

401 

402 def get_scenario_source_generator(self): 

403 return scenario.sources.RectangularSourceGenerator( 

404 magnitude_min=self.magnitude_min, 

405 magnitude_max=self.magnitude_max, 

406 depth_min=5*km, 

407 depth_max=15*km, 

408 radius=self.radius, 

409 nevents=self.nevents) 

410 

411 def get_grond_problem_config(self): 

412 return grond.RectangularProblemConfig( 

413 name_template='rect_source_${event_name}', 

414 decimation_factor=1, 

415 ranges=dict( 

416 north_shift=gf.Range(-15*km, 15*km), 

417 east_shift=gf.Range(-15*km, 15*km), 

418 depth=gf.Range(5*km, 15*km), 

419 length=gf.Range(20*km, 40*km), 

420 width=gf.Range(5*km, 10*km), 

421 dip=gf.Range(0, 90), 

422 strike=gf.Range(0, 180), 

423 rake=gf.Range(0, 90), 

424 slip=gf.Range(.5, 10), 

425 nucleation_x=gf.Range(-1., 1.), 

426 nucleation_y=gf.Range(-1., 1.), 

427 time=gf.Range(-5.0, 5.0, relative='add')) 

428 )