Importing Libraries

import pandas as pd
import numpy as np
import time

from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.feature_selection import RFE
from sklearn.metrics import roc_auc_score, recall_score
from sklearn.model_selection import train_test_split

Loading data and define X / y

# Load cleaned dataset
df = pd.read_pickle("../data/cleaned/NSDUH_2023_clean.pkl")

TARGET = "IRAMDEYR"

# If ANALWT2_C is present, drop it from features (RFE doesn't use weights)
cols_to_drop = [TARGET]
if "ANALWT2_C" in df.columns:
    cols_to_drop.append("ANALWT2_C")

X = df.drop(columns=cols_to_drop)
y = df[TARGET]

print("X shape:", X.shape)
print("y distribution:\n", y.value_counts())

X shape: (45133, 2315)
y distribution:
 IRAMDEYR
0.0    39948
1.0     5185
Name: count, dtype: int64

Ensure numeric & handle NaNs (features only)

# Make sure everything in X is numeric; convert errors to NaN then fill with 0
X = X.apply(pd.to_numeric, errors="coerce").fillna(0)

Train/Test split (NO LEAKAGE)

X_train, X_test, y_train, y_test = train_test_split(
    X,
    y,
    test_size=0.2,
    stratify=y,
    random_state=42
)

print("Train:", X_train.shape, "Test:", X_test.shape)

Train: (36106, 2315) Test: (9027, 2315)

Standardize (fit on train only)

scaler = StandardScaler()

X_train_scaled = scaler.fit_transform(X_train)  # fit only on train
X_test_scaled  = scaler.transform(X_test)       # transform test with same scaler

RFE (AUC)

auc_model = LogisticRegression(
    max_iter=2000,
    n_jobs=-1,
    # no class_weight here → more balanced between precision/recall
)

n_features_to_select = 100
step_size = 0.1  # drop 10% per iteration

start_time = time.time()
print("Starting RFE (AUC)")

rfe_auc = RFE(
    estimator=auc_model,
    n_features_to_select=n_features_to_select,
    step=step_size
)

rfe_auc.fit(X_train_scaled, y_train)

end_time = time.time()
print(f"RFE (AUC) completed in {(end_time - start_time)/60:.2f} minutes")

Starting RFE (AUC)
RFE (AUC) completed in 1.69 minutes

Evaluate AUC and save top-100 list

# Predict probabilities on test using only selected features
X_test_auc = X_test_scaled[:, rfe_auc.support_]
pred_auc   = rfe_auc.estimator_.predict_proba(X_test_auc)[:, 1]

auc_score = roc_auc_score(y_test, pred_auc)
print(f"AUC Score using RFE_AUC features: {auc_score:.4f}")

# Get feature names for the selected mask
feature_names = X.columns
rfe_auc_features = feature_names[rfe_auc.support_].tolist()

# Save top-100 list
pd.Series(rfe_auc_features).to_csv(
    "../results/rfe_auc_top100.csv",
    index=False
)

print("\nRFE_AUC Top 100 Features:")
for f in rfe_auc_features:
    print(f)

AUC Score using RFE_AUC features: 0.9989

RFE_AUC Top 100 Features:
AGE3
IRMARIT
CATAG7
DRVINDETAG
WRKDHRSWK2
WRK35WKUS
WRKDPSTYR
WRKDRGHLP
WRKOKRAND
IRHH65_2
IRPINC3
BKDRVINF
UADOTHM
IITRQANYREC
CIGYR
TOBMON
MJSKNYR
MJSKNMON
CNSANYYR
ILLEMMON
PSILCYEVER
RXHYDMANY
FUCGR18
FUNICVAP21
FULSD18
PYUD5MRJ
EDUD5PNRUNM
UD5ILALANY
SVYRDUDANY
SUTOUTHER
IRSUTINRHAB
SUTOUTCNSPY
SUTNEEDPY
SUNTNOHLP
RSKMRJMON
DIFGETCRK
GRSKMRJWK
DSTNRV30
IMPWORK
COSUITHNK
KSSLR6YRED
WHODASDAED
IIDSTHOP30
IIDSTEFF30
IRIMPGOUT
IIIMPPEOP
IIIMPPEOPM
IRIMPSOC
IRIMPHHLD
IIIMPHHLD
IIIMPHHLDM
IRIMPRESP
IRIMPRESPM
IRIMPWORK
IIIMPWORK
IRSUICTHNK
IRCOSUITHNK
IICOSUIPLNYR
KSSLR6MAX
SPDPSTYR
AKSSLR6WRST
WHODASTOTSC
WHODASDASC
SMIPPPY
SMIPY
AMIPY
SMMIPY
MMIPY
LMIPY
LMMIPY
MICATPY
AMISUD5ANY
LMMISUD5ANY
SMIRSUD5ANY
AMIRSUD5ANY
AMISUD5ANYO
ADSUITPAYR
ADWRELES
ADWREMOR
ADWRSLEP
ADWRSMOR
ADWRENRG
ADWRSLOW
ADWRSLNO
ADWRTHOT
ADPSHMGT
ADRX12MO
ASDSHOM2
ASDSWRK2
ASDSREL2
ASDSSOC2
IRMHTRXMED
MHTRTPY
MHNTENFCV
COPDAGE
CASUPROB
CAMHRCVR
IMFNDLEVER
IIIMFREC
COCLALCUSE

