[examples, plot] Tidy up plotpeaks and plotpartials examples.

3 files changed, 18 insertions, 55 deletions

diff --git a/simpl/examples/plotpartials.py b/simpl/examples/plotpartials.py
index b86b69a..ab71e18 100644
--- a/simpl/examples/plotpartials.py
+++ b/simpl/examples/plotpartials.py
@@ -1,20 +1,15 @@
-import simpl
 import matplotlib.pyplot as plt
-from scipy.io.wavfile import read
+import simpl
 
 input_file = '../../tests/audio/flute.wav'
-audio_in_data = read(input_file)
-audio_in = simpl.asarray(audio_in_data[1]) / 32768.0  # values between -1 and 1
-sample_rate = audio_in_data[0]
+audio = simpl.read_wav(input_file)[0]
+audio = audio[len(audio) / 2:(len(audio) / 2) + 4096]
 
-# take just the first few frames
-audio = audio_in[0:4096]
-# Peak detection and partial tracking using SMS
 pd = simpl.SndObjPeakDetection()
 pd.max_peaks = 60
 peaks = pd.find_peaks(audio)
 pt = simpl.MQPartialTracking()
 pt.max_partials = 60
 partials = pt.find_partials(peaks)
-simpl.plot.plot_partials(partials)
+simpl.plot_partials(partials)
 plt.show()
diff --git a/simpl/examples/plotpeaks.py b/simpl/examples/plotpeaks.py
index 66a538f..01d44ed 100644
--- a/simpl/examples/plotpeaks.py
+++ b/simpl/examples/plotpeaks.py
@@ -4,8 +4,8 @@ import matplotlib.pyplot as plt
 input_file = '../../tests/audio/flute.wav'
 audio = simpl.read_wav(input_file)[0]
 
-# take just the first few frames
-audio = audio[0:4096]
+# take just a few frames
+audio = audio[len(audio) / 2:(len(audio) / 2) + 4096]
 
 # peak detection using the SndObj library
 pd = simpl.SndObjPeakDetection()
diff --git a/simpl/plot/__init__.py b/simpl/plot/__init__.py
index a7aad8e..ea552a9 100644
--- a/simpl/plot/__init__.py
+++ b/simpl/plot/__init__.py
@@ -1,4 +1,3 @@
-import simpl
 import matplotlib.pyplot as plt
 import colours
 
@@ -38,57 +37,26 @@ def plot_peaks(frames):
         _plot_frame_peaks(frame.peaks, frame_number, max_amp)
 
 
-def plot_partials(frames, show_peaks=False):
+def plot_partials(frames, show_peaks=True):
     "Plot partials created by a partial tracking algorithm"
     # Get the maximum peak amplitude, used to select an appropriate
     # colour for each peak.
     max_amp = None
     for frame in frames:
-        if frame.peaks:
-            max_amp = max(max_amp, max([p.amplitude for p in frame.peaks]))
-    # If no max amp then no peaks so return
+        if frame.partials:
+            max_amp = max(max_amp, max([p.amplitude for p in frame.partials]))
+
     if not max_amp:
-        print "Warning: no peaks with an amplitude of > 0 to plot, returning"
+        print "No partial peaks with an amplitude of > 0 to plot"
         return
 
-    # Create Partial objects from frames
-    num_frames = len(frames)
-    partials = []
-    live_partials = [None for i in range(len(frames[0].partials))]
-    for n, f in enumerate(frames):
-        for i, p in enumerate(f.partials):
-            if p.amplitude > 0:
-                # active partial
-                if live_partials[i]:
-                    live_partials[i].add_peak(p)
-                else:
-                    partial = simpl.Partial()
-                    partial.starting_frame = n
-                    partial.add_peak(p)
-                    live_partials[i] = partial
-            else:
-                # currently dead
-                if live_partials[i]:
-                    partials.append(live_partials[i])
-                    live_partials[i] = None
-    for p in live_partials:
-        if p:
-            partials.append(p)
-
-    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))
+    for n in range(len(frames) - 1):
+        for p in range(frame.max_partials):
+            x = [n, n + 1]
+            y = [frames[n].partial(p).frequency,
+                 frames[n + 1].partial(p).frequency]
+            amp = frames[n].partial(p).amplitude
+            plt.plot(x, y, color=colours.pbj(amp / max_amp))
 
     if show_peaks:
         plot_peaks(frames)