##Radiomics model

import os
import pandas as pd
from IPython.display import display
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
from onekey_algo import OnekeyDS as okds
from onekey_algo import get_param_in_cwd

os.makedirs('img', exist_ok=True)
os.makedirs('results', exist_ok=True)
os.makedirs('features', exist_ok=True)
task_type = 'Radiomics_'
mydir = get_param_in_cwd('radio_dir')
group_info = get_param_in_cwd('dataset_column') or 'group'
labelf = get_param_in_cwd('label_file') or os.path.join(mydir, 'label.csv')
labels = [get_param_in_cwd('task_column') or 'label']

from pathlib import Path
from onekey_algo.custom.components.Radiology import diagnose_3d_image_mask_settings, get_image_mask_from_dir

images, masks = get_image_mask_from_dir(mydir, images=f'images', masks=f'masks')

import warnings
import pandas as pd
 
warnings.filterwarnings("ignore")

from onekey_algo.custom.components.Radiology import ConventionalRadiomics

rad_ = None
if os.path.exists(f'features/rad_features.csv'):
    rad_data = pd.read_csv(f'features/rad_features.csv', header=0)
else:
    images, masks = get_image_mask_from_dir(mydir, images=f'images', masks=f'masks')
    param_file = get_param_in_cwd('extractor_settings')
    radiomics = ConventionalRadiomics(param_file, correctMask=True)
    radiomics.extract(images, masks, workers=1, with_fd=get_param_in_cwd('with_fd'))
    rad_data = radiomics.get_label_data_frame(label=1)
    rad_data.columns = [f"intra_{c.replace('-', '_')}" if c != 'ID' else c for c in rad_data.columns]
    rad_data.to_csv(f'features/rad_features.csv', header=True, index=False)
if rad_ is None:
    rad_ = rad_data
else:
    rad_ = pd.merge(rad_, rad_data, on='ID', how='inner')
rad_data = rad_
rad_data

import matplotlib.pyplot as plt
import seaborn as sns
sorted_counts = pd.DataFrame([c.split('_')[-2] for c in rad_data.columns if c !='ID']).value_counts()
sorted_counts = pd.DataFrame(sorted_counts, columns=['count']).reset_index()
sorted_counts = sorted_counts.sort_values(0)
sorted_counts.columns = ['feature_group', 'count']
display(sorted_counts)

plt.figure(figsize=(20, 10))
ax = plt.subplot(121)
plt.pie(sorted_counts['count'], labels=[i for i in sorted_counts['feature_group']], startangle=0,
        counterclock = False, autopct = '%.1f%%')
ax = plt.subplot(122)
ax.spines['top'].set_visible(False)
ax.spines['right'].set_visible(False)
bar = sns.barplot(data=sorted_counts, x='feature_group', y='count', )
plt.savefig(f'img/{task_type}feature_ratio.svg', bbox_inches = 'tight')

from onekey_algo.custom.components.comp1 import fillna
label_data = pd.read_csv(labelf, dtype={'ID': str})
label_data['ID'] = label_data['ID'].map(lambda x: f"{x}.nii.gz" if not (f"{x}".endswith('.nii.gz') or  f"{x}".endswith('.nii')) else x)
label_data

from onekey_algo.custom.utils import print_join_info
event_col = get_param_in_cwd('event_col')
duration_col= get_param_in_cwd('duration_col')

print_join_info(rad_data, label_data)
combined_data = pd.merge(rad_data, label_data[['ID', event_col, duration_col, 'group']], on=['ID'], how='inner')
print(combined_data[event_col].value_counts())
combined_data

from onekey_algo.custom.components.comp1 import normalize_df
data = normalize_df(combined_data, not_norm=['ID', event_col, duration_col], group='group', use_train=True)
data = data.dropna(axis=1)
data.describe()

import pandas as pd
import pingouin as pg
r1a = pd.read_csv('radiomics_r1a.csv', index_col=0)  
r1b = pd.read_csv('radiomics_r1b.csv', index_col=0) 
r2  = pd.read_csv('radiomics_r2.csv',  index_col=0)  
common_cases = r1a.index.intersection(r1b.index).intersection(r2.index)
r1a = r1a.loc[common_cases]
r1b = r1b.loc[common_cases]
r2  = r2.loc[common_cases]
features = r1a.columns.tolist()
icc_summary = pd.DataFrame(index=features,
                           columns=['ICC_intra','ICC_inter'])

