raphtory/algorithms/cores/
k_core.rs

1use crate::{
2    core::{entities::VID, state::compute_state::ComputeStateVec},
3    db::{
4        api::view::{NodeViewOps, StaticGraphViewOps},
5        graph::views::node_subgraph::NodeSubgraph,
6        task::{
7            context::Context,
8            node::eval_node::EvalNodeView,
9            task::{ATask, Job, Step},
10            task_runner::TaskRunner,
11        },
12    },
13};
14use std::collections::HashSet;
15
16#[derive(Clone, Debug)]
17struct KCoreState {
18    alive: bool,
19}
20
21impl Default for KCoreState {
22    fn default() -> Self {
23        Self { alive: true }
24    }
25}
26
27/// Determines which nodes are in the k-core for a given value of k
28///
29/// # Arguments
30///
31/// - `g` - A reference to the graph
32/// - `k` - Value of k such that the returned nodes have degree > k (recursively)
33/// - `iter_count` - The number of iterations to run
34/// - `threads` - number of threads to run on
35///
36/// # Returns
37///
38/// A hash set of nodes in the k core
39///
40pub fn k_core_set<G>(g: &G, k: usize, iter_count: usize, threads: Option<usize>) -> HashSet<VID>
41where
42    G: StaticGraphViewOps,
43{
44    let ctx: Context<G, ComputeStateVec> = g.into();
45
46    let step1 = ATask::new(move |vv| {
47        let deg = vv.degree();
48        let state: &mut KCoreState = vv.get_mut();
49        state.alive = deg >= k;
50        Step::Continue
51    });
52
53    let step2 = ATask::new(move |vv: &mut EvalNodeView<G, KCoreState>| {
54        let prev: bool = vv.prev().alive;
55        if prev {
56            let current = vv
57                .neighbours()
58                .into_iter()
59                .filter(|n| n.prev().alive)
60                .count()
61                >= k;
62            let state: &mut KCoreState = vv.get_mut();
63            if current != prev {
64                state.alive = current;
65                Step::Continue
66            } else {
67                Step::Done
68            }
69        } else {
70            Step::Done
71        }
72    });
73
74    let mut runner: TaskRunner<G, _> = TaskRunner::new(ctx);
75
76    runner.run(
77        vec![Job::new(step1)],
78        vec![Job::read_only(step2)],
79        None,
80        |_, _, _, local| {
81            g.nodes()
82                .iter()
83                .filter(|node| local[node.node.0].alive)
84                .map(|node| node.node)
85                .collect()
86        },
87        threads,
88        iter_count,
89        None,
90        None,
91    )
92}
93
94pub fn k_core<G>(g: &G, k: usize, iter_count: usize, threads: Option<usize>) -> NodeSubgraph<G>
95where
96    G: StaticGraphViewOps,
97{
98    let v_set = k_core_set(g, k, iter_count, threads);
99    g.subgraph(v_set)
100}
101
102#[cfg(test)]
103mod k_core_test {
104    use crate::{algorithms::cores::k_core::k_core_set, prelude::*, test_storage};
105    use std::collections::HashSet;
106
107    #[test]
108    fn k_core_2() {
109        let graph = Graph::new();
110
111        let edges = vec![
112            (1, 2, 1),
113            (1, 3, 2),
114            (1, 4, 3),
115            (3, 1, 4),
116            (3, 4, 5),
117            (3, 5, 6),
118            (4, 5, 7),
119            (5, 6, 8),
120            (5, 8, 9),
121            (7, 5, 10),
122            (8, 5, 11),
123            (1, 9, 12),
124            (9, 1, 13),
125            (6, 3, 14),
126            (4, 8, 15),
127            (8, 3, 16),
128            (5, 10, 17),
129            (10, 5, 18),
130            (10, 8, 19),
131            (1, 11, 20),
132            (11, 1, 21),
133            (9, 11, 22),
134            (11, 9, 23),
135        ];
136
137        for (src, dst, ts) in edges {
138            graph.add_edge(ts, src, dst, NO_PROPS, None).unwrap();
139        }
140
141        test_storage!(&graph, |graph| {
142            let result = k_core_set(graph, 2, usize::MAX, None);
143            let subgraph = graph.subgraph(result.clone());
144            let actual = vec!["1", "3", "4", "5", "6", "8", "9", "10", "11"]
145                .into_iter()
146                .map(|k| k.to_string())
147                .collect::<HashSet<String>>();
148
149            assert_eq!(actual, subgraph.nodes().name().collect::<HashSet<String>>());
150        });
151    }
152}
raphtory/algorithms/cores/k_core.rs

raphtory/algorithms/cores/
k_core.rs
