Set up and use phase layers to represent unphased alleles and pooled design parts

Author: dkhofer

migrations/core/01-initial/up.sql

@@ -129,14 +129,25 @@ CREATE TABLE path_edges (
 ) STRICT;
 CREATE UNIQUE INDEX path_edges_uidx ON path_edges(path_id, edge_id, index_in_path);
 
+CREATE TABLE phase_layers (
+  id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
+  chromosome_index INTEGER NOT NULL,
+  is_reference INTEGER NOT NULL DEFAULT 0
+) STRICT;
+CREATE UNIQUE INDEX phase_layer_uidx ON phase_layers(chromosome_index, is_reference) WHERE is_reference = 1;
+
 CREATE TABLE block_group_edges (
   id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
   block_group_id INTEGER NOT NULL,
   edge_id INTEGER NOT NULL,
   chromosome_index INTEGER,
   phased INTEGER NOT NULL,
+  source_phase_layer_id INTEGER,
+  target_phase_layer_id INTEGER,
   FOREIGN KEY(block_group_id) REFERENCES block_groups(id),
-  FOREIGN KEY(edge_id) REFERENCES edges(id)
+  FOREIGN KEY(edge_id) REFERENCES edges(id),
+  FOREIGN KEY(source_phase_layer_id) REFERENCES phase_layers(id),
+  FOREIGN KEY(target_phase_layer_id) REFERENCES phase_layers(id)
 ) STRICT;
 CREATE UNIQUE INDEX block_group_edges_uidx ON block_group_edges(block_group_id, edge_id, chromosome_index, phased);
 

src/models.rs

@@ -5,10 +5,12 @@ pub mod collection;
 pub mod edge;
 pub mod file_types;
 pub mod metadata;
+pub mod new_block_group_edge;
 pub mod node;
 pub mod operations;
 pub mod path;
 pub mod path_edge;
+pub mod phase_layer;
 pub mod sample;
 pub mod sequence;
 pub mod strand;

src/models/new_block_group_edge.rs

