perf: rewrite not_m_separated_for_all_subsets() in Rustq

Author: jolars

R/extendr-wrappers.R

@@ -134,6 +134,8 @@ rs_minimal_d_separator <- function(session, xs, ys, include, restrict) .Call(wra
 
 rs_m_separated <- function(session, xs, ys, z) .Call(wrap__rs_m_separated, session, xs, ys, z)
 
+rs_not_m_separated_for_all_subsets <- function(session, node_a, node_b, other_nodes, cond_vars) .Call(wrap__rs_not_m_separated_for_all_subsets, session, node_a, node_b, other_nodes, cond_vars)
+
 rs_adjustment_set_parents <- function(session, xs, ys) .Call(wrap__rs_adjustment_set_parents, session, xs, ys)
 
 rs_adjustment_set_backdoor <- function(session, xs, ys) .Call(wrap__rs_adjustment_set_backdoor, session, xs, ys)

R/operations.R

@@ -643,45 +643,25 @@ condition_marginalize <- function(cg, cond_vars = NULL, marg_vars = NULL) {
   other_nodes,
   cond_vars
 ) {
-  n_other <- length(other_nodes)
-
-  # Generate all subsets of other_nodes
-  subsets <- if (n_other == 0L) {
-    list(NULL)
+  node_idx <- .nodes_to_indices(cg, c(node_a, node_b))
+  other_idx <- if (length(other_nodes) == 0L) {
+    integer(0)
   } else {
-    # Build subsets from largest to smallest (often finds separation faster)
-    c(
-      list(other_nodes),
-      if (n_other > 1L) {
-        unlist(
-          lapply(
-            seq_len(n_other - 1L),
-            function(k) combn(other_nodes, n_other - k, simplify = FALSE)
-          ),
-          recursive = FALSE
-        )
-      },
-      list(NULL)
-    )
+    .nodes_to_indices(cg, other_nodes)
   }
-
-  # Check each conditioning set
-
-  for (subset in subsets) {
-    conditioning_set <- c(cond_vars, subset)
-    if (length(conditioning_set) == 0L) {
-      conditioning_set <- NULL
-    }
-
-    if (m_separated(cg, X = node_a, Y = node_b, Z = conditioning_set)) {
-      # Found a set that m-separates them: no edge needed
-      return(FALSE)
-    }
+  cond_idx <- if (length(cond_vars) == 0L) {
+    integer(0)
+  } else {
+    .nodes_to_indices(cg, cond_vars)
   }
 
-  # Not m-separated for any conditioning set: edge is required
-
-  TRUE
+  rs_not_m_separated_for_all_subsets(
+    cg@session,
+    node_idx[[1]],
+    node_idx[[2]],
+    other_idx,
+    cond_idx
+  )
 }
 
 #' @title

src/rust/src/lib.rs

@@ -1898,6 +1898,90 @@ fn rs_m_separated(
         .unwrap_or_else(|e| throw_r_error(e))
 }
 
+#[extendr]
+fn rs_not_m_separated_for_all_subsets(
+    mut session: ExternalPtr<GraphSession>,
+    node_a: i32,
+    node_b: i32,
+    other_nodes: Integers,
+    cond_vars: Integers,
+) -> bool {
+    let a = rint_to_u32(Rint::from(node_a), "node_a");
+    let b = rint_to_u32(Rint::from(node_b), "node_b");
+
+    if a >= session.as_ref().n() {
+        throw_r_error(format!("Index {} is out of bounds", a));
+    }
+    if b >= session.as_ref().n() {
+        throw_r_error(format!("Index {} is out of bounds", b));
+    }
+
+    let mut z_base: Vec<u32> = cond_vars
+        .iter()
+        .map(|ri| rint_to_u32(ri, "cond_vars"))
+        .collect();
+    let other_u: Vec<u32> = other_nodes
+        .iter()
+        .map(|ri| rint_to_u32(ri, "other_nodes"))
+        .collect();
+
+    for &i in &z_base {
+        if i >= session.as_ref().n() {
+            throw_r_error(format!("Index {} is out of bounds", i));
+        }
+    }
+    for &i in &other_u {
+        if i >= session.as_ref().n() {
+            throw_r_error(format!("Index {} is out of bounds", i));
+        }
+    }
+
+    if session
+        .as_mut()
+        .m_separated(&[a], &[b], &z_base)
+        .unwrap_or_else(|e| throw_r_error(e))
+    {
+        return false;
+    }
+
+    let m = other_u.len();
+    if m == 0 {
+        return true;
+    }
+
+    for k in (1..=m).rev() {
+        let mut idx: Vec<usize> = (0..k).collect();
+        loop {
+            z_base.truncate(cond_vars.len());
+            for &ii in &idx {
+                z_base.push(other_u[ii]);
+            }
+
+            if session
+                .as_mut()
+                .m_separated(&[a], &[b], &z_base)
+                .unwrap_or_else(|e| throw_r_error(e))
+            {
+                return false;
+            }
+
+            let mut i = k;
+            while i > 0 && idx[i - 1] == i - 1 + (m - k) {
+                i -= 1;
+            }
+            if i == 0 {
+                break;
+            }
+            idx[i - 1] += 1;
+            for j in i..k {
+                idx[j] = idx[j - 1] + 1;
+            }
+        }
+    }
+
+    true
+}
+
 #[extendr]
 fn rs_adjustment_set_parents(
     mut session: ExternalPtr<GraphSession>,
@@ -2086,6 +2170,7 @@ extendr_module! {
     fn rs_d_separated;
     fn rs_minimal_d_separator;
     fn rs_m_separated;
+    fn rs_not_m_separated_for_all_subsets;
     fn rs_adjustment_set_parents;
     fn rs_adjustment_set_backdoor;
     fn rs_adjustment_set_optimal;

tests/testthat/test-operations.R

@@ -970,15 +970,24 @@ test_that("condition_marginalize and helper branches are covered", {
     "must be the same length"
   )
 
-  expect_type(
+  expect_true(
     caugi:::.not_m_separated_for_all_subsets(
       cg = cg2,
       node_a = "A",
       node_b = "C",
       other_nodes = character(0),
       cond_vars = character(0)
-    ),
-    "logical"
+    )
+  )
+
+  expect_false(
+    caugi:::.not_m_separated_for_all_subsets(
+      cg = cg2,
+      node_a = "A",
+      node_b = "C",
+      other_nodes = "B",
+      cond_vars = character(0)
+    )
   )
 
   edge_rev <- caugi:::.edge_type_from_anteriors(