Merge pull request #167 from EBI-Metagenomics/bugfix/write-gff-missing-contigs-when-user-provided-proteins

Bugfix/write gff missing contigs when user provided proteins

Author: mberacochea

.editorconfig

@@ -19,4 +19,4 @@ indent_size = 2
 indent_size = 2
 
 [*.nf]
-indent_size = 2
+indent_size = 4

bin/contig_taxonomic_assign.py

@@ -2,6 +2,7 @@
 
 import argparse
 import csv
+import logging
 import os
 import re
 import math
@@ -10,38 +11,45 @@
 import pandas as pd
 from ete3 import NCBITaxa
 
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s %(levelname)s %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+)
+
 
 # Some taxa are discontinued, we should exclude them (https://github.com/EBI-Metagenomics/emg-viral-pipeline/issues/113):
-EXCLUDE_TAXA = ["Allolevivirus",
-                "Autographivirinae",
-                "Buttersvirus",
-                "Caudovirales",
-                "Chungbukvirus",
-                "Incheonvirus",
-                "Leviviridae",
-                "Levivirus",
-                "Mandarivirus",
-                "Pbi1virus",
-                "Phicbkvirus",
-                "Radnorvirus",
-                "Sitaravirus",
-                "Vidavervirus",
-                "Myoviridae",
-                "Siphoviridae",
-                "Podoviridae",
-                "Viunavirus",
-                "Orthohepevirus",
-                "Klosneuvirus",
-                "Hendrixvirus",
-                "Rubulavirus",
-                "Avulavirus",
-                "Catovirus",
-                "Nucleorhabdovirus",
-                "Viunavirus",
-                "Gammalipothrixvirus",
-                "Peduovirinae",
-                "Sedoreovirinae"
-                ]
+EXCLUDE_TAXA = [
+    "Allolevivirus",
+    "Autographivirinae",
+    "Buttersvirus",
+    "Caudovirales",
+    "Chungbukvirus",
+    "Incheonvirus",
+    "Leviviridae",
+    "Levivirus",
+    "Mandarivirus",
+    "Pbi1virus",
+    "Phicbkvirus",
+    "Radnorvirus",
+    "Sitaravirus",
+    "Vidavervirus",
+    "Myoviridae",
+    "Siphoviridae",
+    "Podoviridae",
+    "Viunavirus",
+    "Orthohepevirus",
+    "Klosneuvirus",
+    "Hendrixvirus",
+    "Rubulavirus",
+    "Avulavirus",
+    "Catovirus",
+    "Nucleorhabdovirus",
+    "Viunavirus",
+    "Gammalipothrixvirus",
+    "Peduovirinae",
+    "Sedoreovirinae",
+]
 
 
 def main(args):
@@ -62,18 +70,40 @@ def main(args):
             for name, *_, avg_cds, std_cds, _, mult_factor in csv_reader
         }
 
-    file_header = ["contig_ID", "superkingdom", "kingdom", "phylum", "subphylum", "class", "order", "suborder", "family", "subfamily", "genus"]
+    file_header = [
+        "contig_ID",
+        "superkingdom",
+        "kingdom",
+        "phylum",
+        "subphylum",
+        "class",
+        "order",
+        "suborder",
+        "family",
+        "subfamily",
+        "genus",
+    ]
 
     exclude_deprecated_taxa = False
     if args.version4:
         exclude_deprecated_taxa = True
 
-    output_gen = contig_tax(input_df, args.ncbi_db, args.tax_thres, factor_dict, file_header, exclude_deprecated_taxa)
+    output_gen = contig_tax(
+        input_df,
+        args.ncbi_db,
+        args.tax_thres,
+        factor_dict,
+        file_header,
+        exclude_deprecated_taxa,
+    )
 
-    print(args.input_file)
+    logging.info(f"Processing input file: {args.input_file}")
+    logging.info(
+        f"Settings: tax_thres={args.tax_thres:.2f}, version4={args.version4}, exclude_deprecated={exclude_deprecated_taxa}"
+    )
 
     out_file = re.split(r"\.[a-z]+$", os.path.basename(args.input_file))[0]
