1import re 

2import os.path as op 

3from pyrocko import guts, gf 

4from pyrocko.guts import Bool, List, String 

5 

6from .meta import Path, HasPaths, GrondError 

7from .dataset import DatasetConfig 

8from .analysers.base import AnalyserConfig 

9from .analysers.target_balancing import TargetBalancingAnalyserConfig 

10from .problems.base import ProblemConfig 

11from .optimisers.base import OptimiserConfig 

12from .targets.base import TargetGroup 

13from .version import __version__ 

14 

15guts_prefix = 'grond' 

16 

17 

18def color_diff(diff): 

19 green = '\x1b[32m' 

20 red = '\x1b[31m' 

21 blue = '\x1b[34m' 

22 dim = '\x1b[2m' 

23 reset = '\x1b[0m' 

24 

25 for line in diff: 

26 if line.startswith('+'): 

27 yield green + line + reset 

28 elif line.startswith('-'): 

29 yield red + line + reset 

30 elif line.startswith('^'): 

31 yield blue + line + reset 

32 elif line.startswith('@'): 

33 yield dim + line + reset 

34 else: 

35 yield line 

36 

37 

38class EngineConfig(HasPaths): 

39 gf_stores_from_pyrocko_config = Bool.T( 

40 default=True, 

41 help='Load the GF stores from ~/.pyrocko/config') 

42 gf_store_superdirs = List.T( 

43 Path.T(), 

44 help='List of path hosting collection of Green\'s function stores.') 

45 gf_store_dirs = List.T( 

46 Path.T(), 

47 help='List of Green\'s function stores') 

48 

49 def __init__(self, *args, **kwargs): 

50 HasPaths.__init__(self, *args, **kwargs) 

51 self._engine = None 

52 

53 def get_engine(self): 

54 if self._engine is None: 

55 fp = self.expand_path 

56 self._engine = gf.LocalEngine( 

57 use_config=self.gf_stores_from_pyrocko_config, 

58 store_superdirs=fp(self.gf_store_superdirs), 

59 store_dirs=fp(self.gf_store_dirs)) 

60 

61 return self._engine 

62 

63 

64class Config(HasPaths): 

65 rundir_template = Path.T( 

66 help='Rundir for the optimisation, supports templating' 

67 ' (eg. ${event_name})') 

68 dataset_config = DatasetConfig.T( 

69 help='Dataset configuration object') 

70 target_groups = List.T( 

71 TargetGroup.T(), 

72 help='List of ``TargetGroup``s') 

73 problem_config = ProblemConfig.T( 

74 help='Problem config') 

75 analyser_configs = List.T( 

76 AnalyserConfig.T(), 

77 default=[TargetBalancingAnalyserConfig.D()], 

78 help='List of problem analysers') 

79 optimiser_config = OptimiserConfig.T( 

80 help='The optimisers configuration') 

81 engine_config = EngineConfig.T( 

82 default=EngineConfig.D(), 

83 help=':class:`pyrocko.gf.LocalEngine` configuration') 

84 event_names = List.T( 

85 String.T(), 

86 help='Restrict application to given event names. If empty, all events ' 

87 'found through the dataset configuration are considered.') 

88 event_names_exclude = List.T( 

89 String.T(), 

90 help='Event names to be excluded') 

91 

92 def __init__(self, *args, **kwargs): 

93 HasPaths.__init__(self, *args, **kwargs) 

94 

95 def get_event_names(self): 

96 if self.event_names: 

97 names = self.event_names 

98 else: 

99 names = self.dataset_config.get_event_names() 

100 

101 return [name for name in names if name not in self.event_names_exclude] 

102 

103 @property 

104 def nevents(self): 

105 return len(self.dataset_config.get_events()) 

106 

107 def get_dataset(self, event_name): 

108 return self.dataset_config.get_dataset(event_name) 

109 

110 def get_targets(self, event): 

111 ds = self.get_dataset(event.name) 

112 

113 targets = [] 

114 for igroup, target_group in enumerate(self.target_groups): 

115 targets.extend(target_group.get_targets( 

116 ds, event, 'target.%i' % igroup)) 

117 

118 return targets 

119 

120 def setup_modelling_environment(self, problem): 

121 problem.set_engine(self.engine_config.get_engine()) 

122 ds = self.get_dataset(problem.base_source.name) 

