Arviz: Weird `plot_ppc()` for numpyro model with custom mixture distribution

[derekpowell]

Short Description

I'm trying to estimate a finite mixture beta regression model using numpyro. The actual model I'm trying to fit is more complex, but I've reproduced my issue with a simpler model. y_obs is produced by two mixture components: one related to another variable x and another just centered at .50. This is roughly meant to represent how lots of human probability judgments look, where people are sometimes drawn to respond 50% when they are unsure.

Reproducible example and outputs

import numpyro
import numpyro.distributions as dist

from jax import numpy as jnp
from jax import random
from numpyro.infer import MCMC, NUTS, Predictive

import numpy as np
from numpy.random import default_rng

from jax.scipy.special import logit, expit

## simulate some data
n_obs = 2000
rg = default_rng(12345)

x = rg.normal(0,1,n_obs)
y_mean = expit(-.2 + .8*x)

y = rg.beta(y_mean*50, (1-y_mean)*50, n_obs)
p_drawn = .25
drawn = rg.binomial(1, p_drawn, n_obs)
dist50 = rg.beta(.5*300, .5*300, n_obs) # k = 300
y_obs = np.where(drawn==1, dist50, y)

The simulated data looks like this:

## define model

class MixtureBeta(dist.distribution.Distribution):
    '''Mixture of Beta distribtions with marginalized latents

    Takes a and b parameters for beta distributions along with a simplex mixing_probs parameter that defines the mixture probabilities. 
    '''
    arg_constraints = {'concentration1': constraints.positive, 'concentration0': constraints.positive, 'mixing_probs': constraints.simplex}
    reparametrized_params = ['concentration1', 'concentration0', 'mixing_probs']
    support = constraints.unit_interval

    def __init__(self, concentration1, concentration0, mixing_probs, validate_args=None):
        self.concentration1, self.concentration0 = promote_shapes(concentration1, concentration0)
        self.mixing_probs = jnp.reshape(mixing_probs, (1,-1))
        batch_shape = lax.broadcast_shapes(jnp.shape(concentration1), jnp.shape(concentration0))
        concentration1 = jnp.broadcast_to(concentration1, batch_shape)
        concentration0 = jnp.broadcast_to(concentration0, batch_shape)
        self._dirichlet = dist.Dirichlet(jnp.stack([concentration1, concentration0],
                                      axis=-1))
        super(MixtureBeta, self).__init__(batch_shape=batch_shape, validate_args=validate_args)

    def sample(self, key, sample_shape=()):
        assert is_prng_key(key)
        samples = self._dirichlet.sample(key, sample_shape)[..., 0]
        mixing_logits = jax.scipy.special.logit(self.mixing_probs)
        ind = random.categorical(key, mixing_logits)
        return samples[ind]

    @validate_sample
    def log_prob(self, value):
        log_mixing_probs = jnp.log(self.mixing_probs.T)
        dirichlet_probs = self._dirichlet.log_prob(jnp.stack([value, 1. - value], -1))
        sum_probs = jnp.add(log_mixing_probs, dirichlet_probs)
        return jax.nn.logsumexp(sum_probs, axis=0)

def mixmodel_beta_custom(x, y_obs=None):

    alpha = numpyro.sample("alpha", dist.Normal(0,3))
    beta = numpyro.sample("beta", dist.Normal(0,3))
    mixing_probs = numpyro.sample("mixing", dist.Dirichlet(jnp.ones(2)))

    k = numpyro.sample("k", dist.HalfCauchy(10))
    mix_kval = numpyro.sample("mix_kval", dist.HalfCauchy(100))
    mix_muval = numpyro.sample("mix_muval", dist.Beta(1,1))

    y_true = expit(alpha + beta*x)
    n_obs = y_true.shape[0]

    with numpyro.plate("data", x.shape[0]):
        yhat = jnp.stack([y_true, jnp.ones(n_obs)*mix_muval])
        mix_k = jnp.stack([jnp.ones(n_obs)*k, jnp.ones(n_obs)*mix_kval])
        numpyro.sample("yhat", MixtureBeta(yhat*mix_k, (1-yhat)*mix_k, mixing_probs), obs=y_obs)

## fit the model

kernel = NUTS(mixmodel_beta_custom)
mcmc_test = MCMC(kernel, 2000, 2000, num_chains=1)
mcmc_test.run(random.PRNGKey(0), x, y_obs)

## get posterior predictive

posterior_samples = mcmc_test.get_samples()

