--- a
+++ b/AE_run.py
@@ -0,0 +1,179 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# @Time    : 2021/8/7 14:43
+# @Author  : Li Xiao
+# @File    : AE_run.py
+import pandas as pd
+import numpy as np
+import argparse
+from tqdm import tqdm
+import autoencoder_model
+import torch
+import torch.utils.data as Data
+
+def setup_seed(seed):
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+
+def work(data, in_feas, lr=0.001, bs=32, epochs=100, device=torch.device('cpu'), a=0.4, b=0.3, c=0.3, mode=0, topn=100):
+    #name of sample
+    sample_name = data['Sample'].tolist()
+
+    #change data to a Tensor
+    X,Y = data.iloc[:,1:].values, np.zeros(data.shape[0])
+    TX, TY = torch.tensor(X, dtype=torch.float, device=device), torch.tensor(Y, dtype=torch.float, device=device)
+    #train a AE model
+    if mode == 0 or mode == 1:
+        print('Training model...')
+        Tensor_data = Data.TensorDataset(TX, TY)
+        train_loader = Data.DataLoader(Tensor_data, batch_size=bs, shuffle=True)
+
+        #initialize a model
+        mmae = autoencoder_model.MMAE(in_feas, latent_dim=100, a=a, b=b, c=c)
+        mmae.to(device)
+        mmae.train()
+        mmae.train_MMAE(train_loader, learning_rate=lr, device=device, epochs=epochs)
+        mmae.eval()       #before save and test, fix the variables
+        torch.save(mmae, 'model/AE/MMAE_model.pkl')
+
+    #load saved model, used for reducing dimensions
+    if mode == 0 or mode == 2:
+        print('Get the latent layer output...')
+        mmae = torch.load('model/AE/MMAE_model.pkl')
+        omics_1 = TX[:, :in_feas[0]]
+        omics_2 = TX[:, in_feas[0]:in_feas[0]+in_feas[1]]
+        omics_3 = TX[:, in_feas[0]+in_feas[1]:in_feas[0]+in_feas[1]+in_feas[2]]
+        latent_data, decoded_omics_1, decoded_omics_2, decoded_omics_3 = mmae.forward(omics_1, omics_2, omics_3)
+        latent_df = pd.DataFrame(latent_data.detach().cpu().numpy())
+        latent_df.insert(0, 'Sample', sample_name)
+        #save the integrated data(dim=100)
+        latent_df.to_csv('result/latent_data.csv', header=True, index=False)
+
+    print('Extract features...')
+    extract_features(data, in_feas, epochs, topn)
+    return
+
+def extract_features(data, in_feas, epochs, topn=100):
+    # extract features
+    #get each omics data
+    data_omics_1 = data.iloc[:, 1: 1+in_feas[0]]
+    data_omics_2 = data.iloc[:, 1+in_feas[0]: 1+in_feas[0]+in_feas[1]]
+    data_omics_3 = data.iloc[:, 1+in_feas[0]+in_feas[1]: 1+in_feas[0]+in_feas[1]+in_feas[2]]
+
+    #get all features of each omics data
+    feas_omics_1 = data_omics_1.columns.tolist()
+    feas_omics_2 = data_omics_2.columns.tolist()
+    feas_omics_3 = data_omics_3.columns.tolist()
+
+    #calculate the standard deviation of each feature
+    std_omics_1 = data_omics_1.std(axis=0)
+    std_omics_2 = data_omics_2.std(axis=0)
+    std_omics_3 = data_omics_3.std(axis=0)
+
+    #record top N features every 10 epochs
+    topn_omics_1 = pd.DataFrame()
+    topn_omics_2 = pd.DataFrame()
+    topn_omics_3 = pd.DataFrame()
+
+    #used for feature extraction, epoch_ls = [10,20,...], if epochs % 10 != 0, add the last epoch
+    epoch_ls = list(range(10, epochs+10,10))
+    if epochs %10 != 0:
+        epoch_ls.append(epochs)
+    for epoch in tqdm(epoch_ls):
+        #load model
+        mmae = torch.load('model/AE/model_{}.pkl'.format(epoch))
+        #get model variables
+        model_dict = mmae.state_dict()
+
+        #get the absolute value of weights, the shape of matrix is (n_features, latent_layer_dim)
+        weight_omics1 = np.abs(model_dict['encoder_omics_1.0.weight'].detach().cpu().numpy().T)
+        weight_omics2 = np.abs(model_dict['encoder_omics_2.0.weight'].detach().cpu().numpy().T)
+        weight_omics3 = np.abs(model_dict['encoder_omics_3.0.weight'].detach().cpu().numpy().T)
+
+        weight_omics1_df = pd.DataFrame(weight_omics1, index=feas_omics_1)
+        weight_omics2_df = pd.DataFrame(weight_omics2, index=feas_omics_2)
+        weight_omics3_df = pd.DataFrame(weight_omics3, index=feas_omics_3)
+
+        #calculate the weight sum of each feature --> sum of each row
+        weight_omics1_df['Weight_sum'] = weight_omics1_df.apply(lambda x:x.sum(), axis=1)
+        weight_omics2_df['Weight_sum'] = weight_omics2_df.apply(lambda x:x.sum(), axis=1)
+        weight_omics3_df['Weight_sum'] = weight_omics3_df.apply(lambda x:x.sum(), axis=1)
+        weight_omics1_df['Std'] = std_omics_1
+        weight_omics2_df['Std'] = std_omics_2
+        weight_omics3_df['Std'] = std_omics_3
+
+        #importance = Weight * Std
+        weight_omics1_df['Importance'] = weight_omics1_df['Weight_sum']*weight_omics1_df['Std']
+        weight_omics2_df['Importance'] = weight_omics2_df['Weight_sum']*weight_omics2_df['Std']
+        weight_omics3_df['Importance'] = weight_omics3_df['Weight_sum']*weight_omics3_df['Std']
+
+        #select top N features
+        fea_omics_1_top = weight_omics1_df.nlargest(topn, 'Importance').index.tolist()
+        fea_omics_2_top = weight_omics2_df.nlargest(topn, 'Importance').index.tolist()
+        fea_omics_3_top = weight_omics3_df.nlargest(topn, 'Importance').index.tolist()
+
+        #save top N features in a dataframe
+        col_name = 'epoch_'+str(epoch)
+        topn_omics_1[col_name] = fea_omics_1_top
+        topn_omics_2[col_name] = fea_omics_2_top
+        topn_omics_3[col_name] = fea_omics_3_top
+
+    #all of top N features
+    topn_omics_1.to_csv('result/topn_omics_1.csv', header=True, index=False)
+    topn_omics_2.to_csv('result/topn_omics_2.csv', header=True, index=False)
+    topn_omics_3.to_csv('result/topn_omics_3.csv', header=True, index=False)
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--mode', '-m', type=int, choices=[0,1,2], default=0,
+                        help='Mode 0: train&intagrate, Mode 1: just train, Mode 2: just intagrate, default: 0.')
+    parser.add_argument('--seed', '-s', type=int, default=0, help='Random seed, default=0.')
+    parser.add_argument('--path1', '-p1', type=str, required=True, help='The first omics file name.')
+    parser.add_argument('--path2', '-p2', type=str, required=True, help='The second omics file name.')
+    parser.add_argument('--path3', '-p3', type=str, required=True, help='The third omics file name.')
+    parser.add_argument('--batchsize', '-bs', type=int, default=32, help='Training batchszie, default: 32.')
+    parser.add_argument('--learningrate', '-lr', type=float, default=0.001, help='Learning rate, default: 0.001.')
+    parser.add_argument('--epoch', '-e', type=int, default=100, help='Training epochs, default: 100.')
+    parser.add_argument('--latent', '-l', type=int, default=100, help='The latent layer dim, default: 100.')
+    parser.add_argument('--device', '-d', type=str, choices=['cpu', 'gpu'], default='cpu', help='Training on cpu or gpu, default: cpu.')
+    parser.add_argument('--a', '-a', type=float, default=0.6, help='[0,1], float, weight for the first omics data')
+    parser.add_argument('--b', '-b', type=float, default=0.1, help='[0,1], float, weight for the second omics data.')
+    parser.add_argument('--c', '-c', type=float, default=0.3, help='[0,1], float, weight for the third omics data.')
+    parser.add_argument('--topn', '-n', type=int, default=100, help='Extract top N features every 10 epochs, default: 100.')
+    args = parser.parse_args()
+
+    #read data
+    omics_data1 = pd.read_csv(args.path1, header=0, index_col=None)
+    omics_data2 = pd.read_csv(args.path2, header=0, index_col=None)
+    omics_data3 = pd.read_csv(args.path3, header=0, index_col=None)
+
+    #Check whether GPUs are available
+    device = torch.device('cpu')
+    if args.device == 'gpu':
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    #set random seed
+    setup_seed(args.seed)
+
+    if args.a + args.b + args.c != 1.0:
+        print('The sum of weights must be 1.')
+        exit(1)
+
+    #dims of each omics data
+    in_feas = [omics_data1.shape[1] - 1, omics_data2.shape[1] - 1, omics_data3.shape[1] - 1]
+    omics_data1.rename(columns={omics_data1.columns.tolist()[0]: 'Sample'}, inplace=True)
+    omics_data2.rename(columns={omics_data2.columns.tolist()[0]: 'Sample'}, inplace=True)
+    omics_data3.rename(columns={omics_data3.columns.tolist()[0]: 'Sample'}, inplace=True)
+
+    omics_data1.sort_values(by='Sample', ascending=True, inplace=True)
+    omics_data2.sort_values(by='Sample', ascending=True, inplace=True)
+    omics_data3.sort_values(by='Sample', ascending=True, inplace=True)
+
+    #merge the multi-omics data, calculate on common samples
+    Merge_data = pd.merge(omics_data1, omics_data2, on='Sample', how='inner')
+    Merge_data = pd.merge(Merge_data, omics_data3, on='Sample', how='inner')
+    Merge_data.sort_values(by='Sample', ascending=True, inplace=True)
+
+    #train model, reduce dimensions and extract features
+    work(Merge_data, in_feas, lr=args.learningrate, bs=args.batchsize, epochs=args.epoch, device=device, a=args.a, b=args.b, c=args.c, mode=args.mode, topn=args.topn)
+    print('Success! Results can be seen in result file')