SleepGuardAI / find_threshold_recall.py
ElMETRICO's picture
Upload 5 files
b703a33 verified
Raw
History Blame Contribute Delete
2.25 kB
import joblib
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import recall_score, fbeta_score, roc_auc_score
MODEL_PATH = "best_model_v2_calibrated.joblib"
CSV_PATH = "nhanes_sleep_extended_v2.csv"
TARGET = "Sleep_Risk"
FEATURES = [
"Age",
"Gender",
"BMI",
"BP_SYS",
"BP_DIA",
"Phys_Activity_Days",
"DPQ_Score",
"Smoking_Indicator",
"Alcohol_Feature",
"Diabetes_Indicator",
"Cycle",
]
def main():
df = pd.read_csv(CSV_PATH)
# drop rows with missing values in required columns
df = df.dropna(subset=FEATURES + [TARGET]).copy()
X = df[FEATURES]
y = df[TARGET].astype(int)
# Train/test split (keep same logic as training if possible)
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.20, random_state=42, stratify=y
)
bundle = joblib.load(MODEL_PATH)
model = bundle.get("model") if isinstance(bundle, dict) and "model" in bundle else bundle
calibrator = bundle.get("calibrator") if isinstance(bundle, dict) and "calibrator" in bundle else None
# Get probabilities
p_test_raw = model.predict_proba(X_test)[:, 1]
if calibrator is not None:
p_test = calibrator.predict(p_test_raw.reshape(-1, 1))
else:
p_test = p_test_raw
# AUC (threshold-free)
auc = roc_auc_score(y_test, p_test)
print(f"TEST AUC: {auc:.4f}")
# Sweep thresholds
thresholds = np.round(np.arange(0.05, 0.90, 0.01), 2)
best_f2 = (-1, None)
th_recall_085 = None
for th in thresholds:
pred = (p_test >= th).astype(int)
rec = recall_score(y_test, pred)
f2 = fbeta_score(y_test, pred, beta=2)
if f2 > best_f2[0]:
best_f2 = (f2, th)
if th_recall_085 is None and rec >= 0.85:
th_recall_085 = th
print(f"Best threshold by F2: {best_f2[1]} (F2={best_f2[0]:.4f})")
if th_recall_085 is not None:
print(f"Threshold achieving recall ≥ 0.85 : {th_recall_085}")
else:
print("No threshold achieved recall ≥ 0.85 in tested range.")
if __name__ == "__main__":
main()