imf-pgm

imf-pgm.py

#!/usr/bin/env python

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

import daft

pgm = daft.PGM([3.5, 3], origin=[0, -0.5])

# Stars.
pgm.add_node(daft.Node("theta", r"$\theta$", 1, 2))
pgm.add_node(daft.Node("mass", r"$M_i$", 1, 1))
pgm.add_node(daft.Node("mass-obs", r"$\vec{d_i}$", 2, 1,
                       observed=True))
pgm.add_node(daft.Node("obs", r"$\mathrm{obs}_i$", 1.5, 0.5, observed=True))
pgm.add_node(daft.Node("f", r"$f$", 2, 2, fixed=True))
pgm.add_node(daft.Node("selection", r"$s$", 0.25, 0.5, fixed=True))

pgm.add_edge("theta", "mass")
pgm.add_edge("mass", "mass-obs")
pgm.add_edge("f", "mass-obs")
pgm.add_edge("mass", "obs")
pgm.add_edge("selection", "obs")

pgm.add_plate(daft.Plate([0.5, 0, 2, 1.5], shift=-0.1,
                         label=r"$\mathrm{stars:}\,i=1,\cdots,N$"))

# Numbers.
pgm.add_node(daft.Node("npred", r"$N_\mathrm{pred}$", 3, 2))
pgm.add_node(daft.Node("nobs", r"$N_\mathrm{obs}$", 3, 1, observed=True))
pgm.add_edge("npred", "nobs")

pgm.render()
pgm.figure.savefig("imf-pgm.pdf")