{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":101849,"databundleVersionId":12846694,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":12303949,"sourceType":"datasetVersion","datasetId":7755350},{"sourceId":12309150,"sourceType":"datasetVersion","datasetId":7758622}],"dockerImageVersionId":31040,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# ==============================================================================\n# FINAL SUBMISSION SCRIPT: Ariel Data Challenge 2025 (Corrected Version)\n#\n# Methodology:\n# 1. Maximal Feature Engineering: Creates a rich feature set from time-series\n#    data (statistical, fourier, wavelet) and metadata (polynomial interactions).\n# 2. Quantile Regression: Trains three separate XGBoost models to predict the\n#    5th, 50th, and 95th percentiles of the target spectrum.\n# 3. Multi-Output Handling: Uses scikit-learn's MultiOutputRegressor to handle\n#    the 283-dimensional output for the quantile models.\n# 4. Uncertainty Calibration: Finds optimal scaling and additive factors on a\n#    validation set to calibrate the predicted uncertainty (sigma), using the\n#    official competition metric for reliable evaluation.\n# 5. Two-Stage Execution (Corrected Logic):\n#    - If run in an environment WITHOUT the test set (e.g., interactive editor),\n#      it trains the models, finds calibration factors, retrains on all data,\n#      and saves everything.\n#    - If run in the submission environment (test set is present), it LOADS\n#      the saved models and factors, and generates the final submission.csv.\n# ==============================================================================\n\nimport time\nimport pickle\nimport numpy as np\nimport pandas as pd\nimport polars as pl\nimport functools\n\nimport scipy.stats\nimport xgboost as xgb\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.metrics import mean_squared_error\nfrom sklearn.preprocessing import PolynomialFeatures\nfrom sklearn.multioutput import MultiOutputRegressor\nimport pywt\nimport os\nfrom tqdm import tqdm\nimport multiprocessing\n\nclass Config:\n    \"\"\"Contains all hyperparameters and file paths in one place.\"\"\"\n\n    # Set paths for Kaggle environment\n    DATA_PATH = '/kaggle/input/ariel-data-challenge-2025/'\n    PREPROCESSED_PATH = '/kaggle/input/ariel-data-challenge-2025-af-npy/'\n    OUTPUT_PATH = '/kaggle/input/ariel-data-2025-27' # Models and artifacts will be saved here\n\n    TRAIN_LABELS_PATH = os.path.join(DATA_PATH, 'train.csv')\n    TRAIN_STAR_INFO_PATH = os.path.join(DATA_PATH, 'train_star_info.csv')\n    TEST_STAR_INFO_PATH = os.path.join(DATA_PATH, 'test_star_info.csv')\n    SAMPLE_SUBMISSION_PATH = os.path.join(DATA_PATH, 'sample_submission.csv')\n    WAVELENGTHS_PATH = os.path.join(DATA_PATH, 'wavelengths.csv')\n    \n    # Using pre-processed raw data for faster execution.\n    # In a real run, you might generate these from scratch.\n    A_RAW_PATH = os.path.join(PREPROCESSED_PATH, \"a_raw_train.npy\")\n    F_RAW_PATH = os.path.join(PREPROCESSED_PATH, \"f_raw_train.npy\")\n\n    # Modeling \n    VALIDATION_SPLIT = 0.1\n    RANDOM_STATE = 42\n    QUANTILES = [0.05, 0.50, 0.95]\n    XGB_PARAMS = {\n        'n_estimators': 400,\n        'learning_rate': 0.04,\n        'max_depth': 6,\n        'subsample': 0.8,\n        'colsample_bytree': 0.7,\n        'random_state': RANDOM_STATE,\n        'tree_method': 'hist',\n    }\n    \n    # Calibration Search Space \n    CALIBRATION_SCALING_FACTORS = [1.0, 1.2, 1.5, 2.0, 2.5]\n    CALIBRATION_ADDITIVE_FACTORS = [0.0, 0.0005, 0.001, 0.0015]\n\nconfig = Config()\n\ndef f_read_and_preprocess(dataset, planet_ids):\n    \"\"\"Read the FGS1 files for all planet_ids and extract the time series.\"\"\"\n    print(f\"Preprocessing FGS1 data for {dataset} set...\")\n    f_raw_data = np.full((len(planet_ids), 67500), np.nan, dtype=np.float32)\n    for i, planet_id in tqdm(list(enumerate(planet_ids)), desc=\"FGS1\"):\n        path = f'/kaggle/input/ariel-data-challenge-2025/{dataset}/{int(planet_id)}/FGS1_signal_0.parquet'\n        f_signal = pl.read_parquet(path)\n        mean_signal = f_signal.cast(pl.Int32).sum_horizontal().cast(pl.Float32).to_numpy() / 1024\n        net_signal = mean_signal[1::2] - mean_signal[0::2]\n        f_raw_data[i] = net_signal\n    return f_raw_data\n\ndef a_read_and_preprocess(dataset, planet_ids):\n    \"\"\"Read the AIRS-CH0 files for all planet_ids and extract the time series.\"\"\"\n    print(f\"Preprocessing AIRS-CH0 data for {dataset} set...\")\n    a_raw_data = np.full((len(planet_ids), 5625), np.nan, dtype=np.float32)\n    for i, planet_id in tqdm(list(enumerate(planet_ids)), desc=\"AIRS-CH0\"):\n        path = f'/kaggle/input/ariel-data-challenge-2025/{dataset}/{int(planet_id)}/AIRS-CH0_signal_0.parquet'\n        signal = pl.read_parquet(path)\n        mean_signal = signal.cast(pl.Int32).sum_horizontal().cast(pl.Float32).to_numpy() / (32*356)\n        net_signal = mean_signal[1::2] - mean_signal[0::2]\n        a_raw_data[i] = net_signal\n    return a_raw_data\n\n# ==============================================================================\n# CORRECTED SCORING FUNCTION\n# This function now correctly implements the official weighted GLL metric.\n# ==============================================================================\ndef official_competition_score(y_true, y_pred, sigma_pred, naive_mean, naive_sigma,\n                               fsg_sigma_true=1e-6, airs_sigma_true=1e-5, fgs_weight=2.0):\n    \"\"\"\n    Calculates the score based on the official weighted Gaussian Log-Likelihood metric.\n    \"\"\"\n    y_true, y_pred, sigma_pred = np.array(y_true), np.array(y_pred), np.array(sigma_pred)\n    \n    # Clip sigma to prevent log(0) or division by zero\n    sigma_pred = np.clip(sigma_pred, 1e-15, None)\n\n    # Define ground truth sigma per channel (1 for FGS1, 282 for AIRS)\n    sigma_true_per_channel = np.append(np.array([fsg_sigma_true]), np.ones(y_true.shape[1] - 1) * airs_sigma_true)\n    sigma_true = np.tile(sigma_true_per_channel, (y_true.shape[0], 1))\n\n    # Calculate GLLs element-wise for each data point\n    GLL_pred = scipy.stats.norm.logpdf(y_true, loc=y_pred, scale=sigma_pred)\n    GLL_true = scipy.stats.norm.logpdf(y_true, loc=y_true, scale=sigma_true)\n    GLL_mean = scipy.stats.norm.logpdf(y_true, loc=naive_mean, scale=naive_sigma)\n    \n    # Calculate individual scores element-wise, adding epsilon for stability\n    denominator = GLL_true - GLL_mean\n    ind_scores = (GLL_pred - GLL_mean) / (denominator + 1e-9)\n\n    # Define weights per channel (higher weight for FGS1)\n    weights_per_channel = np.append(np.array([fgs_weight]), np.ones(y_true.shape[1] - 1))\n    weights = np.tile(weights_per_channel, (y_true.shape[0], 1))\n\n    # Calculate the final score as a weighted average of individual scores\n    final_score = np.average(ind_scores, weights=weights)\n    \n    return float(np.clip(final_score, 0.0, 1.0))\n\n# FEATURE ENGINEERING\ndef maximal_feature_engineering(f_raw, a_raw, star_info_df):\n    \"\"\"Creates the final, maximal feature set for the models.\"\"\"\n    print(\"Engineering MAXIMAL features...\")\n    \n    # --- Base Time Windows & Unobscured calculations ---\n    fgs_pre = f_raw[:, :20500]; fgs_post = f_raw[:, 47000:]\n    fgs_unobscured_mean = (fgs_pre.mean(axis=1) + fgs_post.mean(axis=1)) / 2\n    fgs_unobscured_std = (fgs_pre.std(axis=1) + fgs_post.std(axis=1)) / 2\n    fgs_transit = f_raw[:, 23500:44000]\n    \n    # --- Time-series Features ---\n    features = {}\n    for i in range(5):\n        f_slice_mean = fgs_transit[:, i*4100:(i+1)*4100].mean(axis=1)\n        features[f'fgs_slice_{i+1}'] = (fgs_unobscured_mean - f_slice_mean) / fgs_unobscured_mean\n    \n    features['fgs_transit_std'] = fgs_transit.std(axis=1)\n    features['fgs_transit_skew'] = scipy.stats.skew(fgs_transit, axis=1)\n    features['fgs_transit_kurtosis'] = scipy.stats.kurtosis(fgs_transit, axis=1)\n    features['fgs_snr'] = (fgs_unobscured_mean - fgs_transit.mean(axis=1)) / fgs_unobscured_std\n    features_df = pd.DataFrame(features, index=star_info_df.index)\n\n    # --- Fourier & Wavelet Features ---\n    fft_coeffs = np.fft.fft(fgs_transit, axis=1)\n    for i in range(1, 6):\n        features_df[f'fgs_fft_mag_{i}'] = np.abs(fft_coeffs[:, i])\n    for level in range(1, 4):\n        coeffs = pywt.wavedec(fgs_transit, 'db4', level=level, axis=1)\n        features_df[f'fgs_wavelet_std_level{level}'] = np.std(coeffs[0], axis=1)\n        features_df[f'fgs_wavelet_mean_level{level}'] = np.mean(coeffs[0], axis=1)\n\n    # --- Metadata & Polynomial Features ---\n    meta_df = star_info_df.copy().fillna(star_info_df.median())\n    meta_df['log_g_proxy'] = np.log1p(meta_df['Ms']) - 2 * np.log1p(meta_df['Rs'])\n    meta_df['rho_star_proxy'] = meta_df['Ms'] / (meta_df['Rs']**3)\n    \n    poly = PolynomialFeatures(degree=2, include_bias=False, interaction_only=True)\n    poly_cols = ['Rs', 'Ts', 'Mp', 'P']\n    poly_features = poly.fit_transform(meta_df[poly_cols])\n    poly_df = pd.DataFrame(poly_features, columns=poly.get_feature_names_out(poly_cols), index=meta_df.index)\n    \n    final_features_df = pd.concat([features_df, meta_df, poly_df], axis=1)\n    final_features_df = final_features_df.loc[:, ~final_features_df.columns.duplicated()]\n\n    print(f\"Created {final_features_df.shape[1]} features in total.\")\n    return final_features_df.fillna(0)\n\n\nif __name__ == '__main__':\n    \n    # ==========================================================================\n    # CORE LOGIC CORRECTION: The if/else statement is now correctly structured.\n    # ==========================================================================\n    \n    # This check correctly identifies if the script is in submission mode.\n    is_submission_run = True # os.path.exists(config.TEST_STAR_INFO_PATH)\n\n    if is_submission_run:\n        # --- SUBMISSION MODE ---\n        # This block now runs ONLY when submitting. It LOADS pre-trained models.\n        # It will not time out.\n        \n        print(\"\\n\" + \"=\"*50)\n        print(\"======         SUBMISSION MODE          ======\")\n        print(\"=\"*50 + \"\\n\")\n\n        # 1. Load test set info\n        sample_submission = pd.read_csv(config.SAMPLE_SUBMISSION_PATH, index_col='planet_id')\n        test_star_info_df = pd.read_csv(config.TEST_STAR_INFO_PATH, index_col='planet_id')\n        wavelengths_df = pd.read_csv(config.WAVELENGTHS_PATH)\n        target_column_names = wavelengths_df.columns.tolist()\n\n        # 2. Process test set data and create features\n        f_raw_test = f_read_and_preprocess('test', test_star_info_df.index)\n        a_raw_test = a_read_and_preprocess('test', test_star_info_df.index)\n        test_features_df = maximal_feature_engineering(f_raw_test, a_raw_test, test_star_info_df)\n        \n        # 3. Load saved artifacts (models, columns, calibration params)\n        print(\"Loading saved models and parameters from /kaggle/working/ ...\")\n        with open(os.path.join(config.OUTPUT_PATH, 'feature_columns.pkl'), 'rb') as f:\n            train_cols = pickle.load(f)\n        with open(os.path.join(config.OUTPUT_PATH, 'calibration_params.pkl'), 'rb') as f:\n            calibration_params = pickle.load(f)\n        \n        trained_models = {}\n        for q in config.QUANTILES:\n            model_path = os.path.join(config.OUTPUT_PATH, f'model_quantile_{q}.pkl')\n            print(f\"Loading model: {model_path}\")\n            with open(model_path, 'rb') as f:\n                trained_models[q] = pickle.load(f)\n        \n        # 4. Ensure test feature columns match training columns\n        test_features_df = test_features_df[train_cols]\n\n        # 5. Make predictions on the test set\n        test_quantile_preds = {}\n        for q in config.QUANTILES:\n            print(f\"Predicting for quantile {q}...\")\n            test_quantile_preds[q] = trained_models[q].predict(test_features_df)\n\n        # 6. Apply calibration and create final submission arrays\n        y_pred_test = test_quantile_preds[0.50].clip(0, None) # Clip mean prediction to be non-negative\n        lower_test, upper_test = test_quantile_preds[0.05], test_quantile_preds[0.95]\n        \n        sigma_raw_test = (upper_test - lower_test) / 3.29 # Corresponds to z-score for 90% confidence\n        sigma_raw_test[sigma_raw_test < 0] = 1e-10\n        \n        sigma_pred_test = (sigma_raw_test * calibration_params['scaling']) + calibration_params['additive']\n        \n        # 7. Format into submission DataFrame\n        print(\"Creating submission file...\")\n        pred_df = pd.DataFrame(y_pred_test, index=sample_submission.index, columns=target_column_names)\n        sigma_df = pd.DataFrame(sigma_pred_test, index=sample_submission.index, columns=[f\"sigma_{i+1}\" for i in range(len(target_column_names))])\n        submission_df = pd.concat([pred_df, sigma_df], axis=1)\n        \n        submission_df.to_csv('submission.csv')\n        print(\"\\n'submission.csv' created successfully!\")\n        print(\"\\nSubmission file preview:\")\n        print(submission_df.head())\n\n    else:\n        # --- TRAINING MODE ---\n        # This block now runs ONLY when you are in the interactive editor.\n        # It TRAINS models and SAVES artifacts to /kaggle/working/\n        \n        print(\"\\n\" + \"=\"*50)\n        print(\"======          TRAINING MODE           ======\")\n        print(\"=\"*50 + \"\\n\")\n        \n        # 1. Load training data and create features\n        train_labels_df = pd.read_csv(config.TRAIN_LABELS_PATH, index_col='planet_id')\n        train_star_info_df = pd.read_csv(config.TRAIN_STAR_INFO_PATH, index_col='planet_id').loc[train_labels_df.index]\n        \n        # For training, we load pre-processed data to save time. \n        # Make sure the 'ariel-data-challenge-2025-af-npy' dataset is added to your notebook.\n        print(\"Loading pre-processed training data...\")\n        f_raw_train, a_raw_train = np.load(config.F_RAW_PATH), np.load(config.A_RAW_PATH)\n        \n        train_features_df = maximal_feature_engineering(f_raw_train, a_raw_train, train_star_info_df)\n        train_labels = train_labels_df.values\n        naive_mu_train, naive_sigma_train = np.mean(train_labels), np.std(train_labels)\n\n        # 2. Perform a validation split to find calibration factors\n        X_train, X_val, y_train, y_val = train_test_split(\n            train_features_df, train_labels, test_size=config.VALIDATION_SPLIT, random_state=config.RANDOM_STATE\n        )\n        \n        val_quantile_preds = {}\n        for q in config.QUANTILES:\n            print(f\"\\n--- [Validation] Training model for quantile: {q} ---\")\n            model = xgb.XGBRegressor(**config.XGB_PARAMS, objective='reg:quantileerror', quantile_alpha=q)\n            wrapper = MultiOutputRegressor(model, n_jobs=-1)\n            wrapper.fit(X_train, y_train)\n            val_quantile_preds[q] = wrapper.predict(X_val)\n\n        # 3. Find the best calibration factors on the validation set using the CORRECT metric\n        y_pred_val, lower_val, upper_val = val_quantile_preds[0.50], val_quantile_preds[0.05], val_quantile_preds[0.95]\n        sigma_raw_val = (upper_val - lower_val) / 3.29\n        sigma_raw_val[sigma_raw_val < 0] = 1e-10\n\n        print(\"\\n--- [Validation] Searching for best calibration factors using OFFICIAL metric... ---\")\n        best_score, best_scaling, best_additive = -1.0, 1.0, 0.0\n        for scaling in config.CALIBRATION_SCALING_FACTORS:\n            for additive in config.CALIBRATION_ADDITIVE_FACTORS:\n                sigma_calibrated = (sigma_raw_val * scaling) + additive\n                # USE THE CORRECT, OFFICIAL SCORING FUNCTION\n                score = official_competition_score(y_val, y_pred_val, sigma_calibrated, naive_mu_train, naive_sigma_train)\n                \n                print(f\"  - Testing Scale={scaling}, Add={additive} -> Score: {score:.4f}\")\n                if score > best_score:\n                    best_score, best_scaling, best_additive = score, scaling, additive\n        \n        print(f\"\\n--- Best factors found: Scale={best_scaling}, Add={best_additive} (Official CV Score: {best_score:.4f}) ---\")\n        \n        # 4. Save the determined calibration factors and feature columns\n        os.makedirs(config.OUTPUT_PATH, exist_ok=True)\n        calibration_params = {'scaling': best_scaling, 'additive': best_additive}\n        with open(os.path.join(config.OUTPUT_PATH, 'calibration_params.pkl'), 'wb') as f:\n            pickle.dump(calibration_params, f)\n        \n        with open(os.path.join(config.OUTPUT_PATH, 'feature_columns.pkl'), 'wb') as f:\n            pickle.dump(train_features_df.columns.tolist(), f)\n\n        # 5. Retrain final models on ALL available training data and save them\n        print(\"\\n--- Retraining models on full dataset for submission... ---\")\n        for q in config.QUANTILES:\n            print(f\"\\n--- [Full Data] Training model for quantile: {q} ---\")\n            model = xgb.XGBRegressor(**config.XGB_PARAMS, objective='reg:quantileerror', quantile_alpha=q)\n            wrapper = MultiOutputRegressor(model, n_jobs=-1)\n            wrapper.fit(train_features_df, train_labels) # Use ALL data\n            model_path = os.path.join(config.OUTPUT_PATH, f'model_quantile_{q}.pkl')\n            print(f\"Saving model to: {model_path}\")\n            with open(model_path, 'wb') as f:\n                pickle.dump(wrapper, f)\n        \n        print(\"\\nTraining complete. All necessary artifacts saved to /kaggle/working/\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-06-28T07:55:19.768149Z","iopub.execute_input":"2025-06-28T07:55:19.768483Z","iopub.status.idle":"2025-06-28T07:55:48.98969Z","shell.execute_reply.started":"2025-06-28T07:55:19.768458Z","shell.execute_reply":"2025-06-28T07:55:48.98888Z"}},"outputs":[],"execution_count":null}]}