Fix argument handling for 'mother'; add cross-hatching for COI; clean up code formatting

Author: Chris Torrence

wave_python/waveletAnalysis.py

@@ -1,30 +1,34 @@
-import numpy as np
-from waveletFunctions import wavelet, wave_signif
 import matplotlib.pylab as plt
-from matplotlib.gridspec import GridSpec
 import matplotlib.ticker as ticker
-from mpl_toolkits.axes_grid1 import make_axes_locatable
+from matplotlib.gridspec import GridSpec
+
+import numpy as np
+
+# from mpl_toolkits.axes_grid1 import make_axes_locatable
+
+from waveletFunctions import wave_signif, wavelet
 
 __author__ = 'Evgeniya Predybaylo'
 
 
 # WAVETEST Example Python script for WAVELET, using NINO3 SST dataset
 #
 # See "http://paos.colorado.edu/research/wavelets/"
-# The Matlab code written January 1998 by C. Torrence is modified to Python by Evgeniya Predybaylo, December 2014
+# The Matlab code written January 1998 by C. Torrence
+# modified to Python by Evgeniya Predybaylo, December 2014
 #
 # Modified Oct 1999, changed Global Wavelet Spectrum (GWS) to be sideways,
 #   changed all "log" to "log2", changed logarithmic axis on GWS to
 #   a normal axis.
-# ------------------------------------------------------------------------------------------------------------------
+# ---------------------------------------------------------------------------
 
 # READ THE DATA
 sst = np.loadtxt('sst_nino3.dat')  # input SST time series
 sst = sst - np.mean(sst)
 variance = np.std(sst, ddof=1) ** 2
 print("variance = ", variance)
 
-#----------C-O-M-P-U-T-A-T-I-O-N------S-T-A-R-T-S------H-E-R-E------------------------------------------------------
+# ----------C-O-M-P-U-T-A-T-I-O-N------S-T-A-R-T-S------H-E-R-E---------------
 
 # normalize by standard deviation (not necessary, but makes it easier
 # to compare with plot on Interactive Wavelet page, at
@@ -52,7 +56,8 @@
 # Significance levels:
 signif = wave_signif(([variance]), dt=dt, sigtest=0, scale=scale,
     lag1=lag1, mother=mother)
-sig95 = signif[:, np.newaxis].dot(np.ones(n)[np.newaxis, :])  # expand signif --> (J+1)x(N) array
+# expand signif --> (J+1)x(N) array
+sig95 = signif[:, np.newaxis].dot(np.ones(n)[np.newaxis, :])
 sig95 = power / sig95  # where ratio > 1, power is significant
 
 # Global wavelet spectrum & significance levels:
@@ -63,65 +68,72 @@
 # Scale-average between El Nino periods of 2--8 years
 avg = np.logical_and(scale >= 2, scale < 8)
 Cdelta = 0.776  # this is for the MORLET wavelet
-scale_avg = scale[:, np.newaxis].dot(np.ones(n)[np.newaxis, :])  # expand scale --> (J+1)x(N) array
+# expand scale --> (J+1)x(N) array
+scale_avg = scale[:, np.newaxis].dot(np.ones(n)[np.newaxis, :])
 scale_avg = power / scale_avg  # [Eqn(24)]
 scale_avg = dj * dt / Cdelta * sum(scale_avg[avg, :])  # [Eqn(24)]
 scaleavg_signif = wave_signif(variance, dt=dt, scale=scale, sigtest=2,
     lag1=lag1, dof=([2, 7.9]), mother=mother)
 
-#------------------------------------------------------ Plotting
+# ------------------------------------------------------ Plotting
 
-#--- Plot time series
+# --- Plot time series
 fig = plt.figure(figsize=(9, 10))
 gs = GridSpec(3, 4, hspace=0.4, wspace=0.75)
