Ordinal regression with ImproperUniform distribution in NumPyro

ordinal-regression-with-improperuniform-in-numpyro.py

# Ordinal regression with ImproperUniform

from jax import numpy as np, random
import numpyro
from numpyro import sample
from numpyro.distributions import constraints, Normal, ImproperUniform, Categorical, OrderedLogistic
from numpyro.infer.mcmc import NUTS, MCMC
import pandas as pd

num_chains = 4
numpyro.set_host_device_count(num_chains)
numpyro.set_platform('cpu')


# Generate data with ordinal structure
simkeys  = random.split(random.PRNGKey(1), 2)
nsim     = 50
nclasses = 3
Y        = Categorical(logits=np.zeros((nclasses,))).sample(simkeys[0], sample_shape=(nsim,))
X        = Normal().sample(simkeys[1], sample_shape = (nsim,))
X       += Y

print("value counts of Y:")
print(pd.Series(Y).value_counts())

for i in range(nclasses):
    print(f"mean(X) for Y == {i}: {X[np.where(Y==i)].mean():.3f}")


# Create models using the `OrderedLogistic` distribution.
# This requires cutpoints that are ordered, so we need to introduce this as a constraint.
# We can use the `ImproperUnifrom` distribution to introduce a parameter
# with an arbitrary support that is otherwise completely uninformative

def model1(X, Y, nclasses=3):
    b_X_eta = sample('b_X_eta', Normal(0, 5))
    c_y     = sample('c_y',     ImproperUniform(support=constraints.ordered_vector, event_shape=(nclasses-1,)))
    with numpyro.plate('obs', X.shape[0]):
        eta = X * b_X_eta
        sample('Y', OrderedLogistic(eta, c_y), obs=Y)

mcmc_key = random.PRNGKey(1234)
kernel = NUTS(model1)
mcmc   = MCMC(kernel, num_warmup=250, num_samples=750, num_chains=num_chains)
mcmc.run(mcmc_key, X,Y, nclasses)
mcmc.print_summary()


# If we have additional information on the parameters that we want to add,
# we can add a `sample` statement that uses a 'regular' prior with an `obs` argument,
# conditioning the parameter drawn from the `ImproperUniform` on the chosen prior

def model2(X, Y, nclasses=3):
    b_X_eta = sample('b_X_eta', Normal(0, 5))
    c_y     = sample('c_y', ImproperUniform(support=constraints.ordered_vector, event_shape=(nclasses-1,)))
    sample('c_y_smp', Normal(0,1), obs=c_y)
    with numpyro.plate('obs', X.shape[0]):
        eta = X * b_X_eta
        sample('Y', OrderedLogistic(eta, c_y), obs=Y)

kernel = NUTS(model2)
mcmc   = MCMC(kernel, num_warmup=250, num_samples=750, num_chains=num_chains)
mcmc.run(mcmc_key, X,Y, nclasses)
mcmc.print_summary()