Merge pull request #98 from pangenome/impg-partition-fewer-files

`impg partition`: emit a single BED file by default

Author: AndreaGuarracino

README.md

@@ -103,9 +103,12 @@ impg query -p alignments.paf -r chr1:1000-2000 --transitive-dfs
 Partition the alignment into smaller pieces:
 
 ```bash
-# Basic partitioning with 1Mb windows
+# Basic partitioning with 1Mb windows (outputs single partitions.bed file with partition number in 4th column)
 impg partition -p alignments.paf -w 1000000
 
+# Output separate files for each partition
+impg partition -p alignments.paf -w 1000000 --separate-files
+
 # Specify output folder for partition files (directory will be created if it doesn't exist)
 impg partition -p alignments.paf -w 1000000 --output-folder results
 
@@ -124,13 +127,13 @@ impg partition -p alignments.paf -w 1000000 --selection-mode haplotype      # by
 # Control transitive search depth and minimum sizes
 impg partition -p alignments.paf -w 1000000 -m 2 -l 10000
 
-# Output as GFA, MAF or FASTA requires sequence files (--sequence-files or --sequence-list)
-impg partition -p alignments.paf -w 1000000 -o gfa --sequence-files *.fa --output-folder gfa_partitions
-impg partition -p alignments.paf -w 1000000 -o maf --sequence-list fastas.txt --output-folder maf_partitions
-impg partition -p alignments.paf -w 1000000 -o fasta --sequence-files *.fa --output-folder fasta_partitions
+# Output as GFA, MAF or FASTA requires sequence files and --separate-files flag
+impg partition -p alignments.paf -w 1000000 -o gfa --sequence-files *.fa --separate-files --output-folder gfa_partitions
+impg partition -p alignments.paf -w 1000000 -o maf --sequence-list fastas.txt --separate-files --output-folder maf_partitions
+impg partition -p alignments.paf -w 1000000 -o fasta --sequence-files *.fa --separate-files --output-folder fasta_partitions
 
 # Works with AGC archives too
-impg partition -p alignments.paf -w 1000000 -o gfa --sequence-files genomes.agc --output-folder gfa_partitions
+impg partition -p alignments.paf -w 1000000 -o gfa --sequence-files genomes.agc --separate-files --output-folder gfa_partitions
 ```
 
 ### Similarity

src/main.rs

