Skip to content

Instantly share code, notes, and snippets.

@kbarbary
Created September 23, 2015 00:36
Show Gist options
  • Save kbarbary/83765ad6532cba17cd2c to your computer and use it in GitHub Desktop.
Save kbarbary/83765ad6532cba17cd2c to your computer and use it in GitHub Desktop.
PGM for Rubin et al 2015 (http://arxiv.org/abs/1507.01602)
#!/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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment