{"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":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# Cell 1 – mount & peek\nimport pandas as pd, os\nROOT = '/kaggle/input/ariel-data-challenge-2025'\nfor f in sorted(os.listdir(ROOT)):\n    print(f)\n\n# quick shapes\nground_truth_spectra = pd.read_csv(f'{ROOT}/train.csv')\nstars  = pd.read_csv(f'{ROOT}/train_star_info.csv')\nprint(stars.columns)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-23T08:25:16.167734Z","iopub.execute_input":"2025-07-23T08:25:16.168321Z","iopub.status.idle":"2025-07-23T08:25:16.269523Z","shell.execute_reply.started":"2025-07-23T08:25:16.1683Z","shell.execute_reply":"2025-07-23T08:25:16.268887Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# For any one planet id, specific job\n\nimport pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport os\nfrom pathlib import Path\n\n# Step 1: Set up paths and load metadata\ndata_dir = \"/kaggle/input/ariel-data-challenge-2025\"\nadc_info = pd.read_csv(f\"{data_dir}/adc_info.csv\")\nwavelengths = pd.read_csv(f\"{data_dir}/wavelengths.csv\")\ntrain_labels = pd.read_csv(f\"{data_dir}/train.csv\")\naxis_info = pd.read_parquet(f\"{data_dir}/axis_info.parquet\")\n\n# Step 2: Load one planet's data (AIRS-CH0)\n# List all planet IDs in 'train' directory\nplanet_ids = sorted(x for x in os.listdir(f\"{data_dir}/train/\") if not x.startswith('.'))\nprint(planet_ids)\n# Use the first planet as an example\nplanet_id = planet_ids[0]  # or change the index for other planets\nsignal_path = os.path.join(data_dir, \"train\", planet_id, \"AIRS-CH0_signal_0.parquet\")\ncalib_dir = os.path.join(data_dir, \"train\", planet_id, \"AIRS-CH0_calibration_0\")\n\ndark_path = os.path.join(calib_dir, \"dark.parquet\")\nflat_path = os.path.join(calib_dir, \"flat.parquet\")\ndead_path = os.path.join(calib_dir, \"dead.parquet\")\nlinear_path = os.path.join(calib_dir, \"linear_corr.parquet\")\nread_path = os.path.join(calib_dir, \"read.parquet\")\n\n\nsignal_data = pd.read_parquet(signal_path)\ndark_data = pd.read_parquet(dark_path)\nflat_data = pd.read_parquet(flat_path)\ndead_data = pd.read_parquet(dead_path)\nlinear_data = pd.read_parquet(linear_path)\nread_data = pd.read_parquet(read_path)\n\n\n# Step 3: Unflatten to 2D images\nimages = signal_data.values.reshape(-1, 32, 356)  # Shape: (11250, 32, 356)\ndark = dark_data.values.reshape(-1, 32, 356).mean(axis=0)  # Average dark frame\nflat = flat_data.values.reshape(-1, 32, 356).mean(axis=0)  # Average flat frame\ndead = dead_data.values.reshape(-1, 32, 356).mean(axis=0)  # Dead pixel mask\nread = read_data.values.reshape(-1, 32, 356).mean(axis=0) # the noise frames\n\ngain = adc_info.iat[0, 3]\noffset = adc_info.iat[0, 2]\n\nimages_corrected = (images / gain) + offset\nimages_corrected = images_corrected.astype(np.float64)\nimages_corrected = images_corrected - dark\nimages_corrected = images_corrected / (flat + 1e-10) # avoid the division by zero error\n# print(images_corrected.shape, dark.shape, flat.shape, dead.shape, read.shape, linear_data.shape)\n# Per wavelength correction mechanism\nimages_linear = np.zeros_like(images_corrected)\nfor x in range(356):  # Loop over wavelengths\n    col_name = f'column_{x}'\n    if col_name in linear_data.columns:\n        corrections = linear_data[col_name].iloc[0:192]  # Get all 192 rows\n        # Check if corrections are arrays or scalars\n        if isinstance(corrections.iloc[0], (list, np.ndarray)):  # Polynomial coefficients\n            # Average coefficients across rows\n            coeffs_mean = np.mean([np.array(c) for c in corrections], axis=0)[::-1]  # Reverse for np.polyval\n            for i in range(images_corrected.shape[0]):\n                for y in range(32):\n                    pixel_value = images_corrected[i, y, x]\n                    images_linear[i, y, x] = np.polyval(coeffs_mean, pixel_value)\n        else:  # Scalar correction\n            scalar_mean = np.mean(corrections)  # Average scalars\n            images_linear[:, :, x] = images_corrected[:, :, x] * scalar_mean\n    else:\n        images_linear[:, :, x] = images_corrected[:, :, x]  # No correction\n\n# Handle dead pixels\ndead_mask = dead > 0\nfor i in range(images_linear.shape[0]):\n    image = images_linear[i]\n    image[dead_mask] = np.median(image[~dead_mask])\n    images_linear[i] = image\n\n# Step 5: Baseline Model - Average AIRS-CH0 to Get Spectrum\nspectrum = np.mean(images_linear, axis=0)  # Shape: (32, 356)\nspectrum_1d = np.mean(spectrum, axis=0)    # Shape: (356,)\n\n# Step 6: Uncertainty Estimation with Read Noise\ntemporal_std = np.std(images_linear, axis=0)  # Temporal variation\ntemporal_std_1d = np.mean(temporal_std, axis=0)  # Shape: (356,)\nread_std = np.std(read, axis=0)  # Read noise contribution\nuncertainty_1d = np.sqrt(temporal_std_1d**2 + read_std**2)  # Combine variances\n\n# Step 7: Visualize the Spectrum\nplt.figure(figsize=(10, 5))\nplt.plot(wavelengths.iloc[0, 0:9], spectrum_1d[0:9], label='Estimated Spectrum')  # Use wl_1 to wl_9\nplt.fill_between(wavelengths.iloc[0, 0:9],\n                 spectrum_1d[0:9] - uncertainty_1d[0:9],\n                 spectrum_1d[0:9] + uncertainty_1d[0:9],\n                 alpha=0.3, label='Uncertainty')\nplt.xlabel('Wavelength (µm)')\nplt.ylabel('Flux')\nplt.title(f'Spectrum for Planet {planet_id}')\nplt.legend()\nplt.show()\n\n# Step 8: Prepare Submission for One Planet\nsubmission = pd.DataFrame({\n    'wavelength': wavelengths.iloc[0, 0:9],  # Use wl_1 to wl_9\n    'spectrum': spectrum_1d[0:9],\n    'sigma': uncertainty_1d[0:9]\n})\nprint(submission)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-23T14:22:03.253798Z","iopub.execute_input":"2025-07-23T14:22:03.254387Z","iopub.status.idle":"2025-07-23T14:22:09.723673Z","shell.execute_reply.started":"2025-07-23T14:22:03.254363Z","shell.execute_reply":"2025-07-23T14:22:09.722881Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\n\n# Load spectral intensity data (train.csv)\nROOT = '/kaggle/input/ariel-data-challenge-2025'\ntrain_spectra = pd.read_csv(f'{ROOT}/train.csv')\nstar_info_df = pd.read_csv(f'{ROOT}/train_star_info.csv')\nprint(f\"train_spectra shape: {train_spectra.shape}\")\nprint(f\"Sample columns: {train_spectra.columns[:10]}\")  # Verify spectral column names\n\n# Load wavelength grid (wavelengths.csv)\nwavelengths = pd.read_csv(f'{ROOT}/wavelengths.csv')\nprint(f\"wavelengths shape: {wavelengths.shape}\")\nprint(f\"Wavelength columns: {wavelengths.columns[:10]}\")\n\n# Convert the first row of wavelengths.csv to a numpy array of floats (wavelength grid)\nwl_values = wavelengths.iloc[0].values.astype(float)\n\n# Select a planet_id to plot\n# For example, the first planet in train_spectra\nplanet_id = train_spectra.loc[0, 'planet_id']\n\n# Extract the spectrum for this planet (all columns except planet_id)\n# Assuming columns after 'planet_id' are ordered spectral intensities corresponding to wl_values\nspectrum_cols = [col for col in train_spectra.columns if col != 'planet_id']\nplanet_spectrum = train_spectra.loc[train_spectra['planet_id'] == planet_id, spectrum_cols].values.flatten()\nprint(f\"Planet ID: {planet_id}\")\nprint(f\"Wavelengths (first 10): {wl_values[:10]}\")\nprint(f\"Spectrum intensities (first 10): {planet_spectrum[:10]}\")\n\n# Verify length matches\nassert len(wl_values) == len(planet_spectrum), \"Wavelengths and spectra length mismatch!\"\n\n# Plot spectrum vs wavelength\nplt.figure(figsize=(10, 5))\nplt.plot(wl_values, planet_spectrum, marker='o', linestyle='-', color='b')\nplt.xlabel('Wavelength (microns)')\nplt.ylabel('Spectral Intensity')\nplt.title(f'Spectrum for Planet {planet_id}')\nplt.grid(True)\nplt.tight_layout()\nplt.show()\ncombined_df = pd.merge(train_spectra, star_info_df, on='planet_id', how='left')\n\n# Check the merged dataframe\nprint(combined_df.shape)\nprint(combined_df.head())\n# Load wavelength grid (shared for spectra)\nwavelengths = pd.read_csv(f'{ROOT}/wavelengths.csv')\nwl_values = wavelengths.iloc[0].values.astype(float)\n\n# Base data directory containing per-planet subfolders\ndata_dir = f'{ROOT}/train'  # adjust to your path","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-23T14:24:56.107897Z","iopub.execute_input":"2025-07-23T14:24:56.108215Z","iopub.status.idle":"2025-07-23T14:24:56.399271Z","shell.execute_reply.started":"2025-07-23T14:24:56.108193Z","shell.execute_reply":"2025-07-23T14:24:56.398674Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport numpy as np\nimport os\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom sklearn.decomposition import PCA\nfrom joblib import Parallel, delayed\n\n# Set device to GPU if available\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"Using device: {device}\")\n\n# Helper: load parquet and convert to GPU tensor with known shape\ndef load_parquet_tensor(path, shape):\n    df = pd.read_parquet(path)\n    arr = df.values.astype(np.float64).reshape(shape)\n    return torch.tensor(arr, dtype=torch.float32, device=device)\n\ndef apply_linear_correction(images, linear_df):\n    corrected = torch.empty_like(images)\n    n_frames, height, wavelengths = images.shape\n\n    for x in range(wavelengths):\n        col_name = f'column_{x}'\n        if col_name in linear_df.columns:\n            corrections = linear_df[col_name].iloc[:192]\n            if isinstance(corrections.iloc[0], (list, np.ndarray)):\n                coeffs_np = np.mean([np.array(c) for c in corrections], axis=0)[::-1]\n                coeffs = torch.tensor(coeffs_np, dtype=images.dtype, device=device)\n\n                pixels = images[:, :, x].reshape(-1)\n                y = torch.zeros_like(pixels)\n                for c in coeffs:\n                    y = y * pixels + c\n                corrected[:, :, x] = y.reshape(n_frames, height)\n            else:\n                scalar = float(np.mean(corrections))\n                corrected[:, :, x] = images[:, :, x] * scalar\n        else:\n            corrected[:, :, x] = images[:, :, x]\n    return corrected\n\ndef load_and_calibrate(planet_id):\n    base = os.path.join(data_dir, str(planet_id))\n    signal_path = os.path.join(base, 'AIRS-CH0_signal_0.parquet')\n    calib_dir = os.path.join(base, 'AIRS-CH0_calibration_0')\n\n    signal = load_parquet_tensor(signal_path, (-1, 32, 356))\n    dark = load_parquet_tensor(os.path.join(calib_dir, 'dark.parquet'), (-1, 32, 356)).mean(dim=0)\n    flat = load_parquet_tensor(os.path.join(calib_dir, 'flat.parquet'), (-1, 32, 356)).mean(dim=0)\n    dead = load_parquet_tensor(os.path.join(calib_dir, 'dead.parquet'), (-1, 32, 356)).mean(dim=0)\n    linear_corr = pd.read_parquet(os.path.join(calib_dir, 'linear_corr.parquet'))\n\n    images_corrected = signal / gain + offset\n    images_corrected = images_corrected - dark.unsqueeze(0)\n    images_corrected = images_corrected / (flat.unsqueeze(0) + 1e-10)\n\n    images_linear = apply_linear_correction(images_corrected, linear_corr)\n\n    dead_mask = dead > 0\n    for i in range(images_linear.shape[0]):\n        img = images_linear[i]\n        median_val = img[~dead_mask].median()\n        img[dead_mask] = median_val\n        images_linear[i] = img\n\n    return images_linear\n\ndef extract_features(images_linear):\n    mean_spatial = images_linear.mean(dim=1).cpu().numpy()\n    pca = PCA(n_components=10)\n    pca_features = pca.fit_transform(mean_spatial)\n    avg_pca_features = pca_features.mean(axis=0)\n\n    avg_spectrum = mean_spatial.mean(axis=0)\n    gradients = np.gradient(avg_spectrum)\n    grad_features = np.array([gradients.mean(), gradients.std()])\n\n    return np.concatenate([avg_pca_features, grad_features])\n# Main execution\navailable_planets = {d for d in os.listdir(data_dir) if not d.startswith('.')}\nmetadata_planets = set(merged_df['planet_id'].astype(str))\nplanets_to_process = list(metadata_planets.intersection(available_planets))\n\nprint(f\"Planets in metadata: {len(metadata_planets)}\")\nprint(f\"Planets in training directory: {len(available_planets)}\")\nprint(f\"Planets to process: {len(planets_to_process)}\")\n    ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-23T14:25:03.657348Z","iopub.execute_input":"2025-07-23T14:25:03.657636Z","iopub.status.idle":"2025-07-23T14:25:03.674863Z","shell.execute_reply.started":"2025-07-23T14:25:03.657617Z","shell.execute_reply":"2025-07-23T14:25:03.674236Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(f\"Starting processing {len(planets_to_process)} planets sequentially on GPU...\")\n\nX, y, planet_ids = [], [], []\n\nfor count, pid_str in enumerate(planets_to_process, 1):\n    print(f\"[{count}/{len(planets_to_process)}] Processing planet {pid_str} ...\")\n    try:\n        images_linear = load_and_calibrate(pid_str)\n        features = extract_features(images_linear)\n\n        target_row = combined_df[combined_df['planet_id'] == int(pid_str)]\n        if target_row.empty:\n            print(f\"Warning: No spectral data for planet {pid_str}\")\n            continue\n\n        spectrum = target_row[spectrum_cols].values.flatten()\n        aux = target_row.drop(columns=['planet_id'] + spectrum_cols).values.flatten()\n        X.append(np.concatenate([features, aux]))\n        y.append(spectrum)\n        planet_ids.append(int(pid_str))\n    except Exception as e:\n        print(f\"Failed planet {pid_str}: {e}\")\n\nX = np.array(X)\ny = np.array(y)\nplanet_ids = np.array(planet_ids)\nprint(f\"Processing complete. Processed {len(X)} planets.\")\nprint(f\"Feature shape: {X.shape}, Target shape: {y.shape}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-23T14:25:06.658933Z","iopub.execute_input":"2025-07-23T14:25:06.659522Z","iopub.status.idle":"2025-07-23T16:40:02.109948Z","shell.execute_reply.started":"2025-07-23T14:25:06.659499Z","shell.execute_reply":"2025-07-23T16:40:02.108596Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom sklearn.decomposition import PCA\nfrom sklearn.ensemble import GradientBoostingRegressor\nfrom sklearn.multioutput import MultiOutputRegressor\n\n# ----------------------------\n# 1. Setup device and directories\n# ----------------------------\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"Using device: {device}\")\n\nBASE_DIR = '/kaggle/input/ariel-data-challenge-2025'\nTRAIN_DIR = os.path.join(BASE_DIR, 'train')\nTEST_DIR = os.path.join(BASE_DIR, 'test')\n\n# Load ADC calibration (common to train and test)\nadc_info = pd.read_csv(os.path.join(BASE_DIR, 'adc_info.csv'))\nGAIN = adc_info.iat[0, 3]\nOFFSET = adc_info.iat[0, 2]\n\n# Load metadata files\ntrain_spectra_df = pd.read_csv(os.path.join(BASE_DIR, 'train.csv'))\ntrain_aux_df = pd.read_csv(os.path.join(BASE_DIR, 'train_star_info.csv'))\ntest_aux_df = pd.read_csv(os.path.join(BASE_DIR, 'test_star_info.csv'))\n\n# Merge train spectra with aux info\ntrain_merged_df = pd.merge(train_spectra_df, train_aux_df, on='planet_id', how='left')\nspectrum_cols = [c for c in train_spectra_df.columns if c != 'planet_id']\n\n# ----------------------------\n# 2. Define GPU-based calibration and feature extraction functions\n# ----------------------------\n\ndef load_parquet_tensor(path, shape):\n    df = pd.read_parquet(path)\n    arr = df.values.astype(np.float64).reshape(shape)\n    return torch.tensor(arr, dtype=torch.float32, device=device)\n\ndef apply_linear_correction(images, linear_df):\n    corrected = torch.empty_like(images)\n    n_frames, height, wavelengths = images.shape\n\n    for x in range(wavelengths):\n        col_name = f'column_{x}'\n        if col_name in linear_df.columns:\n            corrections = linear_df[col_name].iloc[:192]\n            if isinstance(corrections.iloc[0], (list, np.ndarray)):\n                coeffs_np = np.mean([np.array(c) for c in corrections], axis=0)[::-1]\n                coeffs = torch.tensor(coeffs_np, dtype=images.dtype, device=device)\n\n                pixels = images[:, :, x].reshape(-1)\n                y = torch.zeros_like(pixels)\n                for c in coeffs:\n                    y = y * pixels + c\n                corrected[:, :, x] = y.reshape(n_frames, height)\n            else:\n                scalar = float(np.mean(corrections))\n                corrected[:, :, x] = images[:, :, x] * scalar\n        else:\n            corrected[:, :, x] = images[:, :, x]\n    return corrected\n\ndef load_and_calibrate(data_dir, planet_id):\n    base = os.path.join(data_dir, str(planet_id))\n    signal_path = os.path.join(base, 'AIRS-CH0_signal_0.parquet')\n    calib_dir = os.path.join(base, 'AIRS-CH0_calibration_0')\n\n    signal = load_parquet_tensor(signal_path, (-1, 32, 356))\n    dark = load_parquet_tensor(os.path.join(calib_dir, 'dark.parquet'), (-1, 32, 356)).mean(dim=0)\n    flat = load_parquet_tensor(os.path.join(calib_dir, 'flat.parquet'), (-1, 32, 356)).mean(dim=0)\n    dead = load_parquet_tensor(os.path.join(calib_dir, 'dead.parquet'), (-1, 32, 356)).mean(dim=0)\n    linear_corr_df = pd.read_parquet(os.path.join(calib_dir, 'linear_corr.parquet'))\n\n    images_corrected = signal / GAIN + OFFSET\n    images_corrected = images_corrected - dark.unsqueeze(0)\n    images_corrected = images_corrected / (flat.unsqueeze(0) + 1e-10)\n\n    images_linear = apply_linear_correction(images_corrected, linear_corr_df)\n\n    dead_mask = dead > 0\n    for i in range(images_linear.shape[0]):\n        img = images_linear[i]\n        median_val = img[~dead_mask].median()\n        img[dead_mask] = median_val\n        images_linear[i] = img\n\n    return images_linear\n\ndef extract_features(images_linear):\n    # images_linear shape: (frames, 32, 356)\n    mean_over_height = images_linear.mean(dim=1).cpu().numpy()  # shape (frames, 356)\n\n    pca = PCA(n_components=10)\n    pca_features = pca.fit_transform(mean_over_height)  # shape (frames, 10)\n    avg_pca = pca_features.mean(axis=0)\n\n    avg_spectrum = mean_over_height.mean(axis=0)\n    gradient = np.gradient(avg_spectrum)\n    grad_features = np.array([gradient.mean(), gradient.std()])\n    \n    return np.concatenate([avg_pca, grad_features])\n\n# ----------------------------\n# 3. Prepare Training Features and Targets\n# ----------------------------\n\ntrain_planets = list(set(train_merged_df['planet_id'].astype(str)))\nprint(f\"Training on {len(train_planets)} planets\")\n\nX_train, y_train, aux_train = [], [], []\n\nfor i, pid_str in enumerate(train_planets, 1):\n    print(f\"Processing train planet {pid_str} ({i}/{len(train_planets)})\")\n    try:\n        images_linear = load_and_calibrate(TRAIN_DIR, pid_str)\n        feat = extract_features(images_linear)\n\n        row = train_merged_df[train_merged_df['planet_id'] == int(pid_str)]\n        spectrum = row[spectrum_cols].values.flatten()\n        aux = row.drop(columns=['planet_id'] + spectrum_cols).values.flatten()\n\n        X_train.append(np.concatenate([feat, aux]))\n        y_train.append(spectrum)\n    except Exception as e:\n        print(f\"Skipping planet {pid_str} due to error: {e}\")\n\nX_train = np.array(X_train)\ny_train = np.array(y_train)\nprint(f\"Prepared training data: X={X_train.shape}, y={y_train.shape}\")\n\n# ----------------------------\n# 4. Prepare Test Features\n# ----------------------------\n\ntest_planets = [d for d in os.listdir(TEST_DIR) if not d.startswith('.')]\nprint(f\"Processing {len(test_planets)} test planets\")\n\nX_test = []\ntest_planet_ids = []\n\nfor i, pid_str in enumerate(test_planets, 1):\n    print(f\"Processing test planet {pid_str} ({i}/{len(test_planets)})\")\n    try:\n        images_linear = load_and_calibrate(TEST_DIR, pid_str)\n        feat = extract_features(images_linear)\n\n        row = test_aux_df[test_aux_df['planet_id'] == int(pid_str)]\n        if not row.empty:\n            aux = row.drop(columns=['planet_id']).values.flatten()\n            feat_combined = np.concatenate([feat, aux])\n        else:\n            feat_combined = feat\n        \n        X_test.append(feat_combined)\n        test_planet_ids.append(int(pid_str))\n    except Exception as e:\n        print(f\"Skipping test planet {pid_str} due to error: {e}\")\n\nX_test = np.array(X_test)\ntest_planet_ids = np.array(test_planet_ids)\nprint(f\"Prepared test data: X_test={X_test.shape}, number of test planets={len(test_planet_ids)}\")\n\n# ----------------------------\n# 5. Train Model (e.g., Gradient Boosting for mean and uncertainty)\n# ----------------------------\n\n# For demonstration, we train two models:\n# - One for the mean spectral values (283 outputs)\n# - One for uncertainty (dummy here - fixed small value, or train a second model similarly)\n\nprint(\"Training mean spectrum model...\")\nmean_model = MultiOutputRegressor(GradientBoostingRegressor(n_estimators=100, random_state=42))\nmean_model.fit(X_train, y_train)\n\n# For uncertainty, simplest baseline: fixed small uncertainties (e.g., 0.01)\nfixed_uncertainty = 0.01\nuncertainty_pred = np.full_like(y_train, fixed_uncertainty)  # shape = y_train.shape\n\n# ----------------------------\n# 6. Predict on Test Data\n# ----------------------------\n\nprint(\"Predicting test mean spectra...\")\ntest_mean_preds = mean_model.predict(X_test)  # shape (n_test, 283)\n\nprint(\"Using fixed uncertainties for test predictions...\")\ntest_uncertainty_preds = np.full(test_mean_preds.shape, fixed_uncertainty)\n\n# ----------------------------\n# 7. Prepare Submission File\n# ----------------------------\n\n# Concatenate mean and uncertainty predictions\nsubmission_array = np.hstack([test_mean_preds, test_uncertainty_preds])  # shape (n_test, 2*283 = 566)\n\n# Build DataFrame with planet_id and predicted data\nmean_cols = [f'spec_{i+1}' for i in range(283)]\nunc_cols = [f'unc_{i+1}' for i in range(283)]\n\ncolumns = ['planet_id'] + mean_cols + unc_cols\nsubmission_df = pd.DataFrame(\n    np.hstack([test_planet_ids.reshape(-1,1), submission_array]),\n    columns=columns\n)\nsubmission_df['planet_id'] = submission_df['planet_id'].astype(int)\n\nsubmission_path = 'submission.csv'\nsubmission_df.to_csv(submission_path, index=False)\nprint(f\"Submission file saved to {submission_path}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-24T02:43:42.365247Z","iopub.execute_input":"2025-07-24T02:43:42.365782Z","iopub.status.idle":"2025-07-24T04:57:54.407909Z","shell.execute_reply.started":"2025-07-24T02:43:42.365754Z","shell.execute_reply":"2025-07-24T04:57:54.40708Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom sklearn.decomposition import PCA\nfrom sklearn.ensemble import GradientBoostingRegressor\nfrom sklearn.multioutput import MultiOutputRegressor\n\n# ----------------------------\n# 1. Setup device and directories\n# ----------------------------\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"Using device: {device}\")\n\nBASE_DIR = '/kaggle/input/ariel-data-challenge-2025'\nTRAIN_DIR = os.path.join(BASE_DIR, 'train')\nTEST_DIR = os.path.join(BASE_DIR, 'test')\n\n# Load ADC calibration (common to train and test)\nadc_info = pd.read_csv(os.path.join(BASE_DIR, 'adc_info.csv'))\nGAIN = adc_info.iat[0, 3]\nOFFSET = adc_info.iat[0, 2]\n\n# Load metadata files\ntrain_spectra_df = pd.read_csv(os.path.join(BASE_DIR, 'train.csv'))\ntrain_aux_df = pd.read_csv(os.path.join(BASE_DIR, 'train_star_info.csv'))\ntest_aux_df = pd.read_csv(os.path.join(BASE_DIR, 'test_star_info.csv'))\n\n# Merge train spectra with aux info\ntrain_merged_df = pd.merge(train_spectra_df, train_aux_df, on='planet_id', how='left')\nspectrum_cols = [c for c in train_spectra_df.columns if c != 'planet_id']\n\n# ----------------------------\n# 2. Define GPU-based calibration and feature extraction functions\n# ----------------------------\n\ndef load_parquet_tensor(path, shape):\n    df = pd.read_parquet(path)\n    arr = df.values.astype(np.float64).reshape(shape)\n    return torch.tensor(arr, dtype=torch.float32, device=device)\n\ndef apply_linear_correction(images, linear_df):\n    corrected = torch.empty_like(images)\n    n_frames, height, wavelengths = images.shape\n\n    for x in range(wavelengths):\n        col_name = f'column_{x}'\n        if col_name in linear_df.columns:\n            corrections = linear_df[col_name].iloc[:192]\n            if isinstance(corrections.iloc[0], (list, np.ndarray)):\n                coeffs_np = np.mean([np.array(c) for c in corrections], axis=0)[::-1]\n                coeffs = torch.tensor(coeffs_np, dtype=images.dtype, device=device)\n\n                pixels = images[:, :, x].reshape(-1)\n                y = torch.zeros_like(pixels)\n                for c in coeffs:\n                    y = y * pixels + c\n                corrected[:, :, x] = y.reshape(n_frames, height)\n            else:\n                scalar = float(np.mean(corrections))\n                corrected[:, :, x] = images[:, :, x] * scalar\n        else:\n            corrected[:, :, x] = images[:, :, x]\n    return corrected\n\ndef load_and_calibrate(data_dir, planet_id):\n    base = os.path.join(data_dir, str(planet_id))\n    signal_path = os.path.join(base, 'AIRS-CH0_signal_0.parquet')\n    calib_dir = os.path.join(base, 'AIRS-CH0_calibration_0')\n\n    signal = load_parquet_tensor(signal_path, (-1, 32, 356))\n    dark = load_parquet_tensor(os.path.join(calib_dir, 'dark.parquet'), (-1, 32, 356)).mean(dim=0)\n    flat = load_parquet_tensor(os.path.join(calib_dir, 'flat.parquet'), (-1, 32, 356)).mean(dim=0)\n    dead = load_parquet_tensor(os.path.join(calib_dir, 'dead.parquet'), (-1, 32, 356)).mean(dim=0)\n    linear_corr_df = pd.read_parquet(os.path.join(calib_dir, 'linear_corr.parquet'))\n\n    images_corrected = signal / GAIN + OFFSET\n    images_corrected = images_corrected - dark.unsqueeze(0)\n    images_corrected = images_corrected / (flat.unsqueeze(0) + 1e-10)\n\n    images_linear = apply_linear_correction(images_corrected, linear_corr_df)\n\n    dead_mask = dead > 0\n    for i in range(images_linear.shape[0]):\n        img = images_linear[i]\n        median_val = img[~dead_mask].median()\n        img[dead_mask] = median_val\n        images_linear[i] = img\n\n    return images_linear\n\ndef extract_features(images_linear):\n    mean_over_height = images_linear.mean(dim=1).cpu().numpy()  # shape (frames, 356)\n\n    pca = PCA(n_components=10)\n    pca_features = pca.fit_transform(mean_over_height)  # shape (frames, 10)\n    avg_pca = pca_features.mean(axis=0)\n\n    avg_spectrum = mean_over_height.mean(axis=0)\n    gradient = np.gradient(avg_spectrum)\n    grad_features = np.array([gradient.mean(), gradient.std()])\n    \n    return np.concatenate([avg_pca, grad_features])\n\n# ----------------------------\n# 3. Prepare Training Features and Targets\n# ----------------------------\n\ntrain_planets = list(set(train_merged_df['planet_id'].astype(str)))\nprint(f\"Training on {len(train_planets)} planets\")\n\nX_train, y_train, aux_train = [], [], []\n\nfor i, pid_str in enumerate(train_planets, 1):\n    print(f\"Processing train planet {pid_str} ({i}/{len(train_planets)})\")\n    try:\n        images_linear = load_and_calibrate(TRAIN_DIR, pid_str)\n        feat = extract_features(images_linear)\n\n        row = train_merged_df[train_merged_df['planet_id'] == int(pid_str)]\n        spectrum = row[spectrum_cols].values.flatten()\n        aux = row.drop(columns=['planet_id'] + spectrum_cols).values.flatten()\n\n        X_train.append(np.concatenate([feat, aux]))\n        y_train.append(spectrum)\n    except Exception as e:\n        print(f\"Skipping planet {pid_str} due to error: {e}\")\n\nX_train = np.array(X_train)\ny_train = np.array(y_train)\nprint(f\"Prepared training data: X={X_train.shape}, y={y_train.shape}\")\n\n# ----------------------------\n# 4. Prepare Test Features\n# ----------------------------\n\ntest_planets = [d for d in os.listdir(TEST_DIR) if not d.startswith('.')]\nprint(f\"Processing {len(test_planets)} test planets\")\n\nX_test = []\ntest_planet_ids = []\n\nfor i, pid_str in enumerate(test_planets, 1):\n    print(f\"Processing test planet {pid_str} ({i}/{len(test_planets)})\")\n    try:\n        images_linear = load_and_calibrate(TEST_DIR, pid_str)\n        feat = extract_features(images_linear)\n\n        row = test_aux_df[test_aux_df['planet_id'] == int(pid_str)]\n        if not row.empty:\n            aux = row.drop(columns=['planet_id']).values.flatten()\n            feat_combined = np.concatenate([feat, aux])\n        else:\n            feat_combined = feat\n        \n        X_test.append(feat_combined)\n        test_planet_ids.append(int(pid_str))\n    except Exception as e:\n        print(f\"Skipping test planet {pid_str} due to error: {e}\")\n\nX_test = np.array(X_test)\ntest_planet_ids = np.array(test_planet_ids)\nprint(f\"Prepared test data: X_test={X_test.shape}, number of test planets={len(test_planet_ids)}\")\n\n# ----------------------------\n# 5. Train Model (e.g., Gradient Boosting for mean; uncertainty is fixed)\n# ----------------------------\n\nprint(\"Training mean spectrum model...\")\nmean_model = MultiOutputRegressor(GradientBoostingRegressor(n_estimators=100, random_state=42))\nmean_model.fit(X_train, y_train)\n\nfixed_uncertainty = 0.01\n\n# ----------------------------\n# 6. Predict on Test Data\n# ----------------------------\n\nprint(\"Predicting test mean spectra...\")\ntest_mean_preds = mean_model.predict(X_test)  # shape (n_test, 283)\n\nprint(\"Using fixed uncertainties for test predictions...\")\ntest_uncertainty_preds = np.full(test_mean_preds.shape, fixed_uncertainty)\n\n# ----------------------------\n# 7. Prepare Submission File (MATCHING sample_submission.csv format!)\n# ----------------------------\n\n# === CHANGE: Use correct column names as in sample_submission.csv ===\nwl_cols = [f'wl_{i+1}' for i in range(283)]\nsigma_cols = [f'sigma_{i+1}' for i in range(283)]\nsubmission_columns = ['planet_id'] + wl_cols + sigma_cols\n\n# Stack mean and uncertainty predictions and combine with planet_id\nsubmission_array = np.hstack([test_mean_preds, test_uncertainty_preds]) # (n_test, 566)\nfull_array = np.hstack([test_planet_ids.reshape(-1, 1), submission_array])  # (n_test, 567)\n\n# Construct DataFrame\nsubmission_df = pd.DataFrame(full_array, columns=submission_columns)\nsubmission_df['planet_id'] = submission_df['planet_id'].astype(int)\n\nsubmission_path = 'submission.csv'\nsubmission_df.to_csv(submission_path, index=False)\nprint(f\"Submission file saved to {submission_path} with correct format.\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-24T05:10:45.600625Z","iopub.execute_input":"2025-07-24T05:10:45.60093Z","iopub.status.idle":"2025-07-24T07:25:23.894543Z","shell.execute_reply.started":"2025-07-24T05:10:45.600907Z","shell.execute_reply":"2025-07-24T07:25:23.893835Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pdf = pd.read_csv('submission.csv')\nprint(pdf.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-07-24T07:55:33.388326Z","iopub.execute_input":"2025-07-24T07:55:33.388879Z","iopub.status.idle":"2025-07-24T07:55:33.412629Z","shell.execute_reply.started":"2025-07-24T07:55:33.388856Z","shell.execute_reply":"2025-07-24T07:55:33.412057Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}