Merge branch 'master' of github.com:AroneyS/ibis

AroneyS · AroneyS · commit f53c157a42c0 · 2023-05-18T11:12:35.000+10:00
diff --git a/ibis/workflow/coassemble.smk b/ibis/workflow/coassemble.smk
@@ -211,6 +211,7 @@ rule query_processing:
     params:
         sequence_identity=config["appraise_sequence_identity"],
         window_size=60,
+        taxa_of_interest=config["taxa_of_interest"],
     threads:
         64
     script:
diff --git a/ibis/workflow/scripts/query_processing.py b/ibis/workflow/scripts/query_processing.py
@@ -20,12 +20,19 @@ def processing(
     query_read,
     pipe_read,
     SEQUENCE_IDENTITY=0.86,
-    WINDOW_SIZE=60):
+    WINDOW_SIZE=60,
+    TAXA_OF_INTEREST=None):
 
     if len(query_read) == 0:
         empty_output = pl.DataFrame(schema=OUTPUT_COLUMNS)
         return empty_output, empty_output
 
+    # Filter TAXA_OF_INTEREST
+    if TAXA_OF_INTEREST:
+        pipe_read = pipe_read.filter(
+            pl.col("taxonomy").str.contains(TAXA_OF_INTEREST, literal=True)
+        )
+
     appraised = query_read.rename(
         # Rename to match appraise output
         # Query output: query_name, query_sequence, divergence, num_hits, coverage, sample, marker, hit_sequence, taxonomy
@@ -64,7 +71,8 @@ def pipeline(
     query_reads,
     pipe_reads,
     SEQUENCE_IDENTITY=0.86,
-    WINDOW_SIZE=60):
+    WINDOW_SIZE=60,
+    TAXA_OF_INTEREST=None):
 
     print(f"Polars using {str(pl.threadpool_size())} threads")
 
@@ -73,7 +81,8 @@ def pipeline(
             pl.read_csv(query, separator="\t"),
             pl.read_csv(pipe, separator="\t"),
             SEQUENCE_IDENTITY,
-            WINDOW_SIZE)
+            WINDOW_SIZE,
+            TAXA_OF_INTEREST)
         yield binned, unbinned
 
 if __name__ == "__main__":
