primer_check.py

#!/usr/bin/env python3
# Check that all predicted variants are contained in designed primers ('.tsv')
# RUN: python check_primers.py ./output/60mer/k3l_config.yaml

import argparse
import yaml
import pandas as pd
import dms_primer_design
from Bio import SeqIO
from Bio.Seq import Seq
import os
import sys

parser = argparse.ArgumentParser()
parser.add_argument("config", help="configuration file", type=str)
args = parser.parse_args()

print("CONFIGURATION:")
cfg = dms_primer_design.primer_design.merge_config_files(args.config, None)
for key,value in cfg.items(): print(f'{key}:', value)
print()

# parse original .gb wt file and predict missense variants
print(f"Predicting variants from file: {cfg['wt_seq']}")
wt_file_path = cfg['wt_seq']
wt_file = SeqIO.read(wt_file_path, 'genbank')
wt_seq = str(wt_file.seq.upper())

for feature in wt_file.features:
    if feature.type == 'gene':
        gene_start = feature.location.start.position
        gene_end = feature.location.end.position

# identify "windows" in gene where variants are being made
targ_windows = cfg['variant_windows']

predicted_variants = set()
for feature in wt_file.features:
    if feature.type not in targ_windows:
        continue

    start = feature.location.start.position
    end = feature.location.end.position

    win_seq = wt_seq[start:end]
    codon_list = dms_primer_design.codon_table.codons_list(win_seq)
    for i,codon in enumerate(codon_list):
        wt_aa = str(Seq(codon).translate())
        position = int((((start-gene_start)/3)+1)+i)
        missense_aas = dms_primer_design.codon_table.unique_missense_variants(codon)
        for missense_aa in missense_aas:
            missense_variant = wt_aa + str(position) + missense_aa
            predicted_variants.add(missense_variant)

# get designed variants
out_file = cfg['output_path'] + cfg['output_prefix'] + '.tsv'
print(f"Extracting designed variants from file: {out_file}\n")
if not os.path.exists(out_file):
    sys.exit(f"ERROR: primer output file '{out_file}' generated by config does not exist")

df = pd.read_csv(out_file, sep='\t')

designed_variants = set()
for index,row in df.iterrows():
    wt_aa = row['wt_aa']
    position = row['position']
    for missense_aa in row['iupac_aa']:
        missense_variant = wt_aa + str(position) + missense_aa
        designed_variants.add(missense_variant)

# predicted minus designed (should get stop codons)
missing_designed = list(predicted_variants - designed_variants)

if len(missing_designed) == 0:
    print("All predicted variants are in designed primers")
elif {i[-1] for i in missing_designed} == {'*'}:
    print(f"Designed primers are missing {len(missing_designed)} stop ('*') variants")
else:
    print("Designed primers are missing the following predicted variants:")
    for i in missing_designed: print(i)
print()

# get additionally designed variants (should be synonymous)
additional_designed = list(designed_variants - predicted_variants)

if len(additional_designed) == 0:
    print("No additional variants included in designed primers")
elif {i[0] == i[-1] for i in additional_designed} == {True}:
    print(f"{len(additional_designed)} synonymous variants indluded in designed primers:")
    for i in additional_designed: print(i)
else:
    print(f"{len(additional_designed)} variants indluded in designed primers:")
    for i in additional_designed: print(i)
print()