Sera concentrations

Choosing sera concentrations is an important consideration for DMS selections. We’ll use simulated data to estimate the optimal set of sera concentrations that provides the best performance while minimizing the number of concentrations (or selection experiments) required.

import os
import pickle

import altair as alt
import numpy as np
import pandas as pd
import polyclonal

First, we read in a simulated “noisy” dataset measured at six sera concentrations. The variants in this library were simulated to contain a Poisson-distributed number of mutations, with an average of three mutations per gene. The variants also generally span a wide range of escape fractions across the different sera concentrations.

noisy_data = (
    pd.read_csv("RBD_variants_escape_noisy.csv", na_filter=None)
    .query("library == 'avg3muts'")
    .reset_index(drop=True)
)
noisy_data

	library	barcode	concentration	prob_escape	aa_substitutions	IC90
0	avg3muts	AAAACTGCTGAGGAGA	0.125	0.013540		0.08212
1	avg3muts	AAAAGCAGGCTACTCT	0.125	0.066420		0.08212
2	avg3muts	AAAAGCTATAGGTGCC	0.125	0.023120		0.08212
3	avg3muts	AAAAGGTATTAGTGGC	0.125	0.005456		0.08212
4	avg3muts	AAAAGTGCCTTCGTTA	0.125	0.054760		0.08212
...	...	...	...	...	...	...
239995	avg3muts	CTTAAAATAGCTGGTC	0.250	0.000000	Y508W	0.08212
239996	avg3muts	CTTAAAATAGCTGGTC	0.500	0.026480	Y508W	0.08212
239997	avg3muts	CTTAAAATAGCTGGTC	1.000	0.012260	Y508W	0.08212
239998	avg3muts	CTTAAAATAGCTGGTC	2.000	0.000000	Y508W	0.08212
239999	avg3muts	CTTAAAATAGCTGGTC	4.000	0.000000	Y508W	0.08212

240000 rows × 6 columns

While the concentrations are in arbitrary units (0.125, 0.25, 0.5, 1, 2, 4), we can describe them in terms of their ICXX’s against wildtype virus. So, we can interpret each as a sera concentration that neutralizes XX % of wildtype viruses.

wt_data_icxx_df = (
    pd.read_csv("RBD_variants_escape_exact.csv", na_filter=None)
    .query("aa_substitutions == ''")
    .reset_index(drop=True)
    .drop(columns="library")
    .drop_duplicates()
    .assign(ICxx_against_wt=lambda x: round((1 - x["prob_escape"]) * 100, 3))
    .drop(columns=["aa_substitutions", "prob_escape", "IC90"])
    .set_index("concentration")
)

display(wt_data_icxx_df)
wt_data_icxx = wt_data_icxx_df["ICxx_against_wt"].to_dict()

	barcode	ICxx_against_wt
concentration
0.125	AAAAAATGTTCTATCC	95.141
0.125	AAAAACAATCCGGACT	95.141
0.125	AAAAACGCGGTCACTT	95.141
0.125	AAAAACTTGGCTAGCT	95.141
0.125	AAAAAGCAAGGCCCAG	95.141
...	...	...
4.000	TTTGTATGGTCCATAT	99.999
4.000	TTTGTCTCGAATGGTG	99.999
4.000	TTTTAAGCTCATACGC	99.999
4.000	TTTTATCCACCGAACT	99.999
4.000	TTTTCGATCCTTGTCA	99.999

137868 rows × 2 columns

Now, we’ll fit multiple Polyclonal models to data measured at single concentrations. We’ll initialize each Polyclonal model with the same values. 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

Additionally, we’ll store fit models as pickle files, so that we can conveniently load them in the future without having to fit again.

conc_sets = [[0.125], [0.25], [0.5], [1], [2], [4]]

# Make a directory to house pickled models
os.makedirs("fit_polyclonal_models", exist_ok=True)


