Update causal.py

causal.py

@@ -1,4 +1,17 @@
 from __future__ import annotations
+import os
+
+# ---- PyTensor(=PyMCの内部) のコンパイルキャッシュ先を /tmp に固定 ----
+# 環境変数で上書き可: PYTENSOR_BASE=/data/.pytensor（永続化したい場合）
+base = os.environ.get("PYTENSOR_BASE", "/tmp/.pytensor")
+os.makedirs(base, exist_ok=True)
+if "PYTENSOR_FLAGS" in os.environ and "base_compiledir=" not in os.environ["PYTENSOR_FLAGS"]:
+    os.environ["PYTENSOR_FLAGS"] = os.environ["PYTENSOR_FLAGS"] + f",base_compiledir={base}"
+elif "PYTENSOR_FLAGS" not in os.environ:
+    os.environ["PYTENSOR_FLAGS"] = f"base_compiledir={base}"
+os.environ.setdefault("XDG_CACHE_HOME", "/tmp/.cache")
+
+# ---- 以降は従来どおり ----
 import numpy as np
 import pandas as pd
 import pymc as pm
@@ -45,7 +58,7 @@ def fit_uplift_binary(
         p = 0
 
     with pm.Model() as model:
-        n_medium = int(d["medium_idx"].max()) + 1 if hasattr(d["medium_idx"], "max") else int(np.max(d["medium_idx"])) + 1
+        n_medium = int(np.max(d["medium_idx"])) + 1
         mu_re = pm.Normal("mu_re", 0.0, 1.0)
         sd_re = pm.HalfNormal("sd_re", 1.0)
         z_re = pm.Normal("z_re", 0.0, 1.0, shape=n_medium)
@@ -71,9 +84,11 @@ def fit_uplift_binary(
         else:
             pm.Bernoulli("y_obs", p=p_click, observed=d["y"].values)
 
-        idata = pm.sample(draws=draws, tune=draws, chains=2,
-                          target_accept=target_accept, random_seed=random_seed,
-                          progressbar=False)
+        idata = pm.sample(
+            draws=draws, tune=draws, chains=2,
+            target_accept=target_accept, random_seed=random_seed,
+            progressbar=False
+        )
 
     post = idata.posterior
     b0_s = post["intercept"].stack(sample=("chain", "draw"))