{"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":"gpu","dataSources":[{"sourceId":114135,"databundleVersionId":14604229,"sourceType":"competition"}],"dockerImageVersionId":31194,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"#!pip install tensorflow flax scikit-learn matplotlib seaborn polars","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#!/usr/bin/env python3\n\"\"\"\nfinal_ga_optimized.py\nFixes:\n1. XlaRuntimeError (OOM): Removed by using precomputed Gram matrices.\n2. Slowness: Optimized using Subspace Solving (solving 256x256 instead of 3000x3000).\n\"\"\"\n\nimport os\nimport logging\nfrom typing import List, Dict\nfrom functools import partial\n\nimport numpy as np\nimport pandas as pd\nimport polars as pl\nfrom tqdm.auto import tqdm\n\nimport jax\nimport jax.numpy as jnp\nfrom jax import random, jit, vmap, lax\nfrom sklearn.preprocessing import LabelEncoder, StandardScaler\n\n# -------------------------\n# CONFIG\n# -------------------------\nTRAIN_PATH = \"/kaggle/input/detecting-reversal-points-in-us-equities/new_comptetition_data/train.csv\"\nPOP_SIZE = 1024\n# We can increase batch size now because memory usage is tiny with subspace solving\nEVAL_BATCH_SIZE = 128 \nISLANDS = 4\nMAX_TEMPORAL_CANDIDATES = 3000 \nRIDGE_ALPHA = 1.0\nSPARSITY = 0.05\nCXPB = 0.5\nMUTPB = 0.05\nELITISM = 8\n\n# Hard cap for the solver. \n# 5% of 3000 features is 150, so 256 gives plenty of room for mutation growth.\nMAX_SOLVE_CAPACITY = 256 \n\njax.config.update('jax_default_matmul_precision', 'high')\nlogging.basicConfig(level=logging.INFO, format=\"%(asctime)s %(levelname)s: %(message)s\")\nlogger = logging.getLogger(\"ga_pipeline\")\n\n# -------------------------\n# Feature Grouping\n# -------------------------\ndef categorize_features_precise(feature_names: List[str]) -> Dict[str, List[str]]:\n    groups = {'base': [], 'events': [], 'temporal_relations': [], 'other': []}\n    for f in feature_names:\n        if f in {'ratio', 'momentum', 'sm_ratio', 'sm_momentum'}:\n            groups['base'].append(f)\n        elif any(f.startswith(p) for p in ['troughs_', 'peaks_', 'cross_threshold_', 'trending_', 'zone_']):\n            groups['events'].append(f)\n        elif f.startswith(('occurs_within_', 'happens_within_')):\n            groups['temporal_relations'].append(f)\n        else:\n            groups['other'].append(f)\n    return groups\n\n# -------------------------\n# JAX Numerical Core\n# -------------------------\n@jit\ndef compute_feature_scores_jax(X_f32: jnp.ndarray, y_i32: jnp.ndarray) -> jnp.ndarray:\n    \"\"\"Computes correlation scores for initial feature pruning.\"\"\"\n    X_centered = X_f32 - jnp.mean(X_f32, axis=0, keepdims=True)\n    y_f32 = y_i32.astype(jnp.float32)\n    y_centered = y_f32 - jnp.mean(y_f32)\n    cov = jnp.mean(X_centered * y_centered[:, None], axis=0)\n    x_std = jnp.std(X_f32, axis=0)\n    y_std = jnp.std(y_f32)\n    return jnp.abs(cov / (jnp.where(x_std == 0, 1.0, x_std) * jnp.where(y_std == 0, 1.0, y_std)))\n\n@partial(jit, static_argnums=(2,))\ndef compute_f1_macro_jit(y_true: jnp.ndarray, y_pred: jnp.ndarray, num_classes: int):\n    y_true_oh = jax.nn.one_hot(y_true, num_classes)\n    y_pred_oh = jax.nn.one_hot(y_pred, num_classes)\n    tp = jnp.sum(y_true_oh * y_pred_oh, axis=0)\n    fp = jnp.sum((1 - y_true_oh) * y_pred_oh, axis=0)\n    fn = jnp.sum(y_true_oh * (1 - y_pred_oh), axis=0)\n    precision = jnp.where(tp + fp == 0, 0.0, tp / (tp + fp))\n    recall = jnp.where(tp + fn == 0, 0.0, tp / (tp + fn))\n    f1 = jnp.where(precision + recall == 0, 0.0, 2 * precision * recall / (precision + recall))\n    return jnp.mean(f1)\n\n@partial(jit, static_argnums=(4,5))\ndef solve_subspace_jit(XtX_full, Xty_full, mask, alpha, num_classes, max_cap):\n    \"\"\"\n    Solves Ridge Regression ONLY on the active feature subspace (e.g., 256 features).\n    This is ~1000x faster than masking the full 3000 matrix.\n    \"\"\"\n    # 1. Identify active indices (padded to fixed size for JIT)\n    # argsort sorts False (0) before True (1), so we want descending to put 1s at top\n    # Note: jnp.argsort is stable.\n    sort_idx = jnp.argsort(mask, descending=True)\n    active_indices = sort_idx[:max_cap]\n    \n    # 2. Determine which of these indices are actually 'True' in the mask\n    # (If an individual has < 256 genes, the rest are dummy padding)\n    is_active = mask[active_indices] \n    \n    # 3. Gather the Sub-Matrix (256x256) instead of (3000x3000)\n    # JAX handles this gather efficiently\n    XtX_sub = XtX_full[active_indices][:, active_indices]\n    Xty_sub = Xty_full[active_indices]\n    \n    # 4. Apply masking to the sub-system to handle the dummy padding\n    # If a feature is inactive (0), we set its row/col to Identity and target to 0\n    # This forces its weight to 0.0 without breaking the solver with Singular Matrix errors.\n    sub_mask_mat = is_active[:, None] * is_active[None, :]\n    \n    # Zero out rows/cols of dummy features\n    XtX_active = XtX_sub * sub_mask_mat\n    \n    # Add regularization + Identity for dummy slots\n    # diag_reg: alpha for active features, 1.0 for dummy features (to invert safely)\n    diag_reg = jnp.where(is_active > 0, alpha, 1.0)\n    \n    # Create the regularized matrix\n    # The term (jnp.eye * (1-is_active)) puts 1.0 on the diagonal of dummy rows\n    XtX_reg = XtX_active + jnp.diag(diag_reg) + (jnp.eye(max_cap) * (1.0 - is_active[:, None]))\n    \n    Xty_active = Xty_sub * is_active[:, None]\n    \n    # 5. Solve small system\n    W_sub = jnp.linalg.solve(XtX_reg, Xty_active)\n    \n    return W_sub, active_indices\n\ndef make_evaluate_batch_optimized(precomputed_folds, num_classes, alpha, max_cap):\n    \"\"\"\n    Creates the evaluation function.\n    precomputed_folds: List of tuples (XtX_tr, Xty_tr, X_val, y_val)\n    \"\"\"\n    \n    def evaluate_individual(mask):\n        f1_scores = []\n        \n        # Loop over folds (usually 1 or 5)\n        # Because this list is static python list, JAX unrolls it during compilation\n        for (XtX_tr, Xty_tr, X_val, y_val) in precomputed_folds:\n            \n            # Solve only for active genes (Fast!)\n            W_sub, active_indices = solve_subspace_jit(XtX_tr, Xty_tr, mask, alpha, num_classes, max_cap)\n            \n            # Predict\n            # Gather relevant columns from Validation set matching the solved weights\n            X_val_sub = X_val[:, active_indices]\n            \n            logits = X_val_sub @ W_sub\n            y_pred = jnp.argmax(logits, axis=1)\n            \n            f1_scores.append(compute_f1_macro_jit(y_val, y_pred, num_classes))\n            \n        return jnp.mean(jnp.stack(f1_scores))\n\n    # We use vmap here. Since matrices are small (256x256), memory is low.\n    return jit(vmap(evaluate_individual))\n\n# -------------------------\n# GA Engine\n# -------------------------\nclass GAIsland:\n    def __init__(self, rng_key, pop_size, n_genes, sparsity, cxpb, mutpb, elitism):\n        self.rng = rng_key\n        self.pop_size, self.n_genes = pop_size, n_genes\n        self.cxpb, self.mutpb, self.elitism = cxpb, mutpb, elitism\n        sub, self.rng = random.split(self.rng)\n        self.pop = random.bernoulli(sub, p=sparsity, shape=(pop_size, n_genes))\n        self.fitness = jnp.zeros(pop_size)\n\n    def evaluate_full_pop(self, eval_batch_fn):\n        \"\"\"Evaluates population in batches.\"\"\"\n        fits = []\n        for i in range(0, self.pop_size, EVAL_BATCH_SIZE):\n            batch = self.pop[i : i + EVAL_BATCH_SIZE]\n            fits.append(eval_batch_fn(batch))\n        self.fitness = jnp.concatenate(fits)\n\n    def epoch(self, eval_batch_fn):\n        sub, self.rng = random.split(self.rng)\n        idx = random.choice(sub, self.pop_size, (self.pop_size, 3))\n        winners = jnp.argmax(self.fitness[idx], axis=1)\n        parents = self.pop[idx[jnp.arange(self.pop_size), winners]]\n        \n        sub, self.rng = random.split(self.rng)\n        mask = random.bernoulli(sub, 0.5, parents.shape)\n        offspring = jnp.where(random.uniform(sub, (self.pop_size, 1)) < self.cxpb, \n                              jnp.roll(parents, 1, axis=0), parents)\n        offspring = jnp.where(mask, parents, offspring)\n        \n        sub, self.rng = random.split(self.rng)\n        offspring = jnp.logical_xor(offspring, random.bernoulli(sub, self.mutpb, offspring.shape))\n        \n        elite_idx = jnp.argsort(self.fitness)[-self.elitism:]\n        elites = self.pop[elite_idx]\n        self.pop = offspring.at[:self.elitism].set(elites)\n        self.evaluate_full_pop(eval_batch_fn)\n\ndef run_multi_island_ga(eval_batch_fn, n_genes):\n    base_key = random.PRNGKey(42)\n    islands = [GAIsland(random.fold_in(base_key, i), POP_SIZE, n_genes, SPARSITY, CXPB, MUTPB, ELITISM) \n               for i in range(ISLANDS)]\n    \n    logger.info(\"Initial population evaluation...\")\n    for isl in islands: isl.evaluate_full_pop(eval_batch_fn)\n    \n    global_best_score, global_best_ind = -1.0, None\n    \n    pbar = tqdm(range(200), desc=\"GA Evolution\")\n    for gen in pbar:\n        current_gen_best = -1.0\n        for isl in islands: \n            isl.epoch(eval_batch_fn)\n            best_idx = jnp.argmax(isl.fitness)\n            score = float(isl.fitness[best_idx])\n            \n            if score > current_gen_best:\n                current_gen_best = score\n                \n            if score > global_best_score:\n                global_best_score = score\n                global_best_ind = np.array(isl.pop[best_idx])\n                \n        pbar.set_postfix({\"BestF1\": f\"{global_best_score:.4f}\", \"GenBest\": f\"{current_gen_best:.4f}\"})\n        \n    return global_best_ind, global_best_score\n\n# -------------------------\n# Main Execution\n# -------------------------\ndef main():\n    logger.info(\"Initializing Data Pipeline...\")\n    \n    # 1. Load Data\n    schema = pl.scan_csv(TRAIN_PATH).collect_schema()\n    numeric_feats = [c for c, dt in schema.items() if dt in (pl.Float32, pl.Float64, pl.Int32, pl.Int64, pl.Boolean) and c not in ['class_label', 'ticker_id']]\n    \n    df = pl.read_csv(TRAIN_PATH, columns=['class_label'] + numeric_feats)\n    df = df.with_columns(pl.col('class_label').fill_null('None').cast(pl.Utf8))\n    \n    le = LabelEncoder()\n    y_encoded = le.fit_transform(df['class_label'].to_numpy())\n    num_classes = int(len(le.classes_))\n    \n    df_proc = df.with_columns([pl.col(c).cast(pl.Float32).fill_null(0.0) for c in numeric_feats])\n    X_all_jax = jnp.array(df_proc.drop('class_label').to_numpy())\n    y_jax = jnp.array(y_encoded, dtype=jnp.int32)\n    \n    # 2. Pruning Features\n    logger.info(\"Pruning initial feature set...\")\n    scores = np.array(compute_feature_scores_jax(X_all_jax, y_jax))\n    groups = categorize_features_precise(numeric_feats)\n    score_ser = pd.Series(scores, index=df_proc.drop('class_label').columns)\n    \n    essential = groups['base'] + groups['events'] + groups['other']\n    top_temporal = score_ser[score_ser.index.isin(groups['temporal_relations'])].sort_values(ascending=False).head(MAX_TEMPORAL_CANDIDATES).index.tolist()\n    all_star = essential + top_temporal\n    \n    # 3. Prepare Final Dataset\n    logger.info(f\"Selected {len(all_star)} candidate features.\")\n    X_ga_jax = jnp.array(StandardScaler().fit_transform(df_proc.select(all_star).to_numpy()))\n    \n    # 4. Precompute Gram Matrices for Speed\n    # Instead of passing raw data to GA, we calculate XtX and Xty ONCE.\n    logger.info(\"Precomputing Gram Matrices (XtX, Xty)...\")\n    \n    n = X_ga_jax.shape[0]\n    split = int(n * 0.8)\n    \n    # Indices\n    tr_idx = jnp.arange(0, split)\n    val_idx = jnp.arange(split, n)\n    \n    # Train Data (Precompute Matrices)\n    X_tr = X_ga_jax[tr_idx]\n    y_tr = y_jax[tr_idx]\n    y_tr_oh = jax.nn.one_hot(y_tr, num_classes)\n    \n    # The heavy lifting is done here, once.\n    XtX_train = X_tr.T @ X_tr\n    Xty_train = X_tr.T @ y_tr_oh\n    \n    # Validation Data (Keep Raw)\n    X_val = X_ga_jax[val_idx]\n    y_val = y_jax[val_idx]\n    \n    # Package into a list (supports multiple folds if you extend logic later)\n    precomputed_data = [(XtX_train, Xty_train, X_val, y_val)]\n    \n    # 5. Run GA\n    logger.info(\"Starting Evolutionary Search...\")\n    eval_batch_fn = make_evaluate_batch_optimized(precomputed_data, num_classes, RIDGE_ALPHA, MAX_SOLVE_CAPACITY)\n    \n    best_mask, best_f1 = run_multi_island_ga(eval_batch_fn, len(all_star))\n    \n    # 6. Save Results\n    final_features = [all_star[i] for i in range(len(all_star)) if best_mask[i]]\n    pd.Series(final_features).to_csv(\"final_selected_features.csv\", index=False)\n    logger.info(f\"Done. Selected {len(final_features)} features. Max F1: {best_f1:.6f}\")\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-30T11:14:58.342831Z","iopub.execute_input":"2025-12-30T11:14:58.343627Z","iopub.status.idle":"2025-12-30T11:16:28.211724Z","shell.execute_reply.started":"2025-12-30T11:14:58.343595Z","shell.execute_reply":"2025-12-30T11:16:28.211109Z"}},"outputs":[],"execution_count":null}]}