@@ -82,6 +91,7 @@ def pipeline(
 
     SEQUENCE_IDENTITY = snakemake.params.sequence_identity
     WINDOW_SIZE = snakemake.params.window_size
+    TAXA_OF_INTEREST = snakemake.params.taxa_of_interest
     query_reads = snakemake.input.query_reads
     pipe_reads = snakemake.input.pipe_reads
     binned_path = snakemake.output.binned
@@ -91,7 +101,8 @@ def pipeline(
         query_reads,
         pipe_reads,
         SEQUENCE_IDENTITY=SEQUENCE_IDENTITY,
-        WINDOW_SIZE=WINDOW_SIZE
+        WINDOW_SIZE=WINDOW_SIZE,
+        TAXA_OF_INTEREST=TAXA_OF_INTEREST
         )
 
     with open(binned_path, "ab") as binned_file, open(unbinned_path, "ab") as unbinned_file:
diff --git a/test/test_coassemble.py b/test/test_coassemble.py
@@ -334,6 +334,111 @@ def test_coassemble_query_input(self):
             with open(cluster_path) as f:
                 self.assertEqual(expected, f.read())
 
+    def test_coassemble_query_input_taxa_of_interest(self):
+        with in_tempdir():
+            cmd = (
+                f"ibis coassemble "
+                f"--forward {SAMPLE_READS_FORWARD} "
+                f"--reverse {SAMPLE_READS_REVERSE} "
+                f"--genomes {GENOMES} "
+                f"--genome-transcripts {GENOME_TRANSCRIPTS} "
+                f"--sample-query {SAMPLE_QUERY} "
+                f"--sample-singlem {SAMPLE_QUERY_SINGLEM} "
+                f"--sample-read-size {SAMPLE_READ_SIZE} "
+                f"--taxa-of-interest \"p__Actinobacteriota\" "
+                f"--output test "
+                f"--conda-prefix {path_to_conda} "
+                f"--snakemake-args \"cluster_graph\" "
+            )
+            extern.run(cmd)
+
+            config_path = os.path.join("test", "config.yaml")
+            self.assertTrue(os.path.exists(config_path))
+
+            binned_path = os.path.join("test", "coassemble", "appraise", "binned.otu_table.tsv")
+            self.assertTrue(os.path.exists(binned_path))
+            expected = "\n".join(
+                [
+                    "\t".join(["gene", "sample", "sequence", "num_hits", "coverage", "taxonomy", "found_in"]),
+                    ""
+                ]
+            )
+            with open(binned_path) as f:
+                self.assertEqual(expected, f.read())
+
+            unbinned_path = os.path.join("test", "coassemble", "appraise", "unbinned.otu_table.tsv")
+            self.assertTrue(os.path.exists(unbinned_path))
+            expected = "\n".join(
+                [
+                    "\t".join([
+                        "gene",
+                        "sample",
+                        "sequence",
+                        "num_hits",
+                        "coverage",
+                        "taxonomy",
+                        "found_in",
+                    ]),
+                    "\t".join([
+                        "S3.7.ribosomal_protein_S7",
+                        "sample_1",
+                        "TACCAGGTCCCGGTCGAGGTCCGTCCGATCCGCCAGACGACGCTCGCCCTGCGCTGGCTC",
+                        "5",
+                        "8.21",
+                        "Root; d__Bacteria; p__Actinobacteriota; c__Actinomycetia; o__Mycobacteriales; f__Mycobacteriaceae; g__Nocardia; s__Nocardia grenadensis",
+                        "",
+                    ]),
+                    "\t".join([
+                        "S3.7.ribosomal_protein_S7",
+                        "sample_1",
+                        "TATCAGGTGCCTATTGAGGTAAGACCTGAAAGAAGACAGACTTTAGCGCTTCGCTGGATA",
+                        "1",
+                        "1.64",
+                        "Root; d__Bacteria; p__Actinobacteriota; c__Actinomycetia; o__Mycobacteriales; f__Mycobacteriaceae; g__Nocardia; s__Nocardia grenadensis2",
+                        "",
+                    ]),
+                    "\t".join([
+                        "S3.7.ribosomal_protein_S7",
+                        "sample_1",
+                        "TATCAGGTGCCTATTGAGGTAAGACCTGAAAGAAGACAGACTTTAGCGCTTCGCTGGATC",
+                        "5",
+                        "8.21",
+                        "Root; d__Bacteria; p__Actinobacteriota; c__Actinomycetia; o__Mycobacteriales; f__Mycobacteriaceae; g__Nocardia; s__Nocardia grenadensis3",
+                        "",
+                    ]),
+                    "\t".join([
+                        "S3.7.ribosomal_protein_S7",
+                        "sample_2",
+                        "TACCAGGTCCCGGTCGAGGTCCGTCCGATCCGCCAGACGACGCTCGCCCTGCGCTGGCTC",
+                        "3",
+                        "4.92",
+                        "Root; d__Bacteria; p__Actinobacteriota; c__Actinomycetia; o__Mycobacteriales; f__Mycobacteriaceae; g__Nocardia; s__Nocardia grenadensis",
+                        "",
+                    ]),
+                    "\t".join([
+                        "S3.7.ribosomal_protein_S7",
+                        "sample_3",
+                        "TATCAGGTGCCTATTGAGGTAAGACCTGAAAGAAGACAGACTTTAGCGCTTCGCTGGATA",
+                        "6",
+                        "9.85",
+                        "Root; d__Bacteria; p__Actinobacteriota; c__Actinomycetia; o__Mycobacteriales; f__Mycobacteriaceae; g__Nocardia; s__Nocardia grenadensis2",
+                        "",
+                    ]),
+                    ""
+                ]
+            )
+            with open(unbinned_path) as f:
+                self.assertEqual(expected, f.read())
+
+            edges_path = os.path.join("test", "coassemble", "target", "targets.tsv")
+            self.assertTrue(os.path.exists(edges_path))
+
+            edges_path = os.path.join("test", "coassemble", "target", "elusive_edges.tsv")
+            self.assertTrue(os.path.exists(edges_path))
+
+            cluster_path = os.path.join("test", "coassemble", "target", "elusive_clusters.tsv")
+            self.assertTrue(os.path.exists(cluster_path))
+
     def test_coassemble_single_assembly(self):
         with in_tempdir():
             cmd = (
diff --git a/test/test_query_processing.py b/test/test_query_processing.py
@@ -224,6 +224,31 @@ def test_query_processing_remove_EIF(self):
         self.assertDataFrameEqual(expected_binned, observed_binned)
         self.assertDataFrameEqual(expected_unbinned, observed_unbinned)
 
+    def test_query_processing_target_taxa(self):
+        query = pl.DataFrame([
+            ["sample_1", "AAA", 1, 5, 10, "genome_1", "S3.1", "AAA", "Root; d__Bacteria"],
+            ["sample_1", "AAB", 10, 5, 10, "genome_1", "S3.1", "AAA", "Root; d__Bacteria"],
+            ["sample_1", "CCC", 1, 5, 10, "genome_1", "S3.1", "AAA", "Root"],
+            ["sample_1", "CCD", 10, 5, 10, "genome_1", "S3.1", "AAA", "Root"],
+        ], schema=QUERY_COLUMNS)
+        pipe = pl.DataFrame([
+            ["S3.1", "sample_1", "AAA", 5, 10, "Root; d__Bacteria; p__Planctomycetota"],
+            ["S3.1", "sample_1", "AAB", 5, 10, "Root; d__Bacteria; p__Planctomycetota"],
+            ["S3.1", "sample_1", "CCC", 5, 10, "Root"],
+            ["S3.1", "sample_1", "CCD", 5, 10, "Root"],
+        ], schema=PIPE_COLUMNS)
+
+        expected_binned = pl.DataFrame([
+            ["S3.1", "sample_1", "AAA", 5, 10, "Root; d__Bacteria; p__Planctomycetota", "genome_1"],
+        ], schema=APPRAISE_COLUMNS)
+        expected_unbinned = pl.DataFrame([
+            ["S3.1", "sample_1", "AAB", 5, 10, "Root; d__Bacteria; p__Planctomycetota", None],
+        ], schema=APPRAISE_COLUMNS)
+
+        observed_binned, observed_unbinned = processing(query, pipe, TAXA_OF_INTEREST="p__Planctomycetota")
+        self.assertDataFrameEqual(expected_binned, observed_binned)
+        self.assertDataFrameEqual(expected_unbinned, observed_unbinned)
+
 
 if __name__ == '__main__':
     unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -211,6 +211,7 @@ rule query_processing:`
`211`	`211`	`params:`
`212`	`212`	`sequence_identity=config["appraise_sequence_identity"],`
`213`	`213`	`window_size=60,`
	`214`	`+ taxa_of_interest=config["taxa_of_interest"],`
`214`	`215`	`threads:`
`215`	`216`	`64`
`216`	`217`	`script:`