Calculating the waic with Arviz

1

I used a probit model to compute something and now I want to calculate the WAIC and the log likelihood for this case. I was thinking of using arviz for this however when comparing it seems poutine scale is not taken into account when calculating the log likelihood (weighted) and thus not in the WAIC. My model is the following:

def model_pyro_probit_weighted(x_c, y=None, w=None):
    c = x_c.shape[1]

    alpha = pyro.sample('alpha', dist.Normal(0.0, 1.0))
    beta_c = pyro.sample('lambda', dist.Normal(torch.zeros(c), torch.ones(c)).to_event(1)


    y_loc = alpha + x_c.matmul(beta_c.unsqueeze(-1)).squeeze(-1) 
    probs = dist.Normal(1.0 * torch.zeros(n), 1.0 * torch.ones(n)).cdf(y_loc)

    with pyro.plate('data'), pyro.poutine.scale(scale=w):
        pyro.sample('y', dist.Bernoulli(probs=probs), obs=y)

and I run the model as follows:

nuts_kernel = NUTS(model_pyro_probit_weighted)
    mcmc2 = MCMC(kernel=nuts_kernel, num_samples=num_samples_NUTS, warmup_steps=num_warm_NUTS, num_chains=5)
    mcmc2.run(x_c, x_m, y, w)

Then I calculate everything as shown below:

 posterior_predictive = \
            Predictive(model_spec, posterior_samples)(x_c, x_m, y=None, w=w)
        prior = \
            Predictive(model_spec,
                       num_samples=warm_up)(x_c, y=None, w=w)

    pyro_data = az.from_pyro(
        mcmc,
        prior=prior,
        posterior_predictive=posterior_predictive,
        coords={"resp": np.arange(x_c.shape[0])},
        dims={"data": ["resp"]}
    )

az.waic(pyro_data)

so when I do a quick calculation myself without taking the weight into account. Do note the weights are not equal to 1 as that would indeed render the same results.

  output = (alpha + torch.matmul(x_c, beta_c.transpose(1, 0)))

    prob_zero = norm.cdf(x=0, loc=output, scale=1)
    y_dup = y.repeat(prob_zero.shape[1], 1).T
    likelihood = np.where(y_dup == 1, 1 - prob_zero, prob_zero)
    log_likely = np.log(likelihood)
    lppd = np.sum(np.log(np.mean(np.exp(log_likely), axis=1)))
    p_waic = np.sum(np.var(log_likely, axis=1))

I get the same p_waic which should not be the case. So I am wondering where my derivation goes wrong or whether arviz cannot handle poutine scale properly. Any ideas?

summary: the WAIC of arviz is equal to the self calculated unweighted WAIC, even when the weights are not equal to 1

2

I think this is the reason. The log likelihood there is calculated from the raw log_prob which does not take account scale. I guess you can replace obs_site["fn"].log_prob(obs_site["value"]) by

obs_site["fn"].log_prob(obs_site["value"]) * obs_site.get("scale", 1.)

If it works, could you create a PR or issue in arviz? I am happy to review or address it.