def fit_polyclonal(conc_set):
    """
    Fit `Polyclonal` model with data measured for a specific concentration set.
    Returns fit `Polyclonal` object.
    """
    poly_abs = polyclonal.Polyclonal(
        data_to_fit=noisy_data.query(f"concentration in {conc_set}"),
        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",
    )
    poly_abs.fit(reg_escape_weight=0.01, reg_uniqueness2_weight=0)
    return poly_abs


# Store all fit models in a dictionary for future lookup
fit_models = {}

for s in conc_sets:
    # These are the keys for fit models
    model_string = f"noisy_{s}conc_3muts"

    # If the pickled model exists in fit_polyclonal_models directory,
    # load it and update fit_models
    if os.path.exists(f"fit_polyclonal_models/{model_string}.pkl") is True:
        model = pickle.load(open(f"fit_polyclonal_models/{model_string}.pkl", "rb"))
        fit_models.update({model_string: model})
        print(f"Model with {s} was already fit.")
    else:
        # Else, fit a model using fit_polyclonal(), save it to the
        # fit_polyclonal_models directory, and update fit_models
        model = fit_polyclonal(s)
        fit_models.update({model_string: model})
        pickle.dump(model, open(f"fit_polyclonal_models/{model_string}.pkl", "wb"))
        print(f"Model with {s} fit and saved.")

Model with [0.125] was already fit.
Model with [0.25] was already fit.
Model with [0.5] was already fit.
Model with [1] was already fit.
Model with [2] was already fit.
Model with [4] was already fit.

We can look at the correlation between the “true” and inferred mutation-escape values, \(\beta_{m,e}\), for the fit models. These mutation-escape values represent the extent to which mutations mediate escape from specific epitopes.

all_corrs = pd.DataFrame()

for s in conc_sets:
    model = fit_models[f"noisy_{s}conc_3muts"]

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

    corr = (
        mut_escape_pred.groupby("epitope")
        .apply(lambda x: x["escape"].corr(x["predicted escape"]) ** 2)
        .rename("correlation (R^2)")
        .reset_index()
    )

    all_corrs = pd.concat(
        [
            all_corrs,
            corr.assign(
                icxx_set=[" ".join([f"IC{wt_data_icxx[c]}" for c in s])]
                * len(corr.index)
            ),
        ]
    )

# NBVAL_IGNORE_OUTPUT
alt.Chart(all_corrs).mark_circle(size=125).encode(
    x=alt.X("icxx_set:O", sort=alt.EncodingSortField("x", order="descending")),
    y="correlation (R^2):Q",
    column="epitope:N",
    tooltip=["icxx_set:O", alt.Tooltip("correlation (R^2)", format=".3f")],
    color=alt.Color(
        "epitope", scale=alt.Scale(range=["#0072B2", "#CC79A7", "#4C3549"]), legend=None
    ),
).properties(width=200, height=200, title="inferred vs. true mutation escape values")

/home/tyu2/.local/lib/python3.8/site-packages/altair/utils/core.py:317: FutureWarning: iteritems is deprecated and will be removed in a future version. Use .items instead.
  for col_name, dtype in df.dtypes.iteritems():

Additionally, we’ll look at the correlation between “true” and predicted IC90’s for each of the fit models. To do this, we’ll predict the IC90’s of variants in a separate library with a with a different (higher) mutation rate. We therefore read in the “exact” simulated data from a library containing variants with an average of four mutations per gene.

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

We’ll make the comparison on a log scale, and clip IC90s at values >50 as that is likely to be way outside the dynamic range given the concentrations used.

ic90_corrs = pd.DataFrame()

