Automated tiled amplicons design
#!/usr/bin/env python3
# -*- coding: utf-8 -*-


from olivar import build, tiling, save, validate


if __name__ == "__main__":
    build(
        fasta_path = 'example_input/EPI_ISL_402124.fasta', # Path to the fasta reference sequence.

        var_path = 'example_input/delta_omicron_loc.csv', # Optional, path to the csv file of SNP coordinates and frequencies.
        # Required columns: "START", "STOP", "FREQ". "FREQ" is considered as 1.0 if empty. Coordinates are 1-based.

        BLAST_db = 'example_input/Human/GRCh38_primary', # Optional, path to the BLAST database. 
        # Note that this path should end with the name of the BLAST database (e.g., "example_input/Human/GRCh38_primary").

        out_path = './', # Output directory.
        title = 'olivar-ref', # Name of the Olivar reference file [olivar-ref].
        threads = 1 # Number of threads.
    )


    tiling(
        ref_path = 'example_output/olivar-ref.olvr', # Path to the Olivar reference file (.olvr).
        out_path = './', # Output directory.
        title = 'olivar-design', # Name of this design [olivar-design].
        
        max_amp_len = 420, # Maximum amplicon length [420].
        min_amp_len = 252, # Minimum amplicon length. 0.9*{max-amp-len} if set as None. Minimum 120.
        w_egc = 1, # Weight for extreme GC content [1.0].
        w_lc = 1, # Weight for low sequence complexity [1.0].
        w_ns = 1, # Weight for non-specificity [1.0].
        w_var = 1, # Weight for variations [1.0].
        
        temperature = 60, # PCR annealing temperature [60.0].
        salinity = 0.18, # Concentration of monovalent ions in units of molar [0.18].
        dG_max = -11.8, # Maximum free energy change of a primer in kcal/mol [-11.8].
        min_GC = 0.2, # Minimum GC content of a primer [0.2].
        max_GC = 0.75, # Maximum GC content of a primer [0.75].
        min_complexity = 0.4, # Minimum sequence complexity of a primer [0.4].
        max_len = 36, # Maximum length of a primer [36].
        check_var = True, # Filter out primer candidates with variations within 5nt of 3' end [False]. 
        # Setting check_var=True is not recommended when a lot of variations are provided, 
        # since this would significantly reduce the number of primer candidates. 

        fP_prefix = '', # Prefix of forward primer. Empty string '' by default.
        rP_prefix = '', # Prefix of reverse primer. Empty string '' by default.
        
        seed = 10, # Random seed for optimizing primer design regions and primer dimer [10].
        threads = 1, # Number of threads.
    )


    # save(
    #     design_out = 'example_output/olivar-design.olvd', # Path to the Olivar design file (.olvd)
    #     out_path = './', # Output directory.
    # )


    validate(
        primer_pool = 'example_output/olivar-design.csv', # Path to the csv file of a primer pool. 
        # Required columns: "amp_id" (amplicon name), "fP" (sequence of forward primer), "rP" (sequence of reverse primer), "pool" (pool number, e.g., 1).

        pool = 1, # Primer pool number [1]. 
        
        BLAST_db = 'example_input/Human/GRCh38_primary', # Optional, path to the BLAST database. 
        # Note that this path should end with the name of the BLAST database (e.g., "example_input/Human/GRCh38_primary").
        
        out_path = './', # Output directory. 
        title = 'olivar-val', # Name of validation.

        max_amp_len = 1500, # Maximum length of predicted non-specific amplicon [1500].
        temperature = 60, # PCR annealing temperature [60.0].
        threads = 1 # Number of threads.
    )


    validate(
        primer_pool = 'example_output/olivar-design.csv', # Path to the csv file of a primer pool. 
        # Required columns: "amp_id" (amplicon name), "fP" (sequence of forward primer), "rP" (sequence of reverse primer), "pool" (pool number, e.g., 1).

        pool = 2, # Primer pool number [1]. 
        
        BLAST_db = 'example_input/Human/GRCh38_primary', # Optional, path to the BLAST database. 
        # Note that this path should end with the name of the BLAST database (e.g., "example_input/Human/GRCh38_primary").
        
        out_path = './', # Output directory. 
        title = 'olivar-val', # Name of validation.

        max_amp_len = 1500, # Maximum length of predicted non-specific amplicon [1500].
        temperature = 60, # PCR annealing temperature [60.0].
        threads = 1 # Number of threads.
    )

example.py
treangenlab/Olivar
import setuptools

with open("README.md", "r", encoding="utf-8") as fh:
    long_description = fh.read()

setuptools.setup(
    name="olivar",
    version="1.1.3",
    author="Michael X. Wang",
    author_email="xw66@rice.edu",
    description="Olivar PCR tiling design",
    long_description=long_description,
    long_description_content_type="text/markdown",
    url="https://gitlab.com/treangenlab/olivar",
    classifiers=[
        "Programming Language :: Python :: 3",
        "License :: OSI Approved :: MIT License",
        "Operating System :: OS Independent",
    ],
    package_dir={"": "src"},
    packages=['olivar'],
    python_requires=">=3.8",
)

setup.py
treangenlab/Olivar
#!/usr/bin/env python3
# -*- coding: utf-8 -*-


'''
Olivar is a Python3 software for multiplex PCR tiling design. 
Olivar implements a novel algorithm to reduce non-specific amplifications in PCR, 
while avoiding primer design at SNPs and other undesired regions at the same time. 
Olivar also optimize for primer dimers with the SADDLE algorithm (https://www.nature.com/articles/s41467-022-29500-4). 
'''


__author__ = 'Michael X. Wang'
__version__ = '1.1.1'


import os
import sys
CURRENT_DIR = os.path.dirname(__file__)
sys.path.append(CURRENT_DIR)

from main import build, tiling, save, validate

src/olivar/__init__.py
treangenlab/Olivar
#!/usr/bin/env python3
# -*- coding: utf-8 -*-


'''
Some basic tools for processing DNA sequences and coordinates
'''


__author__ = 'Michael X. Wang'


import os
CURRENT_DIR = os.path.dirname(__file__)


from collections import Counter

import numpy as np


def revcomp(seq):
    '''
    give the reverse complement of input sequence
    base & number conversion:
        {'A':0, 'T':1, 'C':2, 'G':3}
    input:
        string or array sequence
    output:
        reverse complement of input
    '''
    if isinstance(seq, str):
        complement = {'A':'T', 'T':'A', 'C':'G', 'G':'C', 'N':'N', 
        'a':'t', 't':'a', 'c':'g', 'g':'c', 'n':'n', 
        '0':'1', '1':'0', '2':'3', '3':'2', '4':'4'}
        bases = [complement[base] for base in seq]
        bases.reverse()
        return ''.join(bases)
    elif isinstance(seq, list):
        complement = {'A':'T', 'T':'A', 'C':'G', 'G':'C', 'N':'N', 
        'a':'t', 't':'a', 'c':'g', 'g':'c', 'n':'n', 
        '0':'1', '1':'0', '2':'3', '3':'2', '4':'4', 
        0:1, 1:0, 2:3, 3:2, 4:4}
        bases = [complement[base] for base in seq]
        bases.reverse()
        return bases
    else:
        return -1


def get_GC(seq):
    '''
    get the GC ratio of input sequence
    base & number conversion:
        {'A':0, 'T':1, 'C':2, 'G':3}
    input:
        string or array sequence
    output:
        GC ratio of input sequence
    '''
    return len([base for base in list(seq) if base in ['G','g',3,'C','c',2]])/len(seq)


def seq2arr(seq_str):
    '''
    convert string sequence to array (list of numbers)
    base & number conversion:
        {'A':0, 'T':1, 'C':2, 'G':3}
    input: 
        string sequence
    output:
        array sequence
    '''
    dic = {'A':0, 'T':1, 'C':2, 'G':3, 'a':0, 't':1, 'c':2, 'g':3}
    seq_list = [dic[s] for s in list(seq_str)]
    return seq_list


def arr2seq(seq_array):
    '''
    convert array to string sequence
    base & number conversion:
        {'A':0, 'T':1, 'C':2, 'G':3}
    input: 
        array sequence
    output:
        string sequence
    '''
    dic = {0:'A', 1:'T', 2:'C', 3:'G'}
    seq_list = [dic[s] for s in seq_array]
    return ''.join(seq_list)


def seq2num(seq):
    '''
    use quaternary to convert array or string sequence to a number
    CAUTION: diffent sequence may output the same number (e.g. 'AACT' and 'AAACT')
    base & number conversion:
        {'A':0, 'T':1, 'C':2, 'G':3}
    input:
        string or array sequence
    output:
        an integer
    '''
    bases = list(seq)
    dic = {'A':0, 'T':1, 'C':2, 'G':3, 
    'a':0, 't':1, 'c':2, 'g':3, 
    0:0, 1:1, 2:2, 3:3}
    bases = [dic[base] for base in bases]
    len_bases = len(seq)
    out_num = 0
    for i in range(len_bases):
        out_num = out_num + 4**(len_bases-i-1)*bases[i]
    return out_num


def num2seq(num, l=None):
    '''
    The reverse function of seq2num
    CAUTION: a single number represents multiple sequences if l=None
    input:
        num: int
        l: length of output sequence; if none, 'A' at sequence beginning is ignored
    output:
        sequence
    '''
    base = ['A', 'T', 'C', 'G']
    seq = []
    if l is None:
        while True:
            new_num = num // 4
            b = num % 4
            seq.append(base[b])
            if new_num == 0:
                break
            num = new_num
        return ''.join(seq[::-1])
    elif num < 4**l:
        for _ in range(l):
            new_num = num // 4
            b = num % 4
            seq.append(base[b])
            num = new_num
        return ''.join(seq[::-1])
    else:
        print('Input number out of range according to input sequence length.')
        return None


def seq2complex(seq_str):
    '''
    Convert string sequence to a list of complex
    '''
    dic = {'A':1, 'T':-1, 'C':1j, 'G':-1j, 'N': 0, 'a':1, 't':-1, 'c':1j, 'g':-1j, 'n': 0}
    seq_list = [dic[s] for s in list(seq_str)]
    return seq_list


