new subcommand: qc-polarization

apps/grond

@@ -5,7 +5,7 @@ import sys
 import logging
 from optparse import OptionParser
 
-from pyrocko import util
+from pyrocko import util, marker
 from pyrocko.gf import Range
 
 import grond
@@ -327,7 +327,6 @@ def command_harvest(args):
 
 
 def command_plot(args):
-    from grond import plot
 
     def setup(parser):
         parser.add_option(
@@ -342,7 +341,17 @@ def command_plot(args):
             '--dpi', '--dpi', dest='dpi', type=float, default=120.,
             help='DPI setting for raster formats (default=120)')
 
-    plotnames_avail = plot.available_plotnames()
+    # plotnames_avail = plot.available_plotnames()
+    # made explicit to avoid import of pyplot before backend can be chosen
+    plotnames_avail = [
+        'bootstrap',
+        'sequence',
+        'contributions',
+        'jointpar',
+        'hudson',
+        'fits',
+        'solution',
+        'location']
 
     details = '''Available <plotnames> are: %s, or "all". Multiple plots are
 selected by specifying a comma-separated list.''' % (
@@ -361,6 +370,12 @@ selected by specifying a comma-separated list.''' % (
     formats = options.formats.split(',')
     dirname = args[1]
 
+    if options.save:
+        import matplotlib
+        matplotlib.use('Agg')
+
+    from grond import plot
+
     try:
         plot.plot_result(
             dirname, plotnames,
@@ -426,15 +441,53 @@ def command_export(args):
 
 def command_qc_polarization(args):
 
-    import grond.qc
-
     def setup(parser):
-        pass
+        parser.add_option(
+            '--time-factor', dest='time_factor', type=float,
+            metavar='NUMBER',
+            default=0.5,
+            help='set duration to extract before and after synthetic P phase '
+                 'arrival, relative to 1/fmin. fmin is taken from the '
+                 'selected target group in the config file (default=%default)')
+        parser.add_option(
+            '--distance-min', dest='distance_min', type=float,
+            metavar='NUMBER',
+            help='minimum event-station distance [m]')
+        parser.add_option(
+            '--distance-max', dest='distance_max', type=float,
+            metavar='NUMBER',
+            help='maximum event-station distance [m]')
+        parser.add_option(
+            '--depth-min', dest='depth_min', type=float,
+            metavar='NUMBER',
+            help='minimum station depth [m]')
+        parser.add_option(
+            '--depth-max', dest='depth_max', type=float,
+            metavar='NUMBER',
+            help='maximum station depth [m]')
+        parser.add_option(
+            '--picks', dest='picks_filename',
+            metavar='FILENAME',
+            help='add file with P picks in Snuffler marker format')
+        parser.add_option(
+            '--save', dest='output_filename',
+            metavar='FILENAME.FORMAT',
+            help='save output to file FILENAME.FORMAT')
+        parser.add_option(
+            '--dpi', dest='output_dpi', type=float, default=120.,
+            metavar='NUMBER',
+            help='DPI setting for raster formats (default=120)')
 
     parser, options, args = cl_parse('qc-polarization', args, setup)
     if len(args) != 3:
         help_and_die(parser, 'missing arguments')
 
+    if options.output_filename:
+        import matplotlib
+        matplotlib.use('Agg')
+
+    import grond.qc
+
     config_path, event_name, target_config_name = args
 
     config = grond.read_config(config_path)
@@ -443,6 +496,24 @@ def command_qc_polarization(args):
 
     engine = config.engine_config.get_engine()
 
+    nsl_to_time = None
+    if options.picks_filename:
+        markers = marker.load_markers(options.picks_filename)
+        marker.associate_phases_to_events(markers)
+
+        nsl_to_time = {}
+        for m in markers:
+            if isinstance(m, marker.PhaseMarker):
+                ev = m.get_event()
+                if ev is not None and ev.name == event_name:
+                    nsl_to_time[m.one_nslc()[:3]] = m.tmin
+
+        if not nsl_to_time:
+            help_and_die(
+                parser,
+                'no markers associated with event "%s" found in file "%s"' % (
+                    event_name, options.picks_filename))
+
     target_config_names_avail = []
     for target_config in config.target_configs:
         name = '%s.%s' % (target_config.super_group, target_config.group or '')
@@ -456,7 +527,15 @@ def command_qc_polarization(args):
             timing = '{cake:P|cake:p|cake:P\\|cake:p\\}'
 
             grond.qc.polarization(
-                ds, store, timing, fmin=fmin, fmax=fmax, ffactor=ffactor)
+                ds, store, timing, fmin=fmin, fmax=fmax, ffactor=ffactor,
+                time_factor=options.time_factor,
+                distance_min=options.distance_min,
+                distance_max=options.distance_max,
+                depth_min=options.depth_min,
+                depth_max=options.depth_max,
+                nsl_to_time=nsl_to_time,
+                output_filename=options.output_filename,
+                output_dpi=options.output_dpi)
 
             return
 

src/qc.py

@@ -6,7 +6,7 @@ import numpy as num
 from matplotlib import pyplot as plt
 from mpl_toolkits.axes_grid1 import ImageGrid
 
-from pyrocko import gf, trace, orthodrome as od, plot
+from pyrocko import gf, orthodrome as od, plot
 from grond import dataset
 
 km = 1000.
@@ -14,7 +14,37 @@ km = 1000.
 logger = logging.getLogger('grond.qc')
 
 
-def polarization(ds, store, timing, fmin, fmax, ffactor=1.5, tfactor=2.):
+def darken(c):
+    return (c[0]*0.5, c[1]*0.5, c[2]*0.5)
+
+
+def plot_color_line(axes, x, y, t, color, tmin, tmax):
+    from matplotlib.colors import LinearSegmentedColormap
+    from matplotlib.collections import LineCollection
+
+    cmap = LinearSegmentedColormap.from_list(
+        'noname', [color, (1., 1., 1.)], 100)
+
+    points = num.array([x, y], dtype=num.float).T.reshape(-1, 1, 2)
+    segments = num.concatenate([points[:-1], points[1:]], axis=1)
+    lc = LineCollection(segments, cmap=cmap, norm=plt.Normalize(tmin, tmax))
+    lc.set_array(t)
+    axes.add_collection(lc, autolim=True)
+
+
+def polarization(
+        ds, store, timing, fmin, fmax, ffactor,
+        time_factor=2.,
+        distance_min=None,
+        distance_max=None,
+        depth_min=None,
+        depth_max=None,
+        size_factor=0.05,
+        nsl_to_time=None,
+        output_filename=None,
+        output_format=None,
+        output_dpi=None):
+
     event = ds.get_event()
     stations = ds.get_stations()
 
@@ -26,14 +56,37 @@ def polarization(ds, store, timing, fmin, fmax, ffactor=1.5, tfactor=2.):
     trs = []
     for station in stations:
 
-        for component in 'ZNE':
+        nsl = station.nsl()
+
+        dist = source.distance_to(station)
+
+        if distance_min is not None and dist < distance_min:
+            continue
+
+        if distance_max is not None and distance_max < dist:
+            continue
+
+        if depth_min is not None and station.depth < depth_min:
+            continue
+
+        if depth_max is not None and depth_max < station.depth:
+            continue
+
+        if nsl_to_time is None:
+            tp = event.time + store.t(timing, source, station)
+
+        else:
+            if nsl not in nsl_to_time:
+                continue
+
+            tp = nsl_to_time[nsl]
 
-            ttp = store.t(timing, source, station)
+        for component in 'ZNE':
 
-            tmin = event.time + ttp - tfactor / fmin
-            tmax = event.time + ttp + tfactor / fmin
+            tmin = tp - time_factor / fmin
+            tmax = tp + time_factor / fmin
 
-            nslc = station.nsl() + (component,)
+            nslc = nsl + (component,)
 
             freqlimits = [
                 fmin / ffactor,
@@ -43,9 +96,6 @@ def polarization(ds, store, timing, fmin, fmax, ffactor=1.5, tfactor=2.):
 
             tfade = 1.0 / (fmin / ffactor)
 
-            tlong = 5.0 / fmin
-            tshort = 0.2 * tlong
-
             try:
                 trs_projected, trs_restituted, trs_raw = \
                     ds.get_waveform(
@@ -56,9 +106,6 @@ def polarization(ds, store, timing, fmin, fmax, ffactor=1.5, tfactor=2.):
                         freqlimits=freqlimits,
                         debug=True)
 
-                # for tr in trs_projected:
-                #     tr.sta_lta_right(tshort, tlong)
-
                 trs.extend(trs_projected)
 
             except dataset.NotFound, e:
@@ -69,9 +116,9 @@ def polarization(ds, store, timing, fmin, fmax, ffactor=1.5, tfactor=2.):
     for tr in trs:
         d_trs[tr.nslc_id[:3]][tr.nslc_id[3]] = tr
 
-    fontsize = 11.
+    fontsize = 10.
     plot.mpl_init(fontsize=fontsize)
-    fig = plt.figure(figsize=(8., 8.))
+    fig = plt.figure(figsize=plot.mpl_papersize('a4', 'landscape'))
     plot.mpl_margins(fig, w=7., h=6., units=fontsize)
 
     grid = ImageGrid(
@@ -81,8 +128,6 @@ def polarization(ds, store, timing, fmin, fmax, ffactor=1.5, tfactor=2.):
         label_mode='L',
         aspect=True)
 
-    print dir(grid)
-
     axes_en = grid[0]
     axes_en.set_ylabel('Northing [km]')
 
@@ -101,13 +146,37 @@ def polarization(ds, store, timing, fmin, fmax, ffactor=1.5, tfactor=2.):
 
     grid[3].set_axis_off()
 
-    factor = 5000.
-
+    locations = []
     for nsl in sorted(d_trs.keys()):
+        station = nsl_to_station[nsl]
+        n, e = od.latlon_to_ne(
+            event.lat, event.lon, station.lat, station.lon)
+
+        locations.append((n, e))
+
+    ns, es = num.array(locations, dtype=num.float).T
+
+    n_min = num.min(ns)
+    n_max = num.max(ns)
+    e_min = num.min(es)
+    e_max = num.max(es)
 
-        tr_e = d_trs[nsl].get('E', None)
-        tr_n = d_trs[nsl].get('N', None)
-        tr_z = d_trs[nsl].get('Z', None)
+    factor = max((n_max - n_min) * size_factor, (e_max - e_min) * size_factor)
+
+    fontsize_annot = fontsize * 0.7
+
+    data = {}
+    for insl, nsl in enumerate(sorted(d_trs.keys())):
+
+        color = plot.mpl_graph_color(insl)
+
+        try:
+            tr_e = d_trs[nsl]['E']
+            tr_n = d_trs[nsl]['N']
+            tr_z = d_trs[nsl]['Z']
+
+        except KeyError:
+            continue
 
         station = nsl_to_station[nsl]
 
@@ -116,26 +185,63 @@ def polarization(ds, store, timing, fmin, fmax, ffactor=1.5, tfactor=2.):
 
         d = station.depth
 
-        axes_en.plot(e/km, n/km, '^')
-        axes_dn.plot(d/km, n/km, '^')
-        axes_ed.plot(e/km, d/km, '^')
-
-        arr_e = tr_e.ydata if tr_e else None
-        arr_n = tr_n.ydata if tr_n else None
-        arr_z = tr_z.ydata if tr_z else None
-
-        amax = num.max(num.abs(num.sqrt(arr_e**2 + arr_n**2 + arr_z**2)))
-        axes_en.plot(
-            (e + arr_e/amax * factor)/km, (n + arr_n/amax * factor)/km,
-            color=plot.mpl_color('skyblue2'))
-        axes_dn.plot(
+        axes_en.annotate(
+            '.'.join(x for x in nsl if x),
+            xy=(e/km, n/km),
+            xycoords='data',
+            xytext=(fontsize_annot/3., fontsize_annot/3.),
+            textcoords='offset points',
+            verticalalignment='bottom',
+            horizontalalignment='left',
+            rotation=0.,
+            size=fontsize_annot)
+
+        axes_en.plot(e/km, n/km, '^', mfc=color, mec=darken(color))
+        axes_dn.plot(d/km, n/km, '^', mfc=color, mec=darken(color))
+        axes_ed.plot(e/km, d/km, '^', mfc=color, mec=darken(color))
+
+        arr_e = tr_e.ydata
+        arr_n = tr_n.ydata
+        arr_z = tr_z.ydata
+        arr_t = tr_z.get_xdata()
+
+        data[nsl] = (arr_e, arr_n, arr_z, arr_t, n, e, d, color)
+
+    amaxs = []
+    amax_hors = []
+    for nsl in sorted(data.keys()):
+        arr_e, arr_n, arr_z, arr_t, n, e, d, color = data[nsl]
+        amaxs.append(
+            num.max(num.abs(num.sqrt(arr_e**2 + arr_n**2 + arr_z**2))))
+        amax_hors.append(
+            num.max(num.abs(num.sqrt(arr_e**2 + arr_n**2))))
+
+    amax = num.median(amaxs)
+    amax_hor = num.median(amax_hors)
+
+    for nsl in sorted(data.keys()):
+        arr_e, arr_n, arr_z, arr_t, n, e, d, color = data[nsl]
+        tmin = arr_t.min()
+        tmax = arr_t.max()
+        plot_color_line(
+            axes_en,
+            (e + arr_e/amax_hor * factor)/km, (n + arr_n/amax_hor * factor)/km,
+            arr_t, color, tmin, tmax)
+        plot_color_line(
+            axes_dn,
             (d - arr_z/amax * factor)/km, (n + arr_n/amax * factor)/km,
-            color=plot.mpl_color('skyblue2'))
-        axes_ed.plot(
+            arr_t, color, tmin, tmax)
+        plot_color_line(
+            axes_ed,
             (e + arr_e/amax * factor)/km, (d - arr_z/amax * factor)/km,
-            color=plot.mpl_color('skyblue2'))
+            arr_t, color, tmin, tmax)
 
     axes_ed.invert_yaxis()
 
+    for axes in (axes_dn, axes_ed, axes_en):
+        axes.autoscale_view(tight=True)
+
+    if output_filename is None:
         plt.show()
-    # trace.snuffle(trs)
+    else:
+        fig.savefig(output_filename, format=output_format, dpi=output_dpi)