{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":101849,"databundleVersionId":13093295,"sourceType":"competition"},{"sourceId":13130542,"sourceType":"datasetVersion","datasetId":8318122}],"dockerImageVersionId":31089,"isInternetEnabled":false,"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\n\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,"execution":{"iopub.status.busy":"2025-09-24T23:40:02.424517Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install /kaggle/input/optuna/optuna-4.5.0-py3-none-any.whl\n\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import optuna\nprint(optuna.__version__)  # should print 4.5.0\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\npath = \"/kaggle/input/ariel-data-challenge-2025/train\"  # replace with your dataset name\nprint(os.listdir(path))","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"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\nclass Config:\n    # 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    # Training hyperparams\n    batch_size = 8\n    learning_rate = 1e-4\n    num_epochs = 20\n    hidden_dim = 64\n    num_layers = 2\n    dropout = 0.1\n\n    airs_feature_dim = 356\n    fgs1_feature_dim = 32\n    # Instrument feature dimensions\n    \n\n    # Dataloader/Device\n    num_workers = 2\n    accelerator = 'gpu' if torch.cuda.is_available() else 'cpu'\n    devices = 1\n\n\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\ntrain_star_info, train_spectra, wavelengths, axis_info, adc_info = load_metadata()\n\nconfig = Config()\n\npd.set_option('display.max_columns', None)  # Show all columns\npd.set_option('display.width', 120)         # Set display width\n\nprint(\"\\n===== Train Star Info =====\")\nprint(train_star_info.head().to_string(index=False))\n\nprint(\"\\n===== Train Spectra Info =====\")\nprint(train_spectra.head().to_string(index=False))\n\nprint(\"\\n===== Wavelength Info =====\")\nprint(wavelengths.head().to_string(index=False))","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\nimport os\nimport glob\nimport pandas as pd\nimport numpy as np\nimport torch\nfrom torch.utils.data import Dataset\nfrom torch.utils.data import DataLoader\nfrom tqdm import tqdm\n\nclass MinimalArielDataset(Dataset):\n    def __init__(self, data_path, config,\n                 star_info=None, spectra=None,\n                 instruments=['AIRS-CH0', 'FGS1'],\n                 light_curve_length=256, mode=\"train\"):\n        \"\"\"\n        mode: 'train' or 'test'\n        \"\"\"\n        self.data_path = data_path\n        self.spectra = spectra\n        self.config = config\n        self.instruments = instruments\n        self.light_curve_length = light_curve_length\n        self.mode = mode\n\n        # Always initialize\n        self.planets = []\n\n        if self.mode == 'train':\n            for pid in star_info['planet_id'].values:\n                planet_folder = os.path.join(data_path, str(pid))\n                if not os.path.exists(planet_folder):\n                    continue\n\n                instrument_data = []\n                for inst in instruments:\n                    sig_files = glob.glob(os.path.join(planet_folder, f\"{inst}_signal_*.parquet\"))\n                    cal_folder = os.path.join(planet_folder, f\"{inst}_calibration_0\")\n                    if sig_files and os.path.exists(cal_folder):\n                        instrument_data.append({\n                            'name': inst,\n                            'signal': sig_files[0],\n                            'calibration': cal_folder\n                        })\n                if instrument_data:\n                    self.planets.append({'planet_id': pid, 'instruments': instrument_data})\n\n        elif self.mode == 'test':\n            planet_folder = data_path\n            pid = os.path.basename(os.path.normpath(data_path))\n\n            instrument_data = []\n            for inst in instruments:\n                sig_files = glob.glob(os.path.join(planet_folder, f\"{inst}_signal_*.parquet\"))\n                cal_folder = os.path.join(planet_folder, f\"{inst}_calibration_0\")\n                if sig_files and os.path.exists(cal_folder):\n                    instrument_data.append({\n                        'name': inst,\n                        'signal': sig_files[0],\n                        'calibration': cal_folder\n                    })\n            if instrument_data:\n                self.planets.append({'planet_id': pid, 'instruments': instrument_data})\n\n        print(f\"Loaded {len(self.planets)} planets.\")\n\n    def __len__(self):\n        return len(self.planets)\n\n    def __getitem__(self, idx):\n        pkg = self.planets[idx]\n        pid = pkg['planet_id']\n\n        # Process instruments → average light curves\n        light_curves = []\n        cache_file = os.path.join('cache', f'{pid}.npy')\n        light_curve = np.load(cache_file)  # (256,)\n        light_curve = torch.FloatTensor(light_curve).unsqueeze(0)\n        \n\n        if self.mode == 'train':\n            \n            target = self.spectra[self.spectra['planet_id']==pid].drop('planet_id', axis=1).values[0]\n            target = target.astype(np.float32)  # simple fixed length\n\n            return {\n            'light_curve': torch.FloatTensor(light_curve),\n            'target': torch.FloatTensor(target)\n            }\n\n        else:\n            return {\n                'light_curve': torch.FloatTensor(light_curve),\n                'planet_id': pid\n            }\n\n    def _process_instrument(self, inst, cal_folder):\n        # Load signal, flatten\n        if inst['name'] == 'AIRS-CH0':\n            signal = pd.read_parquet(inst['signal']).values.reshape(11250, 32, 356)\n        else:\n            signal = pd.read_parquet(inst['signal']).values.reshape(135000, 32, 32)\n\n        dark   = pd.read_parquet(os.path.join(cal_folder, \"dark.parquet\")).values\n        flat   = pd.read_parquet(os.path.join(cal_folder, \"flat.parquet\")).values\n        dead   = pd.read_parquet(os.path.join(cal_folder, \"dead.parquet\")).values\n        \n        linear = pd.read_parquet(os.path.join(cal_folder, \"linear_corr.parquet\")).values\n        feature_dim = self.config.airs_feature_dim if inst['name']=='AIRS-CH0' else self.config.fgs1_feature_dim\n        \n\n        linear = linear.reshape(-1, 6, linear.shape[1]).mean(axis=1)\n\n\n        signal = (signal - dark) / flat  # simplified\n        signal *= linear\n        dead_mask = ((dead > 0) | np.isnan(dead) | (dead < 0))\n        signal = np.where(dead_mask, 0, signal)\n    \n        if dead_mask.shape != signal.shape[1:]:\n            dead_mask = dead_mask.reshape((1,) + dead_mask.shape)\n            signal = np.where(dead_mask, 0, signal)\n    \n        # Collapse to 1D light curve\n        light_curve = signal.sum(axis=(1, 2))  # sum over rows and columns\n        if not isinstance(light_curve, np.ndarray):\n            raise ValueError(\"Light curve is None or not a NumPy array\")\n\n    # Resize/pad to fixed length\n        return self._resize_light_curve(light_curve)\n        \n    def _resize_light_curve(self, lc):\n        # Resize to fixed length\n        target_len = self.light_curve_length\n        if len(lc) > target_len:\n            step = len(lc) // target_len\n            lc = lc[::step][:target_len]\n        elif len(lc) < target_len:\n            pad = target_len - len(lc)\n            lc = np.pad(lc, (0, pad))\n        return lc\n\n\ndef preprocess_dataset(dataset, cache_dir=\"cache\"):\n    os.makedirs(cache_dir, exist_ok=True)\n\n    for pkg in tqdm(dataset.planets, desc=\"Precomputing light curves\"):\n        pid = pkg['planet_id']\n        cache_file = os.path.join(cache_dir, f\"{pid}.npy\")\n\n        # Skip if already cached\n        if os.path.exists(cache_file):\n            continue\n\n        # Process all instruments for this planet\n        light_curves = []\n        for inst in pkg['instruments']:\n            lc = dataset._process_instrument(inst, cal_folder=inst['calibration'])\n            light_curves.append(lc)\n\n        # Average across instruments\n        light_curve = np.mean(light_curves, axis=0).astype(np.float32)\n\n        # Save to cache\n        np.save(cache_file, light_curve)\n\n    print(f\"✅ Preprocessing complete. Cached files in {cache_dir}\")\n\n\n\n\ndataset = MinimalArielDataset(\n    data_path=\"/kaggle/input/ariel-data-challenge-2025/train\",\n    star_info=train_star_info,\n    spectra=train_spectra,\n    config=config,\n    light_curve_length=256,\n    mode='train'\n)\n\npreprocess_dataset(dataset)\n\ntrain_loader = DataLoader(dataset, batch_size=2, shuffle=True, num_workers=1)\n    \n\n# Test one batch\nfor batch in train_loader:\n    print('Light Curve')\n    print(batch['light_curve'].shape)   # [batch_size, 256]\n    print('Target')\n    print(batch['target'].shape)        # [batch_size, 1024]\n    break\n\n         \n       \n\n     \n\n    \n\n\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nclass ArielNet(nn.Module):\n    def __init__(self, input_dim=256, hidden_dim=512, output_dim=283):\n        super().__init__()\n        self.backbone = nn.Sequential(\n            nn.Linear(input_dim, hidden_dim),\n            nn.ReLU(),\n            nn.Linear(hidden_dim, hidden_dim // 2),\n            nn.ReLU(),\n        )\n\n        self.spectrum_head = nn.Linear(hidden_dim // 2, output_dim)\n        self.uncertainty_head = nn.Linear(hidden_dim // 2, output_dim)\n\n    def forward(self, x):\n        features = self.backbone(x)\n        spectra = self.spectrum_head(features)\n        log_var = self.uncertainty_head(features)\n        uncertainty = F.softplus(log_var) + 1e-6  # ensure positivity\n        return spectra, uncertainty\n\n\n        \n        ","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ArielNetFiltrationPipeline:\n    def __init__(self, train_loader, val_loader, test_loader):\n        self.model = ArielNet()\n        self.train_loader = train_loader\n        self.val_loader = val_loader\n        self.test_loader = test_loader\n        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=0.001)\n    def objective(self, trial):\n        # ----- Hyperparameters -----\n        hidden_dim = trial.suggest_int(\"hidden_dim\", 128, 1024, step=128)\n        lr = trial.suggest_loguniform(\"lr\", 1e-4, 1e-2)\n        batch_size = trial.suggest_categorical(\"batch_size\", [16, 32, 64])\n    \n        # ----- Initialize model and optimizer with trial params -----\n        self.model = ArielNet(input_dim=256, hidden_dim=hidden_dim, output_dim=283)\n        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=lr)\n        \n    \n        # ----- Train for a few epochs -----\n        epochs = 5\n        for epoch in range(epochs):\n            train_loss = self.train_one_epoch()\n            print('Train Loss:', train_loss)\n    \n            val_loss = self.validate_one_epoch()\n            print('Validation Loss:', val_loss)\n    \n        # ----- Return validation loss for Optuna -----\n        return val_loss\n\n\n    def fit(self):\n        study = optuna.create_study(\n            direction=\"minimize\",\n            study_name=\"ariel_hyperparam\",\n            storage=\"sqlite:///optuna_ariel.db\",\n            load_if_exists=True\n        )\n    \n        study.optimize(self.objective, n_trials=20)\n    \n        print(\"Best trial:\")\n        trial = study.best_trial\n        print(\"  Value (val_loss):\", trial.value)\n        print(\"  Params: \")\n        for key, value in trial.params.items():\n            print(f\"    {key}: {value}\")\n\n\n\n    def gaussian_nll_loss(self, mu, sigma, target):\n        return (0.5 * torch.log(sigma**2) + 0.5 * ((target - mu)**2 / (sigma**2))).mean()\n\n    def train_one_epoch(self, device='cpu'):\n        self.model.train()\n        total_loss = 0\n        num_samples = 0\n\n        for batch in self.train_loader:\n            light_curve = batch.get('light_curve').to(device)\n            target = batch.get('target').to(device)\n\n            self.optimizer.zero_grad()\n            mu, sigma = self.model(light_curve)\n        \n        # Gaussian NLL loss\n            loss = self.gaussian_nll_loss(mu, sigma, target)\n            loss.backward()\n            self.optimizer.step()\n        \n            total_loss += loss.item() * light_curve.size(0)\n            num_samples += light_curve.size(0)\n    \n        return total_loss / num_samples\n\n    def validate_one_epoch(self, device='cpu'):\n        self.model.eval()\n        total_loss = 0\n        num_samples = 0\n        with torch.no_grad():\n            for batch in self.val_loader:\n                light_curve = batch.get('light_curve').to(device)\n                target = batch.get('target').to(device)\n    \n                mu, sigma = self.model(light_curve)\n            \n            # Gaussian NLL loss\n                loss = self.gaussian_nll_loss(mu, sigma, target)\n                \n            \n                total_loss += loss.item() * light_curve.size(0)\n                num_samples += light_curve.size(0)\n        \n        return total_loss / num_samples\n\n\n    def predict(self, device='cpu'):\n\n        sample_sub = pd.read_csv('/kaggle/input/ariel-data-challenge-2025/sample_submission.csv')\n        print('Sample Submission Head:')\n        print(sample_sub.head())\n        submission_rows = []\n        self.model.eval()\n        with torch.no_grad(): \n            for batch in self.test_loader:\n                print(batch)\n                light_curve = batch.get('light_curve').to(device)\n\n                \n\n                planet_id = batch['planet_id'][0]\n\n                mu, sigma = self.model(light_curve)\n\n                mu = mu.cpu().numpy()\n                sigma = sigma.cpu().numpy()\n\n                row = [planet_id]\n                row.extend(mu[0].tolist())\n                row.extend(sigma[0].tolist())\n                submission_rows.append(row)\n\n            spectra_columns = [f\"wl_{j+1}\" for j in range(283)]\n            uncertainty_columns = [f\"sigma_{j+1}\" for j in range(283)]\n\n            columns = ['planet_id'] + spectra_columns + uncertainty_columns\n            \n            submission_df = pd.DataFrame(submission_rows, columns=columns\n             \n            )\n            print('Submission Dataframe:')\n            print(submission_df.head())\n        submission_df.to_csv(\"submission.csv\", index=False)\n        print(\"Saved submission.csv\")\n                    \n\n                \n\n\n    ","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from torch.utils.data import Subset\nfrom sklearn.model_selection import train_test_split\n\nif __name__ == '__main__':\n    \n    full_dataset = MinimalArielDataset(\n        data_path=\"/kaggle/input/ariel-data-challenge-2025/train\",\n        star_info=train_star_info,\n        spectra=train_spectra,\n        config=config,\n        light_curve_length=256,\n        mode='train'\n    )\n    \n    test_dataset = MinimalArielDataset(\n        data_path=\"/kaggle/input/ariel-data-challenge-2025/test/1103775\",\n        config=config,\n        light_curve_length=256,\n        mode='test'\n    )\n\n\n    \n    test_dataloader = DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=2)\n\n    \n    N = len(full_dataset)\n    indices = list(range(N))\n\n    if N > 1:\n        # Make sure val set has at least 1 sample\n        val_size = max(1, int(0.2 * N))\n        train_size = N - val_size\n        if train_size == 0:  # if dataset is too small, give 1 to train\n            train_size, val_size = 1, 0\n\n        train_idx, val_idx = train_test_split(range(N), test_size=val_size, random_state=42)\n\n        train_dataset = Subset(full_dataset, train_idx)\n        preprocess_dataset(train_dataset)\n        val_dataset   = Subset(full_dataset, val_idx)\n        preprocess_dataset(val_dataset)\n        \n        train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=2)\n        val_loader   = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=2)\n\n    else:\n        # Only 1 sample → no validation\n        train_dataset = full_dataset\n        preprocess_dataset(train_dataset)\n        val_dataset = None\n        preprocess_dataset(val_dataset)\n\n        train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True)\n        val_loader = None\n        \n    # Split indices into train and validation\n    \n    \n    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)\n    val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)\n    \n    \n    \n    filtration_pipeline = ArielNetFiltrationPipeline(\n                            train_loader=train_loader,\n                            val_loader=val_loader,\n                            test_loader=test_dataloader)\n\n    filtration_pipeline.fit()\n    filtration_pipeline.predict()\n","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}