Lines
91.5 %
Functions
32.35 %
Branches
100 %
// use std::io::Read;
use std::io::Write;
use std::path::PathBuf;
use color_eyre::{
eyre::{eyre, WrapErr},
Result,
};
use ndarray::{Array2, Axis};
use petgraph::dot::Dot;
use petgraph::{graph::NodeIndex, Direction, Graph};
use rayon::prelude::*;
use crate::{args::{GraphArgs, Selection, StandardArgs}, structs::HapVariant};
use crate::graphs::DecayGraph;
use crate::io::{read_variable_data_file, read_sample_ids};
use crate::read_vcf::read_vcf_to_matrix;
use crate::structs::{HaploNode, PhasedMatrix};
use crate::utils::{
push_to_output, select_carrier_haplotypes, select_only_longest_haplotypes,
shared_lengths_by_majority,
use crate::core::{open_csv_writer, parse_snp_coord};
#[doc(hidden)]
#[tracing::instrument]
pub fn run(
args: StandardArgs,
graph_args: GraphArgs,
variable_data: Option<PathBuf>,
variable_names: Option<Vec<String>>,
decoy_samples: Option<PathBuf>,
save_branch_haplotypes: bool,
) -> Result<()> {
if args.selection == Selection::Unphased {
return Err(eyre!("Running with unphased data is not supported."))
}
let (contig, variant_pos) = parse_snp_coord(&args.coords)?;
let mut vcf = read_vcf_to_matrix(&args, contig, variant_pos, None)?;
//These are options
let decoy_samples = read_sample_ids(&decoy_samples)?;
// Set clinical data
if let Some(path) = variable_data {
vcf.set_variable_data(read_variable_data_file(path)?)?;
// CSV header
let mut header = vec![];
if let Some(variables) = variable_names.clone() {
let h: Vec<String> = vec!["side", "pos", "1st_n", "2nd_n"]
.iter()
.map(|s| s.to_string())
.collect();
header.extend(h);
for var in variables {
header.push(format!("{var}_1st"));
header.push(format!("{var}_2nd"));
header.push(format!("{var}_t_test"));
tracing::debug!("{header:?}");
if args.selection == Selection::OnlyAlts || args.selection == Selection::OnlyRefs {
vcf = select_carrier_haplotypes(vcf, variant_pos, &args.coords, &args.selection)
if args.selection == Selection::OnlyLongest {
let shared_lengths = shared_lengths_by_majority(&vcf, vcf.variant_idx());
vcf = select_only_longest_haplotypes(&shared_lengths, vcf);
let vec = vec![LocDirection::Left, LocDirection::Right];
vec.par_iter()
.map(|direction| -> Result<()> {
let g = tree_petgraph(&vcf, direction, &decoy_samples, save_branch_haplotypes)?;
if save_branch_haplotypes {
let mut nodes = g.neighbors_directed(NodeIndex::new(0), Direction::Outgoing);
if let Some(node) = nodes.next() {
let haplonode = g.node_weight(node).unwrap();
let mut common_haplo_csv = args.output.clone();
push_to_output(
&args,
&mut common_haplo_csv,
&format!("recursive_mbah_{direction}_common_haplotype"),
"txt",
);
let mut file = std::fs::File::create(&common_haplo_csv)
.wrap_err(eyre!("Path {common_haplo_csv:?}"))?;
file.write_all(haplonode.haplotype().as_bytes())?;
// Write to dot
let mut decay_dot = args.output.clone();
&mut decay_dot,
&format!("recursive_mbah_{direction}"),
"dot",
let mut f = std::fs::File::create(decay_dot)?;
f.write_all(format!("{}", Dot::new(&g)).as_bytes())?;
// Write to json
let mut decay_json = args.output.clone();
&mut decay_json,
"json",
let f = std::fs::File::create(&decay_json).wrap_err(eyre!("Path {decay_json:?}"))?;
serde_json::to_writer(f, &g)?;
// Debug
// let mut file = std::fs::File::open(format!("{out}_{direction}.json"))?;
// let mut data = String::new();
// file.read_to_string(&mut data)?;
// let g: Graph<HaploNode, i64> = serde_json::from_str(&data)?;
// end debug
// Create decay graph svg
let mut dg = DecayGraph::new(&g, &vcf, graph_args.clone());
if let Some(variables) = &variable_names {
dg.variables = Some(variables.clone());
write_clinical_data(
&vcf,
&dg,
variables,
direction,
args.output.clone(),
&header,
)?;
dg.draw_graph()?;
let mut decay_graph = args.output.clone();
&mut decay_graph,
"svg",
svg::save(decay_graph, &dg.document)?;
Ok(())
})
.collect::<Result<()>>()?;
fn write_clinical_data(
args: &StandardArgs,
vcf: &PhasedMatrix,
dg: &DecayGraph,
variables: &[String],
direction: &LocDirection,
mut output: PathBuf,
header: &Vec<String>,
args,
&mut output,
&format!("recursive_mbah_{direction}_clinical_data"),
"csv",
let mut writer = open_csv_writer(output)?;
writer.write_record(header)?;
let nodes = dg.get_majority_nodes().unwrap();
for node in nodes {
let children: Vec<NodeIndex> = dg.g.neighbors_directed(node, Direction::Outgoing).collect();
if children.len() == 2 {
let c1 = dg.g.node_weight(children[0]).unwrap();
let c2 = dg.g.node_weight(children[1]).unwrap();
if c1.indexes_len > 5 && c2.indexes_len > 5 {
let v1 = vcf.get_variable_data_vecs(&c1.indexes, variables)?.unwrap();
let v2 = vcf.get_variable_data_vecs(&c2.indexes, variables)?.unwrap();
let node1 = vcf.get_variable_data_mean(&c1.indexes, variables)?.unwrap();
let node2 = vcf.get_variable_data_mean(&c2.indexes, variables)?.unwrap();
let mut csv_row = vec![
direction.to_string(),
c1.pos.to_string(),
c1.indexes_len.to_string(),
c2.indexes_len.to_string(),
];
for i in 0..node1.len() {
csv_row.push(format!("{:.2}", node1[i]));
csv_row.push(format!("{:.2}", node2[i]));
let t_test = crate::utils::two_tail_welch_t_test(&v1[i], &v2[i]);
csv_row.push(format!("{t_test:.5}"));
writer.write_record(csv_row)?;
#[derive(PartialEq, Eq, PartialOrd, Ord)]
pub enum LocDirection {
Left,
Right,
impl std::fmt::Display for LocDirection {
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
match self {
LocDirection::Right => write!(f, "right"),
LocDirection::Left => write!(f, "left"),
pub fn tree_petgraph(
decoy_samples: &Option<Vec<String>>,
) -> Result<Graph<HaploNode, i64>> {
let indexes: Vec<usize> = (0..vcf.nsamples()).collect();
let mut g = Graph::<HaploNode, i64>::new();
g.add_node(HaploNode {
samples: vcf.get_sample_names(&indexes),
indexes_len: indexes.len(),
pos: 0,
haplotype: vcf.find_haplotype_for_sample("", vcf.variant_idx..vcf.variant_idx, 0),
decoys: match decoy_samples {
Some(decoys) => find_decoys(vcf, &indexes, decoys),
None => None,
},
indexes,
});
match direction {
LocDirection::Right => {
for pos in vcf.variant_idx..vcf.matrix.ncols() {
let node_idxs = g.node_indices();
for idx in node_idxs {
let node = g.node_weight(idx).unwrap();
let count = g.neighbors_directed(idx, Direction::Outgoing).count();
if count == 0 && node.indexes.len() > 1 {
g = check_if_break(vcf, idx, pos, g, direction, decoy_samples, save_branch_haplotypes)?;
LocDirection::Left => {
for pos in (0..vcf.variant_idx).rev() {
Ok(g)
fn check_if_break(
idx: NodeIndex,
pos: usize,
mut g: Graph<HaploNode, i64>,
let submatrix = vcf.matrix.select(Axis(0), &node.indexes);
if let Some((node1, node2)) = find_breaks(submatrix, pos, node, vcf, direction, decoy_samples, save_branch_haplotypes)?
{
let node_idx1 = g.add_node(node1);
let node_idx2 = g.add_node(node2);
let distance = vcf.variant_idx_pos() - vcf.get_pos(pos);
g.add_edge(idx, node_idx1, distance.abs());
g.add_edge(idx, node_idx2, distance.abs());
fn find_breaks(
matrix: Array2<u8>,
node: &HaploNode,
) -> Result<Option<(HaploNode, HaploNode)>> {
let zeroes: Vec<usize> = matrix
.slice(ndarray::s![.., pos])
.zip(node.indexes.clone())
.filter_map(|(n, i)| if *n == 0 { Some(i) } else { None })
let ones: Vec<usize> = matrix
.filter_map(|(n, i)| if *n == 1 { Some(i) } else { None })
if !zeroes.is_empty() && !ones.is_empty() {
let node1 = HaploNode {
samples: vcf.get_sample_names(&zeroes),
pos: distance.unsigned_abs() as usize,
haplotype: match save_branch_haplotypes {
true => find_haplotype_for_haplonode(vcf, zeroes[0], pos, direction),
false => vec![]
indexes_len: zeroes.len(),
Some(decoys) => find_decoys(vcf, &zeroes, decoys),
indexes: zeroes,
let node2 = HaploNode {
samples: vcf.get_sample_names(&ones),
true => find_haplotype_for_haplonode(vcf, ones[0], pos, direction),
indexes_len: ones.len(),
Some(decoys) => find_decoys(vcf, &ones, decoys),
indexes: ones,
return Ok(Some((node1, node2)));
Ok(None)
fn find_haplotype_for_haplonode(
vcf: &PhasedMatrix, idx: usize, pos: usize, direction: &LocDirection
) -> Vec<HapVariant> {
vcf.find_haplotype_for_sample(
vcf.get_contig(),
match &direction {
LocDirection::Right => vcf.variant_idx..pos,
LocDirection::Left => 1 + pos..vcf.variant_idx,
idx)
fn find_decoys(
indexes: &[usize],
decoy_samples: &[String],
) -> Option<Vec<String>> {
let names = vcf.get_sample_names(indexes);
let count = names.iter().filter(|n| decoy_samples.contains(n)).count();
if count == 0 {
None
} else {
Some(names)