1{
2  "op": "replace",
3  "filename": "scikit-learn/custom_cate.py",
4  "start_line": 345,
5  "end_line": 418,
6  "content": "class CATEEstimator(BaseCATEEstimator):\n    \"\"\"Overlap-Stabilized DR + R-Learner Ensemble (OSDRE).\n\n    Novel: Power-law propensity stabilization of DR augmentation terms\n    via ow^beta shrinkage interpolates between standard DR-learner\n    (unbiased, high variance) and X-learner (plug-in, low variance).\n    Combined with weighted R-learner for cross-strategy robustness.\n    \"\"\"\n    def __init__(self):\n        self.n_folds = 5\n        self.beta = 0.5  # shrinkage power: 0=DR, 0.5=sqrt, 1=X-learner\n    def fit(self, X, T, Y):\n        n = len(Y)\n        kf = KFold(n_splits=self.n_folds, shuffle=True, random_state=42)\n        mu0, mu1, e_hat = np.zeros(n), np.zeros(n), np.zeros(n)\n        for tr, va in kf.split(X):\n            c0, c1 = T[tr] == 0, T[tr] == 1\n            kw = dict(n_estimators=200, max_depth=4, learning_rate=0.1,\n                     min_samples_leaf=20, subsample=0.8, random_state=42)\n            m0 = GradientBoostingRegressor(**kw)\n            m1 = GradientBoostingRegressor(**kw)\n            m0.fit(X[tr[c0]], Y[tr[c0]])\n            m1.fit(X[tr[c1]], Y[tr[c1]])\n            mu0[va] = m0.predict(X[va])\n            mu1[va] = m1.predict(X[va])\n            mt = GradientBoostingClassifier(n_estimators=200, max_depth=3,\n                learning_rate=0.1, min_samples_leaf=20, subsample=0.8,\n                random_state=42)\n            mt.fit(X[tr], T[tr])\n            e_hat[va] = mt.predict_proba(X[va])[:, 1]\n        e_hat = np.clip(e_hat, 0.05, 0.95)\n        tau_plug = mu1 - mu0\n        # Stabilized DR: shrink augmentation by own-group propensity^beta\n        # beta=0 gives standard DR, beta=1 gives X-learner, beta=0.5 is sqrt\n        aug = T * (Y - mu1) / e_hat - (1 - T) * (Y - mu0) / (1 - e_hat)\n        ow = T * e_hat + (1 - T) * (1 - e_hat)  # own-group propensity\n        dr_pseudo = tau_plug + (ow ** self.beta) * aug\n        q = np.percentile(np.abs(dr_pseudo), 95)\n        dr_pseudo = np.clip(dr_pseudo, -q, q)\n        # R-learner: use marginal outcome from nuisance models\n        m_hat = e_hat * mu1 + (1 - e_hat) * mu0\n        T_res = T - e_hat\n        safe = np.where(np.abs(T_res) > 0.01, T_res,\n                        np.sign(T_res) * 0.01 + 1e-8)\n        r_pseudo = (Y - m_hat) / safe\n        r_w = T_res ** 2\n        qr = np.percentile(np.abs(r_pseudo), 95)\n        r_pseudo = np.clip(r_pseudo, -qr, qr)\n        # Final CATE models\n        ckw = dict(n_estimators=300, max_depth=3, learning_rate=0.05,\n                  min_samples_leaf=15, subsample=0.8, random_state=42)\n        self._m_dr = GradientBoostingRegressor(**ckw)\n        self._m_dr.fit(X, dr_pseudo)\n        self._m_rl = GradientBoostingRegressor(**ckw)\n        self._m_rl.fit(X, r_pseudo, sample_weight=r_w)\n        return self\n    def predict(self, X):\n        return 0.5 * self._m_dr.predict(X) + 0.5 * self._m_rl.predict(X)\n\n\n"
7}

