Merge pull request #137 from MicrobialDarkMatter/update-epymetheus

Update epymetheus

Author: SorenHeidelbach

flake.lock

@@ -38,7 +38,7 @@
             pkgs.glib
 
             # For python package
-            pkgs.python311Full
+            pkgs.python3
           ];
           shellHook = ''
             export LD_LIBRARY_PATH="${pkgs.libGL}/lib/:${pkgs.stdenv.cc.cc.lib}/lib/:${pkgs.glib.out}/lib/:$LD_LIBRARY_PATH"

flake.nix

@@ -1,7 +1,8 @@
 # Data loading functionalities
 import polars as pl
+
 from nanomotif.seq import Assembly
-from epymetheus import query_pileup_records
+from epymetheus import query_pileup_records, PileupColumn
 import sys
 import pysam
 import pyfastx
@@ -105,7 +106,38 @@ def load_contigs_pileup_bgzip(path: str, contigs: list[str]):
     """
     # Step 2: query pileup from Rust
     log.debug(f"Querying pileup for {len(contigs)} contigs")
-    pileup = query_pileup_records(path, contigs)
+    pileup = query_pileup_records(
+        path,
+        contigs,
+        columns = [
+            PileupColumn.Contig,
+            PileupColumn.Start,
+            PileupColumn.ModType,
+            PileupColumn.Strand,
+            PileupColumn.FractionModified,
+            PileupColumn.NValidCov
+        ]
+    )
+
+    if pileup.height == 0:
+        log.warning(f"No pileup data found for contigs: {contigs}")
+        return pl.DataFrame({
+            "contig": [],
+            "position": [],
+            "mod_type": [],
+            "strand": [],
+            "fraction_mod": [],
+            "Nvalid_cov": [],
+        }, schema = {
+            "contig": pl.String,
+            "position": pl.Int64,
+            "mod_type": pl.Utf8,
+            "strand": pl.Utf8,
+            "fraction_mod": pl.Float64,
+            "Nvalid_cov": pl.Int64,
+        })
+
+    
     log.debug(f"Renaming and removing unnecessary columns")
     pileup = pileup.rename({
         "contig": "contig",
@@ -115,8 +147,7 @@ def load_contigs_pileup_bgzip(path: str, contigs: list[str]):
         "fraction_modified": "fraction_mod",
         "n_valid_cov": "Nvalid_cov",
     }) \
-    .with_columns(pl.col("fraction_mod") / 100) \
-    .select(["contig", "position", "mod_type", "strand", "fraction_mod", "Nvalid_cov"])
+    .with_columns(pl.col("fraction_mod") / 100)
 
     return pileup
 
@@ -213,4 +244,4 @@ def filter_pileup_adjacency_filter(
         .drop("roll_max")
         )
 
-    return pileup
+    return pileup

nanomotif/dataload.py

@@ -141,46 +141,53 @@ def binnary(args, pl):
 
     if args.methylation_output_type == "median":
         output_type = MethylationOutput.Median
-    else:
+    elif args.methylation_output_type == "weighted-mean":
         output_type = MethylationOutput.WeightedMean
+    else:
+        log.error(f"Output type must be either median or weighted mean, got: {args.methylation_output_type}")
+        exit()
 
     contig_methylation_file = f"motifs-scored-read-methylation_{args.methylation_output_type}.tsv"
     if not args.force:
         log.info(f"Check if {contig_methylation_file} exists")
 
     if os.path.isfile(os.path.join(args.out,contig_methylation_file)) and not args.force:
         log.info("motifs-scored-read-methylation.tsv exists. Using existing file! Use --force to override this.")
+        # Load motifs-scored-read-methylation.tsv
+        contig_methylation = pl.read_csv(
+            os.path.join(args.out, contig_methylation_file), separator="\t", has_header = True, schema = {
+                'contig': pl.String(), 'motif': pl.String(), 'mod_type': pl.String(), 'mod_position': pl.Int8(), 'methylation_value': pl.Float64(),  'mean_read_cov': pl.Float64(), 'n_motif_obs': pl.Int32(),
+            }
+        )
 
     elif not os.path.isfile(os.path.join(args.out, contig_methylation_file)) or args.force:
         log.info(f"Running epymetheus to create {contig_methylation_file}")
-        # Create motifs-scored-read-methylation
-        return_code = methylation_pattern(
-            pileup = args.pileup,
-            assembly = args.assembly,
-            motifs = motifs_in_bin_consensus,
-            threads = args.threads,
-            min_valid_read_coverage = args.min_valid_read_coverage,
-            batch_size=1000,
-            min_valid_cov_to_diff_fraction=0.8,
-            output = os.path.join(args.out,contig_methylation_file),
-            allow_assembly_pileup_mismatch=False,
-            output_type = output_type
-        )
-
-        if return_code != 0:
-            log.error("Error running epymetheus")
-
+        try:
+            # Create motifs-scored-read-methylation
+            contig_methylation = methylation_pattern(
+                pileup = args.pileup,
+                assembly = args.assembly,
+                motifs = motifs_in_bin_consensus,
+                threads = args.threads,
+                min_valid_read_coverage = args.min_valid_read_coverage,
+                batch_size=1000,
+                min_valid_cov_to_diff_fraction=0.8,
+                output = os.path.join(args.out,contig_methylation_file),
+                allow_assembly_pileup_mismatch=True,
+                output_type = output_type
+            )
+            if contig_methylation is None or len(contig_methylation) == 0:
+                log.warning("epymetheus returned empty result")
+                sys.exit(1)
+        except Exception as e:
+            log.error(f"Failed to run epymetheus: {e}")
+            sys.exit(1)
+            
 
     log.info("Loading assembly file...")
     assembly = data_processing.read_fasta(args.assembly)
     contig_lengths = data_processing.find_contig_lengths(assembly) 
 
-    # Load motifs-scored-read-methylation.tsv
-    contig_methylation = pl.read_csv(
-        os.path.join(args.out, contig_methylation_file), separator="\t", has_header = True, schema = {
-            'contig': pl.String(), 'motif': pl.String(), 'mod_type': pl.String(), 'mod_position': pl.Int8(), 'methylation_value': pl.Float64(),  'mean_read_cov': pl.Float64(), 'n_motif_obs': pl.Int32(),
-        }
-    )
 
     contig_methylation = contig_methylation\
         .filter((pl.col("n_motif_obs").cast(pl.Float64) * pl.col("mean_read_cov")) >= args.methylation_threshold)

nanomotif/datasets/geobacillus-plasmids.pileup.bed.gz

@@ -35,7 +35,7 @@
         "scikit-learn>=1.5.2",
         "networkx>=3.1",
         "pyarrow>=15.0.2",
-        "epimetheus-py==0.7.4",
+        "epimetheus-py==0.7.5",
         "hdbscan",
         "Bio>=1.6.2",
         "snakemake>=7.32.4",

nanomotif/datasets/geobacillus-plasmids.pileup.bed.gz.tbi

@@ -27,6 +27,25 @@ def test_motif_discovery():
     # Check that the CLI tool executed successfully
     assert result.returncode == 0, "CLI tool did not exit successfully"
 
+def test_motif_discovery_gzip():
+    """
+    """
+    outdir = "tests/cli_test_motif_discovery"
+    
+    cmd = [
+        "nanomotif", "motif_discovery",
+        "-t", "1",
+        "nanomotif/datasets/geobacillus-plasmids.assembly.fasta",
+        "nanomotif/datasets/geobacillus-plasmids.pileup.bed.gz",
+        "-c", "nanomotif/datasets/geobacillus-contig-bin.tsv",
+        "--out", outdir
+    ]
+    
+    result = subprocess.run(cmd)
+    shutil.rmtree(outdir)
+    
+    # Check that the CLI tool executed successfully
+    assert result.returncode == 0, "CLI tool did not exit successfully"
 
 def test_check_installation():
     """
