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.")
def explore_object(self, object_id, input_lc_file=None, time_units=1, auto_detrend_periodic_signals=False, auto_detrend_ratio=1 / 4, detrend_method="cosine", smooth=False, sectors=None, cadence=300): lc_builder = LcBuilder() object_info = lc_builder.build_object_info(object_id, None, sectors, input_lc_file, cadence, None, [{ "P": 0.59925, "D": 200, "T0": 1354.57 }], None, None, None) lc_build = lc_builder.build(object_info, ".") lc = lc_build.lc lc = lc.remove_outliers(sigma_lower=float('inf'), sigma_upper=3) # remove outliers over 3sigma flux = lc.flux.value flux_err = lc.flux_err.value if time_units == 0: time = lc.astropy_time.jd else: time = lc.time.value if smooth: flux = savgol_filter(flux, 11, 2) flux = uniform_filter1d(flux, 11) lc_df = pd.DataFrame(columns=['#time', 'flux', 'flux_err']) lc_df['#time'] = lc.time.value lc_df['flux'] = lc.flux.value lc_df['flux_err'] = lc.flux_err.value lc_df = pd.DataFrame(columns=['flux']) lc_df['flux'] = acf(lc.flux.value, nlags=7200) test_df = pd.DataFrame(lc_df) ax = test_df.plot() ax.set_ylim(lc_df["flux"].min(), lc_df["flux"].max()) #ax.set_xticks(list(range(0, len(ax.get_xticks()))), ax.get_xticks() / 30 / 24) plt.show() #periodogram = lc.to_periodogram(oversample_factor=10) #fig = px.line(x=periodogram.period.astype('<f4'), y=periodogram.power.astype('<f4'), log_x=True) #fig.show() fig = px.scatter(x=lc.time.value, y=lc.flux.value) fig.show() if auto_detrend_periodic_signals: detrend_period = self.__calculate_max_significant_period( lc, periodogram) if detrend_period is not None: flatten_lc, trend_lc = self.__detrend_by_period( lc, detrend_period * auto_detrend_ratio, detrend_method) fig = go.Figure() fig.add_trace( go.Scatter(x=lc.time.value, y=flux, mode='markers', name='Flux')) fig.add_trace( go.Scatter(x=lc.time.value, y=trend_lc, mode='lines+markers', name='Main Trend')) fig.show() if auto_detrend_periodic_signals: fig = px.line(x=lc.time.value, y=flatten_lc) fig.show() lc.flux = flatten_lc bin_means, bin_edges, binnumber = stats.binned_statistic( lc.time.value, lc.flux.value, statistic='mean', bins=len(time) / 5) bin_stds, _, _ = stats.binned_statistic(time, flux, statistic='std', bins=len(time) / 3) bin_width = (bin_edges[1] - bin_edges[0]) bin_centers = bin_edges[1:] - bin_width / 2 time_binned = bin_centers flux_binned = bin_means lc_binned = lk.LightCurve(time=time_binned, flux=flux_binned) fig = px.scatter(x=lc.time.value, y=lc.flux.value) fig.show() while True: try: user_input = input( "Select the period to fold the light curve: ") if user_input.startswith("q"): break period = float(user_input) except ValueError: print("Wrong number.") continue try: user_input = input("Select the t0 to fold the light curve: ") if user_input.startswith("q"): break t0 = float(user_input) except ValueError: print("Wrong number.") continue try: user_input = input("Select the duration (minutes): ") if user_input.startswith("q"): break duration = float(user_input) except ValueError: print("Wrong number.") continue # flux = lc.flux # args_flux_outer = np.argwhere(flux > 1 + depth * 0.001) # flux[args_flux_outer] = np.nan # time[args_flux_outer] = np.nan # args_flux_outer = np.argwhere(flux < 1 + depth * 0.001) # flux[args_flux_outer] = np.nan # time[args_flux_outer] = np.nan j = 0 fig_transit, axs = plt.subplots(4, 4, figsize=(16, 16)) #fig_transit.suptitle("Transit centered plots for TIC " + str(object_id)) for i in np.arange(t0, t0 + period * 16, period)[0:16]: args_plot = np.argwhere((i - 0.1 < time) & (time < i + 0.1)).flatten() plot_time = time[args_plot] plot_flux = flux[args_plot] axs[j // 4][j % 4].scatter(plot_time, plot_flux, color='gray', alpha=1, rasterized=True, label="Flux Transit " + str(j)) j = j + 1 folded_lc = lk.LightCurve(time=lc.time, flux=lc.flux, flux_err=lc.flux_err) folded_lc.remove_nans() # folded_lc.flux_err = np.nan # folded_lc = folded_lc.bin(bins=500) folded_lc.fold(period=period).scatter() plt.title("Phase-folded period: " + format(period, ".2f") + " days") plt.show() fig_transit.show() fig = px.line(x=folded_lc.time.value, y=folded_lc.flux.value) fig.show() mask = wotan.transit_mask(time, period, duration * 2 / 24 / 60, t0) fig = px.scatter(x=time[~mask], y=flux[~mask]) fig.show()
def flatten(lc, t0=None, period=None, duration=1. / 24., window_length=3. / 24., **kwargs): """ Flatten a light curve using the `wotan` package. Parameters ---------- lc : `~lightkurve.LightCurve` A light curve object with the data. t0 : float or iterable, optional Mid-transit time of transit signal(s) to mask. period : float or iterable, optional Period of transit signal(s) to mask. duration : float, optional Duration of transit signal to mask. Defaults to 1 hr. window_length : float, optional Length of the filter window for `wotan.flatten()`. Defaults to 3 hr. kwargs : dict, optional Any extra keyword arguments to pass to `wotan.flatten()`. Returns ------- lc : `~lightkurve.LightCurve` A light curve object with the flattened light curve. trend : ndarray The removed flux trend. Only returned if ``return_trend`` is `True`. """ # Mask transits if any ephemerides are given mask = np.zeros_like(lc.time, dtype=bool) if t0 is not None: t0s = np.array(t0).flatten() periods = np.array(period).flatten() for t0, period in zip(t0s, periods): mask += wotan.transit_mask(time=lc.time, T0=t0, period=period, duration=duration) # Flatten the light curve flux_flat = wotan.flatten(lc.time, lc.flux, window_length=window_length, mask=mask, **kwargs) # Return trend if return_trend=True return_trend = False if isinstance(flux_flat, tuple): return_trend = True flux_flat, trend = flux_flat lcflat = lc.copy() lcflat.flux = flux_flat if return_trend: return (lcflat, trend) else: return lcflat
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.')