implement constant Potts model modularity

raphtory/src/algorithms/community_detection/modularity.rs

@@ -409,3 +409,251 @@ impl ModularityFunction for ModularityUnDir {
         Box::new((0..self.partition.num_nodes()).map(VID))
     }
 }
+
+/// Constant Potts model modularity from https://arxiv.org/pdf/1104.3083
+pub struct ConstModularity {
+    resolution: f64,
+    partition: Partition,
+    adj: Vec<Vec<(VID, f64)>>,
+    self_loops: Vec<f64>,
+    adj_com: Vec<HashMap<ComID, f64>>,
+    n: Vec<i64>,
+    n_com: Vec<i64>,
+    m2: f64,
+    tol: f64,
+}
+
+impl ModularityFunction for ConstModularity {
+    fn new<'graph, G: GraphViewOps<'graph>>(
+        graph: G,
+        weight_prop: Option<&str>,
+        resolution: f64,
+        partition: Partition,
+        tol: f64,
+    ) -> Self {
+        let num_nodes = graph.count_nodes();
+        let n = vec![1; num_nodes];
+        let nodes = graph.nodes();
+        let local_id_map: HashMap<_, _> =
+            nodes.iter().enumerate().map(|(i, n)| (n, VID(i))).collect();
+        let adj: Vec<_> = nodes
+            .iter()
+            .map(|node| {
+                node.edges()
+                    .iter()
+                    .filter(|e| e.dst() != e.src())
+                    .map(|e| {
+                        let w = weight_prop
+                            .map(|w| e.properties().get(w).unwrap_f64())
+                            .unwrap_or(1.0);
+                        let dst_id = local_id_map[&e.nbr().cloned()];
+                        (dst_id, w)
+                    })
+                    .filter(|(_, w)| w >= &tol)
+                    .collect::<Vec<_>>()
+            })
+            .collect();
+        let self_loops: Vec<_> = graph
+            .nodes()
+            .iter()
+            .map(|node| {
+                graph
+                    .edge(node.node, node.node)
+                    .map(|e| {
+                        weight_prop
+                            .map(|w| e.properties().get(w).unwrap_f64())
+                            .unwrap_or(1.0)
+                    })
+                    .filter(|w| w >= &tol)
+                    .unwrap_or(0.0)
+            })
+            .collect();
+        let m2: f64 = adj
+            .iter()
+            .flat_map(|neighbours| neighbours.iter().map(|(_, w)| w))
+            .sum();
+        let adj_com: Vec<_> = adj
+            .iter()
+            .enumerate()
+            .map(|(index, neighbours)| {
+                let mut com_neighbours = HashMap::new();
+                for (n, w) in neighbours {
+                    com_neighbours
+                        .entry(partition.com(n))
+                        .and_modify(|old_w| *old_w += *w)
+                        .or_insert(*w);
+                }
+                if self_loops[index] != 0.0 {
+                    *com_neighbours
+                        .entry(partition.com(&VID(index)))
+                        .or_insert(0.0) += self_loops[index];
+                }
+                com_neighbours
+            })
+            .collect();
+
+        let n_com = partition.coms().map(|(_, com)| com.len() as i64).collect();
+        Self {
+            partition,
+            adj,
+            self_loops,
+            adj_com,
+            resolution,
+            n,
+            n_com,
+            m2,
+            tol,
+        }
+    }
+
+    fn move_delta(&self, node: &VID, new_com: ComID) -> f64 {
+        let old_com = self.partition.com(node);
+        if old_com == new_com {
+            0.0
+        } else {
+            let a = 2.0
+                * (self.adj_com[node.index()].get(&new_com).unwrap_or(&0.0)
+                    - self.adj_com[node.index()].get(&old_com).unwrap_or(&0.0)
+                    + self.self_loops[node.index()]);
+            let p = 2
+                * self.n[node.index()]
+                * (self.n_com[new_com.index()] - self.n_com[old_com.index()]);
+
+            (a - self.resolution * p as f64 / self.m2) / self.m2
+        }
+    }
+
+    fn move_node(&mut self, node: &VID, new_com: ComID) {
+        let old_com = self.partition.com(node);
+        if old_com != new_com {
+            let w_self = self.self_loops[node.index()];
+            match self.adj_com[node.index()]
+                .entry(old_com)
+                .and_modify(|v| *v -= w_self)
+            {
+                Entry::Occupied(v) => {
+                    if *v.get() < self.tol {
+                        v.remove();
+                    }
+                }
+                _ => {
+                    // should only be possible for small values due to tolerance above
+                    debug_assert!(w_self < self.tol)
+                }
+            }
+            if w_self != 0.0 {
+                *self.adj_com[node.index()].entry(new_com).or_insert(0.0) += w_self;
+            }
+
+            for (n, w) in &self.adj[node.index()] {
+                match self.adj_com[n.index()]
+                    .entry(old_com)
+                    .and_modify(|v| *v -= w)
+                {
+                    Entry::Occupied(v) => {
+                        if *v.get() < self.tol {
+                            v.remove();
+                        }
+                    }
+                    _ => {
+                        // should only be possible for small values due to tolerance above
+                        debug_assert!(*w < self.tol)
+                    }
+                }
+                match self.adj_com[node.index()]
+                    .entry(self.partition.com(n))
+                    .and_modify(|v| *v -= w)
+                {
+                    Entry::Occupied(v) => {
+                        if *v.get() < self.tol {
+                            v.remove();
+                        }
+                    }
+                    _ => {
+                        // should only be possible for small values due to tolerance above
+                        debug_assert!(*w < self.tol)
+                    }
+                }
+                *self.adj_com[n.index()].entry(new_com).or_insert(0.0) += w;
+                *self.adj_com[node.index()]
+                    .entry(self.partition.com(n))
+                    .or_insert(0.0) += w;
+            }
+            self.n_com[old_com.index()] -= self.n[node.index()];
+            self.n_com[new_com.index()] += self.n[node.index()];
+        }
+        self.partition.move_node(node, new_com);
+    }
+
+    fn candidate_moves(&self, node: &VID) -> Box<dyn Iterator<Item = ComID> + '_> {
+        Box::new(self.adj_com[node.index()].keys().copied())
+    }
+
+    fn aggregate(&mut self) -> Partition {
+        let old_partition = mem::take(&mut self.partition);
+        let (new_partition, new_to_old, old_to_new) = old_partition.compact();
+        let adj_com: Vec<_> = new_partition
+            .coms()
+            .map(|(_c_new, com)| {
+                let mut neighbours = HashMap::new();
+                for n in com {
+                    for (c_old, w) in &self.adj_com[n.index()] {
+                        *neighbours.entry(old_to_new[c_old]).or_insert(0.0) += w;
+                    }
+                }
+                neighbours
+            })
+            .collect();
+        let adj: Vec<_> = adj_com
+            .iter()
+            .enumerate()
+            .map(|(index, neighbours)| {
+                neighbours
+                    .iter()
+                    .filter(|(ComID(c), _)| c != &index)
+                    .map(|(ComID(index), w)| (VID(*index), *w))
+                    .collect::<Vec<_>>()
+            })
+            .collect();
+        let self_loops: Vec<_> = adj_com
+            .iter()
+            .enumerate()
+            .map(|(index, neighbours)| neighbours.get(&ComID(index)).copied().unwrap_or(0.0))
+            .collect();
+        let n: Vec<_> = new_to_old
+            .into_iter()
+            .map(|ComID(index)| self.n_com[index])
+            .collect();
+        let n_com = n.clone();
+        let partition = Partition::new_singletons(new_partition.num_coms());
+        self.adj = adj;
+        self.adj_com = adj_com;
+        self.self_loops = self_loops;
+        self.n = n;
+        self.n_com = n_com;
+        self.partition = partition;
+        new_partition
+    }
+
+    fn value(&self) -> f64 {
+        let e: f64 = self
+            .partition
+            .coms()
+            .map(|(cid, com)| {
+                com.iter()
+                    .flat_map(|n| self.adj_com[n.index()].get(&cid))
+                    .sum::<f64>()
+            })
+            .sum();
+        let k: i64 = self.n_com.iter().map(|n| n.pow(2)).sum();
+        e / self.m2 - k as f64 / self.m2
+    }
+
+    fn partition(&self) -> &Partition {
+        &self.partition
+    }
+
+    fn nodes(&self) -> Box<dyn Iterator<Item = VID>> {
+        Box::new((0..self.partition.num_nodes()).map(VID))
+    }
+}

