Merge pull request #48 from stellarpower/stellarpower.Fix.PythonPlotting

Fix Issues and Improve Plotting in Python

Author: fastlib

src/fcwt/boilerplate.py

@@ -33,12 +33,7 @@ def cwt(input, fs, f0, f1, fn, nthreads=1, scaling="lin", fast=False, norm=True)
 
     return freqs, output
 
-def plot(input, fs, f0=0, f1=0, fn=0, nthreads=1, scaling="lin", fast=False, norm=True):
-    
-    f0 = f0 if f0 else 1/(input.size/fs)
-    f1 = f1 if f1 else fs/2
-    fn = fn if fn else 100
-    freqs, output = cwt(input, fs, f0, f1, fn, nthreads=nthreads, scaling=scaling, fast=fast, norm=norm)
+def _plot(input, freqs, output, fs, f0, f1, fn):
 
     #create two subplots
     fig, (ax1, ax2) = plt.subplots(2, 1, sharex=True)
@@ -49,14 +44,48 @@ def plot(input, fs, f0=0, f1=0, fn=0, nthreads=1, scaling="lin", fast=False, nor
 
     #plot spectrogram
     ax2.imshow(np.abs(output),aspect='auto')
+    
 
     #set time on x-axis every 10s and frequency on y-axis for 10 ticks
+    # So we have fixed 10 ticks on the Y-axis; 10s occupies one X tick
+
+    XTickInterval       = 10 # Seconds
+    YTickCount          = 9
+    YLabelDecimalPlaces = 1
+
     ax2.set_xlabel('Time (s)')
     ax2.set_ylabel('Frequency (Hz)')
     ax2.set_title('CWT')
-    ax2.set_xticks(np.arange(0,input.size,fs*10),np.arange(0,input.size/fs,10))
-    ax2.set_yticks(np.arange(0,fn,fn/10),np.round(freqs[::int(fn/10)]))
+
+
+    
+    yFrequencies = np.linspace(freqs[0], freqs[-1], num = YTickCount) # Select evenly-spaced frequencies, including the highest and lowest
+
+
+    xTickPositions = np.  arange(0, input.size, fs * XTickInterval)   # For X, just stepping in increments is suitable
+    yTickPositions = np.linspace(0, fn,         num = YTickCount  )   # For Y, we want to ensure f0 and f1 are included as ticks.
+    
+    xLabels  = np.arange(0, input.size/fs, XTickInterval)
+
+    # Round the Y labels to be 1 d.p.
+    yLabels  = np. round(yFrequencies, decimals = YLabelDecimalPlaces) # Will be co-erced to a string for us later.
+
+    
+    ax2.set_xticks( ticks = xTickPositions, labels = xLabels )
+    ax2.set_yticks( ticks = yTickPositions, labels = yLabels )
 
 
     plt.show()
+
+
+def plot(input, fs, f0=0, f1=0, fn=0, nthreads=1, scaling="lin", fast=False, norm=True):
+    
+    f0 = f0 if f0 else 1/(input.size/fs)
+    f1 = f1 if f1 else fs/2
+    fn = fn if fn else 100
+    freqs, output = cwt(input, fs, f0, f1, fn, nthreads=nthreads, scaling=scaling, fast=fast, norm=norm)
+
+    _plot(input, output, fs, f0, f1, fn)
+
+