pedigree/
dmatrix.rs

1use petgraph::{algo::astar, prelude::UnGraph};
2
3use ndarray::{array, Array2, Axis};
4
5use crate::{Edge, Node, Pedigree};
6
7#[derive(Debug)]
8pub struct DMatrix(Array2<f64>);
9
10impl DMatrix {
11    pub fn from(nodes: &Vec<Node>, posterior_max: f64) -> Self {
12        println!();
13        println!();
14        println!("Calculating divergences between all pairs of nodes...");
15        println!();
16        let mut divergences = Array2::<f64>::zeros((nodes.len(), nodes.len()));
17
18        let pb = progress_bars::Progress::new(
19            "Calculating Divergences ",
20            binomial_coefficient(nodes.len(), 2),
21        );
22
23        // Go over all pairs of nodes, excluding self-pairs
24        for (i, first) in nodes.iter().enumerate() {
25            for (j, second) in nodes.iter().skip(i + 1).enumerate() {
26                pb.inc(1);
27                assert!(first.sites.as_ref().is_some());
28                assert!(second.sites.as_ref().is_some());
29                if first.sites.as_ref().unwrap().len() != second.sites.as_ref().unwrap().len() {
30                    println!(
31                        "Lengths do not match, all bets are off: {} vs {}",
32                        first.sites.as_ref().unwrap().len(),
33                        second.sites.as_ref().unwrap().len()
34                    );
35                    divergences[[i, j]] = 0.0;
36                    continue;
37                }
38                assert_eq!(
39                    first.sites.as_ref().unwrap().len(),
40                    second.sites.as_ref().unwrap().len()
41                );
42
43                let mut divergence = 0;
44                let mut compared_sites = 0;
45
46                // Go over all sites in the first sample
47                // IMPORTANT: It is assumed that the same sites are included in the datasets and that the sites are sorted by position
48                for (k, f) in first.sites.as_ref().unwrap().iter().enumerate() {
49                    let s = second
50                        .sites
51                        .as_ref()
52                        .unwrap()
53                        .get(k)
54                        .expect("Partner methylation site must exists");
55
56                    // TODO: remove, is duplicate
57                    if f.posteriormax < posterior_max || s.posteriormax < posterior_max {
58                        continue;
59                    }
60
61                    divergence += f.status_numeric().abs_diff(s.status_numeric());
62                    compared_sites += 1;
63                }
64
65                let divergence = divergence as f64 / (2.0 * compared_sites as f64);
66                divergences[[i, j]] = divergence;
67            }
68        }
69        pb.finish();
70        DMatrix(divergences)
71    }
72    /// Convert graph of divergences to pedigree
73    pub fn convert(&self, nodes: &[Node], edges: &[Edge]) -> Pedigree {
74        let e = edges
75            .iter()
76            .map(|e| {
77                (
78                    e.from.id,
79                    e.to.id,
80                    e.from.generation.abs_diff(e.to.generation) as usize,
81                    // self.0.get((e.from.id, e.to.id)).unwrap(),
82                )
83            })
84            .collect::<Vec<(usize, usize, usize)>>();
85
86        let graph = UnGraph::<usize, usize, usize>::from_edges(e);
87
88        let mut pedigree = Pedigree(Array2::<f64>::default((0, 4)));
89
90        for (i, source) in nodes.iter().enumerate() {
91            for (j, target) in nodes.iter().skip(i + 1).enumerate() {
92                if source.id == target.id {
93                    // Explicitly skip self-pairs
94                    continue;
95                }
96
97                let path = astar(
98                    &graph,
99                    source.id.into(),
100                    |finish| finish == target.id.into(),
101                    |e| *e.weight(),
102                    |_| 0,
103                );
104
105                match path {
106                    None => continue,
107                    Some(path) => {
108                        let distance = path.0;
109                        let visited = path.1;
110
111                        let t0: f64 = visited
112                            .iter()
113                            .map(|n| {
114                                edges
115                                    .iter()
116                                    .find_map(|e| {
117                                        if e.from.id == n.index() {
118                                            return Some(e.from.generation);
119                                        }
120                                        if e.to.id == n.index() {
121                                            return Some(e.to.generation);
122                                        }
123                                        None
124                                    })
125                                    .unwrap()
126                            })
127                            .min()
128                            .unwrap() as f64;
129
130                        let t1 = source.generation as f64;
131                        let t2 = target.generation as f64;
132
133                        let div = self.0.get((i, j)).unwrap().to_owned();
134
135                        assert_eq!(distance as f64, t1 - t0 + t2 - t0);
136                        pedigree
137                            .0
138                            .push(Axis(0), array![t0, t1, t2, div].view())
139                            .expect("Could not insert row into pedigree");
140                    }
141                }
142            }
143        }
144
145        pedigree
146    }
147}
148
149fn binomial_coefficient(n: usize, k: usize) -> usize {
150    let mut result = 1;
151    let k = k.min(n - k); // take advantage of symmetry
152    for i in 0..k {
153        result *= n - i;
154        result /= i + 1;
155    }
156    result
157}
158
159#[cfg(test)]
160mod test {
161    use super::*;
162    #[test]
163    fn test_binomial_coefficient() {
164        assert_eq!(binomial_coefficient(5, 0), 1);
165        assert_eq!(binomial_coefficient(5, 1), 5);
166        assert_eq!(binomial_coefficient(5, 2), 10);
167        assert_eq!(binomial_coefficient(5, 3), 10);
168        assert_eq!(binomial_coefficient(5, 4), 5);
169        assert_eq!(binomial_coefficient(5, 5), 1);
170    }
171}
pedigree/dmatrix.rs

pedigree/
dmatrix.rs