1{
2  "op": "replace",
3  "filename": "scikit-learn/custom_cate.py",
4  "start_line": 345,
5  "end_line": 418,
6  "content": "class CATEEstimator(BaseCATEEstimator):\n    \"\"\"Ensemble-Nuisance Orthogonal CATE (ENOC) estimator.\n\n    Novel contributions:\n    (1) Ensemble nuisance estimation: averages S-learner and T-learner\n        predictions for mu0(X), mu1(X). S-learner pools all data for\n        data-efficient estimation; T-learner captures group-specific\n        patterns. Their average is more robust than either alone.\n    (2) Multi-strategy CATE ensemble: combines DR-Learner (doubly robust),\n        R-Learner (orthogonal debiasing), and S-Learner (regularized\n        stability) via uniform averaging for cross-DGP robustness.\n    \"\"\"\n    def __init__(self):\n        pass\n\n    def fit(self, X, T, Y):\n        n = len(Y)\n        kf = KFold(n_splits=5, shuffle=True, random_state=42)\n        mu0h, mu1h, mh, eh = [np.zeros(n) for _ in range(4)]\n        gk = dict(n_estimators=200, max_depth=4, learning_rate=0.1,\n                  min_samples_leaf=20, subsample=0.8, random_state=42)\n        gc = dict(n_estimators=200, max_depth=3, learning_rate=0.1,\n                  min_samples_leaf=20, subsample=0.8, random_state=42)\n        for tr, va in kf.split(X):\n            nv = len(va)\n            c0, c1 = T[tr] == 0, T[tr] == 1\n            # T-learner nuisance (separate models per treatment group)\n            mu0t = GradientBoostingRegressor(**gk).fit(\n                X[tr[c0]], Y[tr[c0]]).predict(X[va])\n            mu1t = GradientBoostingRegressor(**gk).fit(\n                X[tr[c1]], Y[tr[c1]]).predict(X[va])\n            # S-learner nuisance (pooled model with T as feature)\n            ms = GradientBoostingRegressor(**gk).fit(\n                np.c_[X[tr], T[tr]], Y[tr])\n            mu0s = ms.predict(np.c_[X[va], np.zeros(nv)])\n            mu1s = ms.predict(np.c_[X[va], np.ones(nv)])\n            # Ensemble nuisance: average for robustness\n            mu0h[va] = 0.5 * mu0t + 0.5 * mu0s\n            mu1h[va] = 0.5 * mu1t + 0.5 * mu1s\n            # Marginal outcome E[Y|X] for R-learner\n            mh[va] = GradientBoostingRegressor(**gk).fit(\n                X[tr], Y[tr]).predict(X[va])\n            # Propensity score\n            eh[va] = GradientBoostingClassifier(**gc).fit(\n                X[tr], T[tr]).predict_proba(X[va])[:, 1]\n        eh = np.clip(eh, 0.05, 0.95)\n        # DR pseudo-outcomes with winsorization\n        dr = (mu1h - mu0h + T * (Y - mu1h) / eh\n              - (1 - T) * (Y - mu0h) / (1 - eh))\n        q = np.percentile(np.abs(dr), 97)\n        dr = np.clip(dr, -q, q)\n        # R-learner pseudo-outcomes with safe division\n        Yr, Tr = Y - mh, T - eh\n        sT = np.where(np.abs(Tr) > 0.01, Tr,\n                      np.sign(Tr + 1e-10) * 0.01)\n        rp, rw = Yr / sT, Tr ** 2\n        q2 = np.percentile(np.abs(rp), 95)\n        rp = np.clip(rp, -q2, q2)\n        # Final CATE models\n        fk = dict(n_estimators=300, max_depth=3, learning_rate=0.05,\n                  min_samples_leaf=20, subsample=0.8, random_state=42)\n        self._dr = GradientBoostingRegressor(**fk).fit(X, dr)\n        self._r = GradientBoostingRegressor(**fk).fit(\n            X, rp, sample_weight=rw)\n        self._s = GradientBoostingRegressor(**gk).fit(\n            np.c_[X, T], Y)\n        return self\n\n    def predict(self, X):\n        n = X.shape[0]\n        dp = self._dr.predict(X)\n        rp = self._r.predict(X)\n        sp = (self._s.predict(np.c_[X, np.ones(n)])\n              - self._s.predict(np.c_[X, np.zeros(n)]))\n        return (dp + rp + sp) / 3.0\n\n"
7}

