<s> from __future__ import division <EOL> import warnings <EOL> import numpy as np <EOL> from numpy . testing import ( assert_allclose , assert_equal , <EOL> assert_warns , assert_no_warnings ) <EOL> from . . lomb_scargle_fast import ( extirpolate , bitceil , trig_sum , <EOL> lomb_scargle_fast ) <EOL> from . . import LombScargle , LombScargleFast <EOL> def _generate_data ( N = <NUM_LIT:100> , period = <NUM_LIT:1> , theta = [ <NUM_LIT:10> , <NUM_LIT:2> , <NUM_LIT:3> ] , dy = <NUM_LIT:1> , rseed = <NUM_LIT:0> ) : <EOL> """<STR_LIT>""" <EOL> rng = np . random . RandomState ( rseed ) <EOL> t = <NUM_LIT:20> * period * rng . rand ( N ) <EOL> omega = <NUM_LIT:2> * np . pi / period <EOL> y = theta [ <NUM_LIT:0> ] + theta [ <NUM_LIT:1> ] * np . sin ( omega * t ) + theta [ <NUM_LIT:2> ] * np . cos ( omega * t ) <EOL> dy = dy * ( <NUM_LIT:0.5> + rng . rand ( N ) ) <EOL> y += dy * rng . randn ( N ) <EOL> return t , y , dy <EOL> def test_extirpolate ( ) : <EOL> rng = np . random . RandomState ( <NUM_LIT:0> ) <EOL> x = <NUM_LIT:100> * rng . rand ( <NUM_LIT:50> ) <EOL> y = np . sin ( x ) <EOL> f = lambda x : np . sin ( x / <NUM_LIT:10> ) <EOL> def check_result ( N , M = <NUM_LIT:5> ) : <EOL> y_hat = extirpolate ( x , y , N , M ) <EOL> x_hat = np . arange ( len ( y_hat ) ) <EOL> assert_allclose ( np . dot ( f ( x ) , y ) , np . dot ( f ( x_hat ) , y_hat ) ) <EOL> for N in [ <NUM_LIT:100> , None ] : <EOL> yield check_result , N <EOL> def test_extirpolate_with_integers ( ) : <EOL> rng = np . random . RandomState ( <NUM_LIT:0> ) <EOL> x = <NUM_LIT:100> * rng . rand ( <NUM_LIT:50> ) <EOL> x [ : <NUM_LIT> ] = x [ : <NUM_LIT> ] . astype ( int ) <EOL> y = np . sin ( x ) <EOL> f = lambda x : np . sin ( x / <NUM_LIT:10> ) <EOL> def check_result ( N , M = <NUM_LIT:5> ) : <EOL> y_hat = extirpolate ( x , y , N , M ) <EOL> x_hat = np . arange ( len ( y_hat ) ) <EOL> assert_allclose ( np . dot ( f ( x ) , y ) , np . dot ( f ( x_hat ) , y_hat ) ) <EOL> for N in [ <NUM_LIT:100> , None ] : <EOL> yield check_result , N <EOL> def test_bitceil ( ) : <EOL> slow_bitceil = lambda N : int ( <NUM_LIT:2> ** np . ceil ( np . log2 ( N ) ) ) <EOL> for N in ( <NUM_LIT:2> ** np . arange ( <NUM_LIT:1> , <NUM_LIT:12> ) ) : <EOL> for offset in ( - <NUM_LIT:1> , <NUM_LIT:0> , <NUM_LIT:1> ) : <EOL> assert_equal ( slow_bitceil ( N + offset ) , bitceil ( N + offset ) ) <EOL> def test_trig_sum ( ) : <EOL> rng = np . random . RandomState ( <NUM_LIT:0> ) <EOL> t = <NUM_LIT:10> * rng . rand ( <NUM_LIT:50> ) <EOL> h = np . sin ( t ) <EOL> def check_result ( f0 , adjust_t , freq_factor , df = <NUM_LIT> ) : <EOL> tfit = t - t . min ( ) if adjust_t else t <EOL> S1 , C1 = trig_sum ( tfit , h , df , N = <NUM_LIT:1000> , use_fft = True , <EOL> f0 = f0 , freq_factor = freq_factor , oversampling = <NUM_LIT:10> ) <EOL> S2 , C2 = trig_sum ( tfit , h , df , N = <NUM_LIT:1000> , use_fft = False , <EOL> f0 = f0 , freq_factor = freq_factor , oversampling = <NUM_LIT:10> ) <EOL> assert_allclose ( S1 , S2 , atol = <NUM_LIT> ) <EOL> assert_allclose ( C1 , C2 , atol = <NUM_LIT> ) <EOL> for f0 in [ <NUM_LIT:0> , <NUM_LIT:1> ] : <EOL> for adjust_t in [ True , False ] : <EOL> for freq_factor in [ <NUM_LIT:1> , <NUM_LIT:2> ] : <EOL> yield check_result , f0 , adjust_t , freq_factor <EOL> def test_lomb_scargle_fast ( ) : <EOL> rng = np . random . RandomState ( <NUM_LIT:0> ) <EOL> t = <NUM_LIT:30> * rng . rand ( <NUM_LIT:100> ) <EOL> y = np . sin ( t ) <EOL> dy = <NUM_LIT:0.1> + <NUM_LIT:0.1> * rng . rand ( len ( t ) ) <EOL> y += dy * rng . randn ( len ( t ) ) <EOL> def check_results ( center_data , fit_offset ) : <EOL> freq1 , P1 = lomb_scargle_fast ( t , y , dy , <EOL> center_data = center_data , <EOL> fit_offset = fit_offset , use_fft = True ) <EOL> freq2 , P2 = lomb_scargle_fast ( t , y , dy , <EOL> center_data = center_data , <EOL> fit_offset = fit_offset , use_fft = False ) <EOL> assert_allclose ( freq1 , freq2 ) <EOL> assert_allclose ( P1 , P2 , atol = <NUM_LIT> ) <EOL> for center_data in [ True , False ] : <EOL> for fit_offset in [ True , False ] : <EOL> yield check_results , center_data , fit_offset <EOL> def test_find_best_period ( ) : <EOL> t , y , dy = _generate_data ( ) <EOL> def check_result ( use_fft , fit_offset , center_data ) : <EOL> with warnings . catch_warnings ( ) : <EOL> warnings . simplefilter ( '<STR_LIT:ignore>' ) <EOL> model = LombScargle ( fit_offset = fit_offset , <EOL> center_data = center_data ) . fit ( t , y , dy ) <EOL> fast = LombScargleFast ( fit_offset = fit_offset , <EOL> center_data = center_data , <EOL> use_fft = use_fft ) . fit ( t , y , dy ) <EOL> model . optimizer . period_range = ( <NUM_LIT> , <NUM_LIT:1.0> ) <EOL> fast . optimizer . period_range = ( <NUM_LIT> , <NUM_LIT:1.0> ) <EOL> assert_allclose ( model . best_period , fast . best_period , atol = <NUM_LIT> ) <EOL> for use_fft in [ True , False ] : <EOL> for fit_offset in [ True , False ] : <EOL> for center_data in [ True , False ] : <EOL> yield check_result , use_fft , fit_offset , center_data <EOL> def test_power ( ) : <EOL> t , y , dy = _generate_data ( ) <EOL> f0 = <NUM_LIT> <EOL> df = <NUM_LIT> <EOL> N = <NUM_LIT> <EOL> def check_result ( use_fft , fit_offset , center_data ) : <EOL> model = LombScargle ( fit_offset = fit_offset , <EOL> center_data = center_data ) . fit ( t , y , dy ) <EOL> fast = LombScargleFast ( fit_offset = fit_offset , <EOL> center_data = center_data , <EOL> use_fft = use_fft ) . fit ( t , y , dy ) <EOL> assert_allclose ( model . score_frequency_grid ( f0 , df , N ) , <EOL> fast . score_frequency_grid ( f0 , df , N ) , <EOL> atol = <NUM_LIT> ) <EOL> for use_fft in [ True , False ] : <EOL> for fit_offset in [ True , False ] : <EOL> for center_data in [ True , False ] : <EOL> yield check_result , use_fft , fit_offset , center_data <EOL> def check_warn_on_small_data ( ) : <EOL> t , y , dy = _generate_data ( <NUM_LIT:20> ) <EOL> model = LombScargleFast ( ) <EOL> assert_warns ( UserWarning , model . score_frequency_grid , <NUM_LIT> , <NUM_LIT> , <NUM_LIT> ) <EOL> model = LombScargleFast ( silence_warnings = True ) <EOL> assert_no_warnings ( model . score_frequency_grid , <NUM_LIT> , <NUM_LIT> , <NUM_LIT> ) </s>
