{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.14","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":70367,"databundleVersionId":9188054,"sourceType":"competition"}],"dockerImageVersionId":30762,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport polars as pl\nimport numba as nb\nfrom astropy.stats import sigma_clip\n\nimport scipy.stats\nfrom scipy.signal import fftconvolve\nfrom scipy.signal.windows import get_window\nfrom bisect import bisect\n\nfrom scipy.optimize import curve_fit, least_squares\nimport jax\nimport jax.numpy as jnp\nfrom jax import jacfwd\n\nfrom tqdm import tqdm\nimport matplotlib.pyplot as plt","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-11-02T21:39:00.317304Z","iopub.execute_input":"2024-11-02T21:39:00.318179Z","iopub.status.idle":"2024-11-02T21:39:00.324962Z","shell.execute_reply.started":"2024-11-02T21:39:00.31813Z","shell.execute_reply":"2024-11-02T21:39:00.323926Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Data loading and calibration\nI use the organizers' calibration guidelines and accelerate the process with GPU using CuPy.","metadata":{}},{"cell_type":"code","source":"def load_planet(idx, data_type, engine=\"fastparquet\"):\n    if data_type == \"FGS1\":\n        img_shape = (32, 32)\n    elif data_type == \"AIRS-CH0\":\n        img_shape = (32, 356)\n    else:\n        raise ValueError\n    signal = pl.read_parquet(f'/kaggle/input/ariel-data-challenge-2024/{dataset}/{idx}/{data_type}_signal.parquet').to_numpy().astype(np.float32)\n    return signal\n\ndef load_axis_info():\n    return pd.read_parquet(\"/kaggle/input/ariel-data-challenge-2024/axis_info.parquet\")\n\ndef load_adc_info(dataset):\n    return pd.read_csv(f'/kaggle/input/ariel-data-challenge-2024/{dataset}_adc_info.csv')#, index_col='planet_id')\n\ndef load_calibration(idx, data_type,dataset, cal):\n    calibration = pd.read_parquet(f'/kaggle/input/ariel-data-challenge-2024/{dataset}/{idx}/{data_type}_calibration/{cal}.parquet').values.astype(np.float32)\n    return calibration\n\ndef load_label():\n    Y= pd.read_csv(\"/kaggle/input/ariel-data-challenge-2024/train_labels.csv\")\n    return Y.values[:,1:]\n\n#@nb.njit()\ndef ADC_convert(signal, gain, offset):\n    signal /= gain\n    signal += offset\n    return signal\n\ndef mask_hot_dead(signal, dead, dark):\n    hot = sigma_clip(dark, sigma=5, maxiters=5).mask\n    hot_mask = hot.reshape((-1,))\n    dead_mask = (dead == 1.0).reshape((-1,))\n\n    signal[:, hot_mask] = np.nan\n    signal[:, dead_mask] = np.nan\n    return signal\n\n#@nb.njit()\ndef clean_dark(signal, dark, dt):\n    dark_current = (dt[:, np.newaxis] * dark)\n    signal -= dark_current\n    return signal\n\n#@nb.njit()\ndef clean_flat(signal, flat):\n    signal = (signal) / (flat)\n    return signal\n\ndef apply_linear_corr(linear_corr,signal):\n    for i in range(signal.shape[1]):\n        poli = np.poly1d(np.flip(linear_corr[:, i]))\n        signal[:, i] = poli(signal[:, i])\n    return signal\n\n\n#@nb.njit\ndef apply_linear_corr_np(c, signal):\n\n    assert c.shape[0] == 6  # Ensure the polynomial is of degree 5\n\n    return (\n        (((c[5] * signal + c[4]) * signal + c[3]) * signal + c[2]) * signal + c[1]\n    ) * signal + c[0]\n\nimport cupy as cp\n\ndef apply_linear_corr_gpu(linear_corr, clean_signal_gpu):\n    # Convert the input arrays to CuPy arrays\n    #linear_corr_gpu = cp.asarray(linear_corr)\n    linear_corr_gpu=linear_corr\n    corrected_signal_gpu = (\n        (\n            (\n                (linear_corr_gpu[5] * clean_signal_gpu + linear_corr_gpu[4])\n                * clean_signal_gpu\n                + linear_corr_gpu[3]\n            )\n            * clean_signal_gpu\n            + linear_corr_gpu[2]\n        )\n        * clean_signal_gpu\n        + linear_corr_gpu[1]\n    ) * clean_signal_gpu + linear_corr_gpu[0]\n\n    # Convert the result back to a NumPy array (if needed)\n    #corrected_signal = cp.asnumpy(corrected_signal_gpu)\n    \n    return corrected_signal_gpu\n\n\n\ndef bin_obs(signal ,binning):\n    signal_binned = np.zeros((signal.shape[0]//binning, signal.shape[1], signal.shape[2]))\n    for i in range(signal.shape[0]//binning):\n        signal_binned[i, :, :] = np.mean(signal[i*binning:(i+1)*binning, :, :], axis=0)\n    return signal_binned\n\n\n\ndef numpy_to_cupy(data):\n    data_g = cp.empty_like(data)\n    data_g.set(data)\n    return data_g\n\n# This function load and calibrate the planet idx. All the preprocessing is done on gpu \n# and the function return a cupy array.\ndef load_calibrate(idx,\n                        data_type,\n                        axis_info,\n                        adc_info,\n                        dataset,\n                        do_no_negative = False,\n                        do_gain_offset = True,\n                        do_mask_hot_dead = True,\n                        do_linear_correction = True,\n                        do_black_noise = True,\n                        do_clean_flat = True,\n                        gpu=True,\n                        do_cds=True):\n    \n    if data_type == \"FGS1\":\n        img_shape = (32, 32)\n    elif data_type == \"AIRS-CH0\":\n        img_shape = (32, 356)\n    else:\n        raise ValueError\n    \n    # load signal\n    \n    idx = str(idx)\n    \n\n    signal = load_planet(idx, data_type, dataset)\n\n    \n    # load correction data\n    lin_corr = load_calibration(idx, data_type, dataset, \"linear_corr\").reshape(6, -1)\n    dark = load_calibration(idx, data_type, dataset, \"dark\").reshape((1,-1))\n    dead = load_calibration(idx, data_type, dataset, \"dead\").reshape(1,-1)\n    flat = load_calibration(idx, data_type, dataset, \"flat\").reshape(1,-1)\n\n\n    \n    # 1. gain offset\n\n    \n    if do_mask_hot_dead:\n        signal = mask_hot_dead(signal, dead, dark) # with NaN, avoiding use of masked arrays\n    \n    dt = np.ones(len(signal))*0.1\n    if gpu:\n        dt = numpy_to_cupy(dt)\n        lin_corr = numpy_to_cupy(lin_corr)\n        dark = numpy_to_cupy(dark)\n        dead = numpy_to_cupy(dead)\n        flat = numpy_to_cupy(flat)\n        signal = numpy_to_cupy(signal)\n\n    if do_gain_offset:\n        gain = adc_info[adc_info[\"planet_id\"] == int(idx)][f'{data_type}_adc_gain'].iloc[0]\n        offset = adc_info[adc_info[\"planet_id\"] == int(idx)][f'{data_type}_adc_offset'].iloc[0]\n        signal = ADC_convert(signal, gain, offset)\n    \n\n\n    if do_no_negative:\n        signal = np.where(signal < 0, 0., signal)\n    # 2. Correction linéaire\n    if do_linear_correction:\n        signal = apply_linear_corr_np(lin_corr,signal)\n\n    # 3. Soustraction bruit noir\n    if do_black_noise:\n\n        if data_type == \"FGS1\":\n            dt[1::2] += 0.1\n        else:\n            dt[1::2] += 4.5\n        signal = clean_dark(signal, dark, dt)\n    \n    # CDS\n    if do_cds:\n        signal = signal[1::2] - signal[0::2]\n    \n    if do_clean_flat:\n        signal = clean_flat(signal, flat)\n\n\n    \n    signa_np = signal.reshape((-1, *img_shape)).astype(np.float32)#.get()\n    return signa_np","metadata":{"execution":{"iopub.status.busy":"2024-11-02T21:39:00.327206Z","iopub.execute_input":"2024-11-02T21:39:00.327823Z","iopub.status.idle":"2024-11-02T21:39:00.358143Z","shell.execute_reply.started":"2024-11-02T21:39:00.327778Z","shell.execute_reply":"2024-11-02T21:39:00.357201Z"},"trusted":true,"jupyter":{"source_hidden":true},"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"@nb.njit\ndef bin_signal(signal, binning):\n    signal_binned = np.zeros(signal.shape[0]//binning)\n    for i in range(signal.shape[0]//binning):\n        signal_binned[i] = np.mean(signal[i*binning:(i+1)*binning])\n    return signal_binned\n\ndef get_win_type(win_type, win_size):\n    if win_type is None or win_type == \"none\":\n        filtre = np.ones(win_size)\n    else:\n        filtre = get_window(win_type, win_size)\n        \n    filtre= filtre.astype(np.float32)\n    filtre /= filtre.sum()\n    return filtre\n\n\ndef rolling(data, win_size, win_type=None):\n    filtre = get_win_type(win_type, win_size)\n    \n    # check if there is no nan value\n    isnan = np.isnan(data)\n    number_of_nan = np.sum(isnan)\n    assert number_of_nan == 0\n    \n    if isinstance(data, np.ndarray):\n        data_roll = fftconvolve(data, filtre, mode=\"valid\")\n    elif isinstance(data, cp.ndarray):\n        data_roll = fftconvolve_gpu(data, cp.asarray(filtre), mode=\"valid\")\n    \n    assert data.shape[0] - (filtre.shape[0] -1)  == data_roll.shape[0]\n\n    return data_roll\n\n\n\ndef rolling_2D(data, win_size, win_type=None):\n    filtre = get_win_type(win_type, win_size)\n    filtre = filtre.reshape(1,-1)\n    if isinstance(data, np.ndarray):\n        data_roll = fftconvolve(data, filtre, mode=\"valid\")\n    elif isinstance(data, cp.ndarray):\n        data_roll = fftconvolve_gpu(data, cp.asarray(filtre), mode=\"valid\")\n    return data_roll","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2024-11-02T21:39:00.361985Z","iopub.execute_input":"2024-11-02T21:39:00.362391Z","iopub.status.idle":"2024-11-02T21:39:00.374847Z","shell.execute_reply.started":"2024-11-02T21:39:00.362349Z","shell.execute_reply":"2024-11-02T21:39:00.373966Z"},"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from tqdm import tqdm\nwavelengths = pd.read_csv('/kaggle/input/ariel-data-challenge-2024/wavelengths.csv')\ndataset = \"train\"   # Change this to test for test data\n\naxis_info = load_axis_info()\nadc_info = load_adc_info(dataset)\nids = adc_info.planet_id\n\nsignals_airs = np.zeros((len(ids), 356, 5625), dtype=np.float64)\nsignals_fgs = np.zeros((len(ids), 67500), dtype=np.float64)\n\n# LOAD AIRS | we take only the 15 brightest pixels\nfor i, idx in tqdm(enumerate(ids), total=len(ids)):\n    data = load_calibrate(idx, \"AIRS-CH0\", axis_info, adc_info, dataset, do_no_negative=True)\n    \n    mm = np.median(data, 0)\n    data = data[:,np.argsort(mm,0),np.arange(356)]\n    data_airs_mean = np.nanmean(data[:,-15:,:], 1).T\n\n    signals_airs[i,:,:] = data_airs_mean.get()\n\n\n# LOAD FGS | we take the 100 brightest pixels\nfor i, idx in tqdm(enumerate(ids), total=len(ids)):\n    data_fgs = load_calibrate(idx, \"FGS1\", axis_info, adc_info, dataset, do_no_negative=True)\n    \n    data_fgs_flatten = data_fgs.reshape(-1,32*32)\n    mean_per_pixel = data_fgs_flatten.mean(0)\n    data_fgs_mean = np.nanmean(data_fgs_flatten[:,np.argsort(mean_per_pixel)[-100:]],1)\n    \n    signals_fgs[i, :] = data_fgs_mean.get()","metadata":{"execution":{"iopub.status.busy":"2024-11-02T21:39:00.377178Z","iopub.execute_input":"2024-11-02T21:39:00.37773Z","iopub.status.idle":"2024-11-02T22:11:01.742772Z","shell.execute_reply.started":"2024-11-02T21:39:00.377698Z","shell.execute_reply":"2024-11-02T22:11:01.741812Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Find the transit boundaries using curve fitting\nThe signal can be modeled as the product of a polynomial function $p(x)$, that represents the background noise, and a step function $s(x)$, that represents the transit zone. The definition of the two functions:\n\n\n\\begin{equation}\np(x) = \\sum_{k=0}^{n}a_kx^k\n\\end{equation}\n\n\\begin{equation}\ns(x) = \\left\\{\\begin{array}{ll}\n            1 & \\quad x <  p_1 \\\\\n            \\frac{d-1}{p_2-p_1}(x-p_1) -1 & \\quad z_1 \\leq x \\leq  z_2 \\quad \\\\\n            d & \\quad p_2 < x < p_1 \\\\\n            \\frac{1-d}{p_3-p_4}(x-p_3) -d & \\quad z_3 \\leq x \\leq  z_4 \\quad \\\\\n            1 & \\quad p_4 < x\n        \\end{array}\\right.\n\\end{equation}\n\n\\begin{equation}\n f(x) = p(x) * s(x)\n\\end{equation}\n\n$a_i$ are the coeficients of the polynomial, $z_i$ are the boundaries of the transit zone and $d$ the depth.\n\nThe goal is to fit the function to the data, using least-squares method, to find the boundaries. I used `scipy.optimize.least-squares`, that provides an implementation of 3 gradient-based optimization methods. I used dogbox, but the other two offers similar performances on my experimentations.","metadata":{}},{"cell_type":"code","source":"def lin_eq(start, end, y1, y2, x):\n    return ((y2-y1)/(end-start))*(x-end)+y2\n\ndef relative_pos_to_peak(x, peak):\n    mid = x[len(x)//2]\n    \n    p1 = mid - peak[0]\n    a = p1-peak[1]\n    b = p1+peak[1]\n    \n    p2 = mid + peak[2]\n    c = p2-peak[3]\n    d = p2+peak[3]\n\n    return [bisect(x, a), bisect(x, b), bisect(x, c), bisect(x, d)]\n\ndef multiplicator(x, a, b, c, d, O):\n    out_phase = np.where((x-a<0) | (x-d>0), 1,0)\n    in_phase = np.where((x-b>0) & (x-c<0), O,0)\n\n    descente = lin_eq(a,b,1,O,x)\n    descente = np.where((x-a>=0) & (x-b<=0), descente,0)\n\n    monte = lin_eq(c,d,O,1,x)\n    monte = np.where((x-c>=0) & (x-d<=0), monte,0)\n\n    mult = out_phase + in_phase + monte + descente\n    return mult\n\ndef model_poly_complete(X,coefs,O, peak):\n    x=X[-2]\n    mid = x[len(x)//2]\n\n    p1 = mid - peak[0]\n    a = p1-peak[1]\n    b = p1+peak[1]\n    p2 = mid + peak[2]\n    c = p2-peak[3]\n    d = p2+peak[3]\n\n    p = np.matmul(coefs, X) # polynomial function\n    s = multiplicator(x, a,b,c,d,O) # step function\n    y = p*s\n\n    return y\n\ndef residual_complete(v, x, y):\n    return y - model_poly_complete(x, v[:-5], v[-5], v[-4:])#+np.random.random(len(y))/100000\n\ndef modelling_signal_complete(signal, degs, peak=None):\n    X = np.arange(len(signal), dtype=float)#[4:-5]\n    X = (X-X.mean())/X.std()\n    \n    all_res = []\n    scores=[]\n    for deg in degs:\n        XX = np.vstack([X**i for i in range(deg,-1,-1)])\n        yy=signal\n        x=XX[-2] \n        x-= x.min()\n        #print(x)\n        \n        bound_inf = np.zeros(deg+6) - np.inf\n        bound_sup = np.zeros(deg+6) + np.inf\n    \n        transit_size = 150\n        \n        bound_inf[-4:] = [x[elt] for elt in [transit_size,10,transit_size,10]]\n        bound_inf[-5] = 0\n        bound_sup[-4:] = [x[elt] for elt in [len(x)//2 - transit_size, transit_size-1, len(x)//2 - transit_size, transit_size-1]]\n        bound_sup[-5] = 1\n    \n        res = np.zeros(deg+6)#+1\n        res[-4:] = [x[elt] for elt in [900, 75, 900, 75]]\n        res[-5] = 0.95\n        \n        if peak is not None:\n            a,b,c,d = peak\n            Y_polyfit = np.concatenate((signal[:a],signal[d:]))\n            X_polyfit = XX[-2]\n            X_polyfit = np.concatenate((X_polyfit[:a],X_polyfit[d:]))\n            coefs = np.polyfit(X_polyfit, Y_polyfit,deg)\n            res[:deg+1] = coefs\n    \n        res = least_squares(residual_complete, res,\n                            args=(XX, yy),\n                            bounds = (bound_inf,bound_sup),\n                            method=\"dogbox\",\n                            jac=\"3-point\")\n                            #x_scale=\"jac\").x#.x#[0]\n        scores.append(res.cost)\n        res=res.x\n        all_res.append(\n            {\n             \"peak\":relative_pos_to_peak(x,res[-4:]),\n             \"model\":model_poly_complete(XX, res[:-5], res[-5], res[-4:]),\n             \"model_wo_step\":model_poly_complete(XX, res[:-5], 1, res[-4:]), \n             \"depth\":res[-5]\n             })\n    return all_res[np.argmin(scores)]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-02T22:11:01.745382Z","iopub.execute_input":"2024-11-02T22:11:01.745684Z","iopub.status.idle":"2024-11-02T22:11:01.771255Z","shell.execute_reply.started":"2024-11-02T22:11:01.745652Z","shell.execute_reply":"2024-11-02T22:11:01.770403Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"i=15\nmean_airs = signals_airs[i,:,:].mean(0)\nres = modelling_signal_complete(mean_airs, [2,3,4])\npoly = res[\"model_wo_step\"]\nstep = multiplicator(np.arange(mean_airs.shape[0]), *res[\"peak\"], res[\"depth\"])\n\nplt.figure(figsize=(16,4))\nplt.title(\"polynomial model\")\nplt.subplot(131)\nplt.title(\"polynomial function $p(x)$\")\n#plt.plot(rolling(mean_airs, 10))\nplt.plot(poly)\n\nplt.subplot(132)\nplt.title(\"step function $s(x)$\")\n\nplt.plot(step)\n\nplt.subplot(133)\nplt.title(\"complete model $f(x) = p(x)*s(x)$\")\n\nplt.plot(rolling(mean_airs, 10))\nplt.plot(poly*step)\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-02T22:11:01.772398Z","iopub.execute_input":"2024-11-02T22:11:01.772755Z","iopub.status.idle":"2024-11-02T22:11:03.456855Z","shell.execute_reply.started":"2024-11-02T22:11:01.772706Z","shell.execute_reply":"2024-11-02T22:11:03.455934Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"transit = np.zeros((len(ids),4), dtype=int)\nfor i in tqdm(range(len(ids))):\n    transit[i, :] = modelling_signal_complete(signals_airs[i,:,:].mean(0), [2,3,4])[\"peak\"]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-02T22:11:03.458031Z","iopub.execute_input":"2024-11-02T22:11:03.458376Z","iopub.status.idle":"2024-11-02T22:20:35.924915Z","shell.execute_reply.started":"2024-11-02T22:11:03.458342Z","shell.execute_reply":"2024-11-02T22:20:35.923732Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Estimate transit depth\nSame as transit zone estimation, but $z_i$ is fixed and we do it for each wavelength.","metadata":{}},{"cell_type":"code","source":"@nb.njit\ndef get_power_features(X, deg):\n    assert len(X.shape) == 1\n    X_f = np.zeros((deg+1, X.shape[0]))\n    for i in range(deg+1):\n        X_f[deg-i] = X**i\n    return X_f\n\ndef get_x_normalize(N):\n    X = np.arange(N, dtype=float)\n    return (X-X.mean())/X.std()\n\n# Do the same as linspace bust faster\ndef fast_linspace(start, end, N):\n    delta = end-start\n    return jnp.arange(start, end + delta / (2 * N), delta / (N-1))\n\ndef line_equ(X, x1, x2, y1, y2):\n    a=(y2-y1)/(x2-x1)\n    b = y1-a*x1\n    return a*X + b\n\ndef modelling_signal_wo_peak(signal, peak, degs=[4], p0=None, binning = None, O_estimate=0.995):\n    #signal : (n_signal, ts)\n\n    X = np.arange(signal.shape[0], dtype=float)\n    X=(X-X.mean())/X.std()\n    \n    X_fin = X.copy()\n    signal_fin = signal.copy()\n    peak_ini = peak.copy()\n    X_spacing = abs(X[1] - X[0])\n    X_spacing_ini = X_spacing.copy()\n\n    if binning is not None:\n        X = bin_signal(X, binning)\n        signal = bin_signal(signal, binning)\n        peak = peak//binning\n        X_spacing = X_spacing_ini*binning\n    a,b,c,d=peak\n    \n    all_res = []\n    error = []\n    for degree in degs:\n        # precompute x^i for performance\n        XX = get_power_features(X, degree)\n\n        # This type of optimization method is very sensitive to the initial point.\n        # We try to find a good starting point.\n        p0 = np.zeros(degree+2)\n        p0[-1] = O_estimate\n        # Guess background noise using polyfit\n        Y_polyfit = np.concatenate((signal[:a],signal[d:]))\n        X_polyfit = XX[-2]\n        X_polyfit = np.concatenate((X_polyfit[:a],X_polyfit[d:]))\n        coefs = np.polyfit(X_polyfit, Y_polyfit,degree)\n        p0[:degree+1] = coefs\n\n        f = lambda X, *v: model_poly_wo_peak(X, np.array(v[:-1]), v[-1], (peak-XX.shape[1]/2)*X_spacing)\n        j = lambda X, *v: model_poly_wo_peak_jacobian(X, np.array(v), (peak-XX.shape[1]/2)*X_spacing)\n\n        res, pocv = curve_fit(f, XX, signal, p0 = p0, method=\"trf\", jac=j, ftol=0.00005)\n\n        XX_fin = get_power_features(X_fin, degree)\n        modeled_signal = model_poly_wo_peak(XX_fin, res[:-1], res[-1], (peak_ini-XX_fin.shape[1]/2)*X_spacing_ini)\n        modeled_signal_wo_step = model_poly_wo_peak(XX_fin, res[:-1], 1, (peak_ini-XX_fin.shape[1]/2)*X_spacing_ini)\n\n        er = ((signal_fin - modeled_signal)**2).mean()\n        error.append(er)\n        all_res.append({\"model\" : modeled_signal,\n                       \"model_wo_step\" : modeled_signal_wo_step,\n                       \"depth\" : res[-1],\n                       \"degree\": degree,\n                       \"error\": er})\n        \n    return all_res[np.argmin(error)]\n\ndef residual_wo_peak(v, X, y, peak):\n    return y - model_poly_wo_peak(X, v[:-1], v[-1], peak)\n\ndef line_equ(X, x1, x2, y1, y2):\n    a=(y2-y1)/(x2-x1)\n    b = y1-a*x1\n    #b=0\n    return a*X + b\n\n@jax.jit\ndef model_poly_wo_peak_jacobian(X,v, peak):\n    f = lambda v : model_poly_wo_peak(X, jnp.array(v[:-1]), v[-1], peak)\n    J = jacfwd(f)\n    return J(v)\n\n@jax.jit\ndef model_poly_wo_peak(X,coefs,O, peak):\n    a,b,c,d = peak\n    y=jnp.matmul(coefs, X)\n    \n    O=jnp.minimum(O, 0.9999999999999)\n    t_in = line_equ(X[-2], a, b, 1, O)\n    t_out = line_equ(X[-2], c, d, O, 1)\n    mid = t_in.shape[0]//2\n    t_all = jnp.concatenate((t_in[:mid], t_out[mid:]))\n    t_all = jnp.clip(t_all, O, 1)\n    y*=t_all\n    \n    return y","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-02T22:20:35.926988Z","iopub.execute_input":"2024-11-02T22:20:35.927675Z","iopub.status.idle":"2024-11-02T22:20:35.979316Z","shell.execute_reply.started":"2024-11-02T22:20:35.927616Z","shell.execute_reply":"2024-11-02T22:20:35.978187Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**FGS**","metadata":{}},{"cell_type":"code","source":"depth_fgs = np.zeros(len(ids))\nfor i in tqdm(range(len(ids))):\n    peak = transit[i]*12\n    signal = signals_fgs[i]\n\n    model = modelling_signal_wo_peak(signal,\n                                     peak=peak,\n                                     degs=[2,3],\n                                     binning = 15)\n    \n    depth_estimate = 1 - model[\"depth\"]\n    depth_fgs[i] = depth_estimate","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-02T22:20:35.984472Z","iopub.execute_input":"2024-11-02T22:20:35.985369Z","iopub.status.idle":"2024-11-02T22:21:05.733308Z","shell.execute_reply.started":"2024-11-02T22:20:35.985307Z","shell.execute_reply":"2024-11-02T22:21:05.732346Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"**AIRS-CH0**","metadata":{}},{"cell_type":"code","source":"depth_airs = np.zeros((len(ids),356))\nerror_airs = np.zeros((len(ids),356))\n\nfor i in tqdm(range(signals_airs.shape[0])):\n    # Denoise airs global\n    # Degree 4 maximum\n    si_modeled = modelling_signal_wo_peak(signals_airs[i,:, :].mean(0),\n                                          peak=transit[i],\n                                          degs=[2,3,4],\n                                          binning = 5)\n\n    si_modeled_wo_step = si_modeled[\"model_wo_step\"]\n    si_modeled_wo_step /= si_modeled_wo_step.mean()\n    signals_airs[i,:,:] /= si_modeled_wo_step\n\n    # Estimate depth for each wavelength\n    # Degree 2 maximum\n    for j in range(0,signals_airs.shape[1],4):\n        si_modeled = modelling_signal_wo_peak(signals_airs[i,j:j+4, :].mean(0),\n                                                                      peak=transit[i],\n                                                                      degs=[1,2],\n                                              O_estimate = depth_fgs[i],\n                                             binning = 5)\n        depth_airs[i,j:j+4] = 1 - si_modeled[\"depth\"]\n        error_airs[i,j:j+4] = si_modeled[\"error\"]\n        \n        si_modeled_wo_step = si_modeled[\"model_wo_step\"]\n        si_modeled_wo_step /= si_modeled_wo_step.mean()\n        signals_airs[i,j:j+4,:] = si_modeled[\"model\"] / si_modeled_wo_step","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-02T22:21:05.734631Z","iopub.execute_input":"2024-11-02T22:21:05.735286Z","iopub.status.idle":"2024-11-02T22:38:46.815425Z","shell.execute_reply.started":"2024-11-02T22:21:05.735238Z","shell.execute_reply":"2024-11-02T22:38:46.814435Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"mean_airs = signals_airs.mean(-1)\ndepth_fgs = depth_fgs.reshape(-1,1)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-02T22:38:46.816871Z","iopub.execute_input":"2024-11-02T22:38:46.817604Z","iopub.status.idle":"2024-11-02T22:38:47.926283Z","shell.execute_reply.started":"2024-11-02T22:38:46.817547Z","shell.execute_reply":"2024-11-02T22:38:47.925486Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Feature smooth\nWe smooth the extracted features using a moving average windows.","metadata":{}},{"cell_type":"code","source":"def extract_depth_smooth_airs(depth,mean,error, win, win_type=\"none\", er_factor = 35):\n    factor = mean\n    depth = depth * factor\n    \n    rolling_depth = rolling_2D(depth[:,38-win//2: 320+win//2], win, win_type=win_type)\n    rolling_factor = rolling_2D(factor[:,38-win//2: 320+win//2], win, win_type=win_type)\n    \n    rolling_depth /= rolling_factor\n    return rolling_depth[:,::-1]\n    \ndef extract_smooth_airs_adawin(depth, mean):\n    depth_smooth_airs_adawin =np.zeros((len(depth),282))\n    rolled_depth = extract_depth_smooth_airs(depth, mean, [], 15, win_type=\"taylor\", er_factor = 35)\n    std_depth = rolled_depth[:,:-70].std(1)\n    depth = depth*mean\n\n    win1 = 73\n    win2 = 43\n    win3 = 13\n    \n    for p in range(depth.shape[0]):\n        std = std_depth[p]\n        if std < 0.00017:\n            win = win1\n        elif std <0.00025:\n            win = win2\n        else:\n            win = win3\n        res = rolling(depth[p, 38-win//2: 320+win//2], win, win_type=\"taylor\")\n        res /= rolling(mean[p, 38-win//2: 320+win//2], win, win_type=\"taylor\")\n        \n        depth_smooth_airs_adawin[p,:] = res\n        \n    return depth_smooth_airs_adawin[:,::-1]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-02T22:38:47.927449Z","iopub.execute_input":"2024-11-02T22:38:47.927777Z","iopub.status.idle":"2024-11-02T22:38:47.938609Z","shell.execute_reply.started":"2024-11-02T22:38:47.927725Z","shell.execute_reply":"2024-11-02T22:38:47.937636Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"depth_smooth_airs = extract_smooth_airs_adawin(depth_airs, mean_airs)\n\nfeatures = np.hstack((depth_fgs.reshape(-1,1),\n                      depth_smooth_airs\n                     ))+ 2.4e-5","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-02T22:38:47.9398Z","iopub.execute_input":"2024-11-02T22:38:47.940101Z","iopub.status.idle":"2024-11-02T22:38:48.489291Z","shell.execute_reply.started":"2024-11-02T22:38:47.940054Z","shell.execute_reply":"2024-11-02T22:38:48.488516Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"Y = load_label()\n\nplt.plot(Y[65,1:])\nplt.plot(features[65,1:], label=\"features with smoothing\")\nplt.legend()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-02T22:38:48.490373Z","iopub.execute_input":"2024-11-02T22:38:48.49065Z","iopub.status.idle":"2024-11-02T22:38:48.745511Z","shell.execute_reply.started":"2024-11-02T22:38:48.49062Z","shell.execute_reply":"2024-11-02T22:38:48.744486Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"rmse all :\", ((features - Y)**2).mean()**0.5)\nprint(\"rmse fgs :\", ((depth_fgs[:,0] - Y[:, 0])**2).mean()**0.5)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-02T22:38:48.746804Z","iopub.execute_input":"2024-11-02T22:38:48.747136Z","iopub.status.idle":"2024-11-02T22:38:48.75349Z","shell.execute_reply.started":"2024-11-02T22:38:48.747105Z","shell.execute_reply":"2024-11-02T22:38:48.75257Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class ParticipantVisibleError(Exception):\n    pass\n\ndef competition_score(\n        solution: pd.DataFrame,\n        submission: pd.DataFrame,\n        naive_mean: float,\n        naive_sigma: float,\n        sigma_true: float,\n        row_id_column_name='planet_id',\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        sigma_true: (float) essentially sets the scale of the outputs.\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 pd.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    y_true = solution.values\n\n    GLL_pred = np.sum(scipy.stats.norm.logpdf(y_true, loc=y_pred, scale=sigma_pred))\n    GLL_true = np.sum(scipy.stats.norm.logpdf(y_true, loc=y_true, scale=sigma_true * np.ones_like(y_true)))\n    GLL_mean = np.sum(scipy.stats.norm.logpdf(y_true, loc=naive_mean * np.ones_like(y_true), scale=naive_sigma * np.ones_like(y_true)))\n\n    submit_score = (GLL_pred - GLL_mean)/(GLL_true - GLL_mean)\n    return float(np.clip(submit_score, 0.0, 1.0))\n\ndef postprocessing(pred_array, index, sigma_pred, wavelengths):\n    \"\"\"Create a submission dataframe from its components\n    \n    Parameters:\n    pred_array: ndarray of shape (n_samples, 283)\n    index: pandas.Index of length n_samples with name 'planet_id'\n    sigma_pred: float\n    \n    Return value:\n    df: DataFrame of shape (n_samples, 566) with planet_id as index\n    \"\"\"\n    return pd.concat([pd.DataFrame(pred_array.clip(0, None), index=index, columns=wavelengths.columns),\n                      pd.DataFrame(sigma_pred, index=index, columns=[f\"sigma_{i}\" for i in range(1, 284)])],\n                     axis=1)","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2024-11-02T22:38:48.754869Z","iopub.execute_input":"2024-11-02T22:38:48.755227Z","iopub.status.idle":"2024-11-02T22:38:48.769552Z","shell.execute_reply.started":"2024-11-02T22:38:48.75519Z","shell.execute_reply":"2024-11-02T22:38:48.768582Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"sigma_global = 5.8e-05\nerror=np.tile(np.where(adc_info[['star']] <= 1, sigma_global, sigma_global*1), (1, 283))\n\n#predictions = model.predict(features)\npredictions = features\npredictions[np.where(predictions<0)] = 0\nsub_df = postprocessing(predictions,\n                        adc_info[\"planet_id\"],\n                        sigma_pred=error,\n                       wavelengths=wavelengths)\n\ndisplay(sub_df)\nsub_df.to_csv('submission.csv')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-11-02T22:40:55.989545Z","iopub.execute_input":"2024-11-02T22:40:55.990295Z","iopub.status.idle":"2024-11-02T22:40:56.624091Z","shell.execute_reply.started":"2024-11-02T22:40:55.990245Z","shell.execute_reply":"2024-11-02T22:40:56.623103Z"}},"outputs":[],"execution_count":null}]}