Lines
83.89 %
Functions
20 %
Branches
100 %
use std::collections::HashMap;
use std::path::PathBuf;
use color_eyre::Result;
use rand::{seq::IteratorRandom, thread_rng};
use rayon::prelude::*;
use rust_htslib::bcf::Read;
use crate::args::MergeRule;
use crate::read_vcf::{get_reader, get_samples};
use crate::subcommands::rwc;
use crate::{io::read_sample_ids, utils::filter_samples};
use crate::core::{read_tabix, open_csv_writer};
// Contig, start pos, stop pos, marker count, length in bp
type Shared = Vec<(String, i64, i64, usize, i64)>;
type Location = (String, i64);
type Data<'a> = (i64, usize, i64, Vec<&'a String>);
#[doc(hidden)]
#[tracing::instrument]
#[allow(clippy::too_many_arguments)]
pub fn run(
file: PathBuf,
samples_path: Option<PathBuf>,
mut output: PathBuf,
info_limit: Option<f32>,
sample_size: usize,
niters: i64,
sort: bool,
min_marker_len: usize,
top_markers_n: usize,
merge_rule: MergeRule,
prefix: Option<String>,
) -> Result<()> {
let tabix = read_tabix(&file)?;
let contigs: Vec<String> = tabix.keys().cloned().collect();
if let Some(prefix) = &prefix {
match merge_rule {
MergeRule::Start => output.push(format!("{prefix}_random_sampled_rwc_start.csv")),
MergeRule::Stop => output.push(format!("{prefix}_random_sampled_rwc_stop.csv")),
};
} else {
MergeRule::Start => output.push("random_sampled_rwc_start.csv"),
MergeRule::Stop => output.push("random_sampled_rwc_stop.csv"),
}
let mut writer = open_csv_writer(output)?;
let shared = (0..niters)
.par_bridge()
.map(|_| {
contigs
.par_iter()
.map(|contig| {
// Fetch readers
let mut reader = get_reader(&file, contig, None)?;
let samples = get_samples(reader.header())?;
let wanted = read_sample_ids(&samples_path)?;
let sample_indexes = filter_samples(&samples, wanted);
// Randomly subsample sample_indexes with size n
let mut rng = thread_rng();
let sample_indexes = sample_indexes
.iter()
.copied()
.choose_multiple(&mut rng, sample_size);
let records = reader.records();
// Program logic
let (merged_gt_vec, pos_vec) =
rwc::read_vcf_to_rwc_vec(records, &sample_indexes, info_limit)?;
let vec = find_shared_sequence_lengths(
contig,
merged_gt_vec,
pos_vec,
min_marker_len,
top_markers_n,
)?;
Ok((vec, sample_indexes))
})
.collect::<Result<Vec<(Shared, Vec<usize>)>>>()
.collect::<Result<Vec<Vec<(Shared, Vec<usize>)>>>>()?;
let reader = crate::read_vcf::get_reader(&file, &contigs[0], None)?;
let samples = crate::read_vcf::get_samples(reader.header())?;
let hm = merge_runs(shared, &samples, &merge_rule);
write_matches_to_csv(hm, &mut writer, sort)?;
Ok(())
fn merge_runs<'a>(
shared: Vec<Vec<(Shared, Vec<usize>)>>,
samples: &'a [String],
merge_rule: &MergeRule,
) -> HashMap<Location, Data<'a>> {
let mut hm = HashMap::<Location, Data>::new();
for iteration in shared {
for (shared, indexes) in iteration {
let shared_samples: Vec<&String> = indexes.iter().map(|s| &samples[*s]).collect();
for (contig, start, stop, snp_count, bp_length) in shared {
let merge_value = match merge_rule {
MergeRule::Start => start,
MergeRule::Stop => stop,
let (count, snp_count_avg, bp_length_avg, samples) = hm
.entry((contig, merge_value))
.or_insert((0, snp_count, bp_length, shared_samples.clone()));
*count += 1;
*snp_count_avg = (*snp_count_avg + snp_count) / 2;
*bp_length_avg = (*bp_length_avg + bp_length) / 2;
samples.extend(shared_samples.clone());
samples.sort();
samples.dedup();
hm
fn write_matches_to_csv(
hm: HashMap<Location, Data>,
writer: &mut csv::Writer<Box<dyn std::io::Write>>,
let hm = match sort {
true => {
let mut data: Vec<_> = hm.into_iter().collect();
data.sort_by(|(a, _), (b, _)| b.1.cmp(&a.1));
data
false => hm.into_iter().collect(),
writer.write_record(vec![
"chr",
"pos",
"count",
"snp_count_avg",
"bp_length_avg",
"samples",
])?;
for ((contig, start), (count, snp_count_avg, bp_length_avg, samples)) in hm {
start.to_string(),
count.to_string(),
snp_count_avg.to_string(),
bp_length_avg.to_string(),
samples
.into_iter()
.cloned()
.collect::<Vec<String>>()
.join("&"),
// Every iteration looks for ROH of over 20 markers and takes the top 10 hits
pub fn find_shared_sequence_lengths(
contig: &str,
merged_gt_vec: Vec<u8>,
pos_vec: Vec<i64>,
) -> Result<Shared> {
let mut counter = 1;
let mut shared_sequences = vec![];
merged_gt_vec
.enumerate()
.fold(1, |last, (index, curr)| {
// The last and current element are 0, no contrahomozygotes are present
if last == 0 && *curr == 0 {
counter += 1;
// If the last element is 0 and the current is 1, a shared sequence has ended
} else if last == 0 && *curr == 1 {
if counter >= min_marker_len {
let start = &pos_vec[index - counter];
let stop = &pos_vec[index];
shared_sequences.push((
contig.to_string(),
*start,
*stop,
counter,
stop - start,
));
counter = 1;
*curr
});
counter -= 1;
let start = &pos_vec[merged_gt_vec.len() - 1 - counter];
let stop = &pos_vec[merged_gt_vec.len() - 1];
shared_sequences.sort_by(|a, b| b.3.cmp(&a.3));
Ok(shared_sequences.into_iter().take(top_markers_n).collect())