<s> """<STR_LIT>""" <EOL> from __future__ import print_function , division <EOL> __all__ = [ '<STR_LIT>' ] <EOL> import warnings <EOL> import numpy as np <EOL> from . lomb_scargle import LombScargle <EOL> FACTORIALS = [ <NUM_LIT:1> , <NUM_LIT:1> , <NUM_LIT:2> , <NUM_LIT:6> , <NUM_LIT> , <NUM_LIT> , <NUM_LIT> , <NUM_LIT> , <NUM_LIT> , <NUM_LIT> , <NUM_LIT> ] <EOL> def factorial ( N ) : <EOL> """<STR_LIT>""" <EOL> if N < len ( FACTORIALS ) : <EOL> return FACTORIALS [ N ] <EOL> else : <EOL> from scipy import special <EOL> return int ( special . factorial ( N ) ) <EOL> def bitceil ( N ) : <EOL> """<STR_LIT>""" <EOL> N = int ( N ) - <NUM_LIT:1> <EOL> for i in [ <NUM_LIT:1> , <NUM_LIT:2> , <NUM_LIT:4> , <NUM_LIT:8> , <NUM_LIT:16> , <NUM_LIT:32> ] : <EOL> N |= N >> i <EOL> return N + <NUM_LIT:1> <EOL> def extirpolate ( x , y , N = None , M = <NUM_LIT:4> ) : <EOL> """<STR_LIT>""" <EOL> x , y = map ( np . ravel , np . broadcast_arrays ( x , y ) ) <EOL> if N is None : <EOL> N = int ( np . max ( x ) + <NUM_LIT:0.5> * M + <NUM_LIT:1> ) <EOL> result = np . zeros ( N , dtype = y . dtype ) <EOL> integers = ( x % <NUM_LIT:1> == <NUM_LIT:0> ) <EOL> np . add . at ( result , x [ integers ] . astype ( int ) , y [ integers ] ) <EOL> x , y = x [ ~ integers ] , y [ ~ integers ] <EOL> ilo = np . clip ( ( x - M // <NUM_LIT:2> ) . astype ( int ) , <NUM_LIT:0> , N - M ) <EOL> numerator = y * np . prod ( x - ilo - np . arange ( M ) [ : , np . newaxis ] , <NUM_LIT:0> ) <EOL> denominator = factorial ( M - <NUM_LIT:1> ) <EOL> for j in range ( M ) : <EOL> if j > <NUM_LIT:0> : <EOL> denominator *= j / ( j - M ) <EOL> ind = ilo + ( M - <NUM_LIT:1> - j ) <EOL> np . add . at ( result , ind , numerator / ( denominator * ( x - ind ) ) ) <EOL> return result <EOL> def trig_sum ( t , h , df , N , f0 = <NUM_LIT:0> , freq_factor = <NUM_LIT:1> , <EOL> oversampling = <NUM_LIT:5> , use_fft = True , Mfft = <NUM_LIT:4> ) : <EOL> """<STR_LIT>""" <EOL> df *= freq_factor <EOL> f0 *= freq_factor <EOL> assert df > <NUM_LIT:0> <EOL> t , h = map ( np . ravel , np . broadcast_arrays ( t , h ) ) <EOL> if use_fft : <EOL> Mfft = int ( Mfft ) <EOL> assert ( Mfft > <NUM_LIT:0> ) <EOL> Nfft = bitceil ( N * oversampling ) <EOL> t0 = t . min ( ) <EOL> if f0 > <NUM_LIT:0> : <EOL> h = h * np . exp ( <NUM_LIT> * np . pi * f0 * ( t - t0 ) ) <EOL> tnorm = ( ( t - t0 ) * Nfft * df ) % Nfft <EOL> grid = extirpolate ( tnorm , h , Nfft , Mfft ) <EOL> fftgrid = np . fft . ifft ( grid ) <EOL> if t0 != <NUM_LIT:0> : <EOL> f = f0 + df * np . arange ( Nfft ) <EOL> fftgrid *= np . exp ( <NUM_LIT> * np . pi * t0 * f ) <EOL> fftgrid = fftgrid [ : N ] <EOL> C = Nfft * fftgrid . real <EOL> S = Nfft * fftgrid . imag <EOL> else : <EOL> f = f0 + df * np . arange ( N ) <EOL> C = np . dot ( h , np . cos ( <NUM_LIT:2> * np . pi * f * t [ : , np . newaxis ] ) ) <EOL> S = np . dot ( h , np . sin ( <NUM_LIT:2> * np . pi * f * t [ : , np . newaxis ] ) ) <EOL> return S , C <EOL> def lomb_scargle_fast ( t , y , dy = <NUM_LIT:1> , f0 = <NUM_LIT:0> , df = None , Nf = None , <EOL> center_data = True , fit_offset = True , <EOL> use_fft = True , freq_oversampling = <NUM_LIT:5> , nyquist_factor = <NUM_LIT:2> , <EOL> trig_sum_kwds = None ) : <EOL> """<STR_LIT>""" <EOL> t , y , dy = map ( np . ravel , np . broadcast_arrays ( t , y , dy ) ) <EOL> w = <NUM_LIT:1.> / ( dy ** <NUM_LIT:2> ) <EOL> w /= w . sum ( ) <EOL> if df is None : <EOL> peak_width = <NUM_LIT:1.> / ( t . max ( ) - t . min ( ) ) <EOL> df = peak_width / freq_oversampling <EOL> if Nf is None : <EOL> avg_Nyquist = <NUM_LIT:0.5> * len ( t ) / ( t . max ( ) - t . min ( ) ) <EOL> Nf = max ( <NUM_LIT:16> , ( nyquist_factor * avg_Nyquist - f0 ) / df ) <EOL> Nf = int ( Nf ) <EOL> assert ( df > <NUM_LIT:0> ) <EOL> assert ( Nf > <NUM_LIT:0> ) <EOL> freq = f0 + df * np . arange ( Nf ) <EOL> if center_data or fit_offset : <EOL> y = y - np . dot ( w , y ) <EOL> kwargs = dict . copy ( trig_sum_kwds or { } ) <EOL> kwargs . update ( f0 = f0 , df = df , use_fft = use_fft , N = Nf ) <EOL> Sh , Ch = trig_sum ( t , w * y , ** kwargs ) <EOL> S2 , C2 = trig_sum ( t , w , freq_factor = <NUM_LIT:2> , ** kwargs ) <EOL> if fit_offset : <EOL> S , C = trig_sum ( t , w , ** kwargs ) <EOL> with warnings . catch_warnings ( ) : <EOL> if f0 == <NUM_LIT:0> : <EOL> warnings . simplefilter ( "<STR_LIT:ignore>" ) <EOL> tan_2omega_tau = ( S2 - <NUM_LIT:2> * S * C ) / ( C2 - ( C * C - S * S ) ) <EOL> if np . isnan ( tan_2omega_tau [ <NUM_LIT:0> ] ) : <EOL> tan_2omega_tau [ <NUM_LIT:0> ] = <NUM_LIT:0> <EOL> else : <EOL> tan_2omega_tau = S2 / C2 <EOL> S2w = tan_2omega_tau / np . sqrt ( <NUM_LIT:1> + tan_2omega_tau * tan_2omega_tau ) <EOL> C2w = <NUM_LIT:1> / np . sqrt ( <NUM_LIT:1> + tan_2omega_tau * tan_2omega_tau ) <EOL> Cw = np . sqrt ( <NUM_LIT:0.5> ) * np . sqrt ( <NUM_LIT:1> + C2w ) <EOL> Sw = np . sqrt ( <NUM_LIT:0.5> ) * np . sign ( S2w ) * np . sqrt ( <NUM_LIT:1> - C2w ) <EOL> YY = np . dot ( w , y ** <NUM_LIT:2> ) <EOL> YC = Ch * Cw + Sh * Sw <EOL> YS = Sh * Cw - Ch * Sw <EOL> CC = <NUM_LIT:0.5> * ( <NUM_LIT:1> + C2 * C2w + S2 * S2w ) <EOL> SS = <NUM_LIT:0.5> * ( <NUM_LIT:1> - C2 * C2w - S2 * S2w ) <EOL> if fit_offset : <EOL> CC -= ( C * Cw + S * Sw ) ** <NUM_LIT:2> <EOL> SS -= ( S * Cw - C * Sw ) ** <NUM_LIT:2> <EOL> with warnings . catch_warnings ( ) : <EOL> if fit_offset and f0 == <NUM_LIT:0> : <EOL> warnings . simplefilter ( "<STR_LIT:ignore>" ) <EOL> power = ( YC * YC / CC + YS * YS / SS ) / YY <EOL> if np . isnan ( power [ <NUM_LIT:0> ] ) or np . isinf ( power [ <NUM_LIT:0> ] ) : <EOL> power [ <NUM_LIT:0> ] = <NUM_LIT:0> <EOL> return freq , power <EOL> class LombScargleFast ( LombScargle ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , optimizer = None , center_data = True , fit_offset = True , <EOL> use_fft = True , ls_kwds = None , Nterms = <NUM_LIT:1> , <EOL> fit_period = False , optimizer_kwds = None , <EOL> silence_warnings = False ) : <EOL> self . use_fft = use_fft <EOL> self . ls_kwds = ls_kwds <EOL> self . silence_warnings = silence_warnings <EOL> if Nterms != <NUM_LIT:1> : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> LombScargle . __init__ ( self , optimizer = optimizer , <EOL> center_data = center_data , fit_offset = fit_offset , <EOL> Nterms = <NUM_LIT:1> , regularization = None , <EOL> fit_period = fit_period , <EOL> optimizer_kwds = optimizer_kwds ) <EOL> def _score_frequency_grid ( self , f0 , df , N ) : <EOL> if not self . silence_warnings and self . t . size < <NUM_LIT:50> : <EOL> warnings . warn ( "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> freq , P = lomb_scargle_fast ( self . t , self . y , self . dy , <EOL> f0 = f0 , df = df , Nf = N , <EOL> center_data = self . center_data , <EOL> fit_offset = self . fit_offset , <EOL> use_fft = self . use_fft , <EOL> ** ( self . ls_kwds or { } ) ) <EOL> return P <EOL> def _score ( self , periods ) : <EOL> if not self . silence_warnings : <EOL> warnings . warn ( "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> return LombScargle . _score ( self , periods ) </s>
