def main(): print("Starting tests for wotan...") numpy.testing.assert_almost_equal(t14(R_s=1, M_s=1, P=365), 0.6490025258902046) numpy.testing.assert_almost_equal( t14(R_s=1, M_s=1, P=365, small_planet=True), 0.5403690143737738) print("Transit duration correct.") numpy.random.seed(seed=0) # reproducibility print("Slide clipper...") points = 1000 time = numpy.linspace(0, 30, points) flux = 1 + numpy.sin(time) / points noise = numpy.random.normal(0, 0.0001, points) flux += noise for i in range(points): if i % 75 == 0: flux[i:i + 5] -= 0.0004 # Add some transits flux[i + 50:i + 52] += 0.0002 # and flares clipped = slide_clip(time, flux, window_length=0.5, low=3, high=2, method='mad', center='median') numpy.testing.assert_almost_equal(numpy.nansum(clipped), 948.9926368754939) """ import matplotlib.pyplot as plt plt.scatter(time, flux, s=3, color='black') plt.scatter(time, clipped, s=3, color='orange') plt.show() """ # TESS test print('Loading TESS data from archive.stsci.edu...') path = 'https://archive.stsci.edu/hlsps/tess-data-alerts/' filename = "hlsp_tess-data-alerts_tess_phot_00062483237-s01_tess_v1_lc.fits" #path = 'P:/P/Dok/tess_alarm/' #filename = "hlsp_tess-data-alerts_tess_phot_00062483237-s01_tess_v1_lc.fits" #filename = 'P:/P/Dok/tess_alarm/hlsp_tess-data-alerts_tess_phot_00077031414-s02_tess_v1_lc.fits' #filename = 'tess2018206045859-s0001-0000000201248411-111-s_llc.fits' time, flux = load_file(path + filename) window_length = 0.5 print("Detrending 1 (biweight)...") flatten_lc, trend_lc = flatten(time, flux, window_length, edge_cutoff=1, break_tolerance=0.1, return_trend=True, cval=5.0) numpy.testing.assert_equal(len(trend_lc), 20076) numpy.testing.assert_almost_equal(numpy.nanmax(trend_lc), 28754.985299070882) numpy.testing.assert_almost_equal(numpy.nanmin(trend_lc), 28615.108124724477) numpy.testing.assert_almost_equal(trend_lc[500], 28671.686308143515) numpy.testing.assert_equal(len(flatten_lc), 20076) numpy.testing.assert_almost_equal(numpy.nanmax(flatten_lc), 1.0034653549250616) numpy.testing.assert_almost_equal(numpy.nanmin(flatten_lc), 0.996726610702177) numpy.testing.assert_almost_equal(flatten_lc[500], 1.000577429565131) print("Detrending 2 (andrewsinewave)...") flatten_lc, trend_lc = flatten(time, flux, window_length, method='andrewsinewave', return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18119.15471987987, decimal=2) print("Detrending 3 (welsch)...") flatten_lc, trend_lc = flatten(time, flux, window_length, method='welsch', return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18119.16764691235, decimal=2) print("Detrending 4 (hodges)...") flatten_lc, trend_lc = flatten(time[:1000], flux[:1000], window_length, method='hodges', return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 994.0110525909206, decimal=2) print("Detrending 5 (median)...") flatten_lc, trend_lc = flatten(time, flux, window_length, method='median', return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18119.122065014355, decimal=2) print("Detrending 6 (mean)...") flatten_lc, trend_lc = flatten(time, flux, window_length, method='mean', return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18119.032473037714, decimal=2) print("Detrending 7 (trim_mean)...") flatten_lc, trend_lc = flatten(time, flux, window_length, method='trim_mean', return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18119.095164910334, decimal=2) print("Detrending 8 (supersmoother)...") flatten_lc, trend_lc = flatten(time, flux, window_length, method='supersmoother', return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18123.00632204841, decimal=2) print("Detrending 9 (hspline)...") flatten_lc, trend_lc = flatten(time, flux, window_length, method='hspline', return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18123.07625225313, decimal=2) print("Detrending 10 (cofiam)...") flatten_lc, trend_lc = flatten(time[:2000], flux[:2000], window_length, method='cofiam', return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 1948.9999999987976, decimal=2) print("Detrending 11 (savgol)...") flatten_lc, trend_lc = flatten(time, flux, window_length=301, method='savgol', return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18123.003465539354, decimal=2) print("Detrending 12 (medfilt)...") flatten_lc, trend_lc = flatten(time, flux, window_length=301, method='medfilt', return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18123.22609806557, decimal=2) print("Detrending 12 (gp squared_exp)...") flatten_lc, trend_lc1 = flatten(time[:2000], flux[:2000], method='gp', kernel='squared_exp', kernel_size=10, return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 1948.99958552324, decimal=2) print("Detrending 13 (gp squared_exp robust)...") flatten_lc, trend_lc1 = flatten(time[:2000], flux[:2000], method='gp', kernel='squared_exp', kernel_size=10, robust=True, return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 1948.8820772313468, decimal=2) print("Detrending 14 (gp matern)...") flatten_lc, trend_lc2 = flatten(time[:2000], flux[:2000], method='gp', kernel='matern', kernel_size=10, return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 1949.0001583058202, decimal=2) print("Detrending 15 (gp periodic)...") flatten_lc, trend_lc2 = flatten(time[:2000], flux[:2000], method='gp', kernel='periodic', kernel_size=1, kernel_period=10, return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 1948.9999708985608, decimal=2) time_synth = numpy.linspace(0, 30, 200) flux_synth = numpy.sin(time_synth) + numpy.random.normal(0, 0.1, 200) flux_synth = 1 + flux_synth / 100 time_synth *= 1.5 print("Detrending 16 (gp periodic_auto)...") flatten_lc, trend_lc2 = flatten(time_synth, flux_synth, method='gp', kernel='periodic_auto', kernel_size=1, return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 200, decimal=1) print("Detrending 17 (rspline)...") flatten_lc, trend_lc2 = flatten(time, flux, method='rspline', window_length=1, return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18121.812790732245, decimal=2) print("Detrending 18 (huber)...") flatten_lc, trend_lc = flatten(time[:1000], flux[:1000], method='huber', window_length=0.5, edge_cutoff=0, break_tolerance=0.4, return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 994.01102, decimal=2) print("Detrending 19 (winsorize)...") flatten_lc, trend_lc2 = flatten(time, flux, method='winsorize', window_length=0.5, edge_cutoff=0, break_tolerance=0.4, proportiontocut=0.1, return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18119.064587196448, decimal=2) print("Detrending 20 (pspline)...") flatten_lc, trend_lc = flatten(time, flux, method='pspline', return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18121.832133916843, decimal=2) print("Detrending 21 (hampelfilt)...") flatten_lc, trend_lc5 = flatten(time, flux, method='hampelfilt', window_length=0.5, cval=3, return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18119.158072498867, decimal=2) print("Detrending 22 (lowess)...") flatten_lc, trend_lc1 = flatten(time, flux, method='lowess', window_length=1, return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18123.08085676265, decimal=2) print("Detrending 23 (huber_psi)...") flatten_lc, trend_lc1 = flatten(time, flux, method='huber_psi', window_length=0.5, return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18119.122065014355, decimal=2) print("Detrending 24 (tau)...") flatten_lc, trend_lc2 = flatten(time, flux, method='tau', window_length=0.5, return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18119.02772621119, decimal=2) import numpy as np points = 1000 time = np.linspace(0, 30, points) flux = 1 + np.sin(time) / points noise = np.random.normal(0, 0.0001, points) flux += noise for i in range(points): if i % 75 == 0: flux[i:i + 5] -= 0.0004 # Add some transits flux[i + 50:i + 52] += 0.0002 # and flares print("Detrending 25a (hampel 17A)...") flatten_lc, trend_lc1 = flatten(time, flux, method='hampel', cval=(1.7, 3.4, 8.5), window_length=0.5, return_trend=True) print("Detrending 25b (hampel 25A)...") flatten_lc, trend_lc2 = flatten(time, flux, method='hampel', cval=(2.5, 4.5, 9.5), window_length=0.5, return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 997.9994362858843, decimal=2) print("Detrending 26 (ramsay)...") flatten_lc, trend_lc3 = flatten(time, flux, method='ramsay', cval=0.3, window_length=0.5, return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 997.9974021484584, decimal=2) """ import matplotlib.pyplot as plt plt.scatter(time, flux, s=1, color='black') plt.plot(time[:len(trend_lc1)], trend_lc1, color='blue', linewidth=2) plt.plot(time[:len(trend_lc1)], trend_lc2, color='red', linewidth=2, linestyle='dashed') plt.show() plt.close() #plt.scatter(time, flatten_lc, s=1, color='black') #plt.show() """ print('All tests completed.')
def main(): print("Starting tests for wotan...") numpy.testing.assert_almost_equal(t14(R_s=1, M_s=1, P=365), 0.6490025258902046) numpy.testing.assert_almost_equal( t14(R_s=1, M_s=1, P=365, small_planet=True), 0.5403690143737738 ) print("Transit duration correct.") numpy.random.seed(seed=0) # reproducibility print("Slide clipper...") points = 1000 time = numpy.linspace(0, 30, points) flux = 1 + numpy.sin(time) / points noise = numpy.random.normal(0, 0.0001, points) flux += noise for i in range(points): if i % 75 == 0: flux[i : i + 5] -= 0.0004 # Add some transits flux[i + 50 : i + 52] += 0.0002 # and flares clipped = slide_clip( time, flux, window_length=0.5, low=3, high=2, method="mad", center="median" ) numpy.testing.assert_almost_equal(numpy.nansum(clipped), 948.9926368754939) """ import matplotlib.pyplot as plt plt.scatter(time, flux, s=3, color='black') plt.scatter(time, clipped, s=3, color='orange') plt.show() """ # TESS test print("Loading TESS data from archive.stsci.edu...") path = "https://archive.stsci.edu/hlsps/tess-data-alerts/" # path = 'P:/P/Dok/tess_alarm/' filename = "hlsp_tess-data-alerts_tess_phot_00062483237-s01_tess_v1_lc.fits" time, flux = load_file(path + filename) window_length = 0.5 print("Detrending 1 (biweight)...") flatten_lc, trend_lc = flatten( time, flux, window_length, edge_cutoff=1, break_tolerance=0.1, return_trend=True, cval=5.0, ) numpy.testing.assert_equal(len(trend_lc), 20076) numpy.testing.assert_almost_equal( numpy.nanmax(trend_lc), 28755.03811866676, decimal=2 ) numpy.testing.assert_almost_equal( numpy.nanmin(trend_lc), 28615.110229935075, decimal=2 ) numpy.testing.assert_almost_equal(trend_lc[500], 28671.650565730513, decimal=2) numpy.testing.assert_equal(len(flatten_lc), 20076) numpy.testing.assert_almost_equal( numpy.nanmax(flatten_lc), 1.0034653549250616, decimal=2 ) numpy.testing.assert_almost_equal( numpy.nanmin(flatten_lc), 0.996726610702177, decimal=2 ) numpy.testing.assert_almost_equal(flatten_lc[500], 1.000577429565131, decimal=2) print("Detrending 2 (andrewsinewave)...") flatten_lc, trend_lc = flatten( time, flux, window_length, method="andrewsinewave", return_trend=True ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 18123.15456308313, decimal=2 ) print("Detrending 3 (welsch)...") flatten_lc, trend_lc = flatten( time, flux, window_length, method="welsch", return_trend=True ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 18123.16770590837, decimal=2 ) print("Detrending 4 (hodges)...") flatten_lc, trend_lc = flatten( time[:1000], flux[:1000], window_length, method="hodges", return_trend=True ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 996.0113241694287, decimal=2 ) print("Detrending 5 (median)...") flatten_lc, trend_lc = flatten( time, flux, window_length, method="median", return_trend=True ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 18123.12166552401, decimal=2 ) print("Detrending 6 (mean)...") flatten_lc, trend_lc = flatten( time, flux, window_length, method="mean", return_trend=True ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 18123.032058753546, decimal=2 ) print("Detrending 7 (trim_mean)...") flatten_lc, trend_lc = flatten( time, flux, window_length, method="trim_mean", return_trend=True ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 18123.094751124332, decimal=2 ) print("Detrending 8 (supersmoother)...") flatten_lc, trend_lc = flatten( time, flux, window_length, method="supersmoother", return_trend=True ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 18123.00632204841, decimal=2 ) print("Detrending 9 (hspline)...") flatten_lc, trend_lc = flatten( time, flux, window_length, method="hspline", return_trend=True ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 18123.082601463717, decimal=1 ) print("Detrending 10 (cofiam)...") flatten_lc, trend_lc = flatten( time[:2000], flux[:2000], window_length, method="cofiam", return_trend=True ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 1948.9999999987976, decimal=1 ) print("Detrending 11 (savgol)...") flatten_lc, trend_lc = flatten( time, flux, window_length=301, method="savgol", return_trend=True ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 18123.003465539354, decimal=1 ) print("Detrending 12 (medfilt)...") flatten_lc, trend_lc = flatten( time, flux, window_length=301, method="medfilt", return_trend=True ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 18123.22609806557, decimal=1 ) print("Detrending 12 (gp squared_exp)...") flatten_lc, trend_lc1 = flatten( time[:2000], flux[:2000], method="gp", kernel="squared_exp", kernel_size=10, return_trend=True, ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 1948.9672036416687, decimal=2 ) print("Detrending 13 (gp squared_exp robust)...") flatten_lc, trend_lc1 = flatten( time[:2000], flux[:2000], method="gp", kernel="squared_exp", kernel_size=10, robust=True, return_trend=True, ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 1948.8820772313468, decimal=2 ) print("Detrending 14 (gp matern)...") flatten_lc, trend_lc2 = flatten( time[:2000], flux[:2000], method="gp", kernel="matern", kernel_size=10, return_trend=True, ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 1948.9672464898367, decimal=2 ) print("Detrending 15 (gp periodic)...") flatten_lc, trend_lc2 = flatten( time[:2000], flux[:2000], method="gp", kernel="periodic", kernel_size=1, kernel_period=10, return_trend=True, ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 1948.9999708985608, decimal=2 ) time_synth = numpy.linspace(0, 30, 200) flux_synth = numpy.sin(time_synth) + numpy.random.normal(0, 0.1, 200) flux_synth = 1 + flux_synth / 100 time_synth *= 1.5 print("Detrending 16 (gp periodic_auto)...") flatten_lc, trend_lc2 = flatten( time_synth, flux_synth, method="gp", kernel="periodic_auto", kernel_size=1, return_trend=True, ) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 200, decimal=1) print("Detrending 17 (rspline)...") flatten_lc, trend_lc2 = flatten( time, flux, method="rspline", window_length=1, return_trend=True ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 18121.812790732245, decimal=2 ) """ print("Detrending 18 (huber)...") flatten_lc, trend_lc = flatten( time[:1000], flux[:1000], method='huber', window_length=0.5, edge_cutoff=0, break_tolerance=0.4, return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 996.0112964009066, decimal=2) """ print("Detrending 19 (winsorize)...") flatten_lc, trend_lc2 = flatten( time, flux, method="winsorize", window_length=0.5, edge_cutoff=0, break_tolerance=0.4, proportiontocut=0.1, return_trend=True, ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 18123.064149766662, decimal=2 ) print("Detrending 20 (pspline)...") flatten_lc, trend_lc = flatten(time, flux, method="pspline", return_trend=True) # import matplotlib.pyplot as plt # plt.scatter(time, flux, s=3, color='black') # plt.plot(time, trend_lc) # plt.show() numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 18122.740535799767, decimal=2 ) print("Detrending 21 (hampelfilt)...") flatten_lc, trend_lc5 = flatten( time, flux, method="hampelfilt", window_length=0.5, cval=3, return_trend=True ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 18123.157973016467, decimal=2 ) print("Detrending 22 (lowess)...") flatten_lc, trend_lc1 = flatten( time, flux, method="lowess", window_length=1, return_trend=True ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 18123.039744125545, decimal=2 ) print("Detrending 23 (huber_psi)...") flatten_lc, trend_lc1 = flatten( time, flux, method="huber_psi", window_length=0.5, return_trend=True ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 18123.110893063527, decimal=2 ) print("Detrending 24 (tau)...") flatten_lc, trend_lc2 = flatten( time, flux, method="tau", window_length=0.5, return_trend=True ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 18123.026005725977, decimal=2 ) print("Detrending 25 (cosine)...") flatten_lc, trend_lc2 = flatten( time, flux, method="cosine", window_length=0.5, return_trend=True ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 18122.999999974905, decimal=2 ) print("Detrending 25 (cosine robust)...") flatten_lc, trend_lc2 = flatten( time, flux, method="cosine", robust=True, window_length=0.5, return_trend=True ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 18122.227938535038, decimal=2 ) import numpy as np points = 1000 time = numpy.linspace(0, 30, points) flux = 1 + numpy.sin(time) / points noise = numpy.random.normal(0, 0.0001, points) flux += noise for i in range(points): if i % 75 == 0: flux[i : i + 5] -= 0.0004 # Add some transits flux[i + 50 : i + 52] += 0.0002 # and flares print("Detrending 26 (hampel 17A)...") flatten_lc, trend_lc1 = flatten( time, flux, method="hampel", cval=(1.7, 3.4, 8.5), window_length=0.5, return_trend=True, ) print("Detrending 27 (hampel 25A)...") flatten_lc, trend_lc2 = flatten( time, flux, method="hampel", cval=(2.5, 4.5, 9.5), window_length=0.5, return_trend=True, ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 999.9992212031945, decimal=2 ) print("Detrending 28 (ramsay)...") flatten_lc, trend_lc3 = flatten( time, flux, method="ramsay", cval=0.3, window_length=0.5, return_trend=True ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 999.9970566765148, decimal=2 ) print("Detrending 29 (ridge)...") flatten_lc, trend_lc1 = flatten( time, flux, window_length=0.5, method="ridge", return_trend=True, cval=1 ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 999.9999958887022, decimal=1 ) print("Detrending 30 (lasso)...") flatten_lc, trend_lc2 = flatten( time, flux, window_length=0.5, method="lasso", return_trend=True, cval=1 ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 999.9999894829843, decimal=1 ) print("Detrending 31 (elasticnet)...") flatten_lc, trend_lc3 = flatten( time, flux, window_length=0.5, method="elasticnet", return_trend=True, cval=1 ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 999.9999945063338, decimal=1 ) # Test of transit mask print("Testing transit_mask") filename = "hlsp_tess-data-alerts_tess_phot_00207081058-s01_tess_v1_lc.fits" time, flux = load_file(path + filename) from wotan import transit_mask mask = transit_mask(time=time, period=14.77338, duration=0.21060, T0=1336.141095) numpy.testing.assert_almost_equal(numpy.sum(mask), 302, decimal=1) print("Detrending 32 (transit_mask cosine)") flatten_lc1, trend_lc1 = flatten( time, flux, method="cosine", window_length=0.4, return_trend=True, robust=True, mask=mask, ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc1), 18119.281265446625, decimal=1 ) print("Detrending 33 (transit_mask lowess)") flatten_lc2, trend_lc2 = flatten( time, flux, method="lowess", window_length=0.8, return_trend=True, robust=True, mask=mask, ) # print(numpy.nansum(flatten_lc2)) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc2), 18119.30865711536, decimal=1 ) print("Detrending 34 (transit_mask GP)") mask = transit_mask(time=time[:2000], period=100, duration=0.3, T0=1327.4) flatten_lc2, trend_lc2 = flatten( time[:2000], flux[:2000], method="gp", kernel="matern", kernel_size=0.8, return_trend=True, robust=True, mask=mask, ) # print(numpy.nansum(flatten_lc2)) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc2), 1948.9000170463796, decimal=1 ) print("Detrending 35 (pspline full features)") flatten_lc, trend_lc, nsplines = flatten( time, flux, method="pspline", max_splines=100, edge_cutoff=0.5, return_trend=True, return_nsplines=True, verbose=True, ) # print('lightcurve was split into', len(nsplines), 'segments') # print('chosen number of splines', nsplines) """ import matplotlib.pyplot as plt plt.scatter(time, flux, s=3, color='black') plt.plot(time, trend_lc) plt.show() plt.scatter(time, flatten_lc, s=3, color='black') plt.show() """ numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 16678.312693036027, decimal=1 ) print("Detrending 36 (pspline variable PSPLINES_STDEV_CUT)") flatten_lc, trend_lc, nsplines = flatten( time, flux, method="pspline", max_splines=100, edge_cutoff=0.5, stdev_cut=3, return_trend=True, return_nsplines=True, verbose=True, ) numpy.testing.assert_almost_equal( numpy.nansum(flatten_lc), 16678.292210380347, decimal=2 ) """ import matplotlib.pyplot as plt plt.scatter(time, flux, s=1, color='black') plt.plot(time, trend_lc, color='red', linewidth=2, linestyle='dashed') plt.show() plt.close() """ print("All tests completed.")
t0 = time.time() lcf = lightkurve.search_lightcurve("TIC 251848941", mission="TESS", cadence="short", sector=[2], author="SPOC")\ .download_all() lc = lcf.stitch().remove_nans() lc = lc.remove_outliers(sigma_lower=float('inf'), sigma_upper=3) lc_time = lc.time.value flux = lc.flux.value cadence = 2 window_length = 25 / cadence flux = savgol_filter(flux, 11, 3) R_s = 1.1 M_s = 1.3 P_min = 0.5 P_max = 22 ld_coefficients = [0.2, 0.1] min_duration = wotan.t14(R_s, M_s, P_min, True) max_duration = wotan.t14(R_s, M_s, P_max, True) duration_grid = np.arange(min_duration * 24 * 60 // cadence, max_duration * 24 * 60 // cadence, 1) def calculate_semi_major_axis(period, star_mass): G = 6.674e-11 period_seconds = period * 24. * 3600. mass_kg = star_mass * 2.e30 a1 = (G * mass_kg * period_seconds**2 / 4. / (np.pi**2))**(1. / 3.) return a1 / 1.496e11 a_au = calculate_semi_major_axis(0.5, M_s) a_Rs = a_au / (R_s * 0.00465047)
def analyse(det_met, time_i, flatten_i, n_tra, ab, mass, mass_min, mass_max, radius, radius_min, radius_max, id_run): # --------------------------------------------------------------------------------------------------------------------------------------------------------- # Detrending of the data using WOTAN # --------------------------------------------------------------------------------------------------------------------------------------------------------- # By default is used a 'biweight' method. # It will detrend the lightcurve 'N_detrends' independent times from wotan import flatten if det_met == 2: wl_min = 1 wl_max = 12 else: tdur = t14( R_s=radius, M_s=mass, P=P_protec, small_planet=True ) # we define the typical duration of a small planet in this star wl_min = 3 * tdur # minimum transit duration wl_max = 20 * tdur # maximum transit duration wl_step = (wl_max - wl_min) / N_detrends wl = np.arange( wl_min, wl_max, wl_step) # we define all the posibles window_length that we apply global_flatten_lc = np.zeros((len(wl), len(flatten_i))) global_trend_lc = np.zeros((len(wl), len(flatten_i))) for i in range(0, len(wl)): if det_met == 2: flatten_lc, trend_lc = flatten(time_i, flatten_i, method='gp', kernel='matern', kernel_size=wl[i], return_trend=True, break_tolerance=0.5) else: flatten_lc, trend_lc = flatten(time_i, flatten_i, window_length=wl[i], return_trend=True, method='biweight', break_tolerance=0.5) global_flatten_lc[i] = flatten_lc global_trend_lc[i] = trend_lc ## save in the log file all the information concerning the detrendins applied warnings.filterwarnings("ignore") global_final = np.zeros((len(wl), len(flatten_i))) logprint('\n MODELS IN THE DETRENDING - Run ' + str(id_run)) logprint('========================================\n') if det_met == 2: det_model = 'kernel_size:' else: det_model = 'window_size:' bins = len(time_i) * 2 / minutes bin_means, bin_edges, binnumber = stats.binned_statistic(time_i, flatten_i, statistic='mean', bins=bins) bin_stds, _, _ = stats.binned_statistic(time_i, flatten_i, statistic='std', bins=bins) bin_width = (bin_edges[1] - bin_edges[0]) bin_centers = bin_edges[1:] - bin_width / 2 logprint('PDCSAP_FLUX_' + str(id_run), '\t ', ' ------ ', '\t', '------', '\t', '\t', 'RMS (ppm):', np.std(flatten_i) * 1e6, '\t', 'RMS_10min (ppm):', np.std(bin_means[~np.isnan(bin_means)]) * 1e6) for i in range(len(wl)): flatten = sigma_clip(global_flatten_lc[i], sigma_lower=20, sigma_upper=3) global_final[i] = flatten bins_i = len(time_i) * 2 / minutes bin_means_i, bin_edges_i, binnumber_i = stats.binned_statistic( time_i, flatten, statistic='mean', bins=bins_i) bin_stds_i, _, _ = stats.binned_statistic(time_i, flatten, statistic='std', bins=bins_i) bin_width_i = (bin_edges_i[1] - bin_edges_i[0]) bin_centers_i = bin_edges_i[1:] - bin_width_i / 2 logprint('flatten_lc%s' % i, '\t ', 'trend_lc%s' % i, '\t', det_model, format(wl[i], '0.4f'), '\t', 'RMS (ppm):', np.std(flatten) * 1e6, '\t', 'RMS_10min (ppm):', np.std(bin_means_i[~np.isnan(bin_means_i)]) * 1e6) ## save in a plot all the detrendings and all the data to inspect visually. cases = np.zeros((len(wl), 1)) for i in range(len(wl)): cases[i] = wl[i] figsize = (8, 8) #x,y cols = 3 rows = len(cases) // cols shift = 2 * (1.0 - (np.min(flatten_i)) ) #shift in the between the raw and detrended data ylim_max = 1.0 + 3 * (np.max(flatten_i) - 1.0 ) #shift in the between the raw and detrended data ylim_min = 1.0 - 2.0 * shift #shift in the between the raw and detrended data fig1, axs = plt.subplots(rows, cols, figsize=figsize, constrained_layout=True) axs = trim_axs(axs, len(cases)) for ax, case in zip(axs, cases): if det_met == 2: ax.set_title('ks=%s' % str(np.around(case, decimals=4))) else: ax.set_title('ws=%s' % str(np.around(case, decimals=4))) bins_i = len(time_i) * 2 / minutes bin_means_i, bin_edges_i, binnumber_i = stats.binned_statistic( time_i, global_final[np.nonzero(cases == case)[0][0]], statistic='mean', bins=bins_i) bin_stds_i, _, _ = stats.binned_statistic( time_i, global_final[np.nonzero(cases == case)[0][0]], statistic='std', bins=bins_i) bin_width_i = (bin_edges_i[1] - bin_edges_i[0]) bin_centers_i = bin_edges_i[1:] - bin_width_i / 2 ax.plot(time_i, flatten_i, linewidth=0.05, color='black', alpha=0.75, rasterized=True) ax.plot(time_i, global_trend_lc[np.nonzero(cases == case)[0][0]], linewidth=1, color='orange', alpha=1.0) ax.plot(time_i, global_final[np.nonzero(cases == case)[0][0]] - shift, linewidth=0.05, color='teal', alpha=0.75, rasterized=True) ax.plot(bin_centers_i, bin_means_i - shift, marker='.', markersize=2, color='firebrick', alpha=0.5, linestyle='none', rasterized=True) ax.set_ylim(ylim_min, ylim_max) plt.savefig('Detrends_' + 'run_' + str(id_run) + '_TIC' + str(TIC_ID) + '.png', dpi=200) # --------------------------------------------------------------------------------------------------------------------------------------------------------- # Search of signals in the raw data # --------------------------------------------------------------------------------------------------------------------------------------------------------- logprint('\n SEARCH OF SIGNALS - Run', id_run) logprint('=================================\n') logprint('PDCSAP_FLUX_' + str(id_run), '\t', 'Period', '\t', 'Per_err', '\t', 'N.Tran', '\t', 'Mean Depth(ppt)', '\t', 'T. dur(min)', '\t', 'T0', '\t', 'SNR', '\t', 'SDE', '\t', 'FAP\n') model = transitleastsquares(time_i, flatten_i) results_pdcsap = model.power(u=ab, M_star=mass, M_star_min=mass_min, M_star_max=mass_max, R_star=radius, R_star_min=radius_min, R_star_max=radius_max, period_min=Pmin, period_max=Pmax, n_transits_min=n_tra, show_progress_bar=False) x = [] if results_pdcsap.T0 != 0: for j in range(0, len(results_pdcsap.transit_depths)): # print (results.transit_depths[i]) x = np.append(x, results_pdcsap.transit_depths[j]) x = x[~np.isnan(x)] depth = (1. - np.mean(x)) * 100 / 0.1 #we change to ppt units else: depth = results_pdcsap.transit_depths logprint('-----', '\t ', format(results_pdcsap.period, '.5f'), '\t ', format(results_pdcsap.period_uncertainty, '.6f'), '\t ', results_pdcsap.distinct_transit_count, '\t', format(depth, '.3f'), '\t', format(results_pdcsap.duration * 24 * 60, '.1f'), '\t', results_pdcsap.T0, '\t ', format(results_pdcsap.snr, '.3f'), '\t ', format(results_pdcsap.SDE, '.3f'), '\t ', results_pdcsap.FAP) logprint('\n') # --------------------------------------------------------------------------------------------------------------------------------------------------------- # Saving a plot of the best signal from raw data # --------------------------------------------------------------------------------------------------------------------------------------------------------- fig = save_plot(time_i, global_final, results_pdcsap, i) fig.suptitle('PDCSAP_FLUX_' + str(id_run) + ' ## SNR:' + str(format(results_pdcsap.snr, '.3f')) + ' ## SDE:' + str(format(results_pdcsap.SDE, '.3f')) + ' ## FAP:' + str(results_pdcsap.FAP)) plt.savefig('Run_' + str(id_run) + '_PDCSAP-FLUX_' + 'TIC' + str(TIC_ID) + '.png', bbox_inches='tight', dpi=200) # --------------------------------------------------------------------------------------------------------------------------------------------------------- # Search of signals in the detrended data # --------------------------------------------------------------------------------------------------------------------------------------------------------- if det_met == 2: logprint('Kernel_size', '\t', 'Period', '\t', 'Per_err', '\t', 'N.Tran', '\t', 'Mean Depth(ppt)', '\t', 'T. dur(min)', '\t', 'T0', '\t', 'SNR', '\t', 'SDE', '\t', 'FAP\n') else: logprint('Window_size', '\t', 'Period', '\t', 'Per_err', '\t', 'N.Tran', '\t', 'Mean Depth(ppt)', '\t', 'T. dur(min)', '\t', 'T0', '\t', 'SNR', '\t', 'SDE', '\t', 'FAP\n') SNR = np.zeros((len(wl), 1)) SDE = np.zeros((len(wl), 1)) FAP = np.zeros((len(wl), 1)) periods = np.zeros((len(wl), 1)) per_err = np.zeros((len(wl), 1)) durations = np.zeros((len(wl), 1)) tos = np.zeros((len(wl), 1)) for i in range(len(wl)): print(i) model = transitleastsquares(time_i, global_final[i]) results = model.power(u=ab, M_star=mass, M_star_min=mass_min, M_star_max=mass_max, R_star=radius, R_star_min=radius_min, R_star_max=radius_max, period_min=Pmin, period_max=Pmax, n_transits_min=n_tra, show_progress_bar=False) SNR[i] = results.snr SDE[i] = results.SDE FAP[i] = results.FAP periods[i] = results.period per_err[i] = results.period_uncertainty durations[i] = results.duration tos[i] = results.T0 x = [] if results.T0 != 0: for j in range(0, len(results.transit_depths)): # print (results.transit_depths[i]) x = np.append(x, results.transit_depths[j]) x = x[~np.isnan(x)] depth = (1. - np.mean(x)) * 100 / 0.1 #we change to ppt units else: depth = results.transit_depths logprint(format(wl[i], '.4f'), '\t ', format(results.period, '.5f'), '\t ', format(results.period_uncertainty, '.6f'), '\t ', results.distinct_transit_count, '\t', format(depth, '.3f'), '\t', format(results.duration * 24 * 60, '.1f'), '\t', results.T0, '\t ', format(results.snr, '.3f'), '\t ', format(results.SDE, '.3f'), '\t ', results.FAP) # ----------------------------------------------------------------------------------------------------------------------------------------------------- # Saving a plot of the best signal from detrended data # ----------------------------------------------------------------------------------------------------------------------------------------------------- fig = save_plot(time_i, global_final, results, i) if det_met == 2: fig.suptitle('Run ' + str(id_run) + ' ## kernel_size:' + str(format(wl[i], '.4f')) + ' ## SNR:' + str(format(results.snr, '.3f')) + ' ## SDE:' + str(format(results.SDE, '.3f')) + ' ## FAP:' + str(results.FAP)) #plt.savefig('TIC'+str(TIC_ID)+'_ks='+str(format(wl[i],'.4f'))+'_run_'+str(id_run)+'.png', bbox_inches='tight', dpi=200) plt.savefig('Run_' + str(id_run) + '_ks=' + str(format(wl[i], '.4f')) + '_TIC' + str(TIC_ID) + '.png', bbox_inches='tight', dpi=200) else: fig.suptitle('Run ' + str(id_run) + ' ## window_size:' + str(format(wl[i], '.4f')) + ' ## SNR:' + str(format(results.snr, '.3f')) + ' ## SDE:' + str(format(results.SDE, '.3f')) + ' ## FAP:' + str(results.FAP)) #plt.savefig('TIC'+str(TIC_ID)+'_ws='+str(format(wl[i],'.4f'))+'_run_'+str(id_run)+'.png', bbox_inches='tight', dpi=200) plt.savefig('Run_' + str(id_run) + '_ws=' + str(format(wl[i], '.4f')) + '_TIC' + str(TIC_ID) + '.png', bbox_inches='tight', dpi=200) SNR = np.nan_to_num(SNR) a = np.nanargmax(SNR) #check the maximum SRN if results_pdcsap.snr > SNR[a]: logprint('\nBest Signal -->', '\t', 'PDCSAP_FLUX', '\t', 'SNR:', format(results_pdcsap.snr, '.3f')) if ( results_pdcsap.snr > SNR_min ): # and results_pdcsap.SDE > SDE_min and results_pdcsap.FAP < FAP_max): logprint( '\nBest Signal is good enough to keep searching. Going to the next run.' ) key = 1 else: logprint('\nBest Signal does not look very promising. End') key = 0 else: if det_met == 2: logprint('\nBest Signal -->', '\t', 'flatten_lc' + str(a), '\t', 'kernel_size:', format(wl[a], '0.4f'), '\t', 'SNR:', format(SNR[a][0], '.3f')) else: logprint('\nBest Signal -->', '\t', 'flatten_lc' + str(a), '\t', 'window_size:', format(wl[a], '0.4f'), '\t', 'SNR:', format(SNR[a][0], '.3f')) if (SNR[a] > SNR_min): #and SDE[a] > SDE_min and FAP[a] < FAP_max): logprint( '\nBest Signal is good enough to keep searching. Going to the next run.' ) key = 1 else: logprint('\nBest Signal does not look very promising. End') key = 0 print("### SNR :", str(format(results.snr, '.3f')), " SDE :", str(format(results.SDE, '.3f')), " FAP :", str(results.FAP)) return results_pdcsap, SNR, key, periods, durations, tos
def main(): print("Starting tests for wotan...") numpy.testing.assert_almost_equal( t14(R_s=1, M_s=1, P=365), 0.6490025258902046) numpy.testing.assert_almost_equal( t14(R_s=1, M_s=1, P=365, small_planet=True), 0.5403690143737738) print("Transit duration correct.") numpy.random.seed(seed=0) # reproducibility print("Slide clipper...") points = 1000 time = numpy.linspace(0, 30, points) flux = 1 + numpy.sin(time) / points noise = numpy.random.normal(0, 0.0001, points) flux += noise for i in range(points): if i % 75 == 0: flux[i:i+5] -= 0.0004 # Add some transits flux[i+50:i+52] += 0.0002 # and flares clipped = slide_clip( time, flux, window_length=0.5, low=3, high=2, method='mad', center='median' ) numpy.testing.assert_almost_equal(numpy.nansum(clipped), 948.9926368754939) """ import matplotlib.pyplot as plt plt.scatter(time, flux, s=3, color='black') plt.scatter(time, clipped, s=3, color='orange') plt.show() """ # TESS test print('Loading TESS data from archive.stsci.edu...') path = 'https://archive.stsci.edu/hlsps/tess-data-alerts/' #path = 'P:/P/Dok/tess_alarm/' filename = "hlsp_tess-data-alerts_tess_phot_00062483237-s01_tess_v1_lc.fits" time, flux = load_file(path + filename) window_length = 0.5 print("Detrending 1 (biweight)...") flatten_lc, trend_lc = flatten( time, flux, window_length, edge_cutoff=1, break_tolerance=0.1, return_trend=True, cval=5.0) numpy.testing.assert_equal(len(trend_lc), 20076) numpy.testing.assert_almost_equal(numpy.nanmax(trend_lc), 28755.03811866676, decimal=2) numpy.testing.assert_almost_equal(numpy.nanmin(trend_lc), 28615.110229935075, decimal=2) numpy.testing.assert_almost_equal(trend_lc[500], 28671.650565730513, decimal=2) numpy.testing.assert_equal(len(flatten_lc), 20076) numpy.testing.assert_almost_equal(numpy.nanmax(flatten_lc), 1.0034653549250616, decimal=2) numpy.testing.assert_almost_equal(numpy.nanmin(flatten_lc), 0.996726610702177, decimal=2) numpy.testing.assert_almost_equal(flatten_lc[500], 1.000577429565131, decimal=2) print("Detrending 2 (andrewsinewave)...") flatten_lc, trend_lc = flatten(time, flux, window_length, method='andrewsinewave', return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18119.15471987987, decimal=2) print("Detrending 3 (welsch)...") flatten_lc, trend_lc = flatten(time, flux, window_length, method='welsch', return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18119.16764691235, decimal=2) print("Detrending 4 (hodges)...") flatten_lc, trend_lc = flatten(time[:1000], flux[:1000], window_length, method='hodges', return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 994.0110525909206, decimal=2) print("Detrending 5 (median)...") flatten_lc, trend_lc = flatten(time, flux, window_length, method='median', return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18119.122065014355, decimal=2) print("Detrending 6 (mean)...") flatten_lc, trend_lc = flatten(time, flux, window_length, method='mean', return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18119.032473037714, decimal=2) print("Detrending 7 (trim_mean)...") flatten_lc, trend_lc = flatten(time, flux, window_length, method='trim_mean', return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18119.095164910334, decimal=2) print("Detrending 8 (supersmoother)...") flatten_lc, trend_lc = flatten(time, flux, window_length, method='supersmoother', return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18123.00632204841, decimal=2) print("Detrending 9 (hspline)...") flatten_lc, trend_lc = flatten(time, flux, window_length, method='hspline', return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18123.082601463717, decimal=1) print("Detrending 10 (cofiam)...") flatten_lc, trend_lc = flatten(time[:2000], flux[:2000], window_length, method='cofiam', return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 1948.9999999987976, decimal=1) print("Detrending 11 (savgol)...") flatten_lc, trend_lc = flatten(time, flux, window_length=301, method='savgol', return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18123.003465539354, decimal=1) print("Detrending 12 (medfilt)...") flatten_lc, trend_lc = flatten(time, flux, window_length=301, method='medfilt', return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18123.22609806557, decimal=1) print("Detrending 12 (gp squared_exp)...") flatten_lc, trend_lc1 = flatten( time[:2000], flux[:2000], method='gp', kernel='squared_exp', kernel_size=10, return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 1948.9672036416687, decimal=2) print("Detrending 13 (gp squared_exp robust)...") flatten_lc, trend_lc1 = flatten( time[:2000], flux[:2000], method='gp', kernel='squared_exp', kernel_size=10, robust=True, return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 1948.8820772313468, decimal=2) print("Detrending 14 (gp matern)...") flatten_lc, trend_lc2 = flatten( time[:2000], flux[:2000], method='gp', kernel='matern', kernel_size=10, return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 1948.9672464898367, decimal=2) print("Detrending 15 (gp periodic)...") flatten_lc, trend_lc2 = flatten( time[:2000], flux[:2000], method='gp', kernel='periodic', kernel_size=1, kernel_period=10, return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 1948.9999708985608, decimal=2) time_synth = numpy.linspace(0, 30, 200) flux_synth = numpy.sin(time_synth) + numpy.random.normal(0, 0.1, 200) flux_synth = 1 + flux_synth / 100 time_synth *= 1.5 print("Detrending 16 (gp periodic_auto)...") flatten_lc, trend_lc2 = flatten( time_synth, flux_synth, method='gp', kernel='periodic_auto', kernel_size=1, return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 200, decimal=1) print("Detrending 17 (rspline)...") flatten_lc, trend_lc2 = flatten( time, flux, method='rspline', window_length=1, return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18121.812790732245, decimal=2) print("Detrending 18 (huber)...") flatten_lc, trend_lc = flatten( time[:1000], flux[:1000], method='huber', window_length=0.5, edge_cutoff=0, break_tolerance=0.4, return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 994.01102, decimal=2) print("Detrending 19 (winsorize)...") flatten_lc, trend_lc2 = flatten( time, flux, method='winsorize', window_length=0.5, edge_cutoff=0, break_tolerance=0.4, proportiontocut=0.1, return_trend=True) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18119.064587196448, decimal=2) print("Detrending 20 (pspline)...") flatten_lc, trend_lc = flatten( time, flux, method='pspline', return_trend=True ) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18121.832133916843, decimal=2) print("Detrending 21 (hampelfilt)...") flatten_lc, trend_lc5 = flatten( time, flux, method='hampelfilt', window_length=0.5, cval=3, return_trend=True ) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18119.158072498867, decimal=2) print("Detrending 22 (lowess)...") flatten_lc, trend_lc1 = flatten( time, flux, method='lowess', window_length=1, return_trend=True ) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18123.039744125545, decimal=2) print("Detrending 23 (huber_psi)...") flatten_lc, trend_lc1 = flatten( time, flux, method='huber_psi', window_length=0.5, return_trend=True ) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18119.122065014355, decimal=2) print("Detrending 24 (tau)...") flatten_lc, trend_lc2 = flatten( time, flux, method='tau', window_length=0.5, return_trend=True ) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18119.02772621119, decimal=2) print("Detrending 25 (cosine)...") flatten_lc, trend_lc2 = flatten( time, flux, method='cosine', window_length=0.5, return_trend=True ) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18122.999999974905, decimal=2) print("Detrending 25 (cosine robust)...") flatten_lc, trend_lc2 = flatten( time, flux, method='cosine', robust=True, window_length=0.5, return_trend=True ) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 18122.227938535038, decimal=2) import numpy as np points = 1000 time = numpy.linspace(0, 30, points) flux = 1 + numpy.sin(time) / points noise = numpy.random.normal(0, 0.0001, points) flux += noise for i in range(points): if i % 75 == 0: flux[i:i+5] -= 0.0004 # Add some transits flux[i+50:i+52] += 0.0002 # and flares print("Detrending 26 (hampel 17A)...") flatten_lc, trend_lc1 = flatten( time, flux, method='hampel', cval=(1.7, 3.4, 8.5), window_length=0.5, return_trend=True ) print("Detrending 27 (hampel 25A)...") flatten_lc, trend_lc2 = flatten( time, flux, method='hampel', cval=(2.5, 4.5, 9.5), window_length=0.5, return_trend=True ) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 997.9994362858843, decimal=2) print("Detrending 28 (ramsay)...") flatten_lc, trend_lc3 = flatten( time, flux, method='ramsay', cval=0.3, window_length=0.5, return_trend=True ) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 997.9974021484584, decimal=2) print("Detrending 29 (ridge)...") flatten_lc, trend_lc1 = flatten( time, flux, window_length=0.5, method='ridge', return_trend=True, cval=1) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 999.9999958887022, decimal=1) print("Detrending 30 (lasso)...") flatten_lc, trend_lc2 = flatten( time, flux, window_length=0.5, method='lasso', return_trend=True, cval=1) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 999.9999894829843, decimal=1) print("Detrending 31 (elasticnet)...") flatten_lc, trend_lc3 = flatten( time, flux, window_length=0.5, method='elasticnet', return_trend=True, cval=1) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc), 999.9999945063338, decimal=1) # Test of transit mask print('Testing transit_mask') filename = 'hlsp_tess-data-alerts_tess_phot_00207081058-s01_tess_v1_lc.fits' time, flux = load_file(path + filename) from wotan import transit_mask mask = transit_mask( time=time, period=14.77338, duration=0.21060, T0=1336.141095 ) numpy.testing.assert_almost_equal(numpy.sum(mask), 302, decimal=1) print('Detrending 32 (transit_mask cosine)') flatten_lc1, trend_lc1 = flatten( time, flux, method='cosine', window_length=0.4, return_trend=True, robust=True, mask=mask ) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc1), 18119.281265446625, decimal=1) print('Detrending 33 (transit_mask lowess)') flatten_lc2, trend_lc2 = flatten( time, flux, method='lowess', window_length=0.8, return_trend=True, robust=True, mask=mask ) #print(numpy.nansum(flatten_lc2)) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc2), 18119.30865711536, decimal=1) print('Detrending 34 (transit_mask GP)') mask = transit_mask( time=time[:2000], period=100, duration=0.3, T0=1327.4 ) flatten_lc2, trend_lc2 = flatten( time[:2000], flux[:2000], method='gp', kernel='matern', kernel_size=0.8, return_trend=True, robust=True, mask=mask ) #print(numpy.nansum(flatten_lc2)) numpy.testing.assert_almost_equal(numpy.nansum(flatten_lc2), 1948.9000170463796, decimal=1) """ import matplotlib.pyplot as plt plt.scatter(time[:2000], flux[:2000], s=1, color='black') plt.plot(time[:2000], trend_lc2, color='red', linewidth=2, linestyle='dashed') plt.show() plt.close() """ print('All tests completed.')