Fit polyclonal model¶
Here we fit polyclonal models to the data.
First, import Python modules:
[1]:
import pickle
import altair as alt
import pandas as pd
import polyclonal
import yaml
[2]:
# allow more rows for Altair
_ = alt.data_transformers.disable_max_rows()
Read input data¶
Get parameterized variable from papermill
[3]:
# papermill parameters cell (tagged as `parameters`)
prob_escape_csv = None
n_threads = None
pickle_file = None
antibody = None
[4]:
# Parameters
prob_escape_csv = (
"results/prob_escape/Lib-2_2022-03-14_thaw-3_2130-1-0114-112_1_prob_escape.csv"
)
pickle_file = "results/polyclonal_fits/Lib-2_2022-03-14_thaw-3_2130-1-0114-112_1.pickle"
n_threads = 2
Read the probabilities of escape, and filter for those with sufficient no-antibody counts:
[5]:
print(f"\nReading probabilities of escape from {prob_escape_csv}")
prob_escape = pd.read_csv(
prob_escape_csv, keep_default_na=False, na_values="nan"
).query("`no-antibody_count` >= no_antibody_count_threshold")
assert prob_escape.notnull().all().all()
Reading probabilities of escape from results/prob_escape/Lib-2_2022-03-14_thaw-3_2130-1-0114-112_1_prob_escape.csv
Read the rest of the configuration and input data:
[6]:
# get information from config
with open("config.yaml") as f:
config = yaml.safe_load(f)
antibody = prob_escape["antibody"].unique()
assert len(antibody) == 1, antibody
antibody = antibody[0]
# get site numbering map and the reference sites in order
site_numbering_map = pd.read_csv(config["site_numbering_map"])
reference_sites = site_numbering_map.sort_values("sequential_site")[
"reference_site"
].tolist()
# get the polyclonal configuration for this antibody
with open(config["polyclonal_config"]) as f:
polyclonal_config = yaml.safe_load(f)
if antibody not in polyclonal_config:
raise ValueError(f"`polyclonal_config` lacks configuration for {antibody=}")
antibody_config = polyclonal_config[antibody]
# print names of variables and settings
print(f"{antibody=}")
print(f"{n_threads=}")
print(f"{pickle_file=}")
print(f"{antibody_config=}")
antibody='2130-1-0114-112'
n_threads=2
pickle_file='results/polyclonal_fits/Lib-2_2022-03-14_thaw-3_2130-1-0114-112_1.pickle'
antibody_config={'min_epitope_activity_to_include': 0.2, 'plot_kwargs': {'addtl_slider_stats': {'times_seen': 3, 'functional effect': -1.38}, 'slider_binding_range_kwargs': {'n_models': {'step': 1}, 'times_seen': {'step': 1, 'min': 1, 'max': 25}}, 'heatmap_max_at_least': 2, 'heatmap_min_at_least': -2}, 'max_epitopes': 1, 'fit_kwargs': {'reg_escape_weight': 0.1, 'reg_spread_weight': 0.25, 'reg_activity_weight': 1.0}}
Read spatial distances if relevant:
[7]:
if ("spatial_distances" in config) and (config["spatial_distances"] is not None):
print(f"Reading spatial distances from {config['spatial_distances']}")
spatial_distances = pd.read_csv(config["spatial_distances"])
print(f"Read spatial distances for {len(spatial_distances)} residue pairs")
else:
print("No spatial distances")
spatial_distances = None
No spatial distances
Some summary statistics¶
Note that these statistics are only for the variants that passed upstream filtering in the pipeline.
Number of variants per concentration:
[8]:
display(
prob_escape.groupby("antibody_concentration").aggregate(
n_variants=pd.NamedAgg("barcode", "nunique")
)
)
n_variants | |
---|---|
antibody_concentration | |
0.1599 | 90191 |
0.7995 | 90191 |
3.9974 | 90191 |
Plot mean probability of escape across all variants with the indicated number of mutations. Note that this plot weights each variant the same in the means regardless of how many barcode counts it has. We plot means for both censored (set to between 0 and 1) and uncensored probabilities of escape. Also, note it uses a symlog scale for the y-axis. Mouseover points for values:
[9]:
max_aa_subs = 4 # group if >= this many substitutions
mean_prob_escape = (
prob_escape.assign(
n_subs=lambda x: (
x["aa_substitutions_reference"]
.str.split()
.map(len)
.clip(upper=max_aa_subs)
.map(lambda n: str(n) if n < max_aa_subs else f">{max_aa_subs - 1}")
)
)
.groupby(["antibody_concentration", "n_subs"], as_index=False)
.aggregate({"prob_escape": "mean", "prob_escape_uncensored": "mean"})
.rename(
columns={
"prob_escape": "censored to [0, 1]",
"prob_escape_uncensored": "not censored",
}
)
.melt(
id_vars=["antibody_concentration", "n_subs"],
var_name="censored",
value_name="probability escape",
)
)
mean_prob_escape_chart = (
alt.Chart(mean_prob_escape)
.encode(
x=alt.X("antibody_concentration"),
y=alt.Y(
"probability escape",
scale=alt.Scale(type="symlog", constant=0.05),
),
column=alt.Column("censored", title=None),
color=alt.Color("n_subs", title="n substitutions"),
tooltip=[
alt.Tooltip(c, format=".3g") if mean_prob_escape[c].dtype == float else c
for c in mean_prob_escape.columns
],
)
.mark_line(point=True, size=0.5)
.properties(width=200, height=125)
.configure_axis(grid=False)
)
mean_prob_escape_chart
[9]:
Fit polyclonal
model¶
First, get the fitting related keyword arguments from the configuration passed by snakemake
:
[10]:
max_epitopes = antibody_config["max_epitopes"]
print(f"{max_epitopes=}")
fit_kwargs = antibody_config["fit_kwargs"]
print(f"{fit_kwargs=}")
min_epitope_activity_to_include = antibody_config["min_epitope_activity_to_include"]
print(f"{min_epitope_activity_to_include=}")
max_epitopes=1
fit_kwargs={'reg_escape_weight': 0.1, 'reg_spread_weight': 0.25, 'reg_activity_weight': 1.0}
min_epitope_activity_to_include=0.2
Fit a model to all the data, and keep adding epitopes until we either reach the maximum specified or the new epitope has negative activity. Note that that we fit using the reference based-site-numbering scheme, so results are shown with those numbers:Z
[11]:
models = []
for n_epitopes in range(1, max_epitopes + 1):
print(f"\nFitting model with {n_epitopes=}")
# create model
model = polyclonal.Polyclonal(
n_epitopes=n_epitopes,
data_to_fit=prob_escape.rename(
columns={
"antibody_concentration": "concentration",
"aa_substitutions_reference": "aa_substitutions",
}
),
alphabet=polyclonal.AAS_WITHSTOP_WITHGAP,
sites=reference_sites,
spatial_distances=spatial_distances,
)
# fit model
opt_res = model.fit(logfreq=200, **fit_kwargs)
# display activities
print("Activities of epitopes:")
display(model.activity_wt_df.round(1))
print("Max and mean absolute-value escape at each epitope:")
display(
model.mut_escape_df.groupby("epitope")
.aggregate(
max_escape=pd.NamedAgg("escape", "max"),
mean_abs_escape=pd.NamedAgg("escape", lambda s: s.abs().mean()),
)
.round(1)
)
# stop if activity below threshold for any epitope and fit at least one epitope
if len(models) and any(
model.activity_wt_df["activity"] <= min_epitope_activity_to_include
):
print(f"Stop fitting, epitope has activity <={min_epitope_activity_to_include}")
model = models[-1] # get previous model
break
else:
models.append(model)
print(f"\nThe selected model has {len(model.epitopes)} epitopes")
Fitting model with n_epitopes=1
# First fitting site-level model.
# Starting optimization of 1178 parameters at Fri Feb 17 21:18:48 2023.
step time_sec loss fit_loss reg_escape reg_spread reg_spatial reg_uniqueness reg_uniqueness2 reg_activity
0 0.054559 75861 75860 0 0 0 0 0 0.68264
200 11.89 1303.5 1256.5 41.14 0 0 0 0 5.8421
255 14.997 1303.4 1256.3 41.25 0 0 0 0 5.8433
# Successfully finished at Fri Feb 17 21:19:03 2023.
# Starting optimization of 7014 parameters at Fri Feb 17 21:19:04 2023.
step time_sec loss fit_loss reg_escape reg_spread reg_spatial reg_uniqueness reg_uniqueness2 reg_activity
0 0.073214 2376.7 2080 290.89 2.0617e-30 0 0 0 5.8433
200 15.129 1841.5 1753.9 65.533 15.903 0 0 0 6.1452
209 15.8 1841.4 1753.9 65.505 15.916 0 0 0 6.1434
# Successfully finished at Fri Feb 17 21:19:19 2023.
Activities of epitopes:
epitope | activity | |
---|---|---|
0 | 1 | 6.5 |
Max and mean absolute-value escape at each epitope:
max_escape | mean_abs_escape | |
---|---|---|
epitope | ||
1 | 10.7 | 0.1 |
The selected model has 1 epitopes
Epitope activities:
[12]:
model.activity_wt_barplot()
[12]:
Plot of escape values:
[13]:
df_to_merge = site_numbering_map.rename(columns={"reference_site": "site"})
plot_kwargs = antibody_config["plot_kwargs"]
if "plot_title" not in plot_kwargs:
plot_kwargs["plot_title"] = str(antibody)
if "region" in site_numbering_map:
plot_kwargs["site_zoom_bar_color_col"] = "region"
if "addtl_slider_stats" not in plot_kwargs:
plot_kwargs["addtl_slider_stats"] = {"times_seen": 1}
elif "times_seen" not in plot_kwargs["addtl_slider_stats"]:
plot_kwargs["addtl_slider_stats"]["times_seen"] = 1
if "functional effect" in plot_kwargs["addtl_slider_stats"]:
del plot_kwargs["addtl_slider_stats"]["functional effect"] # only antibody averages
if any(site_numbering_map["sequential_site"] != site_numbering_map["reference_site"]):
if "addtl_tooltip_stats" not in plot_kwargs:
plot_kwargs["addtl_tooltip_stats"] = ["sequential_site"]
else:
plot_kwargs["addtl_tooltip_stats"].append("sequential_site")
model.mut_escape_plot(df_to_merge=df_to_merge, **plot_kwargs)
[13]:
Pickle and save model:
[14]:
print(f"Saving {len(models)} models to {pickle_file=}")
with open(pickle_file, "wb") as f:
pickle.dump(models, f)
Saving 1 models to pickle_file='results/polyclonal_fits/Lib-2_2022-03-14_thaw-3_2130-1-0114-112_1.pickle'