1{
2  "op": "replace",
3  "filename": "scikit-learn/custom_cate.py",
4  "content": "class CATEEstimator(BaseCATEEstimator):\n    \"\"\"Overlap-adaptive orthogonal ensemble for CATE.\n\n    Combines: (i) cross-fitted T-learner nuisances, (ii) a stabilized\n    orthogonal Newton-style pseudo-outcome, (iii) a DR pseudo-outcome,\n    and (iv) R-risk based aggregation of nonlinear and polynomial heads.\n    \"\"\"\n\n    def __init__(self):\n        self.k = 3\n        self.clip = 0.03\n\n    def _y(self):\n        return GradientBoostingRegressor(\n            n_estimators=160, max_depth=3, learning_rate=0.05,\n            min_samples_leaf=15, subsample=0.8, loss=\"huber\", random_state=42,\n        )\n\n    def _t(self):\n        return GradientBoostingClassifier(\n            n_estimators=160, max_depth=3, learning_rate=0.05,\n            min_samples_leaf=15, subsample=0.8, random_state=42,\n        )\n\n    def _tau(self):\n        return GradientBoostingRegressor(\n            n_estimators=220, max_depth=3, learning_rate=0.03,\n            min_samples_leaf=20, subsample=0.8, loss=\"huber\", random_state=42,\n        )\n\n    def fit(self, X, T, Y):\n        X = np.asarray(X)\n        T = np.asarray(T).ravel()\n        Y = np.asarray(Y).ravel()\n        n = len(Y)\n        kf = KFold(n_splits=self.k, shuffle=True, random_state=123)\n        mu0_hat, mu1_hat, e_hat = np.zeros(n), np.zeros(n), np.zeros(n)\n\n        for tr, va in kf.split(X):\n            tr0, tr1 = tr[T[tr] == 0], tr[T[tr] == 1]\n            m0, m1 = self._y(), self._y()\n            m0.fit(X[tr0], Y[tr0])\n            m1.fit(X[tr1], Y[tr1])\n            mu0_hat[va], mu1_hat[va] = m0.predict(X[va]), m1.predict(X[va])\n            gb = self._t().fit(X[tr], T[tr])\n            lr = Pipeline([\n                (\"s\", StandardScaler()),\n                (\"l\", LogisticRegression(C=1.0, max_iter=500)),\n            ]).fit(X[tr], T[tr])\n            e_hat[va] = 0.6 * gb.predict_proba(X[va])[:, 1] + 0.4 * lr.predict_proba(X[va])[:, 1]\n\n        e_hat = np.clip(e_hat, self.clip, 1 - self.clip)\n        tau_t = mu1_hat - mu0_hat\n        m_hat = mu0_hat + e_hat * tau_t\n        r = T - e_hat\n        overlap = e_hat * (1 - e_hat)\n        phi_dr = tau_t + T * (Y - mu1_hat) / e_hat - (1 - T) * (Y - mu0_hat) / (1 - e_hat)\n        denom = r ** 2 + 0.1 * overlap + 1e-4\n        phi_nr = tau_t + r * ((Y - m_hat) - r * tau_t) / denom\n        q = np.percentile(np.abs(np.r_[phi_dr, phi_nr]), 99)\n        phi_dr = np.clip(phi_dr, -q, q)\n        phi_nr = np.clip(phi_nr, -q, q)\n        w = overlap + 0.05\n\n        self._scaler = StandardScaler().fit(X)\n        Xs = self._scaler.transform(X)\n        self._poly = PolynomialFeatures(degree=2, include_bias=False).fit(Xs)\n        Z = self._poly.transform(Xs)\n        oof_gb, oof_rg = np.zeros(n), np.zeros(n)\n        for tr, va in kf.split(X):\n            oof_gb[va] = self._tau().fit(X[tr], phi_nr[tr], sample_weight=w[tr]).predict(X[va])\n            oof_rg[va] = Ridge(alpha=3.0).fit(Z[tr], phi_dr[tr], sample_weight=w[tr]).predict(Z[va])\n\n        self._m0, self._m1 = self._y(), self._y()\n        self._m0.fit(X[T == 0], Y[T == 0])\n        self._m1.fit(X[T == 1], Y[T == 1])\n        self._gb = self._tau().fit(X, phi_nr, sample_weight=w)\n        self._rg = Ridge(alpha=3.0).fit(Z, phi_dr, sample_weight=w)\n\n        resid = Y - m_hat\n        H = np.column_stack([tau_t, oof_gb, oof_rg])\n        risk = np.mean(w[:, None] * (resid[:, None] - r[:, None] * H) ** 2, axis=0)\n        a = np.exp(-(risk - risk.min()) / (np.var(resid) + 1e-8))\n        self._alpha = a / a.sum()\n        full_H = np.column_stack([\n            self._m1.predict(X) - self._m0.predict(X),\n            self._gb.predict(X),\n            self._rg.predict(Z),\n        ])\n        self._offset = phi_dr.mean() - full_H.dot(self._alpha).mean()\n        return self\n\n    def predict(self, X):\n        X = np.asarray(X)\n        Z = self._poly.transform(self._scaler.transform(X))\n        H = np.column_stack([\n            self._m1.predict(X) - self._m0.predict(X),\n            self._gb.predict(X),\n            self._rg.predict(Z),\n        ])\n        return H.dot(self._alpha) + self._offset\n",
5  "after_line": 344,
6  "start_line": 345,
7  "end_line": 418
8}

