Make some paper figures

The next cell is tagged parameters for papermill parameterization:

# tagged parameters for `papermill`

# Parameters
params = {
    "func_scores_293T-Mxra8-E3E2-B-241028-func-1": "results/func_scores/293T-Mxra8-E3E2-B-241028-func-1_func_scores.csv",
    "func_scores_293T-Mxra8-E3E2-B-241028-func-2": "results/func_scores/293T-Mxra8-E3E2-B-241028-func-2_func_scores.csv",
    "func_scores_293T-TIM1-E3E2-B-241028-func-1": "results/func_scores/293T-TIM1-E3E2-B-241028-func-1_func_scores.csv",
    "func_scores_293T-TIM1-E3E2-B-241028-func-2": "results/func_scores/293T-TIM1-E3E2-B-241028-func-2_func_scores.csv",
    "func_scores_293T-Mxra8-6KE1-B-241028-func-1": "results/func_scores/293T-Mxra8-6KE1-B-241028-func-1_func_scores.csv",
    "func_scores_293T-Mxra8-6KE1-B-241028-func-2": "results/func_scores/293T-Mxra8-6KE1-B-241028-func-2_func_scores.csv",
    "func_scores_293T-TIM1-6KE1-B-241028-func-1": "results/func_scores/293T-TIM1-6KE1-B-241028-func-1_func_scores.csv",
    "func_scores_293T-TIM1-6KE1-B-241028-func-2": "results/func_scores/293T-TIM1-6KE1-B-241028-func-2_func_scores.csv",
    "func_scores_293T-Mxra8-E3E2-A-241113-func-1": "results/func_scores/293T-Mxra8-E3E2-A-241113-func-1_func_scores.csv",
    "func_scores_293T-Mxra8-E3E2-A-241113-func-2": "results/func_scores/293T-Mxra8-E3E2-A-241113-func-2_func_scores.csv",
    "func_scores_293T-TIM1-E3E2-A-241113-func-1": "results/func_scores/293T-TIM1-E3E2-A-241113-func-1_func_scores.csv",
    "func_scores_293T-TIM1-E3E2-A-241113-func-2": "results/func_scores/293T-TIM1-E3E2-A-241113-func-2_func_scores.csv",
    "func_scores_293T-Mxra8-6KE1-A-241113-func-1": "results/func_scores/293T-Mxra8-6KE1-A-241113-func-1_func_scores.csv",
    "func_scores_293T-Mxra8-6KE1-A-241113-func-2": "results/func_scores/293T-Mxra8-6KE1-A-241113-func-2_func_scores.csv",
    "func_scores_293T-TIM1-6KE1-A-241113-func-1": "results/func_scores/293T-TIM1-6KE1-A-241113-func-1_func_scores.csv",
    "func_scores_293T-TIM1-6KE1-A-241113-func-2": "results/func_scores/293T-TIM1-6KE1-A-241113-func-2_func_scores.csv",
    "func_scores_C636-E3E2-B-241122-func-1": "results/func_scores/C636-E3E2-B-241122-func-1_func_scores.csv",
    "func_scores_C636-E3E2-B-241122-func-2": "results/func_scores/C636-E3E2-B-241122-func-2_func_scores.csv",
    "func_scores_C636-6KE1-B-241122-func-1": "results/func_scores/C636-6KE1-B-241122-func-1_func_scores.csv",
    "func_scores_C636-6KE1-B-241122-func-2": "results/func_scores/C636-6KE1-B-241122-func-2_func_scores.csv",
    "func_scores_C636-E3E2-A-241212-func-1": "results/func_scores/C636-E3E2-A-241212-func-1_func_scores.csv",
    "func_scores_C636-E3E2-A-241212-func-2": "results/func_scores/C636-E3E2-A-241212-func-2_func_scores.csv",
    "func_scores_C636-6KE1-A-241212-func-1": "results/func_scores/C636-6KE1-A-241212-func-1_func_scores.csv",
    "func_scores_C636-6KE1-A-241212-func-2": "results/func_scores/C636-6KE1-A-241212-func-2_func_scores.csv",
    "func_effects_293T-Mxra8_entry_293T-Mxra8-E3E2-A-241113-func-1": "results/func_effects/by_selection/293T-Mxra8-E3E2-A-241113-func-1_func_effects.csv",
    "func_effects_293T-Mxra8_entry_293T-Mxra8-E3E2-A-241113-func-2": "results/func_effects/by_selection/293T-Mxra8-E3E2-A-241113-func-2_func_effects.csv",
    "func_effects_293T-Mxra8_entry_293T-Mxra8-E3E2-B-241028-func-1": "results/func_effects/by_selection/293T-Mxra8-E3E2-B-241028-func-1_func_effects.csv",
    "func_effects_293T-Mxra8_entry_293T-Mxra8-E3E2-B-241028-func-2": "results/func_effects/by_selection/293T-Mxra8-E3E2-B-241028-func-2_func_effects.csv",
    "func_effects_293T-Mxra8_entry_293T-Mxra8-6KE1-A-241113-func-1": "results/func_effects/by_selection/293T-Mxra8-6KE1-A-241113-func-1_func_effects.csv",
    "func_effects_293T-Mxra8_entry_293T-Mxra8-6KE1-A-241113-func-2": "results/func_effects/by_selection/293T-Mxra8-6KE1-A-241113-func-2_func_effects.csv",
    "func_effects_293T-Mxra8_entry_293T-Mxra8-6KE1-B-241028-func-1": "results/func_effects/by_selection/293T-Mxra8-6KE1-B-241028-func-1_func_effects.csv",
    "func_effects_293T-Mxra8_entry_293T-Mxra8-6KE1-B-241028-func-2": "results/func_effects/by_selection/293T-Mxra8-6KE1-B-241028-func-2_func_effects.csv",
    "func_effects_293T-TIM1_entry_293T-TIM1-E3E2-A-241113-func-1": "results/func_effects/by_selection/293T-TIM1-E3E2-A-241113-func-1_func_effects.csv",
    "func_effects_293T-TIM1_entry_293T-TIM1-E3E2-A-241113-func-2": "results/func_effects/by_selection/293T-TIM1-E3E2-A-241113-func-2_func_effects.csv",
    "func_effects_293T-TIM1_entry_293T-TIM1-E3E2-B-241028-func-1": "results/func_effects/by_selection/293T-TIM1-E3E2-B-241028-func-1_func_effects.csv",
    "func_effects_293T-TIM1_entry_293T-TIM1-E3E2-B-241028-func-2": "results/func_effects/by_selection/293T-TIM1-E3E2-B-241028-func-2_func_effects.csv",
    "func_effects_293T-TIM1_entry_293T-TIM1-6KE1-A-241113-func-1": "results/func_effects/by_selection/293T-TIM1-6KE1-A-241113-func-1_func_effects.csv",
    "func_effects_293T-TIM1_entry_293T-TIM1-6KE1-A-241113-func-2": "results/func_effects/by_selection/293T-TIM1-6KE1-A-241113-func-2_func_effects.csv",
    "func_effects_293T-TIM1_entry_293T-TIM1-6KE1-B-241028-func-1": "results/func_effects/by_selection/293T-TIM1-6KE1-B-241028-func-1_func_effects.csv",
    "func_effects_293T-TIM1_entry_293T-TIM1-6KE1-B-241028-func-2": "results/func_effects/by_selection/293T-TIM1-6KE1-B-241028-func-2_func_effects.csv",
    "func_effects_C636_entry_C636-E3E2-B-241122-func-1": "results/func_effects/by_selection/C636-E3E2-B-241122-func-1_func_effects.csv",
    "func_effects_C636_entry_C636-E3E2-B-241122-func-2": "results/func_effects/by_selection/C636-E3E2-B-241122-func-2_func_effects.csv",
    "func_effects_C636_entry_C636-6KE1-B-241122-func-1": "results/func_effects/by_selection/C636-6KE1-B-241122-func-1_func_effects.csv",
    "func_effects_C636_entry_C636-6KE1-B-241122-func-2": "results/func_effects/by_selection/C636-6KE1-B-241122-func-2_func_effects.csv",
    "func_effects_C636_entry_C636-E3E2-A-241212-func-1": "results/func_effects/by_selection/C636-E3E2-A-241212-func-1_func_effects.csv",
    "func_effects_C636_entry_C636-E3E2-A-241212-func-2": "results/func_effects/by_selection/C636-E3E2-A-241212-func-2_func_effects.csv",
    "func_effects_C636_entry_C636-6KE1-A-241212-func-1": "results/func_effects/by_selection/C636-6KE1-A-241212-func-1_func_effects.csv",
    "func_effects_C636_entry_C636-6KE1-A-241212-func-2": "results/func_effects/by_selection/C636-6KE1-A-241212-func-2_func_effects.csv",
    "codon_variants": "results/variants/codon_variants.csv",
    "annotated_mut_summary": "results/annotated_summary_csvs/entry_293T-Mxra8_C636_293T-TIM1_Mxra8-binding_annotated.csv",
    "annotated_site_summary": "results/annotated_summary_csvs/entry_293T-Mxra8_C636_293T-TIM1_Mxra8-binding_annotated_site_means.csv",
    "mxra8_binding_effects": "results/summaries/binding_mouse_vs_human_Mxra8.csv",
    "mxra8_validation_curves": "manual_analyses/experimental_data/RVP.mutants.neutralization.by.soluble.mouse.Mxra8.csv",
    "mxra8_cell_binding": "manual_analyses/experimental_data/293T-E_MXRA8_Fc_binding.csv",
    "chikv_titers": "manual_analyses/experimental_data/CHIKV_mutant_titers.csv",
    "rvp_titers": "manual_analyses/experimental_data/RVP_mutant_titers.csv",
    "addtl_site_annotations": "data/addtl_site_annotations.csv",
    "natural_summary_293T_Mxra8": "manual_analyses/compare_experimental_natural/results/plots/per_site_summary_293T_Mxra8.csv",
    "natural_summary_293T_TIM1": "manual_analyses/compare_experimental_natural/results/plots/per_site_summary_293T_TIM1.csv",
    "natural_summary_C636": "manual_analyses/compare_experimental_natural/results/plots/per_site_summary_C636.csv",
    "protein_alignment": "nextstrain/results/alignments/protein_ungapped.fa",
    "mxra8_validation_svg": "results/paper_figures/mxra8_validation.svg",
}
min_times_seen = 2
cell_entry_clip_lower = -6
capsid_len = 261
ref_strain = "181-25_MW473668"

