Fit model to data

We will fit a Polyclonal model to the RBD antibody mix we simulated.

First, we read in that simulated data. Recall that we simulated both “exact” and “noisy” data, with several average per-library mutations rates, and at six different concentrations. Here we analyze the noisy data for the library with an average of 3 mutations per gene, measured at three different concentrations, as this represents a fairly realistic representation of a real experiment:

[1]:

import requests
import tempfile

import pandas as pd

import polyclonal

noisy_data = (
    pd.read_csv("RBD_variants_escape_noisy.csv", na_filter=None)
    .query('library == "avg3muts"')
    .query("concentration in [0.25, 1, 4]")
    .reset_index(drop=True)
)

noisy_data

[1]:
library barcode concentration prob_escape aa_substitutions IC90
0 avg3muts AAAACTGCTGAGGAGA 0.25 0.054700 0.08212
1 avg3muts AAAAGCAGGCTACTCT 0.25 0.000000 0.08212
2 avg3muts AAAAGCTATAGGTGCC 0.25 0.007613 0.08212
3 avg3muts AAAAGGTATTAGTGGC 0.25 0.001363 0.08212
4 avg3muts AAAAGTGCCTTCGTTA 0.25 0.000000 0.08212
... ... ... ... ... ... ...
119995 avg3muts GAGCATGATCGACGAA 1.00 0.000000 Y508V H519I 0.10830
119996 avg3muts GAGCATGATCGACGAA 4.00 0.000000 Y508V H519I 0.10830
119997 avg3muts CTTAAAATAGCTGGTC 0.25 0.000000 Y508W 0.08212
119998 avg3muts CTTAAAATAGCTGGTC 1.00 0.012260 Y508W 0.08212
119999 avg3muts CTTAAAATAGCTGGTC 4.00 0.000000 Y508W 0.08212

120000 rows × 6 columns

For spatial regularization (encouraging epitopes to be structurall proximal residues), we read the inter-residue distances in angstroms from PDB 6m0j:

[2]:

# we read the PDB from the webpage into a temporary file and get the distances from that.
# you could also just download the file manually and then read from it.
r = requests.get("https://files.rcsb.org/download/6XM4.pdb")
with tempfile.NamedTemporaryFile() as tmpf:
    _ = tmpf.write(r.content)
    tmpf.flush()
    spatial_distances = polyclonal.pdb_utils.inter_residue_distances(tmpf.name, ["A"])

spatial_distances

[2]:
site_1 site_2 distance chain_1 chain_2
0 27 28 1.332629 A A
1 27 29 4.612508 A A
2 27 30 8.219518 A A
3 27 31 11.016782 A A
4 27 32 13.087037 A A
... ... ... ... ... ...
548623 1308 1310 30.826773 A A
548624 1308 1311 75.350853 A A
548625 1309 1310 12.374796 A A
548626 1309 1311 115.681534 A A
548627 1310 1311 106.112328 A A

548628 rows × 5 columns

Initialize a Polyclonal model with these data, including three epitopes. We know from prior work the three most important epitopes and a key mutation in each, so we use this prior knowledge to “seed” initial guesses that assign large escape values to a key site in each epitope:

  • site 417 for class 1 epitope, which is often the least important

  • site 484 for class 2 epitope, which is often the dominant one

  • site 444 for class 3 epitope, which is often the second most dominant one

[3]:

poly_abs = polyclonal.Polyclonal(
    data_to_fit=noisy_data,
    activity_wt_df=pd.DataFrame.from_records(
        [
            ("1", 1.0),
            ("2", 3.0),
            ("3", 2.0),
        ],
        columns=["epitope", "activity"],
    ),
    site_escape_df=pd.DataFrame.from_records(
        [
            ("1", 417, 10.0),
            ("2", 484, 10.0),
            ("3", 444, 10.0),
        ],
        columns=["epitope", "site", "escape"],
    ),
    data_mut_escape_overlap="fill_to_data",
    spatial_distances=spatial_distances,
)

Now fit the Polyclonal model, logging output every 100 steps. Note how the fitting first just fits a site level model to estimate the average effects of mutations at each site, and then fits the full model. Here we fix the Hill coefficient at one and the non-neutralized fraction at zero:

[4]:

# NBVAL_IGNORE_OUTPUT
opt_res = poly_abs.fit(
    logfreq=200,
    reg_escape_weight=0.01,
    reg_spatial2_weight=1e-5,
    reg_uniqueness2_weight=0,
    fix_non_neutralized_frac=True,
    fix_hill_coefficient=True,
)