@@ -416,6 +416,10 @@ enum Args {
         /// Minimum distance from sequence start/end - closer regions will be extended to the boundaries
         #[clap(long, value_parser, default_value_t = 3000)]
         min_boundary_distance: i32,
+
+        /// Output separate files for each partition when 'bed'
+        #[clap(long, action)]
+        separate_files: bool,
     },
     /// Query overlaps in the alignment
     Query {
@@ -522,6 +526,7 @@ fn main() -> io::Result<()> {
             selection_mode,
             min_missing_size,
             min_boundary_distance,
+            separate_files,
         } => {
             validate_selection_mode(&selection_mode)?;
             validate_output_format(&output_format, &["bed", "gfa", "maf", "fasta"])?;
@@ -534,6 +539,17 @@ fn main() -> io::Result<()> {
                 gfa_maf_fasta.force_large_region,
             )?;
 
+            // Validate single-file output compatibility
+            if !separate_files && output_format != "bed" {
+                return Err(io::Error::new(
+                    io::ErrorKind::InvalidInput,
+                    format!(
+                        "Single-file output is only supported for BED format. Use --separate-files for {} format.",
+                        output_format.to_uppercase()
+                    ),
+                ));
+            }
+
             // Extract reverse_complement before moving gfa_maf_fasta
             let reverse_complement = gfa_maf_fasta.reverse_complement;
 
@@ -562,6 +578,7 @@ fn main() -> io::Result<()> {
                 scoring_params,
                 reverse_complement,
                 common.verbose > 1,
+                separate_files,
             )?;
         }
         Args::Query {

src/partition.rs

@@ -45,6 +45,7 @@ pub fn partition_alignments(
     scoring_params: Option<(u8, u8, u8, u8, u8, u8)>,
     reverse_complement: bool,
     debug: bool,
+    separate_files: bool,
 ) -> io::Result<()> {
     // Initialize windows from starting sequences if provided
     let mut windows = Vec::<(u32, i32, i32)>::new();
@@ -158,6 +159,9 @@ pub fn partition_alignments(
 
     // Track temporary BED files for GFA/MAF conversion
     let mut temp_bed_files = Vec::new();
+    
+    // Collect partitions for single-file output
+    let mut collected_partitions: Vec<(usize, Vec<Interval<u32>>)> = Vec::new();
 
     while !windows.is_empty() {
         if debug {
@@ -329,45 +333,68 @@ pub fn partition_alignments(
                     .map(|(query_interval, _, _)| query_interval)
                     .collect();
 
-                // Write partition
-                match output_format {
-                    "bed" => {
-                        // Write BED file directly
-                        write_partition_bed(
-                            partition_num,
-                            &query_intervals,
-                            impg,
-                            output_folder,
-                            None,
-                        )?;
-                    }
-                    "gfa" | "maf" => {
-                        // Write temporary BED file with .tmp suffix
-                        write_partition_bed(
-                            partition_num,
-                            &query_intervals,
-                            impg,
-                            output_folder,
-                            Some(".tmp"),
-                        )?;
-                        temp_bed_files.push(partition_num);
-                    }
-                    "fasta" => {
-                        // Write FASTA file directly
-                        write_partition_fasta(
-                            partition_num,
-                            &query_intervals,
-                            impg,
-                            output_folder,
-                            sequence_index.expect("Sequence index not found"),
-                            reverse_complement,
-                        )?;
+                // Write partition or collect for single-file output
+                if separate_files {
+                    // Legacy behavior: write separate files
+                    match output_format {
+                        "bed" => {
+                            // Write BED file directly
+                            write_partition_bed(
+                                partition_num,
+                                &query_intervals,
+                                impg,
+                                output_folder,
+                                None,
+                            )?;
+                        }
+                        "gfa" | "maf" => {
+                            // Write temporary BED file with .tmp suffix
+                            write_partition_bed(
+                                partition_num,
+                                &query_intervals,
+                                impg,
+                                output_folder,
+                                Some(".tmp"),
+                            )?;
+                            temp_bed_files.push(partition_num);
+                        }
+                        "fasta" => {
+                            // Write FASTA file directly
+                            write_partition_fasta(
+                                partition_num,
+                                &query_intervals,
+                                impg,
+                                output_folder,
+                                sequence_index.expect("Sequence index not found"),
+                                reverse_complement,
+                            )?;
+                        }
+                        _ => {
+                            return Err(io::Error::new(
+                                io::ErrorKind::InvalidInput,
+                                format!("Unsupported output format: {}", output_format),
+                            ));
+                        }
                     }
-                    _ => {
-                        return Err(io::Error::new(
-                            io::ErrorKind::InvalidInput,
-                            format!("Unsupported output format: {}", output_format),
-                        ));
+                } else {
+                    // New behavior: collect partitions for single-file output
+                    match output_format {
+                        "bed" => {
+                            // Collect BED partitions
+                            collected_partitions.push((partition_num, query_intervals.clone()));
+                        }
+                        "gfa" | "maf" | "fasta" => {
+                            return Err(io::Error::new(
+                                io::ErrorKind::InvalidInput,
+                                "Single-file output is only supported for BED format. Use --separate-files for GFA, MAF, or FASTA formats.".to_string(),
+                            ));
+                        }
+                        _ => {
+                            return Err(io::Error::new(
+                                io::ErrorKind::InvalidInput,
+                                format!("Unsupported output format: {}", output_format),
+                            ));
+                        }
                     }
                 }
 
@@ -462,6 +489,21 @@ pub fn partition_alignments(
             })?;
     }
 
+    // Write collected partitions as single file if not using separate files
+    if !separate_files && !collected_partitions.is_empty() {
+        info!(
+            "Writing {} partitions to single {} file",
+            collected_partitions.len(),
+            output_format
+        );
+        write_single_partition_file(
+            &collected_partitions,
+            impg,
+            output_format,
+            output_folder,
+        )?;
+    }
+
     // Calculate final percentage
     let final_percentage = (total_partitioned_length as f64 / total_sequence_length as f64) * 100.0;
     // Create formatted percentage string with conditional scientific notation
@@ -1344,6 +1386,45 @@ fn write_partition_fasta(
     Ok(())
 }
 
+fn write_single_partition_file(
+    collected_partitions: &[(usize, Vec<Interval<u32>>)],
+    impg: &Impg,
+    output_format: &str,
+    output_folder: Option<&str>,
+) -> io::Result<()> {
+    match output_format {
+        "bed" => {
+            // Create single BED file with partition column
+            let filename = "partitions.bed";
+            let full_path = create_output_path(output_folder, filename)?;
+            let file = File::create(full_path)?;
+            let mut writer = BufWriter::new(file);
+
+            for (partition_num, query_intervals) in collected_partitions {
+                for query_interval in query_intervals {
+                    let name = impg.seq_index.get_name(query_interval.metadata).unwrap();
+                    let (start, end) = if query_interval.first <= query_interval.last {
+                        (query_interval.first, query_interval.last)
+                    } else {
+                        (query_interval.last, query_interval.first)
+                    };
+
+                    writeln!(writer, "{}	{}	{}	{}", name, start, end, partition_num)?;
+                }
+            }
+
+            writer.flush()?;
+            Ok(())
+        }
+        _ => {
+            Err(io::Error::new(
+                io::ErrorKind::InvalidInput,
+                format!("Single-file output not supported for format: {}", output_format),
+            ))
+        }
+    }
+}
+
 pub fn parse_bed_file(bed_file: &str) -> io::Result<Vec<(String, (i32, i32), String)>> {
     let file = File::open(bed_file)?;
     let reader = BufReader::new(file);