Python imports:

import copy
import itertools
import os
from collections import Counter

import altair as alt

from Bio import SeqIO

import dms_variants.codonvarianttable

import matplotlib.cm
import matplotlib.pyplot as plt

import neutcurve

import numpy

import pandas as pd

import scipy.optimize

import scipy.stats

_ = alt.data_transformers.disable_max_rows()

Plot validation assays with RVP for Mxra8 binding

These plot correlations of experimental data for RVP validations and DMS.

First, read the validation data for the curves, fit curves, and plot them:

mxra8_curves = (
    pd.read_csv(params["mxra8_validation_curves"])
    .rename(columns={"Concentration (ug/mL)": "concentration"})
    .set_index("concentration")
)

# get and sort variants
mxra8_validation_variants = [
    c for c in mxra8_curves.columns if not c.startswith("Unnamed:")
]
assert mxra8_validation_variants[0] == "unmutated"
mxra8_validation_variants = ["unmutated"] + [
    tup[1]
    for tup in sorted(((v.split("-")[0], int(v.split("-")[1][1: -1]), v[-1]), v) for v in mxra8_validation_variants[1:])
]

mxra8_curves_tidy = []
for variant in mxra8_validation_variants:
    variant_index = mxra8_curves.columns.tolist().index(variant)
    cols = mxra8_curves.columns[variant_index: variant_index + 3].tolist()
    mxra8_curves_tidy.append(
        mxra8_curves[cols].assign(variant=variant).rename(
            columns={c: f"replicate {i + 1}" for (i, c) in enumerate(cols)}
        )
    )
mxra8_curves_tidy = (
    pd.concat(mxra8_curves_tidy)
    .reset_index()
    .melt(
        id_vars=["concentration", "variant"],
        var_name="replicate",
        value_name="fraction_infectivity",
    )
    .query("concentration != 0")
    .assign(
        serum="serum",
        variant=lambda x: pd.Categorical(x["variant"], mxra8_validation_variants, ordered=True),
    )
    .sort_values("variant")
)

mxra8_curve_fits = neutcurve.CurveFits(
    mxra8_curves_tidy,
    conc_col="concentration",
    fracinf_col="fraction_infectivity",
    virus_col="variant",
    replicate_col="replicate",
    fixslope=[0.5, 2],
)

nviruses = len(mxra8_validation_variants)
variant_colors = ["black"] + [
    tuple(c) for c in matplotlib.cm.get_cmap("tab20", nviruses - 1)(numpy.arange(nviruses - 1))
]
variant_markers = [
    m for m in matplotlib.markers.MarkerStyle.markers.keys()
    if m not in {".", ",", "1", "2", "3", "4"}
][: nviruses]

_, mxra8_curve_fits_axes = mxra8_curve_fits.plotSera(
    titles=[""],
    max_viruses_per_subplot=nviruses,
    colors=variant_colors,
    markers=variant_markers,
    ylabel="fraction infectivity",
    xlabel="mouse Mxra8 (ug/ml)",
    heightscale=0.9,
)
_ = mxra8_curve_fits_axes[0][0].set_yticks([0, 0.5, 1])

/tmp/ipykernel_3610274/2402822684.py:53: MatplotlibDeprecationWarning: The get_cmap function was deprecated in Matplotlib 3.7 and will be removed in 3.11. Use ``matplotlib.colormaps[name]`` or ``matplotlib.colormaps.get_cmap()`` or ``pyplot.get_cmap()`` instead.
/fh/fast/bloom_j/computational_notebooks/xju/2024/CHIKV_181-25_E_DMS/.snakemake/conda/4edaebae31b5476cccce3cc68a120e2f_/lib/python3.13/site-packages/neutcurve/hillcurve.py:1179: RuntimeWarning: invalid value encountered in power

Plot IC50s versus DMS data:

mxra8_ic50s = (
    pd.read_csv(params["mxra8_binding_effects"])
    .assign(
        mutation=lambda x: x["region"] + "-" + x["wildtype"] + x["site"].str.split("(").str[0] + x["mutant"],
    )
    .rename(columns={"mutation": "virus", "binding to mouse Mxra8": "DMS_effect"})
    [["virus", "DMS_effect"]]
    .merge(
        mxra8_curve_fits.fitParams()[["virus", "ic50"]], on="virus", how="right", validate="1:1",
    )
    .assign(
        DMS_effect=lambda x: x["DMS_effect"].where(x["virus"] != "unmutated", 0),
        inv_ic50=lambda x: 1 / x["ic50"],
    )
)

r = mxra8_ic50s["DMS_effect"].corr(numpy.log(mxra8_ic50s["inv_ic50"]))

# per-variant color and marker, matched to the neutralization-curve panel
color_of = dict(zip(mxra8_validation_variants, variant_colors))
marker_of = dict(zip(mxra8_validation_variants, variant_markers))

# --- cell-based Mxra8-Fc binding titration (293T cells expressing E, stained with Mxra8-Fc) ---
mxra8_cell_binding = pd.read_csv(params["mxra8_cell_binding"])
binding_conc_col = mxra8_cell_binding.columns[0]
binding_variants = [
    c for c in mxra8_cell_binding.columns if not c.startswith("Unnamed:") and c != binding_conc_col
]

mxra8_binding_tidy = []
for variant in binding_variants:
    i = mxra8_cell_binding.columns.tolist().index(variant)
    reps = mxra8_cell_binding.iloc[:, [0, i, i + 1]].copy()
    reps.columns = ["concentration", "rep1", "rep2"]
    reps["variant"] = variant
    mxra8_binding_tidy.append(reps)
mxra8_binding_tidy = (
    pd.concat(mxra8_binding_tidy, ignore_index=True)
    .query("concentration > 0")
    .assign(
        mean=lambda x: x[["rep1", "rep2"]].mean(axis=1),
        err=lambda x: (x["rep1"] - x["rep2"]).abs() / 2,
    )
)

# increasing Hill curve with a fixed slope, analogous to the fixslope used for the neut curves
binding_hill_slope = 1.0


def binding_hill(c, bottom, top, ec50):
    return bottom + (top - bottom) / (1 + (ec50 / c) ** binding_hill_slope)


mxra8_validation_fig, mxra8_validation_axes = plt.subplots(
    1, 3, figsize=(10.5, 2.8), gridspec_kw={"width_ratios": [1.5, 1, 1.1], "wspace": 0.45}
)

# panel 1: DMS effect versus RVP 1/IC50
for v, x, y in zip(mxra8_ic50s["virus"], mxra8_ic50s["DMS_effect"], mxra8_ic50s["inv_ic50"]):
    mxra8_validation_axes[1].scatter([x], [y], color=color_of[v], marker=marker_of[v], label=v)
mxra8_validation_axes[1].set_yscale("log")
mxra8_validation_axes[1].set_xlabel("DMS effect on mouse Mxra8 binding")
mxra8_validation_axes[1].set_ylabel("effect in RVP validation (1/IC50)")
mxra8_validation_axes[1].text(
    0.05, 0.9, f"R = {r:.2f}", transform=mxra8_validation_axes[1].transAxes, fontsize=12
)

# panel 2: cell-based Mxra8-Fc binding titration
binding_xgrid = numpy.logspace(
    numpy.log10(mxra8_binding_tidy["concentration"].min()),
    numpy.log10(mxra8_binding_tidy["concentration"].max()),
    200,
)
for variant in binding_variants:
    d = mxra8_binding_tidy.query("variant == @variant").sort_values("concentration")
    mxra8_validation_axes[2].errorbar(
        d["concentration"], d["mean"], yerr=d["err"], fmt=marker_of[variant], color=color_of[variant],
        markersize=6, capsize=2, elinewidth=0.8, linewidth=0,
    )
    xy = numpy.r_[d["concentration"].values, d["concentration"].values]
    yy = numpy.r_[d["rep1"].values, d["rep2"].values]
    popt, _ = scipy.optimize.curve_fit(
        binding_hill, xy, yy, p0=[0, max(yy.max(), 1), numpy.median(d["concentration"])],
        maxfev=20000, bounds=([-5, 0, 1e-4], [20, 200, 1e4]),
    )
    mxra8_validation_axes[2].plot(binding_xgrid, binding_hill(binding_xgrid, *popt), color=color_of[variant], lw=1)
