1
use std::collections::BTreeMap;
2
use std::path::PathBuf;
3

            
4
use color_eyre::{eyre::eyre, Result};
5
use statrs::distribution::ContinuousCDF;
6
use statrs::distribution::Gamma;
7
use statrs::statistics::Statistics;
8

            
9
use crate::args::Selection;
10
use crate::args::StandardArgs;
11
use crate::io::read_recombination_file;
12
use crate::read_vcf::read_vcf_to_matrix;
13
use crate::structs::SharedLength;
14
use crate::utils::{
15
    select_carrier_haplotypes, select_only_longest_alleles, shared_lengths_by_majority,
16
};
17
use crate::core::parse_snp_coord;
18

            
19
pub type Age = (f64, f64, f64);
20

            
21
#[doc(hidden)]
22
6
pub fn run(args: StandardArgs, rec_rates: PathBuf) -> Result<()> {
23
6
    if args.selection == Selection::Unphased {
24
        return Err(eyre!("Running with unphased data is not supported."))
25
6
    }
26

            
27
6
    let (contig, variant_pos) = parse_snp_coord(&args.coords)?;
28

            
29
6
    let mut output = args.output.clone();
30
6
    output.push("mrca_gamma_method.txt");
31

            
32
6
    let rates = read_recombination_file(rec_rates)?;
33

            
34
6
    let mut vcf = read_vcf_to_matrix(&args, contig, variant_pos, None)?;
35

            
36
6
    if args.selection == Selection::OnlyAlts || args.selection == Selection::OnlyRefs {
37
        vcf = select_carrier_haplotypes(vcf, variant_pos, &args.coords, &args.selection)
38
6
    };
39

            
40
6
    let shared_lengths = shared_lengths_by_majority(&vcf, vcf.variant_idx());
41

            
42
6
    let ((i_tau_hat, i_l, i_u), (c_tau_hat, c_l, c_u)) = if args.selection == Selection::OnlyLongest
43
    {
44
6
        let only_longest_lengths = select_only_longest_alleles(&shared_lengths);
45
6
        mrca_gamma_method(only_longest_lengths, variant_pos, &rates)?
46
    } else {
47
        mrca_gamma_method(shared_lengths, variant_pos, &rates)?
48
    };
49

            
50
6
    let data = format!("Independent genealogy:\nage: {i_tau_hat:.3} CI ({i_l:.3}, {i_u:.3})\nCorrelated genealogy:\nage: {c_tau_hat:.3} CI ({c_l:.3}, {c_u:.3})");
51
6

            
52
6
    std::fs::write(&output, data).expect(&format!("Unable to write to {output:?}"));
53
6

            
54
6
    Ok(())
55
6
}
56

            
57
///
58
/// The original R algorithm by Gandolfo et al translated to Rust.
59
/// <https://github.com/bahlolab/DatingRareMutations>
60
///
61
6
pub fn mrca_gamma_method<T: AsRef<SharedLength>>(
62
6
    shared_lengths: Vec<T>,
63
6
    variant_pos: i64,
64
6
    rec_rates: &BTreeMap<u64, f32>,
65
6
) -> Result<(Age, Age)> {
66
6
    let mut l_lengths = shared_lengths
67
6
        .iter()
68
84
        .map(|v| {
69
84
            // Transform from pos to centimorgans by selecting nearest value to the left in the
70
84
            // BTreeMap
71
84
            let variant_cm = rec_rates.range((variant_pos as u64)..).next().unwrap();
72
84
            let break_cm =
73
84
                if let Some(cm) = rec_rates.range(..(v.as_ref().start_pos as u64)).next_back() {
74
84
                    cm
75
                } else {
76
                    rec_rates
77
                        .range((v.as_ref().start_pos as u64)..)
78
                        .next()
79
                        .unwrap()
80
                };
81

            
82
84
            (variant_cm.1 - break_cm.1) as f64 / 100.
83
84
        })
84
6
        .collect::<Vec<f64>>();
85
6

            
86
6
    let mut r_lengths = shared_lengths
87
6
        .iter()
88
84
        .map(|v| {
89
84
            let variant_cm = rec_rates.range((variant_pos as u64)..).next().unwrap();
90
84
            let break_cm = if let Some(cm) = rec_rates.range((v.as_ref().end_pos as u64)..).next() {
91
84
                cm
92
            } else {
93
                rec_rates
94
                    .range(..(v.as_ref().end_pos as u64))
95
                    .next_back()
96
                    .unwrap()
97
            };
98

            
99
84
            (break_cm.1 - variant_cm.1) as f64 / 100.
100
84
        })
101
6
        .collect::<Vec<f64>>();
102
6

            
103
6
    let cc = 0.95;
104
6
    let n = l_lengths.len() as f64;
105
6
    let cs_corr = 0.0;
106
6
    let l_sum: f64 = l_lengths.iter().sum();
107
6
    let r_sum: f64 = r_lengths.iter().sum();
108
6

            
109
6
    // Independent genealogy
110
6
    let length_corr = (l_sum + r_sum - 2.0 * (n - 1.0) * cs_corr) / (2.0 * n);
111
6

            
112
6
    // Sum of ancestral segment lengths with corrections
113
6
    let sum: f64 = l_sum + r_sum + 2.0 * length_corr - 2.0 * (n - 1.0) * cs_corr;
114
6

            
115
6
    // Gamma function MLE bias correction factor
116
6
    let b_c = (2.0 * n - 1.0) / (2.0 * n);
117
6

            
118
6
    // Minimum variance unbiased estimate of T
119
6
    // Compare it with the pure MLE i.e. the mean of the gamma distribution = α / λ
120
6
    let i_tau_hat = (b_c * 2.0 * n) / sum;
121

            
122
6
    let gamma = Gamma::new(2.0 * n, 2.0 * n * b_c)?;
123
6
    let g_l = gamma.inverse_cdf((1.0 - cc) / 2.0);
124
6
    let g_u = gamma.inverse_cdf(cc + (1.0 - cc) / 2.0);
125
6

            
126
6
    let i_tau_hat_l = i_tau_hat * g_l;
127
6
    let i_tau_hat_u = i_tau_hat * g_u;
128
6

            
129
6
    tracing::debug!(
130
        "Independent genealogy:\nage: {i_tau_hat:.3} CI ({i_tau_hat_l:.3}, {i_tau_hat_u:.3})",
131
    );
132

            
133
    // Correlated genealogy
134
6
    let length_corr: f64 = (l_sum + r_sum - 2.0 * (n - 1.0) * cs_corr) / (2.0 * n);
135
6

            
136
78
    let highest_l = l_lengths.iter_mut().max_by(|a, b| a.total_cmp(b)).unwrap();
137
6
    *highest_l = *highest_l + length_corr + cs_corr;
138
6

            
139
78
    let highest_r = r_lengths.iter_mut().max_by(|a, b| a.total_cmp(b)).unwrap();
140
6
    *highest_r = *highest_r + length_corr + cs_corr;
141
6

            
142
6
    let lengths = &l_lengths
143
6
        .iter()
144
6
        .zip(r_lengths.iter())
145
84
        .map(|(a, b)| a + b - 2.0 * cs_corr)
146
6
        .collect::<Vec<f64>>();
147
6

            
148
6
    let sum: f64 = lengths.iter().sum();
149
6

            
150
6
    let term1 = n * lengths.mean().powi(2);
151
6
    let term2 = (lengths.variance() * (1.0 + 2.0 * 2.0))
152
6
        / (n * lengths.mean().powi(2) + lengths.variance() * (n - 1.0));
153
6
    let term3 = lengths.variance() * n - 1.0;
154
6

            
155
6
    let rho_hat = term1 - term2 + term3;
156
6

            
157
6
    let mut n_star = n / (1.0 + (n - 1.0) * rho_hat);
158
6

            
159
6
    if n_star > n {
160
        n_star = n;
161
6
    }
162
6
    if n_star < -n {
163
6
        n_star = -n;
164
6
    }
165

            
166
6
    let b_c = (2.0 * n_star - 1.0) / (2.0 * n_star);
167
6
    let c_tau_hat = (b_c * 2.0 * n) / sum;
168
6

            
169
6
    if -2.0 / (n - 1.0) <= rho_hat && rho_hat < -1.0 / (n - 1.0) {
170
6
        n_star = n;
171
6
    } else if rho_hat < -2.0 / (n - 1.0) {
172
        n_star = n / (1.0 + (n - 1.0) * rho_hat.abs());
173
    };
174

            
175
6
    let gamma = Gamma::new(2.0 * n_star, 2.0 * n_star * b_c)?;
176
6
    let c_l = gamma.inverse_cdf((1.0 - cc) / 2.0);
177
6
    let c_u = gamma.inverse_cdf(cc + (1.0 - cc) / 2.0);
178
6

            
179
6
    let c_tau_hat_l = c_tau_hat * c_l;
180
6
    let c_tau_hat_u = c_tau_hat * c_u;
181
6
    tracing::debug!(
182
        "Correlated genealogy:\nage: {c_tau_hat:.3} CI ({c_tau_hat_l:.3}, {c_tau_hat_u:.3})",
183
    );
184

            
185
6
    let independent = (i_tau_hat, i_tau_hat * g_l, i_tau_hat * g_u);
186
6
    let correlated = (c_tau_hat, c_tau_hat_l, c_tau_hat_u);
187
6

            
188
6
    Ok((independent, correlated))
189
6
}
190

            
191
///
192
/// The original R algorithm by Gandolfo et al translated to Rust.
193
/// <https://github.com/bahlolab/DatingRareMutations>
194
///
195
156
pub fn mrca_only_independent<T: AsRef<SharedLength>>(
196
156
    shared_lengths: Vec<T>,
197
156
    variant_pos: i64,
198
156
    rec_rates: &BTreeMap<u64, f32>,
199
156
) -> Result<(f64, f64, f64)> {
200
156
    let l_lengths = shared_lengths
201
156
        .iter()
202
3276
        .map(|v| {
203
3276
            let variant_cm = if let Some(cm) = rec_rates.range((variant_pos as u64)..).next() {
204
3276
                cm
205
            } else {
206
                rec_rates.range(..(variant_pos as u64)).next_back().unwrap()
207
            };
208

            
209
3276
            let break_cm =
210
3276
                if let Some(cm) = rec_rates.range(..(v.as_ref().start_pos as u64)).next_back() {
211
2466
                    cm
212
                } else {
213
810
                    rec_rates
214
810
                        .range((v.as_ref().start_pos as u64)..)
215
810
                        .next()
216
810
                        .unwrap()
217
                };
218

            
219
3276
            (variant_cm.1 - break_cm.1) as f64 / 100.
220
3276
        })
221
156
        .collect::<Vec<f64>>();
222
156

            
223
156
    let r_lengths = shared_lengths
224
156
        .iter()
225
3276
        .map(|v| {
226
3276
            let variant_cm = if let Some(cm) = rec_rates.range((variant_pos as u64)..).next() {
227
3276
                cm
228
            } else {
229
                rec_rates.range(..(variant_pos as u64)).next_back().unwrap()
230
            };
231

            
232
3276
            let break_cm = if let Some(cm) = rec_rates.range((v.as_ref().end_pos as u64)..).next() {
233
3276
                cm
234
            } else {
235
                rec_rates
236
                    .range(..(v.as_ref().end_pos as u64))
237
                    .next_back()
238
                    .unwrap()
239
            };
240
3276
            assert!(
241
3276
                variant_cm.1 <= break_cm.1,
242
                "centimorgans are nto in order {variant_cm:?} {break_cm:?} {variant_pos} {}",
243
                v.as_ref().end_pos
244
            );
245

            
246
3276
            (break_cm.1 - variant_cm.1) as f64 / 100.
247
3276
        })
248
156
        .collect::<Vec<f64>>();
249
156

            
250
156
    let cc = 0.95;
251
156
    let n = l_lengths.len() as f64;
252
156
    let cs_corr = 0.0;
253
156
    let l_sum: f64 = l_lengths.iter().sum();
254
156
    let r_sum: f64 = r_lengths.iter().sum();
255
156

            
256
156
    // Independent genealogy
257
156
    let length_corr = (l_sum + r_sum - 2.0 * (n - 1.0) * cs_corr) / (2.0 * n);
258
156

            
259
156
    // Sum of ancestral segment lengths with corrections
260
156
    let sum: f64 = l_sum + r_sum + 2.0 * length_corr - 2.0 * (n - 1.0) * cs_corr;
261
156

            
262
156
    // Gamma function MLE bias correction factor
263
156
    let b_c = (2.0 * n - 1.0) / (2.0 * n);
264
156

            
265
156
    // Minimum variance unbiased estimate of T
266
156
    // Compare it with the pure MLE i.e. the mean of the gamma distribution = α / λ
267
156
    let i_tau_hat = (b_c * 2.0 * n) / sum;
268

            
269
156
    let gamma = Gamma::new(2.0 * n, 2.0 * n * b_c)?;
270
156
    let g_l = gamma.inverse_cdf((1.0 - cc) / 2.0);
271
156
    let g_u = gamma.inverse_cdf(cc + (1.0 - cc) / 2.0);
272
156

            
273
156
    let independent = (i_tau_hat, i_tau_hat * g_l, i_tau_hat * g_u);
274
156

            
275
156
    Ok(independent)
276
156
}