from math import ceil, exp import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import pandas as pd from re import match import seaborn as sns import time from threading import Thread from model import ModelFactory class Data: "Container for input and output data" AA = "ACDEFGHIKLMNPQRSTVWY" def parse_seq(self, src: str): "Parse input sequence (plain or FASTA)" lines = src.strip().splitlines() if lines and lines[0].startswith('>'): lines = lines[1:] self.seq = ''.join(lines).upper().replace('\n', '').replace(' ', '') if not self.seq: raise RuntimeError("Sequence is empty") if not all(x in self.model.alphabet for x in self.seq): raise RuntimeError(f"Unsupported characters in sequence: {''.join(x for x in self.seq if x not in self.model.alphabet)}") 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]): self.mode = 'MUT' for resi, (src_c, trg_c) in enumerate(zip(self.seq, self.trg[0]), 1): if src_c != trg_c: self.sub.append(f"{src_c}{resi}{trg_c}") self.resi.append(resi) elif all(match(r'\d+', x) for x in self.trg): self.mode = 'SMS' trh = [] for resi in map(int, self.trg): if resi < 1 or resi > len(self.seq): raise RuntimeError(f"Position {resi} out of range (sequence length: {len(self.seq)})") src_c = self.seq[resi-1] for trg_c in self.AA.replace(src_c, ''): self.sub.append(f"{src_c}{resi}{trg_c}") trh.append(self.seq[:resi-1]+trg_c+self.seq[resi:]) self.resi.append(resi) self.trg = trh elif all(match(r'[A-Z]\d+[A-Z]', x) for x in self.trg): self.mode = 'MUT' self.sub = list(self.trg) trh = [] for x in self.trg: idx = int(x[1:-1]) self.resi.append(idx) trh.append(self.seq[:idx-1]+x[-1]+self.seq[idx:]) for s, *resi_str, _ in self.trg: if self.seq[int(''.join(resi_str))-1] != s: raise RuntimeError(f"Unrecognised input substitution: {self.seq[int(''.join(resi_str))]}{int(''.join(resi_str))} /= {s}{int(''.join(resi_str))}") self.trg = trh else: self.mode = 'DMS' self.trg = [] for resi, src_c in enumerate(self.seq, 1): for trg_c in self.AA.replace(src_c, ''): self.sub.append(f"{src_c}{resi}{trg_c}") self.trg.append(self.seq[:resi-1]+trg_c+self.seq[resi:]) 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" self.model_name = model_name self.model = ModelFactory(model_name) self.parse_seq(src) self.parse_sub(trg) self.scoring_strategy = scoring_strategy self.progress = None self.out = pd.DataFrame(self.sub, columns=['0', self.model_name]) self.out_img_path = f"{out_file}.png" self.out_table = None self.out_csv = f"{out_file}.csv" def parse_output(self) -> None: "Format output data for visualisation" if self.mode == "DMS": self._render_dms() self.out.to_csv(self.out_csv, float_format='%.2f') elif self.mode == "SMS": self._sort_sms() self._style_and_save() elif self.mode == "MUT": self.out = self.out.sort_values(self.model_name, ascending=False) self._style_and_save() else: raise RuntimeError(f"Unrecognised mode {self.mode}") def _style_and_save(self): "Apply table styling and write CSV" self.out_table = (self.out.style .format({col: "{:.2f}" for col in self.out.select_dtypes(float).columns}) .background_gradient(cmap="RdYlGn", vmax=8, vmin=-8).hide(axis=0).hide(axis=1)) self.out.to_csv(self.out_csv, float_format='%.2f', index=False, header=False) def _sort_sms(self): "Sort SMS output by residue then score, top 19 per position, reshaped to columns" 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 _render_dms(self): "Build DMS heatmap from scored mutations" # Group by residue, keep top 19 (all alternatives) grouped = (self.out.assign(resi=self.resi_cycle()) .groupby("resi").head(19)) # Reshape: rows = amino acids + wild-type header, cols = positions blocks = [] for x in range(grouped.shape[0]//19): chunk = grouped.iloc[19*x:19*(x+1)] wt_label = chunk.iloc[0, 0][:-1]+chunk.iloc[0, 0][0] # e.g. "RA" header = pd.Series([wt_label, 0, 0], index=chunk.columns) block = pd.concat([header.to_frame().T, chunk], axis=0, ignore_index=True) block = block.sort_values(['0']).drop(["resi", '0'], axis=1).astype(float) block = block.set_axis(list(self.AA)) blocks.append(block) self.out = (pd.concat(blocks, axis=1) .set_axis([f'{a}{i}' for i, a in enumerate(self.seq, 1)], axis="columns")) self.out /= self.out.abs().max().max() # Layout ncols = min([d for d in range(1, self.out.shape[1]+1) if self.out.shape[1] % d == 0 and 30 <= d <= 60] or [60] , key=lambda x: abs(x-60)) nrows = ceil(self.out.shape[1]/ncols) while self.out.shape[1]/ncols < nrows and ncols > 45 and ncols*nrows >= self.out.shape[1]: ncols -= 1 ncols += 1 self._plot_heatmap(ncols, nrows) def resi_cycle(self): "Repeat self.resi to match out row count" return (self.resi * (len(self.out)//len(self.resi) + 1))[:len(self.out)] def _plot_heatmap(self, ncols, nrows): "Render DMS heatmap to PNG" kw = dict(cmap="RdBu", cbar=False, square=True, xticklabels=1, yticklabels=1 , center=0, fmt='s', annot_kws={"size": "xx-large"}) annotate = lambda df: df.map(lambda x: ' ' if x != 0 else '\u00b7') if nrows < 2: fig = plt.figure(figsize=(12, 6)) sns.heatmap(self.out, annot=annotate(self.out), **kw) fig.tight_layout() else: fig, axes = plt.subplots(nrows=nrows, figsize=(12, 6*nrows)) for i in range(nrows): tmp = self.out.iloc[:, i*ncols:(i+1)*ncols] sns.heatmap(tmp, ax=axes[i], annot=annotate(tmp), **kw) axes[i].set_yticklabels(axes[i].get_yticklabels(), rotation=0) axes[i].set_xticklabels(axes[i].get_xticklabels(), rotation=90) fig.tight_layout() plt.savefig(self.out_img_path, format="png", dpi=150) plt.close(fig) def calculate(self, progress): "run model and parse output" self.progress = progress self.model.run_model(self) done = [False] def _render(): self.parse_output() done[0] = True Thread(target=_render, daemon=True).start() t = 0 while not done[0]: time.sleep(0.33) t += 1 progress(min(1 - exp(-0.1 * t), 0.99), desc="Rendering") return self @property def csv(self): "return output CSV path" return self.out_csv @property def image(self): "return PNG path (DMS) or Styler object (SMS/MUT)" if self.out_table is not None: return self.out_table return self.out_img_path