File size: 5,663 Bytes
c223e92 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 | """
STEP 1 — Preprocessing Module
Load, clean, detrend, and normalize TESS light curves.
"""
from pathlib import Path
from typing import Tuple, Union
import numpy as np
from scipy.signal import savgol_filter
from lightkurve import search_lightcurve
_SAMPLE_DIR = Path(__file__).resolve().parent.parent / "data" / "sample_light_curves"
def _resolve_local_path(tic_id_or_path: Union[str, Path]) -> Union[Path, None]:
"""Resolve a TIC/TOI name to a local .npz file path, if one exists.
Checks: (1) exact given path, (2) ``data/sample_light_curves/`` with
the name as-is, and (3) ``data/sample_light_curves/`` with underscores
substituted for dashes (e.g. *TOI-270* → *TOI_270.npz*).
"""
p = Path(tic_id_or_path)
# (1) Exact path
if p.exists():
return p
# (2) Search in sample directory
candidates = [
_SAMPLE_DIR / p.name,
_SAMPLE_DIR / (p.name + ".npz"),
_SAMPLE_DIR / (str(p.name).replace("-", "_")),
_SAMPLE_DIR / (str(p.name).replace("-", "_") + ".npz"),
]
for c in candidates:
if c.exists():
return c
return None
def load_and_clean(tic_id_or_path: Union[str, Path]) -> Tuple[np.ndarray, np.ndarray]:
"""
Load a TESS light curve for a given TIC ID or local path,
clean it, detrend it, and normalize it.
Args:
tic_id_or_path: TIC ID string (e.g., "TIC 123456789", "TOI-270")
or Path to a local .fits or .npz file
Returns:
Tuple of (time, flux) arrays - cleaned, detrended, normalized
"""
local = _resolve_local_path(tic_id_or_path)
if local is not None:
if local.suffix == ".npz":
data = np.load(local)
time = data["time"]
flux = data["flux"]
elif local.suffix in (".fits", ".fit"):
from lightkurve import read
lc = read(local)
time = lc.time.value
flux = lc.flux.value
else:
raise ValueError(f"Unsupported file format: {local.suffix}")
else:
# Search by TIC ID or TOI name (online fallback)
tic_str = str(tic_id_or_path)
search_result = search_lightcurve(tic_str, mission="TESS")
if len(search_result) == 0:
search_result = search_lightcurve(tic_str, mission="TESS", author="QLP")
if len(search_result) == 0:
raise ValueError(f"No light curves found for {tic_str}")
# Download the first available sector
lc = search_result[0].download()
time = lc.time.value
flux = lc.flux.value
# Remove NaN values
valid_mask = np.isfinite(time) & np.isfinite(flux)
time = time[valid_mask]
flux = flux[valid_mask]
# Remove bad quality cadences if quality array exists
# (This is handled by lightkurve download if quality flag is available)
# Sigma-clip outliers (>3 sigma)
flux_median = np.median(flux)
flux_std = np.std(flux)
outlier_mask = np.abs(flux - flux_median) <= 3 * flux_std
time = time[outlier_mask]
flux = flux[outlier_mask]
# Detrend using Savitzky-Golay filter
# window_length must be odd and <= len(flux)
window_length = min(101, len(flux) - 1)
if window_length % 2 == 0:
window_length -= 1
if window_length >= 5: # Need at least 5 points for polyorder=3
trend = savgol_filter(flux, window_length=window_length, polyorder=3)
flux_detrended = flux / trend
else:
flux_detrended = flux / flux_median
# Normalize by median (baseline = 1.0)
flux_normalized = flux_detrended / np.median(flux_detrended)
return time, flux_normalized
def load_raw_and_clean(tic_id_or_path: Union[str, Path]) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
"""
Load a light curve and return both raw and cleaned versions.
Returns:
Tuple of (time, raw_flux, cleaned_flux)
"""
local = _resolve_local_path(tic_id_or_path)
if local is not None:
if local.suffix == ".npz":
data = np.load(local)
time = data["time"]
raw_flux = data["flux"]
elif local.suffix in (".fits", ".fit"):
from lightkurve import read
lc = read(local)
time = lc.time.value
raw_flux = lc.flux.value
else:
raise ValueError(f"Unsupported file format: {local.suffix}")
else:
tic_str = str(tic_id_or_path)
search_result = search_lightcurve(tic_str, mission="TESS")
if len(search_result) == 0:
search_result = search_lightcurve(tic_str, mission="TESS", author="QLP")
if len(search_result) == 0:
raise ValueError(f"No light curves found for {tic_str}")
lc = search_result[0].download()
time = lc.time.value
raw_flux = lc.flux.value
# Remove NaN
valid_mask = np.isfinite(time) & np.isfinite(raw_flux)
time = time[valid_mask]
raw_flux = raw_flux[valid_mask]
# Sigma-clip
flux_median = np.median(raw_flux)
flux_std = np.std(raw_flux)
outlier_mask = np.abs(raw_flux - flux_median) <= 3 * flux_std
time = time[outlier_mask]
raw_flux = raw_flux[outlier_mask]
# Detrend
window_length = min(101, len(raw_flux) - 1)
if window_length % 2 == 0:
window_length -= 1
if window_length >= 5:
trend = savgol_filter(raw_flux, window_length=window_length, polyorder=3)
flux_detrended = raw_flux / trend
else:
flux_detrended = raw_flux / flux_median
# Normalize
cleaned_flux = flux_detrended / np.median(flux_detrended)
return time, raw_flux, cleaned_flux |