test_model_survival.py

#!/usr/bin/env python
import paths as pt
from tools import data_loader, file_reader
from utility.settings import load_settings
from sksurv.ensemble import RandomSurvivalForest
from sklearn.model_selection import KFold
from sksurv.metrics import concordance_index_censored
from io import StringIO
from pathlib import Path
import shutil

def main():
    data_settings = load_settings(pt.CONFIGS_DIR, "data.yaml")
    target_settings = load_settings(pt.CONFIGS_DIR, "alarm.yaml")
    ats_resolution = data_settings['ats_resolution']

    dl = data_loader.AlarmDataLoader(pt.PROCESSED_DATA_DIR,
                                     "alarm_data.pkl",
                                     target_settings).load_data()
    X, y = dl.get_data()

    ats = {str(i)+'Ats':str for i in range(1, ats_resolution+1)}
    infile = StringIO()
    file_path = pt.PROCESSED_DATA_DIR
    file_name = "alarm_emb.csv"
    with open(Path.joinpath(file_path, file_name), 'r') as fd:
        shutil.copyfileobj(fd, infile)
        infile.seek(0)
        emb_file = file_reader.read_csv(infile, converters=ats)
    X = emb_file

    model = RandomSurvivalForest(n_estimators=200,
                                 max_depth=3,
                                 random_state=0)
    kf = KFold(n_splits=5, shuffle=True, random_state=0)
    c_index_scores = list()

    for train, test in kf.split(X, y): #8 min execution time
        model.fit(X.iloc[train], y[train])
        prediction = model.predict(X.iloc[test])
        c_index = concordance_index_censored(y[test]["Status"],
                                             y[test]["Days_to_alarm"],
                                             prediction)
        c_index_scores.append(c_index[0])

    print(c_index_scores)

if __name__ == '__main__':
    main()