1{
2  "op": "replace",
3  "filename": "scikit-learn/custom_cate.py",
4  "content": "class CATEEstimator(BaseCATEEstimator):\n    \"\"\"Basis-augmented orthogonal learner with overlap-adaptive shrinkage.\n\n    The estimator builds two orthogonal CATE heads from cross-fitted nuisances:\n    a nonlinear Newton-style head and a structured basis DR head. It then\n    learns an overlap-dependent shrinkage rule that interpolates between a\n    stable T-learner and the orthogonal ensemble using R-risk.\n    \"\"\"\n\n    def __init__(self):\n        self.k = 3\n        self.clip = 0.03\n        self.prop_mix = 0.6\n\n    def _y(self):\n        return GradientBoostingRegressor(\n            n_estimators=140, max_depth=3, learning_rate=0.05,\n            min_samples_leaf=12, subsample=0.8, loss=\"huber\", random_state=42,\n        )\n\n    def _t(self):\n        return GradientBoostingClassifier(\n            n_estimators=140, max_depth=3, learning_rate=0.05,\n            min_samples_leaf=12, subsample=0.8, random_state=42,\n        )\n\n    def _tau(self):\n        return GradientBoostingRegressor(\n            n_estimators=180, max_depth=3, learning_rate=0.04,\n            min_samples_leaf=18, subsample=0.8, loss=\"huber\", random_state=42,\n        )\n\n    def _basis(self, X):\n        Xs = self._scaler.transform(np.asarray(X))\n        Xq = Xs[:, :self.q]\n        return np.hstack([\n            Xs,\n            np.abs(Xq),\n            np.maximum(Xq, 0.0),\n            np.sin(np.pi * Xq),\n            np.cos(np.pi * Xq),\n            np.sin(0.5 * np.pi * Xq),\n            self._poly.transform(Xq),\n        ])\n\n    def fit(self, X, T, Y):\n        X = np.asarray(X)\n        T = np.asarray(T).ravel()\n        Y = np.asarray(Y).ravel()\n        n = len(Y)\n        self.y_scale_ = Y.std() + 1e-8\n        Ys = (Y - Y.mean()) / self.y_scale_\n        kf = KFold(n_splits=self.k, shuffle=True, random_state=123)\n        mu0_hat, mu1_hat, e_hat = np.zeros(n), np.zeros(n), np.zeros(n)\n        g0 = Ys[T == 0].mean() if np.any(T == 0) else Ys.mean()\n        g1 = Ys[T == 1].mean() if np.any(T == 1) else Ys.mean()\n\n        for tr, va in kf.split(X):\n            tr0, tr1 = tr[T[tr] == 0], tr[T[tr] == 1]\n            if len(tr0) > 8:\n                m0 = self._y().fit(X[tr0], Ys[tr0])\n                mu0_hat[va] = m0.predict(X[va])\n            else:\n                mu0_hat[va] = g0\n            if len(tr1) > 8:\n                m1 = self._y().fit(X[tr1], Ys[tr1])\n                mu1_hat[va] = m1.predict(X[va])\n            else:\n                mu1_hat[va] = g1\n            gb = self._t().fit(X[tr], T[tr])\n            lr = Pipeline([\n                (\"s\", StandardScaler()),\n                (\"l\", LogisticRegression(C=1.0, max_iter=500)),\n            ]).fit(X[tr], T[tr])\n            e_hat[va] = self.prop_mix * gb.predict_proba(X[va])[:, 1] + (1 - self.prop_mix) * lr.predict_proba(X[va])[:, 1]\n\n        e_hat = np.clip(e_hat, self.clip, 1 - self.clip)\n        tau_t = mu1_hat - mu0_hat\n        m_hat = mu0_hat + e_hat * tau_t\n        r = T - e_hat\n        overlap = e_hat * (1 - e_hat)\n        phi_dr = tau_t + T * (Ys - mu1_hat) / e_hat - (1 - T) * (Ys - mu0_hat) / (1 - e_hat)\n        phi_on = tau_t + r * ((Ys - m_hat) - r * tau_t) / (r ** 2 + 0.25 * overlap + 1e-4)\n        c = np.percentile(np.abs(np.r_[phi_dr, phi_on]), 98)\n        phi_dr = np.clip(phi_dr, -c, c)\n        phi_on = np.clip(phi_on, -c, c)\n        w = 0.1 + 4.0 * overlap\n        resid = Ys - m_hat\n\n        self._scaler = StandardScaler().fit(X)\n        self.q = min(X.shape[1], 10)\n        self._poly = PolynomialFeatures(degree=2, include_bias=False).fit(self._scaler.transform(X)[:, :self.q])\n        B = self._basis(X)\n        oof_tree, oof_lin = np.zeros(n), np.zeros(n)\n        for tr, va in kf.split(X):\n            oof_tree[va] = self._tau().fit(X[tr], phi_on[tr], sample_weight=w[tr]).predict(X[va])\n            oof_lin[va] = Ridge(alpha=1.5).fit(B[tr], phi_dr[tr], sample_weight=w[tr]).predict(B[va])\n\n        H = np.column_stack([oof_tree, oof_lin])\n        risk = np.mean(w[:, None] * (resid[:, None] - r[:, None] * H) ** 2, axis=0)\n        a = np.exp(-(risk - risk.min()) / (np.var(resid) + 1e-8))\n        self.alpha_ = a / a.sum()\n        orth_oof = H.dot(self.alpha_)\n        base_gate = np.clip(4.0 * overlap, 0.0, 1.0)\n        gammas = np.array([0.0, 0.5, 1.0, 2.0, 4.0])\n        mix_risk = []\n        for g in gammas:\n            gate = base_gate ** g\n            tau_mix = (1 - gate) * tau_t + gate * orth_oof\n            mix_risk.append(np.mean(w * (resid - r * tau_mix) ** 2))\n        self.gamma_ = gammas[int(np.argmin(mix_risk))]\n\n        self._m0, self._m1 = self._y(), self._y()\n        self._m0.fit(X[T == 0], Ys[T == 0])\n        self._m1.fit(X[T == 1], Ys[T == 1])\n        self._pgb = self._t().fit(X, T)\n        self._plr = Pipeline([\n            (\"s\", StandardScaler()),\n            (\"l\", LogisticRegression(C=1.0, max_iter=500)),\n        ]).fit(X, T)\n        self._tree = self._tau().fit(X, phi_on, sample_weight=w)\n        self._lin = Ridge(alpha=1.5).fit(B, phi_dr, sample_weight=w)\n\n        e_full = np.clip(self.prop_mix * self._pgb.predict_proba(X)[:, 1] + (1 - self.prop_mix) * self._plr.predict_proba(X)[:, 1], self.clip, 1 - self.clip)\n        gate = np.clip(4.0 * e_full * (1 - e_full), 0.0, 1.0) ** self.gamma_\n        orth = self.alpha_[0] * self._tree.predict(X) + self.alpha_[1] * self._lin.predict(B)\n        tau_full = (1 - gate) * (self._m1.predict(X) - self._m0.predict(X)) + gate * orth\n        self.offset_ = phi_dr.mean() - tau_full.mean()\n        return self\n\n    def predict(self, X):\n        X = np.asarray(X)\n        B = self._basis(X)\n        e = np.clip(self.prop_mix * self._pgb.predict_proba(X)[:, 1] + (1 - self.prop_mix) * self._plr.predict_proba(X)[:, 1], self.clip, 1 - self.clip)\n        gate = np.clip(4.0 * e * (1 - e), 0.0, 1.0) ** self.gamma_\n        orth = self.alpha_[0] * self._tree.predict(X) + self.alpha_[1] * self._lin.predict(B)\n        tau = (1 - gate) * (self._m1.predict(X) - self._m0.predict(X)) + gate * orth + self.offset_\n        return tau * self.y_scale_\n",
5  "after_line": 344,
6  "start_line": 345,
7  "end_line": 445
8}

