1
use std::path::PathBuf;
2

            
3
use color_eyre::{eyre::eyre, Result};
4
use itertools::Itertools;
5
use plotters::prelude::*;
6
use rayon::prelude::*;
7

            
8
use crate::args::{GraphArgs, Selection, StandardArgs};
9
use crate::io::read_recombination_file;
10
use crate::read_vcf::read_vcf_to_matrix;
11
use crate::structs::PhasedMatrix;
12
use crate::subcommands::mrca::mrca_only_independent;
13
use crate::utils::{push_to_output, select_only_longest_alleles, shared_lengths_by_majority};
14
use crate::core::{open_csv_writer, centromeres_hg38, parse_snp_only_contig};
15

            
16
#[doc(hidden)]
17
#[allow(clippy::too_many_arguments)]
18
12
pub fn run(
19
12
    args: StandardArgs,
20
12
    rec_rates: PathBuf,
21
12
    start: Option<i64>,
22
12
    stop: Option<i64>,
23
12
    step_size: usize,
24
12
    no_csv: bool,
25
12
    yes_plot: bool,
26
12
    graph_args: GraphArgs,
27
12
    mark_centromere: bool,
28
12
) -> Result<()> {
29
12
    match args.selection {
30
        Selection::OnlyAlts | Selection::OnlyRefs | Selection::Unphased => 
31
        return Err(eyre!("Running with selection {:?} is not supported.", args.selection)),
32
12
        _ => ()
33
12
    };
34
12
    if args.selection == Selection::OnlyAlts {
35
        panic!("Selection with only alts is not supported for mrca scan")
36
12
    }
37
12

            
38
12
    let coords = args.coords.clone();
39

            
40
12
    let contig = parse_snp_only_contig(&coords)?;
41

            
42
12
    let rates = read_recombination_file(rec_rates)?;
43

            
44
12
    let vcf = read_vcf_to_matrix(&args, contig, 0, None)?;
45

            
46
12
    let start = match start {
47
        Some(start) => vcf.get_nearest_idx_by_pos(start),
48
12
        None => 0,
49
    };
50

            
51
12
    let stop = match stop {
52
        Some(stop) => vcf.get_nearest_idx_by_pos(stop),
53
12
        None => vcf.ncoords(),
54
    };
55

            
56
12
    let ages = find_ages(&vcf, &args, rates, start, stop, step_size)?;
57

            
58
12
    if !no_csv {
59
12
        write_ages_to_csv(&args, &ages, args.output.clone(), &vcf)?;
60
    }
61

            
62
12
    if yes_plot {
63
12
        draw_plot(
64
12
            &vcf,
65
12
            &args,
66
12
            graph_args,
67
12
            args.output.clone(),
68
12
            ages,
69
12
            contig,
70
12
            start,
71
12
            stop,
72
12
            mark_centromere,
73
12
        )?;
74
    }
75

            
76
12
    Ok(())
77
12
}
78

            
79
12
fn write_ages_to_csv(
80
12
    args: &StandardArgs,
81
12
    ages: &Vec<(usize, f64)>,
82
12
    mut output: PathBuf,
83
12
    vcf: &PhasedMatrix,
84
12
) -> Result<()> {
85
12
    push_to_output(args, &mut output, "mrca_scan", "csv");
86

            
87
12
    let mut writer = open_csv_writer(output)?;
88
12
    writer.write_record(vec!["pos", "mrca"])?;
89
168
    for (idx, age) in ages {
90
156
        writer.write_record(vec![format!("{}", vcf.get_pos(*idx)), format!("{age:.5}")])?;
91
    }
92
12
    Ok(())
93
12
}
94

            
95
12
fn find_ages(
96
12
    vcf: &PhasedMatrix,
97
12
    args: &StandardArgs,
98
12
    rates: std::collections::BTreeMap<u64, f32>,
99
12
    start: usize,
100
12
    stop: usize,
101
12
    step_size: usize,
102
12
) -> Result<Vec<(usize, f64)>> {
103
12
    // Create a Vec because par_iter cannot be used with pure ranges and par_bridge does not
104
12
    // return in ordered fashion with .collect()
105
12
    let range: Vec<usize> = (start..stop).collect();
106
12

            
107
12
    match args.selection == Selection::OnlyLongest {
108
6
        true => range
109
6
            .par_iter()
110
378
            .filter(|n| *n % step_size == 0)
111
78
            .map(|i| -> Result<(usize, f64)> {
112
78
                let shared_lengths = shared_lengths_by_majority(vcf, *i);
113
78
                let only_longest_lengths = select_only_longest_alleles(&shared_lengths);
114
78
                let (i_tau_hat, _, _) =
115
78
                    mrca_only_independent(only_longest_lengths, vcf.get_pos(*i), &rates)?;
116
78
                Ok((*i, i_tau_hat.log10()))
117
78
            })
118
6
            .collect(),
119
6
        false => range
120
6
            .par_iter()
121
378
            .filter(|n| *n % step_size == 0)
122
78
            .map(|i| -> Result<(usize, f64)> {
123
78
                let shared_lengths = shared_lengths_by_majority(vcf, *i);
124
78
                let (i_tau_hat, _, _) =
125
78
                    mrca_only_independent(shared_lengths, vcf.get_pos(*i), &rates)?;
126
78
                Ok((*i, i_tau_hat.log10()))
127
78
            })
128
6
            .collect(),
129
    }
130
12
}
131

            
132
#[allow(clippy::too_many_arguments)]
133
12
fn draw_plot(
134
12
    vcf: &PhasedMatrix,
135
12
    args: &StandardArgs,
136
12
    graph_args: GraphArgs,
137
12
    mut output: PathBuf,
138
12
    data: Vec<(usize, f64)>,
139
12
    contig: &str,
140
12
    start: usize,
141
12
    stop: usize,
142
12
    mark_centromere: bool,
143
12
) -> Result<()> {
144
12
    push_to_output(args, &mut output, "mrca_scan", "png");
145
12

            
146
12
    let root =
147
12
        BitMapBackend::new(&output, (graph_args.width, graph_args.height)).into_drawing_area();
148
12

            
149
144
    let y_start = data.iter().min_by(|a, b| a.1.total_cmp(&b.1)).unwrap().1;
150
144
    let y_stop = data.iter().max_by(|a, b| a.1.total_cmp(&b.1)).unwrap().1;
151
12

            
152
12
    root.fill(&WHITE)?;
153
12
    let mut chart = ChartBuilder::on(&root)
154
12
        .margin(10)
155
12
        .set_label_area_size(LabelAreaPosition::Left, 60)
156
12
        .set_label_area_size(LabelAreaPosition::Right, 60)
157
12
        .set_label_area_size(LabelAreaPosition::Bottom, 40)
158
12
        .build_cartesian_2d(start..stop, y_start - 0.2..y_stop + 0.2)?;
159

            
160
12
    chart
161
12
        .configure_mesh()
162
12
        .disable_x_mesh()
163
12
        .disable_y_mesh()
164
12
        .x_labels(30)
165
12
        .max_light_lines(4)
166
12
        .y_desc("MRCA")
167
12
        .x_desc("POS")
168
12
        .draw()?;
169

            
170
    // Raw log10 data line
171
    // chart.draw_series(LineSeries::new(
172
    // data.iter().map(|(pos, mrca)|(*pos,*mrca)),
173
    // &BLUE,
174
    // ))?;
175

            
176
    // Last 10 elements moving average
177
12
    let mut moving_average = vec![];
178
12
    for (a, b, c, d, e, f, g, h, i, j) in data.iter().tuples() {
179
12
        let sum = a.1 + b.1 + c.1 + d.1 + e.1 + f.1 + g.1 + h.1 + i.1 + j.1;
180
12
        let mean = sum / 10.0;
181
12
        moving_average.push((j.0, mean));
182
12
    }
183

            
184
12
    chart.draw_series(LineSeries::new(
185
12
        moving_average.iter().map(|(pos, mrca)| (*pos, *mrca)),
186
12
        ShapeStyle {
187
12
            color: BLACK.mix(0.9),
188
12
            filled: true,
189
12
            stroke_width: graph_args.stroke_width,
190
12
        },
191
12
    ))?;
192

            
193
    // Centromere location
194
12
    if mark_centromere {
195
6
        let (centro_start, centro_stop) = centromeres_hg38(contig);
196
6
        let centro_start = vcf.get_nearest_idx_by_pos(centro_start as i64);
197
6
        let centro_stop = vcf.get_nearest_idx_by_pos(centro_stop as i64);
198
6

            
199
6
        chart.draw_series(LineSeries::new(
200
6
            [(centro_start, y_start), (centro_start, y_stop)].into_iter(),
201
6
            ShapeStyle {
202
6
                color: GREEN.mix(0.6),
203
6
                filled: true,
204
6
                stroke_width: 2,
205
6
            },
206
6
        ))?;
207

            
208
6
        chart.draw_series(LineSeries::new(
209
6
            [(centro_stop, y_start), (centro_stop, y_stop)].into_iter(),
210
6
            ShapeStyle {
211
6
                color: GREEN.mix(0.6),
212
6
                filled: true,
213
6
                stroke_width: 2,
214
6
            },
215
6
        ))?;
216
6
    }
217

            
218
    // Variant location
219
12
    if graph_args.mark_locus {
220
        chart.draw_series(LineSeries::new(
221
            [(vcf.variant_idx(), y_start), (vcf.variant_idx(), y_stop)].into_iter(),
222
            ShapeStyle {
223
                color: BLUE.mix(0.6),
224
                filled: true,
225
                stroke_width: 2,
226
            },
227
        ))?;
228
12
    }
229

            
230
    // Mean line
231
156
    let sum: f64 = data.iter().map(|(_pos, mrca)| mrca).sum();
232
12
    let mean = sum / data.len() as f64;
233
12
    chart.draw_series(LineSeries::new(
234
12
        [(start, mean), (stop, mean)].into_iter(),
235
12
        ShapeStyle {
236
12
            color: RED.mix(0.6),
237
12
            filled: true,
238
12
            stroke_width: graph_args.stroke_width / 2,
239
12
        },
240
12
    ))?;
241

            
242
12
    root.present().expect("Unable to write result to file");
243
12

            
244
12
    Ok(())
245
12
}