@@ -0,0 +1,142 @@
+use crate::models::edge::{Edge, EdgeData};
+use crate::models::traits::*;
+use rusqlite;
+use rusqlite::types::Value;
+use rusqlite::{Connection, Row};
+use std::collections::HashMap;
+
+#[derive(Clone, Debug)]
+pub struct NewBlockGroupEdge {
+    pub id: i64,
+    pub block_group_id: i64,
+    pub edge_id: i64,
+    pub chromosome_index: i64,
+    pub phased: i64,
+    pub source_phase_layer_id: i64,
+    pub target_phase_layer_id: i64,
+}
+
+#[derive(Clone, Debug, Eq, Hash, PartialEq)]
+pub struct NewBlockGroupEdgeData {
+    pub block_group_id: i64,
+    pub edge_id: i64,
+    pub chromosome_index: i64,
+    pub phased: i64,
+    pub source_phase_layer_id: i64,
+    pub target_phase_layer_id: i64,
+}
+
+#[derive(Clone, Debug, Eq, Hash, PartialEq)]
+pub struct NewAugmentedEdge {
+    pub edge: Edge,
+    pub chromosome_index: i64,
+    pub phased: i64,
+    pub source_phase_layer_id: i64,
+    pub target_phase_layer_id: i64,
+}
+
+#[derive(Clone, Debug, Eq, Hash, PartialEq)]
+pub struct NewAugmentedEdgeData {
+    pub edge_data: EdgeData,
+    pub chromosome_index: i64,
+    pub phased: i64,
+    pub source_phase_layer_id: i64,
+    pub target_phase_layer_id: i64,
+}
+
+impl Query for NewBlockGroupEdge {
+    type Model = NewBlockGroupEdge;
+    fn process_row(row: &Row) -> Self::Model {
+        NewBlockGroupEdge {
+            id: row.get(0).unwrap(),
+            block_group_id: row.get(1).unwrap(),
+            edge_id: row.get(2).unwrap(),
+            chromosome_index: row.get(3).unwrap(),
+            phased: row.get(4).unwrap(),
+            source_phase_layer_id: row.get(5).unwrap(),
+            target_phase_layer_id: row.get(6).unwrap(),
+        }
+    }
+}
+
+impl NewBlockGroupEdge {
+    pub fn bulk_create(conn: &Connection, block_group_edges: &[NewBlockGroupEdgeData]) {
+        for chunk in block_group_edges.chunks(100000) {
+            let mut rows_to_insert = vec![];
+            for block_group_edge in chunk {
+                let row = format!(
+                    "({0}, {1}, {2}, {3}, {4}, {5})",
+                    block_group_edge.block_group_id,
+                    block_group_edge.edge_id,
+                    block_group_edge.chromosome_index,
+                    block_group_edge.phased,
+                    block_group_edge.source_phase_layer_id,
+                    block_group_edge.target_phase_layer_id,
+                );
+                rows_to_insert.push(row);
+            }
+
+            let formatted_rows_to_insert = rows_to_insert.join(", ");
+
+            let insert_statement = format!(
+                "INSERT OR IGNORE INTO block_group_edges (block_group_id, edge_id, chromosome_index, phased, source_phase_layer_id, target_phase_layer_id) VALUES {0};",
+                formatted_rows_to_insert
+            );
+            let _ = conn.execute(&insert_statement, ());
+        }
+    }
+
+    pub fn edges_for_block_group(conn: &Connection, block_group_id: i64) -> Vec<NewAugmentedEdge> {
+        let block_group_edges = NewBlockGroupEdge::query(
+            conn,
+            "select * from block_group_edges where block_group_id = ?1;",
+            rusqlite::params!(Value::from(block_group_id)),
+        );
+        let edge_ids = block_group_edges
+            .clone()
+            .into_iter()
+            .map(|block_group_edge| block_group_edge.edge_id)
+            .collect::<Vec<i64>>();
+        let chromosome_index_by_edge_id = block_group_edges
+            .clone()
+            .into_iter()
+            .map(|block_group_edge| (block_group_edge.edge_id, block_group_edge.chromosome_index))
+            .collect::<HashMap<i64, i64>>();
+        let phased_by_edge_id = block_group_edges
+            .clone()
+            .into_iter()
+            .map(|block_group_edge| (block_group_edge.edge_id, block_group_edge.phased))
+            .collect::<HashMap<i64, i64>>();
+        let source_phase_layer_id_by_edge_id = block_group_edges
+            .clone()
+            .into_iter()
+            .map(|block_group_edge| {
+                (
+                    block_group_edge.edge_id,
+                    block_group_edge.source_phase_layer_id,
+                )
+            })
+            .collect::<HashMap<i64, i64>>();
+        let target_phase_layer_id_by_edge_id = block_group_edges
+            .clone()
+            .into_iter()
+            .map(|block_group_edge| {
+                (
+                    block_group_edge.edge_id,
+                    block_group_edge.target_phase_layer_id,
+                )
+            })
+            .collect::<HashMap<i64, i64>>();
+        let edges = Edge::bulk_load(conn, &edge_ids);
+        edges
+            .into_iter()
+            .map(|edge| NewAugmentedEdge {
+                edge: edge.clone(),
+                chromosome_index: *chromosome_index_by_edge_id.get(&edge.id).unwrap(),
+                phased: *phased_by_edge_id.get(&edge.id).unwrap(),
+                source_phase_layer_id: *source_phase_layer_id_by_edge_id.get(&edge.id).unwrap(),
+                target_phase_layer_id: *target_phase_layer_id_by_edge_id.get(&edge.id).unwrap(),
+            })
+            .collect()
+    }
+}

src/models/phase_layer.rs

