1#![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#[repr(transparent)]
42#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
43pub struct ExprId(u32);
44
45impl ExprId {
46 #[must_use]
48 pub const fn from_raw(raw: u32) -> Self {
49 Self(raw)
50 }
51
52 #[must_use]
54 pub const fn raw(self) -> u32 {
55 self.0
56 }
57}
58
59#[repr(transparent)]
61#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
62pub struct VarSetId(u32);
63
64impl VarSetId {
65 #[must_use]
67 pub const fn from_raw(raw: u32) -> Self {
68 Self(raw)
69 }
70
71 #[must_use]
73 pub const fn raw(self) -> u32 {
74 self.0
75 }
76}
77
78#[repr(transparent)]
80#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
81pub struct InterventionSetId(u32);
82
83impl InterventionSetId {
84 #[must_use]
86 pub const fn from_raw(raw: u32) -> Self {
87 Self(raw)
88 }
89
90 #[must_use]
92 pub const fn raw(self) -> u32 {
93 self.0
94 }
95}
96
97#[derive(Clone, Debug, Eq, PartialEq, Hash)]
99pub struct InterventionAssignment {
100 pub variable: VariableId,
102 pub value: Value,
104}
105
106#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
108pub enum ContrastOp {
109 Difference,
111}
112
113#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
115pub enum DomainRef {
116 Observational,
118 Interventional,
120}
121
122#[repr(transparent)]
124#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
125pub struct OutcomeExprId(VariableId);
126
127impl OutcomeExprId {
128 #[must_use]
130 pub const fn identity(variable: VariableId) -> Self {
131 Self(variable)
132 }
133
134 #[must_use]
136 pub const fn variable(self) -> VariableId {
137 self.0
138 }
139}
140
141#[repr(transparent)]
143#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
144pub struct ExprListId(u32);
145
146impl ExprListId {
147 #[must_use]
149 pub const fn from_raw(raw: u32) -> Self {
150 Self(raw)
151 }
152
153 #[must_use]
155 pub const fn raw(self) -> u32 {
156 self.0
157 }
158}
159
160#[derive(Clone, Debug, Eq, PartialEq, Hash)]
162pub enum ExprNode {
163 Distribution {
165 variables: VarSetId,
167 conditioned_on: VarSetId,
169 intervention: InterventionSetId,
171 domain: DomainRef,
173 },
174 Product(ExprListId),
176 SumOut {
178 variables: VarSetId,
180 expr: ExprId,
182 },
183 IntegralOut {
185 variables: VarSetId,
187 expr: ExprId,
189 },
190 Ratio {
192 numerator: ExprId,
194 denominator: ExprId,
196 },
197 Expectation {
199 function: OutcomeExprId,
201 distribution: ExprId,
203 },
204 Contrast {
206 left: ExprId,
208 right: ExprId,
210 op: ContrastOp,
212 },
213}
214
215#[derive(Clone, Debug, Default, Eq, PartialEq)]
217pub struct DerivationMeta {
218 pub rule: Arc<str>,
220 pub note: Option<Arc<str>>,
222}
223
224#[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 node_index: HashMap<ExprNode, ExprId>,
236 derivation: HashMap<u32, DerivationMeta>,
238}
239
240impl CausalExprArena {
241 #[must_use]
243 pub fn new() -> Self {
244 Self::default()
245 }
246
247 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 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 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 pub fn empty_var_set(&mut self) -> VarSetId {
293 self.intern_var_set([])
294 }
295
296 pub fn empty_intervention_set(&mut self) -> InterventionSetId {
298 self.intern_intervention_assignments([])
299 }
300
301 #[must_use]
303 pub fn var_set(&self, id: VarSetId) -> &[VariableId] {
304 &self.var_sets[id.0 as usize]
305 }
306
307 #[must_use]
309 pub fn intervention_assignments(&self, id: InterventionSetId) -> &[InterventionAssignment] {
310 &self.interventions[id.0 as usize]
311 }
312
313 #[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 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 #[must_use]
333 pub fn list(&self, id: ExprListId) -> &[ExprId] {
334 &self.lists[id.0 as usize]
335 }
336
337 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 pub fn set_derivation(&mut self, id: ExprId, meta: DerivationMeta) {
350 self.derivation.insert(id.0, meta);
351 }
352
353 pub fn set_derivation_if_absent(&mut self, id: ExprId, meta: DerivationMeta) {
355 self.derivation.entry(id.0).or_insert(meta);
356 }
357
358 pub fn simplify(&mut self, root: ExprId) -> Result<ExprId, SimplifyError> {
365 simplify::simplify(self, root)
366 }
367
368 #[must_use]
370 pub fn derivation(&self, id: ExprId) -> Option<&DerivationMeta> {
371 self.derivation.get(&id.0)
372 }
373
374 #[must_use]
376 pub fn node(&self, id: ExprId) -> &ExprNode {
377 &self.nodes[id.0 as usize]
378 }
379
380 #[must_use]
382 pub fn len(&self) -> usize {
383 self.nodes.len()
384 }
385
386 #[must_use]
388 pub fn is_empty(&self) -> bool {
389 self.nodes.is_empty()
390 }
391
392 #[must_use]
394 pub fn var_set_count(&self) -> usize {
395 self.var_sets.len()
396 }
397
398 #[must_use]
400 pub fn intervention_set_count(&self) -> usize {
401 self.interventions.len()
402 }
403
404 #[must_use]
406 pub fn list_count(&self) -> usize {
407 self.lists.len()
408 }
409
410 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 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 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 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 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 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 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 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 #[must_use]
654 pub fn pretty(&self, id: ExprId) -> String {
655 pretty_expr(self, id)
656 }
657
658 #[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}