1{
2  "n": 2
3}

1{
2  "op": "replace",
3  "filename": "scikit-learn/custom_cate.py",
4  "content": "class CATEEstimator(BaseCATEEstimator):\n    \"\"\"Cross-fitted orthogonal stack with overlap-aware calibration.\"\"\"\n\n    def __init__(self):\n        self.k = 3\n        self.clip = 0.03\n\n    def _y(self, d=3, lin=False):\n        return Pipeline([(\"s\", StandardScaler()), (\"m\", Ridge(alpha=1.0))]) if lin else GradientBoostingRegressor(\n            n_estimators=120, max_depth=d, learning_rate=0.05,\n            min_samples_leaf=20, subsample=0.8, random_state=42,\n        )\n\n    def _t(self, lin=False):\n        return Pipeline([(\"s\", StandardScaler()), (\"m\", LogisticRegression(C=1.0, max_iter=1000))]) if lin else GradientBoostingClassifier(\n            n_estimators=120, max_depth=3, learning_rate=0.05,\n            min_samples_leaf=20, subsample=0.8, random_state=42,\n        )\n\n    def _fitb(self, X, y, d=3):\n        return [self._y(d).fit(X, y), self._y(lin=True).fit(X, y)]\n\n    def _predb(self, models, X):\n        return 0.8 * models[0].predict(X) + 0.2 * models[1].predict(X)\n\n    def fit(self, X, T, Y):\n        X, T, Y = np.asarray(X), np.asarray(T).astype(int), np.asarray(Y)\n        n = len(Y)\n        kf = KFold(n_splits=min(self.k, max(2, n // 200)), shuffle=True, random_state=42)\n        mu0, mu1, m, e = np.zeros(n), np.zeros(n), np.zeros(n), np.zeros(n)\n\n        for tr, va in kf.split(X):\n            xtr, xva, ttr, ytr = X[tr], X[va], T[tr], Y[tr]\n            m0 = self._fitb(xtr[ttr == 0], ytr[ttr == 0])\n            m1 = self._fitb(xtr[ttr == 1], ytr[ttr == 1])\n            mu0[va], mu1[va] = self._predb(m0, xva), self._predb(m1, xva)\n            mall = self._fitb(xtr, ytr, 2)\n            m[va] = self._predb(mall, xva)\n            tg, tl = self._t().fit(xtr, ttr), self._t(True).fit(xtr, ttr)\n            e[va] = 0.8 * tg.predict_proba(xva)[:, 1] + 0.2 * tl.predict_proba(xva)[:, 1]\n\n        e = np.clip(e, self.clip, 1 - self.clip)\n        r = T - e\n        w = r ** 2 + 1e-3\n        safe = np.where(np.abs(r) > 0.05, r, 0.05 * np.sign(r + 1e-8))\n        pdr = mu1 - mu0 + T * (Y - mu1) / e - (1 - T) * (Y - mu0) / (1 - e)\n        pr = (Y - m) / safe\n        pdr = np.clip(pdr, -np.percentile(np.abs(pdr), 98), np.percentile(np.abs(pdr), 98))\n        pr = np.clip(pr, -np.percentile(np.abs(pr), 98), np.percentile(np.abs(pr), 98))\n\n        gate = np.clip(4 * e * (1 - e), 0.0, 1.0)\n        target, sw = gate * pr + (1 - gate) * pdr, gate + w\n        P = np.zeros((n, 3))\n        P[:, 0] = mu1 - mu0\n        for tr, va in kf.split(X):\n            P[va, 1] = self._y(2).fit(X[tr], pdr[tr], sample_weight=sw[tr]).predict(X[va])\n            P[va, 2] = self._y(3).fit(X[tr], pr[tr], sample_weight=w[tr]).predict(X[va])\n\n        stack = Ridge(alpha=0.5, fit_intercept=False)\n        stack.fit(P, target, sample_weight=sw)\n        a = np.maximum(stack.coef_, 0)\n        self.a_ = a / a.sum() if a.sum() > 1e-8 else np.array([1 / 3, 1 / 3, 1 / 3])\n\n        self.mu0_, self.mu1_ = self._fitb(X[T == 0], Y[T == 0]), self._fitb(X[T == 1], Y[T == 1])\n        self.dr_ = self._y(2).fit(X, pdr, sample_weight=sw)\n        self.rr_ = self._y(3).fit(X, pr, sample_weight=w)\n        fullP = np.column_stack([\n            self._predb(self.mu1_, X) - self._predb(self.mu0_, X),\n            self.dr_.predict(X), self.rr_.predict(X),\n        ])\n        self.off_ = pdr.mean() - fullP.dot(self.a_).mean()\n        return self\n\n    def predict(self, X):\n        P = np.column_stack([\n            self._predb(self.mu1_, X) - self._predb(self.mu0_, X),\n            self.dr_.predict(X), self.rr_.predict(X),\n        ])\n        return P.dot(self.a_) + self.off_\n",
5  "after_line": 418,
6  "start_line": 345,
7  "end_line": 418
8}