mxra8_validation_axes[2].set_xscale("log")
mxra8_validation_axes[2].set_xlabel("mouse Mxra8 concentration (ug/ml)")
mxra8_validation_axes[2].set_ylabel("cell-surface Mxra8-Fc binding (%)")

# shared legend to the right of the whole figure
handles, labels = mxra8_validation_axes[1].get_legend_handles_labels()
mxra8_validation_fig.legend(
    handles, labels, loc="center left", bbox_to_anchor=(0.92, 0.5),
    borderaxespad=0.0, labelspacing=0.1,
)


# move the neutralization curves to the first subplot axis
def move_artists(ax_from, ax_to):
    """Move plot content from ax_from to ax_to."""
    frame_parts = (
        *ax_from.spines.values(), ax_from.patch,
        ax_from.xaxis, ax_from.yaxis,
        ax_from.title, ax_from._left_title, ax_from._right_title,
    )

    # Collect artists to move (can't modify during iteration)
    artists_to_move = [art for art in ax_from.get_children() if art not in frame_parts]

    # Remove from source and add to destination
    for art in artists_to_move:
        art.remove()  # This properly cleans up the figure association
        # Update transform if needed
        if art.get_transform() == ax_from.transData:
            art.set_transform(ax_to.transData)
        # Add to destination
        ax_to.add_artist(art)

    # copy limits / labels, etc.
    ax_to.set_xscale(ax_from.get_xscale())
    ax_to.set_yscale(ax_from.get_yscale())
    ax_to.set_xlim(ax_from.get_xlim())
    ax_to.set_ylim(ax_from.get_ylim())
    ax_to.set_xlabel(ax_from.get_xlabel())
    ax_to.set_ylabel(ax_from.get_ylabel())

move_artists(mxra8_curve_fits_axes[0][0], mxra8_validation_axes[0])
mxra8_validation_axes[0].set_xlabel("mouse Mxra8 concentration (ug/ml)")
mxra8_validation_axes[0].set_ylabel("fraction infectivity")
mxra8_validation_axes[0].set_yticks([0, 0.5, 1])

# final formatting for figure
_ = mxra8_validation_fig.suptitle(
    "validations for mouse Mxra8 binding",
    fontsize=13,
)

print(f"Saving to {params['mxra8_validation_svg']=}")
mxra8_validation_fig.savefig(params["mxra8_validation_svg"], bbox_inches="tight")

Saving to params['mxra8_validation_svg']='results/paper_figures/mxra8_validation.svg'

Distribution of variant functional scores

We want the distribution of variant functional scores, similar to as made by this notebook but including both E3-E2 and 6K-E1 fragments and not including deletions (since those are rare in our libraries).

First read all the functional scores, ignoring those for deletions since those are rare and not reported in paper:

def classify_selection(sel):
    sels = {"293T-Mxra8": "293T-Mxra8", "293T-TIM1":"293T-TIM1", "C636":"C6/36"}
    assert sum(s in sel for s in sels) == 1, sel
    for s in sels:
        if s in sel:
            label = [sels[s]]
    libs = {"-A-": "library A", "-B-": "library B"}
    assert sum(l in sel for l in libs) == 1, sel
    for l in libs:
        if l in sel:
            label.append(libs[l])
    return " ".join(label)


func_scores_df = (
    pd.concat(
        [
            pd.read_csv(f).assign(selection=sel)
            for (sel, f) in params.items() if sel.startswith("func_scores_")
        ]
    )
    .assign(selection=lambda x: x["selection"].map(classify_selection))
    .pipe(dms_variants.codonvarianttable.CodonVariantTable.classifyVariants)
    .query("variant_class != 'deletion'")
)

(
    func_scores_df
    .groupby(["selection", "variant_class"])
    .aggregate(n_variants=pd.NamedAgg("barcode", "count"))
)

		n_variants
selection	variant_class
293T-Mxra8 library A	1 nonsynonymous	142665
	>1 nonsynonymous	79744
	stop	5778
	synonymous	4245
	wildtype	29406
293T-Mxra8 library B	1 nonsynonymous	132342
	>1 nonsynonymous	73868
	stop	5855
	synonymous	4048
	wildtype	27337
293T-TIM1 library A	1 nonsynonymous	142805
	>1 nonsynonymous	79859
	stop	5789
	synonymous	4237
	wildtype	29411
293T-TIM1 library B	1 nonsynonymous	132443
	>1 nonsynonymous	73929
	stop	5863
	synonymous	4048
	wildtype	27349
C6/36 library A	1 nonsynonymous	140546
	>1 nonsynonymous	78516
	stop	5670
	synonymous	4153
	wildtype	28913
C6/36 library B	1 nonsynonymous	131133
	>1 nonsynonymous	73139
	stop	5797
	synonymous	3999
	wildtype	27056

Make the plot:

def ridgeplot(df):
    variant_classes = list(
        reversed(
            [
                c
                for c in [
                    "wildtype",
                    "synonymous",
                    "1 nonsynonymous",
                    ">1 nonsynonymous",
                    "deletion",
                    "stop",
                ]
                if c in set(df["variant_class"])
            ]
        )
    )

    assert set(df["variant_class"]) == set(variant_classes)

    # get smoothed distribution of functional scores
    bins = numpy.linspace(
        df["func_score"].min(),
        df["func_score"].max(),
        num=50,
    )
    smoothed_dist = pd.concat(
        [
            pd.DataFrame(
                {
                    "selection": sel,
                    "variant_class": var,
                    "func_score": bins,
                    "count": scipy.stats.gaussian_kde(df["func_score"])(bins),
                    "mean_func_score": df["func_score"].mean(),
                    "number of variants": len(df),
                }
            )
            for (sel, var), df in df.groupby(["selection", "variant_class"])
        ]
    )

    # assign y / y2 for plotting
    facet_overlap = 0.7  # maximal facet overlap
    max_count = (smoothed_dist["count"]).max()
    smoothed_dist = smoothed_dist.assign(
        y=lambda x: x["variant_class"].map(lambda v: variant_classes.index(v)),
        y2=lambda x: x["y"] + x["count"] / max_count / facet_overlap,
    )

    # ridgeline plot, based on this but using y / y2 rather than row:
    # https://altair-viz.github.io/gallery/ridgeline_plot.html
    ridgeline_chart = (
        alt.Chart(smoothed_dist)
        .encode(
            x=alt.X(
                "func_score", title="functional score for cell entry", scale=alt.Scale(nice=False)
            ),
            y=alt.Y(
                "y",
                scale=alt.Scale(nice=False),
                title=None,
                axis=alt.Axis(
                    ticks=False,
                    domain=False,
                    # set manual labels https://stackoverflow.com/a/64106056
                    values=[v + 0.5 for v in range(len(variant_classes))],
                    labelExpr=f"{str(variant_classes)}[round(datum.value - 0.5)]",
                ),
            ),
            y2=alt.Y2("y2"),
            fill=alt.Fill(
                "mean_func_score:Q",
                title="mean functional score",
                legend=alt.Legend(direction="horizontal"),
                scale=alt.Scale(scheme="yellowgreenblue"),
            ),
            facet=alt.Facet(
                "selection",
                columns=2,
                title=None,
                header=alt.Header(
                    labelFontWeight="bold",
                    labelPadding=0,
                ),
            ),
            tooltip=[
                "selection",
                "variant_class",
                alt.Tooltip(
                    "mean_func_score", format=".2f", title="mean functional score"
                ),
            ],
        )
        .mark_area(
            interpolate="monotone",
            smooth=True,
            fillOpacity=0.8,
            stroke="lightgray",
            strokeWidth=0.5,
        )
        .configure_view(stroke=None)
        .configure_axis(grid=False)
        .properties(width=180, height=22 * len(variant_classes))
    )

    ridgeline_chart = ridgeline_chart.properties(
        autosize=alt.AutoSizeParams(resize=True),
    )

    return ridgeline_chart


func_scores_chart = ridgeplot(func_scores_df)

func_scores_chart

Number of mutations per variant

Plot number of mutations per variant as here but with better axis limits and labels, only keeping the combined libraries:

codon_variants = pd.read_csv(params["codon_variants"])

lib_rename = {"E_A": "library A", "E_B": "library B"}

display(
    codon_variants
    .groupby("library")
    .aggregate(n_variants=pd.NamedAgg("barcode", "count"))
)

max_muts = 4

print(f"Only keeping {lib_rename=}, and clipping at {max_muts=}")

