Add additional level in hierarchical model

Hi there, first post here :relaxed:
I’ve previously been using PyMC3, but wanted to give numpyro a go.

I’ve looked at this example for hierarchical modeling:
http://num.pyro.ai/en/stable/tutorials/bayesian_hierarchical_linear_regression.html

But what I want to do is to add another layer, could be seen as something like country in the above example. But how would that look in regards to the plates and the formula: alpha[patient_id, country_id] + beta[patient_id, country_id]*Weeks

Hi @mattiasthalen, welcome! I think if we have country level, the code would be something like

    with numpyro.plate("patients", n_patients, dim=-2):
        with numpyro.plate("countries", n_countries, dim=-1):
            α = numpyro.sample("α", dist.Normal(μ_α, σ_α))
            β = numpyro.sample("β", dist.Normal(μ_β, σ_β))

    FVC_est = α[PatientID, country_id] + β[PatientID, country_id] * Weeks
1 Like

Interesting!
Is there no need for intermediate parameters, is only the global & local parameters needed?

In PyMC3 I would have something like:
global_alpha_mu -> country_alpha_mu -> alpha

The way we construct a model corresponds 1-1 to a graphical model. So if there is an “intermediate” variable in your graphical model, then you will need to define it in (num)pyro. If you want to record any deterministic value, you can use numpyro.deterministic. Could you post your simplified pymc3 model here so I can get better understanding of what you are looking for?

Thanks for the help!

I made it work in the end :relaxed:

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)
2 Likes