ENH: Added option to run annotate_reads_card in parallel with parsl

README.md

@@ -15,7 +15,7 @@ To install _q2-amr_, follow the steps described below.
 mamba create -yn q2-amr \
   -c https://packages.qiime2.org/qiime2/2024.2/shotgun/released/  \
   -c qiime2 -c conda-forge -c bioconda -c defaults \
-  qiime2 q2cli q2templates q2-types rgi
+  qiime2 q2cli q2templates q2-types q2-feature-table q2-demux rgi
 
 conda activate q2-amr
 
@@ -40,7 +40,7 @@ qiime info
 CONDA_SUBDIR=osx-64 mamba create -yn q2-amr \
   -c https://packages.qiime2.org/qiime2/2024.2/shotgun/released/ \
   -c qiime2 -c conda-forge -c bioconda -c defaults \
-  qiime2 q2cli q2templates q2-types rgi
+  qiime2 q2cli q2templates q2-types q2-feature-table q2-demux rgi
 
 conda activate q2-amr
 conda config --env --set subdir osx-64

ci/recipe/meta.yaml

@@ -19,6 +19,8 @@ requirements:
   run:
     - python=3.8.*
     - qiime2 {{ qiime2_epoch }}.*
+    - q2-demux {{ qiime2_epoch }}.*
+    - q2-feature-table {{ qiime2_epoch }}.*
     - q2-types {{ qiime2_epoch }}.*
     - q2templates {{ qiime2_epoch }}.*
     - q2cli {{ qiime2_epoch }}.*

q2_amr/card/reads.py

@@ -6,9 +6,12 @@
 
 import pandas as pd
 from q2_types.per_sample_sequences import (
+    PairedEndSequencesWithQuality,
+    SequencesWithQuality,
     SingleLanePerSamplePairedEndFastqDirFmt,
     SingleLanePerSampleSingleEndFastqDirFmt,
 )
