New subgraph implementation

Author: frederikfabriciusbjerre

R/extendr-wrappers.R

@@ -46,6 +46,8 @@ markov_blanket_of_ptr <- function(g, idxs) .Call(wrap__markov_blanket_of_ptr, g,
 
 exogenous_nodes_of_ptr <- function(g, undirected_as_parents) .Call(wrap__exogenous_nodes_of_ptr, g, undirected_as_parents)
 
+induced_subgraph_ptr <- function(g, keep) .Call(wrap__induced_subgraph_ptr, g, keep)
+
 is_simple_ptr <- function(g) .Call(wrap__is_simple_ptr, g)
 
 graph_class_ptr <- function(g) .Call(wrap__graph_class_ptr, g)

R/queries.R

@@ -615,52 +615,70 @@ subgraph <- function(cg, nodes = NULL, index = NULL) {
     stop("Supply one of `nodes` or `index`.", call. = FALSE)
   }
 
+  # resolve -> keep_ids0 (0-based), keep_names, preserving order
   if (index_supplied) {
+    if (!is.numeric(index) || anyNA(index)) {
+      stop("`index` must be numeric without NA.", call. = FALSE)
+    }
     idx1 <- as.integer(index)
     if (any(idx1 < 1L) || any(idx1 > nrow(cg@nodes))) {
       stop("`index` out of range (1..n).", call. = FALSE)
     }
+    keep_ids0 <- idx1 - 1L
     keep_names <- cg@nodes$name[idx1]
   } else {
     if (!is.character(nodes)) {
       stop("`nodes` must be a character vector.", call. = FALSE)
     }
+    if (anyNA(nodes)) {
+      stop("`nodes` contains NA.", call. = FALSE)
+    }
     missing <- setdiff(nodes, cg@nodes$name)
     if (length(missing)) {
       stop("Unknown node(s): ", paste(missing, collapse = ", "), call. = FALSE)
     }
     keep_names <- nodes
+    keep_ids0 <- vapply(nodes, cg@name_index_map$get, integer(1))
   }
 
-  if (any(duplicated(keep_names))) {
-    dups <- unique(keep_names[duplicated(keep_names)])
-    stop("`nodes`/`index` contains duplicates: ",
-      paste(dups, collapse = ", "),
-      call. = FALSE
-    )
+  # duplicates are an error
+  if (any(duplicated(keep_ids0))) {
+    dups <- unique(keep_names[duplicated(keep_ids0) | duplicated(keep_ids0, fromLast = TRUE)])
+    stop("`nodes`/`index` contains duplicates: ", paste(dups, collapse = ", "), call. = FALSE)
   }
 
-  # filter edges to the kept nodes; keep constructor’s sort
+  # call Rust (always reindexed)
+  ptr_sub <- induced_subgraph_ptr(cg@ptr, as.integer(keep_ids0))
+
+  # nodes table in input order
+  nodes_sub <- tibble::tibble(name = keep_names)
+
+  # filter edges to kept names and sort like constructor
   keep_set <- fastmap::fastmap()
   for (nm in keep_names) keep_set$set(nm, TRUE)
   edges_sub <- cg@edges |>
     dplyr::filter(keep_set$has(.data$from) & keep_set$has(.data$to)) |>
     dplyr::arrange(.data$from, .data$to, .data$edge)
 
-  # rebuild via constructor: declared nodes contain all edge nodes
-  caugi_graph(
-    from   = edges_sub$from,
-    edge   = edges_sub$edge,
-    to     = edges_sub$to,
-    nodes  = keep_names,
+  # rebuild name_index_map
+  name_index_map_sub <- fastmap::fastmap()
+  for (i in seq_len(nrow(nodes_sub))) name_index_map_sub$set(nodes_sub$name[i], i - 1L)
+
+  state_sub <- .cg_state(
+    nodes = nodes_sub,
+    edges = edges_sub,
+    ptr = ptr_sub,
+    built = TRUE,
     simple = cg@simple,
-    build  = TRUE,
-    class  = cg@graph_class
+    class = cg@graph_class,
+    name_index_map = name_index_map_sub
   )
+  caugi_graph(state = state_sub)
 }
 
 
 
+
 # ──────────────────────────────────────────────────────────────────────────────
 # ──────────────────────────── Relations helpers ───────────────────────────────
 # ──────────────────────────────────────────────────────────────────────────────

src/rust/src/graph/mod.rs

@@ -83,6 +83,83 @@ impl CaugiGraph {
     }
 }
 
+impl CaugiGraph {
+    /// Node-induced subgraph on `keep` (new ids are 0..k-1, in the SAME order as `keep`).
+    /// Returns: (new_core, new_to_old, old_to_new).
+    pub fn induced_subgraph(
+        &self,
+        keep: &[u32],
+    ) -> Result<(CaugiGraph, Vec<u32>, Vec<u32>), String> {
+        let n = self.n() as usize;
+
+        // validate + deduplicate while preserving order
+        let mut seen = vec![false; n];
+        let mut new_to_old: Vec<u32> = Vec::with_capacity(keep.len());
+        for &u in keep {
+            if (u as usize) >= n {
+                return Err("node id out of range".into());
+            }
+            if std::mem::replace(&mut seen[u as usize], true) {
+                return Err("duplicate node id in `keep`".into());
+            }
+            new_to_old.push(u);
+        }
+
+        // old -> new map
+        let mut old_to_new = vec![u32::MAX; n];
+        for (new, &old) in new_to_old.iter().enumerate() {
+            old_to_new[old as usize] = new as u32;
+        }
+
+        // row counts
+        let k = new_to_old.len();
+        let mut row_index: Vec<u32> = Vec::with_capacity(k + 1);
+        row_index.push(0);
+        for &old_u in &new_to_old {
+            let mut cnt = 0u32;
+            for kk in self.row_range(old_u) {
+                let ov = self.col_index[kk] as usize;
+                if ov < n && old_to_new[ov] != u32::MAX {
+                    cnt += 1;
+                }
+            }
+            row_index.push(row_index.last().unwrap() + cnt);
+        }
+
+        // allocate + scatter
+        let nnz = *row_index.last().unwrap() as usize;
+        let mut col_index = vec![0u32; nnz];
+        let mut etype = vec![0u8; nnz];
+        let mut side = vec![0u8; nnz];
+        let mut cur = row_index[..k].to_vec();
+
+        for (new_u, &old_u) in new_to_old.iter().enumerate() {
+            for kk in self.row_range(old_u) {
+                let ov = self.col_index[kk] as usize;
+                let nv = old_to_new[ov];
+                if nv == u32::MAX {
+                    continue;
+                }
+                let p = cur[new_u] as usize;
+                col_index[p] = nv;
+                etype[p] = self.etype[kk];
+                side[p] = self.side[kk];
+                cur[new_u] += 1;
+            }
+        }
+
+        let out = CaugiGraph::from_csr(
+            row_index,
+            col_index,
+            etype,
+            side,
+            self.simple,
+            self.registry.clone(),
+        )?;
+        Ok((out, new_to_old, old_to_new))
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

src/rust/src/graph/view.rs

@@ -131,6 +131,25 @@ impl GraphView {
             _ => Err("all_backdoor_sets is only defined for DAGs".into()),
         }
     }
+
+    pub fn induced_subgraph(
+        &self,
+        keep: &[u32],
+    ) -> Result<GraphView, String> {
+        let (core2, _new_to_old, _old_to_new) = self.core().induced_subgraph(keep)?;
+        let gv = match self {
+            GraphView::Dag(_) => {
+                let d = super::dag::Dag::new(std::sync::Arc::new(core2))?;
+                GraphView::Dag(std::sync::Arc::new(d))
+            }
+            GraphView::Pdag(_) => {
+                let p = super::pdag::Pdag::new(std::sync::Arc::new(core2))?;
+                GraphView::Pdag(std::sync::Arc::new(p))
+            }
+            GraphView::Raw(_) => GraphView::Raw(std::sync::Arc::new(core2)),
+        };
+        Ok(gv)
+    }
 }
 
 #[cfg(test)]

src/rust/src/lib.rs

@@ -611,6 +611,28 @@ fn all_backdoor_sets_ptr(
     extendr_api::prelude::List::from_values(robjs).into_robj()
 }
 
+// ── Subgraph ────────────────────────────────────────────────────────────────
+
+#[extendr]
+fn induced_subgraph_ptr(g: ExternalPtr<GraphView>, keep: Integers) -> Robj {
+    let mut ks: Vec<u32> = Vec::with_capacity(keep.len());
+    for ri in keep.iter() {
+        let u = rint_to_u32(ri, "keep");
+        if u >= g.as_ref().n() {
+            throw_r_error(format!("node id {} out of bounds", u));
+        }
+        ks.push(u);
+    }
+
+    let  sub= g
+        .as_ref()
+        .induced_subgraph(&ks)
+        .unwrap_or_else(|e| throw_r_error(e));
+
+    let sub_ptr = ExternalPtr::new(sub);
+    sub_ptr.into_robj()
+}
+
 extendr_module! {
     mod caugi;
     // registry
@@ -638,6 +660,7 @@ extendr_module! {
     fn descendants_of_ptr;
     fn markov_blanket_of_ptr;
     fn exogenous_nodes_of_ptr;
+    fn induced_subgraph_ptr;
 
     // graph properties
     fn is_simple_ptr;