Refactor ssp_process_traces

CHANGELOG.md

@@ -11,6 +11,10 @@ This release requires at least Python 3.7.
 ### Bugfixes
 
 - Do not ignore picks labeled with lowercase "p" or "s"
+- Fixed: config parameter `p_arrival_tolerance` was used also for S waves,
+  instead of `s_arrival_tolerance` (see [#35])
+- Fix bug where signal and noise windows where plotted with the wrong length,
+  under certain circumstances (see [#35])
 
 ### Requirements
 
@@ -534,3 +538,4 @@ Initial Python port.
 [#28]: https://github.com/SeismicSource/sourcespec/issues/28
 [#30]: https://github.com/SeismicSource/sourcespec/issues/30
 [#31]: https://github.com/SeismicSource/sourcespec/issues/31
+[#35]: https://github.com/SeismicSource/sourcespec/issues/35

sourcespec/ssp_build_spectra.py

@@ -193,7 +193,8 @@ def _cut_signal_noise(config, trace):
             msg = f'{tr_noise_id}: truncating signal window to noise length!'
             trace_signal.data = trace_signal.data[:npts]
             # recompute signal window end time, so that it's plotted correctly
-            _recompute_time_window(trace, config.wave_type[0], npts)
+            _recompute_time_window(
+                trace, config.wave_type[0], npts, keep='start')
         else:
             msg = f'{tr_noise_id}: zero-padding noise window to signal length'
             # Notes:
@@ -206,36 +207,25 @@ def _cut_signal_noise(config, trace):
             pad_len = len(trace_signal) - len(trace_noise)
             trace_noise.data = np.pad(trace_noise.data, (0, pad_len))
             # recompute noise window start time, so that it's plotted correctly
-            _recompute_time_window(
-                trace, 'N', len(trace_signal), from_start=True)
+            _recompute_time_window(trace, 'N', len(trace_signal), keep='end')
         logger.warning(msg)
-
-    # ...and zero pad to spectral_win_length
-    if config.spectral_win_length is not None:
-        trace_signal.trim(starttime=t1,
-                          endtime=t1+config.spectral_win_length,
-                          pad=True,
-                          fill_value=0)
-        trace_noise.trim(starttime=noise_t1,
-                         endtime=noise_t1+config.spectral_win_length,
-                         pad=True,
-                         fill_value=0)
-
     return trace_signal, trace_noise
 
 
-def _recompute_time_window(trace, wave_type, npts, from_start=False):
+def _recompute_time_window(trace, wave_type, npts, keep='start'):
     """Recompute start or end time of signal or noise window,
     based on new number of points"""
     length = npts * trace.stats.delta
-    if from_start:
+    if keep == 'end':
         label, _ = trace.stats.arrivals[f'{wave_type}1']
         t1 = trace.stats.arrivals[f'{wave_type}2'][1] - length
         trace.stats.arrivals[f'{wave_type}1'] = (label, t1)
-    else:
+    elif keep == 'start':
         label, _ = trace.stats.arrivals[f'{wave_type}2']
         t2 = trace.stats.arrivals[f'{wave_type}1'][1] + length
         trace.stats.arrivals[f'{wave_type}2'] = (label, t2)
+    else:
+        raise ValueError('keep must be "start" or "end"')
 
 
 def _check_noise_level(trace_signal, trace_noise, config):
@@ -568,7 +558,15 @@ def _build_signal_and_noise_streams(config, st):
             # RuntimeError is for skipped traces
             logger.warning(msg)
             continue
-    # trim components of the same instrument to the same number of samples
+    return signal_st, noise_st
+
+
+def _trim_components(config, signal_st, noise_st, st):
+    """
+    Trim components of the same instrument to the same number of samples.
+
+    Recompute time window of the signal and noise traces for correct plotting.
+    """
     for id in sorted({tr.id[:-1] for tr in signal_st}):
         st_sel = signal_st.select(id=f'{id}*') + noise_st.select(id=f'{id}*')
         all_npts = {tr.stats.npts for tr in st_sel}
@@ -584,9 +582,8 @@ def _build_signal_and_noise_streams(config, st):
             elif tr.stats.type == 'noise':
                 tr.data = tr.data[-npts:]
         for tr in st.select(id=f'{id}*'):
-            _recompute_time_window(tr, config.wave_type[0], npts)
-            _recompute_time_window(tr, 'N', npts, from_start=True)
-    return signal_st, noise_st
+            _recompute_time_window(tr, config.wave_type[0], npts, keep='start')
+            _recompute_time_window(tr, 'N', npts, keep='end')
 
 
 def _build_signal_and_noise_spectral_streams(config, signal_st, noise_st):
@@ -624,6 +621,19 @@ def _build_signal_and_noise_spectral_streams(config, signal_st, noise_st):
     return spec_st, specnoise_st
 
 
+def _zero_pad(config, trace):
+    """Zero-pad trace to spectral_win_length"""
+    spec_win_len = config.spectral_win_length
+    if spec_win_len is None:
+        return
+    wtype = config.wave_type[0]
+    if trace.stats.type == 'signal':
+        t1 = trace.stats.arrivals[f'{wtype}1'][1]
+    elif trace.stats.type == 'noise':
+        t1 = trace.stats.arrivals['N1'][1]
+    trace.trim(starttime=t1, endtime=t1+spec_win_len, pad=True, fill_value=0)
+
+
 def build_spectra(config, st):
     """
     Build spectra and the ``spec_st`` object.
@@ -634,6 +644,9 @@ def build_spectra(config, st):
     """
     logger.info('Building spectra...')
     signal_st, noise_st = _build_signal_and_noise_streams(config, st)
+    _trim_components(config, signal_st, noise_st, st)
+    for trace in signal_st + noise_st:
+        _zero_pad(config, trace)
     spec_st, specnoise_st = \
         _build_signal_and_noise_spectral_streams(config, signal_st, noise_st)
     weight_st = _build_weight_spectral_stream(config, spec_st, specnoise_st)

sourcespec/ssp_process_traces.py

@@ -21,8 +21,9 @@
 from obspy.core import Stream
 from obspy.core.util import AttribDict
 from sourcespec.ssp_setup import ssp_exit
-from sourcespec.ssp_util import remove_instr_response, hypo_dist
-from sourcespec.ssp_wave_arrival import add_arrivals_to_trace
+from sourcespec.ssp_util import (
+    remove_instr_response, station_to_event_position)
+from sourcespec.ssp_wave_arrival import add_arrival_to_trace
 from sourcespec.clipping_detection import clipping_score, clipping_peaks
 logger = logging.getLogger(__name__.split('.')[-1])
 
@@ -210,35 +211,92 @@ def _process_trace(config, trace):
     return trace_process
 
 
-def _merge_stream(config, st):
+def _add_station_to_event_position(config, trace):
     """
-    Check for gaps and overlaps; remove mean; merge stream.
+    Add to ``trace.stats`` station-to-event distance (hypocentral and
+    epicentral), great-circle distance, azimuth and backazimuth.
     """
-    traceid = st[0].id
-    # First, compute gap/overlap statistics for the whole trace.
-    gaps_olaps = st.get_gaps()
-    gaps = [g for g in gaps_olaps if g[6] >= 0]
-    overlaps = [g for g in gaps_olaps if g[6] < 0]
-    gap_duration = sum(g[6] for g in gaps)
-    if gap_duration > 0:
-        logger.info(
-            f'{traceid}: trace has {gap_duration:.3f} seconds of gaps.')
-        gap_max = config.gap_max
-        if gap_max is not None and gap_duration > gap_max:
-            raise RuntimeError(
-                f'{traceid}: Gap duration larger than gap_max '
-                f'({gap_max:.1f} s): skipping trace')
-    overlap_duration = -1 * sum(g[6] for g in overlaps)
-    if overlap_duration > 0:
-        logger.info(
-            f'{traceid}: trace has {overlap_duration:.3f} seconds of '
-            'overlaps.')
-        overlap_max = config.overlap_max
-        if overlap_max is not None and overlap_duration > overlap_max:
+    try:
+        station_to_event_position(trace)
+    except Exception as e:
+        raise RuntimeError(
+            f'{trace.id}: Unable to compute hypocentral distance: '
+            'skipping trace'
+        ) from e
+    if (
+        config.max_epi_dist is not None and
+        trace.stats.epi_dist > config.max_epi_dist
+    ):
+        raise RuntimeError(
+            f'{trace.id}: Epicentral distance '
+            f'({trace.stats.epi_dist:.1f} km) '
+            f'larger than max_epi_dist ({config.max_epi_dist:.1f} km): '
+            'skipping trace')
+
+
+def _add_arrivals(config, trace):
+    """Add to trace P and S arrival times, travel times and angles."""
+    for phase in 'P', 'S':
+        try:
+            add_arrival_to_trace(trace, phase, config)
+        except Exception as e:
+            for line in str(e).splitlines():
+                logger.warning(line)
             raise RuntimeError(
-                f'{traceid}: overlap duration larger than overlap_max '
-                f'({overlap_max:.1f} s): skipping trace')
-    # Then, compute the same statisics for the signal window.
+                f'{trace.id}: Unable to get {phase} arrival time: '
+                'skipping trace'
+            ) from e
+
+
+def _define_signal_and_noise_windows(config, trace):
+    """Define signal and noise windows for spectral analysis."""
+    p_arrival_time = trace.stats.arrivals['P'][1]
+    if config.wave_type[0] == 'P' and p_arrival_time < trace.stats.starttime:
+        raise RuntimeError(f'{trace.id}: P-window incomplete: skipping trace')
+    s_arrival_time = trace.stats.arrivals['S'][1]
+    if config.wave_type[0] == 'S' and s_arrival_time < trace.stats.starttime:
+        raise RuntimeError(f'{trace.id}: S-window incomplete: skipping trace')
+    # Signal window for spectral analysis (S phase)
+    s_minus_p = s_arrival_time-p_arrival_time
+    s_pre_time = config.signal_pre_time
+    if s_minus_p/2 < s_pre_time:
+        # use (Ts-Tp)/2 if it is smaller than signal_pre_time
+        # (for short-distance records with short S-P interval)
+        s_pre_time = s_minus_p/2
+        logger.warning(
+            f'{trace.id}: signal_pre_time is larger than (Ts-Tp)/2. '
+            'Using (Ts-Tp)/2 instead')
+    t1 = s_arrival_time - s_pre_time
+    t1 = max(trace.stats.starttime, t1)
+    t2 = t1 + config.win_length
+    trace.stats.arrivals['S1'] = ('S1', t1)
+    trace.stats.arrivals['S2'] = ('S2', t2)
+    # Signal window for spectral analysis (P phase)
+    t1 = p_arrival_time - config.signal_pre_time
+    t1 = max(trace.stats.starttime, t1)
+    t2 = t1 + min(config.win_length, s_minus_p)
+    trace.stats.arrivals['P1'] = ('P1', t1)
+    trace.stats.arrivals['P2'] = ('P2', t2)
+    # Noise window for spectral analysis
+    t1 = max(trace.stats.starttime, p_arrival_time - config.noise_pre_time)
+    t2 = t1 + config.win_length
+    if t2 >= p_arrival_time:
+        logger.warning(f'{trace.id}: noise window ends after P-wave arrival')
+        # Note: maybe we should also take into account signal_pre_time here
+        t2 = p_arrival_time
+        t1 = min(t1, t2)
+    trace.stats.arrivals['N1'] = ('N1', t1)
+    trace.stats.arrivals['N2'] = ('N2', t2)
+
+
+def _check_signal_window(config, st):
+    """
+    Check if the signal window has sufficient amount of signal
+    (i.e., not too many gaps).
+
+    This is done on the stream, before merging.
+    """
+    traceid = st[0].id
     st_cut = st.copy()
     if config.wave_type[0] == 'S':
         t1 = st[0].stats.arrivals['S1'][1]
@@ -266,6 +324,36 @@ def _merge_stream(config, st):
         logger.info(
             f'{traceid}: signal window has {overlap_duration:.3f} seconds '
             'of overlaps.')
+
+
+def _merge_stream(config, st):
+    """
+    Check for gaps and overlaps; remove mean; merge stream.
+    """
+    traceid = st[0].id
+    # First, compute gap/overlap statistics for the whole trace.
+    gaps_olaps = st.get_gaps()
+    gaps = [g for g in gaps_olaps if g[6] >= 0]
+    overlaps = [g for g in gaps_olaps if g[6] < 0]
+    gap_duration = sum(g[6] for g in gaps)
+    if gap_duration > 0:
+        logger.info(
+            f'{traceid}: trace has {gap_duration:.3f} seconds of gaps.')
+        gap_max = config.gap_max
+        if gap_max is not None and gap_duration > gap_max:
+            raise RuntimeError(
+                f'{traceid}: Gap duration larger than gap_max '
+                f'({gap_max:.1f} s): skipping trace')
+    overlap_duration = -1 * sum(g[6] for g in overlaps)
+    if overlap_duration > 0:
+        logger.info(
+            f'{traceid}: trace has {overlap_duration:.3f} seconds of '
+            'overlaps.')
+        overlap_max = config.overlap_max
+        if overlap_max is not None and overlap_duration > overlap_max:
+            raise RuntimeError(
+                f'{traceid}: overlap duration larger than overlap_max '
+                f'({overlap_max:.1f} s): skipping trace')
     # Finally, demean (only if trace has not be already preprocessed)
     if config.trace_units == 'auto':
         # Since the count value is generally huge, we need to demean twice
@@ -284,74 +372,6 @@ def _merge_stream(config, st):
     return st[0]
 
 
-def _add_hypo_dist_and_arrivals(config, st):
-    for trace in st:
-        if hypo_dist(trace) is None:
-            raise RuntimeError(
-                f'{trace.id}: Unable to compute hypocentral distance: '
-                'skipping trace')
-        if config.max_epi_dist is not None and \
-                trace.stats.epi_dist > config.max_epi_dist:
-            raise RuntimeError(
-                f'{trace.id}: Epicentral distance '
-                f'({trace.stats.epi_dist:.1f} km) '
-                f'larger than max_epi_dist ({config.max_epi_dist:.1f} km): '
-                'skipping trace')
-        add_arrivals_to_trace(trace, config)
-        try:
-            p_arrival_time = trace.stats.arrivals['P'][1]
-        except KeyError as e:
-            raise RuntimeError(
-                f'{trace.id}: Unable to get P arrival time: skipping trace'
-            ) from e
-        if (config.wave_type[0] == 'P' and
-                p_arrival_time < trace.stats.starttime):
-            raise RuntimeError(
-                f'{trace.id}: P-window incomplete: skipping trace')
-        try:
-            s_arrival_time = trace.stats.arrivals['S'][1]
-        except KeyError as e:
-            raise RuntimeError(
-                f'{trace.id}: Unable to get S arrival time: skipping trace'
-            ) from e
-        if (config.wave_type[0] == 'S' and
-                s_arrival_time < trace.stats.starttime):
-            raise RuntimeError(
-                f'{trace.id}: S-window incomplete: skipping trace')
-        # Signal window for spectral analysis (S phase)
-        s_minus_p = s_arrival_time-p_arrival_time
-        s_pre_time = config.signal_pre_time
-        if s_minus_p/2 < s_pre_time:
-            # use (Ts-Tp)/2 if it is smaller than signal_pre_time
-            # (for short-distance records with short S-P interval)
-            s_pre_time = s_minus_p/2
-            logger.warning(
-                f'{trace.id}: signal_pre_time is larger than (Ts-Tp)/2. '
-                'Using (Ts-Tp)/2 instead')
-        t1 = s_arrival_time - s_pre_time
-        t1 = max(trace.stats.starttime, t1)
-        t2 = t1 + config.win_length
-        trace.stats.arrivals['S1'] = ('S1', t1)
-        trace.stats.arrivals['S2'] = ('S2', t2)
-        # Signal window for spectral analysis (P phase)
-        t1 = p_arrival_time - config.signal_pre_time
-        t1 = max(trace.stats.starttime, t1)
-        t2 = t1 + min(config.win_length, s_minus_p)
-        trace.stats.arrivals['P1'] = ('P1', t1)
-        trace.stats.arrivals['P2'] = ('P2', t2)
-        # Noise window for spectral analysis
-        t1 = max(trace.stats.starttime, p_arrival_time - config.noise_pre_time)
-        t2 = t1 + config.win_length
-        if t2 >= p_arrival_time:
-            logger.warning(
-                f'{trace.id}: noise window ends after P-wave arrival')
-            # Note: maybe we should also take into account signal_pre_time here
-            t2 = p_arrival_time
-            t1 = min(t1, t2)
-        trace.stats.arrivals['N1'] = ('N1', t1)
-        trace.stats.arrivals['N2'] = ('N2', t2)
-
-
 def _skip_ignored(config, id):
     """Skip traces ignored from config."""
     network, station, location, channel = id.split('.')
@@ -382,11 +402,15 @@ def process_traces(config, st):
     logger.info('Processing traces...')
     out_st = Stream()
     for id in sorted({tr.id for tr in st}):
-        # We still use a stream, since the trace can have gaps or overlaps
-        st_sel = st.select(id=id)
         try:
             _skip_ignored(config, id)
-            _add_hypo_dist_and_arrivals(config, st_sel)
+            # We still use a stream, since the trace can have gaps or overlaps
+            st_sel = st.select(id=id)
+            for _trace in st_sel:
+                _add_station_to_event_position(config, _trace)
+                _add_arrivals(config, _trace)
+                _define_signal_and_noise_windows(config, _trace)
+            _check_signal_window(config, st_sel)
             trace = _merge_stream(config, st_sel)
             trace.stats.ignore = False
             trace_process = _process_trace(config, trace)