{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":101849,"databundleVersionId":13093295,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":589384,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":396955,"modelId":415389}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import polars as pl\nimport numpy as np\nimport gc\nfrom joblib import Parallel, delayed\nfrom tqdm.notebook import tqdm\nimport itertools\nfrom sklearn.linear_model import LinearRegression, BayesianRidge\nfrom astropy.stats import sigma_clip\nimport jax\nimport jax.numpy as jnp\nimport optax\nimport os","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-24T14:06:59.608384Z","iopub.execute_input":"2025-09-24T14:06:59.609355Z","iopub.status.idle":"2025-09-24T14:06:59.615523Z","shell.execute_reply.started":"2025-09-24T14:06:59.609319Z","shell.execute_reply":"2025-09-24T14:06:59.614446Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"N = 283\nprefix = '/kaggle/input/ariel-data-challenge-2025'\ntrain = pl.read_csv(f'{prefix}/train.csv')[:1]\ntrain_info = pl.read_csv(f'{prefix}/train_star_info.csv')[:1]\ntest_info = pl.read_csv(f'{prefix}/test_star_info.csv')\ntest_info = test_info.with_columns(\n    pl.col('planet_id').cast(pl.Int64)\n)\ninfo = pl.concat((train_info, test_info))\ninfo","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-24T14:06:59.61709Z","iopub.execute_input":"2025-09-24T14:06:59.617508Z","iopub.status.idle":"2025-09-24T14:06:59.681391Z","shell.execute_reply.started":"2025-09-24T14:06:59.617467Z","shell.execute_reply":"2025-09-24T14:06:59.680471Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"adc_info = pl.read_csv(f'{prefix}/adc_info.csv')\naxis_info = pl.read_parquet(f'{prefix}/axis_info.parquet')\nadc = {\n    'airs': (\n        adc_info['AIRS-CH0_adc_gain'].to_numpy().item(), \n        adc_info['AIRS-CH0_adc_offset'].to_numpy().item(),\n        axis_info['AIRS-CH0-integration_time'][:11250].to_numpy(),\n    ),\n    'fgs1': (\n        adc_info['FGS1_adc_gain'].to_numpy().item(), \n        adc_info['FGS1_adc_offset'].to_numpy().item(),\n        np.ones(135000) * 0.1,\n    ),\n}\nadc","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-24T14:06:59.682214Z","iopub.execute_input":"2025-09-24T14:06:59.682458Z","iopub.status.idle":"2025-09-24T14:06:59.700159Z","shell.execute_reply.started":"2025-09-24T14:06:59.682437Z","shell.execute_reply":"2025-09-24T14:06:59.699301Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"idx_test = len(train_info)\ncal_types = ['dark', 'dead', 'flat', 'linear_corr', 'read']\n\nCv = (((4 * 2 + 2) * 2 + 2) * 2 + 2) * 2 + 2\n\ntarget_by_pid = {}\nfor row in train.rows():\n    target_by_pid[row[0]] = row[1:]\n\ninfo = info.with_columns(\n    (pl.col('Rs') ** 2).alias('Rs2'),\n    (pl.col('sma') / pl.col('Rs')).alias('sma/Rs'),\n    (pl.col('Mp') / pl.col('Ms')).alias('Mp/Ms'),\n    (pl.col('Ms') * pl.col('P') ** 2).alias('a3'),\n)\n\ninfo = info.with_columns(\n    (pl.col('a3') ** 0.333).alias('a'),\n)\n\ndef moving_average(x, k):   \n    c = np.cumsum(x)\n    return (c[k:] - c[:-k]) / k\n\ndef get_phases(x):\n    long = len(x) > 20000\n    k = 655 if long else 88\n    x = moving_average(x, k)\n    W = int(5.8 * k)\n    grad = moving_average(moving_average(np.diff(x), W // 2), W // 2)\n    p0 = np.argmin(grad)\n    p1 = np.argmax(grad) + W\n    l, r = np.quantile(grad, [0.25, 0.75])\n    q = (r - l) * 0.5\n    normal = (l - q <= grad) & (grad <= r + q)\n    while p0 >= 1 and not normal[max(0, p0 - W // 2)]: p0 -= 1\n    while p1 < len(x) and not normal[min(max(p1 - W // 2, 0), len(normal) - 1)]: p1 += 1\n    return (p0, p1), x\n\ndef calibrate(signal, cal, adc):\n    gain, offset, dt = adc\n    signal = signal.astype(np.float64) / gain + offset\n    flat = cal['flat'].copy()\n    linear_corr_flipped = np.flip(cal['linear_corr'], axis=0)\n    for x, y in itertools.product(range(signal.shape[1]), range(signal.shape[2])):\n        poly = np.poly1d(linear_corr_flipped[:, x, y])\n        signal[:, x, y] = poly(signal[:, x, y])\n    hot = sigma_clip(cal['dark'], sigma=8, maxiters=5).mask\n    signal -= cal['dark'] * dt[:, None, None]\n    flat[cal['dead'] | hot] = np.nan\n    signal = signal / flat\n    mean = np.nanmean(signal, 0)[None, :, :]\n    std = np.nanstd(signal, 0)[None, :, :]\n    signal[(signal > mean + 8.0 * std) | (signal < mean - 8.0 * std)] = np.nan\n    return signal\n\ndef get_means(i, pid, j):\n    dir = 'train' if i < idx_test else 'test'\n\n    if os.path.exists(f'{prefix}/{dir}/{pid}/x-airs_mean{j}.npy'):\n        ass = []\n        for ai in range(8):\n            ass.append(\n                np.load(f'{prefix}/{dir}/{pid}/x-airs_mean{j}-{ai}.npy')\n            )\n        return (\n            np.load(f'{prefix}/{dir}/{pid}/x-airs_mean{j}.npy'),\n            np.load(f'{prefix}/{dir}/{pid}/x-fgs1_mean{j}.npy'),\n            np.load(f'{prefix}/{dir}/{pid}/x-airs_std{j}.npy'),\n            np.load(f'{prefix}/{dir}/{pid}/x-fgs1_std{j}.npy'),\n            np.load(f'{prefix}/{dir}/{pid}/x-airs_time{j}.npy'),\n            np.load(f'{prefix}/{dir}/{pid}/x-fgs1_time{j}.npy'),\n            ass,\n        )\n\n    if not os.path.exists(f'{prefix}/{dir}/{pid}/AIRS-CH0_signal_{j}.parquet'):\n        return None, None, None, None, None, None, None\n\n    airs = pl.read_parquet(f'{prefix}/{dir}/{pid}/AIRS-CH0_signal_{j}.parquet').to_numpy().reshape(-1, 32, 356)\n    fgs1 = pl.read_parquet(f'{prefix}/{dir}/{pid}/FGS1_signal_{j}.parquet').to_numpy().reshape(-1, 32, 32)\n    airs_calibration = {}\n    fgs1_calibration = {}\n    for t in cal_types:\n        airs_calibration[t] = pl.read_parquet(f'{prefix}/{dir}/{pid}/AIRS-CH0_calibration_{j}/{t}.parquet').to_numpy().reshape(-1, 32, 356)\n        fgs1_calibration[t] = pl.read_parquet(f'{prefix}/{dir}/{pid}/FGS1_calibration_{j}/{t}.parquet').to_numpy().reshape(-1, 32, 32)\n    airs = calibrate(airs, airs_calibration, adc['airs'])\n    fgs1 = calibrate(fgs1, fgs1_calibration, adc['fgs1'])\n    airs_mean = np.nanmean(airs, (1, 2))\n    fgs1_mean = np.nanmean(fgs1, (1, 2))\n    airs_std = np.nanstd(airs, (1, 2))\n    fgs1_std = np.nanstd(fgs1, (1, 2))\n    airs_time = np.nanmean(airs, (0))\n    fgs1_time = np.nanmean(fgs1, (0))\n    a = np.array_split(np.nanmean(airs, 1), 32, axis=1)\n    ass = [np.nanmean(x, 1) for x in a]\n    return airs_mean, fgs1_mean, airs_std, fgs1_std, airs_time, fgs1_time, ass\n\ndef oclip(x):\n    return x\n\ndef af2f(airs_mean, fgs1_mean, airs_time, fgs1_time, ass, i, pid):\n    airs_time = airs_time / np.mean(airs_mean)\n    fgs1_time = fgs1_time / np.mean(fgs1_mean)\n    airs_mean = airs_mean / np.mean(airs_mean)\n    fgs1_mean = fgs1_mean / np.mean(fgs1_mean)\n    (a0, a1), ac = get_phases(airs_mean)\n    (f0, f1), fc = get_phases(fgs1_mean)\n    x1 = np.array(list(map(np.mean, np.array_split(ac, Cv))))\n    x2 = np.array(list(map(np.mean, np.array_split(fc, Cv))))\n    xx = np.stack([x1, x2], 1)\n    xx = np.log(xx)\n    def stats(x, m, g=1, h=1):\n        try:\n            P1 = np.polyfit(np.arange(len(x)) / 1e4, x, 8)\n            P2 = np.polyfit(np.arange(len(x)) / 1e4, x, 2)\n            P3 = np.polyfit(np.arange(len(x)) / 1e4, x, 1)\n        except:\n            P1 = [0] * 9\n            P2 = [0] * 3\n            P3 = [0] * 2\n        k = max(int(len(x) * 0.06), 300)\n        A = x[:k].mean()\n        B = x[-k:].mean()\n        Q = B - A\n        d = np.diff(x)\n        return [\n            x.min() / A,\n            x.min() / B,\n            x.min() / g,\n            x.min() / h,\n            x.max() / A,\n            x.max() / B,\n            x.max() / g,\n            x.max() / h,\n            x.max() - x.min(),\n            x.mean(),\n            x.max(),\n            x.min(),\n            Q,\n            x.std(),\n            (d ** 2).mean(),\n            *P1,\n            P2[0], P3[0],\n        ]\n    def ss(x):\n        if len(x) == 0:\n            x = np.array([1, 1])\n        k = max(int(len(x) * 0.06), 300)\n        A = x[:k].mean()\n        B = x[-k:].mean()\n        return [\n            x.min() / A,\n            x.min() / B,\n            x.max() - x.min(),\n        ]\n\n    res = []\n    for ai in range(32):\n        xy = ass[ai][1::2]\n        xy = xy / xy.mean()\n        (xy0, xy1), xyc = get_phases(xy)\n        xyl, xym, xyr = xyc[:xy0], xyc[xy0:xy1], xyc[xy1:]\n        res += ss(xym)\n\n    al, am, ar = oclip(ac[:a0]), oclip(ac[a0:a1]), oclip(ac[a1:])\n    fl, fm, fr = oclip(fc[:f0]), oclip(fc[f0:f1]), oclip(fc[f1:])\n    if len(al) == 0: al = ar[::-1]\n    if len(fl) == 0: fl = fr[::-1]\n    if len(ar) == 0: ar = al[::-1]\n    if len(fr) == 0: fr = fl[::-1]\n    acm = ac.mean()\n    fcm = fc.mean()\n    if len(al) == 0:\n        al = np.array([1, 1])\n    if len(ar) == 0:\n        ar = np.array([1, 1])\n    if len(am) == 0:\n        am = np.array([1, 1])\n    if len(fm) == 0:\n        fm = np.array([1, 1])\n    if len(fl) == 0:\n        fl = np.array([1, 1])\n    if len(fr) == 0:\n        fr = np.array([1, 1])\n    return [\n        a0, a1 - a0, f0, f1 - f0, len(ac) - a1, len(fc) - f1,\n        *stats(al, acm), *stats(fl, fcm),\n        *stats(am, acm, al.mean(), ar.mean()),\n        *stats(fm, fcm, fl.mean(), fr.mean()),\n        *stats(ar, acm), *stats(fr, fcm),\n        acm, fcm,\n        *stats(ac, acm), *stats(fc, fcm),\n        np.nanstd(airs_time),\n        np.nanstd(fgs1_time),\n    ] + res, xx\n\ndef get_features_partial(i, pid, j):\n    airs_mean, fgs1_mean, airs_std, fgs1_std, airs_time, fgs1_time, ass = get_means(i, pid, j)\n    if airs_mean is None:\n        return []\n    airs_mean = airs_mean[1::2]\n    fgs1_mean = fgs1_mean[1::2]\n    r = [(pid, af2f(airs_mean, fgs1_mean, airs_time, fgs1_time, ass, i, pid))]\n    return r\n\n\ndef get_features(i, pid):\n    try:\n        a = get_features_partial(i, pid, 0)\n        return a\n    except Exception as e:\n        print(i, pid, e, flush=True)\n        return [0] * (136 * 2)\n\nsignal_features = list(get_features(i, pid) for i, pid in enumerate(tqdm(info['planet_id'])))\nsignal_features = [y for x in signal_features for y in x]\ninfo = info.rows()\nX, y = [], []\nX1 = []\nX2 = []\nsf_by_pid = {}\nidx_by_pos = []\nfor pid, b in signal_features:\n    if pid not in sf_by_pid:\n        sf_by_pid[pid] = []\n    sf_by_pid[pid].append(b)\nfor idx, ai in enumerate(tqdm(info)):\n    for b in sf_by_pid[ai[0]]:\n        X.append(list(ai[1:]) + b[0])\n        X1.append(b[1])\n        if ai[0] in target_by_pid:\n            y.append(target_by_pid[ai[0]])\n        idx_by_pos.append(idx)\n\nX = np.array(X)\nX1 = np.array(X1)\ndef smoothie(x, k=1.0):\n    return x / jnp.sqrt(18.0 + x * x / k)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-24T14:06:59.797706Z","iopub.execute_input":"2025-09-24T14:06:59.798049Z","iopub.status.idle":"2025-09-24T14:07:50.04254Z","shell.execute_reply.started":"2025-09-24T14:06:59.798022Z","shell.execute_reply":"2025-09-24T14:07:50.041596Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"scale = [(0.965606510639190673828125, 1.23688781261444091796875, 1.4539191722869873046875),(0.93791639804840087890625, 1.1010532379150390625, 1.25286996364593505859375),(5563.33154296875, 5890.576171875, 6244.93212890625),(0.57544529438018798828125, 0.9295232295989990234375, 1.62257087230682373046875),(0.0, 0.0, 0.0),(3.9515621662139892578125, 5.20556163787841796875, 6.350939273834228515625),(8.65383434295654296875, 10.5174160003662109375, 12.75456905364990234375),(87.6969146728515625, 88.495147705078125, 89.2470550537109375),(0.932395994663238525390625, 1.52989161014556884765625, 2.113880634307861328125),(6.0701808929443359375, 8.67484188079833984375, 12.66816043853759765625),(0.5052020549774169921875, 0.83054029941558837890625, 1.5174677371978759765625),(16.3290004730224609375, 28.9388275146484375, 43.833316802978515625),(2.5346353054046630859375, 3.06671047210693359375, 3.5214443206787109375),(794.25, 1091.0, 1468.75),(2822.25, 3256.5, 3772.0),(10857.75, 14452.0, 18998.25),(31652.25, 36771.0, 42869.75),(783.0, 1118.5, 1444.75),(10745.75, 14696.5, 18662.75),(0.9998569488525390625, 0.999916136264801025390625, 0.99993956089019775390625),(0.99987208843231201171875, 0.999921500682830810546875, 0.999941885471343994140625),(1.00353085994720458984375, 1.0052769184112548828125, 1.00751078128814697265625),(1.00353085994720458984375, 1.0052769184112548828125, 1.00751078128814697265625),(1.0000574588775634765625, 1.000078678131103515625, 1.0001289844512939453125),(1.0000607967376708984375, 1.0000851154327392578125, 1.00014197826385498046875),(1.00378048419952392578125, 1.0055062770843505859375, 1.00778293609619140625),(1.00378048419952392578125, 1.0055062770843505859375, 1.00778293609619140625),(0.0001520453370176255702972412109375, 0.00019523504306562244892120361328125, 0.00029236203408800065517425537109375),(1.003662109375, 1.005384922027587890625, 1.007622241973876953125),(1.00378048419952392578125, 1.0055062770843505859375, 1.00778293609619140625),(1.00353085994720458984375, 1.0052769184112548828125, 1.00751078128814697265625),(-7.3670627898536622524261474609375e-05, -2.4901282813516445457935333251953125e-06, 5.76951788389123976230621337890625e-05),(3.157885294058360159397125244140625e-05, 4.22052617068402469158172607421875e-05, 7.13277186150662600994110107421875e-05),(1.886164661957234756073376047424972057342529296875e-11, 2.28027770277083874361778725869953632354736328125e-11, 2.6701092725733843735724803991615772247314453125e-11),(-6738094.0, 0.0, 8385781.0),(-4531067.0, -20333.66796875, 2977510.0),(-631669.625, 2527.75634765625, 769912.9375),(-65584.046875, -1371.5577392578125, 53558.96875),(-3204.550048828125, 58.921985626220703125, 4141.21142578125),(-130.0842437744140625, -3.59402942657470703125, 105.3831634521484375),(-1.8121950626373291015625, 0.035912610590457916259765625, 2.0251491069793701171875),(-0.013572975061833858489990234375, -0.000139630050398409366607666015625, 0.012972756288945674896240234375),(1.003692150115966796875, 1.0054035186767578125, 1.0076820850372314453125),(-0.0072089261375367641448974609375, -0.0002693767310120165348052978515625, 0.006282416172325611114501953125),(-0.0010437555611133575439453125, -7.36599540687166154384613037109375e-05, 0.000934020732529461383819580078125),(0.99968230724334716796875, 0.999775111675262451171875, 0.99984633922576904296875),(0.99971091747283935546875, 0.999801158905029296875, 0.99986469745635986328125),(1.003561496734619140625, 1.00529384613037109375, 1.00769507884979248046875),(1.003561496734619140625, 1.00529384613037109375, 1.00769507884979248046875),(1.00013101100921630859375, 1.0001842975616455078125, 1.00026881694793701171875),(1.0001418590545654296875, 1.00020754337310791015625, 1.00030791759490966796875),(1.0040562152862548828125, 1.00579345226287841796875, 1.008228302001953125),(1.0040562152862548828125, 1.00579345226287841796875, 1.008228302001953125),(0.0003844773746095597743988037109375, 0.0004480662173591554164886474609375, 0.0005521043203771114349365234375),(1.003802776336669921875, 1.0055387020111083984375, 1.0079586505889892578125),(1.0040562152862548828125, 1.00579345226287841796875, 1.008228302001953125),(1.003561496734619140625, 1.00529384613037109375, 1.00769507884979248046875),(-0.000146729624248109757900238037109375, -2.00644717551767826080322265625e-05, 9.0521178208291530609130859375e-05),(6.8692301283590495586395263671875e-05, 8.1271384260617196559906005859375e-05, 0.00011033753980882465839385986328125),(5.0973857079750217735636397264897823333740234375e-11, 5.59322831106268125722635886631906032562255859375e-11, 5.8831099714051759974609012715518474578857421875e-11),(-0.018680848181247711181640625, -7.2028606155072338879108428955078125e-06, 0.014713163487613201141357421875),(-0.101960711181163787841796875, 0.0002486536395736038684844970703125, 0.10218577086925506591796875),(-0.2777082920074462890625, -0.001591189764440059661865234375, 0.2229290306568145751953125),(-0.2933330833911895751953125, -0.000864111003465950489044189453125, 0.34948194026947021484375),(-0.254680156707763671875, -0.00181365921162068843841552734375, 0.19789414107799530029296875),(-0.088806100189685821533203125, 8.65631955093704164028167724609375e-05, 0.10127760469913482666015625),(-0.02230098657310009002685546875, -0.000823482754640281200408935546875, 0.021423988044261932373046875),(-0.002111823298037052154541015625, 6.0440839888542541302740573883056640625e-08, 0.0020584422163665294647216796875),(1.0038287639617919921875, 1.00561463832855224609375, 1.00801753997802734375),(-7.881209603510797023773193359375e-05, -7.8633529483340680599212646484375e-06, 9.14051124709658324718475341796875e-05),(-0.000117739444249309599399566650390625, -1.598216840648092329502105712890625e-05, 7.16989816282875835895538330078125e-05),(0.9822652339935302734375, 0.988465845584869384765625, 0.992388665676116943359375),(0.98227560520172119140625, 0.98857009410858154296875, 0.992562115192413330078125),(0.98153650760650634765625, 0.98784220218658447265625, 0.991920530796051025390625),(0.981441974639892578125, 0.987992346286773681640625, 0.99203479290008544921875),(1.00056362152099609375, 1.000821590423583984375, 1.001268863677978515625),(1.0006139278411865234375, 1.0009243488311767578125, 1.0013902187347412109375),(0.99999773502349853515625, 1.00008594989776611328125, 1.0005047321319580078125),(0.99989593029022216796875, 1.000138759613037109375, 1.000759601593017578125),(0.00843309052288532257080078125, 0.01262041740119457244873046875, 0.01915666647255420684814453125),(0.9940059185028076171875, 0.99636423587799072265625, 0.9978039264678955078125),(1.0040633678436279296875, 1.0056612491607666015625, 1.00800418853759765625),(0.988826930522918701171875, 0.993462502956390380859375, 0.995959699153900146484375),(-0.00053122336976230144500732421875, -7.8880562796257436275482177734375e-05, 0.0003507588407956063747406005859375),(0.0027683712542057037353515625, 0.00433591939508914947509765625, 0.007059781812131404876708984375),(1.273519345534879221304436214268207550048828125e-10, 2.54227250362504264558083377778530120849609375e-10, 5.56764800929698822073987685143947601318359375e-10),(-625066.875, -221900.90625, -44245.86328125),(71065.640625, 295810.6875, 746151.375),(-375619.3125, -161772.4375, -45903.171875),(16698.4140625, 47577.52734375, 98884.1875),(-15265.7041015625, -7948.43505859375, -3409.74169921875),(366.40362548828125, 749.53375244140625, 1329.2276611328125),(-59.84078216552734375, -36.130413055419921875, -19.677791595458984375),(0.23587892949581146240234375, 0.491549313068389892578125, 0.853663027286529541015625),(1.00274050235748291015625, 1.00391662120819091796875, 1.0059680938720703125),(0.266791403293609619140625, 0.499116718769073486328125, 1.04555034637451171875),(-0.00203187204897403717041015625, -0.00037915210123173892498016357421875, 0.001116312108933925628662109375),(0.980322360992431640625, 0.98712158203125, 0.991786539554595947265625),(0.98023164272308349609375, 0.9870765209197998046875, 0.99173450469970703125),(0.97943556308746337890625, 0.986421167850494384765625, 0.99126565456390380859375),(0.979411184787750244140625, 0.9865436553955078125, 0.99104869365692138671875),(1.00044906139373779296875, 1.000770092010498046875, 1.001287937164306640625),(1.0004808902740478515625, 1.000820636749267578125, 1.001293182373046875),(1.00004565715789794921875, 1.0001580715179443359375, 1.0004370212554931640625),(0.99996507167816162109375, 1.00020694732666015625, 1.000680446624755859375),(0.009241799823939800262451171875, 0.01392530836164951324462890625, 0.02097549848258495330810546875),(0.992947995662689208984375, 0.99565982818603515625, 0.9973328113555908203125),(1.0041744709014892578125, 1.00592577457427978515625, 1.0082619190216064453125),(0.987090587615966796875, 0.9923651218414306640625, 0.9952399730682373046875),(-0.0004111927119083702564239501953125, -2.2379657821147702634334564208984375e-05, 0.000315364333800971508026123046875),(0.0026756585575640201568603515625, 0.004311509430408477783203125, 0.0069082747213542461395263671875),(4.96226705559710268289563828147947788238525390625e-11, 5.4960307138696151696422020904719829559326171875e-11, 5.986229567600531709103961475193500518798828125e-11),(-0.00248640147037804126739501953125, -0.000625512213446199893951416015625, -0.000115352493594400584697723388671875),(0.00202553975395858287811279296875, 0.0092739574611186981201171875, 0.0325946398079395294189453125),(-0.1727159917354583740234375, -0.0580757446587085723876953125, -0.015051476657390594482421875),(0.058254621922969818115234375, 0.1882747113704681396484375, 0.5035064220428466796875),(-0.8650834560394287109375, -0.34466540813446044921875, -0.120868586003780364990234375),(0.13400883972644805908203125, 0.3634996116161346435546875, 0.81183850765228271484375),(-0.3759687244892120361328125, -0.17266690731048583984375, -0.0613427422940731048583984375),(-0.0048559033311903476715087890625, 0.0121289230883121490478515625, 0.0436977557837963104248046875),(1.0038816928863525390625, 1.00532233715057373046875, 1.00754296779632568359375),(0.0021074297837913036346435546875, 0.00392611511051654815673828125, 0.008411760441958904266357421875),(-0.000137654147692956030368804931640625, -8.11960126156918704509735107421875e-06, 0.000102159450761973857879638671875),(0.99961936473846435546875, 0.999880611896514892578125, 0.9999287128448486328125),(0.999657809734344482421875, 0.999848544597625732421875, 0.99992001056671142578125),(1.0033376216888427734375, 1.005006313323974609375, 1.00731003284454345703125),(1.0033376216888427734375, 1.005006313323974609375, 1.00731003284454345703125),(1.00006997585296630859375, 1.000122547149658203125, 1.0003087520599365234375),(1.00008261203765869140625, 1.0001618862152099609375, 1.00037348270416259765625),(1.00394904613494873046875, 1.0056211948394775390625, 1.00798475742340087890625),(1.00394904613494873046875, 1.0056211948394775390625, 1.00798475742340087890625),(0.000224146511754952371120452880859375, 0.00043076343717984855175018310546875, 0.000847735442221164703369140625),(1.00363028049468994140625, 1.0053212642669677734375, 1.00767338275909423828125),(1.00394904613494873046875, 1.0056211948394775390625, 1.00798475742340087890625),(1.0033376216888427734375, 1.005006313323974609375, 1.00731003284454345703125),(-0.0002868114970624446868896484375, -2.735092721195542253553867340087890625e-06, 0.00021726565319113433361053466796875),(5.161978697287850081920623779296875e-05, 0.00011267312220297753810882568359375, 0.0002404688275419175624847412109375),(1.9471542433979749375794199295341968536376953125e-11, 2.3130775073099130167975090444087982177734375e-11, 2.70929910450856681336517794989049434661865234375e-11),(-4733659.0, 38440.3515625, 10270854.0),(-4893724.0, -42167.171875, 2440465.25),(-569813.125, 4109.56201171875, 753943.625),(-69602.9453125, -1815.5389404296875, 45813.2421875),(-2834.110595703125, 149.56011962890625, 4017.40185546875),(-126.8582763671875, -4.37116146087646484375, 96.3449859619140625),(-1.800997257232666015625, 0.106830179691314697265625, 2.0490779876708984375),(-0.014647297561168670654296875, -0.00037040290771983563899993896484375, 0.0123480148613452911376953125),(1.0036566257476806640625, 1.0052936077117919921875, 1.00756740570068359375),(-0.0154985971748828887939453125, -0.00050850003026425838470458984375, 0.0156521461904048919677734375),(-0.0037785801105201244354248046875, -3.385083618923090398311614990234375e-05, 0.003373198211193084716796875),(0.99949395656585693359375, 0.999776661396026611328125, 0.999872982501983642578125),(0.99954116344451904296875, 0.99979937076568603515625, 0.999880850315093994140625),(1.00345432758331298828125, 1.005062580108642578125, 1.00755393505096435546875),(1.00345432758331298828125, 1.005062580108642578125, 1.00755393505096435546875),(1.00012314319610595703125, 1.0002079010009765625, 1.00044155120849609375),(1.0001213550567626953125, 1.000225067138671875, 1.0005109310150146484375),(1.00422191619873046875, 1.005876064300537109375, 1.0083062648773193359375),(1.00422191619873046875, 1.005876064300537109375, 1.0083062648773193359375),(0.00044901599176228046417236328125, 0.00062108063139021396636962890625, 0.00093199126422405242919921875),(1.0038444995880126953125, 1.0054423809051513671875, 1.0078909397125244140625),(1.00422191619873046875, 1.005876064300537109375, 1.0083062648773193359375),(1.00345432758331298828125, 1.005062580108642578125, 1.00755393505096435546875),(-0.0003697703941725194454193115234375, -4.8345909817726351320743560791015625e-06, 0.0003091360558755695819854736328125),(8.3129518316127359867095947265625e-05, 0.00013093560119159519672393798828125, 0.0002285355585627257823944091796875),(4.993751245851996145574958063662052154541015625e-11, 5.6536601789058948952515493147075176239013671875e-11, 5.9386419393181455461672157980501651763916015625e-11),(-0.019305668771266937255859375, 9.55541618168354034423828125e-05, 0.02262304536998271942138671875),(-0.144001483917236328125, -0.00149716413579881191253662109375, 0.10318274796009063720703125),(-0.2309080064296722412109375, 0.0071408241055905818939208984375, 0.34901034832000732421875),(-0.462748050689697265625, -0.0167715884745121002197265625, 0.285645544528961181640625),(-0.2191917002201080322265625, 0.016096055507659912109375, 0.32568359375),(-0.1237230002880096435546875, -0.0057229991070926189422607421875, 0.09419758617877960205078125),(-0.02122240327298641204833984375, 0.00107232877053320407867431640625, 0.02586633525788784027099609375),(-0.00234811683185398578643798828125, -0.0001880693598650395870208740234375, 0.00208151130937039852142333984375),(1.0038735866546630859375, 1.005471706390380859375, 1.0078756809234619140625),(-0.0001534659531898796558380126953125, -2.298575054737739264965057373046875e-05, 0.000111647226731292903423309326171875),(-0.000293638673610985279083251953125, -1.371003236272372305393218994140625e-05, 0.000248984782956540584564208984375),(0.99987924098968505859375, 0.999914944171905517578125, 0.999941647052764892578125),(0.999922454357147216796875, 0.99994623661041259765625, 0.999962329864501953125),(0.98152530193328857421875, 0.987762510776519775390625, 0.99200165271759033203125),(0.981498658657073974609375, 0.98801624774932861328125, 0.99211919307708740234375),(0.988826930522918701171875, 0.993462502956390380859375, 0.995959699153900146484375),(0.988826930522918701171875, 0.993462502956390380859375, 0.995959699153900146484375),(1.00007259845733642578125, 1.00019323825836181640625, 1.00095999240875244140625),(1.00011336803436279296875, 1.000315189361572265625, 1.0010716915130615234375),(1.0043432712554931640625, 1.00604259967803955078125, 1.0083510875701904296875),(1.0043432712554931640625, 1.00604259967803955078125, 1.0083510875701904296875),(0.008660540916025638580322265625, 0.01284508593380451202392578125, 0.019301511347293853759765625),(0.99987924098968505859375, 0.999914944171905517578125, 0.999941647052764892578125),(1.0043432712554931640625, 1.00604259967803955078125, 1.0083510875701904296875),(0.988826930522918701171875, 0.993462502956390380859375, 0.995959699153900146484375),(-0.0009414298110641539096832275390625, -4.432224886841140687465667724609375e-05, 0.0007620626129209995269775390625),(0.0035870256833732128143310546875, 0.0055523957125842571258544921875, 0.008395728655159473419189453125),(8.9151290516564785093578393571078777313232421875e-11, 1.614589850706948936931439675390720367431640625e-10, 3.16646764364492128152051009237766265869140625e-10),(-1064.030517578125, 524.6649169921875, 2940.323486328125),(-6707.80859375, -1250.09130859375, 2345.743896484375),(-2144.9765625, 1227.9920654296875, 6242.3681640625),(-3066.67626953125, -755.8671875, 1074.249755859375),(-305.932861328125, 250.5075531005859375, 840.5770263671875),(-130.8189239501953125, -43.930370330810546875, 43.29341888427734375),(-3.670526027679443359375, 3.594226837158203125, 9.754825592041015625),(-0.2821333110332489013671875, -0.11013717949390411376953125, 0.12174580991268157958984375),(1.00356388092041015625, 1.0057811737060546875, 1.008780002593994140625),(0.1246194541454315185546875, 0.1903819739818572998046875, 0.283118903636932373046875),(-0.00586759857833385467529296875, -5.795103425043635070323944091796875e-05, 0.0056580291129648685455322265625),(0.97941672801971435546875, 0.9864466190338134765625, 0.991291821002960205078125),(0.979636847972869873046875, 0.98654270172119140625, 0.991127073764801025390625),(0.987090587615966796875, 0.9923651218414306640625, 0.9952399730682373046875),(0.987090587615966796875, 0.9923651218414306640625, 0.9952399730682373046875),(1.0001862049102783203125, 1.00030612945556640625, 1.000925540924072265625),(1.0002081394195556640625, 1.00040328502655029296875, 1.0010459423065185546875),(1.0044887065887451171875, 1.00630700588226318359375, 1.0086116790771484375),(1.0044887065887451171875, 1.00630700588226318359375, 1.0086116790771484375),(0.00957783125340938568115234375, 0.01419599913060665130615234375, 0.021289266645908355712890625),(0.999922454357147216796875, 0.99994623661041259765625, 0.999962329864501953125),(1.0044887065887451171875, 1.00630700588226318359375, 1.0086116790771484375),(0.987090587615966796875, 0.9923651218414306640625, 0.9952399730682373046875),(-0.0007695899694226682186126708984375, -4.68101570731960237026214599609375e-05, 0.0006262971437536180019378662109375),(0.0038623264990746974945068359375, 0.0058696432970464229583740234375, 0.00879232771694660186767578125),(5.1296543401857519484110525809228420257568359375e-11, 5.3586225734481018889709957875311374664306640625e-11, 5.6816225335598602441677940078079700469970703125e-11),(-3.01837144434102810919284820556640625e-06, 5.39597522219992242753505706787109375e-07, 5.371306542656384408473968505859375e-06),(-0.0001451428397558629512786865234375, -1.522519232821650803089141845703125e-05, 8.24686649139039218425750732421875e-05),(-0.00090312235988676548004150390625, 0.00020634685643017292022705078125, 0.001586596830748021602630615234375),(-0.009602242149412631988525390625, -0.001571238040924072265625, 0.0054256878793239593505859375),(-0.01714961230754852294921875, 0.00840328074991703033447265625, 0.033091478049755096435546875),(-0.063647449016571044921875, -0.0193626247346401214599609375, 0.0315130390226840972900390625),(-0.030337996780872344970703125, 0.01895163953304290771484375, 0.058653362095355987548828125),(-0.021292440593242645263671875, -0.0077933068387210369110107421875, 0.011271898634731769561767578125),(1.00367832183837890625, 1.00568330287933349609375, 1.00882565975189208984375),(0.0009131464757956564426422119140625, 0.001374476589262485504150390625, 0.00203430838882923126220703125),(-0.000490535399876534938812255859375, 9.2786140157841145992279052734375e-06, 0.00045905049773864448070526123046875),(1.66531598567962646484375, 1.6668131351470947265625, 1.668218135833740234375),(5.1406402587890625, 5.1502933502197265625, 5.1668376922607421875),(0.881124973297119140625, 0.909497737884521484375, 0.95149767398834228515625),(0.87954533100128173828125, 0.908425867557525634765625, 0.949905574321746826171875),(0.094068743288516998291015625, 0.15674114227294921875, 0.20247180759906768798828125),(0.9552438259124755859375, 0.9632720947265625, 0.973178923130035400390625),(0.954933583736419677734375, 0.96353852748870849609375, 0.97314703464508056640625),(0.04763944447040557861328125, 0.060005672276020050048828125, 0.07052020728588104248046875),(0.980545938014984130859375, 0.98533153533935546875, 0.98838055133819580078125),(0.9806425571441650390625, 0.985458195209503173828125, 0.9886319637298583984375),(0.01650502346456050872802734375, 0.019709713757038116455078125, 0.02531256712973117828369140625),(0.98228085041046142578125, 0.98811972141265869140625, 0.991652429103851318359375),(0.982372343540191650390625, 0.98833620548248291015625, 0.991684436798095703125),(0.010925694368779659271240234375, 0.014745347201824188232421875, 0.0208804793655872344970703125),(0.98240816593170166015625, 0.98832988739013671875, 0.992068111896514892578125),(0.982561528682708740234375, 0.988461017608642578125, 0.992286860942840576171875),(0.0097517259418964385986328125, 0.01385040767490863800048828125, 0.01990083418786525726318359375),(0.98246705532073974609375, 0.98836612701416015625, 0.99236404895782470703125),(0.98266923427581787109375, 0.98856842517852783203125, 0.992420613765716552734375),(0.00946566648781299591064453125, 0.013545830734074115753173828125, 0.01960341818630695343017578125),(0.982517540454864501953125, 0.988407909870147705078125, 0.992321312427520751953125),(0.982415616512298583984375, 0.988537609577178955078125, 0.99253571033477783203125),(0.009298401884734630584716796875, 0.01333428360521793365478515625, 0.0195384323596954345703125),(0.98249256610870361328125, 0.9882113933563232421875, 0.992268979549407958984375),(0.9824306964874267578125, 0.98849523067474365234375, 0.99241411685943603515625),(0.00927531160414218902587890625, 0.01343806274235248565673828125, 0.01950648240745067596435546875),(0.982256352901458740234375, 0.98838317394256591796875, 0.9923717975616455078125),(0.98242437839508056640625, 0.988523423671722412109375, 0.99245202541351318359375),(0.0091733038425445556640625, 0.013449822552502155303955078125, 0.01951968669891357421875),(0.98236405849456787109375, 0.98832762241363525390625, 0.992283046245574951171875),(0.982276976108551025390625, 0.988368928432464599609375, 0.99245321750640869140625),(0.009198659099638462066650390625, 0.013385479338467121124267578125, 0.0196067579090595245361328125),(0.982221126556396484375, 0.98836624622344970703125, 0.99229013919830322265625),(0.982289731502532958984375, 0.9884493350982666015625, 0.99245512485504150390625),(0.0091405212879180908203125, 0.01331843249499797821044921875, 0.0196491889655590057373046875),(0.982248246669769287109375, 0.988262176513671875, 0.99221789836883544921875),(0.982291281223297119140625, 0.98830950260162353515625, 0.99245035648345947265625),(0.0091288425028324127197265625, 0.01335029304027557373046875, 0.01960541866719722747802734375),(0.98217284679412841796875, 0.9881803989410400390625, 0.992271482944488525390625),(0.982253015041351318359375, 0.98841559886932373046875, 0.992407619953155517578125),(0.00902984105050563812255859375, 0.01333371363580226898193359375, 0.01962261833250522613525390625),(0.98210966587066650390625, 0.98822557926177978515625, 0.99226856231689453125),(0.98213994503021240234375, 0.988332450389862060546875, 0.9924404621124267578125),(0.0090467371046543121337890625, 0.01336535997688770294189453125, 0.01962383277714252471923828125),(0.982111990451812744140625, 0.988250434398651123046875, 0.992212116718292236328125),(0.98211848735809326171875, 0.98832643032073974609375, 0.99240553379058837890625),(0.00909706763923168182373046875, 0.01333350129425525665283203125, 0.019597113132476806640625),(0.98209559917449951171875, 0.98815643787384033203125, 0.99219858646392822265625),(0.982152521610260009765625, 0.98837363719940185546875, 0.992403805255889892578125),(0.009084253571927547454833984375, 0.01335233636200428009033203125, 0.0196754522621631622314453125),(0.9821488857269287109375, 0.988239109516143798828125, 0.99223387241363525390625),(0.98220121860504150390625, 0.98833835124969482421875, 0.99252021312713623046875),(0.0090438835322856903076171875, 0.013290318660438060760498046875, 0.019496314227581024169921875),(0.982137620449066162109375, 0.9882965087890625, 0.9922997951507568359375),(0.9821207523345947265625, 0.988362491130828857421875, 0.99243927001953125),(0.008992967195808887481689453125, 0.013115148060023784637451171875, 0.01947176270186901092529296875),(0.982116997241973876953125, 0.9882276058197021484375, 0.9923050403594970703125),(0.98224639892578125, 0.988410413265228271484375, 0.99252545833587646484375),(0.008907833136618137359619140625, 0.013175733387470245361328125, 0.0194257684051990509033203125),(0.9820654392242431640625, 0.988229095935821533203125, 0.99234545230865478515625),(0.98216354846954345703125, 0.98838436603546142578125, 0.9924061298370361328125),(0.00887246243655681610107421875, 0.013100658543407917022705078125, 0.01944480836391448974609375),(0.98209178447723388671875, 0.988293170928955078125, 0.992326557636260986328125),(0.98220622539520263671875, 0.98842871189117431640625, 0.99246311187744140625),(0.00886188447475433349609375, 0.01302433572709560394287109375, 0.0193523205816745758056640625),(0.98218405246734619140625, 0.988317906856536865234375, 0.99236333370208740234375),(0.982213914394378662109375, 0.98846375942230224609375, 0.9925000667572021484375),(0.008833455853164196014404296875, 0.0129585675895214080810546875, 0.01940115354955196380615234375),(0.982161104679107666015625, 0.9882552623748779296875, 0.992323696613311767578125),(0.98219406604766845703125, 0.988368570804595947265625, 0.992477953433990478515625),(0.00880087353289127349853515625, 0.013057641685009002685546875, 0.0194358341395854949951171875),(0.98207128047943115234375, 0.988318026065826416015625, 0.9923174381256103515625),(0.98213100433349609375, 0.988400757312774658203125, 0.99255001544952392578125),(0.0088330470025539398193359375, 0.01300776563584804534912109375, 0.0195406638085842132568359375),(0.98213958740234375, 0.988362789154052734375, 0.992331087589263916015625),(0.98217594623565673828125, 0.9884679317474365234375, 0.99247848987579345703125),(0.008773647248744964599609375, 0.0129812844097614288330078125, 0.0194581411778926849365234375),(0.9822032451629638671875, 0.9883797168731689453125, 0.99232399463653564453125),(0.982149779796600341796875, 0.98850917816162109375, 0.9924213886260986328125),(0.0087333656847476959228515625, 0.01295315660536289215087890625, 0.01940130628645420074462890625),(0.98217427730560302734375, 0.98837578296661376953125, 0.992328464984893798828125),(0.982125341892242431640625, 0.98849427700042724609375, 0.9925410747528076171875),(0.008683544583618640899658203125, 0.0129940025508403778076171875, 0.0192582309246063232421875),(0.9820759296417236328125, 0.98830687999725341796875, 0.9922978878021240234375),(0.982209503650665283203125, 0.988489687442779541015625, 0.9925048351287841796875),(0.0087809748947620391845703125, 0.0129632689058780670166015625, 0.01941337622702121734619140625),(0.982011854648590087890625, 0.98824656009674072265625, 0.992315590381622314453125),(0.982123911380767822265625, 0.988315105438232421875, 0.992425739765167236328125),(0.00885523669421672821044921875, 0.01299333013594150543212890625, 0.0195400752127170562744140625),(0.982150495052337646484375, 0.98834550380706787109375, 0.9922769069671630859375),(0.98215007781982421875, 0.9885485172271728515625, 0.9924411773681640625),(0.0091262497007846832275390625, 0.01324565894901752471923828125, 0.01960206590592861175537109375),(0.98209059238433837890625, 0.98824214935302734375, 0.992188394069671630859375),(0.98202311992645263671875, 0.9883811473846435546875, 0.99234282970428466796875),(0.009241458959877490997314453125, 0.01346961222589015960693359375, 0.01984346844255924224853515625),(0.9818947315216064453125, 0.98797738552093505859375, 0.991794168949127197265625),(0.9819591045379638671875, 0.98822534084320068359375, 0.99204623699188232421875),(0.0102880783379077911376953125, 0.01449439488351345062255859375, 0.0206482894718647003173828125),(-0.0047516324557363986968994140625, 0.001572452485561370849609375, 0.005282062105834484100341796875),(-0.0049221892841160297393798828125, 0.001868290477432310581207275390625, 0.0054603121243417263031005859375),]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-24T14:07:50.044165Z","iopub.execute_input":"2025-09-24T14:07:50.044429Z","iopub.status.idle":"2025-09-24T14:07:50.106287Z","shell.execute_reply.started":"2025-09-24T14:07:50.044407Z","shell.execute_reply":"2025-09-24T14:07:50.105269Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"for i in range(X.shape[1]):\n    l, m, r = scale[i]\n    X[:, i] = smoothie((X[:, i] - m) / (r - l + 1e-8), 9.0)\n\nfor i in range(X1.shape[-1]):\n    l, m, r = scale[i + X.shape[1]]\n    X1[:, :, i] = (X1[:, :, i] - m) / (r - l + 1e-8)\n\nX = jnp.array(X)\nX1 = jnp.array(X1)\ny = jnp.array(y)\nX.shape, y.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-24T14:07:50.107238Z","iopub.execute_input":"2025-09-24T14:07:50.107548Z","iopub.status.idle":"2025-09-24T14:07:50.150811Z","shell.execute_reply.started":"2025-09-24T14:07:50.107512Z","shell.execute_reply":"2025-09-24T14:07:50.149929Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"@jax.jit\ndef glob_linear(key: jax.Array, params, x: jnp.ndarray, x1: jnp.ndarray, *, training: bool):\n    key, *subkeys = jax.random.split(key, 1 + params['R'].shape[0])\n    def head(x, x1, R, S, A, W, O, bh, M, bo, mask, c1, c2, key):\n        #\n        conv = jax.vmap(lambda a, b: jnp.convolve(a, b, mode='valid'), (0, None), 0)\n        conv = jax.vmap(conv, (None, -1), -1)\n        conv2 = jax.vmap(conv, (-1, -2), -1)\n        conv = lambda a, b: conv2(a, b).sum(-1)\n\n        C = conv(x1, c1) / c1.shape[0] ** 1.0 / 1.5\n        C = C.reshape(C.shape[0], -1, 2, C.shape[-1]).max(2)\n\n        C = conv(C, c2) / c2.shape[0] ** 1.0 / 4\n        C = C.reshape(C.shape[0], -1, 2, C.shape[-1]).max(2)\n\n        C = conv(C, c2) / c2.shape[0] ** 1.0 / 4\n        C = C.reshape(C.shape[0], -1, 2, C.shape[-1]).max(2)\n\n        C = conv(C, c2) / c2.shape[0] ** 1.0 / 4\n        C = C.reshape(C.shape[0], -1, 2, C.shape[-1]).max(2)\n\n        C = C.max(-2)\n\n        x = jnp.concatenate([x, C], 1)\n        #\n        mask = jax.nn.softmax(mask)\n        mask = mask / (mask ** 2).sum() ** 0.5 * jnp.sqrt(mask.shape[0])\n        x = x * mask\n        A = A * mask[None, :]\n        def fun(r, a):\n            diff = x - a[None, :]\n            return ((diff ** 2) * r).sum(1)\n        # D = jnp.exp(-jax.vmap(fun, (0, 0), 1)(R, A))\n        D = jnp.pow(1.0 + jax.vmap(fun, (0, 0), 1)(R, A) * 0.5 / S, -S)\n        V = (x @ W) / jnp.sqrt(W.shape[0])\n        V = D * (V + bh)\n        V = (V @ O / jnp.sqrt(O.shape[0]) + bo) * M\n        #\n        V = V.reshape(V.shape[0], -1, 2)\n        V = jnp.stack([\n            V[:, :, 0],\n            V[:, :, 1],\n        ], 2)\n        return V\n    return key, jax.vmap(head, (None, None, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0), 0)(\n        x,\n        x1,\n        jax.nn.softplus(params['R']),\n        jax.nn.softplus(params['S']),\n        params['A'],\n        params['W'],\n        params['O'],\n        params['bh'],\n        jax.nn.softplus(params['M']),\n        params['bo'],\n        params['mask'],\n        params['c1'],\n        params['c2'],\n        jnp.array(subkeys),\n    )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-24T14:07:50.152976Z","iopub.execute_input":"2025-09-24T14:07:50.153243Z","iopub.status.idle":"2025-09-24T14:07:50.176148Z","shell.execute_reply.started":"2025-09-24T14:07:50.153219Z","shell.execute_reply":"2025-09-24T14:07:50.175178Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"loaded_npz = np.load(\"/kaggle/input/ariel00/jax/default/18/params.npz\")\nparams = {k: jnp.array(loaded_npz[k]) for k in loaded_npz.files}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-24T14:07:50.177471Z","iopub.execute_input":"2025-09-24T14:07:50.17781Z","iopub.status.idle":"2025-09-24T14:07:50.663622Z","shell.execute_reply.started":"2025-09-24T14:07:50.177784Z","shell.execute_reply":"2025-09-24T14:07:50.662621Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"key = jax.random.PRNGKey(0)\np = glob_linear(key, params, X, X1, training=False)[1]\nB = jnp.exp(p[:, :, :, 1])\nA = p[:, :, :, 0].mean(0)\nB = jnp.log(1.2 * B.mean(0) + 1.2 * ((p[:, :, :, 0] - A) ** 2).mean(0))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-24T14:07:50.664491Z","iopub.execute_input":"2025-09-24T14:07:50.664808Z","iopub.status.idle":"2025-09-24T14:07:51.941428Z","shell.execute_reply.started":"2025-09-24T14:07:50.664776Z","shell.execute_reply":"2025-09-24T14:07:51.93927Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission = pl.DataFrame(test_info['planet_id'])\nfor i in tqdm(range(N)):\n    submission = submission.with_columns(\n        pl.Series(np.array(A[idx_test:, i])).alias(f'wl_{i+1}')\n    )\nfor i in tqdm(range(N)):\n    submission = submission.with_columns(\n        pl.Series(np.array(jnp.sqrt(jnp.exp(B[idx_test:, i])))).alias(f'sigma_{i+1}')\n    )\nsubmission.write_csv('submission.csv')\nsubmission","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-24T14:07:51.942192Z","iopub.execute_input":"2025-09-24T14:07:51.942478Z","iopub.status.idle":"2025-09-24T14:07:52.22233Z","shell.execute_reply.started":"2025-09-24T14:07:51.942452Z","shell.execute_reply":"2025-09-24T14:07:52.221596Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"params","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-09-24T14:07:52.223195Z","iopub.execute_input":"2025-09-24T14:07:52.223542Z","iopub.status.idle":"2025-09-24T14:07:52.247457Z","shell.execute_reply.started":"2025-09-24T14:07:52.223519Z","shell.execute_reply":"2025-09-24T14:07:52.246549Z"}},"outputs":[],"execution_count":null}]}