Merge branch 'Development' of https://github.com/LinkedEarth/Pyleocli…

…m_util into Development

pyleoclim/core/ui.py

@@ -1209,15 +1209,16 @@ def summary_plot(self, psd=None, scalogram=None, figsize=[8, 10], title=None,
                     time_lim=None, value_lim=None, period_lim=None, psd_lim=None, n_signif_test=100,
                     time_label=None, value_label=None, period_label=None, psd_label='PSD', 
                     wavelet_kwargs = None, psd_kwargs = None, ts_plot_kwargs = None, wavelet_plot_kwargs = None, 
-                    psd_plot_kwargs = None, trunc_series = None, savefig_settings=None, mute=False):
+                    psd_plot_kwargs = None, trunc_series = None, preprocess = True, savefig_settings=None, mute=False):
         ''' Generate a plot of the timeseries and its frequency content through spectral and wavelet analyses.
 
 
         Parameters
         ----------
 
         psd : PSD
-            the PSD object of a Series. If None, will be calculated. This process can be slow as it will be using the WWZ method.
+            the PSD object of a Series. If None, and psd_kwargs is empty, the PSD from the calculated Scalogram will be used. 
+            Otherwise it will be calculated based on specifications in psd_kwargs.
 
         scalogram : Scalogram
             the Scalogram object of a Series. If None, will be calculated. This process can be slow as it will be using the WWZ method.
@@ -1241,7 +1242,7 @@ def summary_plot(self, psd=None, scalogram=None, figsize=[8, 10], title=None,
             the limitation of the psd axis
 
         n_signif_test=100 : int
-            Number of Monte-Carlo simulations to perform for significance testing. Used when psd=None or scalogram=None
+            Number of Monte-Carlo simulations to perform for significance testing. Used when psd=None or scalogram=None.
 
         time_label : str
             the label for the time axis
@@ -1276,15 +1277,22 @@ def summary_plot(self, psd=None, scalogram=None, figsize=[8, 10], title=None,
                 These will be overriden by summary plot to prevent formatting errors
                 
         trunc_series : list or tuple
-            the limitation of the time axis. This will slice the actual time series into one contained within the passed boundaries and as such effect the resulting scalogram and psd objects (assuming said objects are to be generated by summary_plot).
+            the limitation of the time axis. This will slice the actual time series into one contained within the 
+            passed boundaries and as such effect the resulting scalogram and psd objects 
+            (assuming said objects are to be generated by summary_plot).
+        
+        preprocess : bool
+            if True, the series will be standardized and detrended using pyleoclim defaults 
+            prior to the calculation of the scalogram and psd. The unedited series will be used in the plot,
+            while the edited series will be used to calculate the psd and scalogram.
 
         savefig_settings : dict
             the dictionary of arguments for plt.savefig(); some notes below:
             - "path" must be specified; it can be any existed or non-existed path,
               with or without a suffix; if the suffix is not given in "path", it will follow "format"
             - "format" can be one of {"pdf", "eps", "png", "ps"}
 
-        mute : {True,False}
+        mute : bool
             if True, the plot will not show;
             recommend to turn on when more modifications are going to be made on ax
 
@@ -1383,6 +1391,9 @@ def summary_plot(self, psd=None, scalogram=None, figsize=[8, 10], title=None,
         ax['ts'] = self.plot(ax=ax['ts'], **ts_plot_kwargs)
         ax['ts'].xaxis.set_visible(False)
 
+        if preprocess:
+            self = self.standardize().detrend()
+
         if time_lim is not None:
             ax['ts'].set_xlim(time_lim)
             if 'xlim' in ts_plot_kwargs:
@@ -1410,10 +1421,16 @@ def summary_plot(self, psd=None, scalogram=None, figsize=[8, 10], title=None,
         ax['psd'] = plt.subplot(gs[1:4, -3:], sharey=ax['scal'])
 
         if psd is None:
-            if n_signif_test > 0:
-                psd = self.spectral(**psd_kwargs).signif_test(number=n_signif_test)
+            if psd_kwargs:
+                if n_signif_test > 0:
+                    psd = self.spectral(**psd_kwargs).signif_test(number=n_signif_test)
+                else:
+                    psd = self.spectral(**psd_kwargs)
             else:
-                psd = self.spectral(**psd_kwargs)
+                if n_signif_test > 0:
+                    psd = self.spectral(method='wwz').signif_test(number=n_signif_test)
+                else:
+                    psd = self.spectral(method='wwz', scalogram = scalogram)
 
         ax['psd'] = psd.plot(ax=ax['psd'], transpose=True, **psd_plot_kwargs)
 
@@ -1806,7 +1823,7 @@ def detrend(self, method='emd', **kwargs):
         new.value = v_mod
         return new
 
-    def spectral(self, method='lomb_scargle', freq_method='log', freq_kwargs=None, settings=None, label=None, verbose=False):
+    def spectral(self, method='lomb_scargle', freq_method='log', freq_kwargs=None, settings=None, label=None, scalogram=None, verbose=False):
         ''' Perform spectral analysis on the timeseries
 
         Parameters
