Business#

• Executive Summary: This repository contains the accompanying source code for the blog post "Finally! Bayesian Hierarchical Modelling at Scale", demonstrating how to build and train a Bayesian hierarchical Gamma-Poisson model for large-scale daily sales forecasting across many retail stores. The project uses JAX and NumPyro to implement a scalable probabilistic model that shares statistical strength across stores while handling missing observations and padded time series. Data is based on the Kaggle Rossmann Store Sales dataset, with preprocessing, model fitting, and evaluation orchestrated through Jupyter notebooks and a small Python package under src/bhm_at_scale. The code focuses on experiment reproducibility via conda environments, and on comparing hierarchical Bayesian predictions against a conventional regression baseline. Artifacts such as encoded feature matrices, learned parameters, and prediction summaries are written to disk for downstream analysis and visualization.
• Intended Use: The repository is intended for data scientists and researchers who want to study and prototype scalable Bayesian hierarchical models for multi-store retail demand forecasting using Python, JAX, and NumPyro. It is designed to support experimentation, educational use, and reproducible demonstrations that accompany the referenced blog post, rather than serving as a turnkey production forecasting system.
• Non Goals
  • Providing a production-grade, fully tested, and supported forecasting service or generic time-series framework; the code is a research and teaching example that accompanies a blog post and uses a PyScaffold-based project template.
• Use Case: The main use case is forecasting daily sales for many retail stores simultaneously by fitting a Gamma-Poisson Bayesian hierarchical model that shares information across stores while modeling store-specific effects, using the Rossmann Store Sales dataset as a concrete example.
• User Populations
  • Data scientists, machine learning practitioners, and researchers familiar with Python and probabilistic programming who are interested in scalable Bayesian hierarchical modeling for retail sales forecasting.
  • Data Scientist
  • Domain Expert
  • Governance, Compliance & Ethics Officer
  • ML Engineer
  • Product Manager
  • Project Manager
  • Software Developer
  • UX Researcher

Dev Insight#

• Architecture
  • Data Flow
    • Notebook 01-preprocessing downloads the Kaggle 'rossmann-store-sales' dataset, merges store metadata, encodes categorical variables, writes a cleaned feature table to data/preprocessed/edf.csv, and converts it into a 3D feature cube X with separate X_train and X_test NumPy arrays saved under data/preprocessed/.
    • Notebook 02-model loads X_train, uses ModelHandler together with the NumPyro model and guide to fit the hierarchical Gamma-Poisson model via SVI, checkpoints optimizer state, and writes latent parameter summaries and train-set predictions as CSV files under data/result/.
    • Notebook 02-model then loads X_test, fits local store-level parameters for partially observed series using local_guide, and generates test-set predictions which are saved to data/result/test_preds.csv.
    • Notebook 03-evaluation reads the original cleaned data, statistics, and prediction CSVs, and uses plotting utilities from bhm_at_scale.plot and bhm_at_scale.utils to visualize store-level predictions, weekday effects, and promotion effects.
  • Deps Summary
    • Core scientific Python stack: numpy, scipy, pandas, scikit-learn, matplotlib, seaborn.
    • Probabilistic programming and accelerated computation: jax, jaxlib, numpyro (using GammaPoisson, TransformReparam, SVI, Adam optimizer).
    • Experiment tooling and visualization: JupyterLab, daft (for graphical models), seaborn, matplotlib.
    • Data access and environment management: kaggle API for dataset download, conda environment defined in environment.yaml and environment.lock.yaml.
  • Public Apis
    • bhm_at_scale.model.model(X): Gamma-Poisson hierarchical model definition for daily sales forecasting.
    • bhm_at_scale.model.predictive_model(model_params): factory returning a model that uses learned global and local parameters for prediction.
    • bhm_at_scale.model.guide(X): global variational guide defining posterior parameters for all model sites.
    • bhm_at_scale.model.local_guide(model_params)(X): guide that conditions on global parameters and learns local store-level parameters.
    • bhm_at_scale.model.check_model_guide(X, model, guide): utility to sanity-check that a model and guide are syntactically compatible.
    • bhm_at_scale.handler.ModelHandler: class wrapping NumPyro SVI for initializing, fitting, checkpointing optimizer state, and running Predictive for inference.
    • bhm_at_scale.preprocess.encode(df, cols, drop_first): helper to dummy-encode specified categorical columns.
    • bhm_at_scale.preprocess.make_cube(df, idx_cols): converts a long-format DataFrame into a dense multi-dimensional NumPy array cube indexed by given columns.
    • bhm_at_scale.plot.plot_sales_preds(df, hlines, lines, intervals): function to visualize sales, customers, promotions, holidays, and prediction intervals for a single store.
    • bhm_at_scale.plot.PlotStore: functor for plotting history and prediction intervals of a specific store from aggregated prediction DataFrames.
    • bhm_at_scale.plot.plot_densities(df, xlim): helper to draw faceted density plots for coefficient or effect distributions.
    • bhm_at_scale.utils.summary(samples, poisson): computes posterior summary statistics (mean, std, percentiles) for sampled sites.
    • bhm_at_scale.utils.stats_to_df(stats, col_names): transforms summary statistics into a long-form pandas DataFrame.
    • bhm_at_scale.utils.preds_to_df(preds): converts summarized prediction arrays into a store-by-time-step DataFrame.
    • bhm_at_scale.utils.make_intervals(arr, offset): computes contiguous True intervals from a boolean array, used for promo/holiday spans.
    • bhm_at_scale.utils.reorder_cols(df, first, last): utility to reorder DataFrame columns with specified columns first and last.