-plt.subplots_adjust(left=0.1, bottom=0.05, right=0.9, top=0.95, wspace=0, hspace=0)
+plt.subplots_adjust(left=0.1, bottom=0.05, right=0.9, top=0.95,
+                    wspace=0, hspace=0)
 plt.subplot(gs[0, 0:3])
 plt.plot(time, sst, 'k')
 plt.xlim(xlim[:])
 plt.xlabel('Time (year)')
 plt.ylabel('NINO3 SST (\u00B0C)')
 plt.title('a) NINO3 Sea Surface Temperature (seasonal)')
 
-plt.text(time[-1] + 35, 0.5,'Wavelet Analysis\nC. Torrence & G.P. Compo\n' +
+plt.text(time[-1] + 35, 0.5, 'Wavelet Analysis\nC. Torrence & G.P. Compo\n'
     'http://paos.colorado.edu/\nresearch/wavelets/',
     horizontalalignment='center', verticalalignment='center')
 
-#--- Contour plot wavelet power spectrum
+# --- Contour plot wavelet power spectrum
 # plt3 = plt.subplot(3, 1, 2)
 plt3 = plt.subplot(gs[1, 0:3])
 levels = [0, 0.5, 1, 2, 4, 999]
-CS = plt.contourf(time, period, power, len(levels))  #*** or use 'contour'
-im = plt.contourf(CS, levels=levels, colors=['white','bisque','orange','orangered','darkred'])
+# *** or use 'contour'
+CS = plt.contourf(time, period, power, len(levels))
+im = plt.contourf(CS, levels=levels,
+    colors=['white', 'bisque', 'orange', 'orangered', 'darkred'])
 plt.xlabel('Time (year)')
 plt.ylabel('Period (years)')
 plt.title('b) Wavelet Power Spectrum (contours at 0.5,1,2,4\u00B0C$^2$)')
 plt.xlim(xlim[:])
 # 95# significance contour, levels at -99 (fake) and 1 (95# signif)
 plt.contour(time, period, sig95, [-99, 1], colors='k')
 # cone-of-influence, anything "below" is dubious
+plt.fill_between(time, coi * 0 + period[-1], coi, facecolor="none",
+    edgecolor="#00000040", hatch='x')
 plt.plot(time, coi, 'k')
 # format y-scale
-plt3.set_yscale('log', basey=2, subsy=None)
+plt3.set_yscale('log', base=2, subs=None)
 plt.ylim([np.min(period), np.max(period)])
 ax = plt.gca().yaxis
 ax.set_major_formatter(ticker.ScalarFormatter())
 plt3.ticklabel_format(axis='y', style='plain')
 plt3.invert_yaxis()
 # set up the size and location of the colorbar
-# position=fig.add_axes([0.5,0.36,0.2,0.01]) 
-# plt.colorbar(im, cax=position, orientation='horizontal') #, fraction=0.05, pad=0.5)
+# position=fig.add_axes([0.5,0.36,0.2,0.01])
+# plt.colorbar(im, cax=position, orientation='horizontal')
+#   , fraction=0.05, pad=0.5)
 
 # plt.subplots_adjust(right=0.7, top=0.9)
 
-#--- Plot global wavelet spectrum
+# --- Plot global wavelet spectrum
 plt4 = plt.subplot(gs[1, -1])
 plt.plot(global_ws, period)
 plt.plot(global_signif, period, '--')
 plt.xlabel('Power (\u00B0C$^2$)')
 plt.title('c) Global Wavelet Spectrum')
 plt.xlim([0, 1.25 * np.max(global_ws)])
 # format y-scale
-plt4.set_yscale('log', basey=2, subsy=None)
+plt4.set_yscale('log', base=2, subs=None)
 plt.ylim([np.min(period), np.max(period)])
 ax = plt.gca().yaxis
 ax.set_major_formatter(ticker.ScalarFormatter())
@@ -138,6 +150,3 @@
 plt.plot(xlim, scaleavg_signif + [0, 0], '--')
 
 plt.show()
-
-# end of code
-