def stack(stream, length=None, move=None): """ Stack traces in stream by correlation id :param stream: |Stream| object with correlations :param length: time span of one trace in the stack in seconds (alternatively a string consisting of a number and a unit -- ``'d'`` for days and ``'h'`` for hours -- can be specified, i.e. ``'3d'`` stacks together all traces inside a three days time window, default: None, which stacks together all traces) :param move: define a moving stack, float or string, default: None -- no moving stack, if specified move usually is smaller than length to get an overlap in the stacked traces :return: |Stream| object with stacked correlations """ stream.sort() stream_stack = obspy.Stream() ids = {_corr_id(tr) for tr in stream} ids.discard(None) for id_ in ids: traces = [tr for tr in stream if _corr_id(tr) == id_] if length is None: data = np.mean([tr.data for tr in traces], dtype=float, axis=0) tr_stack = obspy.Trace(data, header=traces[0].stats) tr_stack.stats.key = tr_stack.stats.key + '_s' if 'num' in traces[0].stats: tr_stack.stats.num = sum(tr.stats.num for tr in traces) else: tr_stack.stats.num = len(traces) stream_stack.append(tr_stack) else: t1 = traces[0].stats.starttime lensec = _time2sec(length) movesec = _time2sec(move) if move else lensec if (lensec % (24 * 3600) == 0 or isinstance(length, str) and 'd' in length): t1 = UTC(t1.year, t1.month, t1.day) elif (lensec % 3600 == 0 or isinstance(length, str) and 'm' in length): t1 = UTC(t1.year, t1.month, t1.day, t1.hour) t2 = max(t1, traces[-1].stats.endtime - lensec) for t in IterTime(t1, t2, dt=movesec): sel = [ tr for tr in traces if -0.1 <= tr.stats.starttime - t <= lensec + 0.1 ] if len(sel) == 0: continue data = np.mean([tr.data for tr in sel], dtype=float, axis=0) tr_stack = obspy.Trace(data, header=sel[0].stats) key_add = '_s%s' % length + (move is not None) * ('m%s' % move) tr_stack.stats.key = tr_stack.stats.key + key_add tr_stack.stats.starttime = t if 'num' in traces[0].stats: tr_stack.stats.num = sum(tr.stats.num for tr in sel) else: tr_stack.stats.num = len(sel) stream_stack.append(tr_stack) return stream_stack
def correlate(io, day, outkey, edge=60, length=3600, overlap=1800, demean_window=True, discard=None, only_auto_correlation=False, station_combinations=None, component_combinations=('ZZ', ), max_lag=100, keep_correlations=False, stack='1d', njobs=0, **preprocessing_kwargs): """ Correlate data of one day :param io: io config dictionary :param day: |UTC| object with day :param outkey: the output key for the HDF5 index :param edge: additional time span requested from day before and after in seconds :param length: length of correlation in seconds (string possible) :param overlap: length of overlap in seconds (string possible) :param demean_window: demean each window individually before correlating :param discard: discard correlations with less data coverage (float from interval [0, 1]) :param only_auto_correlations: Only correlate stations with itself (different components possible) :param station_combinations: specify station combinations (e.g. ``'CX.PATCX-CX.PB01``, network code can be omitted, e.g. ``'PATCX-PB01'``, default: all) :param component_combinations: component combinations to calculate, tuple of strings with length two, e.g. ``('ZZ', 'ZN', 'RR')``, if ``'R'`` or ``'T'`` is specified, components will be rotated after preprocessing, default: only ZZ components :param max_lag: max time lag in correlations in seconds :param keep_correlatons: write correlations into HDF5 file (dafault: False) :param stack: stack correlations and write stacks into HDF5 file (default: ``'1d'``, must be smaller than one day or one day) .. note:: If you want to stack larger time spans use the separate stack command on correlations or stacked correlations. :param njobs: number of jobs used. Some tasks will run parallel (preprocessing and correlation). :param \*\*preprocessing_kwargs: all other kwargs are passed to `preprocess` """ inventory = io['inventory'] length = _time2sec(length) overlap = _time2sec(overlap) if not keep_correlations and stack is None: msg = ('keep_correlation is False and stack is None -> correlations ' ' would not be saved') raise ValueError(msg) components = set(''.join(component_combinations)) if 'R' in components or 'T' in components: load_components = components - {'R', 'T'} | {'N', 'E'} else: load_components = components if station_combinations is not None: load_stations = set(sta for comb in station_combinations for sta in comb.split('-')) else: load_stations = None # load data stream = obspy.Stream() for smeta in _iter_station_meta(inventory, load_components): if (load_stations is not None and smeta['station'] not in load_stations and '.'.join((smeta['network'], smeta['station'])) not in load_stations): continue stream2 = get_data(smeta, io['data'], io['data_format'], day, overlap=overlap, edge=edge) if stream2: stream += stream2 if len(stream) == 0: log.warning('empty stream for day %s', str(day)[:10]) return preprocess(stream, day, inventory, overlap=overlap, njobs=njobs, **preprocessing_kwargs) # collect trace pairs for correlation next_day = day + 24 * 3600 stations = sorted({tr.id[:-1] for tr in stream}) tasks = [] for station1, station2 in itertools.combinations_with_replacement( stations, 2): if only_auto_correlation and station1 != station2: continue if station_combinations and not any( set(station_comb.split('-')) == ({station1.rsplit('.', 2)[0], station2.rsplit('.', 2)[0]} if '.' in (station_comb) else {station1. split('.')[1], station2.split('.')[1]}) for station_comb in station_combinations): continue stream1 = Stream([tr for tr in stream if tr.id[:-1] == station1]) stream2 = Stream([tr for tr in stream if tr.id[:-1] == station2]) datetime1 = _midtime(stream1[0].stats) datetime2 = _midtime(stream2[0].stats) msg = 'Cannot get coordinates for channel %s datetime %s' try: c1 = inventory.get_coordinates(stream1[0].id, datetime=datetime1) except Exception as ex: raise RuntimeError(msg % (stream1[0].id, datetime1)) from ex try: c2 = inventory.get_coordinates(stream2[0].id, datetime=datetime2) except Exception as ex: raise RuntimeError(msg % (stream2[0].id, datetime2)) from ex args = (c1['latitude'], c1['longitude'], c2['latitude'], c2['longitude']) dist, azi, baz = gps2dist_azimuth(*args) if ('R' in components or 'T' in components) and station1 != station2: stream1 = stream1.copy() stream1b = stream1.copy().rotate('NE->RT', azi) stream1.extend(stream1b.select(component='R')) stream1.extend(stream1b.select(component='T')) stream2 = stream2.copy() stream2b = stream2.copy().rotate('NE->RT', azi) stream2.extend(stream2b.select(component='R')) stream2.extend(stream2b.select(component='T')) it_ = (itertools.product(stream1, stream2) if station1 != station2 else itertools.combinations_with_replacement(stream1, 2)) for tr1, tr2 in it_: comps = tr1.stats.channel[-1] + tr2.stats.channel[-1] if component_combinations and (comps not in component_combinations and comps[::-1] not in component_combinations): continue tasks.append((tr1, tr2, dist, azi, baz)) # start correlation do_work = partial(_slide_and_correlate_traces, day, next_day, length, overlap, discard, max_lag, outkey, demean_window) streams = start_parallel_jobs_inner_loop(tasks, do_work, njobs) xstream = Stream() xstream.traces = [tr for s_ in streams for tr in s_] if len(xstream) > 0: res = {} if keep_correlations: res['corr'] = xstream if stack: res['stack'] = yam.stack.stack(xstream, stack) return res