def complex2seq(seq_complex):
    '''
    Convert a list of complex back to string
    '''
    seq_str = []
    for c in seq_complex:
        if c.real > c.imag:
            if -c.real < c.imag:
                seq_str.append('A')
            elif -c.real > c.imag:
                seq_str.append('G')
            else:
                seq_str.append('N')
        elif c.real < c.imag:
            if -c.real < c.imag:
                seq_str.append('C')
            elif -c.real > c.imag:
                seq_str.append('T')
            else:
                seq_str.append('N')
        else:
            seq_str.append('N')
    return ''.join(seq_str)


def randseq(l, seed=None):
    '''
    Generate random sequence of length l
    input:
        l: sequence length
        seed: random seed for numpy
    output:
        random sequence of length l
    '''
    if seed is not None:
        np.random.seed(seed)
    return arr2seq(np.random.randint(4, size=l).tolist())


def randarr(l, seed=None):
    '''
    Generate random array of length l
    input:
        l: array length
        seed: random seed for numpy
    output:
        random array of length l
    '''
    if seed is not None:
        np.random.seed(seed)
    return np.random.randint(4, size=l).tolist()


def get_complexity(seq):
    # calculate Shannon entropy for word size of 1, 2 and 3
    hash1 = Counter()
    hash2 = Counter()
    hash3 = Counter()

    # count occurance
    for j in range(len(seq)-2):
        hash1[seq[j]] += 1
        hash2[seq[j:j+2]] += 1
        hash3[seq[j:j+3]] += 1
    hash1[seq[-2]] += 1
    hash1[seq[-1]] += 1
    hash2[seq[-2:]] += 1

    # calculate probability mass function
    p1 = np.array(list(hash1.values()))/(len(seq))
    p2 = np.array(list(hash2.values()))/(len(seq)-1)
    p3 = np.array(list(hash3.values()))/(len(seq)-2)

    SE1 = -np.sum(p1 * np.log2(p1))
    SE2 = -np.sum(p2 * np.log2(p2))
    SE3 = -np.sum(p3 * np.log2(p3))
    return min([SE1/2, SE2/4, SE3/6])

src/olivar/basic.py
treangenlab/Olivar
#!/usr/bin/env python3
# -*- coding: utf-8 -*-


'''
Functions for primer design
Architecture:
design.py
    primer_generator(object)
        WAlignScore()
    PrimerSetBadnessFast()
'''


__author__ = 'Michael X. Wang'


import os
CURRENT_DIR = os.path.dirname(__file__)


import random
import json
from collections import Counter

import numpy as np
import pandas as pd
from Bio import SeqIO

import basic
from basic import seq2num
from basic import num2seq
from basic import revcomp


GC_LETTER = ['C', 'G', 'c', 'g']


def WAlignScore(seq1, seq2=None):
    '''
    Calculate primer self badness
    '''
    PENALTY_OFFSET = -0.75 # Formulat is (PO + 1) / (PO + d+1)
    # if =1, then d=1 gives 2/3 socre, if =0, then d=1 gives 1/2 score, if =-0.5, then d=1 gives 1/3 score
    indel_penalty = 3
    replace_penalty = 2
    GC_score = 2
    AT_score = 1
    
    seq1 = seq1.lower()
    len1 = len(seq1)
    
    if seq2 is not None:
        seq2 = seq2.lower()
        len2 = len(seq2)
    else:
        seq2 = basic.revcomp(seq1)
        len2 = len1
        
    if len1 == 0:
        max_score = 0
    else:
        score = np.zeros([len1, len2])
        if seq1[0] == seq2[0]:
            if seq1[0] == 'c' or seq1[0] == 'g':
                score[0][0] = GC_score
            else:
                score[0][0] = AT_score
                
        for j in range(1, len2):
            if seq1[0] == seq2[j]:
                if seq1[0] == 'c' or seq1[0] == 'g':
                    score[0][j] = GC_score
                else:
                    score[0][j] = AT_score
            else:
                score[0][j] = max(score[0][j-1] - indel_penalty, 0)
                
        for i in range(1, len1):
            if seq1[i] == seq2[0]:
                if seq1[i] == 'c' or seq1[i] == 'g':
                    score[i][0] = GC_score
                else:
                    score[i][0] = AT_score
            else:
                score[i][0] = max(score[i-1][0] - indel_penalty, 0)
                
            for j in range(1, len2):
                if seq1[i] == seq2[j]:
                    if seq1[i] == 'c' or seq1[i] == 'g':
                        score[i][j] = GC_score + score[i-1][j-1]
                    else:
                        score[i][j] = AT_score + score[i-1][j-1]
                else:
                    score[i][j] = max([score[i-1][j] - indel_penalty, 
                                       score[i][j-1] - indel_penalty, 
                                       score[i-1][j-1] - replace_penalty, 0])
                    
        # weight by offset
        for i in range(len1):
            for j in range(len2):
                distance = len1 - i - 1
                if seq1[i] == seq2[j]:
                    score[i][j] -= 3*distance
                else:
                    score[i][j] = 0
                    
        max_score = 1.5**(score.max() - 6)
    return max_score


class primer_generator(object):
    '''
    Usage:
        g = primer_generator(temperature, salinity)
        retult_1 = g.get(seq1, dG_max=-11.8, min_GC=0.2, max_GC=0.75, min_complexity=0.4, max_len=36, check_SNP=True, prefix='')
        retult_2 = g.get(seq2, dG_max=-11.8, min_GC=0.1, max_GC=0.9, min_complexity=0.3, max_len=36, check_SNP=True)
    '''
    def __init__(self, temperature=60, salinity=0.18, end_pool_path=None):
        # Parameter = np.array([[11, -7.6, -21.3],
        #                     [12, -7.2, -20.4],
        #                     [13, -8.4, -22.4],
        #                     [14, -7.8, -21.0],
        #                     [21, -7.2, -21.3],
        #                     [22, -7.6, -21.3],
        #                     [23, -8.2, -22.2],
        #                     [24, -8.5, -22.7],
        #                     [31, -8.5, -22.7],
        #                     [32, -7.8, -21.0],
        #                     [33, -8.0, -19.9],
        #                     [34, -10.6, -27.2],
        #                     [41, -8.2, -22.2],
        #                     [42, -8.4, -22.4],
        #                     [43, -9.8, -24.4],
        #                     [44, -8.0, -19.9],])
        paraH = np.array([[-7.6, -7.2, -8.4, -7.8],
                          [-7.2, -7.6, -8.2, -8.5],
                          [-8.5, -7.8, -8.0, -10.6],
                          [-8.2, -8.4, -9.8, -8.0]])
        paraS = np.array([[-21.3, -20.4, -22.4, -21.0],
                          [-21.3, -21.3, -22.2, -22.7],
                          [-22.7, -21.0, -19.9, -27.2],
                          [-22.2, -22.4, -24.4, -19.9]])
        self.paraS = paraS + 0.368 * np.log(salinity)
        self.paraG = paraH - (temperature + 273.15) * self.paraS / 1000
        self.dG_init = 0.2 - (temperature + 273.15) * (-5.7)/1000

        # load pre-generated pool of primer end sequences
        if not end_pool_path is None:
            with open(end_pool_path) as f:
                end_pool = json.load(f)
            self.check_end = True
            self.l_end = len(end_pool[0])
            self.end_pool = [s.lower() for s in end_pool]
        else:
            self.check_end = False
            self.l_end = None
            self.end_pool = None
        
    def StacksDG(self, seq):
        rep = basic.seq2arr(seq)
        dG_stacks = 0
        for i, _ in enumerate(rep[:-1]):
            dG_stacks += self.paraG[rep[i]][rep[i+1]]
        return dG_stacks
    
    def get(self, seq, prefix='', dG_max=-11.8, min_GC=0.2, max_GC=0.75, 
            min_complexity=0.4, min_len=15, max_len=36, check_SNP=True, check_BLAST=False):
        '''
        Input:
            Sequence to generate primers from; restrictions (dG, GC, etc.)
        Output:
            DataFrame of primers; dictionary of failure cases
        '''
        INTR_SCORE = 2

        # input should be lowercase
        if seq == seq.upper(): # all uppercase
            print('input sequence should be in lowercase, with SNPs in uppercase. ')
            seq = seq.lower()
            if check_SNP:
                print('ignore check_SNP')
                check_SNP = False

        primer_seq = [] # primer sequence
        primer_arr = [] # array format
        primer_len = [] # primer length without prefix
        primer_dist = [] # distance to mother sequence end
        badness = []
        primer_dG = []
        BLAST_hits = []
        end_fail = 0 # number of primers failed end check
        SNP_fail = 0 # number of primers failed SNP check
        GC_fail = 0 # number of primers failed GC ratio check
        complexity_fail = 0 # number of primers failed complexity check
        for end_pos in range(len(seq)-1, 8, -1): # end_pos is 0-based
            # check end
            if self.check_end:
                if not seq[end_pos-self.l_end+1:end_pos+1] in self.end_pool:
                    end_fail += 1
                    continue
            # check SNP in last 5 bases
            if check_SNP:
                if seq[end_pos-4:end_pos+1] != seq[end_pos-4:end_pos+1].lower():
                    SNP_fail += 1
                    continue
            start_pos = end_pos - 1
            curr_dG = self.dG_init + self.StacksDG(seq[start_pos:end_pos+1])
            # satisfy minimum length
            while (start_pos > end_pos - min_len + 1) and (start_pos > 0):
                start_pos -= 1
                curr_dG += self.StacksDG(seq[start_pos:(start_pos+2)])
            # long enough to satisfy energy but not exceeding max_len
            while (curr_dG > dG_max) and (start_pos > end_pos - max_len + 1) and (start_pos > 0):
                start_pos -= 1
                curr_dG += self.StacksDG(seq[start_pos:(start_pos+2)])
            # check restrictions
            if (curr_dG <= dG_max) or (end_pos - start_pos + 1 == max_len):
                GC_ratio = basic.get_GC(seq[start_pos:end_pos+1])
                if min_GC <= GC_ratio <= max_GC:
                    complexity = basic.get_complexity(seq[start_pos:end_pos+1])
                    if complexity >= min_complexity:
                        s_raw = seq[start_pos:end_pos+1]
                        s = prefix + s_raw
                        primer_seq.append(s)
                        primer_arr.append(basic.seq2arr(s))
                        primer_len.append(end_pos - start_pos + 1)
                        primer_dist.append(len(seq) - end_pos - 1)
                        primer_dG.append(curr_dG)
                        if check_BLAST:
                            #count = BLAST_seq(s_raw, **BLAST_config)
                            if count > 0:
                                b = INTR_SCORE * WAlignScore(s) + 1000 * np.log10(count)**4
                            else:
                                b = INTR_SCORE * WAlignScore(s)
                        else:
                            count = -1
                            b = INTR_SCORE * WAlignScore(s)
                        BLAST_hits.append(count)
                        badness.append(b)
                    else:
                        complexity_fail += 1
                else:
                    GC_fail += 1
        primers = pd.DataFrame({'seq': primer_seq, 
                                'arr': primer_arr, 
                                'primer_len': primer_len, 
                                'dist': primer_dist, 
                                'badness': badness, 
                                'dG': primer_dG, 
                                'BLAST_hits': BLAST_hits})
        return primers, {'end_fail': end_fail, 'SNP_fail': SNP_fail, 
        'GC_fail': GC_fail, 'complexity_fail': complexity_fail}


