GFA export

.gitignore

@@ -20,3 +20,6 @@ fixtures/*.ipynb
 
 # Mac droppings
 .DS_Store
+
+# Gen tracking dir
+.gen/

Cargo.toml

@@ -12,6 +12,7 @@ itertools = "0.13.0"
 rusqlite = { version = "0.31.0", features = ["bundled", "array", "session"] }
 rusqlite_migration = { version = "1.2.0" , features = ["from-directory"]}
 sha2 = "0.10.8"
+tempfile = "3.12.0"
 noodles = { version = "0.78.0", features = ["core", "vcf", "fasta", "async"] }
 petgraph = "0.6.5"
 chrono = "0.4.38"

fixtures/no_path.gfa

@@ -0,0 +1,7 @@
+S	0	GGGG	*
+S	1	AAAA	*
+S	2	TTTT	*
+S	3	CCCC	*
+L	1	+	2	+	*
+L	2	+	0	+	*
+L	0	+	3	+	*

src/exports.rs

@@ -0,0 +1 @@
+pub mod gfa;

src/exports/gfa.rs

@@ -0,0 +1,212 @@
+use itertools::Itertools;
+use rusqlite::Connection;
+use std::collections::{HashMap, HashSet};
+use std::fs;
+use std::fs::File;
+use std::io::{BufWriter, Write};
+use std::path::PathBuf;
+
+use crate::models::{
+    self,
+    block_group::BlockGroup,
+    block_group_edge::BlockGroupEdge,
+    collection::Collection,
+    edge::{Edge, GroupBlock},
+    path::Path,
+    sequence::Sequence,
+    strand::Strand,
+};
+
+pub fn export_gfa(conn: &Connection, collection_name: &str, filename: &PathBuf) {
+    let block_groups = Collection::get_block_groups(conn, collection_name);
+
+    let mut edge_set = HashSet::new();
+    for block_group in block_groups {
+        let block_group_edges = BlockGroupEdge::edges_for_block_group(conn, block_group.id);
+        edge_set.extend(block_group_edges.into_iter());
+    }
+
+    let mut edges = edge_set.into_iter().collect();
+    let (blocks, boundary_edges) = Edge::blocks_from_edges(conn, &edges);
+    edges.extend(boundary_edges.clone());
+
+    let (graph, edges_by_node_pair) = Edge::build_graph(conn, &edges, &blocks);
+
+    let mut file = File::create(filename).unwrap();
+    let mut writer = BufWriter::new(file);
+
+    let mut terminal_block_ids = HashSet::new();
+    for block in &blocks {
+        if block.sequence_hash == Sequence::PATH_START_HASH
+            || block.sequence_hash == Sequence::PATH_END_HASH
+        {
+            terminal_block_ids.insert(block.id);
+            continue;
+        }
+        writer
+            .write_all(&segment_line(&block.sequence, block.id as usize).into_bytes())
+            .unwrap_or_else(|_| {
+                panic!(
+                    "Error writing segment with sequence {} to GFA stream",
+                    block.sequence,
+                )
+            });
+    }
+
+    let blocks_by_id = blocks
+        .clone()
+        .into_iter()
+        .map(|block| (block.id, block))
+        .collect::<HashMap<i32, GroupBlock>>();
+
+    for (source, target, ()) in graph.all_edges() {
+        if terminal_block_ids.contains(&source) || terminal_block_ids.contains(&target) {
+            continue;
+        }
+        let edge = edges_by_node_pair.get(&(source, target)).unwrap();
+        writer
+            .write_all(
+                &link_line(source, edge.source_strand, target, edge.target_strand).into_bytes(),
+            )
+            .unwrap_or_else(|_| {
+                panic!(
+                    "Error writing link from segment {} to {} to GFA stream",
+                    source, target,
+                )
+            });
+    }
+}
+
+fn segment_line(sequence: &str, index: usize) -> String {
+    format!("S\t{}\t{}\t{}\n", index, sequence, "*")
+}
+
+fn link_line(
+    source_index: i32,
+    source_strand: Strand,
+    target_index: i32,
+    target_strand: Strand,
+) -> String {
+    format!(
+        "L\t{}\t{}\t{}\t{}\t*\n",
+        source_index, source_strand, target_index, target_strand
+    )
+}
+
+mod tests {
+    use rusqlite::Connection;
+    // Note this useful idiom: importing names from outer (for mod tests) scope.
+    use super::*;
+
+    use crate::imports::gfa::import_gfa;
+    use crate::models::{
+        block_group::BlockGroup, block_group_edge::BlockGroupEdge, collection::Collection,
+        edge::Edge, sequence::Sequence,
+    };
+    use crate::test_helpers::get_connection;
+    use tempfile::tempdir;
+
+    #[test]
+    fn test_simple_export() {
+        // Sets up a basic graph and then exports it to a GFA file
+        let conn = get_connection(None);
+
+        let collection_name = "test collection";
+        let collection = Collection::create(&conn, collection_name);
+        let block_group = BlockGroup::create(&conn, collection_name, None, "test block group");
+        let sequence1 = Sequence::new()
+            .sequence_type("DNA")
+            .sequence("AAAA")
+            .save(&conn);
+        let sequence2 = Sequence::new()
+            .sequence_type("DNA")
+            .sequence("TTTT")
+            .save(&conn);
+        let sequence3 = Sequence::new()
+            .sequence_type("DNA")
+            .sequence("GGGG")
+            .save(&conn);
+        let sequence4 = Sequence::new()
+            .sequence_type("DNA")
+            .sequence("CCCC")
+            .save(&conn);
+
+        let edge1 = Edge::create(
+            &conn,
+            Sequence::PATH_START_HASH.to_string(),
+            0,
+            Strand::Forward,
+            sequence1.hash.clone(),
+            0,
+            Strand::Forward,
+            0,
+            0,
+        );
+        let edge2 = Edge::create(
+            &conn,
+            sequence1.hash,
+            4,
+            Strand::Forward,
+            sequence2.hash.clone(),
+            0,
+            Strand::Forward,
+            0,
+            0,
+        );
+        let edge3 = Edge::create(
+            &conn,
+            sequence2.hash,
+            4,
+            Strand::Forward,
+            sequence3.hash.clone(),
+            0,
+            Strand::Forward,
+            0,
+            0,
+        );
+        let edge4 = Edge::create(
+            &conn,
+            sequence3.hash,
+            4,
+            Strand::Forward,
+            sequence4.hash.clone(),
+            0,
+            Strand::Forward,
+            0,
+            0,
+        );
+        let edge5 = Edge::create(
+            &conn,
+            sequence4.hash,
+            4,
+            Strand::Forward,
+            Sequence::PATH_END_HASH.to_string(),
+            0,
+            Strand::Forward,
+            0,
+            0,
+        );
+
+        BlockGroupEdge::bulk_create(
+            &conn,
+            block_group.id,
+            &[edge1.id, edge2.id, edge3.id, edge4.id, edge5.id],
+        );
+        let all_sequences = BlockGroup::get_all_sequences(&conn, block_group.id);
+
+        let temp_dir = tempdir().expect("Couldn't get handle to temp directory");
+        let mut gfa_path = PathBuf::from(temp_dir.path());
+        gfa_path.push("intermediate.gfa");
+
+        export_gfa(&conn, collection_name, &gfa_path);
+        // NOTE: Not directly checking file contents because segments are written in random order
+        import_gfa(&gfa_path, "test collection 2", &conn);
+
+        let block_group2 = Collection::get_block_groups(&conn, "test collection 2")
+            .pop()
+            .unwrap();
+        let all_sequences2 = BlockGroup::get_all_sequences(&conn, block_group2.id);
+
+        assert_eq!(all_sequences, all_sequences2);
+    }
+}