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

1# https://pyrocko.org/grond - GPLv3 

2# 

3# The Grond Developers, 21st Century 

4import time 

5import logging 

6import shutil 

7import os 

8 

9from grond.config import read_config, write_config 

10from grond import meta, run_info 

11from grond.problems.base import load_optimiser_info, load_problem_info, \ 

12 ModelHistory 

13 

14op = os.path 

15 

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

17 

18 

19class GrondEnvironmentError(meta.GrondError): 

20 pass 

21 

22 

23class EventSelectionFailed(GrondEnvironmentError): 

24 pass 

25 

26 

27class NoCurrentEventAvailable(GrondEnvironmentError): 

28 def __init__(self, message='no current event available'): 

29 GrondEnvironmentError.__init__(self, message) 

30 

31 

32class NoEventSelectionAvailable(GrondEnvironmentError): 

33 def __init__(self, message='no event selection available'): 

34 GrondEnvironmentError.__init__(self, message) 

35 

36 

37class NoRundirAvailable(GrondEnvironmentError): 

38 def __init__(self, message='no rundir available'): 

39 GrondEnvironmentError.__init__(self, message) 

40 

41 

42class NoPlotCollectionManagerAvailable(GrondEnvironmentError): 

43 def __init__(self, message='no plot collection manager available'): 

44 GrondEnvironmentError.__init__(self, message) 

45 

46 

47class Environment(object): 

48 

49 def __init__(self, args=None, config=None, event_names=None): 

50 

51 self._current_event_name = None 

52 self._selected_event_names = None 

53 self._config = None 

54 self._plot_collection_manager = None 

55 self._dataset = None 

56 

57 if isinstance(args, str): 

58 args = [args] 

59 

60 if not args and not config: 

61 raise GrondEnvironmentError('missing arguments') 

62 

63 if config and event_names: 

64 self._config_path = None 

65 self._rundir_path = None 

66 self._config = config 

67 

68 if isinstance(event_names, str): 

69 event_names = [event_names] 

70 self.set_selected_event_names(event_names) 

71 

72 elif op.isdir(args[0]): 

73 self._rundir_path = args[0] 

74 self._config_path = op.join(self._rundir_path, 'config.yaml') 

75 

76 else: 

77 self._rundir_path = None 

78 self._config_path = args[0] 

79 self.set_selected_event_names(args[1:]) 

80 

81 self.reset() 

82 

83 @classmethod 

84 def discover(cls, rundir): 

85 running_fn = op.join(rundir, '.running') 

86 while op.exists(running_fn): 

87 try: 

88 cls.verify_rundir(rundir) 

89 return cls([rundir]) 

90 except GrondEnvironmentError: 

91 time.sleep(.25) 

92 raise GrondEnvironmentError('could not discover rundir') 

93 

94 @staticmethod 

95 def verify_rundir(rundir_path): 

96 files = [ 

97 'config.yaml', 

98 'problem.yaml', 

99 'optimiser.yaml', 

100 'misfits' 

101 ] 

102 for fn in files: 

103 if not op.exists(op.join(rundir_path, fn)): 

104 raise GrondEnvironmentError('inconsistent rundir') 

105 

106 def copy(self, destination, force=False): 

107 ''' Copy the environment and return it ''' 

108 files = [ 

109 'config.yaml', 

110 'problem.yaml', 

111 'optimiser.yaml', 

112 'misfits', 

113 'models', 

114 'choices', 

115 'chains' 

116 ] 

117 

118 if op.exists(destination) and not force: 

119 raise OSError('Directory %s already exists' % destination) 

120 

121 destination = op.abspath(destination) 

122 os.makedirs(destination, exist_ok=True) 

123 

124 for file in files: 

125 src = op.join(self._rundir_path, file) 

126 dest = op.join(destination, file) 

127 

128 if not op.isfile(src): 

129 logger.debug('Cannot find file %s', src) 

130 continue 

131 logger.debug('Copying %s to %s', src, dest) 

132 

133 shutil.copy(src, dest) 

134 

135 cls = self.__class__ 

136 return cls(destination) 

137 

138 def reset(self): 

139 if self._dataset: 

140 self._dataset.close() 

141 self._histories = {} 

142 self._dataset = None 

143 self._optimiser = None 

144 self._problem = None 

145 

146 def get_config(self): 

147 if self._config is None: 

148 self._config = read_config(self._config_path) 

149 

150 return self._config 

151 

152 def write_config(self): 

153 write_config(self.get_config(), self.get_config_path()) 

154 

155 def get_available_event_names(self): 

156 return self.get_config().get_event_names() 

157 

158 def set_current_event_name(self, event_name): 

159 self._current_event_name = event_name 

160 self.reset() 

161 

162 def get_current_event_name(self): 

163 if self._current_event_name is None: 

164 try: 

165 self.get_rundir_path() 

166 self._current_event_name = self.get_problem().base_source.name 

167 except NoRundirAvailable: 

168 try: 

169 event_names = self.get_selected_event_names() 

170 if len(event_names) == 1: 

171 self._current_event_name = event_names[0] 

172 else: 

173 raise NoCurrentEventAvailable() 

174 

175 except NoEventSelectionAvailable: 

176 raise NoCurrentEventAvailable() 

177 

178 return self._current_event_name 

179 

