# Example: Horseshoe Regression¶

We demonstrate how to use NUTS to do sparse regression using the Horseshoe prior  for both continuous- and binary-valued responses. For a more complex modeling and inference approach that also supports quadratic interaction terms in a way that is efficient in high dimensions see examples/sparse_regression.py.

References:

 “Handling Sparsity via the Horseshoe,”

Carlos M. Carvalho, Nicholas G. Polson, James G. Scott.

```import argparse
import os
import time

import numpy as np
from scipy.special import expit

import jax.numpy as jnp
import jax.random as random

import numpyro
from numpyro.diagnostics import summary
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS

# regression model with continuous-valued outputs/responses
def model_normal_likelihood(X, Y):
D_X = X.shape

# sample from horseshoe prior
lambdas = numpyro.sample("lambdas", dist.HalfCauchy(jnp.ones(D_X)))
tau = numpyro.sample("tau", dist.HalfCauchy(jnp.ones(1)))

# note that in practice for a normal likelihood we would probably want to
# integrate out the coefficients (as is done for example in sparse_regression.py).
# however, this trick wouldn't be applicable to other likelihoods
# (e.g. bernoulli, see below) so we don't make use of it here.
unscaled_betas = numpyro.sample("unscaled_betas", dist.Normal(0.0, jnp.ones(D_X)))
scaled_betas = numpyro.deterministic("betas", tau * lambdas * unscaled_betas)

# compute mean function using linear coefficients
mean_function = jnp.dot(X, scaled_betas)

prec_obs = numpyro.sample("prec_obs", dist.Gamma(3.0, 1.0))
sigma_obs = 1.0 / jnp.sqrt(prec_obs)

# observe data
numpyro.sample("Y", dist.Normal(mean_function, sigma_obs), obs=Y)

# regression model with binary-valued outputs/responses
def model_bernoulli_likelihood(X, Y):
D_X = X.shape

# sample from horseshoe prior
lambdas = numpyro.sample("lambdas", dist.HalfCauchy(jnp.ones(D_X)))
tau = numpyro.sample("tau", dist.HalfCauchy(jnp.ones(1)))

# note that this reparameterization (i.e. coordinate transformation) improves
# posterior geometry and makes NUTS sampling more efficient
unscaled_betas = numpyro.sample("unscaled_betas", dist.Normal(0.0, jnp.ones(D_X)))
scaled_betas = numpyro.deterministic("betas", tau * lambdas * unscaled_betas)

# compute mean function using linear coefficients
mean_function = jnp.dot(X, scaled_betas)

# observe data
numpyro.sample("Y", dist.Bernoulli(logits=mean_function), obs=Y)

# helper function for HMC inference
def run_inference(model, args, rng_key, X, Y):
start = time.time()
kernel = NUTS(model)
mcmc = MCMC(
kernel,
num_warmup=args.num_warmup,
num_samples=args.num_samples,
num_chains=args.num_chains,
progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True,
)

mcmc.run(rng_key, X, Y)
mcmc.print_summary(exclude_deterministic=False)

samples = mcmc.get_samples()
summary_dict = summary(samples, group_by_chain=False)

print("\nMCMC elapsed time:", time.time() - start)

return summary_dict

# create artificial regression dataset with 3 non-zero regression coefficients
def get_data(N=50, D_X=3, sigma_obs=0.05, response="continuous"):
assert response in ["continuous", "binary"]
assert D_X >= 3

np.random.seed(0)
X = np.random.randn(N, D_X)

# the response only depends on X_0, X_1, and X_2
W = np.array([2.0, -1.0, 0.50])
Y = jnp.dot(X[:, :3], W)
Y -= jnp.mean(Y)

if response == "continuous":
Y += sigma_obs * np.random.randn(N)
elif response == "binary":
Y = np.random.binomial(1, expit(Y))

assert X.shape == (N, D_X)
assert Y.shape == (N,)

return X, Y

def main(args):
N, D_X = args.num_data, 32

print("[Experiment with continuous-valued responses]")
# first generate and analyze data with continuous-valued responses
X, Y = get_data(N=N, D_X=D_X, response="continuous")

# do inference
rng_key, rng_key_predict = random.split(random.PRNGKey(0))
summary = run_inference(model_normal_likelihood, args, rng_key, X, Y)

# lambda should only be large for the first 3 dimensions, which
# correspond to relevant covariates (see get_data)
print("Posterior median over lambdas (leading 5 dimensions):")
print(summary["lambdas"]["median"][:5])
print("Posterior mean over betas (leading 5 dimensions):")
print(summary["betas"]["mean"][:5])

print("[Experiment with binary-valued responses]")
# next generate and analyze data with binary-valued responses
# (note we use more data for the case of binary-valued responses,
# since each response carries less information than a real number)
X, Y = get_data(N=4 * N, D_X=D_X, response="binary")

# do inference
rng_key, rng_key_predict = random.split(random.PRNGKey(0))
summary = run_inference(model_bernoulli_likelihood, args, rng_key, X, Y)

# lambda should only be large for the first 3 dimensions, which
# correspond to relevant covariates (see get_data)
print("Posterior median over lambdas (leading 5 dimensions):")
print(summary["lambdas"]["median"][:5])
print("Posterior mean over betas (leading 5 dimensions):")
print(summary["betas"]["mean"][:5])

if __name__ == "__main__":
assert numpyro.__version__.startswith("0.8.0")
parser = argparse.ArgumentParser(description="Horseshoe regression example")