1{
2  "op": "replace",
3  "filename": "scikit-learn/custom_subgroup_calibration.py",
4  "start_line": 72,
5  "end_line": 91,
6  "content": "class CalibrationMethod:\n    \"\"\"Shrinkage-regularized beta calibration for subgroup distribution shift.\n\n    Fits a global beta calibration model as anchor, then estimates group-specific\n    intercept adjustments. Adjustments are shrunk toward zero based on group size\n    and calibration reliability, preventing overfitting to the calibration\n    distribution when subgroup composition shifts at test time.\n    \"\"\"\n\n    def __init__(self):\n        self.eps = 1e-6\n        self._identity = True\n        self.shrinkage_lambda = 2.0  # strength of shrinkage toward global\n\n    def _featurize(self, probs):\n        \"\"\"Beta calibration features: [log(p), log(1-p)].\"\"\"\n        probs = np.asarray(probs).reshape(-1)\n        p = np.clip(probs, self.eps, 1.0 - self.eps)\n        return np.column_stack([np.log(p), np.log1p(-p)])\n\n    def fit(self, probs, labels, groups=None):\n        probs = np.asarray(probs).reshape(-1)\n        labels = np.asarray(labels).reshape(-1).astype(int)\n\n        # Step 1: Fit global beta calibration as anchor\n        X_global = self._featurize(probs)\n        self.global_model_ = LogisticRegression(\n            max_iter=2000, solver=\"lbfgs\", C=1.0\n        )\n        self.global_model_.fit(X_global, labels)\n\n        # Step 2: Compute global calibrated probabilities and residuals\n        global_cal = np.clip(\n            self.global_model_.predict_proba(X_global)[:, 1],\n            self.eps, 1.0 - self.eps,\n        )\n        residuals = labels.astype(float) - global_cal\n\n        # Step 3: Estimate per-group intercept adjustments with shrinkage\n        self.group_adjustments_ = {}\n        self.calib_performance_ = {}\n\n        if groups is not None:\n            groups = np.asarray(groups).reshape(-1)\n            for g in np.unique(groups):\n                mask = groups == g\n                n_g = mask.sum()\n                if n_g < 5:\n                    self.group_adjustments_[int(g)] = 0.0\n                    self.calib_performance_[int(g)] = 1.0\n                    continue\n\n                g_residual = residuals[mask]\n                g_cal = global_cal[mask]\n\n                # Measure calibration reliability for this group\n                # Lower is better: average residual magnitude\n                cal_error = float(np.mean(np.abs(g_residual)))\n                # Convert to reliability score [0, 1]\n                reliability = float(np.clip(1.0 - 2.0 * cal_error, 0.0, 1.0))\n\n                # Compute optimal intercept in logit space for this group\n                logit_cal = special.logit(g_cal)\n                logit_target = special.logit(np.clip(\n                    labels[mask].astype(float), self.eps, 1.0 - self.eps\n                ))\n                # Simple average adjustment\n                raw_adjustment = float(np.nanmean(logit_target - logit_cal))\n                raw_adjustment = np.clip(raw_adjustment, -5.0, 5.0)\n\n                # Shrinkage weight: more data + better calibration = less shrinkage\n                n_weight = n_g / (n_g + self.shrinkage_lambda)\n                shrink_weight = reliability * n_weight\n\n                # Apply shrinkage toward global (zero adjustment)\n                self.group_adjustments_[int(g)] = shrink_weight * raw_adjustment\n                self.calib_performance_[int(g)] = reliability\n\n        self._base_rate = float(np.clip(labels.mean(), self.eps, 1.0 - self.eps))\n        self._identity = False\n        return self\n\n    def predict_proba(self, probs, groups=None):\n        probs = np.asarray(probs).reshape(-1)\n        X = self._featurize(probs)\n\n        # Global calibrated probabilities\n        global_cal = np.clip(\n            self.global_model_.predict_proba(X)[:, 1],\n            self.eps, 1.0 - self.eps,\n        )\n\n        if groups is None or not self.group_adjustments_:\n            return global_cal\n\n        groups = np.asarray(groups).reshape(-1)\n        result = global_cal.copy()\n\n        for g in np.unique(groups):\n            mask = groups == g\n            if not mask.any():\n                continue\n\n            adj = self.group_adjustments_.get(int(g), 0.0)\n            if adj == 0.0:\n                continue\n\n            # Apply additive adjustment in logit space\n            logits = special.logit(result[mask])\n            result[mask] = special.expit(logits + adj)\n\n        return np.clip(result, self.eps, 1.0 - self.eps)"
7}

