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    /// Cached id of the interned empty variable set. The empty sets are
239    /// requested by nearly every builder and evaluator; caching skips the
240    /// intern-table lookup on repeat calls.
241    empty_var_set_id: Option<VarSetId>,
242    /// Cached id of the interned empty intervention set (see above).
243    empty_intervention_set_id: Option<InterventionSetId>,
244}
245
246impl CausalExprArena {
247    /// Empty arena.
248    #[must_use]
249    pub fn new() -> Self {
250        Self::default()
251    }
252
253    /// Intern a sorted variable set (sorts and dedups input).
254    pub fn intern_var_set(&mut self, vars: impl IntoIterator<Item = VariableId>) -> VarSetId {
255        let mut v: Vec<VariableId> = vars.into_iter().collect();
256        v.sort_unstable();
257        v.dedup();
258        // Borrow-based lookup: allocate the Arc key only on a cache miss.
259        if let Some(id) = self.var_set_index.get(v.as_slice()) {
260            return *id;
261        }
262        let key: Arc<[VariableId]> = Arc::from(v);
263        let id = VarSetId(u32::try_from(self.var_sets.len()).expect("var set id"));
264        self.var_sets.push(Arc::clone(&key));
265        self.var_set_index.insert(key, id);
266        id
267    }
268
269    /// Intern a hard-intervention assignment set (sorted by variable id).
270    pub fn intern_intervention_assignments(
271        &mut self,
272        assignments: impl IntoIterator<Item = InterventionAssignment>,
273    ) -> InterventionSetId {
274        let mut v: Vec<InterventionAssignment> = assignments.into_iter().collect();
275        v.sort_by_key(|a| a.variable.raw());
276        v.dedup_by_key(|a| a.variable.raw());
277        if let Some(id) = self.intervention_index.get(v.as_slice()) {
278            return *id;
279        }
280        let key: Arc<[InterventionAssignment]> = Arc::from(v);
281        let id = InterventionSetId(u32::try_from(self.interventions.len()).expect("id"));
282        self.interventions.push(Arc::clone(&key));
283        self.intervention_index.insert(key, id);
284        id
285    }
286
287    /// Intern an intervention over variables only (value unspecified / placeholder).
288    pub fn intern_intervention_set(
289        &mut self,
290        vars: impl IntoIterator<Item = VariableId>,
291    ) -> InterventionSetId {
292        self.intern_intervention_assignments(
293            vars.into_iter()
294                .map(|variable| InterventionAssignment { variable, value: Value::f64(f64::NAN) }),
295        )
296    }
297
298    /// Empty var set.
299    pub fn empty_var_set(&mut self) -> VarSetId {
300        if let Some(id) = self.empty_var_set_id {
301            return id;
302        }
303        let id = self.intern_var_set([]);
304        self.empty_var_set_id = Some(id);
305        id
306    }
307
308    /// Empty intervention set.
309    pub fn empty_intervention_set(&mut self) -> InterventionSetId {
310        if let Some(id) = self.empty_intervention_set_id {
311            return id;
312        }
313        let id = self.intern_intervention_assignments([]);
314        self.empty_intervention_set_id = Some(id);
315        id
316    }
317
318    /// Look up a var set.
319    #[must_use]
320    pub fn var_set(&self, id: VarSetId) -> &[VariableId] {
321        &self.var_sets[id.0 as usize]
322    }
323
324    /// Look up intervention assignments.
325    #[must_use]
326    pub fn intervention_assignments(&self, id: InterventionSetId) -> &[InterventionAssignment] {
327        &self.interventions[id.0 as usize]
328    }
329
330    /// Variables appearing in an intervention set (legacy helper).
331    #[must_use]
332    pub fn intervention_set(&self, id: InterventionSetId) -> Vec<VariableId> {
333        self.intervention_assignments(id).iter().map(|a| a.variable).collect()
334    }
335
336    /// Intern an expression list.
337    pub fn intern_list(&mut self, exprs: impl IntoIterator<Item = ExprId>) -> ExprListId {
338        let v: Vec<ExprId> = exprs.into_iter().collect();
339        if let Some(id) = self.list_index.get(v.as_slice()) {
340            return *id;
341        }
342        let key: Arc<[ExprId]> = Arc::from(v);
343        let id = ExprListId(u32::try_from(self.lists.len()).expect("list id"));
344        self.lists.push(Arc::clone(&key));
345        self.list_index.insert(key, id);
346        id
347    }
348
349    /// Borrow an interned expression list.
350    #[must_use]
351    pub fn list(&self, id: ExprListId) -> &[ExprId] {
352        &self.lists[id.0 as usize]
353    }
354
355    /// Hash-cons an expression node.
356    pub fn intern(&mut self, node: ExprNode) -> ExprId {
357        if let Some(id) = self.node_index.get(&node) {
358            return *id;
359        }
360        let id = ExprId(u32::try_from(self.nodes.len()).expect("expr id"));
361        self.nodes.push(node.clone());
362        self.node_index.insert(node, id);
363        id
364    }
365
366    /// Attach derivation metadata (does not affect semantic equality).
367    pub fn set_derivation(&mut self, id: ExprId, meta: DerivationMeta) {
368        self.derivation.insert(id.0, meta);
369    }
370
371    /// Attach derivation metadata only when absent (never overwrites ID rules).
372    pub fn set_derivation_if_absent(&mut self, id: ExprId, meta: DerivationMeta) {
373        self.derivation.entry(id.0).or_insert(meta);
374    }
375
376    /// Simplify `root` with worklist-style bottom-up rewrite + memoization.
377    ///
378    /// # Errors
379    ///
380    /// [`SimplifyError`] if a `SumOut`/`IntegralOut` binds a variable absent from its
381    /// body's free variables — an ill-formed estimand. See [`SimplifyError`] docs.
382    pub fn simplify(&mut self, root: ExprId) -> Result<ExprId, SimplifyError> {
383        simplify::simplify(self, root)
384    }
385
386    /// Borrow derivation metadata.
387    #[must_use]
388    pub fn derivation(&self, id: ExprId) -> Option<&DerivationMeta> {
389        self.derivation.get(&id.0)
390    }
391
392    /// Borrow a node.
393    #[must_use]
394    pub fn node(&self, id: ExprId) -> &ExprNode {
395        &self.nodes[id.0 as usize]
396    }
397
398    /// Number of nodes.
399    #[must_use]
400    pub fn len(&self) -> usize {
401        self.nodes.len()
402    }
403
404    /// Whether empty.
405    #[must_use]
406    pub fn is_empty(&self) -> bool {
407        self.nodes.is_empty()
408    }
409
410    /// Number of interned variable sets (for serialization).
411    #[must_use]
412    pub fn var_set_count(&self) -> usize {
413        self.var_sets.len()
414    }
415
416    /// Number of interned intervention sets (for serialization).
417    #[must_use]
418    pub fn intervention_set_count(&self) -> usize {
419        self.interventions.len()
420    }
421
422    /// Number of interned expression lists (for serialization).
423    #[must_use]
424    pub fn list_count(&self) -> usize {
425        self.lists.len()
426    }
427
428    /// Build the backdoor adjustment functional for ATE:
429    /// `E[Y | do(T=active)] − E[Y | do(T=control)]` under adjustment by Z.
430    pub fn backdoor_ate(
431        &mut self,
432        treatment: VariableId,
433        outcome: VariableId,
434        adjustment: &[VariableId],
435        active: Value,
436        control: Value,
437    ) -> ExprId {
438        let left = self.backdoor_potential_outcome(treatment, outcome, adjustment, active);
439        let right = self.backdoor_potential_outcome(treatment, outcome, adjustment, control);
440        let contrast = self.intern(ExprNode::Contrast { left, right, op: ContrastOp::Difference });
441        self.set_derivation(
442            contrast,
443            DerivationMeta {
444                rule: Arc::from("backdoor.adjustment"),
445                note: Some(Arc::from(format!("ATE adjustment set size {}", adjustment.len()))),
446            },
447        );
448        contrast
449    }
450
451    fn backdoor_potential_outcome(
452        &mut self,
453        treatment: VariableId,
454        outcome: VariableId,
455        adjustment: &[VariableId],
456        level: Value,
457    ) -> ExprId {
458        let z = self.intern_var_set(adjustment.iter().copied());
459        let y = self.intern_var_set([outcome]);
460        let empty = self.empty_var_set();
461        let empty_i = self.empty_intervention_set();
462        let do_t = self.intern_intervention_assignments([InterventionAssignment {
463            variable: treatment,
464            value: level,
465        }]);
466
467        let dist_body = self.intern(ExprNode::Distribution {
468            variables: y,
469            conditioned_on: z,
470            intervention: do_t,
471            domain: DomainRef::Interventional,
472        });
473        let z_marg = self.intern(ExprNode::Distribution {
474            variables: z,
475            conditioned_on: empty,
476            intervention: empty_i,
477            domain: DomainRef::Observational,
478        });
479        let product = {
480            let list = self.intern_list([dist_body, z_marg]);
481            self.intern(ExprNode::Product(list))
482        };
483        let summed = self.intern(ExprNode::SumOut { variables: z, expr: product });
484        self.intern(ExprNode::Expectation {
485            function: OutcomeExprId::identity(outcome),
486            distribution: summed,
487        })
488    }
489
490    /// Build the front-door functional for ATE:
491    /// `E[Y | do(T=active)] − E[Y | do(T=control)]`, mediated through `M` via
492    /// `sum_m P(m | t) * sum_t' P(y | m, t') P(t')` (FD condition 2 reduces
493    /// `P(m | do(t))` to the observational `P(m | t)`).
494    pub fn frontdoor_ate(
495        &mut self,
496        treatment: VariableId,
497        outcome: VariableId,
498        mediators: &[VariableId],
499        active: Value,
500        control: Value,
501    ) -> ExprId {
502        let left = self.frontdoor_potential_outcome(treatment, outcome, mediators, active);
503        let right = self.frontdoor_potential_outcome(treatment, outcome, mediators, control);
504        let contrast = self.intern(ExprNode::Contrast { left, right, op: ContrastOp::Difference });
505        self.set_derivation(
506            contrast,
507            DerivationMeta {
508                rule: Arc::from("frontdoor"),
509                note: Some(Arc::from(format!("front-door mediator set size {}", mediators.len()))),
510            },
511        );
512        contrast
513    }
514
515    /// Linear temporal-mediation path-product ATE contrast (same product-of-coefficients
516    /// geometry as front-door under a linear SEM, tagged `temporal_mediation` — not front-door).
517    pub fn temporal_mediation_ate(
518        &mut self,
519        treatment: VariableId,
520        outcome: VariableId,
521        mediators: &[VariableId],
522        active: Value,
523        control: Value,
524    ) -> ExprId {
525        let left = self.frontdoor_potential_outcome(treatment, outcome, mediators, active);
526        let right = self.frontdoor_potential_outcome(treatment, outcome, mediators, control);
527        let contrast = self.intern(ExprNode::Contrast { left, right, op: ContrastOp::Difference });
528        self.set_derivation(
529            contrast,
530            DerivationMeta {
531                rule: Arc::from("temporal_mediation"),
532                note: Some(Arc::from(format!(
533                    "linear temporal mediation path-product; mediator set size {}",
534                    mediators.len()
535                ))),
536            },
537        );
538        contrast
539    }
540
541    fn frontdoor_potential_outcome(
542        &mut self,
543        treatment: VariableId,
544        outcome: VariableId,
545        mediators: &[VariableId],
546        level: Value,
547    ) -> ExprId {
548        let m = self.intern_var_set(mediators.iter().copied());
549        let y = self.intern_var_set([outcome]);
550        let t = self.intern_var_set([treatment]);
551        let m_and_t = self.intern_var_set(mediators.iter().copied().chain([treatment]));
552        let empty = self.empty_var_set();
553        let empty_i = self.empty_intervention_set();
554        let do_t = self.intern_intervention_assignments([InterventionAssignment {
555            variable: treatment,
556            value: level,
557        }]);
558
559        // P(m | t): observational under FD condition 2; treatment level bound so
560        // the evaluator treats it as fixed (not free).
561        let m_given_t = self.intern(ExprNode::Distribution {
562            variables: m,
563            conditioned_on: t,
564            intervention: do_t,
565            domain: DomainRef::Observational,
566        });
567        // P(y | m, t').
568        let y_given_m_t = self.intern(ExprNode::Distribution {
569            variables: y,
570            conditioned_on: m_and_t,
571            intervention: empty_i,
572            domain: DomainRef::Observational,
573        });
574        // P(t').
575        let t_marginal = self.intern(ExprNode::Distribution {
576            variables: t,
577            conditioned_on: empty,
578            intervention: empty_i,
579            domain: DomainRef::Observational,
580        });
581        let inner_product = {
582            let list = self.intern_list([y_given_m_t, t_marginal]);
583            self.intern(ExprNode::Product(list))
584        };
585        let inner_summed = self.intern(ExprNode::SumOut { variables: t, expr: inner_product });
586        let outer_product = {
587            let list = self.intern_list([m_given_t, inner_summed]);
588            self.intern(ExprNode::Product(list))
589        };
590        let outer_summed = self.intern(ExprNode::SumOut { variables: m, expr: outer_product });
591        self.intern(ExprNode::Expectation {
592            function: OutcomeExprId::identity(outcome),
593            distribution: outer_summed,
594        })
595    }
596
597    /// Build the Wald IV functional for binary instrument `Z`:
598    /// `(E[Y|Z=1] − E[Y|Z=0]) / (E[T|Z=1] − E[T|Z=0])`.
599    ///
600    /// `active` / `control` are recorded in derivation metadata (treatment contrast
601    /// scaling); the ratio itself conditions on instrument levels 1 and 0.
602    pub fn iv_wald(
603        &mut self,
604        treatment: VariableId,
605        outcome: VariableId,
606        instruments: &[VariableId],
607        active: &Value,
608        control: &Value,
609    ) -> ExprId {
610        let z = instruments.first().copied().unwrap_or(treatment);
611        let z1 = Value::f64(1.0);
612        let z0 = Value::f64(0.0);
613        let outcome_given_z1 = self.observational_conditional_mean(outcome, z, z1.clone());
614        let outcome_given_z0 = self.observational_conditional_mean(outcome, z, z0.clone());
615        let treatment_given_z1 = self.observational_conditional_mean(treatment, z, z1);
616        let treatment_given_z0 = self.observational_conditional_mean(treatment, z, z0);
617        let num = self.intern(ExprNode::Contrast {
618            left: outcome_given_z1,
619            right: outcome_given_z0,
620            op: ContrastOp::Difference,
621        });
622        let den = self.intern(ExprNode::Contrast {
623            left: treatment_given_z1,
624            right: treatment_given_z0,
625            op: ContrastOp::Difference,
626        });
627        let ratio = self.intern(ExprNode::Ratio { numerator: num, denominator: den });
628        self.set_derivation(
629            ratio,
630            DerivationMeta {
631                rule: Arc::from("iv.wald"),
632                note: Some(Arc::from(format!(
633                    "Wald IV ratio using {} instrument(s); treatment contrast [{active:?}, {control:?}]",
634                    instruments.len()
635                ))),
636            },
637        );
638        ratio
639    }
640
641    /// Observational `E[outcome | conditioner = level]`.
642    ///
643    /// The conditioning level is bound via an intervention assignment so the
644    /// evaluator treats it as fixed (not free), while the factor remains
645    /// observational `P(outcome | conditioner)`.
646    fn observational_conditional_mean(
647        &mut self,
648        outcome: VariableId,
649        conditioner: VariableId,
650        level: Value,
651    ) -> ExprId {
652        let y = self.intern_var_set([outcome]);
653        let z = self.intern_var_set([conditioner]);
654        let bind = self.intern_intervention_assignments([InterventionAssignment {
655            variable: conditioner,
656            value: level,
657        }]);
658        let dist = self.intern(ExprNode::Distribution {
659            variables: y,
660            conditioned_on: z,
661            intervention: bind,
662            domain: DomainRef::Observational,
663        });
664        self.intern(ExprNode::Expectation {
665            function: OutcomeExprId::identity(outcome),
666            distribution: dist,
667        })
668    }
669
670    /// Pretty-print an expression (diagnostics only; not an equality key).
671    #[must_use]
672    pub fn pretty(&self, id: ExprId) -> String {
673        pretty_expr(self, id)
674    }
675
676    /// Render an expression as LaTeX (diagnostics only; not an equality key).
677    #[must_use]
678    pub fn latex(&self, id: ExprId) -> String {
679        latex_expr(self, id)
680    }
681}
682
683impl fmt::Display for ExprId {
684    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
685        write!(f, "E{}", self.0)
686    }
687}
688
689#[cfg(test)]
690mod tests {
691    use super::*;
692
693    #[test]
694    fn var_sets_are_sorted_and_interned() {
695        let mut a = CausalExprArena::new();
696        let s1 = a.intern_var_set([VariableId::from_raw(2), VariableId::from_raw(1)]);
697        let s2 = a.intern_var_set([VariableId::from_raw(1), VariableId::from_raw(2)]);
698        assert_eq!(s1, s2);
699        assert_eq!(a.var_set(s1), &[VariableId::from_raw(1), VariableId::from_raw(2)]);
700    }
701
702    #[test]
703    fn hash_cons_reuses_nodes() {
704        let mut a = CausalExprArena::new();
705        let empty = a.empty_var_set();
706        let empty_i = a.empty_intervention_set();
707        let n1 = a.intern(ExprNode::Distribution {
708            variables: empty,
709            conditioned_on: empty,
710            intervention: empty_i,
711            domain: DomainRef::Observational,
712        });
713        let n2 = a.intern(ExprNode::Distribution {
714            variables: empty,
715            conditioned_on: empty,
716            intervention: empty_i,
717            domain: DomainRef::Observational,
718        });
719        assert_eq!(n1, n2);
720        assert_eq!(a.len(), 1);
721    }
722
723    #[test]
724    fn backdoor_ate_contrasts_distinct_levels() {
725        let mut a = CausalExprArena::new();
726        let id = a.backdoor_ate(
727            VariableId::from_raw(0),
728            VariableId::from_raw(1),
729            &[VariableId::from_raw(2)],
730            Value::f64(1.0),
731            Value::f64(0.0),
732        );
733        let meta = a.derivation(id).unwrap();
734        assert_eq!(&*meta.rule, "backdoor.adjustment");
735        let ExprNode::Contrast { left, right, .. } = a.node(id) else {
736            panic!("expected contrast");
737        };
738        assert_ne!(left, right);
739        let pretty = a.pretty(id);
740        assert!(pretty.contains('−') || pretty.contains("E["));
741        let latex = a.latex(id);
742        assert!(latex.contains("\\mathbb{E}") || latex.contains("\\mathrm{do}"));
743        assert!(latex.contains('-'));
744    }
745
746    #[test]
747    fn frontdoor_ate_contrasts_distinct_levels() {
748        let mut a = CausalExprArena::new();
749        let id = a.frontdoor_ate(
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, "frontdoor");
758        let ExprNode::Contrast { left, right, .. } = a.node(id) else {
759            panic!("expected contrast");
760        };
761        assert_ne!(left, right);
762    }
763
764    #[test]
765    fn iv_wald_is_ratio_of_instrument_contrasts() {
766        let mut a = CausalExprArena::new();
767        let id = a.iv_wald(
768            VariableId::from_raw(0),
769            VariableId::from_raw(1),
770            &[VariableId::from_raw(2)],
771            &Value::f64(1.0),
772            &Value::f64(0.0),
773        );
774        let meta = a.derivation(id).unwrap();
775        assert_eq!(&*meta.rule, "iv.wald");
776        let ExprNode::Ratio { numerator, denominator } = a.node(id) else {
777            panic!("expected Wald ratio");
778        };
779        assert!(matches!(a.node(*numerator), ExprNode::Contrast { .. }));
780        assert!(matches!(a.node(*denominator), ExprNode::Contrast { .. }));
781        assert_ne!(*numerator, *denominator);
782    }
783}
antecedent_expr/lib.rs

antecedent_expr/
lib.rs
