

In [96]:

    
pgm = daft.PGM(shape=(6,8), node_unit=1, )

# TODO: boldify vectors

###################################################################
# Add nodes
#

# progenitor
pgm.add_node(daft.Node("x_p", r"$x_p$", 3, 6))
pgm.add_node(daft.Node("w_p", r"$w_p$", 2, 6))
pgm.add_node(daft.Node("D_p", r"$D_p$", 1, 6, observed=True))
pgm.add_node(daft.Node("cov_p", r"$\Sigma_p$", 1, 7, observed=True))
pgm.add_node(daft.Node("m_0", r"$m_0$", 1, 5, observed=True))
pgm.add_node(daft.Node("mdot", r"$\dot{m}$", 1, 4))

# potential
pgm.add_node(daft.Node("Phi", r"$\Phi$", 3, 7))

# other
pgm.add_node(daft.Node("t_int", r"$t_{\rm int}$", 4, 7))
pgm.add_node(daft.Node("lambda", r"$\lambda$", 5, 6))

# star plate
pgm.add_plate(daft.Plate([2,1,3,4], label="$n=1...N$"))

# stars
pgm.add_node(daft.Node("x_n", r"$x_n$", 3, 4))
pgm.add_node(daft.Node("t_n", r"$\tau_n$", 4, 4))
pgm.add_node(daft.Node("w_n", r"$w_n$", 3, 3))
pgm.add_node(daft.Node("D_n", r"$D_n$", 3, 2, observed=True))
pgm.add_node(daft.Node("cov_n", r"$\Sigma$", 4, 2, observed=True))

###################################################################
# Add edges
#

# progenitor
pgm.add_edge('Phi', 'x_p')
pgm.add_edge('t_int', 'x_p')
pgm.add_edge('x_p', 'w_p')
pgm.add_edge('w_p', 'D_p')
pgm.add_edge('cov_p', 'D_p')

# stars
pgm.add_edge('lambda', 'x_n')
pgm.add_edge('x_p', 'x_n')
pgm.add_edge('m_0', 'x_n')
pgm.add_edge('mdot', 'x_n')
pgm.add_edge('m_0', 'mdot')
pgm.add_edge('t_n', 'x_n')
pgm.add_edge('x_n', 'w_n')
pgm.add_edge('w_n', 'D_n')
pgm.add_edge('cov_n', 'D_n')

ax = pgm.render()
fig = ax.figure
fig.savefig("pgm.pdf")