{"id":82,"date":"2025-12-08T03:32:44","date_gmt":"2025-12-07T19:32:44","guid":{"rendered":"https:\/\/connectword.dpdns.org\/?p=82"},"modified":"2025-12-08T03:32:44","modified_gmt":"2025-12-07T19:32:44","slug":"a-coding-implementation-of-a-complete-hierarchical-bayesian-regression-workflow-in-numpyro-using-jax-powered-inference-and-posterior-predictive-analysis","status":"publish","type":"post","link":"https:\/\/connectword.dpdns.org\/?p=82&lang=en","title":{"rendered":"A Coding Implementation of a Complete Hierarchical Bayesian Regression Workflow in NumPyro Using JAX-Powered Inference and Posterior Predictive Analysis"},"content":{"rendered":"

In this tutorial, we explore hierarchical Bayesian regression with NumPyro<\/strong><\/a> and walk through the entire workflow in a structured manner. We start by generating synthetic data, then we define a probabilistic model that captures both global patterns and group-level variations. Through each snippet, we set up inference using NUTS, analyze posterior distributions, and perform posterior predictive checks to understand how well our model captures the underlying structure. By approaching the tutorial step by step, we build an intuitive understanding of how NumPyro enables flexible, scalable Bayesian modeling. Check out the\u00a0Full Codes here<\/a><\/strong>.<\/p>\n

\n
\n
\n
\n
<\/div>\n
<\/div>\n
<\/div>\n<\/div>\n
Copy Code<\/span>Copied<\/span>Use a different Browser<\/span><\/a><\/div>\n<\/div>\n
try:\n   import numpyro\nexcept ImportError:\n   !pip install -q \"llvmlite>=0.45.1\" \"numpyro[cpu]\" matplotlib pandas\n\n\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport jax\nimport jax.numpy as jnp\nfrom jax import random\nimport numpyro\nimport numpyro.distributions as dist\nfrom numpyro.infer import MCMC, NUTS, Predictive\nfrom numpyro.diagnostics import hpdi\n\n\nnumpyro.set_host_device_count(1)<\/code><\/pre>\n<\/div>\n<\/div>\n
We set up our environment by installing NumPyro and importing all required libraries. We prepare JAX, NumPyro, and plotting tools so we have everything ready for Bayesian inference. As we run this cell, we ensure our Colab session is fully equipped for hierarchical modeling. Check out the\u00a0Full Codes here<\/a><\/strong>.<\/p>\n
\n
\n
\n
\n
<\/div>\n
<\/div>\n
<\/div>\n<\/div>\n
Copy Code<\/span>Copied<\/span>Use a different Browser<\/span><\/a><\/div>\n<\/div>\n
def generate_data(key, n_groups=8, n_per_group=40):\n   k1, k2, k3, k4 = random.split(key, 4)\n   true_alpha = 1.0\n   true_beta = 0.6\n   sigma_alpha_g = 0.8\n   sigma_beta_g = 0.5\n   sigma_eps = 0.7\n   group_ids = np.repeat(np.arange(n_groups), n_per_group)\n   n = n_groups * n_per_group\n   alpha_g = random.normal(k1, (n_groups,)) * sigma_alpha_g\n   beta_g = random.normal(k2, (n_groups,)) * sigma_beta_g\n   x = random.normal(k3, (n,)) * 2.0\n   eps = random.normal(k4, (n,)) * sigma_eps\n   a = true_alpha + alpha_g[group_ids]\n   b = true_beta + beta_g[group_ids]\n   y = a + b * x + eps\n   df = pd.DataFrame({\"y\": np.array(y), \"x\": np.array(x), \"group\": group_ids})\n   truth = dict(true_alpha=true_alpha, true_beta=true_beta,\n                sigma_alpha_group=sigma_alpha_g, sigma_beta_group=sigma_beta_g,\n                sigma_eps=sigma_eps)\n   return df, truth\n\n\nkey = random.PRNGKey(0)\ndf, truth = generate_data(key)\nx = jnp.array(df[\"x\"].values)\ny = jnp.array(df[\"y\"].values)\ngroups = jnp.array(df[\"group\"].values)\nn_groups = int(df[\"group\"].nunique())<\/code><\/pre>\n<\/div>\n<\/div>\n
We generate synthetic hierarchical data that mimics real-world group-level variation. We convert this data into JAX-friendly arrays so NumPyro can process it efficiently. By doing this, we lay the foundation for fitting a model that learns both global trends and group differences. Check out the\u00a0Full Codes here<\/a><\/strong>.<\/p>\n
\n
\n
\n
\n
<\/div>\n
<\/div>\n
<\/div>\n<\/div>\n
Copy Code<\/span>Copied<\/span>Use a different Browser<\/span><\/a><\/div>\n<\/div>\n
def hierarchical_regression_model(x, group_idx, n_groups, y=None):\n   mu_alpha = numpyro.sample(\"mu_alpha\", dist.Normal(0.0, 5.0))\n   mu_beta = numpyro.sample(\"mu_beta\", dist.Normal(0.0, 5.0))\n   sigma_alpha = numpyro.sample(\"sigma_alpha\", dist.HalfCauchy(2.0))\n   sigma_beta = numpyro.sample(\"sigma_beta\", dist.HalfCauchy(2.0))\n   with numpyro.plate(\"group\", n_groups):\n       alpha_g = numpyro.sample(\"alpha_g\", dist.Normal(mu_alpha, sigma_alpha))\n       beta_g = numpyro.sample(\"beta_g\", dist.Normal(mu_beta, sigma_beta))\n   sigma_obs = numpyro.sample(\"sigma_obs\", dist.Exponential(1.0))\n   alpha = alpha_g[group_idx]\n   beta = beta_g[group_idx]\n   mean = alpha + beta * x\n   with numpyro.plate(\"data\", x.shape[0]):\n       numpyro.sample(\"y\", dist.Normal(mean, sigma_obs), obs=y)\n\n\nnuts = NUTS(hierarchical_regression_model, target_accept_prob=0.9)\nmcmc = MCMC(nuts, num_warmup=1000, num_samples=1000, num_chains=1, progress_bar=True)\nmcmc.run(random.PRNGKey(1), x=x, group_idx=groups, n_groups=n_groups, y=y)\nsamples = mcmc.get_samples()<\/code><\/pre>\n<\/div>\n<\/div>\n
We define our hierarchical regression model and launch the NUTS-based MCMC sampler. We allow NumPyro to explore the posterior space and learn parameters such as group intercepts and slopes. As this sampling completes, we obtain rich posterior distributions that reflect uncertainty at every level. Check out the\u00a0Full Codes here<\/a><\/strong>.<\/p>\n
\n
\n
\n
\n
<\/div>\n
<\/div>\n
<\/div>\n<\/div>\n
Copy Code<\/span>Copied<\/span>Use a different Browser<\/span><\/a><\/div>\n<\/div>\n
def param_summary(arr):\n   arr = np.asarray(arr)\n   mean = arr.mean()\n   lo, hi = hpdi(arr, prob=0.9)\n   return mean, float(lo), float(hi)\n\n\nfor name in [\"mu_alpha\", \"mu_beta\", \"sigma_alpha\", \"sigma_beta\", \"sigma_obs\"]:\n   m, lo, hi = param_summary(samples[name])\n   print(f\"{name}: mean={m:.3f}, HPDI=[{lo:.3f}, {hi:.3f}]\")\n\n\npredictive = Predictive(hierarchical_regression_model, samples, return_sites=[\"y\"])\nppc = predictive(random.PRNGKey(2), x=x, group_idx=groups, n_groups=n_groups)\ny_rep = np.asarray(ppc[\"y\"])\n\n\ngroup_to_plot = 0\nmask = df[\"group\"].values == group_to_plot\nx_g = df.loc[mask, \"x\"].values\ny_g = df.loc[mask, \"y\"].values\ny_rep_g = y_rep[:, mask]\n\n\norder = np.argsort(x_g)\nx_sorted = x_g[order]\ny_rep_sorted = y_rep_g[:, order]\ny_med = np.median(y_rep_sorted, axis=0)\ny_lo, y_hi = np.percentile(y_rep_sorted, [5, 95], axis=0)\n\n\nplt.figure(figsize=(8, 5))\nplt.scatter(x_g, y_g)\nplt.plot(x_sorted, y_med)\nplt.fill_between(x_sorted, y_lo, y_hi, alpha=0.3)\nplt.show()<\/code><\/pre>\n<\/div>\n<\/div>\n
We analyze our posterior samples by computing summaries and performing posterior predictive checks. We visualize how well the model recreates observed data for a selected group. This step helps us understand how accurately our model captures the underlying generative process. Check out the\u00a0Full Codes here<\/a><\/strong>.<\/p>\n
\n
\n
\n
\n
<\/div>\n
<\/div>\n
<\/div>\n<\/div>\n
Copy Code<\/span>Copied<\/span>Use a different Browser<\/span><\/a><\/div>\n<\/div>\n
alpha_g = np.asarray(samples[\"alpha_g\"]).mean(axis=0)\nbeta_g = np.asarray(samples[\"beta_g\"]).mean(axis=0)\n\n\nfig, axes = plt.subplots(1, 2, figsize=(12, 4))\naxes[0].bar(range(n_groups), alpha_g)\naxes[0].axhline(truth[\"true_alpha\"], linestyle=\"--\")\naxes[1].bar(range(n_groups), beta_g)\naxes[1].axhline(truth[\"true_beta\"], linestyle=\"--\")\nplt.tight_layout()\nplt.show()<\/code><\/pre>\n<\/div>\n<\/div>\n
We plot the estimated group-level intercepts and slopes to compare their learned patterns with the true values. We explore how each group behaves and how the model adapts to their differences. This final visualization brings together the complete picture of hierarchical inference.<\/p>\n
In conclusion, we implemented how NumPyro allows us to model hierarchical relationships with clarity, efficiency, and strong expressive power. We observed how the posterior results reveal meaningful global and group-specific effects, and how predictive checks validate the model\u2019s fit to the generated data. As we put everything together, we gain confidence in constructing, fitting, and interpreting hierarchical models using JAX-powered inference. This process strengthens our ability to apply Bayesian thinking to richer, more realistic datasets where multilevel structure is essential.<\/p>\n
\n
Check out the\u00a0Full Codes here<\/a><\/strong>.\u00a0Feel free to check out our\u00a0GitHub Page for Tutorials, Codes and Notebooks<\/a><\/mark><\/strong>.\u00a0Also,\u00a0feel free to follow us on\u00a0Twitter<\/mark><\/a><\/strong>\u00a0and don\u2019t forget to join our\u00a0100k+ ML SubReddit<\/a><\/strong>\u00a0and Subscribe to\u00a0our Newsletter<\/a><\/strong>. Wait! are you on telegram?\u00a0now you can join us on telegram as well.<\/a><\/strong><\/p>\n
The post A Coding Implementation of a Complete Hierarchical Bayesian Regression Workflow in NumPyro Using JAX-Powered Inference and Posterior Predictive Analysis<\/a> appeared first on MarkTechPost<\/a>.<\/p>","protected":false},"excerpt":{"rendered":"
In this tutorial, we explore h…<\/p>\n","protected":false},"author":1,"featured_media":29,"comment_status":"open","ping_status":"open","sticky":false,"template":"","format":"standard","meta":{"footnotes":""},"categories":[1],"tags":[],"class_list":["post-82","post","type-post","status-publish","format-standard","has-post-thumbnail","hentry","category-uncategorized"],"_links":{"self":[{"href":"https:\/\/connectword.dpdns.org\/index.php?rest_route=\/wp\/v2\/posts\/82","targetHints":{"allow":["GET"]}}],"collection":[{"href":"https:\/\/connectword.dpdns.org\/index.php?rest_route=\/wp\/v2\/posts"}],"about":[{"href":"https:\/\/connectword.dpdns.org\/index.php?rest_route=\/wp\/v2\/types\/post"}],"author":[{"embeddable":true,"href":"https:\/\/connectword.dpdns.org\/index.php?rest_route=\/wp\/v2\/users\/1"}],"replies":[{"embeddable":true,"href":"https:\/\/connectword.dpdns.org\/index.php?rest_route=%2Fwp%2Fv2%2Fcomments&post=82"}],"version-history":[{"count":0,"href":"https:\/\/connectword.dpdns.org\/index.php?rest_route=\/wp\/v2\/posts\/82\/revisions"}],"wp:featuredmedia":[{"embeddable":true,"href":"https:\/\/connectword.dpdns.org\/index.php?rest_route=\/wp\/v2\/media\/29"}],"wp:attachment":[{"href":"https:\/\/connectword.dpdns.org\/index.php?rest_route=%2Fwp%2Fv2%2Fmedia&parent=82"}],"wp:term":[{"taxonomy":"category","embeddable":true,"href":"https:\/\/connectword.dpdns.org\/index.php?rest_route=%2Fwp%2Fv2%2Fcategories&post=82"},{"taxonomy":"post_tag","embeddable":true,"href":"https:\/\/connectword.dpdns.org\/index.php?rest_route=%2Fwp%2Fv2%2Ftags&post=82"}],"curies":[{"name":"wp","href":"https:\/\/api.w.org\/{rel}","templated":true}]}}