for feat in features:
    df_intra = pd.DataFrame({
        'case' : list(common_cases)*2,
        'rater': ['R1a']*len(common_cases) + ['R1b']*len(common_cases),
        'value': r1a[feat].tolist() + r1b[feat].tolist()
    })
    icc_i = pg.intraclass_corr(data=df_intra,
                               targets='case',
                               raters='rater',
                               ratings='value')
    icc_intra = icc_i.loc[icc_i['Type']=='ICC3','ICC'].values[0]
    
    df_inter = pd.DataFrame({
        'case' : list(common_cases)*2,
        'rater': ['R1a']*len(common_cases) + ['R2']*len(common_cases),
        'value': r1a[feat].tolist() + r2[feat].tolist()
    })
    icc_e = pg.intraclass_corr(data=df_inter,
                               targets='case',
                               raters='rater',
                               ratings='value')
    icc_inter = icc_e.loc[icc_e['Type']=='ICC3','ICC'].values[0]
    
    icc_summary.loc[feat,'ICC_intra']  = icc_intra
    icc_summary.loc[feat,'ICC_inter']  = icc_inter

good_feats = icc_summary[
    (icc_summary['ICC_intra']  >= 0.8) &
    (icc_summary['ICC_inter']  >= 0.8)
]
good_feats.to_csv('rad_features_sel.csv')    

from onekey_algo.custom.components.comp1 import normalize_df, select_feature

corr_name = get_param_in_cwd('corr_name', 'pearson')
if os.path.exists(f'features/{task_type}rad_features_sel.csv') and False:
    data = pd.read_csv(f'features/{task_type}rad_features_sel.csv', header=0)
else:
    tgroup = data[data['group'] == 'train']
    sel_feature = select_feature(tgroup[[c for c in tgroup.columns if c not in [event_col, duration_col]]].corr(corr_name), 
                                 threshold=0.8, topn=134, verbose=False)
    data = data[['ID'] + sel_feature + [event_col, duration_col, 'group']]
    data.to_csv(f'features/{task_type}rad_features_sel.csv', header=True, index=False)
data

import numpy as np
import onekey_algo.custom.components as okcomp
from collections import OrderedDict

train_data = data[(data[group_info] == 'train')]
subsets = get_param_in_cwd('subsets')
val_datasets = OrderedDict()
for subset in subsets:
    val_data = data[data[group_info] == subset]
    val_datasets[subset] = val_data
    val_data.to_csv(f'features/{task_type}{subset}_features_norm.csv', index=False)

print('，'.join([f"{subset}样本数：{d_.shape}" for subset, d_ in val_datasets.items()]))

from onekey_algo.custom.components.survival import uni_cox

if os.path.exists(f'features/{task_type}rad_features_unisel.csv') and False:
    train_data = pd.read_csv(f'features/{task_type}rad_features_unisel.csv')
else:
    sel_features = uni_cox(train_data, duration_col=duration_col, event_col=event_col,
                           cols=[c for c in train_data.columns if c not in [event_col, duration_col, 'ID', 'group']], 
                           verbose=False, pvalue_thres=get_param_in_cwd('pvalue', 0.05))
    train_data = train_data[['ID'] + sel_features + [event_col, duration_col, 'group']]
    train_data.to_csv(f'features/{task_type}rad_features_unisel.csv', header=True, index=False)
train_data

from onekey_algo.custom.components.survival import get_x_y_survival, lasso_cox_cv
COEF_THRESHOLD = 1e-6

X, y = get_x_y_survival(train_data, val_outcome=1, event_col=event_col, duration_col=duration_col)
sel_features = lasso_cox_cv(X, y, max_iter=100,  norm_X=False, prefix=f"{task_type}", l1_ratio=0.5, cv=10, weights_fig_size=(10, 15))

train_data = train_data[['ID'] + list(sel_features.index) + [event_col, duration_col]]
for subset in subsets:
    val_datasets[subset] = val_datasets[subset][['ID'] + list(sel_features.index) + [event_col, duration_col]]
    val_datasets[subset].to_csv(f'features/{task_type}{subset}_cox.csv', index=False)

import seaborn as sns
import matplotlib.pyplot as plt

if train_data.shape[1] < 150:
    pp = sns.clustermap(train_data[[c for c in train_data.columns if c not in [event_col, duration_col]]].corr(corr_name), 
                        linewidths=.5, figsize=(20.0, 16.0), cmap='YlGnBu')
    plt.setp(pp.ax_heatmap.get_yticklabels(), rotation=0)
    plt.savefig(f'img/{task_type}feature_cluster.svg', bbox_inches = 'tight')

