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}