Skip to main content

xlog_prob/mc/
evidence.rs

1//! Evidence forcing for Monte Carlo sampling.
2
3use xlog_core::{Result, XlogError};
4
5use super::{McProgram, McSamplingMethod};
6
7/// Why evidence may or may not be forceable to root Bernoulli variables.
8#[derive(Debug, Clone, PartialEq, Eq)]
9pub enum ForceabilityReason {
10    /// Every evidence literal maps directly to a sampled Bernoulli variable.
11    AllForceable,
12    /// At least one evidence literal is derived rather than sampled directly.
13    ContainsDerivedEvidence,
14    /// A negated annotated-disjunction head cannot be represented by direct clamping.
15    ContainsNegativeAdHeadEvidence,
16    /// The program contains no evidence to clamp.
17    NoEvidence,
18}
19
20/// Compiled evidence forcing for the MC sampler.
21#[derive(Debug, Clone)]
22pub struct EvidenceForcing {
23    /// Per-variable flag indicating that the sampler must force the variable.
24    pub force_mask: Vec<u8>,
25    /// Per-variable Boolean value used when the corresponding mask entry is set.
26    pub forced_value: Vec<u8>,
27    /// Whether all evidence can be enforced by direct variable clamping.
28    pub forceable: bool,
29    /// Explanation for the forceability decision.
30    pub reason: ForceabilityReason,
31}
32
33impl McProgram {
34    /// Resolve the sampling method from config + evidence forceability.
35    pub(super) fn resolve_sampling_method(
36        &self,
37        requested: Option<McSamplingMethod>,
38    ) -> Result<(McSamplingMethod, EvidenceForcing)> {
39        let forcing = self.compile_evidence_forcing()?;
40        let method = match requested {
41            Some(McSamplingMethod::EvidenceClamping) => {
42                if !forcing.forceable {
43                    return Err(XlogError::Execution(format!(
44                        "Cannot use EvidenceClamping: {:?}",
45                        forcing.reason
46                    )));
47                }
48                McSamplingMethod::EvidenceClamping
49            }
50            Some(McSamplingMethod::Rejection) => McSamplingMethod::Rejection,
51            None => {
52                if forcing.forceable {
53                    McSamplingMethod::EvidenceClamping
54                } else {
55                    McSamplingMethod::Rejection
56                }
57            }
58        };
59        Ok((method, forcing))
60    }
61
62    /// Compiles program evidence into per-variable sampler constraints.
63    pub fn compile_evidence_forcing(&self) -> Result<EvidenceForcing> {
64        let num_vars = self.bernoulli_probs.len();
65        let mut force_mask = vec![0u8; num_vars];
66        let mut forced_value = vec![0u8; num_vars];
67
68        if self.evidence.is_empty() {
69            return Ok(EvidenceForcing {
70                force_mask,
71                forced_value,
72                forceable: false,
73                reason: ForceabilityReason::NoEvidence,
74            });
75        }
76
77        for (atom, expected) in &self.evidence {
78            // Try to match against prob fact specs
79            if let Some(spec) = self.prob_facts.iter().find(|s| &s.atom == atom) {
80                force_mask[spec.var_idx] = 1;
81                forced_value[spec.var_idx] = if *expected { 1 } else { 0 };
82                continue;
83            }
84
85            // Try to match against AD choice atoms (positive evidence only)
86            let mut found_ad = false;
87            for ad in &self.annotated_disjunctions {
88                if let Some(choice_idx) = ad.choices.iter().position(|c| c == atom) {
89                    if !*expected {
90                        // evidence(ad_head, false) — not forceable in v0.5.1
91                        return Ok(EvidenceForcing {
92                            force_mask: vec![0u8; num_vars],
93                            forced_value: vec![0u8; num_vars],
94                            forceable: false,
95                            reason: ForceabilityReason::ContainsNegativeAdHeadEvidence,
96                        });
97                    }
98
99                    let num_decision_vars = ad.decision_vars.len();
100                    if choice_idx < num_decision_vars {
101                        // Force d_i = 0 for all i < choice_idx, d_{choice_idx} = 1
102                        for i in 0..choice_idx {
103                            force_mask[ad.decision_vars[i]] = 1;
104                            forced_value[ad.decision_vars[i]] = 0;
105                        }
106                        force_mask[ad.decision_vars[choice_idx]] = 1;
107                        forced_value[ad.decision_vars[choice_idx]] = 1;
108                    } else {
109                        // Last head (no none branch): force all decision vars to 0
110                        for &dv in &ad.decision_vars {
111                            force_mask[dv] = 1;
112                            forced_value[dv] = 0;
113                        }
114                    }
115                    found_ad = true;
116                    break;
117                }
118            }
119            if found_ad {
120                continue;
121            }
122
123            // Evidence atom not found in prob facts or AD choices → derived/deterministic
124            return Ok(EvidenceForcing {
125                force_mask: vec![0u8; num_vars],
126                forced_value: vec![0u8; num_vars],
127                forceable: false,
128                reason: ForceabilityReason::ContainsDerivedEvidence,
129            });
130        }
131
132        Ok(EvidenceForcing {
133            force_mask,
134            forced_value,
135            forceable: true,
136            reason: ForceabilityReason::AllForceable,
137        })
138    }
139}