causal_hub/inference/topological_order.rs
1use std::collections::VecDeque;
2
3use ndarray::prelude::*;
4
5use crate::{
6 models::{DiGraph, Graph},
7 set,
8};
9
10/// Topological sort trait.
11pub trait TopologicalOrder {
12 /// Returns the topological sort of the graph.
13 ///
14 /// # Returns
15 ///
16 /// A vector of vertex indices in topological order,
17 /// or `None` if the order does not exists.
18 ///
19 fn topological_order(&self) -> Option<Vec<usize>>;
20}
21
22impl TopologicalOrder for DiGraph {
23 fn topological_order(&self) -> Option<Vec<usize>> {
24 // Compute the in-degrees of the vertices.
25 let mut in_degree = self
26 .to_adjacency_matrix()
27 .mapv(|x| x as usize)
28 .sum_axis(Axis(0));
29 // Initialize queue with vertices having in-degree 0
30 let mut to_be_visited: VecDeque<usize> = in_degree
31 .iter()
32 .enumerate()
33 .filter_map(|(i, &d)| if d == 0 { Some(i) } else { None })
34 .collect();
35
36 // Initialize the order vector.
37 let mut order = Vec::with_capacity(in_degree.len());
38 // While there are vertices to be visited ...
39 while let Some(i) = to_be_visited.pop_front() {
40 // Add the vertex to the order.
41 order.push(i);
42 // For each neighbor, reduce its in-degree.
43 self.children(&set![i]).into_iter().for_each(|y| {
44 // Decrement the in-degree of the child.
45 in_degree[y] -= 1;
46 // If the in-degree becomes 0, add it to the queue.
47 if in_degree[y] == 0 {
48 to_be_visited.push_back(y);
49 }
50 });
51 }
52
53 // Check if the order contains all vertices.
54 if in_degree.len() == order.len() {
55 Some(order)
56 } else {
57 None // The graph is not a DAG.
58 }
59 }
60}