{"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":101849,"databundleVersionId":13093295,"sourceType":"competition"}],"dockerImageVersionId":31090,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"### This code includes data loading, preprocessing, model architecture, training, and submission generation.","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nimport pytorch_lightning as pl\nfrom pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint\nimport glob\nimport os\nfrom sklearn.preprocessing import StandardScaler\nimport warnings\nwarnings.filterwarnings('ignore')\n\n# Set random seeds for reproducibility\npl.seed_everything(42)\n\n# Configuration\nclass Config:\n    # Data paths\n    train_data_path = \"/kaggle/input/ariel-data-challenge-2025/train\"\n    test_data_path = \"/kaggle/input/ariel-data-challenge-2025/test\"\n    metadata_path = \"/kaggle/input/ariel-data-challenge-2025\"\n    \n    # Model parameters\n    batch_size = 2\n    learning_rate = 1e-4\n    num_epochs = 3\n    hidden_dim = 32\n    num_layers = 1\n    dropout = 0.1\n    \n    # Instrument parameters\n    airs_shape = (32, 356)\n    fgs1_shape = (32, 32)\n    \n    # Training parameters\n    num_workers = 1\n    accelerator = 'gpu' if torch.cuda.is_available() else 'cpu'\n    devices = 1\n\nconfig = Config()\n\n# Load metadata\ndef load_metadata():\n    train_star_info = pd.read_csv(\"/kaggle/input/ariel-data-challenge-2025/train_star_info.csv\")\n    train_spectra = pd.read_csv(\"/kaggle/input/ariel-data-challenge-2025/train.csv\")\n    wavelengths = pd.read_csv(\"/kaggle/input/ariel-data-challenge-2025/wavelengths.csv\")\n    axis_info = pd.read_parquet(\"/kaggle/input/ariel-data-challenge-2025/axis_info.parquet\")\n    adc_info = pd.read_csv(\"/kaggle/input/ariel-data-challenge-2025/adc_info.csv\")\n    \n    return train_star_info, train_spectra, wavelengths, axis_info, adc_info\n\n# Data preprocessing functions\ndef preprocess_images(data, gain, offset, instrument):\n    data = data.astype(np.float64) / gain + offset\n    return data\n\ndef extract_light_curve(data, instrument):\n    light_curve = np.nansum(data, axis=(1, 2))\n    return light_curve\n\n# Dataset class\nclass ArielDataset(Dataset):\n    def __init__(self, data_path, star_info, spectra, adc_info, is_train=True, max_samples=2):\n        self.data_path = data_path\n        self.star_info = star_info\n        self.spectra = spectra\n        self.adc_info = adc_info\n        self.is_train = is_train\n        self.planet_ids = star_info['planet_id'].values\n        \n        self.file_paths = []\n        for planet_id in self.planet_ids:\n            planet_id_str = str(planet_id)\n            planet_path = os.path.join(data_path, planet_id_str)\n            \n            if not os.path.exists(planet_path):\n                continue\n                \n            airs_files = glob.glob(os.path.join(planet_path, \"AIRS-CH0_signal_*.parquet\"))\n            fgs1_files = glob.glob(os.path.join(planet_path, \"FGS1_signal_*.parquet\"))\n            \n            for airs_file in airs_files:\n                self.file_paths.append((planet_id_str, 'AIRS-CH0', airs_file))\n            \n            for fgs1_file in fgs1_files:\n                self.file_paths.append((planet_id_str, 'FGS1', fgs1_file))\n            \n            if len(self.file_paths) >= max_samples:\n                break\n        \n        print(f\"Found {len(self.file_paths)} observation files\")\n    \n    def __len__(self):\n        return len(self.file_paths)\n    \n    def __getitem__(self, idx):\n        planet_id, instrument, file_path = self.file_paths[idx]\n        \n        try:\n            signal_data = pd.read_parquet(file_path).values\n            \n            if instrument == 'AIRS-CH0':\n                gain = self.adc_info['AIRS-CH0_adc_gain'].values[0]\n                offset = self.adc_info['AIRS-CH0_adc_offset'].values[0]\n            else:\n                gain = self.adc_info['FGS1_adc_gain'].values[0]\n                offset = self.adc_info['FGS1_adc_offset'].values[0]\n            \n            signal_data = preprocess_images(signal_data, gain, offset, instrument)\n            \n            if instrument == 'AIRS-CH0':\n                signal_data = signal_data.reshape(-1, config.airs_shape[0], config.airs_shape[1])\n            else:\n                signal_data = signal_data.reshape(-1, config.fgs1_shape[0], config.fgs1_shape[1])\n            \n            light_curve = extract_light_curve(signal_data, instrument)\n            \n            # Downsample to fixed length\n            target_length = 256  # Fixed length for all sequences\n            if len(light_curve) > target_length:\n                step = len(light_curve) // target_length\n                light_curve = light_curve[::step]\n                if len(light_curve) > target_length:\n                    light_curve = light_curve[:target_length]\n            elif len(light_curve) < target_length:\n                # Pad with zeros if shorter\n                pad_length = target_length - len(light_curve)\n                light_curve = np.concatenate([light_curve, np.zeros(pad_length)])\n            \n            star_params = self.star_info[self.star_info['planet_id'] == int(planet_id)].drop('planet_id', axis=1).values[0]\n            \n            if self.is_train:\n                target = self.spectra[self.spectra['planet_id'] == int(planet_id)].drop('planet_id', axis=1).values[0]\n            else:\n                target = np.zeros(283)\n                \n        except Exception as e:\n            print(f\"Error loading data for {planet_id}: {e}\")\n            light_curve = np.zeros(256)\n            star_params = np.zeros(8)\n            target = np.zeros(283)\n        \n        light_curve = torch.FloatTensor(light_curve)\n        star_params = torch.FloatTensor(star_params)\n        target = torch.FloatTensor(target)\n        \n        return {\n            'light_curve': light_curve,\n            'star_params': star_params,\n            'target': target,\n            'planet_id': planet_id,\n            'instrument': instrument\n        }\n\n# Custom collate function\ndef collate_fn(batch):\n    light_curves = [item['light_curve'] for item in batch]\n    star_params = torch.stack([item['star_params'] for item in batch])\n    targets = torch.stack([item['target'] for item in batch])\n    planet_ids = [item['planet_id'] for item in batch]\n    instruments = [item['instrument'] for item in batch]\n    \n    # All sequences should be same length now (256)\n    light_curves = torch.stack(light_curves)\n    attention_masks = torch.ones_like(light_curves)\n    \n    return {\n        'light_curves': light_curves,\n        'attention_masks': attention_masks,\n        'star_params': star_params,\n        'targets': targets,\n        'planet_ids': planet_ids,\n        'instruments': instruments\n    }\n\n# Simplified Model architecture\nclass ArielModel(nn.Module):\n    def __init__(self, input_dim, star_param_dim, output_dim, hidden_dim=32, dropout=0.1):\n        super(ArielModel, self).__init__()\n        \n        # Light curve encoder (1D CNN)\n        self.light_curve_encoder = nn.Sequential(\n            nn.Conv1d(1, 8, kernel_size=5, stride=2, padding=2),\n            nn.ReLU(),\n            nn.BatchNorm1d(8),\n            nn.Conv1d(8, 16, kernel_size=5, stride=2, padding=2),\n            nn.ReLU(),\n            nn.BatchNorm1d(16),\n            nn.Conv1d(16, 32, kernel_size=5, stride=2, padding=2),\n            nn.ReLU(),\n            nn.BatchNorm1d(32),\n            nn.AdaptiveAvgPool1d(16)  # Fixed size output\n        )\n        \n        # Star parameter encoder\n        self.star_encoder = nn.Sequential(\n            nn.Linear(star_param_dim, hidden_dim),\n            nn.ReLU(),\n            nn.Dropout(dropout)\n        )\n        \n        # Combined feature processing\n        self.combined_processor = nn.Sequential(\n            nn.Linear(32 * 16 + hidden_dim, hidden_dim),\n            nn.ReLU(),\n            nn.Dropout(dropout),\n            nn.Linear(hidden_dim, hidden_dim),\n            nn.ReLU(),\n            nn.Dropout(dropout),\n            nn.Linear(hidden_dim, output_dim)\n        )\n        \n        # Uncertainty estimation head\n        self.uncertainty_head = nn.Sequential(\n            nn.Linear(32 * 16 + hidden_dim, hidden_dim),\n            nn.ReLU(),\n            nn.Dropout(dropout),\n            nn.Linear(hidden_dim, output_dim),\n            nn.Softplus()\n        )\n    \n    def forward(self, light_curves, attention_masks, star_params):\n        batch_size, seq_len = light_curves.shape\n        \n        # Reshape for CNN\n        light_curves = light_curves.unsqueeze(1)\n        \n        # Encode light curves with CNN\n        cnn_features = self.light_curve_encoder(light_curves)\n        cnn_features = cnn_features.view(batch_size, -1)  # Flatten\n        \n        # Encode star parameters\n        star_features = self.star_encoder(star_params)\n        \n        # Combine features\n        combined_features = torch.cat([cnn_features, star_features], dim=1)\n        \n        # Predict spectrum\n        spectrum_pred = self.combined_processor(combined_features)\n        \n        # Predict uncertainty\n        uncertainty = self.uncertainty_head(combined_features)\n        \n        return spectrum_pred, uncertainty\n\n# PyTorch Lightning Module\nclass ArielLightningModule(pl.LightningModule):\n    def __init__(self, model, learning_rate=1e-4):\n        super().__init__()\n        self.model = model\n        self.learning_rate = learning_rate\n        self.loss_fn = self.gaussian_log_likelihood_loss\n    \n    def gaussian_log_likelihood_loss(self, pred_mean, pred_std, target):\n        variance = pred_std ** 2 + 1e-6\n        log_likelihood = -0.5 * (torch.log(2 * torch.tensor(np.pi)) + torch.log(variance) + (target - pred_mean) ** 2 / variance)\n        return -log_likelihood.mean()\n    \n    def training_step(self, batch, batch_idx):\n        light_curves = batch['light_curves']\n        attention_masks = batch['attention_masks']\n        star_params = batch['star_params']\n        targets = batch['targets']\n        \n        pred_mean, pred_std = self.model(light_curves, attention_masks, star_params)\n        loss = self.loss_fn(pred_mean, pred_std, targets)\n        \n        self.log('train_loss', loss, prog_bar=True)\n        return loss\n    \n    def validation_step(self, batch, batch_idx):\n        light_curves = batch['light_curves']\n        attention_masks = batch['attention_masks']\n        star_params = batch['star_params']\n        targets = batch['targets']\n        \n        pred_mean, pred_std = self.model(light_curves, attention_masks, star_params)\n        loss = self.loss_fn(pred_mean, pred_std, targets)\n        \n        self.log('val_loss', loss, prog_bar=True)\n        return loss\n    \n    def configure_optimizers(self):\n        return optim.Adam(self.parameters(), lr=self.learning_rate)\n\n# Main training function\ndef train_model():\n    train_star_info, train_spectra, wavelengths, axis_info, adc_info = load_metadata()\n    \n    train_dataset = ArielDataset(config.train_data_path, train_star_info, train_spectra, adc_info, is_train=True, max_samples=2)\n    \n    train_loader = DataLoader(\n        train_dataset, \n        batch_size=config.batch_size, \n        shuffle=True, \n        num_workers=config.num_workers,\n        collate_fn=collate_fn\n    )\n    \n    val_loader = DataLoader(\n        train_dataset, \n        batch_size=config.batch_size, \n        shuffle=False, \n        num_workers=config.num_workers,\n        collate_fn=collate_fn\n    )\n    \n    input_dim = 1\n    star_param_dim = train_star_info.drop('planet_id', axis=1).shape[1]\n    output_dim = train_spectra.drop('planet_id', axis=1).shape[1]\n    \n    model = ArielModel(input_dim, star_param_dim, output_dim, \n                      hidden_dim=config.hidden_dim, \n                      dropout=config.dropout)\n    \n    lightning_model = ArielLightningModule(model, learning_rate=config.learning_rate)\n    \n    trainer = pl.Trainer(\n        max_epochs=config.num_epochs,\n        accelerator=config.accelerator,\n        devices=config.devices,\n        log_every_n_steps=1,\n        enable_progress_bar=True\n    )\n    \n    trainer.fit(lightning_model, train_loader, val_loader)\n    \n    return lightning_model, trainer\n\n# Prediction function\ndef predict(model, data_loader):\n    model.eval()\n    all_preds = []\n    all_uncertainties = []\n    all_planet_ids = []\n    \n    with torch.no_grad():\n        for batch in data_loader:\n            light_curves = batch['light_curves'].to(model.device)\n            attention_masks = batch['attention_masks'].to(model.device)\n            star_params = batch['star_params'].to(model.device)\n            \n            pred_mean, pred_std = model(light_curves, attention_masks, star_params)\n            \n            all_preds.append(pred_mean.cpu().numpy())\n            all_uncertainties.append(pred_std.cpu().numpy())\n            all_planet_ids.extend(batch['planet_ids'])\n    \n    unique_planet_ids = list(set(all_planet_ids))\n    predictions = {pid: {'mean': [], 'std': []} for pid in unique_planet_ids}\n    \n    for i, pid in enumerate(all_planet_ids):\n        predictions[pid]['mean'].append(all_preds[i])\n        predictions[pid]['std'].append(all_uncertainties[i])\n    \n    final_predictions = {}\n    for pid in unique_planet_ids:\n        mean_vals = np.mean(predictions[pid]['mean'], axis=0)\n        std_vals = np.mean(predictions[pid]['std'], axis=0)\n        final_predictions[pid] = {'mean': mean_vals, 'std': std_vals}\n    \n    return final_predictions\n\n# Create submission file\ndef create_submission(predictions, sample_submission_path, output_path='/kaggle/working/submission.csv'):\n    sample_submission = pd.read_csv(sample_submission_path)\n    \n    submission_data = []\n    for planet_id in sample_submission['planet_id']:\n        planet_id_str = str(planet_id)\n        if planet_id_str in predictions:\n            pred = predictions[planet_id_str]\n            row = [planet_id] + list(pred['mean']) + list(pred['std'])\n            submission_data.append(row)\n        else:\n            row = [planet_id] + [1e-6] * (len(sample_submission.columns) - 1)\n            submission_data.append(row)\n    \n    submission_df = pd.DataFrame(submission_data, columns=sample_submission.columns)\n    submission_df.to_csv(output_path, index=False)\n    print(f\"Submission saved to {output_path}\")\n    \n    return submission_df\n\n# Main execution\nif __name__ == \"__main__\":\n    print(\"Starting NeurIPS - Ariel Data Challenge 2025 solution...\")\n    \n    print(\"Training model...\")\n    try:\n        lightning_model, trainer = train_model()\n        model = lightning_model.model\n    except Exception as e:\n        print(f\"Training failed: {e}\")\n        sample_submission = pd.read_csv(\"/kaggle/input/ariel-data-challenge-2025/sample_submission.csv\")\n        dummy_preds = {str(pid): {'mean': np.ones(283) * 1e-6, 'std': np.ones(283) * 1e-6} \n                      for pid in sample_submission['planet_id']}\n        create_submission(dummy_preds, \"/kaggle/input/ariel-data-challenge-2025/sample_submission.csv\")\n        exit()\n    \n    print(\"Loading test data...\")\n    test_star_info = pd.read_csv(\"/kaggle/input/ariel-data-challenge-2025/test_star_info.csv\")\n    adc_info = pd.read_csv(\"/kaggle/input/ariel-data-challenge-2025/adc_info.csv\")\n    \n    test_spectra = pd.DataFrame({'planet_id': test_star_info['planet_id']})\n    for i in range(283):\n        test_spectra[f'wavelength_{i}'] = 0.0\n    \n    test_dataset = ArielDataset(config.test_data_path, test_star_info, test_spectra, adc_info, is_train=False, max_samples=5)\n    test_loader = DataLoader(\n        test_dataset, \n        batch_size=config.batch_size, \n        shuffle=False, \n        num_workers=config.num_workers,\n        collate_fn=collate_fn\n    )\n    \n    print(\"Making predictions...\")\n    predictions = predict(model, test_loader)\n    \n    print(\"Creating submission...\")\n    sample_submission_path = \"/kaggle/input/ariel-data-challenge-2025/sample_submission.csv\"\n    submission = create_submission(predictions, sample_submission_path)\n    \n    print(\"Done!\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-04T13:15:07.971585Z","iopub.execute_input":"2025-09-04T13:15:07.972102Z","iopub.status.idle":"2025-09-04T13:16:08.951604Z","shell.execute_reply.started":"2025-09-04T13:15:07.972069Z","shell.execute_reply":"2025-09-04T13:16:08.950855Z"}},"outputs":[],"execution_count":null}]}