Sampling a hierarchical ordered logistic regression

Questions
4

I realise I’m two years too late, but I’m posting more for visibility as I was running into issues myself and I’ve managed to get a hierarchical ordered logistic regression working and sampling in PyMC (PyMC 4.0.0b1), so it may be helpful for others. I’ve been working through a problem in the book noted above, chapter 12 problem 12H2. Which asks the reader to fit a hierarchical ordered logistic regression model onto the Trolley data set in the book. See code below:

d = pd.read_csv('Data/Trolley.csv', delimiter=';')
d['id_factorised'] = pd.factorize(d.id)[0]
action = shared(d['action'].values)
intention = shared(d['intention'].values)
contact = shared(d['contact'].values)
N_ids = d['id'].unique().shape[0]

with pm.Model() as m_12h2_null:
    a = pm.Normal(
        'a', 0., 10.,
        transform=pm.distributions.transforms.ordered,
        shape=6, testval=np.arange(6) - 2.5)
    
    resp_obs = pm.OrderedLogistic(
        'resp_obs', 0., a,
        observed=d.response.values - 1
    )
    trace_m12h2_null = pm.sample(2000)

with pm.Model() as m_12h2_flexible:
    sigma = pm.HalfCauchy('sigma', 1.)
    a_id = pm.Normal('a_id', 0, sigma, shape=N_ids)
    
    a = pm.Normal(
        'a', 0., 10.,
        transform=pm.distributions.transforms.ordered,
        shape=6,
        initval=trace_m12h2_null['posterior']['a'].mean(dim=['chain', 'draw'])
    )
    
    bA = pm.Normal('bA', 0., 10.)
    bI = pm.Normal('bI', 0., 10.)
    bC = pm.Normal('bC', 0., 10.)
    phi  = a_id[d.id_factorised.values] + bA * action + bI * intention + bC * contact 
 
    resp_obs = pm.OrderedLogistic(
        'resp_obs', phi, a,
        observed=d.response - 1
    )
    trace_m12h2_flexible = pm.sample(4000, tune=1500)

The key thing to get it sampling was to set a sensible initval parameter in the prior for a. This was done by fitting a null model onto the data and using the trace of a from that null model as an initial value for initval. Sampling takes about 9 mins for this set-up and I think the model probably needs to be tweaked to improve sampling effeciency, but otherwise this will hopefully be a good start for anyone wanting to try and fit one of these models.