Add new required inputs to PipelineStepBlastContigs

chanzuckerberg · Apr 17, 2024 · eb75a2b · eb75a2b
1 parent 593c7c7
commit eb75a2b
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 13 deletions.
diff --git a/lib/idseq-dag/idseq_dag/steps/blast_contigs.py b/lib/idseq-dag/idseq_dag/steps/blast_contigs.py
@@ -94,8 +94,15 @@ def run(self):
             command.write_text_to_file('[]', contig_summary_json)
             return  # return in the middle of the function
 
+        lineage_db = s3.fetch_reference(
+            self.additional_files["lineage_db"],
+            self.ref_dir_local,
+            allow_s3mi=False)  # Too small to waste s3mi
+
+        accession2taxid_dict = s3.fetch_reference(self.additional_files["accession2taxid"], self.ref_dir_local)
+
         (read_dict, accession_dict, _selected_genera) = m8.summarize_hits(hit_summary)
-        PipelineStepBlastContigs.run_blast(db_type, blast_m8, assembled_contig, reference_fasta, blast_top_m8)
+        PipelineStepBlastContigs.run_blast(db_type, blast_m8, assembled_contig, reference_fasta, blast_top_m8, lineage_db, accession2taxid_dict)
         read2contig = {}
         generate_info_from_sam(bowtie_sam, read2contig, duplicate_cluster_sizes_path=duplicate_cluster_sizes_path)
 
@@ -106,11 +113,6 @@ def run(self):
                                          refined_hit_summary, refined_m8)
 
         # Generating taxon counts based on updated results
