1
use color_eyre::Result;
2
use ndarray::{s, Array1, Array2, Axis};
3
use ndarray_rand::{RandomExt, SamplingStrategy};
4
use rayon::prelude::*;
5

            
6
use crate::args::{Selection, StandardArgs};
7
use crate::read_vcf::{read_vcf_to_contra_matrix, read_vcf_to_matrix};
8
use crate::utils::{
9
    select_carrier_haplotypes, select_only_longest_haplotypes, shared_lengths_by_majority, push_to_output,
10
};
11
use crate::core::{open_csv_writer, parse_snp_coord};
12

            
13
#[doc(hidden)]
14
24
#[tracing::instrument]
15
12
pub fn run(
16
    args: StandardArgs,
17
    niters: i64,
18
    sample_size: usize,
19
) -> Result<()> {
20
    let mut output = args.output.clone();
21
    push_to_output(&args, &mut output, "random_sampled_shared_haplotype", "csv");
22
    let mut writer = open_csv_writer(output)?;
23
    let (contig, variant_pos) = parse_snp_coord(&args.coords)?;
24

            
25
    let (variant_matrix, pos_vec) = match args.selection {
26
        Selection::OnlyAlts | Selection::OnlyRefs => {
27
            let vcf = read_vcf_to_matrix(&args, contig, variant_pos, None)?;
28
            let vcf = select_carrier_haplotypes(vcf, variant_pos, &args.coords, &args.selection);
29
378
            let pos_vec = vcf.coords().iter().map(|c| c.pos).collect::<Vec<i64>>();
30
            (vcf.matrix, pos_vec)
31
        }
32
        Selection::OnlyLongest => {
33
            let mut vcf = read_vcf_to_matrix(&args, contig, variant_pos, None)?;
34
            let shared_lengths = shared_lengths_by_majority(&vcf, vcf.variant_idx());
35
            vcf = select_only_longest_haplotypes(&shared_lengths, vcf);
36
            let pos_vec = vcf.coords().iter().map(|c| c.pos).collect::<Vec<i64>>();
37
            (vcf.matrix, pos_vec)
38
        }
39
        Selection::All => {
40
            let vcf = read_vcf_to_matrix(&args, contig, variant_pos, None)?;
41
            let pos_vec = vcf.coords().iter().map(|c| c.pos).collect::<Vec<i64>>();
42
            (vcf.matrix, pos_vec)
43
        }
44
        Selection::Unphased => read_vcf_to_contra_matrix(&args, contig, None)?
45
    };
46
    tracing::debug!("Finished reading the matrix with selection {:?}.", args.selection);
47

            
48
    writer.write_record(vec!["start", "snp_count", "bp_length"])?;
49

            
50
    // Multithreaded random sampling of the variant matrix
51
    let lengths = (0..niters)
52
        .into_par_iter()
53
1194
        .flat_map(|_| {
54
1194
            let sampled_matrix = variant_matrix.sample_axis(
55
1194
                Axis(0),
56
1194
                sample_size,
57
1194
                SamplingStrategy::WithoutReplacement,
58
1194
            );
59
1194
            let shared_sequence =
60
1194
                shared_sequences(&sampled_matrix, args.selection != Selection::Unphased);
61
1194
            find_shared_sequence_lengths(shared_sequence, &pos_vec, variant_pos)
62
1194
        })
63
        .collect::<Vec<(i64, i64, i64)>>();
64
    tracing::debug!("Finished processing lengths.");
65

            
66
    for (start, snp_count, bp_length) in lengths {
67
        writer.write_record(vec![
68
            start.to_string(),
69
            snp_count.to_string(),
70
            bp_length.to_string(),
71
        ])?;
72
    }
73

            
74
    Ok(())
75
}
76

            
77
// Create a 1-D array from 2-D a genotype matrix that denotes all the positions where the samples
78
// have a contrahomozygote in between them. A 0 denotes no contrahomozygotes are present and 1
79
// denotes the opposite
80
//
81
#[doc(hidden)]
82
1200
pub fn shared_sequences(matrix: &Array2<u8>, phased: bool) -> Array1<u8> {
83
1200
    let mut shared_sequence = Array1::<u8>::from_elem(matrix.ncols(), 0);
84

            
85
    // Flatten a 2-D matrix into an array. If both 0 and 2 exist denote as 1 else remain as 0
86
74778
    for i in 0..matrix.ncols() {
87
74778
        let column = matrix.slice(ndarray::s![.., i]);
88
74778

            
89
74778
        if phased {
90
135726
            if column.into_iter().any(|v| *v == 0) && column.into_iter().any(|v| *v == 1) {
91
4794
                shared_sequence.slice_mut(s![i]).fill(1);
92
32712
            }
93
129846
        } else if column.into_iter().any(|v| *v == 0) && column.into_iter().any(|v| *v == 2) {
94
            shared_sequence.slice_mut(s![i]).fill(1);
95
36624
        }
96
    }
97
1200
    shared_sequence
98
1200
}
99

            
100
#[doc(hidden)]
101
1200
pub fn find_shared_sequence_lengths(
102
1200
    shared_sequence: Array1<u8>,
103
1200
    pos_coords: &[i64],
104
1200
    variant_pos: i64,
105
1200
) -> Option<(i64, i64, i64)> {
106
1200
    let mut shared_lengths_and_pos = None;
107
1200
    let mut counter = 1;
108
1200

            
109
1200
    shared_sequence
110
1200
        .iter()
111
1200
        .enumerate()
112
1200
        .fold(1, |last, (index, curr)| {
113
75534
            if last == 0 && *curr == 0 {
114
65718
                counter += 1;
115
65718
            } else if last == 0 && *curr == 1 {
116
3816
                let start = pos_coords.get(index - counter).unwrap();
117
3816
                let stop = pos_coords.get(index).unwrap();
118
3816
                // Find the sequence that overlaps the targeted variant
119
3816
                if *start < variant_pos && *stop > variant_pos {
120
600
                    shared_lengths_and_pos = Some((*start, counter as i64, stop - start));
121
3216
                }
122
3816
                counter = 1;
123
6000
            }
124
75534
            *curr
125
75534
        });
126
1200
    counter -= 1;
127
1200
    if counter > 0 {
128
1200
        let index =  shared_sequence.len() - 1;
129
1200
        let start = pos_coords.get(index - counter).unwrap();
130
1200
        let stop = pos_coords.get(index).unwrap();
131
1200
        // Find the sequence that overlaps the targeted variant
132
1200
        if *start < variant_pos && *stop > variant_pos {
133
600
            shared_lengths_and_pos = Some((*start, counter as i64, stop - start));
134
600
        }
135
    }
136
1200
    shared_lengths_and_pos
137
1200
}