annotate 0.4.2/bin/process_centrifuge_output.py @ 143:620bffa66bbb tip

planemo upload
author kkonganti
date Thu, 11 Jul 2024 14:48:07 -0400
parents 52045ea4679d
children
rev   line source
kkonganti@105 1 #!/usr/bin/env python
kkonganti@105 2
kkonganti@105 3 import os
kkonganti@105 4 import argparse
kkonganti@105 5 import logging as log
kkonganti@105 6 import pandas as pd
kkonganti@105 7 import numpy as np
kkonganti@105 8 from Bio import SeqIO
kkonganti@105 9
kkonganti@105 10
kkonganti@105 11 def main():
kkonganti@105 12 # READ IN ARGUMENTS
kkonganti@105 13 desc = """
kkonganti@105 14 This script is part of the centriflaken pipeline: It processes centrifuge
kkonganti@105 15 output and produces either a filtered FASTQ or a text file of FASTQ IDs based
kkonganti@105 16 on the supplied taxa/bug
kkonganti@105 17 """
kkonganti@105 18 parser = argparse.ArgumentParser(prog='process_centrifuge_output.py', description=desc)
kkonganti@105 19 parser.add_argument("-v", dest='verbose', action="store_true", help="For more verbose output")
kkonganti@105 20 parser.add_argument("-i", dest='input_fastq', required=False,
kkonganti@105 21 help="Path to input FASTQ file (same as input to centrifuge). If not mentioned, \
kkonganti@105 22 a text file of sequence IDs are produced instead of a FASTQ file")
kkonganti@105 23 parser.add_argument("-t", dest='taxa_filtered_fastq_file', required=True,
kkonganti@105 24 help="Path to output FASTQ or output text file filtered by the taxa specified")
kkonganti@105 25 parser.add_argument("-r", dest='cent_report', required=True, help="Path to centrifuge report")
kkonganti@105 26 parser.add_argument("-o", dest='cent_output', required=True, help="Path to centrifuge output")
kkonganti@105 27 parser.add_argument("-b", dest='bug', required=True,
kkonganti@105 28 help="Name or fragment of name of the bug by which reads are extracted")
kkonganti@105 29 args = parser.parse_args()
kkonganti@105 30
kkonganti@105 31 # MORE INFO IF VERBOSE
kkonganti@105 32 if args.verbose:
kkonganti@105 33 log.basicConfig(format="%(levelname)s: %(message)s", level=log.DEBUG)
kkonganti@105 34 else:
kkonganti@105 35 log.basicConfig(format="%(levelname)s: %(message)s")
kkonganti@105 36
kkonganti@105 37 # ASSIGN VARIABLES
kkonganti@105 38 input_fastq = args.input_fastq
kkonganti@105 39 taxa_filtered_fastq_file = args.taxa_filtered_fastq_file
kkonganti@105 40 cent_report = args.cent_report
kkonganti@105 41 cent_output = args.cent_output
kkonganti@105 42 bug = args.bug
kkonganti@105 43 report_col_list = ["name", "taxID"]
kkonganti@105 44 output_col_list = ["taxID", "readID"]
kkonganti@105 45
kkonganti@105 46 # Match and filter taxa names and ids from centrifuge report file
kkonganti@105 47 report_df = pd.read_csv(cent_report, delimiter="\t", usecols=report_col_list)
kkonganti@105 48 report_df['name'] = report_df['name'].str.lower()
kkonganti@105 49 filt_report_df = report_df[report_df['name'].str.contains(bug.lower())]
kkonganti@105 50 #print("\nMatching taxa names and ids:\n",filt_report_df)
kkonganti@105 51 taxID_list = filt_report_df['taxID']
kkonganti@105 52
kkonganti@105 53 # Match the above tax ids to read ids from centrifuge output file and deduplicate
kkonganti@105 54 output_df = pd.read_csv(cent_output, delimiter="\t", usecols=output_col_list)
kkonganti@105 55 filt_output_df = output_df.loc[output_df['taxID'].isin(taxID_list)]
kkonganti@105 56 readID_list = filt_output_df['readID']
kkonganti@105 57 readID_dedup_list = np.unique(readID_list)
kkonganti@105 58 TF=open(taxa_filtered_fastq_file, "w")
kkonganti@105 59
kkonganti@105 60 if (not input_fastq):
kkonganti@105 61 # print("\nFILTERED READ ID LIST:\n", readID_dedup_list)
kkonganti@105 62 for ID in readID_dedup_list:
kkonganti@105 63 TF.write(f"{ID}\n")
kkonganti@105 64 else:
kkonganti@105 65 # Extract filtered reads from input fastq and write to output fastq
kkonganti@105 66 print ("Indexing reads..")
kkonganti@105 67 rec = SeqIO.index(input_fastq,"fastq")
kkonganti@105 68 for i in readID_dedup_list:
kkonganti@105 69 if i in rec:
kkonganti@105 70 SeqIO.write(rec[i], TF, "fastq")
kkonganti@105 71
kkonganti@105 72 TF.close()
kkonganti@105 73
kkonganti@105 74 if __name__ == "__main__":
kkonganti@105 75 main()