Refactoring ccftime plot to use objects

msnoise/plots/ccftime.py

@@ -17,13 +17,17 @@
 """
 # plot interferogram
 
+import datetime
+
 import matplotlib.dates as mdates
 import matplotlib.pyplot as plt
+import numpy as np
 from matplotlib.widgets import Cursor
-
-from obspy.signal.filter import envelope as obspy_envelope
 from obspy.signal.filter import bandpass
-from ..api import *
+from obspy.signal.filter import envelope as obspy_envelope
+
+from msnoise.api import connect, get_config, get_maxlag_samples, \
+    build_movstack_datelist, get_filters, get_results
 
 
 def main(sta1, sta2, filterid, components, mov_stack=1, ampli=5, seismic=False,
@@ -33,11 +37,13 @@ def main(sta1, sta2, filterid, components, mov_stack=1, ampli=5, seismic=False,
     samples = get_maxlag_samples(db)
     cc_sampling_rate = float(get_config(db, 'cc_sampling_rate'))
     start, end, datelist = build_movstack_datelist(db)
-    base = mdates.date2num(start) 
-    plt.figure(figsize=(12, 9))
+    base = mdates.date2num(start)
+
+    fig = plt.figure(figsize=(12, 9))
+
     sta1 = sta1.replace('.', '_')
     sta2 = sta2.replace('.', '_')
-    t = np.arange(samples)/cc_sampling_rate - maxlag
+    t = np.arange(samples) / cc_sampling_rate - maxlag
 
     if refilter:
         freqmin, freqmax = refilter.split(':')
@@ -46,49 +52,54 @@ def main(sta1, sta2, filterid, components, mov_stack=1, ampli=5, seismic=False,
 
     if sta2 >= sta1:
         pair = "%s:%s" % (sta1, sta2)
-        
+
         print("New Data for %s-%s-%i-%i" % (pair, components, filterid,
                                             mov_stack))
         nstack, stack_total = get_results(db, sta1, sta2, filterid, components,
                                           datelist, mov_stack, format="matrix")
-        ax = plt.subplot(111)
+        ax = fig.add_subplot(111)
         for i, line in enumerate(stack_total):
             if np.all(np.isnan(line)):
                 continue
+
             if refilter:
                 line = bandpass(line, freqmin, freqmax, cc_sampling_rate,
                                 zerophase=True)
             if envelope:
                 line = obspy_envelope(line)
+
             line /= line.max()
-            plt.plot(t, line * ampli + i + base, c='k')
+            ax.plot(t, line * ampli + i + base, c='k')
+
             if seismic:
                 y1 = np.ones(len(line)) * i
-                y2 = line*ampli + i + base
-                plt.fill_between(t, y1, y2, where=y2 >= y1, facecolor='k',
-                                 interpolate=True)
+                y2 = line * ampli + i + base
+                ax.fill_between(t, y1, y2, where=y2 >= y1, facecolor='k',
+                                interpolate=True)
 
         for filterdb in get_filters(db, all=True):
             if filterid == filterdb.ref:
                 low = float(filterdb.low)
                 high = float(filterdb.high)
                 break
-
-        plt.xlabel("Lag Time (s)")
-        plt.axhline(0, lw=0.5, c='k')
-        plt.grid()
-        title = '%s : %s, %s, Filter %d (%.2f - %.2f Hz), Stack %d' %\
+
+        ax.set_xlabel("Lag Time (s)")
+        ax.axhline(0, lw=0.5, c='k')
+        ax.grid()
+        ax.scatter(0, [start, ], alpha=0)
+        ax.set_ylim(start - datetime.timedelta(days=ampli),
+                    end + datetime.timedelta(days=ampli))
+        ax.set_xlim(-maxlag, maxlag)
+        ax.fmt_ydata = mdates.DateFormatter('%Y-%m-%d')
+        cursor = Cursor(ax, useblit=True, color='red', linewidth=1.2)
+
+        title = '%s : %s, %s, Filter %d (%.2f - %.2f Hz), Stack %d' % \
                 (sta1.replace('_', '.'), sta2.replace('_', '.'), components,
                  filterid, low, high, mov_stack)
+
         if refilter:
             title += ", Re-filtered (%.2f - %.2f Hz)" % (freqmin, freqmax)
-        plt.title(title)
-        plt.scatter(0, [start, ], alpha=0)
-        plt.ylim(start-datetime.timedelta(days=ampli),
-                 end+datetime.timedelta(days=ampli))
-        plt.xlim(-maxlag, maxlag)
-        ax.fmt_ydata = mdates.DateFormatter('%Y-%m-%d')
-        cursor = Cursor(ax, useblit=True, color='red', linewidth=1.2)
+        ax.set_title(title)
 
         if outfile:
             if outfile.startswith("?"):
@@ -99,6 +110,9 @@ def main(sta1, sta2, filterid, components, mov_stack=1, ampli=5, seismic=False,
                                                                   mov_stack))
             outfile = "ccftime " + outfile
             print("output to:", outfile)
-            plt.savefig(outfile)
+            fig.savefig(outfile)
+
         if show:
-            plt.show()
+            fig.show()
+        else:
+            plt.close(fig)