import pickle
import jax

import matplotlib.pyplot as plt
import jax.numpy as jnp
import tensorflow_probability.substrates.jax as tfp
from scipy.stats import gaussian_kde
import plotly.express as px
import pandas as pd
import pickle
tfd = tfp.distributions
import plotly
plotly.offline.init_notebook_mode()

x = jnp.linspace(8,12,5000)
all_pdf = []

Ajax

with open('./results_data/weibull_Ajax','rb') as f:
    posterior  = pickle.load(f)
ajax_samples = posterior.sample(seed = jax.random.PRNGKey(1), sample_shape = (10000,))
kde_ajax = gaussian_kde(ajax_samples["theta"])
pdf_ajax = kde_ajax(x)
all_pdf.append(pdf_ajax)

MCMC Blackjax

import numpy as np
with open('./results_data/MCMC_Blackjax','rb') as f:
    mcmc_model  = pickle.load(f)
mcmc_samples= np.array(mcmc_model.position[300:])
kde_black = gaussian_kde(mcmc_samples.reshape((-1,)))
pdf_black = kde_black(x)
all_pdf.append(pdf_black)

HMC Stan

with open('results_data/HMC_Stan','rb') as f:
    stan_hmc  = pickle.load(f)
stan_hmc_kde=gaussian_kde(stan_hmc)
stan_hmc_pdf=stan_hmc_kde(x)
all_pdf.append(stan_hmc_pdf)

all_label = ["Ajax VI estimate"]*x.shape[0] + ["Blackjax RMH estimate"]*x.shape[0]+ ["Stan HMC"]*x.shape[0]
x_repeated = jnp.tile(x,3)
all_pdf = jnp.array(all_pdf).reshape((-1))
print( x_repeated.shape, all_pdf.shape)
to_df = {
    "theta":x_repeated,
    "PDF":all_pdf,
    "label": all_label

}
df = pd.DataFrame(to_df)
fig = px.line(to_df,"theta","PDF",color="label",title="Weibull Poisson results") 
fig.show()
fig.write_html("weibull_poisson_result_plotly.html")

(15000,) (15000,)

!jupyter nbconvert --to HTML weibull_poisson_results.ipynb

[NbConvertApp] Converting notebook weibull_poisson_results.ipynb to HTML
[NbConvertApp] Writing 8367001 bytes to weibull_poisson_results.html