{"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}],"dockerImageVersionId":30746,"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') # 135000, 4\nstar = adc_info['star'].values\ndt_a = axis_info['AIRS-CH0-integration_time'].dropna().values\naxis_info = pd.read_parquet(dir1 + 'axis_info.parquet') # 135000, 4\n#labels = pd.read_csv(dir1 + 'train_labels.csv', index_col='planet_id')\n#y = labels.values\nprint('end init', int(time()-t0), 'sec')","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-10-30T01:21:32.492088Z","iopub.execute_input":"2024-10-30T01:21:32.492467Z","iopub.status.idle":"2024-10-30T01:21:37.324373Z","shell.execute_reply.started":"2024-10-30T01:21:32.492435Z","shell.execute_reply":"2024-10-30T01:21:37.323113Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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 = \"v20\"\n\nBINNING = 1 # 15\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            # set pixels that are close to 0 to 0. For dimension of 32 only - leave 356 alone: CV-=0.0008 - worse, undo\n            #signal[:,(signal.nanmean(axis=0).nanmean(axis=-1) < 40),:] = 0 # here limit is 40, 56% of pixels are reset. So, 18 out of 32.\n            signal = torch.nanmean(signal, axis=1).cpu().numpy()\n\n        # save the processed signal\n        #np.save(\n        #    f\"{planet_id}_{sensor}_signal_{VERSION}.npz\",\n        #    signal.astype(np.float64),\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-30T01:21:37.327118Z","iopub.execute_input":"2024-10-30T01:21:37.327916Z","iopub.status.idle":"2024-10-30T01:21:48.357342Z","shell.execute_reply.started":"2024-10-30T01:21:37.327877Z","shell.execute_reply":"2024-10-30T01:21:48.35626Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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. 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.07, 2.32, 2.73, 3.26, 3.59]\nstds = [0.18, 0.09, 0.08, 0.19, 0.20]\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-30T01:21:48.36328Z","iopub.execute_input":"2024-10-30T01:21:48.363636Z","iopub.status.idle":"2024-10-30T01:21:48.426546Z","shell.execute_reply.started":"2024-10-30T01:21:48.363607Z","shell.execute_reply":"2024-10-30T01:21:48.425384Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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] # use mean2(A), blended - same as all_s. 5625\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-30T01:21:48.427777Z","iopub.execute_input":"2024-10-30T01:21:48.428094Z","iopub.status.idle":"2024-10-30T01:21:48.463875Z","shell.execute_reply.started":"2024-10-30T01:21:48.428068Z","shell.execute_reply":"2024-10-30T01:21:48.462567Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# calc all_s *******************************************************************************************************\ndef 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() # MAE - minimize it. Seems better than MSE\n    return q * 1000 # this increases solver precision\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) # centered on 0\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) # smooth signal first - really helps\n\n    # excl wider transition period from polyfit - really helps\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) # centered on 0\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) # smooth signal first - really helps\n\n    # excl wider transition period from polyfit - really helps\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-30T01:21:48.465243Z","iopub.execute_input":"2024-10-30T01:21:48.465795Z","iopub.status.idle":"2024-10-30T01:21:48.659758Z","shell.execute_reply.started":"2024-10-30T01:21:48.465756Z","shell.execute_reply":"2024-10-30T01:21:48.65872Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# diff of obscured\nw = 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-30T01:21:48.661266Z","iopub.execute_input":"2024-10-30T01:21:48.661636Z","iopub.status.idle":"2024-10-30T01:21:48.671022Z","shell.execute_reply.started":"2024-10-30T01:21:48.661607Z","shell.execute_reply":"2024-10-30T01:21:48.669734Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# smoothing\nN = 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 # degree of polyfit for dd. param. 2 or 3 are the same\nfor i in range(x.shape[0]):\n    xx0 = np.arange(-x.shape[1]//2, x.shape[1]//2) # centered on 0\n    mult = np.zeros(x.shape[1], dtype=np.float32)\n    mult[bp[i]:bp2[i]] = 1\n\n    # subtract poly from x **************************************\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    # excl wider transition period from polyfit - really helps\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) # to make this even. Has to be >1 to fully capture the transition. Making it <2 does not help\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) # this is a better fit than min/max or median\n    d2i = -x[i,bp[i] - N2//2 - R:bp[i] - N2//2 + R,:].mean(axis=0)\n    di = (d1i + d2i) / 2 / unobscured[i] # ratio\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\n# use 5 best features: e=25.7\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#if log: print('all_s_blend', np.round(np.sqrt(((y.mean(-1)-all_s)**2).mean())*1e6, 1), int(time()-t0), 'sec')\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.\n# Then coeffs should be cleaner and L2 will work better\nind = [4,5,6,7,8] # 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 two approaches\nx = np.concatenate((d, dd), axis=1)\n\n# drop some features: 9,5,6,7,8\nind = [0,1,2,3,4,10,11,12,13,14,15,16,17]\nx = x[:,ind]\nif log: print('end filter', x.shape, int(time()-t0), 'sec')","metadata":{"execution":{"iopub.status.busy":"2024-10-30T01:21:48.67246Z","iopub.execute_input":"2024-10-30T01:21:48.672833Z","iopub.status.idle":"2024-10-30T01:21:48.974261Z","shell.execute_reply.started":"2024-10-30T01:21:48.672797Z","shell.execute_reply":"2024-10-30T01:21:48.973141Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# inference\nxs0 = x[star==0,:]\nxs1 = x[star==1,:]\nc = np.load(dir2 + 'c.npy') # use c12 here instead of c: reduces score by 0.002 - undo. V138(0.664) vs V130(0.666)\npred1 = np.matmul(x, c.T).astype(np.float32)\nc0 = np.load(dir2 + 'c0.npy')\npreds0 = np.matmul(xs0, c0.T).astype(np.float32)\nc1 = np.load(dir2 + 'c1.npy')\npreds1 = np.matmul(xs1, c1.T).astype(np.float32)\n\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 # =score * -1, so minimize this\n        \nmodel0 = tf.keras.models.load_model(dir2 + 'model_0.keras', custom_objects={'customloss1': customloss1})\nmodel1 = tf.keras.models.load_model(dir2 + 'model_1.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=34, verbose=0) + model1.predict(x, batch_size=34, verbose=0)) / 2\n\nprint('finished predicting', int(time()-t0), 'sec')","metadata":{"execution":{"iopub.status.busy":"2024-10-30T01:21:48.975692Z","iopub.execute_input":"2024-10-30T01:21:48.975988Z","iopub.status.idle":"2024-10-30T01:22:03.723964Z","shell.execute_reply.started":"2024-10-30T01:21:48.975962Z","shell.execute_reply":"2024-10-30T01:22:03.722893Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# assign sigma\npred1 = pred1 + all_s.reshape(-1, 1)\npred = np.concatenate((pred1, pred1), axis=-1)\n\n# do not use different models for stars 0 and 1:V139(0....) vs V130(0.666)\n# star 0\nidx0 = star == 0\npreds0 = preds0 + all_s[idx0].reshape(-1, 1)\n#pred[idx0,:283] = preds0\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)\n#pred[idx1,:283] = preds1\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 # 200 for star 2+\n\n\n\n# sigma by planet\ns = pred[:,:283].std(-1)\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') # is v0 better than v1? Yes, by 0.001\n# 1.3: ...\nstd_mult0 = std_mult0 * 1.3 - 0.3\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 # tune this later\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\n# vary sigma by star\n# star 0:\n# *0.9 =-1 (143 vs 139)\n# *1.1 =+1 (144 vs 139) - test increasing even more (to do)\npred[star == 0,283:] *= 1.1\n\n# star 1:\n# *1.15 =-3 (145 vs 144)\n# *0.85 =+2 (147 vs 144) - test decreasing even more (to do)\n# *0.75=??? (152 vs 151)\npred[star == 1,283:] *= 0.75\n\n# star 2:\n# *1.15 =+1 (148 vs 147) - test increasing even more (to do)\npred[star == 2,283:] *= 1.15\n\n# star 3: best is prob around 1.025\n# *1.15 =-1 (149 vs 148)\n# *0.85 =-2 (151 vs 148)\npred[star == 3,283:] *= 1.025\n\n\nprint('finished std', pred.shape, int(time()-t0), 'sec')","metadata":{"execution":{"iopub.status.busy":"2024-10-30T01:22:03.725411Z","iopub.execute_input":"2024-10-30T01:22:03.726016Z","iopub.status.idle":"2024-10-30T01:22:03.754503Z","shell.execute_reply.started":"2024-10-30T01:22:03.725986Z","shell.execute_reply":"2024-10-30T01:22:03.753042Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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-30T01:22:03.757587Z","iopub.execute_input":"2024-10-30T01:22:03.757926Z","iopub.status.idle":"2024-10-30T01:22:03.782699Z","shell.execute_reply.started":"2024-10-30T01:22:03.757897Z","shell.execute_reply":"2024-10-30T01:22:03.78147Z"},"trusted":true},"execution_count":null,"outputs":[]}]}