posterior_predictive = Predictive(mixmodel_beta_custom, posterior_samples)(
    random.PRNGKey(1), x
)
prior = Predictive(mixmodel_beta_custom, num_samples=500)(
    random.PRNGKey(10), x
)

az_data = az.from_numpyro(mcmc_test, prior=prior, posterior_predictive=posterior_predictive)

az.plot_ppc(data=az_data, var_names = "yhat", data_pairs={"yhat":"yhat"}, num_pp_samples=500)

This model recovers the correct parameters in the simulated data

                 mean       std    median      5.0%     95.0%     n_eff     r_hat
      alpha     -0.20      0.01     -0.20     -0.22     -0.19   2477.07      1.00
       beta      0.79      0.01      0.79      0.78      0.81   2464.89      1.00
          k     49.54      2.00     49.48     46.17     52.77   2153.06      1.00
   mix_kval    317.33     29.35    317.21    268.25    365.02   1972.80      1.00
  mix_muval      0.50      0.00      0.50      0.50      0.50   2539.01      1.00
  mixing[0]      0.75      0.01      0.75      0.73      0.77   2612.54      1.00
  mixing[1]      0.25      0.01      0.25      0.23      0.27   2612.54      1.00

but produces the funky plot below:

I'm wondering what i'm doing wrong? Unfortunately, I'm not sure if this is a numpyro or arviz issue, but thought I'd ask here first since the estimation seems to be working correctly but the plotting doesn't.

Appreciate any help!

[ahartikainen]

What is the shape for the yhat? Is it a vector?

I think our ppc does some kind of flattening which causes your draws to be plot under the same var. If you check ppc (solid blue line) you can see that there are two distinct groups (in densities). This is probably due to flattening your mixture model.

cc @OriolAbril might have an idea what to do

[OriolAbril]

I ran the code to see what was going on, and it looks like the draws either sample from one of the mixture components or from the other, is that what should happen?

The mean therefore looks a bit closer to the observed kde, but the other kde which are done will the pp samples for a single draw are completely different.

import matplotlib.pyplot as plt
fig, ax = plt.subplots()

where = np.empty(az_data.posterior_predictive.dims["draw"])
for d in range(az_data.posterior_predictive.dims["draw"]):
    draw_data = az_data.posterior_predictive["yhat"].sel(draw=d).values.flatten()
    az.plot_kde(
        draw_data, 
        #label=f"{d}", 
        ax=ax, 
        plot_kwargs=dict(color="k" if draw_data.min() < .25 else "r")
    )
    where[d] = draw_data.min() < .25
print(where.mean())
# 0.899

[derekpowell]

Thanks @ahartikainen and @OriolAbril for the replies: Yes the yhat should be a vector drawn from the mixture distribution not just its individual components. Your diagnosis of what's happening seems exactly right but I can't figure out why it's happening. For comparison I also wrote a model using a mixture of normals and it has no problem with the pp checks.

class MixtureNormal(dist.distribution.Distribution):
    '''Mixture of Normal distribtions with marginalized latents

    Takes loc and scale parameters for normal distributions along with a simplex mixing_probs parameter that defines the mixture probabilities. 
    '''
    arg_constraints = {
        'loc': constraints.real, 
        'scale': constraints.positive, 
        'mixing_probs': constraints.simplex
    }
    support = constraints.real
    reparametrized_params = ['loc', 'scale', 'mixing_probs']

    def __init__(self, loc=0., scale=1., mixing_probs=1., validate_args=None):
        self.loc, self.scale = promote_shapes(loc, scale)
        self.mixing_probs = jnp.reshape(mixing_probs, (1,-1))
        batch_shape = lax.broadcast_shapes(jnp.shape(loc), jnp.shape(scale))
        super(MixtureNormal, self).__init__(batch_shape=batch_shape, event_shape=1, validate_args=validate_args)

    def sample(self, key, sample_shape=()):
        assert is_prng_key(key)
        mixing_logits = jax.scipy.special.logit(self.mixing_probs)
        ind = random.categorical(key, mixing_logits)
        eps = random.normal(key, shape=sample_shape + self.batch_shape + self.event_shape)
        return self.loc[ind] + eps * self.scale[ind]

    @validate_sample
    def log_prob(self, value):
        log_mixing_probs = jnp.log(self.mixing_probs.T)
        probs = dist.Normal(self.loc, self.scale).log_prob(value)
        sum_probs = jnp.add(log_mixing_probs, probs)

        return jax.nn.logsumexp(sum_probs, axis=0)

