Thomas 1 år sedan
förälder
incheckning
244b12e4dc
2 ändrade filer med 66 tillägg och 50 borttagningar
  1. 55 2
      src/assembler/mod.rs
  2. 11 48
      src/assembler/wtdbg2.rs

+ 55 - 2
src/assembler/mod.rs

@@ -3,10 +3,10 @@ use std::{
 };
 
 use anyhow::{Context, Ok};
-use log::warn;
+use log::{info, warn};
 use thiserror::Error;
 
-use crate::io::bam::{all_bam_paths, create_bam_leave_two_out};
+use crate::io::{bam::{all_bam_paths, cp_mod_tags, create_bam_leave_two_out, remap_bam, run_modkit}, fasta::{fai, write_fasta}};
 
 use self::{
     flye::{Flye, FlyeConfig},
@@ -130,6 +130,59 @@ pub enum AssembleError {
     MaxRescueDepthReached(String),
 }
 
+pub fn default_save(
+    assembler: &str,
+    contigs: Option<Vec<String>>,
+    input_id: String,
+    output_dir: PathBuf,
+    input_fq: String,
+    input_records: Vec<bam::Record>,
+    tmp_dir: String,
+) -> anyhow::Result<()> {
+    if contigs.is_none() {
+        anyhow::bail!(AssembleError::NoContig(input_id));
+    }
+
+    for (i, contig) in contigs.unwrap().iter().enumerate() {
+        let suffixe = if i == 0 {
+            "".to_string()
+        } else {
+            format!("_{i}")
+        };
+
+        if !output_dir.exists() {
+            fs::create_dir_all(&output_dir)?;
+        }
+
+        let contig_id = format!("{input_id}{suffixe}_{assembler}");
+        let contig_fa = format!("{}/{contig_id}.fa", output_dir.display());
+
+        info!("Saving contig {contig_id} in {contig_fa}");
+        write_fasta(&contig_fa, &vec![(contig_id.clone(), contig.clone())])?;
+        fai(&contig_fa)?;
+
+        // Remaping input bam to contig
+        info!("Mapping input reads to {contig_id}");
+        let new_bam = format!("{}/{contig_id}.bam", output_dir.display());
+        remap_bam(&contig_fa, &input_fq, &new_bam)?;
+
+        // clean bwa indices
+        // remove_bwa_indices(&contig_fa)?;
+
+        // Copy modified base tags to new bam
+        cp_mod_tags(&input_records, &new_bam)?;
+
+        // Run modkit
+        let modkit_pileup = format!("{}/{contig_id}_mod.bed", output_dir.display());
+        run_modkit(&new_bam, &contig_fa, &modkit_pileup)?;
+    }
+
+    // Cleaning
+    fs::remove_dir_all(tmp_dir)?;
+
+    Ok(())
+}
+
 pub fn calculate_shannon_entropy(sequence: &str) -> f64 {
     let total_length = sequence.len() as f64;
 

+ 11 - 48
src/assembler/wtdbg2.rs

@@ -11,14 +11,10 @@ use seq_io::fasta::Record;
 
 use crate::{
     assembler::AssembleError,
-    io::{
-        bam::{cp_mod_tags, read_bam, remap_bam, run_modkit},
-        fasta::{fai, write_fasta},
-        fastq::records_to_fastq,
-    },
+    io::{bam::read_bam, fastq::records_to_fastq},
 };
 
-use super::{calculate_shannon_entropy, Assemble, AssembleConfig};
+use super::{calculate_shannon_entropy, default_save, Assemble, AssembleConfig};
 
 #[derive(Debug, Clone)]
 pub struct Wtdbg2Config {
@@ -128,48 +124,15 @@ impl Assemble for Wtdbg2 {
     }
 
     fn save(self) -> anyhow::Result<()> {
-        if self.contigs.is_none() {
-            anyhow::bail!(AssembleError::NoContig(self.input_id));
-        }
-
-        for (i, contig) in self.contigs.unwrap().iter().enumerate() {
-            let suffixe = if i == 0 {
-                "".to_string()
-            } else {
-                format!("_{i}")
-            };
-
-            if !self.config.output_dir.exists() {
-                fs::create_dir_all(&self.config.output_dir)?;
-            }
-
-            let contig_id = format!("{}{suffixe}_wtdbg2", self.input_id);
-            let contig_fa = format!("{}/{contig_id}.fa", self.config.output_dir.display());
-
-            info!("Saving contig {contig_id} in {contig_fa}");
-            write_fasta(&contig_fa, &vec![(contig_id.clone(), contig.clone())])?;
-            fai(&contig_fa)?;
-
-            // Remaping input bam to contig
-            info!("Mapping  input reads to {contig_id}");
-            let new_bam = format!("{}/{contig_id}.bam", self.config.output_dir.display());
-            remap_bam(&contig_fa, &self.input_fq, &new_bam)?;
-
-            // clean bwa indices
-            // remove_bwa_indices(&contig_fa)?;
-
-            // Copy modified base tags to new bam
-            cp_mod_tags(&self.input_records, &new_bam)?;
-
-            // Run modkit
-            let modkit_pileup = format!("{}/{contig_id}_mod.bed", self.config.output_dir.display());
-            run_modkit(&new_bam, &contig_fa, &modkit_pileup)?;
-        }
-
-        // Cleaning
-        fs::remove_dir_all(self.tmp_dir)?;
-
-        Ok(())
+        default_save(
+            "wtdbg2",
+            self.contigs,
+            self.input_id,
+            self.config.output_dir,
+            self.input_fq,
+            self.input_records,
+            self.tmp_dir,
+        )
     }
 }