Cleanup and bug fixes in plot_waveform

git-svn-id: https://svn.ligo.caltech.edu/svn/bayeswave/trunk@55 c56465c9-8126-4a4f-9d7d-ac845eff4865

postprocess/plot_waveform.py

@@ -7,17 +7,47 @@ import glob
 import argparse
 import os
 
+
+# -------------------------------------------------
+# Define helper function to extract median waveform
+# -------------------------------------------------
+def get_waveform(fileList, ifo):
+    if not len(fileList) == 0:
+        count=0
+        for filename in fileList:
+            time, strain = np.loadtxt(filename, unpack=True)
+            if count == 0:
+                strain_master = np.zeros((200, strain.size))
+            strain_master[count] = strain
+            count += 1
+        print "   "+str(count)+" waveforms found for IFO "+str(int(ifo)+1)
+        median_waveform = np.median(strain_master, axis=0)
+
+        down_vec = np.zeros(median_waveform.size)
+        up_vec = np.zeros(median_waveform.size)
+        for sample in range(median_waveform.size):
+            vector = np.transpose(strain_master)[sample]
+            sort_vec = np.sort(vector)
+            up_quart = sort_vec[ int(0.84*sort_vec.size) ]
+            down_quart = sort_vec[ int(0.16*sort_vec.size) ]
+            down_vec[sample] = down_quart
+            up_vec[sample] = up_quart
+            
+    return (time, median_waveform, up_vec, down_vec)
+
+
+# -- Parse command line arguments
 parser = argparse.ArgumentParser(description='Produce html page for a BayesWaveBurst run.')
 parser.add_argument('--mdc', default=False, help='Is this this an MDC? True/False (Default is false)')
 args = parser.parse_args()
 mdc = args.mdc
 
-jobname = glob.glob('*bayeswave.run')[0]
-bayeswave = open(jobname, 'r')
-
 # ---------------------------------------
 # Extract information from bayeswave.run
 # ---------------------------------------
+jobname = glob.glob('*bayeswave.run')[0]
+bayeswave = open(jobname, 'r')
+
 ifoNames = []
 for line in bayeswave:
     spl = line.split()
@@ -67,33 +97,14 @@ for ifo in ifoList:
 
     # -- Read Signal model
     fileList = glob.glob('waveforms/'+str(jobname)+'signal_1_wave_*.dat.{0}'.format(ifo))
-    if not len(fileList) == 0:
-        count=0
-        for filename in fileList:
-            time, strain = np.loadtxt(filename, unpack=True)
-            if count == 0:
-                strain_master = np.zeros((200, strain.size))
-            strain_master[count] = strain
-            count += 1
-        print "   "+str(count)+" waveforms found for IFO "+str(int(ifo)+1)
-        median_waveform = np.median(strain_master, axis=0)
-
-        down_vec = np.zeros(median_waveform.size)
-        up_vec = np.zeros(median_waveform.size)
-        for sample in range(median_waveform.size):
-            vector = np.transpose(strain_master)[sample]
-            sort_vec = np.sort(vector)
-            up_quart = sort_vec[ int(0.84*sort_vec.size) ]
-            down_quart = sort_vec[ int(0.16*sort_vec.size) ]
-            down_vec[sample] = down_quart
-            up_vec[sample] = up_quart
-
-        # -- Get axis info
-        wave = up_vec
-        # wave = median_waveform
-        wave_max = np.maximum(wave.max(), d_s.max())
-        sig_times = time[wave > 0.05*wave.max()]
-        axisList.append([sig_times.min(), sig_times.max(), -wave_max*1.1, wave_max*1.1])
+    (time, median_waveform, up_vec, down_vec) = get_waveform(fileList, ifo)
+    
+    # -- Get axis info
+    wave = up_vec
+    # wave = median_waveform
+    wave_max = np.maximum(wave.max(), d_s.max())
+    sig_times = time[wave > 0.05*wave.max()]
+    axisList.append([sig_times.min(), sig_times.max(), -wave_max*1.1, wave_max*1.1])
 
 # -- Select Axis with biggest y-value
 ymax = 0
@@ -110,137 +121,80 @@ for ifo in ifoList:
     # Read in signal model and data 
     # ----------------------------------
     fileList = glob.glob('waveforms/'+str(jobname)+'signal_1_wave_*.dat.{0}'.format(ifo))
