Thanks for the help!
I made it work in the end 
I was after this:
alpha = numpyro.sample('alpha', alpha_prior)
beta0 = numpyro.sample('beta0', beta0_prior)
beta1 = numpyro.sample('beta1', beta1_prior)
sigma = numpyro.sample('sigma', sigma_prior)
alpha_sigma = numpyro.sample('alpha_sigma', numpyro.distributions.HalfNormal(1))
beta0_sigma = numpyro.sample('beta0_sigma', numpyro.distributions.HalfNormal(1))
beta1_sigma = numpyro.sample('beta1_sigma', numpyro.distributions.HalfNormal(1))
with numpyro.plate('level1', n_level1_ids, dim = -1):
level1_alpha_offset = numpyro.sample('level1_alpha_offset', numpyro.distributions.Normal(0, 1))
level1_alpha = numpyro.deterministic('level1_alpha', alpha + level1_alpha_offset*alpha_sigma)
level1_beta0_offset = numpyro.sample('level1_beta0_offset', numpyro.distributions.Normal(0, 1))
level1_beta0 = numpyro.deterministic('level1_beta0', beta0 + level1_beta0_offset*beta0_sigma)
level1_beta1_offset = numpyro.sample('level1_beta1_offset', numpyro.distributions.Normal(0, 1))
level1_beta1 = numpyro.deterministic('level1_beta1', beta1 + level1_beta1_offset*beta1_sigma)
level1_alpha_sigma = numpyro.sample('level1_alpha_sigma', numpyro.distributions.HalfNormal(1))
level1_beta0_sigma = numpyro.sample('level1_beta0_sigma', numpyro.distributions.HalfNormal(1))
level1_beta1_sigma = numpyro.sample('level1_beta1_sigma', numpyro.distributions.HalfNormal(1))
with numpyro.plate('level2', n_level2_ids, dim = -2):
level2_alpha_offset = numpyro.sample('level2_alpha_offset', numpyro.distributions.Normal(0, 1))
level2_alpha = numpyro.deterministic('level2_alpha', level1_alpha + level2_alpha_offset*level1_alpha_sigma)
level2_beta0_offset = numpyro.sample('level2_beta0_offset', numpyro.distributions.Normal(0, 1))
level2_beta0 = numpyro.deterministic('level2_beta0', level1_beta0 + level2_beta0_offset*level1_beta0_sigma)
level2_beta1_offset = numpyro.sample('level2_beta1_offset', numpyro.distributions.Normal(0, 1))
level2_beta1 = numpyro.deterministic('level2_beta1', level1_beta1 + level2_beta0_offset*level1_beta1_sigma)
Full model function:
def quadratic_model(x,
y = None,
level1_id = None,
level2_id = None,
level2_beta = False,
inverse = False,
standardize = True,
priors = None,
**kwargs):
'''
Bayesian model to estimate the parameters used in the quadratic regression `y = alpha + beta0*x + beta1*x^2`.
Up to two additional levels are supported to make a hierarchical model.
Parameters
----------
x : list of float
The independent variable.
y : list of float, optional
Dependent variable.
level1_id : list of int, optional
Optional list of integers for additional level in the model.
level2_id : list of int, optional
Optional list of integers for additional level in the model.
level2_name : str, default `level2`
Set the display name of level 2.
level2_beta : bool, default False
Set to `True` if `beta0` & `beta1` should be inferred for level 2.
inverse : bool, default False
Flip the independent & dependent variables.
standardize : bool, default True
Standardize variables using `sklearn.preprocessing.StandardScaler()`.
Unscaled parameters are returned in addition to the scaled ones.
priors : dict, default None
Provide a dict of priors for the parameters.
E.g. `{'alpha': numpyro.distributions.Normal(0, 1)}`.
'''
x_obs, y_obs = [x, y] if not inverse else [y, x]
# Standardize inputs
if standardize:
xy_scaler = sklearn.preprocessing.StandardScaler()
xy = np.array(list(zip(x_obs, y_obs)))
xy_scaled = xy_scaler.fit_transform(xy)
x_obs, y_obs = np.array(list(zip(*xy_scaled)))
x_mean, y_mean = xy_scaler.mean_
x_sd, y_sd = xy_scaler.scale_
inverse_alpha = lambda alpha, beta0, beta1: y_sd*(alpha - beta0*x_mean/x_sd + beta1*x_mean**2/x_sd**2) + y_mean
inverse_beta0 = lambda beta0, beta1: beta0*y_sd/x_sd - 2*beta1*x_mean*y_sd/x_sd**2
inverse_beta1 = lambda beta1: beta1*y_sd/x_sd**2
level1_check = type(level1_id) is np.ndarray
level2_check = type(level2_id) is np.ndarray
# Hyperparameters
alpha_prior = numpyro.distributions.Normal(0, 1)
beta0_prior = numpyro.distributions.Normal(0, 1)
beta1_prior = numpyro.distributions.Normal(0, 1)
sigma_prior = numpyro.distributions.HalfNormal(1)
if type(priors) is dict:
alpha_prior = priors.get('alpha', alpha_prior)
beta0_prior = priors.get('beta0', beta0_prior)
beta1_prior = priors.get('beta1', beta1_prior)
sigma_prior = priors.get('sigma', sigma_prior)
alpha = numpyro.sample('alpha', alpha_prior)
beta0 = numpyro.sample('beta0', beta0_prior)
beta1 = numpyro.sample('beta1', beta1_prior)
sigma = numpyro.sample('sigma', sigma_prior)
if standardize:
numpyro.deterministic('alpha_unscaled', inverse_alpha(alpha, beta0, beta1))
numpyro.deterministic('beta0_unscaled', inverse_beta0(beta0, beta1))
numpyro.deterministic('beta1_unscaled', inverse_beta1(beta1))
# Level 1 Parameters
if level1_check:
n_level1_ids = len(np.unique(level1_id))
alpha_sigma = numpyro.sample('alpha_sigma', numpyro.distributions.HalfNormal(1))
beta0_sigma = numpyro.sample('beta0_sigma', numpyro.distributions.HalfNormal(1))
beta1_sigma = numpyro.sample('beta1_sigma', numpyro.distributions.HalfNormal(1))
with numpyro.plate('level1', n_level1_ids, dim = -1):
level1_alpha_offset = numpyro.sample('level1_alpha_offset', numpyro.distributions.Normal(0, 1))
level1_alpha = numpyro.deterministic('level1_alpha', alpha + level1_alpha_offset*alpha_sigma)
level1_beta0_offset = numpyro.sample('level1_beta0_offset', numpyro.distributions.Normal(0, 1))
level1_beta0 = numpyro.deterministic('level1_beta0', beta0 + level1_beta0_offset*beta0_sigma)
level1_beta1_offset = numpyro.sample('level1_beta1_offset', numpyro.distributions.Normal(0, 1))
level1_beta1 = numpyro.deterministic('level1_beta1', beta1 + level1_beta1_offset*beta1_sigma)
if standardize:
numpyro.deterministic('level1_alpha_unscaled', inverse_alpha(level1_alpha, level1_beta0, level1_beta1))
numpyro.deterministic('level1_beta0_unscaled', inverse_beta0(level1_beta0, level1_beta1))
numpyro.deterministic('level1_beta1_unscaled', inverse_beta1(level1_beta1))
# Level 2 Parameters
if level2_check:
n_level2_ids = len(np.unique(level2_id))
level1_alpha_sigma = numpyro.sample('level1_alpha_sigma', numpyro.distributions.HalfNormal(1))
if level2_beta:
level1_beta0_sigma = numpyro.sample('level1_beta0_sigma', numpyro.distributions.HalfNormal(1))
level1_beta1_sigma = numpyro.sample('level1_beta1_sigma', numpyro.distributions.HalfNormal(1))
with numpyro.plate('level2', n_level2_ids, dim = -2):
level2_alpha_offset = numpyro.sample('level2_alpha_offset', numpyro.distributions.Normal(0, 1))
level2_alpha = numpyro.deterministic('level2_alpha', level1_alpha + level2_alpha_offset*level1_alpha_sigma)
if level2_beta:
level2_beta0_offset = numpyro.sample('level2_beta0_offset', numpyro.distributions.Normal(0, 1))
level2_beta0 = numpyro.deterministic('level2_beta0', level1_beta0 + level2_beta0_offset*level1_beta0_sigma)
level2_beta1_offset = numpyro.sample('level2_beta1_offset', numpyro.distributions.Normal(0, 1))
level2_beta1 = numpyro.deterministic('level2_beta1', level1_beta1 + level2_beta0_offset*level1_beta1_sigma)
if standardize:
numpyro.deterministic('level2_alpha_unscaled', inverse_alpha(level2_alpha, level1_beta0, level1_beta1))
if (standardize) & (level2_beta):
numpyro.deterministic('level2_beta0_unscaled', inverse_beta0(level2_beta0, level2_beta1))
numpyro.deterministic('level2_beta1_unscaled', inverse_beta1(level2_beta1))
# Update parameter assignments
if level1_check:
alpha = level1_alpha[level1_id]
beta0 = level1_beta0[level1_id]
beta1 = level1_beta1[level1_id]
if (level1_check) & (level2_check):
alpha = level2_alpha[level2_id, level1_id]
if (level1_check) & (level2_check) & (level2_beta):
beta0 = level2_beta0[level2_id, level1_id]
beta1 = level2_beta1[level2_id, level1_id]
# Likelihood
mu = alpha + beta0*x_obs + beta1*x_obs**2
with numpyro.plate('observations', len(x_obs)):
numpyro.sample('y_obs', numpyro.distributions.Normal(mu, sigma), obs = y_obs)