{"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":"gpu","dataSources":[{"sourceId":101849,"databundleVersionId":12846694,"sourceType":"competition"},{"sourceId":12303949,"sourceType":"datasetVersion","datasetId":7755350},{"sourceId":12309150,"sourceType":"datasetVersion","datasetId":7758622},{"sourceId":12388933,"sourceType":"datasetVersion","datasetId":7812079}],"dockerImageVersionId":31040,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport time\nimport pickle\nimport warnings\nwarnings.filterwarnings('ignore')\n\nimport numpy as np\nimport pandas as pd\nimport polars as pl\nimport scipy.stats\nimport pywt\nfrom sklearn.model_selection import GroupKFold, train_test_split\nfrom sklearn.isotonic import IsotonicRegression\nfrom sklearn.linear_model import Ridge\nfrom sklearn.multioutput import MultiOutputRegressor\nfrom sklearn.metrics import mean_squared_error\nimport xgboost as xgb\nimport lightgbm as lgb\nfrom tqdm import tqdm\n\nclass Config:\n    # Paths\n    DATA_PATH = '/kaggle/input/ariel-data-challenge-2025'\n    PREP_PATH = '/kaggle/input/ariel-data-challenge-2025-af-npy/'\n    OUT_PATH  = '/kaggle/working/'\n    # Files\n    TRAIN_LABELS = os.path.join(DATA_PATH, 'train.csv')\n    STAR_INFO    = os.path.join(DATA_PATH, 'train_star_info.csv')\n    TEST_INFO    = os.path.join(DATA_PATH, 'test_star_info.csv')\n    SAMPLE_SUB   = os.path.join(DATA_PATH, 'sample_submission.csv')\n    WAVES        = os.path.join(DATA_PATH, 'wavelengths.csv')\n    # CV\n    N_FOLDS      = 5\n    RANDOM_STATE = 42\n    # Model params\n    XGB_PARAMS = {\n        'objective': 'reg:squarederror',\n        'n_estimators': 1000,\n        'learning_rate': 0.03,\n        'max_depth': 6,\n        'subsample': 0.8,\n        'colsample_bytree': 0.7,\n        'tree_method': 'hist',\n        'random_state': RANDOM_STATE,\n        'verbosity': 0\n    }\n    LGB_PARAMS = {\n        'objective': 'regression',\n        'learning_rate': 0.02,\n        'num_leaves': 128,\n        'feature_fraction': 0.8,\n        'bagging_fraction': 0.8,\n        'bagging_freq': 5,\n        'min_data_in_leaf': 30,\n        'n_estimators': 2000,\n        'verbose': -1,\n        'seed': RANDOM_STATE\n    }\n    EARLY_STOPPING = 50\n\ncfg = Config()\n\ndef load_raw(path, dataset, planet_ids, signal_file, normalize_factor):\n    arr = np.zeros((len(planet_ids), normalize_factor.size), dtype=np.float32)\n    for i, pid in enumerate(tqdm(planet_ids, desc=f\"Loading {signal_file}\")):\n        df = pl.read_parquet(f\"{path}/{dataset}/{int(pid)}/{signal_file}\")\n        sig = df.cast(pl.Int32).sum_horizontal().to_numpy().astype(np.float32)\n        diff = sig[1::2] - sig[0::2]\n        arr[i] = diff\n    return arr\n\ndef make_features(f_raw, a_raw, star_info):\n    N = f_raw.shape[1]\n    transit = f_raw[:, 23500:44000]\n    feats = {}\n    # rolling mean/std on transit in windows of 1000\n    window = 1000\n    for i in range((transit.shape[1] // window)):\n        seg = transit[:, i*window:(i+1)*window]\n        feats[f'roll_mean_{i}'] = seg.mean(1)\n        feats[f'roll_std_{i}']  = seg.std(1)\n    # wavelength‑weighted means\n    waves = pd.read_csv(cfg.WAVES)\n    weights = waves.iloc[:, 1].to_numpy()\n    feats['fgs_wmean'] = (transit * weights[:transit.shape[1]]).sum(1)/ weights[:transit.shape[1]].sum()\n    # FFT & wavelet as before\n    fft = np.fft.rfft(transit, axis=1)\n    for k in range(1, 6):\n        feats[f'fft_{k}'] = np.abs(fft[:, k])\n    for lvl in range(1, 4):\n        c = pywt.wavedec(transit, 'db4', level=lvl, axis=1)[0]\n        feats[f'wav_std_{lvl}'] = c.std(1)\n    # AIRS block as before\n    a_tr = a_raw[:, 500:4500]\n    feats['airs_mean'] = a_tr.mean(1)\n    feats['airs_skew'] = scipy.stats.skew(a_tr, axis=1)\n    fft_a = np.fft.rfft(a_tr, axis=1)\n    for k in range(1, 6):\n        feats[f'a_fft_{k}'] = np.abs(fft_a[:, k])\n    # metadata\n    meta = star_info.fillna(star_info.median())\n    meta['Mp_P'] = meta['Mp']/(meta['P']+1e-6)\n    meta['logRs_logTs'] = np.log1p(meta['Rs'])*np.log1p(meta['Ts'])\n    X = pd.DataFrame(feats, index=star_info.index)\n    X = pd.concat([X, meta], axis=1).fillna(0)\n    return X\n\ndef official_score(y_true, y_pred, sigma, mu0, sigma0):\n    y_true,y_pred,sigma = map(np.array, (y_true,y_pred,sigma))\n    base = scipy.stats.norm.logpdf(y_true, loc=y_true, scale=sigma0)\n    glp  = scipy.stats.norm.logpdf(y_true, loc=y_pred, scale=sigma)\n    gm   = scipy.stats.norm.logpdf(y_true, loc=mu0,     scale=sigma0)\n    return ((glp - gm)/(base - gm + 1e-9)).mean()\n\n# -------------------------\n# 1) TRAIN + CALIBRATE\n# -------------------------\nlabels = pd.read_csv(cfg.TRAIN_LABELS, index_col='planet_id')\nstars  = pd.read_csv(cfg.STAR_INFO,    index_col='planet_id').loc[labels.index]\nf_raw  = np.load(cfg.PREP_PATH + 'f_raw_train.npy')\na_raw  = np.load(cfg.PREP_PATH + 'a_raw_train.npy')\nX      = make_features(f_raw, a_raw, stars)\ny      = labels.values\nmu0, sigma0 = y.mean(), y.std()\n\ngroups = pd.qcut(stars['Rs'], cfg.N_FOLDS, labels=False).values\noof_med = np.zeros_like(y)\noof_sig = np.zeros_like(y)\n\nkf = GroupKFold(n_splits=cfg.N_FOLDS)\nfor fold, (tr, val) in enumerate(kf.split(X, groups=groups), 1):\n    print(f\"Fold {fold}/{cfg.N_FOLDS}\")\n    X_tr, X_val = X.iloc[tr], X.iloc[val]\n    y_tr, y_val = y[tr], y[val]\n\n    # --- XGBoost with early stopping ---\n    xgb_model = xgb.XGBRegressor(**cfg.XGB_PARAMS)\n    xgb_model.fit(\n        X_tr, y_tr,\n        eval_set=[(X_val, y_val)],\n        early_stopping_rounds=cfg.EARLY_STOPPING,\n        verbose=False\n    )\n    xgb_pred = xgb_model.predict(X_val)\n\n    # --- LightGBM without early stopping in MultiOutputRegressor ---\n    lgbm = MultiOutputRegressor(lgb.LGBMRegressor(**cfg.LGB_PARAMS))\n    lgbm.fit(X_tr, y_tr)\n    lgb_pred = lgbm.predict(X_val)\n\n    # --- stacking for median ---\n    stacked = np.stack([xgb_pred, lgb_pred], axis=2).mean(2)\n    ridge = Ridge(random_state=cfg.RANDOM_STATE).fit(stacked, y_val)\n    med = ridge.predict(stacked)\n    oof_med[val] = med\n\n    # --- raw sigma from quantile LightGBM (also without early stopping) ---\n    lower = MultiOutputRegressor(\n        lgb.LGBMRegressor(**{**cfg.LGB_PARAMS, 'objective':'quantile', 'alpha':0.05})\n    ).fit(X_tr, y_tr).predict(X_val)\n    upper = MultiOutputRegressor(\n        lgb.LGBMRegressor(**{**cfg.LGB_PARAMS, 'objective':'quantile', 'alpha':0.95})\n    ).fit(X_tr, y_tr).predict(X_val)\n    oof_sig[val] = (upper - lower) / 3.29\n\n# calibrate per output dimension\ncalibrators = []\nfor dim in range(y.shape[1]):\n    ir = IsotonicRegression(out_of_bounds='clip')\n    ir.fit(oof_sig[:,dim], np.abs(y[:,dim]-oof_med[:,dim]))\n    calibrators.append(ir)\n# save calibrators & feature list\nwith open(cfg.OUT_PATH + 'calibrators.pkl', 'wb') as f:\n    pickle.dump(calibrators, f)\npickle.dump(X.columns.tolist(), open(cfg.OUT_PATH + 'features.pkl','wb'))\n\nscore = official_score(y, oof_med, np.column_stack([c.predict(oof_sig[:,d]) for d,c in enumerate(calibrators)]), mu0, sigma0)\nprint(f\"\\nOOF official score: {score:.6f}\")\n\n# -------------------------\n# 2) MAKE SUBMISSION\n# -------------------------\nsample = pd.read_csv(cfg.SAMPLE_SUB, index_col='planet_id')\nstars_test = pd.read_csv(cfg.TEST_INFO, index_col='planet_id')\nf_test = load_raw(cfg.DATA_PATH, 'test', stars_test.index, 'FGS1_signal_0.parquet', np.arange(67500))\na_test = load_raw(cfg.DATA_PATH, 'test', stars_test.index, 'AIRS-CH0_signal_0.parquet', np.arange(5625))\nX_test = make_features(f_test, a_test, stars_test)\n\n# retrain on full data\nxgb_full = xgb.XGBRegressor(**cfg.XGB_PARAMS).fit(X, y)\nlgb_full = MultiOutputRegressor(lgb.LGBMRegressor(**cfg.LGB_PARAMS)).fit(X, y)\npred_x = xgb_full.predict(X_test)\npred_l = lgb_full.predict(X_test)\nstack_full = np.stack([pred_x, pred_l], axis=2).mean(2)\nridge_full = Ridge(random_state=cfg.RANDOM_STATE).fit(\n    np.stack([xgb_full.predict(X), lgb_full.predict(X)], axis=2).mean(2), y\n)\nmed_full = ridge_full.predict(stack_full)\n\n# quantile sigmas\nlow_full = MultiOutputRegressor(\n    lgb.LGBMRegressor(**{**cfg.LGB_PARAMS, 'objective':'quantile','alpha':0.05})\n).fit(X, y).predict(X_test)\nup_full  = MultiOutputRegressor(\n    lgb.LGBMRegressor(**{**cfg.LGB_PARAMS, 'objective':'quantile','alpha':0.95})\n).fit(X, y).predict(X_test)\nraw_full = (up_full - low_full)/3.29\n\n# apply calibrators\nwith open(cfg.OUT_PATH + 'calibrators.pkl','rb') as f:\n    calibrators = pickle.load(f)\nsig_full = np.column_stack([calibrators[d].predict(raw_full[:,d]) for d in range(raw_full.shape[1])])\n\n# build submission\nwaves = pd.read_csv(cfg.WAVES)\ndf_med = pd.DataFrame(np.clip(med_full, 0, None), index=sample.index, columns=waves.columns)\ndf_sig = pd.DataFrame(sig_full, index=sample.index, columns=[f'sigma_{i+1}' for i in range(waves.shape[1])])\nsubmission = pd.concat([df_med, df_sig], axis=1)\nsubmission.to_csv('submission.csv')\nprint(\"Done! submission.csv created.\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}