{"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":"none","dataSources":[{"sourceId":70367,"databundleVersionId":9188054,"sourceType":"competition"},{"sourceId":9548483,"sourceType":"datasetVersion","datasetId":5817588}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install /kaggle/input/astropy-install/*","metadata":{"execution":{"iopub.status.busy":"2024-10-29T14:21:11.945217Z","iopub.execute_input":"2024-10-29T14:21:11.945646Z","iopub.status.idle":"2024-10-29T14:21:49.53822Z","shell.execute_reply.started":"2024-10-29T14:21:11.945606Z","shell.execute_reply":"2024-10-29T14:21:49.536776Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport itertools\nimport os\nimport gc\nimport time\nfrom IPython.utils import io\nfrom astropy.stats import sigma_clip\nfrom multiprocessing import Pool\nimport statsmodels.api as sm\nfrom tqdm import tqdm\nimport scipy as sp\nimport math\nimport warnings\nwarnings.filterwarnings('ignore', category=pd.errors.PerformanceWarning)\nwarnings.filterwarnings('ignore', category=UserWarning)\nwarnings.filterwarnings('ignore', category=RuntimeWarning)","metadata":{"execution":{"iopub.status.busy":"2024-10-29T14:21:49.540511Z","iopub.execute_input":"2024-10-29T14:21:49.540897Z","iopub.status.idle":"2024-10-29T14:21:52.396908Z","shell.execute_reply.started":"2024-10-29T14:21:49.540857Z","shell.execute_reply":"2024-10-29T14:21:52.395479Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Ver\n* final version","metadata":{}},{"cell_type":"markdown","source":"# CFG","metadata":{}},{"cell_type":"code","source":"class CFG:\n    ### data processing cfg\n    mode='test'\n    train_dataset_path=None #'/kaggle/input/adc-dataset-v4'\n    process_batch_size=100\n    \n    ### phase split\n    bound_extending_step=1\n    \n    ### solver\n    win_avg_k=10\n    shift_avg=True\n    iqr_k=2.0\n    iqr_avg=1\n    poly_k=5\n    smooth_r=20\n    smooth_weight_decay=0.45\n    svd_dim=140\n    keep_fgs1=True\n    svd_snr_power=0.6\n    smooth_snr_power=1.0\n    \n    ### postprocesing\n    low_rank_dim=4\n    ","metadata":{"execution":{"iopub.status.busy":"2024-10-29T14:21:52.398383Z","iopub.execute_input":"2024-10-29T14:21:52.398915Z","iopub.status.idle":"2024-10-29T14:21:52.406Z","shell.execute_reply.started":"2024-10-29T14:21:52.398873Z","shell.execute_reply":"2024-10-29T14:21:52.404643Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Metric","metadata":{}},{"cell_type":"code","source":"def GLL(mu,sigma,y,reduce=True):\n    var=sigma**2\n    L= -0.5*(np.log(2*np.pi)+np.log(var)+(y-mu)**2/var)\n    if reduce:\n        return  L.mean()\n    else:\n        return L\n\ndef metric(pred_mu,pred_sigma,target,reduce=True,compute_ref=False):\n    L=GLL(pred_mu,pred_sigma,target,reduce)\n    L_ideal=GLL(0,1e-5,0)\n    if compute_ref:\n        L_ref=GLL(target.mean(),target.std(),target)\n    else:\n        L_ref=4.942895\n    return (L-L_ref)/(L_ideal-L_ref)","metadata":{"execution":{"iopub.status.busy":"2024-10-29T14:21:52.409176Z","iopub.execute_input":"2024-10-29T14:21:52.409709Z","iopub.status.idle":"2024-10-29T14:21:52.423807Z","shell.execute_reply.started":"2024-10-29T14:21:52.409655Z","shell.execute_reply":"2024-10-29T14:21:52.422488Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Data Processing\n* almost same as public notebooks, clip 2 more dims for last 60 wavelens","metadata":{}},{"cell_type":"code","source":"def ADC_convert(signal, gain, offset):\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, dark, dt):\n    dark = np.tile(dark, (signal.shape[0], 1, 1))\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":{"execution":{"iopub.status.busy":"2024-10-29T14:21:52.425295Z","iopub.execute_input":"2024-10-29T14:21:52.425771Z","iopub.status.idle":"2024-10-29T14:21:52.442304Z","shell.execute_reply.started":"2024-10-29T14:21:52.425706Z","shell.execute_reply":"2024-10-29T14:21:52.44108Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_planet_data(adc_info,dataset='train',sensor='AIRS-CH0',planet_id=785834,binning = 15):\n    data_dict={}\n    cut_inf, cut_sup = 39, 321\n    sensor_sizes_dict = {\"AIRS-CH0\":[[11250, 32, 356], [1, 32, cut_sup-cut_inf]], \"FGS1\":[[135000, 32, 32], [1, 32, 32]]}\n    binned_dict = {\"AIRS-CH0\":[11250 // binning // 2, 282], \"FGS1\":[135000 // binning // 2]}\n    linear_corr_dict = {\"AIRS-CH0\":(6, 32, 356), \"FGS1\":(6, 32, 32)}\n\n    signal = pd.read_parquet(f'/kaggle/input/ariel-data-challenge-2024/{dataset}/{planet_id}/{sensor}_signal.parquet').to_numpy()\n    dark_frame = pd.read_parquet(f'/kaggle/input/ariel-data-challenge-2024/{dataset}/' + str(planet_id) + '/' + sensor + '_calibration/dark.parquet', engine='pyarrow').to_numpy()\n    dead_frame = pd.read_parquet(f'/kaggle/input/ariel-data-challenge-2024/{dataset}/' + str(planet_id) + '/' + sensor + '_calibration/dead.parquet', engine='pyarrow').to_numpy()\n    flat_frame = pd.read_parquet(f'/kaggle/input/ariel-data-challenge-2024/{dataset}/' + str(planet_id) + '/' + sensor + '_calibration/flat.parquet', engine='pyarrow').to_numpy()\n    linear_corr = pd.read_parquet(f'/kaggle/input/ariel-data-challenge-2024/{dataset}/' + str(planet_id) + '/' + sensor + '_calibration/linear_corr.parquet').values.astype(np.float64).reshape(linear_corr_dict[sensor])\n\n    #### recover signal\n    signal = signal.reshape(sensor_sizes_dict[sensor][0]) \n    cond=adc_info['planet_id']==planet_id\n    gain = adc_info.loc[cond,f'{sensor}_adc_gain'].item()\n    offset = adc_info.loc[cond,f'{sensor}_adc_offset'].item()\n    signal = signal / gain + offset\n\n    ####\n    hot = sigma_clip(\n        dark_frame, sigma=5, maxiters=5\n    ).mask\n\n    if sensor == 'AIRS-CH0':\n        signal = signal[:, :, cut_inf:cut_sup] #11250 * 32 * 282\n        dt = np.ones(len(signal))*0.1 \n        dt[1::2] += 4.5 #@bilzard idea\n        linear_corr = linear_corr[:, :, cut_inf:cut_sup]\n        dark_frame = dark_frame[:, cut_inf:cut_sup]\n        dead_frame = dead_frame[:, cut_inf:cut_sup]\n        flat_frame = flat_frame[:, cut_inf:cut_sup]\n        hot = hot[:, cut_inf:cut_sup]\n    elif sensor == 'FGS1':\n        dt = np.ones(len(signal))*0.1\n        dt[1::2] += 0.1\n\n    signal = signal.clip(0) #@graySnow idea\n    linear_corr_signal = apply_linear_corr(linear_corr, signal)\n    signal = clean_dark(linear_corr_signal, dark_frame, dt)\n    flat = flat_frame.reshape(sensor_sizes_dict[sensor][1])\n    flat[dead_frame.reshape(sensor_sizes_dict[sensor][1])] = np.nan\n    flat[hot.reshape(sensor_sizes_dict[sensor][1])] = np.nan\n    signal = signal / flat\n    \n    if sensor == \"FGS1\":\n        signal = signal[:,10:22,10:22] \n        signal = signal.reshape(sensor_sizes_dict[sensor][0][0],144)\n\n    if sensor != \"FGS1\":\n        signal[:,:11,:60]=0\n        signal[:,21:,:60]=0\n        signal = signal[:,10:22,:] \n\n    mean_signal = np.nanmean(signal, axis=1) # mean over the 32*32(FGS1) or 32(CH0) pixels\n    cds_signal = (mean_signal[1::2] - mean_signal[0::2])\n\n    binned = np.zeros((binned_dict[sensor]))\n    for j in range(cds_signal.shape[0] // binning):\n        binned[j] = cds_signal[j*binning:j*binning+binning].mean(axis=0)\n\n    if sensor == \"FGS1\":\n        binned = binned.reshape((binned.shape[0],1))\n    return binned","metadata":{"execution":{"iopub.status.busy":"2024-10-29T14:21:52.443873Z","iopub.execute_input":"2024-10-29T14:21:52.444403Z","iopub.status.idle":"2024-10-29T14:21:52.464049Z","shell.execute_reply.started":"2024-10-29T14:21:52.444304Z","shell.execute_reply":"2024-10-29T14:21:52.462737Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_data_wrapper(inputs):\n    adc_info,dataset,sensor,planet_id,binning = inputs\n    X=get_planet_data(adc_info,dataset,sensor,planet_id,binning)\n    return X\n\ndef get_data_with_pool(adc_info,dataset='train',sensor='AIRS-CH0',binning = 15,pool_size=4):\n    tasks=[]\n    for i in adc_info.index:\n        planet_id=adc_info.loc[i,'planet_id']\n        tasks.append((adc_info.loc[[i]],dataset,sensor,planet_id,binning))\n    with Pool(pool_size) as p:\n        X = list(tqdm(p.imap(get_data_wrapper, tasks), total=len(tasks)))\n    X=np.stack(X)\n    return X","metadata":{"execution":{"iopub.status.busy":"2024-10-29T14:21:52.465471Z","iopub.execute_input":"2024-10-29T14:21:52.466132Z","iopub.status.idle":"2024-10-29T14:21:52.478032Z","shell.execute_reply.started":"2024-10-29T14:21:52.466081Z","shell.execute_reply":"2024-10-29T14:21:52.476764Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Signal Processing","metadata":{}},{"cell_type":"code","source":"def win_avg(X,k=1):\n    X=np.cumsum(np.concatenate([np.ones((1,X.shape[1])),X],axis=0),axis=0)\n    X=(X[k:]-X[:-k])/k\n    return X\n\ndef svd_denoise(signals,dim=1,weight=None):\n    signals_=signals.copy()\n    if weight is not None:\n        signals_*=weight.reshape((1,-1))\n    U,D,V=np.linalg.svd(signals_.T@signals_)\n    U=(signals_@V.T)/D**0.5\n    pcs=U[:,:dim]\n    signals_denoised=pcs@(np.linalg.inv(pcs.T@pcs)@(pcs.T@signals))\n    return signals_denoised\n\ndef get_trend(signals):\n    return signals.mean(axis=1)\n\ndef smoothing_over_wavelens(X,r=0,boundary='stuck',weight_decay=0,wave_weight=None):\n    X_=np.zeros(X.shape)\n    for i in range(X.shape[1]):\n        start=i-r\n        end=i+1+r\n        if start<0:\n            if boundary=='stuck':\n                start=0\n                end=2*r+1\n            elif boundary=='clip':\n                start=0\n        elif end>X.shape[1]:\n            if boundary=='stuck':\n                end=X.shape[1]\n                start=end-2*r-1\n            elif boundary=='clip':\n                end=X.shape[1]\n        # centered weight decay \n        weight=np.arange(start,end)\n        weight=(1-weight_decay)**np.abs(weight-i)\n        if wave_weight is not None:\n            weight*=wave_weight[start:end]\n        weight=weight.reshape((1,-1))\n        X_[:,i]=(X[:,start:end]*weight).mean(axis=1)\n    return X_\n\n\ndef get_snr_fusion_weight(phase_mark,signals,degree=3):\n    signal_mag=np.abs(signals).mean(axis=0)\n    signals=signals[phase_mark==1]\n    m=len(signals)\n    poly_matrix=np.ones((m,degree+1))*np.arange(1,m+1).reshape((-1,1))/m\n    for i in range(degree+1):\n        poly_matrix[:,i]**=i\n    signals_hat=poly_matrix@np.linalg.inv(poly_matrix.T@poly_matrix)@poly_matrix.T@signals\n    resi=signals-signals_hat\n    noise_var=(resi**2).sum(axis=0)/(m-degree-1)+1e-12\n    weight=signal_mag/noise_var\n    weight/=weight.mean()\n    return weight","metadata":{"execution":{"iopub.status.busy":"2024-10-29T14:21:52.479567Z","iopub.execute_input":"2024-10-29T14:21:52.479916Z","iopub.status.idle":"2024-10-29T14:21:52.496485Z","shell.execute_reply.started":"2024-10-29T14:21:52.47988Z","shell.execute_reply":"2024-10-29T14:21:52.495201Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Phase Split\n* find center of gap region based on smoothed gradient\n* find boundary of gap region based on monotonic condition","metadata":{}},{"cell_type":"code","source":"def get_phase_mark(trend,grad_accum_window=10,detrend_window=-1,gap_len=10,bound_extending_step=1,min_gap=0,detrend=False,detrend_degree=3,detrend_bound_gap=20):\n    # detrending\n    if detrend:\n        m=len(trend)\n        poly_matrix=np.ones((m,detrend_degree+1))*np.arange(1,m+1).reshape((-1,1))/m\n        for i in range(detrend_degree+1):\n            poly_matrix[:,i]**=i\n        fit_index=list(range(detrend_bound_gap))+list(range(m-detrend_bound_gap,m))\n        X_fit=poly_matrix[fit_index]\n        Y_fit=trend[fit_index]\n        trend_fit=poly_matrix@np.linalg.inv(X_fit.T@X_fit)@(X_fit.T@Y_fit)\n        trend=trend-trend_fit\n    \n    df=pd.DataFrame()\n    df['trend']=trend\n    df['grad']=df['trend'].diff()\n    df['grad_accum']=df['grad'].rolling(window=grad_accum_window+1,center=True,win_type='bartlett').mean()\n    if detrend_window>0:\n        df['grad_accum']-=df['grad'].rolling(window=detrend_window+1,center=True).mean()\n    left_center=df['grad_accum'].argmin()\n    right_center=df['grad_accum'].argmax()\n    if left_center>=right_center:\n        left_center=df.loc[gap_len:len(df)//2-gap_len,'grad_accum'].argmin()\n        right_center=df.loc[len(df)//2+gap_len:len(df)-gap_len,'grad_accum'].argmax()\n        \n    # extend boundary based on monotonic condition\n    start1=left_center-bound_extending_step-min_gap\n    last_gap=trend[start1:start1+bound_extending_step].mean()\n    while abs(start1-left_center)<gap_len:\n        new_gap=trend[start1-bound_extending_step:start1].mean()\n        if new_gap>=last_gap:\n            start1-=bound_extending_step\n            last_gap=new_gap\n        else:\n            break\n\n    end1=left_center+bound_extending_step+min_gap\n    last_gap=trend[end1-bound_extending_step:end1].mean()\n    while abs(end1-left_center)<gap_len:\n        new_gap=trend[end1:end1+bound_extending_step].mean()\n        if new_gap<=last_gap:\n            end1+=bound_extending_step\n            last_gap=new_gap\n        else:\n            break\n            \n    start2=right_center-bound_extending_step-min_gap\n    last_gap=trend[start2:start2+bound_extending_step].mean()\n    while abs(start2-right_center)<gap_len:\n        new_gap=trend[start2-bound_extending_step:start2].mean()\n        if new_gap<=last_gap:\n            start2-=bound_extending_step\n            last_gap=new_gap\n        else:\n            break\n\n    end2=right_center+bound_extending_step+min_gap\n    last_gap=trend[end2-bound_extending_step:end2].mean()\n    while abs(end2-right_center)<gap_len:\n        new_gap=trend[end2:end2+bound_extending_step].mean()\n        if new_gap>=last_gap:\n            end2+=bound_extending_step\n            last_gap=new_gap\n        else:\n            break\n    \n    phase_mark=2*np.ones(len(df))\n    phase_mark[:start1]=0\n    phase_mark[end2+1:]=0\n    phase_mark[end1+1:start2]=1\n    return phase_mark","metadata":{"execution":{"iopub.status.busy":"2024-10-29T14:21:52.498454Z","iopub.execute_input":"2024-10-29T14:21:52.499145Z","iopub.status.idle":"2024-10-29T14:21:52.517817Z","shell.execute_reply.started":"2024-10-29T14:21:52.499062Z","shell.execute_reply":"2024-10-29T14:21:52.516547Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Base Models\n* robust_ols_solver gives base coef\n* all_split_solver when signal is clear","metadata":{}},{"cell_type":"code","source":"def robust_ols_solver(phase_mark,signal,degree=4,iqr_k=None,iqr_avg=1,create_time_gap=True,prior_keep_cond=None,weight=None,return_details=False,res_mult=283):\n    details={\n        'keep_cond':None,\n        'resi':None,\n    }\n    oot_cond=phase_mark==0\n    non_gap_cond=phase_mark!=2\n    \n    Y=signal[non_gap_cond]\n    X=[]\n    \n    predictor=signal.reshape((-1,1)).copy()\n    predictor[~oot_cond]=0\n    predictor=predictor[non_gap_cond]\n    \n    if create_time_gap:\n        m=len(signal)\n    else:\n        m=non_gap_cond.sum()\n    poly_matrix=np.ones((m,degree+1))*np.arange(1,m+1).reshape((-1,1))/m\n    for i in range(degree+1):\n        poly_matrix[:,i]**=i\n    if create_time_gap:\n        poly_matrix=poly_matrix[non_gap_cond]\n        \n    X=np.concatenate([predictor,poly_matrix],axis=1)\n    if prior_keep_cond is not None:\n        X=X[prior_keep_cond]\n        Y=Y[prior_keep_cond]\n        if weight is not None:\n            weight=weight[prior_keep_cond]\n        details['keep_cond']=prior_keep_cond\n    elif iqr_k is not None:\n        model=sm.QuantReg(Y,X)\n        fit_res=model.fit(q=0.5)\n        resi=fit_res.resid\n        if iqr_avg>1:\n            resi=np.convolve(resi,np.ones(iqr_avg),mode='same')/np.convolve(np.ones(resi.shape),np.ones(iqr_avg),mode='same')\n        q1=np.quantile(resi,0.25)\n        q3=np.quantile(resi,0.75)\n        iqr=q3-q1\n        keep_cond=resi<(q3+iqr_k*iqr)\n        keep_cond&=resi>(q1-iqr_k*iqr)\n        X=X[keep_cond]\n        Y=Y[keep_cond]\n        details['keep_cond']=keep_cond\n        if weight is not None:\n            weight=weight[keep_cond]\n    if weight is not None:\n        X*=weight.reshape((-1,1))\n        Y*=weight\n    model=sm.OLS(Y,X)\n    fit_res=model.fit()\n    theta=fit_res.params\n    std=fit_res.bse\n    coef=theta[0]\n    coef_std=std[0]\n    details['resi']=fit_res.resid\n    if return_details:\n        return coef,coef_std,details\n    return coef,coef_std","metadata":{"execution":{"iopub.status.busy":"2024-10-29T14:21:52.519716Z","iopub.execute_input":"2024-10-29T14:21:52.52024Z","iopub.status.idle":"2024-10-29T14:21:52.535739Z","shell.execute_reply.started":"2024-10-29T14:21:52.520188Z","shell.execute_reply":"2024-10-29T14:21:52.534392Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def all_split_solver(phase_mark,signals,smooth_r=0,smooth_weight_decay=0,svd_dim=None,degree=4,prior_keep_cond=None,keep_fgs1=True,svd_snr_power=0.2,smooth_snr_power=0.5):\n    if keep_fgs1:\n        vec_copy=signals[:,0].copy()\n        \n    if svd_dim is not None:\n        snr_wave_weight=get_snr_fusion_weight(phase_mark,signals,degree=3)\n        signals=svd_denoise(signals,dim=svd_dim,weight=snr_wave_weight**svd_snr_power)\n        \n    if smooth_r>0:\n        snr_wave_weight=get_snr_fusion_weight(phase_mark,signals,degree=3)\n        signals=smoothing_over_wavelens(signals,r=smooth_r,boundary='stuck',weight_decay=smooth_weight_decay,wave_weight=snr_wave_weight**smooth_snr_power)\n        \n    if keep_fgs1:\n        signals[:,0]=vec_copy\n        \n    ### create eqn\n    oot_cond=phase_mark==0\n    non_gap_cond=phase_mark!=2\n\n    Y=signals[non_gap_cond]\n    X=signals.copy()\n    X[~oot_cond]=0\n    X=X[non_gap_cond]\n    \n    m=len(signals)\n    P=np.ones((m,degree+1))*np.arange(1,m+1).reshape((-1,1))/m\n    for i in range(degree+1):\n        P[:,i]**=i\n    P=P[non_gap_cond]\n    m=len(P)\n\n    if prior_keep_cond is not None:\n        Y=Y[prior_keep_cond]\n        X=X[prior_keep_cond]\n        P=P[prior_keep_cond]\n    \n    coef_map=np.linalg.inv(P.T@P)@P.T\n    X_=X-P@(coef_map@X)\n    Y_=Y-P@(coef_map@Y)\n    coef=(X_*Y_).sum(axis=0)/(X_*X_).sum(axis=0)\n    \n    resi=Y_-X_*coef.reshape((1,-1))\n    var=(resi**2).sum(axis=0)/(m-degree-2)\n    coef_std=np.sqrt(var/(X_*X_).sum(axis=0))\n    \n    # cov based on residual correlation\n    cov=resi/(resi**2).mean(axis=0)**0.5\n    cov*=coef_std\n    cov=(cov.T@cov)/len(cov)\n\n    autocorr=(resi[1:]*resi[:-1]).mean(axis=0)/((resi[1:]*resi[1:]).mean(axis=0)*(resi[:-1]*resi[:-1]).mean(axis=0))**0.5\n    pzr=2*(resi>0).astype('int32')-1\n    pzr=((pzr[1:]*pzr[:-1])<0).mean(axis=0)\n    return coef,coef_std,cov,autocorr,pzr","metadata":{"execution":{"iopub.status.busy":"2024-10-29T14:21:52.539132Z","iopub.execute_input":"2024-10-29T14:21:52.539545Z","iopub.status.idle":"2024-10-29T14:21:52.554743Z","shell.execute_reply.started":"2024-10-29T14:21:52.539494Z","shell.execute_reply":"2024-10-29T14:21:52.553488Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Solve Planets","metadata":{}},{"cell_type":"code","source":"def solve_planet(inputs):\n    planet_id,signals,wave_lens=inputs\n    signals=win_avg(signals,k=CFG.win_avg_k)\n    \n    res=pd.DataFrame()\n    res['wave_len']=wave_lens\n    res['planet_id']=planet_id\n    \n    res['coef']=0.0\n    res['coef_std']=0.0\n    res['coef2']=0.0\n    res['coef_std2']=0.0\n    cov=0\n    autocorr=0\n    pzr=0\n    count=0\n    oot_pts=0\n    it_pts=0\n    # average k model results with bin=k (k=CFG.win_avg_k)\n    for i in range(CFG.win_avg_k):\n        signals_=signals[i::CFG.win_avg_k]\n        trend=get_trend(signals_)\n        scale=30/CFG.win_avg_k\n        phase_mark=get_phase_mark(\n            trend,\n            grad_accum_window=int(10*scale),\n            gap_len=int(4*scale),\n            bound_extending_step=CFG.bound_extending_step,\n            detrend=True,\n            detrend_degree=3,\n            detrend_bound_gap=int(20*scale)\n        )\n        coef1,coef_std1,details=robust_ols_solver(\n            phase_mark,\n            trend,\n            degree=CFG.poly_k,\n            iqr_k=CFG.iqr_k,\n            iqr_avg=CFG.iqr_avg,\n            return_details=True,\n        )\n        coef2,coef_std2,cov2,autocorr2,pzr2=all_split_solver(\n            phase_mark,\n            signals_,\n            smooth_r=CFG.smooth_r,\n            smooth_weight_decay=CFG.smooth_weight_decay,\n            svd_dim=CFG.svd_dim,\n            degree=CFG.poly_k,\n            prior_keep_cond=details['keep_cond'],\n            keep_fgs1=CFG.keep_fgs1,\n            svd_snr_power=CFG.svd_snr_power,\n            smooth_snr_power=CFG.smooth_snr_power,\n        )\n        res['coef']+=coef1\n        res['coef_std']+=coef_std1**2\n        res['coef2']+=coef2\n        res['coef_std2']+=coef_std2**2\n        cov+=cov2\n        autocorr+=autocorr2\n        pzr+=pzr2\n        count+=1\n        oot_pts+=(phase_mark==0).sum()\n        it_pts+=(phase_mark==1).sum()\n        if not CFG.shift_avg:\n            break\n            \n    res['coef']/=count\n    res['coef_std']/=count\n    res['coef_std']**=0.5\n    res['coef2']/=count\n    # since noise in highly correlated among shifted results, here I just use simple average\n    res['coef_std2']/=count\n    res['coef_std2']**=0.5\n    res['cov2']=(cov/count).tolist()\n    res['autocorr']=autocorr/count\n    res['pzr']=pzr/count\n    res['oot_pts']=oot_pts/count\n    res['it_pts']=it_pts/count\n    return res","metadata":{"execution":{"iopub.status.busy":"2024-10-29T14:21:52.556843Z","iopub.execute_input":"2024-10-29T14:21:52.557399Z","iopub.status.idle":"2024-10-29T14:21:52.571404Z","shell.execute_reply.started":"2024-10-29T14:21:52.557345Z","shell.execute_reply":"2024-10-29T14:21:52.570186Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"mode=CFG.mode\nadc_info=pd.read_csv(f'/kaggle/input/ariel-data-challenge-2024/{mode}_adc_info.csv')\nwave_df=pd.read_csv('/kaggle/input/ariel-data-challenge-2024/wavelengths.csv')\nif mode=='train':\n    labels=pd.read_csv('/kaggle/input/ariel-data-challenge-2024/train_labels.csv')\nelse:\n    labels=None","metadata":{"execution":{"iopub.status.busy":"2024-10-29T14:21:52.573413Z","iopub.execute_input":"2024-10-29T14:21:52.574027Z","iopub.status.idle":"2024-10-29T14:21:52.606451Z","shell.execute_reply.started":"2024-10-29T14:21:52.573986Z","shell.execute_reply":"2024-10-29T14:21:52.605325Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"feature_df=[]\nbatchs=math.ceil(len(adc_info)/CFG.process_batch_size)\nfor i in range(batchs):\n    adc_info_=adc_info.iloc[CFG.process_batch_size*i:CFG.process_batch_size*(i+1)]\n    if CFG.mode=='train' and CFG.train_dataset_path is not None:\n        fgs1=None\n        ch0=None\n        planet_data=np.load(os.path.join(CFG.train_dataset_path,f'planet_data_batch{i}.npy'))\n    else:\n        fgs1=get_data_with_pool(adc_info_,dataset=mode,sensor='FGS1',binning = 12,pool_size=4)\n        ch0=get_data_with_pool(adc_info_,dataset=mode,sensor='AIRS-CH0',binning = 1,pool_size=4)\n        planet_data=np.concatenate(\n            [fgs1,ch0[:,:,::-1]],\n            axis=2\n        )\n    tasks=[]\n    for j,idx in enumerate(adc_info_.index):\n        planet_id=adc_info_.loc[idx,'planet_id']\n        signals=planet_data[j].copy()\n        tasks.append((planet_id,signals,wave_df.columns))\n    with Pool(4) as p:\n        feature_df_ = list(tqdm(p.imap(solve_planet, tasks), total=len(tasks)))\n    feature_df_=pd.concat(feature_df_,axis=0).reset_index(drop=True)\n    feature_df_=feature_df_.merge(adc_info[['planet_id','star']],on='planet_id',how='left')\n    feature_df.append(feature_df_)\n    ######\n    del fgs1\n    del ch0\n    del planet_data\n    del tasks\n    gc.collect()\n    #break\nfeature_df=pd.concat(feature_df,axis=0).reset_index(drop=True)\nif CFG.mode!='test':\n    feature_df.to_parquet('feature_df.pqt')","metadata":{"execution":{"iopub.status.busy":"2024-10-29T14:21:52.607819Z","iopub.execute_input":"2024-10-29T14:21:52.608155Z","iopub.status.idle":"2024-10-29T14:22:07.518451Z","shell.execute_reply.started":"2024-10-29T14:21:52.608122Z","shell.execute_reply":"2024-10-29T14:22:07.516914Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Postprocessing","metadata":{}},{"cell_type":"code","source":"def postprocesing(feature_df):\n    coef_ub=0.95\n    sigma_lb=5e-6\n    \n    # constant prediction as base coef, it has best mse error in general\n    feature_df['mu']=feature_df['coef'].copy()\n    feature_df['sigma']=feature_df['coef_std'].copy()\n    \n    # cases we don't know what's happening, never happen in training set and public LB\n    misspecified_cond=feature_df['mu']>coef_ub\n    misspecified_cond|=feature_df['coef2']>coef_ub\n    misspecified_cond|=feature_df['mu'].isna()\n    misspecified_cond|=feature_df['coef2'].isna()\n    misspecified_cond|=feature_df['sigma'].isna()\n    misspecified_cond|=feature_df['coef_std2'].isna()\n    misspecified_cond|=feature_df['sigma']==0\n    misspecified_cond|=feature_df['coef_std2']==0\n    \n    # denoise solved coef with population information, assuming underlying factors are low ranked\n    mu2=feature_df.pivot(index='planet_id',columns='wave_len',values='coef2')[wave_df.columns].values\n    U,_,_=np.linalg.svd(mu2.T)\n    U=U[:,:CFG.low_rank_dim]\n    for _,sub_df in tqdm(feature_df.groupby('planet_id')):\n        coef=sub_df['coef2'].values\n        cov=np.array(sub_df['cov2'].tolist())\n        # GLS with cov shrinkage\n        prc=np.linalg.inv(cov+4e-8*np.eye(len(coef)))\n        P=U@np.linalg.inv(U.T@prc@U)@U.T@prc\n        P2=U@np.linalg.inv(U.T@U)@U.T\n        coef_trf=P@coef\n        # use OLS cov, simple GLS cov is under estimated\n        cov_trf=P2@cov@P2.T\n        feature_df.loc[sub_df.index,'coef2']=coef_trf\n        feature_df.loc[sub_df.index,'coef_std2']=cov_trf.diagonal()**0.5\n    # some ensemble effect\n    alpha=0.6 \n    feature_df['mu']=alpha*feature_df['mu']+(1-alpha)*feature_df['coef2']\n    # correct for base coef, it is biased wrt energy&target distribution\n    feature_df['mu']*=feature_df.loc[~misspecified_cond,'coef2'].mean()/feature_df.loc[~misspecified_cond,'coef'].mean()\n    \n    # 'observe' highly biased coef from less biased coef2\n    feature_df['coef_gap']=feature_df.eval('abs(coef-coef2)')\n    feature_df['sigma']=feature_df[['coef_gap','coef_std2']].max(axis=1)\n    # use population estimation when elementwise accurate estimation is infeasible\n    feature_df['bound']=(feature_df['coef_gap']**2).mean()**0.5\n    feature_df['sigma']=feature_df[['sigma','bound']].min(axis=1)\n    # at least with its own coef_std\n    feature_df['sigma']=feature_df[['sigma','coef_std']].max(axis=1)\n    \n    # when significant bias exists less biased one is better\n    feature_df['planet_gap^2']=feature_df['coef_gap']**2\n    feature_df['planet_gap^2']=feature_df.groupby('planet_id')['planet_gap^2'].transform('mean')\n    feature_df['planet_var']=feature_df.eval('coef_std**2+coef_std2**2')\n    feature_df['planet_var']=feature_df.groupby('planet_id')['planet_var'].transform('mean')\n    sig_bias_cond=np.sqrt(feature_df['planet_gap^2']/feature_df['planet_var'])>1.5\n    sig_bias_cond|=feature_df.eval('coef_gap/sqrt(coef_std**2+coef_std2**2)')>2.0\n    sig_bias_cond&=~misspecified_cond\n    feature_df.loc[sig_bias_cond,'mu']=feature_df.loc[sig_bias_cond,'coef2']\n    \n    # some tricks to extend sigma, since coef2 is still biased and seems there are also unknown bias in data\n    feature_df['left_gap']=(feature_df['coef2']-feature_df.groupby('planet_id')['coef2'].shift(1)).abs().fillna(0)\n    feature_df['right_gap']=(feature_df['coef2']-feature_df.groupby('planet_id')['coef2'].shift(-1)).abs().fillna(0)\n    feature_df['coef_nbr_max_gap']=feature_df[['left_gap','right_gap']].max(axis=1)\n    feature_df.loc[sig_bias_cond,'sigma']=feature_df.loc[sig_bias_cond].eval('2.0*coef_std2+1.0*coef_nbr_max_gap')\n    \n    # cases we don't know what's happening\n    # check mu,sigma again after fusion\n    misspecified_cond|=feature_df['mu']>coef_ub\n    misspecified_cond|=feature_df['mu'].isna()\n    misspecified_cond|=feature_df['sigma'].isna()\n    if misspecified_cond.mean()>0:\n        print('misspecified rate:',misspecified_cond.mean().round(3))\n    feature_df.loc[misspecified_cond,'mu']=2.5e-3\n    feature_df.loc[misspecified_cond,'sigma']=1.7e-3\n    \n    # safe cutting\n    feature_df['mu']=feature_df['mu'].clip(lower=0)\n    feature_df['sigma']=feature_df['sigma'].clip(lower=sigma_lb) \n    return feature_df","metadata":{"execution":{"iopub.status.busy":"2024-10-29T14:22:07.520258Z","iopub.execute_input":"2024-10-29T14:22:07.520659Z","iopub.status.idle":"2024-10-29T14:22:07.539721Z","shell.execute_reply.started":"2024-10-29T14:22:07.52062Z","shell.execute_reply":"2024-10-29T14:22:07.538611Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"* less biased score","metadata":{}},{"cell_type":"code","source":"if labels is not None:\n    for star in [[0],[1],[0,1]]:\n        star_cond=feature_df['star'].isin(star)\n        if star_cond.sum()==0:\n            continue\n        mu=feature_df[star_cond].pivot(index='planet_id',columns='wave_len',values='coef2')[wave_df.columns].values\n        sigma=feature_df[star_cond].pivot(index='planet_id',columns='wave_len',values='coef_std2')[wave_df.columns].values\n        target=labels[labels['planet_id'].isin(feature_df[star_cond]['planet_id'])].values[:,1:]\n        score=metric(mu,sigma,target)\n        best_score=metric(mu,np.abs(target-mu),target)\n        print(f'star{star} score/best_score:',round(score,5),'/',round(best_score,5))","metadata":{"execution":{"iopub.status.busy":"2024-10-29T14:22:07.541743Z","iopub.execute_input":"2024-10-29T14:22:07.542253Z","iopub.status.idle":"2024-10-29T14:22:07.553651Z","shell.execute_reply.started":"2024-10-29T14:22:07.542203Z","shell.execute_reply":"2024-10-29T14:22:07.552448Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"* fusion score","metadata":{"execution":{"iopub.status.busy":"2024-10-19T07:16:13.656707Z","iopub.execute_input":"2024-10-19T07:16:13.657835Z","iopub.status.idle":"2024-10-19T07:16:13.665347Z","shell.execute_reply.started":"2024-10-19T07:16:13.657745Z","shell.execute_reply":"2024-10-19T07:16:13.66388Z"}}},{"cell_type":"code","source":"feature_df=postprocesing(feature_df)\nif labels is not None:\n    for star in [[0],[1],[0,1]]:\n        star_cond=feature_df['star'].isin(star)\n        if star_cond.sum()==0:\n            continue\n        mu=feature_df[star_cond].pivot(index='planet_id',columns='wave_len',values='mu')[wave_df.columns].values\n        sigma=feature_df[star_cond].pivot(index='planet_id',columns='wave_len',values='sigma')[wave_df.columns].values\n        target=labels[labels['planet_id'].isin(feature_df[star_cond]['planet_id'])].values[:,1:]\n        score=metric(mu,sigma,target)\n        best_score=metric(mu,np.abs(target-mu),target)\n        print(f'star{star} score/best_score:',round(score,5),'/',round(best_score,5))","metadata":{"execution":{"iopub.status.busy":"2024-10-29T14:22:07.555189Z","iopub.execute_input":"2024-10-29T14:22:07.555591Z","iopub.status.idle":"2024-10-29T14:22:07.727625Z","shell.execute_reply.started":"2024-10-29T14:22:07.555531Z","shell.execute_reply":"2024-10-29T14:22:07.726084Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"* less biased score after svd","metadata":{}},{"cell_type":"code","source":"if labels is not None:\n    for star in [[0],[1],[0,1]]:\n        star_cond=feature_df['star'].isin(star)\n        if star_cond.sum()==0:\n            continue\n        mu=feature_df[star_cond].pivot(index='planet_id',columns='wave_len',values='coef2')[wave_df.columns].values\n        sigma=feature_df[star_cond].pivot(index='planet_id',columns='wave_len',values='coef_std2')[wave_df.columns].values\n        target=labels[labels['planet_id'].isin(feature_df[star_cond]['planet_id'])].values[:,1:]\n        score=metric(mu,sigma,target)\n        best_score=metric(mu,np.abs(target-mu),target)\n        print(f'star{star} score/best_score:',round(score,5),'/',round(best_score,5))","metadata":{"execution":{"iopub.status.busy":"2024-10-29T14:22:07.729268Z","iopub.execute_input":"2024-10-29T14:22:07.73007Z","iopub.status.idle":"2024-10-29T14:22:07.739011Z","shell.execute_reply.started":"2024-10-29T14:22:07.729996Z","shell.execute_reply":"2024-10-29T14:22:07.737761Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Submit","metadata":{}},{"cell_type":"code","source":"if mode=='test':\n    mu=feature_df.pivot(index='planet_id',columns='wave_len',values='mu')[wave_df.columns]\n    sigma=feature_df.pivot(index='planet_id',columns='wave_len',values='sigma')[wave_df.columns]\n    submission = pd.read_csv('/kaggle/input/ariel-data-challenge-2024/sample_submission.csv')\n    submission.iloc[:,1:1+283]=mu.loc[submission['planet_id'].values].values\n    submission.iloc[:,1+283:1+283*2]=sigma.loc[submission['planet_id'].values].values\n    submission.to_csv('submission.csv',index=False)\n    display(submission)","metadata":{"execution":{"iopub.status.busy":"2024-10-29T14:22:07.740424Z","iopub.execute_input":"2024-10-29T14:22:07.740846Z","iopub.status.idle":"2024-10-29T14:22:07.814674Z","shell.execute_reply.started":"2024-10-29T14:22:07.740801Z","shell.execute_reply":"2024-10-29T14:22:07.813483Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}