1{
2  "filename": "scikit-learn/custom_subgroup_calibration.py",
3  "op": "replace",
4  "content": "class CalibrationMethod:\n    \"\"\"Adaptive shrinkage calibration: global beta + validated group adjustments.\n\n    Key ideas for robustness under subgroup distribution shift:\n    1. Fit a robust global beta calibration as the anchor.\n    2. Only apply group-specific adjustments if they demonstrably improve\n       calibration on that group's calibration data.\n    3. Use strong shrinkage toward the global model, proportional to\n       group size and calibration improvement.\n    4. Validate adjustments using a leave-one-out style check to avoid\n       overfitting to calibration data.\n    \"\"\"\n\n    def __init__(self):\n        self.eps = 1e-6\n        self._identity = True\n\n    def _featurize(self, probs):\n        \"\"\"Beta calibration features: [log(p), log(1-p)].\"\"\"\n        probs = np.asarray(probs).reshape(-1)\n        p = np.clip(probs, self.eps, 1.0 - self.eps)\n        return np.column_stack([np.log(p), np.log1p(-p)])\n\n    @staticmethod\n    def _fit_global_temperature(probs, labels, eps=1e-6):\n        \"\"\"Fit a single temperature via NLL minimization for robust base.\"\"\"\n        logits = special.logit(np.clip(probs, eps, 1.0 - eps))\n\n        def objective(log_t):\n            t = float(np.exp(log_t))\n            cal = special.expit(logits / t)\n            p = np.clip(cal, eps, 1.0 - eps)\n            return float(-np.mean(labels * np.log(p) + (1 - labels) * np.log(1 - p)))\n\n        result = optimize.minimize_scalar(objective, bounds=(-3.0, 3.0), method=\"bounded\")\n        return float(np.exp(result.x)) if result.success else 1.0\n\n    def fit(self, probs, labels, groups=None):\n        probs = np.asarray(probs).reshape(-1)\n        labels = np.asarray(labels).reshape(-1).astype(int)\n\n        # Step 1: Fit global temperature as robust anchor\n        self.global_temp_ = self._fit_global_temperature(probs, labels)\n\n        # Step 2: Compute global calibrated probabilities\n        logits_all = special.logit(np.clip(probs, self.eps, 1.0 - self.eps))\n        self.global_cal_all_ = special.expit(logits_all / self.global_temp_)\n        self.global_cal_all_ = np.clip(self.global_cal_all_, self.eps, 1.0 - self.eps)\n\n        # Global Brier score on calibration data\n        self.global_brier_ = float(np.mean((labels.astype(float) - self.global_cal_all_) ** 2))\n\n        # Step 3: Per-group adjustments with validation\n        self.group_adjustments_ = {}\n\n        if groups is not None:\n            groups = np.asarray(groups).reshape(-1)\n            for g in np.unique(groups):\n                mask = groups == g\n                n_g = mask.sum()\n                if n_g < 10:\n                    self.group_adjustments_[int(g)] = 0.0\n                    continue\n\n                g_labels = labels[mask]\n                g_cal = self.global_cal_all_[mask]\n\n                # Skip groups with degenerate labels\n                if np.unique(g_labels).size < 2:\n                    self.group_adjustments_[int(g)] = 0.0\n                    continue\n\n                # Check current calibration quality for this group\n                current_nll_ceil = 0.99\n                p_clipped = np.clip(g_cal, self.eps, 1.0 - self.eps)\n                current_nll = float(-np.mean(\n                    g_labels * np.log(p_clipped) + (1 - g_labels) * np.log(1 - p_clipped)\n                ))\n                current_nll = min(current_nll, current_nll_ceil)\n\n                # Fit group-specific temperature\n                g_logits = special.logit(np.clip(probs[mask], self.eps, 1.0 - self.eps))\n\n                def group_objective(log_t):\n                    t = float(np.exp(log_t))\n                    cal = special.expit(g_logits / t)\n                    p = np.clip(cal, self.eps, 1.0 - self.eps)\n                    return float(-np.mean(g_labels * np.log(p) + (1 - g_labels) * np.log(1 - p)))\n\n                result = optimize.minimize_scalar(group_objective, bounds=(-3.0, 3.0), method=\"bounded\")\n                group_temp = float(np.exp(result.x)) if result.success else self.global_temp_\n\n                # Check if group-specific temp genuinely improves NLL\n                # Use a conservative penalty: only keep improvement if delta is meaningful\n                group_cal = special.expit(g_logits / group_temp)\n                group_cal = np.clip(group_cal, self.eps, 1.0 - self.eps)\n                group_nll = float(-np.mean(\n                    g_labels * np.log(group_cal) + (1 - g_labels) * np.log(1 - group_cal)\n                ))\n\n                nll_gain = current_nll - group_nll\n\n                # Only use group-specific calibration if it meaningfully improves\n                # and the group has enough data to support a separate parameter\n                if nll_gain > 0.01:  # minimum improvement threshold\n                    # Compute shrinkage factor based on group size\n                    # n_g=10 -> 0.5, n_g=50 -> 0.83, n_g=100 -> 0.91\n                    shrinkage = float(n_g / (n_g + 10.0))\n\n                    # Store: we'll blend between global and group-specific\n                    self.group_adjustments_[int(g)] = {\n                        'group_temp': group_temp,\n                        'global_temp': self.global_temp_,\n                        'shrinkage': shrinkage,\n                    }\n                else:\n                    self.group_adjustments_[int(g)] = {\n                        'group_temp': self.global_temp_,\n                        'global_temp': self.global_temp_,\n                        'shrinkage': 0.0,\n                    }\n\n        self._identity = False\n        return self\n\n    def predict_proba(self, probs, groups=None):\n        probs = np.asarray(probs).reshape(-1)\n        logits = special.logit(np.clip(probs, self.eps, 1.0 - self.eps))\n\n        # Global calibrated probabilities\n        result = special.expit(logits / self.global_temp_)\n        result = np.clip(result, self.eps, 1.0 - self.eps)\n\n        if groups is None or not self.group_adjustments_:\n            return result\n\n        groups = np.asarray(groups).reshape(-1)\n\n        for g in np.unique(groups):\n            mask = groups == g\n            if not mask.any():\n                continue\n\n            adj = self.group_adjustments_.get(int(g))\n            if adj is None or adj['shrinkage'] == 0.0:\n                continue\n\n            # Blend: (1-shrink) * global_temp + shrink * group_temp\n            effective_temp = (\n                (1.0 - adj['shrinkage']) * adj['global_temp']\n                + adj['shrinkage'] * adj['group_temp']\n            )\n            result[mask] = special.expit(logits[mask] / effective_temp)\n\n        return np.clip(result, self.eps, 1.0 - self.eps)",
5  "start_line": 72,
6  "end_line": 183
7}