nmuts_dist = (
    codon_variants
    .query("library in @lib_rename")
    .assign(
        library=lambda x: x["library"].map(lib_rename),
        n_muts=lambda x: x["n_aa_substitutions"].clip(upper=max_muts),
    )
    .groupby(["library", "n_muts"], as_index=False)
    .aggregate(n_variants=pd.NamedAgg("barcode", "count"))
    .assign(
        n_muts_label=lambda x: x["n_muts"].map(
            lambda n: str(n) if n < max_muts else f">{n - 1}"
        )
    )
)

nmuts_dist_chart = (
    alt.Chart(nmuts_dist)
    .encode(
        alt.X(
            "n_muts_label",
            sort=alt.SortField("n_muts"),
            title="number amino-acid mutations",
            axis=alt.Axis(labelAngle=0),
        ),
        alt.Y("n_variants", title="number of barcoded variants"),
        alt.Column(
            "library",
            title=None,
            header=alt.Header(labelFontStyle="bold", labelFontSize=11, labelPadding=0),
        ),
    )
    .mark_bar(color="black")
    .configure_axis(grid=False)
    .properties(height=150, width=150)
)

nmuts_dist_chart

	n_variants
library
6KE1_A	69359
6KE1_B	62446
E3E2_A	65426
E3E2_B	62495
E_A	134757
E_B	124915

Only keeping lib_rename={'E_A': 'library A', 'E_B': 'library B'}, and clipping at max_muts=4

Plot correlation among estimated functional effects on each cell

Plot similar to here but aggregating across libraries:

func_effects_by_lib = (
    pd.concat(
        [
            pd.read_csv(f).assign(
                name=name,
                cell=name.split("_")[2],
                replicate=name.split("-")[-1],
                library="library A" if "-A-" in name else "library B",
                region="E3E2" if "E3E2" in name else "6KE1",
            )
            for (name, f) in params.items() if name.startswith("func_effects_")
        ],
        ignore_index=True,
    )
    .query("wildtype != mutant")
    .query("times_seen >= @min_times_seen")
    .assign(
        mut_in_region=lambda x: x.apply(
            lambda r: (
                ("6K" in r["site"] or "E1" in r["site"]) and (r["region"] == "6KE1")
                or ("E2" in r["site"] or "E2" in r["site"]) and (r["region"] == "E3E2")
            ),
            axis=1,
        ),
    )
    .query("mut_in_region")
    .assign(
        library_replicate=lambda x: x["library"] + ", replicate " + x["replicate"],
        cell=lambda x: x["cell"].map(
            {
                "C636": "entry in C6/36 cells",
                "293T-Mxra8": "entry in 293T-Mxra8 cells",
                "293T-TIM1": "entry in 293T-TIM1 cells",
            }
        ),
    )
    [["site", "wildtype", "mutant", "functional_effect", "library_replicate", "cell"]]
)

func_effects_by_lib

	site	wildtype	mutant	functional_effect	library_replicate	cell
1346	1(E2)	S	-	-0.9849	library A, replicate 1	entry in 293T-Mxra8 cells
1347	1(E2)	S	A	0.2444	library A, replicate 1	entry in 293T-Mxra8 cells
1348	1(E2)	S	C	-1.9140	library A, replicate 1	entry in 293T-Mxra8 cells
1349	1(E2)	S	D	-0.1799	library A, replicate 1	entry in 293T-Mxra8 cells
1350	1(E2)	S	E	-0.7129	library A, replicate 1	entry in 293T-Mxra8 cells
...	...	...	...	...	...	...
306391	439(E1)	H	Y	-0.3790	library A, replicate 2	entry in C6/36 cells
306393	440(E1)	*	K	0.0617	library A, replicate 2	entry in C6/36 cells
306394	440(E1)	*	Q	-2.4410	library A, replicate 2	entry in C6/36 cells
306395	440(E1)	*	W	-0.4964	library A, replicate 2	entry in C6/36 cells
306396	440(E1)	*	Y	-0.2349	library A, replicate 2	entry in C6/36 cells

210846 rows × 6 columns

Now make the plot:

for cell, cell_df in func_effects_by_lib.groupby("cell"):
    corr_panels = []
    for sel1, sel2 in itertools.combinations(sorted(cell_df["library_replicate"].unique()), 2):
        corr_df = (
            cell_df.query("library_replicate == @sel1")[["functional_effect", "site", "mutant"]]
            .rename(columns={"functional_effect": sel1})
            .merge(
                cell_df.query("library_replicate == @sel2")[["functional_effect", "site", "mutant"]].rename(
                    columns={"functional_effect": sel2}
                ),
                validate="one_to_one",
            )
            .drop(columns=["site", "mutant"])
        )
        n = len(corr_df)
        r = corr_df[[sel1, sel2]].corr().values[1, 0]
        corr_panels.append(
            alt.Chart(corr_df)
            .encode(
                alt.X(sel1, scale=alt.Scale(nice=False, padding=4)),
                alt.Y(sel2, scale=alt.Scale(nice=False, padding=4)),
            )
            .mark_circle(color="black", size=25, opacity=0.15)
            .properties(
                width=135,
                height=135,
                title=alt.TitleParams(
                    f"R = {r:.2f}, N = {n}", fontSize=11, fontWeight="normal", dy=2
                ),
            )
        )
    
    corr_chart = alt.hconcat(*corr_panels, spacing=5).configure_axis(grid=False).properties(
        title=alt.TitleParams(f"correlations among libraries and replicates for mutations effects on {cell}", anchor="middle")
    )
    display(corr_chart)

Summary of cell entry effects

First read the effects, apply a floor, and compute the site mean and number of effective amino acids at each site:

cell_entry_types = {
    "entry in 293T_Mxra8 cells": "293T-Mxra8",
    "entry in C636 cells": "C6/36",
    "entry in 293T_TIM1 cells": "293T-TIM1",
}

cell_entry_mut = (
    pd.read_csv(params["annotated_mut_summary"])
    .rename(columns=cell_entry_types)
    .melt(
        id_vars=["site", "sequential_site", "wildtype", "mutant", "region", "domain", "contacts"],
        value_vars=cell_entry_types.values(),
        var_name="cell",
        value_name="cell entry",
    )
    .assign(**{"cell entry": lambda x: x["cell entry"].clip(lower=cell_entry_clip_lower)})
)

cell_entry_site = (
    cell_entry_mut
    .assign(
        p_unnorm=lambda x: numpy.exp(x["cell entry"]),
        p=lambda x: x["p_unnorm"] / x.groupby(["cell", "site"])["p_unnorm"].transform("sum"),
        entry_no_wt=lambda x: x["cell entry"].where(x["wildtype"] != x["mutant"], pd.NA),
    )
    .groupby(
        ["cell", "site", "sequential_site", "wildtype", "region", "domain", "contacts"],
        dropna=False,
        as_index=False,
    )
    .aggregate(
        n_effective=pd.NamedAgg("p", lambda p: numpy.exp((-p * numpy.log(p)).sum())),
        site_mean_entry=pd.NamedAgg("entry_no_wt", lambda s: s.dropna().mean()),
    )
    .sort_values("sequential_site")
)

Plot various site statistics showing mutational tolerance. Plot as the mean effect of mutations and the number of effective amino acids assigning probability weights proportional to the exponential of the cell entry effect:

metrics = {"n_effective": "effective amino acids", "site_mean_entry": "mean cell entry effect"}

title_prefix = {
    "n_effective": "mutational tolerance at each site for entry in ", 
    "site_mean_entry": "effects of mutations at each site for entry in ",
}

width = 600

for cell, metric in itertools.product(cell_entry_site["cell"].unique(), metrics):

    df = cell_entry_site.query("cell == @cell")[["site", "sequential_site", "region", metric]]
    
    site_entry_lines = (
        alt.Chart(df)
        .encode(
            alt.X(
                "site",
                sort=alt.SortField("sequential_site"),
                axis=alt.Axis(values=cell_entry_site["site"].iloc[90::242], labelAngle=0),
            ),
            alt.Y(
                metric,
                title=metrics[metric],
                scale=alt.Scale(nice=False, padding=2),
                axis=alt.Axis(grid=False)
            ),
            alt.Color("region", scale=alt.Scale(range=["gray", "darkgreen", "darkblue", "teal"])),
            tooltip=df.columns.tolist(),
        )
        .mark_rect(opacity=1)
        .properties(height=110, width=width)
    )

    text_df = df.groupby("region", as_index=False).aggregate(x=pd.NamedAgg("sequential_site", "mean"))

    text_chart = (
        alt.Chart(text_df)
        .encode(
            alt.X(
                "x:Q",
                title=None,
                scale=alt.Scale(domain=(df["sequential_site"].min(), df["sequential_site"].max())),
                axis=None,
            ),
            alt.Text("region"),
            alt.Color("region", legend=None),
        )
        .mark_text(fontWeight="bold", fontSize=12)
        .properties(width=width, height=1)
    )

    site_entry_chart = (
        alt.vconcat(text_chart, site_entry_lines, spacing=0)
        .configure_view(stroke=None)
        .properties(title=alt.TitleParams(title_prefix[metric] + cell + " cells", anchor="middle", fontWeight="normal"))
    )

    display(site_entry_chart)