-        lineage_db = s3.fetch_reference(
-            self.additional_files["lineage_db"],
-            self.ref_dir_local,
-            allow_s3mi=False)  # Too small to waste s3mi
-
         deuterostome_db = None
         if self.additional_files.get("deuterostome_db"):
             deuterostome_db = s3.fetch_reference(self.additional_files["deuterostome_db"],
@@ -264,7 +266,7 @@ def update_read_dict(read2contig, blast_top_blastn_6_path, read_dict, accession_
             return (consolidated_dict, read2blastm8, contig2lineage, added_reads)
 
     @staticmethod
-    def run_blast_nt(blast_index_path, blast_m8, assembled_contig, reference_fasta, blast_top_m8):
+    def run_blast_nt(blast_index_path, blast_m8, assembled_contig, reference_fasta, blast_top_m8, lineage_db, accession2taxid_dict):
         blast_type = 'nucl'
         blast_command = 'blastn'
         command.execute(
@@ -308,10 +310,10 @@ def run_blast_nt(blast_index_path, blast_m8, assembled_contig, reference_fasta,
             )
         )
         # further processing of getting the top m8 entry for each contig.
-        get_top_m8_nt(blast_m8, blast_top_m8)
+        get_top_m8_nt(blast_m8, lineage_db, accession2taxid_dict, blast_top_m8)
 
     @staticmethod
-    def run_blast_nr(blast_index_path, blast_m8, assembled_contig, reference_fasta, blast_top_m8):
+    def run_blast_nr(blast_index_path, blast_m8, assembled_contig, reference_fasta, blast_top_m8, lineage_db, accession2taxid_dict):
         blast_type = 'prot'
         blast_command = 'blastx'
         command.execute(
@@ -349,4 +351,4 @@ def run_blast_nr(blast_index_path, blast_m8, assembled_contig, reference_fasta,
             )
         )
         # further processing of getting the top m8 entry for each contig.
-        get_top_m8_nr(blast_m8, blast_top_m8)
+        get_top_m8_nr(blast_m8, lineage_db, accession2taxid_dict, blast_top_m8)
diff --git a/workflows/short-read-mngs/postprocess.wdl b/workflows/short-read-mngs/postprocess.wdl
@@ -157,6 +157,7 @@ task BlastContigs_refined_gsnap_out {
     File assembly_nt_refseq_fasta
     File duplicate_cluster_sizes_tsv
     File lineage_db
+    File accession2taxid
     File taxon_blacklist
     File deuterostome_db
     Boolean use_deuterostome_filter
@@ -171,7 +172,7 @@ task BlastContigs_refined_gsnap_out {
     --input-files '[["~{gsnap_out_gsnap_m8}", "~{gsnap_out_gsnap_deduped_m8}", "~{gsnap_out_gsnap_hitsummary_tab}", "~{gsnap_out_gsnap_counts_with_dcr_json}"], ["~{assembly_contigs_fasta}", "~{assembly_scaffolds_fasta}", "~{assembly_read_contig_sam}", "~{assembly_contig_stats_json}"], ["~{assembly_nt_refseq_fasta}"], ["~{duplicate_cluster_sizes_tsv}"]]' \
     --output-files '["assembly/gsnap.blast.m8", "assembly/gsnap.reassigned.m8", "assembly/gsnap.hitsummary2.tab", "assembly/refined_gsnap_counts_with_dcr.json", "assembly/gsnap_contig_summary.json", "assembly/gsnap.blast.top.m8"]' \
     --output-dir-s3 '~{s3_wd_uri}' \
-    --additional-files '{"lineage_db": "~{lineage_db}", "taxon_blacklist": "~{taxon_blacklist}", "deuterostome_db": "~{if use_deuterostome_filter then '~{deuterostome_db}' else ''}"}' \
+    --additional-files '{"lineage_db": "~{lineage_db}", "accession2taxid": "~{accession2taxid}", "taxon_blacklist": "~{taxon_blacklist}", "deuterostome_db": "~{if use_deuterostome_filter then '~{deuterostome_db}' else ''}"}' \
     --additional-attributes '{"db_type": "nt", "use_taxon_whitelist": ~{use_taxon_whitelist}}'
   >>>
   output {
@@ -204,6 +205,7 @@ task BlastContigs_refined_rapsearch2_out {
     File assembly_nr_refseq_fasta
     File duplicate_cluster_sizes_tsv
     File lineage_db
+    File accession2taxid
     File taxon_blacklist
     Boolean use_taxon_whitelist
   }
@@ -216,7 +218,7 @@ task BlastContigs_refined_rapsearch2_out {
     --input-files '[["~{rapsearch2_out_rapsearch2_m8}", "~{rapsearch2_out_rapsearch2_deduped_m8}", "~{rapsearch2_out_rapsearch2_hitsummary_tab}", "~{rapsearch2_out_rapsearch2_counts_with_dcr_json}"], ["~{assembly_contigs_fasta}", "~{assembly_scaffolds_fasta}", "~{assembly_read_contig_sam}", "~{assembly_contig_stats_json}"], ["~{assembly_nr_refseq_fasta}"], ["~{duplicate_cluster_sizes_tsv}"]]' \
     --output-files '["assembly/rapsearch2.blast.m8", "assembly/rapsearch2.reassigned.m8", "assembly/rapsearch2.hitsummary2.tab", "assembly/refined_rapsearch2_counts_with_dcr.json", "assembly/rapsearch2_contig_summary.json", "assembly/rapsearch2.blast.top.m8"]' \
     --output-dir-s3 '~{s3_wd_uri}' \
-    --additional-files '{"lineage_db": "~{lineage_db}", "taxon_blacklist": "~{taxon_blacklist}"}' \
+    --additional-files '{"lineage_db": "~{lineage_db}", "accession2taxid": "~{accession2taxid}", "taxon_blacklist": "~{taxon_blacklist}"}' \
     --additional-attributes '{"db_type": "nr", "use_taxon_whitelist": ~{use_taxon_whitelist}}'
   >>>
   output {
@@ -489,6 +491,7 @@ workflow czid_postprocess {
     String nr_db = "s3://czid-public-references/ncbi-sources/2021-01-22/nr"
     File nr_loc_db = "s3://czid-public-references/alignment_data/2021-01-22/nr_loc.db"
     File lineage_db = "s3://czid-public-references/taxonomy/2021-01-22/taxid-lineages.db"
+    File accession2taxid_db = "s3://czid-public-references/ncbi-indexes-prod/2021-01-22/index-generation-2/accession2taxid.marisa"
     File taxon_blacklist = "s3://czid-public-references/taxonomy/2021-01-22/taxon_blacklist.txt"
     File deuterostome_db = "s3://czid-public-references/taxonomy/2021-01-22/deuterostome_taxids.txt"
     Boolean use_deuterostome_filter = true
@@ -556,6 +559,7 @@ workflow czid_postprocess {
       assembly_nt_refseq_fasta = DownloadAccessions_gsnap_accessions_out.assembly_nt_refseq_fasta,
       duplicate_cluster_sizes_tsv = duplicate_cluster_sizes_tsv,
       lineage_db = lineage_db,
+      accession2taxid = accession2taxid_db,
       taxon_blacklist = taxon_blacklist,
       deuterostome_db = deuterostome_db,
       use_deuterostome_filter = use_deuterostome_filter,
@@ -577,6 +581,7 @@ workflow czid_postprocess {
       assembly_nr_refseq_fasta = DownloadAccessions_rapsearch2_accessions_out.assembly_nr_refseq_fasta,
       duplicate_cluster_sizes_tsv = duplicate_cluster_sizes_tsv,
       lineage_db = lineage_db,
+      accession2taxid = accession2taxid_db,
       taxon_blacklist = taxon_blacklist,
       use_taxon_whitelist = use_taxon_whitelist
   }

diff --git a/workflows/short-read-mngs/test/local_test_viral.yml b/workflows/short-read-mngs/test/local_test_viral.yml
@@ -26,6 +26,7 @@ postprocess.nt_db: s3://czid-public-references/test/viral-alignment-indexes/vira
 postprocess.nt_loc_db: s3://czid-public-references/test/viral-alignment-indexes/viral_nt_loc.marisa
 postprocess.nr_db: s3://czid-public-references/test/viral-alignment-indexes/viral_nr
 postprocess.nr_loc_db: s3://czid-public-references/test/viral-alignment-indexes/viral_nr_loc.marisa
+postprocess.accession2taxid_db: s3://czid-public-references/mini-database/alignment_indexes/2020-08-20-viral/viral_accessions2taxid.marisa
 experimental.nt_db: s3://czid-public-references/test/viral-alignment-indexes/viral_nt
 experimental.nt_loc_db: s3://czid-public-references/test/viral-alignment-indexes/viral_nt_loc.marisa
-experimental.nt_info_db: s3://czid-public-references/test/viral-alignment-indexes/viral_nt_info.marisa
+experimental.nt_info_db: s3://czid-public-references/test/viral-alignment-indexes/viral_nt_info.marisa