Averaging models

Probably the best way to ensure robust inferences and estimate errors is to have multiple experimental replicates, ideally on different libraries.

Here we describe how to average model fits across libraries and/or replicates.

Split data into replicates

We will use our data for the RBD as an earlier examples, but split it into several libraries / replicates.

Specifically, we will fit two different libraries: avg2muts and avg3muts, which have different barcodes and also different mutation rates (although of course in real life you might sometimes want to average results from different libraries with the same mutation rates). We will also simulate having two replicates for each library just by sampling each library. To make this example faster, we’ll just use one concentration:

import pandas as pd

import polyclonal.polyclonal
import polyclonal.polyclonal_collection


# read data
all_data = pd.read_csv("RBD_variants_escape_noisy.csv", na_filter=None)

# split by library and replicates
libraries = ["avg2muts", "avg3muts"]  # the two libraries to use
concentrations = [1]  # use just use this concentration
n_replicates = 2  # number of replicates per library

data_by_replicate = {
    (library, replicate + 1): (
        all_data.query("library == @library")
        .query("concentration in @concentrations")
        .sample(frac=0.3, random_state=replicate)
    )
    for library in libraries
    for replicate in range(n_replicates)
}

Fit models to each replicate

We now fit a Polyclonal model to each replicate using just 2 epitopes, as the data don’t seem sufficient to accurately fit all three epitopes. Then we arrange the models in a data frame:

# first create a data frame with all the models
models_by_replicate = {}
for (library, replicate), data in data_by_replicate.items():
    model = polyclonal.Polyclonal(data_to_fit=data, n_epitopes=2)
    models_by_replicate[(library, replicate)] = model
models_df = (
    pd.Series(models_by_replicate, name="model")
    .rename_axis(["library", "replicate"])
    .reset_index()
)

# now fit the models
n_fit, n_failed, models_df["model"] = polyclonal.polyclonal_collection.fit_models(
    models_df["model"],
    n_threads=2,
    reg_escape_weight=0.01,
    reg_uniqueness2_weight=0,
)

Note how the models are arranged in a data frame:

# NBVAL_IGNORE_OUTPUT

models_df

	library	replicate	model
0	avg2muts	1	<polyclonal.polyclonal.Polyclonal object at 0x...
1	avg2muts	2	<polyclonal.polyclonal.Polyclonal object at 0x...
2	avg3muts	1	<polyclonal.polyclonal.Polyclonal object at 0x...
3	avg3muts	2	<polyclonal.polyclonal.Polyclonal object at 0x...

Average the models

Now we create a PolyclonalAverage model with the models to average. Note that by default the “average” used by PolyclonalAverage is the median rather than the mean between epitopes, although this is a parameter that can also be set to mean.

If your epitopes are too different or poorly defined (e.g., you are trying to fit more epitopes than can be consistently inferred from the data), then you may get an epitope harmonization error:

model_avg = polyclonal.PolyclonalAverage(models_df)

Let’s look at the correlation among the escape at each epitope across models:

# NBVAL_IGNORE_OUTPUT

model_avg.mut_escape_corr_heatmap()

Look at the activities of the epitopes and the rest of the curves. Note how a dark line is shown for the average, and thin lines for individual replicates. It should generally be the case that the epitope with greater activity (more left shifted in plot below) should also be better correlated among replicates (heatmap above) as it can be inferred more reliably:

# NBVAL_IGNORE_OUTPUT

model_avg.curves_plot()

We can access the average escape values. Note that there is also a column escape_min_magnitude that gives the lowest magnitude value:

# NBVAL_IGNORE_OUTPUT

model_avg.mut_escape_df

	epitope	site	wildtype	mutant	mutation	escape_mean	escape_median	escape_min_magnitude	escape_std	n_models	times_seen	frac_models
0	1	331	N	A	N331A	-0.066612	-0.031718	0.026679	0.324873	4	17.75	1.0
1	1	331	N	D	N331D	-0.098052	-0.100283	-0.071111	0.022076	4	11.25	1.0
2	1	331	N	E	N331E	-0.061359	-0.008731	-0.001655	0.112576	4	10.25	1.0
3	1	331	N	F	N331F	0.270243	0.170869	0.014552	0.330315	4	10.00	1.0
4	1	331	N	G	N331G	0.186271	0.164941	0.031369	0.232272	4	25.00	1.0
...	...	...	...	...	...	...	...	...	...	...	...	...
3859	2	531	T	R	T531R	-0.176471	-0.112467	-0.033899	0.187640	4	27.00	1.0
3860	2	531	T	S	T531S	-0.075136	-0.011741	0.000405	0.155037	4	31.75	1.0
3861	2	531	T	V	T531V	-0.041719	-0.005589	0.012383	0.094055	4	19.50	1.0
3862	2	531	T	W	T531W	0.169718	0.100366	-0.004133	0.266617	4	5.25	1.0
3863	2	531	T	Y	T531Y	-0.051762	-0.050223	-0.032461	0.020582	4	11.75	1.0

