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

antecedent_expr/
lib.rs
