|
import pandas as pd |
|
import numpy as np |
|
import seaborn as sns |
|
import matplotlib.pyplot as plt |
|
from functools import partial |
|
from utils.metric_utils import spearman, topk_mean, hit_rate, aucroc |
|
|
|
|
|
def get_stratified_metrics(df, model_name, max_n_mut, metric_fn): |
|
strat_metrics = np.zeros(max_n_mut+1) |
|
for i in range(0, max_n_mut+1): |
|
|
|
tmp = df |
|
if i > 0: |
|
tmp = df[df.n_mut == i] |
|
strat_metrics[i] = metric_fn(tmp[model_name], tmp.log_fitness) |
|
return strat_metrics |
|
|
|
|
|
def plot_stratified_metrics(ax, df, models, max_n_mut, metric_fn, vmin, vmax): |
|
strat_matrix = np.zeros((len(models), 1+max_n_mut)) |
|
xticklabels=['All'] + list(range(1, max_n_mut+1)) |
|
for i, m in enumerate(models): |
|
strat_matrix[i] = get_stratified_metrics(df, m, max_n_mut, metric_fn) |
|
sns.heatmap(strat_matrix, yticklabels=models, xticklabels=xticklabels, |
|
vmin=vmin, vmax=vmax, ax=ax, cmap='viridis') |
|
ax.set_xlabel('# Mutations') |
|
ax.vlines([1], *ax.get_ylim(), colors='black') |
|
|
|
|
|
def plot_auc_and_corr(df, models, functional_threshold, wt_log_fitness, |
|
max_n_mut=5, vmin=[None, None, None], vmax=[None, None, None], topk=96): |
|
fig, axes = plt.subplots(1, 3, figsize=(14, 4), sharex=True, sharey=True) |
|
|
|
ax = axes[0] |
|
fn = partial(aucroc, y_cutoff=functional_threshold) |
|
plot_stratified_metrics(ax, df, models, max_n_mut, fn, vmin[0], vmax[0]) |
|
ax.set_title(f'Functional vs Non-Functional AUC-ROC') |
|
|
|
ax = axes[1] |
|
fn = partial(aucroc, y_cutoff=wt_log_fitness) |
|
plot_stratified_metrics(ax, df[df.log_fitness >= functional_threshold], |
|
models, max_n_mut, fn, vmin[1], vmax[1]) |
|
ax.set_title(f'Functional, <WT vs >=WT AUC-ROC') |
|
|
|
ax = axes[2] |
|
plot_stratified_metrics(ax, df[df.log_fitness >= functional_threshold], |
|
models, max_n_mut, spearman, vmin[2], vmax[2]) |
|
ax.set_title('Rank Correlation (Functional)') |
|
|
|
fig.suptitle('Model performance, stratified by # mutations') |
|
plt.subplots_adjust(wspace=0.1, top=0.85) |
|
plt.show() |
|
|
