Add a default colormap argument to rainbowgram

zounds/spectral/functional.py

@@ -8,6 +8,7 @@
 from zounds.loudness import log_modulus, unit_scale
 import numpy as np
 from scipy.signal import resample
+from matplotlib import cm
 
 
 def fft(x, axis=-1, padding_samples=0):
@@ -74,28 +75,14 @@ def phase_shift(coeffs, samplerate, time_shift, axis=-1, frequency_band=None):
     return new_coeffs
 
 
-# def apply_scale(short_time_fft, scale, reducer=np.sum, window=None):
-#     magnitudes = np.abs(short_time_fft.real)
-#     output = np.zeros(
-#         short_time_fft.shape[:-1] + (len(scale),), dtype=magnitudes.dtype)
-#     output = ArrayWithUnits(
-#         output, short_time_fft.dimensions[:-1] + (FrequencyDimension(scale),))
-#     window = window or IdentityWindowingFunc()
-#     for i, freq_band in enumerate(scale):
-#         reduced_band = reducer(magnitudes[..., freq_band] * window, axis=-1)
-#         output[..., i] = reduced_band
-#     return output
-
-
 def apply_scale(short_time_fft, scale, window=None):
     magnitudes = np.abs(short_time_fft.real)
     spectrogram = scale.apply(magnitudes, window)
     dimensions = short_time_fft.dimensions[:-1] + (FrequencyDimension(scale),)
     return ArrayWithUnits(spectrogram, dimensions)
 
 
-def rainbowgram(time_frequency_repr, colormap):
-
+def rainbowgram(time_frequency_repr, colormap=cm.rainbow):
     # magnitudes on a log scale, and shifted and
     # scaled to the unit interval
     magnitudes = np.abs(time_frequency_repr.real)