Leonardo0711/exoplanet-api
0
1#!/usr/bin/env python32# -*- coding: utf-8 -*-3"""4predict_part.py — Predicción por PARTES (sin stitch) usando el modelo "parts"5 6• Usa por defecto:7 models/baseline_parts.pkl8 models/baseline_parts_thresholds.json9 10• Flujo: (opcional) curate → features por PARTE → predicción → (opcional) agregado por target.11• Metadata por PARTE (segment/author) opcional, igual que en el training:12 columnas esperadas: mission, label, target_id, segment, author,13 stellar_radius_rsun, stellar_mass_msun, stellar_teff_k, stellar_logg_cgs, period_days14 15Novedades:16 - Genera artefactos/figuras por PARTE en out/artifacts_parts/<key>/:17 * lc_raw.png (curva normalizada)18 * lc_dtr.png (curva aplanada)19 * periodogram.png (BLS agregado)20 * phase.png (fase con binning)21 * odd_even.png (puntos impares/pares en ventana)22 * marks.png (ventanas de tránsito sobre la serie)23 * manifest.json (metadatos de la parte)24 - Si metadata trae period_days ∈ [0.5, 30], BLS se restringe a [0.9P, 1.1P].25 - Preprocesamiento igual al training (filtro de calidad, sigma-clip, flatten, pmin ajustado por span).26 - Orden de features desde baseline_parts_thresholds.json ("feature_names").27 28Ejemplos:29 30# A) ***Una sola parte*** (ej. KIC 4680772, part003)31python predict_part.py --raw-root "prueba/raw" --out "prueba/out_4680772_part003" --only-target 'KIC[_ ]4680772.*part003' --metadata "prueba/metadata.part.csv"32 33# B) ***Todas las partes*** de un target (ej. KIC 4680772)34python predict_part.py --raw-root "prueba/raw" --out "prueba/out_4680772_allparts" --only-target "4680772" --metadata "prueba/metadata.part.csv"35 36# C) ***Varias partes específicas*** (ej. 003, 004 y 007)37python predict_part.py --raw-root "prueba/raw" --out "prueba/out_4680772_part_003_004_007" --only-target 'KIC[_ ]4680772.*part(003|004|007)' --metadata "prueba/metadata.part.csv"38 39# D) (Opcional) Agregar por target con MAX cuando corriste varias partes40python predict_part.py --raw-root "prueba/raw" --out "prueba/out_4680772_allparts" --only-target "4680772" --aggregate-by-target max --metadata "prueba/metadata.part.csv"41"""42 43from __future__ import annotations44import argparse, json, math, re, sys, warnings45from dataclasses import dataclass, asdict46from concurrent.futures import ThreadPoolExecutor, as_completed47from pathlib import Path48from typing import Optional, Dict, Any, List, Tuple49 50import numpy as np51import pandas as pd52from tqdm import tqdm53from astropy.timeseries import BoxLeastSquares54import joblib55 56# Matplotlib (modo headless para servidores/CI)57import matplotlib58matplotlib.use("Agg")59import matplotlib.pyplot as plt60 61# ========================= CONFIG / DEFAULTS =========================62MODELS_DIR = Path("models")63DEFAULT_PKL = MODELS_DIR / "baseline_parts.pkl"64DEFAULT_THR = MODELS_DIR / "baseline_parts_thresholds.json"65 66P_MIN, P_MAX = 0.5, 30.0 # días (rango BLS)67DUR_FRACS = np.array([0.008, 0.015, 0.03, 0.05], dtype=float)68ROLL_WINDOW = 40169BINS = 20070EPS = 1e-971 72QUALITY_CANDS = ["quality","quality_flags","dq","flag","flags"]73 74# ========================= NORMALIZADORES (IDs/segment/author) =========================75_ID_CANON_RE = re.compile(r'(?i)\b(KIC|TIC)\b[ _-]?(\d+)')76 77def canonical_tid(s: Optional[str], prefer_prefix: Optional[str] = None) -> str:78 if not s:79 return ""80 s2 = str(s).strip()81 m = _ID_CANON_RE.search(s2)82 if m:83 return f"{m.group(1).upper()} {int(m.group(2))}"84 if s2.isdigit() and prefer_prefix:85 return f"{prefer_prefix.upper()} {int(s2)}"86 s2 = s2.replace("_", " ").replace("-", " ")87 s2 = re.sub(r"\s+", " ", s2).strip()88 return s2.upper()89 90def canonical_segment(s: Optional[str]) -> str:91 if not s or str(s).strip() == "":92 return ""93 ss = str(s).strip().upper().replace("SEG", "").strip()94 m = re.match(r'^(Q|S|C)\s*0*?(\d+)$', ss)95 if m:96 tag = m.group(1)97 num = int(m.group(2))98 if tag == "S":99 return f"S{num:03d}"100 else:101 return f"{tag}{num:02d}"102 if ss.isdigit():103 return f"Q{int(ss):02d}"104 return ss105 106def canonical_author(s: Optional[str]) -> str:107 if not s:108 return ""109 return str(s).strip().upper()110 111def sanitize_key(s: str) -> str:112 s2 = re.sub(r'\s+', '_', str(s))113 s2 = re.sub(r'[^A-Za-z0-9_\-\.]+', '', s2)114 return s2[:120] if len(s2) > 120 else s2115 116# ========================= UTILS FS / MATCH =========================117def ensure_dir(p: Path): p.mkdir(parents=True, exist_ok=True)118 119_ID_RE = re.compile(r'(KIC|TIC)[ _-]?(\d+)', re.IGNORECASE)120def match_only_target(path: Path, only: Optional[str]) -> bool:121 """Permite substring/regex simple sobre la ruta/stem; también acepta ID numérico."""122 if not only:123 return True124 s = str(path)125 try:126 if re.search(only, s, re.IGNORECASE):127 return True128 except re.error:129 if only.lower() in s.lower():130 return True131 if str(only).isdigit() and re.search(rf"{int(only)}", s):132 return True133 m = _ID_RE.search(path.stem) or _ID_RE.search(s)134 if m and re.search(str(only), f"{m.group(1)} {m.group(2)}", re.IGNORECASE):135 return True136 return False137 138# ============================== CURATE ==============================139KEPLER_RE = re.compile(r"^(KIC_\d+)(?:_Q(\d{1,2}))?(?:_author-([A-Za-z0-9\-]+))?_part(\d{3})$")140TESS_RE = re.compile(r"^(TIC_\d+)(?:_S(\d{1,3}))?(?:_author-([A-Za-z0-9\-]+))?_part(\d{3})$")141 142def parse_from_path(csv_path: Path) -> Dict[str, Optional[str]]:143 try:144 label = csv_path.parent.name145 mission = csv_path.parent.parent.name146 except Exception:147 label = None; mission = None148 stem = csv_path.stem149 m = KEPLER_RE.match(stem)150 if m:151 tid = m.group(1).replace("_"," ")152 q = int(m.group(2)) if m.group(2) else None153 auth = m.group(3) or "Kepler"154 return {"mission": mission or "Kepler","label": label,"target_id": tid,155 "segment": (f"Q{q:02d}" if q is not None else ""), "author": auth, "part": int(m.group(4))}156 m = TESS_RE.match(stem)157 if m:158 tid = m.group(1).replace("_"," ")159 s = int(m.group(2)) if m.group(2) else None160 auth= m.group(3) or "SPOC"161 return {"mission": mission or "TESS","label": label,"target_id": tid,162 "segment": (f"S{s:03d}" if s is not None else ""), "author": auth, "part": int(m.group(4))}163 return {"mission": mission, "label": label, "target_id": stem.split("_",1)[0].replace("_"," "),164 "segment":"", "author":None, "part":None}165 166def lower_map(df: pd.DataFrame) -> Dict[str, str]: return {c.lower(): c for c in df.columns}167TIME_CANDS = ["time","bkjd","btjd","jd","bjd","time_bjd","time_tjd","mjd"]168FLUX_PREF = ["pdcsap_flux","kspsap_flux","sap_flux","flux","psf_flux","aperture_flux","flux_raw","aperture_sum","net_flux"]169FLUX_ERRS = ["pdcsap_flux_err","kspsap_flux_err","sap_flux_err","flux_err","psf_flux_err","aperture_flux_err"]170QUALITYS = QUALITY_CANDS171 172def pick_time_col(df: pd.DataFrame) -> Optional[str]:173 lm = lower_map(df)174 for k in TIME_CANDS:175 if k in lm: return lm[k]176 return None177 178def pick_flux_cols(df: pd.DataFrame) -> Tuple[Optional[str], Optional[str], Optional[str]]:179 lm = lower_map(df)180 fcol = None181 for k in FLUX_PREF:182 if k in lm: fcol = lm[k]; break183 if fcol is None:184 for c in df.columns:185 if c.lower() not in TIME_CANDS and pd.api.types.is_numeric_dtype(df[c]):186 fcol = c; break187 ferr = None188 for k in FLUX_ERRS:189 if k in lm: ferr = lm[k]; break190 qcol = None191 for k in QUALITYS:192 if k in lm: qcol = lm[k]; break193 return fcol, ferr, qcol194 195def load_meta_for(csv_path: Path) -> Optional[Dict]:196 meta = csv_path.with_suffix(".meta.json")197 if meta.exists():198 try:199 return json.loads(meta.read_text(encoding="utf-8"))200 except Exception:201 return None202 return None203 204def reconstruct_time_from_meta(n_rows: int, meta: Dict) -> Optional[np.ndarray]:205 if n_rows <= 0 or meta is None: return None206 def _get(k):207 v = meta.get(k)208 try: return float(v) if v is not None else None209 except Exception: return None210 tstart = _get("TSTART") or _get("LC_START")211 tstop = _get("TSTOP") or _get("LC_END")212 dt = _get("TIMEDEL")213 if n_rows == 1 and tstart is not None: return np.array([tstart], float)214 if tstart is not None and tstop is not None and tstop >= tstart:215 return np.linspace(tstart, tstop, n_rows, dtype=float)216 if tstart is not None and dt is not None:217 return tstart + dt * np.arange(n_rows, dtype=float)218 return None219 220@dataclass221class CurateResult:222 file: str; out_file: Optional[str]; status: str; reason: str223 mission: Optional[str]; label: Optional[str]; target_id: Optional[str]224 n_rows_in: int; n_rows_out: int; flux_col_used: Optional[str]; time_reconstructed: bool225 226def curate_one(csv_path: Path, raw_root: Path, curated_root: Path, min_rows: int) -> CurateResult:227 meta_info = parse_from_path(csv_path)228 try:229 df = pd.read_csv(csv_path)230 except UnicodeDecodeError:231 df = pd.read_csv(csv_path, encoding="latin-1")232 except Exception as e:233 return CurateResult(csv_path.as_posix(), None, "READ_ERROR", str(e),234 meta_info["mission"], meta_info["label"], meta_info["target_id"], 0, 0, None, False)235 n_in = int(len(df))236 time_col = pick_time_col(df)237 fcol, ferr, qcol = pick_flux_cols(df)238 t_reco = False239 if time_col is None:240 meta = load_meta_for(csv_path)241 tvec = reconstruct_time_from_meta(n_in, meta) if meta else None242 if tvec is None or len(tvec) != n_in:243 return CurateResult(csv_path.as_posix(), None, "NO_TIME", "sin 'time' y no se pudo reconstruir",244 meta_info["mission"], meta_info["label"], meta_info["target_id"],245 n_in, 0, fcol, False)246 df.insert(0, "time", tvec); time_col = "time"; t_reco = True247 if fcol is None:248 return CurateResult(csv_path.as_posix(), None, "NO_FLUX", "no se encontró columna de flujo útil",249 meta_info["mission"], meta_info["label"], meta_info["target_id"], n_in, 0, None, t_reco)250 251 out = pd.DataFrame()252 out["time"] = pd.to_numeric(df[time_col], errors="coerce")253 flux_raw = pd.to_numeric(df[fcol], errors="coerce")254 out["flux_raw"] = flux_raw255 med = float(np.nanmedian(flux_raw)) if np.isfinite(np.nanmedian(flux_raw)) else np.nan256 out["flux"] = flux_raw / med if (np.isfinite(med) and med != 0) else flux_raw257 out["flux_err"] = (pd.to_numeric(df[ferr], errors="coerce")/med) if (ferr and np.isfinite(med) and med!=0) else (pd.to_numeric(df[ferr], errors="coerce") if ferr else np.nan)258 out["quality"] = pd.to_numeric(df[qcol], errors="coerce") if qcol else 0259 out["mission"] = meta_info["mission"] or ""260 out["label"] = meta_info["label"] or ""261 out["target_id"]=meta_info["target_id"] or ""262 out["segment"] = meta_info["segment"] or ""263 out["author"] = meta_info["author"] or ""264 out["flux_source"] = fcol265 out = out.replace([np.inf,-np.inf], np.nan).dropna(subset=["time","flux"]).sort_values("time").reset_index(drop=True)266 n_out = int(len(out))267 if n_out < min_rows:268 return CurateResult(csv_path.as_posix(), None, "TOO_FEW_ROWS", f"{n_out} < min_rows={min_rows}",269 meta_info["mission"], meta_info["label"], meta_info["target_id"], n_in, n_out, fcol, t_reco)270 try: rel = csv_path.relative_to(raw_root)271 except Exception: rel = Path(out["mission"].iloc[0] or "UNK")/ (out["label"].iloc[0] or "UNK")/ csv_path.name272 out_path = (curated_root / rel); ensure_dir(out_path.parent); out.to_csv(out_path, index=False)273 return CurateResult(csv_path.as_posix(), out_path.as_posix(), "OK", "", meta_info["mission"], meta_info["label"],274 meta_info["target_id"], n_in, n_out, fcol, t_reco)275 276def run_curate(raw_root: Path, curated_root: Path, only_target: Optional[str],277 workers: int, min_rows: int, limit: Optional[int]) -> Path:278 ensure_dir(curated_root)279 files = [p for p in raw_root.glob("**/*.csv") if match_only_target(p, only_target)]280 files = sorted(files)281 if limit: files = files[:limit]282 if not files:283 print(f"[WARN] No hay CSV en {raw_root} (filtro only-target={only_target}).")284 return curated_root285 results: List[CurateResult] = []286 with ThreadPoolExecutor(max_workers=max(1,workers)) as ex:287 futs = [ex.submit(curate_one, p, raw_root, curated_root, min_rows) for p in files]288 for f in tqdm(as_completed(futs), total=len(futs), desc="Curando"):289 try: results.append(f.result())290 except Exception as e:291 results.append(CurateResult("(unknown)", None, "EXC", str(e), None,None,None,0,0,None,False))292 rep = pd.DataFrame([asdict(r) for r in results])293 ensure_dir(curated_root / "reports")294 rep.to_csv(curated_root / "reports" / "curation_report.csv", index=False)295 status_counts = rep["status"].value_counts(dropna=False).to_dict()296 print("[CURATE] status:", status_counts)297 return curated_root298 299# ====================== FEATURES (como training) ======================300def robust_median_trend(y: np.ndarray, window: int) -> np.ndarray:301 if window % 2 == 0: window += 1302 s = pd.Series(y)303 tr = s.rolling(window=window, center=True, min_periods=max(5, window//5)).median().to_numpy()304 med = np.nanmedian(y)305 return np.where(np.isfinite(tr), tr, med)306 307def flatten(time: np.ndarray, flux: np.ndarray, window: int = ROLL_WINDOW) -> Tuple[np.ndarray, np.ndarray]:308 trend = robust_median_trend(flux, window)309 trend = np.where(trend == 0, 1.0, trend)310 return flux / trend, trend311 312def fold_phase(time: np.ndarray, t0: float, period: float):313 return ((time - t0 + 0.5*period) % period) / period - 0.5314 315def bin_phase(phase: np.ndarray, flux: np.ndarray, n_bins: int):316 edges = np.linspace(-0.5, 0.5, n_bins + 1)317 idx = np.digitize(phase, edges) - 1318 x = np.array([np.nanmean(phase[idx==i]) for i in range(n_bins)])319 y = np.array([np.nanmedian(flux[idx==i]) for i in range(n_bins)])320 m = np.isfinite(x) & np.isfinite(y)321 return x[m], y[m]322 323def cadence_minutes(time: np.ndarray) -> Optional[float]:324 if time is None or len(time) < 3: return None325 dt = np.diff(np.sort(time)); dt = dt[np.isfinite(dt)]326 if len(dt) == 0: return None327 return float(np.nanmedian(dt) * 24 * 60)328 329def bls_best_for_periods(bls: BoxLeastSquares, periods: np.ndarray, fracs: np.ndarray) -> Dict[str, float]:330 best = {"power": -np.inf}331 for frac in fracs:332 durations = periods * frac333 m = durations < periods334 if not np.any(m): continue335 res = bls.power(periods[m], durations[m])336 pwr = np.asarray(res.power)337 k = int(np.nanargmax(pwr))338 if pwr[k] > best["power"]:339 best = {"period": float(np.asarray(res.period)[k]),340 "duration": float(np.asarray(res.duration)[k]),341 "t0": float(np.asarray(res.transit_time)[k]),342 "power": float(pwr[k])}343 return best344 345def run_bls_two_stage(t: np.ndarray, f: np.ndarray, pmin: float, pmax: float, fracs: np.ndarray) -> Dict[str, float]:346 # pmin ajustado por span/2.5 (como en training para partes)347 span = float(np.nanmax(t) - np.nanmin(t)) if t.size else 0.0348 pmin_eff = max(pmin, min(pmax, span/2.5) if span>0 else pmin)349 pmax_eff = max(pmin_eff, pmax)350 351 bls = BoxLeastSquares(t, f)352 grid = np.geomspace(pmin_eff, pmax_eff, 900)353 stacks = []354 for frac in fracs:355 durations = grid * frac356 m = durations < grid357 if not np.any(m): continue358 stacks.append(np.asarray(bls.power(grid[m], durations[m]).power))359 if not stacks: raise RuntimeError("BLS: sin combinaciones válidas (parte).")360 agg = np.sum(np.stack(stacks, axis=0), axis=0)361 top_idx = np.argsort(agg)[-2:][::-1]362 best = {"power": -np.inf}363 for idx in top_idx:364 p0 = grid[idx]; pf = np.linspace(p0*0.98, p0*1.02, 700)365 cand = bls_best_for_periods(bls, pf, fracs)366 if cand and cand["power"] > best["power"]: best = cand367 if best["power"] == -np.inf: raise RuntimeError("BLS no encontró señal (parte).")368 return best369 370def choose_best_alias(t, flat, best, n_bins=BINS) -> Dict[str, Any]:371 baseP = best["period"]; cand=[]372 for tag, P in [("P/2",baseP/2.0), ("P",baseP), ("2P",baseP*2.0)]:373 if P<=0: continue374 frac = max(EPS, best["duration"]/baseP); dur= frac*P375 phase = fold_phase(t, best["t0"], P)376 xb,yb = bin_phase(phase, flat, n_bins)377 w = max(frac*1.2, 2/n_bins)378 in0 = np.abs(xb-0.0)<=w; inS = np.abs(np.abs(xb)-0.5)<=w379 med0 = float(np.nanmedian(yb[in0])) if np.any(in0) else 1.0380 medS = float(np.nanmedian(yb[inS])) if np.any(inS) else 1.0381 depth = max(0.0, (1.0 - med0)*1e6); sec = max(0.0, (1.0 - medS)*1e6)382 score = depth - 5.0*sec383 cand.append({"tag":tag,"period":P,"duration":dur,"depth_ppm":depth,"secondary_ppm":sec,"score":score})384 return max(cand, key=lambda d: d["score"])385 386def estimate_depth_snr(phase: np.ndarray, flat: np.ndarray, period: float, duration: float) -> Tuple[Optional[float], Optional[float], Optional[float]]:387 if period is None or duration is None: return None, None, None388 frac = max(0.01, min(0.2, duration/period)); win = 1.5*frac389 in_tr = np.abs(phase) <= win/2; oot = np.abs(phase) >= win390 if np.sum(in_tr)<10 or np.sum(oot)<20: return None, None, None391 depth = np.nanmedian(flat[oot]) - np.nanmedian(flat[in_tr])392 rms = np.nanstd(flat[oot])393 snr = depth / (rms / np.sqrt(np.sum(in_tr))) if (rms and rms>0) else None394 return float(depth*1e6), (float(snr) if snr is not None else None), float(rms)395 396# ====================== METADATA (por PARTE) ======================397STAR_COLS = ["stellar_radius_rsun","stellar_mass_msun","stellar_teff_k","stellar_logg_cgs"]398 399def load_meta_star_props(meta_csv: Optional[str]) -> Optional[pd.DataFrame]:400 if not meta_csv:401 return None402 p = Path(meta_csv)403 if not p.exists():404 print(f"[WARN] No encuentro metadata: {p.as_posix()}")405 return None406 try:407 df = pd.read_csv(p)408 except Exception as e:409 print(f"[WARN] Error leyendo metadata: {e}")410 return None411 412 for c in ["mission", "label", "target_id", "segment", "author", "quarter", "sector", "campaign"]:413 if c in df.columns:414 df[c] = df[c].astype(str).str.strip()415 416 if "segment" not in df.columns:417 def _seg(row):418 for k, pref in [("quarter","Q"), ("sector","S"), ("campaign","C")]:419 v = row.get(k)420 try:421 if pd.notnull(v):422 val = int(float(v))423 return f"{pref}{val:03d}" if pref == "S" else f"{pref}{val:02d}"424 except Exception:425 pass426 return ""427 df["segment"] = df.apply(_seg, axis=1)428 429 base_keys = ["mission", "label", "target_id", "segment"]430 if "author" in df.columns:431 base_keys.append("author")432 433 keep = [c for c in (base_keys + ["period_days"] + STAR_COLS) if c in df.columns]434 if "target_id" not in keep:435 return None436 437 out = df[keep].copy()438 439 dedup_keys = [c for c in ["mission","label","target_id","segment","author"] if c in out.columns]440 out = out.sort_values(dedup_keys).drop_duplicates(dedup_keys, keep="first")441 442 out["_tid"] = out["target_id"].apply(lambda x: canonical_tid(x))443 out["_seg"] = out["segment"].apply(lambda x: canonical_segment(x) if pd.notnull(x) else "")444 if "author" in out.columns:445 out["_auth"] = out["author"].apply(lambda x: canonical_author(x) if pd.notnull(x) else "")446 else:447 out["_auth"] = ""448 return out449 450def lookup_meta_row(enrich: pd.DataFrame, tid_real: str, seg_real: str, auth_real: str) -> Optional[pd.Series]:451 if enrich is None or enrich.empty:452 return None453 groups = []454 if auth_real is not None:455 g1 = enrich[(enrich["_tid"] == tid_real) & (enrich["_seg"] == seg_real) & (enrich["_auth"] == auth_real)]456 groups.append(g1)457 g2 = enrich[(enrich["_tid"] == tid_real) & (enrich["_seg"] == seg_real)]458 groups.append(g2)459 g3 = enrich[(enrich["_tid"] == tid_real)]460 groups.append(g3)461 for g in groups:462 if len(g):463 if "period_days" in g.columns:464 notnull = g[g["period_days"].notnull()]465 if len(notnull):466 return notnull.iloc[0]467 return g.iloc[0]468 return None469 470# ======================= LECTURA CSV por parte =======================471META_BASIC = ["mission","label","target_id","segment","author"]472 473def read_part_csv(p: Path) -> Optional[pd.DataFrame]:474 try:475 df = pd.read_csv(p)476 except Exception:477 return None478 if not {"time","flux"}.issubset(df.columns):479 return None480 qcol = next((c for c in QUALITY_CANDS if c in df.columns), None)481 if qcol and qcol != "quality":482 df["quality"] = pd.to_numeric(df[qcol], errors="coerce")483 elif "quality" not in df.columns:484 df["quality"] = 0485 need = set(META_BASIC) - set(df.columns)486 if need:487 try:488 label = p.parent.name489 mission = p.parent.parent.name490 except Exception:491 label = ""; mission = ""492 stem = p.stem493 m = KEPLER_RE.match(stem) or TESS_RE.match(stem)494 tid, seg, auth = "", "", ""495 if m:496 tid = m.group(1).replace("_"," ")497 if m.re is KEPLER_RE:498 seg = f"Q{int(m.group(2)):02d}" if m.group(2) else ""499 auth = m.group(3) or ("Kepler" if "Kepler" in mission else "")500 else:501 seg = f"S{int(m.group(2)):03d}" if m.group(2) else ""502 auth = m.group(3) or ("SPOC" if "TESS" in mission else "")503 if "mission" not in df: df["mission"] = mission504 if "label" not in df: df["label"] = label505 if "target_id" not in df: df["target_id"] = tid506 if "segment" not in df: df["segment"] = seg507 if "author" not in df: df["author"] = auth508 return df509 510# ===================== FEATURES + ARTEFACTOS (parte) =====================511def compute_features_and_artifacts_per_part(p: Path,512 enrich: Optional[pd.DataFrame],513 sigma_clip: float, qfilter: str, qmask: int,514 pmin: float, pmax: float,515 artifacts_dir: Optional[Path]) -> Optional[Dict[str, Any]]:516 df = read_part_csv(p)517 if df is None or df.empty:518 return None519 520 if "quality" in df.columns:521 q = df["quality"].fillna(0).astype(float)522 if qfilter == "good":523 df = df[q <= 0]524 elif qfilter == "bitmask" and qmask > 0:525 df = df[(q.fillna(0).astype(int) & int(qmask)) == 0]526 527 t = pd.to_numeric(df["time"], errors="coerce").to_numpy(float)528 f = pd.to_numeric(df["flux"], errors="coerce").to_numpy(float)529 m = np.isfinite(t) & np.isfinite(f)530 t, f = t[m], f[m]531 if t.size < 300:532 return None533 534 med = np.nanmedian(f)535 if np.isfinite(med) and med != 0:536 f = f / med537 538 if sigma_clip and sigma_clip > 0:539 med = np.nanmedian(f); mad = 1.4826*np.nanmedian(np.abs(f-med))540 if np.isfinite(mad) and mad > 0:541 lo, hi = med - sigma_clip*mad, med + sigma_clip*mad542 keep = np.isfinite(f) & (f>=lo) & (f<=hi)543 t, f = t[keep], f[keep]544 if t.size < 300: return None545 546 f_flat, trend = flatten(t, f, ROLL_WINDOW)547 548 mission = str(df["mission"].iloc[0]) if "mission" in df.columns else ""549 label = str(df["label"].iloc[0]) if "label" in df.columns else ""550 tid = str(df["target_id"].iloc[0]) if "target_id" in df.columns else ""551 seg = str(df["segment"].iloc[0]) if "segment" in df.columns else ""552 auth = str(df["author"].iloc[0]) if "author" in df.columns else ""553 554 # Restricción por metadata555 pmin_eff, pmax_eff = pmin, pmax556 rstar = mstar = teff = logg = None557 if enrich is not None:558 prefer_prefix = "KIC" if tid.upper().startswith("KIC") else ("TIC" if tid.upper().startswith("TIC") else None)559 ct = canonical_tid(tid, prefer_prefix=prefer_prefix)560 cs = canonical_segment(seg)561 ca = canonical_author(auth)562 row = lookup_meta_row(enrich, ct, cs, ca)563 if row is not None:564 if "period_days" in row and pd.notnull(row.get("period_days")):565 try:566 Pcat = float(row["period_days"])567 if P_MIN <= Pcat <= P_MAX:568 pmin_eff, pmax_eff = 0.9*Pcat, 1.1*Pcat569 except Exception:570 pass571 teff = float(row.get("stellar_teff_k")) if pd.notnull(row.get("stellar_teff_k")) else None572 logg = float(row.get("stellar_logg_cgs")) if pd.notnull(row.get("stellar_logg_cgs")) else None573 rstar = float(row.get("stellar_radius_rsun")) if pd.notnull(row.get("stellar_radius_rsun")) else None574 mstar = float(row.get("stellar_mass_msun")) if pd.notnull(row.get("stellar_mass_msun")) else None575 576 # BLS (con heurística span/2.5)577 try:578 best = run_bls_two_stage(t, f_flat, pmin_eff, pmax_eff, DUR_FRACS)579 except Exception:580 return None581 582 alias = choose_best_alias(t, f_flat, best, BINS)583 P, D, power = alias["period"], alias["duration"], best["power"]584 phase = fold_phase(t, best["t0"], P)585 depth_ppm, snr, rms_oot = estimate_depth_snr(phase, f_flat, P, D)586 587 frac = max(EPS, D/P if (P and D) else 0.02); w=frac/2588 in1 = np.abs(phase) <= w; in2 = np.abs(np.abs(phase)-0.5)<=w589 odd_even_ppm = secondary_ppm = None590 if np.sum(in1)>=5 and np.sum(in2)>=5:591 d1 = 1 - np.nanmedian(f_flat[in1]); d2 = 1 - np.nanmedian(f_flat[in2])592 odd_even_ppm = float((d1-d2)*1e6)593 secondary_ppm = float((1-np.nanmedian(f_flat[in2]))*1e6)594 595 span_d = float(np.nanmax(t) - np.nanmin(t))596 ntr = int(span_d/P) if (P and P>0) else None597 rp_rs = math.sqrt(max(depth_ppm,0.0)/1e6) if (depth_ppm is not None and np.isfinite(depth_ppm)) else None598 planet_re = (rp_rs * rstar * 109.1) if (rp_rs is not None and rstar is not None) else None599 600 # ---------- Artefactos ----------601 if artifacts_dir is not None:602 # clave: <tid>__<segment>__<author>__<stem>603 key = sanitize_key("__".join([x for x in [tid, seg, auth, p.stem] if x]))604 out_k = artifacts_dir / key605 ensure_dir(out_k)606 607 plt.style.use("dark_background")608 609 # 1) lc_raw.png610 fig = plt.figure(figsize=(9,4.5), dpi=120)611 ax = plt.gca()612 ax.plot(t, f, lw=0.6)613 ax.set_xlabel("Tiempo (días)")614 ax.set_ylabel("Flujo normalizado")615 ax.set_title(f"{tid} {seg} {auth} — Parte")616 ax.grid(alpha=.2)617 fig.tight_layout()618 fig.savefig(out_k / "lc_raw.png")619 plt.close(fig)620 621 # 2) lc_dtr.png622 fig = plt.figure(figsize=(9,4.5), dpi=120)623 ax = plt.gca()624 ax.plot(t, f_flat, lw=0.5)625 ax.set_xlabel("Tiempo (días)")626 ax.set_ylabel("Flujo aplanado")627 ax.set_title(f"{tid} {seg} {auth} — Detrended")628 ax.grid(alpha=.2)629 fig.tight_layout()630 fig.savefig(out_k / "lc_dtr.png")631 plt.close(fig)632 633 # 3) periodogram.png (agregado usado internamente)634 # volvemos a calcular el agregado para dibujar635 bls = BoxLeastSquares(t, f_flat)636 grid = np.geomspace(max(P_MIN, min(P_MAX, (span_d/2.5) if span_d>0 else P_MIN)), max(P_MAX, P_MAX), 900)637 stacks = []638 for frac_d in DUR_FRACS:639 durations = grid * frac_d640 msk = durations < grid641 if np.any(msk):642 stacks.append(np.asarray(bls.power(grid[msk], durations[msk]).power))643 if stacks:644 agg = np.sum(np.stack(stacks, axis=0), axis=0)645 fig = plt.figure(figsize=(9,4.5), dpi=120)646 ax = plt.gca()647 ax.plot(grid, agg, lw=0.9)648 ax.axvline(P, ls="--", lw=1.0)649 ax.set_xscale("log")650 ax.set_xlabel("Período (días)")651 ax.set_ylabel("Potencia (BLS agregada)")652 ax.set_title(f"{tid} {seg} {auth} — Periodograma")653 ax.grid(alpha=.2)654 fig.tight_layout()655 fig.savefig(out_k / "periodogram.png")656 plt.close(fig)657 658 # 4) phase.png659 xb, yb = bin_phase(phase, f_flat, BINS)660 fig = plt.figure(figsize=(9,4.5), dpi=120)661 ax = plt.gca()662 ax.scatter(phase, f_flat, s=2, alpha=0.25)663 ax.plot(xb, yb, lw=1.2)664 ax.set_xlabel("Fase")665 ax.set_ylabel("Flujo aplanado")666 ax.set_title(f"{tid} {seg} — Fase (P={P:.5f} d, dur={D:.4f} d)")667 ax.grid(alpha=.2)668 fig.tight_layout()669 fig.savefig(out_k / "phase.png")670 plt.close(fig)671 672 # 5) odd_even.png673 if np.sum(in1) > 0:674 k = np.floor((t - best["t0"])/P + 0.5).astype(int)675 odd = (k % 2) != 0676 fig = plt.figure(figsize=(9,4.5), dpi=120)677 ax = plt.gca()678 ax.scatter(phase[in1 & (~odd)], f_flat[in1 & (~odd)], s=5, alpha=0.7, label="pares")679 ax.scatter(phase[in1 & odd], f_flat[in1 & odd], s=5, alpha=0.7, label="impares")680 ax.set_xlabel("Fase")681 ax.set_ylabel("Flujo aplanado")682 ax.set_title(f"{tid} {seg} {auth} — Odd/Even")683 ax.legend()684 ax.grid(alpha=.2)685 fig.tight_layout()686 fig.savefig(out_k / "odd_even.png")687 plt.close(fig)688 689 # 6) marks.png (ventanas de tránsito en el dominio del tiempo)690 fig = plt.figure(figsize=(9,4.5), dpi=120)691 ax = plt.gca()692 ax.plot(t, f_flat, lw=0.5)693 if P and D:694 tmin, tmax = float(np.min(t)), float(np.max(t))695 n_lo = int(np.floor((tmin - best["t0"]) / P)) - 1696 n_hi = int(np.ceil((tmax - best["t0"]) / P)) + 1697 n_total = n_hi - n_lo + 1698 step = max(1, int(np.ceil(n_total / 120)))699 for n in range(n_lo, n_hi+1, step):700 center = best["t0"] + n*P701 ax.axvspan(center - D/2, center + D/2, color="tab:blue", alpha=0.15, lw=0)702 ax.set_xlabel("Tiempo (días)")703 ax.set_ylabel("Flujo aplanado")704 ax.set_title(f"{tid} {seg} — Ventanas de tránsito")705 ax.grid(alpha=.2)706 fig.tight_layout()707 fig.savefig(out_k / "marks.png")708 plt.close(fig)709 710 manifest = {711 "mission": mission, "label": label,712 "target_id": tid, "segment": seg, "author": auth,713 "csv": str(p.as_posix()),714 "n_rows": int(len(t)),715 "period_days": P, "duration_days": D, "t0": float(best["t0"]),716 }717 (out_k / "manifest.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")718 719 # ---------- Features ----------720 return {721 "out_file": p.as_posix(),722 "mission": mission, "label": label, "target_id": tid,723 "segment": seg, "author": auth,724 "n_rows": int(len(t)),725 "cadence_min": cadence_minutes(t),726 "span_days": span_d,727 "bls_power": power,728 "period_days": P,729 "duration_days": D,730 "depth_ppm": depth_ppm,731 "secondary_ppm": secondary_ppm,732 "snr_transit": snr,733 "odd_even_delta_ppm": odd_even_ppm,734 "rms_oot": rms_oot,735 "n_transits_est": ntr,736 "rp_over_rs": rp_rs,737 "planet_radius_re": planet_re,738 "teff_k": teff,739 "logg_cgs": logg,740 "stellar_radius_rsun": rstar,741 "stellar_mass_msun": mstar742 }743 744def extract_features_per_part(p: Path, enrich: Optional[pd.DataFrame],745 sigma_clip: float, qfilter: str, qmask: int,746 pmin: float, pmax: float) -> Optional[Dict[str, Any]]:747 # wrapper sin artefactos (por compatibilidad si lo necesitas)748 return compute_features_and_artifacts_per_part(749 p, enrich, sigma_clip, qfilter, qmask, pmin, pmax, artifacts_dir=None750 )751 752# =========================== MODELO / CLI ===========================753DEFAULT_FEATURE_ORDER = [754 "n_rows","cadence_min","span_days","bls_power","period_days","duration_days",755 "depth_ppm","secondary_ppm","snr_transit","odd_even_delta_ppm","rms_oot",756 "n_transits_est","rp_over_rs","planet_radius_re","teff_k","logg_cgs",757 "stellar_radius_rsun","stellar_mass_msun"758]759 760def load_model(pkl_path: Path, thr_json: Path):761 clf = joblib.load(pkl_path)762 thr = json.loads(thr_json.read_text(encoding="utf-8")) if thr_json.exists() else {"default_threshold_value": 0.5}763 feat_names = thr.get("feature_names")764 order = feat_names if isinstance(feat_names, list) and len(feat_names)>0 else DEFAULT_FEATURE_ORDER765 return clf, thr, order766 767def predict_one_row(model, row: Dict[str,Any], feature_order: List[str]) -> float:768 X = pd.DataFrame([{k: row.get(k, np.nan) for k in feature_order}])769 return float(model.predict_proba(X)[:,1][0])770 771def build_cli():772 p = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)773 # Entradas774 p.add_argument("--raw-root", type=str, default=None, help="Raíz de CSV crudos (se curan).")775 p.add_argument("--curated-root", type=str, default=None, help="Si lo das, saltas el curado.")776 p.add_argument("--out", type=str, default="PredOutParts", help="Raíz de salidas.")777 p.add_argument("--only-target", type=str, default=None, help="Filtro substring/regex sobre ruta/stem (p. ej. '11336883', 'part003', 'S031.*part005').")778 p.add_argument("--limit", type=int, default=None, help="Máximo de CSV crudos a curar o partes a procesar (ordenado).")779 p.add_argument("--workers", type=int, default=8)780 p.add_argument("--min-rows", type=int, default=100, help="Mínimo de filas tras curado.")781 # Filtros/feats (como training)782 p.add_argument("--sigma-clip", type=float, default=6.0)783 p.add_argument("--quality-filter", type=str, default="good", choices=["none","good","bitmask"])784 p.add_argument("--quality-bitmask", type=int, default=0)785 # Modelo786 p.add_argument("--model", type=str, default=str(DEFAULT_PKL))787 p.add_argument("--thresholds", type=str, default=str(DEFAULT_THR))788 p.add_argument("--aggregate-by-target", type=str, default="max", choices=["max","mean"],789 help="Agregación de probabilidades por target (post-predicción).")790 # Metadata por PARTE791 p.add_argument("--metadata", type=str, default=None, help="CSV por parte (mission,label,target_id,segment,author,period_days,props estelares)")792 # Artefactos793 p.add_argument("--no-artifacts", action="store_true", help="Si se pasa, NO se guardan PNGs/manifest por parte.")794 return p.parse_args()795 796# =============================== MAIN ===============================797def main():798 warnings.filterwarnings("ignore")799 800 args = build_cli()801 out_root = Path(args.out).resolve(); ensure_dir(out_root)802 803 # 1) CURATE (si hay raw-root) → curated_dir804 if args.curated_root:805 curated_dir = Path(args.curated_root).resolve()806 if not curated_dir.exists():807 print(f"[ERROR] --curated-root no existe: {curated_dir}"); sys.exit(1)808 print(f"[INFO] Usando curated existente: {curated_dir.as_posix()}")809 else:810 if not args.raw_root:811 print("[ERROR] Debes indicar --raw-root o --curated-root."); sys.exit(1)812 raw_root = Path(args.raw_root).resolve()813 curated_dir = out_root / "curated"814 curated_dir = run_curate(raw_root, curated_dir, args.only_target, args.workers, args.min_rows, args.limit)815 816 # 2) Cargar modelo + thresholds + orden de features817 model, thr, feature_order = load_model(Path(args.model), Path(args.thresholds))818 thr_val = float(thr.get("thresholds",{}).get("best_f1",{}).get("thr", thr.get("default_threshold_value", 0.5)))819 820 # 3) Metadata por parte821 enrich = load_meta_star_props(args.metadata)822 823 # 4) Listar partes a procesar824 files = sorted([p for p in curated_dir.glob("**/*.csv") if match_only_target(p, args.only_target)])825 if args.limit:826 files = files[:int(args.limit)]827 if not files:828 print(f("[WARN] No hay CSV curados en {curated_dir} (filtro only-target={args.only_target})."))829 sys.exit(0)830 831 # Carpeta de artefactos por parte832 artifacts_dir = None if args.no_artifacts else (out_root / "artifacts_parts")833 if artifacts_dir is not None:834 ensure_dir(artifacts_dir)835 836 # 5) Features + Predicción837 rows_feat, preds, man = [], [], []838 for p in tqdm(files, desc="Features (per-part)", unit="file"):839 f = compute_features_and_artifacts_per_part(840 p, enrich,841 sigma_clip=args.sigma_clip,842 qfilter=args.quality_filter,843 qmask=args.quality_bitmask,844 pmin=P_MIN, pmax=P_MAX,845 artifacts_dir=artifacts_dir846 )847 if f is None:848 continue849 850 proba = predict_one_row(model, f, feature_order)851 rows_feat.append(f)852 preds.append({853 "mission": f["mission"], "label": f["label"], "target_id": f["target_id"],854 "segment": f["segment"], "author": f["author"],855 "file": f["out_file"], "proba_confirmed": proba,856 "pred_label": ("CONFIRMED" if proba>=thr_val else "FALSE_POSITIVE"),857 "period_days": f.get("period_days"),858 "duration_days": f.get("duration_days"),859 "depth_ppm": f.get("depth_ppm"),860 "secondary_ppm": f.get("secondary_ppm"),861 "snr_transit": f.get("snr_transit"),862 "odd_even_delta_ppm": f.get("odd_even_delta_ppm"),863 "rms_oot": f.get("rms_oot"),864 "n_transits_est": f.get("n_transits_est"),865 "rp_over_rs": f.get("rp_over_rs"),866 "planet_radius_re": f.get("planet_radius_re"),867 "teff_k": f.get("teff_k"),868 "logg_cgs": f.get("logg_cgs"),869 "stellar_radius_rsun": f.get("stellar_radius_rsun"),870 "stellar_mass_msun": f.get("stellar_mass_msun")871 })872 man.append({"csv": p.as_posix(), "kind":"curated"})873 874 if rows_feat:875 pd.DataFrame(rows_feat).to_csv(out_root / "features_part.csv", index=False)876 pd.DataFrame(preds).to_csv(out_root / "predictions_part.csv", index=False)877 pd.DataFrame(man).to_csv(out_root / "manifest_part.csv", index=False)878 print(f"[OK] per-part: {len(preds)} predicciones → { (out_root/'predictions_part.csv').as_posix() }")879 880 # 6) Agregado por target (max/mean)881 agg = (pd.DataFrame(preds)882 .groupby(["mission","label","target_id"])["proba_confirmed"]883 .agg(args.aggregate_by_target).reset_index()884 .rename(columns={"proba_confirmed": f"proba_{args.aggregate_by_target}"}))885 agg["pred_label"] = np.where(agg[f"proba_{args.aggregate_by_target}"]>=thr_val,"CONFIRMED","FALSE_POSITIVE")886 agg.to_csv(out_root / "predictions_part_by_target.csv", index=False)887 print(f"[OK] agregado por target ({args.aggregate_by_target}) → { (out_root/'predictions_part_by_target.csv').as_posix() }")888 else:889 print("[WARN] per-part: no se obtuvieron features de ningún archivo.")890 891 print(f"\n[UMBRAL USADO] {thr_val:.3f}")892 print("[FIN]")893 894if __name__ == "__main__":895 main()896