Merging Wenjie's edits from years ago

pyflex/__init__.py

@@ -38,7 +38,7 @@ class PyflexWarning(UserWarning):
 
 # Setup the logger.
 logger = logging.getLogger("pyflex")
-logger.setLevel(logging.WARNING)
+logger.setLevel(logging.DEBUG)
 # Prevent propagating to higher loggers.
 logger.propagate = 0
 # Console log handler.

pyflex/config.py

@@ -20,17 +20,21 @@ class Config(object):
     def __init__(self, min_period, max_period, stalta_waterlevel=0.07,
                  tshift_acceptance_level=10.0, tshift_reference=0.0,
                  dlna_acceptance_level=1.3, dlna_reference=0.0,
-                 cc_acceptance_level=0.7, s2n_limit=1.5, earth_model="ak135",
+                 cc_acceptance_level=0.7, earth_model="ak135",
+                 s2n_limit=1.5, s2n_limit_energy=1.5,
                  min_surface_wave_velocity=3.0,
-                 max_time_before_first_arrival=50.0,
+                 max_surface_wave_velocity=4.2,
+                 max_time_before_first_arrival=None,
+                 max_time_after_last_arrival=None,
                  c_0=1.0, c_1=1.5, c_2=0.0,
                  c_3a=4.0, c_3b=2.5, c_4a=2.0, c_4b=6.0,
                  check_global_data_quality=False, snr_integrate_base=3.5,
                  snr_max_base=3.0, noise_start_index=0, noise_end_index=None,
-                 signal_start_index=None, signal_end_index=-1,
+                 signal_start_index=None, signal_end_index=None,
                  window_weight_fct=None,
                  window_signal_to_noise_type="amplitude",
-                 resolution_strategy="interval_scheduling"):
+                 resolution_strategy="interval_scheduling",
+                 selection_mode=None):
         """
         Central configuration object for Pyflex.
 
@@ -105,28 +109,51 @@ def __init__(self, min_period, max_period, stalta_waterlevel=0.07,
             same number of samples as the data.
         :type cc_acceptance_level: float or :class:`numpy.ndarray`
 
-        :param s2n_limit: Limit of the signal to noise ratio per window. If
-            the maximum amplitude of the window over the maximum amplitude
-            of the global noise of the waveforms is smaller than this
+        :param s2n_limit: Limit of the amplitude signal-to-noise
+            ratio per window. If the maximum amplitude of the window over
+            the maximum amplitude of the global noise of the waveforms is
+            smaller than this window, then it will be rejected. Can be
+            either a single value or an array with the same number of
+            samples as the data.
+        :type s2n_limit: float or :class:`numpy.ndarray`
+
+         :param s2n_limit_energy: Limit of the energy signal-to-noise ratio
+            per window. If the average energy of the window over the average
+            energy of the global noise of the waveforms is smaller than this
             window, then it will be rejected. Can be either a single value
             or an array with the same number of samples as the data.
-        :type s2n_limit: float or :class:`numpy.ndarray`
+        :type s2n_limit_energy: float or :class:`numpy.ndarray`
 
         :param earth_model: The earth model used for the traveltime
             calculations. Either ``"ak135"`` or ``"iasp91"``.
         :type earth_model: str
 
         :param min_surface_wave_velocity: The minimum surface wave velocity
-            in km/s. All windows containing data later then this velocity
-            will be rejected. Only used if station and event information is
-            available.
+            in km/s. The two values, min_surface_wave_velocity and
+            max_surface_wave_velocity are used to calculate surface wave time
+            region. All windows containing data earlier than the min velocity
+            and later than the max velocity will be treated as surface waves.
+            Only used if station and event information is available.
         :type min_surface_wave_velocity: float
 
+        :param max_surface_wave_velocity: The maxium surface wave velocity
+            in km/s. The two values, min_surface_wave_velocity and
+            max_surface_wave_velocity are used to calculate surface wave time
+            region. All windows containing data earlier than the min velocity
+            and later than the max velocity will be treated as surface waves.
+            Only used if station and event information is available.
+        :type max_surface_wave_velocity: float
+
         :param max_time_before_first_arrival: This is the minimum starttime
             of any window in seconds before the first arrival. No windows will
             have a starttime smaller than this.
         :type max_time_before_first_arrival: float
 
+        :param max_time_after_last_arrival: This is the max endtime
+            of any window in seconds after the last arrival. No windows will
+            have a endtime larger then this.
+        :type max_time_after_last_arrival: float
+
         :param c_0: Fine tuning constant for the rejection of windows based
             on the height of internal minima. Any windows with internal
             minima lower then this value times the STA/LTA water level at
@@ -198,12 +225,19 @@ def __init__(self, min_period, max_period, stalta_waterlevel=0.07,
         :type window_weight_fct: function
 
         :param window_signal_to_noise_type: The type of signal to noise
-            ratio used to reject windows. If ``"amplitude"``, then the
-            largest amplitude before the arrival is the noise amplitude and
-            the largest amplitude in the window is the signal amplitude. If
-            ``"energy"`` the time normalized energy is used in both cases.
-            The later one is a bit more stable when having random wiggles
-            before the first arrival.
+            ratio used to reject windows
+            Possiblities are:
+            * ``"amplitude"``: then the largest amplitude before the arrival
+                is the noise amplitude and the largest amplitude in the
+                window is the signal amplitude. If the value is smaller than
+                "s2n_limit" then the windows will be kept.
+            * ``"energy"``: the time normalized energy is used. This one is
+                a bit more stable when having random wiggles before the
+                first arrival. If the value is smaller than
+                "s2n_limit_energy", then the window will be kept.
+            * ``"amplitude_and_energy"``: then both checks(amplitude and
+                energy) are applied to the trace. But please remember to
+                assign values for "s2n_limit" and "s2n_limit_energy".
         :type window_signal_to_noise_type: str
 
         :param resolution_strategy: Strategy used to resolve overlaps.
@@ -212,6 +246,23 @@ def __init__(self, min_period, max_period, stalta_waterlevel=0.07,
             non-overlapping windows. Merging will simply merge overlapping
             windows.
         :type resolution_strategy: str
+
+        :param selection_mode: Strategy to select different types of phases.
+            For most cases, event and station information are required to
+            calculate the arrival time information.
+            Possibilities are:
+            * ``"all_waves"``: all windows after first arrival
+            * ``"body_waves"``: windows in the body wave region(after
+                first arrival and before surface wave arrival).
+            * ``"surface_waves"``: windows in surface wave region(velocity
+                between min_surface_wave_velocity and
+                max_surface_wave_velocity).
+            * ``"body_and_surface_waves"``: windows inside body and surface
+                wave regions.
+            * ``mantle_waves``: pyflex will only accept windows after the
+                surface wave region.
+            * ``custom``: all windows will pass the check(nothing rejected)
+        :type selection_mode: str
         """
         self.min_period = min_period
         self.max_period = max_period