def PrimerSetBadnessFast(all_fP: list, all_rP=[], existing=[]):
    '''
    Input:
        all_fP: list of strings
        all_rP: list of strings
        existing: list of strings
    Output:
        Badness: total Loss of the primer set (all fP and rP)
        Badness_component: Badness of each primer
    '''
    endhash4 = Counter()
    endhash5 = Counter()
    endhash6 = Counter()
    middlehash7 = Counter()
    middlehash8 = Counter()

    PENALTY_OFFSET = 0 # if=1, then d=1 gives 2/3 socre, if =0, then d=1 gives 1/2 score, if =-0.5, then d=1 gives 1/3 score

    END4 = 1
    END5 = 4
    END6 = 20
    MIDDLE7 = 100
    MIDDLE8 = 500

    # NOTE: Intramolecular checks done at GeneratePrimerCandidates to be efficient

    # set all sequences as lowercase
    all_fP = [p.lower() for p in all_fP]
    all_rP = [p.lower() for p in all_rP]
    existing = [p.lower() for p in existing]

    # Set up end hash tables (check end 4 nt to middle binding)
    for p in all_fP+all_rP+existing:
        endhash4[p[-4:]] += 1
        endhash5[p[-5:]] += 1
        endhash6[p[-6:]] += 1
        
    # Set up middle hash table
    # Middlehash penalizes based on closeness to 3' end, Badness = 2 / (distance to 3' + 2);
    # So absolute last 7 is worth 1, 1nt away is worth 0.67, 2nt away is worth 0.5, etc.
    for p in all_fP+all_rP+existing:
        l = len(p)
        for j in range(l-6):
            middlehash7[p[j:j+7]] += (PENALTY_OFFSET + 1) / (l - j - 6 + PENALTY_OFFSET)
        for j in range(l-7):
            middlehash8[p[j:j+8]] += (PENALTY_OFFSET + 1) / (l - j - 7 + PENALTY_OFFSET)

    # Run through each sequence's reverse complement to add up badness
    Badness = 0
    Badness_component = []
    for one_side_primer in [all_fP, all_rP]:
        one_side_badness = []
        for p in one_side_primer:
            p_badness = 0
            c = basic.revcomp(p)
            l = len(c)
            for j in range(l-3):
                k = c[j:j+4]
                try:
                    endscore4 = endhash4[k]
                    numGC = len([b for b in k if b in GC_LETTER])
                    p_badness += (endscore4 * END4 * (PENALTY_OFFSET+1)/(j+1+PENALTY_OFFSET)) * (2**numGC)
                except KeyError:
                    pass
            for j in range(l-4):
                k = c[j:j+5]
                try:
                    endscore5 = endhash5[k]
                    numGC = len([b for b in k if b in GC_LETTER])
                    p_badness += (endscore5 * END5 * (PENALTY_OFFSET+1)/(j+1+PENALTY_OFFSET)) * (2**numGC)
                except KeyError:
                    pass
            for j in range(l-5):
                k = c[j:j+6]
                try:
                    endscore6 = endhash6[k]
                    numGC = len([b for b in k if b in GC_LETTER])
                    p_badness += (endscore6 * END6 * (PENALTY_OFFSET+1)/(j+1+PENALTY_OFFSET)) * (2**numGC)
                except KeyError:
                    pass
            for j in range(l-6):
                k = c[j:j+7]
                try:
                    midscore7 = middlehash7[k]
                    numGC = len([b for b in k if b in GC_LETTER])
                    p_badness += (midscore7 * MIDDLE7 * (PENALTY_OFFSET+1)/(j+1+PENALTY_OFFSET)) * (2**numGC)
                except KeyError:
                    pass
            for j in range(l-7):
                k = c[j:j+8]
                try:
                    midscore8 = middlehash8[k]
                    numGC = len([b for b in k if b in GC_LETTER])
                    p_badness += (midscore8 * MIDDLE8 * (PENALTY_OFFSET+1)/(j+1+PENALTY_OFFSET)) * (2**numGC)
                except KeyError:
                    pass
            Badness += p_badness
            one_side_badness.append(p_badness)
        Badness_component.append(one_side_badness)

    return Badness, Badness_component

src/olivar/design.py
treangenlab/Olivar
#!/usr/bin/env python3
# -*- coding: utf-8 -*-


'''
Main workflow of Olivar tiling.
Architecture:
main.py
    build()
    tiling()
        design_context_seq()
            generate_context()
                find_min_loc()
        get_primer()
        optimize()
        to_df()
        tiling_save()
    validate()
'''


__author__ = 'Michael X. Wang'


import os
CURRENT_DIR = os.path.dirname(__file__)

import json
import pickle # needs python>=3.8
import random
import copy
import multiprocessing
from time import time
from math import ceil, floor
from copy import deepcopy

import pandas as pd
import numpy as np
from numpy.random import rand, default_rng
from tqdm import tqdm

from Bio import SeqIO

import plotly
import plotly.graph_objects as go
from plotly.subplots import make_subplots

# internal modules
import basic
import design # SADDLE
from ncbi_tools import BLAST_batch_short, ns_simulation
from design import PrimerSetBadnessFast

N_POOLS = 2 # number of primer pool (not to be changed)

# algorithmic parameters
PRIMER_DESIGN_LEN = 40 # length of primer design region [40]
UPPER_GC = 0.75 # extreme GC upper bond [0.75]
LOWER_GC = 0.25 # extreme GC lower bond [0.75]
LOW_COMPLEXITY = 0.4 # low complexity lower bond [0.4]
CHOICE_RATE = 0.3 # bottom CHOICE_RATE lowest risk primer design regions are choosen from candidate region
RISK_TH = 0.1 # top RISK_TH highest risk primer design regions are considered as loss


