1
use std::path::PathBuf;
2

            
3
use ndarray::{s, Array2};
4
use statrs::distribution::ContinuousCDF;
5
use statrs::distribution::StudentsT;
6
use statrs::statistics::Statistics;
7

            
8
use crate::args::{Selection, StandardArgs};
9
use crate::structs::{PhasedMatrix, SharedLength};
10

            
11
pub fn push_to_output(args: &StandardArgs, output: &mut PathBuf, name: &str, suffix: &str) {
12
357
    if let Some(prefix) = &args.prefix {
13
1
        match args.selection {
14
1
            Selection::All => output.push(format!("{prefix}_{name}.{suffix}")),
15
            Selection::OnlyAlts => output.push(format!("{prefix}_{name}_only_alts.{suffix}")),
16
            Selection::OnlyRefs => output.push(format!("{prefix}_{name}_only_refs.{suffix}")),
17
            Selection::OnlyLongest => output.push(format!("{prefix}_{name}_only_longest.{suffix}")),
18
            Selection::Unphased => output.push(format!("{prefix}_{name}_rwc.{suffix}")),
19
        }
20
    } else {
21
356
        match args.selection {
22
140
            Selection::All => output.push(format!("{name}.{suffix}")),
23
24
            Selection::OnlyAlts => output.push(format!("{name}_only_alts.{suffix}")),
24
24
            Selection::OnlyRefs => output.push(format!("{name}_only_refs.{suffix}")),
25
156
            Selection::OnlyLongest => output.push(format!("{name}_only_longest.{suffix}")),
26
12
            Selection::Unphased => output.push(format!("{name}_rwc.{suffix}")),
27
        }
28
    }
29
357
}
30

            
31
pub fn push_to_output_mbah(args: &StandardArgs, output: &mut PathBuf, name: &str, suffix: &str) {
32
12
    if let Some(prefix) = &args.prefix {
33
        match args.selection {
34
            Selection::All | Selection::OnlyLongest => output.push(format!("{prefix}_{name}.{suffix}")),
35
            Selection::OnlyAlts => output.push(format!("{prefix}_{name}_only_alts.{suffix}")),
36
            Selection::OnlyRefs => output.push(format!("{prefix}_{name}_only_refs.{suffix}")),
37
            Selection::Unphased => output.push(format!("{prefix}_{name}_unphased.{suffix}")),
38
        }
39
    } else {
40
12
        match args.selection {
41
12
            Selection::All | Selection::OnlyLongest => output.push(format!("{name}.{suffix}")),
42
            Selection::OnlyAlts => output.push(format!("{name}_only_alts.{suffix}")),
43
            Selection::OnlyRefs => output.push(format!("{name}_only_refs.{suffix}")),
44
            Selection::Unphased => output.push(format!("{name}_unphased.{suffix}")),
45
        }
46
    }
47
12
}
48

            
49
pub fn filter_samples(samples: &Vec<String>, wanted: Option<Vec<String>>) -> Vec<usize> {
50
868
    if let Some(wanted) = wanted {
51
11
        for i in &wanted {
52
8
            if !samples.contains(i) {
53
2
                tracing::warn!("Wanted sample {i} is not in the vcf");
54
6
            }
55
        }
56

            
57
3
        samples
58
3
            .iter()
59
3
            .enumerate()
60
7
            .filter(|(_, s)| wanted.contains(s))
61
6
            .map(|(i, _)| i)
62
3
            .collect()
63
    } else {
64
865
        (0..samples.len()).collect()
65
    }
66
868
}
67

            
68
48
pub fn select_carrier_haplotypes(
69
48
    mut vcf: PhasedMatrix,
70
48
    variant_pos: i64,
71
48
    coords: &str,
72
48
    selection: &Selection
73
48
) -> PhasedMatrix {
74
48
    assert_eq!(
75
48
        vcf.variant_idx_pos(),
76
        variant_pos,
77
        "Cannot select variant carriers: no variant at locus {coords}"
78
    );
79
48
    let samples_len = vcf.matrix.nrows() / 2;
80
48
    let mut haplotype_idxs = vec![];
81
756
    for first_allele_idx in 0..samples_len {
82
756
        let second_allele_idx = first_allele_idx + samples_len;
83
756
        let a = *vcf.matrix.slice(s![first_allele_idx, vcf.variant_idx]).into_scalar();
84
756
        let b = *vcf.matrix.slice(s![second_allele_idx, vcf.variant_idx]).into_scalar();
85
756

            
86
756
        match selection {
87
            Selection::OnlyAlts => {
88
504
                match (a, b) {
89
126
                    (1, 0) => haplotype_idxs.push(first_allele_idx),
90
378
                    (0, 1) => haplotype_idxs.push(second_allele_idx),
91
                    (1, 1) => {
92
                        haplotype_idxs.push(first_allele_idx);
93
                        haplotype_idxs.push(second_allele_idx);
94
                    }
95
                    (0, 0) => {}
96
                    _ => panic!(),
97
                }
98
            }
99
            Selection::OnlyRefs => {
100
252
                match (a, b) {
101
42
                    (1, 0) => haplotype_idxs.push(second_allele_idx),
102
210
                    (0, 1) => haplotype_idxs.push(first_allele_idx),
103
                    (1, 1) => {}
104
                    (0, 0) => {
105
                        haplotype_idxs.push(first_allele_idx);
106
                        haplotype_idxs.push(second_allele_idx);
107
                    }
108
                    _ => panic!(),
109
                }
110
            }
111
            _ => panic!("Invalid selection method for alleles")
112
        }
113
    }
114
54
    vcf.select_rows(haplotype_idxs);
115
54
    vcf
116
54
}
117

            
118
282
pub fn shared_lengths_by_majority(vcf: &PhasedMatrix, variant_idx: usize) -> Vec<SharedLength> {
119
282
    let mut lengths: Vec<SharedLength> = (0..vcf.nsamples())
120
7896
        .map(|i| SharedLength {
121
7896
            index: i,
122
7896
            start: None,
123
7896
            end: None,
124
7896
            start_pos: 0,
125
7896
            end_pos: 0,
126
7896
        })
127
282
        .collect();
128
282

            
129
282
    // Right side
130
282
    let mut mutating_indexes: Vec<usize> = (0..vcf.nsamples()).collect();
131
8022
    for pos in variant_idx..vcf.matrix.ncols() {
132
8022
        let submatrix = vcf.matrix.select(ndarray::Axis(0), &mutating_indexes);
133
8022
        // Returns the majority of indexes
134
8022
        let majority_idxs = return_majority(&submatrix, pos, &mutating_indexes);
135
8022

            
136
8022
        if majority_idxs.len() == 1 {
137
5544
            lengths.iter_mut().for_each(|i| {
138
5544
                if i.end.is_none() {
139
396
                    i.end = Some(pos)
140
5148
                }
141
5544
            });
142
198
            break;
143
        } else {
144
            // If the given indexes equals majority of indexes, no indexes were removed
145
7824
            if mutating_indexes.len() != majority_idxs.len() {
146
4362
                // Add end positions to all indexes present in the original indexes, but not in the
147
4362
                // majority of indexes
148
61896
                mutating_indexes.iter().for_each(|i| {
149
61896
                    if !majority_idxs.contains(i) {
150
6288
                        lengths[*i].end = Some(pos)
151
55608
                    }
152
61896
                });
153
4362
            }
154
        }
155
7824
        mutating_indexes = majority_idxs;
156
    }
157

            
158
    // Left side
159
282
    let mut mutating_indexes: Vec<usize> = (0..vcf.nsamples()).collect();
160
7440
    for pos in (0..variant_idx).rev() {
161
7440
        let submatrix = vcf.matrix.select(ndarray::Axis(0), &mutating_indexes);
162
7440
        let majority_idxs = return_majority(&submatrix, pos, &mutating_indexes);
163
7440

            
164
7440
        if majority_idxs.len() == 1 {
165
5544
            lengths.iter_mut().for_each(|i| {
166
5544
                if i.start.is_none() {
167
396
                    i.start = Some(pos)
168
5148
                }
169
5544
            });
170
198
            break;
171
7242
        } else if mutating_indexes.len() != majority_idxs.len() {
172
94914
            mutating_indexes.iter().for_each(|i| {
173
94914
                if !majority_idxs.contains(i) {
174
6096
                    lengths[*i].start = Some(pos)
175
88818
                }
176
94914
            });
177
5880
        }
178

            
179
7242
        mutating_indexes = majority_idxs;
180
    }
181

            
182
7896
    lengths.iter_mut().for_each(|i| {
183
7896
        if i.start.is_none() {
184
1404
            i.start = Some(0)
185
6492
        }
186
7896
        if i.end.is_none() {
187
1212
            i.end = Some(vcf.matrix.ncols() - 1)
188
6684
        }
189
7896
        i.start_pos = vcf.get_pos(i.start());
190
7896
        i.end_pos = vcf.get_pos(i.end());
191
7896
    });
192
282

            
193
282
    lengths
194
282
}
195

            
196
15469
fn return_majority(matrix: &Array2<u8>, pos: usize, indexes: &[usize]) -> Vec<usize> {
197
15469
    let zeroes: Vec<usize> = matrix
198
15469
        .slice(ndarray::s![.., pos])
199
15469
        .iter()
200
15469
        .zip(indexes.to_owned())
201
219408
        .filter_map(|(n, i)| if *n == 0 { Some(i) } else { None })
202
15469
        .collect();
203
15469
    let ones: Vec<usize> = matrix
204
15469
        .slice(ndarray::s![.., pos])
205
15469
        .iter()
206
15469
        .zip(indexes.to_owned())
207
219402
        .filter_map(|(n, i)| if *n == 1 { Some(i) } else { None })
208
15469
        .collect();
209
15469

            
210
15469
    match zeroes.len().cmp(&ones.len()) {
211
484
        std::cmp::Ordering::Less => ones,
212
14449
        std::cmp::Ordering::Greater => zeroes,
213
        std::cmp::Ordering::Equal => {
214
536
            if ones.len() != 1 {
215
138
                tracing::warn!(
216
                "At {} an equal amount of genotypes present: {} {}, selecting alts as the majority",
217
                pos,
218
                ones.len(),
219
                zeroes.len()
220
            );
221
398
            }
222
536
            ones
223
        }
224
    }
225
15469
}
226

            
227
90
pub fn select_only_longest_haplotypes<T: AsRef<SharedLength>>(
228
90
    lengths: &Vec<T>,
229
90
    mut vcf: PhasedMatrix,
230
90
) -> PhasedMatrix {
231
90
    let mut only_longest_lengths = vec![];
232

            
233
1260
    for index in 0..lengths.len() / 2 {
234
1260
        let index2 = index + lengths.len() / 2;
235
1260
        let allele1 = &lengths[index];
236
1260
        let allele2 = &lengths[index2];
237
1260
        let allele1_length = allele1.as_ref().end_pos - allele1.as_ref().start_pos;
238
1260
        let allele2_length = allele2.as_ref().end_pos - allele2.as_ref().start_pos;
239
1260

            
240
1260
        match allele1_length.cmp(&allele2_length) {
241
1176
            std::cmp::Ordering::Less => only_longest_lengths.push(index2),
242
84
            std::cmp::Ordering::Greater => only_longest_lengths.push(index),
243
            std::cmp::Ordering::Equal => {
244
                tracing::warn!(
245
                "Sample {} has two equally long haplotypes. Selecting the first one in the matrix.",
246
                vcf.get_sample_name(index),
247
            );
248
                only_longest_lengths.push(index)
249
            }
250
        }
251
    }
252

            
253
90
    vcf.select_rows(only_longest_lengths);
254
90

            
255
90
    vcf
256
90
}
257

            
258
108
pub fn select_only_longest_alleles(shared_lengths: &Vec<SharedLength>) -> Vec<&SharedLength> {
259
108
    let mut only_longest_lengths = vec![];
260

            
261
1512
    for index in 0..&shared_lengths.len() / 2 {
262
1512
        let allele1 = &shared_lengths[index];
263
1512
        let allele2 = &shared_lengths[index + shared_lengths.len() / 2];
264
1512
        let allele1_length = allele1.end_pos - allele1.start_pos;
265
1512
        let allele2_length = allele2.end_pos - allele2.start_pos;
266
1512

            
267
1512
        match allele1_length.cmp(&allele2_length) {
268
1134
            std::cmp::Ordering::Less => only_longest_lengths.push(allele2),
269
378
            std::cmp::Ordering::Greater => only_longest_lengths.push(allele1),
270
            std::cmp::Ordering::Equal => {
271
                tracing::warn!(
272
                    "Sample index {} has two equally long haplotypes. Selecting the first one in the matrix.",
273
                    index
274
                );
275
                only_longest_lengths.push(allele1)
276
            }
277
        }
278
    }
279
108
    only_longest_lengths
280
108
}
281

            
282
9
pub fn two_tail_welch_t_test(x: &[f64], y: &[f64]) -> f64 {
283
9
    let n1 = x.len() as f64;
284
9
    let n2 = y.len() as f64;
285
9
    let mean1 = x.mean();
286
9
    let mean2 = y.mean();
287
9
    let variance1 = f64::powi(x.std_dev(), 2);
288
9
    let variance2 = f64::powi(y.std_dev(), 2);
289
9
    let t = (mean1 - mean2) / (variance1 / n1 + variance2 / n2).sqrt();
290
9
    let var_by_n1 = variance1 / n1;
291
9
    let var_by_n2 = variance2 / n2;
292
9
    let upper_part = f64::powi(var_by_n1 + var_by_n2, 2);
293
9
    let lower_part = f64::powi(var_by_n1, 2) / (n1 - 1.0) + f64::powi(var_by_n2, 2) / (n2 - 1.0);
294
9
    let df = upper_part / lower_part;
295
9
    let tdist = StudentsT::new(0.0, 1.0, df).unwrap();
296
9
    tdist.cdf(t) * 2.0
297
9
}
298

            
299
#[cfg(test)]
300
#[rustfmt::skip]
301
mod tests {
302
    use super::*;
303

            
304
1
    #[test]
305
1
    fn sample_filtering() {
306
1
        let samples = vec!["foo".to_string(), "fii".to_string()];
307
1
        let wanted = vec!["foo".to_string(), "fii".to_string()];
308
1
        let sample_indexes = filter_samples(&samples, Some(wanted));
309
1
        assert_eq!(sample_indexes, vec![0,1]);
310

            
311
1
        let samples = vec!["foo".to_string(), "fii".to_string()];
312
1
        let wanted = vec!["fee".to_string(), "foo".to_string(), "fii".to_string()];
313
1
        let sample_indexes = filter_samples(&samples, Some(wanted));
314
1
        assert_eq!(sample_indexes, vec![0,1]);
315

            
316
1
        let samples = vec!["faa".to_string(), "foo".to_string(), "fii".to_string()];
317
1
        let wanted = vec!["fee".to_string(), "foo".to_string(), "fii".to_string()];
318
1
        let sample_indexes = filter_samples(&samples, Some(wanted));
319
1
        assert_eq!(sample_indexes, vec![1,2]);
320

            
321
1
        let sample_indexes = filter_samples(&samples, None);
322
1
        assert_eq!(sample_indexes, vec![0,1,2]);
323
1
    }
324

            
325
1
    #[test]
326
1
    fn test_two_tail_welch_t_test() {
327
1
        let vec1 = vec![1.0, 2.0, 1.0, 3.0];
328
1
        let vec2 = vec![1.0, 2.0, 1.0, 5.0];
329
1
        let result = two_tail_welch_t_test(&vec1, &vec2);
330
1
        assert_eq!("0.6596", &format!("{result:.4}"));
331

            
332
1
        let vec1 = vec![1.0, -2.0, 1.0, 3.0];
333
1
        let vec2 = vec![-1.0, 2.0, 1.0, 5.0];
334
1
        let result = two_tail_welch_t_test(&vec1, &vec2);
335
1
        assert_eq!("0.5606", &format!("{result:.4}"));
336

            
337
1
        let vec1 = vec![1.0, -2.0, 1.0, 3.0];
338
1
        let vec2 = vec![-1.0, 2.0, 1.0, 5.0, 5.0, 5.0];
339
1
        let result = crate::utils::two_tail_welch_t_test(&vec1, &vec2);
340
1
        assert_eq!("0.1959", &format!("{result:.4}"));
341
1
    }
342

            
343
1
    #[test]
344
1
    fn test_return_majority() {
345
1
        let matrix = ndarray::arr2(&[
346
1
                [ 1,  1,  1],
347
1
                [ 0,  1,  0]
348
1
        ]);
349
1

            
350
1
        let maj = return_majority(&matrix, 0, &[5, 8]);
351
1
        assert_eq!(maj, vec![5]);
352

            
353
1
        let maj = return_majority(&matrix, 1, &[5, 8]);
354
1
        assert_eq!(maj, vec![5, 8]);
355

            
356
1
        let maj = return_majority(&matrix, 2, &[5, 8]);
357
1
        assert_eq!(maj, vec![5]);
358

            
359
1
        let matrix = ndarray::arr2(&[
360
1
                [ 1,  0,  1, 1],
361
1
                [ 0,  0,  1, 1],
362
1
                [ 1,  0,  0, 1]
363
1
        ]);
364
1
        let maj = return_majority(&matrix, 0, &[5, 8, 100]);
365
1
        assert_eq!(maj, vec![5, 100]);
366

            
367
1
        let maj = return_majority(&matrix, 1, &[5, 8, 100]);
368
1
        assert_eq!(maj, vec![5, 8, 100]);
369

            
370
1
        let maj = return_majority(&matrix, 2, &[5, 8, 100]);
371
1
        assert_eq!(maj, vec![5, 8]);
372

            
373
1
        let maj = crate::utils::return_majority(&matrix, 3, &[5, 8, 100]);
374
1
        assert_eq!(maj, vec![5, 8, 100]);
375
1
    }
376

            
377
1
    #[test]
378
1
    fn test_push_to_output() {
379
1
        let mut output = std::path::PathBuf::new();
380
1
        let args = crate::args::StandardArgs::default();
381
1
        push_to_output(&args, &mut output, "picture", "png");
382
1
        assert_eq!(output, std::path::PathBuf::from("picture.png"));
383

            
384
1
        let mut output = std::path::PathBuf::from("./foo");
385
1
        let args = crate::args::StandardArgs::default();
386
1
        push_to_output(&args, &mut output, "picture", "png");
387
1
        assert_eq!(output, std::path::PathBuf::from("./foo/picture.png"));
388

            
389
1
        let mut output = std::path::PathBuf::from("./foo");
390
1
        let args = crate::args::StandardArgs { 
391
1
            prefix: Some("nice".to_string()), 
392
1
            ..Default::default()
393
1
        };
394
1
        push_to_output(&args, &mut output, "picture", "png");
395
1
        assert_eq!(output, std::path::PathBuf::from("./foo/nice_picture.png"));
396
1
    }
397
}