3864 rows × 12 columns

We can do the same for the ICXX values:

# NBVAL_IGNORE_OUTPUT

model_avg.mut_icXX_df(
    x=0.9,
    icXX_col="IC90",
    log_fold_change_icXX_col="log2 fold change IC90",
)

	site	wildtype	mutant	log2 fold change IC90 mean	log2 fold change IC90 median	log2 fold change IC90 min_magnitude	log2 fold change IC90 std	n_models	times_seen	frac_models
0	331	N	A	-0.159587	-0.206684	-0.135967	0.257279	4	17.75	1.0
1	331	N	D	-0.086087	-0.054268	-0.002643	0.101813	4	11.25	1.0
2	331	N	E	0.034466	0.016051	0.015768	0.067786	4	10.25	1.0
3	331	N	F	0.106040	0.142322	0.141842	0.287263	4	10.00	1.0
4	331	N	G	0.283214	0.282422	0.088191	0.163483	4	25.00	1.0
...	...	...	...	...	...	...	...	...	...	...
2100	531	T	S	-0.068568	-0.016541	-0.004995	0.144948	4	31.75	1.0
2101	531	T	T	0.000000	0.000000	0.000000	0.000000	4	NaN	1.0
2102	531	T	V	0.013326	-0.002602	0.039361	0.236514	4	19.50	1.0
2103	531	T	W	0.199249	0.119866	-0.009133	0.287019	4	5.25	1.0
2104	531	T	Y	-0.021049	-0.048014	-0.044700	0.124192	4	11.75	1.0

2105 rows × 10 columns

Or the per-replicate escape values:

# NBVAL_IGNORE_OUTPUT

model_avg.mut_escape_df_replicates

	epitope	site	wildtype	mutant	mutation	escape	times_seen	library	replicate
0	1	331	N	A	N331A	-0.090115	19	avg2muts	1
1	1	331	N	D	N331D	-0.110795	10	avg2muts	1
2	1	331	N	E	N331E	-0.001655	11	avg2muts	1
3	1	331	N	F	N331F	0.037305	10	avg2muts	1
4	1	331	N	G	N331G	0.031369	18	avg2muts	1
...	...	...	...	...	...	...	...	...	...
15429	2	531	T	R	T531R	-0.033899	24	avg3muts	2
15430	2	531	T	S	T531S	0.000405	42	avg3muts	2
15431	2	531	T	V	T531V	-0.023562	24	avg3muts	2
15432	2	531	T	W	T531W	0.532476	6	avg3muts	2
15433	2	531	T	Y	T531Y	-0.074142	9	avg3muts	2

15434 rows × 9 columns

Now let’s plot the escape. See how you can select mutations based not only on how many times they are seen (averaged over all models in average), but also the number of models in which they are seen.

# NBVAL_IGNORE_OUTPUT

model_avg.mut_escape_plot(addtl_slider_stats={"times_seen": 2})

Here is the same plot plotting the lowest magnitude escape value across libraries for each mutation:

# NBVAL_IGNORE_OUTPUT

model_avg.mut_escape_plot(
    addtl_slider_stats={"times_seen": 2}, avg_type="min_magnitude"
)

We can also plot the log fold change in IC90 caused by each mutation:

# NBVAL_IGNORE_OUTPUT

model_avg.mut_icXX_plot(
    addtl_slider_stats={"times_seen": 2},
    avg_type="median",
)

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Cell In[12], line 3
      1 # NBVAL_IGNORE_OUTPUT
----> 3 model_avg.mut_icXX_plot(
      4     addtl_slider_stats={"times_seen": 2},
      5     avg_type="median",
      6 )

File ~/polyclonal/polyclonal/polyclonal_collection.py:1174, in PolyclonalCollection.mut_icXX_plot(self, x, icXX_col, log_fold_change_icXX_col, min_c, max_c, logbase, check_wt_icXX, biochem_order_aas, df_to_merge, positive_color, negative_color, avg_type, init_n_models, per_model_tooltip, scale_stat_col, **kwargs)
   1172     max_escape_std = kwargs["data_df"][std_col].max()
   1173     kwargs["addtl_slider_stats"][std_col] = max_escape_std
-> 1174     kwargs["addtl_slider_stats_as_max"].append(std_col)
   1176 if "heatmap_min_at_least" not in kwargs:
   1177     kwargs["heatmap_min_at_least"] = -2

KeyError: 'addtl_slider_stats_as_max'