from lifelines import CoxPHFitter

cph = CoxPHFitter(penalizer=0.1)
cph.fit(train_data[[c for c in train_data.columns if c != 'ID']], duration_col=duration_col, event_col=event_col)
cph.print_summary()

print(cph.concordance_index_)
su = cph.summary[['exp(coef)', 'exp(coef) lower 95%', 'exp(coef) upper 95%', 'p']]
su.columns = ['HR', 'HR lower 95%', 'HR upper 95%', 'pvalue']
su.reset_index().to_csv(f'features/{task_type}features_HR.csv', index=False)
su

import matplotlib.pyplot as plt

plt.figure(figsize=(10, 8))
cph.plot(hazard_ratios=True)
plt.savefig(f'img/{task_type}feature_pvalue.svg')
plt.show()

from lifelines import CoxPHFitter
from lifelines.statistics import logrank_test
from lifelines import KaplanMeierFitter
from lifelines.plotting import add_at_risk_counts

thres = 1e-4
bst_split = {'train': 2.0, 'val':1, 'test': 1}
loc = {'train': 0.2, 'val':0.5, 'test': 0.5}
for subset, test_data in val_datasets.items():
    c_index = cph.score(test_data[[c for c in test_data.columns if c != 'ID']], scoring_method="concordance_index")
    y_pred = cph.predict_partial_hazard(test_data[[c for c in test_data.columns if c != 'ID']])
    cox_data = pd.concat([test_data, y_pred], axis=1)
    mean = cox_data.describe()[0]['50%']
    cox_data['HR'] = cox_data[0] > mean
        cox_data['HR'] = cox_data[0] > 1
    
    cox_data['HR'] = cox_data[0] > bst_split[subset]

    dem = (cox_data["HR"] == True)
    results = logrank_test(cox_data[duration_col][dem], cox_data[duration_col][~dem], 
                           event_observed_A=cox_data[event_col][dem], event_observed_B=cox_data[event_col][~dem])
    p_value = f"={results.p_value:.3f}" if results.p_value > thres else f'<{thres}'
    plt.title(f"Cohort {subset} C-index:{c_index:.3f}")
    plt.ylabel('Probability')
    if sum(dem):
        kmf_high = KaplanMeierFitter()
        kmf_high.fit(cox_data[duration_col][dem], event_observed=cox_data[event_col][dem], label="High Risk")
        kmf_high.plot_survival_function(color='r')
    if sum(~dem):
        kmf_low = KaplanMeierFitter()
        kmf_low.fit(cox_data[duration_col][~dem], event_observed=cox_data[event_col][~dem], label="Low Risk")
        kmf_low.plot_survival_function(color='g')
    plt.text(0.5, loc[subset] if subset in loc else 0.2, f"P{p_value}")
    plt.xlabel('Time(months)')
    add_at_risk_counts(kmf_high, kmf_low, rows_to_show=['At risk'])
    plt.savefig(f'img/{task_type}KM_{subset}.svg', bbox_inches='tight')
    plt.show()

import os
import numpy as np

def get_prediction(model: CoxPHFitter, data, ID=None, **kwargs):
    hr = model.predict_partial_hazard(data)
    expectation = model.predict_expectation(data)
    
    predictions = pd.concat([hr, expectation], axis=1)
    predictions.columns = ['HR', 'expectation']
    if ID is not None:
        predictions = pd.concat([ID, hr, expectation], axis=1)
        predictions.columns = ['ID', 'HR', 'expectation']
    else:
        predictions = pd.concat([hr, expectation], axis=1)
        predictions.columns = ['HR', 'expectation']
    return predictions

os.makedirs('results', exist_ok=True)
info = []
for subset, test_data in val_datasets.items():
    if subset in get_param_in_cwd('subsets'):
        results = get_prediction(cph, test_data, ID=test_data['ID'])
        results.to_csv(f'results/{task_type}cox_predictions_{subset}.csv', index=False)
        results['group'] = subset
        info.append(results)
        pd.merge(results, label_data[['ID', event_col, duration_col]], on='ID', how='inner').to_csv(f'features/{task_type}4xtile_{subset}.txt', 
                                                                                                    index=False, sep='\t')
info = pd.concat(info, axis=0)
info

##Clinical mode

