{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.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":70367,"databundleVersionId":9188054,"sourceType":"competition"},{"sourceId":9722998,"sourceType":"datasetVersion","datasetId":5559910},{"sourceId":9774816,"sourceType":"datasetVersion","datasetId":5548325}],"dockerImageVersionId":30747,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import random, gc, warnings\nimport numpy as np\nimport pandas as pd\nfrom time import time\nimport scipy.stats\nfrom scipy.signal import savgol_filter\nimport cupy as cp\nfrom astropy.stats import sigma_clip\nfrom scipy.optimize import minimize\nt0 = time()\n\ndir1 = '/kaggle/input/ariel-data-challenge-2024/'\ndir2 = '/kaggle/input/ariel-ym1/'\nprefix = 'test/'\nlog = True\n\ndef apply_linear_corr_gpu(linear_corr, signal):\n    linear_corr = linear_corr.reshape(1, 6, linear_corr.shape[1], linear_corr.shape[2])\n    linear_corr = cp.asarray(linear_corr)\n    signal = cp.asarray(signal)   \n    signal = linear_corr[0,0,:,:] + signal * (linear_corr[0,1,:,:] + signal * (linear_corr[0,2,:,:] + signal * (linear_corr[0,3,:,:] + signal * (linear_corr[0,4,:,:] + signal * linear_corr[0,5,:,:]))))\n    return cp.asnumpy(signal)\n\nadc_info = pd.read_csv(dir1 + 'test_adc_info.csv', index_col='planet_id')\naxis_info = pd.read_parquet(dir1 + 'axis_info.parquet')\nstar = adc_info['star'].values\ndt_a = axis_info['AIRS-CH0-integration_time'].dropna().values\naxis_info = pd.read_parquet(dir1 + 'axis_info.parquet')\nprint('end init', int(time()-t0), 'sec')","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-10-27T22:57:27.348391Z","iopub.execute_input":"2024-10-27T22:57:27.348752Z","iopub.status.idle":"2024-10-27T22:57:32.067101Z","shell.execute_reply.started":"2024-10-27T22:57:27.348722Z","shell.execute_reply":"2024-10-27T22:57:32.06592Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# read data\nprint('start read data', int(time()-t0), 'sec')\n\nimport itertools\nfrom tqdm import tqdm\nimport torch\nimport torch.nn.functional as F\n\nROOT = \"/kaggle/input/ariel-data-challenge-2024/\"\nVERSION = \"v26\"\n\nBINNING = 1\n\nMODE = \"test\"\n\n\nsensor_sizes_dict = {\n    \"AIRS-CH0\": [[11250, 32, 356], [32, 356]],\n    \"FGS1\": [[135000, 32, 32], [32, 32]],\n}  # input, mask\n\n# 16 center pixels, rest is just noise\ncl = 8\ncr = 24\n\n\ndef get_gain_offset(planet_id, sensor, mode):\n    \"\"\"\n    Get the gain and offset for a given planet and sensor\n    \"\"\"\n    gain_offset_csv = pd.read_csv(f\"{ROOT}/{mode}_adc_info.csv\")\n    planet_gain_offset = gain_offset_csv[gain_offset_csv[\"planet_id\"] == planet_id]\n\n    gain = planet_gain_offset[sensor + \"_adc_gain\"].values[0]\n    offset = planet_gain_offset[sensor + \"_adc_offset\"].values[0]\n    return gain, offset\n\n\ndef read_data(planet_id, sensor, mode):\n    \"\"\"\n    Read the data for a given planet and sensor\n    \"\"\"\n    # get all noise correction frames and signal\n    signal = pd.read_parquet(\n        f\"{ROOT}/{mode}/{planet_id}/{sensor}_signal.parquet\",\n        engine=\"pyarrow\",\n    )\n    dark_frame = pd.read_parquet(\n        f\"{ROOT}/{mode}/{planet_id}/{sensor}_calibration/dark.parquet\",\n        engine=\"pyarrow\",\n    )\n    dead_frame = pd.read_parquet(\n        f\"{ROOT}/{mode}/{planet_id}/{sensor}_calibration/dead.parquet\",\n        engine=\"pyarrow\",\n    )\n    linear_corr_frame = pd.read_parquet(\n        f\"{ROOT}/{mode}/{planet_id}/{sensor}_calibration/linear_corr.parquet\",\n        engine=\"pyarrow\",\n    )\n    flat_frame = pd.read_parquet(\n        f\"{ROOT}/{mode}/{planet_id}/{sensor}_calibration/flat.parquet\",\n        engine=\"pyarrow\",\n    )\n    # read_frame = pd.read_parquet(\n    #     f\"{ROOT}/{mode}/{planet_id}/{sensor}_calibration/read.parquet\",\n    #     engine=\"pyarrow\",\n    # )\n\n    # reshape to sensor shape and cast to float64\n    signal = signal.values.astype(np.float64).reshape(sensor_sizes_dict[sensor][0])[\n        :, cl:cr, :\n    ]\n    dark_frame = dark_frame.values.astype(np.float64).reshape(\n        sensor_sizes_dict[sensor][1]\n    )[cl:cr, :]\n    dead_frame = dead_frame.values.reshape(sensor_sizes_dict[sensor][1])[cl:cr, :]\n    flat_frame = flat_frame.values.astype(np.float64).reshape(\n        sensor_sizes_dict[sensor][1]\n    )[cl:cr, :]\n    # read_frame = read_frame.values.reshape(sensor_sizes_dict[sensor][1])\n    linear_corr = linear_corr_frame.values.astype(np.float64).reshape(\n        [6] + sensor_sizes_dict[sensor][1]\n    )[:, cl:cr, :]\n\n    return (\n        signal,\n        dark_frame,\n        dead_frame,\n        linear_corr,\n        flat_frame,\n        # read_frame,\n    )\n\n\ndef ADC_convert(signal, gain, offset):\n    \"\"\"\n    Step 1: Analog-to-Digital Conversion (ADC) correction\n\n    The Analog-to-Digital Conversion (adc) is performed by the detector to convert the\n    pixel voltage into an integer number. We revert this operation by using the gain\n    and offset for the calibration files 'train_adc_info.csv'.\n    \"\"\"\n\n    return signal / gain + offset\n\n\ndef mask_hot_dead(signal, dead, dark):\n    \"\"\"\n    Step 2: Mask hot/dead pixel\n\n    The dead pixels map is a map of the pixels that do not respond to light and, thus,\n    can't be accounted for any calculation. In all these frames the dead pixels are\n    masked using python masked arrays. The bad pixels are thus masked but left\n    uncorrected. Some methods can be used to correct bad-pixels but this task,\n    if needed, is left to the participants.\n    \"\"\"\n\n    hot = sigma_clip(dark, sigma=5, maxiters=5).mask\n    hot = np.tile(hot, (signal.shape[0], 1, 1))\n    dead = np.tile(dead, (signal.shape[0], 1, 1))\n\n    # Set values to np.nan where dead or hot pixels are found\n    signal[dead] = np.nan\n    signal[hot] = np.nan\n    return signal\n\n\ndef apply_linear_corr(c, signal):\n    \"\"\"\n    Step 3: linearity Correction\n\n    The non-linearity of the pixels' response can be explained as capacitive leakage\n    on the readout electronics of each pixel during the integration time. The number\n    of electrons in the well is proportional to the number of photons that hit the\n    pixel, with a quantum efficiency coefficient. However, the response of the pixel\n    is not linear with the number of electrons in the well. This effect can be\n    described by a polynomial function of the number of electrons actually in the well.\n    The data is provided with calibration files linear_corr.parquet that are the\n    coefficients of the inverse polynomial function and can be used to correct this\n    non-linearity effect.\n    Using horner's method to evaluate the polynomial\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\n\ndef clean_dark(signal, dark, dt):\n    \"\"\"\n    Step 4: dark current subtraction\n\n    The data provided include calibration for dark current estimation, which can be\n    used to pre-process the observations. Dark current represents a constant signal\n    that accumulates in each pixel during the integration time, independent of the\n    incoming light. To obtain the corrected image, the following conventional approach\n    is applied: The data provided include calibration files such as dark frames or\n    dead pixels' maps. They can be used to pre-process the observations. The dark frame\n    is a map of the detector response to a very short exposure time, to correct for the\n    dark current of the detector.\n\n    image - (dark * dt)\n\n    The corrected image is conventionally obtained via the following: where the dark\n    current map is first corrected for the dead pixel.\n    \"\"\"\n\n    dark = torch.tile(dark, (signal.shape[0], 1, 1))\n    signal -= dark * dt[:, None, None]\n    return signal\n\n\ndef get_cds(signal):\n    \"\"\"\n    Step 5: Get Correlated Double Sampling (CDS)\n\n    The science frames are alternating between the start of the exposure and the end of\n    the exposure. The lecture scheme is a ramp with a double sampling, called\n    Correlated Double Sampling (CDS), the detector is read twice, once at the start\n    of the exposure and once at the end of the exposure. The final CDS is the\n    difference (End of exposure) - (Start of exposure).\n    \"\"\"\n\n    return torch.subtract(signal[1::2, :, :], signal[::2, :, :])\n\n\ndef bin_obs(signal, binning):\n    \"\"\"\n    Step 5.1: Bin Observations\n\n    The data provided are binned in the time dimension. The binning is performed by\n    summing the signal over the time dimension.\n    \"\"\"\n\n    assert signal.shape[0] % binning == 0  # Ensure the binning is possible\n\n    # cds_transposed = signal.transpose(0, 2, 1)\n    cds_binned = torch.zeros(\n        (\n            signal.shape[0] // binning,\n            signal.shape[1],\n            signal.shape[2],\n        ),\n        device=\"cuda:0\",\n    )\n    for i in range(signal.shape[0] // binning):\n        cds_binned[i, :, :] = torch.sum(\n            signal[i * binning : (i + 1) * binning, :, :], axis=0\n        )\n    return cds_binned\n\n\ndef correct_flat_field(flat, signal):\n    \"\"\"\n    Step 6: Flat Field Correction\n\n    The flat field is a map of the detector response to uniform illumination, to\n    correct for the pixel-to-pixel variations of the detector, for example the\n    different quantum efficiencies of each pixel.\n    \"\"\"\n\n    return signal / flat\n\n\ndef nan_interpolation(tensor):\n    # Assume tensor is of shape (batch, height, width)\n    nan_mask = torch.isnan(tensor)\n\n    # Replace NaNs with zero temporarily\n    tensor_filled = torch.where(\n        nan_mask, torch.tensor(0.0, device=tensor.device), tensor\n    )\n\n    # Create a binary mask (0 where NaNs were and 1 elsewhere)\n    ones = torch.ones_like(tensor, device=tensor.device)\n    weight = torch.where(nan_mask, torch.tensor(0.0, device=tensor.device), ones)\n\n    # Perform interpolation by convolving with a kernel\n    # using bilinear interpolation\n    kernel = torch.ones(1, 1, 1, 3, device=tensor.device, dtype=tensor.dtype)\n\n    # Apply padding to the tensor and weight to prevent boundary issues\n    tensor_padded = F.pad(\n        tensor_filled.unsqueeze(1), (1, 1, 0, 0), mode=\"replicate\"\n    ).squeeze(1)\n    weight_padded = F.pad(weight.unsqueeze(1), (1, 1, 0, 0), mode=\"replicate\").squeeze(\n        1\n    )\n\n    # Convolve the filled tensor and the weight mask\n    tensor_conv = F.conv2d(tensor_padded.unsqueeze(1), kernel, stride=1)\n    weight_conv = F.conv2d(weight_padded.unsqueeze(1), kernel, stride=1)\n\n    # Compute interpolated values (normalized by weights)\n    interpolated_tensor = tensor_conv / weight_conv\n\n    # Apply the interpolated values only to the positions of NaNs\n    result = torch.where(nan_mask, interpolated_tensor.squeeze(1), tensor)\n\n    return result\n\n\ndef process_planet(planet_id):\n    \"\"\"\n    Process a single planet's data\n    \"\"\"\n    axis_info = pd.read_parquet(ROOT + \"axis_info.parquet\")\n    dt_airs = axis_info[\"AIRS-CH0-integration_time\"].dropna().values\n\n    for sensor in [\"FGS1\", \"AIRS-CH0\"]:\n        # load all data for this planet and sensor\n        signal, dark_frame, dead_frame, linear_corr, flat_frame = read_data(\n            planet_id, sensor, mode=MODE\n        )\n        gain, offset = get_gain_offset(planet_id, sensor, mode=MODE)\n\n        # Step 1: ADC correction\n        signal = ADC_convert(signal, gain, offset)\n\n        # Step 2: Mask hot/dead pixel\n        signal = mask_hot_dead(signal, dead_frame, dark_frame)\n\n        # clip at 0\n        signal = signal.clip(0)\n\n        # Step 3: linearity Correction\n        signal = apply_linear_corr(\n            torch.tensor(linear_corr).to(\"cuda:0\"), torch.tensor(signal).to(\"cuda:0\")\n        )\n\n        # Step 4: dark current subtraction\n        if sensor == \"FGS1\":\n            dt = torch.ones(len(signal), device=\"cuda:0\") * 0.1\n            dt[1::2] += 4.5\n        elif sensor == \"AIRS-CH0\":\n            dt = torch.tensor(dt_airs).to(\"cuda:0\")\n            dt[1::2] += 0.1\n\n        signal = clean_dark(signal, torch.tensor(dark_frame).to(\"cuda:0\"), dt)\n\n        # Step 5: Get Correlated Double Sampling (CDS)\n        signal = get_cds(signal)\n\n        # Step 5.1: Bin Observations\n        if sensor == \"FGS1\":\n            signal = bin_obs(signal, binning=BINNING * 12)\n        elif sensor == \"AIRS-CH0\":\n            signal = bin_obs(signal, binning=BINNING)\n\n        # Step 6: Flat Field Correction\n        signal = correct_flat_field(torch.tensor(flat_frame).to(\"cuda:0\"), signal)\n\n        # Step 7: Interpolate NaNs (twice!)\n        signal = nan_interpolation(signal)\n        signal = nan_interpolation(signal)\n\n        # Step 8: Sum over spatial axis\n        if sensor == \"FGS1\":\n            signal = torch.nanmean(signal, axis=[1, 2]).cpu().numpy()\n        elif sensor == \"AIRS-CH0\":\n            signal = torch.nanmean(signal, axis=1).cpu().numpy()\n\n        if sensor == \"FGS1\":\n            s1 = signal.astype(np.float32).copy()\n    signal = signal.astype(np.float32) \n    # combine. F first\n    signal = np.concatenate((s1.reshape(-1,1), signal), axis=-1)\n    return signal\n\n\n\nadc_info = pd.read_csv(f\"{ROOT}/{MODE}_adc_info.csv\", index_col=\"planet_id\")\nplanet_ids = adc_info.index.tolist()\n\nx = np.zeros([len(planet_ids), 5625, 357], dtype=np.float32)\nfor i, planet_id in enumerate(planet_ids):\n    x[i,:,:] = process_planet(planet_id)\n    if i%5 == 0:\n        print('    ', i, len(planet_ids), int(time()-t0), 'sec')\n\nprint('end read data', x.shape, int(time()-t0), 'sec')","metadata":{"execution":{"iopub.status.busy":"2024-10-27T22:57:32.07323Z","iopub.execute_input":"2024-10-27T22:57:32.073542Z","iopub.status.idle":"2024-10-27T22:57:42.339328Z","shell.execute_reply.started":"2024-10-27T22:57:32.073516Z","shell.execute_reply":"2024-10-27T22:57:42.338355Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# group data\nsh = (x.shape[0], x.shape[1], 1)\nx1 = np.concatenate((x[:,:,0].reshape(sh), x[:,:,1:].mean(-1).reshape(sh)), axis=-1) # 1=F, 2=mean(A)\nx1 = np.concatenate((x1, x[:,:,1+37:-37].mean(-1).reshape(sh)), axis=-1) # 3=mean2(A) - true spectrum\n\n# scale all stars to the same (uniform) spectrum.\n# After totals - unscaled data is a lot more stable, use it for bp/all_s\nstar_list = pd.Series(star).unique()\nfor s in star_list:\n    x[star==s,:,:] /= x[star==s,:,:].mean(0).mean(0).reshape(1,1,-1)\n\n# add 3 components to capture the slope: level, and 2 linear ones\n# 4=level - Am3\nx1 = np.concatenate((x1, x[:,:,1+37:-37].mean(-1).reshape(sh)), axis=-1).astype(np.float32)\n\nx = np.concatenate((x[:,:,0].reshape(sh), np.flip(x[:,:,1+37:-37], -1)), -1) # flip and drop tails\nwl = pd.read_csv(dir1 + 'wavelengths.csv').values.astype(np.float32).ravel() # 283, 0.71 + 1.95 to 3.90\n\n# co2=2.03, co=2.35, h2o=2.69, 4.25(main), nh3=3.0(main), ch4=3.3(main)\ncc = [2.00, 2.38, 2.58, 2.76, 3.00, 3.32, 3.50, 3.98]\nstds = [0.08, 0.12, 0.13, 0.11, 0.16, 0.16, 0.19, 0.19]\nfor i in range(len(cc)): # loop over components\n    c = cc[i]\n    std = stds[i]\n    # construct gaussian component with mean and std\n    comp = np.exp(-(wl-c)**2/2/std**2)\n    comp /= comp.max() # JIC\n    # blend\n    x2 = (x * comp.reshape(1,1,-1)).sum(-1)\n    x1 = np.concatenate((x1, x2.reshape(sh)), axis=-1).astype(np.float32)\n\nx = x1 # grouped data\nif log: print('end grouping', x.shape, int(time()-t0), 'sec')","metadata":{"execution":{"iopub.status.busy":"2024-10-27T22:57:42.341702Z","iopub.execute_input":"2024-10-27T22:57:42.341999Z","iopub.status.idle":"2024-10-27T22:57:42.399101Z","shell.execute_reply.started":"2024-10-27T22:57:42.341974Z","shell.execute_reply":"2024-10-27T22:57:42.398078Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# find transit zones: use the derivative\ndef phase_detector(signal):\n    MIN = signal.shape[0] // 2\n    signal1 = signal[:MIN ]\n    signal2 = signal[MIN :]\n\n    first_derivative1 = np.gradient(signal1)\n    first_derivative2 = np.gradient(signal2)\n    \n    phase1 = np.argmin(first_derivative1)\n    phase2 = np.argmax(first_derivative2) + MIN\n    return phase1, phase2\n\nbp = np.zeros(x.shape[0], dtype=np.int32)\nbp2 = np.zeros(x.shape[0], dtype=np.int32)\nfor i in range(x.shape[0]):\n    signal = x[i,:,2] - 0.2 * x[i,:,1]\n    signal = signal.reshape(-1,15).mean(-1) # collapse by 15; 375\n    signal = savgol_filter(signal, 11, 2) # smooth\n    p1, p2 = phase_detector(signal)\n    bp[i] = 15 * p1\n    bp2[i] = 15 * p2\nif log: print('end bp', x.shape, int(0.49+np.mean(np.minimum(1842,np.maximum(1800,bp)))), int(time()-t0), 'sec')","metadata":{"execution":{"iopub.status.busy":"2024-10-27T22:57:42.400267Z","iopub.execute_input":"2024-10-27T22:57:42.400624Z","iopub.status.idle":"2024-10-27T22:57:42.433875Z","shell.execute_reply.started":"2024-10-27T22:57:42.400598Z","shell.execute_reply":"2024-10-27T22:57:42.432748Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def try_s(s):\n    y0 = signal * (1 + s * mult)\n    z = np.polyfit(xx0, y0, deg)\n    p = np.poly1d(z)\n    d = p(xx0) - y0\n    q = np.abs(d).mean()\n    return q * 1000\n\ndeg =  4 # degree of polyfit for all_s. param.\nNS2  = 70 # numper of periods used for smoothing. param.\nw    = 70 # width of transition period; wider to exclude uneven transition. param.\nall_s = []\nfor i in range(x.shape[0]):\n    xx0 = np.arange(-x.shape[1]//2, x.shape[1]//2)\n    mult = np.zeros(x.shape[1], dtype=np.float32)\n    mult[bp[i]:bp2[i]] = 1\n    \n    signal = x[i,:,2] - 0.2 * x[i,:,1] # use mean2(A), blended\n    signal = savgol_filter(signal, NS2, 2)\n\n    idx = [j for j in range(x.shape[1]) if j < bp[i] - w or j > bp2[i] + w or (j > bp[i] + w and j < bp2[i] - w)]\n    xx0 = xx0[idx]\n    signal = signal[idx]\n    mult = mult[idx]\n    \n    r = minimize(try_s, [0.0003], method = 'Nelder-Mead')\n    all_s.append(r.x[0])\nall_s = np.array(all_s).astype(np.float32)\n\nall_s2 = []\nfor i in range(x.shape[0]):\n    xx0 = np.arange(-x.shape[1]//2, x.shape[1]//2)\n    mult = np.zeros(x.shape[1], dtype=np.float32)\n    mult[bp[i]:bp2[i]] = 1\n    \n    signal = x[i,:,3] # use mean3(A)\n    signal = savgol_filter(signal, NS2, 2)\n\n    idx = [j for j in range(x.shape[1]) if j < bp[i] - w or j > bp2[i] + w or (j > bp[i] + w and j < bp2[i] - w)]\n    xx0 = xx0[idx]\n    signal = signal[idx]\n    mult = mult[idx]\n    \n    r = minimize(try_s, [0.0003], method = 'Nelder-Mead')\n    all_s2.append(r.x[0])\nall_s2 = np.array(all_s2).astype(np.float32)\nif log: print('end all_s', int(time()-t0), 'sec')\n","metadata":{"execution":{"iopub.status.busy":"2024-10-27T22:57:42.435553Z","iopub.execute_input":"2024-10-27T22:57:42.436417Z","iopub.status.idle":"2024-10-27T22:57:42.56755Z","shell.execute_reply.started":"2024-10-27T22:57:42.43638Z","shell.execute_reply":"2024-10-27T22:57:42.566614Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"w = 10\nunobscured = []\nobscured = []\nfor i in range(x.shape[0]):\n    obscuredi = x[i,bp[i] + w:bp2[i] - w,:].mean(axis=0)\n    unobscuredi = (x[i,:bp[i] - w,:].mean(axis=0) + x[i,bp2[i] + w:,:].mean(axis=0)) / 2\n    obscured.append(obscuredi)\n    unobscured.append(unobscuredi)\nunobscured = np.array(unobscured)\ndd = (unobscured - np.array(obscured)) / unobscured # ratio\nif log: print('end diff of obscured', int(time()-t0), 'sec')","metadata":{"execution":{"iopub.status.busy":"2024-10-27T22:57:42.568717Z","iopub.execute_input":"2024-10-27T22:57:42.569031Z","iopub.status.idle":"2024-10-27T22:57:42.577714Z","shell.execute_reply.started":"2024-10-27T22:57:42.569006Z","shell.execute_reply":"2024-10-27T22:57:42.576547Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"N = 250\nx = savgol_filter(x, N, 2, axis=1) # 673, 5625, 6. Here 3 vs 2 does not matter\n\n# fit polynomial to x\ndeg2 = 2\nfor i in range(x.shape[0]):\n    xx0 = np.arange(-x.shape[1]//2, x.shape[1]//2)\n    mult = np.zeros(x.shape[1], dtype=np.float32)\n    mult[bp[i]:bp2[i]] = 1\n\n    signal = x[i,:,:] # all components\n    signal = savgol_filter(signal, NS2, 2, axis=0)\n    y0 = signal * (1 + all_s[i] * mult.reshape(-1,1))\n\n    idx = [j for j in range(x.shape[1]) if j < bp[i] - w or j > bp2[i] + w or (j > bp[i] + w and j < bp2[i] - w)]\n    xx0 = xx0[idx]\n    y0 = y0[idx]\n    \n    z = np.polyfit(xx0, y0, deg2) # use deg2 here\n    pred = 0\n    xx0 = np.arange(-x.shape[1]//2, x.shape[1]//2) # centered on 0\n    for j in range(deg2): # excl const\n        pred = xx0.reshape(-1,1) * (z[j,:] + pred)\n    x[i,:,:] -= pred\n\nN2 = int(((N * 2.5) // 2) * 2) # fully capture the transition\nx = x[:,N2:,:] - x[:,:-N2,:] # diff: 673, 5625-2*N, 17 (cuts N from each end)\nd = []\nfor i in range(x.shape[0]):\n    R = int(N * .75) # window for finding peaks\n    d1i = x[i,bp2[i] - N2//2 - R:bp2[i] - N2//2 + R,:].mean(axis=0)\n    d2i = -x[i,bp[i] - N2//2 - R:bp[i] - N2//2 + R,:].mean(axis=0)\n    di = (d1i + d2i) / 2 / unobscured[i]\n    d.append(di)\ndel x\ngc.collect()\nd = np.array(d)\n\n# define all_s as best POSITIVE blend of all vars incl max/min\nall_s = np.maximum(all_s, d[:,3]) * .333 + np.minimum(all_s, d[:,2]) * .204 + np.minimum(all_s2, dd[:,3]) * .035 + np.maximum(all_s2, d[:,1]) * .152 + np.maximum(d[:,1], d[:,3]) * .28\n\n# redefine d/dd (0=F, 1=Am1, 2=Am2, 3=Am3, 4=c1, 5=c2, 6=c3, 7=c4, 8=c5) in terms of all_s/Am3.\nind = [4,5,6,7,8,9,10,11] # center c1-c5 on Am3\nd[:,ind]  = d[:,ind] - d[:,3].reshape(-1,1)\ndd[:,ind] = dd[:,ind] - dd[:,3].reshape(-1,1)\nind = [0,1,2,3] # center F/Am on all_s\nd[:,ind]  = d[:,ind] - all_s.reshape(-1,1)\ndd[:,ind] = dd[:,ind] - all_s.reshape(-1,1)\n\n# combine the two\nx = np.concatenate((d, dd), axis=1)\n\n# drop some features\nind = [0, 1, 4, 5, 7, 11, 13, 14, 15, 17, 18, 19, 20, 21, 22, 23]\nx = x[:,ind]\nif log: print('end filter', x.shape, int(time()-t0), 'sec')","metadata":{"execution":{"iopub.status.busy":"2024-10-27T22:57:42.579289Z","iopub.execute_input":"2024-10-27T22:57:42.5797Z","iopub.status.idle":"2024-10-27T22:57:42.873412Z","shell.execute_reply.started":"2024-10-27T22:57:42.579668Z","shell.execute_reply":"2024-10-27T22:57:42.872444Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# inference\nxs0 = x[star==0,:]\nxs1 = x[star==1,:]\nc = np.load('/kaggle/input/ariel24-models/c_10.npy')\nc_2 = np.load('/kaggle/input/ariel24-models/c_9.npy')\npred1 = (np.matmul(x, c.T).astype(np.float32) + np.matmul(x, c_2.T).astype(np.float32)) / 2\nc0 = np.load('/kaggle/input/ariel24-models/c0_10.npy')\nc0_2 = np.load('/kaggle/input/ariel24-models/c0_9.npy')\npreds0 = (np.matmul(xs0, c0.T).astype(np.float32) + np.matmul(xs0, c0_2.T).astype(np.float32)) / 2\nc1 = np.load('/kaggle/input/ariel24-models/c1_10.npy')\nc1_2 = np.load('/kaggle/input/ariel24-models/c1_9.npy')\npreds1 = (np.matmul(xs1, c1.T).astype(np.float32) + np.matmul(xs1, c1_2.T).astype(np.float32)) / 2\n\n\n# NN to predict sigma\nimport tensorflow as tf\nfrom tensorflow import keras\nfrom tensorflow.keras.optimizers import Adam\nfrom tensorflow.keras.optimizers.schedules import CosineDecayRestarts\nfrom tensorflow.keras.layers import Input, Dense\n\ndef customloss1(mse, p):\n    var = p * p\n    d2 = (tf.math.log(var) + mse / var + 11.724) / 11.3017\n    loss = tf.math.reduce_mean(d2, axis=-1)\n    return loss\n        \nmodel0 = tf.keras.models.load_model('/kaggle/input/ariel24-models/model_0_10.keras', custom_objects={'customloss1': customloss1})\nmodel1 = tf.keras.models.load_model('/kaggle/input/ariel24-models/model_1_10.keras', custom_objects={'customloss1': customloss1})\nmodel2 = tf.keras.models.load_model('/kaggle/input/ariel24-models/model_0_9.keras', custom_objects={'customloss1': customloss1})\nmodel3 = tf.keras.models.load_model('/kaggle/input/ariel24-models/model_1_9.keras', custom_objects={'customloss1': customloss1})\n\n# norm x\nx = np.concatenate((x, 0.1 * all_s.reshape(-1,1)), axis=1)\nx *= 6000\nwith tf.device('/GPU:0'):\n    p2 = (model0.predict(x, batch_size=64, verbose=0) + model1.predict(x, batch_size=64, verbose=0) + model2.predict(x, batch_size=64, verbose=0) + model3.predict(x, batch_size=64, verbose=0)) / 4\n\nprint('finished predicting', int(time()-t0), 'sec')","metadata":{"execution":{"iopub.status.busy":"2024-10-27T22:57:42.874808Z","iopub.execute_input":"2024-10-27T22:57:42.875498Z","iopub.status.idle":"2024-10-27T22:57:56.67744Z","shell.execute_reply.started":"2024-10-27T22:57:42.875462Z","shell.execute_reply":"2024-10-27T22:57:56.676443Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# assign sigma\npred1 = pred1 + all_s.reshape(-1, 1)\npred = np.concatenate((pred1, pred1), axis=-1)\n\n# star 0\nidx0 = star == 0\npreds0 = preds0 + all_s[idx0].reshape(-1, 1)\npred[idx0,283:] = 0.000042 * 3 * .9 * .9 * .9 * .9 * .9 * .9\n\n# star 1\nidx1 = star == 1\npreds1 = preds1 + all_s[idx1].reshape(-1, 1)\npred[idx1,283:] = 0.000034 * 3 * .9 * .9 * .9 * .9 * .9 * .9\n\n# star 2+\nidx2 = star > 1\npred[idx2,283:] = 0.000200 * .9 * .9 * .9 * .9 * .9 * .9 * .9\n\n\n\n# sigma by planet\ns = pred[:,:283].std(-1) + 10e-6\n# scale to the same total\nsig0 = pred[:,283:].mean()\ns *= sig0 / s.mean()\n\n# assign std\npred[:,283:] = s.reshape(-1, 1)\n\n# apply mult by WL\nstd_mult0 = np.load(dir2 + 'std_mult0.npy')\npred[:,283:] = pred[:,283:] * std_mult0.reshape(1, -1)\n\n\n\n# NN sigma\n# scale to the same total\np2 *= sig0 / p2.mean()\n# scale p2 differently; 1.3 seems best\np2 *= 1.3\n\n\n# blend old(pred std) and new(NN)\npred[:,283:] = p2 * 0.1 + 0.9 * pred[:,283:]\n\n# combine prior best prediction with const * std_mult0. Specifically: sigma_new = W * std_mult0 * sigma_const + (1-W) * sigma_old.\nsigma_const = 50e-6\nW = 0.3\nsigma_old = pred[:,283:].copy()\nsigma_new = W * std_mult0.reshape(1, -1) * sigma_const + (1-W) * sigma_old\npred[:,283:] = sigma_new\n\n\n# vary sigma by star\n# star 0:\npred[star == 0,283:] *= 1.1\n# star 1:\npred[star == 1,283:] *= 0.70\n# star 2:\npred[star == 2,283:] *= 1.15\n# star 3: best is prob around 1.025\npred[star == 3,283:] *= 1.025\n\n# final tune\npred[:, 283:] *= 0.98\n\nprint('finished std', pred.shape, int(time()-t0), 'sec')","metadata":{"execution":{"iopub.status.busy":"2024-10-27T22:57:56.680385Z","iopub.execute_input":"2024-10-27T22:57:56.680968Z","iopub.status.idle":"2024-10-27T22:57:56.701585Z","shell.execute_reply.started":"2024-10-27T22:57:56.680941Z","shell.execute_reply":"2024-10-27T22:57:56.700312Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# save submission\nsub = pd.DataFrame(pred.astype(np.float32))\nsub.index = planet_ids\nsub.index.name = \"planet_id\"\ncol1 = ['wl_'+str(i+1) for i in range(283)]\ncol2 = ['sigma_'+str(i+1) for i in range(283)]\nsub.columns = col1 + col2\nsub.to_csv('submission.csv')\nprint('finished sub', int(time()-t0), 'sec')","metadata":{"execution":{"iopub.status.busy":"2024-10-27T22:57:56.702649Z","iopub.execute_input":"2024-10-27T22:57:56.703023Z","iopub.status.idle":"2024-10-27T22:57:56.732456Z","shell.execute_reply.started":"2024-10-27T22:57:56.702984Z","shell.execute_reply":"2024-10-27T22:57:56.731418Z"},"trusted":true},"outputs":[],"execution_count":null}]}