Lines
80.45 %
Functions
12.5 %
Branches
100 %
use std::path::PathBuf;
use color_eyre::{eyre::eyre, Result};
use itertools::Itertools;
use plotters::prelude::*;
use rayon::prelude::*;
use crate::args::{GraphArgs, Selection, StandardArgs};
use crate::io::read_recombination_file;
use crate::read_vcf::read_vcf_to_matrix;
use crate::structs::PhasedMatrix;
use crate::subcommands::mrca::mrca_only_independent;
use crate::utils::{push_to_output, select_only_longest_alleles, shared_lengths_by_majority};
use crate::core::{open_csv_writer, centromeres_hg38, parse_snp_only_contig};
#[doc(hidden)]
#[allow(clippy::too_many_arguments)]
pub fn run(
args: StandardArgs,
rec_rates: PathBuf,
start: Option<i64>,
stop: Option<i64>,
step_size: usize,
no_csv: bool,
yes_plot: bool,
graph_args: GraphArgs,
mark_centromere: bool,
) -> Result<()> {
match args.selection {
Selection::OnlyAlts | Selection::OnlyRefs | Selection::Unphased =>
return Err(eyre!("Running with selection {:?} is not supported.", args.selection)),
_ => ()
};
if args.selection == Selection::OnlyAlts {
panic!("Selection with only alts is not supported for mrca scan")
}
let coords = args.coords.clone();
let contig = parse_snp_only_contig(&coords)?;
let rates = read_recombination_file(rec_rates)?;
let vcf = read_vcf_to_matrix(&args, contig, 0, None)?;
let start = match start {
Some(start) => vcf.get_nearest_idx_by_pos(start),
None => 0,
let stop = match stop {
Some(stop) => vcf.get_nearest_idx_by_pos(stop),
None => vcf.ncoords(),
let ages = find_ages(&vcf, &args, rates, start, stop, step_size)?;
if !no_csv {
write_ages_to_csv(&args, &ages, args.output.clone(), &vcf)?;
if yes_plot {
draw_plot(
&vcf,
&args,
graph_args,
args.output.clone(),
ages,
contig,
start,
stop,
mark_centromere,
)?;
Ok(())
fn write_ages_to_csv(
args: &StandardArgs,
ages: &Vec<(usize, f64)>,
mut output: PathBuf,
vcf: &PhasedMatrix,
push_to_output(args, &mut output, "mrca_scan", "csv");
let mut writer = open_csv_writer(output)?;
writer.write_record(vec!["pos", "mrca"])?;
for (idx, age) in ages {
writer.write_record(vec![format!("{}", vcf.get_pos(*idx)), format!("{age:.5}")])?;
fn find_ages(
rates: std::collections::BTreeMap<u64, f32>,
start: usize,
stop: usize,
) -> Result<Vec<(usize, f64)>> {
// Create a Vec because par_iter cannot be used with pure ranges and par_bridge does not
// return in ordered fashion with .collect()
let range: Vec<usize> = (start..stop).collect();
match args.selection == Selection::OnlyLongest {
true => range
.par_iter()
.filter(|n| *n % step_size == 0)
.map(|i| -> Result<(usize, f64)> {
let shared_lengths = shared_lengths_by_majority(vcf, *i);
let only_longest_lengths = select_only_longest_alleles(&shared_lengths);
let (i_tau_hat, _, _) =
mrca_only_independent(only_longest_lengths, vcf.get_pos(*i), &rates)?;
Ok((*i, i_tau_hat.log10()))
})
.collect(),
false => range
mrca_only_independent(shared_lengths, vcf.get_pos(*i), &rates)?;
fn draw_plot(
data: Vec<(usize, f64)>,
contig: &str,
push_to_output(args, &mut output, "mrca_scan", "png");
let root =
BitMapBackend::new(&output, (graph_args.width, graph_args.height)).into_drawing_area();
let y_start = data.iter().min_by(|a, b| a.1.total_cmp(&b.1)).unwrap().1;
let y_stop = data.iter().max_by(|a, b| a.1.total_cmp(&b.1)).unwrap().1;
root.fill(&WHITE)?;
let mut chart = ChartBuilder::on(&root)
.margin(10)
.set_label_area_size(LabelAreaPosition::Left, 60)
.set_label_area_size(LabelAreaPosition::Right, 60)
.set_label_area_size(LabelAreaPosition::Bottom, 40)
.build_cartesian_2d(start..stop, y_start - 0.2..y_stop + 0.2)?;
chart
.configure_mesh()
.disable_x_mesh()
.disable_y_mesh()
.x_labels(30)
.max_light_lines(4)
.y_desc("MRCA")
.x_desc("POS")
.draw()?;
// Raw log10 data line
// chart.draw_series(LineSeries::new(
// data.iter().map(|(pos, mrca)|(*pos,*mrca)),
// &BLUE,
// ))?;
// Last 10 elements moving average
let mut moving_average = vec![];
for (a, b, c, d, e, f, g, h, i, j) in data.iter().tuples() {
let sum = a.1 + b.1 + c.1 + d.1 + e.1 + f.1 + g.1 + h.1 + i.1 + j.1;
let mean = sum / 10.0;
moving_average.push((j.0, mean));
chart.draw_series(LineSeries::new(
moving_average.iter().map(|(pos, mrca)| (*pos, *mrca)),
ShapeStyle {
color: BLACK.mix(0.9),
filled: true,
stroke_width: graph_args.stroke_width,
},
))?;
// Centromere location
if mark_centromere {
let (centro_start, centro_stop) = centromeres_hg38(contig);
let centro_start = vcf.get_nearest_idx_by_pos(centro_start as i64);
let centro_stop = vcf.get_nearest_idx_by_pos(centro_stop as i64);
[(centro_start, y_start), (centro_start, y_stop)].into_iter(),
color: GREEN.mix(0.6),
stroke_width: 2,
[(centro_stop, y_start), (centro_stop, y_stop)].into_iter(),
// Variant location
if graph_args.mark_locus {
[(vcf.variant_idx(), y_start), (vcf.variant_idx(), y_stop)].into_iter(),
color: BLUE.mix(0.6),
// Mean line
let sum: f64 = data.iter().map(|(_pos, mrca)| mrca).sum();
let mean = sum / data.len() as f64;
[(start, mean), (stop, mean)].into_iter(),
color: RED.mix(0.6),
stroke_width: graph_args.stroke_width / 2,
root.present().expect("Unable to write result to file");