from lifelines import CoxPHFitter
import pandas as pd
from onekey_algo.custom.components.comp1 import normalize_df
from sklearn.model_selection import train_test_split
from onekey_algo import get_param_in_cwd
from onekey_algo.custom.components.comp1 import fillna

event_col = get_param_in_cwd('event_col')
group_info = 'group'
task_type = 'Clinical_'
duration_col= get_param_in_cwd('duration_col')
data = pd.read_csv('data/clinical.csv', dtype={'ID': str})
data = data[data.columns[:-3]]
label_data = pd.read_csv(get_param_in_cwd('label_file'), dtype={'ID': str})
label_data['ID'] = label_data['ID'].map(lambda x: f"{x}.nii.gz" if not (f"{x}".endswith('.nii.gz') or  f"{x}".endswith('.nii')) else x)
data = pd.merge(data, label_data[['ID', event_col, duration_col, 'group']], on='ID', how='inner')
data

import numpy as np
import onekey_algo.custom.components as okcomp
from collections import OrderedDict

train_data = data[(data[group_info] == 'train')]
subsets = get_param_in_cwd('subsets')
val_datasets = OrderedDict()
for subset in subsets:
    val_data = data[data[group_info] == subset]
    val_datasets[subset] = val_data
    val_data.to_csv(f'features/{task_type}{subset}_features_norm.csv', index=False)

import os
from onekey_algo.custom.components.survival import uni_cox

if os.path.exists(f'features/{task_type}features_unisel.csv') and False:
    train_data = pd.read_csv(f'features/{task_type}features_unisel.csv')
else:
    sel_features, uni_info = uni_cox(train_data, duration_col=duration_col, event_col=event_col,
                                     cols=[c for c in train_data.columns if c not in [event_col, duration_col, 'ID', 'group']], 
                                     pvalue_thres=0.05, verbose=True)
    train_data = train_data[['ID'] + sel_features + [event_col, duration_col, 'group']]
    data[['ID'] + sel_features + [event_col, duration_col, 'group']].to_csv(f'features/{task_type}features_unisel.csv', header=True, index=False)
train_data

from lifelines import CoxPHFitter

cph = CoxPHFitter(penalizer=0.5)
cph.fit(train_data[[c for c in train_data.columns if c not in ['ID', 'group']]], duration_col=duration_col, event_col=event_col)
cph.print_summary()

multi_info = cph.summary[['exp(coef)', 'exp(coef) lower 95%', 'exp(coef) upper 95%', 'p']].reset_index()
i = pd.merge(uni_info, multi_info, on='covariate', how='left', suffixes=['_Uni', '_Multi'])
multi_data = data[['ID'] + list(multi_info[multi_info['p'] < 0.05]['covariate'])]
multi_data.to_csv(f'features/{task_type}features_mulsel.csv', header=True, index=False)
i['p_Uni'] = i['p_Uni'].map(lambda x: '<0.05' if not pd.isna(x) and x < 0.05 else f"{x:.3f}")
i['p_Multi'] = i['p_Multi'].map(lambda x: '<0.05' if not pd.isna(x) and x < 0.05 else f"{x:.3f}")
i.columns = ['features_name', 'HR', 'lower 95%CI', 'upper 95%CI', 'pvalue', 'HR', 'lower 95%CI', 'upper 95%CI', 'pvalue']
i = i.applymap(lambda x: '' if pd.isna(x) else x)
i.to_csv('results/unimulti.csv', index=False)
i

import matplotlib.pyplot as plt

plt.figure(figsize=(10, train_data.shape[1]-4))
cph.plot(hazard_ratios=True)
plt.savefig('img/Clinic_feature_pvalue.svg')
plt.show()

from lifelines import CoxPHFitter
from lifelines.statistics import logrank_test
from lifelines import KaplanMeierFitter
from lifelines.plotting import add_at_risk_counts

thres = 1e-4
bst_split = {'train': 1.0, 'val':1, 'test': 1}
loc = {'train': 0.5, 'val':0.5, 'test': 0.5}

