def sample_plot_gibbs(x0, z0, kv, π, μ, σ, n_iterations, xmin, xmax): x_hist, z_hist = gibbs_sample(x0, z0, kv, π, μ, σ, n_iterations) colors = ["tab:blue" if z else "tab:red" for z in z_hist] fig, axs = plt.subplots() axs.scatter(np.arange(n_iterations), x_hist, s=20, facecolors="none", edgecolors=colors) pml.savefig("gibbs_scatter.pdf") fig = plt.figure() axs = plt.axes(projection="3d") plot_gmm_3d_trace(x_hist, π, μ, σ, "Gibbs sampling", xmin, xmax, axs) pml.style3d(axs, 1.5, 1, 0.8) plt.subplots_adjust(left=0.001, bottom=0.208, right=0.7) pml.savefig("gibbs_trace.pdf", pad_inches=0, bbox_inches="tight") fig, axs = plt.subplots() sm.graphics.tsa.plot_acf(x_hist, lags=45, alpha=None, title="Gibbs", ax=axs) pml.savefig("gibbs_autocorrelation.pdf")
def sample_plot_mh(x0, τ, π, μ, σ, n_iterations, xmin, xmax): x_hist = metropolis_sample(x0, τ, π, μ, σ, n_iterations) fig = plt.figure() axs = plt.axes(projection="3d") plot_gmm_3d_trace(x_hist, π, μ, σ, f"MH with $N(0,{τ}^2)$ proposal", xmin, xmax, axs) pml.style3d(axs, 1.5, 1, 0.8) plt.subplots_adjust(left=0.001, bottom=0.208) pml.savefig(f"mh_trace_{τ}tau.pdf", pad_inches=0, bbox_inches="tight") fig, axs = plt.subplots() sm.graphics.tsa.plot_acf(x_hist, lags=45, alpha=None, title=f"MH with $N(0,{τ}^2)$ proposal", ax=axs) pml.savefig(f"mh_autocorrelation_{τ}tau.pdf")
def plot_3d_belief_state(mu_hist, dim, ax, skip=3, npoints=2000, azimuth=-30, elevation=30): nsteps = len(mu_hist) xmin, xmax = mu_hist[..., dim].min(), mu_hist[..., dim].max() xrange = jnp.linspace(xmin, xmax, npoints).reshape(-1, 1) res = np.apply_along_axis(lambda X: pml.kdeg(xrange, X[..., None], 0.5), 1, mu_hist) densities = res[..., dim] for t in range(0, nsteps, skip): tloc = t * np.ones(npoints) px = densities[t] ax.plot(tloc, xrange, px, c="tab:blue", linewidth=1) ax.set_zlim(0, 1) pml.style3d(ax, 1.8, 1.2, 0.7, 0.8) ax.view_init(elevation, azimuth) ax.set_xlabel(r"$t$", fontsize=13) ax.set_ylabel(r"$x_{" f"d={dim}" ",t}$", fontsize=13) ax.set_zlabel(r"$p(x_{d, t} \vert y_{1:t})$", fontsize=13)