Average mutational effects for an assay (eg, antibody)¶
This notebook averages selections that measure escape from neutralization in an assay. Below in the code, "antibody" is used as a generic term for the agent that neutralizes the virus.
Import Python modules.
We use polyclonal
for the averaging and plotting:
import copy
import pickle
import altair as alt
import pandas as pd
import polyclonal
import polyclonal.utils
_ = alt.data_transformers.disable_max_rows()
This notebook is parameterized by papermill
.
The next cell is tagged as parameters
to get the passed parameters.
# this cell is tagged parameters for `papermill` parameterization
assay = None
site_numbering_map_csv = None
prob_escape_mean_csvs = None
pickles = None
assay_config = None
avg_pickle_file = None
effect_csv = None
icXX_csv = None
effect_html = None
icXX_html = None
params = None
# Parameters
params = {
"icXX": 90,
"escape_plot_kwargs": {
"addtl_slider_stats": {"times_seen": 3, "escape_std": 1.5},
"addtl_tooltip_stats": ["sequential_site"],
"heatmap_max_at_least": 2,
"heatmap_min_at_least": -2,
"init_floor_at_zero": False,
"init_site_statistic": "sum",
"site_zoom_bar_color_col": "region",
"avg_type": "median",
"addtl_slider_stats_as_max": ["escape_std"],
"per_model_tooltip": True,
},
"plot_hide_stats": {
"functional effect": {
"csv": "results/func_effects/averages/293T_high_ACE2_entry_func_effects.csv",
"csv_col": "effect",
"init": -1.5,
"min_filters": {"times_seen": 3},
},
"ACE2 binding": {
"csv": "results/ACE2_binding/averages/monomeric_ACE2_mut_effect.csv",
"csv_col": "ACE2 binding_mean",
"init": -1.5,
"min_filters": {"times_seen": 3},
},
},
"legend": "Use the site zoom bar at the top to zoom in on specific sites. The line plot shows a summary statistic indicating of the mutation effects at each site. The heat map shows effects for individual mutations, with parental amino-acid identities indicated by x and gray indicating non-measured mutations. Mouse over points for details.\n\nOptions at the bottom of the plot let you modify the display, such as by selecting how many different variants a mutation must be seen in to be shown (*minimum times_seen*), how many different experimental selections the mutation was measured in (*minimum n_models*), what site summary statistic to show, etc.\n\nYou can filter by the functional effects of mutations. Mutations removed by this filter are shown as dark gray squares in the heat map to distinguish unmeasured mutations from ones measured to be deleterious.\n\nThe *minimum max* slider is useful to select the sites where mutations have the largest effect.\n",
"title": "Sera 500C escape in medium ACE2 cells",
"selections": [
"Lib1-230805-sera-500C_mediumACE2",
"Lib2-230805-sera-500C_mediumACE2",
],
}
prob_escape_mean_csvs = [
"results/antibody_escape/by_selection/Lib1-230805-sera-500C_mediumACE2_prob_escape_mean.csv",
"results/antibody_escape/by_selection/Lib2-230805-sera-500C_mediumACE2_prob_escape_mean.csv",
]
pickles = [
"results/antibody_escape/by_selection/Lib1-230805-sera-500C_mediumACE2_polyclonal_model.pickle",
"results/antibody_escape/by_selection/Lib2-230805-sera-500C_mediumACE2_polyclonal_model.pickle",
]
assay_config = {
"title": "Antibody/serum escape",
"selections": "antibody_selections",
"averages": "avg_antibody_escape",
"prob_escape_scale": {"type": "symlog", "constant": 0.04},
"scale_stat": 1,
"stat_name": "escape",
}
assay = "antibody_escape"
site_numbering_map_csv = "data/site_numbering_map.csv"
avg_pickle_file = (
"results/antibody_escape/averages/sera_500C_mediumACE2_polyclonal_model.pickle"
)
effect_csv = "results/antibody_escape/averages/sera_500C_mediumACE2_mut_effect.csv"
icXX_csv = "results/antibody_escape/averages/sera_500C_mediumACE2_mut_icXX.csv"
effect_html = (
"results/antibody_escape/averages/sera_500C_mediumACE2_mut_effect_nolegend.html"
)
icXX_html = (
"results/antibody_escape/averages/sera_500C_mediumACE2_mut_icXX_nolegend.html"
)
print(f"Analyzing results for {assay=}")
Analyzing results for assay='antibody_escape'
Read the input data and parameters:
site_numbering_map = pd.read_csv(site_numbering_map_csv).rename(
columns={"reference_site": "site"}
)
assert site_numbering_map[["site", "sequential_site"]].notnull().all().all()
assert len(params["selections"]) == len(set(params["selections"]))
# read Polyclonal models into a data frame that can be passed to PolyclonalAverage
models_df = pd.DataFrame(
[(s, pickle.load(open(f, "rb"))) for s, f in zip(params["selections"], pickles)],
columns=["selection", "model"],
)
# read prob_escape means all into one data frame
prob_escape_means = pd.concat(
[
pd.read_csv(f).assign(selection=s)
for s, f in zip(params["selections"], prob_escape_mean_csvs)
],
ignore_index=True,
)
# get the plot kwargs
escape_plot_kwargs = params["escape_plot_kwargs"]
Neutralization at concentrations used for each selection¶
For each selection going into the average, plot the average fraction neutralization (probability of escape) of variants with different numbers of mutations, both for the censored values used to fit the models and the uncensored values. Note the concentrations not used in the model fits are shown fainter and in a different shape:
mean_prob_escape_chart = (
alt.Chart(prob_escape_means)
.encode(
x=alt.X(
"concentration",
**(
{"title": assay_config["concentration_title"]}
if "concentration_title" in assay_config
else {}
),
scale=alt.Scale(
**(
assay_config["concentration_scale"]
if "concentration_scale" in assay_config
else {"type": "log"}
)
),
),
y=alt.Y(
"probability escape",
scale=alt.Scale(**assay_config["prob_escape_scale"]),
),
column=alt.Column(
"censored",
title=None,
header=alt.Header(labelFontWeight="bold", labelFontSize=10),
),
row=alt.Row(
"selection",
title=None,
header=alt.Header(labelFontWeight="bold", labelFontSize=10),
),
color=alt.Color("n_substitutions"),
tooltip=[
alt.Tooltip(c, format=".3g") if c == "probability escape" else c
for c in prob_escape_means.columns
],
shape=alt.Shape("use_in_fit", scale=alt.Scale(domain=[True, False])),
opacity=alt.Opacity(
"use_in_fit", scale=alt.Scale(domain=[True, False], range=[0.9, 0.3])
),
)
.mark_line(point=True, size=0.75, opacity=0.8)
.properties(width=230, height=145)
.configure_axis(grid=False)
.configure_point(size=50)
)
mean_prob_escape_chart
Average mutation effects¶
First build a PolyclonalAverage
:
avg_model = polyclonal.PolyclonalAverage(models_df)
print(f"Saving the average model to {avg_pickle_file}")
with open(avg_pickle_file, "wb") as f:
pickle.dump(avg_model, f)
Saving the average model to results/antibody_escape/averages/sera_500C_mediumACE2_polyclonal_model.pickle
Neutralization curves against unmutated protein (which reflect the wildtype activities, Hill coefficients, and non-neutralizable fractions):
avg_model.curves_plot()
Get the mutation escape values, and apply any filters also being applied to the final average plot:
mut_escape_df = avg_model.mut_escape_df_replicates
# apply `times_seen` filter
try:
times_seen = escape_plot_kwargs["addtl_slider_stats"]["times_seen"]
except KeyError:
times_seen = 1
n = len(mut_escape_df)
mut_escape_df = mut_escape_df.query("times_seen >= @times_seen")
print(
f"Filtering for `times_seen` >= {times_seen} removes "
f"{n - len(mut_escape_df)} of {n} mutations"
)
# apply any other filters
if "plot_hide_stats" in params:
for stat, stat_d in params["plot_hide_stats"].items():
n = len(mut_escape_df)
min_stat = stat_d["init"]
df_to_merge = pd.read_csv(stat_d["csv"])
if not avg_model.sequential_integer_sites:
df_to_merge["site"] = df_to_merge["site"].astype(str)
if "min_filters" in stat_d:
for col, min_col in stat_d["min_filters"].items():
df_to_merge = df_to_merge.query(f"{col} >= @min_col")
col = stat_d["csv_col"]
assert col not in mut_escape_df.columns, f"{col=}, {mut_escape_df.columns=}"
mut_escape_df = mut_escape_df.merge(
df_to_merge[["site", "mutant", col]],
how="left",
validate="many_to_one",
)
mut_escape_df = mut_escape_df[
(mut_escape_df[col] >= min_stat) | mut_escape_df[col].isnull()
].drop(columns=col)
print(
f"Filtering for {stat} >= {min_stat} removes "
f"{n - len(mut_escape_df)} of {n} mutations"
)
Filtering for `times_seen` >= 3 removes 3088 of 13781 mutations Filtering for functional effect >= -1.5 removes 244 of 10693 mutations Filtering for ACE2 binding >= -1.5 removes 80 of 10449 mutations
Correlation of escape across different selections:
mut_escape_corr = polyclonal.utils.tidy_to_corr(
mut_escape_df,
sample_col="selection",
label_col="mutation",
value_col="escape",
group_cols="epitope",
method="pearson",
).rename(columns={"correlation": "r"})
polyclonal.plot.corr_heatmap(
corr_df=mut_escape_corr,
corr_col="r",
sample_cols="selection",
group_col=None if mut_escape_corr["epitope"].nunique() == 1 else "epitope",
)
Site line plots for the site values for each individual selection (model) in the average. This makes it easier to tell if one selection is an outlier before we plot the full averages below, and how correlated the selections are. Note the plot is interactive: you can mouseover points and change the site metric shown.
site_stats = ["sum", "mean", "max", "min"]
addtl_site_cols = [
c for c in site_numbering_map.columns if c != "site" and c.endswith("site")
]
site_escape_df = (
pd.concat(
[
mut_escape_df.assign(
floor_at_zero=floor_at_zero,
escape=lambda x: (x["escape"] * assay_config["scale_stat"]).clip(
lower=0 if floor_at_zero else None,
),
)
for floor_at_zero in [True, False]
]
)
.groupby(["epitope", "selection", "floor_at_zero", "site"], as_index=False)
.aggregate(**{stat: pd.NamedAgg("escape", stat) for stat in site_stats})
.melt(
id_vars=["epitope", "selection", "floor_at_zero", "site"],
value_vars=site_stats,
var_name="site statistic",
value_name="site_val",
)
.merge(
site_numbering_map[["site", *addtl_site_cols]].assign(
site=lambda x: x["site"].astype(type(avg_model.sites[0]))
),
validate="many_to_one",
)
)
try:
init_site_stat = params["escape_plot_kwargs"]["init_site_statistic"]
except KeyError:
init_site_stat = "sum"
site_statistic_selection = alt.selection_point(
fields=["site statistic"],
bind=alt.binding_radio(
name="site statistic",
options=site_stats,
),
value=init_site_stat,
)
try:
init_floor_at_zero = params["escape_plot_kwargs"]["init_floor_at_zero"]
except KeyError:
init_floor_at_zero = True
floor_selection = alt.selection_point(
fields=["floor_at_zero"],
bind=alt.binding_radio(
name="floor at zero",
options=[True, False],
),
value=init_floor_at_zero,
)
site_selection = alt.selection_point(fields=["site"], on="mouseover", empty=False)
per_selection_site_escape_chart_base = (
alt.Chart(site_escape_df)
.encode(
x=alt.X(
"site",
sort=alt.SortField("sequential_site"),
axis=alt.Axis(labelOverlap=True),
scale=alt.Scale(nice=False, zero=False),
),
y=alt.Y("site_val", title=assay),
color="epitope",
tooltip=[
"site",
alt.Tooltip("site_val", format=".2f", title=assay),
],
)
.properties(width=750, height=85)
.add_params(site_statistic_selection, floor_selection, site_selection)
.transform_filter(site_statistic_selection)
.transform_filter(floor_selection)
)
per_selection_site_escape_chart_lines = per_selection_site_escape_chart_base.mark_line(
size=0.75
)
per_selection_site_escape_chart_points = per_selection_site_escape_chart_base.encode(
size=alt.condition(site_selection, alt.value(75), alt.value(30)),
strokeWidth=alt.condition(site_selection, alt.value(2), alt.value(0)),
).mark_circle(filled=True, stroke="orange")
per_selection_escape_chart = (
(per_selection_site_escape_chart_lines + per_selection_site_escape_chart_points)
.facet(
facet=alt.Facet(
"selection",
title=None,
header=alt.Header(labelPadding=0),
),
columns=1,
spacing=5,
)
.configure_axis(grid=False)
)
per_selection_escape_chart
Plot standard deviation in escape versus escape. This is useful to look at if you have imposed a filter on the standard deviation on the escape to filter variants with large variation.
escape_plot_kwargs = params["escape_plot_kwargs"]
try:
avg_type = escape_plot_kwargs["avg_type"]
except KeyError:
avg_type = avg_model.default_avg_to_plot
std_col = f"{assay} std"
avg_col = f"{assay} {avg_type}"
std_df = avg_model.mut_escape_df_w_model_values.rename(
columns={f"escape_{avg_type}": avg_col, "escape_std": std_col}
)
# replace dots in column names, which cause problems in tooltips, with underscore
std_df = std_df.rename(columns={c: c.replace(".", "_") for c in std_df.columns})
try:
init_times_seen = escape_plot_kwargs["addtl_slider_stats"]["times_seen"]
except KeyError:
init_times_seen = 1
times_seen_slider = alt.param(
value=init_times_seen,
bind=alt.binding_range(
name="minimum times_seen",
min=1,
max=std_df[std_col].quantile(0.9),
),
)
try:
init_std = escape_plot_kwargs["addtl_slider_stats"]["escape_std"]
except KeyError:
try:
init_std = escape_plot_kwargs["addtl_slider_stats"][f"{assay}_std"]
except KeyError:
init_std = std_df[std_col].max()
std_slider = alt.param(
value=init_std,
bind=alt.binding_range(
name=f"maximum {std_col}",
min=0,
max=std_df[std_col].max(),
),
)
std_chart = (
alt.Chart(std_df)
.add_params(std_slider, times_seen_slider)
.transform_filter(alt.datum["times_seen"] >= times_seen_slider)
.transform_calculate(above_max_std=alt.datum[std_col] > std_slider)
.encode(
alt.X(std_col),
alt.Y(avg_col),
alt.Color(
"above_max_std:N",
scale=alt.Scale(domain=[False, True]),
legend=alt.Legend(orient="bottom", symbolOpacity=1),
),
tooltip=[
alt.Tooltip(c, format=".3g") if std_df[c].dtype == float else c
for c in std_df.columns
],
)
.mark_circle(opacity=0.2, strokeOpacity=1, stroke="black", strokeWidth=0.5)
.resolve_scale(x="independent", y="independent")
.properties(width=250, height=250)
.configure_axis(grid=False)
)
std_chart
Plot and save the mutation effects for the average model:
# build up arguments used to format plot
plot_hide_stats = params["plot_hide_stats"]
if "addtl_slider_stats" not in escape_plot_kwargs:
escape_plot_kwargs["addtl_slider_stats"] = {}
if "addtl_slider_stats_hide_not_filter" not in escape_plot_kwargs:
escape_plot_kwargs["addtl_slider_stats_hide_not_filter"] = []
if "escape_std" in escape_plot_kwargs["addtl_slider_stats"]:
if "addtl_slider_stats_as_max" not in escape_plot_kwargs:
raise ValueError(
"You specified `escape_std` in `addtl_slider_stats` but did not add it to "
"`addtl_slider_stats_as_max`. If you really do not want `escape_std` in "
"`addtl_slider_stats_as_max`, then specify that list without it."
)
escape_plot_kwargs["df_to_merge"] = []
for stat, stat_d in plot_hide_stats.items():
escape_plot_kwargs["addtl_slider_stats"][stat] = stat_d["init"]
escape_plot_kwargs["addtl_slider_stats_hide_not_filter"].append(stat)
merge_df = pd.read_csv(stat_d["csv"]).rename(columns={stat_d["csv_col"]: stat})
if "min_filters" in stat_d:
for col, col_min in stat_d["min_filters"].items():
if col not in merge_df.columns:
raise ValueError(f"{stat=} CSV lacks {col=}\n{merge_df.columns=}")
merge_df = merge_df[merge_df[col] >= col_min]
escape_plot_kwargs["df_to_merge"].append(merge_df[["site", "mutant", stat]])
escape_plot_kwargs["df_to_merge"].append(
site_numbering_map[["site", *addtl_site_cols, "region"]]
)
if "addtl_tooltip_stats" not in escape_plot_kwargs:
escape_plot_kwargs["addtl_tooltip_stats"] = []
for c in addtl_site_cols:
if c not in escape_plot_kwargs["addtl_tooltip_stats"]:
escape_plot_kwargs["addtl_tooltip_stats"].append(c)
if "init_n_models" not in escape_plot_kwargs:
escape_plot_kwargs["init_n_models"] = len(avg_model.models) // 2 + 1
escape_plot_kwargs["scale_stat_col"] = assay_config["scale_stat"]
if assay_config["stat_name"] != "escape":
escape_plot_kwargs["rename_stat_col"] = assay_config["stat_name"]
print(f"Writing mutation values to {effect_csv}")
effects_for_csv = avg_model.mut_escape_df_w_model_values.drop(
columns="escape_min_magnitude"
)
for c in ["escape_mean", "escape_median"] + params["selections"]:
effects_for_csv[c] = assay_config["scale_stat"] * effects_for_csv[c]
effects_for_csv["escape_std"] = (
abs(assay_config["scale_stat"]) * effects_for_csv["escape_std"]
)
effects_for_csv = effects_for_csv.rename(
columns={
f"escape_{stat}": f"{assay_config['stat_name']}_{stat}"
for stat in ["mean", "median", "std"]
}
)
effects_for_csv.to_csv(effect_csv, index=False, float_format="%.4g")
escape_chart = avg_model.mut_escape_plot(**copy.deepcopy(escape_plot_kwargs))
print(f"Writing chart to {effect_html}")
escape_chart.save(effect_html)
display(escape_chart)
Writing mutation values to results/antibody_escape/averages/sera_500C_mediumACE2_mut_effect.csv Writing chart to results/antibody_escape/averages/sera_500C_mediumACE2_mut_effect_nolegend.html
Plot and save the predicted change in ICXX induced by each mutation:
icXX = params["icXX"]
print(f"Getting predicted changes in IC{icXX}")
icXX_col = f"IC{icXX}"
log_fold_change_icXX_col = f"log2 fold change {icXX_col}"
print(f"Writing changes in ICXX to {icXX_csv}")
icXX_for_csv = avg_model.mut_icXX_df_w_model_values(
x=icXX / 100.0,
icXX_col=icXX_col,
log_fold_change_icXX_col=log_fold_change_icXX_col,
).drop(columns=[log_fold_change_icXX_col + s for s in [" min_magnitude", " std"]])
for c in [log_fold_change_icXX_col + s for s in [" mean", " median"]] + params[
"selections"
]:
icXX_for_csv[c] = assay_config["scale_stat"] * icXX_for_csv[c]
icXX_for_csv.to_csv(icXX_csv, index=False, float_format="%.4g")
if "rename_stat_col" in escape_plot_kwargs:
del escape_plot_kwargs["rename_stat_col"]
if "escape_std" in escape_plot_kwargs["addtl_slider_stats"]:
del escape_plot_kwargs["addtl_slider_stats"]["escape_std"]
if "addtl_slider_stats_as_max" in escape_plot_kwargs:
if "escape_std" in escape_plot_kwargs["addtl_slider_stats_as_max"]:
escape_plot_kwargs["addtl_slider_stats_as_max"].remove("escape_std")
icXX_chart = avg_model.mut_icXX_plot(
x=icXX / 100.0,
icXX_col=icXX_col,
log_fold_change_icXX_col=log_fold_change_icXX_col,
**copy.deepcopy(escape_plot_kwargs),
)
print(f"Writing ICXX chart to {icXX_html}")
icXX_chart.save(icXX_html)
display(icXX_chart)
Getting predicted changes in IC90 Writing changes in ICXX to results/antibody_escape/averages/sera_500C_mediumACE2_mut_icXX.csv Writing ICXX chart to results/antibody_escape/averages/sera_500C_mediumACE2_mut_icXX_nolegend.html