max_ic90 = 50
for s in conc_sets:
    model = fit_models[f"noisy_{s}conc_3muts"]

    ic90s = (
        exact_data[["aa_substitutions", "IC90"]]
        .assign(IC90=lambda x: x["IC90"].clip(upper=max_ic90))
        .drop_duplicates()
    )
    ic90s = model.filter_variants_by_seen_muts(ic90s)
    ic90s = model.icXX(ic90s, x=0.9, col="predicted_IC90", max_c=max_ic90)

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

    corr = ic90s["log_IC90"].corr(ic90s["predicted_log_IC90"]) ** 2

    ic90_corrs = pd.concat(
        [
            ic90_corrs,
            pd.DataFrame(
                {
                    "correlation (R^2)": corr,
                    "icxx_set": [f"IC{wt_data_icxx[c]}" for c in s],
                }
            ),
        ]
    )

# NBVAL_IGNORE_OUTPUT
alt.Chart(ic90_corrs).mark_circle(size=125).encode(
    x="icxx_set:O",
    y="correlation (R^2):Q",
    tooltip=["icxx_set", alt.Tooltip("correlation (R^2)", format=".3f")],
).properties(width=200, height=200, title="predicted vs. true IC90")

/home/tyu2/.local/lib/python3.8/site-packages/altair/utils/core.py:317: FutureWarning: iteritems is deprecated and will be removed in a future version. Use .items instead.
  for col_name, dtype in df.dtypes.iteritems():

Across all models fit on data with single concentrations, the correlation was very strong.

Next, we’ll fit a couple more Polyclonal models to determine if adding a second or third concentration to the best performing single concentration will improve inference of the mutation-escape values. As a reference, we also fit a Polyclonal model on all six concentrations.

conc_sets = [
    [0.25, 1],
    [0.5, 1],
    [1, 2],
    [1, 4],
    [0.5, 1, 2],
    [0.25, 1, 4],
    [0.125, 0.25, 0.5, 1, 2, 4],
]

for s in conc_sets:
    # These are the keys for fit models
    model_string = f"noisy_{s}conc_3muts"

    # If the pickled model exists in fit_polyclonal_models directory,
    # load it and add to fit_models
    if os.path.exists(f"fit_polyclonal_models/{model_string}.pkl") is True:
        model = pickle.load(open(f"fit_polyclonal_models/{model_string}.pkl", "rb"))
        fit_models.update({model_string: model})
        print(f"Model with {s} was already fit.")
    else:
        # Else, fit a model using fit_polyclonal(), save it to the
        # fit_polyclonal_models directory, and add to fit_models
        model = fit_polyclonal(s)
        fit_models.update({model_string: model})
        pickle.dump(model, open(f"fit_polyclonal_models/{model_string}.pkl", "wb"))
        print(f"Model with {s} fit and saved.")

Model with [0.25, 1] was already fit.
Model with [0.5, 1] was already fit.
Model with [1, 2] was already fit.
Model with [1, 4] was already fit.
Model with [0.5, 1, 2] was already fit.
Model with [0.25, 1, 4] was already fit.
Model with [0.125, 0.25, 0.5, 1, 2, 4] was already fit.

Again, lets look at the correlation between “true” and predicted mutation-escape values for each of the fit models.

# add best performing single concentration to concentration sets, so we can plot as reference
conc_sets_to_plot = [[1]] + conc_sets

all_corrs = pd.DataFrame()

for s in conc_sets_to_plot:
    model = fit_models[f"noisy_{s}conc_3muts"]

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

    corr = (
        mut_escape_pred.groupby("epitope")
        .apply(lambda x: x["escape"].corr(x["predicted escape"]) ** 2)
        .rename("correlation (R^2)")
        .reset_index()
    )

    all_corrs = pd.concat(
        [
            all_corrs,
            corr.assign(
                icxx_set=[",".join([f"IC{wt_data_icxx[c]}" for c in s])]
                * len(corr.index),
                num_concs=len([c for c in s]),
            ),
        ]
    ).reset_index(drop=True)

