I am running a hierarchical linear regression model with “small effect sizes” and relatively large variability of the effect size across different groups. In addition there is also a very large variance of the observation around its expectation.
When sampling with a large number of groups I receive the following warning
...lib/python3.11/site-packages/nutpie/compile_pymc.py:483: NumbaWarning: Cannot cache compiled function "numba_funcified_fgraph" as it uses dynamic globals (such as ctypes pointers and large global arrays)
return inner(x, *_shared_tuple)
I am wondering if I can change something in the definition of my model (cf. code full example below) in order to improve the efficiency of sampling and get rid off the warning.
Specifically I am wondering if I
- use integers instead of dates to label groups in the definition of
coords - explictly convert
pandas.Seriestonumpy.ndarray, e.g. useframe["y"].to_numpy()instead offrame["y"]. - Remove the mutable groups, but then how would I make predictions with the model?
- Change this line
mu = (beta_t[t] * x).sum(axis=-1)to a different expression that computes the expression (vector product) for all date groups t simultaneously - using mu[t] in the likelihood instead?
If yes, how can I do that? - Use additional changes/transformations of my problem to speed up things
In addition I am observing that the posterior (summary) for variables tau[x_4] does not seem to “cover” the “true” value, although the model specification closely follows the data generating process. The same is true for tau[x_5] and tau[x_6] also.
In general the posterior seems to “over-estimate” the values for the across group variability of my regression slopes, tau. Is this just a result of inefficient sampling or am I overlooking something else here.
I am omitting any intercept (on population or group level) in my model since the generating process has a zero intercept. When I add an intercept to the model the sampling efficiency gets worse.
I am using nutpie with a compiled model in order to speed-up sampling but the results are qualitatively identical when using the python based NUTS sampler.
I haven’t been able to reproduce this using a more concise example, so I am posting the rather lengthy full example here.
The full example can be seen below.
import time
import nutpie
import pytensor
pytensor.config.gcc__cxxflags += " -isysroot /Library/Developer/CommandLineTools/SDKs/MacOSX.sdk"
import numpy as np
import pandas as pd
import pymc as pm
import arviz as az
BETA = [0.0008, 0.0011, -0.0012, 0.0012, -0.0010, -0.0040]
TAU = [0.008, 0.010, 0.008, 0.006, 0.010, 0.019]
SIGMA = 0.07
def generate_data(num_groups: int = 5, num_samples_per_group: int = 20) -> pd.DataFrame:
rng = np.random.default_rng(seed=42)
dates = pd.date_range("2001-01-01", freq="B", periods=num_groups)
ids = [f"id_{i+1}" for i in range(num_samples_per_group)]
num_predictors = len(BETA)
X = []
y = []
for _ in range(num_groups):
# Use unit random normal predictors for each group
Xg = rng.normal(loc=0.0, scale=1.0, size=(num_samples_per_group, num_predictors))
beta_g = rng.normal(loc=BETA, scale=TAU)
yg = Xg.dot(beta_g) + SIGMA * rng.normal(size=num_samples_per_group)
X.append(Xg)
y.append(yg)
X_ = pd.DataFrame(np.concatenate(X, axis=0), columns=[f"x_{i+1}" for i in range(num_predictors)])
y_ = pd.DataFrame(np.concatenate(y, axis=0), columns=["y"])
frame = pd.concat([X_, y_], axis=1)
dates_ = np.repeat(dates, num_samples_per_group)
frame["date"] = dates_
ids_ = ids * num_groups
frame["id"] = ids_
return frame.set_index(["date", "id"])
def make_model(frame: pd.DataFrame) -> pm.Model:
unique_dates = [t.date() for t in frame.index.unique("date")]
predictors = [col for col in frame.columns if col.startswith("x")]
date_idx = frame.index.get_level_values("date").map(
dict(zip(unique_dates, range(len(unique_dates))))
)
coords = {
"date": unique_dates,
"predictor": predictors
}
with pm.Model(coords=coords) as model:
# Data
x = pm.MutableData("x", frame[predictors])
t = pm.MutableData("t", date_idx)
y = pm.ConstantData("y", frame["y"])
# Panel/population level
beta = pm.Normal("beta", dims="predictor")
tau = pm.HalfNormal("tau", 0.01, dims="predictor")
sigma = pm.HalfNormal("sigma", 0.15)
# Group level
epsilon_beta_t = pm.ZeroSumNormal("epsilon_beta_t", dims=["date", "predictor"])
beta_t = beta + tau * epsilon_beta_t
# Linear model, is this implementation efficient?
mu = (beta_t[t] * x).sum(axis=-1)
#Likelihood
pm.Normal("y_obs", mu=mu, sigma=sigma, observed=y, shape=mu.shape)
return model
def add_truth(summary):
ind = summary.index.str.startswith("beta")
summary.loc[ind, "true"] = BETA
ind = summary.index.str.startswith("tau")
summary.loc[ind, "true"] = TAU
summary.loc["sigma", "true"] = SIGMA
tmp = summary.pop("true")
summary.insert(0, "true", tmp)
if __name__ == "__main__":
frame = generate_data(num_groups=500, num_samples_per_group=270)
model = make_model(frame)
t0 = time.time()
trace = nutpie.sample(nutpie.compile_pymc_model(model))
t = time.time() - t0
print(f"Time for sampling is {t=:.3f}s.")
summary = az.summary(trace, var_names=["beta", "tau", "sigma"]).loc[:, ["mean", "hdi_3%", "hdi_97%", "ess_bulk", "ess_tail", "r_hat"]]
add_truth(summary)
print(summary)
>>>
Time for sampling is t=507.438s.
true mean hdi_3% hdi_97% ess_bulk ess_tail r_hat
beta[x_1] 0.0008 0.000 -0.000 0.001 1324.0 2626.0 1.00
beta[x_2] 0.0011 0.001 0.000 0.002 805.0 1612.0 1.01
beta[x_3] -0.0012 -0.002 -0.002 -0.001 1248.0 1956.0 1.00
beta[x_4] 0.0012 0.001 0.001 0.002 12081.0 4090.0 1.00
beta[x_5] -0.0010 -0.001 -0.002 -0.000 852.0 1444.0 1.01
beta[x_6] -0.0040 -0.003 -0.005 -0.001 580.0 1115.0 1.01
tau[x_1] 0.0080 0.009 0.008 0.009 1670.0 2868.0 1.00
tau[x_2] 0.0100 0.012 0.011 0.012 1462.0 2502.0 1.00
tau[x_3] 0.0080 0.009 0.008 0.010 1687.0 2902.0 1.00
tau[x_4] 0.0060 0.000 0.000 0.001 1337.0 1899.0 1.00
tau[x_5] 0.0100 0.012 0.011 0.013 1299.0 2637.0 1.00
tau[x_6] 0.0190 0.026 0.025 0.028 1015.0 2409.0 1.01
sigma 0.0700 0.070 0.070 0.071 11587.0 4065.0 1.00
If you have any suggestions/ideas for debugging this I’d be very glad.