@@ -152,6 +171,104 @@ def test_detect_contamination():
         else:
             print(f"File not found: {outfile}")
 
+
+
+def test_detect_contamination_bgzip():
+    """
+    """
+    import polars as pl
+    from Bio import SeqIO
+    from Bio.Seq import Seq
+    from Bio.SeqRecord import SeqRecord
+    from epymetheus.epymetheus import bgzf_pileup
+
+    infile = "nanomotif/datasets/geobacillus-plasmids.assembly.fasta"
+    outfile_a = "nanomotif/datasets/geobacillus-plasmids.assembly.duplicated.fasta"
+
+    # Read all records from the original file
+    records = list(SeqIO.parse(infile, "fasta"))
+
+    # We'll store our new records here
+    new_records = []
+
+    # For each record, if it's contig_2 or contig_3, create 5 duplicates
+    # with IDs appended by _1.._5. Otherwise, keep it as is.
+    for record in records:
+        if record.id in ["contig_2", "contig_3"]:
+            for i in range(1, 6):
+                new_id = f"{record.id}_{i}"
+                # Create a new SeqRecord with the same sequence
+                new_record = SeqRecord(
+                    record.seq,
+                    id=new_id,
+                    description=""
+                )
+                new_records.append(new_record)
+        else:
+            # For non-contig_2/3, simply keep the original record
+            new_records.append(record)
+
+    # Write out the new FASTA file
+    SeqIO.write(new_records, outfile_a, "fasta")
+
+    print(f"Duplicated contigs written to: {outfile_a}")
+
+
+    infile = "nanomotif/datasets/geobacillus-plasmids.pileup.bed"
+    outfile_p = "nanomotif/datasets/geobacillus-plasmids.pileup.duplicated.bed"
+    
+    p = pl.read_csv(infile, has_header = False, separator = "\t")
+    
+    p_dup = pl.DataFrame()
+    for contig in ["contig_3", "contig_2"]:
+        p_tmp = p.filter(pl.col("column_1") == contig)
+
+        for i in range(1, 6):
+            p_i = p_tmp.with_columns(
+                (pl.col("column_1") + f"_{i}").alias("column_1")
+            )
+
+            p_dup = pl.concat([p_dup, p_i])
+
+    p_dup.write_csv(outfile_p, separator = "\t", include_header = False)
+    bgzf_pileup(outfile_p)
+
+    
+
+    outfile_b = "nanomotif/datasets/geobacillus-plasmids.contig_bin.tmp.tsv"
+    contig_bin = pl.DataFrame({
+                                  "contig": [f"contig_{i}_{j}" for i in [2, 3] for j in range(1, 6)],
+                                  "bin": ["bin_1"] * 10,
+                              })
+
+    contig_bin.write_csv(outfile_b, separator ="\t", include_header = False)
+    
+    outdir = "tests/cli_test_detect_contamination"
+
+    cmd = [
+        "nanomotif", "detect_contamination",
+        "-t", "1",
+        "--force",
+        "--assembly", outfile_a,
+        "--pileup", outfile_p + ".gz",
+        "--contig_bins", outfile_b,
+        "--bin_motifs", "nanomotif/datasets/geobacillus-plasmids.bin-motifs.tsv",
+        "--out", outdir
+    ]
+    result = subprocess.run(cmd)
+    
+    # Check that the CLI tool executed successfully
+    shutil.rmtree(outdir)
+    assert result.returncode == 0, "CLI tool did not exit successfully"
+    for outfile in [outfile_a, outfile_p, outfile_b, outfile_p + ".gz", outfile_p + "..gz.tbi"]:
+        if os.path.exists(outfile):
+            os.remove(outfile)
+            print(f"Deleted: {outfile}")
+        else:
+            print(f"File not found: {outfile}")
+
+
+
 def test_detect_contamination_weighted_mean():
     """
     """