1use crate::admg::Admg;
13use crate::dsep::DSeparationWorkspace;
14use crate::error::GraphError;
15use crate::pag::Pag;
16use crate::types::{DenseNodeId, Endpoint};
17use crate::workspace::GraphWorkspace;
18
19#[derive(Clone, Debug)]
21pub struct PagCompletion {
22 pub graph: Pag,
24 pub index: usize,
26}
27
28#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
30pub struct CompletionValidationReport {
31 pub assignments_examined: u64,
33 pub rejected_non_ancestral: u64,
35 pub rejected_nonmaximal: u64,
37 pub rejected_local_incompatible: u64,
39 pub ambiguous_global_class: bool,
41 pub equivalence_audit_skipped: bool,
47 pub represented_completions: usize,
49}
50
51#[derive(Clone, Debug)]
53pub struct CompletionSampler {
54 completions: Vec<Pag>,
55 max_completions: usize,
56 next_index: usize,
57 n_circle_sites: usize,
58 report: CompletionValidationReport,
59}
60
61const MAX_AUDITED_CIRCLE_SITES: usize = 16;
62const MAX_EQUIVALENCE_NODES: usize = 12;
63const MAX_EQUIVALENCE_QUERIES: u128 = 2_000_000;
64
65impl CompletionSampler {
66 #[allow(clippy::needless_pass_by_value)] pub fn new(pag: Pag, max_completions: usize) -> Result<Self, GraphError> {
76 let mut sites = Vec::new();
77 let n = pag.node_count();
78 for i in 0..n {
79 let a = DenseNodeId::try_from_usize(i)?;
80 for (b, at_a, at_b) in pag.neighbors(a) {
81 if b.raw() < a.raw() {
82 continue;
83 }
84 if matches!(at_a, Endpoint::Circle) {
85 sites.push((a, b, true));
86 }
87 if matches!(at_b, Endpoint::Circle) {
88 sites.push((a, b, false));
89 }
90 }
91 }
92 if sites.len() > MAX_AUDITED_CIRCLE_SITES {
93 return Err(GraphError::InvalidEndpoints {
94 message: "PAG completion audit supports at most 16 circle endpoints",
95 });
96 }
97 let total = 1u64 << sites.len();
98 let audit_class = sites.is_empty() || {
105 let n = pag.node_count() as u128;
106 let conditioning_sets = 1u128 << pag.node_count().saturating_sub(2);
107 let queries = u128::from(total)
108 .saturating_mul(n.saturating_mul(n.saturating_sub(1)) / 2)
109 .saturating_mul(conditioning_sets);
110 pag.node_count() <= MAX_EQUIVALENCE_NODES && queries <= MAX_EQUIVALENCE_QUERIES
111 };
112 let mut report = CompletionValidationReport {
113 assignments_examined: total,
114 equivalence_audit_skipped: !audit_class,
115 ..Default::default()
116 };
117 let mut completions = Vec::new();
118 let mut represented_completions = 0usize;
119 let mut reference_signature = None;
120 for mask in 0..total {
121 let Some(candidate) = orient_assignment(&pag, &sites, mask) else {
122 report.rejected_non_ancestral += 1;
123 continue;
124 };
125 if !is_maximal_ancestral_graph(&candidate) {
126 report.rejected_nonmaximal += 1;
127 continue;
128 }
129 if !preserves_unshielded_colliders(&pag, &candidate) {
130 report.rejected_local_incompatible += 1;
131 continue;
132 }
133 if audit_class {
134 let signature =
135 if sites.is_empty() { Vec::new() } else { m_separation_signature(&candidate) };
136 if let Some(reference) = &reference_signature {
137 if reference != &signature {
138 report.ambiguous_global_class = true;
139 continue;
140 }
141 } else {
142 reference_signature = Some(signature);
143 }
144 }
145 represented_completions += 1;
146 if completions.len() < max_completions {
147 completions.push(candidate);
148 }
149 }
150 if report.ambiguous_global_class {
151 completions.clear();
154 }
155 report.represented_completions =
156 if report.ambiguous_global_class { 0 } else { represented_completions };
157 Ok(Self {
158 completions,
159 max_completions,
160 next_index: 0,
161 n_circle_sites: sites.len(),
162 report,
163 })
164 }
165
166 #[must_use]
168 pub fn max_completions(&self) -> usize {
169 self.max_completions
170 }
171
172 #[must_use]
174 pub fn n_circle_sites(&self) -> usize {
175 self.n_circle_sites
176 }
177
178 #[must_use]
180 pub const fn validation_report(&self) -> CompletionValidationReport {
181 self.report
182 }
183
184 #[must_use]
191 pub const fn class_audit_incomplete(&self) -> bool {
192 self.report.equivalence_audit_skipped
193 }
194
195 #[must_use]
202 pub fn hit_cap(&self) -> bool {
203 self.report.represented_completions > self.max_completions
204 }
205}
206
207fn orient_assignment(
208 base: &Pag,
209 sites: &[(DenseNodeId, DenseNodeId, bool)],
210 mask: u64,
211) -> Option<Pag> {
212 let mut graph = base.clone();
213 for (i, &(a, b, at_a_circle)) in sites.iter().enumerate() {
214 let endpoint = if ((mask >> i) & 1) == 1 { Endpoint::Arrow } else { Endpoint::Tail };
215 let edge = graph.edge_between(a, b)?;
216 let marks = if at_a_circle { (endpoint, edge.at_b) } else { (edge.at_a, endpoint) };
217 graph.set_marks(a, b, marks.0, marks.1).ok()?;
218 }
219 is_ancestral_orientation(&graph).then_some(graph)
220}
221
222#[must_use]
225pub fn is_mag_completion(g: &Pag) -> bool {
226 is_maximal_ancestral_graph(g)
227}
228
229fn is_ancestral_orientation(g: &Pag) -> bool {
230 let n = g.node_count();
231 let mut ws = GraphWorkspace::default();
232 for i in 0..n {
233 let a = DenseNodeId::try_from_usize(i).expect("node fit");
234 for (b, at_a, at_b) in g.neighbors(a) {
235 if b.raw() < a.raw() {
236 continue;
237 }
238 if matches!(at_a, Endpoint::Circle | Endpoint::Conflict)
239 || matches!(at_b, Endpoint::Circle | Endpoint::Conflict)
240 {
241 return false;
242 }
243 if matches!((at_a, at_b), (Endpoint::Tail, Endpoint::Tail)) {
245 return false;
246 }
247 if matches!((at_a, at_b), (Endpoint::Arrow, Endpoint::Arrow)) {
248 if g.reaches_directed_with(&mut ws, a, b) || g.reaches_directed_with(&mut ws, b, a)
250 {
251 return false;
252 }
253 }
254 }
255 }
256 true
257}
258
259#[must_use]
261pub fn is_maximal_ancestral_graph(g: &Pag) -> bool {
262 if !is_ancestral_orientation(g) {
263 return false;
264 }
265 let admg = as_admg(g);
269 let n = g.node_count();
270 let mut graph_ws = GraphWorkspace::default();
271 let mut sep_ws = DSeparationWorkspace::default();
272 for i in 0..n {
273 let x = DenseNodeId::from_raw(u32::try_from(i).expect("node fit"));
274 for j in (i + 1)..n {
275 let y = DenseNodeId::from_raw(u32::try_from(j).expect("node fit"));
276 if g.has_edge(x, y) {
277 continue;
278 }
279 let separating: Vec<_> = (0..n)
280 .map(|k| DenseNodeId::from_raw(u32::try_from(k).expect("node fit")))
281 .filter(|&node| {
282 node != x
283 && node != y
284 && (g.reaches_directed_with(&mut graph_ws, node, x)
285 || g.reaches_directed_with(&mut graph_ws, node, y))
286 })
287 .collect();
288 if !admg.is_m_separated(x, y, &separating, &mut sep_ws).expect("known nodes") {
289 return false;
290 }
291 }
292 }
293 true
294}
295
296fn preserves_unshielded_colliders(pag: &Pag, mag: &Pag) -> bool {
297 let n = pag.node_count();
298 for middle_i in 0..n {
299 let middle = DenseNodeId::from_raw(u32::try_from(middle_i).expect("node fit"));
300 let neighbors: Vec<_> = pag.neighbors(middle).map(|(node, mark, _)| (node, mark)).collect();
301 for i in 0..neighbors.len() {
302 for j in (i + 1)..neighbors.len() {
303 let (left, pag_left_mark) = neighbors[i];
304 let (right, pag_right_mark) = neighbors[j];
305 if pag.has_edge(left, right) {
306 continue;
307 }
308 let mag_left_mark = mag
309 .neighbors(middle)
310 .find(|(node, _, _)| *node == left)
311 .map(|(_, mark, _)| mark)
312 .expect("same skeleton");
313 let mag_right_mark = mag
314 .neighbors(middle)
315 .find(|(node, _, _)| *node == right)
316 .map(|(_, mark, _)| mark)
317 .expect("same skeleton");
318 let pag_collider = matches!(pag_left_mark, Endpoint::Arrow)
319 && matches!(pag_right_mark, Endpoint::Arrow);
320 let mag_collider = matches!(mag_left_mark, Endpoint::Arrow)
321 && matches!(mag_right_mark, Endpoint::Arrow);
322 if pag_collider != mag_collider {
323 return false;
324 }
325 }
326 }
327 }
328 true
329}
330
331fn as_admg(g: &Pag) -> Admg {
332 let mut admg = Admg::with_variables(u32::try_from(g.node_count()).expect("node count fits"));
333 for i in 0..g.node_count() {
334 let a = DenseNodeId::from_raw(u32::try_from(i).expect("node fit"));
335 for (b, at_a, at_b) in g.neighbors(a) {
336 if b.raw() < a.raw() {
337 continue;
338 }
339 match (at_a, at_b) {
340 (Endpoint::Tail, Endpoint::Arrow) => {
341 admg.insert_directed(a, b).expect("validated MAG");
342 }
343 (Endpoint::Arrow, Endpoint::Tail) => {
344 admg.insert_directed(b, a).expect("validated MAG");
345 }
346 (Endpoint::Arrow, Endpoint::Arrow) => {
347 admg.insert_bidirected(a, b).expect("validated MAG");
348 }
349 _ => unreachable!("validated directed MAG marks"),
350 }
351 }
352 }
353 admg
354}
355
356fn m_separation_signature(g: &Pag) -> Vec<bool> {
357 let admg = as_admg(g);
358 let n = g.node_count();
359 let mut signature = Vec::new();
360 let mut ws = DSeparationWorkspace::default();
361 for i in 0..n {
362 for j in (i + 1)..n {
363 let others: Vec<_> = (0..n).filter(|&k| k != i && k != j).collect();
364 for mask in 0..(1usize << others.len()) {
365 let z: Vec<_> = others
366 .iter()
367 .enumerate()
368 .filter(|(bit, _)| ((mask >> bit) & 1) == 1)
369 .map(|(_, &k)| DenseNodeId::from_raw(u32::try_from(k).expect("node fit")))
370 .collect();
371 signature.push(
372 admg.is_m_separated(
373 DenseNodeId::from_raw(u32::try_from(i).expect("node fit")),
374 DenseNodeId::from_raw(u32::try_from(j).expect("node fit")),
375 &z,
376 &mut ws,
377 )
378 .expect("known nodes"),
379 );
380 }
381 }
382 }
383 signature
384}
385
386impl Iterator for CompletionSampler {
387 type Item = PagCompletion;
388
389 fn next(&mut self) -> Option<Self::Item> {
390 if self.next_index >= self.max_completions {
391 return None;
392 }
393 let graph = self.completions.get(self.next_index)?.clone();
394 let index = self.next_index;
395 self.next_index += 1;
396 Some(PagCompletion { graph, index })
397 }
398}
399
400#[cfg(test)]
401mod tests {
402 use super::*;
403 use crate::pag::Pag;
404
405 #[test]
406 fn respects_max_completions_bound() {
407 let mut pag = Pag::with_variables(2);
408 pag.insert_circle_circle(DenseNodeId::from_raw(0), DenseNodeId::from_raw(1)).unwrap();
409 let sampler = CompletionSampler::new(pag, 2).unwrap();
410 assert_eq!(sampler.n_circle_sites(), 2);
411 let collected: Vec<_> = sampler.collect();
412 assert!(collected.len() <= 2);
413 assert!(!collected.is_empty());
414 for c in &collected {
415 assert!(is_mag_completion(&c.graph));
416 let e =
417 c.graph.edge_between(DenseNodeId::from_raw(0), DenseNodeId::from_raw(1)).unwrap();
418 assert!(!matches!(e.at_a, Endpoint::Circle));
419 assert!(!matches!(e.at_b, Endpoint::Circle));
420 assert!(!matches!((e.at_a, e.at_b), (Endpoint::Tail, Endpoint::Tail)));
422 }
423 }
424
425 #[test]
426 fn no_circle_yields_single_completion() {
427 let mut pag = Pag::with_variables(2);
428 pag.insert_directed(DenseNodeId::from_raw(0), DenseNodeId::from_raw(1)).unwrap();
429 let collected: Vec<_> = CompletionSampler::new(pag, 10).unwrap().collect();
430 assert_eq!(collected.len(), 1);
431 assert!(is_mag_completion(&collected[0].graph));
432 }
433
434 #[test]
435 fn rejects_almost_directed_cycle() {
436 let mut g = Pag::with_variables(3);
438 let a = DenseNodeId::from_raw(0);
439 let b = DenseNodeId::from_raw(1);
440 let c = DenseNodeId::from_raw(2);
441 g.insert_directed(a, b).unwrap();
442 g.insert_directed(b, c).unwrap();
443 g.insert_bidirected(a, c).unwrap();
444 assert!(!is_mag_completion(&g));
445 }
446
447 #[test]
448 fn accepts_bidirected_without_directed_path() {
449 let mut g = Pag::with_variables(2);
450 g.insert_bidirected(DenseNodeId::from_raw(0), DenseNodeId::from_raw(1)).unwrap();
451 assert!(is_mag_completion(&g));
452 }
453
454 #[test]
455 fn rejects_nonmaximal_ancestral_orientation() {
456 let mut graph = Pag::with_variables(4);
460 let endpoint_x = DenseNodeId::from_raw(0);
461 let inner_a = DenseNodeId::from_raw(1);
462 let inner_b = DenseNodeId::from_raw(2);
463 let endpoint_y = DenseNodeId::from_raw(3);
464 graph.insert_bidirected(endpoint_x, inner_a).unwrap();
465 graph.insert_bidirected(inner_a, inner_b).unwrap();
466 graph.insert_bidirected(inner_b, endpoint_y).unwrap();
467 graph.insert_directed(inner_a, endpoint_y).unwrap();
468 graph.insert_directed(inner_b, endpoint_x).unwrap();
469 assert!(is_ancestral_orientation(&graph), "graph is ancestral but not maximal");
470 assert!(!is_maximal_ancestral_graph(&graph));
471
472 let sampler = CompletionSampler::new(graph, 8).unwrap();
473 assert_eq!(sampler.validation_report().rejected_nonmaximal, 1);
474 assert_eq!(sampler.count(), 0, "nonmaximal orientation must never be yielded");
475 }
476
477 #[test]
478 fn rejects_unshielded_collider_not_encoded_by_pag() {
479 let mut pag = Pag::with_variables(3);
483 let x = DenseNodeId::from_raw(0);
484 let m = DenseNodeId::from_raw(1);
485 let y = DenseNodeId::from_raw(2);
486 pag.insert_circle_circle(x, m).unwrap();
487 pag.insert_circle_circle(m, y).unwrap();
488 let sampler = CompletionSampler::new(pag, 32).unwrap();
489 assert!(sampler.validation_report().rejected_local_incompatible > 0);
490 for completion in sampler {
491 let xm = completion.graph.edge_between(x, m).unwrap();
492 let my = completion.graph.edge_between(m, y).unwrap();
493 let at_m_from_x = if xm.a == m { xm.at_a } else { xm.at_b };
494 let at_m_from_y = if my.a == m { my.at_a } else { my.at_b };
495 assert!(!(at_m_from_x == Endpoint::Arrow && at_m_from_y == Endpoint::Arrow));
496 }
497 }
498
499 #[test]
500 fn fails_closed_when_marks_admit_multiple_global_classes() {
501 let mut pag = Pag::with_variables(4);
506 let x = DenseNodeId::from_raw(0);
507 let q = DenseNodeId::from_raw(1);
508 let b = DenseNodeId::from_raw(2);
509 let y = DenseNodeId::from_raw(3);
510 pag.insert_directed(x, q).unwrap();
511 pag.insert_bidirected(q, b).unwrap();
512 pag.insert_directed(q, y).unwrap();
513 pag.insert_circle_circle(b, y).unwrap();
514
515 let sampler = CompletionSampler::new(pag, 8).unwrap();
516 assert!(sampler.validation_report().ambiguous_global_class);
517 assert_eq!(sampler.count(), 0, "an underoriented PAG must fail closed");
518 }
519
520 #[test]
521 fn skips_rather_than_refuses_equivalence_audits_above_work_bound() {
522 let mut pag = Pag::with_variables(12);
527 for (left, right) in [(0, 1), (2, 3), (4, 5)] {
528 pag.insert_circle_circle(DenseNodeId::from_raw(left), DenseNodeId::from_raw(right))
529 .unwrap();
530 }
531 let sampler = CompletionSampler::new(pag, 8).unwrap();
532 assert!(sampler.class_audit_incomplete());
533 assert!(sampler.validation_report().equivalence_audit_skipped);
534 assert!(!sampler.validation_report().ambiguous_global_class);
535 assert!(sampler.count() > 0, "locally valid completions are still yielded");
536 }
537
538 #[test]
539 fn wide_pags_with_circle_marks_are_not_refused_outright() {
540 let mut pag = Pag::with_variables(13);
543 pag.insert_circle_circle(DenseNodeId::from_raw(0), DenseNodeId::from_raw(1)).unwrap();
544 let sampler = CompletionSampler::new(pag, 8).unwrap();
545 assert!(sampler.class_audit_incomplete());
546 assert!(sampler.count() > 0);
547 }
548
549 fn random_pag_with_circles(rng: &mut antecedent_core::CausalRng, n: u32) -> Pag {
550 let mut pag = Pag::with_variables(n);
551 let mut order: Vec<u32> = (0..n).collect();
553 for i in (1..usize::try_from(n).unwrap_or(0)).rev() {
554 let bound = u64::try_from(i + 1).unwrap_or(1);
555 let j = usize::try_from(rng.next_u64() % bound).unwrap_or(0);
556 order.swap(i, j);
557 }
558 let n_usize = usize::try_from(n).unwrap_or(0);
559 for i in 0..n_usize {
560 for j in (i + 1)..n_usize {
561 if rng.next_u64() % 3 != 0 {
562 continue;
563 }
564 let a = DenseNodeId::from_raw(order[i]);
565 let b = DenseNodeId::from_raw(order[j]);
566 let kind = rng.next_u64() % 4;
567 let _ = match kind {
568 0 => pag.insert_directed(a, b),
569 1 => pag.insert_circle_arrow(a, b),
570 2 => pag.insert_circle_circle(a, b),
571 _ => pag.insert_bidirected(a, b),
572 };
573 }
574 }
575 pag
576 }
577
578 #[test]
580 fn property_completions_respect_bound_and_no_circles() {
581 use antecedent_core::CausalRng;
582
583 let mut rng = CausalRng::from_seed(23);
584 for _ in 0..40 {
585 let n = 2 + u32::try_from(rng.next_u64() % 3).unwrap_or(0); let pag = random_pag_with_circles(&mut rng, n);
587 let max_c = 1 + usize::try_from(rng.next_u64() % 4).unwrap_or(0); let Ok(sampler) = CompletionSampler::new(pag, max_c) else {
589 continue; };
591 let collected: Vec<_> = sampler.collect();
592 assert!(collected.len() <= max_c, "exceeded max_completions");
593 for (i, c) in collected.iter().enumerate() {
594 assert_eq!(c.index, i);
595 assert!(is_mag_completion(&c.graph));
596 for i in 0..c.graph.node_count() {
597 let a = DenseNodeId::from_raw(u32::try_from(i).unwrap());
598 for (b, at_a, at_b) in c.graph.neighbors(a) {
599 if b.raw() < a.raw() {
600 continue;
601 }
602 assert!(!matches!(at_a, Endpoint::Circle | Endpoint::Conflict));
603 assert!(!matches!(at_b, Endpoint::Circle | Endpoint::Conflict));
604 }
605 }
606 }
607 }
608 }
609
610 #[test]
613 fn property_definite_msep_stable_across_completions() {
614 use antecedent_core::CausalRng;
615
616 let mut rng = CausalRng::from_seed(29);
617 for _ in 0..30 {
618 let n = 3u32;
619 let pag = random_pag_with_circles(&mut rng, n);
620 let Ok(sampler) = CompletionSampler::new(pag.clone(), 8) else {
621 continue;
622 };
623 if sampler.n_circle_sites() > 4 {
624 continue; }
626 let completions: Vec<_> = sampler.collect();
627 if completions.is_empty() {
628 continue;
629 }
630 for x in 0..n {
631 for y in 0..n {
632 if x == y {
633 continue;
634 }
635 let xi = DenseNodeId::from_raw(x);
636 let yi = DenseNodeId::from_raw(y);
637 let Ok(pag_sep) = pag.is_m_separated(xi, yi, &[], 32, 6) else {
639 continue; };
641 if pag_sep {
642 continue; }
644 for c in &completions {
645 let Ok(comp_sep) = c.graph.is_m_separated(xi, yi, &[], 32, 6) else {
646 continue;
647 };
648 assert!(
649 !comp_sep,
650 "PAG m-connected but completion {} separated {}–{}",
651 c.index, x, y
652 );
653 }
654 }
655 }
656 }
657 }
658}