raphtory/src/python/packages/algorithms.rs

@@ -12,7 +12,8 @@ use crate::{
         },
         community_detection::{
             label_propagation::label_propagation as label_propagation_rs,
-            louvain::louvain as louvain_rs, modularity::ModularityUnDir,
+            louvain::louvain as louvain_rs,
+            modularity::{ConstModularity, ModularityUnDir},
         },
         components,
         cores::k_core::k_core_set,
@@ -73,7 +74,11 @@ use crate::{
 };
 #[cfg(feature = "storage")]
 use pometry_storage::algorithms::connected_components::connected_components as connected_components_rs;
-use pyo3::{prelude::*, types::PyList};
+use pyo3::{
+    exceptions::{PyTypeError, PyValueError},
+    prelude::*,
+    types::PyList,
+};
 use rand::{prelude::StdRng, SeedableRng};
 use raphtory_api::core::Direction;
 use raphtory_storage::core_ops::CoreGraphOps;
@@ -842,25 +847,31 @@ pub fn temporal_SEIR(
     )
 }
 
-/// Louvain algorithm for community detection
+/// Louvain algorithm for community detection with configuration model
 ///
 /// Arguments:
 ///     graph (GraphView): the graph view
 ///     resolution (float): the resolution parameter for modularity. Defaults to 1.0.
 ///     weight_prop (str | None): the edge property to use for weights (has to be float)
 ///     tol (None | float): the floating point tolerance for deciding if improvements are significant (default: 1e-8)
