test_final_multithreading.py

import pysam
import concurrent.futures
import sys
import getopt
import logging


#Options
output_file = ""
try:
    opts, args = getopt.getopt(sys.argv[1:], 'b:w:s:t:o:e:')
    for opt, arg in opts:
        if opt in ("-b"):
            bamfile = arg
        if opt in ("-w"):
            window_size = int(arg)
        if opt in ("-s"):
            step_size = int(arg)
        if opt in ("-t"):
            num_threads = int(arg)
        if opt in ("-o"):
            output_file = open(arg, "a")
        if opt in ("-e"):
            logfile = arg
except getopt.GetoptError:
    print('Invalid argument')
    sys.exit(1)
    
logging.basicConfig(filename =  '%s' % (logfile), filemode = 'a', level = logging.INFO, format = '%(asctime)s %(levelname)s - %(message)s')
logging.info('start')

def calcul_depth_seq(bamfile, seq, window_size, step_size, numero_de_thread):
    sys.stderr.write("\t entering calcul_depth_seq : thread numero %s\n" % numero_de_thread)
    samfile = pysam.AlignmentFile(bamfile, "rb")
    seq_length = samfile.lengths[samfile.references.index(seq)] #Permet d'obtenir la longueur de la sequence

    for pos_start in range(0, seq_length - window_size +1, step_size): #Itère sur la sequence par "fenêtre" en fonction de la taille
        pos_end = pos_start + window_size #pos_start + window_size donne la position de début de la fenêtre, puis grâce à min() et seq_length je peux définir position de fin. le min me permet d'être certain que pos_end < seq_lenght. C-a-d que si la fenêtre dépasse la taille de la sequence, j'aurais quand même pos_end
        window_depths = [] #liste pour stocker la profondeur de chaque fenêtre

        for pileupcolumn in samfile.pileup(seq, start=pos_start, stop=pos_end): #itère sur chaque colone de pileup
            window_depths.append(pileupcolumn.n)#ajoute le nb de lecture couvrant cette position à la liste window_depths

        if window_depths:#calcul la profondeur moyenne pour chaque sequence
            sys.stderr.write("\t\t entering wd thread numero %s\n" % numero_de_thread)
            avg_depth = sum(window_depths) / len(window_depths)
            sys.stdout.write(f"{seq}:{pos_start}-{pos_end}: {avg_depth}\n")            
            output_file.write(f"{seq}:{pos_start}-{pos_end}: {avg_depth}\n")
           
    sys.stderr.write("\t leaving calcul_depth_seq thread numero %s\n" % numero_de_thread)
    
def calcul_depth_all(bamfile, window_size, step_size, num_threads):
    sys.stderr.write("entering calcul_depth_all \n")
    samfile=pysam.AlignmentFile(bamfile, "rb")
    
    with concurrent.futures.ThreadPoolExecutor(max_workers = num_threads) as executor:
        futures = [] #liste pour stocker les futures
        n=0
        for seq in samfile.references:
            n+=1
            future = executor.submit(calcul_depth_seq, bamfile, seq, window_size, step_size, n)
            futures.append(future)
            
        concurrent.futures.wait(futures)
        
        samfile.close()

#if output_file != sys.stdout:
 #   output_file.close()

#Programme principal
if len(sys.argv) <2 or bamfile is None:
    print("Usage : File alignement .bam")
    sys.exit(1)

bamfile_path = bamfile
calcul_depth_all(bamfile_path, window_size, step_size, num_threads)
logging.info('end')
samfile.close()