Bayesian Hierarchical Linear Model — Mixture Heterogeneity

Model

This extends the hierarchical linear model by replacing its single-normal coefficient distribution with a mixture of normals: each unit $i$ still has its own regression $\beta_i$, but the population of $\beta_i$ is now a $K$-component mixture, with optional unit covariates $z_i$ shifting the means through $\Delta$. The single-normal hierarchy is the special case $K=1$. It replicates bayesm's rhierLinearMixture (Rossi, Allenby & McCulloch 2005, Ch. 5).

• $\beta_i$ — unit $i$'s regression coefficients; the population is a mixture of $K$ normals
• $(\pi_k, \mu_k, \Sigma_k)$ — weight, mean, covariance of component $k$
• $\sigma_i^2$ — per-unit error variance; $\Delta$ — optional centred unit-covariate effects

As in the multinomial-logit mixture, the central question is whether the mixture functions as a flexible heterogeneity density (skew / fat tails a single normal misses) or as a segmentation device (clean discrete clusters). Both applications here land on the same answer — the flexible-density role earns its keep; discrete segments generally do not survive scrutiny.

6-Block Gibbs sampler

Every block is conjugate (the linear outcome makes $\beta_i$ Gaussian directly — no Metropolis needed, unlike the MNL mixture): a 6-block Gibbs scan over coefficients, component assignments, per-component Normal/Inverse-Wishart parameters, mixing weights, per-unit error variances, and the optional $\Delta$ covariate block. Components are sorted by intercept each sweep to resolve label switching. Model comparison uses a marginal-over-$\beta$ WAIC.

Notebooks

