PGM for Rubin et al 2015 (http://arxiv.org/abs/1507.01602)

rubin_pgm.py

#!/usr/bin/env python

from matplotlib import rc
from daft import PGM, Node, Plate
rc("font", family="serif", size=12)
rc("text", usetex=True)


pgm = PGM([7.65, 5.3], origin=[0., -0.2], observed_style='inner')

# -----------------------------------------------------------------------
# Inner Plate

# True SN parameters
pgm.add_node(Node("x1_true", r"$x_{1i}^\mathrm{true}$", 3, 3))
pgm.add_node(Node("mb_true", r"$m_{Bi}^\mathrm{true}$", 4, 3))
pgm.add_node(Node("c_true", r"$c_i^\mathrm{true}$", 5, 3))

# Obs SN parameters
pgm.add_node(Node("x1_fit", r"$x_{1i}^\mathrm{fit}$", 3, 1.9, observed=True))
pgm.add_node(Node("mb_fit", r"$m_{Bi}^\mathrm{fit}$", 4, 1.9, observed=True))
pgm.add_node(Node("c_fit", r"$c_i^\mathrm{fit}$", 5, 1.9, observed=True))

# true -> fit edges
pgm.add_edge("x1_true", "x1_fit")
pgm.add_edge("mb_true", "mb_fit")
pgm.add_edge("c_true", "c_fit")

pgm.add_plate(Plate([2.25, 0.75, 3.25, 2.75],
                    label=r"SNe $i = 1, \dots, N_{SN, j}$",
                    shift=-0.1))

# -----------------------------------------------------------------------------
# Outer Plate

# x1 and c distributions
pgm.add_node(Node("x1_dist", r"$R_{x_1 j}, \vec{x}_{1j}^\ast$",
                  1.25, 3.0, aspect=2.0))
pgm.add_node(Node("c_dist", r"$R_{cj}, \vec{c}_j^\ast, \vec{\alpha}_{cj}$",
                  6.5, 3.0, aspect=2.5))
pgm.add_edge("x1_dist", "x1_true")
pgm.add_edge("c_dist", "c_true")


pgm.add_node(Node("outliers",
                  "$f_{\\mathrm{outl}}, \\sigma_j^\\mathrm{int}, "
                  "\\rho^{\mathrm{int}},$\n$f_{m_B}, f_{x_1}, f_c$",
                  1.25, 0.75, scale=1.6, aspect=2.0))
pgm.add_edge("outliers", "x1_fit")
pgm.add_edge("outliers", "mb_fit")
pgm.add_edge("outliers", "c_fit")

# selection effects
pgm.add_node(Node("selection",
                  r"$m^\mathrm{cut}_j, \sigma^\mathrm{cut}_j, "
                  "a^\mathrm{cut}_j$",
                  6.5, 0.75, scale=1.1, aspect=2.7))
pgm.add_edge("selection", "mb_fit")
pgm.add_edge("selection", "c_fit")

pgm.add_plate(Plate([0.2, 0., 7.25, 3.7], label=r"Samples $j = 1, \dots, N_{samples}$"))

# -----------------------------------------------------------------------------
# global parameters

# systematics
pgm.add_node(Node("sys", r"$\vec{\Delta}_\mathrm{sys}$", 2.5, 4.5))
pgm.add_edge("sys", "x1_true")
pgm.add_edge("sys", "mb_true")
pgm.add_edge("sys", "c_true")

# all the others
pgm.add_node(Node("global",
                  "$\\Omega_m, M_B, \\delta(0), \\delta(\\infty),$\n"
                  "$\\beta_B,\\beta_R, \\alpha_S, \\alpha_L$",
                  5.5, 4.5, aspect=2.7, scale=1.7))
pgm.add_edge("global", "mb_true")


# -----------------------------------------------------------------------------
# Render and save.
pgm.render()
pgm.figure.savefig("rubinpgm.png", dpi=150)