@@ -1827,6 +1844,9 @@ def spectral(self, method='lomb_scargle', freq_method='log', freq_kwargs=None, s
         label : str
             Label for the PSD object
 
+        scalogram : pyleoclim.core.ui.Series.Scalogram
+            The return of the wavelet analysis; effective only when the method is 'wwz'
+
         verbose : bool
             If True, will print warning messages if there is any
 
@@ -1931,6 +1951,19 @@ def spectral(self, method='lomb_scargle', freq_method='log', freq_kwargs=None, s
             fig, ax = psd_wwz_signif.plot(title='PSD using WWZ method')
             pyleo.closefig(fig)
 
+        We may take advantage of a pre-calculated scalogram using WWZ to accelerate the spectral analysis
+        (although note that the default parameters for spectral and wavelet analysis using WWZ are different):
+
+        .. ipython:: python
+            :okwarning:
+            :okexcept:    
+
+            scal_wwz = ts_std.wavelet(method='wwz')  # wwz is the default method
+            psd_wwz_fast = ts_std.spectral(method='wwz', scalogram=scal_wwz)
+            @savefig spec_wwz_fast.png
+            fig, ax = psd_wwz_fast.plot(title='PSD using WWZ method w/ pre-calculated scalogram')
+            pyleo.closefig(fig)
+
         - Periodogram
 
         .. ipython:: python
@@ -1992,6 +2025,16 @@ def spectral(self, method='lomb_scargle', freq_method='log', freq_kwargs=None, s
         args['welch'] = {}
         args['periodogram'] = {}
         args[method].update(settings)
+
+        if method == 'wwz' and scalogram is not None:
+            args['wwz'].update(
+                {
+                    'wwa': scalogram.amplitude,
+                    'wwz_Neffs': scalogram.wwz_Neffs,
+                    'wwz_freq': scalogram.frequency,
+                }
+            )
+
         spec_res = spec_func[method](self.value, self.time, **args[method])
         if type(spec_res) is dict:
             spec_res = dict2namedtuple(spec_res)
@@ -2110,6 +2153,11 @@ def wavelet(self, method='wwz', settings=None, freq_method='log', ntau=None, fre
 
         args[method].update(settings)
         wave_res = wave_func[method](self.value, self.time, **args[method])
+        if method == 'wwz':
+            wwz_Neffs = wave_res.Neffs
+        else:
+            wwz_Neffs = None
+
         scal = Scalogram(
             frequency=wave_res.freq,
             time=wave_res.time,
@@ -2121,6 +2169,7 @@ def wavelet(self, method='wwz', settings=None, freq_method='log', ntau=None, fre
             freq_method=freq_method,
             freq_kwargs=freq_kwargs,
             wave_args=args[method],
+            wwz_Neffs=wwz_Neffs,
         )
 
         return scal
@@ -3073,7 +3122,7 @@ def plot(self, in_loglog=True, in_period=True, label=None, xlabel=None, ylabel='
             return ax
 
 class Scalogram:
-    def __init__(self, frequency, time, amplitude, coi=None, label=None, Neff=3, timeseries=None,
+    def __init__(self, frequency, time, amplitude, coi=None, label=None, Neff=3, wwz_Neffs=None, timeseries=None,
                  wave_method=None, wave_args=None, signif_qs=None, signif_method=None, freq_method=None, freq_kwargs=None,
                  period_unit=None, time_label=None):
         '''
@@ -3086,6 +3135,10 @@ def __init__(self, frequency, time, amplitude, coi=None, label=None, Neff=3, tim
             amplitude : array
                 the amplitude at each (frequency, time) point;
                 note the dimension is assumed to be (frequency, time)
+            Neff : int
+                the threshold of the number of effective samples
+            wwz_Neffs : array
+                the matrix of effective number of points in the time-scale coordinates obtained from wwz
         '''
         self.frequency = np.array(frequency)
         self.time = np.array(time)
@@ -3102,6 +3155,8 @@ def __init__(self, frequency, time, amplitude, coi=None, label=None, Neff=3, tim
         self.signif_method = signif_method
         self.freq_method = freq_method
         self.freq_kwargs = freq_kwargs
+        if wave_method == 'wwz':
+            self.wwz_Neffs = wwz_Neffs
 
         if period_unit is not None:
             self.period_unit = period_unit
@@ -3633,7 +3688,7 @@ def signif_test(self, number=200, method='ar1', seed=None, qs=[0.95], settings=N
 
         cohs = []
         for i in tqdm(range(number), desc='Performing wavelet coherence on surrogate pairs', total=number, disable=mute_pbar):
-            coh_tmp = surr1.series_list[i].wavelet_coherence(surr2.series_list[i], tau=self.time, freq_method=self.freq_method, freq_kwargs=self.freq_kwargs)
+            coh_tmp = surr1.series_list[i].wavelet_coherence(surr2.series_list[i], settings={'tau': self.time, 'freq': self.frequency})
             cohs.append(coh_tmp.coherence)
 
         cohs = np.array(cohs)

pyleoclim/tests/test_ui_Series.py

@@ -433,18 +433,20 @@ class TestUiSeriesSummaryPlot:
     def test_summary_plot_t0(self):
         ''' Generate a colored noise and run the summary_plot() function.
         Note that we should avoid pyleo.showfig() in tests.
+        
+        Passing pre generated scalogram and psd.
         '''
         alpha = 1
         t, v = gen_colored_noise(nt=100, alpha=alpha)
         ts = pyleo.Series(time=t, value=v)
-        psd = ts.spectral()
         scal = ts.wavelet()
+        psd = ts.spectral()
         period_label='Period'
         psd_label='PSD'
         time_label='Time'
         value_label='Value'
         fig, ax = ts.summary_plot(
-            psd=psd, scalogram=scal, figsize=[4, 5], title='Test',
+            psd = psd, scalogram=scal, figsize=[4, 5], title='Test',
             period_label=period_label, psd_label=psd_label,
             value_label=value_label, time_label=time_label,
             mute=True,
@@ -456,6 +458,34 @@ def test_summary_plot_t0(self):
         assert ax['ts'].properties()['ylabel'] == value_label, 'Value label is not being passed properly'
 
         plt.close(fig)
+
+    def test_summary_plot_t1(self):
+        ''' Generate a colored noise and run the summary_plot() function.
+        Note that we should avoid pyleo.showfig() in tests.
+    
+        Passing just a pre generated psd.
+        '''
+        alpha = 1
+        t, v = gen_colored_noise(nt=100, alpha=alpha)
+        ts = pyleo.Series(time=t, value=v)
+        psd = ts.spectral()
+        period_label='Period'
+        psd_label='PSD'
+        time_label='Time'
+        value_label='Value'
+        fig, ax = ts.summary_plot(
+            psd = psd, figsize=[4, 5], title='Test',
+            period_label=period_label, psd_label=psd_label,
+            value_label=value_label, time_label=time_label,
+            n_signif_test = 1, mute=True,
+        )
+
+        assert ax['scal'].properties()['ylabel'] == period_label, 'Period label is not being passed properly'
+        assert ax['psd'].properties()['xlabel'] == psd_label, 'PSD label is not being passed properly'
+        assert ax['scal'].properties()['xlabel'] == time_label, 'Time label is not being passed properly'
+        assert ax['ts'].properties()['ylabel'] == value_label, 'Value label is not being passed properly'
+
+        plt.close(fig)
 
 class TestUiSeriesCorrelation:
     ''' Test Series.correlation()

pyleoclim/utils/spectral.py

@@ -615,7 +615,7 @@ def wwz_psd(ys, ts, freq=None, freq_method='log', freq_kwargs=None,
             tau=None, c=1e-3, nproc=8,
             detrend=False, sg_kwargs=None, gaussianize=False,
             standardize=False, Neff=3, anti_alias=False, avgs=2,
-            method='Kirchner_numba'):
+            method='Kirchner_numba', wwa=None, wwz_Neffs=None, wwz_freq=None):
     ''' Return the psd of a timeseries using wwz method.
 
     Parameters
@@ -678,6 +678,16 @@ def wwz_psd(ys, ts, freq=None, freq_method='log', freq_kwargs=None,
         flag for whether spectrum is derived from instantaneous point measurements (avgs<>1)
         OR from measurements averaged over each sampling interval (avgs==1)
 
+    wwa : array
+        the weighted wavelet amplitude, returned from pyleoclim.utils.wavelet.wwz
+
+    wwz_Neffs : array
+        the matrix of effective number of points in the time-scale coordinates,
+        returned from pyleoclim.utils.wavelet.wwz
+
+    wwz_freq : array
+        the returned frequency vector from pyleoclim.utils.wavelet.wwz
+
     Returns
     -------
 
@@ -710,11 +720,15 @@ def wwz_psd(ys, ts, freq=None, freq_method='log', freq_kwargs=None,
 
     # get wwa but AR1_q is not needed here so set nMC=0
     #  wwa, _, _, coi, freq, _, Neffs, _ = wwz(ys_cut, ts_cut, freq=freq, tau=tau, c=c, nproc=nproc, nMC=0,
-    res_wwz = wwz(ys_cut, ts_cut, freq=freq, tau=tau, c=c, nproc=nproc, nMC=0,
-              detrend=detrend, sg_kwargs=sg_kwargs,
-              gaussianize=gaussianize, standardize=standardize, method=method)
-
-    psd = wwa2psd(res_wwz.amplitude, ts_cut, res_wwz.Neffs, freq=res_wwz.freq, Neff=Neff, anti_alias=anti_alias, avgs=avgs)
+    if wwa is None or wwz_Neffs is None or wwz_freq is None:
+        res_wwz = wwz(ys_cut, ts_cut, freq=freq, tau=tau, c=c, nproc=nproc, nMC=0,
+                  detrend=detrend, sg_kwargs=sg_kwargs,
+                  gaussianize=gaussianize, standardize=standardize, method=method)
+        wwa = res_wwz.amplitude
+        wwz_Neffs = res_wwz.Neffs
+        wwz_freq = res_wwz.freq
+
+    psd = wwa2psd(wwa, ts_cut, wwz_Neffs, freq=wwz_freq, Neff=Neff, anti_alias=anti_alias, avgs=avgs)
     #  psd[1/freqs > np.max(coi)] = np.nan  # cut off the unreliable part out of the coi
     #  psd = psd[1/freqs <= np.max(coi)] # cut off the unreliable part out of the coi
     #  freqs = freqs[1/freqs <= np.max(coi)]

pyleoclim/utils/wavelet.py

@@ -1103,7 +1103,7 @@ def wwa2psd(wwa, ts, Neffs, freq=None, Neff=3, anti_alias=False, avgs=2):
     ts : array
         the time points, should be pre-truncated so that the span is exactly what is used for wwz
     Neffs : array
-        the matrix of effective number of points in the time-scale coordinates obtained from wwz from wwz
+        the matrix of effective number of points in the time-scale coordinates obtained from wwz
     freq : array
         vector of frequency from wwz
     Neff : int