Plot showing distribution of mean-mutation effects at each site. Do this for different proteins (E3, E2, 6K, E1) and different subdomains, as well as different types of Mxra8 contact sites:

for domain_col, cell in itertools.product(
    ["domain", "contacts", "region"],
    cell_entry_site["cell"].unique(),
):
    df = (
        cell_entry_site
        [cell_entry_site[domain_col].notnull()]
        .query("cell == @cell")
        [["site_mean_entry", "sequential_site", domain_col]]
    )

    if domain_col == "contacts":
        df = df.query("contacts != 'no'")

    dist_chart = (
        alt.Chart(df)
        .encode(
            alt.X(
                "site_mean_entry",
                axis=alt.Axis(titleFontWeight="normal", titleFontSize=12),
                scale=alt.Scale(nice=False, padding=4),
                title="mean effect of mutations at each site",
            ),
            alt.Y(
                domain_col,
                sort=df[domain_col].unique(),
                title=None,
                axis=alt.Axis(labelFontSize=12),
            ),
        )
        .mark_boxplot(
            outliers=False,
            box=alt.MarkConfig(fill="darkgray", stroke="black"),
            median=alt.MarkConfig(color="black"),
            rule=alt.MarkConfig(color='black'),
            size=13,
        )
        .properties(
            width=170,
            height=alt.Step(17),
            title=f"entry in {cell} cells"
        )
        .configure_axis(grid=False)
    )
    
    display(dist_chart)

Plot the same for the effect of mutations on mouse Mxra8 binding at each type of Mxra8 contact site:

mxra8_binding_contacts = (
    pd.read_csv(params["annotated_mut_summary"])
    .query("wildtype != mutant")
    .dropna(subset=["binding to mouse Mxra8", "contacts"])
    .query("contacts != 'no'")
    .groupby(["site", "sequential_site", "contacts"], as_index=False)
    .aggregate(site_mean_binding=pd.NamedAgg("binding to mouse Mxra8", "mean"))
)

contacts_order = ["intraspike", "wrapped", "interspike"]

mxra8_binding_contacts_chart = (
    alt.Chart(mxra8_binding_contacts)
    .encode(
        alt.X(
            "site_mean_binding",
            axis=alt.Axis(titleFontWeight="normal", titleFontSize=12),
            scale=alt.Scale(nice=False, padding=4),
            title="mean effect of mutations at each site",
        ),
        alt.Y("contacts", sort=contacts_order, title=None, axis=alt.Axis(labelFontSize=12)),
    )
    .mark_boxplot(
        outliers=False,
        box=alt.MarkConfig(fill="darkgray", stroke="black"),
        median=alt.MarkConfig(color="black"),
        rule=alt.MarkConfig(color="black"),
        size=13,
    )
    .properties(width=170, height=alt.Step(17), title="binding to mouse Mxra8")
    .configure_axis(grid=False)
)

mxra8_binding_contacts_chart

Statistically compare how much the cell-entry effects diverge among the three cell lines (293T-Mxra8, 293T-TIM1, C6/36) for each protein (E3, E2, 6K, E1). For each site, compute the variance in site_mean_entry across the three cell lines, then compare the distribution of this per-site divergence metric between proteins:

cell_cols = list(cell_entry_site["cell"].unique())

cell_entry_divergence = (
    cell_entry_site
    .pivot_table(
        index=["site", "sequential_site", "region"],
        columns="cell",
        values="site_mean_entry",
    )
    .reset_index()
    .dropna(subset=cell_cols)
)

cell_entry_divergence["divergence"] = cell_entry_divergence[cell_cols].var(axis=1, ddof=1)

regions = cell_entry_divergence.sort_values("sequential_site")["region"].unique().tolist()

print("Number of sites with divergence measurements, by protein:")
print(cell_entry_divergence["region"].value_counts())

divergence_by_region = {
    region: cell_entry_divergence.query("region == @region")["divergence"]
    for region in regions
}

h, p_kw = scipy.stats.kruskal(*divergence_by_region.values())
print(f"\nKruskal-Wallis test across all {len(regions)} proteins: H = {h:.2f}, P = {p_kw:.2e}")

print("\nE2 versus each other protein (one-sided Mann-Whitney U, alternative='greater'):")
for region in regions:
    if region == "E2":
        continue
    u, p = scipy.stats.mannwhitneyu(
        divergence_by_region["E2"], divergence_by_region[region], alternative="greater",
    )
    print(f"  E2 vs {region}: U = {u:.1f}, P = {p:.2e}")

Number of sites with divergence measurements, by protein:
region
E1    439
E2    423
E3     65
6K     61
Name: count, dtype: int64

Kruskal-Wallis test across all 4 proteins: H = 113.74, P = 1.72e-24

E2 versus each other protein (one-sided Mann-Whitney U, alternative='greater'):
  E2 vs E3: U = 17892.0, P = 4.52e-05
  E2 vs 6K: U = 21081.0, P = 5.77e-16
  E2 vs E1: U = 123230.0, P = 4.63e-17

divergence_chart = (
    alt.Chart(cell_entry_divergence[["divergence", "region"]])
    .encode(
        alt.X(
            "divergence",
            axis=alt.Axis(titleFontWeight="normal", titleFontSize=12),
            scale=alt.Scale(nice=False, padding=4),
            title="variance in site-mean entry effect across cell lines",
        ),
        alt.Y(
            "region",
            sort=regions,
            title=None,
            axis=alt.Axis(labelFontSize=12),
        ),
    )
    .mark_boxplot(
        outliers=False,
        box=alt.MarkConfig(fill="darkgray", stroke="black"),
        median=alt.MarkConfig(color="black"),
        rule=alt.MarkConfig(color='black'),
        size=13,
    )
    .properties(
        width=170,
        height=alt.Step(17),
        title="variance among cell lines in entry effects, by protein",
    )
    .configure_axis(grid=False)
)

divergence_chart

Compare DMS constraint to natural sequence constraint

We want to put the DMS results in the context of natural evolution.

For each site, we compare two site-level entropy measures already computed by the separate manual_analyses/compare_experimental_natural pipeline: the Shannon entropy of the amino acid observed at that site in an alignment of natural CHIKV sequences (alignment_entropy; low = conserved in nature), and the Shannon entropy of the DMS-derived amino-acid preferences at that site for a given entry phenotype (preference_entropy; low = constrained in the experiment). In the plots below these are labeled in plain language as variability in nature and mutational tolerance in DMS; the exact quantities are as defined here, and should be spelled out in the methods or figure legend. Note that the two have the same units but very different scales, so they are not directly comparable to each other. We show the two entropy distributions first by protein, then by finer-grained structural domain:

natural_summary_phenotypes = {
    "293T-Mxra8": params["natural_summary_293T_Mxra8"],
    "293T-TIM1": params["natural_summary_293T_TIM1"],
    "C6/36": params["natural_summary_C636"],
}

site_domains = pd.read_csv(params["addtl_site_annotations"])[["sequential_site", "domain"]]

# the fine-grained annotation has 18 domains, too many to read in a boxplot; group them by
# protein, splitting out only the E2 cytoplasmic tail
domain_groups = {
    "E3": "E3",
    "E2-n-linker": "E2-non-cytoplasmic",
    "E2-arch1": "E2-non-cytoplasmic",
    "E2-arch2": "E2-non-cytoplasmic",
    "E2-domain a": "E2-non-cytoplasmic",
    "E2-domain b": "E2-non-cytoplasmic",
    "E2-domain c": "E2-non-cytoplasmic",
    "E2-stem": "E2-non-cytoplasmic",
    "E2-transmembrane": "E2-non-cytoplasmic",
    "E2-cytoplasmic": "E2-cytoplasmic",
    "6K": "6K",
    "E1-domain I": "E1",
    "E1-domain II": "E1",
    "E1-domain III": "E1",
    "E1-fusion loop": "E1",
    "E1-stem": "E1",
    "E1-transmembrane": "E1",
    "E1-cytoplasmic": "E1",
}

natural_vs_experimental = (
    pd.concat(
        [
            pd.read_csv(f).assign(cell=cell)
            for (cell, f) in natural_summary_phenotypes.items()
        ],
        ignore_index=True,
    )
    .merge(site_domains, on="sequential_site", validate="many_to_one")
    .dropna(subset=["alignment_entropy", "preference_entropy", "domain"])
)

unmapped_domains = set(natural_vs_experimental["domain"]) - set(domain_groups)
assert not unmapped_domains, f"domains missing from `domain_groups`: {unmapped_domains}"
natural_vs_experimental["domain_group"] = natural_vs_experimental["domain"].map(domain_groups)

print("Number of sites compared, by grouped domain:")
print(natural_vs_experimental.groupby("cell")["domain_group"].value_counts())

print("\nOverall (all domains, all cells) natural versus DMS entropy:")
print(
    natural_vs_experimental[["alignment_entropy", "preference_entropy"]]
    .median()
    .rename("median entropy (nats)")
)

Number of sites compared, by grouped domain:
cell        domain_group      
293T-Mxra8  E1                    428
            E2-non-cytoplasmic    391
            E3                     64
            6K                     61
            E2-cytoplasmic         32
