PyMC multilevel hierarchical model #pymc #hierarchical #censored

pymc_multilevel_hierarchical_model.py

def titer_hierarchy_model(
    df: pd.DataFrame,
    l1_name: str = "day",
    l2_name: str = "day_size",
    transform: Callable = day_size_to_day,
) -> pm.Model:
    """
    Args:
        df: DataFrame
        l1_name: Column that contains the first level.
        l2_name: Column that contains the second level.
        transform: Function that takes second level items and maps them to the first level.
            If the second level is "day_size", which is encoded in items like "d1_4.5ug" 
            in the "day_size" column then this function would want to extract a value of 
            `1` as the day (assuming that "day" is the first level and values of `1` occur
            in the "day" column)
    """

    l1 = pd.Categorical(df[l1_name])
    l2 = pd.Categorical(df[l2_name])

    # index that maps each value in level 2 to its correct level 1
    idx = [l1.categories.get_loc(transform(value)) for value in l2.categories]

    coords = {l1_name: l1.categories, l2_name: l2.categories}

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

        # Hyper priors
        mu = pm.Normal("mu", 2.5, 2.5)
        sigma = pm.Exponential("sigma", 1.0)

        # Level 1
        _j = pm.Normal("_j", 0, 1, dims=l1_name)  # non-centered
        mu_1 = pm.Deterministic(f"mu_{l1_name}", _j * sigma + mu, dims=l1_name)
        sigma_1 = pm.Exponential(f"sigma_{l1_name}", 1.0)

        # Level 2
        _k = pm.Normal("_k", 0, 1, dims=l2_name)  # non-centered
        mu_2 = pm.Deterministic(f"mu_{l2_name}", _k * sigma_1 + mu_1[idx], dims=l2_name)

        # Observed titers
        titer_latent = pm.StudentT.dist(
            mu=mu_2[l2.codes],
            sigma=pm.Exponential("sigma_obs", 1.0),
            nu=pm.Exponential("nu_obs", 1),
        )
        pm.Censored("obs", titer_latent, lower=0.0, observed=df["log10_titer"])

    return model