@@ -0,0 +1,62 @@
+use rusqlite::Connection;
+
+#[derive(Clone, Debug)]
+pub struct PhaseLayer {
+    pub id: i64,
+    pub chromosome_index: i64,
+    pub is_reference: i64,
+}
+
+pub const UNPHASED_CHROMOSOME_INDEX: i64 = -1;
+
+impl PhaseLayer {
+    pub fn get_or_create(
+        conn: &Connection,
+        chromosome_index: i64,
+        is_reference: i64,
+    ) -> Result<i64, &'static str> {
+        let phase_layer_id: i64 = match conn.query_row(
+            "select id from phase_layers where chromosome_index = ?1 AND is_reference = ?2",
+            (chromosome_index, is_reference),
+            |row| row.get(0),
+        ) {
+            Ok(res) => res,
+            Err(rusqlite::Error::QueryReturnedNoRows) => 0,
+            Err(_e) => {
+                panic!("Error querying the database: {_e}");
+            }
+        };
+        if phase_layer_id != 0 {
+            return Ok(phase_layer_id);
+        }
+
+        let new_phase_layer_id = PhaseLayer::create(conn, chromosome_index, is_reference);
+
+        Ok(new_phase_layer_id)
+    }
+
+    pub fn create(conn: &Connection, chromosome_index: i64, is_reference: i64) -> i64 {
+        let query = "INSERT INTO phase_layers (chromosome_index, is_reference) VALUES (?1, ?2) RETURNING (id)";
+        let mut stmt = conn.prepare(query).unwrap();
+        match stmt.query_row((chromosome_index, is_reference), |row| row.get(0)) {
+            Ok(res) => res,
+            Err(rusqlite::Error::SqliteFailure(err, details)) => {
+                if err.code == rusqlite::ErrorCode::ConstraintViolation {
+                    println!("{err:?} {details:?}");
+                    conn
+                        .query_row(
+                            "select id from phase_layers where chromosome_index = ?1 and is_reference = ?2",
+                            (chromosome_index, is_reference),
+                            |row| row.get(0),
+                        )
+                        .unwrap()
+                } else {
+                    panic!("something bad happened querying the database")
+                }
+            }
+            Err(_) => {
+                panic!("something bad happened querying the database")
+            }
+        }
+    }
+}

src/updates/library.rs

@@ -8,11 +8,12 @@ use std::io::BufReader;
 use std::str;
 
 use crate::models::block_group::BlockGroup;
-use crate::models::block_group_edge::{BlockGroupEdge, BlockGroupEdgeData};
 use crate::models::edge::{Edge, EdgeData};
 use crate::models::file_types::FileTypes;
+use crate::models::new_block_group_edge::{NewBlockGroupEdge, NewBlockGroupEdgeData};
 use crate::models::node::Node;
 use crate::models::operations::OperationInfo;
+use crate::models::phase_layer::{PhaseLayer, UNPHASED_CHROMOSOME_INDEX};
 use crate::models::sample::Sample;
 use crate::models::sequence::Sequence;
 use crate::models::strand::Strand;
@@ -130,6 +131,21 @@ pub fn update_with_library(
     let node_end_coordinate = end_coordinate - end_block.start + end_block.sequence_start;
 
     let mut new_edges = HashSet::new();
+    let mut phase_layers_by_node_id = HashMap::new();
+    let unphased_layer_id1 = PhaseLayer::create(conn, UNPHASED_CHROMOSOME_INDEX, 0);
+    let unphased_layer_id2 = PhaseLayer::create(conn, UNPHASED_CHROMOSOME_INDEX, 0);
+    // TODO: Change path to use block group edges instead of edges,
+    // then use the phase layer of the start and end nodes here instead of the unphased one
+    phase_layers_by_node_id.insert(start_block.node_id, unphased_layer_id1);
+    phase_layers_by_node_id.insert(end_block.node_id, unphased_layer_id2);
+
+    for parts in parts_list.iter() {
+        for part in parts.iter() {
+            let phased_layer_id = PhaseLayer::create(conn, UNPHASED_CHROMOSOME_INDEX, 0);
+            phase_layers_by_node_id.insert(**part, phased_layer_id);
+        }
+    }
+
     let start_parts = parts_list.first().unwrap();
     for start_part in *start_parts {
         let edge = EdgeData {
@@ -179,16 +195,19 @@ pub fn update_with_library(
     path_changes_count *= end_parts.len();
 
     let new_edge_ids = Edge::bulk_create(conn, &new_edges.iter().cloned().collect());
-    let new_block_group_edges = new_edge_ids
+    let new_edges = Edge::bulk_load(conn, &new_edge_ids);
+    let new_block_group_edges = new_edges
         .iter()
-        .map(|edge_id| BlockGroupEdgeData {
+        .map(|edge| NewBlockGroupEdgeData {
             block_group_id: path.block_group_id,
-            edge_id: *edge_id,
+            edge_id: edge.id,
             chromosome_index: 0,
             phased: 0,
+            source_phase_layer_id: *phase_layers_by_node_id.get(&edge.source_node_id).unwrap(),
+            target_phase_layer_id: *phase_layers_by_node_id.get(&edge.target_node_id).unwrap(),
         })
         .collect::<Vec<_>>();
-    BlockGroupEdge::bulk_create(conn, &new_block_group_edges);
+    NewBlockGroupEdge::bulk_create(conn, &new_block_group_edges);
 
     let summary_str = format!("{region_name}: {path_changes_count} changes.\n");
     operation_management::end_operation(