baogorek · July 15, 2021 23:33 · Jul 15, 2021
diff --git a/ananke.py b/ananke.py
@@ -0,0 +1,60 @@
+# pip install ananke-causal
+from ananke import graphs
+from ananke import identification
+from ananke.estimation import CausalEffect
+
+import numpy as np
+import pandas as pd
+
+# Simulate front-door situation with confounder Z
+N = 100000
+z = np.random.normal(size=N)
+x = .8 * z + np.random.normal(size=N) > 0
+x = x.astype(int)
+
+m = .5 * x + np.random.normal(size=N)
+y = .7 * z + 1.2 * m + np.random.normal(size=N)
+
+df = pd.DataFrame({'X': x, 'Y': y, "M": m, "Z": z})
+
+# Note that this is a front door graph where M is the mediator
+# DAG example with a single confounder and a front-door path
+vertices = ['X', 'Z', 'Y', 'M']
+edges = [
+    ('X', 'M'), ('M', 'Y'),  # Mediation path
+    ('Z', 'X'), ('Z', 'Y'),  # Confounding path
+]
+dag = graphs.DAG(vertices, edges)
+dag_graph = dag.draw(direction='LR')  # Need Graphviz installed
+dag_graph.view(filename="front_door")  # Wait for the browser to open (20 seconds)
+id_pya = identification.OneLineID(
+    graph=dag,
+    treatments=['X'],
+    outcomes=['Y']
+)
+id_pya.id()  # Is it identified?
+id_pya.functional()  # The Functional (have no idea what this is)
+
+ACE_estimand = CausalEffect(graph=dag, treatment='X', outcome='Y')
+ace = ACE_estimand.compute_effect(df, "eff-aipw")
+print(f"truth = {1.2 * .5} vs est = {np.round(ace, 4)} for {N=}")
+
+# ADMG front-door example with covariance instead of observed Z
+vertices = ['X', 'Y', 'M']
+di_edges = [('X', 'M'), ('M', 'Y')]  # Mediation path
+bi_edges = [('X', 'Y')]  # Confounding path
+admg = graphs.ADMG(vertices, di_edges=di_edges, bi_edges=bi_edges)
+digraph = admg.draw(direction='LR')
+digraph.view(filename="front_door_ADMG")  # Wait for the browser to open (20 seconds)
+id_pya = identification.OneLineID(graph=admg, treatments=['X'],
+                                  outcomes=['Y'])
+id_pya.id()  # Is it identified?
+id_pya.functional()
+
+ACE_estimand2 = CausalEffect(graph=admg, treatment='X', outcome='Y')
+
+ace = ACE_estimand2.compute_effect(df, "eff-aipw")  # This doesn't work
+ace2 = ACE_estimand2.compute_effect(df, "p-ipw")
+ace3 = ACE_estimand2.compute_effect(df, "d-ipw")
+ace4 = ACE_estimand2.compute_effect(df, "apipw")
+print(f"truth = {1.2 * .5} vs est = {np.round(ace3, 3)} for {N=}")