Hierarchical Causal Models

.github/workflows/tests.yml

@@ -104,7 +104,7 @@ jobs:
 
       - name: Install dependencies
         run: |
-          sudo apt-get install graphviz
+          sudo apt-get install graphviz libgraphviz-dev
       - name: Test with pytest and generate coverage file
         run:
           uvx -p ${{ matrix.python-version }} --with tox-uv tox -e py

README.md

@@ -108,6 +108,7 @@ unimplemented for the last 15 years of publications including:
 | IDC Star                        | [Shpitser and Pearl, 2012](https://arxiv.org/abs/1206.5294)                         |
 | Surrogate Outcomes              | [Tikka and Karvanen, 2018](https://arxiv.org/abs/1806.07172)                        |
 | Counterfactual Transportability | [Correia, Lee, Bareinboim, 2022](https://proceedings.mlr.press/v162/correa22a.html) |
+| Hierarchical Causal Models      | [Weinstein and Blei, 2024](https://arxiv.org/abs/2401.05330)                        |
 
 Apply an algorithm to an Acyclic Directed Mixed Graph (ADMG) and a causal query
 to generate an estimand represented in the DSL like:

docs/source/hierarchical.rst

@@ -0,0 +1,3 @@
+Hierarchical Graphs
+===================
+.. automodapi:: y0.hierarchical

docs/source/index.rst

@@ -18,6 +18,7 @@ y0 |release| Documentation
     tian_id
     transport
     lvdags
+    hierarchical
     struct
     cli
 

src/y0/examples/hierarchical.py

@@ -0,0 +1,227 @@
+"""An example directory of hierarchical causal models."""
+
+from y0.dsl import A, B, C, D, U, Y, Z
+from y0.graph import NxMixedGraph
+from y0.hierarchical import HierarchicalCausalModel, HierarchicalStructuralCausalModel, QVariable
+
+__all__ = [
+    "confounder_hcm",
+]
+
+Q_A = QVariable(name="A")
+Q_Y = QVariable(name="Y")
+Q_Y_A = QVariable(name="Y", parents=frozenset([A]))
+Q_A_Z = QVariable(name="A", parents=frozenset([Z]))
+Q_Z = QVariable(name="Z")
+Q_Y_A_Z = QVariable(name="Y", parents=frozenset([A, Z]))
+
+
+def _make_confounder_hcm() -> HierarchicalCausalModel:
+    """Pytest fixture for the Confounder HCM in Figure 2 (a)."""
+    graph = HierarchicalCausalModel()
+    graph.add_observed_node(A)
+    graph.add_observed_node(Y)
+    graph.add_edge(U, A)
+    graph.add_edge(A, Y)
+    graph.add_edge(U, Y)
+    graph.add_subunits([A, Y])
+    return graph
+
+
+confounder_hcm = _make_confounder_hcm()
+
+
+def get_confounder_hcgm() -> HierarchicalCausalModel:
+    """Pytest fixture for the Confounder HCGM in Figure 2 (b)."""
+    hcm = HierarchicalCausalModel()
+    hcm.add_observed_node(A)
+    hcm.add_observed_node(Y)
+    hcm.add_observed_node(Q_A)
+    hcm.add_observed_node(Q_Y_A)
+    hcm.add_edge(U, Q_A)
+    hcm.add_edge(Q_A, A)
+    hcm.add_edge(A, Y)
+    hcm.add_edge(U, Q_Y_A)
+    hcm.add_edge(Q_Y_A, Y)
+    hcm.add_subunits([A, Y])
+    return hcm
+
+
+def get_confounder_hscm() -> HierarchicalStructuralCausalModel:
+    """Pytest fixture for Confounder HSCM in Section 2.3 Eqn. (8)."""
+    hscm = HierarchicalStructuralCausalModel()
+    hscm.add_observed_node(A)
+    hscm.add_observed_node(Y)
+    hscm.add_unobserved_node(U)
+    hscm.add_edge(U, A)
+    hscm.add_edge(A, Y)
+    hscm.add_edge(U, Y)
+    hscm.add_subunits([A, Y])
+    return hscm
+
+
+def get_school_confounder_interference_hcm() -> HierarchicalCausalModel:
+    """Pytest fixture for the Confounder Interference HCM in Figure 2 (e)."""
+    hcm = HierarchicalCausalModel()
+    hcm.add_observed_node(A)
+    hcm.add_observed_node(Y)
+    hcm.add_observed_node(Z)
+    hcm.add_edge(U, A)
+    hcm.add_edge(A, Y)
+    hcm.add_edge(U, Y)
+    hcm.add_edge(A, Z)
+    hcm.add_edge(Z, Y)
+    hcm.add_subunits([A, Y])
+    return hcm
+
+
+def get_school_confounder_interference_hcgm() -> HierarchicalCausalModel:
+    """Pytest School fixture for the Confounder Interference HCGM in Hierarchical Causal Models notebook."""
+    hcm = HierarchicalCausalModel()
+    hcm.add_observed_node(A)
+    hcm.add_observed_node(Y)
+    hcm.add_observed_node(Z)
+    hcm.add_observed_node(Q_A)
+    hcm.add_observed_node(Q_Y_A)
+    hcm.add_edge(U, Q_A)
+    hcm.add_edge(Q_A, A)
+    hcm.add_edge(A, Y)
+    hcm.add_edge(U, Q_Y_A)
+    hcm.add_edge(Q_Y_A, Y)
+    hcm.add_edge(A, Z)
+    hcm.add_edge(Z, Q_Y_A)
+    hcm.add_subunits([A, Y])
+    return hcm
+
+
+def get_school_confounder_interference_hscm() -> HierarchicalStructuralCausalModel:
+    """Pytest fixture for the Confounder Interference HSCM in Sec. 2.3 Eqn. (9)."""
+    student_tutoring, student_test_score, school_budget = A, Y, U
+    hscm = HierarchicalStructuralCausalModel()
+    hscm.add_observed_node(student_tutoring)
+    hscm.add_observed_node(student_tutoring)
+    hscm.add_observed_node(student_test_score)
+    hscm.add_unobserved_node(school_budget)
+    hscm.add_edge(school_budget, student_tutoring)
+    hscm.add_edge(student_tutoring, student_tutoring)
+    hscm.add_edge(school_budget, student_tutoring)
+    hscm.add_edge(student_tutoring, student_test_score)
+    hscm.add_edge(student_test_score, student_tutoring)
+    hscm.add_subunits([student_tutoring, student_tutoring])
+    return hscm
+
+
+def get_instrument_hcm() -> HierarchicalCausalModel:
+    """Pytest fixture for the Instrument HCM in Figure 2 (i)."""
+    hcm = HierarchicalCausalModel()
+    hcm.add_observed_node(A)
+    hcm.add_observed_node(Y)
+    hcm.add_observed_node(Z)
+    hcm.add_edge(U, A)
+    hcm.add_edge(A, Y)
+    hcm.add_edge(U, Y)
+    hcm.add_edge(Z, A)
+    hcm.add_subunits([A, Z])
+    return hcm
+
+
+def get_instrument_subunit_graph() -> HierarchicalCausalModel:
+    """Pytest fixture for the Instrument HCM subunit graph."""
+    subg = HierarchicalCausalModel()
+    subg.add_observed_node(A)
+    subg.add_observed_node(Z)
+    subg.add_edge(Z, A)
+    return subg
+
+
+def get_instrument_hcgm() -> HierarchicalCausalModel:
+    """Pytest fixture for the Instrument HCGM in Figure 2(j)."""
+    hcm = HierarchicalCausalModel()
+    hcm.add_observed_node(A)
+    hcm.add_observed_node(Y)
+    hcm.add_observed_node(Z)
+    hcm.add_observed_node(Q_Z)
+    hcm.add_observed_node(Q_A_Z)
+    hcm.add_edge(U, Q_A_Z)
+    hcm.add_edge(Q_A_Z, A)
+    hcm.add_edge(Q_Z, Z)
+    hcm.add_edge(A, Y)
+    hcm.add_edge(U, Y)
+    hcm.add_edge(Z, A)
+    hcm.add_subunits([A, Z])
+    return hcm
+
+
+def get_instrument_hscm() -> HierarchicalStructuralCausalModel:
+    """Pytest fixture for the Instrument HSCM in Section 2.3 Eqn. (10)."""
+    hscm = HierarchicalStructuralCausalModel()
+    hscm.add_observed_node(A)
+    hscm.add_observed_node(Y)
+    hscm.add_observed_node(Z)
+    hscm.add_unobserved_node(U)
+    hscm.add_edge(U, A)
+    hscm.add_edge(A, Y)
+    hscm.add_edge(U, Y)
+    hscm.add_edge(Z, A)
+    hscm.add_subunits([A, Z])
+    return hscm
+
+
+def get_compl_subgraph_hcm() -> HierarchicalCausalModel:
+    """Pytest fixture for HCM with complicated subgraph structure."""
+    hcm = HierarchicalCausalModel.from_lists(
+        observed_subunits=[A, B, C, Y],
+        observed_units=["D"],
+        unobserved_units=[U],
+        edges=[
+            (U, A),
+            (U, B),
+            (U, C),
+            (A, B),
+            (B, Y),
+            (C, Y),
+            (A, D),
+            (D, C),
+        ],
+    )
+    return hcm
+
+
+confounder_augmented_admg = NxMixedGraph.from_edges(
+    undirected=[(Q_A, Q_Y_A)], directed=[(Q_A, Q_Y), (Q_Y_A, Q_Y)]
+)
+"""Pytest fixture for augmented Confounder HCM in Figure 2 (d)."""
+
+
+confounder_interference_augmented_admg = NxMixedGraph.from_edges(
+    undirected=[(Q_A, Q_Y_A)], directed=[(Q_A, Q_Y), (Q_Y_A, Q_Y), (Q_A, Z), (Z, Q_Y_A)]
+)
+"""Pytest fixture for augmented Confounder Interference HCM in Figure 2 (h)."""
+
+
+instrument_augmented_admg = NxMixedGraph.from_edges(
+    undirected=[(Q_A_Z, Y)], directed=[(Q_A_Z, Q_A), (Q_Z, Q_A), (Q_A, Y)]
+)
+"""Pytest fixture for augmented Instrument HCM in Figure A2."""
+
+
+instrument_marginalized_admg = NxMixedGraph.from_edges(
+    undirected=[(Q_A_Z, Y)], directed=[(Q_A_Z, Q_A), (Q_A, Y)]
+)
+"""Pytest fixture for augmented Instrument HCM in Figure 2 (l)."""
+
+
+confounder_collapsed_admg = NxMixedGraph.from_edges(undirected=[(Q_A, Q_Y_A)])
+"""Pytest fixture for collapsed Confounder HCM in Figure 2 (c)."""
+
+
+confounder_interference_collapsed_admg = NxMixedGraph.from_edges(
+    undirected=[(Q_A, Q_Y_A)], directed=[(Q_A, Z), (Z, Q_Y_A)]
+)
+"""Pytest fixture for collapsed Confounder Interference HCM in Figure 2 (g)."""
+
+
+instrument_collapsed_admg = NxMixedGraph.from_edges(
+    undirected=[(Q_A_Z, Y)], directed=[(Q_A_Z, Y), (Q_Z, Y)]
+)
+"""Pytest fixture for collapsed Instrument HCM in Figure 2 (k)."""