ml_pipeline/read_data_split.py at main · ambercheny/ml_pipeline · GitHub

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
import os

import pandas as pd


def _read_table(path: str) -> pd.DataFrame:
    ext = os.path.splitext(path)[1].lower()
    if ext == ".parquet":
        return pd.read_parquet(path)
    if ext == ".pkl":
        return pd.read_pickle(path)
    if ext == ".csv":
        return pd.read_csv(path)
    raise ValueError(f"Unsupported file extension: {ext}. Use parquet/pkl/csv.")


def _read_split_ids(id_dir: str, sample_id: str):
    split_to_ids = {}
    for split_name in ["train", "val", "test"]:
        split_path = os.path.join(id_dir, f"{split_name}_ids.csv")
        split_df = pd.read_csv(split_path)
        if sample_id in split_df.columns:
            split_to_ids[split_name] = split_df[sample_id].astype(str)
        else:
            # Fallback for legacy files where the first column may have a wrong name.
            split_to_ids[split_name] = split_df.iloc[:, 0].astype(str)
    return split_to_ids


def _split_feature_df(logger, feature_df, split_to_ids, sample_id):
    feature_df = feature_df.copy()
    feature_df[sample_id] = feature_df[sample_id].astype(str)

    train_ids = set(split_to_ids["train"].unique())
    val_ids = set(split_to_ids["val"].unique())
    test_ids = set(split_to_ids["test"].unique())

    df_train = feature_df[feature_df[sample_id].isin(train_ids)]
    df_valid = feature_df[feature_df[sample_id].isin(val_ids)]
    df_test = feature_df[feature_df[sample_id].isin(test_ids)]

    logger.info(f"df_train: {df_train.shape}, sample_id: {df_train[sample_id].nunique()}")
    logger.info(f"df_valid: {df_valid.shape}, sample_id: {df_valid[sample_id].nunique()}")
    logger.info(f"df_test: {df_test.shape}, sample_id: {df_test[sample_id].nunique()}")

    train_valid_overlap = set(df_train[sample_id]).intersection(df_valid[sample_id])
    train_test_overlap = set(df_train[sample_id]).intersection(df_test[sample_id])
    valid_test_overlap = set(df_valid[sample_id]).intersection(df_test[sample_id])
    logger.info("check overlapping sample ids")
    if len(train_valid_overlap) + len(train_test_overlap) + len(valid_test_overlap) > 0:
        if len(train_valid_overlap) > 0:
            logger.info(f"Train-Valid overlap: {len(train_valid_overlap)}")
        if len(train_test_overlap) > 0:
            logger.info(f"Train-Test overlap: {len(train_test_overlap)}")
        if len(valid_test_overlap) > 0:
            logger.info(f"Valid-Test overlap: {len(valid_test_overlap)}")
    else:
        logger.info("  >>>> No overlapping sample ids")

    return df_train, df_valid, df_test


def read_data_split(
    logger,
    data_seed,
    feature_folder,
    id_folder,
    outcome_folder,
    outcome, # outcome string
    outcome_file, # outcome file name
    feature_file,
    sample_id="sample_id",
    outcome_col="event", # outcome label
):
    # 1. get all features
    feature_path = os.path.join(feature_folder, feature_file)
    feature_df = _read_table(feature_path)
    logger.info(f"feature read from {feature_path}, shape={feature_df.shape}")


    # 2. get id path
    id_dir = os.path.join(id_folder, f"seed{data_seed}/{outcome}")

    if not os.path.exists(id_dir):
        # Fallback: search for the exact seed split folder with train_ids.csv
        candidate_root = os.path.join(id_folder, f"seed{data_seed}")
        candidates = []
        if os.path.isdir(candidate_root):
            for name in os.listdir(candidate_root):
                c = os.path.join(candidate_root, name)
                if os.path.isfile(os.path.join(c, "train_ids.csv")):
                    candidates.append(c)
        if len(candidates) == 1:
            id_dir = candidates[0]
        elif len(candidates) > 1:
            raise ValueError(f"Multiple split folders found under {candidate_root}: {candidates}")
        else:
            raise FileNotFoundError(f"Could not find split folder under {candidate_root}")

    # 3. split ids
    split_to_ids = _read_split_ids(id_dir, sample_id)
    df_train, df_valid, df_test = _split_feature_df(logger, feature_df, split_to_ids, sample_id)


    # 4. get outcomes
    outcome_path = os.path.join(outcome_folder, outcome_file)
    outcome_df = _read_table(outcome_path)
    outcome_df = outcome_df[[sample_id, outcome_col]].dropna().copy()
    outcome_df[sample_id] = outcome_df[sample_id].astype(str)
    ### get outcomes of train, valid, test
    y_train = df_train[[sample_id]].merge(outcome_df, on=sample_id, how="inner")
    y_valid = df_valid[[sample_id]].merge(outcome_df, on=sample_id, how="inner")
    y_test = (
        df_test[[sample_id]]
        .merge(outcome_df, on=sample_id, how="inner")
        .drop_duplicates(subset=sample_id, keep="first")
    )
    logger.info(f"y shape: train {y_train.shape}, valid {y_valid.shape}, test {y_test.shape}")

    # 4. get features of train, valid, test
    df_train = df_train[df_train[sample_id].isin(y_train[sample_id].unique())]
    df_valid = df_valid[df_valid[sample_id].isin(y_valid[sample_id].unique())]
    df_test = df_test[df_test[sample_id].isin(y_test[sample_id].unique())]
    logger.info(f"X shape: train {df_train.shape}, valid {df_valid.shape}, test {df_test.shape}")

    feature_cols = [c for c in df_train.columns if c != sample_id]
    X_train = df_train[feature_cols]
    X_valid = df_valid[feature_cols]
    X_test = df_test[feature_cols]

    y_train = y_train[outcome_col]
    y_valid = y_valid[outcome_col]
    y_test = y_test[outcome_col]

    logger.info(f"X isnull: train {X_train.isnull().sum().sum()}, valid {X_valid.isnull().sum().sum()}, test {X_test.isnull().sum().sum()}")
    logger.info(f"y isnull: train {y_train.isnull().sum()}, valid {y_valid.isnull().sum()}, test {y_test.isnull().sum()}")
    logger.info(f"y outcome: train {y_train.mean()}, valid {y_valid.mean()}, test {y_test.mean()}")

    return X_train, X_valid, X_test, y_train, y_valid, y_test