# Imports
import os
import pandas as pd
import altair as alt

# Plotting colors
# re-arranged for plot
tol_muted_adjusted = [
    "#AA4499",
    "#88CCEE",
    "#EE7733",
    "#44AA99",
    "#1f78b4",
    "#CC6677",
    "#117733",
    "#999933",
    "#DDCC77",
    "#CC3311",
    "#882255",
    "#000000",
    "#DDDDDD",
]

# Allow more rows for Altair
_ = alt.data_transformers.disable_max_rows()

# this cell is tagged as `parameters` for papermill parameterization
variant_data = None

out_dir = None
out_file = None

# Parameters
variant_data = "results/variants/codon_variants.csv"
out_dir = "results/summary_of_libraries/"
out_file = (
    "results/summary_of_libraries/present_and_absent_mutations_in_DMS_libraries.csv"
)

# # Uncomment for running interactive
# variant_data = "../results/variants/codon_variants.csv"

# out_dir = "../results/summary_of_libraries/"
# out_file = "../results/summary_of_libraries/present_and_absent_mutations_in_DMS_libraries.csv"

# Load data
variant_df = pd.read_csv(variant_data)

# Group all variants with >= 8 amino-acid mutations
variant_df["n_aa_substitutions"] = variant_df["n_aa_substitutions"].apply(lambda x: 8 if x >= 8 else x)

distribution_plot = alt.Chart(variant_df).mark_bar(color="#000000", size=7).encode(
    x=alt.X(
        "n_aa_substitutions", 
        axis=alt.Axis(
            title="AA muts",
            values=[0, 1, 2, 3, 4, 5, 6, 7, 8],
            domainWidth=1,
            domainColor="black",
            tickColor="black",
            labelFontSize=8,
            labelFontWeight="normal",
            titleFontWeight="normal",
        ),
    ),
    y=alt.Y(
        "count()",
        axis=alt.Axis(
            title="number of variants",
            values=[0, 5000, 10000, 15000],
            domainWidth=1,
            domainColor="black",
            tickColor="black",
            labelFontSize=8,
            labelFontWeight="normal",
            titleFontWeight="normal",
        ),
        scale=alt.Scale(domain=[0,15000]),
    ),
    facet=alt.Facet(
        "library", 
        title=None, 
        columns=2,
        header=alt.Header(
            labelFontSize=8,
            labelFontWeight="bold",
        ),
    ),
).properties(
    width=75, 
    height=100,
).configure_axis(
    grid=False,
    labelFontSize=8,
    titleFontSize=8,
    labelFontWeight="normal",
) 

distribution_plot

# Single mutant variant counts
libA_single_mutants = (
    variant_df.query("library == 'LibA' & n_aa_substitutions == 1")
    .reset_index(drop=True)
    ["aa_substitutions"].unique()
)
libA_single_mutants = len([x for x in libA_single_mutants if (x[-1] != "-" and x[1:-1] != "492")])
print(f"Number of amino-acid mutations measured in single mutant variants in library A: {libA_single_mutants} ({libA_single_mutants/9820 * 100:.1f} % possible mutations)")

libB_single_mutants = (
    variant_df.query("library == 'LibB' & n_aa_substitutions == 1")
    .reset_index(drop=True)
    ["aa_substitutions"].unique()
)
libB_single_mutants = len([x for x in libB_single_mutants if (x[-1] != "-" and x[1:-1] != "492")])
print(f"Number of amino-acid mutations measured in single mutant variants in library B: {libB_single_mutants} ({libB_single_mutants/9820 * 100:.1f} % possible mutations)")

# All mutant variant counts
libA_all_mutants = (
    variant_df.query("library == 'LibA' & n_aa_substitutions >= 1")
    .reset_index(drop=True)
)
libA_all_mutants["list_mutations"] = libA_all_mutants.apply(lambda x: x["aa_substitutions"].split(" "), axis=1)
libA_all_mutants = (
    libA_all_mutants.explode("list_mutations")
    ["list_mutations"].unique()
)
libA_all_mutants = len([x for x in libA_all_mutants if (x[-1] != "-" and x[1:-1] != "492")])
print(f"Number of amino-acid mutations measured in all barcoded variants in library A: {libA_all_mutants} ({libA_all_mutants/9820 * 100:.1f} % possible mutations)")

