Hierarchical Model - Slow Sampling ( 1.8 sec per draw ) - Primary suspect is how the multi-level hierarchy is coded

Questions
3

I have a solution that reduced the training time tremendously. I’ll provide the code along with some today data so that it’s easier to follow:

import arviz as az
import numpy as np
import pandas as pd
import pymc3 as pm

toy_data = pd.DataFrame({"region":["West"]*25 + ["East"]*10, "state":["CA"]*15 + ["NV"]*10 + ["NY"]*10,
               "county":["Los Angeles"]*10 + ["San Diego"]*3 + ["Orange"]*2 + ["Clark"]*5 + ["Washoe"]*5 + ["Kings"]*5 + ["Queens"]*5,
               "y": list(2+ 0.1*np.random.randn(10)) + list(2.5 + 0.1*np.random.randn(3)) + list(3. + 0.1*np.random.randn(2)) + list(5. + 0.1*np.random.randn(5)) + list(5.5 + 0.1*np.random.randn(5)) + list(8. + 0.1*np.random.randn(5)) + list(8.5 + 0.1*np.random.randn(5))})

toy_data["state_region"] = toy_data["state"] + "_" + toy_data["region"]
toy_data["state_region"] = toy_data["state_region"].astype("category")
toy_data["county_state_region"] = toy_data["county"] + "_" + toy_data["state"] + "_" + toy_data["region"]
toy_data["county_state_region"] = toy_data["county_state_region"].astype("category")


toy_data["region"] = toy_data["region"].astype("category")
toy_data["state"] = toy_data["state"].astype("category")
toy_data["county"] = toy_data["county"].astype("category")

region_index = (
    toy_data[["region", "state", "state_region"]]
    .drop_duplicates()["region"]
    .cat.codes
)

state_index = (
    toy_data[[ "region", "state", "county", "county_state_region"]]
    .drop_duplicates()["state"]
    .cat.codes
)
county = toy_data["county"].cat.codes

number_of_regions = toy_data["region"].nunique()
number_of_states = toy_data["state"].nunique()
number_of_counties = toy_data["county"].nunique()

with pm.Model() as toy_model:

    # Hyper Priors: Region
    sigma_a_region = pm.HalfNormal("sigma_a_region", 1.)
    mu_a_region = pm.Normal("mu_a_region", mu = 5., sigma = 2.5,)
    a_region_offset = pm.Normal("a_region_offset", 0., 1., shape = number_of_regions)

    # Country coefficients
    a_region = pm.Deterministic("a_region", mu_a_region + sigma_a_region*a_region_offset)

    # Hyper Priors: State
    sigma_a_state = pm.HalfNormal("sigma_a_state", 2.)
    a_state_offset = pm.Normal("a_state_offset", 0., 1., shape = number_of_states)

    # State coefficients
    a_state = pm.Deterministic("a_state", a_region[region_index.values] + sigma_a_state*a_state_offset)    

    # Hyper Priors: County
    sigma_a_county = pm.HalfNormal("sigma_a_county", 2.)
    a_county_offset = pm.Normal("a_county_offset", 0., 1., shape = number_of_counties)

    # County coefficients
    a_county = pm.Deterministic("a_county", a_state[state_index.values] + sigma_a_county*a_county_offset)

    # Model error
    epsilon = pm.HalfNormal("eps", 5.)

    yield_estimate = (a_county[county.values]) # + pm.math.dot(X[modeling_columns].values, beta) )

    # Data likelihood
    yield_like = pm.Normal("yield_likelihood", 
                           mu=yield_estimate, 
                           sigma=epsilon, 
                           observed = toy_data.y,
                           )

model_graph
model_graph1068×776 45.2 KB

With this model structure, I no longer have the a_state and a_county down at the observation level, which seems to tremendously boosted the sampling rate. Also, I changed the eps HalfCauchy to a HalfNormal, which has also sped up the rate of sampling.