+///     modularity (Literal("configuration", "constant")): the modularity function to use. Default to "configuration".
 ///
 /// Returns:
 ///     NodeStateUsize: Mapping of nodes to their community assignment
 #[pyfunction]
-#[pyo3[signature=(graph, resolution=1.0, weight_prop=None, tol=None)]]
+#[pyo3[signature=(graph, resolution=1.0, weight_prop=None, tol=None, modularity="configuration")]]
 pub fn louvain(
     graph: &PyGraphView,
     resolution: f64,
     weight_prop: Option<&str>,
     tol: Option<f64>,
-) -> NodeState<'static, usize, DynamicGraph> {
-    louvain_rs::<ModularityUnDir, _>(&graph.graph, resolution, weight_prop, tol)
+    modularity: &str,
+) -> PyResult<NodeState<'static, usize, DynamicGraph>> {
+    match modularity {
+        "configuration" => Ok(louvain_rs::<ModularityUnDir, _>(&graph.graph, resolution, weight_prop, tol)),
+        "constant" => Ok(louvain_rs::<ConstModularity, _>(&graph.graph, resolution, weight_prop, tol)),
+        other => Err(PyValueError::new_err(format!("'{other}' not a valid value for modularity, should be one of 'configuration' or 'constant'")))
+    }
 }
 
 /// Fruchterman Reingold layout algorithm

raphtory/tests/algo_tests/community_detection.rs

@@ -57,6 +57,7 @@ fn lpa_test() {
 }
 
 use proptest::prelude::*;
+use raphtory::algorithms::community_detection::modularity::ConstModularity;
 
 #[test]
 fn test_louvain() {
@@ -90,6 +91,12 @@ fn test_all_nodes_assigned_inner(edges: Vec<(u64, u64, f64)>) {
             .nodes()
             .iter()
             .all(|n| result.get_by_node(n).is_some()));
+
+        let result = louvain::<ConstModularity, _>(graph, 1.0, Some("weight"), None);
+        assert!(graph
+            .nodes()
+            .iter()
+            .all(|n| result.get_by_node(n).is_some()));
     });
 }
 
@@ -106,6 +113,12 @@ fn test_all_nodes_assigned_inner_unweighted(edges: Vec<(u64, u64)>) {
             .nodes()
             .iter()
             .all(|n| result.get_by_node(n).is_some()));
+
+        let result = louvain::<ConstModularity, _>(graph, 1.0, None, None);
+        assert!(graph
+            .nodes()
+            .iter()
+            .all(|n| result.get_by_node(n).is_some()));
     });
 }