File size: 6,687 Bytes
bae5726
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
"""
Traditional machine learning models
"""


import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import joblib

from sklearn.linear_model import Ridge
from sklearn.svm import SVR
from sklearn.ensemble import RandomForestRegressor
from sklearn.neighbors import KNeighborsRegressor
import xgboost

from tqdm import tqdm
import os, subprocess, sys
import importlib
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT))
from model.ephod.training import trainutils
importlib.reload(trainutils);




WEIGHT_TYPES = ['bin_inv', 'bin_inv_sqrt', 'LDS_inv', 'LDS_inv_sqrt', 'LDS_extreme']




def get_target_data(path=None):
    '''Return pHopt data as a dictionary of arrays'''

    if path is None: # Else specify to correct path
        path = PROJECT_ROOT / "conf" / "training_data" / "pHopt_data.csv"
    df = pd.read_csv(path, index_col=0)
    data = {}
    for key in ['Training', 'Validation', 'Testing']:
        dfsel = df[df['Split']==key]
        keydata = {'accession': dfsel['Accession'].values, 
                    'sequence': dfsel['Sequence'].values,
                    'y': dfsel['pHopt'].values}
        keydata['weights'] = {name : trainutils.get_sample_weights(keydata['y'], name) for name in WEIGHT_TYPES}
        data[key] = keydata
    
    return data



def get_features_data(data, path=None):
    '''Fetch features data and add to data dictionary'''
    
    if path is None:
        path = PROJECT_ROOT / "conf" / "training_data" / "aac.csv"
    X = pd.read_csv(path, index_col=0)
    for key in data.keys():
        print(key)
        data[key]['X'] = X.loc[data[key]['accession'], :].values    
    
    return data
    


    
class Trainer():
    
    def __init__(self, data, modelname='ridge'):

        self.data = data
        self.modelname = modelname
        self.param_space = self.get_param_space()
        self.sampled_params = trainutils.sample_hyperparameters(self.param_space, n=200)

    
    def get_model(self):
        
        if self.modelname == 'ridge':
            model = Ridge()
        elif self.modelname == 'svr':
            model = SVR()
        elif self.modelname == 'rforest':
            model = RandomForestRegressor(n_jobs=-1)
        elif self.modelname == 'knr':
            model = KNeighborsRegressor(n_jobs=-1)
        elif self.modelname == 'xgboost':
            model = xgboost.XGBRegressor(n_jobs=-1)
        else:
            raise ValueError(f'Unrecognized model: {self.modelname}')
        
        return model

    
    def get_param_space(self):
        '''Define space for hyperparameter search'''
        
        if self.modelname == 'ridge':
            param_space = {'alpha': 10.0 ** np.arange(-8,8+1)}
        elif self.modelname == 'svr':
            param_space = {'kernel': ['poly', 'rbf'], 'gamma': ['scale', 'auto'], 'C': 10. ** np.arange(-5, 5)}
        elif self.modelname == 'rforest':
            param_space = {'n_estimators': [10, 20, 50, 100, 200, 500, 1000], 'criterion': ['mse', 'mae'],  
                           'max_features': [0.25, 0.5, 0.75, None], 'max_samples': [0.25, 0.5, 0.75, None],
                           'max_depth': [5,10,None]}
        elif self.modelname == 'knr':
            param_space = {'n_neighbors':[1, 2, 5, 10, 20, 50, 100], 'weights':['distance', 'uniform']}
        elif self.modelname == 'xgboost': 
            param_space = {'learning_rate': [0.01, 0.5, 0.1, 0.2, None], 'reg_alpha': [0, 0.01, 0.1, 1, 10, 100, None],
                           'reg_lambda': [0.01, 0.1, 1, 10, 100, None], 'max_depth': [3, 5, 7, 9, 11, None],
                           'max_delta_step': [0, 1, 5, 10, None], 'min_child_weight': [1, 3, 5, 7, 10], 
                           'n_estimators':[20, 50, 100, 200, 500, None]}
        param_space['weight_type'] = WEIGHT_TYPES # Add reweighting technique as hyperparameter
        
        return param_space


    def normalize_data(self):
    
        Xtrain, Xval, Xtest = [data[key]['X'] for key in ['Training', 'Validation', 'Testing']]
        ytrain, yval, ytest = [data[key]['y'] for key in ['Training', 'Validation', 'Testing']]
        means, stds = np.mean(Xtrain, axis=0), np.std(Xtrain, axis=0)
        Xtrain, Xval, Xtest = [(item - means) / (stds + 1e-8) for item in [Xtrain, Xval, Xtest]] 
        
        return ((Xtrain, Xval, Xtest), (ytrain, yval, ytest))

    
    def train(self):
        '''Train multiple models with all sampled hyperparameters'''
        
        (Xtrain, Xval, Xtest), (ytrain, yval, ytest) = self.normalize_data()
        best_rmse = 1e10 # Initialize to a large number
        
        for params in tqdm(self.sampled_params):
            
            # Sample weights for training (always use bin_inv in validation/testing)
            trainweights = self.data['Training']['weights'][params['weight_type']]
            valweights, testweights = [self.data[key]['weights']['bin_inv'] for key in ('Validation', 'Testing')]
            # Train
            modelparams = {k:v for k,v in params.items() if k!='weight_type'}
            model = self.get_model()
            model = model.set_params(**modelparams)
            model = model.fit(Xtrain, ytrain, sample_weight=trainweights)
            # Validate
            yvalpred = model.predict(Xval)
            valperf = trainutils.performance(yval, yvalpred, valweights)
            # Test
            ytestpred = model.predict(Xtest)
            testperf = trainutils.performance(ytest, ytestpred, testweights)
            # Store results
            rmse = valperf['rmse']
            if rmse < best_rmse:
                best_rmse = rmse
                best_results = params.copy()
                best_results.update({'valperf': valperf, 'testperf': testperf, 'yvalpred': yvalpred,
                                     'ytestpred': ytestpred, 'model': model})
                self.best_results = best_results
        
        return best_results




if __name__ == '__main__': 

    modelnames = ['ridge', 'svr', 'knr', 'rforest', 'xgboost']
    repnames = ['aac', 'ifeaturepca', 'ifeaturerfe', 'onehot', 'esm1v', 'esm1b', 'prott5', 'progen2', 
                'tranception', 'carp']
    data = get_target_data() # First, download from zenodo
    for modelname in modelnames:
        for repname in repnames:
            # Data for training
            data = get_features_data(data, PROJECT_ROOT / "conf" / "training_data" / f"{repname}.csv")
            # Train model
            trainer = Trainer(data, modelname)
            results = trainer.train()
            joblib.dump(results, PROJECT_ROOT / "weight" / f"{repname}-{modelname}.pkl")