Graph tools for causal inference

Sjoerd Smit, Wolfram Research Europe Ltd

(RowDefault[, (Spacer1), , (Spacer1), ])  #WolframTechConf

ABSTRACT

We know that correlation is not causation, but that doesn’t mean that data science cannot say anything substantial about causation at all. In this talk we will look at graph-based tools that we can use to understand spurious correlations in concrete scientific problems and how to deal with them.

What is causality?

Philosophy: is causality real and can we learn anything about it?

I argue “yes” because we clearly have some degree of mastery over the world around us

Causal inference is closely related to scientific reasoning

Correlation ≠ causation

Causality is never “just in the data”

Need causal hypotheses to test

Directed Acyclic Graph (DAG) is a common tool for this

How to investigate causality?

Experimentation

Very useful when it can be done with great control over relevant parameters (e.g., physics)

Not always possible (e.g. in astrophysics) or ethical (e.g., behavioral sciences)

Graphically:

With[{s="exposure", g="confounder", c="outcome"},
Labeled[CausalGraph[{Exposure[s]->Outcome[c], (*t-> c,*)g->s,g->c},
VertexCoordinates->{s->{0,0},
c->{2,0},g->{1,1}
}
],"Observational data", Top]
]

With[{d="coin flip",s="exposure", g="confounder", c="outcome"},
Labeled[CausalGraph[{d -> s, s->c, (*t-> c,*)g->c},
VertexCoordinates->{s->{0,0},
t->{1,0},c->{2,0},g->{1,1},d->{0, 1}
}],"Experimental data", Top]
]

Observation

Data is often readily available

It is essential to have a holistic and scientific view of how the data came to be

Sometimes we have no information about certain parts of the data-generation process (uncontrolled confounders etc.)

Basic causal graphs

The Fork (common cause)

CausalGraph[{Exposure["x"]->Outcome["y"],"c"->"x","c"->"y"}]

c=RandomInteger[{1,5},100];
x= 1+0.9c+RandomReal[1,Length[c]];
y=2+0.9c + RandomReal[1,Length[c]]; (* no direct causal link x -> y *)
data=Transpose[Dataset@AssociationThread[{"x","c","y"},{x,c,y}]];
ListPlot[data[[All,{"x","y"}]],AxesLabel->Automatic]

ListPlot[colourPoints[data,{"x","y"}->"c","Rainbow",Automatic]]

Impact of x and c on y:

