"""Robustness checks: re-run the deficit analysis with alternative choices (final cut set), and the astrometric-quality diagnostics for both cut sets. Three kinds of variant recompute : the calibration is refitted with a different setting (neighbour number, neighbour scales); subset : a subset of a clean sample (calibration NOT refitted): stricter astrometry, lower extinction, distance, latitude, apparent magnitude; also "one cut at a time" subsets of the sample selected with the ORIGINAL astrometric cuts (read from the stored origcuts tables); threshold : the same deficits with different flagging rules (thresholds, single band, colour cut, which of the three criteria are required). For every variant, survey and dimming fraction f we report the clean sample size, the number of sensitive stars, the number flagged, the injection recovery, how many flagged stars are unresolved (a flagged star is resolved only if the vetting resolved it; stars newly flagged by a variant were never vetted and are counted as unresolved) and the 95% upper limit (population limit, with the selection efficiency eps per Teff class, the end-to-end recovery correction C and the vetting efficiency V of the matching cut set). The full pipeline under the original cuts (own selection model, own vetting) is not a variant here: it is run by run_all.py with DIM_STARS_CUTSET=orig and its limits are in data/origcuts/limits.csv. Diagnostics data/astrometry_diagnostic.csv three-way candidate rate under the ORIGINAL cuts (RUWE < 1.4, parallax/error >= 20) in bins of RUWE and parallax/error: the evidence that motivated the tightening data/astrometry_diagnostic_final.csv the same inside the FINAL cuts, plus bins of apparent magnitude Writes data/robustness.csv and the two diagnostics. """ import os import warnings import numpy as np import pandas as pd import dimlib as L from analysis import CLASSES, local_faint_sigma, teff_class from common import DATA_MAIN, DATA_ORIG, WORK, dpkl, log warnings.filterwarnings("ignore") SURVEYS = ("apogee", "lamost", "gspspec") ALL3 = ("spec", "ms", "cmd") def prepare(d, survey, k=L.KNN_K, scale=None, poe_min=None, ruwe_max=None, floors=None, precut=False): """Cuts + calibration + deficits for all stars passing the cuts, with sigma_i for D, D_G and D_K.""" if precut: m = np.ones(len(d), bool) else: m, _ = L.apply_cuts(d, survey, verbose=False, poe_min=poe_min, ruwe_max=ruwe_max) s, cal = L.compute_deficits(d[m], k=k, scale=scale, floors=floors) for band, col in (("max", "D"), ("G", "D_G"), ("K", "D_K")): s["sigma_" + band] = np.maximum(local_faint_sigma(s[col].values, cal.idx)[0], 0.015) s["cls"] = teff_class(s.teff.values) s["popok"] = cal.rk < L.KNN_RMAX return s def evaluate(s, band="max", nsig_thr=None, nsig_sens=None, colour=True, zmax=L.COLOUR_ZMAX, require=ALL3, extra=None, edge_mult=None): ok = s.popok.values.copy() if colour: ok &= ((np.abs(s.z_bprp) < zmax) & (np.abs(s.z_gk) < zmax)).values if extra is not None: ok &= np.asarray(extra(s), bool) c = s[ok].reset_index(drop=True) if edge_mult is not None: L.set_floor(c, edge_mult) c["D"] = c["D" if band == "max" else "D_" + band] c["sigma_i"] = c["sigma_" + band] return c, nsig_thr, nsig_sens, require # variants that refit the calibration (the main-sequence floors of the first four are those of the baseline: # they do not depend on k, and a common rescaling of the scales leaves the 2-D neighbour sets unchanged) RECOMPUTE = [ ("baseline", {}, False), ("k = 25 neighbours", dict(k=25), True), ("k = 100 neighbours", dict(k=100), True), ("neighbour scales x0.6", dict(scale={a: 0.6 * b for a, b in L.KNN_SCALE.items()}), True), ("neighbour scales x1.6", dict(scale={a: 1.6 * b for a, b in L.KNN_SCALE.items()}), True), ] # evaluated on the baseline computation ON_BASE = [ ("threshold 1 sigma above expected deficit", dict(nsig_thr=1.0)), ("threshold 3 sigma above expected deficit", dict(nsig_thr=3.0)), ("sensitivity requirement 4 sigma", dict(nsig_sens=4.0)), ("sensitivity requirement 7 sigma", dict(nsig_sens=7.0)), ("Ks band only", dict(band="K")), ("G band only", dict(band="G")), ("main-sequence floor multiplier 1.0", dict(edge_mult=1.0)), ("main-sequence floor multiplier 2.5", dict(edge_mult=2.5)), ("no colour-consistency cut", dict(colour=False)), ("colour cut at 3 sigma", dict(zmax=3.0)), ("spectroscopic criterion only (no main-sequence floors)", dict(require=("spec",))), ("spectroscopic + logg-free floor only", dict(require=("spec", "ms"))), ("spectroscopic + colour-based floor only", dict(require=("spec", "cmd"))), ("subset: RUWE < 1.1", dict(extra=lambda d: d.ruwe < 1.1)), ("subset: parallax/error >= 100", dict(extra=lambda d: d.poe >= 100)), ("subset: E(B-V)_SFD < 0.05", dict(extra=lambda d: d.sfd < 0.05)), ("subset: distance < 300 pc", dict(extra=lambda d: d.dist < 300)), ("subset: distance >= 300 pc", dict(extra=lambda d: d.dist >= 300)), ("subset: |b| >= 40 deg", dict(extra=lambda d: np.abs(d.glat) >= 40)), ("subset: brighter half (G < survey median)", dict(extra=lambda d: d.G < np.median(d.G))), ("subset: fainter half (G >= survey median)", dict(extra=lambda d: d.G >= np.median(d.G))), ("subset: main sequence only (logg >= 4.1)", dict(extra=lambda d: d.logg >= L.LOGG_MS)), ] # subsets of the sample selected with the ORIGINAL astrometric cuts (stored origcuts tables; calibration from # that sample). The last column says which cut set's selection efficiency applies (it follows the parallax cut). ON_ORIG = [ ("original cuts, all (RUWE < 1.4, parallax/error >= 20)", {}, "orig"), ("original sample: RUWE < 1.4 with parallax/error >= 50", dict(extra=lambda d: d.poe >= 50), "final"), ("original sample: parallax/error >= 20 with RUWE < 1.2", dict(extra=lambda d: d.ruwe < 1.2), "orig"), ("original sample: RUWE < 1.2 and parallax/error >= 50", dict(extra=lambda d: (d.ruwe < 1.2) & (d.poe >= 50)), "final"), ] # columns used by the cuts, the calibration and the subset definitions NEEDED = ["source_id", "teff", "logg", "feh", "snr", "flag_ok", "ncomp", "offsets", "e_teff", "e_logg", "flags_gspspec", "solved", "G", "BP", "RP", "dup", "plx", "poe", "sfd", "ruwe", "ipd_fmp", "ipd_gofha", "bprp_excess", "K", "e_K", "d2m", "n2m", "qfl", "MG", "MK", "bprp0", "gk0", "dist", "glat"] def flags_full(s, f, inject=False): """(sensitive, flagged) at level f for every row of a prepare() frame (False outside the clean sample).""" c = s[["D", "m_ms", "m_cmd"]].copy() c["sigma_i"] = s.sigma_max.values sens, fl = L.flag_mask(c, f, inject=inject) ok = s.popok.values & (np.abs(s.z_bprp.values) < L.COLOUR_ZMAX) & (np.abs(s.z_gk.values) < L.COLOUR_ZMAX) return sens & ok, fl & ok RUWE_BINS = (0.0, 1.0, 1.1, 1.2, 1.3, 1.4) POE_BINS = (20, 30, 50, 70, 100, 200, np.inf) def from_table(sv, cutset): """Stored per-star table of a cut set, in the form evaluate() expects.""" s = pd.read_pickle(dpkl(sv, cutset)) s["popok"] = s["pop"].values s["sigma_max"] = s.sigma_i.values return s def diagnostic(c, sv, cutset): """Three-way candidate rate (before vetting) in bins of astrometric quality and apparent magnitude.""" rows = [] gmed = float(np.median(c.G.values)) bins = dict(ruwe=[(a, b) for a, b in zip(RUWE_BINS[:-1], RUWE_BINS[1:]) if a < L.CUTSETS[cutset]["ruwe_max"]], poe=[(a, b) for a, b in zip(POE_BINS[:-1], POE_BINS[1:]) if b > L.CUTSETS[cutset]["poe_min"]], G=[(0.0, gmed), (gmed, 99.0)]) for f in L.F_GRID: sens, fl = L.flag_mask(c, f) for qty, bb in bins.items(): x = c[qty].values for lo, hi in bb: q = (x >= lo) & (x < hi) ns, nf = int((sens & q).sum()), int((fl & q).sum()) rows.append(dict(cutset=cutset, survey=sv, f=f, quantity=qty, lo=lo, hi=hi, n_sensitive=ns, n_flagged=nf, rate=nf / ns if ns else np.nan, rate_UL95=L.poisson_ul(nf) / ns if ns else np.nan)) if cutset == "orig": # RUWE bins separately below and above parallax/error = 50 for plo, phi in ((20.0, 50.0), (50.0, np.inf)): for lo, hi in bins["ruwe"]: q = (c.ruwe.values >= lo) & (c.ruwe.values < hi) & (c.poe.values >= plo) & (c.poe.values < phi) ns, nf = int((sens & q).sum()), int((fl & q).sum()) rows.append(dict(cutset=cutset, survey=sv, f=f, quantity="ruwe | poe %g-%g" % (plo, phi), lo=lo, hi=hi, n_sensitive=ns, n_flagged=nf, rate=nf / ns if ns else np.nan, rate_UL95=L.poisson_ul(nf) / ns if ns else np.nan)) return rows def run(): from injection import correction ctx = {} for cs, dd in (("final", DATA_MAIN), ("orig", DATA_ORIG)): if not os.path.exists(os.path.join(dd, "limits.csv")): continue cand = pd.read_csv(os.path.join(dd, "candidates_all.csv")) ctx[cs] = dict(eps=pd.read_csv(os.path.join(dd, "selection_eff.csv")), lim=pd.read_csv(os.path.join(dd, "limits.csv")), cand=cand) rows, diag = [], {"final": [], "orig": []} def resolved(cs, sv): cc = ctx[cs]["cand"] cc = cc[cc.survey == sv] res = cc[["status_%g" % g for g in L.F_GRID]].apply(lambda col: col.astype(str).str.startswith("resolved")) return set(cc.source_id[res.any(axis=1)]) def record(sv, vname, kind, c, nthr, nsens, req, any_res, eps_cs="final", v_cs="final"): eps, lim0 = ctx[eps_cs]["eps"], ctx[v_cs]["lim"] for f in L.F_GRID: sens, fl = L.flag_mask(c, f, nsig_thr=nthr, nsig_sens=nsens, require=req) _, rec = L.flag_mask(c, f, nsig_thr=nthr, nsig_sens=nsens, require=req, inject=True) ids = set(c.source_id.values[fl]) n_unres = len([i for i in ids if i not in any_res]) expo = 0.0 for cname, _, _ in CLASSES: q = sens & (c.cls.values == cname) if q.sum() == 0: continue e = eps[(eps.survey == sv) & (eps.cls == cname) & (eps.f == f)].eps.iloc[0] expo += q.sum() * rec[q].mean() * e expo *= correction(sv, f) V = lim0[(lim0["sample"] == sv) & (lim0.f == f) & (lim0.tier == "three-way")].mean_V.iloc[0] in_faint = c.G.values >= np.median(c.G.values) rows.append(dict(survey=sv, variant=vname, kind=kind, f=f, n_clean=len(c), n_sens=int(sens.sum()), n_flagged=int(fl.sum()), n_unresolved=n_unres, n_flagged_fainter_half=int((fl & in_faint).sum()), recovery=float(rec[sens].mean()) if sens.any() else np.nan, sigma_core=L.robust_sigma(c.D.values), UL95_pop=L.poisson_ul(n_unres) / (expo * V) if expo > 0 else np.nan, UL95_pop_unvetted=L.poisson_ul(int(fl.sum())) / expo if expo > 0 else np.nan)) for sv in SURVEYS: d = pd.read_pickle(os.path.join(WORK, "%s_stars.pkl" % sv)) m, _ = L.apply_cuts(d, sv, verbose=False) # final cut set (this stage always runs with it) d = d.loc[m, [c for c in NEEDED if c in d.columns]].reset_index(drop=True) any_res = resolved("final", sv) base = None for vname, kw, reuse_floors in RECOMPUTE: log("robustness: %s / %s" % (sv, vname)) s = prepare(d, sv, floors=base if reuse_floors else None, precut=True, **kw) if vname == "baseline": base = s[["source_id"] + L.FLOOR_COLS] record(sv, vname, "recompute", *evaluate(s), any_res) if vname == "baseline": for v2, kw2 in ON_BASE: kind = "subset" if v2.startswith("subset") else "threshold" record(sv, v2, kind, *evaluate(s, **kw2), any_res) diag["final"] += diagnostic(evaluate(s)[0], sv, "final") del d, s if "orig" in ctx and os.path.exists(dpkl(sv, "orig")): log("robustness: %s / subsets of the original-cut sample" % sv) so = from_table(sv, "orig") res_o = resolved("orig", sv) for v2, kw2, eps_cs in ON_ORIG: record(sv, v2, "subset", *evaluate(so, **kw2), res_o, eps_cs=eps_cs, v_cs="orig") diag["orig"] += diagnostic(evaluate(so)[0], sv, "orig") del so out = pd.DataFrame(rows) out.to_csv(os.path.join(DATA_MAIN, "robustness.csv"), index=False) d_o, d_f = pd.DataFrame(diag["orig"]), pd.DataFrame(diag["final"]) d_o.to_csv(os.path.join(DATA_MAIN, "astrometry_diagnostic.csv"), index=False) d_f.to_csv(os.path.join(DATA_MAIN, "astrometry_diagnostic_final.csv"), index=False) pd.set_option("display.width", 250) pd.set_option("display.max_rows", 900) piv = out.pivot_table(index=["survey", "variant"], columns="f", values=["n_flagged", "n_unresolved", "UL95_pop"], sort=False) print(piv.to_string(float_format=lambda x: "%.3g" % x)) for t in (d_o, d_f): if len(t): print(t[t.f <= 0.5].to_string(index=False, float_format=lambda x: "%.3g" % x)) return out if __name__ == "__main__": run()