adds ability to plot misfit windows and read ASDFDataSets

pysep/recsec.py

@@ -95,6 +95,7 @@
 from glob import glob
 from datetime import datetime
 from matplotlib.ticker import MultipleLocator
+from matplotlib.patches import Rectangle
 from obspy import read, Stream
 from obspy.geodetics import (kilometers2degrees, gps2dist_azimuth)
 
@@ -126,7 +127,7 @@ class RecordSection:
     3) produces record section waveform figures.
     """
     def __init__(self, pysep_path=None, syn_path=None, 
-                 stations=None, source=None, st=None, st_syn=None, 
+                 stations=None, source=None, st=None, st_syn=None, windows=None,
                  sort_by="default", scale_by=None, time_shift_s=None, 
                  zero_pad_s=None, move_out=None, 
                  min_period_s=None, max_period_s=None,
@@ -197,6 +198,15 @@ def __init__(self, pysep_path=None, syn_path=None,
         :type st_syn: obspy.core.stream.Stream
         :param st_syn: Stream objects containing synthetic time series to be
             plotted on the record section. Must be same length as `st`
+        :type windows: dict of lists of tuples
+        :param windows: EXPERIMENTAL FEATURE -- plot misfit windows collected
+            by windowing algorithm like Pyflex. Essentially these are provided
+            as start and end times for each trace in the Stream. The dictionary
+            should be formated where keys are the Trace IDs of observed
+            waveforms, and values are lists of tuples, where each tuple 
+            represents a window start and end time in seconds. See 
+            `pysep.utils.io.read_asdfdataset` for code on how windows are 
+            structured
 
         .. note::
             Waveform plotting organization parameters
@@ -519,6 +529,9 @@ def __init__(self, pysep_path=None, syn_path=None,
         self.xlim_syn = []
         self.sorted_idx = []
 
+        # Experimental Feature
+        self.windows = windows
+
     def read_data(self, path, data_type, wildcard="*", source=None,
                   stations=None, srcfmt=None):
         """
@@ -611,7 +624,7 @@ def read_data(self, path, data_type, wildcard="*", source=None,
         else:
             raise NotImplementedError("`data_type` must be 'data' or 'syn'")
 
-        return st
+        return st                
 
     def check_parameters(self):
         """
@@ -1683,6 +1696,8 @@ def _plot_trace(self, idx, y_index, choice="st"):
             to plot the observed or synthetic waveforms
         """
         linewidth = self.kwargs.get("linewidth", .25)
+        window_alpha = self.kwargs.get("window_alpha", .1)
+        window_color = self.kwargs.get("window_color", "orange")
 
         # Used to differentiate the two types of streams for plotting diffs
         choices = ["st", "st_syn"]
@@ -1715,7 +1730,7 @@ def _plot_trace(self, idx, y_index, choice="st"):
         else:
             sign = 1
 
-        # Ampplitude scaling may change between observed and synthetic if e.g.,
+        # Amplitude scaling may change between observed and synthetic if e.g.,
         # we are doing trace-wise normalization
         if choice == "st":
             amplitude_scaling = self.amplitude_scaling
@@ -1726,6 +1741,15 @@ def _plot_trace(self, idx, y_index, choice="st"):
 
         y = sign * tr.data / amplitude_scaling[idx] + self.y_axis[y_index]
 
+        # Experimental: Plot windows over the record section if provided by User
+        if self.windows and tr.id in self.windows:
+            for window in self.windows[tr.id]:
+                rc = Rectangle(xy=(window[0] + tshift, y.min()), 
+                               width=window[1] - window[0], 
+                               height=y.max() - y.min(), alpha=window_alpha, 
+                               color=window_color, zorder=2)
+                self.ax.add_patch(rc)
+
         # Truncate waveforms to get figure scaling correct. Make the distinction
         # between data and synthetics which may have different start and end 
         # times natively

pysep/utils/io.py

@@ -14,7 +14,7 @@
 
 from pysep import logger
 from pysep.utils.mt import moment_magnitude, seismic_moment
-from pysep.utils.fmt import format_event_tag_legacy, channel_code
+from pysep.utils.fmt import channel_code
 from pysep.utils.cap_sac import append_sac_headers, append_sac_headers_cartesian
 
 
@@ -498,6 +498,67 @@ def read_event_file(fid):
     return list_out
 
 
+def read_asdfdataset(path, evaluation):
+    """
+    Read an ASDFDataSet created by a SeisFlows Pyaflowa inversion run. 
+    The dataset should contain observed and synthetic waveforms, and 
+    optionally contains misfit windows. Will return Streams with SAC headers
+
+    .. note::
+
+        This function makes assumptions about the PyASDF data structure which
+        is defined by the external package Pyatoa. If Pyatoa changes that
+        structure, this function will break.
+
+    :type path: str
+    :param path: Path to the ASDF dataset to be read
+    :type evaluation: str
+    :param evaluation: evaluation to take synthetics from. These are saved
+        following a format specified by Pyatoa, but usually follow the form
+        iteration/step_count, e.g., i01s00 gives iteration 1, step count 0.
+        Take a look at the waveform tags in `ASDFDataSet.waveforms[<station>]`
+        for tags following the 'synthetic_' prefix
+    """
+    # PySEP, by default will not require PyASDF to be installed
+    try:
+        from pyasdf import ASDFDataSet  # NOQA
+    except ImportError:
+        logger.critical("pyasdf is not installed. Please install pyasdf "
+                        "to read ASDF datasets")
+        return None, None
+
+    with ASDFDataSet(path) as ds:
+        event = ds.events[0]
+        st_out = Stream()
+        st_syn_out = Stream()
+        for station in ds.waveforms.list():
+            inv = ds.waveforms[station].StationXML
+            st = ds.waveforms[station].observed
+            st_syn = ds.waveforms[station][f"synthetic_{evaluation}"]
+
+            st_out += append_sac_headers(st, event, inv)
+            st_syn_out += append_sac_headers(st_syn, event, inv)
+
+        # Read windows from the dataset
+        windows = {}
+        if hasattr(ds.auxiliary_data, "MisfitWindows"):
+            iter_ = evaluation[:3]  # 'i01s00' -> 'i01'
+            step = evaluation[3:]
+            for station in ds.auxiliary_data.MisfitWindows[iter_][step].list():
+                parameters = ds.auxiliary_data.MisfitWindows[iter_][step][
+                    station].parameters
+                trace_id = parameters["channel_id"]
+                starttime = parameters["relative_starttime"]
+                endtime = parameters["relative_endtime"]    
+                # initialize empty window array
+                if trace_id not in windows:
+                    windows[trace_id] = []
+
+                windows[trace_id].append((starttime, endtime))
+
+    return st_out, st_syn_out, windows
+
+
 def write_sem(st, unit, path="./", time_offset=0):
     """
     Write an ObsPy Stream in the two-column ASCII format that Specfem uses