import pymc as pm
import numpy as np
import arviz as az
from pytensor.tensor import subtensor as st

17. Multinomial regression*#

Adapted from Unit 7: NHANESmulti.odc

Problem statement#

The National Health and Nutrition Examination Survey (NHANES) is a program of studies designed to assess the health and nutritional status of adults and children in the United States. The survey is unique in that it combines interviews and physical examinations.

Assume that N subjects select a choice form K categories. The i-th subject is characterized by 3 covariates x[i,1], x[i,2], and x[i,3]. Given the covariates, model the probability of a subject selecting the category k, k=1,…,K.

# data
# fmt: off
y = np.array([[1, 0, 0, 0, 0],
              [0, 1, 0, 0, 0],
              [1, 0, 0, 0, 0],
              [0, 0, 1, 0, 0],
              [0, 1, 0, 0, 0],
              [0, 0, 1, 0, 0],
              [0, 0, 0, 1, 0],
              [0, 0, 0, 0, 1],
              [0, 0, 0, 0, 1],
              [0, 0, 0, 1, 0]])

X =np.array([[2, 4, 9],
             [1, 5, 10],
             [1, 6, 14],
             [2, 4, 21],
             [2, 4, 22],
             [2, 6, 30],
             [3, 3, 33],
             [3, 2, 36],
             [3, 1, 40],
             [4, 1, 44]])
# fmt: on

y.shape

(10, 5)

X_aug = np.concatenate((np.ones((X.shape[0], 1)), X), axis=1)
X_aug.shape

(10, 4)

with pm.Model() as m:
    y_data = pm.Data("y", y)
    X_data = pm.Data("X", X_aug)

    _beta = pm.Normal("_beta", 0, tau=0.1, shape=(4, 5))
    # first col in BUGS example is assigned 0
    # might be better to prepend a column of zeroes rather than overwrite
    beta = pm.Deterministic("beta", st.set_subtensor(_beta[:, 0], 0))
    eta = pm.math.exp(pm.math.dot(X_data, beta))
    p = eta / pm.math.sum(eta, axis=1)

    pm.Multinomial("likelihood", n=1, p=p, observed=y_data, shape=X_data.shape)

    trace = pm.sample(10000)

Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [_beta]



Sampling 4 chains for 1_000 tune and 10_000 draw iterations (4_000 + 40_000 draws total) took 25 seconds.













X_new = np.array([1, 3, 3, 30]).reshape((1, 4))

with m:
    pm.set_data({"X": X_new})
    ppc = pm.sample_posterior_predictive(trace, predictions=True)









Sampling: [likelihood]















az.summary(ppc.predictions)













  
    

      
      mean
      sd
      hdi_3%
      hdi_97%
      mcse_mean
      mcse_sd
      ess_bulk
      ess_tail
      r_hat
    

  
  
    

      likelihood[0, 0]
      0.008
      0.090
      0.0
      0.0
      0.000
      0.000
      38988.0
      38988.0
      1.0
    

    

      likelihood[0, 1]
      0.039
      0.194
      0.0
      0.0
      0.001
      0.001
      40973.0
      40973.0
      1.0
    

    

      likelihood[0, 2]
      0.119
      0.324
      0.0
      1.0
      0.002
      0.001
      40603.0
      40000.0
      1.0
    

    

      likelihood[0, 3]
      0.683
      0.465
      0.0
      1.0
      0.002
      0.002
      40085.0
      40000.0
      1.0
    

    

      likelihood[0, 4]
      0.150
      0.357
      0.0
      1.0
      0.002
      0.001
      40403.0
      40000.0
      1.0
    

  











%load_ext watermark
%watermark -n -u -v -iv -p pytensor









Last updated: Sun Mar 09 2025

Python implementation: CPython
Python version       : 3.12.7
IPython version      : 8.29.0

pytensor: 2.26.4

arviz   : 0.20.0
pymc    : 5.19.1
pytensor: 2.26.4
numpy   : 1.26.4