-
-    if not len(fileList) == 0:
-        count=0
-        for filename in fileList:
-            time, strain = np.loadtxt(filename, unpack=True)
-            if count == 0:
-                strain_master = np.zeros((200, strain.size))
-            strain_master[count] = strain
-            count += 1
-
-        print "   "+str(count)+" waveforms found for IFO "+str(int(ifo)+1)
-
-        median_waveform = np.median(strain_master, axis=0)
-
-        down_vec = np.zeros(median_waveform.size)
-        up_vec = np.zeros(median_waveform.size)
-        for sample in range(median_waveform.size):
-        
-            vector = np.transpose(strain_master)[sample]
-            sort_vec = np.sort(vector)
-            up_quart = sort_vec[ int(0.84*sort_vec.size) ]
-            down_quart = sort_vec[ int(0.16*sort_vec.size) ]
-
-            down_vec[sample] = down_quart
-            up_vec[sample] = up_quart
+    (time, median_waveform, up_vec, down_vec) = get_waveform(fileList, ifo)
         
-        
-        # -- Read whitened data
-        d_t, d_s = np.loadtxt('waveforms/'+str(jobname)+'clean_1_data_199.dat.{0}'.format(ifo), unpack=True)
+    # -- Read whitened data
+    d_t, d_s = np.loadtxt('waveforms/'+str(jobname)+'clean_1_data_199.dat.{0}'.format(ifo), unpack=True)
   
