| 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() |
|
|
| |
| 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" |
| |
| grouped = (self.out.assign(resi=self.resi_cycle()) |
| .groupby("resi").head(19)) |
|
|
| |
| 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] |
| 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() |
|
|
| |
| 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 |
|
|
