Test alignment parsing by Targets¶

import contextlib
import os
import random
import re
import tempfile

from IPython.display import display

import pandas as pd

import yaml

import alignparse.minimap2
import alignparse.targets

random.seed(1)  # for reproducible output
nts = "ACGT"  # valid nucleotides


def randseq(length):
    """Random nucleotide sequence of given length."""
    return "".join(random.choices(nts, k=length))


def mutseq(wtseq):
    """Random mutant sequence from a wildtype sequence."""
    return "".join(random.choice([nt for nt in nts if nt != wt]) for wt in wtseq)


def get_wt_query(
    target, ambiguous_features=("barcode", "variant_tag5", "variant_tag3")
):
    """`(description, seq)` for wildtype query, ambiguous features assigned."""
    seq = target.seq
    description = []
    for fname in ambiguous_features:
        f = target.get_feature(fname)
        fseq = randseq(f.length)
        seq = seq[: f.start] + fseq + seq[f.end :]
        description.append(f"{f.name}={fseq}")
    assert len(seq) == len(target.seq)
    assert re.fullmatch(f"[{nts}]+", seq)
    return ",".join(description), seq


queries = []

# get two random (unmappable) queries
queries.append(("unmapped_1", randseq(target.length)))
queries.append(("unmapped_2", randseq(target.length // 2)))

# get a fully wildtype query
queries.append(get_wt_query(target))

# query with query clipping at both ends and a codon substitution in gene
desc, seq = get_wt_query(target)
# add substitution to gene
gene = target.get_feature("gene")
geneseq = gene.seq
mutstart = 6
mutlength = 3
mutcodon = mutseq(gene.seq[mutstart : mutstart + mutlength])
sub_desc = []
for i in range(mutlength):
    wt = geneseq[i + mutstart]
    mut = mutcodon[i]
    sub_desc.append(f"{wt}{mutstart + i}{mut}")
    geneseq = geneseq[: mutstart + i] + mut + geneseq[mutstart + i + 1 :]
seq = seq[: gene.start] + geneseq + seq[gene.end :]
desc += f",gene_{'-'.join(sub_desc)}"
# add clipping
desc += ",query_clip5=9"
seq = randseq(9) + seq
desc += ",query_clip3=7"
seq += randseq(7)
# add to list of queries
queries.append((desc, seq))

# query with a deletion in gene and in insertion in spacer
desc, seq = get_wt_query(target)
delstart = 7
dellength = 15
geneseq = gene.seq[:delstart] + gene.seq[delstart + dellength :]
spacer = target.get_feature("spacer")
insstart = 4
ins = "G"
spacerseq = spacer.seq[:insstart] + ins + spacer.seq[insstart:]
seq = (
    seq[: gene.start]
    + geneseq
    + seq[gene.end : spacer.start]
    + spacerseq
    + seq[spacer.end :]
)
desc += f",gene_del{delstart}to{delstart + dellength}"
desc += f",spacer_ins{insstart}{ins}"
queries.append((desc, seq))

# query with deletion spanning gene and spacer
desc, seq = get_wt_query(target)
delstartgene = gene.length - 7
delendgene = delstartgene + 7
geneseq = gene.seq[:delstartgene]
delendspacer = 9
spacerseq = spacer.seq[delendspacer:]
seq = seq[: gene.start] + geneseq + spacerseq + seq[spacer.end :]
desc += f",gene_del{delstartgene}to{delendgene},spacer_del1to{delendspacer}"
queries.append((desc, seq))

# query with insertion at boundary of gene and spacer, note how
# it is assigned as being at end of gene
desc, seq = get_wt_query(target)
ins = randseq(14)
seq = seq[: gene.end] + ins + seq[spacer.start :]
desc += f",gene_ins{gene.length}{ins}"
queries.append((desc, seq))

# query with target clipping at both ends
desc, seq = get_wt_query(target)
target_clip5 = target.get_feature("termini5").length + 4
target_clip3 = 2
seq = seq[target_clip5:-target_clip3]
desc = desc.split(",")
desc = ",".join([desc[0], "variant_tag5=None", desc[2]])
desc += f",target_clip5={target_clip5},target_clip3={target_clip3}"
desc += ",termini5=None,gene=<clip4>,termini3=<clip2>"
queries.append((desc, seq))

# NBVAL_IGNORE_OUTPUT

mapper = alignparse.minimap2.Mapper(alignparse.minimap2.OPTIONS_CODON_DMS)

with contextlib.ExitStack() as stack:

    # write queries to FASTA file
    queryfile = tempfile.NamedTemporaryFile("r+", suffix=".fasta")
    queryfile.write("\n".join(f">{desc}\n{seq}" for desc, seq in queries))
    queryfile.flush()

    # align the queries to the target
    alignmentfile = tempfile.NamedTemporaryFile("r+", suffix=".sam")
    targets.align(queryfile.name, alignmentfile.name, mapper)

    # get the alignment cs data frame
    alignments_cs = targets._parse_alignment_cs(alignmentfile.name)

    # get the alignment results
    readstats, aligned, filtered = targets.parse_alignment(alignmentfile.name)