• Code Overview
  • Components

• Entrypoints
  • Jupyter notebooks in the notebooks/ directory (01-preprocessing.ipynb, 02-model.ipynb, 03-evaluation.ipynb) that orchestrate data download, preprocessing, model training, prediction, and evaluation.
  • Python package modules under src/bhm_at_scale, which are imported from the notebooks (e.g., bhm_at_scale.model, bhm_at_scale.handler, bhm_at_scale.preprocess, bhm_at_scale.plot, bhm_at_scale.utils).
• Languages
  • Python
• Quality Signals
  • Coverage Hint: A coverage configuration (.coveragerc) is present and points coverage.py at the bhm_at_scale package, but the tests/ directory only contains an empty conftest.py and no actual unit tests.
  • Tests Present: false

Governance#

• Policies
  • The code is licensed under the GNU General Public License, version 2 or later (GPL-2.0-or-later).

Integration#

• Api: No network or REST API is provided; integration is via direct Python imports of the bhm_at_scale package and execution of the Jupyter notebooks.
• Ux Notes: The primary user experience is notebook-driven: users set up the conda environment, install the package, and then run the numbered Jupyter notebooks, which import the bhm_at_scale modules, produce inline plots, and write CSV and image outputs for inspection.

Meta#

• Created At: 2020
• Language: Python
• Links
  • Documentation: https://github.com/FlorianWilhelm/bhm-at-scale
  • Repo: https://github.com/FlorianWilhelm/bhm-at-scale
  • Repository: https://github.com/FlorianWilhelm/bhm-at-scale
• Maturity: beta
• Owners

• Tags
  • Bayesian hierarchical modelling
  • Gamma-Poisson model
  • retail sales forecasting
  • Rossmann Store Sales
  • time series
  • JAX
  • NumPyro
  • probabilistic programming
• Title: Bayesian Hierarchical Modelling at Scale

Ml Core#

• Artifact U R Is
  • data/preprocessed/edf.csv
  • data/preprocessed/X_train.npz
  • data/preprocessed/X_test.npz
  • data/result/df.csv
  • data/result/optim_state.pickle
  • data/result/coef_mus.csv
  • data/result/coef_sigmas.csv
  • data/result/stats.csv
  • data/result/train_preds.csv
  • data/result/test_preds.csv
  • docs/bhm_model.png
• Baselines
  • A conventional baseline uses scikit-learn's LinearRegression on log-transformed sales for a single store, fitted on a small history window and compared against the hierarchical Bayesian model's predictions.
• Datasets

