#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@author: Zhi Huang
"""
import sys, os
from pathlib import Path
project_folder = Path("..").resolve()
model_folder = project_folder / "model"
sys.path.append(model_folder.absolute().as_posix())
import SALMON
import pandas as pd
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import transforms
from torch.autograd import Variable
from collections import Counter
import pandas as pd
import matplotlib.pyplot as plt
import math
import random
from imblearn.over_sampling import RandomOverSampler
from lifelines.statistics import logrank_test
import json
import tables
import logging
import csv
import numpy as np
import optunity
import pickle
import time
from sklearn.model_selection import KFold
from sklearn import preprocessing
import matplotlib.pyplot as plt
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument('--num_epochs', type=int, default=100, help="Number of epochs to train for. Default: 100")
parser.add_argument('--measure_while_training', action='store_true', default=False, help='disables measure while training (make program faster)')
parser.add_argument('--batch_size', type=int, default=256, help="Number of batches to train/test for. Default: 256")
parser.add_argument('--dataset', type=int, default=7)
parser.add_argument('--nocuda', action='store_true', default=False, help='disables CUDA training')
parser.add_argument('--verbose', default=1, type=int)
parser.add_argument('--results_dir', default=(project_folder / "experiments/Results").absolute().as_posix(), help="results dir")
return parser.parse_args()
if __name__=='__main__':
torch.cuda.empty_cache()
args = parse_args()
plt.ioff()
# model file
num_epochs = args.num_epochs
batch_size = args.batch_size
learning_rate_range = 10**np.arange(-4,-1,0.3)
cuda = True
verbose = 0
measure_while_training = True
dropout_rate = 0
lambda_1 = 1e-5 # L1
if args.dataset == 1:
dataset_subset = "1_RNAseq"
elif args.dataset == 2:
dataset_subset = "2_miRNAseq"
elif args.dataset == 3:
dataset_subset = "3_RNAseq+miRNAseq"
elif args.dataset == 4:
dataset_subset = "4_RNAseq+miRNAseq+cnv+tmb"
elif args.dataset == 5:
dataset_subset = "5_RNAseq+miRNAseq+clinical"
elif args.dataset == 6:
dataset_subset = "6_cnv+tmb+clinical"
elif args.dataset == 7:
dataset_subset = "7_RNAseq+miRNAseq+cnv+tmb+clinical"
datasets_5folds = pickle.load( open( (project_folder / "data/BRCA_583_new/datasets_5folds.pickle").absolute().as_posix(), "rb" ) )
for i in range(5):
print("5 fold CV -- %d/5" % (i+1))
# dataset
TIMESTRING = time.strftime("%Y%m%d-%H.%M.%S", time.localtime())
results_dir_dataset = args.results_dir + '/' + dataset_subset + '/run_' + TIMESTRING + '_fold_' + str(i+1)
if not os.path.exists(results_dir_dataset):
os.makedirs(results_dir_dataset)
logging.basicConfig(filename=results_dir_dataset+'/mainlog.log',level=logging.DEBUG)
# print("Arguments:",args)
# logging.info("Arguments: %s" % args)
datasets = datasets_5folds[str(i+1)]
len_of_RNAseq = 57
len_of_miRNAseq = 12
len_of_cnv = 1
len_of_tmb = 1
len_of_clinical = 3
length_of_data = {}
length_of_data['mRNAseq'] = len_of_RNAseq
length_of_data['miRNAseq'] = len_of_miRNAseq
length_of_data['CNB'] = len_of_cnv
length_of_data['TMB'] = len_of_tmb
length_of_data['clinical'] = len_of_clinical
if args.dataset == 1:
#### RNAseq Only
datasets['train']['x'] = datasets['train']['x'][:, 0:len_of_RNAseq]
datasets['test']['x'] = datasets['test']['x'][:, 0:len_of_RNAseq]
elif args.dataset == 2:
#### miRNAseq Only
datasets['train']['x'] = datasets['train']['x'][:, len_of_RNAseq:(len_of_RNAseq + len_of_miRNAseq)]
datasets['test']['x'] = datasets['test']['x'][:, len_of_RNAseq:(len_of_RNAseq + len_of_miRNAseq)]
elif args.dataset == 3:
#### RNAseq + miRNAseq
datasets['train']['x'] = datasets['train']['x'][:, 0:(len_of_RNAseq + len_of_miRNAseq)]
datasets['test']['x'] = datasets['test']['x'][:, 0:(len_of_RNAseq + len_of_miRNAseq)]
elif args.dataset == 4:
#### RNAseq + miRNAseq + CNB + all TMB
datasets['train']['x'] = datasets['train']['x'][:, 0:(len_of_RNAseq + len_of_miRNAseq + len_of_cnv + len_of_tmb)]
datasets['test']['x'] = datasets['test']['x'][:, 0:(len_of_RNAseq + len_of_miRNAseq + len_of_cnv + len_of_tmb)]
elif args.dataset == 5:
#### RNAseq + miRNAseq + clinical (age+ER+PR)
datasets['train']['x'] = np.concatenate((datasets['train']['x'][:, 0:(len_of_RNAseq + len_of_miRNAseq)], \
datasets['train']['x'][:, (len_of_RNAseq + len_of_miRNAseq + len_of_cnv + len_of_tmb):]),1)
datasets['test']['x'] = np.concatenate((datasets['test']['x'][:, 0:(len_of_RNAseq + len_of_miRNAseq)], \
datasets['test']['x'][:, (len_of_RNAseq + len_of_miRNAseq + len_of_cnv + len_of_tmb):]),1)
elif args.dataset == 6:
#### CNB + all TMB + clinical (age+ER+PR)
datasets['train']['x'] = datasets['train']['x'][:, (len_of_RNAseq + len_of_miRNAseq):]
datasets['test']['x'] = datasets['test']['x'][:, (len_of_RNAseq + len_of_miRNAseq):]
elif args.dataset == 7:
#### RNAseq + miRNAseq + CNB + all TMB + clinical (age+ER+PR)
datasets['train']['x'] = datasets['train']['x']
datasets['test']['x'] = datasets['test']['x']
# =============================================================================
# # Finding optimal learning rate w.r.t. concordance index
# =============================================================================
ci_list = []
for j, lr in enumerate(learning_rate_range):
print("[%d/%d] current lr: %.4E" %((j+1), len(learning_rate_range), lr))
logging.info("[%d/%d] current lr: %.4E" %((j+1), len(learning_rate_range), lr))
model, loss_nn_all, pvalue_all, c_index_all, c_index_list, acc_train_all, code_output = \
SALMON.train(datasets, num_epochs, batch_size, lr, dropout_rate,\
lambda_1, length_of_data, cuda, measure_while_training, verbose)
epochs_list = range(num_epochs)
plt.figure(figsize=(8,4))
plt.plot(epochs_list, c_index_list['train'], "b--",linewidth=1)
plt.plot(epochs_list, c_index_list['test'], "g-",linewidth=1)
plt.legend(['train', 'test'])
plt.xlabel("epochs")
plt.ylabel("Concordance index")
plt.savefig(results_dir_dataset + "/convergence_%02d_lr=%.2E.png" % (j, lr),dpi=300)
code_test, loss_nn_sum, acc_test, pvalue_pred, c_index_pred, lbl_pred_all, OS_event_test, OS_test = \
SALMON.test(model, datasets, 'test', length_of_data, batch_size, cuda, verbose)
ci_list.append(c_index_pred)
print("current concordance index: ", c_index_pred,"\n")
logging.info("current concordance index: %.10f\n" % c_index_pred)
optimal_lr = learning_rate_range[np.argmax(ci_list)]
print("Optimal learning rate: %.4E, optimal c-index: %.10f" % (optimal_lr, np.max(ci_list)))
logging.info("Optimal learning rate: %.4E, optimal c-index: %.10f" % (optimal_lr, np.max(ci_list)))
# =============================================================================
# # Training
# =============================================================================
model, loss_nn_all, pvalue_all, c_index_all, c_index_list, acc_train_all, code_output = \
SALMON.train(datasets, num_epochs, batch_size, optimal_lr, dropout_rate,\
lambda_1, length_of_data, cuda, measure_while_training, verbose)
code_train, loss_nn_sum, acc_train, pvalue_pred, c_index_pred, lbl_pred_all_train, OS_event_train, OS_train = \
SALMON.test(model, datasets, 'train', length_of_data, batch_size, cuda, verbose)
print("[Final] Apply model to training set: c-index: %.10f, p-value: %.10e" % (c_index_pred, pvalue_pred))
logging.info("[Final] Apply model to training set: c-index: %.10f, p-value: %.10e" % (c_index_pred, pvalue_pred))
code_test, loss_nn_sum, acc_test, pvalue_pred, c_index_pred, lbl_pred_all_test, OS_event_test, OS_test = \
SALMON.test(model, datasets, 'test', length_of_data, batch_size, cuda, verbose)
print("[Final] Apply model to testing set: c-index: %.10f, p-value: %.10e" % (c_index_pred, pvalue_pred))
logging.info("[Final] Apply model to testing set: c-index: %.10f, p-value: %.10e" % (c_index_pred, pvalue_pred))
with open(results_dir_dataset + '/model.pickle', 'wb') as handle:
pickle.dump(model, handle, protocol=pickle.HIGHEST_PROTOCOL)
with open(results_dir_dataset + '/c_index_list_by_epochs.pickle', 'wb') as handle:
pickle.dump(c_index_list, handle, protocol=pickle.HIGHEST_PROTOCOL)
with open(results_dir_dataset + '/hazard_ratios_lbl_pred_all_train.pickle', 'wb') as handle:
pickle.dump(lbl_pred_all_train, handle, protocol=pickle.HIGHEST_PROTOCOL)
with open(results_dir_dataset + '/OS_event_train.pickle', 'wb') as handle:
pickle.dump(OS_event_train, handle, protocol=pickle.HIGHEST_PROTOCOL)
with open(results_dir_dataset + '/OS_train.pickle', 'wb') as handle:
pickle.dump(OS_train, handle, protocol=pickle.HIGHEST_PROTOCOL)
with open(results_dir_dataset + '/code_train.pickle', 'wb') as handle:
pickle.dump(code_train, handle, protocol=pickle.HIGHEST_PROTOCOL)
with open(results_dir_dataset + '/hazard_ratios_lbl_pred_all_test.pickle', 'wb') as handle:
pickle.dump(lbl_pred_all_test, handle, protocol=pickle.HIGHEST_PROTOCOL)
with open(results_dir_dataset + '/OS_event_test.pickle', 'wb') as handle:
pickle.dump(OS_event_test, handle, protocol=pickle.HIGHEST_PROTOCOL)
with open(results_dir_dataset + '/OS_test.pickle', 'wb') as handle:
pickle.dump(OS_test, handle, protocol=pickle.HIGHEST_PROTOCOL)
with open(results_dir_dataset + '/code_test.pickle', 'wb') as handle:
pickle.dump(code_test, handle, protocol=pickle.HIGHEST_PROTOCOL)
epochs_list = range(num_epochs)
plt.figure(figsize=(8,4))
plt.plot(epochs_list, c_index_list['train'], "b--",linewidth=1)
plt.plot(epochs_list, c_index_list['test'], "g-",linewidth=1)
plt.legend(['train', 'test'])
plt.xlabel("epochs")
plt.ylabel("Concordance index")
plt.savefig(results_dir_dataset + "/convergence.png",dpi=300)
Datasets
Models