180 def set_selected_event_names(self, args): 

181 event_names = self.get_available_event_names() 

182 if len(args) == 0: 

183 if len(event_names) == 1: 

184 self._selected_event_names = event_names 

185 else: 

186 if not event_names: 

187 raise EventSelectionFailed( 

188 'No event file found, check your config!') 

189 raise EventSelectionFailed( 

190 'Ambiguous event selection. Select from available events:' 

191 '\n %s\n or \'all\' to use all available events' 

192 % '\n '.join(event_names)) 

193 

194 elif len(args) == 1 and args[0] == 'all': 

195 self._selected_event_names = event_names 

196 

197 else: 

198 self._selected_event_names = [] 

199 for event_name in args: 

200 if event_name not in event_names: 

201 self._selected_event_names = None 

202 raise EventSelectionFailed( 

203 'No such event: %s' % event_name) 

204 

205 self._selected_event_names.append(event_name) 

206 

207 @property 

208 def nevents_selected(self): 

209 return len(self.get_selected_event_names()) 

210 

211 def get_selected_event_names(self): 

212 if self._selected_event_names is None: 

213 raise NoEventSelectionAvailable() 

214 

215 return self._selected_event_names 

216 

217 def get_dataset(self): 

218 if self._dataset is None: 

219 event_name = self.get_current_event_name() 

220 self._dataset = self.get_config().get_dataset(event_name) 

221 

222 return self._dataset 

223 

224 def set_rundir_path(self, path): 

225 self._rundir_path = path 

226 

227 def get_rundir_path(self): 

228 if self._rundir_path is None: 

229 raise NoRundirAvailable() 

230 

231 return self._rundir_path 

232 

233 def have_rundir(self): 

234 return self._rundir_path is not None 

235 

236 def get_run_info_path(self): 

237 return op.join(self.get_rundir_path(), 'run_info.yaml') 

238 

239 def get_run_info(self): 

240 run_info_path = self.get_run_info_path() 

241 if not op.exists(run_info_path): 

242 info = run_info.RunInfo() 

243 return info 

244 else: 

245 return run_info.read_info(run_info_path) 

246 

247 def set_run_info(self, info): 

248 run_info_path = self.get_run_info_path() 

249 run_info.write_info(info, run_info_path) 

250 

251 def get_optimiser(self): 

252 if self._optimiser is None: 

253 try: 

254 self._optimiser = load_optimiser_info(self.get_rundir_path()) 

255 except NoRundirAvailable: 

256 self._optimiser = \ 

257 self.get_config().optimiser_config.get_optimiser() 

258 

259 return self._optimiser 

260 

261 def get_problem(self): 

262 if self._problem is None: 

263 try: 

264 self._problem = load_problem_info(self.get_rundir_path()) 

265 except NoRundirAvailable: 

266 self._problem = \ 

267 self.get_config().get_problem( 

268 self.get_dataset().get_event()) 

269 

270 return self._problem 

271 

272 def get_history(self, subset=None): 

273 if subset not in self._histories: 

274 self._histories[subset] = \ 

275 ModelHistory( 

276 self.get_problem(), 

277 nchains=self.get_optimiser().nchains, 

278 path=meta.xjoin(self.get_rundir_path(), subset)) 

279 

280 self._histories[subset].ensure_bootstrap_misfits( 

281 self.get_optimiser()) 

282 

283 return self._histories[subset] 

284 

285 def set_plot_collection_manager(self, pcm): 

286 self._plot_collection_manager = pcm 

287 

288 def get_plot_collection_manager(self): 

289 if self._plot_collection_manager is None: 

290 raise NoPlotCollectionManagerAvailable() 

291 

292 return self._plot_collection_manager 

293 

294 def setup_modelling(self): 

295 '''Must be called before any modelling can be done.''' 

296 logger.debug('Setting up modelling...') 

297 self.get_config().setup_modelling_environment(self.get_problem()) 

298 ds = self.get_dataset() 

299 for target in self.get_problem().targets: 

300 target.set_dataset(ds) 

301 

302 def get_plot_classes(self): 

303 '''Discover all plot classes relevant for the setup.''' 

304 

305 plots = set() 

306 try: 

307 plots.update(self.get_problem().get_plot_classes()) 

308 except GrondEnvironmentError: 

309 pass 

310 

311 try: 

312 plots.update(self.get_optimiser().get_plot_classes()) 

313 except GrondEnvironmentError: 

314 pass 

315 

316 try: 

317 for target in self.get_problem().targets: 

318 plots.update(target.get_plot_classes()) 

319 except GrondEnvironmentError: 

320 pass 

321 

322 return sorted(list(plots), key=lambda plot: plot.name) 

323 

324 def get_plots_path(self): 

325 try: 

326 return op.join(self.get_rundir_path(), 'plots') 

327 except NoRundirAvailable: 

328 return 'plots' 

329 

330 def get_config_path(self): 

331 return self._config_path 

332 

333 def is_running(self): 

334 return op.exists(self.get_rundir_path, '.running') 

335 

336 

337__all__ = [ 

338 'GrondEnvironmentError', 

339 'EventSelectionFailed', 

340 'NoCurrentEventAvailable', 

341 'NoRundirAvailable', 

342 'NoPlotCollectionManagerAvailable', 

343 'Environment', 

344]