123 synt = ds.synthetic_test 

124 if synt: 

125 synt.set_problem(problem) 

126 problem.base_source = problem.get_source(synt.get_x()) 

127 

128 def get_problem(self, event): 

129 targets = self.get_targets(event) 

130 problem = self.problem_config.get_problem( 

131 event, self.target_groups, targets) 

132 self.setup_modelling_environment(problem) 

133 return problem 

134 

135 def get_elements(self, ypath): 

136 return list(guts.iter_elements(self, ypath)) 

137 

138 def set_elements(self, ypath, value): 

139 guts.set_elements(self, ypath, value, regularize=True) 

140 

141 def clone(self): 

142 return guts.clone(self) 

143 

144 

145def read_config(path): 

146 try: 

147 config = guts.load(filename=path) 

148 except OSError: 

149 raise GrondError( 

150 'Cannot read Grond configuration file: %s' % path) 

151 

152 if not isinstance(config, Config): 

153 raise GrondError('Invalid Grond configuration in file "%s".' % path) 

154 

155 config.set_basepath(op.dirname(path) or '.') 

156 return config 

157 

158 

159def write_config(config, path): 

160 try: 

161 basepath = config.get_basepath() 

162 dirname = op.dirname(path) or '.' 

163 config.change_basepath(dirname) 

164 guts.dump( 

165 config, 

166 filename=path, 

167 header='Grond configuration file, version %s' % __version__) 

168 

169 config.change_basepath(basepath) 

170 

171 except OSError: 

172 raise GrondError( 

173 'Cannot write Grond configuration file: %s' % path) 

174 

175 

176def diff_configs(path1, path2): 

177 import sys 

178 import difflib 

179 from pyrocko import guts_agnostic as aguts 

180 

181 t1 = aguts.load(filename=path1) 

182 t2 = aguts.load(filename=path2) 

183 

184 s1 = aguts.dump(t1) 

185 s2 = aguts.dump(t2) 

186 

187 result = list(difflib.unified_diff( 

188 s1.splitlines(1), s2.splitlines(1), 

189 'left', 'right')) 

190 

191 if sys.stdout.isatty(): 

192 sys.stdout.writelines(color_diff(result)) 

193 else: 

194 sys.stdout.writelines(result) 

195 

196 

197class YPathError(GrondError): 

198 pass 

199 

200 

201def parse_yname(yname): 

202 ident = r'[a-zA-Z][a-zA-Z0-9_]*' 

203 rint = r'-?[0-9]+' 

204 m = re.match( 

205 r'^(%s)(\[((%s)?(:)(%s)?|(%s))\])?$' 

206 % (ident, rint, rint, rint), yname) 

207 

208 if not m: 

209 raise YPathError('Syntax error in component: "%s"' % yname) 

210 

211 d = dict( 

212 name=m.group(1)) 

213 

214 if m.group(2): 

215 if m.group(5): 

216 istart = iend = None 

217 if m.group(4): 

218 istart = int(m.group(4)) 

219 if m.group(6): 

220 iend = int(m.group(6)) 

221 

222 d['slice'] = (istart, iend) 

223 else: 

224 d['index'] = int(m.group(7)) 

225 

226 return d 

227 

228 

229def _decend(obj, ynames): 

230 if ynames: 

231 for sobj in iter_get_obj(obj, ynames): 

232 yield sobj 

233 else: 

234 yield obj 

235 

236 

237def iter_get_obj(obj, ynames): 

238 yname = ynames.pop(0) 

239 d = parse_yname(yname) 

240 if d['name'] not in obj.T.propnames: 

241 raise AttributeError(d['name']) 

242 

243 obj = getattr(obj, d['name']) 

244 

245 if 'index' in d: 

246 sobj = obj[d['index']] 

247 for ssobj in _decend(sobj, ynames): 

248 yield ssobj 

249 

250 elif 'slice' in d: 

251 for i in range(*slice(*d['slice']).indices(len(obj))): 

252 sobj = obj[i] 

253 for ssobj in _decend(sobj, ynames): 

254 yield ssobj 

255 else: 

256 for sobj in _decend(obj, ynames): 

257 yield sobj 

258 

259 

260__all__ = ''' 

261 EngineConfig 

262 Config 

263 read_config 

264 write_config 

265 diff_configs 

266'''.split()