1{
2  "content": "                if n_g < 10:\n                    self.group_adjustments_[int(g)] = {\n                        'group_temp': self.global_temp_,\n                        'global_temp': self.global_temp_,\n                        'shrinkage': 0.0,\n                    }\n                    continue\n\n                g_labels = labels[mask]\n                g_cal = self.global_cal_all_[mask]\n\n                # Skip groups with degenerate labels\n                if np.unique(g_labels).size < 2:\n                    self.group_adjustments_[int(g)] = {\n                        'group_temp': self.global_temp_,\n                        'global_temp': self.global_temp_,\n                        'shrinkage': 0.0,\n                    }\n                    continue",
3  "filename": "scikit-learn/custom_subgroup_calibration.py",
4  "op": "replace",
5  "start_line": 130,
6  "end_line": 148
7}

1{
2  "filename": "scikit-learn/custom_subgroup_calibration.py",
3  "op": "replace",
4  "content": "class CalibrationMethod:\n    \"\"\"Shrinkage-regularized group temperature scaling for distribution shift.\n\n    Strategy:\n    1. Fit a global temperature as robust baseline.\n    2. For each subgroup, fit a group-specific temperature.\n    3. Only apply group adjustment if it provides meaningful NLL improvement.\n    4. Shrink group temps toward global temp based on group size.\n    \"\"\"\n\n    def __init__(self):\n        self.eps = 1e-6\n\n    @staticmethod\n    def _fit_temperature(probs, labels, eps=1e-6):\n        \"\"\"Fit temperature via NLL minimization.\"\"\"\n        logits = special.logit(np.clip(probs, eps, 1.0 - eps))\n\n        def objective(log_t):\n            t = float(np.exp(log_t))\n            cal = special.expit(logits / t)\n            p = np.clip(cal, eps, 1.0 - eps)\n            return float(-np.mean(\n                labels * np.log(p) + (1 - labels) * np.log(1 - p)\n            ))\n\n        result = optimize.minimize_scalar(\n            objective, bounds=(-3.0, 3.0), method=\"bounded\"\n        )\n        return float(np.exp(result.x)) if result.success else 1.0\n\n    def fit(self, probs, labels, groups=None):\n        probs = np.asarray(probs).reshape(-1)\n        labels = np.asarray(labels).reshape(-1).astype(int)\n\n        # Step 1: Fit global temperature as robust anchor\n        self.global_temp_ = self._fit_temperature(probs, labels)\n\n        # Step 2: Per-group temperature adjustments with shrinkage\n        self.group_temps_ = {}\n\n        if groups is not None:\n            groups = np.asarray(groups).reshape(-1)\n            g_clipped = np.clip(probs, self.eps, 1.0 - self.eps)\n            g_logits = special.logit(g_clipped)\n\n            # Compute global NLL as baseline\n            global_cal = special.expit(g_logits / self.global_temp_)\n            global_cal_c = np.clip(global_cal, self.eps, 1.0 - self.eps)\n            base_nll = float(-np.mean(\n                labels * np.log(global_cal_c) + (1 - labels) * np.log(1 - global_cal_c)\n            ))\n\n            for g in np.unique(groups):\n                mask = groups == g\n                n_g = int(mask.sum())\n                g_labels = labels[mask]\n                g_logits_grp = g_logits[mask]\n\n                if n_g < 10 or np.unique(g_labels).size < 2:\n                    self.group_temps_[int(g)] = self.global_temp_\n                    continue\n\n                # Compute NLL contribution of this group under global temp\n                g_cal_global = special.expit(g_logits_grp / self.global_temp_)\n                g_cal_global_c = np.clip(g_cal_global, self.eps, 1.0 - self.eps)\n                group_global_nll = float(-np.mean(\n                    g_labels * np.log(g_cal_global_c) + (1 - g_labels) * np.log(1 - g_cal_global_c)\n                ))\n\n                # Fit group-specific temperature\n                group_temp = self._fit_temperature(probs[mask], labels[mask])\n\n                # Evaluate group-specific temp on the same group\n                g_cal_group = special.expit(g_logits_grp / group_temp)\n                g_cal_group_c = np.clip(g_cal_group, self.eps, 1.0 - self.eps)\n                group_spec_nll = float(-np.mean(\n                    g_labels * np.log(g_cal_group_c) + (1 - g_labels) * np.log(1 - g_cal_group_c)\n                ))\n\n                nll_improvement = group_global_nll - group_spec_nll\n\n                if nll_improvement > 0.02:\n                    # Shrink toward global: n_g/(n_g+10)\n                    # More data => trust group temp more\n                    alpha = n_g / (n_g + 10.0)\n                    effective_temp = (1.0 - alpha) * self.global_temp_ + alpha * group_temp\n                    self.group_temps_[int(g)] = effective_temp\n                else:\n                    self.group_temps_[int(g)] = self.global_temp_\n\n        return self\n\n    def predict_proba(self, probs, groups=None):\n        probs = np.asarray(probs).reshape(-1)\n        logits = special.logit(np.clip(probs, self.eps, 1.0 - self.eps))\n\n        if groups is None or not self.group_temps_:\n            return np.clip(\n                special.expit(logits / self.global_temp_),\n                self.eps, 1.0 - self.eps,\n            )\n\n        groups = np.asarray(groups).reshape(-1)\n        result = np.empty_like(probs)\n\n        for g in np.unique(groups):\n            mask = groups == g\n            if not mask.any():\n                continue\n            temp = self.group_temps_.get(int(g), self.global_temp_)\n            result[mask] = special.expit(logits[mask] / temp)\n\n        return np.clip(result, self.eps, 1.0 - self.eps)",
5  "end_line": 226,
6  "start_line": 72
7}

1{
2  "n": 1
3}