Cheese (IRI scanner panel, 88 retailers). A 2-component mixture is only mildly preferred by WAIC over a single normal ($\Delta \approx 4$); $K=3$ does not improve on $K=2$, and no store commits to a component (every store's posterior membership clusters near 0.5, none above $\approx 0.65$). The three engines even find different modes (base-volume split, display split, label-switched), with PyMC's NUTS failing to converge ($\hat r \approx 1.7$) on the multimodal target — the coefficient distribution is continuous and only weakly non-normal, not discrete segments. The robust outputs are the single-normal hierarchy and the $\Delta$ large-market effect. Rats growth data (Gelfand et al. 1990; 30 rats × 5 ages). The canonical BiRats backbone: $K=1$ recovers population intercept $\approx 242$ g (weight at day 22), slope $\approx 6.2$ g/day, and a positive intercept–slope correlation, with textbook partial-pooling shrinkage. WAIC selects $K=1$ — there are no discrete growth types. Four independent engines agree on the BiRats fit (from-scratch Gibbs, PyMC NUTS, bayesm, and lme4), and both Bayesian mixture routes confirm a single bivariate normal suffices.

Downloads

Sampler Module — Source Code

"""
Bayesian Hierarchical Linear Model Mixture -- Gibbs sampler.
Replicates bayesm::rhierLinearMixture (Rossi, Allenby & McCulloch 2005, Ch. 5).

Model
-----
  y_i  = X_i @ beta_i + eps_i,             eps_i ~ N(0, sigma2_i I_Ti)
  beta_i | s_i=k ~ MvN(mu_k + Delta @ z_i, Sigma_k)
  s_i  | pi      ~ Categorical(pi)
  pi             ~ Dirichlet(a * 1_K)
  mu_k           ~ MvN(mubar, (Amu * I_p)^{-1})   [independent of Sigma_k]
  Sigma_k        ~ IW(nu, V)
  sigma2_i       ~ IG(nu_e/2, nu_e * ssq_i / 2)

  If Z=None: Delta = 0  (pure mixture, no store covariates)
  If Z given: vec(Delta) ~ N(deltabar, Ad^{-1})
    Z should be centred (no intercept column), matching bayesm convention.

Identification
--------------
  Components are sorted by mu_k[0] (intercept) ascending after each
  (mu_k, Sigma_k) draw to resolve label switching.

6-block Gibbs
-------------
  (1)  beta_i   | s_i, sigma2_i, mu_{s_i}, Sigma_{s_i}  -- Normal conjugate
  (2)  s_i      | beta_i, pi, {mu_k, Sigma_k}           -- Categorical
  (3a) mu_k     | {beta_i: s_i=k}, Sigma_k              -- Normal conjugate
  (3b) Sigma_k  | {beta_i: s_i=k}, mu_k                 -- IW conjugate
  (4)  pi       | {s_i}                                  -- Dirichlet
  (5)  sigma2_i | beta_i, y_i, X_i                       -- IG conjugate
  (6)  Delta    | {beta_i, s_i, mu_{s_i}, {Sigma_k}}    -- GLS Normal [if Z]

Delta convention
----------------
  Delta is stored as (p, nz).  delta_vec = Delta.ravel() [C/row-major]:
  delta_vec[j*nz + l] = Delta[j, l] = effect of z[:,l] on beta[:,j].
  Deltadraw columns follow this row-major ordering (differs from bayesm's
  column-major vec(); convert with .reshape(p,nz).ravel('F') if needed).
"""

import time
import numpy as np
from scipy.stats import invwishart
from scipy.linalg import cholesky, solve_triangular


# ── private helpers ───────────────────────────────────────────────────────────

def _chol_draw(Prec, rhs, rng):
    """x ~ N(Prec^{-1} rhs, Prec^{-1}) via Cholesky of the precision matrix."""
    L  = cholesky(Prec, lower=True)
    v  = solve_triangular(L,   rhs, lower=True)
    mu = solve_triangular(L.T, v,   lower=False)
    z  = solve_triangular(L.T, rng.standard_normal(len(rhs)), lower=False)
    return mu + z


def _log_mvn(x, mu, L):
    """Log density of MvN(mu, L L') at x.  L is lower Cholesky of Sigma."""
    p   = len(x)
    dev = x - mu
    v   = solve_triangular(L, dev, lower=True)
    return (-0.5 * p * np.log(2.0 * np.pi)
            - np.sum(np.log(np.diag(L)))
            - 0.5 * float(v @ v))


# ── public API ────────────────────────────────────────────────────────────────

def rhier_linear_mixture(regdata, K=2, Z=None, Prior=None, Mcmc=None, seed=42):
    """
    Gibbs sampler for the Hierarchical Linear Model Mixture.

    Parameters
    ----------
    regdata : list of nreg dicts, each with 'y': (T_i,) and 'X': (T_i, p)
    K       : int, number of mixture components (default 2)
    Z       : (nreg, nz) store-level covariates (centred, no intercept).
              If None, pure mixture with component means mu_k only.
    Prior   : dict with optional keys
                a        : float or (K,)  Dirichlet concentration, default 5.
                mubar    : (p,)           prior mean for mu_k,      default 0
                Amu      : float          precision for mu_k prior,  default 0.1
                nu       : int            IW df for Sigma_k,         default p+3
                V        : (p,p)          IW scale for Sigma_k,      default (p+2)*I
                nu_e     : float          IG df for sigma2_i,        default 3.
                ssq      : float|(nreg,)  IG scale for sigma2_i,     default 1.
                Ad       : (p,p)          precision for Delta,        default 0.01*I_p
                            [only if Z provided; for nz>1 use (p*nz, p*nz)]
                deltabar : (p,)|(p*nz,)  prior mean for delta_vec,  default 0
    Mcmc    : dict with optional keys
                R      : int  total draws,     default 10000
                burn   : int  burn-in,         default 0.2*R
                keep   : int  thinning,        default 1
                nprint : int  print frequency, default R//10

    Returns
    -------
    dict with
      betadraw    : (nreg, p, R_kept)   store-level coefficient draws
      sdraw       : (R_kept, nreg)      component assignment, 0-indexed, sorted
      mudraw      : (R_kept, K, p)      component mean draws
      Sigmadraw   : (R_kept, K, p, p)   component covariance draws
      pidraw      : (R_kept, K)         mixing weight draws
      sigmasqdraw : (R_kept, nreg)      error variance draws
      Deltadraw   : (R_kept, p*nz)      [only if Z provided, row-major vec]
    """
    if Prior is None:
        Prior = {}
    if Mcmc is None:
        Mcmc = {'R': 10000}

    rng  = np.random.default_rng(seed)
    nreg = len(regdata)
    p    = regdata[0]['X'].shape[1]

    # ── priors ───────────────────────────────────────────────────────────────
    a_raw  = Prior.get('a', 5.)
    a_dir  = (np.full(K, float(a_raw)) if np.isscalar(a_raw)
              else np.asarray(a_raw, float))
    mubar  = np.asarray(Prior.get('mubar', np.zeros(p)), float)
    Amu    = float(Prior.get('Amu', 0.1))
    A_mu   = Amu * np.eye(p)                                   # (p, p) precision
    nu0    = int(Prior.get('nu',  p + 3))
    V0     = np.asarray(Prior.get('V', float(p + 2) * np.eye(p)), float)
    nu_e   = float(Prior.get('nu_e', 3.0))
    ssq_raw = Prior.get('ssq', 1.0)
    ssq    = (np.full(nreg, ssq_raw, float) if np.isscalar(ssq_raw)
              else np.asarray(ssq_raw, float))

    # ── MCMC settings ────────────────────────────────────────────────────────
    R      = int(Mcmc['R'])
    burn   = int(Mcmc.get('burn',   int(0.2 * R)))
    keep   = int(Mcmc.get('keep',   1))
    nprint = int(Mcmc.get('nprint', max(R // 10, 1)))
    R_kept = (R - burn) // keep

    # ── data ─────────────────────────────────────────────────────────────────
    ys   = [np.asarray(d['y'], float) for d in regdata]
    Xs   = [np.asarray(d['X'], float) for d in regdata]
    Ts   = np.array([len(y) for y in ys])
    XtXs = [X.T @ X for X in Xs]
    Xtys = [X.T @ y for X, y in zip(Xs, ys)]

    # ── Z covariate (optional) ────────────────────────────────────────────────
    use_Z = Z is not None
    if use_Z:
        Z_arr = np.asarray(Z, float)           # (nreg, nz)
        nz    = Z_arr.shape[1]
        dim_d = p * nz                         # dimension of delta_vec
        Ad    = np.asarray(Prior.get('Ad', 0.01 * np.eye(dim_d)), float)
        d0    = np.asarray(Prior.get('deltabar', np.zeros(dim_d)), float)
        Delta = np.zeros((p, nz))              # current Delta matrix
    else:
        nz = 0

    # ── initialise ────────────────────────────────────────────────────────────
    sigma2  = ssq.copy()
    mu_k    = np.tile(mubar, (K, 1)).copy()                  # (K, p)
    Sigma_k = np.stack([V0 / max(nu0 - p - 1, 1)] * K)      # (K, p, p)
    Sig_inv = np.stack([np.linalg.inv(Sigma_k[k]) for k in range(K)])  # (K,p,p)
    pi_k    = np.full(K, 1.0 / K)
    s       = np.arange(nreg, dtype=int) % K                 # round-robin init
    betas   = np.zeros((nreg, p))

    # ── storage ───────────────────────────────────────────────────────────────
    betadraw    = np.zeros((nreg,  p,    R_kept))
    sdraw       = np.zeros((R_kept, nreg), dtype=int)
    mudraw      = np.zeros((R_kept, K, p))
    Sigmadraw   = np.zeros((R_kept, K, p, p))
    pidraw      = np.zeros((R_kept, K))
    sigmasqdraw = np.zeros((R_kept, nreg))
    if use_Z:
        Deltadraw = np.zeros((R_kept, dim_d))

    i_kept = 0
    t0     = time.time()

    for r in range(R):

        # ── block 1: beta_i | s_i, sigma2_i, mu_{s_i}, Sigma_{s_i} ──────────
        for i in range(nreg):
            k_i    = s[i]
            mu_eff = (mu_k[k_i] + Delta @ Z_arr[i]) if use_Z else mu_k[k_i]
            Prec_i = Sig_inv[k_i] + XtXs[i] / sigma2[i]
            rhs_i  = Sig_inv[k_i] @ mu_eff + Xtys[i] / sigma2[i]
            betas[i] = _chol_draw(Prec_i, rhs_i, rng)

        # ── block 2: s_i | beta_i, pi, {mu_k, Sigma_k} ───────────────────────
        L_ks   = [cholesky(Sigma_k[k], lower=True) for k in range(K)]
        log_pi = np.log(np.clip(pi_k, 1e-300, None))
        for i in range(nreg):
            log_w = np.empty(K)
            for k in range(K):
                mu_eff    = (mu_k[k] + Delta @ Z_arr[i]) if use_Z else mu_k[k]
                log_w[k]  = log_pi[k] + _log_mvn(betas[i], mu_eff, L_ks[k])
            log_w -= log_w.max()
            w      = np.exp(log_w)
            w     /= w.sum()
            s[i]   = rng.choice(K, p=w)

        # ── block 3: (mu_k, Sigma_k) | {beta_i: s_i=k}, Delta ───────────────
        for k in range(K):
            mask   = s == k
            n_k    = int(mask.sum())
            beff_k = betas[mask].copy()          # (n_k, p)
            if use_Z:
                beff_k -= (Delta @ Z_arr[mask].T).T   # subtract Delta z_i shift

            # 3a: mu_k | beff_k, Sigma_k  (Normal conjugate, independent prior)
            if n_k > 0:
                Prec_mu = A_mu + n_k * Sig_inv[k]
                rhs_mu  = A_mu @ mubar + Sig_inv[k] @ beff_k.sum(0)
            else:
                Prec_mu = A_mu.copy()
                rhs_mu  = A_mu @ mubar
            mu_k[k] = _chol_draw(Prec_mu, rhs_mu, rng)

            # 3b: Sigma_k | beff_k, mu_k  (IW conjugate)
            if n_k > 0:
                dev     = beff_k - mu_k[k]
                scatter = dev.T @ dev
                Sigma_k[k] = invwishart.rvs(df=nu0 + n_k,
                                             scale=V0 + scatter,
                                             random_state=rng)
            else:
                Sigma_k[k] = invwishart.rvs(df=nu0, scale=V0, random_state=rng)
            Sig_inv[k] = np.linalg.inv(Sigma_k[k])

        # label-switching: sort components by intercept ascending
        order   = np.argsort(mu_k[:, 0])
        mu_k    = mu_k[order]
        Sigma_k = Sigma_k[order]
        Sig_inv = Sig_inv[order]
        s_new   = np.empty_like(s)
        for new_k, old_k in enumerate(order):
            s_new[s == old_k] = new_k
        s = s_new

        # ── block 4: pi | {s_i} ───────────────────────────────────────────────
        counts = np.bincount(s, minlength=K).astype(float)
        pi_k   = rng.dirichlet(a_dir + counts)

        # ── block 5: sigma2_i | beta_i, y_i, X_i ─────────────────────────────
        for i in range(nreg):
            e        = ys[i] - Xs[i] @ betas[i]
            scale_i  = (nu_e * ssq[i] + float(e @ e)) / 2.0
            sigma2[i] = scale_i / rng.gamma((nu_e + Ts[i]) / 2.0)

        # ── block 6: Delta | {beta_i, s_i, mu_{s_i}, Sigma_k} [if Z] ─────────
        if use_Z:
            # Full GLS update for delta_vec = Delta.ravel() [row-major, shape p*nz]
            # Model: beta_i - mu_{s_i} = Delta @ z_i + u_i, u_i ~ N(0, Sigma_{s_i})
            # Precision: P = Ad + sum_i kron(Sig_inv[s_i], outer(z_i, z_i))
            # rhs:       r = Ad @ d0 + sum_i kron(Sig_inv[s_i] @ bi_eff, z_i)
            P_D   = Ad.copy()
            rhs_D = Ad @ d0
            for i in range(nreg):
                k_i    = s[i]
                zi     = Z_arr[i]                        # (nz,)
                bi_eff = betas[i] - mu_k[k_i]           # (p,)
                P_D   += np.kron(Sig_inv[k_i], np.outer(zi, zi))
                rhs_D += np.kron(Sig_inv[k_i] @ bi_eff, zi)
            delta_vec = _chol_draw(P_D, rhs_D, rng)     # (p*nz,) row-major
            Delta = delta_vec.reshape(p, nz)             # (p, nz)

        if nprint and (r + 1) % nprint == 0:
            elapsed = time.time() - t0
            print(f'  Iter {r+1:6d}/{R}  ({elapsed:.0f}s)')

        if r >= burn and (r - burn) % keep == 0:
            betadraw[:, :, i_kept]  = betas
            sdraw[i_kept]           = s
            mudraw[i_kept]          = mu_k
            Sigmadraw[i_kept]       = Sigma_k
            pidraw[i_kept]          = pi_k
            sigmasqdraw[i_kept]     = sigma2
            if use_Z:
                Deltadraw[i_kept]   = delta_vec
            i_kept += 1

    elapsed = time.time() - t0
    print(f'Done in {elapsed / 60:.1f} min  |  kept: {R_kept}')

    out = {
        'betadraw':    betadraw,
        'sdraw':       sdraw,
        'mudraw':      mudraw,
        'Sigmadraw':   Sigmadraw,
        'pidraw':      pidraw,
        'sigmasqdraw': sigmasqdraw,
    }
    if use_Z:
        out['Deltadraw'] = Deltadraw
    return out


def simulate_hlm_mixture(nreg, p, K, true_pi, true_mu, true_Sigma,
                          true_sigma2=1.0, ni=20, seed=0):
    """
    Simulate data from the HLM Mixture model with known parameters.

    Parameters
    ----------
    nreg        : int, number of units (stores)
    p           : int, number of regression coefficients
    K           : int, number of mixture components
    true_pi     : (K,) mixing weights (must sum to 1)
    true_mu     : (K, p) component means
    true_Sigma  : (K, p, p) component covariances
    true_sigma2 : float or (nreg,), error variance per unit
    ni          : int or list, observations per unit
    seed        : int, random seed

    Returns
    -------
    regdata : list of nreg dicts with 'y', 'X'
    betas   : (nreg, p) true store-level coefficients
    comps   : (nreg,) true component assignments (0-indexed)
    """
    rng    = np.random.default_rng(seed)
    L_sigs = [np.linalg.cholesky(true_Sigma[k]) for k in range(K)]
    sig2   = (np.full(nreg, float(true_sigma2)) if np.isscalar(true_sigma2)
              else np.asarray(true_sigma2, float))
    ns     = [ni] * nreg if np.isscalar(ni) else list(ni)

    comps  = rng.choice(K, size=nreg, p=true_pi)
    betas  = np.array([
        true_mu[comps[i]] + L_sigs[comps[i]] @ rng.standard_normal(p)
        for i in range(nreg)
    ])

    regdata = []
    for i in range(nreg):
        X = np.column_stack([np.ones(ns[i]),
                              rng.standard_normal((ns[i], p - 1))])
        y = X @ betas[i] + rng.normal(0., np.sqrt(sig2[i]), ns[i])
        regdata.append({'y': y, 'X': X})

    return regdata, betas, comps


def posterior_summary(draws, param_names=None, true_vals=None):
    """Posterior mean, SD, 95% CrI for each column of a (R, q) draws array."""
    import pandas as pd
    q = draws.shape[1]
    if param_names is None:
        param_names = [f'p{j}' for j in range(q)]
    rows = []
    for j, name in enumerate(param_names):
        d      = draws[:, j]
        lo, hi = np.percentile(d, [2.5, 97.5])
        row    = {'param': name, 'mean': d.mean(), 'sd': d.std(),
                  'q025': lo, 'q975': hi}
        if true_vals is not None:
            row['true']    = float(true_vals[j])
            row['covered'] = bool(lo <= true_vals[j] <= hi)
        rows.append(row)
    return pd.DataFrame(rows).set_index('param')

References

• Rossi, P. E., Allenby, G. M. & McCulloch, R. (2005). Bayesian Statistics and Marketing, Ch. 5. Wiley. — bayesm::rhierLinearMixture, and the IRI cheese panel
• Gelfand, A. E., Hills, S. E., Racine-Poon, A. & Smith, A. F. M. (1990). Illustration of Bayesian inference in normal data models using Gibbs sampling. Journal of the American Statistical Association 85(412), 972–985. — the Rats / BiRats growth data
• Stephens, M. (2000). Dealing with label switching in mixture models. Journal of the Royal Statistical Society: Series B 62(4), 795–809. — label switching — the identification problem the mixture fits must handle
• Watanabe, S. (2010). Asymptotic equivalence of Bayes cross validation and widely applicable information criterion in singular learning theory. Journal of Machine Learning Research 11, 3571–3594. — WAIC, used to select K
• Lewandowski, D., Kurowicka, D. & Joe, H. (2009). Generating random correlation matrices based on vines and extended onion method. Journal of Multivariate Analysis 100(9), 1989–2001. — the LKJ prior used in the PyMC fits
• Bates, D., Mächler, M., Bolker, B. & Walker, S. (2015). Fitting linear mixed-effects models using lme4. Journal of Statistical Software 67(1), 1–48. — the frequentist benchmark the hierarchical fits are checked against