Experiments with plotting in TF plane

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

postprocess/tfhist.py

+import numpy as np
+from matplotlib import use
+use('Agg')
+import matplotlib.pyplot as plt
+from scipy import stats
+from scipy import optimize
+
+plt.close('all')
+# ----------------------------
+# Read run file for basic info
+# ----------------------------
+infile = open('bayeswave.run', 'r')
+for line in infile:
+    if line.find('duration of data') > -1:
+        run_duration = int(line.split()[4])
+        trig_time = run_duration - 2
+    if line.find('sampling rate') > -1:
+        sample_rate = int(line.split()[3])
+
+print "The duration is {0} and sample rate is {1}".format(run_duration, sample_rate)
+
+# infile = ('signal_1_wavechain.data.0', 'r')
+
+#data = np.recfromtxt('chains/signal_1_wavechain.dat.0')
+#print data
+
+infile = open('chains/signal_1_wavechain.dat.0', 'r')
+datastring = infile.read()
+infile.close()
+dataList = datastring.split('\n')
+# -- Remove blank lines
+dataList = filter(None, dataList)
+data = [ np.array(line.split(), dtype=float) for line in dataList ]
+
+
+# -- Only use some fraction of the samples
+ds = len(data)
+data = data[ds/2::100]
+print "Using {0} samples".format(len(data))
+
+wn = np.array([ int(line[0]) for line in data ])
+
+plt.hist(wn)
+plt.savefig('number_hist.png')
+plt.close()
+
+# ----------------------
+# Select chains with right number of wavelets
+# ----------------------
+number = int(stats.mode(wn)[0][0])
+print number
+# data2 = data[ wn == number ]
+data2 = [data[i] for i in range(len(data)) if wn[i] == number]
+print len(data2)
+
+# -----------------------------
+# Read in parameters of wavelets
+# ----------------------------
+wavelets = np.zeros( (number*len(data2), 5) )
+count = 0
+for line in data2:
+    wavedata = [line[indx] for indx in range(8,len(line))]
+    for i in range(number):
+        wavelets[count] = wavedata[i*5:(i+1)*5]
+        count += 1
+
+# print wavelets
+
+# ---------------------------
+# Make scatter plot of pixels  
+# ---------------------------
+tpwavelet = np.transpose(wavelets)
+t0 = tpwavelet[0] - trig_time
+f0 = tpwavelet[1]
+Q  = tpwavelet[2]
+amp = tpwavelet[3]
+
+# -- Calculate deltas
+delta_f = f0 * np.sqrt(11) / Q
+delta_t = 1.0 / delta_f
+
+nyquist = sample_rate / 2
+
+plt.figure()
+plt.scatter(t0, f0)
+plt.axis([-0.5, 0.5, 0, nyquist])
+plt.savefig('tfscatter.png')
+plt.close()
+
+# ---------------
+# Try ellipse plot
+# ----------------
+from matplotlib.patches import Ellipse
+
+print "Calculating TF ellipses..."
+# ells = [Ellipse( xy=(t0[i], f0[i]), width=delta_t[i], height=delta_f[i], angle=0 )for i in range( len(f0) ) ]
+ells = [Ellipse( xy=(t0[i], f0[i]), width=delta_t[i], height=delta_f[i], angle=0 )for i in range( len(f0) ) ]
+fig = plt.figure()
+ax = fig.add_subplot(111)
+print "Plotting TF ellipses...."
+count = 0
+for e in ells:
+    ax.add_artist(e)
+    e.set_clip_box(ax.bbox)
+#   e.set_alpha(1.0/len(f0))
+    e.set_alpha(0.01)
+    e.set_facecolor('black')
+    e.set_edgecolor('none')
+
+    count += 1
+    print "Plotting ellipse {0}".format(count)
+
+ax.set_xlim(-0.5, 0.5)
+ax.set_ylim(0, nyquist)
+plt.savefig('ellipse.png')
+plt.close()
+
+# --------------
+# Try a histogram
+# ---------------
+# t_edges = np.arange(-0.5, 0.5, 1.0/100)
+# f_edges = np.arange(0, nyquist, 2)
+nbins = 100
+trange = [-0.5, 0.5]
+frange = [0.0, nyquist]
+
+fig = plt.figure()
+H, xedge, yedge = np.histogram2d(t0, f0, bins=nbins, range=[trange, frange])
+
+# H needs to be rotated and flipped
+H = np.rot90(H)
+H = np.flipud(H)
+
+# Mask zeros
+Hmasked = np.ma.masked_where(H==0,H) # Mask pixels with a value of zero
+
+# x,y = np.meshgrid(xedge, yedge)
+fig = plt.figure()
+plt.pcolormesh(xedge,yedge,np.log10(Hmasked))
+cbar = plt.colorbar()
+cbar.ax.set_ylabel('log10 Counts')
+plt.xlabel('Time (s)')
+plt.ylabel('Frequency (Hz)')
+plt.savefig('tfhist.png')
+plt.close()
+
+# ------------------------------
+# Fit gaussian to the histogram
+# ------------------------------
+# X, Y = plt.meshgrid(xedge, yedge)
+
+# -- http://wiki.scipy.org/Cookbook/FittingData#head-11870c917b101bb0d4b34900a0da1b7deb613bf7
+#def gaussian(height, center_x, center_y, width_x, width_y):
+#    """Returns a gaussian function with the given parameters"""
+#    width_x = float(width_x)
+#    width_y = float(width_y)
+#    return lambda x,y: height*exp(
+#        -(((center_x-x)/width_x)**2+((center_y-y)/width_y)**2)/2)
+
+
+#def many_gaussian(heightV, x0V, y0V, width_xV, width_yV):
+#    """Returns a sum of gaussians"""