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