{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":96164,"databundleVersionId":11418275,"sourceType":"competition"}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport os\nfrom pathlib import Path\nfrom sklearn.model_selection import KFold, train_test_split\nfrom sklearn.preprocessing import StandardScaler, QuantileTransformer\nfrom sklearn.feature_selection import mutual_info_regression, SelectKBest\nfrom scipy.stats import pearsonr, spearmanr\nfrom tqdm import tqdm\nimport random\nimport warnings\nwarnings.filterwarnings(\"ignore\", category=RuntimeWarning)\n\n# Deep Learning imports\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import DataLoader, TensorDataset\nimport torch.nn.functional as F\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\nimport math\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"Using device: {device}\")\n\n# =========================\n# Configuration\n# =========================\nclass Config:\n    TRAIN_PATH = \"/kaggle/input/drw-crypto-market-prediction/train.parquet\"\n    TEST_PATH = \"/kaggle/input/drw-crypto-market-prediction/test.parquet\"\n    SUBMISSION_PATH = \"/kaggle/input/drw-crypto-market-prediction/sample_submission.csv\"\n    \n    # Extended features for SAINT\n    SAINT_FEATURES = [\n        \"X863\", \"X856\", \"X598\", \"X862\", \"X385\", \"X852\", \"X603\", \"X860\", \"X674\",\n        \"X415\", \"X345\", \"X855\", \"X174\", \"X302\", \"X178\", \"X168\", \"X612\",\n        \"X425\", \"X132\", \"X691\", \"X593\", \"X377\", \"X285\", \"X126\", \"X419\", \"X604\",\n        \"X84\", \"X138\", \"X413\", \"X291\", \"X40\", \"X123\", \"X81\", \"X853\", \"X854\",\n        \"X777\", \"X219\", \"X776\", \"X180\", \"X781\", \"X445\", \"X444\", \"X384\", \"X466\",\n        \"X95\", \"X583\", \"X272\", \"X137\", \"X533\", \"X758\", \"X279\", \"X297\",\n        \"X21\", \"X20\", \"X28\", \"X29\", \"X19\", \"X27\", \"X22\", \"X198\", \"X89\", \"X90\",\n        \"X98\", \"X96\", \"X97\", \"X383\", \"X427\", \"X451\", \"X283\",\n        \"X753\", \"X497\", \"X748\", \"X820\", \"X566\", \"X535\", \"X394\", \"X618\",\n        \"X429\", \"X381\", \"X387\", \"X890\", \"X752\", \"X375\", \"X68\", \"X152\",\n        \"X110\", \"X850\", \"X851\", \"X481\", \"X321\", \"X363\", \"X405\", \"X492\",\n        \"X888\", \"X421\", \"X333\", \"X817\", \"X586\", \"X292\", \"X344\", \"X532\",\n        \"bid_qty\", \"ask_qty\", \"buy_qty\", \"sell_qty\", \"volume\"\n    ]\n\n    LABEL_COLUMN = \"label\"\n    RANDOM_STATE = 42\n\ndef set_seed(seed=42):\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n\n# =========================\n# SAINT Components\n# =========================\nclass FeatureEmbedder(nn.Module):\n    \"\"\"Embeds features with positional encoding\"\"\"\n    def __init__(self, num_features, embed_dim, dropout=0.1):\n        super().__init__()\n        self.num_features = num_features\n        self.embed_dim = embed_dim\n        \n        # Feature embedding\n        self.feature_embed = nn.Linear(num_features, embed_dim)\n        \n        # Positional encoding for features\n        self.pos_encoding = nn.Parameter(torch.randn(1, num_features, embed_dim))\n        nn.init.normal_(self.pos_encoding, std=0.02)\n        \n        # Feature-wise projections\n        self.feature_projections = nn.ModuleList([\n            nn.Linear(1, embed_dim) for _ in range(num_features)\n        ])\n        \n        self.norm = nn.LayerNorm(embed_dim)\n        self.dropout = nn.Dropout(dropout)\n        \n    def forward(self, x):\n        batch_size = x.size(0)\n        \n        # Method 1: Global embedding\n        global_embed = self.feature_embed(x)\n        \n        # Method 2: Feature-wise embeddings\n        feature_embeds = []\n        for i in range(self.num_features):\n            feat = x[:, i:i+1]\n            embed = self.feature_projections[i](feat)\n            feature_embeds.append(embed)\n        \n        feature_stack = torch.stack(feature_embeds, dim=1)  # [batch, num_features, embed_dim]\n        \n        # Add positional encoding\n        feature_stack = feature_stack + self.pos_encoding\n        \n        # Combine global and feature-wise\n        combined = global_embed.unsqueeze(1) + feature_stack\n        \n        # Normalize and dropout\n        output = self.norm(combined)\n        output = self.dropout(output)\n        \n        return output, global_embed\n\nclass IntersampleAttention(nn.Module):\n    \"\"\"Fixed attention mechanism between different samples in a batch\"\"\"\n    def __init__(self, embed_dim, num_heads, dropout=0.1):\n        super().__init__()\n        self.embed_dim = embed_dim\n        self.num_heads = num_heads\n        \n        # Simple self-attention that processes batch dimension\n        self.attention = nn.MultiheadAttention(\n            embed_dim, num_heads, dropout=dropout, batch_first=True\n        )\n        self.norm = nn.LayerNorm(embed_dim)\n        self.dropout = nn.Dropout(dropout)\n        \n    def forward(self, x, mask=None):\n        # x shape: [batch, embed_dim]\n        batch_size = x.size(0)\n        \n        # Simple approach: treat batch as sequence for self-attention\n        # Add a dummy sequence dimension\n        x_seq = x.unsqueeze(1)  # [batch, 1, embed_dim]\n        \n        # Self-attention within the batch\n        attn_out, _ = self.attention(x_seq, x_seq, x_seq, attn_mask=mask)\n        \n        # Remove sequence dimension\n        attn_out = attn_out.squeeze(1)  # [batch, embed_dim]\n        \n        # Residual connection\n        output = self.norm(x + self.dropout(attn_out))\n        \n        return output\n\nclass MixtureOfExperts(nn.Module):\n    \"\"\"Mixture of Experts layer\"\"\"\n    def __init__(self, embed_dim, num_experts, expert_dim, dropout=0.1):\n        super().__init__()\n        self.num_experts = num_experts\n        \n        # Expert networks\n        self.experts = nn.ModuleList([\n            nn.Sequential(\n                nn.Linear(embed_dim, expert_dim),\n                nn.ReLU(),\n                nn.Dropout(dropout),\n                nn.Linear(expert_dim, embed_dim)\n            ) for _ in range(num_experts)\n        ])\n        \n        # Gating network\n        self.gate = nn.Sequential(\n            nn.Linear(embed_dim, num_experts),\n            nn.Softmax(dim=-1)\n        )\n        \n        # Top-k selection\n        self.top_k = min(2, num_experts)  # Use top-2 experts\n        \n    def forward(self, x):\n        # Get gating weights\n        gates = self.gate(x)  # [batch, num_experts]\n        \n        # Select top-k experts\n        top_k_gates, top_k_indices = torch.topk(gates, self.top_k, dim=-1)\n        \n        # Renormalize top-k gates\n        top_k_gates = top_k_gates / top_k_gates.sum(dim=-1, keepdim=True)\n        \n        # Apply selected experts\n        output = torch.zeros_like(x)\n        for i in range(self.top_k):\n            for j in range(self.num_experts):\n                mask = (top_k_indices[:, i] == j)\n                if mask.any():\n                    expert_out = self.experts[j](x[mask])\n                    gate_values = top_k_gates[mask, i].unsqueeze(-1)\n                    output[mask] += gate_values * expert_out\n        \n        return output\n\nclass HierarchicalProcessor(nn.Module):\n    \"\"\"Process features at multiple scales\"\"\"\n    def __init__(self, embed_dim, num_levels, dropout=0.1):\n        super().__init__()\n        self.num_levels = num_levels\n        \n        # Different pooling operations for hierarchy\n        self.pooling_ops = nn.ModuleList([\n            nn.AdaptiveAvgPool1d(2 ** (num_levels - i - 1))\n            for i in range(num_levels)\n        ])\n        \n        # Processing at each level\n        self.level_processors = nn.ModuleList([\n            nn.Sequential(\n                nn.Linear(embed_dim, embed_dim),\n                nn.LayerNorm(embed_dim),\n                nn.GELU(),\n                nn.Dropout(dropout)\n            ) for _ in range(num_levels)\n        ])\n        \n        # Fusion layer\n        self.fusion = nn.Linear(embed_dim * num_levels, embed_dim)\n        \n    def forward(self, x):\n        # x shape: [batch, num_features, embed_dim]\n        batch_size, num_features, embed_dim = x.size()\n        \n        # Process at different scales\n        level_outputs = []\n        for i in range(self.num_levels):\n            # Pool features\n            x_pooled = self.pooling_ops[i](x.transpose(1, 2)).transpose(1, 2)\n            \n            # Global pooling at this level\n            x_global = x_pooled.mean(dim=1)  # [batch, embed_dim]\n            \n            # Process\n            x_processed = self.level_processors[i](x_global)\n            level_outputs.append(x_processed)\n        \n        # Concatenate all levels\n        multi_scale = torch.cat(level_outputs, dim=-1)  # [batch, embed_dim * num_levels]\n        \n        # Fuse\n        output = self.fusion(multi_scale)\n        \n        return output\n\nclass SAINT(nn.Module):\n    \"\"\"SAINT: Self-Attention and INtersample Transformer\"\"\"\n    def __init__(self, config):\n        super().__init__()\n        \n        self.num_features = config['input_dim']\n        self.embed_dim = config['embed_dim']\n        self.num_layers = config['num_layers']\n        self.use_intersample = config.get('use_intersample_attention', True)\n        self.use_moe = config.get('use_mixture_of_experts', True)\n        \n        # Feature embedding\n        self.feature_embedder = FeatureEmbedder(\n            self.num_features,\n            self.embed_dim,\n            config['embed_dropout']\n        )\n        \n        # Self-attention layers\n        self.self_attention_layers = nn.ModuleList([\n            nn.TransformerEncoderLayer(\n                d_model=self.embed_dim,\n                nhead=config['num_heads'],\n                dim_feedforward=config['ff_dim'],\n                dropout=config['attention_dropout'],\n                activation='gelu',\n                batch_first=True\n            ) for _ in range(self.num_layers)\n        ])\n        \n        # Intersample attention (optional)\n        if self.use_intersample:\n            self.intersample_attention = IntersampleAttention(\n                self.embed_dim,\n                config['intersample_heads'],\n                config['intersample_dropout']\n            )\n        \n        # Mixture of experts (optional)\n        if self.use_moe:\n            self.moe = MixtureOfExperts(\n                self.embed_dim,\n                config['num_experts'],\n                config['expert_dim'],\n                config['moe_dropout']\n            )\n        \n        # Hierarchical processor\n        self.hierarchical = HierarchicalProcessor(\n            self.embed_dim,\n            config['num_hierarchy_levels'],\n            config['hierarchy_dropout']\n        )\n        \n        # Final prediction head\n        self.prediction_head = nn.Sequential(\n            nn.Linear(self.embed_dim, config['head_dim']),\n            nn.LayerNorm(config['head_dim']),\n            nn.GELU(),\n            nn.Dropout(config['head_dropout']),\n            nn.Linear(config['head_dim'], config['head_dim'] // 2),\n            nn.GELU(),\n            nn.Dropout(config['head_dropout']),\n            nn.Linear(config['head_dim'] // 2, 1)\n        )\n        \n    def forward(self, x, return_embeddings=False):\n        batch_size = x.size(0)\n        \n        # Embed features\n        feature_embeds, global_embed = self.feature_embedder(x)\n        \n        # Self-attention over features\n        attended_features = feature_embeds\n        \n        for layer in self.self_attention_layers:\n            attended_features = layer(attended_features)\n        \n        # Pool features\n        pooled_features = attended_features.mean(dim=1)  # [batch, embed_dim]\n        \n        # Intersample attention\n        if self.use_intersample and batch_size > 1:\n            pooled_features = self.intersample_attention(pooled_features)\n        \n        # Mixture of experts\n        if self.use_moe:\n            pooled_features = pooled_features + self.moe(pooled_features)\n        \n        # Hierarchical processing\n        hierarchical_features = self.hierarchical(attended_features)\n        \n        # Combine all representations\n        final_features = pooled_features + hierarchical_features + global_embed\n        \n        # Predictions\n        predictions = self.prediction_head(final_features)\n        \n        if return_embeddings:\n            return predictions, final_features\n        \n        return predictions\n\n# =========================\n# Feature Engineering\n# =========================\ndef add_features(df):\n    # Original features\n    df['bid_ask_interaction'] = df['bid_qty'] * df['ask_qty']\n    df['bid_buy_interaction'] = df['bid_qty'] * df['buy_qty']\n    df['bid_sell_interaction'] = df['bid_qty'] * df['sell_qty']\n    df['ask_buy_interaction'] = df['ask_qty'] * df['buy_qty']\n    df['ask_sell_interaction'] = df['ask_qty'] * df['sell_qty']\n\n    df['volume_weighted_sell'] = df['sell_qty'] * df['volume']\n    df['buy_sell_ratio'] = df['buy_qty'] / (df['sell_qty'] + 1e-10)\n    df['selling_pressure'] = df['sell_qty'] / (df['volume'] + 1e-10)\n    df['log_volume'] = np.log1p(df['volume'])\n\n    df['effective_spread_proxy'] = np.abs(df['buy_qty'] - df['sell_qty']) / (df['volume'] + 1e-10)\n    df['bid_ask_imbalance'] = (df['bid_qty'] - df['ask_qty']) / (df['bid_qty'] + df['ask_qty'] + 1e-10)\n    df['order_flow_imbalance'] = (df['buy_qty'] - df['sell_qty']) / (df['buy_qty'] + df['sell_qty'] + 1e-10)\n    df['liquidity_ratio'] = (df['bid_qty'] + df['ask_qty']) / (df['volume'] + 1e-10)\n    \n    # New microstructure features\n    df['net_order_flow'] = df['buy_qty'] - df['sell_qty']\n    df['normalized_net_flow'] = df['net_order_flow'] / (df['volume'] + 1e-10)\n    df['buying_pressure'] = df['buy_qty'] / (df['volume'] + 1e-10)\n    df['volume_weighted_buy'] = df['buy_qty'] * df['volume']\n    \n    df['total_depth'] = df['bid_qty'] + df['ask_qty']\n    df['depth_imbalance'] = (df['bid_qty'] - df['ask_qty']) / (df['total_depth'] + 1e-10)\n    df['relative_spread'] = np.abs(df['bid_qty'] - df['ask_qty']) / (df['total_depth'] + 1e-10)\n    df['log_depth'] = np.log1p(df['total_depth'])\n    \n    df['kyle_lambda'] = np.abs(df['net_order_flow']) / (df['volume'] + 1e-10)\n    df['flow_toxicity'] = np.abs(df['order_flow_imbalance']) * df['volume']\n    df['aggressive_flow_ratio'] = (df['buy_qty'] + df['sell_qty']) / (df['total_depth'] + 1e-10)\n    \n    df['volume_depth_ratio'] = df['volume'] / (df['total_depth'] + 1e-10)\n    df['activity_intensity'] = (df['buy_qty'] + df['sell_qty']) / (df['volume'] + 1e-10)\n    df['log_buy_qty'] = np.log1p(df['buy_qty'])\n    df['log_sell_qty'] = np.log1p(df['sell_qty'])\n    df['log_bid_qty'] = np.log1p(df['bid_qty'])\n    df['log_ask_qty'] = np.log1p(df['ask_qty'])\n    \n    df['realized_spread_proxy'] = 2 * np.abs(df['net_order_flow']) / (df['volume'] + 1e-10)\n    df['price_impact_proxy'] = df['net_order_flow'] / (df['total_depth'] + 1e-10)\n    df['quote_volatility_proxy'] = np.abs(df['depth_imbalance'])\n    \n    df['flow_depth_interaction'] = df['net_order_flow'] * df['total_depth']\n    df['imbalance_volume_interaction'] = df['order_flow_imbalance'] * df['volume']\n    df['depth_volume_interaction'] = df['total_depth'] * df['volume']\n    df['buy_sell_spread'] = np.abs(df['buy_qty'] - df['sell_qty'])\n    df['bid_ask_spread'] = np.abs(df['bid_qty'] - df['ask_qty'])\n    \n    df['trade_informativeness'] = df['net_order_flow'] / (df['bid_qty'] + df['ask_qty'] + 1e-10)\n    df['execution_shortfall_proxy'] = df['buy_sell_spread'] / (df['volume'] + 1e-10)\n    df['adverse_selection_proxy'] = df['net_order_flow'] / (df['total_depth'] + 1e-10) * df['volume']\n    \n    df['fill_probability'] = df['volume'] / (df['buy_qty'] + df['sell_qty'] + 1e-10)\n    df['execution_rate'] = (df['buy_qty'] + df['sell_qty']) / (df['total_depth'] + 1e-10)\n    df['market_efficiency'] = df['volume'] / (df['bid_ask_spread'] + 1e-10)\n    \n    df['sqrt_volume'] = np.sqrt(df['volume'])\n    df['sqrt_depth'] = np.sqrt(df['total_depth'])\n    df['volume_squared'] = df['volume'] ** 2\n    df['imbalance_squared'] = df['order_flow_imbalance'] ** 2\n    \n    df['bid_ratio'] = df['bid_qty'] / (df['total_depth'] + 1e-10)\n    df['ask_ratio'] = df['ask_qty'] / (df['total_depth'] + 1e-10)\n    df['buy_ratio'] = df['buy_qty'] / (df['buy_qty'] + df['sell_qty'] + 1e-10)\n    df['sell_ratio'] = df['sell_qty'] / (df['buy_qty'] + df['sell_qty'] + 1e-10)\n    \n    df['liquidity_consumption'] = (df['buy_qty'] + df['sell_qty']) / (df['total_depth'] + 1e-10)\n    df['market_stress'] = df['volume'] / (df['total_depth'] + 1e-10) * np.abs(df['order_flow_imbalance'])\n    df['depth_depletion'] = df['volume'] / (df['bid_qty'] + df['ask_qty'] + 1e-10)\n    \n    df['net_buying_ratio'] = df['net_order_flow'] / (df['volume'] + 1e-10)\n    df['directional_volume'] = df['net_order_flow'] * np.log1p(df['volume'])\n    df['signed_volume'] = np.sign(df['net_order_flow']) * df['volume']\n    \n    # Replace infinities and NaNs\n    df = df.replace([np.inf, -np.inf], 0).fillna(0)\n    \n    return df\n\ndef load_data():\n    # Load data with all features available\n    all_features = Config.SAINT_FEATURES\n    train_df = pd.read_parquet(Config.TRAIN_PATH, columns=all_features + [Config.LABEL_COLUMN])\n    test_df = pd.read_parquet(Config.TEST_PATH, columns=all_features)\n    submission_df = pd.read_csv(Config.SUBMISSION_PATH)\n    print(f\"Loaded data - Train: {train_df.shape}, Test: {test_df.shape}, Submission: {submission_df.shape}\")\n\n    # Add features\n    train_df = add_features(train_df)\n    test_df = add_features(test_df)\n\n    return train_df.reset_index(drop=True), test_df.reset_index(drop=True), submission_df\n\n# =========================\n# SAINT Training\n# =========================\ndef train_saint_model(model, train_loader, val_loader, config, device):\n    \"\"\"Train SAINT model\"\"\"\n    \n    # Loss and optimizer\n    criterion = nn.HuberLoss(delta=config.get('huber_delta', 1.0))\n    optimizer = optim.AdamW(\n        model.parameters(),\n        lr=config['learning_rate'],\n        weight_decay=config['weight_decay']\n    )\n    \n    # Scheduler\n    scheduler = ReduceLROnPlateau(optimizer, mode='max', factor=0.5, patience=5, min_lr=1e-6)\n    \n    # Training settings\n    best_val_pearson = -np.inf\n    patience_counter = 0\n    patience = config.get('patience', 10)\n    num_epochs = config.get('num_epochs', 50)\n    \n    for epoch in range(num_epochs):\n        # Training\n        model.train()\n        train_loss = 0.0\n        train_batches = 0\n        \n        progress_bar = tqdm(train_loader, desc=f\"Epoch {epoch+1}/{num_epochs}\")\n        for inputs, targets in progress_bar:\n            inputs, targets = inputs.to(device), targets.to(device)\n            \n            # Add noise\n            if config.get('noise_factor', 0) > 0:\n                noise = torch.randn_like(inputs) * config['noise_factor']\n                inputs = inputs + noise\n            \n            optimizer.zero_grad()\n            \n            outputs = model(inputs)\n            loss = criterion(outputs, targets)\n            \n            loss.backward()\n            torch.nn.utils.clip_grad_norm_(model.parameters(), config.get('grad_clip', 1.0))\n            optimizer.step()\n            \n            train_loss += loss.item()\n            train_batches += 1\n            \n            progress_bar.set_postfix({'loss': f'{loss.item():.4f}'})\n        \n        # Validation\n        model.eval()\n        val_loss = 0.0\n        val_preds = []\n        val_targets = []\n        \n        with torch.no_grad():\n            for inputs, targets in val_loader:\n                inputs, targets = inputs.to(device), targets.to(device)\n                \n                outputs = model(inputs)\n                loss = criterion(outputs, targets)\n                \n                val_loss += loss.item()\n                val_preds.extend(outputs.cpu().numpy().flatten())\n                val_targets.extend(targets.cpu().numpy().flatten())\n        \n        # Metrics\n        avg_train_loss = train_loss / train_batches\n        avg_val_loss = val_loss / len(val_loader)\n        val_pearson = pearsonr(val_targets, val_preds)[0]\n        val_spearman = spearmanr(val_targets, val_preds)[0]\n        \n        print(f\"\\nTrain Loss: {avg_train_loss:.4f}, Val Loss: {avg_val_loss:.4f}\")\n        print(f\"Val Pearson: {val_pearson:.4f}, Val Spearman: {val_spearman:.4f}\")\n        \n        # Update scheduler\n        scheduler.step(val_pearson)\n        \n        # Save best model\n        if val_pearson > best_val_pearson:\n            best_val_pearson = val_pearson\n            patience_counter = 0\n            torch.save(model.state_dict(), \"best_saint.pt\")\n            print(f\"✅ New best model saved! Pearson: {best_val_pearson:.4f}\")\n        else:\n            patience_counter += 1\n            if patience_counter >= patience:\n                print(f\"Early stopping triggered after {epoch+1} epochs\")\n                break\n    \n    # Load best model\n    model.load_state_dict(torch.load(\"best_saint.pt\"))\n    \n    return model, best_val_pearson\n\ndef train_saint(train_df, test_df):\n    print(\"\\n=== Training SAINT Model ===\")\n    \n    # Set seed\n    set_seed(42)\n    \n    # Get SAINT features\n    saint_features = Config.SAINT_FEATURES.copy()\n    \n    # Add engineered features to SAINT features\n    engineered_features = [\n        \"log_volume\", 'bid_ask_interaction', 'net_order_flow', 'normalized_net_flow',\n        'buying_pressure', 'total_depth', 'depth_imbalance', 'kyle_lambda', \n        'aggressive_flow_ratio', 'volume_depth_ratio', 'log_buy_qty', 'log_sell_qty',\n        'price_impact_proxy', 'market_stress', 'liquidity_consumption'\n    ]\n    \n    all_saint_features = saint_features + engineered_features\n    all_saint_features = list(set(all_saint_features))  # Remove duplicates\n    \n    # Ensure all features exist\n    all_saint_features = [f for f in all_saint_features if f in train_df.columns]\n    \n    print(f\"Using {len(all_saint_features)} features for SAINT\")\n    \n    # Use recent data (last 85%)\n    train_size = int(0.85 * len(train_df))\n    train_data = train_df.iloc[-train_size:].reset_index(drop=True)\n    \n    # Split for validation\n    split_idx = int(0.8 * len(train_data))\n    train_split = train_data[:split_idx].copy()\n    val_split = train_data[split_idx:].copy()\n    \n    y_train = train_split[Config.LABEL_COLUMN].values\n    y_val = val_split[Config.LABEL_COLUMN].values\n    \n    X_train = train_split[all_saint_features].values\n    X_val = val_split[all_saint_features].values\n    \n    # Feature selection\n    print(\"\\nSelecting features...\")\n    selector = SelectKBest(score_func=mutual_info_regression, k=min(105, len(all_saint_features)))\n    X_train_selected = selector.fit_transform(X_train, y_train)\n    X_val_selected = selector.transform(X_val)\n    \n    selected_features = [all_saint_features[i] for i in selector.get_support(indices=True)]\n    print(f\"Selected {len(selected_features)} features\")\n    \n    # Transform data\n    print(\"\\nTransforming data...\")\n    transformer = QuantileTransformer(output_distribution='normal', random_state=42)\n    X_train_transformed = transformer.fit_transform(X_train_selected)\n    X_val_transformed = transformer.transform(X_val_selected)\n    \n    # Create data loaders\n    train_dataset = TensorDataset(\n        torch.tensor(X_train_transformed, dtype=torch.float32),\n        torch.tensor(y_train, dtype=torch.float32).unsqueeze(1)\n    )\n    val_dataset = TensorDataset(\n        torch.tensor(X_val_transformed, dtype=torch.float32),\n        torch.tensor(y_val, dtype=torch.float32).unsqueeze(1)\n    )\n    \n    # Try different hyperparameter configurations\n    configs = [\n        # Config 1: Basic SAINT\n        {\n            'config_name': 'basic_saint',\n            'embed_dim': 128,\n            'num_layers': 3,\n            'num_heads': 4,\n            'ff_dim': 512,\n            'embed_dropout': 0.2,\n            'attention_dropout': 0.2,\n            'use_intersample_attention': True,\n            'intersample_heads': 4,\n            'intersample_dropout': 0.2,\n            'use_mixture_of_experts': False,\n            'num_experts': 4,\n            'expert_dim': 128,\n            'moe_dropout': 0.3,\n            'num_hierarchy_levels': 3,\n            'hierarchy_dropout': 0.2,\n            'head_dim': 256,\n            'head_dropout': 0.3,\n            'learning_rate': 0.001,\n            'weight_decay': 0.01,\n            'batch_size': 512,\n            'huber_delta': 1.0,\n            'noise_factor': 0.01,\n            'grad_clip': 1.0,\n            'num_epochs': 30,\n            'patience': 10,\n            'input_dim': X_train_transformed.shape[1]\n        },\n        # Config 2: Medium SAINT\n        {\n            'config_name': 'medium_saint',\n            'embed_dim': 192,\n            'num_layers': 4,\n            'num_heads': 6,\n            'ff_dim': 768,\n            'embed_dropout': 0.25,\n            'attention_dropout': 0.25,\n            'use_intersample_attention': True,\n            'intersample_heads': 6,\n            'intersample_dropout': 0.25,\n            'use_mixture_of_experts': True,\n            'num_experts': 4,\n            'expert_dim': 192,\n            'moe_dropout': 0.3,\n            'num_hierarchy_levels': 3,\n            'hierarchy_dropout': 0.25,\n            'head_dim': 384,\n            'head_dropout': 0.35,\n            'learning_rate': 0.0008,\n            'weight_decay': 0.008,\n            'batch_size': 256,\n            'huber_delta': 0.8,\n            'noise_factor': 0.015,\n            'grad_clip': 1.5,\n            'num_epochs': 25,\n            'patience': 8,\n            'input_dim': X_train_transformed.shape[1]\n        },\n        # Config 3: Light SAINT\n        {\n            'config_name': 'light_saint',\n            'embed_dim': 96,\n            'num_layers': 2,\n            'num_heads': 3,\n            'ff_dim': 384,\n            'embed_dropout': 0.15,\n            'attention_dropout': 0.15,\n            'use_intersample_attention': False,\n            'intersample_heads': 3,\n            'intersample_dropout': 0.15,\n            'use_mixture_of_experts': False,\n            'num_experts': 3,\n            'expert_dim': 96,\n            'moe_dropout': 0.2,\n            'num_hierarchy_levels': 2,\n            'hierarchy_dropout': 0.15,\n            'head_dim': 192,\n            'head_dropout': 0.25,\n            'learning_rate': 0.002,\n            'weight_decay': 0.005,\n            'batch_size': 1024,\n            'huber_delta': 1.0,\n            'noise_factor': 0.008,\n            'grad_clip': 1.0,\n            'num_epochs': 35,\n            'patience': 12,\n            'input_dim': X_train_transformed.shape[1]\n        }\n    ]\n    \n    # Train ensemble of models\n    ensemble_models = []\n    ensemble_scores = []\n    \n    for config in configs:\n        print(f\"\\n=== Training SAINT {config['config_name']} ===\")\n        \n        # Create data loaders\n        train_loader = DataLoader(train_dataset, batch_size=config['batch_size'], shuffle=True)\n        val_loader = DataLoader(val_dataset, batch_size=1024, shuffle=False)\n        \n        # Create model\n        model = SAINT(config).to(device)\n        print(f\"Model parameters: {sum(p.numel() for p in model.parameters()):,}\")\n        \n        # Train\n        model, best_score = train_saint_model(model, train_loader, val_loader, config, device)\n        \n        ensemble_models.append(model)\n        ensemble_scores.append(best_score)\n        \n        print(f\"Model {config['config_name']} best validation Pearson: {best_score:.4f}\")\n    \n    # Make test predictions\n    print(\"\\n=== Making SAINT Test Predictions ===\")\n    \n    # Transform test data\n    X_test = test_df[all_saint_features].values\n    X_test_selected = selector.transform(X_test)\n    X_test_transformed = transformer.transform(X_test_selected)\n    \n    # Make predictions with each model\n    all_predictions = []\n    \n    for model in ensemble_models:\n        model.eval()\n        test_dataset = TensorDataset(torch.tensor(X_test_transformed, dtype=torch.float32))\n        test_loader = DataLoader(test_dataset, batch_size=2048, shuffle=False)\n        \n        predictions = []\n        with torch.no_grad():\n            for (inputs,) in test_loader:\n                inputs = inputs.to(device)\n                outputs = model(inputs)\n                predictions.extend(outputs.cpu().numpy().flatten())\n        \n        all_predictions.append(np.array(predictions))\n    \n    # Ensemble predictions\n    weights = np.array(ensemble_scores)\n    weights = weights / weights.sum()\n    \n    final_predictions = np.zeros_like(all_predictions[0])\n    for pred, weight in zip(all_predictions, weights):\n        final_predictions += weight * pred\n    \n    # Post-processing\n    pred_mean = train_df[Config.LABEL_COLUMN].mean()\n    pred_std = train_df[Config.LABEL_COLUMN].std()\n    final_predictions = np.clip(\n        final_predictions,\n        pred_mean - 4 * pred_std,\n        pred_mean + 4 * pred_std\n    )\n    \n    print(f\"\\nSAINT ensemble weights: {weights}\")\n    print(f\"SAINT prediction stats - Mean: {final_predictions.mean():.6f}, Std: {final_predictions.std():.6f}\")\n    \n    return final_predictions\n\n# =========================\n# Main Execution\n# =========================\nif __name__ == \"__main__\":\n    # Load data\n    train_df, test_df, submission_df = load_data()\n    \n    # Train SAINT model\n    saint_predictions = train_saint(train_df, test_df)\n    \n    # Save SAINT submission\n    saint_submission = submission_df.copy()\n    saint_submission[\"prediction\"] = saint_predictions\n    saint_submission.to_csv(\"submission_saint_fixed.csv\", index=False)\n    print(f\"\\nSaved: submission_saint_fixed.csv\")\n    \n    # Show sample predictions\n    print(\"\\nSample predictions (first 10 rows):\")\n    print(saint_submission[['ID', 'prediction']].head(10))","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}