@@ -223,13 +274,24 @@ def __init__(self, min_period, max_period, stalta_waterlevel=0.07,
         self.dlna_reference = dlna_reference
         self.cc_acceptance_level = cc_acceptance_level
         self.s2n_limit = s2n_limit
+        self.s2n_limit_energy = s2n_limit_energy
 
         if earth_model.lower() not in ("ak135", "iasp91"):
             raise PyflexError("Earth model must either be 'ak135' or "
                               "'iasp91'.")
         self.earth_model = earth_model.lower()
         self.min_surface_wave_velocity = min_surface_wave_velocity
-        self.max_time_before_first_arrival = max_time_before_first_arrival
+        self.max_surface_wave_velocity = max_surface_wave_velocity
+
+        if max_time_before_first_arrival is None:
+            self.max_time_before_first_arrival = 2 * min_period
+        else:
+            self.max_time_before_first_arrival = max_time_before_first_arrival
+
+        if max_time_after_last_arrival is None:
+            self.max_time_after_last_arrival = max_period
+        else:
+            self.max_time_after_last_arrival = max_time_after_last_arrival
 
         self.c_0 = c_0
         self.c_1 = c_1
@@ -242,6 +304,7 @@ def __init__(self, min_period, max_period, stalta_waterlevel=0.07,
         self.check_global_data_quality = check_global_data_quality
         self.snr_integrate_base = snr_integrate_base
         self.snr_max_base = snr_max_base
+
         self.noise_start_index = noise_start_index
         self.noise_end_index = noise_end_index
         self.signal_start_index = signal_start_index
@@ -250,9 +313,10 @@ def __init__(self, min_period, max_period, stalta_waterlevel=0.07,
         self.window_weight_fct = window_weight_fct
 
         snr_type = window_signal_to_noise_type.lower()
-        if snr_type not in ["amplitude", "energy"]:
-            raise PyflexError("The window signal to noise type must be either"
-                              "'amplitude' or 'energy'.")
+        if snr_type not in ["amplitude", "energy", "amplitude_and_energy"]:
+            raise PyflexError(
+                    "The window signal to noise type must be"
+                    "'amplitude', 'energy' or 'amplitude_and_energy'.")
         self.window_signal_to_noise_type = snr_type
 
         if resolution_strategy.lower() not in ["interval_scheduling", "merge"]:
@@ -261,14 +325,34 @@ def __init__(self, min_period, max_period, stalta_waterlevel=0.07,
                 "'interval_scheduling' or 'merge'.")
         self.resolution_strategy = resolution_strategy.lower()
 
+        # Add a specific selection mode for different types of phases
+        if selection_mode is not None:
+
+            selection_mode = selection_mode.strip()
+            if selection_mode.split(":")[0].strip().lower() not in [
+                    "all_waves", "body_and_surface_waves", "body_waves",
+                    "surface_waves", "mantle_waves", "custom", "phase_list",
+                    "body_and_mantle_waves"]:
+                raise ValueError(
+                        "Invalide selection_mode({}). Choose: 1) all_waves;"
+                        "2) body_and_surface_waves; 3) body_waves; "
+                        "4) surface_waves; 5) mantle_waves; "
+                        "6) custom; 7) phase_list;".format(selection_mode))
+
+        self.selection_mode = selection_mode
+
+
+
+
     def _convert_to_array(self, npts):
         """
         Internally converts the acceptance and water levels to arrays. Not
         called during initialization as the array length is not yet known.
         """
         attributes = ("stalta_waterlevel", "tshift_acceptance_level",
                       "dlna_acceptance_level", "cc_acceptance_level",
-                      "s2n_limit")
+                      "s2n_limit", "s2n_limit_energy")
+
         for name in attributes:
             attr = getattr(self, name)
 
@@ -278,6 +362,47 @@ def _convert_to_array(self, npts):
                         "Config value '%s' does not have the same number of "
                         "samples as the waveforms." % name)
                 setattr(self, name, np.array(attr))
-                continue
+            else:
+                setattr(self, name, attr * np.ones(npts))
 
-            setattr(self, name, attr * np.ones(npts))
+    def _convert_negative_index(self, npts):
+        """
+        Convert negative index into positive. So we can compare and check
+        """
+        def add_npts(_idx, _npts):
+            if _idx is None:
+                return None
+            _idx = int(_idx)
+            if _idx < 0:
+                return _idx + _npts
+            else:
+                return _idx
+
+        self.noise_start_index = add_npts(self.noise_start_index, npts)
+        self.noise_end_index = add_npts(self.noise_end_index, npts)
+        self.signal_start_index = add_npts(self.signal_start_index, npts)
+        self.signal_end_index = add_npts(self.signal_end_index, npts)
+
+        if self.noise_start_index is not None and \
+                self.noise_end_index is not None:
+            if self.noise_start_index >= self.noise_end_index:
+                raise ValueError(
+                    "noise_start_index is larger than "
+                    "noise_end_idx: %d > %d" % (self.noise_start_index,
+                                                self.noise_end_index))
+
+        if self.signal_start_index is not None and \
+                self.signal_end_index is not None:
+            if self.signal_start_index >= self.signal_end_index:
+                raise ValueError("singal_start_index is larger than "
+                                 "signal_end_index: %d > %d"
+                                 % (self.signal_start_index,
+                                    self.signal_end_index))
+
+        if self.noise_end_index is not None and \
+                self.signal_start_index is not None:
+            if self.noise_end_index > self.signal_start_index:
+                raise ValueError("noise_end_index is larger than "
+                                 "signal_start_index: %d > %d"
+                                 % (self.noise_end_index,
+                                    self.signal_start_index))

pyflex/interval_scheduling.py

@@ -59,18 +59,18 @@ def schedule_weighted_intervals(I):
         OPT[j] = max(I[j].weight + OPT[p[j]], OPT[j - 1])
 
     # given OPT and p, find actual solution intervals in O(n)
-    O = []
+    res = []
 
     def compute_solution(j):
         if j >= 0:  # will halt on OPT[-1]
             if I[j].weight + OPT[p[j]] > OPT[j - 1]:
-                O.append(I[j])
+                res.append(I[j])
                 compute_solution(p[j])
             else:
                 compute_solution(j - 1)
     compute_solution(len(I) - 1)
 
-    # resort, as our O is in reverse order (OPTIONAL)
-    O.sort(key=lambda x: x.right)
+    # resort, as our res is in reverse order (OPTIONAL)
+    res.sort(key=lambda x: x.right)
 
-    return O
+    return res