for subset, test_data in val_datasets.items():
    c_index = cph.score(test_data[[c for c in test_data.columns if c != 'ID']], scoring_method="concordance_index")
    y_pred = cph.predict_partial_hazard(test_data[[c for c in test_data.columns if c != 'ID']])
    cox_data = pd.concat([test_data, y_pred], axis=1)
    mean = cox_data.describe()[0]['50%']
    cox_data['HR'] = cox_data[0] > mean
    dem = (cox_data["HR"] == True)
    results = logrank_test(cox_data[duration_col][dem], cox_data[duration_col][~dem], 
                           event_observed_A=cox_data[event_col][dem], event_observed_B=cox_data[event_col][~dem])
    p_value = f"={results.p_value:.3f}" if results.p_value > thres else f'<{thres}'
    plt.title(f"Cohort {subset} C-index:{c_index:.3f}")
    plt.ylabel('Probability')
    if sum(dem):
        kmf_high = KaplanMeierFitter()
        kmf_high.fit(cox_data[duration_col][dem], event_observed=cox_data[event_col][dem], label="High Risk")
        kmf_high.plot_survival_function(color='r')
    if sum(~dem):
        kmf_low = KaplanMeierFitter()
        kmf_low.fit(cox_data[duration_col][~dem], event_observed=cox_data[event_col][~dem], label="Low Risk")
        kmf_low.plot_survival_function(color='g')
    plt.text(0.5, loc[subset] if subset in loc else 0.2, f"P{p_value}")
    plt.xlabel('Time(months)')
    add_at_risk_counts(kmf_high, kmf_low, rows_to_show=['At risk'])
    plt.savefig(f'img/{task_type}KM_{subset}.svg', bbox_inches='tight')
    plt.show()

import os
import numpy as np

def get_prediction(model: CoxPHFitter, data, ID=None, **kwargs):
    hr = model.predict_partial_hazard(data)
    expectation = model.predict_expectation(data)
    
    predictions = pd.concat([hr, expectation], axis=1)
    predictions.columns = ['HR', 'expectation']
    if ID is not None:
        predictions = pd.concat([ID, hr, expectation], axis=1)
        predictions.columns = ['ID', 'HR', 'expectation']
    else:
        predictions = pd.concat([hr, expectation], axis=1)
        predictions.columns = ['HR', 'expectation']
    return predictions

os.makedirs('results', exist_ok=True)
info = []
for subset, test_data in val_datasets.items():
    if subset in get_param_in_cwd('subsets'):
        results = get_prediction(cph, test_data, ID=test_data['ID'])
        results.to_csv(f'results/{task_type}cox_predictions_{subset}.csv', index=False)
        results['group'] = subset
        info.append(results)
        pd.merge(results, label_data[['ID', event_col, duration_col]], on='ID', how='inner').to_csv(f'features/{task_type}4xtile_{subset}.txt', 
                                                                                                    index=False, sep='\t')
info = pd.concat(info, axis=0)
info

### Combined model

from lifelines import CoxPHFitter
import pandas as pd
from onekey_algo.custom.components.comp1 import normalize_df
from sklearn.model_selection import train_test_split
from onekey_algo import get_param_in_cwd
from onekey_algo.custom.components.comp1 import fillna
from lifelines.utils import concordance_index
event_col = get_param_in_cwd('event_col')
group_info = 'group'
task_type = 'Combined_'
duration_col= get_param_in_cwd('duration_col')
data = pd.read_csv('features/Clinical_features_mulsel.csv')
label_data = pd.read_csv(get_param_in_cwd('label_file'))
label_data['ID'] = label_data['ID'].map(lambda x: f"{x}.nii.gz" if not (f"{x}".endswith('.nii.gz') or  f"{x}".endswith('.nii')) else x)

data = pd.merge(data, label_data[['ID', event_col, duration_col, 'group']], on='ID', how='inner')
data

def get_sig(prefix):
    results = [pd.read_csv(f'results/{prefix}_cox_predictions_{subset}.csv')for subset in get_param_in_cwd('subsets')]
    results = pd.concat(results, axis=0)
    return results

mn = {'Radiomics': f'Radiomics'}
for sig in mn:
    signature = get_sig(mn[sig])
    signature.columns = ['ID', sig, '-']
    data = pd.merge(data, signature, on='ID', how='inner')
data

import numpy as np
import onekey_algo.custom.components as okcomp
from collections import OrderedDict

group_info = 'group'
train_data = data[(data[group_info] == 'train')]

subsets = get_param_in_cwd('subsets')
val_datasets = OrderedDict()
for subset in subsets:
    val_data = data[data[group_info] == subset]
    val_datasets[subset] = val_data
    val_data.to_csv(f'features/{task_type}{subset}_cox.csv', index=False)

print('，'.join([f"{subset}样本数：{d_.shape}" for subset, d_ in val_datasets.items()]))

from lifelines import CoxPHFitter