1{
2  "op": "create",
3  "filename": "scikit-learn/scikit-learn/custom_cate.py",
4  "content": "import os\nimport runpy\n\nif __name__ == \"__main__\":\n    here = os.path.dirname(__file__)\n    target = os.path.abspath(os.path.join(here, \"..\", \"custom_cate.py\"))\n    runpy.run_path(target, run_name=\"__main__\")\n\n",
5  "after_line": 1,
6  "start_line": 1,
7  "end_line": 1
8}

1{
2  "op": "replace",
3  "filename": "scikit-learn/custom_cate.py",
4  "content": "class CATEEstimator(BaseCATEEstimator):\n    \"\"\"Overlap-adaptive orthogonal stacking with disagreement shrinkage.\"\"\"\n\n    def __init__(self):\n        self.k, self.clip = 3, 0.03\n\n    def _y(self, d=3, lin=False):\n        return Pipeline([(\"s\", StandardScaler()), (\"m\", Ridge(alpha=1.0))]) if lin else GradientBoostingRegressor(\n            n_estimators=100, max_depth=d, learning_rate=0.05,\n            min_samples_leaf=20, subsample=0.8, random_state=42)\n\n    def _t(self, lin=False):\n        return Pipeline([(\"s\", StandardScaler()), (\"m\", LogisticRegression(C=1.0, max_iter=1000))]) if lin else GradientBoostingClassifier(\n            n_estimators=100, max_depth=3, learning_rate=0.05,\n            min_samples_leaf=20, subsample=0.8, random_state=42)\n\n    def _fitb(self, X, y, d=3): return [self._y(d).fit(X, y), self._y(lin=True).fit(X, y)]\n    def _predb(self, m, X): return 0.8 * m[0].predict(X) + 0.2 * m[1].predict(X)\n    def _eh(self, mg, ml, X): return 0.8 * mg.predict_proba(X)[:, 1] + 0.2 * ml.predict_proba(X)[:, 1]\n\n    def fit(self, X, T, Y):\n        X, T, Y = np.asarray(X), np.asarray(T).astype(int), np.asarray(Y)\n        n = len(Y)\n        kf = KFold(n_splits=self.k, shuffle=True, random_state=42)\n        mu0, mu1, m, e = np.zeros(n), np.zeros(n), np.zeros(n), np.zeros(n)\n\n        for tr, va in kf.split(X):\n            xtr, xva, ttr, ytr = X[tr], X[va], T[tr], Y[tr]\n            mu0[va] = self._predb(self._fitb(xtr[ttr == 0], ytr[ttr == 0]), xva)\n            mu1[va] = self._predb(self._fitb(xtr[ttr == 1], ytr[ttr == 1]), xva)\n            m[va] = self._predb(self._fitb(xtr, ytr, 2), xva)\n            mg, ml = self._t().fit(xtr, ttr), self._t(True).fit(xtr, ttr)\n            e[va] = self._eh(mg, ml, xva)\n\n        e = np.clip(e, self.clip, 1 - self.clip)\n        r = T - e\n        safe = np.sign(r + 1e-8) * np.maximum(np.abs(r), 0.05)\n        gate, w = np.clip(4 * e * (1 - e), 0.0, 1.0), r ** 2 + 1e-3\n        tau_t = mu1 - mu0\n        pdr = tau_t + T * (Y - mu1) / e - (1 - T) * (Y - mu0) / (1 - e)\n        pr = (Y - m) / safe\n        pto = r * (Y - m) / (e * (1 - e))\n        for z in (pdr, pr, pto):\n            q = np.percentile(np.abs(z), 98)\n            z[:] = np.clip(z, -q, q)\n\n        target, sw = (1 - gate) * pdr + gate * (0.7 * pto + 0.3 * pr), 0.2 + gate + w\n        P = np.zeros((n, 4)); P[:, 0] = tau_t\n        for tr, va in kf.split(X):\n            P[va, 1] = self._y(2).fit(X[tr], pdr[tr], sample_weight=sw[tr]).predict(X[va])\n            P[va, 2] = self._y(2).fit(X[tr], pto[tr], sample_weight=0.1 + gate[tr]).predict(X[va])\n            P[va, 3] = self._y(3).fit(X[tr], pr[tr], sample_weight=w[tr]).predict(X[va])\n\n        st = Ridge(alpha=0.5, fit_intercept=False).fit(P, target, sample_weight=sw)\n        a = np.maximum(st.coef_, 0); self.a_ = a / a.sum() if a.sum() > 1e-8 else np.ones(4) / 4\n        self.mu0_, self.mu1_ = self._fitb(X[T == 0], Y[T == 0]), self._fitb(X[T == 1], Y[T == 1])\n        self.m_, self.tg_, self.tl_ = self._fitb(X, Y, 2), self._t().fit(X, T), self._t(True).fit(X, T)\n        self.dr_ = self._y(2).fit(X, pdr, sample_weight=sw)\n        self.to_ = self._y(2).fit(X, pto, sample_weight=0.1 + gate)\n        self.rr_ = self._y(3).fit(X, pr, sample_weight=w)\n        P = np.column_stack([self._predb(self.mu1_, X) - self._predb(self.mu0_, X), self.dr_.predict(X), self.to_.predict(X), self.rr_.predict(X)])\n        raw, base = P.dot(self.a_), P[:, 0]\n        eh = np.clip(self._eh(self.tg_, self.tl_, X), self.clip, 1 - self.clip)\n        disp = ((P[:, 1:] - P[:, 1:].mean(1, keepdims=True)) ** 2).mean(1)\n        self.ds_, self.ate_ = np.median(disp) + 1e-6, pdr.mean()\n        shrink = (0.25 + 0.75 * np.clip(4 * eh * (1 - eh), 0, 1)) / (1 + disp / self.ds_)\n        self.off_ = self.ate_ - (base + shrink * (raw - base)).mean()\n        return self\n\n    def predict(self, X):\n        eh = np.clip(self._eh(self.tg_, self.tl_, X), self.clip, 1 - self.clip)\n        P = np.column_stack([self._predb(self.mu1_, X) - self._predb(self.mu0_, X), self.dr_.predict(X), self.to_.predict(X), self.rr_.predict(X)])\n        raw, base = P.dot(self.a_), P[:, 0]\n        disp = ((P[:, 1:] - P[:, 1:].mean(1, keepdims=True)) ** 2).mean(1)\n        shrink = (0.25 + 0.75 * np.clip(4 * eh * (1 - eh), 0, 1)) / (1 + disp / self.ds_)\n        return base + shrink * (raw - base) + self.off_\n",
5  "after_line": 423,
6  "start_line": 345,
7  "end_line": 423
8}