Normal @ LinearModelFit[Normal@data[All,{#x,#c}->#y&],{1, \[FormalX],\[FormalC]},{ \[FormalX],\[FormalC]}]

Impact of just x on y:

Normal @ LinearModelFit[Normal@data[All,{#x}->#y&],{1, \[FormalX]},{ \[FormalX]}]

Conclusion you generally need to adjust for common causes

CausalGraph[{Exposure["x"]->Outcome["y"],Adjusted["c"]->"x","c"->"y"}]

Note that there’s a big difference between prediction accuracy and being able to assess causality
It’s perfectly possible (and permissible) to predict x -> y, even if there’s no direct causal link!

Basic causal graphs

The Pipe (mediation)

CausalGraph[{Exposure["x"]->"m", "m"->Outcome["y"]}]

x= RandomReal[1,100];
m=0.8 x+ RandomReal[0.2,Length[x]];
y=1.25 m+RandomReal[0.1, Length[m]];
data=Transpose[Dataset@AssociationThread[{"x","m","y"},{x,m,y}]];
ListPlot[data[[All,{"x","y"}]],AxesLabel->Automatic]

ListPlot[colourPoints[data,{"x","y"}->"m","Rainbow",5]]

Impact of both features onto y:

Normal @ LinearModelFit[Normal@data[All,{#x,#m}->#y&],{1, \[FormalX],\[FormalM]},{ \[FormalX],\[FormalM]}]

Impact of just x:

Normal @ LinearModelFit[Normal@data[All,{#x}->#y&],{1, \[FormalX]},{ \[FormalX]}]

If you want to include m in your analysis, you have to measure the effect x->m and m->y separately and compose the results:

Normal @ LinearModelFit[Normal@data[All,{#x}->#m&],{1, \[FormalX]},{ \[FormalX]}]

Normal @ LinearModelFit[Normal@data[All,{#m}->#y&],{1, \[FormalM]},{ \[FormalM]}]

Coefficient[,\[FormalX]]*Coefficient[,\[FormalM]]

Conclusion you generally should not control for mediators, because that will make the causal effect more difficult to see.

CausalGraph[{Exposure["x"]->Adjusted["m"], "m"->Outcome["y"]}]

However, if you only want to predict y accurately, it is better to include m in your model.

Basic causal graphs

The Collider (common effect)

CausalGraph[
{Exposure["x"]->Outcome["y"], "x"->"c","y"->"c"},
]

x= RandomReal[1,100];
y= RandomReal[1,Length[x]]; (* no direct effect x -> y *)
c=x+y+RandomReal[0.1, Length[x]];
data=Transpose[Dataset@AssociationThread[{"x","c","y"},{x,c,y}]];
ListPlot[data[[All,{"x","y"}]],AxesLabel->Automatic]

ListPlot[colourPoints[data,{"x","y"}->"c","Rainbow",5]]

Impact of x on y if you include c:

Normal @ LinearModelFit[Normal@data[All,{#x,#c}->#y&],{1, \[FormalX],\[FormalC]},{ \[FormalX],\[FormalC]}]

Impact of just x:

Normal @ LinearModelFit[Normal@data[All,{#x}->#y&],{1, \[FormalX]},{ \[FormalX]}]

Conclusion you generally should not control for colliders, because that create spurious correlations

CausalGraph[{Exposure["x"]->Outcome["y"], "x"->Adjusted["c"],"y"->"c"},
]

Basic causal graphs

The Descendant (measure-by-proxy)

CausalGraph[
{"x","m","y","z"},
{Exposure["x"]->"m", "m"->Outcome["y"],"m"->"z"},

]

Controlling for a descendant is similar to controlling for the variable that is influencing it, but it only works well if the link is strong.

Basic causal graphs

The Instrument

If you can’t control for u:

CausalGraph[{Exposure["x"]->Outcome["y"],"u"->"x","u"->"y"}]

There might be something that can nudge x in a way that it sets up a natural “experiment”:

CausalGraph[
{Exposure["t"]->"x", "x"->Outcome["y"],"u"->"x","u"->"y"},

]

Some more advanced examples

graph = CausalGraph[{Exposure["x"],Outcome["y"],"z", "a","b"},{"x"->"y", "a"->"x","a"->"z","b"->"y","b"->"z","z"->"x","z"->"y"},

]

FindIndependenceConditions[graph]

UpdateCausalInformation[graph, "z"->Adjusted]

UpdateCausalInformation[graph, {"z"->Adjusted,"a"->Adjusted}]

BlockBackdoorPaths[graph]

graph2=UpdateCausalInformation[graph,"a"->Unobserved]

BlockBackdoorPaths[graph2]

Application: smoking (see [1])

Problem: how do you exclude the effect of unobserved confounders?

With[{
s="smoking", t="tar in lungs", g="genetics", c="cancer"
},
CausalGraph[{Unobserved[g] -> s, Exposure[s]->Outcome[c],g->c},
VertexCoordinates->{s->{0,0},
t->{1,0},c->{2,0},g->{1,1}
}]
]

Answer: investigate the causal pathway in more detail

The full path is still confounded:

With[{
s="smoking", t="tar in lungs", g="genetics", c="cancer"
},
CausalGraph[{Unobserved[g] -> s, Exposure[s]->t, t->Outcome[c],g->c},
VertexCoordinates->{s->{0,0},
t->{1,0},c->{2,0},g->{1,1}
}]
]

But the two components of the path can be identified!

Unconfounded because of the collider:

With[{
s="smoking", t="tar in lungs", g="genetics", c="cancer"
},
CausalGraph[{Unobserved[g] -> s, Exposure[s]->Outcome[t], t->c,g->c},
VertexCoordinates->{s->{0,0},
t->{1,0},c->{2,0},g->{1,1}
}]
]

Smoking is observable and can be used to block the backdoor path:

With[{
s="smoking", t="tar in lungs", g="genetics", c="cancer"
},
CausalGraph[{Unobserved[g] -> Adjusted[s], s->Exposure[t], t->Outcome[c],g->c},
VertexCoordinates->{s->{0,1/2},
t->{0,0},c->{2,0},g->{1,1}
}]
]

References

[1] Judea Pearl, The Book of Why (2018), ISBN 978-0-241-24263-6

[2] Johannes Textor, Benito van der Zander, Mark K. Gilthorpe, Maciej Liskiewicz, George T.H. Ellison. Robust causal inference using directed acyclic graphs: the R package ‘dagitty’. International Journal of Epidemiology 45(6):1887-1894, 2016. 
https://dagitty.net/

Initialization

Clear[colourPoints];
colourPoints[data_Dataset, rest___] := colourPoints[Normal[data], rest];
colourPoints[data : {} | {{}}, ___] := data;
colourPoints[
	data_, cols : {__} -> colorColumn_, 
	colFun_, nColors : _Integer | Automatic
] := Module[{
   points = Values @ data[[All, cols]],
   fun = Replace[colFun, s_String :> ColorData[s]],
   colors
   },
  colors = data[[All, colorColumn]];
  If[ ! MatchQ[colors, {__?NumericQ}],
   	colors = ArrayComponents[colors]
   ];
  If[MatchQ[colors, {__Integer}] && 
    MatchQ[fun, HoldPattern[ColorDataFunction[_, "Indexed", __]]],
   colors = fun[colors]
   ,
   If[nColors === Automatic,
    colors = Function[n, fun[n], Listable][Rescale[colors]],
    colors = 
     Function[n, fun[n], Listable][
      Ceiling[Rescale[colors, MinMax[colors], {$MachineEpsilon, 1}], 
       1./nColors]]
    ]
   ];
  Thread[Style[points, colors]]
];


Unprotect @@ Names["CausalInference`" ~~ __];
ClearAll @@ Names["CausalInference`" ~~ __];

BeginPackage["CausalInference`"];



GeneralUtilities`SetUsage[CausalGraphQ,
	"CausalGraphQ[gr$] returns true if gr$ is a simple, directed and acyclic graph."
];
GeneralUtilities`SetUsage[CausalGraph,
	"CausalGraph[gr$, props$] constructs a causal graph."
];

Unobserved;
Adjusted;
Outcome;
Exposure;

CausalInformation;
UpdateCausalInformation;
CausalVertexInformation;

GeneralUtilities`SetUsage[CausalDo,
	"CausalDo[gr$, list$] deletes all directed edges in causal graph gr$ that go into any of the vertices in list$."
];

GeneralUtilities`SetUsage[CausalMediatingVertices,
	"CausalMediatingVertices[gr$, s$ -> t$] gives all vertices that are mediating the causal effect of s$ on t$."
];

GeneralUtilities`SetUsage[BlockBackdoorPaths,
	"BlockBackdoorPaths[gr$, s$ -> t$, list$] finds all possible subsets of condition candidates in list$ that block the backdoor paths between s$ and t$ in causal graph gr$."
];
GeneralUtilities`SetUsage[BackdoorPathsClosedQ,
	"BackdoorPathsClosedQ[gr$, s$ -> t$, list$] returns True if all backdoor paths between s$ and t$ in causal graph gr$ are closed given a set of condition variables in list$."
];

CorrectlyAdjustedCausalGraphQ;

ConditionallyIndependentQ;
FindIndependenceConditions;

GeneralUtilities`SetUsage[Causes,
	"Causes[gr$] finds all vertices that are the cause of another vertex in the causal graph gr$.
Causes[gr$, list$] finds the direct causes of vertices in list$."
];
GeneralUtilities`SetUsage[Effects,
	"Effects[gr$] finds all vertices that are the effect of another vertex in the causal graph gr$.
Effects[gr$, list$] finds the direct effects of vertices in list$."
];

GeneralUtilities`SetUsage[RootCauses,
	"RootCauses[gr$] finds all vertices in causal graph gr$ that are not descendants of any other vertex."
];
GeneralUtilities`SetUsage[RootEffects,
	"RootEffects[gr$] finds all vertices in causal graph gr$ that are not ancestors of any other vertex."
];


InformationPaths::usage = "InformationPaths finds all paths between two nodes in a causal graph along which information can flow and annotates the vertices with their causal type along that path (\"Pipe\", \"Fork\", \"Collider\" or \"Descendant\").";

ToUndirectedGraph::usage = "ToUndirectedGraph[gr] replaces all directed edges in a graph to undirected edges.";


InformationPath::usage = "InformationPath[...] represents a path through a causal graph in terms of ForwardArrow and BackwardArrow.";
InformationPathQ;
ForwardArrow::usage = "ForwardArrow represents a directed edge in an information path.";
BackwardArrow::usage = "BackwardArrow represents a directed edge in an information path.";


MarkovDecomposition::usage = "MarkovDecomposition[gr] give the joint probability distribution of the vertices in gr.";
FeatureDependencies::usage = "FeatureDependencies[gr] returns a list of rules corresponding to the feature dependencies encoded by the graph.";

ℙ;

CausalGraphTrain;

Begin["`Private`"];


Scan[
	Function[sym,
		sym[_] := $Failed;
		sym[_, _, __] := $Failed;
	],
	{ForwardArrow, BackwardArrow}
];

InformationPathQ[expr_] := And[
	MatchQ[expr, InformationPath[(_ForwardArrow | _BackwardArrow)..]],
	With[{
		list = List @@ expr
	},
		SameQ[list[[1 ;; -2 , 2]], list[[2 ;; All, 1]]]
	]
];

CausalPathExistsQ[gr_?GraphQ, s_ -> t_] := And[
	s =!= t,
	MatchQ[
		FindPath[gr, s, t, Infinity, 1],
		{{__}..}
	]
];

CausalMediatingVertices[gr_Graph, s_ -> t_] := Replace[
	Quiet @ FindPath[gr, s, t, Infinity, All],
	{
		list_List :> DeleteCases[Union @@ list, s | t],
		_ :> $Failed
	}
];

taggedCausalGraphQ[gr_] := GraphQ[gr] && TrueQ @ AnnotationValue[gr, "CausalGraph"];
untaggedCausalGraphQ[gr_] := GraphQ[gr] && !taggedCausalGraphQ[gr];
tagCausalGraph[gr_Graph?GraphQ] := If[ taggedCausalGraphQ[gr],
	gr,
	If[ CausalGraphQ[gr],
		Annotate[gr, "CausalGraph" -> True],
		$Failed
	]
];

CausalGraphQ[gr_Graph?GraphQ] := Or[
	taggedCausalGraphQ[gr],
	SimpleGraphQ[gr] && DirectedGraphQ[gr] && AcyclicGraphQ[gr]
];

causalWrapper = Unobserved | Adjusted | Outcome | Exposure;
causalWrapperHeads = List @@ causalWrapper;

causalVertexStyles = <|
	Unobserved -> {White},
	Adjusted -> {Blue},
	Exposure -> {Yellow},
	Outcome -> {Orange}
|>;
causalEdgeStyles = <|
	(* "NonCausalBlockedEdges" -> None, *)
	"NonCausalOpenEdges" -> Directive[Red, Dashed],
	"CausalOpenEdges" -> Directive[Blue],
	"CausalBlockedEdges" -> Directive[Blue, Dashed]
|>;

openFilter["Open"][boole_] := boole;
openFilter["Closed"][boole_] := !boole;
openFilter[All][boole_] := True;

pathTypeEdges[type_] := pathTypeEdges[type, "Open"];
pathTypeEdges[type_, open_][gr_Graph] := Module[{
	info = reverseInfo[gr],
	direction,
	adjusted
},
	direction = causalDirections @ info;
	adjusted = allAdjustedVertices @ info;
	Switch[ direction,
		_Rule,
			pathTypeEdges[type, open][gr, direction, adjusted],
		{___Rule},
			DeleteDuplicates @ Flatten @ Map[
				pathTypeEdges[type, open][gr, #, adjusted]&,
				direction
			],
		_,
			{}
	]
];
pathTypeEdges[type_, open_][gr_Graph, type_, dir_Rule, adjusted_List] := pathTypeEdges[type, open][
	InformationPaths[gr, dir, type, adjusted]
];

pathTypeEdges[type_, open_][paths_List] :=  DeleteDuplicates @ Flatten @ Map[EdgeList] @ Lookup[
	Select[paths, Function[openFilter[open][#OpenQ] && MatchQ[#PathType, type]]],
	"Graph",
	{}
];
pathTypeEdges[__][__] := {};

classifyEdges[gr_] := Module[{
	info = reverseInfo[gr],
	direction
},
	direction = causalDirections @ info;
	If[ MatchQ[direction, _Rule | {___Rule}],
		Replace[
			classifyEdges[gr, direction, allAdjustedVertices @ info],
			Except[_Association]?(AllTrue[ListQ]) :> <||>
		],
		<||>
	]
];
classifyEdges[gr_, dir_Rule, adjusted_] := Module[{
	paths = InformationPaths[gr, dir, All, adjusted]
},
	mergeClassifications @ Map[Identity] @ Comap[
		<|
			"CausalOpenEdges" -> pathTypeEdges["Causal", "Open"],
			"CausalBlockedEdges" -> pathTypeEdges["Causal", "Closed"],
			"NonCausalOpenEdges" -> pathTypeEdges[Except["Causal"], "Open"],
			"NonCausalBlockedEdges" -> pathTypeEdges[Except["Causal"], "Closed"]
		|>,
		paths
	]
];
classifyEdges[gr_, dirList_List, adjusted_] := <||>; (* TODO: handle multi exposure/outcome scenarios *)
(* 
classifyEdges[gr_, dirList_List, adjusted_] := mergeClassifications @ Map[
	classifyEdges[gr, #, adjusted]&,
	dirList
];
*)

mergeClassifications[list : {___Association}] := mergeClassifications @ Merge[
	list,
	Apply[Join] /* DeleteDuplicates
];
mergeClassifications[assoc_Association] := Module[{
	backdoors = Lookup[assoc, "NonCausalOpenEdges", {}],
	blockedCausal = Lookup[assoc, "CausalBlockedEdges", {}]
},
	DeleteCases[{}] @ Association[
		Map[
			Complement[#, backdoors, blockedCausal]&,
			KeyTake[assoc, {"CausalOpenEdges", "NonCausalBlockedEdges"}]
		],
		"CausalBlockedEdges" -> Complement[
			Lookup[assoc, "CausalBlockedEdges", {}],
			backdoors
		],
		"NonCausalOpenEdges" -> backdoors
	]
];

styleCausalGraph[gr_Graph] := With[{
	vertexStyles = Normal @ Map[
		Lookup[causalVertexStyles, #, Automatic]&,
		causalInfo[gr]
	],
	edgeStyles = DeleteDuplicatesBy[First] @ Join[
		Flatten @ Values @ Merge[
			KeyIntersection[{causalEdgeStyles, classifyEdges[gr]}],
			Apply[
				Function[{style, edges},
					Map[
						Function[# -> style],
						edges
					]
				]
			]
		],
		Thread[EdgeList[gr] -> None] (* Default *)
	]
},
	Annotate[gr,
		{
			VertexStyle -> vertexStyles,
			If[ edgeStyles === {}, Nothing, EdgeStyle -> edgeStyles]
		}
	]
];
styleCausalGraph[other_] := other;

removeCausalWrappers[causalWrapper[vertex_]] := vertex;
removeCausalWrappers[list_List] := removeCausalWrappers /@ list;
removeCausalWrappers[(edgeHead : DirectedEdge | UndirectedEdge)[a_, b_, rest___]] := edgeHead[
	removeCausalWrappers[a],
	removeCausalWrappers[b],
	rest
];
removeCausalWrappers[vertex_] := vertex;

removeStyling[gr_] := Graph[VertexList[gr], EdgeList[gr]];

edgesPatt = {(Rule | DirectedEdge)[_, __]..};
toDirectedEdges[list : {___DirectedEdge}] := list;
toDirectedEdges[list_List] := DirectedEdge @@@ list;

CausalGraph[list : edgesPatt, rest___] := With[{
	edges = toDirectedEdges[list]
},
	CausalGraph[
		findVertices[edges],
		edges,
		rest
	]
];
CausalGraph[vertices_List, list : edgesPatt, rest___] := Module[{
	edges = toDirectedEdges[list],
	opts = {rest},
	gr, allVerts
},
	Enclose[
		allVerts = Join[vertices, Apply[Join, List @@@ edges]];
		allVerts = findVertices[allVerts];
		gr = Graph[
			Replace[allVerts, (h : causalWrapper)[vert_] :> Annotation[vert, CausalInformation -> h], {1}],
			removeCausalWrappers[edges],
			Sequence @@ opts,
			VertexLabels -> Automatic
		];
		gr = Confirm[tagCausalGraph[gr]];
		styleCausalGraph[gr]
	]
];
CausalGraph[gr_Graph] := styleCausalGraph @ tagCausalGraph[gr];

CausalVertexInformation[gr_Graph, v : _ : All] := With[{
	res = If[ CausalGraphQ[gr],
		Replace[
			causalInfo[gr, v],
			{
				r_Rule :> Keys[r],
				a_Association :> Normal[a],
				_ :> $Failed
			}
		],
		None
	]
},
	res /; !FailureQ[res]
];

causalInfo[gr_] := causalInfo[gr, All]
causalInfo[gr_, All] := causalInfo[gr, VertexList[gr]];
causalInfo[gr_, v_List] := Association @ Map[causalInfo[gr, #]&, v];
causalInfo[gr_, v_] /; VertexQ[gr, v] := v -> Replace[
	AnnotationValue[{gr, v}, CausalInformation],
	_?FailureQ :> None
];

addAllWrapperHeads[assoc_] := Join[Association @ Map[Function[# -> {}], causalWrapperHeads], assoc];

reverseInfo[gr_] := addAllWrapperHeads @ GroupBy[
	Normal @ causalInfo[gr],
	Last -> First
];

allAdjustedVertices[gr_Graph] := allAdjustedVertices[reverseInfo[gr]];
allAdjustedVertices[revInfo_Association] := Join @@ Values[KeyTake[revInfo, {Adjusted}]];

causalDirections[info_] := Switch[{Length @ info[Exposure], Length @ info[Outcome]},
	{0, _} | {_, 0},
		$Failed,
	{1, 1},
		First @ info[Exposure] -> First @ info[Outcome],
	_,
		Flatten @ Outer[Rule, info[Exposure], info[Outcome]]
];

UpdateCausalInformation[gr_?untaggedCausalGraphQ, rest__] := UpdateCausalInformation[
	tagCausalGraph[gr],
	rest
];
UpdateCausalInformation[gr_?CausalGraphQ, v_ -> status : (causalWrapper | None)] /; VertexQ[gr, v] := styleCausalGraph[
	If[status === None,
		AnnotationDelete[{gr, v}, CausalInformation],
		Annotate[{gr, v}, CausalInformation -> status]
	]
];
UpdateCausalInformation[gr_?CausalGraphQ, rules : {___Rule}] := Catch[
	styleCausalGraph @ Block[{
		styleCausalGraph = Function[#]
	},
		Fold[
			Replace[
				UpdateCausalInformation[#1, #2],
				Except[_?GraphQ] :> Throw[$Failed, UpdateCausalInformation]
			]&,
			gr,
			rules
		]
	],
	UpdateCausalInformation
];
UpdateCausalInformation[
	gr_?CausalGraphQ,
	assoc_Association?AssociationQ
] := UpdateCausalInformation[
	gr,
	Flatten[Thread /@ Reverse[Normal[assoc], 2]]
] /; ContainsOnly[Keys[KeyDrop[None] @ assoc], causalWrapperHeads];
UpdateCausalInformation[_, _] := $Failed;

findVertices[edges : edgesPatt] := findVertices[Join @@ MapApply[List, edges]];
findVertices[vertices_] := DeleteDuplicatesBy[
	Join[
		Cases[vertices, causalWrapper[_]],
		vertices
	],
	removeCausalWrappers
];

ToUndirectedGraph[gr_?GraphQ] := With[{
	edges = EdgeList[gr],
	vert = VertexList[gr]
},
	Graph[
		vert,
		UndirectedEdge @@@ edges
	]
];
ToUndirectedGraph[___] := $Failed;

CausalDo[gr_, v_] := Module[{
	doVertices = v,
	allVertices
},
	Enclose[
		ConfirmAssert[CausalGraphQ @ gr];
		allVertices = VertexList[gr];
		If[ MemberQ[allVertices, doVertices],
			doVertices = {doVertices},
			ConfirmAssert[
				SubsetQ[allVertices, doVertices],
				"Invalid vertex specification"
			]
		];
		ConfirmBy[
			EdgeDelete[gr, DirectedEdge[_, Alternatives @@ doVertices]],
			CausalGraphQ
		]
	]
];

MapApply[
	Function[{sym, str},
		MakeBoxes[expr : sym[a_, b_], form : StandardForm] := With[{
			boxes = RowBox[{
				MakeBoxes[a, form],
				str,
				MakeBoxes[b, form]
			}]
		},
			InterpretationBox[boxes, expr]
		];
	],
	{
		{ForwardArrow, "→"},
		{BackwardArrow, "←"}
	}
];


MakeBoxes[path : InformationPath[__]?InformationPathQ, form : StandardForm] := With[{
	boxes = RowBox[Flatten @ {
		MakeBoxes[InformationPath, form], "[",
		Append[ToBoxes[path[[-1, -1]], form]] @ ReplaceAll[
			List @@ path,
			{
				ForwardArrow[a_, _] :> {MakeBoxes[a, form], "→"},
				BackwardArrow[a_, _] :> {MakeBoxes[a, form], "←"}
			}
		],
		 "]"
	}]
},
	InterpretationBox[boxes, path]
];

$pipe = "Pipe";
$collider = "Collider";
$fork = "Fork";
$pathNode = $pipe | $collider | $fork;

Clear[classifyVertices];
classifyVertices[path : InformationPath[__]] := Join[
	{path[[1, 1]] -> {"Type" -> "Source"}},
	Replace[
		Partition[List @@ path, 2, 1],
		{
			{ForwardArrow[_, b_], ForwardArrow[b_, _]} :> Rule @@ {b, {"Type" -> $pipe}},
			{BackwardArrow[_, b_], BackwardArrow[b_, _]} :> Rule @@ {b, {"Type" -> $pipe}},
			{ForwardArrow[_, b_], BackwardArrow[b_, _]} :> Rule @@ {b, {"Type" -> $collider}},
			{BackwardArrow[_, b_], ForwardArrow[b_, _]} :> Rule @@ {b, {"Type" -> $fork}}
		},
		{1}
	],
	{path[[-1, -1]] -> {"Type" -> "Target"}}
];

Clear[augmentedPathGraph];
augmentedPathGraph[fullGraph_, vertexPath_List, rules_, cond_] := Module[{
	innerVertices = vertexPath[[2 ;; -2]],
	pathVertices = Alternatives @@ vertexPath,
	informationPath = directedPath[vertexPath, rules],
	descendantEdges,
	annotations,
	pathEdges,
	pathType,
	vertexClassification,
	assoc
},
	pathType = Switch[informationPath,
		InformationPath[__ForwardArrow],
			"Causal",
		InformationPath[_BackwardArrow, ___],
			"Backdoor",
		_,
			"NonCausal"
	];
	descendantEdges = Flatten @ Map[
		Thread @ DirectedEdge[
			#,
			DeleteCases[
				VertexOutComponent[fullGraph, #],
				pathVertices
			]
		]&,
		innerVertices
	];
	vertexClassification = Join[
		KeyValueMap[
			Function[#1 -> {"Type" -> "Descendant", "From" -> #2}],
			GroupBy[Last -> First] @ descendantEdges
		],
		classifyVertices[informationPath]
	];
	If[ vertexClassification =!= {},
		annotations = MapAt[
			With[{
				type = Lookup[#, "Type"]
			},
				{"Type" -> type, Tooltip -> type}
			]&,
			vertexClassification,
			{All, 2}
		]
		,
		annotations = {}
	];
	pathEdges = directedPathToEdges[informationPath];
	assoc = <|
		"PathType" -> pathType,
		"Graph" -> Graph[
			GraphUnion[
				Graph[pathEdges],
				Graph[descendantEdges, VertexLabels -> Automatic]
			],
			AnnotationRules -> annotations,
			GraphHighlight -> Join[pathEdges, vertexPath[[{1, -1}]]],
			FilterRules[Options[fullGraph], Except[VertexCoordinates]]
		],
		"InformationPath" -> informationPath,
		"VertexClassification" -> MapApply[
			Function[
				Prepend["Vertex" -> #1] @ Association[#2]
			],
			vertexClassification
		]
	|>;
	assoc["OpenQ"] = pathOpenQ[assoc, cond, False];
	assoc
];

makePathRules[DirectedEdge[left_, right_]] := {
	{left, right} :> ForwardArrow[left, right],
	{right, left} :> BackwardArrow[right, left]
};

directedPath[vertexSequence_List, rules_] := InformationPath @@ Replace[
	Partition[vertexSequence, 2, 1],
	rules,
	{1}
];

directedPathToEdges[path_InformationPath] := Replace[
	List @@ path,
	{
		ForwardArrow[a_, b_] :> DirectedEdge[a, b],
		BackwardArrow[a_, b_] :> DirectedEdge[b, a]
	},
	{1}
];


InformationPaths[gr_] := InformationPaths[gr, Automatic];
InformationPaths[gr_?CausalGraphQ, Automatic] := Module[{
	info = reverseInfo[gr]
},
	Enclose[
		Select[AssociationQ] @ Flatten @ Outer[
			Confirm @ InformationPaths[gr, #1 -> #2, All, info]&,
			info[Exposure],
			info[Outcome]
		]
	 ] /; Length @ info[Exposure] >= 1 && Length @ info[Outcome] >= 1
];
InformationPaths[gr_, s_ -> t_] := InformationPaths[gr, s -> t, All];
InformationPaths[gr_, s_ -> t_, type_] := InformationPaths[gr, s -> t, type, Automatic];
InformationPaths[gr_?CausalGraphQ, s_ -> t_, type_, Automatic] := InformationPaths[gr, s -> t, type, reverseInfo[gr]];
InformationPaths[gr_?GraphQ, s_ -> t_, "Backdoor", rest___] := InformationPaths[
	EdgeDelete[gr, DirectedEdge[s, _]] (* Backdoor paths always start with an arrow going into the start point *),
	s -> t,
	All,
	rest
];
InformationPaths[gr_, dir_, type_, condList_List] := InformationPaths[gr, dir, type,
	addAllWrapperHeads @ <|Adjusted -> condList|>
];
InformationPaths[gr_?GraphQ, s_ -> t_, type_, causalAssoc_?AssociationQ] := If[
	s === t || ! SubsetQ[VertexList[gr], {s, t}]
	,
	Enclose[Confirm[$Failed, "Invalid source or target vertex"]]
	,
	Module[{
		graph = gr,
		adjusted = DeleteCases[causalAssoc[Adjusted], s | t],
		replacements,
		directedPaths
	},
		Enclose[
			ConfirmAssert[CausalGraphQ @ graph];
			replacements = Flatten @ Map[makePathRules, EdgeList[gr]];
			directedPaths = Map[
				ConfirmBy[
					augmentedPathGraph[graph, #, replacements, adjusted],
					AssociationQ
				]&,
				FindPath[
					Switch[ type,
						"Causal",
							graph,
						_,
							ConfirmBy[ToUndirectedGraph[graph], GraphQ]
					],
					s,
					t,
					Infinity,
					All
				]
			];
			ConfirmMatch[directedPaths, {___?AssociationQ}];
			If[ type === "NonCausal",
				directedPaths = Select[directedPaths, !MatchQ[#PathType, "Causal"]&]
			];
			directedPaths
		]
	]
];
InformationPaths[___] := $Failed;

(* pathAssoc is one of the associations returned by augmentedPathGraph *)
pathClosedQ[set_][assoc_] := pathClosedQ[assoc, set];
(* TODO: vectorize efficiently over multiple paths for a given condition set *)
pathClosedQ[pathAssoc_?AssociationQ, conditionSet_List, descendantConditionedQ : (_?BooleanQ) : False] := Module[{
	conditionedVerticesOnPath = conditionSet,
	descendants = If[ TrueQ @ descendantConditionedQ,
		Cases[
			pathAssoc["VertexClassification"],
			KeyValuePattern[{
				"Vertex" -> Alternatives @@ conditionSet,
				"Type" -> "Descendant"
			}]
		],
		{}
	],
	pathNodes = Cases[
		pathAssoc["VertexClassification"],
		KeyValuePattern[{"Type" -> $pathNode}]
	]
},
	(* Add the vertices that are conditioned on by proxy through descendants *)
	conditionedVerticesOnPath = DeleteDuplicates @ Join[
		conditionedVerticesOnPath,
		Join @@ Lookup[descendants, "From", {}]
	];
	AnyTrue[pathNodes,
		Function[
			Or[ (* Check if any of the nodes on the path are blocked *)
				#Type === $collider && !MemberQ[conditionedVerticesOnPath, #Vertex], (* colliders block when not conditioned on *)
				MatchQ[#Type, $fork | $pipe] && MemberQ[conditionedVerticesOnPath, #Vertex] (* pipes and forks only block when conditioned on *)
			]
		]
	]
];
pathClosedQ[_, __] := $Failed;

pathOpenQ[set_][assoc_] := pathOpenQ[assoc, set];
pathOpenQ[assoc_, set_, boole___] := Replace[pathClosedQ[assoc, set, boole], b_?BooleanQ :> !b];

ConditionallyIndependentQ[gr_, st : {_, _}] := ConditionallyIndependentQ[gr, st, {}];
ConditionallyIndependentQ[gr_?CausalGraphQ, {s_, t_}, cond_List] := With[{
	paths = InformationPaths[gr, s -> t, All, cond]
},
	If[ ListQ[paths],
		AllTrue[paths, Function[#OpenQ === False]],
		$Failed
	]
];

BackdoorPathsClosedQ[gr_?CausalGraphQ, s_ -> t_, set_List, opts : OptionsPattern[]] := Module[{
	paths,
	conditionSet = set,
	graph = gr
},
	Enclose[
		ConfirmAssert[
			CausalPathExistsQ[graph, s -> t],
			"There is no causal path from between the source and target vertices."
		];
		paths = Confirm @ InformationPaths[graph, s -> t, "NonCausal", conditionSet];
		AllTrue[paths, Function[#OpenQ === False]]
	]
];

CorrectlyAdjustedCausalGraphQ[gr_?CausalGraphQ] := Module[{
	info = reverseInfo[gr]
},
	TrueQ @ Enclose[
		Outer[
			Function[
				ConfirmAssert @ CausalPathExistsQ[gr, #1 -> #2];
				ConfirmAssert[!MemberQ[info[Adjusted], CausalMediatingVertices[gr, #1 -> #2]]];
				ConfirmAssert @ BackdoorPathsClosedQ[gr, #1 -> #2, info[Adjusted]];
			],
			info[Exposure],
			info[Outcome]
		];
		True
	]
];
CorrectlyAdjustedCausalGraphQ[_] := False;

BlockBackdoorPaths[gr_] := BlockBackdoorPaths[gr, Automatic];
BlockBackdoorPaths[gr_?CausalGraphQ, Automatic] := Module[{
	info = reverseInfo[gr]
},
	Enclose[
		Apply[Intersection] @ Apply[Join] @ Outer[
			Confirm @ BlockBackdoorPaths[gr, #1 -> #2, info]&,
			info[Exposure],
			info[Outcome]
		]
	 ] /; Length @ info[Exposure] >= 1 && Length @ info[Outcome] >= 1
];
BlockBackdoorPaths[gr_, s_ -> t_] := BlockBackdoorPaths[gr, s -> t, Automatic];
BlockBackdoorPaths[gr_?CausalGraphQ, s_ -> t_, Automatic] := BlockBackdoorPaths[gr, s -> t, reverseInfo[gr]];
BlockBackdoorPaths[gr_, dir_, condList_List] := BlockBackdoorPaths[gr, dir,
	addAllWrapperHeads @ <|Unobserved -> condList|>
];

BlockBackdoorPaths[gr_?CausalGraphQ, s_ -> t_, vertexInfo_Association] := Module[{
	graph = gr,
	conditionCandidates = Complement[VertexList[gr], vertexInfo[Unobserved], {s, t}, CausalMediatingVertices[gr, s -> t]],
	alreadyAdjusted = vertexInfo[Adjusted],
	paths,
	sets,
	descendantConditionedQ = False
},
	Enclose[
		ConfirmAssert[
			CausalPathExistsQ[graph, s -> t],
			"There is no causal path from between the source and target vertices."
		];
		paths = Confirm @ InformationPaths[graph, s -> t, "NonCausal", {}];
		sets = Select[
			ConfirmMatch[Subsets[conditionCandidates], {__List}],
			ContainsAll[alreadyAdjusted]
		];
		ConfirmMatch[
			Select[
				sets,
				Function[set,
					AllTrue[paths, pathClosedQ[#, set, descendantConditionedQ]&]
				]
			],
			{___List}
		]
	]
];

FindIndependenceConditions[gr_?CausalGraphQ] := Module[{
	infoAssoc = reverseInfo[gr],
	vList = VertexList[gr],
	observables,
	adjustableVertices,
	adjustedVertices,
	conditionSets,
	independencePairs
},
	Enclose[
		observables = Complement[vList, infoAssoc[Unobserved]];
		adjustedVertices = infoAssoc[Adjusted];
		adjustableVertices = observables;
		conditionSets = Select[
			ConfirmMatch[Subsets[adjustableVertices], {__List}],
			ContainsAll[adjustedVertices]
		];
		independencePairs = Subsets[observables, {2}];
		Flatten @ Outer[
			Function[
				If[ IntersectingQ[#1, #2] || !ConditionallyIndependentQ[gr, #1, #2],
					{},
					If[ #2 === {},
						Perpendicular @@ #1,
						Conditioned[Perpendicular @@ #1, #2]
					]
				]
			],
			independencePairs,
			conditionSets,
			1
		]
		
	]
];

(* Define Cause and Effects *)
MapApply[
	Function[{sym, index},
		sym[gr_?CausalGraphQ] := Replace[
			EdgeList[gr],
			{
				list : {__} :> DeleteDuplicates @ list[[All, index]],
				Except[_List] :> $Failed
			}
		];
		sym[___] := $Failed
	]
	,
	{
		{Causes, 1},
		{Effects, 2}
	}
];

(* Define RootCauses and RootEffects in terms of Causes and Effects *)
MapApply[
	Function[{sym, sym2},
		sym[gr_?CausalGraphQ] :=  Enclose[
			DeleteCases[
				VertexList[gr],
				Alternatives @@ Confirm[sym2[gr]]
			]
		];
		sym[___] := $Failed
	]
	,
	{
		{RootCauses, Effects},
		{RootEffects, Causes}
	}
];

MarkovDecomposition[gr_] := Module[{
	graph = gr,
	sourceVertices,
	edges,
	assoc
},
	Enclose[
		ConfirmAssert[
			CausalGraphQ[graph],
			"The graph is not a causal graph."
		];
		edges = EdgeList[gr];
		assoc = GroupBy[edges, Last -> First, Apply[And]];
		sourceVertices = VertexList[gr, Except[Alternatives @@ Keys[assoc]]];
		Times[
			Times @@ (ℙ /@ sourceVertices),
			Times @@ KeyValueMap[
				Function[
					ℙ[Conditioned[#1, #2]]
				],
				assoc
			]
		]
	]
];

FeatureDependencies[graph_?CausalGraphQ] := Module[{
	edges = EdgeList[graph],
	rules
},
	Enclose[
		rules = ConfirmBy[
			GroupBy[
				ConfirmBy[edges, ListQ],
				Last -> First
			],
			AssociationQ
		];
		Reverse[Normal[rules], {2}]

	]
];
FeatureDependencies[___] := $Failed;

CausalGraphTrain[data_, graph_?GraphQ, rest__] := CausalGraphTrain[data, FeatureDependencies[graph], rest];

CausalGraphTrain[
	data : {__Association},
	featureDependencies_List,
	specs_?AssociationQ,
	opts : OptionsPattern[]
] := Module[{
	predictors,
	trainingData = KeyUnion @ data,
	featureSpecs = specs,
	keys
},
	Enclose[
		keys = Keys @ First[trainingData];
		ConfirmAssert[
			And[
				MatchQ[featureDependencies, {(_List -> _String)..}],
				SubsetQ[keys, Keys[featureSpecs]],
				SubsetQ[Keys[featureSpecs], Values[featureDependencies]],
				AllTrue[featureSpecs, StringQ[#] || AssociationQ[#]&]
			]
		];
		featureSpecs = Map[
			Replace[s_String :> <|"FeatureType" -> s|>],
			featureSpecs
		];
		predictors = Association @ MapApply[
			Function[{featuresIn, featureOut},
				(featuresIn -> featureOut) -> If[ featureSpecs[featureOut, "FeatureType"] === "Numerical",
					Predict,
					Classify
				][
					KeyTake[trainingData, Append[featuresIn, featureOut]] -> featureOut,
					Sequence @@ Lookup[featureSpecs[featureOut], "TrainingOptions", {}]
				]
			],
			featureDependencies
		]
	]
];

End[];


EndPackage[];