Identify hierarchical causal model specified in pyro

[djinnome]

create figure 2 models in pyro. use render_model to export to dot so that it can be read into pygraphviz.

import pyro

def figure_2a_confounder():
    """
    Simulates the confounder model for Figure 2(a) in the paper.
    """
    # Define the confounder model
    def confounder_model():
        # Define the prior distribution for the confounder variable
        confounder = pyro.sample("U", dist.Uniform(0.0, 1.0))
        # Define the prior distribution for the treatment variable
        with pyro.plate("students", 10):
            treatment = pyro.sample("A", dist.Bernoulli(confounder))
        # Define the prior distribution for the outcome variable
            outcome = pyro.sample("Y", dist.Normal(confounder + 2.0 * treatment, 1.0))
        return outcome
    return confounder_model

pyro.render_model(figure_2a_confounder(), render_deterministic=True, filename="confounder.dot")

confounder.dot.txt