{"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"},{"sourceId":12943380,"sourceType":"datasetVersion","datasetId":8190942},{"sourceId":13153798,"sourceType":"datasetVersion","datasetId":7868170}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"### library","metadata":{}},{"cell_type":"code","source":"!cp -r '/kaggle/input/nips-adc-package/package' '/kaggle/working/'\n!pip install --no-index --find-links=/kaggle/working/package scipy==1.16.0","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2025-09-24T05:56:20.045877Z","iopub.execute_input":"2025-09-24T05:56:20.046939Z","iopub.status.idle":"2025-09-24T05:56:23.523505Z","shell.execute_reply.started":"2025-09-24T05:56:20.046901Z","shell.execute_reply":"2025-09-24T05:56:23.522716Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\n\nimport numpy as np\nfrom numpy.polynomial.polynomial import Polynomial\n\nimport glob\n\nimport os\n\nfrom tqdm import tqdm\n\nimport matplotlib.pyplot as plt\n\ntry:\n    from astropy.stats import sigma_clip\nexcept:\n    pass\n\nimport itertools\n\nfrom scipy.optimize import minimize, curve_fit, least_squares\nfrom scipy.sparse import lil_matrix\nfrom scipy.ndimage import gaussian_filter1d\n\nimport torch\n\nif torch.cuda.is_available():\n    import cupy as cp\n\nfrom sklearn.decomposition import PCA\nfrom sklearn.model_selection import KFold\nfrom sklearn.ensemble import GradientBoostingRegressor, GradientBoostingClassifier\nfrom sklearn.metrics import r2_score, accuracy_score\n\nfrom multiprocessing import Pool","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2025-09-24T05:56:23.525258Z","iopub.execute_input":"2025-09-24T05:56:23.526051Z","iopub.status.idle":"2025-09-24T05:56:23.53247Z","shell.execute_reply.started":"2025-09-24T05:56:23.525996Z","shell.execute_reply":"2025-09-24T05:56:23.531672Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### config","metadata":{}},{"cell_type":"code","source":"class CustomConfig:\n    root = '/kaggle/input/ariel-data-challenge-2025/'\n\n    seed = 42\n    \n    n_fold = 4\n\n    split = 'train'#'test'\n\n    crop = 8\n\n    cut_inf = 39\n    cut_sup = 321\n\n    binning = {\n        'AIRS-CH0' : 30,\n        'FGS1' : 30 * 12,\n    }\n\n    columns = pd.read_csv(root + 'sample_submission.csv', index_col = 'planet_id').columns\n\n    threshold = 1e-2\n    sigma = 4e-4\n    fgs_sigma = 8e-4\n\n    if '2024' in root:\n        train = pd.read_csv(root + 'train_labels.csv')\n    else:\n        train = pd.read_csv(root + 'train.csv')\n\n    naive_mean = train.values[:, 1:].mean()\n    naive_sigma = train.values[:, 1:].std()\n\n    fgs_weight = (0.4 / 1.95) * 282\n\n    instruments = [\n        'FGS1',\n        'AIRS-CH0',\n    ]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-24T05:56:23.533333Z","iopub.execute_input":"2025-09-24T05:56:23.533579Z","iopub.status.idle":"2025-09-24T05:56:23.641463Z","shell.execute_reply.started":"2025-09-24T05:56:23.533563Z","shell.execute_reply":"2025-09-24T05:56:23.640798Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### utils","metadata":{}},{"cell_type":"code","source":"# ref.: https://www.kaggle.com/code/gordonyip/calibrating-and-binning-ariel-data/notebook\n\ndef ADC_convert(signal, gain=0.4369, offset=-1000):\n    \"\"\"The Analog-to-Digital Conversion (adc) is performed by the detector to convert\n    the pixel voltage into an integer number. Since we are using the same conversion number \n    this year, we have simply hard-coded it inside. \"\"\"\n    signal = signal.astype(np.float64)\n    signal /= gain\n    signal += offset\n    return signal\n\ndef mask_hot_dead(signal, dead, dark):\n    hot = sigma_clip(\n        dark, sigma=5, maxiters=5\n    ).mask\n    hot = np.tile(hot, (signal.shape[0], 1, 1))\n    dead = np.tile(dead, (signal.shape[0], 1, 1))\n    signal = np.ma.masked_where(dead, signal)\n    signal = np.ma.masked_where(hot, signal)\n    return signal\n\ndef apply_linear_corr(linear_corr,clean_signal):\n    linear_corr = np.flip(linear_corr, axis=0)\n    for x, y in itertools.product(\n                range(clean_signal.shape[1]), range(clean_signal.shape[2])\n            ):\n        poli = np.poly1d(linear_corr[:, x, y])\n        clean_signal[:, x, y] = poli(clean_signal[:, x, y])\n    return clean_signal\n\ndef clean_dark(signal, dead, dark, dt):\n\n    dark = np.ma.masked_where(dead, dark)\n    dark = np.tile(dark, (signal.shape[0], 1, 1))\n\n    signal -= dark* dt[:, np.newaxis, np.newaxis]\n    return signal\n\ndef get_cds(signal):\n    cds = signal[:,1::2,:,:] - signal[:,::2,:,:]\n    return cds\n\ndef bin_obs(cds_signal,binning):\n    cds_transposed = cds_signal.transpose(0,1,3,2)\n    cds_binned = np.zeros((cds_transposed.shape[0], cds_transposed.shape[1]//binning, cds_transposed.shape[2], cds_transposed.shape[3]))\n    for i in range(cds_transposed.shape[1]//binning):\n        cds_binned[:,i,:,:] = np.sum(cds_transposed[:,i*binning:(i+1)*binning,:,:], axis=1)\n    return cds_binned\n\ndef correct_flat_field(flat,dead, signal):\n    flat = flat.transpose(1, 0)\n    dead = dead.transpose(1, 0)\n    flat = np.ma.masked_where(dead, flat)\n    flat = np.tile(flat, (signal.shape[0], 1, 1))\n    signal = signal / flat\n    return signal","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2025-09-24T05:56:23.643086Z","iopub.execute_input":"2025-09-24T05:56:23.643358Z","iopub.status.idle":"2025-09-24T05:56:23.653488Z","shell.execute_reply.started":"2025-09-24T05:56:23.643332Z","shell.execute_reply":"2025-09-24T05:56:23.652805Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def apply_linear_corr_gpu(linear_corr, signal):\n    '''\n    horner’s method in-place : y = (((((a5 * x + a4) * x + a3) * x + a2) * x + a1) * x + a0)\n    '''\n    \n    linear_corr = cp.asarray(linear_corr)\n    signal = cp.asarray(signal)    \n    out = cp.full_like(signal, linear_corr[-1])\n\n    for i in range(linear_corr.shape[0] - 2, -1, -1):\n        cp.multiply(out, signal, out = out)    \n        cp.add(out, linear_corr[i], out = out)\n    \n    out = cp.asnumpy(out)\n    return out\n\nif __name__ == '__main__':\n    '''\n    linear_corr = np.random.randn(6, 32, 32)\n    signal = np.random.randn(135000, 32, 32)\n\n    x1 = apply_linear_corr(linear_corr.copy(), signal.copy())\n    x2 = apply_linear_corr_gpu(linear_corr.copy(), signal.copy())\n\n    print(np.allclose(x1, x2))\n    '''","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2025-09-24T05:56:23.654231Z","iopub.execute_input":"2025-09-24T05:56:23.654443Z","iopub.status.idle":"2025-09-24T05:56:23.671626Z","shell.execute_reply.started":"2025-09-24T05:56:23.65442Z","shell.execute_reply":"2025-09-24T05:56:23.671052Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_solution(args, train, planet_ids):\n    rows = []\n    for planet_id in planet_ids:\n        true = train[train.planet_id == planet_id].values[0, 1:]\n        \n        row = {'planet_id' : planet_id}\n        for i, column in enumerate(args.columns[:283]):\n            row[column] = true[i]\n\n        rows.append(row)\n\n    solution = pd.DataFrame(rows)\n    return solution\n\ndef get_score(args, submission, solution = None, print_score = True):            \n    planet_ids = [int(_) for _ in submission.planet_id.values]\n\n    if '2024' in args.root:\n        train = pd.read_csv('/kaggle/input/ariel-data-challenge-2024/train_labels.csv')\n    else:\n        train = pd.read_csv('/kaggle/input/ariel-data-challenge-2025/train.csv')\n\n    if solution is None:\n        solution = get_solution(args, train, planet_ids)\n\n    scores, scores2, score = score_function(\n        solution = solution.copy(),\n        submission = submission.copy(),\n        row_id_column_name = 'planet_id',\n        naive_mean = args.naive_mean,\n        naive_sigma = args.naive_sigma,\n        fgs_weight = args.fgs_weight,\n        print_score = print_score,\n    )\n    return scores, scores2, score","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2025-09-24T05:56:23.672425Z","iopub.execute_input":"2025-09-24T05:56:23.672632Z","iopub.status.idle":"2025-09-24T05:56:23.692826Z","shell.execute_reply.started":"2025-09-24T05:56:23.672609Z","shell.execute_reply":"2025-09-24T05:56:23.692218Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def remove_outlier(data, window = 100, threshold = 5.0):\n    df = pd.DataFrame(data)\n    \n    mean = df.rolling(window = window, center = True, min_periods = 1).mean()\n    std = df.rolling(window = window, center = True, min_periods = 1).std()\n    \n    mask = (df - mean).abs() > (threshold * std)\n    \n    df[mask] = mean[mask]\n    data = df.to_numpy()\n    return data\n\ndef bin_function(data, binning):\n    _data = np.zeros((data.shape[0]//binning, data.shape[1]))\n\n    for i in range(data.shape[0]//binning):\n        _data[i, :] = np.mean(data[i * binning:(i + 1) * binning, :], axis = 0)\n        \n    return _data\n\ndef pca_function(x, n_components):\n    pca = PCA()\n    \n    _x = pca.fit_transform(x)\n    _x[:, n_components:] = 0\n    x = pca.inverse_transform(_x)[:, :]\n    \n    return x","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2025-09-24T05:56:23.693507Z","iopub.execute_input":"2025-09-24T05:56:23.693704Z","iopub.status.idle":"2025-09-24T05:56:23.706253Z","shell.execute_reply.started":"2025-09-24T05:56:23.693689Z","shell.execute_reply":"2025-09-24T05:56:23.705469Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### metric","metadata":{}},{"cell_type":"code","source":"# ref.: https://www.kaggle.com/code/metric/ariel-gaussian-log-likelihood\n\nimport numpy as np\nimport pandas as pd\nimport pandas.api.types\nimport scipy.stats\n\n\nclass ParticipantVisibleError(Exception):\n    pass\n\n\ndef score_function(\n    solution: pd.DataFrame,\n    submission: pd.DataFrame,\n    row_id_column_name: str,\n    naive_mean: float,\n    naive_sigma: float,\n    fsg_sigma_true: float = 1e-6,\n    airs_sigma_true: float = 1e-5,\n    fgs_weight: float = 1,\n    print_score : bool = True,\n) -> float:\n    \"\"\"\n    This is a Gaussian Log Likelihood based metric. For a submission, which contains the predicted mean (x_hat) and variance (x_hat_std),\n    we calculate the Gaussian Log-likelihood (GLL) value to the provided ground truth (x). We treat each pair of x_hat,\n    x_hat_std as a 1D gaussian, meaning there will be 283 1D gaussian distributions, hence 283 values for each test spectrum,\n    the GLL value for one spectrum is the sum of all of them.\n\n    Inputs:\n        - solution: Ground Truth spectra (from test set)\n            - shape: (nsamples, n_wavelengths)\n        - submission: Predicted spectra and errors (from participants)\n            - shape: (nsamples, n_wavelengths*2)\n        naive_mean: (float) mean from the train set.\n        naive_sigma: (float) standard deviation from the train set.\n        fsg_sigma_true: (float) standard deviation from the FSG1 instrument for the test set.\n        airs_sigma_true: (float) standard deviation from the AIRS instrument for the test set.\n        fgs_weight: (float) relative weight of the fgs channel\n    \"\"\"\n\n    del solution[row_id_column_name]\n    del submission[row_id_column_name]\n\n    if submission.min().min() < 0:\n        raise ParticipantVisibleError('Negative values in the submission')\n    for col in submission.columns:\n        if not pandas.api.types.is_numeric_dtype(submission[col]):\n            raise ParticipantVisibleError(f'Submission column {col} must be a number')\n\n    n_wavelengths = len(solution.columns)\n    if len(submission.columns) != n_wavelengths * 2:\n        raise ParticipantVisibleError('Wrong number of columns in the submission')\n\n    y_pred = submission.iloc[:, :n_wavelengths].values\n    # Set a non-zero minimum sigma pred to prevent division by zero errors.\n    sigma_pred = np.clip(submission.iloc[:, n_wavelengths:].values, a_min=10**-15, a_max=None)\n    sigma_true = np.append(\n        np.array(\n            [\n                fsg_sigma_true,\n            ]\n        ),\n        np.ones(n_wavelengths - 1) * airs_sigma_true,\n    )\n    y_true = solution.values\n\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 * np.ones_like(y_true))\n    GLL_mean = scipy.stats.norm.logpdf(y_true, loc=naive_mean * np.ones_like(y_true), scale=naive_sigma * np.ones_like(y_true))\n\n    # normalise the score, right now it becomes a matrix instead of a scalar.\n    ind_scores = (GLL_pred - GLL_mean) / (GLL_true - GLL_mean)\n    if print_score:\n        print(ind_scores.mean(1))\n\n    weights = np.append(np.array([fgs_weight]), np.ones(len(solution.columns) - 1))\n    weights = weights * np.ones_like(ind_scores)\n    submit_score = np.average(ind_scores, weights=weights)\n\n    scores = ((ind_scores * weights).sum(1) / weights.sum(1))\n    if print_score:\n        print(scores)\n    return scores, ind_scores.mean(1), float(submit_score)#float(np.clip(submit_score, 0.0, 1.0))","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2025-09-24T05:56:23.70709Z","iopub.execute_input":"2025-09-24T05:56:23.707289Z","iopub.status.idle":"2025-09-24T05:56:23.723456Z","shell.execute_reply.started":"2025-09-24T05:56:23.707274Z","shell.execute_reply":"2025-09-24T05:56:23.722736Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### preprocess","metadata":{}},{"cell_type":"code","source":"def read_data(args, planet_id, instrument, observation_count = 0):\n    if '2024' in args.root:\n        signal = pd.read_parquet(args.root + f'{args.split}/{planet_id}/{instrument}_signal.parquet').values.astype(np.float64)\n        \n        flat = pd.read_parquet(args.root + f'{args.split}/{planet_id}/{instrument}_calibration/flat.parquet').values.astype(np.float64)    \n        dark = pd.read_parquet(args.root + f'{args.split}/{planet_id}/{instrument}_calibration/dark.parquet').values.astype(np.float64)    \n        dead = pd.read_parquet(args.root + f'{args.split}/{planet_id}/{instrument}_calibration/dead.parquet').values.astype(np.float64)\n        linear_corr = pd.read_parquet(args.root + f'{args.split}/{planet_id}/{instrument}_calibration/linear_corr.parquet').values.astype(np.float64)\n    else:\n        signal = pd.read_parquet(args.root + f'{args.split}/{planet_id}/{instrument}_signal_{observation_count}.parquet').values.astype(np.float64)\n        \n        flat = pd.read_parquet(args.root + f'{args.split}/{planet_id}/{instrument}_calibration_{observation_count}/flat.parquet').values.astype(np.float64)    \n        dark = pd.read_parquet(args.root + f'{args.split}/{planet_id}/{instrument}_calibration_{observation_count}/dark.parquet').values.astype(np.float64)    \n        dead = pd.read_parquet(args.root + f'{args.split}/{planet_id}/{instrument}_calibration_{observation_count}/dead.parquet').values.astype(np.float64)\n        linear_corr = pd.read_parquet(args.root + f'{args.split}/{planet_id}/{instrument}_calibration_{observation_count}/linear_corr.parquet').values.astype(np.float64)\n\n    if instrument == 'AIRS-CH0':\n        signal = signal.reshape(signal.shape[0], 32, 356)[:, args.crop:32-args.crop, args.cut_inf:args.cut_sup]\n\n        dt = axis_info['AIRS-CH0-integration_time'].dropna().values\n        dt[1::2] += 0.1\n        \n        flat = flat.reshape(32, 356)[args.crop:32-args.crop, args.cut_inf:args.cut_sup]    \n        dark = dark.reshape(32, 356)[args.crop:32-args.crop, args.cut_inf:args.cut_sup]     \n        dead = dead.reshape(32, 356)[args.crop:32-args.crop, args.cut_inf:args.cut_sup]     \n        linear_corr = linear_corr.reshape(6, 32, 356)[:, args.crop:32-args.crop, args.cut_inf:args.cut_sup] \n    else:\n        signal = signal.reshape(signal.shape[0], 32, 32)[:, args.crop:32-args.crop, args.crop:32-args.crop]\n\n        dt = np.ones(signal.shape[0]) * 0.1\n        dt[1::2] += 0.1\n        \n        flat = flat.reshape(32, 32)[args.crop:32-args.crop, args.crop:32-args.crop]\n        dark = dark.reshape(32, 32)[args.crop:32-args.crop, args.crop:32-args.crop]\n        dead = dead.reshape(32, 32)[args.crop:32-args.crop, args.crop:32-args.crop]\n        linear_corr = linear_corr.reshape(6, 32, 32)[:, args.crop:32-args.crop, args.crop:32-args.crop]\n    return signal, dt, flat, dark, dead, linear_corr\n\ndef calibrate_data(args, signal, dt, flat, dark, dead, linear_corr, binning):\n    signal = ADC_convert(signal)\n    \n    #signal = mask_hot_dead(signal, dead, dark)\n\n    if torch.cuda.is_available():\n        signal = apply_linear_corr_gpu(linear_corr, signal)\n    else:\n        signal = apply_linear_corr(linear_corr, signal)\n    \n    signal = clean_dark(signal, dead, dark, dt)\n\n    signal = signal[None]\n\n    signal = get_cds(signal)\n    \n    signal = signal.transpose(0, 1, 3, 2)\n\n    signal = signal[0]\n\n    signal = correct_flat_field(flat, dead, signal)\n    signal = signal.filled(fill_value = np.nan)\n\n    signal = signal[:, ::-1]\n    return signal\n\ndef preprocess_function(args, planet_id, observation_count = 0):\n    data = []\n    for instrument in args.instruments:\n        _data, dt, flat, dark, dead, linear_corr = read_data(args, planet_id, instrument, observation_count)\n        _data = calibrate_data(args, _data, dt, flat, dark, dead, linear_corr, args.binning[instrument])           \n        _data = np.nanmean(_data, axis = 2)\n    \n        _data = remove_outlier(_data)\n        if args.binning[instrument] != None:\n            _data = bin_function(_data, args.binning[instrument])\n\n        if instrument == 'FGS1':\n            _data = np.mean(_data, axis = 1, keepdims = True)\n\n        data.append(_data)\n\n    data = np.concatenate(data, axis = 1)\n    return data\n\nif __name__ == '__main__':\n    args = CustomConfig()\n\n    args.root = '/kaggle/input/ariel-data-challenge-2025/'\n \n    axis_info = pd.read_parquet(args.root + 'axis_info.parquet')\n    \n    planet_ids = glob.glob(args.root + f'{args.split}/*')\n    planet_ids = sorted([_.split('/')[-1] for _ in planet_ids])\n\n    #'''\n    if args.split == 'train':\n        cv = pd.read_csv('/kaggle/input/nips-adc-data/cv6.csv')\n        planet_ids = cv[(cv.weighted_score <= 0.44) & (cv.success == True)]['planet_id'].tolist()\n        print('n_planet : ', len(planet_ids))\n    #'''\n\n    index = np.random.randint(0, len(planet_ids))\n    planet_id = planet_ids[index]\n    print('index : ', index, ', planet_id : ', planet_id)\n    \n    data = preprocess_function(args, planet_id)\n    print('data : ', data.shape)\n\n    plt.plot(data.mean(1))\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-24T05:56:23.724351Z","iopub.execute_input":"2025-09-24T05:56:23.724615Z","iopub.status.idle":"2025-09-24T05:56:32.900372Z","shell.execute_reply.started":"2025-09-24T05:56:23.724592Z","shell.execute_reply":"2025-09-24T05:56:32.89961Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### stage1","metadata":{}},{"cell_type":"code","source":"def model1(time, T1, T2, T3, T4, depth, dark, *coeffs):\n    def get_transit(time, T1, T2, T3, T4, depth, dark):\n        transit = np.ones_like(time)\n\n        ingress_mask = (time >= T1) & (time <= T2)\n        egress_mask = (time >= T3) & (time <= T4)\n        transit_mask = (time > T2) & (time < T3)\n\n        transit[ingress_mask] = 1.0 - depth * (time[ingress_mask] - T1) / (T2 - T1)\n        transit[egress_mask] = (1.0 - depth) + depth * (time[egress_mask] - T3) / (T4 - T3)  \n        transit[transit_mask] = 1.0 - depth + dark * (time[transit_mask] - T2) * (time[transit_mask] - T3)\n        return transit\n    \n    system = Polynomial(coeffs)(time)\n    transit = get_transit(time, T1, T2, T3, T4, depth, dark)\n    return system + transit\n\ndef loss_function(params, time, data):\n    pred = model1(time, *params)\n    loss = np.mean((data - pred) ** 2)\n    return loss\n\ndef stage1_function(args, data, degree = 3):\n    if 'FGS1' not in args.instruments:\n        data = data.mean(1)\n    else:\n        data = data[:, 1:].mean(1)\n    data = data - data.mean()\n    data = data / data.std()\n\n    time = np.linspace(0, 1, data.shape[0])\n    \n    x0 = [\n        0.10,               # T1\n        0.25,               # T2\n        0.75,               # T3\n        0.90,               # T4\n        1,                  # depth\n        1,                  # dark\n    ] + [0] * (degree + 1)  # coeffs\n\n    constraints = [\n        {'type' : 'ineq', 'fun' : lambda p : p[1] - p[0] - 1e-2}, # T2 > T1\n        {'type' : 'ineq', 'fun' : lambda p : p[2] - p[1] - 1e-2}, # T3 > T2\n        {'type' : 'ineq', 'fun' : lambda p : p[3] - p[2] - 1e-2}, # T4 > T3\n    ]\n\n    bounds = [\n        (0, 1),                      # T1\n        (0, 1),                      # T2\n        (0, 1),                      # T3\n        (0, 1),                      # T4\n        (0, 1e1),                    # depth\n        (0, 1e1),                    # dark\n    ] + [(-1e1, 1e1)] * (degree + 1) # coeffs\n    \n    result = minimize(\n        fun = loss_function,   \n        x0 = x0,           \n        args = (time, data),\n        method = 'SLSQP',        \n        constraints = constraints,\n        bounds = bounds,\n        options = {\n            'maxiter' : 500,\n            'disp' : False,\n        },\n        tol = 1e-16,\n    )\n\n    params = result.x\n    if args.threshold != None:\n        success = result.fun < args.threshold\n    else:\n        success = True\n\n    success = success and (params[0] > 0.01) and (params[3] < 0.99)\n    return success, params\n\nif __name__ == '__main__':      \n    stage1_success, stage1_params = stage1_function(args, data)\n    print('params : ', [round(_, 4) for _ in stage1_params])\n    print('success : ', stage1_success)\n\n    true = data.mean(1)\n    true -= true.mean()\n    true /= true.std()\n\n    time = np.linspace(0, 1, data.shape[0]) \n    pred1 = model1(time, *stage1_params)\n    pred2 = model1(time, *stage1_params[:4], 0, 0, *stage1_params[6:])\n    plt.plot(true, color = 'g', alpha = 0.5)\n    plt.plot(pred1, color = 'r', alpha = 0.5)\n    plt.plot(pred2, color = 'y', alpha = 0.5)\n    plt.show()","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2025-09-24T05:56:32.902557Z","iopub.execute_input":"2025-09-24T05:56:32.902828Z","iopub.status.idle":"2025-09-24T05:56:33.191976Z","shell.execute_reply.started":"2025-09-24T05:56:32.902811Z","shell.execute_reply":"2025-09-24T05:56:33.191218Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### stage2","metadata":{}},{"cell_type":"code","source":"\"\"\"def area(d, x, R):\n    _area = np.zeros_like(d + x, dtype = float)\n\n    mask1 = (d <= np.abs(R - x))\n    mask2 = (d > np.abs(R - x)) & (d < (R + x))\n\n    _area[mask1] = (np.pi * (np.minimum(x, R) ** 2))[mask1]\n\n    with np.errstate(divide = 'ignore', invalid = 'ignore'):\n        arg1 = (d ** 2 + x ** 2 - R ** 2) / (2 * d * x)\n        arg2 = (d ** 2 + R ** 2 - x ** 2) / (2 * d * R)\n        arg3 = (-d + x + R) * (d + x - R) * (d - x + R) * (d + x + R)\n\n    arg1 = np.clip(arg1, -1, 1)\n    arg2 = np.clip(arg2, -1, 1)\n    arg3 = np.clip(arg3, 0, None)\n\n    _area[mask2] = (x ** 2 * np.arccos(arg1) + R ** 2 * np.arccos(arg2) - 0.5 * np.sqrt(arg3))[mask2]\n    return _area\n\ndef intensity(x, args):\n    c1, c2, c3, c4 = args\n\n    norm = (- c1 / 10 - c2 / 6 - 3 * c3 / 14 - c4 / 4 + 0.5) * 2 * np.pi\n\n    x = np.clip(x, None, 0.99995)\n    sqrtmu = (1 - x ** 2) ** 0.25\n    return (1 - c1 * (1 - sqrtmu) - c2 * (1 - sqrtmu ** 2) - c3 * (1 - sqrtmu ** 3) - c4 * (1 - sqrtmu ** 4)) / norm\n\ndef calc_limb_darkening(d_array, rprs, intensity_args, n_step = 5):\n    f_array = np.ones_like(d_array, dtype = float)\n\n    x_in_array  = np.maximum(d_array - rprs, 0.0)\n    x_out_array = np.minimum(d_array + rprs, 1.0)\n\n    mask = (x_in_array < 1.0) & ((x_out_array - x_in_array) >= 1e-7)\n\n    if not np.any(mask):\n        return f_array\n\n    d = d_array[mask]\n    x_in = x_in_array[mask]\n    x_out = x_out_array[mask]\n\n    x = np.linspace(x_in, x_out, n_step, axis = 1)\n    dx = x[:, 1:] - x[:, :-1]\n\n    Int = intensity(x[:, 1:] - dx / 2, intensity_args)\n\n    A = area(d[:, np.newaxis], x, rprs)\n    delta = np.sum((A[:, 1:] - A[:, :-1]) * Int, axis = 1)\n\n    f_array[mask] = 1.0 - delta\n    return f_array\n\ndef get_transit(t, star_info, rp, c1, c2, c3, c4, P, sma, i, t0):\n    f = 2 * np.pi * (t - t0) / P\n    b = sma * np.cos(np.radians(i))\n    d = np.sqrt((sma * np.sin(f)) ** 2 + (b * np.cos(f)) ** 2)\n\n    transit = calc_limb_darkening(d, rp, [c1, c2, c3, c4])\n    return transit\n\ndef get_system(t, c, *coeffs):\n    system_time = Polynomial(coeffs[:len(coeffs) // 2])(t)\n    system_channel = Polynomial(coeffs[len(coeffs) // 2:])(c)\n\n    system = system_time[:, None] * system_channel[None, :]\n    system = 1 + system\n    return system\n\ndef get_jac_sparsity(param_info, n_time, n_channel):\n    m = n_time * n_channel\n    n = param_info['n_param']\n\n    jac_sparsity = lil_matrix((m, n), dtype = int)\n    for param in ['P', 'sma', 'i', 't0', 'sys']:\n        start = param_info[param][0]\n        end = param_info[param][1]\n        jac_sparsity[:, start:end] = 1\n\n    for param in ['rp', 'c1', 'c2', 'c3', 'c4']:\n        for j in range(n_channel):\n            jac_sparsity[\n                np.arange(j, m, n_channel),\n                param_info[param][0] + j,\n            ] = 1\n            \n    return jac_sparsity\n\ndef model2(params, t, star_info, param_info, n_channel):\n    c = np.linspace(0, 1, n_channel)\n\n    rp = params[param_info['rp'][0]:param_info['rp'][1]]\n\n    c1 = ((np.tanh(params[param_info['c1'][0]:param_info['c1'][1]]) + 1.0) / 2) * param_info['ci_max']\n    c2 = ((np.tanh(params[param_info['c2'][0]:param_info['c2'][1]]) + 1.0) / 2) * param_info['ci_max']\n    c3 = ((np.tanh(params[param_info['c3'][0]:param_info['c3'][1]]) + 1.0) / 2) * param_info['ci_max']\n    c4 = ((np.tanh(params[param_info['c4'][0]:param_info['c4'][1]]) + 1.0) / 2) * param_info['ci_max']\n    \n    P = params[param_info['P'][0]:param_info['P'][1]]\n    sma = params[param_info['sma'][0]:param_info['sma'][1]]\n    i = params[param_info['i'][0]:param_info['i'][1]]\n    t0 = params[param_info['t0'][0]:param_info['t0'][1]]\n\n    sys = params[param_info['sys'][0]:param_info['sys'][1]]\n\n    transit = []\n    for j in range(n_channel):\n        _transit = get_transit(t, star_info, rp[j], c1[j], c2[j], c3[j], c4[j], P, sma, i, t0)\n        transit.append(_transit)\n\n    transit = np.stack(transit, axis = 1)\n    system = get_system(t, c, *sys)\n\n    pred = transit * system\n    return pred\n\ndef fun(params, t, star_info, param_info, true):\n    n_channel = true.shape[1]\n\n    pred = model2(params, t, star_info, param_info, n_channel)\n    loss = true - pred\n    return loss.ravel()\n\ndef process_single(inputs):\n    args, ydata, time, star_info, T1, T2, T3, T4, degree, window, ci_max = inputs\n\n    system_mask = (time < T1) | (time > T4)\n    transit_mask = (time > T2) & (time < T3)\n\n    target = ydata[:, 0]\n    target /= target[system_mask].mean()\n\n    def model(t, star_info, rp, c1, c2, c3, c4, P, sma, i, t0, *coeffs):\n        f = 2 * np.pi * (t - t0) / P\n        b = sma * np.cos(np.radians(i))\n        d = np.sqrt((sma * np.sin(f)) ** 2 + (b * np.cos(f)) ** 2)\n\n        system = Polynomial(coeffs)(t)\n        transit = calc_limb_darkening(d, rp, [c1, c2, c3, c4])\n        return system, transit\n\n    def f(xdata, rp, c1, c2, c3, c4, P, sma, i, t0, *coeffs):\n        system, transit = model(xdata, star_info, rp, c1, c2, c3, c4, P, sma, i, t0, *coeffs)\n        return system * transit\n\n    rp_0 = np.sqrt(1 - target[transit_mask].mean())\n\n    ci_0 = [1e-1] * 4\n\n    P_0 = star_info['P'] * 24 / 7.5\n    sma_0 = star_info['sma']\n    i_0 = star_info['i']\n    t0_0 = (T1 + T4) / 2\n\n    sysi_0 = [0] * (degree + 1)\n\n    p0 = [rp_0] + ci_0 + [P_0, sma_0, i_0, t0_0] + sysi_0\n    bounds = (\n        [rp_0 * (1.0 - 1e-1)] + [0.0] * 4 + [_ * (1.0 - 1e-1) for _ in [P_0, sma_0, i_0, t0_0]] + [-1e1] * (degree + 1),\n        [rp_0 * (1.0 + 1e-1)] + [ci_max] * 4 + [_ * (1.0 + 1e-1) for _ in [P_0, sma_0, i_0, t0_0]] + [+1e1] * (degree + 1),\n    )\n\n    popt, pcov = curve_fit(\n        f = f,\n        xdata = time,\n        ydata = target,\n        p0 = p0,\n        bounds = bounds,\n    )\n    return popt\n\ndef process_multiple(inputs):\n    args, targets, time, star_info, T1, T2, T3, T4, sysi_degree, window, ci_max, plot, params, param_info = inputs\n\n    system_mask = (time < T1) | (time > T4)\n    transit_mask = (time > T2) & (time < T3)\n\n    if params == None:\n        rp_0 = np.sqrt(1 - targets[transit_mask].mean(0)).tolist()\n    \n        ci_0 = [-1.0] * targets.shape[1] * 4\n        \n        P_0 = [star_info['P'] * 24 / 7.5]\n        sma_0 = [star_info['sma']]\n        i_0 = [star_info['i']]\n        t0_0 = [(T1 + T4) / 2]\n    \n        sysi_0 = [1e-3] * (sysi_degree + 1) + [1e-3] * (sysi_degree + 1)\n    \n        x0 = rp_0 + ci_0 + P_0 + sma_0 + i_0 + t0_0 + sysi_0\n    \n        bounds = (\n            [_ * (1.0 - 1e-1) for _ in rp_0] + [-1e1] * len(ci_0) + [_ * (1.0 - 2e-1) for _ in P_0 + sma_0 + i_0 + t0_0] + [-1e1] * len(sysi_0),\n            [_ * (1.0 + 1e-1) for _ in rp_0] + [+1e1] * len(ci_0) + [_ * (1.0 + 2e-1) for _ in P_0 + sma_0 + i_0 + t0_0] + [+1e1] * len(sysi_0),\n        )\n    else:\n        params = params.x\n        \n        rp_0 = params[param_info['rp'][0]:param_info['rp'][1]].tolist()\n    \n        c1_0 = params[param_info['c1'][0]:param_info['c1'][1]].tolist()\n        c2_0 = params[param_info['c2'][0]:param_info['c2'][1]].tolist()\n        c3_0 = params[param_info['c3'][0]:param_info['c3'][1]].tolist()\n        c4_0 = params[param_info['c4'][0]:param_info['c4'][1]].tolist()\n    \n        ci_0 = c1_0 + c2_0 + c3_0 + c4_0\n        \n        P_0 = params[param_info['P'][0]:param_info['P'][1]].tolist()\n        sma_0 = params[param_info['sma'][0]:param_info['sma'][1]].tolist()\n        i_0 = params[param_info['i'][0]:param_info['i'][1]].tolist()\n        t0_0 = params[param_info['t0'][0]:param_info['t0'][1]].tolist()\n    \n        sysi_0 = params[param_info['sys'][0]:param_info['sys'][1]].tolist()\n        sysi_0 = [1e-3] * (sysi_degree - param_info['sysi_degree']) + sysi_0[:param_info['sysi_degree'] + 1] + \\\n                 [1e-3] * (sysi_degree - param_info['sysi_degree']) + sysi_0[param_info['sysi_degree'] + 1:]\n\n        x0 = rp_0 + ci_0 + P_0 + sma_0 + i_0 + t0_0 + sysi_0\n\n        bounds = (\n            [_ * (1.0 - 1e-1) for _ in rp_0] + [-1e1] * len(ci_0) + [_ * (1.0 - 1e-1) for _ in P_0 + sma_0 + i_0 + t0_0] + [-1e1] * len(sysi_0),\n            [_ * (1.0 + 1e-1) for _ in rp_0] + [+1e1] * len(ci_0) + [_ * (1.0 + 1e-1) for _ in P_0 + sma_0 + i_0 + t0_0] + [+1e1] * len(sysi_0),\n        )\n\n    param_info = {\n        'rp' : [0, len(rp_0)],\n        'c1' : [len(rp_0), len(rp_0) + (len(ci_0) // 4)],\n        'c2' : [len(rp_0) + (len(ci_0) // 4), len(rp_0) + 2 * (len(ci_0) // 4)],\n        'c3' : [len(rp_0) + 2 * (len(ci_0) // 4), len(rp_0) + 3 * (len(ci_0) // 4)],\n        'c4' : [len(rp_0) + 3 * (len(ci_0) // 4), len(rp_0) + len(ci_0)],\n        'P' : [len(rp_0) + len(ci_0), len(rp_0) + len(ci_0) + 1],\n        'sma' : [len(rp_0) + len(ci_0) + 1, len(rp_0) + len(ci_0) + 2],\n        'i' : [len(rp_0) + len(ci_0) + 2, len(rp_0) + len(ci_0) + 3],\n        't0' : [len(rp_0) + len(ci_0) + 3, len(rp_0) + len(ci_0) + 4],\n        'sys' : [len(rp_0) + len(ci_0) + 4, len(rp_0) + len(ci_0) + 4 + len(sysi_0)],\n        'n_param' : len(rp_0) + len(ci_0) + 4 + len(sysi_0),\n        'ci_max' : ci_max,\n        'sysi_degree' : sysi_degree,\n    }\n\n    jac_sparsity = get_jac_sparsity(param_info, targets.shape[0], targets.shape[1])\n\n    with Pool(processes = os.cpu_count()) as pool:\n        res = least_squares(\n            fun = fun,\n            x0 = x0,\n            bounds = bounds,\n            method = 'trf',\n            jac_sparsity = jac_sparsity,\n            args = (time, star_info, param_info, targets),\n            workers = pool.map,\n            verbose = 2 if plot else 0,\n        )\n    return res, param_info\n\ndef stage2_function(args, ydata, star_info, T1, T2, T3, T4, ci_max = 0.5, sysi_degree = 3, window = 5, downsample = 4, plot = False):\n    time = np.linspace(0, 1, ydata.shape[0])\n\n    system_mask = (time < T1) | (time > T4)\n    transit_mask = (time > T2) & (time < T3)\n\n    targets = []\n    for j in range(ydata.shape[1]):\n        if j > 200:\n            window = 50\n        \n        if j == 0:\n            target = ydata[:, 0:1]\n        else:\n            target = ydata[:, max(1, j - window):(j + window + 1)]\n\n        target = target.mean(1)\n        target /= target[system_mask].mean()\n        targets.append(target)\n\n    targets = np.stack(targets, axis = 1)\n    targets = targets[:, 1::downsample]\n\n    res, param_info = process_multiple([args, targets, time, star_info, T1, T2, T3, T4, 3, 5, 0.50, plot, None, None])\n    res, param_info = process_multiple([args, targets, time, star_info, T1, T2, T3, T4, 4, 5, 0.50, plot, res, param_info])\n    airs_popt, cost, nfev = res.x, res.cost, res.nfev\n    \n    fgs1_popt = process_single([args, ydata, time, star_info, T1, T2, T3, T4, 3, 5, 0.5])\n\n    means = np.array(airs_popt[param_info['rp'][0]:param_info['rp'][1]]) ** 2\n    x1 = np.linspace(0, 1, 282)\n    x2 = np.linspace(0, 1, means.shape[0])\n    means = np.interp(x1, x2, means).tolist()\n    means = [fgs1_popt[0] ** 2] + means\n\n    sigmas = [args.fgs_sigma] + [args.sigma] * 282\n\n    pred = np.array(means + sigmas)\n    pred = pred.clip(0)\n\n    params = np.concatenate([\n        airs_popt,\n        np.array([cost]),\n        np.array([nfev]),\n        fgs1_popt,\n    ], axis = 0)\n\n    if plot:\n        x1, x2 = np.meshgrid(\n            np.linspace(0, targets.shape[1], targets.shape[1]),\n            np.linspace(0, targets.shape[0], targets.shape[0]),\n        )\n        fig = plt.figure()\n        ax = fig.add_subplot(projection = '3d')\n        ax.plot_surface(x1, x2, targets, cmap = 'viridis', color = 'g', alpha = 0.5)\n        plt.show()\n\n        _pred = model2(airs_popt, time, star_info, param_info, targets.shape[1])\n\n        for j in range(targets.shape[1]):\n            if j % (targets.shape[1] // 4) == 0:\n                plt.plot(targets[:, j], color = 'g', alpha = 0.5)\n                plt.plot(_pred[:, j], color = 'r', alpha = 0.5)\n                plt.show()\n\n        print('airs_popt : ', airs_popt.shape)\n        print('fgs1_popt : ', fgs1_popt.shape)\n\n    return pred, params\n\nif __name__ == '__main__':\n    train_star_info = pd.read_csv('/kaggle/input/ariel-data-challenge-2025/train_star_info.csv')\n    test_star_info = pd.read_csv('/kaggle/input/ariel-data-challenge-2025/test_star_info.csv')\n\n    if args.split == 'train':\n        if '2024' in args.root:\n            train = pd.read_csv('/kaggle/input/ariel-data-challenge-2024/train_labels.csv')\n        else:\n            train = pd.read_csv('/kaggle/input/ariel-data-challenge-2025/train.csv')\n\n        cv5 = pd.read_csv('/kaggle/input/nips-adc-data/cv5.csv')\n        cv6 = pd.read_csv('/kaggle/input/nips-adc-data/cv6.csv')\n\n        true = train[train.planet_id == int(planet_id)].values[0, 1:]\n        star_info = dict(train_star_info[train_star_info.planet_id == int(planet_id)].reset_index(drop = True).loc[0])\n        pred, stage2_params = stage2_function(args, data, star_info, *stage1_params[:4], plot = True)\n\n        plt.plot(true, color = 'g', alpha = 0.5)\n        plt.plot(cv5[cv5.planet_id == int(planet_id)].values[0, 1:1 + 283], color = 'c', alpha = 0.5)\n        plt.plot(cv6[cv6.planet_id == int(planet_id)].values[0, 1:1 + 283], color = 'y', alpha = 0.5)\n        plt.plot(pred[:283], color = 'r', alpha = 0.5)\n        plt.show()\n\n        submission = pd.DataFrame(pred[None], columns = args.columns)\n        submission['planet_id'] = [int(planet_id)]\n        _, _, score = get_score(args, submission)\n\n        print('score : ', score)\n        print('cv5 : ', cv5[cv5.planet_id == int(planet_id)]['weighted_score'].values[0])\n        print('cv6 : ', cv6[cv6.planet_id == int(planet_id)]['weighted_score'].values[0])\n        print('params : ', stage2_params.shape)\n\n        print(train_star_info[train_star_info.planet_id == int(planet_id)])\n    else:\n        star_info = dict(test_star_info[test_star_info.planet_id == int(planet_id)].reset_index(drop = True).loc[0])\n\n        pred, stage2_params = stage2_function(args, data, star_info, *stage1_params[:4], plot = True)\n\n        plt.plot(pred[:283], color = 'r', alpha = 0.5)\n        plt.show()\n\n        print(test_star_info[test_star_info.planet_id == int(planet_id)])\"\"\"\npass","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2025-09-24T05:56:33.193005Z","iopub.execute_input":"2025-09-24T05:56:33.193282Z","iopub.status.idle":"2025-09-24T05:56:33.203904Z","shell.execute_reply.started":"2025-09-24T05:56:33.193263Z","shell.execute_reply":"2025-09-24T05:56:33.20304Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\"\"\"def area(d, x, R):\n    _area = np.zeros_like(d + x, dtype = float)\n\n    mask1 = (d <= np.abs(R - x))\n    mask2 = (d > np.abs(R - x)) & (d < (R + x))\n\n    _area[mask1] = (np.pi * (np.minimum(x, R) ** 2))[mask1]\n\n    with np.errstate(divide = 'ignore', invalid = 'ignore'):\n        arg1 = (d ** 2 + x ** 2 - R ** 2) / (2 * d * x)\n        arg2 = (d ** 2 + R ** 2 - x ** 2) / (2 * d * R)\n        arg3 = (-d + x + R) * (d + x - R) * (d - x + R) * (d + x + R)\n\n    arg1 = np.clip(arg1, -1, 1)\n    arg2 = np.clip(arg2, -1, 1)\n    arg3 = np.clip(arg3, 0, None)\n\n    _area[mask2] = (x ** 2 * np.arccos(arg1) + R ** 2 * np.arccos(arg2) - 0.5 * np.sqrt(arg3))[mask2]\n    return _area\n\ndef intensity(x, args):\n    c1, c2, c3, c4 = args\n\n    norm = (- c1 / 10 - c2 / 6 - 3 * c3 / 14 - c4 / 4 + 0.5) * 2 * np.pi\n\n    x = np.clip(x, None, 0.99995)\n    sqrtmu = (1 - x ** 2) ** 0.25\n    return (1 - c1 * (1 - sqrtmu) - c2 * (1 - sqrtmu ** 2) - c3 * (1 - sqrtmu ** 3) - c4 * (1 - sqrtmu ** 4)) / norm\n\ndef calc_limb_darkening(d_array, rprs, intensity_args, n_step = 5):\n    f_array = np.ones_like(d_array, dtype = float)\n\n    x_in_array  = np.maximum(d_array - rprs, 0.0)\n    x_out_array = np.minimum(d_array + rprs, 1.0)\n\n    mask = (x_in_array < 1.0) & ((x_out_array - x_in_array) >= 1e-7)\n\n    if not np.any(mask):\n        return f_array\n\n    d = d_array[mask]\n    x_in = x_in_array[mask]\n    x_out = x_out_array[mask]\n\n    x = np.linspace(x_in, x_out, n_step, axis = 1)\n    dx = x[:, 1:] - x[:, :-1]\n\n    Int = intensity(x[:, 1:] - dx / 2, intensity_args)\n\n    A = area(d[:, np.newaxis], x, rprs)\n    delta = np.sum((A[:, 1:] - A[:, :-1]) * Int, axis = 1)\n\n    f_array[mask] = 1.0 - delta\n    return f_array\n\ndef get_transit(t, star_info, rp, c1, c2, c3, c4, P, sma, i, t0):\n    f = 2 * np.pi * (t - t0) / P\n    b = sma * np.cos(np.radians(i))\n    d = np.sqrt((sma * np.sin(f)) ** 2 + (b * np.cos(f)) ** 2)\n\n    transit = calc_limb_darkening(d, rp, [c1, c2, c3, c4])\n    return transit\n\ndef get_system(t, c, *coeffs):\n    system_time = Polynomial(coeffs[:len(coeffs) // 2])(t)\n    system_channel = Polynomial(coeffs[len(coeffs) // 2:])(c)\n\n    system = system_time[:, None] * system_channel[None, :]\n    system = 1 + system\n    return system\n\ndef get_jac_sparsity(param_info, n_time, n_channel):\n    m = n_time * n_channel\n    n = param_info['n_param']\n\n    jac_sparsity = lil_matrix((m, n), dtype = int)\n    for param in ['P', 'sma', 'i', 't0', 'sys']:\n        start = param_info[param][0]\n        end = param_info[param][1]\n        jac_sparsity[:, start:end] = 1\n\n    for param in ['rp', 'c1', 'c2', 'c3', 'c4']:\n        for j in range(n_channel):\n            jac_sparsity[\n                np.arange(j, m, n_channel),\n                param_info[param][0] + j,\n            ] = 1\n            \n    return jac_sparsity\n\ndef model2(params, t, star_info, param_info, n_channel):\n    c = np.linspace(0, 1, n_channel)\n\n    rp = params[param_info['rp'][0]:param_info['rp'][1]]\n\n    c1 = ((np.tanh(params[param_info['c1'][0]:param_info['c1'][1]]) + 1.0) / 2) * param_info['ci_max']\n    c2 = ((np.tanh(params[param_info['c2'][0]:param_info['c2'][1]]) + 1.0) / 2) * param_info['ci_max']\n    c3 = ((np.tanh(params[param_info['c3'][0]:param_info['c3'][1]]) + 1.0) / 2) * param_info['ci_max']\n    c4 = ((np.tanh(params[param_info['c4'][0]:param_info['c4'][1]]) + 1.0) / 2) * param_info['ci_max']\n    \n    P = params[param_info['P'][0]:param_info['P'][1]]\n    sma = params[param_info['sma'][0]:param_info['sma'][1]]\n    i = params[param_info['i'][0]:param_info['i'][1]]\n    t0 = params[param_info['t0'][0]:param_info['t0'][1]]\n\n    sys = params[param_info['sys'][0]:param_info['sys'][1]]\n\n    transit = []\n    for j in range(n_channel):\n        _transit = get_transit(t, star_info, rp[j], c1[j], c2[j], c3[j], c4[j], P, sma, i, t0)\n        transit.append(_transit)\n\n    transit = np.stack(transit, axis = 1)\n    system = get_system(t, c, *sys)\n\n    pred = transit * system\n    return pred\n\ndef fun(params, t, star_info, param_info, true):\n    n_channel = true.shape[1]\n\n    pred = model2(params, t, star_info, param_info, n_channel)\n    loss = true - pred\n    return loss.ravel()\n\ndef process_single(inputs):\n    args, ydata, time, star_info, T1, T2, T3, T4, degree, window, ci_max = inputs\n\n    system_mask = (time < T1) | (time > T4)\n    transit_mask = (time > T2) & (time < T3)\n\n    target = ydata[:, 0]\n    target /= target[system_mask].mean()\n\n    def model(t, star_info, rp, c1, c2, c3, c4, P, sma, i, t0, *coeffs):\n        f = 2 * np.pi * (t - t0) / P\n        b = sma * np.cos(np.radians(i))\n        d = np.sqrt((sma * np.sin(f)) ** 2 + (b * np.cos(f)) ** 2)\n\n        system = Polynomial(coeffs)(t)\n        transit = calc_limb_darkening(d, rp, [c1, c2, c3, c4])\n        return system, transit\n\n    def f(xdata, rp, c1, c2, c3, c4, P, sma, i, t0, *coeffs):\n        system, transit = model(xdata, star_info, rp, c1, c2, c3, c4, P, sma, i, t0, *coeffs)\n        return system * transit\n\n    rp_0 = np.sqrt(1 - target[transit_mask].mean())\n\n    ci_0 = [1e-1] * 4\n\n    P_0 = star_info['P'] * 24 / 7.5\n    sma_0 = star_info['sma']\n    i_0 = star_info['i']\n    t0_0 = (T1 + T4) / 2\n\n    sysi_0 = [0] * (degree + 1)\n\n    p0 = [rp_0] + ci_0 + [P_0, sma_0, i_0, t0_0] + sysi_0\n    bounds = (\n        [rp_0 * (1.0 - 1e-1)] + [0.0] * 4 + [_ * (1.0 - 1e-1) for _ in [P_0, sma_0, i_0, t0_0]] + [-1e1] * (degree + 1),\n        [rp_0 * (1.0 + 1e-1)] + [ci_max] * 4 + [_ * (1.0 + 1e-1) for _ in [P_0, sma_0, i_0, t0_0]] + [+1e1] * (degree + 1),\n    )\n\n    popt, pcov = curve_fit(\n        f = f,\n        xdata = time,\n        ydata = target,\n        p0 = p0,\n        bounds = bounds,\n    )\n    return popt\n\ndef process_multiple(inputs):\n    args, targets, time, star_info, T1, T2, T3, T4, sysi_degree, window, ci_max, plot, params, param_info = inputs\n\n    system_mask = (time < T1) | (time > T4)\n    transit_mask = (time > T2) & (time < T3)\n\n    if params == None:\n        rp_0 = np.sqrt(1 - targets[transit_mask].mean(0)).tolist()\n    \n        #ci_0 = [-1.0] * targets.shape[1] * 4\n        ######################################\n        ci_0 = [-0.5] * targets.shape[1] * 4\n        ######################################\n        \n        P_0 = [star_info['P'] * 24 / 7.5]\n        sma_0 = [star_info['sma']]\n        i_0 = [star_info['i']]\n        t0_0 = [(T1 + T4) / 2]\n    \n        sysi_0 = [1e-3] * (sysi_degree + 1) + [1e-3] * (sysi_degree + 1)\n    \n        x0 = rp_0 + ci_0 + P_0 + sma_0 + i_0 + t0_0 + sysi_0\n    \n        bounds = (\n            [_ * (1.0 - 1e-1) for _ in rp_0] + [-1e1] * len(ci_0) + [_ * (1.0 - 2e-1) for _ in P_0 + sma_0 + i_0 + t0_0] + [-1e1] * len(sysi_0),\n            [_ * (1.0 + 1e-1) for _ in rp_0] + [+1e1] * len(ci_0) + [_ * (1.0 + 2e-1) for _ in P_0 + sma_0 + i_0 + t0_0] + [+1e1] * len(sysi_0),\n        )\n    else:\n        params = params.x\n        \n        rp_0 = params[param_info['rp'][0]:param_info['rp'][1]].tolist()\n    \n        c1_0 = params[param_info['c1'][0]:param_info['c1'][1]].tolist()\n        c2_0 = params[param_info['c2'][0]:param_info['c2'][1]].tolist()\n        c3_0 = params[param_info['c3'][0]:param_info['c3'][1]].tolist()\n        c4_0 = params[param_info['c4'][0]:param_info['c4'][1]].tolist()\n    \n        ci_0 = c1_0 + c2_0 + c3_0 + c4_0\n        \n        P_0 = params[param_info['P'][0]:param_info['P'][1]].tolist()\n        sma_0 = params[param_info['sma'][0]:param_info['sma'][1]].tolist()\n        i_0 = params[param_info['i'][0]:param_info['i'][1]].tolist()\n        t0_0 = params[param_info['t0'][0]:param_info['t0'][1]].tolist()\n    \n        sysi_0 = params[param_info['sys'][0]:param_info['sys'][1]].tolist()\n        sysi_0 = [1e-3] * (sysi_degree - param_info['sysi_degree']) + sysi_0[:param_info['sysi_degree'] + 1] + \\\n                 [1e-3] * (sysi_degree - param_info['sysi_degree']) + sysi_0[param_info['sysi_degree'] + 1:]\n\n        x0 = rp_0 + ci_0 + P_0 + sma_0 + i_0 + t0_0 + sysi_0\n\n        bounds = (\n            [_ * (1.0 - 1e-1) for _ in rp_0] + [-1e1] * len(ci_0) + [_ * (1.0 - 1e-1) for _ in P_0 + sma_0 + i_0 + t0_0] + [-1e1] * len(sysi_0),\n            [_ * (1.0 + 1e-1) for _ in rp_0] + [+1e1] * len(ci_0) + [_ * (1.0 + 1e-1) for _ in P_0 + sma_0 + i_0 + t0_0] + [+1e1] * len(sysi_0),\n        )\n\n    param_info = {\n        'rp' : [0, len(rp_0)],\n        'c1' : [len(rp_0), len(rp_0) + (len(ci_0) // 4)],\n        'c2' : [len(rp_0) + (len(ci_0) // 4), len(rp_0) + 2 * (len(ci_0) // 4)],\n        'c3' : [len(rp_0) + 2 * (len(ci_0) // 4), len(rp_0) + 3 * (len(ci_0) // 4)],\n        'c4' : [len(rp_0) + 3 * (len(ci_0) // 4), len(rp_0) + len(ci_0)],\n        'P' : [len(rp_0) + len(ci_0), len(rp_0) + len(ci_0) + 1],\n        'sma' : [len(rp_0) + len(ci_0) + 1, len(rp_0) + len(ci_0) + 2],\n        'i' : [len(rp_0) + len(ci_0) + 2, len(rp_0) + len(ci_0) + 3],\n        't0' : [len(rp_0) + len(ci_0) + 3, len(rp_0) + len(ci_0) + 4],\n        'sys' : [len(rp_0) + len(ci_0) + 4, len(rp_0) + len(ci_0) + 4 + len(sysi_0)],\n        'n_param' : len(rp_0) + len(ci_0) + 4 + len(sysi_0),\n        'ci_max' : ci_max,\n        'sysi_degree' : sysi_degree,\n    }\n\n    jac_sparsity = get_jac_sparsity(param_info, targets.shape[0], targets.shape[1])\n\n    with Pool(processes = os.cpu_count()) as pool:\n        res = least_squares(\n            fun = fun,\n            x0 = x0,\n            bounds = bounds,\n            method = 'trf',\n            jac_sparsity = jac_sparsity,\n            args = (time, star_info, param_info, targets),\n            workers = pool.map,\n            verbose = 2 if plot else 0,\n        )\n    return res, param_info\n\ndef stage2_function(args, ydata, star_info, T1, T2, T3, T4, ci_max = 0.5, sysi_degree = 3, window = 5, downsample = 4, plot = False):\n    time = np.linspace(0, 1, ydata.shape[0])\n\n    system_mask = (time < T1) | (time > T4)\n    transit_mask = (time > T2) & (time < T3)\n\n    targets = []\n    for j in range(ydata.shape[1]):\n        if j > 200:\n            window = 50\n        \n        if j == 0:\n            target = ydata[:, 0:1]\n        else:\n            target = ydata[:, max(1, j - window):(j + window + 1)]\n\n        target = target.mean(1)\n        target /= target[system_mask].mean()\n        targets.append(target)\n\n    targets = np.stack(targets, axis = 1)\n    targets = targets[:, 1::downsample]\n\n    #res, param_info = process_multiple([args, targets, time, star_info, T1, T2, T3, T4, 3, 5, 0.50, plot, None, None])\n    #res, param_info = process_multiple([args, targets, time, star_info, T1, T2, T3, T4, 4, 5, 0.50, plot, res, param_info])\n    ###########################################################################################################################\n    res, param_info = process_multiple([args, targets, time, star_info, T1, T2, T3, T4, 3, 5, 0.25, plot, None, None])\n    res, param_info = process_multiple([args, targets, time, star_info, T1, T2, T3, T4, 4, 5, 0.25, plot, res, param_info])\n    ###########################################################################################################################\n    airs_popt, cost, nfev = res.x, res.cost, res.nfev\n    \n    fgs1_popt = process_single([args, ydata, time, star_info, T1, T2, T3, T4, 3, 5, 0.5])\n\n    means = np.array(airs_popt[param_info['rp'][0]:param_info['rp'][1]]) ** 2\n    x1 = np.linspace(0, 1, 282)\n    x2 = np.linspace(0, 1, means.shape[0])\n    means = np.interp(x1, x2, means).tolist()\n    means = [fgs1_popt[0] ** 2] + means\n\n    sigmas = [args.fgs_sigma] + [args.sigma] * 282\n\n    pred = np.array(means + sigmas)\n    pred = pred.clip(0)\n\n    params = np.concatenate([\n        airs_popt,\n        np.array([cost]),\n        np.array([nfev]),\n        fgs1_popt,\n    ], axis = 0)\n\n    if plot:\n        x1, x2 = np.meshgrid(\n            np.linspace(0, targets.shape[1], targets.shape[1]),\n            np.linspace(0, targets.shape[0], targets.shape[0]),\n        )\n        fig = plt.figure()\n        ax = fig.add_subplot(projection = '3d')\n        ax.plot_surface(x1, x2, targets, cmap = 'viridis', color = 'g', alpha = 0.5)\n        plt.show()\n\n        _pred = model2(airs_popt, time, star_info, param_info, targets.shape[1])\n\n        for j in range(targets.shape[1]):\n            if j % (targets.shape[1] // 4) == 0:\n                plt.plot(targets[:, j], color = 'g', alpha = 0.5)\n                plt.plot(_pred[:, j], color = 'r', alpha = 0.5)\n                plt.show()\n\n        print('airs_popt : ', airs_popt.shape)\n        print('fgs1_popt : ', fgs1_popt.shape)\n\n    return pred, params\n\nif __name__ == '__main__':\n    train_star_info = pd.read_csv('/kaggle/input/ariel-data-challenge-2025/train_star_info.csv')\n    test_star_info = pd.read_csv('/kaggle/input/ariel-data-challenge-2025/test_star_info.csv')\n\n    if args.split == 'train':\n        if '2024' in args.root:\n            train = pd.read_csv('/kaggle/input/ariel-data-challenge-2024/train_labels.csv')\n        else:\n            train = pd.read_csv('/kaggle/input/ariel-data-challenge-2025/train.csv')\n\n        cv6 = pd.read_csv('/kaggle/input/nips-adc-data/cv6.csv')\n        cv7 = pd.read_csv('/kaggle/input/nips-adc-data/cv7.csv')\n\n        true = train[train.planet_id == int(planet_id)].values[0, 1:]\n        star_info = dict(train_star_info[train_star_info.planet_id == int(planet_id)].reset_index(drop = True).loc[0])\n        pred, stage2_params = stage2_function(args, data, star_info, *stage1_params[:4], plot = True)\n\n        plt.plot(true, color = 'g', alpha = 0.5)\n        plt.plot(cv6[cv6.planet_id == int(planet_id)].values[0, 1:1 + 283], color = 'c', alpha = 0.5)\n        plt.plot(cv7[cv7.planet_id == int(planet_id)].values[0, 1:1 + 283], color = 'y', alpha = 0.5)\n        plt.plot(pred[:283], color = 'r', alpha = 0.5)\n        plt.show()\n\n        submission = pd.DataFrame(pred[None], columns = args.columns)\n        submission['planet_id'] = [int(planet_id)]\n        _, _, score = get_score(args, submission)\n\n        print('score : ', score)\n        print('cv6 : ', cv6[cv6.planet_id == int(planet_id)]['weighted_score'].values[0])\n        print('cv7 : ', cv7[cv7.planet_id == int(planet_id)]['weighted_score'].values[0])\n        print('params : ', stage2_params.shape)\n\n        print(train_star_info[train_star_info.planet_id == int(planet_id)])\n    else:\n        star_info = dict(test_star_info[test_star_info.planet_id == int(planet_id)].reset_index(drop = True).loc[0])\n\n        pred, stage2_params = stage2_function(args, data, star_info, *stage1_params[:4], plot = True)\n\n        plt.plot(pred[:283], color = 'r', alpha = 0.5)\n        plt.show()\n\n        print(test_star_info[test_star_info.planet_id == int(planet_id)])\"\"\"\npass","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2025-09-24T05:56:33.204796Z","iopub.execute_input":"2025-09-24T05:56:33.205004Z","iopub.status.idle":"2025-09-24T05:56:33.223497Z","shell.execute_reply.started":"2025-09-24T05:56:33.204987Z","shell.execute_reply":"2025-09-24T05:56:33.22288Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def area(d, x, R):\n    _area = np.zeros_like(d + x, dtype = float)\n\n    mask1 = (d <= np.abs(R - x))\n    mask2 = (d > np.abs(R - x)) & (d < (R + x))\n\n    _area[mask1] = (np.pi * (np.minimum(x, R) ** 2))[mask1]\n\n    with np.errstate(divide = 'ignore', invalid = 'ignore'):\n        arg1 = (d ** 2 + x ** 2 - R ** 2) / (2 * d * x)\n        arg2 = (d ** 2 + R ** 2 - x ** 2) / (2 * d * R)\n        arg3 = (-d + x + R) * (d + x - R) * (d - x + R) * (d + x + R)\n\n    arg1 = np.clip(arg1, -1, 1)\n    arg2 = np.clip(arg2, -1, 1)\n    arg3 = np.clip(arg3, 0, None)\n\n    _area[mask2] = (x ** 2 * np.arccos(arg1) + R ** 2 * np.arccos(arg2) - 0.5 * np.sqrt(arg3))[mask2]\n    return _area\n\ndef intensity(x, args):\n    c1, c2, c3, c4 = args\n\n    norm = (- c1 / 10 - c2 / 6 - 3 * c3 / 14 - c4 / 4 + 0.5) * 2 * np.pi\n\n    x = np.clip(x, None, 0.99995)\n    sqrtmu = (1 - x ** 2) ** 0.25\n    return (1 - c1 * (1 - sqrtmu) - c2 * (1 - sqrtmu ** 2) - c3 * (1 - sqrtmu ** 3) - c4 * (1 - sqrtmu ** 4)) / norm\n\ndef calc_limb_darkening(d_array, rprs, intensity_args, n_step = 5):\n    f_array = np.ones_like(d_array, dtype = float)\n\n    x_in_array  = np.maximum(d_array - rprs, 0.0)\n    x_out_array = np.minimum(d_array + rprs, 1.0)\n\n    mask = (x_in_array < 1.0) & ((x_out_array - x_in_array) >= 1e-7)\n\n    if not np.any(mask):\n        return f_array\n\n    d = d_array[mask]\n    x_in = x_in_array[mask]\n    x_out = x_out_array[mask]\n\n    x = np.linspace(x_in, x_out, n_step, axis = 1)\n    dx = x[:, 1:] - x[:, :-1]\n\n    Int = intensity(x[:, 1:] - dx / 2, intensity_args)\n\n    A = area(d[:, np.newaxis], x, rprs)\n    delta = np.sum((A[:, 1:] - A[:, :-1]) * Int, axis = 1)\n\n    f_array[mask] = 1.0 - delta\n    return f_array\n\ndef get_transit(t, star_info, rp, c1, c2, c3, c4, P, sma, i, t0):\n    f = 2 * np.pi * (t - t0) / P\n    b = sma * np.cos(np.radians(i))\n    d = np.sqrt((sma * np.sin(f)) ** 2 + (b * np.cos(f)) ** 2)\n\n    transit = calc_limb_darkening(d, rp, [c1, c2, c3, c4])\n    return transit\n\ndef get_system(t, c, *coeffs):\n    system_time = Polynomial(coeffs[:len(coeffs) // 2])(t)\n    system_channel = Polynomial(coeffs[len(coeffs) // 2:])(c)\n\n    system = system_time[:, None] * system_channel[None, :]\n    system = 1 + system\n    return system\n\ndef get_jac_sparsity(param_info, n_time, n_channel):\n    m = n_time * n_channel\n    n = param_info['n_param']\n\n    jac_sparsity = lil_matrix((m, n), dtype = int)\n    for param in ['P', 'sma', 'i', 't0', 'sys']:\n        start = param_info[param][0]\n        end = param_info[param][1]\n        jac_sparsity[:, start:end] = 1\n\n    for param in ['rp', 'c1', 'c2', 'c3', 'c4']:\n        for j in range(n_channel):\n            jac_sparsity[\n                np.arange(j, m, n_channel),\n                param_info[param][0] + j,\n            ] = 1\n            \n    return jac_sparsity\n\ndef model2(params, t, star_info, param_info, n_channel):\n    c = np.linspace(0, 1, n_channel)\n\n    rp = params[param_info['rp'][0]:param_info['rp'][1]]\n\n    c1 = ((np.tanh(params[param_info['c1'][0]:param_info['c1'][1]]) + 1.0) / 2) * param_info['ci_max']\n    c2 = ((np.tanh(params[param_info['c2'][0]:param_info['c2'][1]]) + 1.0) / 2) * param_info['ci_max']\n    c3 = ((np.tanh(params[param_info['c3'][0]:param_info['c3'][1]]) + 1.0) / 2) * param_info['ci_max']\n    c4 = ((np.tanh(params[param_info['c4'][0]:param_info['c4'][1]]) + 1.0) / 2) * param_info['ci_max']\n    \n    P = params[param_info['P'][0]:param_info['P'][1]]\n    sma = params[param_info['sma'][0]:param_info['sma'][1]]\n    i = params[param_info['i'][0]:param_info['i'][1]]\n    t0 = params[param_info['t0'][0]:param_info['t0'][1]]\n\n    sys = params[param_info['sys'][0]:param_info['sys'][1]]\n\n    transit = []\n    for j in range(n_channel):\n        _transit = get_transit(t, star_info, rp[j], c1[j], c2[j], c3[j], c4[j], P, sma, i, t0)\n        transit.append(_transit)\n\n    transit = np.stack(transit, axis = 1)\n    system = get_system(t, c, *sys)\n\n    pred = transit * system\n    return pred\n\ndef fun(params, t, star_info, param_info, true):\n    n_channel = true.shape[1]\n\n    pred = model2(params, t, star_info, param_info, n_channel)\n    loss = true - pred\n    return loss.ravel()\n\ndef process_single(inputs):\n    args, ydata, time, star_info, T1, T2, T3, T4, degree, window, ci_max = inputs\n\n    system_mask = (time < T1) | (time > T4)\n    transit_mask = (time > T2) & (time < T3)\n\n    target = ydata[:, 0]\n    target /= target[system_mask].mean()\n\n    def model(t, star_info, rp, c1, c2, c3, c4, P, sma, i, t0, *coeffs):\n        f = 2 * np.pi * (t - t0) / P\n        b = sma * np.cos(np.radians(i))\n        d = np.sqrt((sma * np.sin(f)) ** 2 + (b * np.cos(f)) ** 2)\n\n        system = Polynomial(coeffs)(t)\n        transit = calc_limb_darkening(d, rp, [c1, c2, c3, c4])\n        return system, transit\n\n    def f(xdata, rp, c1, c2, c3, c4, P, sma, i, t0, *coeffs):\n        system, transit = model(xdata, star_info, rp, c1, c2, c3, c4, P, sma, i, t0, *coeffs)\n        return system * transit\n\n    rp_0 = np.sqrt(1 - target[transit_mask].mean())\n\n    ci_0 = [1e-1] * 4\n\n    P_0 = star_info['P'] * 24 / 7.5\n    sma_0 = star_info['sma']\n    i_0 = star_info['i']\n    t0_0 = (T1 + T4) / 2\n\n    sysi_0 = [0] * (degree + 1)\n\n    p0 = [rp_0] + ci_0 + [P_0, sma_0, i_0, t0_0] + sysi_0\n    bounds = (\n        [rp_0 * (1.0 - 1e-1)] + [0.0] * 4 + [_ * (1.0 - 1e-1) for _ in [P_0, sma_0, i_0, t0_0]] + [-1e1] * (degree + 1),\n        [rp_0 * (1.0 + 1e-1)] + [ci_max] * 4 + [_ * (1.0 + 1e-1) for _ in [P_0, sma_0, i_0, t0_0]] + [+1e1] * (degree + 1),\n    )\n\n    popt, pcov = curve_fit(\n        f = f,\n        xdata = time,\n        ydata = target,\n        p0 = p0,\n        bounds = bounds,\n    )\n    return popt\n\ndef process_multiple(inputs):\n    args, targets, time, star_info, T1, T2, T3, T4, sysi_degree, window, ci_max, plot, params, param_info = inputs\n\n    system_mask = (time < T1) | (time > T4)\n    transit_mask = (time > T2) & (time < T3)\n\n    if params == None:\n        rp_0 = np.sqrt(1 - targets[transit_mask].mean(0)).tolist()\n    \n        #ci_0 = [-1.0] * targets.shape[1] * 4\n        ######################################\n        ci_0 = [-0.5] * targets.shape[1] * 4\n        ######################################\n        \n        P_0 = [star_info['P'] * 24 / 7.5]\n        sma_0 = [star_info['sma']]\n        i_0 = [star_info['i']]\n        t0_0 = [(T1 + T4) / 2]\n    \n        sysi_0 = [1e-3] * (sysi_degree + 1) + [1e-3] * (sysi_degree + 1)\n    \n        x0 = rp_0 + ci_0 + P_0 + sma_0 + i_0 + t0_0 + sysi_0\n    \n        bounds = (\n            [_ * (1.0 - 1e-1) for _ in rp_0] + [-1e1] * len(ci_0) + [_ * (1.0 - 2e-1) for _ in P_0 + sma_0 + i_0 + t0_0] + [-1e1] * len(sysi_0),\n            [_ * (1.0 + 1e-1) for _ in rp_0] + [+1e1] * len(ci_0) + [_ * (1.0 + 2e-1) for _ in P_0 + sma_0 + i_0 + t0_0] + [+1e1] * len(sysi_0),\n        )\n    else:\n        params = params.x\n        \n        rp_0 = params[param_info['rp'][0]:param_info['rp'][1]].tolist()\n    \n        c1_0 = params[param_info['c1'][0]:param_info['c1'][1]].tolist()\n        c2_0 = params[param_info['c2'][0]:param_info['c2'][1]].tolist()\n        c3_0 = params[param_info['c3'][0]:param_info['c3'][1]].tolist()\n        c4_0 = params[param_info['c4'][0]:param_info['c4'][1]].tolist()\n    \n        ci_0 = c1_0 + c2_0 + c3_0 + c4_0\n        \n        P_0 = params[param_info['P'][0]:param_info['P'][1]].tolist()\n        sma_0 = params[param_info['sma'][0]:param_info['sma'][1]].tolist()\n        i_0 = params[param_info['i'][0]:param_info['i'][1]].tolist()\n        t0_0 = params[param_info['t0'][0]:param_info['t0'][1]].tolist()\n    \n        sysi_0 = params[param_info['sys'][0]:param_info['sys'][1]].tolist()\n        sysi_0 = [1e-3] * (sysi_degree - param_info['sysi_degree']) + sysi_0[:param_info['sysi_degree'] + 1] + \\\n                 [1e-3] * (sysi_degree - param_info['sysi_degree']) + sysi_0[param_info['sysi_degree'] + 1:]\n\n        x0 = rp_0 + ci_0 + P_0 + sma_0 + i_0 + t0_0 + sysi_0\n\n        bounds = (\n            [_ * (1.0 - 1e-1) for _ in rp_0] + [-1e1] * len(ci_0) + [_ * (1.0 - 1e-1) for _ in P_0 + sma_0 + i_0 + t0_0] + [-1e1] * len(sysi_0),\n            [_ * (1.0 + 1e-1) for _ in rp_0] + [+1e1] * len(ci_0) + [_ * (1.0 + 1e-1) for _ in P_0 + sma_0 + i_0 + t0_0] + [+1e1] * len(sysi_0),\n        )\n\n    param_info = {\n        'rp' : [0, len(rp_0)],\n        'c1' : [len(rp_0), len(rp_0) + (len(ci_0) // 4)],\n        'c2' : [len(rp_0) + (len(ci_0) // 4), len(rp_0) + 2 * (len(ci_0) // 4)],\n        'c3' : [len(rp_0) + 2 * (len(ci_0) // 4), len(rp_0) + 3 * (len(ci_0) // 4)],\n        'c4' : [len(rp_0) + 3 * (len(ci_0) // 4), len(rp_0) + len(ci_0)],\n        'P' : [len(rp_0) + len(ci_0), len(rp_0) + len(ci_0) + 1],\n        'sma' : [len(rp_0) + len(ci_0) + 1, len(rp_0) + len(ci_0) + 2],\n        'i' : [len(rp_0) + len(ci_0) + 2, len(rp_0) + len(ci_0) + 3],\n        't0' : [len(rp_0) + len(ci_0) + 3, len(rp_0) + len(ci_0) + 4],\n        'sys' : [len(rp_0) + len(ci_0) + 4, len(rp_0) + len(ci_0) + 4 + len(sysi_0)],\n        'n_param' : len(rp_0) + len(ci_0) + 4 + len(sysi_0),\n        'ci_max' : ci_max,\n        'sysi_degree' : sysi_degree,\n    }\n\n    jac_sparsity = get_jac_sparsity(param_info, targets.shape[0], targets.shape[1])\n\n    with Pool(processes = os.cpu_count()) as pool:\n        res = least_squares(\n            fun = fun,\n            x0 = x0,\n            bounds = bounds,\n            method = 'trf',\n            jac_sparsity = jac_sparsity,\n            args = (time, star_info, param_info, targets),\n            workers = pool.map,\n            verbose = 2 if plot else 0,\n        )\n    return res, param_info\n\ndef stage2_function(args, ydata, star_info, T1, T2, T3, T4, ci_max = 0.5, sysi_degree = 3, window = 5, downsample = 4, plot = False):\n    time = np.linspace(0, 1, ydata.shape[0])\n\n    system_mask = (time < T1) | (time > T4)\n    transit_mask = (time > T2) & (time < T3)\n\n    targets = []\n    for j in range(ydata.shape[1]):\n        #if j > 200:\n        #    window = 50\n        \n        if j == 0:\n            target = ydata[:, 0:1]\n        else:\n            target = ydata[:, max(1, j - window):(j + window + 1)]\n\n        target = target.mean(1)\n        target /= target[system_mask].mean()\n        targets.append(target)\n\n    targets = np.stack(targets, axis = 1)\n    targets = targets[:, 1::downsample]\n\n    #res, param_info = process_multiple([args, targets, time, star_info, T1, T2, T3, T4, 3, 5, 0.50, plot, None, None])\n    #res, param_info = process_multiple([args, targets, time, star_info, T1, T2, T3, T4, 4, 5, 0.50, plot, res, param_info])\n    ###########################################################################################################################\n    res, param_info = process_multiple([args, targets, time, star_info, T1, T2, T3, T4, 3, 5, 0.25, plot, None, None])\n    res, param_info = process_multiple([args, targets, time, star_info, T1, T2, T3, T4, 4, 5, 0.25, plot, res, param_info])\n    ###########################################################################################################################\n    airs_popt, cost, nfev = res.x, res.cost, res.nfev\n    \n    fgs1_popt = process_single([args, ydata, time, star_info, T1, T2, T3, T4, 3, 5, 0.5])\n\n    means = np.array(airs_popt[param_info['rp'][0]:param_info['rp'][1]]) ** 2\n    x1 = np.linspace(0, 1, 282)\n    x2 = np.linspace(0, 1, means.shape[0])\n    means = np.interp(x1, x2, means).tolist()\n    means = [fgs1_popt[0] ** 2] + means\n\n    sigmas = [args.fgs_sigma] + [args.sigma] * 282\n\n    pred = np.array(means + sigmas)\n    pred = pred.clip(0)\n\n    params = np.concatenate([\n        airs_popt,\n        np.array([cost]),\n        np.array([nfev]),\n        fgs1_popt,\n    ], axis = 0)\n\n    if plot:\n        x1, x2 = np.meshgrid(\n            np.linspace(0, targets.shape[1], targets.shape[1]),\n            np.linspace(0, targets.shape[0], targets.shape[0]),\n        )\n        fig = plt.figure()\n        ax = fig.add_subplot(projection = '3d')\n        ax.plot_surface(x1, x2, targets, cmap = 'viridis', color = 'g', alpha = 0.5)\n        plt.show()\n\n        _pred = model2(airs_popt, time, star_info, param_info, targets.shape[1])\n\n        for j in range(targets.shape[1]):\n            if j % (targets.shape[1] // 4) == 0:\n                plt.plot(targets[:, j], color = 'g', alpha = 0.5)\n                plt.plot(_pred[:, j], color = 'r', alpha = 0.5)\n                plt.show()\n\n        print('airs_popt : ', airs_popt.shape)\n        print('fgs1_popt : ', fgs1_popt.shape)\n\n    return pred, params\n\nif __name__ == '__main__':\n    train_star_info = pd.read_csv('/kaggle/input/ariel-data-challenge-2025/train_star_info.csv')\n    test_star_info = pd.read_csv('/kaggle/input/ariel-data-challenge-2025/test_star_info.csv')\n\n    if args.split == 'train':\n        if '2024' in args.root:\n            train = pd.read_csv('/kaggle/input/ariel-data-challenge-2024/train_labels.csv')\n        else:\n            train = pd.read_csv('/kaggle/input/ariel-data-challenge-2025/train.csv')\n\n        cv6 = pd.read_csv('/kaggle/input/nips-adc-data/cv6.csv')\n        cv7 = pd.read_csv('/kaggle/input/nips-adc-data/cv7.csv')\n\n        true = train[train.planet_id == int(planet_id)].values[0, 1:]\n        star_info = dict(train_star_info[train_star_info.planet_id == int(planet_id)].reset_index(drop = True).loc[0])\n        pred, stage2_params = stage2_function(args, data, star_info, *stage1_params[:4], plot = True)\n\n        plt.plot(true, color = 'g', alpha = 0.5)\n        plt.plot(cv6[cv6.planet_id == int(planet_id)].values[0, 1:1 + 283], color = 'c', alpha = 0.5)\n        plt.plot(cv7[cv7.planet_id == int(planet_id)].values[0, 1:1 + 283], color = 'y', alpha = 0.5)\n        plt.plot(pred[:283], color = 'r', alpha = 0.5)\n        plt.show()\n\n        submission = pd.DataFrame(pred[None], columns = args.columns)\n        submission['planet_id'] = [int(planet_id)]\n        _, _, score = get_score(args, submission)\n\n        print('score : ', score)\n        print('cv6 : ', cv6[cv6.planet_id == int(planet_id)]['weighted_score'].values[0])\n        print('cv7 : ', cv7[cv7.planet_id == int(planet_id)]['weighted_score'].values[0])\n        print('params : ', stage2_params.shape)\n\n        print(train_star_info[train_star_info.planet_id == int(planet_id)])\n    else:\n        star_info = dict(test_star_info[test_star_info.planet_id == int(planet_id)].reset_index(drop = True).loc[0])\n\n        pred, stage2_params = stage2_function(args, data, star_info, *stage1_params[:4], plot = True)\n\n        plt.plot(pred[:283], color = 'r', alpha = 0.5)\n        plt.show()\n\n        print(test_star_info[test_star_info.planet_id == int(planet_id)])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-24T05:56:33.224179Z","iopub.execute_input":"2025-09-24T05:56:33.224363Z","iopub.status.idle":"2025-09-24T05:56:48.056074Z","shell.execute_reply.started":"2025-09-24T05:56:33.224349Z","shell.execute_reply":"2025-09-24T05:56:48.05507Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### postprocess","metadata":{}},{"cell_type":"code","source":"def postprocess_function(args, models, submission, star_info, success, stage1_params, stage2_params):\n    inputs = get_inputs(args, submission, star_info.loc[submission.planet_id], stage1_params, stage2_params)\n\n    pred = []\n    for model in models:\n        pred.append(model.predict(inputs))\n    pred = np.stack(pred, axis = 0).mean(0)\n    pred = pred.clip(args.post_sigma[0])\n\n    submission.loc[success, [f'sigma_{c}' for c in range(2, 283 + 1)]] = np.broadcast_to(pred[:, np.newaxis], (pred.shape[0], 282))[success]\n    submission.loc[success, ['sigma_1']] = 2 * pred[success]\n\n    if args.pca_components != None:\n        submission.loc[success, [f'wl_{c}' for c in range(2, 283 + 1)]] = pca_function(\n            submission.loc[success, [f'wl_{c}' for c in range(2, 283 + 1)]].values, args.pca_components\n        ).clip(0)\n    return submission\n\ndef get_inputs(args, df, star_info, stage1_params = None, stage2_params = None):\n    if isinstance(stage1_params, np.ndarray):\n        T = stage1_params[:, :4]\n    else:\n        T = df[['T1', 'T2', 'T3', 'T4']].values\n\n    param_info = {\n        'airs' : {\n            'rp': [0, 71], \n            'c1': [71, 142], \n            'c2': [142, 213], \n            'c3': [213, 284], \n            'c4': [284, 355], \n            'P': [355, 356], \n            'sma': [356, 357], \n            'i': [357, 358], \n            't0': [358, 359], \n            'sys': [359, 369], \n        },\n        'cost' : [369, 370], \n        'nfev' : [370, 371], \n        'fgs1' : {\n            'rp': [371, 372], \n            'c1': [372, 373], \n            'c2': [373, 374], \n            'c3': [374, 375], \n            'c4': [375, 376], \n            'P': [376, 377], \n            'sma': [377, 378], \n            'i': [378, 379], \n            't0': [379, 380], \n            'sys': [380, 384], \n        },\n    }\n\n    if isinstance(stage2_params, np.ndarray):\n        pass\n\n    else:\n        stage2_params = np.concatenate([\n            df[[f'airs_popt_{c}' for c in range(1, 1 + 369)]].values,\n            df[['cost', 'nfev']].values,\n            df[[f'fgs1_popt_{c}' for c in range(1, 1 + 13)]].values,\n        ], axis = 1)\n\n    rp = np.concatenate([\n        stage2_params[:, param_info['airs']['rp'][0]:param_info['airs']['rp'][1]],\n        stage2_params[:, param_info['fgs1']['rp'][0]:param_info['fgs1']['rp'][1]],\n    ], axis = 1)\n\n    c1 = np.concatenate([\n        stage2_params[:, param_info['airs']['c1'][0]:param_info['airs']['c1'][1]],\n        stage2_params[:, param_info['fgs1']['c1'][0]:param_info['fgs1']['c1'][1]],\n    ], axis = 1)\n\n    c2 = np.concatenate([\n        stage2_params[:, param_info['airs']['c2'][0]:param_info['airs']['c2'][1]],\n        stage2_params[:, param_info['fgs1']['c2'][0]:param_info['fgs1']['c2'][1]],\n    ], axis = 1)\n\n    c3 = np.concatenate([\n        stage2_params[:, param_info['airs']['c3'][0]:param_info['airs']['c3'][1]],\n        stage2_params[:, param_info['fgs1']['c3'][0]:param_info['fgs1']['c3'][1]],\n    ], axis = 1)\n\n    c4 = np.concatenate([\n        stage2_params[:, param_info['airs']['c4'][0]:param_info['airs']['c4'][1]],\n        stage2_params[:, param_info['fgs1']['c4'][0]:param_info['fgs1']['c4'][1]],\n    ], axis = 1)\n\n    P = np.concatenate([\n        stage2_params[:, param_info['airs']['P'][0]:param_info['airs']['P'][1]],\n        stage2_params[:, param_info['fgs1']['P'][0]:param_info['fgs1']['P'][1]],\n    ], axis = 1)\n    \n    sma = np.concatenate([\n        stage2_params[:, param_info['airs']['sma'][0]:param_info['airs']['sma'][1]],\n        stage2_params[:, param_info['fgs1']['sma'][0]:param_info['fgs1']['sma'][1]],\n    ], axis = 1)\n    \n    i = np.concatenate([\n        stage2_params[:, param_info['airs']['i'][0]:param_info['airs']['i'][1]],\n        stage2_params[:, param_info['fgs1']['i'][0]:param_info['fgs1']['i'][1]],\n    ], axis = 1)\n    \n    t0 = np.concatenate([\n        stage2_params[:, param_info['airs']['t0'][0]:param_info['airs']['t0'][1]],\n        stage2_params[:, param_info['fgs1']['t0'][0]:param_info['fgs1']['t0'][1]],\n    ], axis = 1)\n\n    cost = stage2_params[:, param_info['cost'][0]:param_info['cost'][1]]\n    nfev = stage2_params[:, param_info['nfev'][0]:param_info['nfev'][1]]\n\n    inputs = np.concatenate([\n        T,\n        T[:, 3:4] - T[:, 0:1],\n        T[:, 2:3] - T[:, 1:2],\n\n        rp[:, :-1].mean(1, keepdims = True),\n        c1[:, :-1].mean(1, keepdims = True),\n        c2[:, :-1].mean(1, keepdims = True),\n        c3[:, :-1].mean(1, keepdims = True),\n        c4[:, :-1].mean(1, keepdims = True),\n\n        rp[:, :-1].std(1, keepdims = True),\n        c1[:, :-1].std(1, keepdims = True),\n        c2[:, :-1].std(1, keepdims = True),\n        c3[:, :-1].std(1, keepdims = True),\n        c4[:, :-1].std(1, keepdims = True),\n\n        #P,\n        #sma,\n        #i,\n        #t0,\n\n        cost,\n        nfev,\n\n        star_info.values,\n    ], axis = 1)\n    return inputs\n\ndef get_targets(args, df, star_info, label):\n    _targets = df['target'].values\n\n    targets = []\n    for j in range(_targets.shape[0]):\n        targets.append(args.post_sigma[_targets[j]])\n\n    targets = np.stack(targets, axis = 0)\n    return targets\n\ndef train_function(args, df):\n    print('n_train : ', len(df))\n\n    star_info = pd.read_csv(args.root + f'train_star_info.csv')\n    star_info = star_info.set_index('planet_id')\n\n    label = pd.read_csv(args.root + 'train.csv')\n    label = label.set_index('planet_id')\n\n    planet_ids = df.planet_id.unique()\n\n    solution = get_solution(args, label.reset_index(), planet_ids)\n\n    #'''\n    for j in range(len(df)):\n        df.loc[j, [f'wl_{c}' for c in range(2, 283 + 1)]] = gaussian_filter1d(\n            df.loc[j, [f'wl_{c}' for c in range(2, 283 + 1)]].values.astype(float), sigma = 1.0,\n        ).clip(0)\n    #'''\n\n    targets = []\n    for sigma in args.post_sigma:\n        _df = df.copy()\n        _df.loc[:, [f'sigma_{c}' for c in range(2, 283 + 1)]] = sigma\n        _df.loc[:, 'sigma_1'] = 2.0 * sigma\n\n        scores, _, _ = get_score(args, _df[_df.columns[:567]], solution = solution, print_score = False)\n        targets.append(scores)\n\n    targets = np.stack(targets, axis = 1)\n    targets = targets.argmax(1)\n\n    df['target'] = targets\n\n    if args.post_wl != None:\n        df.loc[:, [f'wl_{c}' for c in range(2, 283 + 1)]] *= (1.0 + args.post_wl[0])\n        df.loc[:, ['wl_1']] *= (1.0 + args.post_wl[1])\n\n    outputs = {}\n\n    kf = KFold(n_splits = args.n_fold, shuffle = True, random_state = args.seed)\n    for i, (train_index, test_index) in enumerate(kf.split(planet_ids)):\n        train_planets = planet_ids[train_index]\n        test_planets = planet_ids[test_index]\n\n        train_df = df[(df.planet_id.isin(train_planets)) & (df.success == True)].reset_index(drop = True)\n        test_df = df[df.planet_id.isin(test_planets)].reset_index(drop = True)\n\n        train_star_info = star_info.loc[train_df.planet_id]\n        test_star_info = star_info.loc[test_planets]\n\n        train_label = label.loc[train_df.planet_id]\n        test_label = label.loc[test_planets]\n\n        train_inputs = get_inputs(args, train_df, train_star_info)\n        train_targets = get_targets(args, train_df, train_star_info, train_label)\n\n        test_inputs = get_inputs(args, test_df, test_star_info)\n        test_targets = get_targets(args, test_df, test_star_info, test_label)\n\n        model = GradientBoostingRegressor(\n            random_state = args.seed,\n            n_estimators = 250,\n            max_depth = 3,\n        )\n\n        model.fit(train_inputs, train_targets)\n\n        score = model.score(test_inputs, test_targets)\n\n        print(f'fold{i + 1}', score)\n\n        outputs[f'fold{i + 1}'] = {}\n        outputs[f'fold{i + 1}']['train_planets'] = train_planets\n        outputs[f'fold{i + 1}']['test_planets'] = test_planets\n        outputs[f'fold{i + 1}']['model'] = model\n        outputs[f'fold{i + 1}']['score'] = score\n\n    return outputs\n\ndef get_cv(args, df):\n    outputs = train_function(args, df)\n\n    star_info = pd.read_csv(args.root + f'train_star_info.csv')\n    star_info = star_info.set_index('planet_id')\n\n    cv = []\n    scores = []\n    for i in range(args.n_fold):\n        output = outputs[f'fold{i + 1}']\n\n        planet_ids = output['test_planets']\n\n        _df = df.loc[df.planet_id.isin(planet_ids)]\n        _df = _df.reset_index(drop = True)\n\n        success = _df['success'].values\n        stage1_params = _df[['T1', 'T2', 'T3', 'T4']].values\n        stage2_params = np.concatenate([\n            _df[[f'airs_popt_{c}' for c in range(1, 1 + 369)]].values,\n            _df[['cost', 'nfev']].values,\n            _df[[f'fgs1_popt_{c}' for c in range(1, 1 + 13)]].values,\n        ], axis = 1)\n\n        if args.postprocess:\n            _df = postprocess_function(args, [output['model']], _df, star_info.loc[_df.planet_id], success, stage1_params, stage2_params)\n\n        weighted_scores, _, score = get_score(args, _df[_df.columns[:567]], print_score = False)\n        _df['weighted_score'] = weighted_scores\n        _df['fold'] = [i + 1] * len(_df)\n\n        cv.append(_df)\n        scores.append(score)\n\n    cv = pd.concat(cv, axis = 0)\n    cv = cv.reset_index(drop = True)\n    return cv, scores, outputs\n\nif __name__ == '__main__':\n    args = CustomConfig()\n\n    args.cv_path = '/kaggle/input/nips-adc-data/cv6.csv'\n\n    args.postprocess = True\n    args.post_wl = [1e-3, -5e-3]\n    args.post_sigma = np.linspace(5e-5, 2e-3, 500)\n\n    args.pca_components = 6\n\n    df = pd.read_csv(args.cv_path)\n    df['planet_id'] = df['planet_id'].astype(int)\n    df = df[~df.planet_id.isin([2486733311])]\n    df = df.sort_values(by = 'planet_id')\n    df = df.reset_index(drop = True)\n\n    df, scores, _ = get_cv(args, df)\n    print('scores : ', scores)\n    print('cv : ', np.mean(scores))\n\n    df.to_csv('cv.csv', index = False)\n\n    '''\n    n_train :  1099\n    fold1 0.3815623750338394\n    fold2 0.1702556503990037\n    fold3 0.6332230605037319\n    fold4 0.7112889028015135\n    scores :  [0.5901795315417878, 0.5731021979279475, 0.5937740064245882, 0.5807180647396613]\n    cv :  0.5844434501584962\n    '''","metadata":{"trusted":true,"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\"\"\"def postprocess_function(args, models, submission, star_info, success, stage1_params, stage2_params):\n    inputs = get_inputs(args, submission, star_info.loc[submission.planet_id], stage1_params, stage2_params)\n\n    pred = []\n    for model in models:\n        pred.append(model.predict(inputs))\n    pred = np.stack(pred, axis = 0).mean(0)\n    pred = pred.clip(args.post_sigma[0])\n\n    #################################################################################################\n    if args.post_wl != None:\n        submission.loc[success, [f'wl_{c}' for c in range(2, 283 + 1)]] *= (1.0 + args.post_wl[0])\n        submission.loc[success, ['wl_1']] *= (1.0 + args.post_wl[1])\n    #################################################################################################\n\n    submission.loc[success, [f'sigma_{c}' for c in range(2, 283 + 1)]] = np.broadcast_to(pred[:, np.newaxis], (pred.shape[0], 282))[success]\n    submission.loc[success, ['sigma_1']] = 2 * pred[success]\n\n    if args.pca_components != None:\n        submission.loc[success, [f'wl_{c}' for c in range(2, 283 + 1)]] = pca_function(\n            submission.loc[success, [f'wl_{c}' for c in range(2, 283 + 1)]].values, args.pca_components\n        ).clip(0)\n    return submission\n\ndef get_inputs(args, df, star_info, stage1_params = None, stage2_params = None):\n    if isinstance(stage1_params, np.ndarray):\n        T = stage1_params[:, :4]\n    else:\n        T = df[['T1', 'T2', 'T3', 'T4']].values\n\n    param_info = {\n        'airs' : {\n            'rp': [0, 71], \n            'c1': [71, 142], \n            'c2': [142, 213], \n            'c3': [213, 284], \n            'c4': [284, 355], \n            'P': [355, 356], \n            'sma': [356, 357], \n            'i': [357, 358], \n            't0': [358, 359], \n            'sys': [359, 369], \n        },\n        'cost' : [369, 370], \n        'nfev' : [370, 371], \n        'fgs1' : {\n            'rp': [371, 372], \n            'c1': [372, 373], \n            'c2': [373, 374], \n            'c3': [374, 375], \n            'c4': [375, 376], \n            'P': [376, 377], \n            'sma': [377, 378], \n            'i': [378, 379], \n            't0': [379, 380], \n            'sys': [380, 384], \n        },\n    }\n\n    if isinstance(stage2_params, np.ndarray):\n        pass\n\n    else:\n        stage2_params = np.concatenate([\n            df[[f'airs_popt_{c}' for c in range(1, 1 + 369)]].values,\n            df[['cost', 'nfev']].values,\n            df[[f'fgs1_popt_{c}' for c in range(1, 1 + 13)]].values,\n        ], axis = 1)\n\n    rp = np.concatenate([\n        stage2_params[:, param_info['airs']['rp'][0]:param_info['airs']['rp'][1]],\n        stage2_params[:, param_info['fgs1']['rp'][0]:param_info['fgs1']['rp'][1]],\n    ], axis = 1)\n\n    c1 = np.concatenate([\n        stage2_params[:, param_info['airs']['c1'][0]:param_info['airs']['c1'][1]],\n        stage2_params[:, param_info['fgs1']['c1'][0]:param_info['fgs1']['c1'][1]],\n    ], axis = 1)\n\n    c2 = np.concatenate([\n        stage2_params[:, param_info['airs']['c2'][0]:param_info['airs']['c2'][1]],\n        stage2_params[:, param_info['fgs1']['c2'][0]:param_info['fgs1']['c2'][1]],\n    ], axis = 1)\n\n    c3 = np.concatenate([\n        stage2_params[:, param_info['airs']['c3'][0]:param_info['airs']['c3'][1]],\n        stage2_params[:, param_info['fgs1']['c3'][0]:param_info['fgs1']['c3'][1]],\n    ], axis = 1)\n\n    c4 = np.concatenate([\n        stage2_params[:, param_info['airs']['c4'][0]:param_info['airs']['c4'][1]],\n        stage2_params[:, param_info['fgs1']['c4'][0]:param_info['fgs1']['c4'][1]],\n    ], axis = 1)\n\n    P = np.concatenate([\n        stage2_params[:, param_info['airs']['P'][0]:param_info['airs']['P'][1]],\n        stage2_params[:, param_info['fgs1']['P'][0]:param_info['fgs1']['P'][1]],\n    ], axis = 1)\n    \n    sma = np.concatenate([\n        stage2_params[:, param_info['airs']['sma'][0]:param_info['airs']['sma'][1]],\n        stage2_params[:, param_info['fgs1']['sma'][0]:param_info['fgs1']['sma'][1]],\n    ], axis = 1)\n    \n    i = np.concatenate([\n        stage2_params[:, param_info['airs']['i'][0]:param_info['airs']['i'][1]],\n        stage2_params[:, param_info['fgs1']['i'][0]:param_info['fgs1']['i'][1]],\n    ], axis = 1)\n    \n    t0 = np.concatenate([\n        stage2_params[:, param_info['airs']['t0'][0]:param_info['airs']['t0'][1]],\n        stage2_params[:, param_info['fgs1']['t0'][0]:param_info['fgs1']['t0'][1]],\n    ], axis = 1)\n\n    cost = stage2_params[:, param_info['cost'][0]:param_info['cost'][1]]\n    nfev = stage2_params[:, param_info['nfev'][0]:param_info['nfev'][1]]\n\n    inputs = np.concatenate([\n        T,\n        T[:, 3:4] - T[:, 0:1],\n        T[:, 2:3] - T[:, 1:2],\n\n        rp[:, :-1].mean(1, keepdims = True),\n        c1[:, :-1].mean(1, keepdims = True),\n        c2[:, :-1].mean(1, keepdims = True),\n        c3[:, :-1].mean(1, keepdims = True),\n        c4[:, :-1].mean(1, keepdims = True),\n\n        rp[:, :-1].std(1, keepdims = True),\n        c1[:, :-1].std(1, keepdims = True),\n        c2[:, :-1].std(1, keepdims = True),\n        c3[:, :-1].std(1, keepdims = True),\n        c4[:, :-1].std(1, keepdims = True),\n\n        #P,\n        #sma,\n        #i,\n        #t0,\n\n        cost,\n        nfev,\n\n        star_info.values,\n    ], axis = 1)\n    return inputs\n\ndef get_targets(args, df, star_info, label):\n    _targets = df['target'].values\n\n    targets = []\n    for j in range(_targets.shape[0]):\n        targets.append(args.post_sigma[_targets[j]])\n\n    targets = np.stack(targets, axis = 0)\n    return targets\n\ndef train_function(args, df):\n    print('n_train : ', len(df))\n\n    star_info = pd.read_csv(args.root + f'train_star_info.csv')\n    star_info = star_info.set_index('planet_id')\n\n    label = pd.read_csv(args.root + 'train.csv')\n    label = label.set_index('planet_id')\n\n    planet_ids = df.planet_id.unique()\n\n    solution = get_solution(args, label.reset_index(), planet_ids)\n\n    #'''\n    for j in range(len(df)):\n        df.loc[j, [f'wl_{c}' for c in range(2, 283 + 1)]] = gaussian_filter1d(\n            df.loc[j, [f'wl_{c}' for c in range(2, 283 + 1)]].values.astype(float), sigma = 1.0,\n        ).clip(0)\n    #'''\n\n    targets = []\n    for sigma in args.post_sigma:\n        _df = df.copy()\n        _df.loc[:, [f'sigma_{c}' for c in range(2, 283 + 1)]] = sigma\n        _df.loc[:, 'sigma_1'] = 2.0 * sigma\n\n        scores, _, _ = get_score(args, _df[_df.columns[:567]], solution = solution, print_score = False)\n        targets.append(scores)\n\n    targets = np.stack(targets, axis = 1)\n    targets = targets.argmax(1)\n\n    df['target'] = targets\n\n    '''\n    if args.post_wl != None:\n        df.loc[:, [f'wl_{c}' for c in range(2, 283 + 1)]] *= (1.0 + args.post_wl[0])\n        df.loc[:, ['wl_1']] *= (1.0 + args.post_wl[1])\n    '''\n\n    outputs = {}\n\n    kf = KFold(n_splits = args.n_fold, shuffle = True, random_state = args.seed)\n    for i, (train_index, test_index) in enumerate(kf.split(planet_ids)):\n        train_planets = planet_ids[train_index]\n        test_planets = planet_ids[test_index]\n\n        train_df = df[(df.planet_id.isin(train_planets)) & (df.success == True)].reset_index(drop = True)\n        test_df = df[df.planet_id.isin(test_planets)].reset_index(drop = True)\n\n        train_star_info = star_info.loc[train_df.planet_id]\n        test_star_info = star_info.loc[test_planets]\n\n        train_label = label.loc[train_df.planet_id]\n        test_label = label.loc[test_planets]\n\n        train_inputs = get_inputs(args, train_df, train_star_info)\n        train_targets = get_targets(args, train_df, train_star_info, train_label)\n\n        test_inputs = get_inputs(args, test_df, test_star_info)\n        test_targets = get_targets(args, test_df, test_star_info, test_label)\n\n        model = GradientBoostingRegressor(\n            random_state = args.seed,\n            n_estimators = 250,\n            max_depth = 3,\n        )\n\n        model.fit(train_inputs, train_targets)\n\n        score = model.score(test_inputs, test_targets)\n\n        print(f'fold{i + 1}', score)\n\n        outputs[f'fold{i + 1}'] = {}\n        outputs[f'fold{i + 1}']['train_planets'] = train_planets\n        outputs[f'fold{i + 1}']['test_planets'] = test_planets\n        outputs[f'fold{i + 1}']['model'] = model\n        outputs[f'fold{i + 1}']['score'] = score\n\n    return outputs\n\ndef get_cv(args, df):\n    outputs = train_function(args, df)\n\n    star_info = pd.read_csv(args.root + f'train_star_info.csv')\n    star_info = star_info.set_index('planet_id')\n\n    cv = []\n    scores = []\n    for i in range(args.n_fold):\n        output = outputs[f'fold{i + 1}']\n\n        planet_ids = output['test_planets']\n\n        _df = df.loc[df.planet_id.isin(planet_ids)]\n        _df = _df.reset_index(drop = True)\n\n        success = _df['success'].values\n        stage1_params = _df[['T1', 'T2', 'T3', 'T4']].values\n        stage2_params = np.concatenate([\n            _df[[f'airs_popt_{c}' for c in range(1, 1 + 369)]].values,\n            _df[['cost', 'nfev']].values,\n            _df[[f'fgs1_popt_{c}' for c in range(1, 1 + 13)]].values,\n        ], axis = 1)\n\n        if args.postprocess:\n            _df = postprocess_function(args, [output['model']], _df, star_info.loc[_df.planet_id], success, stage1_params, stage2_params)\n\n        weighted_scores, _, score = get_score(args, _df[_df.columns[:567]], print_score = False)\n        _df['weighted_score'] = weighted_scores\n        _df['fold'] = [i + 1] * len(_df)\n\n        cv.append(_df)\n        scores.append(score)\n\n    cv = pd.concat(cv, axis = 0)\n    cv = cv.reset_index(drop = True)\n    return cv, scores, outputs\n\nif __name__ == '__main__':\n    args = CustomConfig()\n\n    args.cv_path = '/kaggle/input/nips-adc-data/cv6.csv'\n\n    args.postprocess = True\n    args.post_wl = [1e-3, -5e-3]\n    args.post_sigma = np.linspace(5e-5, 2e-3, 500)\n\n    args.pca_components = 6\n\n    df = pd.read_csv(args.cv_path)\n    df['planet_id'] = df['planet_id'].astype(int)\n    df = df[~df.planet_id.isin([2486733311])]\n    df = df.sort_values(by = 'planet_id')\n    df = df.reset_index(drop = True)\n\n    df, scores, _ = get_cv(args, df)\n    print('scores : ', scores)\n    print('cv : ', np.mean(scores))\n\n    df.to_csv('cv.csv', index = False)\n\n    '''\n    n_train :  1099\n    fold1 0.3815623750338394\n    fold2 0.1702556503990037\n    fold3 0.6332230605037319\n    fold4 0.7112889028015135\n    scores :  [0.5901797323680259, 0.5731023016822754, 0.5937743250282688, 0.5807186997366831]\n    cv :  0.5844437647038133\n    '''\"\"\"\npass","metadata":{"trusted":true,"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### run","metadata":{}},{"cell_type":"code","source":"def inference_function(args):\n    planet_ids = glob.glob(args.root + f'{args.split}/*')\n    planet_ids = sorted([int(_.split('/')[-1]) for _ in planet_ids])\n\n    if args.split == 'train':\n        planet_ids = planet_ids#[::11]#[:1]\n\n    star_info = pd.read_csv(args.root + f'{args.split}_star_info.csv')\n    star_info = star_info.set_index('planet_id')\n    star_info = star_info.loc[planet_ids]\n\n    rows = []\n    success = []\n    stage1_params = []\n    stage2_params = []\n    for planet_id in tqdm(planet_ids):\n        data = preprocess_function(args, planet_id)\n\n        _success, _stage1_params = stage1_function(args, data)\n\n        if _success:\n            try:\n                pred, _stage2_params = stage2_function(\n                    args, \n                    data, \n                    dict(star_info.loc[int(planet_id)]),\n                    *_stage1_params[:4],\n                )  \n                assert _stage2_params.shape == args.param_shape\n            except:\n                print(planet_id)\n                pred = np.array([args.naive_mean] * 283 + [args.naive_sigma] * 283)\n                _stage2_params = np.ones(args.param_shape)\n        else:\n            pred = np.array([args.naive_mean] * 283 + [args.naive_sigma] * 283)\n            _stage2_params = np.ones(args.param_shape)\n        \n        row = {'planet_id' : planet_id}\n        for i, column in enumerate(args.columns):\n            row[column] = pred[i]\n\n        rows.append(row)\n        success.append(_success)\n        stage1_params.append(_stage1_params)\n        stage2_params.append(_stage2_params)\n        \n    submission = pd.DataFrame(rows)\n    success = np.array(success)\n    stage1_params = np.stack(stage1_params, axis = 0)\n    stage2_params = np.stack(stage2_params, axis = 0)\n\n    if args.postprocess:         \n        df = pd.read_csv(args.cv_path)\n        df['planet_id'] = df['planet_id'].astype(int)\n        df = df[~df.planet_id.isin([2486733311])]\n        df = df.sort_values(by = 'planet_id')\n        df = df.reset_index(drop = True)\n\n        if args.split == 'train':\n            df = df[~df.planet_id.isin(planet_ids)]\n            df = df.reset_index(drop = True)\n        \n        df, _, outputs = get_cv(args, df)\n        models = [outputs[f'fold{i + 1}']['model'] for i in range(args.n_fold)]\n        submission = postprocess_function(args, models, submission, star_info, success, stage1_params, stage2_params)\n\n    return submission, success, stage1_params, stage2_params\n\nif __name__ == '__main__':\n    args = CustomConfig()\n\n    args.root = '/kaggle/input/ariel-data-challenge-2025/'\n    args.threshold = 1e-2\n    args.sigma = 4e-4\n    args.fgs_sigma = 8e-4\n\n    args.param_shape = (384, )\n\n    args.cv_path = '/kaggle/input/nips-adc-data/cv6.csv'\n\n    args.postprocess = False#True\n    args.post_wl = [1e-3, -5e-3]\n    args.post_sigma = np.linspace(5e-5, 2e-3, 500)\n\n    args.pca_components = 6\n    \n    submission, success, stage1_params, stage2_params = inference_function(args)\n\n    if args.split == 'train':\n        weighted_scores, unweighted_scores, score = get_score(args, submission)\n        print('score : ', score)\n        submission = pd.concat([\n            submission,\n            pd.DataFrame({'weighted_score': weighted_scores}),\n            pd.DataFrame({'unweighted_score': unweighted_scores}),\n            pd.DataFrame({'success': success}),\n\n            # T1, T2, T3, T4\n            pd.DataFrame(stage1_params[:, :4], columns = ['T1', 'T2', 'T3', 'T4']),\n\n            # airs_popt\n            pd.DataFrame(stage2_params[:, :369], columns = [f'airs_popt_{c}' for c in range(1, 1 + 369)]),\n\n            # cost, nfev\n            pd.DataFrame(stage2_params[:, 369:369 + 2], columns = ['cost', 'nfev']),\n\n            # fgs1_popt\n            pd.DataFrame(stage2_params[:, 369 + 2:369 + 2 + 13], columns = [f'fgs1_popt_{c}' for c in range(1, 1 + 13)]),\n        ], axis = 1)\n\n    submission.to_csv('submission.csv', index = False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-24T05:56:48.05709Z","iopub.execute_input":"2025-09-24T05:56:48.057875Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"'''\n[0.47232071]\n[0.44065144]\nscore :  0.4406514400215444\n\n[0.47223889]\n[0.43968711]\nscore :  0.4396871130029494\n'''\npass","metadata":{"trusted":true,"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission","metadata":{"trusted":true,"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null}]}