Spaces:
Running
Running
| from math import ceil | |
| import matplotlib.pyplot as plt | |
| import pandas as pd | |
| from re import match | |
| import seaborn as sns | |
| from model import Model | |
| class Data: | |
| """Container for input and output data""" | |
| # Initialise empty model as static class member for efficiency | |
| model = Model() | |
| def parse_seq(self, src: str): | |
| """Parse input sequence""" | |
| self.seq = src.strip().upper().replace('\n', '') | |
| if not all(x in self.model.alphabet for x in self.seq): | |
| raise RuntimeError("Unrecognised characters in sequence") | |
| def parse_sub(self, trg: str): | |
| """Parse input substitutions""" | |
| self.mode = None | |
| self.sub = list() | |
| self.trg = trg.strip().upper().split() | |
| self.resi = list() | |
| # Identify running mode | |
| if len(self.trg) == 1 and len(self.trg[0]) == len(self.seq) and match(r'^\w+$', self.trg[0]): | |
| # If single string of same length as sequence, seq vs seq mode | |
| self.mode = 'MUT' | |
| for resi, (src, trg) in enumerate(zip(self.seq, self.trg[0]), 1): | |
| if src != trg: | |
| self.sub.append(f"{src}{resi}{trg}") | |
| self.resi.append(resi) | |
| else: | |
| if all(match(r'\d+', x) for x in self.trg): | |
| # If all strings are numbers, deep mutational scanning mode | |
| self.mode = 'DMS' | |
| for resi in map(int, self.trg): | |
| src = self.seq[resi-1] | |
| for trg in "ACDEFGHIKLMNPQRSTVWY".replace(src, ''): | |
| self.sub.append(f"{src}{resi}{trg}") | |
| self.resi.append(resi) | |
| elif all(match(r'[A-Z]\d+[A-Z]', x) for x in self.trg): | |
| # If all strings are of the form X#Y, single substitution mode | |
| self.mode = 'MUT' | |
| self.sub = self.trg | |
| self.resi = [int(x[1:-1]) for x in self.trg] | |
| for s, *resi, _ in self.trg: | |
| if self.seq[int(''.join(resi))-1] != s: | |
| raise RuntimeError(f"Unrecognised input substitution {self.seq[int(''.join(resi))]}{int(''.join(resi))} /= {s}{int(''.join(resi))}") | |
| else: | |
| self.mode = 'TMS' | |
| for resi, src in enumerate(self.seq, 1): | |
| for trg in "ACDEFGHIKLMNPQRSTVWY".replace(src, ''): | |
| self.sub.append(f"{src}{resi}{trg}") | |
| self.resi.append(resi) | |
| self.sub = pd.DataFrame(self.sub, columns=['0']) | |
| def __init__(self, src:str, trg:str, model_name:str='facebook/esm2_t33_650M_UR50D', scoring_strategy:str='masked-marginals', out_file='out'): | |
| "initialise data" | |
| # if model has changed, load new model | |
| if self.model.model_name != model_name: | |
| self.model_name = model_name | |
| self.model = Model(model_name) | |
| self.parse_seq(src) | |
| self.offset = 0 | |
| self.parse_sub(trg) | |
| self.scoring_strategy = scoring_strategy | |
| self.token_probs = None | |
| self.out = pd.DataFrame(self.sub, columns=['0', self.model_name]) | |
| self.out_img = f'{out_file}.png' | |
| self.out_csv = f'{out_file}.csv' | |
| def parse_output(self) -> None: | |
| "format output data for visualisation" | |
| if self.mode == 'TMS': | |
| self.process_tms_mode() | |
| self.out.to_csv(self.out_csv, float_format='%.2f') | |
| else: | |
| if self.mode == 'DMS': | |
| self.sort_by_residue_and_score() | |
| elif self.mode == 'MUT': | |
| self.sort_by_score() | |
| else: | |
| raise RuntimeError(f"Unrecognised mode {self.mode}") | |
| self.out.columns = [str(i) for i in range(self.out.shape[1])] | |
| self.out_img = (self.out.style | |
| .format(lambda x: f'{x:.2f}' if isinstance(x, float) else x) | |
| .hide(axis=0) | |
| .background_gradient(cmap="RdYlGn", vmax=8, vmin=-8)) | |
| self.out.to_csv(self.out_csv, float_format='%.2f', index=False, header=False) | |
| def sort_by_score(self): | |
| self.out = self.out.sort_values(self.model_name, ascending=False) | |
| def sort_by_residue_and_score(self): | |
| self.out = (self.out.assign(resi=self.out['0'].str.extract(r'(\d+)', expand=False).astype(int)) | |
| .sort_values(['resi', self.model_name], ascending=[True,False]) | |
| .groupby(['resi']) | |
| .head(19) | |
| .drop(['resi'], axis=1)) | |
| self.out = pd.concat([self.out.iloc[19*x:19*(x+1)].reset_index(drop=True) for x in range(self.out.shape[0]//19)] | |
| , axis=1).set_axis(range(self.out.shape[0]//19*2), axis='columns') | |
| def process_tms_mode(self): | |
| self.out = self.assign_resi_and_group() | |
| self.out = self.concat_and_set_axis() | |
| self.out /= self.out.abs().max().max() | |
| divs = self.calculate_divs() | |
| ncols = min(divs, key=lambda x: abs(x-60)) | |
| nrows = ceil(self.out.shape[1]/ncols) | |
| ncols = self.adjust_ncols(ncols, nrows) | |
| self.plot_heatmap(ncols, nrows) | |
| def assign_resi_and_group(self): | |
| return (self.out.assign(resi=self.out['0'].str.extract(r'(\d+)', expand=False).astype(int)) | |
| .groupby(['resi']) | |
| .head(19)) | |
| def concat_and_set_axis(self): | |
| return (pd.concat([(self.out.iloc[19*x:19*(x+1)] | |
| .pipe(self.create_dataframe) | |
| .sort_values(['0'], ascending=[True]) | |
| .drop(['resi', '0'], axis=1) | |
| .set_axis(['A', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'K', 'L', | |
| 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'V', 'W', 'Y']) | |
| .astype(float) | |
| ) for x in range(self.out.shape[0]//19)] | |
| , axis=1) | |
| .set_axis([f'{a}{i}' for i, a in enumerate(self.seq, 1)], axis='columns')) | |
| def create_dataframe(self, df): | |
| return pd.concat([pd.Series([df.iloc[0, 0][:-1]+df.iloc[0, 0][0], 0, 0], index=df.columns).to_frame().T, df], axis=0, ignore_index=True) | |
| def calculate_divs(self): | |
| return [x for x in range(1, self.out.shape[1]+1) if self.out.shape[1] % x == 0 and 30 <= x and x <= 60] or [60] | |
| def adjust_ncols(self, ncols, nrows): | |
| while self.out.shape[1]/ncols < nrows and ncols > 45 and ncols*nrows >= self.out.shape[1]: | |
| ncols -= 1 | |
| return ncols + 1 | |
| def plot_heatmap(self, ncols, nrows): | |
| if nrows < 2: | |
| self.plot_single_heatmap() | |
| else: | |
| self.plot_multiple_heatmaps(ncols, nrows) | |
| plt.savefig(self.out_img, format='png', dpi=300) | |
| def plot_single_heatmap(self): | |
| fig = plt.figure(figsize=(12, 6)) | |
| sns.heatmap(self.out | |
| , cmap='RdBu' | |
| , cbar=False | |
| , square=True | |
| , xticklabels=1 | |
| , yticklabels=1 | |
| , center=0 | |
| , annot=self.out.map(lambda x: ' ' if x != 0 else '·') | |
| , fmt='s' | |
| , annot_kws={'size': 'xx-large'}) | |
| fig.tight_layout() | |
| def plot_multiple_heatmaps(self, ncols, nrows): | |
| fig, ax = plt.subplots(nrows=nrows, figsize=(12, 6*nrows)) | |
| for i in range(nrows): | |
| tmp = self.out.iloc[:,i*ncols:(i+1)*ncols] | |
| label = tmp.map(lambda x: ' ' if x != 0 else '·') | |
| sns.heatmap(tmp | |
| , ax=ax[i] | |
| , cmap='RdBu' | |
| , cbar=False | |
| , square=True | |
| , xticklabels=1 | |
| , yticklabels=1 | |
| , center=0 | |
| , annot=label | |
| , fmt='s' | |
| , annot_kws={'size': 'xx-large'}) | |
| ax[i].set_yticklabels(ax[i].get_yticklabels(), rotation=0) | |
| ax[i].set_xticklabels(ax[i].get_xticklabels(), rotation=90) | |
| fig.tight_layout() | |
| def calculate(self): | |
| "run model and parse output" | |
| self.model.run_model(self) | |
| self.parse_output() | |
| return self | |
| def csv(self): | |
| "return output data" | |
| return self.out_csv | |
| def image(self): | |
| "return output data" | |
| return self.out_img | |