def mixmodel_custom(x, y=None):

    alpha = numpyro.sample("alpha", dist.Normal(0,10))
    beta = numpyro.sample("beta", dist.Normal(0,10))
    sigma = numpyro.sample("sigma", dist.HalfCauchy(10))

    mean_outlier = numpyro.sample("mean_o", dist.Normal(0,10))
    sigma_outlier = numpyro.sample("sigma_o", dist.HalfCauchy(10))
    p_outlier = numpyro.sample("p_outlier", dist.Dirichlet(jnp.ones(2)))

    y_fore = alpha + beta*x

    ## have to manually manage shapes of things
    yhat = jnp.stack([y_fore, mean_outlier.repeat(x.shape[0])])
    mix_sigma = jnp.stack([sigma.repeat(x.shape[0]), sigma_outlier.repeat(x.shape[0])])

    with numpyro.plate("data", x.shape[0]):
        numpyro.sample("yhat", MixtureNormal(yhat, mix_sigma, p_outlier), obs=y)

And doing inference just like before yields this nice-looking pp plot:

[OriolAbril]

@fehiepsi can probably help more on that. I don't really know why it's not working.

[fehiepsi]

@derekpowell Here are some lines that might cause the issue:

• The mixing probability should have shape batch_shape + (num_mixtures,)

self.mixing_probs = jnp.reshape(mixing_probs, (1,-1))

• This line does not respect sample_shape. You need to make sure that mixing_logits is broadcasted to have shape sample_shape + self.batch_shape + (num_mixtures,)

ind = random.categorical(key, mixing_logits)

• Those lines do not respect broadcasting

log_mixing_probs = jnp.log(self.mixing_probs.T)  # .T is not broadcastable

jax.nn.logsumexp(sum_probs, axis=0)  # axis=0 is not broadcastable



md5-68857c3f6e4fc2efbfffbcf0d4f90118



yhat = jnp.stack([y_true, jnp.ones(n_obs)*mix_muval])  # jnp.stack with the default axis=0 is not broadcastable
mix_k = jnp.stack([jnp.ones(n_obs)*k, jnp.ones(n_obs)*mix_kval])  # same as above

[fehiepsi]

@derekpowell I just address the above issues and it seems to return the expected result. Could you double-check?

class MixtureBeta(dist.Distribution):
    def __init__(self, concentration1, concentration0, mixing_probs, validate_args=None):
        expand_shape = jax.lax.broadcast_shapes(
            jnp.shape(concentration1), jnp.shape(concentration0), jnp.shape(mixing_probs))
        self._beta = dist.Beta(concentration1, concentration0).expand(expand_shape)
        self._categorical = dist.Categorical(jnp.broadcast_to(mixing_probs, expand_shape))
        super(MixtureBeta, self).__init__(batch_shape=expand_shape[:-1], validate_args=validate_args)

    def sample(self, key, sample_shape=()):
        key, key_idx = random.split(key)
        samples = self._beta.sample(key, sample_shape)
        ind = self._categorical.sample(key_idx, sample_shape)
        return jnp.take_along_axis(samples, ind[..., None], -1)[..., 0]

    def log_prob(self, value):
        dirichlet_probs = self._beta.log_prob(value[..., None])
        sum_probs = self._categorical.logits + dirichlet_probs
        return jax.nn.logsumexp(sum_probs, axis=-1)

def mixmodel_beta_custom(x, y_obs=None):
    ...
    yhat = jnp.stack([y_true, jnp.ones(n_obs)*mix_muval], -1)
    mix_k = jnp.stack([jnp.ones(n_obs)*k, jnp.ones(n_obs)*mix_kval], -1)

I think the takeaways here is to make sure that the implementation follows tensor shape senmatics: tensor/numpy arrays broadcast by aligning on the right. By the way, the above code can be generalized (using similar code) to more general patterns like

class Mixture(dist.Distribution):
    def __init__(self, base_dist, mixing_probs, validate_args=None):
        self.base_dist = ...
        self._categorical = ...

where we can replace base_dist with Beta, Normal,... If you find that this Mixture class is useful, it would be great to open a FR or make a PR in numpyro. :)

[derekpowell]

@fehiepsi Thank you! Works perfectly in both the example and my actual application. I figured I was botching up the tensor shapes but couldn't write the tests to figure out where I was going wrong. I'll see if I can wrap my head around your code well enough to attempt a generalization for a PR myself, but otherwise will try to draft up a clear FR.

Thanks to all for your help, closing this as it wasn't a problem with arviz

[fehiepsi]

attempt a generalization for a PR myself, but otherwise will try to draft up a clear FR.

Thanks! FYI, PyTorch has MixtureSameFamily that you can use as an API reference (the implementation there is a bit complicated to follow).