# NBVAL_IGNORE_OUTPUT
alt.Chart(all_corrs).mark_circle(size=125).encode(
    x=alt.X(
        "icxx_set:N",
        sort=alt.EncodingSortField("x", order="descending"),
    ),
    y="correlation (R^2):Q",
    tooltip=["icxx_set", alt.Tooltip("correlation (R^2)", format=".3f")],
    color=alt.Color(
        "epitope", scale=alt.Scale(range=["#0072B2", "#CC79A7", "#4C3549"]), legend=None
    ),
    row="epitope:N",
).properties(
    width=375, height=200, title="inferred vs. true mutation escape values"
).configure_axis(
    labelLimit=300
)

/home/tyu2/.local/lib/python3.8/site-packages/altair/utils/core.py:317: FutureWarning: iteritems is deprecated and will be removed in a future version. Use .items instead.
  for col_name, dtype in df.dtypes.iteritems():

We see that the performance of models fit on data with two or three concentrations is nearly indistinguishable from that of the model fit on data with all six concentrations. Additionally, adding two or three concentrations does lead to a modest improvement in the correlation, especially for the class 1 epitope, which is expected to be the hardest to predict since it has the lowest wildtype activity. We can also summarize this by plotting the number of concentrations used for fitting on the x-axis.

# NBVAL_IGNORE_OUTPUT
alt.Chart(all_corrs).mark_circle(size=125).encode(
    x=alt.X("num_concs:O", axis=alt.Axis(labelAngle=0)),
    y="correlation (R^2):Q",
    column="epitope:N",
    tooltip=[
        "icxx_set:O",
        "num_concs:O",
        alt.Tooltip("correlation (R^2)", format=".3f"),
    ],
    color=alt.Color(
        "epitope", scale=alt.Scale(range=["#0072B2", "#CC79A7", "#4C3549"]), legend=None
    ),
).properties(width=150, height=200, title="inferred vs. true mutation escape values")

Again, lets look at the correlation between “true” and predicted IC90’s for each of the fit models.

ic90_corrs = pd.DataFrame()

max_ic90 = 50
for s in conc_sets_to_plot:
    model = fit_models[f"noisy_{s}conc_3muts"]

    ic90s = (
        exact_data[["aa_substitutions", "IC90"]]
        .assign(IC90=lambda x: x["IC90"].clip(upper=max_ic90))
        .drop_duplicates()
    )
    ic90s = model.filter_variants_by_seen_muts(ic90s)
    ic90s = model.icXX(ic90s, x=0.9, col="predicted_IC90", max_c=max_ic90)

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

    corr = ic90s["log_IC90"].corr(ic90s["predicted_log_IC90"]) ** 2

    ic90_corrs = pd.concat(
        [
            ic90_corrs,
            pd.DataFrame(
                {
                    "correlation (R^2)": corr,
                    "icxx_set": [[f"IC{wt_data_icxx[c]}" for c in s]],
                }
            ),
        ]
    )

# NBVAL_IGNORE_OUTPUT
alt.Chart(ic90_corrs).mark_circle(size=125).encode(
    x=alt.X(
        "icxx_set:O",
        sort=alt.EncodingSortField("x", order="descending"),
        axis=alt.Axis(labelAngle=0),
    ),
    y="correlation (R^2):Q",
    tooltip=["icxx_set", alt.Tooltip("correlation (R^2)", format=".3f")],
).properties(width=375, height=200, title="predicted vs. true IC90")

/home/tyu2/.local/lib/python3.8/site-packages/altair/utils/core.py:317: FutureWarning: iteritems is deprecated and will be removed in a future version. Use .items instead.
  for col_name, dtype in df.dtypes.iteritems():

Fitting with more concentrations maintained the excellent IC90 prediction exhibited by models fit on data with a single concentration.

Summary

Based on these simulation experiments:

1. Choose a single concentration that is potent enough (ex. IC99.9) to neutralize all wildtype viruses, but not potent enough to neutralize all the variants.

2. Adding a second or third concentration can help refine inference of true mutation-escape values, especially for the most subdominant epitopes. Due to technical difficulties in precisely achieving the IC99.9 in a given experiment, we suggest using concentrations that are 2-4 fold higher or lower than IC99.9 in order to more likely span the correct dynamic range.