!pip install numpy matplotlib astropy wotan transitleastsquares sklearn statsmodels sklearn supersmoother pygam
from ipywidgets import interactive
import ipywidgets as widgets
import numpy as np
import matplotlib.pyplot as plt
from astropy.io import fits
from wotan import flatten
from transitleastsquares import resample
# Load the data and bin down (too make this demo fast)
def load_file(filename):
"""Loads a TESS *spoc* FITS file and returns TIME, PDCSAP_FLUX"""
hdu = fits.open(filename)
time = hdu[1].data['TIME']
flux = hdu[1].data['PDCSAP_FLUX']
flux[flux == 0] = np.nan
return time, flux
print('Loading TESS data from archive.stsci.edu...')
path = 'https://archive.stsci.edu/hlsps/tess-data-alerts/'
filename = 'hlsp_tess-data-alerts_tess_phot_00077031414-s02_tess_v1_lc.fits'
time, flux = load_file(path + filename)
time, flux = resample(time, flux, factor=20)
# Detrending function
def func(method, window_length, break_tolerance, edge_cutoff='0.1', cval=5):
f, ax = plt.subplots(2, sharex=True, figsize=(12, 12))
if method == 'trim_mean' or method == 'winsorize':
cval /= 10 # must be a fraction >0, <0.5
if cval >=0.5:
cval = 0.49
flatten_lc, trend_lc = flatten(
time,
flux,
method=method,
window_length=window_length,
edge_cutoff=edge_cutoff,
break_tolerance=break_tolerance,
return_trend=True,
cval=cval
)
ax[0].plot(time, trend_lc, color='black', linewidth=3)
ax[0].scatter(time, flux, edgecolors='k', c='yellow', s=30)
ax[0].set_xlim(min(time), max(time))
ax[0].set_ylabel('Raw flux')
ax[1].scatter(time, flatten_lc, edgecolors='k', c='black', s=30)
ax[1].set_ylim(0.995, 1.005)
ax[1].set_ylabel('Detrended flux')
plt.xlabel('Time (days)')
f.subplots_adjust(hspace=0)
plt.show();
return time
y1=interactive(
func,
method=["biweight", "hodges", "welsch", "median", "andrewsinewave", "mean", "trim_mean", "winsorize"],
window_length=(0.1, 2, 0.1),
break_tolerance=(0, 1, 0.1),
edge_cutoff=(0, 1, 0.1),
cval=(1, 9, 1)
)
y2=interactive(
func,
method=["hspline", "pspline", "rspline"],
window_length=(0.1, 2, 0.1),
break_tolerance=(0, 1, 0.1),
edge_cutoff=(0, 1, 0.1),
cval=(1, 9, 1)
)
widgets.VBox([widgets.HBox([y1, y2])])