293T-TIM1   E1                    428
            E2-non-cytoplasmic    391
            E3                     64
            6K                     61
            E2-cytoplasmic         32
C6/36       E1                    428
            E2-non-cytoplasmic    391
            E3                     64
            6K                     61
            E2-cytoplasmic         32
Name: count, dtype: int64

Overall (all domains, all cells) natural versus DMS entropy:
alignment_entropy     0.00316
preference_entropy    2.39000
Name: median entropy (nats), dtype: float64

alignment_entropy does not depend on cell line/phenotype, so we show it once, alongside the phenotype-specific preference_entropy for each cell line. First by protein (E3, E2, 6K, E1):

# colors match `manual_analyses/compare_experimental_natural/scripts/comparison_plots.py`
C_ALN = "#1f77b4"  # natural alignment entropy
C_PREF = "#d62728"  # DMS preference entropy

# plain-language axis titles; the underlying quantity (site-level Shannon entropy of the
# amino-acid distribution) is defined in the methods / figure legend
ENTROPY_AXIS_TITLES = {
    "alignment_entropy": "variability in nature",
    "preference_entropy": "mutational tolerance in DMS",
}

# alignment_entropy is identical for every cell (it doesn't depend on phenotype), so just
# take it from one cell rather than showing three identical copies
first_cell = natural_vs_experimental["cell"].unique()[0]


def entropy_by_group_chart(group_col, group_order, box_size, row_step):
    """Natural-sequence entropy alongside per-cell-line DMS entropy, split by `group_col`."""
    natural_box = (
        alt.Chart(
            natural_vs_experimental.query("cell == @first_cell")[[group_col, "alignment_entropy"]]
        )
        .encode(
            alt.X(
                "alignment_entropy",
                title=ENTROPY_AXIS_TITLES["alignment_entropy"],
                scale=alt.Scale(nice=False, padding=4),
            ),
            alt.Y(group_col, sort=group_order, title=None, axis=alt.Axis(labelFontSize=11)),
        )
        .mark_boxplot(
            outliers=False,
            box=alt.MarkConfig(fill=C_ALN, stroke="black"),
            median=alt.MarkConfig(color="black"),
            rule=alt.MarkConfig(color="black"),
            size=box_size,
        )
        .properties(
            width=160,
            height=alt.Step(row_step),
            title=alt.TitleParams(
                "natural sequences", fontSize=13, fontWeight="normal", anchor="middle"
            ),
        )
    )

    dms_boxes = []
    for cell, cell_df in natural_vs_experimental.groupby("cell"):
        dms_boxes.append(
            alt.Chart(cell_df[[group_col, "preference_entropy"]])
            .encode(
                alt.X(
                    "preference_entropy",
                    title=ENTROPY_AXIS_TITLES["preference_entropy"],
                    scale=alt.Scale(nice=False, padding=4),
                ),
                alt.Y(group_col, sort=group_order, title=None, axis=None),
            )
            .mark_boxplot(
                outliers=False,
                box=alt.MarkConfig(fill=C_PREF, stroke="black"),
                median=alt.MarkConfig(color="black"),
                rule=alt.MarkConfig(color="black"),
                size=box_size,
            )
            .properties(
                width=160,
                height=alt.Step(row_step),
                title=alt.TitleParams(cell, fontSize=13, fontWeight="normal", anchor="middle"),
            )
        )

    return (
        alt.hconcat(natural_box, *dms_boxes, spacing=5)
        .configure_axis(grid=False, titleFontWeight="normal")
    )


regions_order = natural_vs_experimental.sort_values("sequential_site")["region"].unique().tolist()

print("Number of sites compared, by protein:")
print(natural_vs_experimental.query("cell == @first_cell")["region"].value_counts().to_string())

entropy_region_chart = entropy_by_group_chart("region", regions_order, box_size=12, row_step=22)

entropy_region_chart

Number of sites compared, by protein:
region
E1    428
E2    423
E3     64
6K     61

Now the same comparison with the E2 cytoplasmic tail separated from the rest of E2:

# order the grouped domains by the most N-terminal site each contains
domain_groups_order = (
    natural_vs_experimental
    .groupby("domain_group")["sequential_site"]
    .min()
    .sort_values()
    .index
    .tolist()
)

entropy_dist_chart = entropy_by_group_chart(
    "domain_group", domain_groups_order, box_size=11, row_step=20
)

entropy_dist_chart

Effects of mutations observed in natural sequences vs. all possible mutations

As another way to relate the DMS measurements to natural evolution, we identify amino-acid mutations actually observed in the natural CHIKV sequence alignment (relative to the 181/25 DMS strain), and compare the distribution of their cell-entry effects to the distribution for all possible measured mutations.

First find the natural mutations by comparing each sequence in the alignment to the DMS strain:

IGNORED_AA = {"-", "X", "*"}  # gaps, ambiguous residues, stop codons

records = {r.id: str(r.seq) for r in SeqIO.parse(params["protein_alignment"], "fasta")}
ref_seq = records[ref_strain]
other_seqs = [seq for rid, seq in records.items() if rid != ref_strain]
print(f"Alignment: {len(records)} sequences (including reference {ref_strain}), length {len(ref_seq)}")

# map alignment column (0-based) to sequential_site; must match
# manual_analyses/compare_experimental_natural/scripts/verify_and_map.py, where
# alignment_column = capsid_len + 1 + i and sequential_site = 2 + i (1-based columns)
natural_muts = []
for col_idx in range(capsid_len, len(ref_seq)):
    sequential_site = col_idx - capsid_len + 2
    ref_aa = ref_seq[col_idx]
    if ref_aa in IGNORED_AA:
        continue
    counts = Counter(
        seq[col_idx] for seq in other_seqs
        if seq[col_idx] != ref_aa and seq[col_idx] not in IGNORED_AA
    )
    for mut_aa, count in counts.items():
        natural_muts.append((sequential_site, ref_aa, mut_aa, count))

natural_muts_df = pd.DataFrame(natural_muts, columns=["sequential_site", "wildtype", "mutant", "mutation_count"])

for thresh in [1, 2, 3, 4, 5]:
    sub = natural_muts_df.query("mutation_count >= @thresh")
    print(
        f"  min_mutation_count >= {thresh}: {len(sub)} distinct natural mutations "
        f"at {sub['sequential_site'].nunique()} sites"
    )

min_mutation_count = 2
natural_muts_df = natural_muts_df.query("mutation_count >= @min_mutation_count")
print(
    f"\nUsing min_mutation_count={min_mutation_count}: {len(natural_muts_df)} distinct natural mutations "
    f"at {natural_muts_df['sequential_site'].nunique()} sites"
)

Alignment: 2833 sequences (including reference 181-25_MW473668), length 1248
  min_mutation_count >= 1: 894 distinct natural mutations at 554 sites
  min_mutation_count >= 2: 393 distinct natural mutations at 308 sites
  min_mutation_count >= 3: 281 distinct natural mutations at 230 sites
  min_mutation_count >= 4: 236 distinct natural mutations at 199 sites
  min_mutation_count >= 5: 206 distinct natural mutations at 177 sites

Using min_mutation_count=2: 393 distinct natural mutations at 308 sites

Merge the natural mutations with the DMS cell-entry effects already computed above (cell_entry_mut, floored at cell_entry_clip_lower), and compare to the distribution for all possible measured mutations:

all_possible_mutations = cell_entry_mut.query("wildtype != mutant")

natural_mut_effects = natural_muts_df.merge(
    cell_entry_mut[["sequential_site", "wildtype", "mutant", "cell", "cell entry"]],
    on=["sequential_site", "wildtype", "mutant"],
    how="inner",
    validate="one_to_many",  # natural_muts_df has 1 row per mutation, cell_entry_mut has 1 per (mutation, cell)
)

for cell in all_possible_mutations["cell"].unique():
    n_all = (all_possible_mutations["cell"] == cell).sum()
    n_nat = (natural_mut_effects["cell"] == cell).sum()
    print(f"{cell}: {n_all} total possible mutations, {n_nat} natural mutations with DMS data")

293T-Mxra8: 18474 total possible mutations, 388 natural mutations with DMS data
C6/36: 18474 total possible mutations, 388 natural mutations with DMS data
293T-TIM1: 18474 total possible mutations, 388 natural mutations with DMS data

ALL_MUTS_LABEL = "all possible mutations"
NATURAL_MUTS_LABEL = "mutations in natural sequences"
C_ALL_MUTS = "#1f77b4"
C_NATURAL_MUTS = "#2ca02c"

mutation_sets = pd.concat(
    [
        all_possible_mutations[["cell", "cell entry"]].assign(mutation_set=ALL_MUTS_LABEL),
        natural_mut_effects[["cell", "cell entry"]].assign(mutation_set=NATURAL_MUTS_LABEL),
    ],
    ignore_index=True,
).dropna(subset=["cell entry"])

