vanAmsterdam · June 2, 2020 19:48
diff --git a/ordinal-regression-with-improperuniform-in-numpyro.py b/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()
	# 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()