+from q2_types.sample_data import SampleData
 
 from q2_amr.card.utils import create_count_table, load_card_db, read_in_txt, run_command
 from q2_amr.types import (
@@ -19,6 +22,64 @@
 
 
 def annotate_reads_card(
+    ctx,
+    reads,
+    card_db,
+    aligner="kma",
+    threads=1,
+    include_wildcard=False,
+    include_other_models=False,
+    num_partitions=None,
+):
+    # Get all actions used by the pipeline
+    if reads.type <= SampleData[SequencesWithQuality]:
+        partition_method = ctx.get_action("demux", "partition_samples_single")
+    elif reads.type <= SampleData[PairedEndSequencesWithQuality]:
+        partition_method = ctx.get_action("demux", "partition_samples_paired")
+
+    annotate = ctx.get_action("amr", "_annotate_reads_card")
+
+    collate_allele_annotations = ctx.get_action(
+        "amr", "collate_reads_allele_annotations"
+    )
+    collate_gene_annotations = ctx.get_action("amr", "collate_reads_gene_annotations")
+    merge_tables = ctx.get_action("feature-table", "merge")
+
+    # Partition the reads
+    (partitioned_seqs,) = partition_method(reads, num_partitions)
+
+    allele_annotations = []
+    gene_annotations = []
+    allele_tables = []
+    gene_tables = []
+
+    # Run _annotate_reads_card for every partition
+    for read in partitioned_seqs.values():
+        (allele_annotation, gene_annotation, allele_table, gene_table) = annotate(
+            read, card_db, aligner, threads, include_wildcard, include_other_models
+        )
+
+        # Append output artifacts to lists
+        allele_annotations.append(allele_annotation)
+        gene_annotations.append(gene_annotation)
+        allele_tables.append(allele_table)
+        gene_tables.append(gene_table)
+
+    # Collate annotation and feature table artifacts
+    (collated_allele_annotations,) = collate_allele_annotations(allele_annotations)
+    (collated_gene_annotations,) = collate_gene_annotations(gene_annotations)
+    (collated_allele_tables,) = merge_tables(allele_tables)
+    (collated_gene_tables,) = merge_tables(gene_tables)
+
+    return (
+        collated_allele_annotations,
+        collated_gene_annotations,
+        collated_allele_tables,
+        collated_gene_tables,
+    )
+
+
+def _annotate_reads_card(
     reads: Union[
         SingleLanePerSamplePairedEndFastqDirFmt, SingleLanePerSampleSingleEndFastqDirFmt
     ],
@@ -36,24 +97,34 @@ def annotate_reads_card(
     paired = isinstance(reads, SingleLanePerSamplePairedEndFastqDirFmt)
     manifest = reads.manifest.view(pd.DataFrame)
     allele_frequency_list, gene_frequency_list = [], []
+
     amr_allele_annotation = CARDAlleleAnnotationDirectoryFormat()
     amr_gene_annotation = CARDGeneAnnotationDirectoryFormat()
+
     with tempfile.TemporaryDirectory() as tmp:
+        # Load CARD database files
         load_card_db(
             card_db=card_db,
             fasta=True,
             include_other_models=include_other_models,
             include_wildcard=include_wildcard,
         )
         for samp in list(manifest.index):
+            # Set paths for forward and reverse reads files
             fwd = manifest.loc[samp, "forward"]
             rev = manifest.loc[samp, "reverse"] if paired else None
+
+            # Create sample directories in the output directories
             samp_allele_dir = os.path.join(str(amr_allele_annotation), samp)
             samp_gene_dir = os.path.join(str(amr_gene_annotation), samp)
             os.makedirs(samp_allele_dir)
             os.makedirs(samp_gene_dir)
-            samp_input_dir = os.path.join(tmp, samp)
-            os.makedirs(samp_input_dir)
+
+            # Create sample directory in the tmp directory
+            samp_tmp_dir = os.path.join(tmp, samp)
+            os.makedirs(samp_tmp_dir)
+
+            # Run annotation
             run_rgi_bwt(
                 cwd=tmp,
                 samp=samp,
@@ -65,13 +136,13 @@ def annotate_reads_card(
                 include_other_models=include_other_models,
             )
 
-            # Create a frequency table and add it to a list, for gene and allele
-            # mapping data
+            # Create a frequency table for gene and allele mapping data and add it to
+            # the frequency table lists
             for map_type, table_list in zip(
                 ["allele", "gene"], [allele_frequency_list, gene_frequency_list]
             ):
                 path_txt = os.path.join(
-                    samp_input_dir, f"output.{map_type}_mapping_data.txt"
+                    samp_tmp_dir, f"output.{map_type}_mapping_data.txt"
                 )
                 frequency_table = read_in_txt(
                     path=path_txt,
@@ -87,15 +158,16 @@ def annotate_reads_card(
             ):
                 files = [f"{map_type}_mapping_data.txt"]
                 # mapping statistics only go to the allele directories
-                files.extend(
-                    ["overall_mapping_stats.txt", "sorted.length_100.bam"]
-                ) if map_type == "allele" else None
+                if map_type == "allele":
+                    files.extend(["overall_mapping_stats.txt", "sorted.length_100.bam"])
+
                 for file in files:
                     shutil.copy(
-                        os.path.join(samp_input_dir, "output." + file),
+                        os.path.join(samp_tmp_dir, "output." + file),
                         os.path.join(des_dir, file),
                     )
 
+    # Merge all frequency tables into one for alleles and genes separately
     allele_feature_table = create_count_table(allele_frequency_list)
     gene_feature_table = create_count_table(gene_frequency_list)
     return (

q2_amr/card/tests/data/partitioned/annotate_mags_output_2/sample2/bin1/amr_annotation.txt

@@ -0,0 +1,3 @@
+ORF_ID	Contig	Start	Stop	Orientation	Cut_Off	Pass_Bitscore	Best_Hit_Bitscore	Best_Hit_ARO	Best_Identities	ARO	Model_type	SNPs_in_Best_Hit_ARO	Other_SNPs	Drug Class	Resistance Mechanism	AMR Gene Family	Predicted_DNA	Predicted_Protein	CARD_Protein_Sequence	Percentage Length of Reference Sequence	ID	Model_ID	Nudged	Note
+k141_1197_2 # 683 # 1345 # 1 # ID=49_2;partial=00;start_type=ATG;rbs_motif=GGA/GAG/AGG;rbs_spacer=5-10bp;gc_cont=0.588	k141_1197_2	683	1345	+	Strict	200	326.635	vanX gene in vanO cluster	75.25	3002954	protein homolog model	n/a	n/a	glycopeptide antibiotic	antibiotic target alteration	vanX; glycopeptide resistance gene cluster	ATGAAGGGTGACTTCGTTTTCGTTGATGAGTGTGTTCCAGGAGTCCGCTGGGATGCCAAATACGCCACATCGGACAACTTCACCGGCAAACCGGTGGAGGGATATCTGGCCAACCGGATTGTCGGGACCAGGGCTTTGTGCTCAGCGCTGGAAAGCGTGCGGCAACGGGCTGCATCCCGCGGTTTCGGGTTGCTGCTGTGGGACGGCTACCGCCCGCAGCGCGCCGTGGATTCGTTCCTGCACTGGGCGAAACAACCAGAGGACGGCGCAACTAAACGCCGCCACTATCCAAATATTTCCCGAGCGGAAATGTTCGAAAGAGGATACGTAGCCTCCAAGTCCGGCCACAGCCGGGGCAGCACCGTCGATTTGACCCTGTATGACCTGGTTACCGGTGACCTCGTTCCCATGGGCGGCGGCCACGACTTGATGGATGAAATTTCGCATCACGGAGCGCCCGGCATCACCCGGGCCGAGACCGGCAACCGCCACACGCTGCGTTCGCTCATGGAGGCCTGCGGTTTCAGTTCCTACGATTCTGAGTGGTGGCATTACACCCTGAAGAACGAACCCTATCCGGACACTTATTTCGATTTTCCCGTTACGGATCCGCTTCCATCAGACGCCGCAACGGCCAGGGACCTTGTCTTCCAGAATGCATAG	MKGDFVFVDECVPGVRWDAKYATSDNFTGKPVEGYLANRIVGTRALCSALESVRQRAASRGFGLLLWDGYRPQRAVDSFLHWAKQPEDGATKRRHYPNISRAEMFERGYVASKSGHSRGSTVDLTLYDLVTGDLVPMGGGHDLMDEISHHGAPGITRAETGNRHTLRSLMEACGFSSYDSEWWHYTLKNEPYPDTYFDFPVTDPLPSDAATARDLVFQNA	MNDDFVYVDDWVPGVRWDAKYATWDNFTGKPVDGYLANRIVGTRALCAALEQAREKAASLGFGLLLWDGYRPRRAVDSFLRWSEQPEDGQTKQRHYPNIDRPEMLEKGYVATQSGHSRGGAVDLTLYHLATGELAPMGGDHDLMDPISHHRARGIKPIESKNRELLRSIMEDCGFDRYDCEWWHYTLKREPYPDVYFDFPIT	108.91	gnl|BL_ORD_ID|1674|hsp_num:0	1699
+k141_10683_1 # 1 # 453 # 1 # ID=423_1;partial=10;start_type=Edge;rbs_motif=None;rbs_spacer=None;gc_cont=0.658	k141_10683_1	1	453	+	Strict	50	90.8929	vanY gene in vanM cluster	38.62	3002961	protein homolog model	n/a	n/a	glycopeptide antibiotic	antibiotic target alteration	vanY; glycopeptide resistance gene cluster	GAGGCTGCAGGGGCCTACCGGCAAATGGCCGCGGAAGCGGGCGCCGCCGGAGTTCCCATGTCCGCGGTGAGCGGCTTTCGGACCGGAGCAGAGCAGGACCAGCTGTACGTCTCCTACACGGAGAACTTTGGGCCGGAGGCAGCCGACGCCATTTCGGCCCGTCCCGGGTACAGCGAGCATCAGACGGGGCTGGCCATCGACATCGCCAACCCGGACGGAACCTGCGCCCTGGAATCCTGCTTCGCCGAAACCTTGGCGGGTTCGTGGGCGGCCGCCAATGCCCAGCACTACGGCTTCATCATCCGTTATCCGGCAGGAGCCGAGCACATCACCGGGTACGCCCATGAACCGTGGCATCTGCGGTACGTGGGGACGGAACATGCCCGGACAATGCACGACGCCGGCACCACCTTGGAAGAATATCTGGGACTTCCTGCCGCGCCGGGTTACTGA	EAAGAYRQMAAEAGAAGVPMSAVSGFRTGAEQDQLYVSYTENFGPEAADAISARPGYSEHQTGLAIDIANPDGTCALESCFAETLAGSWAAANAQHYGFIIRYPAGAEHITGYAHEPWHLRYVGTEHARTMHDAGTTLEEYLGLPAAPGY	MVFQGNLLLVNNEYPVLEESIKTDVVNLFKHDELTKGYELLNREIYLSEKVAREFSEMVDAAEKEGVRHFSINSGFRNFDEQNALYQEMGSDYALPAGYSEHNLGLALDIGSTQMEMSEAPEGKWLEDNAWEYGFILRYPMDKTAITGIQYEPWHFRYVGLPHSAIIEEKNFALEEYLDFLKEQKSISGTIHGENYEISYYPITEKTDIEMPANLHYEISGNNMDGVIVTVYR	64.38	gnl|BL_ORD_ID|1675|hsp_num:0	1713

q2_amr/card/tests/data/partitioned/kmer_analysis_mags_2/sample2/bin1/61mer_analysis.json

@@ -0,0 +1 @@
+{"k141_1617_82 # 105913 # 109473 # 1 # ID=42_82;partial=00;start_type=TTG;rbs_motif=AGGA;rbs_spacer=5-10bp;gc_cont=0.630": {"ORF": "k141_1617_82 # 105913 # 109473 # 1 # ID=42_82;partial=00;start_type=TTG;rbs_motif=AGGA;rbs_spacer=5-10bp;gc_cont=0.630", "contig": "k141_1617_82", "HSP": "gnl|BL_ORD_ID|4713|hsp_num:0", "ARO_model": "Bifidobacterium adolescentis rpoB mutants conferring resistance to rifampicin", "type_hit": "Perfect", "#_of_kmers_in_sequence": 3501, "#_of_AMR_kmers": 2274, "taxonomic_info": {"species": {"Bifidobacterium longum": 75, "Bifidobacterium dentium": 67, "Parascardovia denticolens": 3, "Bifidobacterium animalis": 7, "Bifidobacterium bifidum": 9}, "genus": {}}, "genomic_info": {"chr + plasmid": 0, "plasmid": 0, "chr": 234}}}

q2_amr/card/tests/data/partitioned/kmer_analysis_mags_2/sample2/bin1/61mer_analysis_rgi_summary.txt

@@ -0,0 +1,2 @@
+ORF_ID	Contig	Cut_Off	Best_Hit_ARO	CARD*kmer Prediction	Taxonomic kmers	Genomic kmers
+"NC_000962.3_273 # 314309 # 314854 # -1 # ID=1_273;partial=00;start_type=GTG;rbs_motif=AGG;rbs_spacer=4bp;gc_cont=0.679"	NC_000962.3_273	Perfect	AAC(2')-Ic	Mycobacterium tuberculosis (chromosome)	"Mycobacterium tuberculosis: 486; "	"chr + plasmid: 0; plasmid: 0; chr: 486; "

q2_amr/card/tests/data/reads_paired/MANIFEST

@@ -0,0 +1,2 @@
+sample-id,filename,direction
+sample1,sample1_00_L001_R1_001.fastq.gz,forward

q2_amr/card/tests/data/reads_paired/metadata.yml

@@ -0,0 +1 @@
+phred-offset: 33

q2_amr/card/tests/data/reads_single/MANIFEST

@@ -0,0 +1,2 @@
+sample-id,filename,direction
+sample1,sample1_00_L001_R1_001.fastq.gz,forward

q2_amr/card/tests/data/reads_single/metadata.yml

@@ -0,0 +1 @@
+phred-offset: 33

q2_amr/card/tests/test_reads.py

@@ -7,9 +7,10 @@
     SingleLanePerSamplePairedEndFastqDirFmt,
     SingleLanePerSampleSingleEndFastqDirFmt,
 )
+from qiime2 import Artifact
 from qiime2.plugin.testing import TestPluginBase
 
-from q2_amr.card.reads import annotate_reads_card, run_rgi_bwt
+from q2_amr.card.reads import _annotate_reads_card, annotate_reads_card, run_rgi_bwt
 from q2_amr.types import (
     CARDAlleleAnnotationDirectoryFormat,
     CARDDatabaseDirectoryFormat,
@@ -78,8 +79,8 @@ def annotate_reads_card_test_body(self, read_type):
         ), patch("q2_amr.card.reads.read_in_txt", mock_read_in_txt), patch(
             "q2_amr.card.reads.create_count_table", mock_create_count_table
         ):
-            # Run annotate_reads_card function
-            result = annotate_reads_card(reads, card_db)
+            # Run _annotate_reads_card function
+            result = _annotate_reads_card(reads, card_db)
 
             # Retrieve the path to cwd directory from mock_run_rgi_bwt arguments
             first_call_args = mock_run_rgi_bwt.call_args_list[0]
@@ -197,3 +198,50 @@ def test_exception_raised(self):
             mock_run_command.side_effect = subprocess.CalledProcessError(1, "cmd")
             run_rgi_bwt()
             self.assertEqual(str(cm.exception), expected_message)
+
+    def test_annotate_reads_card_pipeline_paired(self):
+        reads_path = self.get_data_path("reads_paired")
+        reads = SingleLanePerSamplePairedEndFastqDirFmt(reads_path, "r")
+        reads_artifact = Artifact.import_data(
+            "SampleData[PairedEndSequencesWithQuality]", reads
+        )
+        self._test_annotate_reads_card_pipeline(reads_artifact)
+
+    def test_annotate_reads_card_pipeline_single(self):
+        reads_path = self.get_data_path("reads_single")
+        reads = SingleLanePerSampleSingleEndFastqDirFmt(reads_path, "r")
+        reads_artifact = Artifact.import_data("SampleData[SequencesWithQuality]", reads)
+        self._test_annotate_reads_card_pipeline(reads_artifact)
+
+    def _test_annotate_reads_card_pipeline(self, reads):
+        # Mock the get_action method to return MagicMock objects
+        mock_ctx = MagicMock()
+        mock_ctx.get_action.side_effect = [
+            MagicMock(
+                return_value=({"1": "artifact_reads_1", "2": "artifact_reads_2"},)
+            ),
+            MagicMock(
+                return_value=(
+                    "artifact_annotation_allele",
+                    "artifact_annotation_gene",
+                    "artifact_feature_table_allele",
+                    "artifact_feature_table_gene",
+                )
+            ),
+            MagicMock(return_value=("artifact_allele_annotation_collated",)),
+            MagicMock(return_value=("artifact_gene_annotation_collated",)),
+            MagicMock(return_value=("artifact_feature_table_merged",)),
+        ]
+
+        # Call function with mocked ctx
+        result = annotate_reads_card(ctx=mock_ctx, reads=reads, card_db=None)
+
+        self.assertEqual(
+            result,
+            (
+                "artifact_allele_annotation_collated",
+                "artifact_gene_annotation_collated",
+                "artifact_feature_table_merged",
+                "artifact_feature_table_merged",
+            ),
+        )

q2_amr/plugin_setup.py

@@ -42,7 +42,7 @@
     partition_reads_allele_annotations,
     partition_reads_gene_annotations,
 )
-from q2_amr.card.reads import annotate_reads_card
+from q2_amr.card.reads import _annotate_reads_card, annotate_reads_card
 from q2_amr.types import (
     CARDAnnotationJSONFormat,
     CARDAnnotationTXTFormat,
@@ -150,32 +150,77 @@
     citations=[citations["alcock_card_2023"]],
 )
 
-P_aligner, T_allele_annotation, T_gene_annotation = TypeMap(
+P_aligner, T_allele_annotation = TypeMap(
     {
+        Str % Choices("kma"): SampleData[CARDAlleleAnnotation % Properties("kma")],
         Str
-        % Choices("kma"): (
-            SampleData[CARDAlleleAnnotation % Properties("kma")],
-            SampleData[CARDGeneAnnotation % Properties("kma")],
-        ),
-        Str
-        % Choices("bowtie2"): (
-            SampleData[CARDAlleleAnnotation % Properties("bowtie2")],
-            SampleData[CARDGeneAnnotation % Properties("bowtie2")],
-        ),
-        Str
-        % Choices("bwa"): (
-            SampleData[CARDAlleleAnnotation % Properties("bwa")],
-            SampleData[CARDGeneAnnotation % Properties("bwa")],
-        ),
+        % Choices("bowtie2"): SampleData[CARDAlleleAnnotation % Properties("bowtie2")],
+        Str % Choices("bwa"): SampleData[CARDAlleleAnnotation % Properties("bwa")],
     }
 )
 
-plugin.methods.register_function(
+plugin.pipelines.register_function(
     function=annotate_reads_card,
     inputs={
         "reads": SampleData[PairedEndSequencesWithQuality | SequencesWithQuality],
         "card_db": CARDDatabase,
     },
+    parameters={
+        "aligner": P_aligner,
+        "threads": Int % Range(0, None, inclusive_start=False),
+        "include_wildcard": Bool,
+        "include_other_models": Bool,
+        "num_partitions": Int % Range(0, None, inclusive_start=False),
+    },
+    outputs=[
+        ("amr_allele_annotation", T_allele_annotation),
+        ("amr_gene_annotation", SampleData[CARDGeneAnnotation]),
+        ("allele_feature_table", FeatureTable[Frequency]),
+        ("gene_feature_table", FeatureTable[Frequency]),
+    ],
+    input_descriptions={
+        "reads": "Paired or single end reads.",
+        "card_db": "CARD Database.",
+    },
+    parameter_descriptions={
+        "aligner": "Specify alignment tool.",
+        "threads": "Number of threads (CPUs) to use.",
+        "include_wildcard": "Additionally align reads to the in silico predicted "
+        "allelic variants available in CARD's Resistomes & Variants"
+        " data set. This is highly recommended for non-clinical "
+        "samples .",
+        "include_other_models": "The default settings will align reads against "
+        "CARD's protein homolog models. With include_other_"
+        "models set to True, reads are additionally aligned to "
+        "protein variant models, rRNA mutation models, and "
+        "protein over-expression models. These three model "
+        "types require comparison to CARD's curated lists of "
+        "mutations known to confer phenotypic antibiotic "
+        "resistance to differentiate alleles conferring "
+        "resistance from antibiotic susceptible alleles, "
+        "but RGI as of yet does not perform this comparison. "
+        "Use these results with caution.",
+        "num_partitions": "Number of partitions that should run in parallel.",
+    },
+    output_descriptions={
+        "amr_allele_annotation": "AMR annotation mapped on alleles.",
+        "amr_gene_annotation": "AMR annotation mapped on genes.",
+        "allele_feature_table": "Frequency table of ARGs in all samples for allele "
+        "mapping.",
+        "gene_feature_table": "Frequency table of ARGs in all samples for gene "
+        "mapping.",
+    },
+    name="Annotate reads with antimicrobial resistance genes from CARD.",
+    description="Annotate reads with antimicrobial resistance genes from CARD.",
+    citations=[citations["alcock_card_2023"]],
+)
+
+plugin.methods.register_function(
+    function=_annotate_reads_card,
+    inputs={
+        "reads": SampleData[PairedEndSequencesWithQuality | SequencesWithQuality],
+        "card_db": CARDDatabase,
+    },
     parameters={
         "aligner": P_aligner,
         "threads": Int % Range(0, None, inclusive_start=False),
@@ -184,13 +229,13 @@
     },
     outputs=[
         ("amr_allele_annotation", T_allele_annotation),
-        ("amr_gene_annotation", T_gene_annotation),
+        ("amr_gene_annotation", SampleData[CARDGeneAnnotation]),
         ("allele_feature_table", FeatureTable[Frequency]),
         ("gene_feature_table", FeatureTable[Frequency]),
     ],
     input_descriptions={
         "reads": "Paired or single end reads.",
-        "card_db": "CARD Database",
+        "card_db": "CARD Database.",
     },
     parameter_descriptions={
         "aligner": "Specify alignment tool.",

setup.py

@@ -35,6 +35,7 @@
         ],
         "q2_amr.card.tests": [
             "data/*",
+            "data/*/*",
             "data/*/*/*/*",
             "data/*/*/*/*/*",
         ],