libB_all_mutants = (
    variant_df.query("library == 'LibB' & n_aa_substitutions >= 1")
    .reset_index(drop=True)
)
libB_all_mutants["list_mutations"] = libB_all_mutants.apply(lambda x: x["aa_substitutions"].split(" "), axis=1)
libB_all_mutants = (
    libB_all_mutants.explode("list_mutations")
    ["list_mutations"].unique()
)
libB_all_mutants = len([x for x in libB_all_mutants if (x[-1] != "-" and x[1:-1] != "492")])
print(f"Number of amino-acid mutations measured in all barcoded variants in library B: {libB_all_mutants} ({libB_all_mutants/9820 * 100:.1f} % possible mutations)")

Number of amino-acid mutations measured in single mutant variants in library A: 5737 (58.4 % possible mutations)
Number of amino-acid mutations measured in single mutant variants in library B: 6455 (65.7 % possible mutations)
Number of amino-acid mutations measured in all barcoded variants in library A: 9669 (98.5 % possible mutations)

Number of amino-acid mutations measured in all barcoded variants in library B: 9702 (98.8 % possible mutations)

# Create dataframe of all possible mutations
AAs = ['R','K','H','D','E','Q','N','S','T','Y','W','F','A','I','L','M','V','G','P','C','*'] * 491
positions = [v for v in list(range(1,492)) for _ in range(21)]
all_muts = pd.DataFrame({"site" : positions, "mutant" : AAs})

# Extract all mutations and filter deletions
variant_df = (
    variant_df.query("n_aa_substitutions >= 1")
    .reset_index(drop=True)
)
variant_df["list_mutations"] = variant_df.apply(lambda x: x["aa_substitutions"].split(" "), axis=1)
variant_df = (
    variant_df.explode("list_mutations")[["library", "list_mutations"]]
    .drop_duplicates()
    .reset_index(drop=True)
)
variant_df["wildtype"] = variant_df["list_mutations"].str[0]
variant_df["site"] = variant_df["list_mutations"].str[1:-1].astype(int)
variant_df["mutant"] = variant_df["list_mutations"].str[-1]
variant_df = (
    variant_df
    .query("mutant != '-'")
    .reset_index(drop=True)
    .rename(columns={"list_mutations" : "present_mutation"})
)

# Remove all wildtype mutations from df
possible_mutations = (
    all_muts.merge(
        variant_df[["site", "wildtype"]].drop_duplicates(),
        how="left",
        on=["site"],
        validate="many_to_many",
    )
    .query("mutant != wildtype")
    .reset_index(drop=True)
)
possible_mutations["possible_mutation"] = possible_mutations["wildtype"] + possible_mutations["site"].astype(str) + possible_mutations["mutant"]

# Add all mutations present in library A
possible_mutations = (
    possible_mutations.merge(
        variant_df.query("library == 'LibA'")[["site", "wildtype", "mutant", "present_mutation"]],
        how="left",
        on=["site", "wildtype", "mutant"],
        validate="one_to_one",
    )
    .rename(columns={"present_mutation" : "mutation_present_in_libA"})
)
possible_mutations["mutation_present_in_libA"] = possible_mutations["mutation_present_in_libA"].fillna(-1)
possible_mutations["mutation_present_in_libA"] = (
    possible_mutations["mutation_present_in_libA"].apply(lambda x: False if x == -1 else True)
)

# Add all mutations present in library B
possible_mutations = (
    possible_mutations.merge(
        variant_df.query("library == 'LibB'")[["site", "wildtype", "mutant", "present_mutation"]],
        how="left",
        on=["site", "wildtype", "mutant"],
        validate="one_to_one",
    )
    .rename(columns={"present_mutation" : "mutation_present_in_libB"})
)
possible_mutations["mutation_present_in_libB"] = possible_mutations["mutation_present_in_libB"].fillna(-1)
possible_mutations["mutation_present_in_libB"] = (
    possible_mutations["mutation_present_in_libB"].apply(lambda x: False if x == -1 else True)
)

# Make output dir if doesn't exist
if not os.path.exists(out_dir):
    os.mkdir(out_dir)

possible_mutations.to_csv(out_file, index=False)

Visualize distribution of amino-acid variants¶