-        # -- Plot the median waveform and whitened data
-        fignum = plt.figure()
-        plt.plot(d_t, d_s, 'r', linewidth=1, alpha=0.3)
-        plt.plot(time, median_waveform, 'b', linewidth=2)
-        plt.fill_between(time, down_vec, up_vec, facecolor='blue', edgecolor='blue', alpha=0.3)
-        
-        plt.axis(axwinner)
-        axisList.append([sig_times.min(), sig_times.max(), -wave_max*1.1, wave_max*1.1])
-        plt.xlabel('Time (s)')
-        plt.ylabel('Best Fit Waveform (blue) in whitened data (red)')
-        plt.title('IFO {0}'.format(ifo))
-
-        # -- Save the plot
-        plt.savefig('waveform_{0}.png'.format(ifo))
-        plt.close()
-
-        # -- Plot the clean, whitened data and best fit waveform
-        d_t, d_s = np.loadtxt('waveforms/'+str(jobname)+'clean_1_data_199.dat.{0}'.format(ifo), unpack=True)
-        plt.figure()
-        plt.plot(d_t, d_s, 'r', linewidth=3)
-        plt.plot(time, median_waveform, 'b', linewidth=1)
-        vertrange = np.maximum( median_waveform.max(), d_s.max() )
-        plt.axis([sig_times.min(), sig_times.max(), -vertrange*1.1, vertrange*1.1])
-        plt.xlabel('Time (s)')
-        plt.ylabel('Whitened Data (red) and Signal model (blue)')
-        plt.title('IFO {0}'.format(ifo))
-        plt.savefig('data_{0}.png'.format(ifo))
-        plt.close()
-
-        # -- Attempt to make a spectrogram of median waveform
-        ts = time[1] - time[0]
-        fs = int(np.round(1.0/ts))
-        NFFT = fs / 16
-        window = np.blackman(NFFT)
-        holder = np.zeros(1*NFFT)
-        seg = np.append(holder, median_waveform)
-        seg = np.append(seg, holder)
-        spec_power, freqs, bins = mlab.specgram(median_waveform, NFFT=NFFT, Fs=fs, 
-                                        window=window, noverlap=int(0.8*NFFT))
-
-        # -- Try eliminating quiet pixels
-        sort_power = np.sort(spec_power.flatten())[::-1]
-        sum_power = np.cumsum(sort_power)
-        indx = np.nonzero( sum_power < 0.95*sum_power.max() )
-        point90 = sort_power[indx[0].max()]
-        spec_power[ spec_power < point90 ] = spec_power.min()
-        vmin = np.log10(spec_power[ spec_power > spec_power.min()].min())
-
-        # -- Try do it yourself spectrogram
-        plt.figure()
-        ax = plt.subplot(111)
-        ax.pcolormesh(bins-0.5, freqs, np.log10(spec_power), vmin=vmin, cmap='OrRd')
-        plt.title('IFO {0}'.format(ifo))
-        plt.xlabel('Time (s)')
-        plt.ylabel('Frequency (Hz)')
-        plt.axis([axwinner[0], axwinner[1], 0, fs/2])
-        plt.savefig('self_spec_{0}.png'.format(ifo))
-        plt.close()
+    # -- Plot the median waveform and whitened data
+    fignum = plt.figure()
+    plt.plot(d_t, d_s, 'r', linewidth=1, alpha=0.3)
+    plt.plot(time, median_waveform, 'b', linewidth=2)
+    plt.fill_between(time, down_vec, up_vec, facecolor='blue', edgecolor='blue', alpha=0.3)
+    
+    plt.axis(axwinner)
+    axisList.append([sig_times.min(), sig_times.max(), -wave_max*1.1, wave_max*1.1])
+    plt.xlabel('Time (s)')
+    plt.ylabel('Best Fit Waveform (blue) in whitened data (red)')
+    plt.title('IFO {0}'.format(ifo))
+
+    # -- Save the plot
+    plt.savefig('waveform_{0}.png'.format(ifo))
+    plt.close()
+    
+    # -- Attempt to make a spectrogram of median waveform
+    ts = time[1] - time[0]
+    fs = int(np.round(1.0/ts))
+    NFFT = fs / 16
+    window = np.blackman(NFFT)
+    holder = np.zeros(1*NFFT)
+    seg = np.append(holder, median_waveform)
+    seg = np.append(seg, holder)
+    spec_power, freqs, bins = mlab.specgram(median_waveform, NFFT=NFFT, Fs=fs, 
+                                            window=window, noverlap=int(0.8*NFFT))
+
+    # -- Try eliminating quiet pixels
+    sort_power = np.sort(spec_power.flatten())[::-1]
+    sum_power = np.cumsum(sort_power)
+    indx = np.nonzero( sum_power < 0.95*sum_power.max() )
+    point90 = sort_power[indx[0].max()]
+    spec_power[ spec_power < point90 ] = spec_power.min()
+    vmin = np.log10(spec_power[ spec_power > spec_power.min()].min())
+
+    # -- Try do it yourself spectrogram
+    fig = plt.figure()
+    ax = plt.subplot(111)
+    binwidth = (bins[1] - bins[0]) / 2.0
+    cax = ax.pcolormesh(bins-0.5-binwidth, freqs, np.log10(spec_power), vmin=vmin, cmap='OrRd')
+    plt.title('IFO {0}'.format(ifo))
+    plt.xlabel('Time (s)')
+    plt.ylabel('Frequency (Hz)')
+    plt.axis([axwinner[0], axwinner[1], 0, fs/2])
+    cb = plt.colorbar(cax)
+    cb.set_label('Logrithmic Signal Power')
+    plt.savefig('self_spec_{0}.png'.format(ifo))
+    plt.close()
 
     # --------------------------------
     # -- Plot the median glitch model
     # --------------------------------
     fileList = glob.glob('waveforms/'+str(jobname)+'glitch_1_glitch_*.dat.{0}'.format(ifo))
+    (time, median_waveform, up_vec, down_vec) = get_waveform(fileList, ifo)
 
-    if not len(fileList) == 0:
-        count=0
-        for filename in fileList:
-            time, strain = np.loadtxt(filename, unpack=True)
-            if count == 0:
-                strain_master = np.zeros((200, strain.size))
-            strain_master[count] = strain
-            count += 1
-
-        median_waveform = np.median(strain_master, axis=0)
-
-        down_vec = np.zeros(median_waveform.size)
-        up_vec = np.zeros(median_waveform.size)
-        for sample in range(median_waveform.size):
-        
-            vector = np.transpose(strain_master)[sample]
-            sort_vec = np.sort(vector)
-            up_quart = sort_vec[ int(0.84*sort_vec.size) ]
-            down_quart = sort_vec[ int(0.16*sort_vec.size) ]
-
-            down_vec[sample] = down_quart
-            up_vec[sample] = up_quart
+    plt.figure()
+    plt.plot(d_t, d_s, 'r', linewidth=1, alpha=0.3)
+    plt.plot(time, median_waveform, 'k', linewidth=2)
+    plt.fill_between(time, down_vec, up_vec, facecolor='black', edgecolor='black', alpha=0.3)
+    
+    # -- Adjust plot properties
+    plt.axis(axwinner)
+    plt.xlabel('Time (s)')
+    plt.ylabel('Best Fit Glitch Model (whitened)')
+    plt.title('IFO {0}'.format(ifo))
 
