antecedent_expr/
lib.rs

1//! Arena-backed causal-functional IR.
2//!
3//! # Modules
4//!
5//! - [`estimand`] — identified estimand + method tags
6//! - [`eval`] — compiled evaluators over providers
7//! - [`simplify`] — algebraic simplification
8//! - [`pretty`] / [`latex`] — display helpers
9//! - [`provider`] — distribution / table / posterior providers
10//!
11//! SPDX-License-Identifier: MIT OR Apache-2.0
12
13#![forbid(unsafe_code)]
14#![deny(missing_docs)]
15
16pub mod estimand;
17pub mod eval;
18pub mod latex;
19pub mod pretty;
20pub mod provider;
21pub mod simplify;
22
23pub use estimand::{EstimandMethod, IdentifiedEstimand, RdDesignParams};
24pub use eval::CompiledEvaluator;
25pub use provider::{
26    Assignment, DistributionProvider, EmpiricalTableProvider, EvalContext, EvalError, FactorSpec,
27    GaussianDensityProvider, PosteriorDrawProvider, QuadratureNodes,
28};
29pub use simplify::SimplifyError;
30
31use latex::latex_expr;
32use pretty::pretty_expr;
33
34use std::collections::HashMap;
35use std::fmt;
36use std::sync::Arc;
37
38use antecedent_core::{Value, VariableId};
39
40/// Opaque expression node id.
41#[repr(transparent)]
42#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
43pub struct ExprId(u32);
44
45impl ExprId {
46    /// Create from a raw index (tests / deserialization).
47    #[must_use]
48    pub const fn from_raw(raw: u32) -> Self {
49        Self(raw)
50    }
51
52    /// Raw index.
53    #[must_use]
54    pub const fn raw(self) -> u32 {
55        self.0
56    }
57}
58
59/// Interned sorted variable set id.
60#[repr(transparent)]
61#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
62pub struct VarSetId(u32);
63
64impl VarSetId {
65    /// Create from a raw index (deserialization).
66    #[must_use]
67    pub const fn from_raw(raw: u32) -> Self {
68        Self(raw)
69    }
70
71    /// Raw index.
72    #[must_use]
73    pub const fn raw(self) -> u32 {
74        self.0
75    }
76}
77
78/// Interned intervention-set id (hard assignments `do(V := value)`).
79#[repr(transparent)]
80#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
81pub struct InterventionSetId(u32);
82
83impl InterventionSetId {
84    /// Create from a raw index (deserialization).
85    #[must_use]
86    pub const fn from_raw(raw: u32) -> Self {
87        Self(raw)
88    }
89
90    /// Raw index.
91    #[must_use]
92    pub const fn raw(self) -> u32 {
93        self.0
94    }
95}
96
97/// One hard intervention assignment in an interned set.
98#[derive(Clone, Debug, Eq, PartialEq, Hash)]
99pub struct InterventionAssignment {
100    /// Target variable.
101    pub variable: VariableId,
102    /// Assigned value under `do(·)`.
103    pub value: Value,
104}
105
106/// Contrast operator between two expressions.
107#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
108pub enum ContrastOp {
109    /// Left − right.
110    Difference,
111}
112
113/// Domain reference for a distribution .
114#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
115pub enum DomainRef {
116    /// Observational P(·).
117    Observational,
118    /// Interventional P(· | do(·)).
119    Interventional,
120}
121
122/// Outcome function id .
123#[repr(transparent)]
124#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
125pub struct OutcomeExprId(VariableId);
126
127impl OutcomeExprId {
128    /// Identity outcome Y.
129    #[must_use]
130    pub const fn identity(variable: VariableId) -> Self {
131        Self(variable)
132    }
133
134    /// Underlying variable.
135    #[must_use]
136    pub const fn variable(self) -> VariableId {
137        self.0
138    }
139}
140
141/// Expression list id (product children).
142#[repr(transparent)]
143#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
144pub struct ExprListId(u32);
145
146impl ExprListId {
147    /// Create from a raw index (deserialization).
148    #[must_use]
149    pub const fn from_raw(raw: u32) -> Self {
150        Self(raw)
151    }
152
153    /// Raw index.
154    #[must_use]
155    pub const fn raw(self) -> u32 {
156        self.0
157    }
158}
159
160/// Semantic expression node (no derivation metadata).
161#[derive(Clone, Debug, Eq, PartialEq, Hash)]
162pub enum ExprNode {
163    /// Joint / conditional distribution factor.
164    Distribution {
165        /// Variables in the factor.
166        variables: VarSetId,
167        /// Conditioning set.
168        conditioned_on: VarSetId,
169        /// Intervention set (empty for observational).
170        intervention: InterventionSetId,
171        /// Domain.
172        domain: DomainRef,
173    },
174    /// Product of factors.
175    Product(ExprListId),
176    /// Discrete marginalization.
177    SumOut {
178        /// Variables summed out.
179        variables: VarSetId,
180        /// Body.
181        expr: ExprId,
182    },
183    /// Continuous marginalization.
184    IntegralOut {
185        /// Variables integrated out.
186        variables: VarSetId,
187        /// Body.
188        expr: ExprId,
189    },
190    /// Ratio of expressions.
191    Ratio {
192        /// Numerator.
193        numerator: ExprId,
194        /// Denominator.
195        denominator: ExprId,
196    },
197    /// Expectation of an outcome under a distribution.
198    Expectation {
199        /// Outcome function.
200        function: OutcomeExprId,
201        /// Distribution expression.
202        distribution: ExprId,
203    },
204    /// Contrast of two expectations / functionals.
205    Contrast {
206        /// Left side.
207        left: ExprId,
208        /// Right side.
209        right: ExprId,
210        /// Operator.
211        op: ContrastOp,
212    },
213}
214
215/// Separate derivation metadata keyed by expression id.
216#[derive(Clone, Debug, Default, Eq, PartialEq)]
217pub struct DerivationMeta {
218    /// Human-readable rule tag (e.g. `backdoor.adjustment`).
219    pub rule: Arc<str>,
220    /// Optional note.
221    pub note: Option<Arc<str>>,
222}
223
224/// Arena for causal expressions with interned variable sets.
225#[derive(Clone, Debug, Default)]
226pub struct CausalExprArena {
227    nodes: Vec<ExprNode>,
228    var_sets: Vec<Arc<[VariableId]>>,
229    var_set_index: HashMap<Arc<[VariableId]>, VarSetId>,
230    interventions: Vec<Arc<[InterventionAssignment]>>,
231    intervention_index: HashMap<Arc<[InterventionAssignment]>, InterventionSetId>,
232    lists: Vec<Arc<[ExprId]>>,
233    list_index: HashMap<Arc<[ExprId]>, ExprListId>,
234    /// Hash-cons map from node → id.
235    node_index: HashMap<ExprNode, ExprId>,
236    /// Derivation metadata (optional; not part of semantic equality).
237    derivation: HashMap<u32, DerivationMeta>,
238}
239
240impl CausalExprArena {
241    /// Empty arena.
242    #[must_use]
243    pub fn new() -> Self {
244        Self::default()
245    }
246
247    /// Intern a sorted variable set (sorts and dedups input).
248    pub fn intern_var_set(&mut self, vars: impl IntoIterator<Item = VariableId>) -> VarSetId {
249        let mut v: Vec<VariableId> = vars.into_iter().collect();
250        v.sort_unstable();
251        v.dedup();
252        let key: Arc<[VariableId]> = Arc::from(v);
253        if let Some(id) = self.var_set_index.get(&key) {
254            return *id;
255        }
256        let id = VarSetId(u32::try_from(self.var_sets.len()).expect("var set id"));
257        self.var_sets.push(Arc::clone(&key));
258        self.var_set_index.insert(key, id);
259        id
260    }
261
262    /// Intern a hard-intervention assignment set (sorted by variable id).
263    pub fn intern_intervention_assignments(
264        &mut self,
265        assignments: impl IntoIterator<Item = InterventionAssignment>,
266    ) -> InterventionSetId {
267        let mut v: Vec<InterventionAssignment> = assignments.into_iter().collect();
268        v.sort_by_key(|a| a.variable.raw());
269        v.dedup_by_key(|a| a.variable.raw());
270        let key: Arc<[InterventionAssignment]> = Arc::from(v);
271        if let Some(id) = self.intervention_index.get(&key) {
272            return *id;
273        }
274        let id = InterventionSetId(u32::try_from(self.interventions.len()).expect("id"));
275        self.interventions.push(Arc::clone(&key));
276        self.intervention_index.insert(key, id);
277        id
278    }
279
280    /// Intern an intervention over variables only (value unspecified / placeholder).
281    pub fn intern_intervention_set(
282        &mut self,
283        vars: impl IntoIterator<Item = VariableId>,
284    ) -> InterventionSetId {
285        self.intern_intervention_assignments(
286            vars.into_iter()
287                .map(|variable| InterventionAssignment { variable, value: Value::f64(f64::NAN) }),
288        )
289    }
290
291    /// Empty var set.
292    pub fn empty_var_set(&mut self) -> VarSetId {
293        self.intern_var_set([])
294    }
295
296    /// Empty intervention set.
297    pub fn empty_intervention_set(&mut self) -> InterventionSetId {
298        self.intern_intervention_assignments([])
299    }
300
301    /// Look up a var set.
302    #[must_use]
303    pub fn var_set(&self, id: VarSetId) -> &[VariableId] {
304        &self.var_sets[id.0 as usize]
305    }
306
307    /// Look up intervention assignments.
308    #[must_use]
309    pub fn intervention_assignments(&self, id: InterventionSetId) -> &[InterventionAssignment] {
310        &self.interventions[id.0 as usize]
311    }
312
313    /// Variables appearing in an intervention set (legacy helper).
314    #[must_use]
315    pub fn intervention_set(&self, id: InterventionSetId) -> Vec<VariableId> {
316        self.intervention_assignments(id).iter().map(|a| a.variable).collect()
317    }
318
319    /// Intern an expression list.
320    pub fn intern_list(&mut self, exprs: impl IntoIterator<Item = ExprId>) -> ExprListId {
321        let key: Arc<[ExprId]> = Arc::from(exprs.into_iter().collect::<Vec<_>>());
322        if let Some(id) = self.list_index.get(&key) {
323            return *id;
324        }
325        let id = ExprListId(u32::try_from(self.lists.len()).expect("list id"));
326        self.lists.push(Arc::clone(&key));
327        self.list_index.insert(key, id);
328        id
329    }
330
331    /// Borrow an interned expression list.
332    #[must_use]
333    pub fn list(&self, id: ExprListId) -> &[ExprId] {
334        &self.lists[id.0 as usize]
335    }
336
337    /// Hash-cons an expression node.
338    pub fn intern(&mut self, node: ExprNode) -> ExprId {
339        if let Some(id) = self.node_index.get(&node) {
340            return *id;
341        }
342        let id = ExprId(u32::try_from(self.nodes.len()).expect("expr id"));
343        self.nodes.push(node.clone());
344        self.node_index.insert(node, id);
345        id
346    }
347
348    /// Attach derivation metadata (does not affect semantic equality).
349    pub fn set_derivation(&mut self, id: ExprId, meta: DerivationMeta) {
350        self.derivation.insert(id.0, meta);
351    }
352
353    /// Attach derivation metadata only when absent (never overwrites ID rules).
354    pub fn set_derivation_if_absent(&mut self, id: ExprId, meta: DerivationMeta) {
355        self.derivation.entry(id.0).or_insert(meta);
356    }
357
358    /// Simplify `root` with worklist-style bottom-up rewrite + memoization.
359    ///
360    /// # Errors
361    ///
362    /// [`SimplifyError`] if a `SumOut`/`IntegralOut` binds a variable absent from its
363    /// body's free variables — an ill-formed estimand. See [`SimplifyError`] docs.
364    pub fn simplify(&mut self, root: ExprId) -> Result<ExprId, SimplifyError> {
365        simplify::simplify(self, root)
366    }
367
368    /// Borrow derivation metadata.
369    #[must_use]
370    pub fn derivation(&self, id: ExprId) -> Option<&DerivationMeta> {
371        self.derivation.get(&id.0)
372    }
373
374    /// Borrow a node.
375    #[must_use]
376    pub fn node(&self, id: ExprId) -> &ExprNode {
377        &self.nodes[id.0 as usize]
378    }
379
380    /// Number of nodes.
381    #[must_use]
382    pub fn len(&self) -> usize {
383        self.nodes.len()
384    }
385
386    /// Whether empty.
387    #[must_use]
388    pub fn is_empty(&self) -> bool {
389        self.nodes.is_empty()
390    }
391
392    /// Number of interned variable sets (for serialization).
393    #[must_use]
394    pub fn var_set_count(&self) -> usize {
395        self.var_sets.len()
396    }
397
398    /// Number of interned intervention sets (for serialization).
399    #[must_use]
400    pub fn intervention_set_count(&self) -> usize {
401        self.interventions.len()
402    }
403
404    /// Number of interned expression lists (for serialization).
405    #[must_use]
406    pub fn list_count(&self) -> usize {
407        self.lists.len()
408    }
409
410    /// Build the backdoor adjustment functional for ATE:
411    /// `E[Y | do(T=active)] − E[Y | do(T=control)]` under adjustment by Z.
412    pub fn backdoor_ate(
413        &mut self,
414        treatment: VariableId,
415        outcome: VariableId,
416        adjustment: &[VariableId],
417        active: Value,
418        control: Value,
419    ) -> ExprId {
420        let left = self.backdoor_potential_outcome(treatment, outcome, adjustment, active);
421        let right = self.backdoor_potential_outcome(treatment, outcome, adjustment, control);
422        let contrast = self.intern(ExprNode::Contrast { left, right, op: ContrastOp::Difference });
423        self.set_derivation(
424            contrast,
425            DerivationMeta {
426                rule: Arc::from("backdoor.adjustment"),
427                note: Some(Arc::from(format!("ATE adjustment set size {}", adjustment.len()))),
428            },
429        );
430        contrast
431    }
432
433    fn backdoor_potential_outcome(
434        &mut self,
435        treatment: VariableId,
436        outcome: VariableId,
437        adjustment: &[VariableId],
438        level: Value,
439    ) -> ExprId {
440        let z = self.intern_var_set(adjustment.iter().copied());
441        let y = self.intern_var_set([outcome]);
442        let empty = self.empty_var_set();
443        let empty_i = self.empty_intervention_set();
444        let do_t = self.intern_intervention_assignments([InterventionAssignment {
445            variable: treatment,
446            value: level,
447        }]);
448
449        let dist_body = self.intern(ExprNode::Distribution {
450            variables: y,
451            conditioned_on: z,
452            intervention: do_t,
453            domain: DomainRef::Interventional,
454        });
455        let z_marg = self.intern(ExprNode::Distribution {
456            variables: z,
457            conditioned_on: empty,
458            intervention: empty_i,
459            domain: DomainRef::Observational,
460        });
461        let product = {
462            let list = self.intern_list([dist_body, z_marg]);
463            self.intern(ExprNode::Product(list))
464        };
465        let summed = self.intern(ExprNode::SumOut { variables: z, expr: product });
466        self.intern(ExprNode::Expectation {
467            function: OutcomeExprId::identity(outcome),
468            distribution: summed,
469        })
470    }
471
472    /// Build the front-door functional for ATE:
473    /// `E[Y | do(T=active)] − E[Y | do(T=control)]`, mediated through `M` via
474    /// `sum_m P(m | t) * sum_t' P(y | m, t') P(t')` (FD condition 2 reduces
475    /// `P(m | do(t))` to the observational `P(m | t)`).
476    pub fn frontdoor_ate(
477        &mut self,
478        treatment: VariableId,
479        outcome: VariableId,
480        mediators: &[VariableId],
481        active: Value,
482        control: Value,
483    ) -> ExprId {
484        let left = self.frontdoor_potential_outcome(treatment, outcome, mediators, active);
485        let right = self.frontdoor_potential_outcome(treatment, outcome, mediators, control);
486        let contrast = self.intern(ExprNode::Contrast { left, right, op: ContrastOp::Difference });
487        self.set_derivation(
488            contrast,
489            DerivationMeta {
490                rule: Arc::from("frontdoor"),
491                note: Some(Arc::from(format!("front-door mediator set size {}", mediators.len()))),
492            },
493        );
494        contrast
495    }
496
497    /// Linear temporal-mediation path-product ATE contrast (same product-of-coefficients
498    /// geometry as front-door under a linear SEM, tagged `temporal_mediation` — not front-door).
499    pub fn temporal_mediation_ate(
500        &mut self,
501        treatment: VariableId,
502        outcome: VariableId,
503        mediators: &[VariableId],
504        active: Value,
505        control: Value,
506    ) -> ExprId {
507        let left = self.frontdoor_potential_outcome(treatment, outcome, mediators, active);
508        let right = self.frontdoor_potential_outcome(treatment, outcome, mediators, control);
509        let contrast = self.intern(ExprNode::Contrast { left, right, op: ContrastOp::Difference });
510        self.set_derivation(
511            contrast,
512            DerivationMeta {
513                rule: Arc::from("temporal_mediation"),
514                note: Some(Arc::from(format!(
515                    "linear temporal mediation path-product; mediator set size {}",
516                    mediators.len()
517                ))),
518            },
519        );
520        contrast
521    }
522
523    fn frontdoor_potential_outcome(
524        &mut self,
525        treatment: VariableId,
526        outcome: VariableId,
527        mediators: &[VariableId],
528        level: Value,
529    ) -> ExprId {
530        let m = self.intern_var_set(mediators.iter().copied());
531        let y = self.intern_var_set([outcome]);
532        let t = self.intern_var_set([treatment]);
533        let m_and_t = self.intern_var_set(mediators.iter().copied().chain([treatment]));
534        let empty = self.empty_var_set();
535        let empty_i = self.empty_intervention_set();
536        let do_t = self.intern_intervention_assignments([InterventionAssignment {
537            variable: treatment,
538            value: level,
539        }]);
540
541        // P(m | t): observational under FD condition 2; treatment level bound so
542        // the evaluator treats it as fixed (not free).
543        let m_given_t = self.intern(ExprNode::Distribution {
544            variables: m,
545            conditioned_on: t,
546            intervention: do_t,
547            domain: DomainRef::Observational,
548        });
549        // P(y | m, t').
550        let y_given_m_t = self.intern(ExprNode::Distribution {
551            variables: y,
552            conditioned_on: m_and_t,
553            intervention: empty_i,
554            domain: DomainRef::Observational,
555        });
556        // P(t').
557        let t_marginal = self.intern(ExprNode::Distribution {
558            variables: t,
559            conditioned_on: empty,
560            intervention: empty_i,
561            domain: DomainRef::Observational,
562        });
563        let inner_product = {
564            let list = self.intern_list([y_given_m_t, t_marginal]);
565            self.intern(ExprNode::Product(list))
566        };
567        let inner_summed = self.intern(ExprNode::SumOut { variables: t, expr: inner_product });
568        let outer_product = {
569            let list = self.intern_list([m_given_t, inner_summed]);
570            self.intern(ExprNode::Product(list))
571        };
572        let outer_summed = self.intern(ExprNode::SumOut { variables: m, expr: outer_product });
573        self.intern(ExprNode::Expectation {
574            function: OutcomeExprId::identity(outcome),
575            distribution: outer_summed,
576        })
577    }
578
579    /// Build the Wald IV functional for binary instrument `Z`:
580    /// `(E[Y|Z=1] − E[Y|Z=0]) / (E[T|Z=1] − E[T|Z=0])`.
581    ///
582    /// `active` / `control` are recorded in derivation metadata (treatment contrast
583    /// scaling); the ratio itself conditions on instrument levels 1 and 0.
584    pub fn iv_wald(
585        &mut self,
586        treatment: VariableId,
587        outcome: VariableId,
588        instruments: &[VariableId],
589        active: &Value,
590        control: &Value,
591    ) -> ExprId {
592        let z = instruments.first().copied().unwrap_or(treatment);
593        let z1 = Value::f64(1.0);
594        let z0 = Value::f64(0.0);
595        let outcome_given_z1 = self.observational_conditional_mean(outcome, z, z1.clone());
596        let outcome_given_z0 = self.observational_conditional_mean(outcome, z, z0.clone());
597        let treatment_given_z1 = self.observational_conditional_mean(treatment, z, z1);
598        let treatment_given_z0 = self.observational_conditional_mean(treatment, z, z0);
599        let num = self.intern(ExprNode::Contrast {
600            left: outcome_given_z1,
601            right: outcome_given_z0,
602            op: ContrastOp::Difference,
603        });
604        let den = self.intern(ExprNode::Contrast {
605            left: treatment_given_z1,
606            right: treatment_given_z0,
607            op: ContrastOp::Difference,
608        });
609        let ratio = self.intern(ExprNode::Ratio { numerator: num, denominator: den });
610        self.set_derivation(
611            ratio,
612            DerivationMeta {
613                rule: Arc::from("iv.wald"),
614                note: Some(Arc::from(format!(
615                    "Wald IV ratio using {} instrument(s); treatment contrast [{active:?}, {control:?}]",
616                    instruments.len()
617                ))),
618            },
619        );
620        ratio
621    }
622
623    /// Observational `E[outcome | conditioner = level]`.
624    ///
625    /// The conditioning level is bound via an intervention assignment so the
626    /// evaluator treats it as fixed (not free), while the factor remains
627    /// observational `P(outcome | conditioner)`.
628    fn observational_conditional_mean(
629        &mut self,
630        outcome: VariableId,
631        conditioner: VariableId,
632        level: Value,
633    ) -> ExprId {
634        let y = self.intern_var_set([outcome]);
635        let z = self.intern_var_set([conditioner]);
636        let bind = self.intern_intervention_assignments([InterventionAssignment {
637            variable: conditioner,
638            value: level,
639        }]);
640        let dist = self.intern(ExprNode::Distribution {
641            variables: y,
642            conditioned_on: z,
643            intervention: bind,
644            domain: DomainRef::Observational,
645        });
646        self.intern(ExprNode::Expectation {
647            function: OutcomeExprId::identity(outcome),
648            distribution: dist,
649        })
650    }
651
652    /// Pretty-print an expression (diagnostics only; not an equality key).
653    #[must_use]
654    pub fn pretty(&self, id: ExprId) -> String {
655        pretty_expr(self, id)
656    }
657
658    /// Render an expression as LaTeX (diagnostics only; not an equality key).
659    #[must_use]
660    pub fn latex(&self, id: ExprId) -> String {
661        latex_expr(self, id)
662    }
663}
664
665impl fmt::Display for ExprId {
666    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
667        write!(f, "E{}", self.0)
668    }
669}
670
671#[cfg(test)]
672mod tests {
673    use super::*;
674
675    #[test]
676    fn var_sets_are_sorted_and_interned() {
677        let mut a = CausalExprArena::new();
678        let s1 = a.intern_var_set([VariableId::from_raw(2), VariableId::from_raw(1)]);
679        let s2 = a.intern_var_set([VariableId::from_raw(1), VariableId::from_raw(2)]);
680        assert_eq!(s1, s2);
681        assert_eq!(a.var_set(s1), &[VariableId::from_raw(1), VariableId::from_raw(2)]);
682    }
683
684    #[test]
685    fn hash_cons_reuses_nodes() {
686        let mut a = CausalExprArena::new();
687        let empty = a.empty_var_set();
688        let empty_i = a.empty_intervention_set();
689        let n1 = a.intern(ExprNode::Distribution {
690            variables: empty,
691            conditioned_on: empty,
692            intervention: empty_i,
693            domain: DomainRef::Observational,
694        });
695        let n2 = a.intern(ExprNode::Distribution {
696            variables: empty,
697            conditioned_on: empty,
698            intervention: empty_i,
699            domain: DomainRef::Observational,
700        });
701        assert_eq!(n1, n2);
702        assert_eq!(a.len(), 1);
703    }
704
705    #[test]
706    fn backdoor_ate_contrasts_distinct_levels() {
707        let mut a = CausalExprArena::new();
708        let id = a.backdoor_ate(
709            VariableId::from_raw(0),
710            VariableId::from_raw(1),
711            &[VariableId::from_raw(2)],
712            Value::f64(1.0),
713            Value::f64(0.0),
714        );
715        let meta = a.derivation(id).unwrap();
716        assert_eq!(&*meta.rule, "backdoor.adjustment");
717        let ExprNode::Contrast { left, right, .. } = a.node(id) else {
718            panic!("expected contrast");
719        };
720        assert_ne!(left, right);
721        let pretty = a.pretty(id);
722        assert!(pretty.contains('−') || pretty.contains("E["));
723        let latex = a.latex(id);
724        assert!(latex.contains("\\mathbb{E}") || latex.contains("\\mathrm{do}"));
725        assert!(latex.contains('-'));
726    }
727
728    #[test]
729    fn frontdoor_ate_contrasts_distinct_levels() {
730        let mut a = CausalExprArena::new();
731        let id = a.frontdoor_ate(
732            VariableId::from_raw(0),
733            VariableId::from_raw(1),
734            &[VariableId::from_raw(2)],
735            Value::f64(1.0),
736            Value::f64(0.0),
737        );
738        let meta = a.derivation(id).unwrap();
739        assert_eq!(&*meta.rule, "frontdoor");
740        let ExprNode::Contrast { left, right, .. } = a.node(id) else {
741            panic!("expected contrast");
742        };
743        assert_ne!(left, right);
744    }
745
746    #[test]
747    fn iv_wald_is_ratio_of_instrument_contrasts() {
748        let mut a = CausalExprArena::new();
749        let id = a.iv_wald(
750            VariableId::from_raw(0),
751            VariableId::from_raw(1),
752            &[VariableId::from_raw(2)],
753            &Value::f64(1.0),
754            &Value::f64(0.0),
755        );
756        let meta = a.derivation(id).unwrap();
757        assert_eq!(&*meta.rule, "iv.wald");
758        let ExprNode::Ratio { numerator, denominator } = a.node(id) else {
759            panic!("expected Wald ratio");
760        };
761        assert!(matches!(a.node(*numerator), ExprNode::Contrast { .. }));
762        assert!(matches!(a.node(*denominator), ExprNode::Contrast { .. }));
763        assert_ne!(*numerator, *denominator);
764    }
765}
antecedent_expr/lib.rs

antecedent_expr/
lib.rs