cph = CoxPHFitter(penalizer=0.3)
cph.fit(train_data[[c for c in train_data.columns if c not in ['ID', 'group'] and not c.startswith('-')]], 
        duration_col=duration_col, event_col=event_col)
cph.print_summary()

cph.concordance_index_

import matplotlib.pyplot as plt

plt.figure(figsize=(10, train_data.shape[1]-4))
cph.plot(hazard_ratios=True)
plt.savefig(f'img/{task_type}_feature_pvalue.svg')
plt.show()

from onekey_algo.custom.components import nomogram
from onekey_algo import get_param_in_cwd

cox_data = train_data.round(decimals=2)
sur_name = get_param_in_cwd('time_settings')
nomogram.nomogram(train_data, duration=duration_col, result=event_col, 
                  columns=[c for c in train_data.columns if c not in ['ID', 'group', duration_col, event_col] and not c.startswith('-')],
                  survs=list(sur_name.values()), surv_names=list(sur_name.keys()), with_r=False,
                  width=7000, height=5000, save_name=f'img/nomogram.png',
                  x_range='0.01,0.25,0.5,0.75,0.95')

from lifelines import CoxPHFitter
from lifelines.statistics import logrank_test
from lifelines import KaplanMeierFitter
from lifelines.plotting import add_at_risk_counts

c_index_list = []
thres = 1e-4
bst_split = {'train': 1.0, 'val':1, 'test': 1}
loc = {'train': 0.5, 'val':0.5, 'test': 0.5}
for subset, test_data in val_datasets.items():
    c_index = cph.score(test_data[[c for c in test_data.columns if c != 'ID']], scoring_method="concordance_index")
    y_pred = cph.predict_partial_hazard(test_data[[c for c in test_data.columns if c != 'ID']])
    cox_data = pd.concat([test_data, y_pred], axis=1)
    mean = cox_data.describe()[0]['50%']
    cox_data['HR'] = cox_data[0] > mean
    dem = (cox_data["HR"] == True)
    results = logrank_test(cox_data[duration_col][dem], cox_data[duration_col][~dem], 
                           event_observed_A=cox_data[event_col][dem], event_observed_B=cox_data[event_col][~dem])
    p_value = f"={results.p_value:.3f}" if results.p_value > thres else f'<{thres}'
    plt.title(f"Cohort {subset} C-index:{c_index:.3f}")
    plt.ylabel('Probability')
    if sum(dem):
        kmf_high = KaplanMeierFitter()
        kmf_high.fit(cox_data[duration_col][dem], event_observed=cox_data[event_col][dem], label="High Risk")
        kmf_high.plot_survival_function(color='r')
    if sum(~dem):
        kmf_low = KaplanMeierFitter()
        kmf_low.fit(cox_data[duration_col][~dem], event_observed=cox_data[event_col][~dem], label="Low Risk")
        kmf_low.plot_survival_function(color='g')
    plt.text(0.5, loc[subset] if subset in loc else 0.2, f"P{p_value}")
    plt.xlabel('Time(months)')
    add_at_risk_counts(kmf_high, kmf_low, rows_to_show=['At risk'])
    plt.savefig(f'img/{task_type}KM_{subset}.svg', bbox_inches='tight')
    plt.show()

import os
import numpy as np

def get_prediction(model: CoxPHFitter, data, ID=None, **kwargs):
    hr = model.predict_partial_hazard(data)
    expectation = model.predict_expectation(data)
    
    predictions = pd.concat([hr, expectation], axis=1)
    predictions.columns = ['HR', 'expectation']
    if ID is not None:
        predictions = pd.concat([ID, hr, expectation], axis=1)
        predictions.columns = ['ID', 'HR', 'expectation']
    else:
        predictions = pd.concat([hr, expectation], axis=1)
        predictions.columns = ['HR', 'expectation']
    return predictions

os.makedirs('results', exist_ok=True)
info = []
for subset, test_data in val_datasets.items():
    if subset in get_param_in_cwd('subsets'):
        results = get_prediction(cph, test_data, ID=test_data['ID'])
        results.to_csv(f'results/{task_type}cox_predictions_{subset}.csv', index=False)
        results['group'] = subset
        info.append(results)
        pd.merge(results, label_data[['ID', event_col, duration_col]], on='ID', how='inner').to_csv(f'features/{task_type}4xtile_{subset}.txt', 
                                                                                                    index=False, sep='\t')
info = pd.concat(info, axis=0)
info