-        plt.figure()
-        plt.plot(d_t, d_s, 'r', linewidth=1, alpha=0.3)
-        plt.plot(time, median_waveform, 'k', linewidth=2)
-        plt.fill_between(time, down_vec, up_vec, facecolor='black', edgecolor='black', alpha=0.3)
-
-        # -- Adjust plot properties
-        plt.axis(axwinner)
-        plt.xlabel('Time (s)')
-        plt.ylabel('Best Fit Glitch Model (whitened)')
-        plt.title('IFO {0}'.format(ifo))
-
-        # -- Save the plot
-        plt.savefig('glitch_{0}.png'.format(ifo))
-        plt.close()
+    # -- Save the plot
+    plt.savefig('glitch_{0}.png'.format(ifo))
+    plt.close()
 
     # -----------------------
     # Plot the median PSDs
@@ -397,10 +351,19 @@ if os.path.isfile('condorOut.txt'):
 
 # -- If MDC, read SNR info
 if mdc=='True':
-    snrFile = open('*snr/snr.dat', 'r')
-    snrLine = snrFile.read()
+    print "Looking for SNRs..."
+    snrList = []
+    foundSnr=False
+    snrFile = open('condorOut.txt', 'r')
+    for line in snrFile:
+        if line.find('Recompute')>0:
+            print "Found Recompute section!"
+            foundSnr = True
+        if (line.find('Injected SNR')>-1) and foundSnr:
+            print "Found an SNR!!"
+            snrList.append(float(line.split()[2]))
     snrFile.close
-    snrList = [float(snr) for snr in snrLine.split()]
+
 
 info = info + "The detectors are {0} \n".format( ', '.join(ifoNames) )
 info = info + "   The trigger time is GPS {0} \n".format(gps)
@@ -435,12 +398,6 @@ web.write("Reconstructed waveforms in whitened data.  Red is data, blue is recon
 for ifo in ifoList:
     web.write("<img src='./waveform_{0}.png' alt='waveform_{0}.png' width=500>".format(ifo))
 
-# -- Plot the data and best fit waveform
-# web.write("<h3>Whitened Data (red) and Median Reconstructed Waveform (blue) </h3>")
-# web.write("Reconstructed waveforms in whitened data.  Red is data, blue is reconstructed waveform<br/>")
-# for ifo in ifoList:
-#    web.write("<img src='./data_{0}.png' alt='data_{0}.png' width=500>".format(ifo))
-
 # -- Write out Glitch model
 web.write("<h3>Median Glitch Model Waveforms and 1-sigma Uncertainties</h3>")
 web.write("Reconstructed waveforms in whitened data<br/>")
@@ -450,7 +407,6 @@ for ifo in ifoList:
 # -- Write out spectrogram
 web.write("<h3>Spectrogram of median reconstructed signal model waveform</h3>")
 for ifo in ifoList:
-    # web.write("<img src='./spec_{0}.png' alt='spec_{0}.png' width=500>".format(ifo) )
     web.write("<img src='./self_spec_{0}.png' alt='self_spec_{0}.png' width=500>".format(ifo) )
 
 # -- Write out the PSDs
@@ -463,11 +419,6 @@ web.write("<h3>Diagnostic plots of Likelihood and Temperature</h3>")
 web.write("<img src='./likelihood.png' alt='likelihood.png' width=500>")
 web.write("<img src='./TL.png' alt='TL.png' width=500> <br/>" )
     
-# -- Put all plots on web site
-# plotlist = glob.glob('*.png')
-# for plot in plotlist:
-#     web.write("<img src='{0}' width=200>".format(plot))
-
 # -- Include skymap
 web.write('<h3>Skymap</h3>')
 web.write("<img src='./skymap.png' alt='No_skymap_available' width=650>")