| import os | |
| import numpy as np | |
| import pandas as pd | |
| from sklearn.linear_model import Ridge, Lasso, LinearRegression | |
| from utils import seqs_to_onehot, read_fasta, load_rows_by_numbers | |
| from predictors.base_predictors import BaseRegressionPredictor, BaseGPPredictor | |
| class ESMPredictor(BaseRegressionPredictor): | |
| """ESM likelihood as features in regression.""" | |
| def __init__(self, dataset_name, rep_name, reg_coef=1e-8, path_prefix='', | |
| **kwargs): | |
| super(ESMPredictor, self).__init__(dataset_name, reg_coef, Ridge) | |
| seqs_path = path_prefix + os.path.join('data', dataset_name, 'seqs.fasta') | |
| seqs = read_fasta(seqs_path) | |
| id2seq = pd.Series(index=np.arange(len(seqs)), data=seqs, name='seq') | |
| esm_data_path = path_prefix + os.path.join('inference', dataset_name, | |
| 'esm', rep_name, 'pll.csv') | |
| ll = pd.read_csv(esm_data_path, index_col=0) | |
| ll['id'] = ll.index.to_series().apply( | |
| lambda x: int(x.replace('id_', ''))) | |
| ll = ll.join(id2seq, on='id', how='left') | |
| self.seq2score_dict = dict(zip(ll.seq, ll.pll)) | |
| def seq2score(self, seqs): | |
| scores = np.array([self.seq2score_dict.get(s, 0.0) for s in seqs]) | |
| return scores | |
| def seq2feat(self, seqs): | |
| return self.seq2score(seqs)[:, None] | |
| def predict_unsupervised(self, seqs): | |
| return self.seq2score(seqs) | |
| class GlobalESMPredictor(ESMPredictor): | |
| def __init__(self, dataset_name, **kwargs): | |
| super(GlobalESMPredictor, self).__init__( | |
| dataset_name, 'global', **kwargs) | |
| class EvotunedESMPredictor(ESMPredictor): | |
| def __init__(self, dataset_name, rep_name='uniref100', **kwargs): | |
| super(EvotunedESMPredictor, self).__init__( | |
| dataset_name, rep_name, **kwargs) | |
| class ESMRegressionPredictor(BaseRegressionPredictor): | |
| """Regression on ESM representation.""" | |
| def __init__(self, dataset_name, rep_name, reg_coef=1.0, **kwargs): | |
| super(ESMRegressionPredictor, self).__init__(dataset_name, reg_coef, | |
| Ridge, **kwargs) | |
| self.load_rep(dataset_name, rep_name) | |
| def load_rep(self, dataset_name, rep_name): | |
| self.rep_path = os.path.join('inference', dataset_name, 'esm', | |
| rep_name, 'rep.npy*') | |
| self.seq_path = os.path.join('data', dataset_name, 'seqs.fasta') | |
| self.seqs = read_fasta(self.seq_path) | |
| self.seq2id = dict(zip(self.seqs, range(len(self.seqs)))) | |
| def seq2feat(self, seqs): | |
| """Look up representation by sequences.""" | |
| ids = [self.seq2id[s] for s in seqs] | |
| return load_rows_by_numbers(self.rep_path, ids) | |
| class EvotunedESMRegressionPredictor(ESMRegressionPredictor): | |
| def __init__(self, dataset_name, **kwargs): | |
| super(EvotunedESMRegressionPredictor, self).__init__(dataset_name, | |
| 'uniref100', **kwargs) | |
| class GlobalESMRegressionPredictor(ESMRegressionPredictor): | |
| def __init__(self, dataset_name, **kwargs): | |
| super(GlobalESMRegressionPredictor, self).__init__(dataset_name, | |
| 'global', **kwargs) | |