Added colour to default plots (with heatmap) and some minor code tidying

1 files changed, 88 insertions, 0 deletions

diff --git a/plot/__init__.py b/plot/__init__.py
new file mode 100644
index 0000000..4954b07
--- /dev/null
+++ b/plot/__init__.py
@@ -0,0 +1,88 @@
+# Copyright (c) 2009 John Glover, National University of Ireland, Maynooth
+# 
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation; either version 2 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of 
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program; if not, write to the Free Software
+# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
+
+import matplotlib.pyplot as plt
+import colours
+
+def plot_frame_peaks(peaks):
+    "Plot peaks in one frame"
+    x_values = []
+    y_values = []
+    for peak in peaks:
+        x_values.append(int(peak.frequency))
+        y_values.append(peak.amplitude)
+    plt.plot(x_values, y_values, 'ro')
+
+#def _plot_frame_peaks(frame, frame_number, max_amp=0):
+#    "Plot one frame, which is a list of Peak objects"
+#    x_values = [frame_number for x in range(len(frame))]
+#    y_values = [int(peak.frequency) for peak in frame]
+#    plt.plot(x_values, y_values, 'ro')
+
+def _plot_frame_peaks(frame, frame_number, max_amp=0):
+    "Plot one frame, which is a list of Peak objects"
+    for peak in frame:
+        plt.plot(frame_number, int(peak.frequency), linestyle="None",
+                 marker="o", markersize=2, markeredgewidth=None,
+                 markerfacecolor=colours.pbj(peak.amplitude/max_amp))
+    
+def plot_peaks(peaks):
+    "Plot peaks found by a peak detection algorithm"
+    max_amp = max([max([p.amplitude for p in f if p]) for f in peaks if f])
+    for frame_number, frame in enumerate(peaks):
+        _plot_frame_peaks(frame, frame_number, max_amp)
+        
+#def plot_partials(partials, show_peaks=True):
+#    "Plot partials created by a partial tracking algorithm"
+#    max_amp = max([max([peak.amplitude for peak in p.peaks if p.peaks])
+#                   for p in partials if p])
+#    num_frames = max([partial.get_last_frame() for partial in partials])
+
+#    peaks = [[] for f in range(num_frames)]
+#    for partial in partials:
+#        x_values = []
+#        y_values = []
+#        for peak_number, peak in enumerate(partial.peaks):
+#            x_values.append(partial.starting_frame + peak_number)
+#            y_values.append(int(peak.frequency))
+#            peaks[partial.starting_frame + peak_number].append(peak)
+#        plt.plot(x_values, y_values, "b")
+#    if show_peaks:
+#        plot_peaks(peaks)  
+
+def plot_partials(partials, show_peaks=False):
+    "Plot partials created by a partial tracking algorithm"
+    max_amp = max([max([peak.amplitude for peak in p.peaks if p.peaks])
+                   for p in partials if p])
+    num_frames = max([partial.get_last_frame() for partial in partials])
+
+    peaks = [[] for f in range(num_frames)]
+    for partial in partials:
+        x_values = []
+        y_values = []
+        avg_amp = 0.0
+        num_peaks = 0
+        for peak_number, peak in enumerate(partial.peaks):
+            x_values.append(partial.starting_frame + peak_number)
+            y_values.append(int(peak.frequency))
+            avg_amp += peak.amplitude
+            num_peaks += 1
+            peaks[partial.starting_frame + peak_number].append(peak)
+        avg_amp /= num_peaks
+        plt.plot(x_values, y_values, color=colours.pbj(avg_amp/max_amp))
+    if show_peaks:
+        plot_peaks(peaks)  
+