# group sizes differ slightly between cell lines (a few mutations lack a measurement in one
# cell line), so they are reported here rather than put in a shared axis label
for cell in cell_entry_types.values():
    cell_muts = mutation_sets.query("cell == @cell")
    all_effects = cell_muts.query("mutation_set == @ALL_MUTS_LABEL")["cell entry"]
    natural_effects = cell_muts.query("mutation_set == @NATURAL_MUTS_LABEL")["cell entry"]
    u, p = scipy.stats.mannwhitneyu(natural_effects, all_effects, alternative="greater")
    print(
        f"{cell}: all possible n = {len(all_effects)}, median = {all_effects.median():.2f}; "
        f"natural n = {len(natural_effects)}, median = {natural_effects.median():.2f}; "
        f"one-sided Mann-Whitney U = {u:.1f}, P = {p:.2e}"
    )

natural_mut_dist_chart = (
    alt.Chart(mutation_sets)
    .encode(
        alt.X("cell entry:Q", title="cell entry effect", scale=alt.Scale(nice=False, padding=4)),
        alt.Y(
            "mutation_set:N",
            sort=[ALL_MUTS_LABEL, NATURAL_MUTS_LABEL],
            title=None,
            axis=alt.Axis(labelFontSize=11, labelLimit=250),
        ),
        alt.Color(
            "mutation_set:N",
            scale=alt.Scale(
                domain=[ALL_MUTS_LABEL, NATURAL_MUTS_LABEL],
                range=[C_ALL_MUTS, C_NATURAL_MUTS],
            ),
            legend=None,  # each box is labeled directly on the y axis
        ),
    )
    # 1.5 x IQR whiskers with outliers hidden; the lower quartile of the all-possible set sits
    # on the `cell_entry_clip_lower` floor, so those boxes have no left whisker
    .mark_boxplot(
        outliers=False,
        median=alt.MarkConfig(color="black"),
        rule=alt.MarkConfig(color="black"),
        size=16,
    )
    .properties(width=230, height=alt.Step(24))
    .facet(
        row=alt.Facet(
            "cell:N",
            sort=list(cell_entry_types.values()),
            title=None,
            header=alt.Header(
                labelFontSize=12,
                labelFontWeight="normal",
                labelOrient="top",
                labelAnchor="start",
                labelPadding=1,
            ),
        ),
        spacing=10,
    )
    .configure_axis(grid=False, labelFontSize=12, titleFontSize=13, titleFontWeight="normal")
    .configure_view(stroke=None)
)

natural_mut_dist_chart

293T-Mxra8: all possible n = 18466, median = -1.93; natural n = 388, median = -0.20; one-sided Mann-Whitney U = 5497911.0, P = 3.27e-75
C6/36: all possible n = 18439, median = -2.30; natural n = 388, median = -0.22; one-sided Mann-Whitney U = 5444412.5, P = 9.57e-72
293T-TIM1: all possible n = 18464, median = -1.39; natural n = 388, median = -0.14; one-sided Mann-Whitney U = 5482000.0, P = 2.72e-73

Occurrence in nature of mutations tolerated in one or both hosts

CHIKV alternates between mammalian and mosquito hosts, so mutations may be constrained by the requirements for entry into both. We classify every measured mutation by whether it is tolerated for entry in 293T-Mxra8 (mammalian) and C6/36 (mosquito) cells, then ask how often mutations of each class are observed in the natural sequence alignment (in at least min_mutation_count sequences, as above):

entry_tolerated_threshold = -1

mut_entry_by_cell = (
    cell_entry_mut
    .query("wildtype != mutant")
    .pivot_table(
        index=["sequential_site", "wildtype", "mutant"],
        columns="cell",
        values="cell entry",
    )
    .reset_index()
    .dropna(subset=["293T-Mxra8", "C6/36"])
)

# `natural_muts_df` is already restricted to mutations seen in >= min_mutation_count sequences
natural_mutation_set = set(
    zip(
        natural_muts_df["sequential_site"],
        natural_muts_df["wildtype"],
        natural_muts_df["mutant"],
    )
)
mut_entry_by_cell["observed_in_nature"] = [
    key in natural_mutation_set
    for key in zip(
        mut_entry_by_cell["sequential_site"],
        mut_entry_by_cell["wildtype"],
        mut_entry_by_cell["mutant"],
    )
]

mammal_tolerated = mut_entry_by_cell["293T-Mxra8"] >= entry_tolerated_threshold
mosquito_tolerated = mut_entry_by_cell["C6/36"] >= entry_tolerated_threshold

CLASS_BOTH = "tolerated in both"
CLASS_MAMMAL = "tolerated in 293T-Mxra8 only"
CLASS_MOSQUITO = "tolerated in C6/36 only"
CLASS_NEITHER = "tolerated in neither"
mutation_class_order = [CLASS_BOTH, CLASS_MAMMAL, CLASS_MOSQUITO, CLASS_NEITHER]

mut_entry_by_cell["mutation_class"] = numpy.select(
    [
        mammal_tolerated & mosquito_tolerated,
        mammal_tolerated & ~mosquito_tolerated,
        ~mammal_tolerated & mosquito_tolerated,
    ],
    [CLASS_BOTH, CLASS_MAMMAL, CLASS_MOSQUITO],
    default=CLASS_NEITHER,
)

natural_by_class = []
for mutation_class in mutation_class_order:
    class_muts = mut_entry_by_cell.query("mutation_class == @mutation_class")
    n_observed = int(class_muts["observed_in_nature"].sum())
    n_mutations = len(class_muts)
    ci = scipy.stats.binomtest(n_observed, n_mutations).proportion_ci(method="wilson")
    natural_by_class.append(
        {
            "mutation_class": mutation_class,
            "n_mutations": n_mutations,
            "n_observed": n_observed,
            "fraction_label": f"{n_observed}/{n_mutations}",
            "pct_observed": 100 * n_observed / n_mutations,
            "ci_lower": 100 * ci.low,
            "ci_upper": 100 * ci.high,
        }
    )
natural_by_class = pd.DataFrame(natural_by_class)

print(
    f"Mutations measured in both cell lines: {len(mut_entry_by_cell)}; "
    f"tolerated means cell entry >= {entry_tolerated_threshold}"
)
print(natural_by_class[["mutation_class", "n_mutations", "n_observed", "pct_observed"]].round(2).to_string(index=False))

print(f"\nFisher exact test of each class versus '{CLASS_BOTH}':")
both_observed = mut_entry_by_cell.query("mutation_class == @CLASS_BOTH")["observed_in_nature"]
for mutation_class in [CLASS_MAMMAL, CLASS_MOSQUITO, CLASS_NEITHER]:
    class_observed = mut_entry_by_cell.query("mutation_class == @mutation_class")["observed_in_nature"]
    odds, p = scipy.stats.fisher_exact(
        [
            [int(both_observed.sum()), int((~both_observed).sum())],
            [int(class_observed.sum()), int((~class_observed).sum())],
        ]
    )
    print(f"  {mutation_class}: odds ratio = {odds:.2f}, P = {p:.2e}")

Mutations measured in both cell lines: 18436; tolerated means cell entry >= -1
              mutation_class  n_mutations  n_observed  pct_observed
           tolerated in both         6272         301          4.80
tolerated in 293T-Mxra8 only         1042          29          2.78
     tolerated in C6/36 only          406           4          0.99
        tolerated in neither        10716          54          0.50

Fisher exact test of each class versus 'tolerated in both':
  tolerated in 293T-Mxra8 only: odds ratio = 1.76, P = 2.78e-03
  tolerated in C6/36 only: odds ratio = 5.07, P = 6.36e-05
  tolerated in neither: odds ratio = 9.95, P = 7.44e-79

natural_by_class_row_step = 26
natural_by_class_y = alt.Y(
    "mutation_class",
    sort=mutation_class_order,
    title=None,
    axis=alt.Axis(labelFontSize=11, labelLimit=250),
)

natural_by_class_bars = (
    alt.Chart(natural_by_class)
    .mark_bar(color="#4c78a8", size=18)
    .encode(
        alt.X(
            "pct_observed",
            title=f"mutations seen in ≥{min_mutation_count} natural sequences",
            scale=alt.Scale(nice=False, padding=4),
            # the axis title carries no unit, so every tick gets a percent sign
            axis=alt.Axis(labelExpr="datum.label + '%'"),
        ),
        natural_by_class_y,
    )
)

natural_by_class_ci = (
    alt.Chart(natural_by_class)
    .mark_rule(color="black")
    .encode(alt.X("ci_lower"), alt.X2("ci_upper"), natural_by_class_y)
)

# short perpendicular ticks at each end, so the interval does not read as a leader line
natural_by_class_caps = alt.layer(
    *[
        alt.Chart(natural_by_class)
        .mark_tick(color="black", thickness=1, size=9)
        .encode(alt.X(bound), natural_by_class_y)
        for bound in ["ci_lower", "ci_upper"]
    ]
)

# the classes differ in size by >25-fold, so a percentage alone hides how few mutations
# some of them contain; the fractions go in their own column to keep them visually
# separate from the confidence intervals
natural_by_class_labels = (
    alt.Chart(natural_by_class)
    .mark_text(align="left", fontSize=11)
    .encode(
        x=alt.value(0),
        y=alt.Y("mutation_class", sort=mutation_class_order, title=None, axis=None),
        text=alt.Text("fraction_label"),
    )
    .properties(width=62, height=alt.Step(natural_by_class_row_step))
)

