{"cells":[{"cell_type":"markdown","id":"8e7c9352","metadata":{},"source":"# Ariel Data Challenge 2025 - Baseline + Transformer Model"},{"cell_type":"markdown","id":"2feb1199","metadata":{},"source":"## 📥 1. Import & Setup"},{"cell_type":"code","execution_count":null,"id":"04ad4731","metadata":{},"outputs":[],"source":"import numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom sklearn.model_selection import train_test_split\nimport matplotlib.pyplot as plt\nimport os"},{"cell_type":"markdown","id":"c69f6b1d","metadata":{},"source":"## 📊 2. Load Data"},{"cell_type":"markdown","id":"ff2f1525","metadata":{},"source":"# Replace with real paths or uploaded data\ntrain_df = pd.read_csv(\"data/train.csv\")\ntest_df = pd.read_csv(\"data/test.csv\")\n\nplanet_ids_test = test_df['planet_id'].values\nX_test = test_df.drop(\"planet_id\", axis=1).values.astype(np.float32)\n\nX = train_df.drop(columns=[\"planet_id\", \"target_spectrum\"]).values.astype(np.float32)\ny = train_df[\"target_spectrum\"].values.reshape(-1, 283).astype(np.float32)\n\nX_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.1, random_state=42)"},{"cell_type":"markdown","id":"9cda1dff","metadata":{},"source":"## 🧪 3. Dataset & Dataloader"},{"cell_type":"code","execution_count":null,"id":"25e78aa2","metadata":{},"outputs":[],"source":"class SpectraDataset(Dataset):\n    def __init__(self, X, y=None):\n        self.X = X\n        self.y = y\n\n    def __len__(self):\n        return len(self.X)\n\n    def __getitem__(self, idx):\n        x = self.X[idx]\n        if self.y is not None:\n            return torch.tensor(x), torch.tensor(self.y[idx])\n        return torch.tensor(x)\n\ntrain_ds = SpectraDataset(X_train, y_train)\nval_ds = SpectraDataset(X_val, y_val)\ntest_ds = SpectraDataset(X_test)\n\ntrain_loader = DataLoader(train_ds, batch_size=64, shuffle=True)\nval_loader = DataLoader(val_ds, batch_size=64)\ntest_loader = DataLoader(test_ds, batch_size=64)"},{"cell_type":"markdown","id":"41318e35","metadata":{},"source":"## 🧠 4. Model: SimpleSpectraDenoiser"},{"cell_type":"code","execution_count":null,"id":"da37f3af","metadata":{},"outputs":[],"source":"class SimpleSpectraDenoiser(nn.Module):\n    def __init__(self, input_dim=283, hidden_dim=512):\n        super().__init__()\n        self.shared = nn.Sequential(\n            nn.Linear(input_dim, hidden_dim),\n            nn.ReLU(),\n            nn.Linear(hidden_dim, hidden_dim),\n            nn.ReLU()\n        )\n        self.mean_head = nn.Linear(hidden_dim, input_dim)\n        self.std_head = nn.Linear(hidden_dim, input_dim)\n\n    def forward(self, x):\n        h = self.shared(x)\n        mean = self.mean_head(h)\n        std = torch.exp(self.std_head(h))\n        return mean, std"},{"cell_type":"markdown","id":"1059e839","metadata":{},"source":"## 🧠 4B. Transformer-based Model (Advanced)"},{"cell_type":"code","execution_count":null,"id":"b6d9a794","metadata":{},"outputs":[],"source":"import math\n\nclass PositionalEncoding(nn.Module):\n    def __init__(self, d_model, max_len=283):\n        super().__init__()\n        pe = torch.zeros(max_len, d_model)\n        position = torch.arange(0, max_len, dtype=torch.float32).unsqueeze(1)\n        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))\n        pe[:, 0::2] = torch.sin(position * div_term)\n        pe[:, 1::2] = torch.cos(position * div_term)\n        self.pe = pe.unsqueeze(0)\n\n    def forward(self, x):\n        return x + self.pe[:, :x.size(1)].to(x.device)\n\nclass SpectraTransformer(nn.Module):\n    def __init__(self, input_dim=283, d_model=64, nhead=8, num_layers=4):\n        super().__init__()\n        self.input_linear = nn.Linear(1, d_model)\n        self.pos_encoder = PositionalEncoding(d_model)\n        encoder_layers = nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward=256, batch_first=True)\n        self.transformer_encoder = nn.TransformerEncoder(encoder_layers, num_layers)\n        self.output_mean = nn.Linear(d_model, 1)\n        self.output_std = nn.Linear(d_model, 1)\n\n    def forward(self, x):\n        x = x.unsqueeze(-1)\n        x = self.input_linear(x)\n        x = self.pos_encoder(x)\n        h = self.transformer_encoder(x)\n        mean = self.output_mean(h).squeeze(-1)\n        std = torch.exp(self.output_std(h).squeeze(-1))\n        return mean, std"},{"cell_type":"markdown","id":"6038241d","metadata":{},"source":"## 🏋️‍♂️ 5. Training Function"},{"cell_type":"code","execution_count":null,"id":"63003fb6","metadata":{},"outputs":[],"source":"def train_model(model, train_loader, val_loader, epochs=20, lr=1e-3):\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    model = model.to(device)\n    optimizer = torch.optim.Adam(model.parameters(), lr=lr)\n    criterion = nn.MSELoss()\n\n    best_val_loss = float(\"inf\")\n    for epoch in range(epochs):\n        model.train()\n        total_loss = 0\n        for xb, yb in train_loader:\n            xb, yb = xb.to(device), yb.to(device)\n            pred_mean, _ = model(xb)\n            loss = criterion(pred_mean, yb)\n            optimizer.zero_grad()\n            loss.backward()\n            optimizer.step()\n            total_loss += loss.item()\n\n        model.eval()\n        val_loss = 0\n        with torch.no_grad():\n            for xb, yb in val_loader:\n                xb, yb = xb.to(device), yb.to(device)\n                pred_mean, _ = model(xb)\n                loss = criterion(pred_mean, yb)\n                val_loss += loss.item()\n\n        print(f\"Epoch {epoch+1}: Train Loss = {total_loss/len(train_loader):.4f}, Val Loss = {val_loss/len(val_loader):.4f}\")\n        if val_loss < best_val_loss:\n            best_val_loss = val_loss\n            torch.save(model.state_dict(), \"best_model.pth\")"},{"cell_type":"markdown","id":"915858c9","metadata":{},"source":"## 🚀 6. Train Transformer"},{"cell_type":"code","execution_count":null,"id":"3066b0dc","metadata":{},"outputs":[],"source":"model = SpectraTransformer()\ntrain_model(model, train_loader, val_loader)"},{"cell_type":"markdown","id":"50261db7","metadata":{},"source":"## 📤 7. Inference & Submission"},{"cell_type":"code","execution_count":null,"id":"e61bb879","metadata":{},"outputs":[],"source":"model.load_state_dict(torch.load(\"best_model.pth\"))\nmodel.eval()\n\nall_means, all_stds = [], []\nwith torch.no_grad():\n    for xb in test_loader:\n        mean, std = model(xb)\n        all_means.append(mean.numpy())\n        all_stds.append(std.numpy())\n\nmean_array = np.concatenate(all_means, axis=0)\nstd_array = np.concatenate(all_stds, axis=0)\n\nsubmission = pd.DataFrame()\nsubmission[\"planet_id\"] = planet_ids_test\nsubmission = pd.concat([submission, pd.DataFrame(mean_array), pd.DataFrame(std_array)], axis=1)\nsubmission.to_csv(\"submission.csv\", index=False)"}],"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":"none","dataSources":[{"sourceId":12351634,"sourceType":"datasetVersion","datasetId":7787005}],"dockerImageVersionId":31040,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat":4,"nbformat_minor":4}