Multi-level model with nonsense state and county coefficients

Questions
4

I’m not yet familiar with pm.Data usage, so I got rid of it and just used the idx arrays generated by pandas factorize.
Also, I think mu_a_county is one too many intercepts, so I got rid of that as well.
I don’t have graphviz installed so I can’t check that, but this seems to work:

with pm.Model(coords=coords) as hierarchical_model_varying_intercept_county_state:

    # Hyper Priors: State
    sigma_a_state = pm.HalfNormal("sigma_a_state", 3.)
    mu_a_state = pm.Normal("mu_a_state", mu=9., sigma = 2.5)

    # State coefficients
    a_state_offset = pm.Normal("a_state_offset", mu = 0., sigma = 1., dims="state")
    a_state = pm.Deterministic("a_state", mu_a_state + sigma_a_state*a_state_offset) 

    # Hyper Priors: County
    sigma_a_county = pm.HalfNormal("sigma_a_county", 3.)
    
    # County coefficients
    a_county_offset = pm.Normal("a_county_offset", mu = 0., sigma = 1., dims= "county")
    a_county = pm.Deterministic("a_county", a_state[state] +  sigma_a_county*a_county_offset[county])

    # Model error
    epsilon = pm.HalfCauchy("eps", 5.0)

    # Data likelihood
    disease_like = pm.Normal("disease_likelihood", 
                           mu=a_county, 
                           sigma=epsilon, 
                           observed = sample_data.disease_count.values)