{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[],"dockerImageVersionId":31192,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# !pip install jax jaxlib numpy scikit-learn tqdm xgboost matplotlib\n\nimport numpy as np\nimport jax\nimport jax.numpy as jnp\nfrom sklearn.metrics import f1_score\nfrom tqdm.auto import tqdm\nimport warnings\nimport xgboost as xgb\nfrom sklearn.metrics import f1_score\nfrom sklearn.model_selection import train_test_split\nimport time\nimport matplotlib.pyplot as plt\n\n# --- JAX Configuration and Kernel/Loss Functions (Your Original Code) ---\nwarnings.filterwarnings(\"ignore\")\njax.config.update(\"jax_enable_x64\", True)\n\ndef _fit_tree_kernel(X, residuals, max_depth, num_thresholds):\n    \"\"\"JAX-compiled function to fit a single decision tree.\"\"\"\n    num_data_points = X.shape[0]\n    num_nodes_at_max_depth = 2 ** max_depth\n    all_feature_indices = jnp.zeros((max_depth, num_nodes_at_max_depth), dtype=jnp.int32)\n    all_thresholds = jnp.zeros((max_depth, num_nodes_at_max_depth), dtype=jnp.float64)\n    all_left_outputs = jnp.zeros((max_depth, num_nodes_at_max_depth), dtype=jnp.float64)\n    all_right_outputs = jnp.zeros((max_depth, num_nodes_at_max_depth), dtype=jnp.float64)\n    node_mask = jnp.zeros(num_data_points, dtype=jnp.int32)\n\n    def loop_body(depth, carry):\n        node_mask, all_feature_indices, all_thresholds, all_left_outputs, all_right_outputs = carry\n        num_nodes_at_this_depth = 2 ** depth\n\n        def node_loop_body(node_index, node_carry):\n            (afi, ath, alo, aro) = node_carry\n            def return_default(): return (jnp.array(-1, dtype=jnp.int32), 0.0, 0.0, 0.0)\n            data_mask = node_mask == node_index\n            is_valid_node = jnp.sum(data_mask) > 0\n            \n            def calculate_split():\n                def find_best_for_feature(feature_column):\n                    node_vals = jnp.where(data_mask, feature_column, jnp.nan)\n                    thresholds = jnp.nanquantile(node_vals, q=jnp.linspace(0.0, 1.0, num_thresholds))\n                    def inner_loop(state, threshold):\n                        best_score, best_thr, best_left, best_right = state\n                        left_mask = jnp.logical_and(data_mask, feature_column < threshold)\n                        right_mask = jnp.logical_and(data_mask, ~left_mask)\n                        \n                        count_left = jnp.sum(left_mask)\n                        count_right = jnp.sum(right_mask)\n                        sum_left = jnp.sum(residuals * left_mask)\n                        sum_right = jnp.sum(residuals * right_mask)\n                        \n                        gain = jax.lax.cond(\n                            jnp.logical_and(count_left > 0, count_right > 0),\n                            lambda: ((sum_left**2 / count_left) + (sum_right**2 / count_right)),\n                            lambda: -jnp.inf  # Penalize invalid splits\n                        )\n\n                        is_better = gain > best_score\n                        new_state = (\n                            jax.lax.cond(is_better, lambda: gain, lambda: best_score),\n                            jax.lax.cond(is_better, lambda: threshold, lambda: best_thr),\n                            jax.lax.cond(is_better, lambda: sum_left / count_left, lambda: best_left),\n                            jax.lax.cond(is_better, lambda: sum_right / count_right, lambda: best_right),\n                        )\n                        return new_state, None\n\n                    initial_state = (-jnp.inf, 0.0, 0.0, 0.0)\n                    final_state, _ = jax.lax.scan(inner_loop, initial_state, thresholds)\n                    return final_state\n\n                best_scores, best_thresholds, best_lefts, best_rights = jax.vmap(find_best_for_feature, in_axes=1)(X)\n                best_feature_idx = jnp.argmax(best_scores).astype(jnp.int32)\n                return (best_feature_idx, best_thresholds[best_feature_idx], best_lefts[best_feature_idx], best_rights[best_feature_idx])\n            \n            feature_idx, threshold, left_out, right_out = jax.lax.cond(is_valid_node, calculate_split, return_default)\n            afi = afi.at[depth, node_index].set(feature_idx)\n            ath = ath.at[depth, node_index].set(threshold)\n            alo = alo.at[depth, node_index].set(left_out)\n            aro = aro.at[depth, node_index].set(right_out)\n            return (afi, ath, alo, aro)\n        \n        carry_state = (all_feature_indices, all_thresholds, all_left_outputs, all_right_outputs)\n        new_carry_state = jax.lax.fori_loop(0, num_nodes_at_this_depth, node_loop_body, carry_state)\n        (all_feature_indices, all_thresholds, all_left_outputs, all_right_outputs) = new_carry_state\n        \n        def update_mask_body(i, current_mask):\n            is_in_node = current_mask == i\n            split_feature = all_feature_indices[depth, i]\n            split_threshold = all_thresholds[depth, i]\n            valid_split = split_feature >= 0\n            def apply_split():\n                split_condition_mask = X[:, split_feature] < split_threshold\n                goes_left = jnp.logical_and(is_in_node, split_condition_mask)\n                goes_right = jnp.logical_and(is_in_node, ~split_condition_mask)\n                new_mask = jnp.where(goes_left, 2 * i, current_mask)\n                new_mask = jnp.where(goes_right, 2 * i + 1, new_mask)\n                return new_mask\n            return jax.lax.cond(valid_split, apply_split, lambda: current_mask)\n\n        node_mask = jax.lax.fori_loop(0, num_nodes_at_this_depth, update_mask_body, node_mask)\n        return (node_mask, all_feature_indices, all_thresholds, all_left_outputs, all_right_outputs)\n\n    initial_carry = (node_mask, all_feature_indices, all_thresholds, all_left_outputs, all_right_outputs)\n    final_carry = jax.lax.fori_loop(0, max_depth, loop_body, initial_carry)\n    return final_carry[1:]\n\ndef _predict_leaf_index_kernel(X, tree_params, max_depth):\n    \"\"\"JAX-compiled function to find the leaf index for each data point.\"\"\"\n    feature_indices, thresholds, _, _ = tree_params\n    node_indices = jnp.zeros(X.shape[0], dtype=jnp.int32)\n    def body(carry, depth):\n        node_indices = carry\n        f_idx = feature_indices[depth, node_indices]\n        thr = thresholds[depth, node_indices]\n        f_vals = jnp.take_along_axis(X, f_idx[:, None], axis=1).squeeze(axis=1)\n        split_mask = f_vals < thr\n        return jnp.where(split_mask, 2 * node_indices, 2 * node_indices + 1), None\n    final_nodes, _ = jax.lax.scan(body, node_indices, jnp.arange(max_depth))\n    return final_nodes\n\n_fit_tree_kernel_jit = jax.jit(_fit_tree_kernel, static_argnums=(2, 3))\n_predict_leaf_index_kernel_jit = jax.jit(_predict_leaf_index_kernel, static_argnums=2)\n\ndef weighted_bce_loss(y_true, y_pred_logits, weights):\n    return jnp.mean(weights * jnp.log(1 + jnp.exp(-y_true * y_pred_logits)))\n\ndef mse_loss(y_true, y_pred):\n    return jnp.mean((y_true - y_pred)**2)\n\ndef combined_loss(y_true, y_pred_class_logits, y_true_reco, y_pred_reco, weights, lambda_ratio):\n    supervised_loss = lambda_ratio * weighted_bce_loss(y_true, y_pred_class_logits, weights)\n    reconstruction_loss = (1 - lambda_ratio) * mse_loss(y_true_reco, y_pred_reco)\n    return supervised_loss + reconstruction_loss\n\nclass JAXGBMDualHead:\n    def __init__(self, n_estimators=100, learning_rate=0.1, max_depth=3, batch_size=20000, num_thresholds=70, lambda_ratio=0.5):\n        self.n_estimators = n_estimators\n        self.learning_rate = learning_rate\n        self.max_depth = max_depth\n        self.num_thresholds = num_thresholds\n        self.batch_size = batch_size\n        self.lambda_ratio = lambda_ratio\n        self.trees = []\n        self.reco_head_params = None\n\n    def fit(self, X_input, X_reco, y_class, sample_weights):\n        X_jnp = jnp.asarray(X_input, dtype=jnp.float64)\n        X_reco_jnp = jnp.asarray(X_reco, dtype=jnp.float64)\n        y_class_jnp = jnp.asarray(y_class, dtype=jnp.float64).ravel()\n        sample_weights_jnp = jnp.asarray(sample_weights, dtype=jnp.float64).ravel()\n        num_leaves = 2 ** self.max_depth\n        num_samples = X_jnp.shape[0]\n        if self.reco_head_params is None:\n            num_features = X_jnp.shape[1]\n            self.reco_head_params = jnp.zeros((num_leaves, num_features))\n        current_predictions = jnp.zeros_like(y_class_jnp)\n\n        @jax.jit\n        def _update_reco_head(reco_params, leaf_idx, X_reco_jnp, learning_rate):\n            reco_loss_grad_fn = jax.grad(lambda rp: mse_loss(X_reco_jnp, jnp.take(rp, leaf_idx, axis=0)))\n            reco_head_grad = reco_loss_grad_fn(reco_params)\n            return reco_params - learning_rate * reco_head_grad\n\n        for i in tqdm(range(self.n_estimators), desc=\"Training Dual-Head GBM\"):\n            residuals = -jax.grad(lambda preds: weighted_bce_loss(y_class_jnp, preds, sample_weights_jnp))(current_predictions)\n            full_tree_params = None\n            all_leaf_sums = jnp.zeros(num_leaves)\n            all_leaf_counts = jnp.zeros(num_leaves)\n            for j in range(0, num_samples, self.batch_size):\n                X_batch = X_jnp[j:j+self.batch_size]\n                r_batch = residuals[j:j+self.batch_size]\n                if full_tree_params is None:\n                    tree_params_batch = _fit_tree_kernel_jit(X_batch, r_batch, self.max_depth, self.num_thresholds)\n                    full_tree_params = tree_params_batch\n                leaf_idx = _predict_leaf_index_kernel_jit(X_batch, full_tree_params, self.max_depth)\n                all_leaf_sums = all_leaf_sums.at[leaf_idx].add(r_batch)\n                all_leaf_counts = all_leaf_counts.at[leaf_idx].add(1.0)\n            gamma_per_leaf = all_leaf_sums / (all_leaf_counts + 1e-8)\n            leaf_idx_full = _predict_leaf_index_kernel_jit(X_jnp, full_tree_params, self.max_depth)\n            gamma_for_samples = jnp.take(gamma_per_leaf, leaf_idx_full)\n            current_predictions += self.learning_rate * gamma_for_samples\n            self.trees.append((full_tree_params, np.array(gamma_per_leaf)))\n            self.reco_head_params = _update_reco_head(self.reco_head_params, leaf_idx_full, X_reco_jnp, self.learning_rate)\n            reco_preds = jnp.take(self.reco_head_params, leaf_idx_full, axis=0)\n            total_loss = combined_loss(y_class_jnp, current_predictions, X_reco_jnp, reco_preds, sample_weights_jnp, self.lambda_ratio)\n            \n            # --- UPDATED PRINT STATEMENT ---\n            if (i + 1) % 10 == 0 or i == 0:\n                print(f\"Epoch {i+1}/{self.n_estimators}, Total Combined Loss: {total_loss:.4f}\")\n            # --- END UPDATED PRINT STATEMENT ---\n\n    def predict(self, X):\n        X_jnp = jnp.asarray(X, dtype=jnp.float64)\n        predictions = jnp.zeros(X_jnp.shape[0])\n        for tree_params, gamma_per_leaf in self.trees:\n            tree_params_jnp = jax.tree_map(jnp.asarray, tree_params)\n            leaf_idx = _predict_leaf_index_kernel_jit(X_jnp, tree_params_jnp, self.max_depth)\n            gamma_jnp = jnp.asarray(gamma_per_leaf)\n            predictions += self.learning_rate * jnp.take(gamma_jnp, leaf_idx)\n        return 1 / (1 + jnp.exp(-predictions))\n        \n# ----------------------------------------------------------------------\n# MODIFIED run_comparison FUNCTION TO RETURN PLOTTING DATA\n# ----------------------------------------------------------------------\ndef run_comparison():\n    print(\"Generating synthetic data...\")\n    num_samples = 20000\n    np.random.seed(42)\n    \n    # Feature with a predictable pattern\n    time_series_feature = np.sin(np.linspace(0, 10 * np.pi, num_samples))\n    y_class = (np.diff(time_series_feature) > 0).astype(int)\n    X = np.zeros((num_samples, 5))\n    X[:, 0] = time_series_feature\n    X[:, 1:] = np.random.randn(num_samples, 4) * 0.1\n    X = X[:-1]\n    y_class = y_class.ravel()\n    \n    # Create train/test split for a more rigorous comparison\n    X_train, X_test, y_train, y_test = train_test_split(X, y_class, test_size=0.2, random_state=42)\n    \n    X_reco_train = X_train + np.random.randn(X_train.shape[0], X_train.shape[1]) * 0.05\n    sample_weights = np.ones(y_train.shape[0])\n\n    # --- JAXGBMDualHead ---\n    print(\"--- Testing JAXGBMDualHead ---\")\n    gbm_jax = JAXGBMDualHead(n_estimators=50, learning_rate=0.1, max_depth=3, batch_size=2000, num_thresholds=20, lambda_ratio=0.5)\n    start_time_jax = time.time()\n    gbm_jax.fit(X_train, X_reco_train, y_train, sample_weights)\n    end_time_jax = time.time()\n    preds_probs_jax = gbm_jax.predict(X_test)\n    preds_binary_jax = (preds_probs_jax > 0.5).astype(int)\n    f1_jax = f1_score(y_test, preds_binary_jax)\n    training_time_jax = end_time_jax - start_time_jax\n    \n    print(f\"\\n✅ JAXGBMDualHead Results:\")\n    print(f\"   F1 Score: {f1_jax:.4f}\")\n    print(f\"   Training Time: {training_time_jax:.2f} seconds\")\n    \n    # --- XGBoost Classifier ---\n    print(\"\\n--- Testing XGBoost Classifier ---\")\n    gbm_xgb = xgb.XGBClassifier(n_estimators=50, learning_rate=0.1, max_depth=3, use_label_encoder=False, eval_metric='logloss', random_state=42)\n    start_time_xgb = time.time()\n    gbm_xgb.fit(X_train, y_train)\n    end_time_xgb = time.time()\n    preds_probs_xgb = gbm_xgb.predict_proba(X_test)[:, 1]\n    preds_binary_xgb = (preds_probs_xgb > 0.5).astype(int)\n    f1_xgb = f1_score(y_test, preds_binary_xgb)\n    training_time_xgb = end_time_xgb - start_time_xgb\n\n    print(f\"\\n✅ XGBoost Results:\")\n    print(f\"   F1 Score: {f1_xgb:.4f}\")\n    print(f\"   Training Time: {training_time_xgb:.2f} seconds\")\n    \n    # Return all necessary results/data for both plotting functions\n    return {\n        'comparison': {\n            'JAXGBMDualHead': {'F1': f1_jax, 'Time': training_time_jax},\n            'XGBoost': {'F1': f1_xgb, 'Time': training_time_xgb}\n        },\n        'data': {\n            'y_test': y_test,\n            'preds_jax': np.array(preds_probs_jax) \n        }\n    }\n\n# ----------------------------------------------------------------------\n# PLOTTING FUNCTIONS\n# ----------------------------------------------------------------------\ndef plot_results(results_data):\n    \"\"\"Generates and saves a comparison plot of F1 Score and Training Time.\"\"\"\n    results = results_data['comparison']\n    models = list(results.keys())\n    f1_scores = [results[m]['F1'] for m in models]\n    train_times = [results[m]['Time'] for m in models]\n\n    fig, axes = plt.subplots(1, 2, figsize=(12, 5))\n    bar_colors = ['#4CAF50', '#007ACC']\n\n    # --- F1 Score Plot ---\n    axes[0].bar(models, f1_scores, color=bar_colors)\n    axes[0].set_title('Classification Performance (F1 Score)', fontsize=14)\n    axes[0].set_ylabel('F1 Score')\n    axes[0].grid(axis='y', linestyle=':', alpha=0.6)\n    axes[0].set_ylim(0.0, max(f1_scores) * 1.1)\n    for i, v in enumerate(f1_scores):\n        axes[0].text(i, v + 0.01, f\"{v:.4f}\", color='black', ha='center', fontweight='bold')\n\n    # --- Training Time Plot ---\n    axes[1].bar(models, train_times, color=bar_colors)\n    axes[1].set_title('Training Time Comparison (50 Estimators)', fontsize=14)\n    axes[1].set_ylabel('Time (seconds)')\n    axes[1].grid(axis='y', linestyle=':', alpha=0.6)\n    axes[1].set_ylim(0, max(train_times) * 1.1)\n    for i, v in enumerate(train_times):\n        axes[1].text(i, v + max(train_times) * 0.02, f\"{v:.2f}s\", color='black', ha='center', fontweight='bold')\n\n    fig.suptitle('JAX-based Dual-Head GBM vs. Standard XGBoost Comparison', fontsize=16)\n    plt.tight_layout(rect=[0, 0.03, 1, 0.95])\n    plot_filename = 'gbm_comparison_plot.png'\n    plt.savefig(plot_filename, dpi=300)\n    print(f\"✨ Comparison plot saved as: {plot_filename}\")\n    plt.show()\n\ndef plot_residuals(results_data):\n    \"\"\"Generates a residual and probability distribution plot for JAXGBMDualHead.\"\"\"\n    y_test = results_data['data']['y_test']\n    preds_probs_jax = results_data['data']['preds_jax']\n    \n    residuals = y_test - preds_probs_jax\n\n    fig, axes = plt.subplots(1, 2, figsize=(14, 6))\n    \n    # --- Residual Scatter Plot ---\n    axes[0].scatter(preds_probs_jax, residuals, alpha=0.3, s=5, c=y_test, cmap='viridis')\n    axes[0].axhline(0, color='r', linestyle='--', linewidth=1)\n    axes[0].set_title('Residuals ($y - \\hat{p}$) vs. Predicted Probability', fontsize=14)\n    axes[0].set_xlabel('Predicted Probability ($\\hat{p}$)', fontsize=12)\n    axes[0].set_ylabel('Residual ($y - \\hat{p}$)', fontsize=12)\n    axes[0].set_ylim(-1.1, 1.1)\n    axes[0].grid(True, linestyle=':', alpha=0.6)\n\n    # --- Probability Distribution Plot ---\n    axes[1].hist(preds_probs_jax[y_test == 0], bins=20, alpha=0.6, label='True Class 0', color='#007ACC')\n    axes[1].hist(preds_probs_jax[y_test == 1], bins=20, alpha=0.6, label='True Class 1', color='#4CAF50')\n    axes[1].set_title('Distribution of Predicted Probabilities by True Class', fontsize=14)\n    axes[1].set_xlabel('Predicted Probability ($\\hat{p}$)', fontsize=12)\n    axes[1].set_ylabel('Frequency', fontsize=12)\n    axes[1].legend(title='True Class', fontsize=10)\n    axes[1].grid(axis='y', linestyle=':', alpha=0.6)\n\n    fig.suptitle('JAXGBMDualHead Model Diagnostic Plots (Open Source Demonstration)', fontsize=16)\n    plt.tight_layout(rect=[0, 0.03, 1, 0.95])\n    \n    plot_filename = 'gbm_residuals_plot.png'\n    plt.savefig(plot_filename, dpi=300)\n    print(f\"✨ Residuals plot saved as: {plot_filename}\")\n    plt.show()\n\n# ----------------------------------------------------------------------\n# MAIN EXECUTION BLOCK\n# ----------------------------------------------------------------------\nif __name__ == '__main__':\n    # 1. Run the comparison and get the results/data\n    all_results_data = run_comparison()\n    \n    # 2. Plot performance comparison (gbm_comparison_plot.png)\n    plot_results(all_results_data)\n    \n    # 3. Plot residuals and distribution (gbm_residuals_plot.png)\n    plot_residuals(all_results_data)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-02T10:11:57.689957Z","iopub.execute_input":"2026-01-02T10:11:57.690199Z","iopub.status.idle":"2026-01-02T10:12:11.814692Z","shell.execute_reply.started":"2026-01-02T10:11:57.690175Z","shell.execute_reply":"2026-01-02T10:12:11.813539Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install xgboost\n\nimport numpy as np\nimport jax\nimport jax.numpy as jnp\nfrom sklearn.metrics import f1_score\nfrom tqdm.auto import tqdm\nimport warnings\nimport xgboost as xgb\nfrom sklearn.metrics import f1_score\nfrom sklearn.model_selection import train_test_split\nimport time\n# --- ADDED: Import matplotlib for plotting ---\nimport matplotlib.pyplot as plt\nprint(f\"Default backend: {jax.default_backend()}\")\nprint(f\"Available devices: {jax.devices()}\")\nwarnings.filterwarnings(\"ignore\")\n\n# --- Memory & Backend Optimization ---\njax.config.update(\"jax_enable_x64\", False)\ntry:\n    jax.config.update(\"jax_platforms\", \"cpu,cuda\") \nexcept:\n    pass\n\nwarnings.filterwarnings(\"ignore\")\nimport os\nprint(\"CUDA_VISIBLE_DEVICES:\", os.getenv(\"CUDA_VISIBLE_DEVICES\"))\nprint(\"JAX devices:\", jax.devices())\n\n#original # --- JAX Configuration and Loss Functions ---\n#warnings.filterwarnings(\"ignore\")\n#jax.config.update(\"jax_enable_x64\", True)\n\ndef weighted_bce_loss(y_true, y_pred_logits, weights):\n    \"\"\"Calculates weighted Binary Cross-Entropy (for y_true in {-1, 1}).\"\"\"\n    return jnp.mean(weights * jnp.log(1 + jnp.exp(-y_true * y_pred_logits)))\n\ndef mse_loss(y_true, y_pred):\n    \"\"\"Calculates Mean Squared Error.\"\"\"\n    return jnp.mean((y_true - y_pred)**2)\n\ndef combined_loss(y_true, y_pred_class_logits, y_true_reco, y_pred_reco, weights, lambda_ratio):\n    \"\"\"Combines weighted BCE for classification and MSE for reconstruction.\"\"\"\n    supervised_loss = lambda_ratio * weighted_bce_loss(y_true, y_pred_class_logits, weights)\n    reconstruction_loss = (1 - lambda_ratio) * mse_loss(y_true_reco, y_pred_reco)\n    return supervised_loss + reconstruction_loss\n\n# --- JAX Tree Kernels ---\n\n# 💥 UPDATED SIGNATURE: Accepts grad, hess, and lambda_l2\ndef _fit_tree_kernel(X, grad, hess, lambda_l2, max_depth, num_thresholds):\n    \"\"\"\n    JAX-compiled function to fit a single decision tree structure \n    using the standard XGBoost-style Gain formula.\n    \"\"\"\n    num_data_points = X.shape[0]\n    num_nodes_at_max_depth = 2 ** max_depth\n    \n    all_feature_indices = jnp.zeros((max_depth, num_nodes_at_max_depth), dtype=jnp.int32)\n    all_thresholds = jnp.zeros((max_depth, num_nodes_at_max_depth), dtype=jnp.float64)\n    node_mask = jnp.zeros(num_data_points, dtype=jnp.int32)\n\n    def loop_body(depth, carry):\n        node_mask, all_feature_indices, all_thresholds = carry\n        num_nodes_at_this_depth = 2 ** depth\n\n        def node_loop_body(node_index, node_carry):\n            (afi, ath) = node_carry\n            def return_default(): return (jnp.array(-1, dtype=jnp.int32), 0.0)\n            data_mask = node_mask == node_index\n            is_valid_node = jnp.sum(data_mask) > 0\n            \n            # --- Calculate Split Start ---\n            def calculate_split():\n                \n                # Calculate the aggregated G and H for the current node (G_node, H_node)\n                sum_g_node = jnp.sum(grad * data_mask)\n                sum_h_node = jnp.sum(hess * data_mask)\n                \n                # Node gain term (constant for all splits at this node)\n                # Term = - (Sum[G_node]^2 / (Sum[H_node] + lambda_l2))\n                gain_node_term = - (sum_g_node**2 / (sum_h_node + lambda_l2 + 1e-8))\n\n                def find_best_for_feature(feature_column):\n                    node_vals = jnp.where(data_mask, feature_column, jnp.nan)\n                    thresholds = jnp.nanquantile(node_vals, q=jnp.linspace(0.01, 0.99, num_thresholds))\n                    \n                    def inner_loop(state, threshold):\n                        best_score, best_thr = state\n                        left_mask = jnp.logical_and(data_mask, feature_column < threshold)\n                        right_mask = jnp.logical_and(data_mask, ~left_mask)\n                        \n                        # Aggregated G and H for left and right children\n                        sum_g_left = jnp.sum(grad * left_mask)\n                        sum_h_left = jnp.sum(hess * left_mask)\n                        sum_g_right = jnp.sum(grad * right_mask)\n                        sum_h_right = jnp.sum(hess * right_mask)\n                        \n                        # Term_L = (Sum[G_L]^2 / (Sum[H_L] + lambda_l2))\n                        # Term_R = (Sum[G_R]^2 / (Sum[H_R] + lambda_l2))\n                        # Gain = Term_L + Term_R + Gain_Node_Term\n                        gain = jax.lax.cond(\n                            jnp.logical_and(sum_h_left > 0, sum_h_right > 0),\n                            lambda: (\n                                (sum_g_left**2 / (sum_h_left + lambda_l2)) + \n                                (sum_g_right**2 / (sum_h_right + lambda_l2)) + \n                                gain_node_term\n                            ),\n                            lambda: -jnp.inf\n                        )\n\n                        is_better = gain > best_score\n                        new_state = (\n                            jax.lax.cond(is_better, lambda: gain, lambda: best_score),\n                            jax.lax.cond(is_better, lambda: threshold, lambda: best_thr),\n                        )\n                        return new_state, None\n\n                    initial_state = (-jnp.inf, 0.0)\n                    final_state, _ = jax.lax.scan(inner_loop, initial_state, thresholds)\n                    return final_state\n\n                best_scores, best_thresholds = jax.vmap(find_best_for_feature, in_axes=1)(X)\n                best_feature_idx = jnp.argmax(best_scores).astype(jnp.int32)\n                return (best_feature_idx, best_thresholds[best_feature_idx])\n            # --- Calculate Split End ---\n            \n            feature_idx, threshold = jax.lax.cond(is_valid_node, calculate_split, return_default)\n            afi = afi.at[depth, node_index].set(feature_idx)\n            ath = ath.at[depth, node_index].set(threshold)\n            return (afi, ath)\n        \n        carry_state = (all_feature_indices, all_thresholds)\n        new_carry_state = jax.lax.fori_loop(0, num_nodes_at_this_depth, node_loop_body, carry_state)\n        (all_feature_indices, all_thresholds) = new_carry_state\n        \n        def update_mask_body(i, current_mask):\n            is_in_node = current_mask == i\n            split_feature = all_feature_indices[depth, i]\n            split_threshold = all_thresholds[depth, i]\n            valid_split = split_feature >= 0\n            \n            def apply_split():\n                split_condition_mask = X[:, split_feature] < split_threshold\n                goes_left = jnp.logical_and(is_in_node, split_condition_mask)\n                goes_right = jnp.logical_and(is_in_node, ~split_condition_mask)\n                new_mask = jnp.where(goes_left, 2 * i, current_mask)\n                new_mask = jnp.where(goes_right, 2 * i + 1, new_mask)\n                return new_mask\n            \n            return jax.lax.cond(valid_split, apply_split, lambda: current_mask)\n\n        node_mask = jax.lax.fori_loop(0, num_nodes_at_this_depth, update_mask_body, node_mask)\n        return (node_mask, all_feature_indices, all_thresholds)\n\n    initial_carry = (node_mask, all_feature_indices, all_thresholds)\n    final_carry = jax.lax.fori_loop(0, max_depth, loop_body, initial_carry)\n    dummy_outputs = jnp.zeros((max_depth, num_nodes_at_max_depth), dtype=jnp.float64)\n    return final_carry[1], final_carry[2], dummy_outputs, dummy_outputs\n\n# 💥 UPDATED JIT: Includes lambda_l2 as a static argument\n_fit_tree_kernel_jit = jax.jit(_fit_tree_kernel, static_argnums=(3, 4, 5)) \n\ndef _predict_leaf_index_kernel(X, tree_params, max_depth):\n    \"\"\"JAX-compiled function to find the leaf index for each data point.\"\"\"\n    feature_indices, thresholds, _, _ = tree_params\n    node_indices = jnp.zeros(X.shape[0], dtype=jnp.int32)\n    def body(carry, depth):\n        node_indices = carry\n        f_idx = feature_indices[depth, node_indices]\n        thr = thresholds[depth, node_indices]\n        f_vals = jnp.take_along_axis(X, f_idx[:, None], axis=1).squeeze(axis=1) \n        split_mask = f_vals < thr\n        return jnp.where(split_mask, 2 * node_indices, 2 * node_indices + 1), None\n    \n    final_nodes, _ = jax.lax.scan(body, node_indices, jnp.arange(max_depth))\n    return final_nodes\n\n_predict_leaf_index_kernel_jit = jax.jit(_predict_leaf_index_kernel, static_argnums=2)\n\n# --- JAXGBMDualHead Class Definition ---\n\nclass JAXGBMDualHead:\n    def __init__(self, n_estimators=100, learning_rate=0.1, max_depth=3, batch_size=20000, num_thresholds=70, lambda_ratio=0.5, lambda_l2=1.0):\n        self.n_estimators = n_estimators\n        self.learning_rate = learning_rate\n        self.max_depth = max_depth\n        self.num_thresholds = num_thresholds\n        self.batch_size = batch_size\n        self.lambda_ratio = lambda_ratio\n        self.lambda_l2 = lambda_l2\n        self.trees = []\n        self.reco_head_params = None\n\n    def fit(self, X_input, X_reco, y_class, sample_weights):\n        X_jnp = jnp.asarray(X_input, dtype=jnp.float64)\n        X_reco_jnp = jnp.asarray(X_reco, dtype=jnp.float64)\n        y_class_jnp = jnp.asarray(y_class, dtype=jnp.float64).ravel()\n        sample_weights_jnp = jnp.asarray(sample_weights, dtype=jnp.float64).ravel()\n        num_leaves = 2 ** self.max_depth\n        num_samples = X_jnp.shape[0]\n        if self.reco_head_params is None:\n            num_features = X_jnp.shape[1]\n            self.reco_head_params = jnp.zeros((num_leaves, num_features))\n        current_predictions = jnp.zeros_like(y_class_jnp)\n\n        @jax.jit\n        def _update_reco_head(reco_params, leaf_idx, X_reco_jnp, learning_rate):\n            reco_loss_grad_fn = jax.grad(lambda rp: mse_loss(X_reco_jnp, jnp.take(rp, leaf_idx, axis=0)))\n            reco_head_grad = reco_loss_grad_fn(reco_params)\n            return reco_params - learning_rate * reco_head_grad\n\n        for i in tqdm(range(self.n_estimators), desc=\"Training Dual-Head GBM\"):\n            # 1️⃣ Gradient / Hessian calculation\n            clipped_logits = jnp.clip(current_predictions, -20.0, 20.0)\n            grad = -jax.grad(lambda logits: weighted_bce_loss(\n                y_class_jnp, logits, sample_weights_jnp))(clipped_logits)\n\n            prob = jax.nn.sigmoid(clipped_logits)\n            hess = sample_weights_jnp * prob * (1.0 - prob)\n            \n            full_tree_params = None \n\n            # 2️⃣ Fit a fresh tree structure using batch-wise approximation\n            for start in range(0, num_samples, self.batch_size):\n                X_batch = X_jnp[start:start + self.batch_size]\n                g_batch = grad[start:start + self.batch_size]\n                h_batch = hess[start:start + self.batch_size] # 💥 NEW: get hessian batch\n                \n                # 💥 UPDATED CALL: Pass hessian batch and lambda_l2\n                tree_params_batch = _fit_tree_kernel_jit(\n                    X_batch, g_batch, h_batch, self.lambda_l2, self.max_depth, self.num_thresholds\n                )\n                \n                full_tree_params = tree_params_batch\n                \n            # 3️⃣ Calculate full dataset's leaf indices\n            leaf_idx_full = _predict_leaf_index_kernel_jit(\n                X_jnp, full_tree_params, self.max_depth\n            )\n\n            # 4️⃣ Accumulate full Gradient / Hessian\n            leaf_sum_g = jnp.zeros(num_leaves).at[leaf_idx_full].add(grad)\n            leaf_sum_h = jnp.zeros(num_leaves).at[leaf_idx_full].add(hess)\n\n            # 5️⃣ Leaf output (γ) – Newton step with L2 regularizer\n            gamma_leaf = -leaf_sum_g / (leaf_sum_h + self.lambda_l2 + 1e-8)\n\n            # 6️⃣ Update ensemble predictions\n            current_predictions = (\n                current_predictions\n                + self.learning_rate * jnp.take(gamma_leaf, leaf_idx_full)\n            )\n\n            self.trees.append((full_tree_params, np.array(gamma_leaf)))\n\n            # 7️⃣ Reconstruction head update\n            self.reco_head_params = _update_reco_head(\n                self.reco_head_params, leaf_idx_full, X_reco_jnp, self.learning_rate\n            )\n\n            # 8️⃣ Optional loss reporting\n            reco_pred = jnp.take(self.reco_head_params, leaf_idx_full, axis=0)\n            total_loss = combined_loss(\n                y_class_jnp,\n                current_predictions,\n                X_reco_jnp,\n                reco_pred,\n                sample_weights_jnp,\n                self.lambda_ratio,\n            )\n            if (i + 1) % 10 == 0 or i == 0:\n                tqdm.write(\n                    f\"Estimator {i+1}/{self.n_estimators} – Combined loss: {total_loss:.5f}\"\n                )\n\n    def predict(self, X):\n        X_jnp = jnp.asarray(X, dtype=jnp.float64)\n        predictions = jnp.zeros(X_jnp.shape[0])\n        for tree_params, gamma_per_leaf in self.trees:\n            tree_params_jnp = jax.tree_map(jnp.asarray, tree_params)\n            leaf_idx = _predict_leaf_index_kernel_jit(X_jnp, tree_params_jnp, self.max_depth)\n            gamma_jnp = jnp.asarray(gamma_per_leaf)\n            predictions += self.learning_rate * jnp.take(gamma_jnp, leaf_idx)\n        return 1 / (1 + jnp.exp(-predictions))\n\n\n# ----------------------------------------------------------------------\n# PLOTTING FUNCTIONS\n# ----------------------------------------------------------------------\ndef plot_results(results):\n    \"\"\"Generates and saves a comparison plot of F1 Score and Training Time.\"\"\"\n    models = list(results['comparison'].keys())\n    f1_scores = [results['comparison'][m]['F1'] for m in models]\n    train_times = [results['comparison'][m]['Time'] for m in models]\n\n    fig, axes = plt.subplots(1, 2, figsize=(12, 5))\n    \n    bar_colors = ['#4CAF50', '#007ACC']\n\n    # --- F1 Score Plot ---\n    axes[0].bar(models, f1_scores, color=bar_colors)\n    axes[0].set_title('Classification Performance (F1 Score)')\n    axes[0].set_ylabel('F1 Score')\n    axes[0].grid(axis='y', linestyle='--', alpha=0.7)\n    axes[0].set_ylim(0.0, max(f1_scores) * 1.1)\n    \n    for i, v in enumerate(f1_scores):\n        axes[0].text(i, v + 0.01, f\"{v:.4f}\", color='black', ha='center')\n\n    # --- Training Time Plot ---\n    axes[1].bar(models, train_times, color=bar_colors)\n    axes[1].set_title('Training Time Comparison')\n    axes[1].set_ylabel('Time (seconds)')\n    axes[1].grid(axis='y', linestyle='--', alpha=0.7)\n    axes[1].set_ylim(0, max(train_times) * 1.1)\n    \n    for i, v in enumerate(train_times):\n        axes[1].text(i, v + max(train_times) * 0.02, f\"{v:.2f}s\", color='black', ha='center')\n\n    fig.suptitle('JAX-based Dual-Head GBM vs. Standard XGBoost Comparison', fontsize=16)\n    plt.tight_layout(rect=[0, 0.03, 1, 0.95])\n    \n    plot_filename = 'gbm_comparison_plot.png'\n    plt.savefig(plot_filename, dpi=300)\n    print(f\"\\n[Plot saved to {plot_filename}]\")\n\ndef plot_residuals(y_test, preds_probs_jax):\n    \"\"\"Generates a residual and probability distribution plot for JAXGBMDualHead.\"\"\"\n    \n    residuals = y_test - preds_probs_jax\n\n    fig, axes = plt.subplots(1, 2, figsize=(12, 5))\n    \n    # --- Residual Scatter Plot ---\n    axes[0].scatter(preds_probs_jax, residuals, alpha=0.3, s=5)\n    axes[0].axhline(0, color='r', linestyle='--')\n    axes[0].set_title('Residuals ($y - \\hat{p}$) vs. Predicted Probability')\n    axes[0].set_xlabel('Predicted Probability ($\\hat{p}$)')\n    axes[0].set_ylabel('Residual ($y - \\hat{p}$)')\n    axes[0].set_ylim(-1.1, 1.1)\n    axes[0].grid(True, linestyle='--', alpha=0.5)\n\n    # --- Probability Distribution Plot ---\n    axes[1].hist(preds_probs_jax[y_test == 0], bins=20, alpha=0.6, label='True Class 0', color='#007ACC')\n    axes[1].hist(preds_probs_jax[y_test == 1], bins=20, alpha=0.6, label='True Class 1', color='#4CAF50')\n    axes[1].set_title('Distribution of Predicted Probabilities')\n    axes[1].set_xlabel('Predicted Probability ($\\hat{p}$)')\n    axes[1].set_ylabel('Frequency')\n    axes[1].legend()\n    axes[1].grid(axis='y', linestyle='--', alpha=0.5)\n\n    fig.suptitle('JAXGBMDualHead Model Diagnostic Plots', fontsize=16)\n    plt.tight_layout(rect=[0, 0.03, 1, 0.95])\n    \n    plot_filename = 'gbm_residuals_plot.png'\n    plt.savefig(plot_filename, dpi=300)\n    print(f\"[Plot saved to {plot_filename}]\")\n\n# ----------------------------------------------------------------------\n# Run Comparison Function (UPDATED to collect plotting data)\n# ----------------------------------------------------------------------\ndef run_comparison():\n    print(\"Generating synthetic data...\")\n    num_samples = 20000\n    np.random.seed(42)\n    \n    # Data Setup\n    time_series_feature = np.sin(np.linspace(0, 10 * np.pi, num_samples))\n    y_class = (np.diff(time_series_feature) > 0).astype(int)\n    X = np.zeros((num_samples, 5))\n    X[:, 0] = time_series_feature\n    X[:, 1:] = np.random.randn(num_samples, 4) * 0.1\n    X = X[:-1]\n    y_class = y_class.ravel()\n    \n    y_class_jax = np.where(y_class == 0, -1, 1) \n    X_train, X_test, y_train_jax, y_test = train_test_split(X, y_class_jax, test_size=0.2, random_state=42)\n    y_train_xgb = np.where(y_train_jax == -1, 0, 1) \n    y_test_xgb = np.where(y_test == -1, 0, 1)\n    X_reco_train = X_train + np.random.randn(X_train.shape[0], X_train.shape[1]) * 0.05\n    sample_weights = np.ones(y_train_jax.shape[0])\n\n    comparison_results = {}\n    preds_probs_jax = np.zeros_like(y_test_xgb, dtype=np.float64) # Initialize for error handling\n\n    # --- JAXGBMDualHead Training ---\n    print(\"\\n--- Testing JAXGBMDualHead (UPDATED) ---\")\n    gbm_jax = JAXGBMDualHead(n_estimators=50, learning_rate=0.1, max_depth=3, batch_size=2000, num_thresholds=40, lambda_ratio=0.9, lambda_l2=0.001)\n    \n    start_time_jax = time.time()\n    try:\n        gbm_jax.fit(X_train, X_reco_train, y_train_jax, sample_weights)\n        end_time_jax = time.time()\n        preds_probs_jax = gbm_jax.predict(X_test)\n        preds_probs_jax = np.asarray(preds_probs_jax) \n        preds_binary_jax = (preds_probs_jax > 0.5).astype(int)\n        f1_jax = f1_score(y_test_xgb, preds_binary_jax)\n    except Exception as e:\n        print(f\"\\n❌ ERROR during JAXGBMDualHead execution: {e}\")\n        f1_jax = np.nan\n        end_time_jax = start_time_jax # Set time to 0 if fit failed\n        preds_probs_jax = np.zeros_like(y_test_xgb, dtype=np.float64) \n        \n    training_time_jax = end_time_jax - start_time_jax\n    comparison_results['JAXGBMDualHead'] = {'F1': f1_jax, 'Time': training_time_jax}\n\n    # --- XGBoost Training ---\n    print(\"\\n--- Testing XGBoost Classifier ---\")\n    gbm_xgb = xgb.XGBClassifier(n_estimators=50, learning_rate=0.1, max_depth=3, use_label_encoder=False, eval_metric='logloss')\n    start_time_xgb = time.time()\n    try:\n        gbm_xgb.fit(X_train, y_train_xgb)\n        end_time_xgb = time.time()\n        preds_probs_xgb = gbm_xgb.predict_proba(X_test)[:, 1]\n        preds_binary_xgb = (preds_probs_xgb > 0.5).astype(int)\n        f1_xgb = f1_score(y_test_xgb, preds_binary_xgb)\n    except Exception as e:\n        print(f\"\\n❌ ERROR during XGBoost execution: {e}\")\n        f1_xgb = np.nan\n        end_time_xgb = start_time_xgb\n\n    training_time_xgb = end_time_xgb - start_time_xgb\n    comparison_results['XGBoost'] = {'F1': f1_xgb, 'Time': training_time_xgb}\n    \n    # --- Final Print Results ---\n    print(\"\\n\" + \"=\"*40)\n    print(\"🚀 Final Performance Summary\")\n    print(\"=\"*40)\n    \n    print(f\"\\n✅ JAXGBMDualHead Results:\")\n    print(f\"   F1 Score: {f1_jax:.4f}\")\n    print(f\"   Training Time: {training_time_jax:.2f} seconds\")\n    \n    print(f\"\\n✅ XGBoost Results:\")\n    print(f\"   F1 Score: {f1_xgb:.4f}\")\n    print(f\"   Training Time: {training_time_xgb:.2f} seconds\")\n\n    return {\n        'comparison': comparison_results,\n        'data': {\n            'y_test': y_test_xgb, \n            'preds_jax': preds_probs_jax\n        }\n    }\n\n# ----------------------------------------------------------------------\n# MAIN EXECUTION BLOCK (UPDATED to call plotting functions)\n# ----------------------------------------------------------------------\nif __name__ == '__main__':\n    print(\"\\n\" + \"#\"*50)\n    print(\"WARNING: The environment must have JAX, XGBoost, and Matplotlib\")\n    print(\"installed for the script to run and generate plots.\")\n    print(\"The code is provided below with plots enabled.\")\n    print(\"#\"*50 + \"\\n\")\n    \n    comparison_data = run_comparison()\n    plot_results(comparison_data)\n    y_test = comparison_data['data']['y_test']\n    preds_jax = comparison_data['data']['preds_jax']\n    plot_residuals(y_test, preds_jax)\n    ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-23T11:56:20.769141Z","iopub.execute_input":"2025-12-23T11:56:20.769796Z","iopub.status.idle":"2025-12-23T11:56:30.974393Z","shell.execute_reply.started":"2025-12-23T11:56:20.769765Z","shell.execute_reply":"2025-12-23T11:56:30.973655Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport jax\nimport jax.numpy as jnp\nfrom sklearn.metrics import f1_score\nfrom tqdm.auto import tqdm\nimport warnings\nimport xgboost as xgb\nfrom sklearn.model_selection import train_test_split\nimport time\nimport matplotlib.pyplot as plt\nfrom functools import partial\n\n# --- JAX Configuration ---\nwarnings.filterwarnings(\"ignore\")\njax.config.update(\"jax_enable_x64\", True)\n\n# --- Loss Functions ---\ndef weighted_bce_loss(y_true, y_pred_logits, weights):\n    # y_true in {-1, 1}\n    return jnp.mean(weights * jnp.log(1 + jnp.exp(-y_true * y_pred_logits)))\n\ndef mse_loss(y_true, y_pred):\n    return jnp.mean((y_true - y_pred)**2)\n\ndef combined_loss(y_true, y_pred_class_logits, y_true_reco, y_pred_reco, weights, lambda_ratio):\n    supervised_loss = lambda_ratio * weighted_bce_loss(y_true, y_pred_class_logits, weights)\n    reconstruction_loss = (1 - lambda_ratio) * mse_loss(y_true_reco, y_pred_reco)\n    return supervised_loss + reconstruction_loss\n\n# --- High-Performance JAX GBM Kernels (Scan Optimized) ---\n\n@partial(jax.jit, static_argnums=(4, 5))\ndef _fit_tree_kernel(X, grad, hess, lambda_l2, max_depth, num_thresholds):\n    \"\"\"Fits tree structure using XGBoost Gain logic and scan-optimized loops.\"\"\"\n    num_data_points = X.shape[0]\n    num_nodes_total = 2**max_depth\n    \n    all_f_idx = jnp.full((max_depth, num_nodes_total), -1, dtype=jnp.int32)\n    all_thrs = jnp.zeros((max_depth, num_nodes_total), dtype=jnp.float64)\n    node_mask = jnp.zeros(num_data_points, dtype=jnp.int32)\n\n    def depth_scan(carry, depth):\n        node_mask, afi, ath = carry\n        num_nodes_this_depth = 2**depth\n\n        def node_loop(node_idx, node_carry):\n            afi_loc, ath_loc = node_carry\n            in_node = (node_mask == node_idx)\n            \n            def find_split():\n                # Node statistics\n                G_node = jnp.sum(jnp.where(in_node, grad, 0.0))\n                H_node = jnp.sum(jnp.where(in_node, hess, 0.0))\n                node_gain_term = -(G_node**2 / (H_node + lambda_l2 + 1e-8))\n\n                def feat_search(feat_col):\n                    v_min = jnp.nanmin(jnp.where(in_node, feat_col, jnp.nan))\n                    v_max = jnp.nanmax(jnp.where(in_node, feat_col, jnp.nan))\n                    \n                    def thr_scan(best, thr):\n                        best_g, best_t = best\n                        left = in_node & (feat_col < thr)\n                        G_L = jnp.sum(jnp.where(left, grad, 0.0))\n                        H_L = jnp.sum(jnp.where(left, hess, 0.0))\n                        G_R, H_R = G_node - G_L, H_node - H_L\n                        \n                        gain = jax.lax.cond((H_L > 0) & (H_R > 0),\n                            lambda: (G_L**2 / (H_L + lambda_l2)) + (G_R**2 / (H_R + lambda_l2)) + node_gain_term,\n                            lambda: -jnp.inf)\n                        \n                        is_better = gain > best_g\n                        return (jax.lax.select(is_better, gain, best_g), jax.lax.select(is_better, thr, best_t)), None\n\n                    thrs = jnp.linspace(v_min, v_max, num_thresholds)\n                    res, _ = jax.lax.scan(thr_scan, (-jnp.inf, 0.0), thrs)\n                    return res\n\n                scores, thrs = jax.vmap(feat_search, in_axes=1)(X)\n                best_f = jnp.argmax(scores).astype(jnp.int32)\n                return best_f, thrs[best_f]\n\n            # FIXED: Explicitly cast return values to match types exactly\n            res = jax.lax.cond(jnp.sum(in_node) > 1, \n                               find_split, \n                               lambda: (jnp.array(-1, dtype=jnp.int32), jnp.array(0.0, dtype=jnp.float64)))\n            \n            return afi_loc.at[depth, node_idx].set(res[0]), ath_loc.at[depth, node_idx].set(res[1])\n\n        afi, ath = jax.lax.fori_loop(0, num_nodes_this_depth, node_loop, (afi, ath))\n        \n        def mask_update(i, m):\n            f_idx, thr = afi[depth, i], ath[depth, i]\n            is_in = (m == i)\n            cond = is_in & (f_idx >= 0)\n            go_left = cond & (X[:, f_idx] < thr)\n            m = jnp.where(go_left, 2*i, m)\n            m = jnp.where(cond & (~go_left), 2*i + 1, m)\n            return m\n        node_mask_new = jax.lax.fori_loop(0, num_nodes_this_depth, mask_update, node_mask)\n        return (node_mask_new, afi, ath), None\n\n    final_state, _ = jax.lax.scan(depth_scan, (node_mask, all_f_idx, all_thrs), jnp.arange(max_depth))\n    return final_state[1], final_state[2]\n\n@partial(jax.jit, static_argnums=2)\ndef _predict_leaf_indices(X, tree_params, max_depth):\n    f_indices, thresholds = tree_params\n    node_indices = jnp.zeros(X.shape[0], dtype=jnp.int32)\n    def depth_step(idx, depth):\n        f_idx, thr = f_indices[depth, idx], thresholds[depth, idx]\n        # Safe indexing to avoid out-of-bounds errors on non-split nodes (-1)\n        safe_f_idx = jnp.maximum(0, f_idx)\n        f_vals = jnp.take_along_axis(X, safe_f_idx[:, None], axis=1).squeeze(1)\n        next_idx = jnp.where(f_vals < thr, 2*idx, 2*idx + 1)\n        return jnp.where(f_idx >= 0, next_idx, idx), None\n    final_idx, _ = jax.lax.scan(depth_step, node_indices, jnp.arange(max_depth))\n    return final_idx\n\n# --- JAXGBMDualHead Implementation ---\nclass JAXGBMDualHead:\n    def __init__(self, n_estimators=100, learning_rate=0.1, max_depth=3, num_thresholds=32, lambda_ratio=0.5, lambda_l2=1.0):\n        self.n_estimators = n_estimators\n        self.learning_rate = learning_rate\n        self.max_depth = max_depth\n        self.num_thresholds = num_thresholds\n        self.lambda_ratio = lambda_ratio\n        self.lambda_l2 = lambda_l2\n        self.trees = []\n        self.reco_head_params = None\n\n    def fit(self, X_input, X_reco, y_class, sample_weights):\n        X_jnp = jnp.asarray(X_input, dtype=jnp.float64)\n        X_reco_jnp = jnp.asarray(X_reco, dtype=jnp.float64)\n        y_class_jnp = jnp.asarray(y_class, dtype=jnp.float64)\n        weights_jnp = jnp.asarray(sample_weights, dtype=jnp.float64)\n        \n        num_leaves = 2 ** self.max_depth\n        self.reco_head_params = jnp.zeros((num_leaves, X_reco_jnp.shape[1]))\n        current_logits = jnp.zeros_like(y_class_jnp)\n\n        @jax.jit\n        def update_reco(params, leaf_idx, target, lr):\n            grad_reco = jax.grad(lambda p: mse_loss(target, jnp.take(p, leaf_idx, axis=0)))(params)\n            return params - lr * grad_reco\n\n        for i in tqdm(range(self.n_estimators), desc=\"Training\"):\n            p = jax.nn.sigmoid(current_logits)\n            y_bin = (y_class_jnp + 1) / 2\n            grad = weights_jnp * (p - y_bin)\n            hess = weights_jnp * p * (1.0 - p) + 1e-8\n            \n            f_idx, thrs = _fit_tree_kernel(X_jnp, grad, hess, self.lambda_l2, self.max_depth, self.num_thresholds)\n            leaf_idx = _predict_leaf_indices(X_jnp, (f_idx, thrs), self.max_depth)\n            \n            sum_g = jnp.zeros(num_leaves).at[leaf_idx].add(grad)\n            sum_h = jnp.zeros(num_leaves).at[leaf_idx].add(hess)\n            gamma = -sum_g / (sum_h + self.lambda_l2)\n            \n            current_logits += self.learning_rate * gamma[leaf_idx]\n            self.trees.append(((f_idx, thrs), np.array(gamma)))\n            self.reco_head_params = update_reco(self.reco_head_params, leaf_idx, X_reco_jnp, self.learning_rate)\n            \n            if (i+1) % 10 == 0:\n                loss = combined_loss(y_class_jnp, current_logits, X_reco_jnp, jnp.take(self.reco_head_params, leaf_idx, axis=0), weights_jnp, self.lambda_ratio)\n                tqdm.write(f\"Estimator {i+1} Loss: {loss:.5f}\")\n\n    def predict(self, X):\n        X_jnp = jnp.asarray(X, dtype=jnp.float64)\n        logits = jnp.zeros(X_jnp.shape[0])\n        for (struct, gamma) in self.trees:\n            idx = _predict_leaf_indices(X_jnp, struct, self.max_depth)\n            logits += self.learning_rate * jnp.asarray(gamma)[idx]\n        return jax.nn.sigmoid(logits)\n\ndef run_comparison():\n    print(\"Generating synthetic data...\")\n    num_samples = 10000\n    np.random.seed(42)\n    t = np.linspace(0, 10 * np.pi, num_samples)\n    X = np.column_stack([np.sin(t), np.cos(t), np.random.randn(num_samples, 3)])\n    y = (np.roll(X[:, 0], -1) > X[:, 0]).astype(int)[:-1]\n    X = X[:-1]\n    \n    y_jax = np.where(y == 0, -1, 1)\n    X_train, X_test, y_train_jax, y_test_jax = train_test_split(X, y_jax, test_size=0.2, random_state=42)\n    y_train_xgb = (y_train_jax + 1) // 2\n    y_test_xgb = (y_test_jax + 1) // 2\n    \n    # JAX Fit\n    print(\"\\n--- Testing JAX Optimized GBM ---\")\n    model_jax = JAXGBMDualHead(n_estimators=50, learning_rate=0.1, max_depth=3, lambda_l2=0.1)\n    start = time.time()\n    model_jax.fit(X_train, X_train, y_train_jax, np.ones(len(y_train_jax)))\n    jax_time = time.time() - start\n    jax_preds = model_jax.predict(X_test)\n    jax_f1 = f1_score(y_test_xgb, (jax_preds > 0.5).astype(int))\n    \n    # XGBoost Fit\n    print(\"\\n--- Testing Standard XGBoost ---\")\n    model_xgb = xgb.XGBClassifier(n_estimators=50, learning_rate=0.1, max_depth=3, reg_lambda=0.1)\n    start = time.time()\n    model_xgb.fit(X_train, y_train_xgb)\n    xgb_time = time.time() - start\n    xgb_f1 = f1_score(y_test_xgb, model_xgb.predict(X_test))\n\n    print(f\"\\nFinal Results:\")\n    print(f\"JAX GBM: F1={jax_f1:.4f}, Time={jax_time:.2f}s\")\n    print(f\"XGBoost: F1={xgb_f1:.4f}, Time={xgb_time:.2f}s\")\n\nif __name__ == '__main__':\n    run_comparison()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-08T12:07:59.927334Z","iopub.execute_input":"2026-01-08T12:07:59.927835Z","iopub.status.idle":"2026-01-08T12:08:03.676884Z","shell.execute_reply.started":"2026-01-08T12:07:59.927806Z","shell.execute_reply":"2026-01-08T12:08:03.675107Z"}},"outputs":[],"execution_count":null}]}