Repository: AroneyS/binchicken (language: python)

Issue:
Add input to restrict coassembly samples
- [x] Implement coassembly sample filtering in cluster_graph
- [x] Add to coassemble snakemake
- [x] Add `--coassembly-samples` argument
- [x] Add `--coassembly-samples-list` argument


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index c0dfcf6d..8536d29a 100644
--- a/README.md
+++ b/README.md
@@ -41,10 +41,9 @@ pip install ibis-genome
 ## Ibis coassemble
 
 Snakemake pipeline to discover coassembly sample clusters based on co-occurrence of single-copy marker genes, excluding those genes present in reference genomes (e.g. previously recovered genomes).
-Creates graph with samples as nodes and the number of overlapping sequences provided by SingleM.
 The taxa of the considered sequences can be filtered to target a specific taxon (e.g. the phylum Planctomycetota).
-The graph is clustered using the Girvan-Newman algorithm to provide sample groupings.
-Aviary assemble/recover commands are generated based on proposed coassemblies.
+Assembly and recovery can be run directly, or the coassemblies with differential-abudance-binning samples can be run in the tool of your choice.
+Aviary assemble/recover commands are also generated based on proposed coassemblies.
 Optionally, reads can be mapped to the matched bins with only unmapped reads being assembled.
 Paired end reads of form reads_1.1.fq, reads_1_1.fq and reads_1_R1.fq are automatically detected and matched to their basename.
 
@@ -95,7 +94,7 @@ ibis iterate --exclude-coassemblies reads_1,reads_2 --new-genomes new_genome_1.f
 
 ## Ibis update
 