• Features
  • DayOfWeek_1
  • DayOfWeek_2
  • DayOfWeek_3
  • DayOfWeek_4
  • DayOfWeek_5
  • DayOfWeek_6
  • DayOfWeek_7
  • Promo
  • StateHoliday_0
  • StateHoliday_1
  • StateHoliday_2
  • StateHoliday_3
  • SchoolHoliday
  • Promo2
  • StoreVariant_11
  • StoreVariant_13
  • StoreVariant_21
  • StoreVariant_22
  • StoreVariant_23
  • StoreVariant_31
  • StoreVariant_33
  • StoreVariant_41
  • StoreVariant_43
• Problem: Model and infer a Gamma-Poisson Bayesian hierarchical model for daily store-level sales, where store- and feature-specific coefficients and dispersion parameters are shared across many stores to improve forecasts, especially for stores with limited data, using the Rossmann Store Sales dataset.
• Training: Training is performed via stochastic variational inference (SVI) in NumPyro using the Trace_ELBO loss and an Adam optimizer. The ModelHandler class constructs the optimizer, initializes SVI with the model and guide, and runs JAX-jitted training loops over a 3D input tensor X of shape (n_stores, n_days, n_features+1), where the last dimension is the target. The 02-model notebook demonstrates fitting the global model on X_train for thousands of epochs with different learning rates, checkpointing and restoring optimizer state, and then fitting local store-level parameters on short history windows of X_test using local_guide to adapt the model to new stores.

Provenance#

• Changelog
  • Date: 2025-11-24T14:10:12.537Z
  • Summary: ML System Card run live-test-1763993411 observed 35 changed files
  • Files
    • Path: .coveragerc
    • Path: .gitignore
    • Path: .isort.cfg
    • Path: .pre-commit-config.yaml
    • Path: AUTHORS.rst
    • Path: CHANGELOG.rst
    • Path: LICENSE.txt
    • Path: README.md
    • Path: data/.gitignore
    • Path: data/external/.gitignore
    • Path: data/interim/.gitignore
    • Path: data/preprocessed/.gitignore
    • Path: data/raw/.gitignore
    • Path: data/result/.gitignore
    • Path: docs/Makefile
    • Path: docs/_static/.gitignore
    • Path: docs/authors.rst
    • Path: docs/changelog.rst
    • Path: docs/conf.py
    • Path: docs/index.rst
    • Path: docs/license.rst
    • Path: environment.lock.yaml
    • Path: environment.yaml
    • Path: notebooks/01-preprocessing.ipynb
    • Path: notebooks/02-model.ipynb
    • Path: notebooks/03-evaluation.ipynb
    • Path: setup.cfg
    • Path: setup.py
    • Path: src/bhm_at_scale/init.py
    • Path: src/bhm_at_scale/handler.py
    • Path: src/bhm_at_scale/model.py
    • Path: src/bhm_at_scale/plot.py
    • Path: src/bhm_at_scale/preprocess.py
    • Path: src/bhm_at_scale/utils.py
    • Path: tests/conftest.py
  • Run Id: live-test-1763993411
  • Head Sha: 4c5d5cd94a14ba3c523d4a5b7bc2b79a3acfc5dd

Stakeholder Notes#

• Data scientist: Current Analysis: The Gamma-Poisson hierarchical structure is holding up well for long-tail store distributions. However, the DayOfWeek drift suggests we may need to re-evaluate the seasonal priors.

• Prior Check: Re-run check_model_guide with wider sigma on temporal components. * Data Quality: Missing PromoInterval data in 2 regions causing singular matrices in local guide.

• Ml engineer: Deployment Readiness: JAX compilation times are stable, but memory usage during inference on X_test is creeping up.

• Action Item: optimize gen_partitions to allow smaller batch sizes during inference. * Env: environment.lock.yaml was updated; ensure CI pipeline matches new numpyro version.

• Product manager: Business Impact: The pilot for the first 50 stores shows +12% accuracy improvement over the legacy regression baseline.

• Rollout: Ready to expand to "Region North" next sprint. * Risks: Handling of holiday weeks is still manual (see StateHoliday mapping changes).