def build(fasta_path: str, var_path: str, BLAST_db: str, out_path: str, title: str, threads: int):
    '''
    Build the Olivar reference file for tiled amplicon design
    Input:
        fasta_path: Path to the fasta reference sequence.
        var_path: Optional, path to the csv file of SNP coordinates and frequencies. 
            Required columns: "START", "STOP", "FREQ". "FREQ" is considered as 1.0 if empty. Coordinates are 1-based.
        BLAST_db: Optional, path to the BLAST database. 
            Note that this path should end with the name of the BLAST database (e.g., "example_input/Human/GRCh38_primary").
        out_path: Output directory [./]. 
        title: Name of the Olivar reference file [olivar-ref]. 
        threads: Number of threads [1]. 
    '''
    n_cpu = threads

    if not os.path.exists(out_path):
        os.makedirs(out_path)

    # all coordinates are 1-based and inclusive
    word_size = 28
    offset = 14 # word_size should be divisible by offset
    print('Building Olivar reference with word_size=%d, offset=%d' % (word_size, offset))

    n_cycle = word_size/offset
    if n_cycle != int(n_cycle):
        raise ValueError('word_size should be divisible by offset.')
    else:
        n_cycle = int(n_cycle)

    for record in SeqIO.parse(fasta_path, 'fasta'):
        seq_raw = str(record.seq).lower() # SNPs in upper case
        break # read first record
    
    # check ambiguous bases
    ambiguous_bases = set(seq_raw) - {'a', 't', 'c', 'g'}
    if ambiguous_bases:
        raise ValueError(f'Ambiguous bases are not supported. {ambiguous_bases} found in reference fasta file.')

    # load SNP/iSNV coordinates
    var_arr = np.zeros(len(seq_raw))
    if var_path:
        seq_raw = list(seq_raw)
        df = pd.read_csv(var_path, sep=',', index_col=False)

        # variant csv file sanity check
        required_col = {
            'START': 'int64', 
            'STOP': 'int64'
        }
        optional_col = {
            'FREQ': 'float64'
        }
        for cname, expected_dtype in required_col.items():
            try:
                received_dtype = df[cname].dtype
            except KeyError:
                raise ValueError(f'Column {cname} not found in {var_path}')
            if received_dtype != expected_dtype:
                raise ValueError(f'Invalid value found in column {cname} of {var_path}, only {expected_dtype} is accepted (received {received_dtype})')
        for cname, expected_dtype in optional_col.items():
            try:
                received_dtype = df[cname].dtype
            except KeyError:
                # column not provided
                print("Warning: 'FREQ' is not provided")
                df[cname] = np.nan
                continue
            if received_dtype != expected_dtype:
                raise ValueError(f'Invalid value found in column {cname} of {var_path}, only {expected_dtype} is accepted (received {received_dtype})')

        # add variant frequency to var_arr
        for i, row in df.iterrows():
            freq = row['FREQ']
            if np.isnan(freq):
                # FREQ is not provided
                freq = 1.0
            var_arr[int(row['START'])-1:int(row['STOP'])] += freq
            # capitalize SNPs
            for pos in range(int(row['START'])-1, int(row['STOP'])):
                seq_raw[pos] = seq_raw[pos].upper()
        var_arr = var_arr**0.5 # amplify low frequency variants
        seq_raw = ''.join(seq_raw)
    else:
        print('No sequence variation coordinates provided, skipped.')

    # calculate risk of GC, non-specificity, complexity
    seq_len_temp = (len(seq_raw)//offset) * offset # length of temporary sequence
    start_temp = (len(seq_raw) - seq_len_temp)//2 + 1 # start coordinate of sequence to be processed as words
    stop_temp = start_temp + seq_len_temp - 1

    # fetch all words
    all_word = []
    for pos in range(0, seq_len_temp-word_size+1, offset):
        word_start = start_temp+pos-1
        word = seq_raw[word_start : word_start+word_size]
        all_word.append(word)

    # get GC, complexity, BLAST hits of each word
    with multiprocessing.Pool(processes=n_cpu) as pool:
        print('Calculating GC content and sequence complexity...')
        tik = time()
        all_gc = np.array(pool.map(basic.get_GC, all_word))
        all_complexity = np.array(pool.map(basic.get_complexity, all_word))
        print('Finished in %.3fs' % (time()-tik))
    if BLAST_db:
        print('Calculating non-specificity...')
        all_hits, _ = BLAST_batch_short(all_word, db=BLAST_db, n_cpu=n_cpu, mode='rough') # tabular output
    else:
        print('No BLAST database provided, skipped.')
        all_hits = np.zeros(len(all_word))

    # calculate score array
    start = start_temp + (n_cycle-1)*offset # start coordinate of sequence with score array
    stop = stop_temp - (n_cycle-1)*offset
    seq_len = stop - start + 1 # length of final sequence
    gc_arr = np.zeros(seq_len)
    comp_arr = np.zeros(seq_len)
    hits_arr = np.zeros(seq_len)
    for pos in range(seq_len):
        n = pos//offset
        gc_arr[pos] = np.sum(all_gc[n:n+n_cycle])
        comp_arr[pos] = np.sum(all_complexity[n:n+n_cycle])
        hits_arr[pos] = np.sum(all_hits[n:n+n_cycle])
    gc_arr = gc_arr/n_cycle
    hits_arr = hits_arr/n_cycle

    olv_ref = {
        'seq': seq_raw, 
        'seq_record': record, 
        'start': start, 
        'stop': stop, 
        'gc_arr': gc_arr, 
        'comp_arr': comp_arr, 
        'var_arr': var_arr[start-1:stop], 
        'hits_arr': hits_arr, 
        'all_hits': all_hits
    }
    save_path = os.path.join(out_path, '%s.olvr' % title)
    with open(save_path, 'wb') as f:
        pickle.dump(olv_ref, f, protocol=5) # protocol 5 needs python>=3.8
        print('Reference file saved as %s' % save_path)


def find_min_loc(risk_arr, start, stop, rng):
    '''
    find the next primer design region within start and stop.
    '''
    score = []
    loc = np.array(range(start, stop-PRIMER_DESIGN_LEN+2))
    for i in loc:
        score.append(sum(risk_arr[i-1: i-1+PRIMER_DESIGN_LEN]))
    if len(loc) == 1:
        idx_rand = 0
    elif CHOICE_RATE < 1:
        k = int(len(loc)*CHOICE_RATE)
        if k == 0:
            k = 1
        idx_k_lowest = np.argpartition(score, k)[:k]
        idx_rand = rng.choice(idx_k_lowest)
    else:
        idx_rand = rng.choice(range(len(loc)))
    primer_start = loc[idx_rand]
    return primer_start, primer_start+PRIMER_DESIGN_LEN-1, score[idx_rand]


def generate_context(SOI):
    '''
    find a valid set of PDRs that cover the whole sequence of interest (SOI).
    '''
    risk_arr, start, stop, max_amp_len, min_amp_len, seed = SOI
    rng = default_rng(seed)
    
    all_context_seq = [] # coordinates of all context sequences
    all_risk = []

    # generate the first pair of primer design region
    fp_start, fp_stop, fp_risk = find_min_loc(risk_arr, start, start+3*PRIMER_DESIGN_LEN-1, rng)
    rp_stop, rp_start, rp_risk = find_min_loc(risk_arr, fp_start+min_amp_len-PRIMER_DESIGN_LEN, fp_start+max_amp_len-1, rng)
    all_context_seq.append((fp_start, fp_stop, rp_stop, rp_start))
    all_risk.append((fp_risk, rp_risk))

    # generate the second pair of primer design region
    left_lim = max(fp_stop+1, rp_start+2*PRIMER_DESIGN_LEN-max_amp_len+1) # leave space for the next fP design region
    fp_start, fp_stop, fp_risk = find_min_loc(risk_arr, left_lim, rp_stop-1, rng)
    left_lim = max(fp_start+min_amp_len-PRIMER_DESIGN_LEN, rp_start+PRIMER_DESIGN_LEN+1)
    right_lim = fp_start+max_amp_len-1
    if right_lim > stop:
        raise ValueError('sequence provided is too short.')
    rp_stop, rp_start, rp_risk = find_min_loc(risk_arr, left_lim, right_lim, rng)
    all_context_seq.append((fp_start, fp_stop, rp_stop, rp_start))
    all_risk.append((fp_risk, rp_risk))

    # design the rest context sequences
    cs_id = 2 # id of the next context sequence to be designed (0-based)
    while True:
        # fp
        # all_context_seq[cs_id-2][3]: rp_start of the last last context sequence
        left_lim = max(all_context_seq[cs_id-2][3]+1, rp_start+2*PRIMER_DESIGN_LEN-max_amp_len+1)
        fp_start, fp_stop, fp_risk = find_min_loc(risk_arr, left_lim, rp_stop-1, rng)
        # rp
        left_lim = max(fp_start+min_amp_len-PRIMER_DESIGN_LEN, rp_start+PRIMER_DESIGN_LEN+1)
        right_lim = fp_start+max_amp_len-1
        if right_lim > stop:
            break
        rp_stop, rp_start, rp_risk = find_min_loc(risk_arr, left_lim, right_lim, rng)

        all_context_seq.append((fp_start, fp_stop, rp_stop, rp_start))
        all_risk.append((fp_risk, rp_risk))
        cs_id += 1

    all_context_seq = np.array(all_context_seq)
    all_risk = np.array(all_risk)
    all_risk_flatten = all_risk.flatten()
    k = int(len(all_risk_flatten) * (1-RISK_TH))
    loss = sum(np.partition(all_risk_flatten, k)[k:]**2)
    #loss = sum(all_risk_flatten[all_risk_flatten>=RISK_TH])
    return all_context_seq, all_risk, loss


def design_context_seq(config):
    '''
    PDR optimization. 
    Input:
        config: design configurations
    Output:
        all_plex_info
    '''
    ref_path = config['ref_path']
    max_amp_len = config['max_amp_len']
    w_gc = config['w_egc']
    w_complexity = config['w_lc']
    w_hits = config['w_ns']
    w_var = config['w_var']
    seed = config['seed']
    n_cpu = config['threads']
    
    # set random number generator
    rng_parent = default_rng(seed)
    
    with open(ref_path,  'rb') as f:
        olv_ref = pickle.load(f)

    seq_raw = olv_ref['seq'] # SNPs are capitalized in build()
    start = olv_ref['start'] # start of design region, 1-based, closed
    stop = olv_ref['stop'] # stop of design region, 1-based, closed
    gc_arr = olv_ref['gc_arr']
    comp_arr = olv_ref['comp_arr']
    var_arr = olv_ref['var_arr']
    hits_arr = olv_ref['hits_arr']
    print('Successfully loaded reference file %s' % ref_path)
    
    if not config['min_amp_len']:
        min_amp_len = int(max_amp_len*0.9)
    else:
        min_amp_len = config['min_amp_len']
    if min_amp_len < 3*PRIMER_DESIGN_LEN:
        raise ValueError('min_amp_len is too small')

    gc_arr = np.logical_or(gc_arr<LOWER_GC, gc_arr>UPPER_GC).astype(int)
    comp_arr = (comp_arr < LOW_COMPLEXITY).astype(int)
    if max(hits_arr) != 0:
        hits_arr = hits_arr/max(hits_arr) # normalize

    # make arrays the same length as seq_raw
    gc_arr = w_gc*0.5*np.concatenate((np.zeros(start-1), gc_arr, np.zeros(len(seq_raw)-stop)))
    comp_arr = w_complexity*0.5*np.concatenate((np.zeros(start-1), comp_arr, np.zeros(len(seq_raw)-stop)))
    hits_arr = w_hits*np.concatenate((np.zeros(start-1), hits_arr, np.zeros(len(seq_raw)-stop)))
    var_arr = w_var*10*np.concatenate((np.zeros(start-1), var_arr, np.zeros(len(seq_raw)-stop)))

    # construct risk array (first row is risk, second row is coordinate on seq_rawy)
    risk_arr = gc_arr + comp_arr + hits_arr + var_arr

    N = 500*len(risk_arr)//max_amp_len # number of primer sets to generate
    #N = 100
    rand_int = rng_parent.integers(2**32, size=N) # random seeds for each iteration

    # single thread
    # design = [generate_context((risk_arr, start, stop, max_amp_len, min_amp_len, rand_int[i])) for i in tqdm(range(N))]

    # multi threads
    print('reference sequence length: %d' % len(seq_raw))
    print('design region: %d:%d' % (start, stop))
    print('region length: %d' % (stop-start+1))
    print('number of PDR sets to be tested: %d' % N)
    print('Designing PDRs...')
    tik = time()
    with multiprocessing.Pool(processes=n_cpu) as pool:
        batch = [(risk_arr, start, stop, max_amp_len, min_amp_len, rand_int[i]) for i in range(N)]
        design = list(
            tqdm(
                pool.imap(generate_context, batch, chunksize=max(N//n_cpu//10, 1)), # chunksize=1 is slow
                total=N, 
                ascii=' >'
            )
        )
    print('Finished in %.3fs' % (time()-tik))

    # Loss of all iterations
    all_loss = [d[2] for d in design]
    all_loss.sort(reverse=True)

    # find the best arangement of primer design regions
    best_design = sorted(design, key=lambda x:x[2])
    all_context_seq, all_risk, loss = best_design[0]
    print('Best total risk: %.3f' % loss)
    
    # prepare output
    all_plex_info = {}
    for i, (context_seq, risk) in enumerate(zip(all_context_seq, all_risk)):
        cs = seq_raw[context_seq[0]-1:context_seq[3]] # actual context sequence
        plex_info = {
            'tube': i%2 + 1, 
            'context_seq_coords': (context_seq[0], context_seq[3]), 
            'primers_coords': tuple(context_seq), 
            'insert_coords': (context_seq[1]+1, context_seq[2]-1), 
            'context_seq': cs, 
            'context_seq_revcomp': basic.revcomp(cs), 
            'fP_design': seq_raw[context_seq[0]-1:context_seq[1]], 
            'rP_design': basic.revcomp(seq_raw[context_seq[2]-1:context_seq[3]]), 
            'risk': tuple(risk)
        }
        all_plex_info['%s_%d' % (config['title'], i+1)] = plex_info
    print('total amplicons: %d' % (i+1))
    cover_start = all_context_seq[0][1]+1
    cover_stop = all_context_seq[-1][2]-1
    print('covered region: %d:%d' % (cover_start, cover_stop))
    print('coverage of reference sequence: %.3f%%' % (100*(cover_stop-cover_start+1)/len(seq_raw)))
    return all_plex_info, risk_arr, gc_arr, comp_arr, hits_arr, var_arr, all_loss, olv_ref['seq_record']


def get_primer(all_plex_info, config):
    '''
    Generate candidates of fP and rP for each plex
    Input:
        all_plex_info: output of design_context_seq()
        config: design configurations
    Output:
        all_plex_info: add primer candidates to input all_plex_info
    '''
    print('Generating primer candidates...')
    tik = time()

    temperature_fP = config['temperature']
    temperature_rP = config['temperature']
    salinity = config['salinity']
    fP_adapter = config['fP_prefix']
    rP_adapter = config['rP_prefix']
    default_fP_setting = {
        'dG_max': config['dG_max'], 
        'min_GC': config['min_GC'], 
        'max_GC': config['max_GC'], 
        'min_complexity': config['min_complexity'], 
        'max_len': config['max_len'], 
        'check_SNP': config['check_SNP']
    }
    default_rP_setting = deepcopy(default_fP_setting)
    check_BLAST = False

    # generate primer candidates
    fP_generator = design.primer_generator(temperature_fP, salinity)
    rP_generator = design.primer_generator(temperature_rP, salinity)
    n = 0
    for plex_id, plex_info in tqdm(all_plex_info.items(), ascii=' >'):
        # set primer prefix
        fP_prefix = fP_adapter
        rP_prefix = rP_adapter
        n += 1
        
        # generate fP
        fP_design = plex_info['fP_design'] # sequence to generate fP candidates
        fP_setting = deepcopy(default_fP_setting)
        fP, fail = fP_generator.get(fP_design, fP_prefix, check_BLAST=check_BLAST, **fP_setting)
        while len(fP) == 0:
            print('Fail to generate fP, update setting: %s' % plex_id)
            # detect failure mode with the most cases
            fail_mode = max(fail, key=fail.get)
            # update setting
            if fail_mode == 'SNP_fail' and fP_setting['check_SNP'] == True:
                fP_setting['check_SNP'] = False
            elif fail_mode == 'GC_fail':
                fP_setting['min_GC'] -= 0.05
                fP_setting['max_GC'] += 0.05
            elif fail_mode == 'complexity_fail':
                fP_setting['min_complexity'] -= 0.05
            # generate primers again
            fP, fail = fP_generator.get(fP_design, fP_prefix, check_BLAST=check_BLAST, **fP_setting)
        plex_info['fP_candidate'] = fP
        plex_info['fP_setting'] = fP_setting
        
        # generate rP
        rP_design = plex_info['rP_design'] # sequence to generate rP candidates
        rP_setting = deepcopy(default_rP_setting)
        rP, fail = rP_generator.get(rP_design, rP_prefix, check_BLAST=check_BLAST, **rP_setting)
        while len(rP) == 0:
            print('Fail to generate rP, update setting: %s' % plex_id)
            # detect failure mode with the most cases
            fail_mode = max(fail, key=fail.get)
            # update setting
            if fail_mode == 'SNP_fail' and rP_setting['check_SNP'] == True:
                rP_setting['check_SNP'] = False
            elif fail_mode == 'GC_fail':
                rP_setting['min_GC'] -= 0.05
                rP_setting['max_GC'] += 0.05
            elif fail_mode == 'complexity_fail':
                rP_setting['min_complexity'] -= 0.05
            # generate primers again
            rP, fail = rP_generator.get(rP_design, rP_prefix, check_BLAST=check_BLAST, **rP_setting)
        plex_info['rP_candidate'] = rP
        plex_info['rP_setting'] = rP_setting
        
    print('Finished in %.3fs' % (time()-tik))
    return all_plex_info


def optimize(all_plex_info, config):
    '''
    Optimize fP and rP with the SADDLE algorithm.
    Input:
        all_plex_info: output of get_primer()
        config: design configurations
    Output:
        all_plex_info: add optimized fP and rP to each plex of input all_plex_info
    '''
    print('Optimizing primer dimers...')
    # set random seed
    np.random.seed(config['seed'])
    random.seed(config['seed'])

    # automatically determine # SADDLE iterations
    pool_size = max(len(all_plex_info)/N_POOLS - 100, 0)
    InitSATemp = 1000 + 10*pool_size
    NUMSTEPS = 10 + int(pool_size/10)
    ZEROSTEPS = NUMSTEPS
    TimePerStep = 1000
    # InitSATemp = config['InitSATemp']
    # NUMSTEPS = config['NumSteps']
    # ZEROSTEPS = config['ZeroSteps']
    # TimePerStep = config['TimePerStep']

    # generate primer pair candidates (fP-rP)
    n_pair = []
    for plex_id, plex_info in all_plex_info.items():
        optimize = []
        plex_amp_len = []
        for i, fp in plex_info['fP_candidate'].iterrows():
            for j, rp in plex_info['rP_candidate'].iterrows():
                insert_len = plex_info['insert_coords'][1] - plex_info['insert_coords'][0] + 1
                # check amplicon length
                amp_len = fp['primer_len'] + fp['dist'] + insert_len + rp['dist'] + rp['primer_len']
                plex_amp_len.append(amp_len)
                optimize.append([fp, rp])
        plex_info['optimize'] = optimize
        n_pair.append(len(optimize))
    print('total primer pairs %d' % sum(n_pair))
    print('average pairs per plex %.2f' % (sum(n_pair)/len(n_pair)))
    
    # optimize each tube
    all_lc = []
    tik = time()
    for i_tube in range(1, N_POOLS+1):
        print('\npool %d, simulated annealing...' % i_tube)
        # plex_id of current tube
        curr_tube = [plex_id for plex_id, plex_info in all_plex_info.items() if plex_info['tube'] == i_tube]

        # load existing primers
        #existing_primer = config['existing_primer']
        existing_primer = []

        # randomly select one primer set (one pair each plex)
        curr_index = {}
        curr_fp = {}
        curr_rp = {}
        curr_badness = 0
        for plex_id in curr_tube:
            # randomly select one primer pair
            i = floor(rand() * len(all_plex_info[plex_id]['optimize']))
            curr_index[plex_id] = i
            fp = all_plex_info[plex_id]['optimize'][i][0]
            rp = all_plex_info[plex_id]['optimize'][i][1]
            curr_fp[plex_id] = fp['seq']
            curr_rp[plex_id] = rp['seq']
            curr_badness += fp['badness'] + rp['badness']
        inter_badness, comp_badness = PrimerSetBadnessFast(list(curr_fp.values()), list(curr_rp.values()), existing_primer)
        curr_badness += inter_badness
        print('initial loss = %.3f' % curr_badness)

        # optimization
        learning_curve = []
        SATemp = InitSATemp
        for step in range(NUMSTEPS + ZEROSTEPS):
            print('SA temperature = %.3f, SADDLE Loss = %.3f' % (SATemp, curr_badness))
            for t in range(TimePerStep):
                # pick a plex to change primer pair
                mutplex = random.choice(curr_tube) # randomly select a plex_id in current tube
                n_pair = len(all_plex_info[mutplex]['optimize'])
                if n_pair == 1:
                    continue
                mutindex = floor(rand() * n_pair) # index of new primer pair

                # copy current design
                new_index = deepcopy(curr_index)
                new_fp = deepcopy(curr_fp)
                new_rp = deepcopy(curr_rp)
                
                # calculate new badness
                new_index[mutplex] = mutindex
                new_fp[mutplex] = all_plex_info[mutplex]['optimize'][mutindex][0]['seq']
                new_rp[mutplex] = all_plex_info[mutplex]['optimize'][mutindex][1]['seq']
                new_badness = 0
                for plex_id, i in new_index.items():
                    fp = all_plex_info[plex_id]['optimize'][i][0]
                    rp = all_plex_info[plex_id]['optimize'][i][1]
                    new_badness += fp['badness'] + rp['badness']
                inter_badness, new_comp_badness = PrimerSetBadnessFast(list(new_fp.values()), list(new_rp.values()), existing_primer)
                new_badness += inter_badness
                
                # calculate probability of accepting mutation
                if new_badness < curr_badness:
                    accept = True
                else:
                    if SATemp == 0:
                        accept = False
                    elif rand() < 2**((curr_badness - new_badness)/SATemp):
                        accept = True
                    else:
                        accept = False
                        
                # implement acceptance
                if accept:
                    curr_index[mutplex] = new_index[mutplex]
                    curr_fp[mutplex] = new_fp[mutplex]
                    curr_rp[mutplex] = new_rp[mutplex]
                    curr_badness = new_badness
                    comp_badness = new_comp_badness
                    
                learning_curve.append(curr_badness)
            # reduce SATemp
            SATemp = max(0, (SATemp - InitSATemp/NUMSTEPS))
            
        # save optimization result
        for n, (plex_id, i) in enumerate(curr_index.items()):
            all_plex_info[plex_id]['fP'] = all_plex_info[plex_id]['optimize'][i][0]
            all_plex_info[plex_id]['rP'] = all_plex_info[plex_id]['optimize'][i][1]
            all_plex_info[plex_id]['fP_badness'] = comp_badness[0][n]
            all_plex_info[plex_id]['rP_badness'] = comp_badness[1][n]

        all_lc.append(learning_curve)
    
    print('Primer dimer optimization finished in %.3fs' % (time()-tik))
    return all_plex_info, all_lc


def to_df(all_plex_info, config):
    '''
    Format sequences and coordinates as a DataFrame
    Input:
        all_plex_info: output of optimize()
        config: design configurations.
    Output:
        DataFrame of all amplicons
    '''
    print('Formatting design output...')

    l_fP_adp = len(config['fP_prefix']) # length of adapter/prefix
    l_rP_adp = len(config['rP_prefix'])

    plex_name = []
    pool = []

    # sequence specific part of primers
    fp = []
    rp = []

    # coordinates (1-based)
    start = []
    insert_start = []
    insert_end = []
    end = []

    # sequences
    amp = []
    insert = []

    # full primer (with adapters/prefixes)
    fp_full = []
    rp_full = []

    # ARTIC columns
    art_genome = []
    art_start = []
    art_stop = []
    art_primer_name = []
    art_pool = []
    art_strand = []
    art_primer_seq = []

    #print('amplicons containing SNPs or iSNVs in primers:')
    for plex_id, plex_info in all_plex_info.items():
        plex_name.append(plex_id)
        pool.append(plex_info['tube'])
        
        curr_fp = plex_info['fP'] # fP object
        curr_rp = plex_info['rP']

        curr_fp_seq = curr_fp['seq'][l_fP_adp:] # fP sequence
        curr_rp_seq = curr_rp['seq'][l_rP_adp:]
        # if curr_fp_seq != curr_fp_seq.lower() or curr_rp_seq != curr_rp_seq.lower():
        #     print(plex_id)

        fp.append(curr_fp_seq)
        rp.append(curr_rp_seq)

        # calculate coordinates
        pc = plex_info['primers_coords'] # primer design coords ([fp_start, fp_end, rp_start, rp_end])
        curr_start = pc[1] - curr_fp['dist'] - curr_fp['primer_len'] + 1
        curr_insert_start = curr_start + curr_fp['primer_len']
        curr_insert_end = pc[2] + curr_rp['dist'] - 1
        curr_end = curr_insert_end + curr_rp['primer_len']
        start.append(curr_start)
        insert_start.append(curr_insert_start)
        insert_end.append(curr_insert_end)
        end.append(curr_end)

        curr_context_seq = plex_info['context_seq']
        amp.append(curr_context_seq[curr_start-pc[0]:curr_end-pc[0]+1])
        insert.append(curr_context_seq[curr_insert_start-pc[0]:curr_insert_end-pc[0]+1])

        fp_full.append(curr_fp['seq']) # fP with adapter/prefix
        rp_full.append(curr_rp['seq'])

        # ARTIC columns
        # first fP, then rP
        art_genome.extend(['input-seq-1', 'input-seq-1'])
        art_start.extend([curr_start, curr_insert_end+1])
        art_stop.extend([curr_insert_start-1, curr_end])
        art_primer_name.extend([plex_id+'_LEFT', plex_id+'_RIGHT'])
        art_pool.extend([plex_info['tube'], plex_info['tube']])
        art_strand.extend(['+', '-'])
        art_primer_seq.extend([curr_fp_seq.upper(), curr_rp_seq.upper()])

    # format as DataFrame
    df = pd.DataFrame({
        'amplicon_id': plex_name, 
        'pool': pool, 
        'fP': fp, 
        'rP': rp, 
        'start': start, 
        'insert_start': insert_start, 
        'insert_end': insert_end, 
        'end': end, 
        'amplicon': amp, 
        'insert': insert, 
        'fP_full': fp_full, 
        'rP_full': rp_full, 
    })

    # ARTIC DataFrame
    art_df = pd.DataFrame({
        'genome': art_genome, 
        'start': art_start, 
        'stop': art_stop, 
        'primer_name': art_primer_name, 
        'pool': art_pool, 
        'strand': art_strand, 
        'primer_seq': art_primer_seq
    })
    return df, art_df


def tiling(ref_path: str, out_path: str, title: str, max_amp_len: int, min_amp_len: int, 
    w_egc: float, w_lc: float, w_ns: float, w_var: float, temperature: float, salinity: float, 
    dG_max: float, min_GC: float, max_GC: float, min_complexity: float, max_len: int, 
    check_var: bool, fP_prefix: str, rP_prefix: str, seed: int, threads: int):
    '''
    Design tiled amplicons. 
    Input:
        ref_path: Path to the Olivar reference file (.olvr).
        out_path: Output directory [./].
        title: Name of this design [olivar-design].
        max_amp_len: Maximum amplicon length [420].
        min_amp_len: Minimum amplicon length. 0.9*{max-amp-len} if set as None. Minimum 120.
        w_egc: Weight for extreme GC content [1.0].
        w_lc: Weight for low sequence complexity [1.0].
        w_ns: Weight for non-specificity [1.0].
        w_var: Weight for variations [1.0].
        temperature: PCR annealing temperature [60.0].
        salinity: Concentration of monovalent ions in units of molar [0.18].
        dG_max: Maximum free energy change of a primer in kcal/mol [-11.8].
        min_GC: Minimum GC content of a primer [0.2].
        max_GC: Maximum GC content of a primer [0.75].
        min_complexity: Minimum sequence complexity of a primer [0.4].
        max_len: Maximum length of a primer [36].
        check_var: Filter out primer candidates with variations within 5nt of 3' end [False]. 
            Setting check_var=True is not recommended when a lot of variations are provided, 
            since this would significantly reduce the number of primer candidates. 
        fP_prefix: Prefix of forward primer. Empty string '' by default.
        rP_prefix: Prefix of reverse primer. Empty string '' by default.
        seed: Random seed for optimizing primer design regions and primer dimer [10].
        threads: Number of threads [1].
    '''
    config = {
        'ref_path': ref_path, 
        'out_path': out_path, 
        'title': title, 
        'max_amp_len': max_amp_len, 
        'min_amp_len': min_amp_len, 
        'w_egc': w_egc, 
        'w_lc': w_lc, 
        'w_ns': w_ns, 
        'w_var': w_var, 
        'temperature': temperature, 
        'salinity': salinity, 
        'dG_max': dG_max, 
        'min_GC': min_GC, 
        'max_GC': max_GC, 
        'min_complexity': min_complexity, 
        'max_len': max_len, 
        'check_SNP': check_var, 
        'fP_prefix': fP_prefix, 
        'rP_prefix': rP_prefix, 
        'seed': seed, 
        'threads': threads
    }

    if not os.path.exists(config['out_path']):
        os.makedirs(config['out_path'])

    all_plex_info, risk_arr, gc_arr, comp_arr, hits_arr, var_arr, all_loss, seq_record = design_context_seq(config)
    all_plex_info_primer = get_primer(all_plex_info, config)
    all_plex_info_optimize, learning_curve = optimize(all_plex_info_primer, config)
    df, art_df = to_df(all_plex_info_optimize, config)

    # format and save design output
    design_out = {
        'config': config, 
        'risk_arr': risk_arr, 
        'gc_arr': gc_arr, 
        'comp_arr': comp_arr, 
        'hits_arr': hits_arr, 
        'var_arr': var_arr, 
        'all_loss': all_loss, 
        'seq': seq_record, # Bio.SeqRecord
        'df': df, 
        'art_df': art_df, # ARTIC output format
        'all_plex_info': all_plex_info_optimize, 
        'learning_curve': learning_curve
    }
    design_path = os.path.join(config['out_path'], '%s.olvd' % config['title'])
    with open(design_path, 'wb') as f:
        pickle.dump(design_out, f, protocol=5) # protocol 5 needs python>=3.8
        print('Design file saved as %s' % design_path)

    # save human readable files
    save(design_out, config['out_path'])


def save(design_out, out_path: str):
    '''
    Load from a previous Olivar design object (in tiling()) or file (.olvd) and save output files.
    Input:
        design_out: Output of tiling(), or load the olvd file with pickle.
        out_path: Output directory.
    '''
    # load data
    if type(design_out) is str:
        print(f'Loading Olivar design from {design_out}...')
        try:
            with open(design_out,  'rb') as f:
                design_out = pickle.load(f)
                print(f'Successfully loaded Olivar design.')
        except FileNotFoundError:
            raise FileNotFoundError('Olivar design file not found.')
    
    if not os.path.exists(out_path):
        os.makedirs(out_path)
        
    config = design_out['config']
    risk_arr = design_out['risk_arr']
    gc_arr = design_out['gc_arr']
    comp_arr = design_out['comp_arr']
    hits_arr = design_out['hits_arr']
    var_arr = design_out['var_arr']
    all_loss = design_out['all_loss']
    lc = design_out['learning_curve']
    df = design_out['df']

    # column name mapper for backward compatibility
    df.rename(columns={
        'amp_id': 'amplicon_id', 
        'amp': 'amplicon'
    }, inplace=True)

    # save configurations
    save_path = os.path.join(out_path, '%s.json' % config['title'])
    with open(save_path, 'w') as f:
        json.dump(config, f, indent=4)
    print(f'Configurations saved as {save_path}')

    # save sequences and coordinates as csv
    save_path = os.path.join(out_path, '%s.csv' % config['title'])
    df.to_csv(save_path, index=False)
    print(f'Primer pools saved as {save_path}')

    # save ARTIC output format
    try:
        art_df = design_out['art_df']
        save_path = os.path.join(out_path, '%s.scheme.bed' % config['title'])
        art_df.sort_values(by='start', axis=0, ascending=True, inplace=True)
        art_df.to_csv(save_path, sep='\t', header=False, index=False)
        print(f'Primer pools (ARTIC format) saved as {save_path}')
    except KeyError:
        print('ARTIC/PrimalScheme format is not included in older versions of Olivar, skipped.')

    # save reference sequence
    save_path = os.path.join(out_path, '%s.fasta' % config['title'])
    with open(save_path, 'w') as f:
        SeqIO.write([design_out['seq']], f, 'fasta')
    print(f'Reference sequence saved as {save_path}')

    # save risk array
    risk = pd.DataFrame({
        'position': range(1, len(risk_arr)+1), 
        'base': list(design_out['seq'].seq), 
        'extreme GC': gc_arr, 
        'low complexity': comp_arr, 
        'non-specificity': hits_arr, 
        'variations': var_arr, 
        'risk': risk_arr
    })
    save_path = os.path.join(out_path, '%s_risk.csv' % config['title'])
    risk.to_csv(save_path, index=False)
    print(f'Risk scores saved as {save_path}')

    #------------- Loss and SADDLE Loss -------------#
    fig = make_subplots(
        rows=2, cols=2, 
        subplot_titles=('Optimization of primer design regions', 
            '', 
            'Primer dimer optimization (pool-1)', 
            'Primer dimer optimization (pool-2)')
    )

    fig.add_trace(
        go.Scatter(
            y=all_loss, 
            line=dict(color='#1f77b4'), 
            showlegend=False
        ), 
        row=1, col=1
    )
    fig.add_trace(
        go.Scatter(
            y=lc[0], 
            line=dict(color='#1f77b4'), 
            showlegend=False
        ), 
        row=2, col=1
    )
    fig.add_trace(
        go.Scatter(
            y=lc[1], 
            line=dict(color='#1f77b4'), 
            showlegend=False
        ), 
        row=2, col=2
    )

    fig.update_xaxes(title_text='iterations (sorted)', row=1, col=1)
    fig.update_xaxes(title_text='iterations', row=2, col=1)
    fig.update_xaxes(title_text='iterations', row=2, col=2)

    fig.update_yaxes(title_text='Loss', type='log', row=1, col=1)
    fig.update_yaxes(title_text='SADDLE Loss', row=2, col=1)
    fig.update_yaxes(title_text='SADDLE Loss', row=2, col=2)

    # save html figure
    save_path = os.path.join(out_path, '%s_Loss.html' % config['title'])
    with open(save_path, 'w') as f:
        f.write(plotly.io.to_html(fig))
    print(f'Optimization plots saved as {save_path}')
    #------------- Loss and SADDLE Loss -------------#

    #------------- risk array and primers -------------#
    # create figure
    fig = go.Figure()

    # set figure scale
    base_offset = (0.5*config['w_egc'] + 0.5*config['w_lc'] + config['w_ns'])/6

    # plot risk array
    r = np.arange(len(risk_arr))
    fig.add_trace(
        go.Scatter(
            x=r, y=gc_arr,
            hoverinfo='skip',
            mode='lines',
            line=dict(width=0, color='#1f77b4'),
            name='extreme GC',
            stackgroup='one' # define stack group
        )
    )
    fig.add_trace(
        go.Scatter(
            x=r, y=comp_arr,
            hoverinfo='skip',
            mode='lines',
            line=dict(width=0, color='#ff7f0e'),
            name='low complexity',
            stackgroup='one'
        )
    )
    fig.add_trace(
        go.Scatter(
            x=r, y=hits_arr,
            hoverinfo='skip',
            mode='lines',
            line=dict(width=0, color='#2ca02c'),
            name='non-specificity',
            stackgroup='one'
        )
    )
    fig.add_trace(
        go.Scatter(
            x=r, y=var_arr,
            hoverinfo='skip',
            mode='lines',
            line=dict(width=0, color='#d62728'),
            name='variations',
            stackgroup='one'
        )
    )

    # plot primers
    fp_rp_diff = 0.1 # distance between fP and rP
    head_dx = 7 # arrow head length
    head_dy = base_offset*7/90 # arrow head height
    primer_x = [] # x coords for primer plot
    primer_y = [] # y coords for primer plot
    hover_text = []

    # pool 1
    fp_offset = 2 * base_offset
    rp_offset = (2-fp_rp_diff) * base_offset
    for i, row in df[df['pool']==1].iterrows():
        fp_start = row['start']-1
        fp_stop = row['insert_start']-1
        rp_start = row['insert_end']
        rp_stop = row['end']
        # [fp_start, fp_stop] and [fp_offset, fp_offset] is the body of fP
        # [fp_stop-head_dx, fp_stop] and [fp_offset+head_dy, fp_offset] is the head of fP
        primer_x.extend([fp_start, fp_stop, None, fp_stop-head_dx, fp_stop, None, 
            rp_start, rp_stop, None, rp_start, rp_start+head_dx, None])
        primer_y.extend([fp_offset, fp_offset, None, fp_offset+head_dy, fp_offset, None, 
            rp_offset, rp_offset, None, rp_offset, rp_offset-head_dy, None])
        hover_text.extend(['pool-1 %s fP: %d - %d' % (row['amplicon_id'], row['start'], fp_stop)]*6 + \
            ['pool-1 %s rP: %d - %d' % (row['amplicon_id'], rp_start+1, rp_stop)]*6)
    
    # pool 2
    fp_offset = base_offset
    rp_offset = (1-fp_rp_diff) * base_offset
    for i, row in df[df['pool']==2].iterrows():
        fp_start = row['start']-1
        fp_stop = row['insert_start']-1
        rp_start = row['insert_end']
        rp_stop = row['end']
        # [fp_start, fp_stop] and [fp_offset, fp_offset] is the body of fP
        # [fp_stop-head_dx, fp_stop] and [fp_offset+head_dy, fp_offset] is the head of fP
        primer_x.extend([fp_start, fp_stop, None, fp_stop-head_dx, fp_stop, None, 
            rp_start, rp_stop, None, rp_start, rp_start+head_dx, None])
        primer_y.extend([fp_offset, fp_offset, None, fp_offset+head_dy, fp_offset, None, 
            rp_offset, rp_offset, None, rp_offset, rp_offset-head_dy, None])
        hover_text.extend(['pool-2 %s fP: %d - %d' % (row['amplicon_id'], row['start'], fp_stop)]*6 + \
            ['pool-2 %s rP: %d - %d' % (row['amplicon_id'], rp_start+1, rp_stop)]*6)
    
    # plot all primers
    fig.add_trace(
        go.Scatter(
            x=primer_x, y=primer_y, 
            mode='lines', # connect points with lines
            connectgaps=False, # gaps (np.nan or None) in the provided data arrays are not connected
            line=dict(color='rgb(0,0,0)'), # black
            showlegend=False, 
            hovertemplate='%{text}<extra></extra>', # <extra></extra> hides the "trace 4"
            text=hover_text
        )
    )

    # title and axis
    fig.update_layout(
        hoverlabel=dict(
            bgcolor='white'
        ), 
        # title
        title_text="risk components are stacked together", 
        title_x=0.98, 
        titlefont=dict(
            size=18,
        ), 

        # axis
        xaxis=dict(
            tickformat='%d', 
            tickfont=dict(
                size=14, 
            ), 
            showgrid=False, 
            rangeslider=dict(
                visible=True
            )
        ), 
        yaxis=dict(
            range=[0, base_offset*6], 
            tickfont=dict(
                size=14, 
            ), 
            showgrid=False
        ), 

        # axis title
        xaxis_title='position', 
        yaxis_title='risk', 

        # global font
        font=dict(
            size=16,
        )
    )

    # save html figure
    save_path = os.path.join(out_path, '%s.html' % config['title'])
    with open(save_path, 'w') as f:
        f.write(plotly.io.to_html(fig))
    print(f'Risk and primer viewer saved as {save_path}')
    #------------- risk array and primers -------------#


def validate(primer_pool: str, pool: int, BLAST_db: str, out_path: str, 
    title: str, max_amp_len: int, temperature: float, threads: int):
    '''
    Run analysis on a primer pool for multiplex PCR
    Input:
        primer_pool: Path to the csv file of a primer pool. 
            Required columns: "amplicon_id" (amplicon name), "fP" (sequence of forward primer), "rP" (sequence of reverse primer), "pool" (pool number, e.g., 1).
        pool: Primer pool number [1]. 
        BLAST_db: Optional, path to the BLAST database. 
            Note that this path should end with the name of the BLAST database (e.g., "example_input/Human/GRCh38_primary").
        out_path: Output directory [./]. 
        title: Name of validation [olivar-val].
        max_amp_len: Maximum length of predicted non-specific amplicon [1500]. 
        temperature: PCR annealing temperature [60.0]. 
        threads: Number of threads [1]. 
    '''
    if not os.path.exists(out_path):
        os.makedirs(out_path)

    seq_list = []
    seq_names = []
    df = pd.read_csv(primer_pool, sep=',', index_col=False)

    # column name mapper for backward compatibility
    df.rename(columns={
        'amp_id': 'amplicon_id', 
        'amp': 'amplicon'
    }, inplace=True)
    
    for i, row in df[df['pool']==pool].iterrows():
        seq_list.extend([row['fP'], row['rP']])
        seq_names.extend(['%s_fP' % row['amplicon_id'], '%s_rP' % row['amplicon_id']])
    print(f'Successfully loaded {primer_pool}, pool-{pool}.')

    # non-specific simulation
    if BLAST_db:
        df_ns, df_count, all_hits = ns_simulation(BLAST_db, seq_list, seq_names, max_amp_len, threads)
        save_path = os.path.join(out_path, f'{title}_pool-{pool}_ns-amp.csv')
        df_ns.to_csv(save_path, index=False)
        print('Non-specific amplicons saved as %s' % save_path)
        save_path = os.path.join(out_path, f'{title}_pool-{pool}_ns-pair.csv')
        df_count.to_csv(save_path, index=False)
        print('Non-specific primer pairs saved as %s' % save_path)
    else:
        all_hits = np.zeros(len(seq_list)) - 1
        print('No BLAST database provided, skipped.')

    # calculate Badness
    #all_rep = [seq2arr(p) for p in seq_list]
    _, all_bad = PrimerSetBadnessFast(seq_list)
    all_bad = all_bad[0]

    # calculate dG
    gen = design.primer_generator(temperature=temperature, salinity=0.18)
    all_dG = [gen.dG_init+gen.StacksDG(p) for p in seq_list]

    df_val = pd.DataFrame({
        'name': seq_names, 
        'seq': seq_list, 
        'length': [len(p) for p in seq_list], 
        '%GC': [basic.get_GC(p) for p in seq_list], 
        'complexity': [basic.get_complexity(p) for p in seq_list], 
        'dG (%dC)' % temperature: all_dG, 
        'self_align': [design.WAlignScore(p) for p in seq_list], 
        'dimer_score': all_bad, 
        'BLAST_hits': all_hits
    })
    save_path = os.path.join(out_path, f'{title}_pool-{pool}.csv')
    df_val.to_csv(save_path, index=False)
    print('Validation file saved as %s' % save_path)

src/olivar/main.py
treangenlab/Olivar
#!/usr/bin/env python3
# -*- coding: utf-8 -*-


'''
Wrappers for NCBI BLAST+
Architecture:
ncbi_tools.py
    ns_simulation()
        BLAST_batch_short()
'''


__author__ = 'Michael X. Wang'


import os
CURRENT_DIR = os.path.dirname(__file__)


import copy
from time import time

import pandas as pd
from Bio.Blast import NCBIXML
from Bio.Blast.Applications import NcbiblastnCommandline
from tqdm import tqdm


def BLAST_batch_short(seq_list: list, db: str, n_cpu=1, seq_names=None, mode='rough'):
    '''
    Run local BLAST on a list of short sequences and return a corresponding list of hits. 
    Input:
        seq_list: list of sequences
        db: path to BLAST database
        n_cpu: number of threads to use
        seq_names: Default None. If specified, should be a list of sequence names corresponding to seq_list
        mode: 'precise' or 'rough'
    '''
    BATCH_SIZE = 128
    N = len(seq_list)

    BLAST_input = os.path.join(CURRENT_DIR, 'BLAST_query.fasta')
    if mode == 'precise':
        BLAST_output = os.path.join(CURRENT_DIR, 'BLAST_result.xml')
    elif mode == 'rough':
        BLAST_output = os.path.join(CURRENT_DIR, 'BLAST_result.tsv')
    else:
        raise ValueError('Invalid mode. mode should be either "precise" or "rough"')

    if seq_names is None:
        seq_names = ['query_%d' % i for i in range(N)]
    if N != len(seq_names):
        raise ValueError('length of seq_names should be the same as seq_list.')
    
    print(f'Running blastn-short with {n_cpu} CPU(s), batch size {BATCH_SIZE}...')
    tik = time()
    pbar = tqdm(total=N, ascii=' >') # progress bar

    n_batch = 0
    if mode == 'precise':
        all_count_batch = []
    elif mode == 'rough':
        all_count_batch = [0] * N
    all_loc_batch = []
    while BATCH_SIZE*n_batch < N:
        batch_start = BATCH_SIZE*n_batch
        batch_stop = min(BATCH_SIZE*(n_batch+1), N)
        seq_batch = seq_list[batch_start: batch_stop]
        name_batch = seq_names[batch_start: batch_stop]
        n_batch += 1
        #print('%d/%d' % (batch_start, N))

        # prepare fasta file for query
        with open(BLAST_input, 'w') as f:
            for name, seq in zip(name_batch, seq_batch):
                f.write('>%s\n%s\n' % (name, seq))
        
        # run blast
        if mode == 'precise':
            # referred to Primer-BLAST paper
            config = {
                'outfmt': 5, # xml
                'evalue': 5000, 
                'reward': 1, 
                'penalty': -1, 
                'gapopen': 2, 
                'gapextend': 1
            }
        elif mode == 'rough':
            # default of blastn-short
            config = {
                'outfmt': 6, # tabular
                'evalue': 10, 
                'reward': 1, 
                'penalty': -3, 
                'gapopen': 5, 
                'gapextend': 2
            }
        blastn_cline = NcbiblastnCommandline(query=BLAST_input, 
            db=db, 
            out=BLAST_output, 
            task='blastn-short', 
            num_threads=n_cpu, 
            **config)
        stdout, stderr = blastn_cline()
        
        # parse result file
        with open(BLAST_output, 'r') as handle:
            if mode == 'precise':
                all_records = NCBIXML.parse(handle)
                all_count = []
                all_loc = []
                for record in all_records:
                    # single BLAST query
                    count = 0
                    loc_dict = {}
                    qlen = record.query_letters
                    for alignment in record.alignments:
                        loc_list = []
                        for hsp in alignment.hsps:
                            if (hsp.query_end == qlen) and \
                            (hsp.match[-5:].count(' ')<2) and \
                            (hsp.identities>qlen-5):
                                # no 3' overhang
                                # no more than 1 mismatch within 5nt at 3'end
                                # no more than 4 total mismatches
                                loc_list.append((hsp.sbjct_start, hsp.sbjct_end, record.query))
                                count += 1
                        if loc_list:
                            loc_dict[alignment.accession] = loc_list
                    all_count.append(count)
                    all_loc.append(loc_dict)
                all_count_batch.extend(all_count)
                all_loc_batch.extend(all_loc)
            elif mode == 'rough':
                '''
                1. qseqid      query or source (e.g., gene) sequence id
                2. sseqid      subject  or target (e.g., reference genome) sequence id
                3. pident      percentage of identical matches
                4. length      alignment length (sequence overlap)
                5. mismatch    number of mismatches
                6. gapopen     number of gap openings
                7. qstart      start of alignment in query
                8. qend        end of alignment in query
                9. sstart      start of alignment in subject
                10. send        end of alignment in subject
                11. evalue      expect value
                12. bitscore    bit score
                '''
                # get first hsp
                hsp = handle.readline()[:-1].split('\t')
                query_name = hsp[0]
                count = 1
                for line in handle:
                    hsp = line[:-1].split('\t')
                    if hsp[0] == query_name:
                        count += 1
                    else:
                        # next query sequence
                        all_count_batch[int(query_name.split('_')[1])] = count
                        query_name = hsp[0]
                        count = 1
                all_count_batch[int(query_name.split('_')[1])] = count # save count of last record

        pbar.update(batch_stop-batch_start)
    pbar.close()

    print('finished in %.2fs' % (time()-tik))
    return all_count_batch, all_loc_batch


def ns_simulation(BLAST_db: str, seq_list: list, seq_names: list, max_amp_len=1500, n_cpu=1):
    '''
    Predict all possible non-specific amplification for primers in seq_list
    Input:
        BLAST_db: path to local BLAST database
        seq_list: list of primers
        seq_names: list of primer names
        max_amp_len: maximum amplicon length
        n_cpu: number of threads to use
    return: 
        dataframe of ns amplicons
        dataframe of ns primer pairs
    '''
    all_hits, all_loc = BLAST_batch_short(seq_list, db=BLAST_db, n_cpu=n_cpu, seq_names=seq_names, mode='precise')

    # merge hits within each alignment (chromosome)
    loc_merged = {}
    for loc_dict in all_loc:
        for acc, loc_list in loc_dict.items():
            try:
                loc_merged[acc].extend(copy.copy(loc_list))
            except KeyError:
                # copy to prevent modifying all_loc
                loc_merged[acc] = copy.copy(loc_list)

    # count legitimate primer pairs within each alignemnt (chromosome)
    pair_dict = {}
    for acc, loc_list in loc_merged.items():
        loc_list.sort(key=lambda x: x[0])
        list_len = len(loc_list)
        for i, loc in enumerate(loc_list[:-1]):
            j = i + 1
            while (j < list_len) and (loc_list[j][0] - loc[0] < max_amp_len):
                # amplicon length within range
                if (loc[0] < loc[1]) and (loc_list[j][1] < loc_list[j][0]):
                    # primer orientation is correct
                    pair = (loc[2], loc_list[j][2]) # primer names
                    try:
                        pair_dict[pair].append((acc, loc, loc_list[j]))
                    except KeyError:
                        pair_dict[pair] = [(acc, loc, loc_list[j])]
                j += 1
    
    # output each non-specific amplicon
    df_ns = {
        'fP': [], 
        'rP': [], 
        'acc': [], 
        'fP_start': [], 
        'fP_stop': [], 
        'rP_start': [], 
        'rP_stop': []
    }
    df_count = {
        'fP': [], 
        'rP': [], 
        'hits': []
    }
    for pair, loc_list in pair_dict.items():
        df_count['fP'].append(pair[0])
        df_count['rP'].append(pair[1])
        df_count['hits'].append(len(loc_list))
        for loc in loc_list:
            df_ns['fP'].append(pair[0])
            df_ns['rP'].append(pair[1])
            df_ns['acc'].append(loc[0])
            df_ns['fP_start'].append(loc[1][0])
            df_ns['fP_stop'].append(loc[1][1])
            df_ns['rP_start'].append(loc[2][0])
            df_ns['rP_stop'].append(loc[2][1])
    df_ns = pd.DataFrame(df_ns)
    df_ns['amp_len'] = df_ns['rP_start'] - df_ns['fP_start'] + 1
    return df_ns, pd.DataFrame(df_count), all_hits

src/olivar/ncbi_tools.py
treangenlab/Olivar