-Applies further processing to a previous Ibis coassemble run: downloading SRA reads, generating unmapped reads files, and/or running Aviary commands.
+Applies further processing to a previous Ibis coassemble run: downloading SRA reads, generating unmapped reads files, and/or running assembly/recovery commands.
 
 ```bash
 # Example: update previous run to download SRA reads
diff --git a/ibis.yml b/ibis.yml
index 2d83a0eb..e910d5bd 100644
--- a/ibis.yml
+++ b/ibis.yml
@@ -5,9 +5,8 @@ channels:
   - defaults
 dependencies:
   - python=3.10.*
-  - snakemake=7.26.*
+  - snakemake=7.32.*
   - mamba=1.4.*
-  - pandas=2.0.*
   - networkx=3.1
   - bird_tool_utils_python=0.4.*
   - extern=0.4.*
diff --git a/ibis/__init__.py b/ibis/__init__.py
index 792c52c0..f8c6ac7f 100644
--- a/ibis/__init__.py
+++ b/ibis/__init__.py
@@ -1,1 +1,1 @@
-__version__ = "0.9.4.1"
+__version__ = "0.9.5"
diff --git a/ibis/config/template_coassemble.yaml b/ibis/config/template_coassemble.yaml
index 323ac0e3..599848ce 100644
--- a/ibis/config/template_coassemble.yaml
+++ b/ibis/config/template_coassemble.yaml
@@ -1,4 +1,5 @@
 max_threads: 8
+coassembly_samples: []
 singlem_metapackage: 
 genomes:
   genome: 
diff --git a/ibis/ibis.py b/ibis/ibis.py
index 569da4b1..2e64c347 100755
--- a/ibis/ibis.py
+++ b/ibis/ibis.py
@@ -390,6 +390,9 @@ def coassemble(args):
             os.path.join(args.output, "coassemble", "read_size.csv"),
         )
 
+    if args.coassembly_samples_list:
+        args.coassembly_samples = read_list(args.coassembly_samples_list)
+
     if args.exclude_coassemblies_list:
         args.exclude_coassemblies = read_list(args.exclude_coassemblies_list)
 
@@ -438,6 +441,7 @@ def coassemble(args):
         "reads_2": reverse_reads,
         "genomes": genomes if args.genomes else None,
         "singlem_metapackage": metapackage,
+        "coassembly_samples": args.coassembly_samples,
         # Clustering config
         "taxa_of_interest": args.taxa_of_interest if args.taxa_of_interest else None,
         "appraise_sequence_identity": args.appraise_sequence_identity / 100 if args.appraise_sequence_identity > 1 else args.appraise_sequence_identity,
@@ -783,6 +787,7 @@ def build(args):
     args.reverse_list = None
     args.genomes_list = None
     args.new_genomes_list = None
+    args.coassembly_samples_list = None
     args.aviary_gtdbtk_dir = "."
     args.aviary_checkm2_dir = "."
     args.aviary_cores = None
@@ -935,6 +940,8 @@ def add_base_arguments(argument_group):
         argument_group.add_argument("--reverse-list", help="input reverse nucleotide read sequence(s) newline separated")
         argument_group.add_argument("--genomes", nargs='+', help="Reference genomes for read mapping")
         argument_group.add_argument("--genomes-list", help="Reference genomes for read mapping newline separated")
+        argument_group.add_argument("--coassembly-samples", nargs='+', help="Restrict coassembly to these samples. Remaining samples will still be used for recovery [default: use all samples]", default=[])
+        argument_group.add_argument("--coassembly-samples-list", help="Restrict coassembly to these samples, newline separated. Remaining samples will still be used for recovery [default: use all samples]", default=[])
 
     def add_evaluation_options(argument_group):
         argument_group.add_argument("--checkm-version", type=int, help="CheckM version to use to quality cutoffs [default: 2]", default=2)
diff --git a/ibis/workflow/coassemble.smk b/ibis/workflow/coassemble.smk
index c558d1bc..9a2a6cc4 100644
--- a/ibis/workflow/coassemble.smk
+++ b/ibis/workflow/coassemble.smk
@@ -3,11 +3,9 @@
 #############
 ruleorder: no_genomes > query_processing > update_appraise > singlem_appraise
 ruleorder: mock_download_sra > download_sra
-localrules: all, summary, singlem_summarise_genomes, singlem_appraise_filtered, no_genomes, download_sra, mock_download_sra, collect_genomes, finish_mapping, finish_qc, aviary_commands, aviary_recover, aviary_combine
 
 import os
-os.environ['OPENBLAS_NUM_THREADS'] = '1'
-import pandas as pd
+import polars as pl
 from ibis.ibis import FAST_AVIARY_MODE
 
 output_dir = os.path.abspath("coassemble")
@@ -53,12 +51,12 @@ def get_cat(wildcards):
 
 def get_reads_coassembly(wildcards, forward=True, recover=False):
     checkpoint_output = checkpoints.cluster_graph.get(**wildcards).output[0]
-    elusive_clusters = pd.read_csv(checkpoint_output, sep = "\t")
+    elusive_clusters = pl.read_csv(checkpoint_output, separator="\t")
 
     if recover:
-        sample_names = elusive_clusters[elusive_clusters["coassembly"] == wildcards.coassembly]["recover_samples"].iloc[0]
+        sample_names = elusive_clusters.filter(pl.col("coassembly") == wildcards.coassembly).get_column("recover_samples").to_list()[0]
     else:
-        sample_names = elusive_clusters[elusive_clusters["coassembly"] == wildcards.coassembly]["samples"].iloc[0]
+        sample_names = elusive_clusters.filter(pl.col("coassembly") == wildcards.coassembly).get_column("samples").to_list()[0]
 
     sample_names = sample_names.split(",")
 
@@ -74,8 +72,8 @@ def get_reads_coassembly(wildcards, forward=True, recover=False):
 
 def get_coassemblies(wildcards):
     checkpoint_output = checkpoints.cluster_graph.get().output[0]
-    elusive_clusters = pd.read_csv(checkpoint_output, sep = "\t")
-    coassemblies = elusive_clusters["coassembly"].to_list()
+    elusive_clusters = pl.read_csv(checkpoint_output, separator="\t")
+    coassemblies = elusive_clusters.get_column("coassembly").to_list()
 
     return [output_dir + f"/coassemble/{c}/recover.done" for c in coassemblies]
 
@@ -86,6 +84,7 @@ rule all:
         output_dir + "/commands/recover_commands.sh" if not config["run_aviary"] else [],
         output_dir + "/commands/done" if config["run_aviary"] else [],
         output_dir + "/summary.tsv",
+    localrule: True
 
 rule summary:
     input:
@@ -93,6 +92,7 @@ rule summary:
         read_size = output_dir + "/unmapped_read_size.csv" if config["assemble_unmapped"] else [],
     output:
         summary = output_dir + "/summary.tsv",
+    localrule: True
     script:
         "scripts/summarise_coassemblies.py"
 
@@ -179,6 +179,7 @@ rule singlem_summarise_genomes:
         logs_dir + "/summarise/{version,.*}genomes.log"
     params:
         singlem_metapackage = config["singlem_metapackage"]
+    localrule: True
     conda:
         "env/singlem.yml"
     shell:
@@ -231,6 +232,7 @@ rule singlem_appraise_filtered:
         binned = output_dir + "/appraise/binned.otu_table.tsv",
     params:
         bad_package = "S3.18.EIF_2_alpha",
+    localrule: True
     shell:
         "grep -v {params.bad_package} {input.unbinned} > {output.unbinned} && "
         "grep -v {params.bad_package} {input.binned} > {output.binned}"
@@ -303,6 +305,7 @@ rule no_genomes:
     output:
         unbinned = temp(output_dir + "/appraise/unbinned_raw.otu_table.tsv") if config["no_genomes"] else [],
         binned = temp(output_dir + "/appraise/binned_raw.otu_table.tsv") if config["no_genomes"] else [],
+    localrule: True
     log:
         logs_dir + "/appraise/appraise.log"
     script:
@@ -362,6 +365,7 @@ checkpoint cluster_graph:
         num_coassembly_samples = config["num_coassembly_samples"],
         max_coassembly_samples = config["max_coassembly_samples"],
         max_recovery_samples = config["max_recovery_samples"],
+        coassembly_samples = config["coassembly_samples"],
         exclude_coassemblies = config["exclude_coassemblies"],
     threads: 64
     resources:
@@ -405,6 +409,7 @@ rule download_sra:
         expand(output_dir + "/sra/{read}.done", read=config["sra"]) if config["sra"] else [],
     output:
         output_dir + "/sra/all_done"
+    localrule: True
     shell:
         "touch {output}"
 
@@ -416,6 +421,7 @@ rule mock_download_sra:
         sra_u = workflow.basedir + "/../../test/data/sra/" + config["sra"][0] + ".fastq.gz" if config["sra"] else "",
         sra_f = " ".join([workflow.basedir + "/../../test/data/sra/" + s + "_1.fastq.gz" for s in config["sra"][1:]]) if config["sra"] else "",
         sra_r = " ".join([workflow.basedir + "/../../test/data/sra/" + s + "_2.fastq.gz" for s in config["sra"][1:]]) if config["sra"] else "",
+    localrule: True
     conda:
         "env/kingfisher.yml"
     log:
@@ -473,6 +479,7 @@ rule collect_genomes:
         genomes = config["genomes"],
         sample = "{read}",
         min_appraised = config["unmapping_min_appraised"],
+    localrule: True
     script:
         "scripts/collect_reference_bins.py"
 
@@ -562,6 +569,7 @@ rule finish_mapping:
         mapped_reads_2.values(),
     output:
         output_dir + "/mapping/done"
+    localrule: True
     shell:
         "touch {output}"
 
@@ -571,6 +579,7 @@ rule finish_qc:
         qc_reads_2.values(),
     output:
         output_dir + "/qc/done"
+    localrule: True
     shell:
         "touch {output}"
 
@@ -592,6 +601,7 @@ rule aviary_commands:
         memory = config["aviary_memory"],
         threads = config["aviary_threads"],
         speed = config["aviary_speed"],
+    localrule: True
     log:
         logs_dir + "/aviary_commands.log"
     script:
@@ -663,6 +673,7 @@ rule aviary_recover:
         fast = "--workflow recover_mags_no_singlem --skip-binners maxbin concoct rosella --skip-abundances --refinery-max-iterations 0" if config["aviary_speed"] == FAST_AVIARY_MODE else "",
         snakemake_profile = f"--snakemake-profile {config['snakemake_profile']}" if config["snakemake_profile"] else "",
         cluster_retries = f"--cluster-retries {config['cluster_retries']}" if config["cluster_retries"] else "",
+    localrule: True
     threads:
         1
     resources:
@@ -701,5 +712,6 @@ rule aviary_combine:
         mapping = output_dir + "/mapping/done" if config["assemble_unmapped"] else [],
     output:
         output_dir + "/commands/done",
+    localrule: True
     shell:
         "touch {output} "
diff --git a/ibis/workflow/evaluate.smk b/ibis/workflow/evaluate.smk
index 8824ffb1..6cb8622a 100644
--- a/ibis/workflow/evaluate.smk
+++ b/ibis/workflow/evaluate.smk
@@ -2,8 +2,6 @@
 ### Setup ###
 #############
 import os
-os.environ['OPENBLAS_NUM_THREADS'] = '1'
-import pandas as pd
 import re
 
 output_dir = os.path.abspath("evaluate")
diff --git a/ibis/workflow/scripts/cluster_graph.py b/ibis/workflow/scripts/cluster_graph.py
index 9fa0f1a3..5f27483c 100644
--- a/ibis/workflow/scripts/cluster_graph.py
+++ b/ibis/workflow/scripts/cluster_graph.py
@@ -99,6 +99,7 @@ def pipeline(
         MAX_RECOVERY_SAMPLES=20,
         MIN_CLUSTER_TARGETS=1,
         MAX_SAMPLES_COMBINATIONS=100,
+        COASSEMBLY_SAMPLES=[],
         EXCLUDE_COASSEMBLIES=[]):
 
     logging.info(f"Polars using {str(pl.threadpool_size())} threads")
@@ -127,6 +128,17 @@ def pipeline(
                 )
         )
 
+        if COASSEMBLY_SAMPLES:
+            coassembly_edges = (
+                elusive_edges
+                .with_columns(
+                    pl.col("samples").list.eval(pl.element().filter(pl.element().is_in(COASSEMBLY_SAMPLES)))
+                    )
+                .filter(pl.col("samples").list.lengths() >= pl.col("cluster_size"))
+            )
+        else:
+            coassembly_edges = elusive_edges
+
         read_size = (
             read_size
             .lazy()
@@ -154,6 +166,7 @@ def pipeline(
             clusters = [
                 elusive_edges
                 .explode("samples")
+                .filter((not COASSEMBLY_SAMPLES) | pl.col("samples").is_in(COASSEMBLY_SAMPLES))
                 .group_by("samples")
                 .agg(pl.col("target_ids").flatten())
                 .with_columns(
@@ -165,7 +178,7 @@ def pipeline(
         else:
             logging.info("Forming candidate sample clusters")
             clusters = [
-                elusive_edges
+                coassembly_edges
                 .filter(pl.col("style") == "match")
                 .filter(pl.col("cluster_size") >= MIN_COASSEMBLY_SAMPLES)
                 .filter(pl.col("target_ids").list.lengths() >= MIN_CLUSTER_TARGETS)
@@ -174,7 +187,7 @@ def pipeline(
 
             if is_pooled:
                 clusters.append(
-                    elusive_edges
+                    coassembly_edges
                     .filter(pl.col("style") == "pool")
                     .filter(pl.col("cluster_size") >= MIN_COASSEMBLY_SAMPLES)
                     # Prevent combinatorial explosion (also, large clusters are less useful for distinguishing between clusters)
@@ -242,7 +255,7 @@ def pipeline(
             .collect(streaming=True)
             .pipe(
                 join_list_subsets,
-                df2=elusive_edges
+                df2=coassembly_edges
                     .filter(pl.col("style") == "pool")
                     .filter(pl.col("samples").list.lengths() >= MAX_SAMPLES_COMBINATIONS)
                     .collect(streaming=True),
@@ -290,6 +303,7 @@ def pipeline(
     MAX_COASSEMBLY_SAMPLES = snakemake.params.max_coassembly_samples
     MIN_COASSEMBLY_SAMPLES = snakemake.params.num_coassembly_samples
     MAX_RECOVERY_SAMPLES = snakemake.params.max_recovery_samples
+    COASSEMBLY_SAMPLES = snakemake.params.coassembly_samples
     EXCLUDE_COASSEMBLIES = snakemake.params.exclude_coassemblies
     elusive_edges_path = snakemake.input.elusive_edges
     read_size_path = snakemake.input.read_size
@@ -310,6 +324,7 @@ def pipeline(
         MAX_COASSEMBLY_SAMPLES=MAX_COASSEMBLY_SAMPLES,
         MIN_COASSEMBLY_SAMPLES=MIN_COASSEMBLY_SAMPLES,
         MAX_RECOVERY_SAMPLES=MAX_RECOVERY_SAMPLES,
+        COASSEMBLY_SAMPLES=COASSEMBLY_SAMPLES,
         EXCLUDE_COASSEMBLIES=EXCLUDE_COASSEMBLIES,
         MIN_CLUSTER_TARGETS=min_cluster_targets,
         MAX_SAMPLES_COMBINATIONS=100,
diff --git a/ibis/workflow/scripts/separate_SingleM_seq.py b/ibis/workflow/scripts/separate_SingleM_seq.py
deleted file mode 100755
index 9176cda7..00000000
--- a/ibis/workflow/scripts/separate_SingleM_seq.py
+++ /dev/null
@@ -1,24 +0,0 @@
-###############################
-### separate_SingleM_seq.py ###
-###############################
-# Author: Samuel Aroney
-
-import os
-os.environ['OPENBLAS_NUM_THREADS'] = '1'
-import pandas as pd
-import re
-from Bio.Seq import Seq
-from Bio.SeqRecord import SeqRecord
-from Bio import SeqIO
-
-otu_table = pd.read_csv(snakemake.input.otu_table, sep="\t")
-
-otu_table["id"] = otu_table.index
-otu_table["id"] = otu_table["id"].astype(str)
-
-for gene in otu_table["gene"].unique():
-    gene_table = otu_table[otu_table["gene"] == gene]
-    records = [SeqRecord(Seq(re.sub("-", "", gene_table.iloc[i]["sequence"])), id = gene_table.iloc[i]["id"]) for i in range(len(gene_table))]
-    SeqIO.write(records, snakemake.params.output_dir + "/" + gene + ".fasta", 'fasta')
-
-otu_table.to_csv(snakemake.output.id_otu_table, sep="\t", index=False)
diff --git a/ibis/workflow/scripts/summarise_coassemblies.py b/ibis/workflow/scripts/summarise_coassemblies.py
index 4c462eb9..31cfa533 100644
--- a/ibis/workflow/scripts/summarise_coassemblies.py
+++ b/ibis/workflow/scripts/summarise_coassemblies.py
@@ -4,15 +4,41 @@
 # Author: Samuel Aroney
 
 import os
-os.environ['OPENBLAS_NUM_THREADS'] = '1'
-import pandas as pd
+import polars as pl
 
-elusive_clusters = pd.read_csv(snakemake.input.elusive_clusters, sep="\t")
-summary = elusive_clusters[["coassembly", "samples", "length", "total_targets", "total_size"]]
+def processing(elusive_clusters, read_size):
+    print(f"Polars using {str(pl.threadpool_size())} threads")
 
-if snakemake.input.read_size:
-    read_size = pd.read_csv(snakemake.input.read_size, names = ["sample", "read_size"])
-    read_sizes = read_size.set_index("sample").to_dict()["read_size"]
-    summary["unmapped_size"] = summary["samples"].apply(lambda x: sum([read_sizes[sample] for sample in x.split(",")]))
+    summary = (
+        elusive_clusters
+        .select("coassembly", "samples", "length", "total_targets", "total_size")
+    )
 
-summary.to_csv(snakemake.output.summary, sep="\t", index=False)
+    if read_size is not None:
+        summary = (
+            summary
+            .with_columns(sample = pl.col("samples").str.split(","))
+            .explode("sample")
+            .join(read_size, on="sample", how="left")
+            .group_by("coassembly", "samples", "length", "total_targets", "total_size")
+            .agg(unmapped_size = pl.sum("read_size"))
+        )
+
+    return summary
+
+if __name__ == "__main__":
+    os.environ["POLARS_MAX_THREADS"] = str(snakemake.threads)
+    import polars as pl
+    
+    elusive_clusters_path = snakemake.input.elusive_clusters
+    read_size_path = snakemake.input.read_size
+    output_path = snakemake.output.summary
+
+    elusive_clusters = pl.read_csv(elusive_clusters_path, separator="\t")
+    if read_size_path:
+        read_size = pl.read_csv(read_size_path, has_header=False, new_columns=["sample", "read_size"])
+    else:
+        read_size = None
+
+    summary = processing(elusive_clusters, read_size)
+    summary.write_csv(output_path, separator="\t")
diff --git a/pyproject.toml b/pyproject.toml
index f506779b..04defe3d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -14,7 +14,6 @@ license = { file="LICENSE" }
 requires-python = ">=3.10"
 dependencies = ["snakemake>=6.0.5",
   "mamba",
-  "pandas",
   "networkx",
   "bird_tool_utils",
   "extern",


Test patch:
diff --git a/.github/workflows/test-ibis.yml b/.github/workflows/test-ibis.yml
index c8b1ab24..99b41da7 100644
--- a/.github/workflows/test-ibis.yml
+++ b/.github/workflows/test-ibis.yml
@@ -36,6 +36,8 @@ jobs:
           python test/test_target_elusive.py -b
           python test/test_cluster_graph.py -b
           python test/test_collect_reference_bins.py -b
+          python test/test_summarise_coassemblies.py -b
+          python test/test_no_genomes.py -b
           python test/test_aviary_commands.py
           python test/test_is_interleaved.py
           python test/test_build.py
@@ -43,3 +45,4 @@ jobs:
           python test/test_evaluate.py
           python test/test_evaluate_script.py -b
           python test/test_coassemble.py
+          python test/test_iterate.py
diff --git a/test/test_cluster_graph.py b/test/test_cluster_graph.py
index d9e4bea5..89ea3bd6 100644
--- a/test/test_cluster_graph.py
+++ b/test/test_cluster_graph.py
@@ -497,6 +497,92 @@ def test_cluster_exclude_coassemblies(self):
         observed = pipeline(elusive_edges, read_size, EXCLUDE_COASSEMBLIES=["2,3", "5,6"])
         self.assertDataFrameEqual(expected, observed)
 
+    def test_cluster_restrict_coassembly_samples(self):
+        elusive_edges = pl.DataFrame([
+            ["match", 2, "1,2", "1,2,3"],
+            ["match", 2, "1,3", "1,3"],
+            ["match", 2, "1,4", "1,4"],
+            ["match", 2, "2,3", "2,3"],
+            ["match", 2, "2,4", "2,4"],
+            ["match", 2, "3,4", "1,3,4"],
+            ["match", 2, "4,5", "5"],
+            ["match", 2, "4,6", "5"],
+            ["match", 2, "5,6", "5,6,7"],
+        ], schema = ELUSIVE_EDGES_COLUMNS)
+        read_size = pl.DataFrame([
+            ["1", 1000],
+            ["2", 2000],
+            ["3", 3000],
+            ["4", 4000],
+            ["5", 5000],
+            ["6", 6000],
+        ], schema=READ_SIZE_COLUMNS)
+
+        expected = pl.DataFrame([
+            ["1,2", 2, 3, 3000, "1,2,3,4", "coassembly_0"],
+            ["4,5", 2, 1, 9000, "4,5,6", "coassembly_1"],
+        ], schema=ELUSIVE_CLUSTERS_COLUMNS)
+        observed = pipeline(
+            elusive_edges,
+            read_size,
+            MAX_RECOVERY_SAMPLES=4,
+            COASSEMBLY_SAMPLES=["1", "2", "4", "5"],
+            )
+        self.assertDataFrameEqual(expected, observed)
+
+    def test_cluster_restrict_coassembly_samples_no_overlap(self):
+        elusive_edges = pl.DataFrame([
+            ["match", 2, "1,2", "1,2"],
+            ["match", 2, "1,3", "3"],
+            ["match", 2, "2,3", "3"],
+            ["match", 2, "4,5", "1"],
+        ], schema = ELUSIVE_EDGES_COLUMNS)
+        read_size = pl.DataFrame([
+            ["1", 1000],
+            ["2", 2000],
+            ["3", 3000],
+            ["4", 4000],
+            ["5", 5000],
+            ["6", 6000],
+        ], schema=READ_SIZE_COLUMNS)
+
+        expected = pl.DataFrame([
+            ["1,2", 2, 2, 3000, "1,2,4,5", "coassembly_0"],
+        ], schema=ELUSIVE_CLUSTERS_COLUMNS)
+        observed = pipeline(
+            elusive_edges,
+            read_size,
+            MAX_RECOVERY_SAMPLES=4,
+            COASSEMBLY_SAMPLES=["1", "2"],
+            )
+        self.assertDataFrameEqual(expected, observed)
+
+    def test_cluster_restrict_coassembly_samples_single_assembly(self):
+        elusive_edges = pl.DataFrame([
+            ["match", 2, "1,2", "1,2,3,4"],
+            ["match", 2, "3,1", "5"],
+            ["match", 2, "3,2", "6,7"],
+        ], schema = ELUSIVE_EDGES_COLUMNS)
+        read_size = pl.DataFrame([
+            ["1", 1000],
+            ["2", 1000],
+            ["3", 1000],
+        ], schema=READ_SIZE_COLUMNS)
+
+        expected = pl.DataFrame([
+            ["2", 1, 6, 1000, "1,2,3", "coassembly_0"],
+            ["1", 1, 5, 1000, "1,2,3", "coassembly_1"],
+        ], schema=ELUSIVE_CLUSTERS_COLUMNS)
+        observed = pipeline(
+            elusive_edges,
+            read_size,
+            MAX_COASSEMBLY_SAMPLES=1,
+            MIN_COASSEMBLY_SAMPLES=1,
+            MAX_RECOVERY_SAMPLES=4,
+            COASSEMBLY_SAMPLES=["1", "2"],
+            )
+        self.assertDataFrameEqual(expected, observed)
+
     def test_join_list_subsets(self):
         with pl.StringCache():
             df1 = (
diff --git a/test/test_coassemble.py b/test/test_coassemble.py
index b126fe04..e5605a13 100644
--- a/test/test_coassemble.py
+++ b/test/test_coassemble.py
@@ -578,6 +578,7 @@ def test_coassemble_single_assembly(self):
                 f"--reverse {SAMPLE_READS_REVERSE} "
                 f"--singlem-metapackage {METAPACKAGE} "
                 f"--single-assembly "
+                f"--coassembly-samples sample_1 sample_2 "
                 f"--no-genomes "
                 f"--output test "
                 f"--conda-prefix {path_to_conda} "
@@ -617,14 +618,6 @@ def test_coassemble_single_assembly(self):
                         "sample_1,sample_2,sample_3",
                         "coassembly_1"
                     ]),
-                    "\t".join([
-                        "sample_3",
-                        "1",
-                        "2",
-                        "3624",
-                        "sample_1,sample_2,sample_3",
-                        "coassembly_2"
-                    ]),
                     ""
                 ]
             )
@@ -809,6 +802,7 @@ def test_coassemble_default_config(self):
             self.assertEqual(config["assemble_unmapped"], False)
             self.assertEqual(config["aviary_threads"], 64)
             self.assertEqual(config["aviary_memory"], 500)
+            self.assertEqual(config["coassembly_samples"], [])
 
     def test_coassemble_singlem_inputs(self):
         with in_tempdir():
@@ -977,6 +971,7 @@ def test_coassemble_files_of_paths(self):
             write_string_to_file(SAMPLE_READS_REVERSE, "sample_reads_reverse")
             write_string_to_file(GENOMES, "genomes")
             write_string_to_file(GENOME_TRANSCRIPTS, "genome_transcripts")
+            write_string_to_file("sample_1\nsample_2", "coassembly_samples")
             write_string_to_file("sample_1,sample_2", "exclude_coassemblies")
 
             cmd = (
@@ -986,6 +981,7 @@ def test_coassemble_files_of_paths(self):
                 f"--genomes-list genomes "
                 f"--genome-transcripts-list genome_transcripts "
                 f"--singlem-metapackage {METAPACKAGE} "
+                f"--coassembly-samples-list coassembly_samples "
                 f"--exclude-coassemblies-list exclude_coassemblies "
                 f"--assemble-unmapped "
                 f"--output test "
@@ -1013,6 +1009,7 @@ def test_coassemble_files_of_paths(self):
 
             config = load_configfile(os.path.join("test", "config.yaml"))
             self.assertEqual(config["genomes"], {os.path.splitext(os.path.basename(g))[0]: g for g in GENOMES.split(" ")})
+            self.assertEqual(config["coassembly_samples"], ["sample_1", "sample_2"])
             self.assertEqual(config["exclude_coassemblies"], ["sample_1,sample_2"])
             self.assertEqual(config["reads_1"], {os.path.splitext(os.path.splitext(os.path.basename(s))[0])[0]: s for s in SAMPLE_READS_FORWARD.split(" ")})
             self.assertEqual(config["reads_2"], {os.path.splitext(os.path.splitext(os.path.basename(s))[0])[0]: s for s in SAMPLE_READS_REVERSE.split(" ")})
diff --git a/test/test_evaluate.py b/test/test_evaluate.py
index 4bdd4f51..09bd58f1 100644
--- a/test/test_evaluate.py
+++ b/test/test_evaluate.py
@@ -127,17 +127,16 @@ def test_evaluate(self):
 
             summarise_path = os.path.join("test", "evaluate", "summarise", "bins_summarised.otu_table.tsv")
             self.assertTrue(os.path.exists(summarise_path))
-            expected = "\n".join(
-                [
-                    "\t".join(["gene", "sample", "sequence", "num_hits", "coverage", "taxonomy"]),
-                    "\t".join(["S3.7.ribosomal_protein_S7", "coassembly_0-0_transcripts", "TATCAAGTTCCACAAGAAGTTAGAGGAGAAAGAAGAATCTCGTTAGCTATTAGATGGATT", "1", "1.14", "Root; d__Bacteria"]),
-                    "\t".join(["S3.7.ribosomal_protein_S7", "coassembly_0-1_transcripts", "TTCCAGGTGCCTACTGAAGTTCGTCCCGAGCGTAAAATTGCATTGGGTATGAAATGGCTG", "1", "1.18", "Root; d__Bacteria; p__Bacteroidota; c__Bacteroidia; o__Sphingobacteriales; f__Sphingobacteriaceae; g__Mucilaginibacter; s__Mucilaginibacter_sp013286235"]),
-                    "\t".join(["S3.7.ribosomal_protein_S7", "coassembly_0-1_transcripts", "TACCAGGTTCCTGTTGAAGTTCGCCCATCACGCCGTTCTGCTTTGGCAATGCGCTGGGTG", "1", "1.14", "Root; d__Bacteria; p__Proteobacteria; c__Gammaproteobacteria; o__Burkholderiales; f__Burkholderiaceae; g__Polynucleobacter; s__Polynucleobacter_sp018688335"]),
-                    ""
-                ]
-            )
             with open(summarise_path) as f:
-                self.assertEqual(expected, f.read())
+                observed = f.read()
+
+            otu1 = "\t".join(["S3.7.ribosomal_protein_S7", "coassembly_0-0_transcripts", "TATCAAGTTCCACAAGAAGTTAGAGGAGAAAGAAGAATCTCGTTAGCTATTAGATGGATT", "1", "1.14", "Root; d__Bacteria"])
+            otu2 = "\t".join(["S3.7.ribosomal_protein_S7", "coassembly_0-1_transcripts", "TACCAGGTTCCTGTTGAAGTTCGCCCATCACGCCGTTCTGCTTTGGCAATGCGCTGGGTG", "1", "1.14", "Root; d__Bacteria; p__Proteobacteria; c__Gammaproteobacteria; o__Burkholderiales; f__Burkholderiaceae; g__Polynucleobacter; s__Polynucleobacter_sp018688335"])
+            otu3 = "\t".join(["S3.7.ribosomal_protein_S7", "coassembly_0-1_transcripts", "TTCCAGGTGCCTACTGAAGTTCGTCCCGAGCGTAAAATTGCATTGGGTATGAAATGGCTG", "1", "1.18", "Root; d__Bacteria; p__Bacteroidota; c__Bacteroidia; o__Sphingobacteriales; f__Sphingobacteriaceae; g__Mucilaginibacter; s__Mucilaginibacter_sp013286235"])
+
+            self.assertTrue(otu1 in observed)
+            self.assertTrue(otu2 in observed)
+            self.assertTrue(otu3 in observed)
 
     def test_evaluate_specified_files(self):
         with in_tempdir():
@@ -259,17 +258,16 @@ def test_evaluate_genome_input(self):
 
             summarise_path = os.path.join("test", "evaluate", "summarise", "bins_summarised.otu_table.tsv")
             self.assertTrue(os.path.exists(summarise_path))
-            expected = "\n".join(
-                [
-                    "\t".join(["gene", "sample", "sequence", "num_hits", "coverage", "taxonomy"]),
-                    "\t".join(["S3.7.ribosomal_protein_S7", "coassembly_0-bin_1_transcripts", "TATCAAGTTCCACAAGAAGTTAGAGGAGAAAGAAGAATCTCGTTAGCTATTAGATGGATT", "1", "1.14", "Root; d__Bacteria"]),
-                    "\t".join(["S3.7.ribosomal_protein_S7", "coassembly_0-bin_3_transcripts", "TTCCAGGTGCCTACTGAAGTTCGTCCCGAGCGTAAAATTGCATTGGGTATGAAATGGCTG", "1", "1.14", "Root; d__Bacteria; p__Bacteroidota; c__Bacteroidia; o__Sphingobacteriales; f__Sphingobacteriaceae; g__Mucilaginibacter; s__Mucilaginibacter_sp013286235"]),
-                    "\t".join(["S3.7.ribosomal_protein_S7", "coassembly_0-bin_3_transcripts", "TACCAGGTTCCTGTTGAAGTTCGCCCATCACGCCGTTCTGCTTTGGCAATGCGCTGGGTG", "1", "1.14", "Root; d__Bacteria; p__Proteobacteria; c__Gammaproteobacteria; o__Burkholderiales; f__Burkholderiaceae; g__Polynucleobacter; s__Polynucleobacter_sp018688335"]),
-                    ""
-                ]
-            )
             with open(summarise_path) as f:
-                self.assertEqual(expected, f.read())
+                observed = f.read()
+
+            otu1 = "\t".join(["S3.7.ribosomal_protein_S7", "coassembly_0-bin_1_transcripts", "TATCAAGTTCCACAAGAAGTTAGAGGAGAAAGAAGAATCTCGTTAGCTATTAGATGGATT", "1", "1.14", "Root; d__Bacteria"])
+            otu2 = "\t".join(["S3.7.ribosomal_protein_S7", "coassembly_0-bin_3_transcripts", "TACCAGGTTCCTGTTGAAGTTCGCCCATCACGCCGTTCTGCTTTGGCAATGCGCTGGGTG", "1", "1.14", "Root; d__Bacteria; p__Proteobacteria; c__Gammaproteobacteria; o__Burkholderiales; f__Burkholderiaceae; g__Polynucleobacter; s__Polynucleobacter_sp018688335"])
+            otu3 = "\t".join(["S3.7.ribosomal_protein_S7", "coassembly_0-bin_3_transcripts", "TTCCAGGTGCCTACTGAAGTTCGTCCCGAGCGTAAAATTGCATTGGGTATGAAATGGCTG", "1", "1.14", "Root; d__Bacteria; p__Bacteroidota; c__Bacteroidia; o__Sphingobacteriales; f__Sphingobacteriaceae; g__Mucilaginibacter; s__Mucilaginibacter_sp013286235"])
+
+            self.assertTrue(otu1 in observed)
+            self.assertTrue(otu2 in observed)
+            self.assertTrue(otu3 in observed)
 
     def test_evaluate_genome_files_input(self):
         with in_tempdir():
@@ -333,17 +331,16 @@ def test_evaluate_metapackage_env_variable(self):
 
             summarise_path = os.path.join("test", "evaluate", "summarise", "bins_summarised.otu_table.tsv")
             self.assertTrue(os.path.exists(summarise_path))
-            expected = "\n".join(
-                [
-                    "\t".join(["gene", "sample", "sequence", "num_hits", "coverage", "taxonomy"]),
-                    "\t".join(["S3.7.ribosomal_protein_S7", "coassembly_0-0_transcripts", "TATCAAGTTCCACAAGAAGTTAGAGGAGAAAGAAGAATCTCGTTAGCTATTAGATGGATT", "1", "1.14", "Root; d__Bacteria"]),
-                    "\t".join(["S3.7.ribosomal_protein_S7", "coassembly_0-1_transcripts", "TTCCAGGTGCCTACTGAAGTTCGTCCCGAGCGTAAAATTGCATTGGGTATGAAATGGCTG", "1", "1.18", "Root; d__Bacteria; p__Bacteroidota; c__Bacteroidia; o__Sphingobacteriales; f__Sphingobacteriaceae; g__Mucilaginibacter; s__Mucilaginibacter_sp013286235"]),
-                    "\t".join(["S3.7.ribosomal_protein_S7", "coassembly_0-1_transcripts", "TACCAGGTTCCTGTTGAAGTTCGCCCATCACGCCGTTCTGCTTTGGCAATGCGCTGGGTG", "1", "1.14", "Root; d__Bacteria; p__Proteobacteria; c__Gammaproteobacteria; o__Burkholderiales; f__Burkholderiaceae; g__Polynucleobacter; s__Polynucleobacter_sp018688335"]),
-                    ""
-                ]
-            )
             with open(summarise_path) as f:
-                self.assertEqual(expected, f.read())
+                observed = f.read()
+
+            otu1 = "\t".join(["S3.7.ribosomal_protein_S7", "coassembly_0-0_transcripts", "TATCAAGTTCCACAAGAAGTTAGAGGAGAAAGAAGAATCTCGTTAGCTATTAGATGGATT", "1", "1.14", "Root; d__Bacteria"])
+            otu2 = "\t".join(["S3.7.ribosomal_protein_S7", "coassembly_0-1_transcripts", "TACCAGGTTCCTGTTGAAGTTCGCCCATCACGCCGTTCTGCTTTGGCAATGCGCTGGGTG", "1", "1.14", "Root; d__Bacteria; p__Proteobacteria; c__Gammaproteobacteria; o__Burkholderiales; f__Burkholderiaceae; g__Polynucleobacter; s__Polynucleobacter_sp018688335"])
+            otu3 = "\t".join(["S3.7.ribosomal_protein_S7", "coassembly_0-1_transcripts", "TTCCAGGTGCCTACTGAAGTTCGTCCCGAGCGTAAAATTGCATTGGGTATGAAATGGCTG", "1", "1.18", "Root; d__Bacteria; p__Bacteroidota; c__Bacteroidia; o__Sphingobacteriales; f__Sphingobacteriaceae; g__Mucilaginibacter; s__Mucilaginibacter_sp013286235"])
+
+            self.assertTrue(otu1 in observed)
+            self.assertTrue(otu2 in observed)
+            self.assertTrue(otu3 in observed)
 
     def test_evaluate_cluster(self):
         with in_tempdir():
diff --git a/test/test_summarise_coassemblies.py b/test/test_summarise_coassemblies.py
new file mode 100644
index 00000000..fd7054b8
--- /dev/null
+++ b/test/test_summarise_coassemblies.py
@@ -0,0 +1,96 @@
+#!/usr/bin/env python3
+
+import unittest
+import os
+os.environ["POLARS_MAX_THREADS"] = "1"
+import polars as pl
+from polars.testing import assert_frame_equal
+from ibis.workflow.scripts.summarise_coassemblies import processing
+
+ELUSIVE_CLUSTERS_COLUMNS={
+    "coassembly": str,
+    "samples": str,
+    "length": int,
+    "total_targets": int,
+    "total_size": int,
+    }
+
+READ_SIZE_COLUMNS={
+    "sample": str,
+    "read_size": int,
+    }
+
+SUMMARY_COLUMNS={
+    "coassembly": str,
+    "samples": str,
+    "length": int,
+    "total_targets": int,
+    "total_size": int,
+    "unmapped_size": int,
+    }
+
+class Tests(unittest.TestCase):
+    def assertDataFrameEqual(self, a, b):
+        assert_frame_equal(a, b, check_dtype=False)
+
+    def test_summarise_coassemblies(self):
+        elusive_clusters = pl.DataFrame([
+            ["coassembly_1", "sample_1,sample_2", 2, 200, 1000],
+            ["coassembly_2", "sample_1,sample_3", 2, 100, 2000],
+        ], schema=ELUSIVE_CLUSTERS_COLUMNS)
+
+        read_size = pl.DataFrame([
+            ["sample_1", 100],
+            ["sample_2", 200],
+            ["sample_3", 300],
+        ], schema=READ_SIZE_COLUMNS)
+
+        expected = pl.DataFrame([
+            ["coassembly_1", "sample_1,sample_2", 2, 200, 1000, 300],
+            ["coassembly_2", "sample_1,sample_3", 2, 100, 2000, 400],
+        ], schema=SUMMARY_COLUMNS)
+
+        observed = processing(elusive_clusters, read_size)
+        self.assertDataFrameEqual(expected, observed)
+
+    def test_summarise_coassemblies_empty_input(self):
+        elusive_clusters = pl.DataFrame([], schema=ELUSIVE_CLUSTERS_COLUMNS)
+        read_size = pl.DataFrame([], schema=READ_SIZE_COLUMNS)
+
+        expected = pl.DataFrame([], schema=SUMMARY_COLUMNS)
+
+        observed = processing(elusive_clusters, read_size)
+        self.assertDataFrameEqual(expected, observed)
+
+    def test_summarise_coassemblies_empty_elusive_clusters(self):
+        elusive_clusters = pl.DataFrame([], schema=ELUSIVE_CLUSTERS_COLUMNS)
+        read_size = pl.DataFrame([
+            ["sample_1", 100],
+            ["sample_2", 200],
+            ["sample_3", 300],
+        ], schema=READ_SIZE_COLUMNS)
+
+        expected = pl.DataFrame([], schema=SUMMARY_COLUMNS)
+
+        observed = processing(elusive_clusters, read_size)
+        self.assertDataFrameEqual(expected, observed)
+
+    def test_summarise_coassemblies_empty_read_size(self):
+        elusive_clusters = pl.DataFrame([
+            ["coassembly_1", "sample_1,sample_2", 2, 200, 1000],
+            ["coassembly_2", "sample_1,sample_3", 2, 100, 2000],
+        ], schema=ELUSIVE_CLUSTERS_COLUMNS)
+
+        read_size = None
+
+        expected = pl.DataFrame([
+            ["coassembly_1", "sample_1,sample_2", 2, 200, 1000],
+            ["coassembly_2", "sample_1,sample_3", 2, 100, 2000],
+        ], schema={k:v for k,v in SUMMARY_COLUMNS.items() if k != "unmapped_size"})
+
+        observed = processing(elusive_clusters, read_size)
+        self.assertDataFrameEqual(expected, observed)
+
+
+if __name__ == '__main__':
+    unittest.main()