RFE (Recall)

rec_model = LogisticRegression(
    max_iter=2000,
    n_jobs=-1,
    class_weight="balanced"  # tilt toward minority class → better recall
)

start_time = time.time()
print("Starting RFE (Recall)")

rfe_rec = RFE(
    estimator=rec_model,
    n_features_to_select=n_features_to_select,
    step=step_size
)

rfe_rec.fit(X_train_scaled, y_train)

end_time = time.time()
print(f"RFE (Recall) completed in {(end_time - start_time)/60:.2f} minutes")

Starting RFE (Recall)
RFE (Recall) completed in 1.43 minutes

Evaluate Recall and save top-100 list

# Predict class labels on test using only selected features
X_test_rec = X_test_scaled[:, rfe_rec.support_]
pred_rec   = rfe_rec.estimator_.predict(X_test_rec)

rec_score = recall_score(y_test, pred_rec)
print(f"Recall using RFE_Recall features: {rec_score:.4f}")

# Get feature names for the selected mask
rfe_rec_features = feature_names[rfe_rec.support_].tolist()

# Save top-100 list
pd.Series(rfe_rec_features).to_csv(
    "../results/rfe_recall_top100.csv",
    index=False
)

print("\nRFE_Recall Top 100 Features:")
for f in rfe_rec_features:
    print(f)

Recall using RFE_Recall features: 0.9894

RFE_Recall Top 100 Features:
AGE3
IRSEX
CATAG7
DRVINDETAG
SEXAGE
WRKDHRSWK2
WRK35WKUS
WRKDPSTYR
WRKDRGEDU
WRKDRGHLP
IRKI17_2
IIKI17_2
CAIDCHIP
IRFAMSSI
IRPINC3
BKMVTHFT
IIBZONMYR
IISTMNMINIT
CIGYR
TOBMON
TOBVNICFLAG
MJSMKYR
MJSKNYR
MJSKNMON
CNSANYYR
MJONLYFLAG
PSILCYEVER
RXHYDMANY
FUNICVAP21
FULSD18
PYUD5MRJ
UD5ILALANY
SUTOUTHER
SUTOUTCNSPY
SUNTNOHLP
GRSKMRJWK
SNRLDCSN
SNRLFRND
DSTNRV30
DSTHOP30
IMPREMEM
IMPGOUT
IMPWORK
COSUITHNK
KSSLR6YRED
WHODASDAED
IIDSTEFF30
IRDSTNGD30
IRIMPGOUT
IRIMPSOC
IRIMPHHLD
IRIMPRESP
IRIMPRESPM
IRIMPWORK
IRSUICTHNK
IRCOSUITHNK
KSSLR6MAX
AKSSLR6WRST
WHODASTOTSC
WHODASDASC
SMIPPPY
SMIPY
AMIPY
SMMIPY
MMIPY
LMIPY
LMMIPY
MICATPY
AMIRSUD5ANY
AMISUD5ANYO
ADSUITPAYR
ADWRELES
ADWRSLEP
ADWRSMOR
ADWRENRG
ADWRSLOW
ADWRSLNO
ADWRTHOT
ADPSHMGT
ADRX12MO
ASDSHOM2
ASDSWRK2
ASDSREL2
ASDSSOC2
IRMHTRXMED
MHTRTPY
MHTNSEEKPY
MHTSKTHPY
DIABETEAG
COPDAGE
ASTHMAAGE
ASTHMANOW
CASUPROB
CAMHRCVR
SYNSTMEVR
IMFEVER
IISYNSTMREC
CASUPROB2
CAMHPROB2
COCLALCUSE