-    
+
     if not os.path.exists(args.outdir):
         os.mkdir(args.outdir)
     with open(
@@ -85,44 +115,71 @@ def main(args):
             tsv_writer.writerow(item)
 
 
-def contig_tax(annot_df, ncbi_db, tax_thres, taxon_factor_dict, output_taxa_order, exclude_deprecated_taxa=False):
+def contig_tax(
+    annot_df,
+    ncbi_db,
+    tax_thres,
+    taxon_factor_dict,
+    output_taxa_order,
+    exclude_deprecated_taxa=False,
+):
     """This function takes the annotation table generated by viral_contig_maps.py and generates a table that
     provides the taxonomic lineage of each viral contig, based on the corresponding ViPhOG annotations"""
 
     ncbi = NCBITaxa(dbfile=ncbi_db)
     viphog_rank = ["genus", "subfamily", "family", "order"]
     contig_set = set(annot_df["Contig"])
 
+    logging.info(f"Assigning taxonomy for {len(contig_set)} contigs")
+
+    assigned = 0
+    unassigned_no_hits = 0
+    unassigned_below_thres = 0
+
     for contig in contig_set:
         contig_lineage = []
         contig_df = annot_df[annot_df["Contig"] == contig]
         total_prot = len(contig_df)
         annot_prot = sum(contig_df["Best_hit"] != "No hit")
         if annot_prot == 0:
+            logging.debug(f"Contig {contig}: no ViPhOG hits ({total_prot} proteins) - skipping")
+            unassigned_no_hits += 1
             contig_lineage.extend([""] * len(output_taxa_order[1:]))
         else:
+            logging.debug(f"Contig {contig}: {annot_prot}/{total_prot} proteins with ViPhOG hits")
             contig_hits = contig_df[pd.notnull(contig_df["Label"])]["Label"].values
             taxid_list = []
             for item in contig_hits:
                 if len(ncbi.get_name_translator([item])):
                     taxid_list.append(ncbi.get_name_translator([item])[item][0])
                 else:
-                    print(f'No {item} found in NCBI db')
-                    
+                    logging.warning(f"Taxon label '{item}' not found in NCBI db (contig {contig})")
+
             hit_lineages = []
             for item in taxid_list:
                 lineage_dict = {}
                 try:
                     for x, y in ncbi.get_rank(ncbi.get_lineage(item)).items():
                         if y in viphog_rank:
                             taxa_to_check = ncbi.get_taxid_translator([x])[x]
-                            if exclude_deprecated_taxa and taxa_to_check not in EXCLUDE_TAXA:
+                            # Check is some taxa should be excluded
+                            if exclude_deprecated_taxa:
+                                if taxa_to_check in EXCLUDE_TAXA:
+                                    logging.info(f"Taxon '{taxa_to_check}' is deprecated and excluded (contig {contig})")
+                                else:
+                                    lineage_dict[y] = taxa_to_check
+                            else:
                                 lineage_dict[y] = taxa_to_check
                     if lineage_dict:
                         hit_lineages.append(lineage_dict)
                 except ValueError:
-                    print(f'Can not return lineage for {item}')
+                    logging.warning(f"Cannot retrieve lineage for taxid {item} (contig {contig})")
                     pass
+
+            if not hit_lineages:
+                logging.debug(f"Contig {contig}: no valid lineages resolved from {len(taxid_list)} hits")
+
+            contig_assigned = False
             for rank in output_taxa_order[::-1][:-1]:
                 taxon_list = [item.get(rank) for item in hit_lineages]
                 total_hits = sum(pd.notnull(taxon_list))
@@ -142,14 +199,21 @@ def contig_tax(annot_df, ncbi_db, tax_thres, taxon_factor_dict, output_taxa_orde
                             if hit_taxon in taxon_factor_dict.keys()
                             else 1
                         )
-                        if prop_hits < tax_thres * taxon_factor:
-                            hit_bound = math.ceil(tax_thres * taxon_factor * total_hits)
+                        effective_thres = tax_thres * taxon_factor
+                        logging.debug(
+                            f"Contig {contig} | rank={rank} | taxon={hit_taxon} | prop={prop_hits:.2f} | thres={effective_thres:.2f} (factor={taxon_factor:.2f})"
+                        )
+                        if prop_hits < effective_thres:
+                            hit_bound = math.ceil(effective_thres * total_hits)
                             hit_diff = hit_bound - hit_count
                             under_thres.append((hit_taxon, hit_diff))
                         else:
                             over_thres.append((hit_taxon, prop_hits))
                     if len(over_thres) == 0:
                         best_under = sorted(under_thres, key=lambda x: x[1])[0]
+                        logging.debug(
+                            f"Contig {contig} | rank={rank}: all below threshold, best candidate is '{best_under[0]}'"
+                        )
                         contig_lineage.append(best_under[0])
                     else:
                         sorted_over_thres = [
@@ -176,22 +240,39 @@ def contig_tax(annot_df, ncbi_db, tax_thres, taxon_factor_dict, output_taxa_orde
                                 taxon_lineage_list = [
                                     taxon_lineage_dict.get(item, "")
                                     for item in output_taxa_order[
-                                        1:output_taxa_order.index(rank)+1 
+                                        1 : output_taxa_order.index(rank) + 1
                                     ]
                                 ]
+                                logging.info(f"Contig {contig} assigned at rank '{rank}' via taxon '{taxon}'")
+                                contig_assigned = True
                                 break
                         else:
+                            logging.debug(
+                                f"Contig {contig} | rank={rank}: candidates {sorted_over_thres} all failed CDS size filter"
+                            )
                             contig_lineage.append("")
                             continue
                         contig_lineage.reverse()
                         contig_lineage = taxon_lineage_list + contig_lineage
                         break
+
+            if annot_prot > 0:
+                if contig_assigned:
+                    assigned += 1
+                else:
+                    logging.debug(f"Contig {contig}: had hits but could not be assigned (below threshold or no valid lineage)")
+                    unassigned_below_thres += 1
+
         contig_lineage = [contig] + contig_lineage
         yield contig_lineage
 
+    total = len(contig_set)
+    logging.info(
+        f"Summary: {assigned}/{total} contigs assigned | {unassigned_no_hits} no ViPhOG hits | {unassigned_below_thres} hits but unassigned"
+    )
+
 
 if __name__ == "__main__":
-
     parser = argparse.ArgumentParser(
         description="Generate tabular file with taxonomic assignment of viral contigs based on ViPhOG annotations"
     )
@@ -232,7 +313,7 @@ def contig_tax(annot_df, ncbi_db, tax_thres, taxon_factor_dict, output_taxa_orde
         "--version4",
         dest="version4",
         help="Flag for whether to use version 4 or not. This has implications for taxa that are ignored in version 4.",
-        action="store_true"
+        action="store_true",
     )
     args = parser.parse_args()