#
# Fitting site-level model.
# Starting optimization of 522 parameters at Tue Apr  4 16:01:08 2023.
        step    time_sec        loss    fit_loss  reg_escape  reg_spread reg_spatial reg_uniqueness reg_uniqueness2 reg_activity reg_hill_coefficient reg_non_neutralized_frac
           0     0.10483      2991.9      2980.2     0.29701           0           0              0               0       11.418                    0                        0
         200      24.563        1876      1836.9      2.3802           0      25.111              0               0       11.569                    0                        0
         308      37.745      1873.6      1835.4      2.3448           0      24.241              0               0       11.592                    0                        0
# Successfully finished at Tue Apr  4 16:01:45 2023.
#
# Fitting model.
# Starting optimization of 5799 parameters at Tue Apr  4 16:01:45 2023.
        step    time_sec        loss    fit_loss  reg_escape  reg_spread reg_spatial reg_uniqueness reg_uniqueness2 reg_activity reg_hill_coefficient reg_non_neutralized_frac
           0     0.13785      2026.1      1960.5      29.725  3.6371e-30      24.241              0               0       11.592                    0                        0
         200      31.743      492.84      378.86      33.574      20.624      46.484              0               0         13.3                    0                        0
         400      62.818      479.16       366.6      32.471      21.461      45.134              0               0       13.489                    0                        0
         600      93.955      476.55      364.79      31.686      22.005      44.596              0               0       13.467                    0                        0
         800      124.81      469.69      358.68      31.313      22.319      43.888              0               0       13.489                    0                        0
        1000      155.41      466.36      356.22      31.011      22.471      43.178              0               0       13.479                    0                        0
        1200      186.12      462.73      353.63      30.727      22.423      42.489              0               0       13.454                    0                        0
        1400      216.23      462.18      352.96      30.669      22.483      42.602              0               0       13.464                    0                        0
        1600      246.06      457.12      347.99      30.439      22.776      42.496              0               0       13.421                    0                        0
        1800      275.75      456.23      347.31      30.348      22.746      42.382              0               0       13.447                    0                        0
        1826      279.48      456.23      347.31      30.345       22.75      42.375              0               0       13.449                    0                        0
# Successfully finished at Tue Apr  4 16:06:25 2023.

We can now visualize the resulting curve fits and activity values and compare them to the “true” activity values used to simulate the data. Note that the Hill coefficients are one and the non-neutralized fractions are zero because we fixed them as such above:

[5]:

poly_abs.curve_specs_df.round(1)

[5]:
epitope activity hill_coefficient non_neutralized_frac
0 1 1.5 1.0 0.0
1 2 3.2 1.0 0.0
2 3 2.4 1.0 0.0 
[6]:

# NBVAL_IGNORE_OUTPUT
poly_abs.curves_plot()

[6]:

[7]:

# NBVAL_IGNORE_OUTPUT
poly_abs.activity_wt_barplot()

[7]:

[8]:

# NBVAL_IGNORE_OUTPUT

import altair as alt

true_activities = pd.read_csv("RBD_activity_wt_df.csv")


activity_wt_comparison = (
    pd.concat(
        [
            poly_abs.activity_wt_df.rename(columns={"activity": "predicted"}),
            true_activities.rename(columns={"activity": "actual"}).drop(
                columns="epitope"
            ),
        ],
        axis=1,
    )
).melt(
    id_vars=["epitope"],
    value_vars=["actual", "predicted"],
    var_name="value_type",
    value_name="wildtype activity",
)

alt.Chart(activity_wt_comparison).mark_bar(size=35).encode(
    x="value_type:O",
    y="wildtype activity:Q",
    color=alt.Color(
        "epitope:N", scale=alt.Scale(range=list(poly_abs.epitope_colors.values()))
    ),
    column="epitope:N",
    tooltip=["value_type", alt.Tooltip("wildtype activity", format=".3f"), "epitope"],
).properties(width=100, height=125)

[8]:

Similarly, we can visualize the resulting fits for the escape values, and compare them to the “true” escape values used to simulate the data. Note how entries can be filtered by how many variants a mutation is seen in (times_seen):

[9]:

# NBVAL_IGNORE_OUTPUT
poly_abs.mut_escape_plot()

[9]:

For these simulated data, we can also see how well the fit model does on the “true” simulated values from a library with a different (higher) mutation rate. We therefore read in the “exact” simulated data from a library with a different mutation rate:

[10]:

exact_data = (
    pd.read_csv("RBD_variants_escape_exact.csv", na_filter=None)
    .query('library == "avg3muts"')
    .query("concentration in [0.25, 1, 0.5]")
    .reset_index(drop=True)
)