1{
2  "op": "replace",
3  "filename": "scikit-learn/custom_cate.py",
4  "start_line": 345,
5  "end_line": 418,
6  "content": "class CATEEstimator(BaseCATEEstimator):\n    \"\"\"Overlap-Adaptive Doubly Robust Ensemble (OADRE) for CATE.\n    Novel: (1) Overlap-adaptive blending of DR correction and T-learner\n    plug-in via h(x)=min(e,1-e). (2) OADR + R-learner ensemble.\n    (3) Iterative influence trimming for pseudo-outcome denoising.\"\"\"\n\n    def __init__(self):\n        self._nf = 3\n\n    def fit(self, X, T, Y):\n        n = len(Y)\n        kf = KFold(n_splits=self._nf, shuffle=True, random_state=42)\n        mu0, mu1, mY, eX = [np.zeros(n) for _ in range(4)]\n        kw = dict(n_estimators=200, max_depth=4, learning_rate=0.1,\n                  min_samples_leaf=15, subsample=0.8, random_state=42)\n        for tr, va in kf.split(X):\n            t0, t1 = T[tr] == 0, T[tr] == 1\n            g0 = GradientBoostingRegressor(**kw)\n            g1 = GradientBoostingRegressor(**kw)\n            gm = GradientBoostingRegressor(**kw)\n            g0.fit(X[tr[t0]], Y[tr[t0]]); g1.fit(X[tr[t1]], Y[tr[t1]])\n            gm.fit(X[tr], Y[tr])\n            mu0[va] = g0.predict(X[va])\n            mu1[va] = g1.predict(X[va])\n            mY[va] = gm.predict(X[va])\n            gc = GradientBoostingClassifier(n_estimators=200, max_depth=3,\n                learning_rate=0.1, min_samples_leaf=15, subsample=0.8, random_state=42)\n            gc.fit(X[tr], T[tr])\n            eX[va] = gc.predict_proba(X[va])[:, 1]\n        eX = np.clip(eX, 0.05, 0.95)\n        # --- OADR: overlap-adaptive DR pseudo-outcomes ---\n        plugin = mu1 - mu0\n        corr = T * (Y - mu1) / eX - (1 - T) * (Y - mu0) / (1 - eX)\n        h = np.minimum(eX, 1 - eX)\n        alpha = np.minimum(1.0, h / 0.15)\n        psi_dr = plugin + alpha * corr\n        vp = T / eX**2 + (1 - T) / (1 - eX)**2\n        w_dr = 1.0 / (vp + np.median(vp))\n        w_dr = w_dr / w_dr.mean()\n        # --- R-learner pseudo-outcomes ---\n        Tr, Yr = T - eX, Y - mY\n        sT = np.where(np.abs(Tr) > 0.01, Tr, np.sign(Tr) * 0.01 + 1e-8)\n        psi_r = Yr / sT\n        q = np.percentile(np.abs(psi_r), 95)\n        psi_r = np.clip(psi_r, -q, q)\n        w_r = Tr ** 2\n        w_r = w_r / (w_r.mean() + 1e-12)\n        # --- Fit OADR with iterative influence trimming ---\n        fkw = dict(n_estimators=300, max_depth=4, learning_rate=0.05,\n                   min_samples_leaf=15, subsample=0.8, random_state=42)\n        f1 = GradientBoostingRegressor(**fkw)\n        f1.fit(X, psi_dr, sample_weight=w_dr)\n        dev = np.abs(psi_dr - f1.predict(X))\n        thr = np.percentile(dev, 90)\n        rw = w_dr * np.minimum(1.0, thr / (dev + 1e-8))\n        self._f1 = GradientBoostingRegressor(**fkw)\n        self._f1.fit(X, psi_dr, sample_weight=rw)\n        # --- Fit R-learner model ---\n        self._f2 = GradientBoostingRegressor(**fkw)\n        self._f2.fit(X, psi_r, sample_weight=w_r)\n        return self\n\n    def predict(self, X):\n        return 0.5 * self._f1.predict(X) + 0.5 * self._f2.predict(X)\n\n\n"
7}

1{
2  "n": 2
3}