host_tradeoff_chart = (
    alt.hconcat(
        (natural_by_class_bars + natural_by_class_ci + natural_by_class_caps).properties(
            width=260, height=alt.Step(natural_by_class_row_step)
        ),
        natural_by_class_labels,
        spacing=8,
    )
    .configure_axis(grid=False, titleFontWeight="normal")
    .configure_view(stroke=None)
)

host_tradeoff_chart

Validation assays of RVPs and CHIKV mutants with loss-of-function mutations for entry in a cell line

First plot reporter virus particle titers:

# color shades of blue if expected to be attenauted in Mxra8, orange if mosquitoe
viruses = {
    'unmutated': "black",
    'R119K': "#6A5ACD",
    'K120D': "#367588",
    'I121E': "#4682B4",
    'R119K-K120D': "#1E90FF",
    'R119K-I121E': "#4169E1",
    'K120D-I121E': "#545499",
    'A157S': "#FFB300",
    'Q158T': "#ED9121",
    'Q158V': "#FF7518",
    'A157S-Q158T': "#E9963A",
    'A157S-Q158V': "#CC5500",
}

# shapes by number of mutations
shapes = {}
for v in viruses:
    if v == "unmutated":
        shapes[v] = "square"
    elif "-" in v:
        shapes[v] = "diamond"
    else:
        shapes[v] = "circle"

rvp_titers = (
    pd.read_csv(params["rvp_titers"])
    .groupby(["cell", "virus"], as_index=False)
    .aggregate(
        titer=pd.NamedAgg("titer", "mean"),
        titer_sem=pd.NamedAgg("titer", "sem"),
    )
    .assign(
        virus=lambda x: x["virus"].str.replace("E2_", "").str.replace("_", "-"),
        titer_upper=lambda x: x["titer"] + x["titer_sem"],
        titer_lower=lambda x: x["titer"] - x["titer_sem"],
    )
)

assert set(rvp_titers["virus"]) == set(viruses)

cells = ["293T-TIM1", "293T-Mxra8", "Huh7", "HFF-1", "C6/36", "Aag2"]

rvp_titers_base = (
    alt.Chart(rvp_titers)
    .encode(
        alt.Y("virus", title=None, sort=list(viruses)),
        alt.Color(
            "virus",
            scale=alt.Scale(domain=list(viruses), range=list(viruses.values())),
            legend=alt.Legend(title="E2 mutations", titleFontSize=11)
        ),
    )
    .properties(width=200, height=alt.Step(10))
)

rvp_titers_points = (
    rvp_titers_base
    .encode(
        alt.X(
            "titer",
            title="titer relative to unmutated",
            scale=alt.Scale(type="log", nice=False, padding=6),
            axis=alt.Axis(values=[1, 0.1, 0.01, 0.001], grid=False, titleFontWeight="normal", titleFontSize=11)
        ),
        alt.Shape("virus", scale=alt.Scale(domain=list(shapes), range=list(shapes.values()))),
    )
    .mark_point(filled=True, fillOpacity=1, size=60)
)

rvp_titers_errorbars = (
    rvp_titers_base
    .encode(
        alt.X("titer_upper"),
        alt.X2("titer_lower"),
    )
    .mark_rule(size=1)
)

rvp_titers_chart = (
    (rvp_titers_errorbars + rvp_titers_points)
    .facet(
        row=alt.Facet(
            "cell",
            title=None,
            header=alt.Header(orient="right", labelPadding=2, labelFontSize=11, labelFontWeight="bold"),
            sort=cells,
        ),
        spacing=3,
    )
    .resolve_scale(y="independent")
    .properties(
        title=alt.TitleParams("alphavirus RVP titers", anchor="middle", dx=-30),
    )
)

rvp_titers_chart

Now plot CHIKV titers:

chikv_titers = (
    pd.read_csv(params["chikv_titers"])
    .groupby(["cell", "virus", "timepoint"], as_index=False)
    .aggregate(
        titer=pd.NamedAgg("titer", "mean"),
        titer_sem=pd.NamedAgg("titer", "sem"),
    )
    .assign(
        virus=lambda x: x["virus"].str.replace("E2_", "").str.replace("_", "-"),
        titer_upper=lambda x: x["titer"] + x["titer_sem"],
        # can't plot infinitely low error bars, so clip them
        titer_lower=lambda x: (x["titer"] - x["titer_sem"]).clip(lower=x["titer"].min() / 2),
    )
)

chikv_titers_base = (
    alt.Chart(chikv_titers)
    .encode(
        alt.X(
            "timepoint",
            title="hours post-infection",
            axis=alt.Axis(grid=False, titleFontWeight="normal", titleFontSize=11),
        ),
        alt.Color(
            "virus",
            scale=alt.Scale(domain=list(viruses), range=list(viruses.values())),
            legend=None,
        ),
    )
    .properties(width=200, height=134)
)

chikv_titers_points = (
    chikv_titers_base
    .encode(
        alt.Y(
            "titer",
            title="titer (RLU)",
            scale=alt.Scale(type="log", nice=False, padding=8),
            axis=alt.Axis(
                grid=False,
                titleFontWeight="normal",
                titleFontSize=11,
                format="~e",
                values=[1e1, 1e2, 1e3, 1e4, 1e5, 1e6, 1e7],
            ),
        ),
        alt.Shape("virus", scale=alt.Scale(domain=list(shapes), range=list(shapes.values()))),
    )
    .mark_point(filled=True, fillOpacity=1, size=60)
)

chikv_titers_lines = (
    chikv_titers_base
    .encode(
        alt.Y("titer", scale=alt.Scale(type="log"), axis=alt.Axis(grid=False)),
    )
    .mark_line(size=1)
)

chikv_titers_errorbars = (
    chikv_titers_base
    .encode(
        alt.Y("titer_lower", scale=alt.Scale(type="log"), axis=alt.Axis(grid=False)),
        alt.Y2("titer_upper"),
    )
    .mark_rule(size=1)
)

chikv_titers_chart = (
    (chikv_titers_lines + chikv_titers_errorbars + chikv_titers_points)
    .facet(
        row=alt.Facet(
            "cell",
            title=None,
            header=alt.Header(orient="right", labelPadding=2, labelFontSize=11, labelFontWeight="bold"),
            sort=cells,
        ),
        spacing=3,
    )
    .resolve_scale(y="independent")
    .properties(
        title=alt.TitleParams("authentic CHIKV titers", anchor="middle", dx=20),
    )
)

chikv_titers_chart

(
    alt.hconcat(rvp_titers_chart, chikv_titers_chart, spacing=30)
    .resolve_scale(color="independent", shape="independent")
    .configure_legend(offset=30, strokeColor="gray", padding=7)
    .configure_view(stroke="gray")
)

HFF-1 RVP titers

Plot the RVP titers in HFF-1 cells (human foreskin fibroblasts) separately, in the same style as the titer panels above:

hff1_titers = (
    pd.read_csv(params["rvp_titers"])
    .query("cell == 'HFF-1'")
    .groupby(["cell", "virus"], as_index=False)
    .aggregate(
        titer=pd.NamedAgg("titer", "mean"),
        titer_sem=pd.NamedAgg("titer", "sem"),
    )
    .assign(
        virus=lambda x: x["virus"].str.replace("E2_", "").str.replace("_", "-"),
        titer_upper=lambda x: x["titer"] + x["titer_sem"],
        titer_lower=lambda x: x["titer"] - x["titer_sem"],
    )
)

assert set(hff1_titers["virus"]) == set(viruses)

# reuse the `viruses` colors and `shapes` defined for the RVP titer panels above
hff1_titers_base = (
    alt.Chart(hff1_titers)
    .encode(
        alt.Y("virus", title=None, sort=list(viruses)),
        alt.Color(
            "virus",
            scale=alt.Scale(domain=list(viruses), range=list(viruses.values())),
            legend=alt.Legend(title="E2 mutations", titleFontSize=11),
        ),
    )
    .properties(width=200, height=alt.Step(14))
)

hff1_titers_points = (
    hff1_titers_base
    .encode(
        alt.X(
            "titer",
            title="titer relative to unmutated",
            scale=alt.Scale(type="log", nice=False, padding=6),
            axis=alt.Axis(values=[1, 0.1, 0.01, 0.001], grid=False, titleFontWeight="normal", titleFontSize=11),
        ),
        alt.Shape("virus", scale=alt.Scale(domain=list(shapes), range=list(shapes.values()))),
    )
    .mark_point(filled=True, fillOpacity=1, size=60)
)

hff1_titers_errorbars = (
    hff1_titers_base
    .encode(alt.X("titer_upper"), alt.X2("titer_lower"))
    .mark_rule(size=1)
)

hff1_titers_chart = (
    (hff1_titers_errorbars + hff1_titers_points)
    .properties(title=alt.TitleParams("HFF-1 RVP titers", anchor="middle"))
    .configure_axis(grid=False)
)

hff1_titers_chart