We will compare the true simulated IC90 values to those predicted by the fit model. We make the comparison on a log scale, and clip IC90s at values >50 as likely to be way outside the dynamic range given the concentrations used:

[11]:

# NBVAL_IGNORE_OUTPUT

import numpy

from plotnine import *

max_ic90 = 50

# we only need the variants, not the concentration for the IC90 comparison
ic90s = (
    exact_data[["aa_substitutions", "IC90"]]
    .assign(IC90=lambda x: x["IC90"].clip(upper=max_ic90))
    .drop_duplicates()
)

ic90s = poly_abs.icXX(ic90s, x=0.9, col="predicted_IC90", max_c=max_ic90)

ic90s = ic90s.assign(
    log_IC90=lambda x: numpy.log10(x["IC90"]),
    predicted_log_IC90=lambda x: numpy.log10(x["predicted_IC90"]),
)

corr = ic90s["log_IC90"].corr(ic90s["predicted_log_IC90"]) ** 2
print(f"Correlation (R^2) is {corr:.2f}")

ic90_corr_plot = (
    ggplot(ic90s)
    + aes("log_IC90", "predicted_log_IC90")
    + geom_point(alpha=0.1, size=1)
    + theme_classic()
    + theme(figure_size=(3, 3))
)

_ = ic90_corr_plot.draw()
# ic90_corr_plot.save('IC90_pred_vs_actual.pdf', height=3.5, width=3.5)

Correlation (R^2) is 0.97
_images/fit_RBD_18_1.png

Next we see how well the model predicts the variant-level escape probabilities \(p_v\left(c\right)\), by reading in exact data from the simulations, and then making predictions of escape probabilities. We both examine and plot the correlations:

[12]:

# NBVAL_IGNORE_OUTPUT

exact_vs_pred = poly_abs.prob_escape(variants_df=exact_data)

print(f"Correlations (R^2) at each concentration:")
display(
    exact_vs_pred.groupby("concentration")
    .apply(lambda x: x["prob_escape"].corr(x["predicted_prob_escape"]) ** 2)
    .rename("correlation (R^2)")
    .reset_index()
    .round(2)
)

pv_corr_plot = (
    ggplot(exact_vs_pred)
    + aes("prob_escape", "predicted_prob_escape")
    + geom_point(alpha=0.1, size=1)
    + facet_wrap("~ concentration", nrow=1)
    + theme_classic()
    + theme(figure_size=(3 * exact_vs_pred["concentration"].nunique(), 3))
)

_ = pv_corr_plot.draw()

Correlations (R^2) at each concentration:
concentration correlation (R^2)
0 0.25 0.99
1 0.50 1.00
2 1.00 1.00

We also examine the correlation between the “true” and inferred mutation-escape values, \(\beta_{m,e}\). In general, it’s necessary to ensure the epitopes match up for this type of comparison as it is arbitrary which epitope in the model is given which name. But above we seeded the epitopes at the site level using site_effects_df when we initialized the Polyclonal object, so they match up with class 1, 2, and 3:

[13]:

# NBVAL_IGNORE_OUTPUT

import altair as alt

mut_escape_pred = pd.read_csv("RBD_mut_escape_df.csv").merge(
    (
        poly_abs.mut_escape_df.assign(
            epitope=lambda x: "class " + x["epitope"].astype(str)
        ).rename(columns={"escape": "predicted escape"})
    ),
    on=["mutation", "epitope"],
    validate="one_to_one",
)

print("Correlation (R^2) between predicted and true values:")
corr = (
    mut_escape_pred.groupby("epitope")
    .apply(lambda x: x["escape"].corr(x["predicted escape"]) ** 2)
    .rename("correlation (R^2)")
    .reset_index()
)
display(corr.round(2))

# for testing since we nbval ignore cell output
assert (
    numpy.allclose(
        corr["correlation (R^2)"], numpy.array([0.60, 0.96, 0.95]), atol=0.02
    )
    == True
)

corr_chart = (
    alt.Chart(mut_escape_pred)
    .encode(
        x="escape",
        y="predicted escape",
        color=alt.Color(
            "epitope", scale=alt.Scale(range=list(poly_abs.epitope_colors.values()))
        ),
        tooltip=["mutation", "epitope"],
    )
    .mark_point(opacity=0.5)
    .properties(width=250, height=250)
    .facet(column="epitope")
    .resolve_scale(
        x="independent",
        y="independent",
    )
)

corr_chart

Correlation (R^2) between predicted and true values:
epitope correlation (R^2)
0 class 1 0.60
1 class 2 0.96
2 class 3 0.95 
[13]:

The correlations are strongest for the dominant epitope (class 2), which makes sense as this will drive the highest escape signal.