{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"## Configuration parameters\n# MODE = 'train'\nMODE = 'test'\n\n# USE_POLARS = False\nUSE_POLARS = True\n\n# TRAIN_MAX_EVENTS = 20000\nTRAIN_MAX_EVENTS = None\nTRAIN_BATCH_START = 1\nTRAIN_N_BATCHES = 1\n\n## I pulled in one piece of older code to demonstrate \"before and after\".\n## Set to True (and USE_POLARS=False) if interested in seeing the difference.\nUSE_UNOPTIMIZED = False\n\n\n#### HYPERPARAMETERS ####\n\n## For setting auxiliary = False\n## Hand-tuned and hand-validated, I mostly used batches 100-105, and probably early on also touched batch 1.\n## TODO: revisit the deep core logic now that I've learned about the deep veto layer: https://www.kaggle.com/competitions/icecube-neutrinos-in-deep-ice/discussion/381702\nFIND_BEST_POINTS = True\n\nMIN_PRIMARY_DATAPOINTS = 2\n\nMAX_Z = 3\nMAX_DEEP_Z = 1\nMAX_T = 350\nMAX_DEEP_T = 180\n\nMAX_Z = 4\nMAX_DEEP_Z = 1\nMAX_T = 360\nMAX_DEEP_T = 150\n\n## Ensemble and algorithm selection parameters.\n## First weight is center of charge algorithm introduced in this notebook.\n## Second weight is the unweighted version. Which turns out to be the least-squares algorithm: https://www.kaggle.com/competitions/icecube-neutrinos-in-deep-ice/discussion/381747\nUSE_ENSEMBLE = True\n\nWEIGHTS = [0.5, 0.5]\n\nif not USE_ENSEMBLE and not USE_POLARS:\n    ALGORITHM = 'center of charge'\n#     ALGORITHM = 'least squares'\n    if ALGORITHM == 'least squares':\n        USE_WEIGHTED_LEAST_SQUARES = False\n\n## Constants\nINPUT_DIR = '/kaggle/input/icecube-neutrinos-in-deep-ice'\n\n## Basic configuration override logic\nif MODE == 'test':\n    TRAIN_MAX_EVENTS = None\n    \nif USE_ENSEMBLE:\n    USE_WEIGHTED_LEAST_SQUARES = False","metadata":{"execution":{"iopub.status.busy":"2023-01-29T12:13:52.905412Z","iopub.execute_input":"2023-01-29T12:13:52.90597Z","iopub.status.idle":"2023-01-29T12:13:52.938968Z","shell.execute_reply.started":"2023-01-29T12:13:52.905856Z","shell.execute_reply":"2023-01-29T12:13:52.937517Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Imports and Helper Functions","metadata":{}},{"cell_type":"code","source":"if USE_POLARS:\n    try:\n        import polars as pl\n    except:\n        print('Installing polars, please wait about 35 seconds...')\n        !pip install /kaggle/input/polars01516/polars-0.15.16-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl\n        import polars as pl","metadata":{"execution":{"iopub.status.busy":"2023-01-29T12:13:52.94154Z","iopub.execute_input":"2023-01-29T12:13:52.941978Z","iopub.status.idle":"2023-01-29T12:14:27.924246Z","shell.execute_reply.started":"2023-01-29T12:13:52.94194Z","shell.execute_reply":"2023-01-29T12:14:27.92269Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport math\nimport time\nimport gc\nfrom tqdm import tqdm\ntqdm.pandas()\n\nimport pickle","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-01-29T12:14:27.927994Z","iopub.execute_input":"2023-01-29T12:14:27.928389Z","iopub.status.idle":"2023-01-29T12:14:27.936079Z","shell.execute_reply.started":"2023-01-29T12:14:27.928352Z","shell.execute_reply":"2023-01-29T12:14:27.934812Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"## Condensed for space. See here for expanded original version: https://www.kaggle.com/code/sohier/mean-angular-error\ndef angular_dist_score(az_true, zen_true, az_pred, zen_pred):\n    if not (np.all(np.isfinite(az_true)) and\n            np.all(np.isfinite(zen_true)) and\n            np.all(np.isfinite(az_pred)) and\n            np.all(np.isfinite(zen_pred))):\n        raise ValueError(\"All arguments must be finite\")\n    sa1 = np.sin(az_true)\n    ca1 = np.cos(az_true)\n    sz1 = np.sin(zen_true)\n    cz1 = np.cos(zen_true)\n    sa2 = np.sin(az_pred)\n    ca2 = np.cos(az_pred)\n    sz2 = np.sin(zen_pred)\n    cz2 = np.cos(zen_pred)\n    scalar_prod = sz1*sz2*(ca1*ca2 + sa1*sa2) + (cz1*cz2)\n    scalar_prod =  np.clip(scalar_prod, -1, 1)\n    return np.average(np.abs(np.arccos(scalar_prod)))","metadata":{"execution":{"iopub.status.busy":"2023-01-29T12:14:27.941707Z","iopub.execute_input":"2023-01-29T12:14:27.942212Z","iopub.status.idle":"2023-01-29T12:14:27.953096Z","shell.execute_reply.started":"2023-01-29T12:14:27.942169Z","shell.execute_reply":"2023-01-29T12:14:27.95161Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"## TODO: It would be good to benchmark versus other implementations, like arctan2 used here: https://www.kaggle.com/code/shlomoron/icecube-eda-pca-baseline-cv-1-28-lb-1-274 \n\n## This version has a small optimization trick, calculating azimuth without regard for z or zenith\n## This version is suboptimal if the vectors are already unit vectors, or if you need 3d unit vectors again later for some other step.\ndef angles_from_vectors(vectors):\n    v_squared = np.square(vectors)\n    \n    ## Shortcut optimization for azimuth: calculate 2d unit vectors for x and y independent of z\n    xy_sq = np.sum(v_squared[:, 0:2], axis=1)\n    xy_d = np.sqrt(xy_sq)[:, None]\n    np.seterr(divide='ignore', invalid='ignore') ## Turn off the warning temporarily\n    vectors[:, 0:2] = np.where(xy_d == 0, xy_d, vectors[:, 0:2]/xy_d)\n\n    ## For z, use full 3d unit vector\n    d = np.sqrt(xy_sq + v_squared[:, 2])\n    vectors[:, 2] = np.where(d == 0, d, vectors[:, 2]/d)\n    np.seterr(divide='warn', invalid='warn') ## Turn back on\n\n    ## As mentioned by others, clip solely to avoid floating point errors, the unit vectors should already be within this range.\n    vectors =  np.clip(vectors, -1, 1)\n\n    azimuth = np.arccos(vectors[:, 0])\n    ## if y < 0, convert from quadrants 1 and 2 to quadrants 3 and 4\n    azimuth = np.where(vectors[:, 1] >= 0, azimuth, 2*math.pi - azimuth)\n    azimuth = np.where(np.isfinite(azimuth), azimuth, 0.0)\n\n    zenith = np.arccos(vectors[:, 2])\n    ## IMPORTANT: zenith angles are not evenly distributed, so set the error case to pi/2!\n    ## (even though x, y, z might be. It would be a fun exercise to check if random values\n    ##  for x, y, z converted to zenith angles would match the observed distribution of zenith angles in the train labels)\n    zenith = np.where(np.isfinite(zenith), zenith, math.pi/2)\n\n    return np.stack([azimuth, zenith], axis=1)","metadata":{"execution":{"iopub.status.busy":"2023-01-29T12:14:27.955348Z","iopub.execute_input":"2023-01-29T12:14:27.9559Z","iopub.status.idle":"2023-01-29T12:14:27.972958Z","shell.execute_reply.started":"2023-01-29T12:14:27.955856Z","shell.execute_reply":"2023-01-29T12:14:27.971525Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"## Takes a list of azimuth np arrays, a list of zenith np arrays, and an optional list of numerical weights,\n## and ensembles into a final direction.\n##\n## It's not really optimal in terms of lines of code nor performance,\n## since in most or all cases you are converting a unit vector to an angle,\n## converting back to a unit vector, averaging, then converting to the final angle.\n## However, it is quite convenient, because you can always use this at the end\n## to ensemble the results originating from any number of notebooks or sources.\ndef average_angles(az_list, zen_list, weights=None):\n    assert(len(az_list) == len(zen_list))\n    total = az_list[0].shape[0]\n    x = np.zeros(total)\n    y = np.zeros(total)\n    z = np.zeros(total)\n    for i in range(len(az_list)):\n        w = 1\n        if weights is not None:\n            w = weights[i]\n        az = az_list[i]\n        zen = zen_list[i]\n        assert(az.shape[0] == total)\n        assert(zen.shape[0] == total)\n        if not (np.all(np.isfinite(az)) and\n                np.all(np.isfinite(zen))):\n            raise ValueError(\"All arguments must be finite\")\n        sz = np.sin(zen)\n        x += w*np.cos(az)*sz\n        y += w*np.sin(az)*sz\n        z += w*np.cos(zen)\n    tot_w = len(az_list)\n    if weights is not None:\n        tot_w = sum(weights)\n    x = x / tot_w\n    y = y / tot_w\n    z = z / tot_w\n    d = np.sqrt(np.square(x) + np.square(y) + np.square(z))\n    x = x / d\n    y = y / d\n    z = z / d\n    return angles_from_vectors(np.stack([x, y, z], axis=1))","metadata":{"execution":{"iopub.status.busy":"2023-01-29T12:14:27.974988Z","iopub.execute_input":"2023-01-29T12:14:27.975506Z","iopub.status.idle":"2023-01-29T12:14:27.99135Z","shell.execute_reply.started":"2023-01-29T12:14:27.975441Z","shell.execute_reply":"2023-01-29T12:14:27.990189Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Section 1 - Center of Charge with Pandas\nIn this section, we demonstrate the core algorithm. \n\nIf USE_BEST_POINTS is True, we also use our algorithm in Section 2 to get higher quality data points as input.\n\nLater, in Section 3, we will rewrite this section using Polars for lightning fast execution time!\n\n⚡🧊⚡**WARNING**⚡🧊⚡: When substituting Polars in a Pandas Python environment, there is a distinct possibility of emitting Cherenkov radiation due to code execution that's faster than the speed of light in a Pandas medium!","metadata":{}},{"cell_type":"markdown","source":"### Dual time-weighted Center of Charge in more detail\nThe core concept is if the charge - the light - from the neutrino collision travels omni-directionally and also in the direction of the neutrino, then if we center all 'early' time events versus all 'late' time events, the spherical omni-directional attributes will, on average, cancel out, leaving only the directional attributes we are looking for. We could do a naive technique like first half vs last half, but in reality we want to treat all time gaps as equal relative to the length of the time gap. So we can normalize all timestamps as 0 to 1 (0 is earliest to 1 as latest) or the inverse 1 to 0 (1 is earliest). In this way, we calculate two points, draw a line from earliest to latest to know the direction the neutrino is going, then (**IMPORTANT**) flip it to answer the organizer's stated question of where did the neutrino come from.\n\nUsing 'charge' instead of mass, we can trivially calculate the 'center of charge' of an arbitrary number of points. The time weights just become another weighting besides charge in this very simple algorithm.\n\nIf v_i is each observation's x,y,z coordinates, and t0_i (t1_i) is the calculated time weights, then simply:\n\n``` python3\nv0 = sum(v_i * charge_i * t0_i) / sum(charge_i * t0_i)\nv1 = sum(v_i * charge_i * t1_i) / sum(charge_i * t1_i)\n```\nWhich can also be expressed in the form as simply sum(w * v) / sum(w), for arbitrary weights, which in our case we are using dual time-weights: w0 = charge * t0; w1 = charge * t1.\n\nThen from that, v1 - v0 is the direction vector, with the reverse (v0 - v1) being the direction the neutrino is from.","metadata":{}},{"cell_type":"code","source":"def center_of_charge(batch):\n    ## Groupby -> transform is the key trick to avoiding the dreaded 'for each event' loop,\n    ## and thus getting ~30-60x speed boost improvement!\n    ## If you need any min, max, mean or other simple value from the event group,\n    ## you can precalculate it for each group and broadcast it\n    batch['ev_t_min'] = batch.groupby('event_id')['time'].transform('min')\n    batch['ev_t_max'] = batch.groupby('event_id')['time'].transform('max')\n    \n    ## Now we can just implement our formula! w0 and w1 are the time-weighted charge cases.\n    ## Gather the values we need\n    batch['w1'] = batch.charge * (batch.time - batch.ev_t_min) / (batch.ev_t_max - batch.ev_t_min)\n    batch['w0'] = batch.charge - batch.w1\n    batch['wx0'] = batch.x * batch.w0\n    batch['wy0'] = batch.y * batch.w0\n    batch['wz0'] = batch.z * batch.w0\n    batch['wx1'] = batch.x * batch.w1\n    batch['wy1'] = batch.y * batch.w1\n    batch['wz1'] = batch.z * batch.w1\n    df = batch[['w0', 'w1', 'wx0', 'wy0', 'wz0', 'wx1', 'wy1', 'wz1']]\n\n    ## Calculate all the sums!\n    df = df.groupby('event_id').sum()\n    \n    ## Now do the final divide of the weighted center by the sum of the weights.\n    df[['wx0', 'wy0', 'wz0']] = df[['wx0', 'wy0', 'wz0']].div(df.w0, axis=0)\n    df[['wx1', 'wy1', 'wz1']] = df[['wx1', 'wy1', 'wz1']].div(df.w1, axis=0)\n    \n    ## The direction the neutrino is traveling FROM is point0 - point1, instead of point1 - point0.\n    ## Counter-intuitive to me, but fortunately, easy to notice and correct if your score is > 1.57 instead of less.\n    df[['x', 'y', 'z']] = df[['wx0', 'wy0', 'wz0']].values - df[['wx1', 'wy1', 'wz1']].values\n\n    df = df[['x', 'y', 'z']]\n    df[['azimuth', 'zenith']] = angles_from_vectors(df.values)\n\n    return(df[['azimuth', 'zenith']])","metadata":{"execution":{"iopub.status.busy":"2023-01-29T12:14:27.992922Z","iopub.execute_input":"2023-01-29T12:14:27.994042Z","iopub.status.idle":"2023-01-29T12:14:28.012018Z","shell.execute_reply.started":"2023-01-29T12:14:27.993997Z","shell.execute_reply":"2023-01-29T12:14:28.010569Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Weighted least-squares algorithm\nSolverWorld discusses (unweighted) least squares algorithm [here](https://www.kaggle.com/competitions/icecube-neutrinos-in-deep-ice/discussion/381747) and implements [here](https://www.kaggle.com/code/solverworld/icecube-neutrino-path-least-squares-1-214)\n\n$$v_{est} = \\frac{ \\langle ri * ti \\rangle -  \\langle r \\rangle * \\langle t \\rangle}{\\langle t^2 \\rangle - \\langle t \\rangle ^2}$$\n\nUPDATE: This algorithm uses 4 dimensions of data (x,y,z,time). How do we modify the formula to give arbitrary weights to each point?\n\nInstead of <> meaning take the mean, let's generalize that as <> meaning sum(w_i * _ith_value_inside_<>) / sum(w_i). In the normal case, where all w_i = 1, that simplifies to the mean.\n\nNow we have 5 dimensions using this expanded formula.\n\n","metadata":{}},{"cell_type":"code","source":"def least_squares(batch, weighted=False):\n    batch['xt'] = batch.x * batch.time\n    batch['yt'] = batch.y * batch.time\n    batch['zt'] = batch.z * batch.time\n    batch['tt'] = batch.time * batch.time\n    if weighted:\n        df = batch[['x', 'y', 'z', 'time', 'xt', 'yt', 'zt', 'tt']] * batch.charge.values[:, None]\n        df['charge'] = batch.charge\n        df = df.groupby('event_id').sum()\n        df = df.div(df.charge, axis=0)\n    else:\n        df = batch[['x', 'y', 'z', 'time', 'xt', 'yt', 'zt', 'tt']]\n        df = df.groupby('event_id').mean()\n    df[['x', 'y', 'z']] = (\n                              (df[['xt', 'yt', 'zt']].values - (df[['x', 'y', 'z']].values * df['time'].values[:, None]))\n                            / (df['tt'].values - (df.time.values * df.time.values))[:, None]\n                          )\n    ## Reverse it\n    df = -df[['x', 'y', 'z']]\n    df[['azimuth', 'zenith']] = angles_from_vectors(df.values)\n    return df[['azimuth', 'zenith']]\n","metadata":{"execution":{"iopub.status.busy":"2023-01-29T12:14:28.014282Z","iopub.execute_input":"2023-01-29T12:14:28.014765Z","iopub.status.idle":"2023-01-29T12:14:28.030763Z","shell.execute_reply.started":"2023-01-29T12:14:28.014725Z","shell.execute_reply":"2023-01-29T12:14:28.029071Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def process_batch(batch_id, sensor, max_events=None):\n    print('load batch...')\n    batch = pd.read_parquet(f'{INPUT_DIR}/{MODE}/batch_{batch_id}.parquet')\n\n    ## Limit to max_events\n    if max_events is not None:\n        batch_i = batch.reset_index()\n        event_ids = batch_i.event_id.drop_duplicates()\n        end_index = event_ids.index[max_events]\n        batch = batch_i[:end_index].set_index('event_id')\n    print(batch.shape)\n\n    ## Merge in sensor x,y,z data\n    batch = batch.reset_index().merge(sensor, how='left', on='sensor_id', left_index=False).set_index('event_id')\n\n    ## The logic for auxiliary = False is very basic, we can improve on it. Initial discussion here: \n    if FIND_BEST_POINTS:\n        batch = find_best_points_pandas(batch)\n\n    ## Limit to primary (aux=False) datapoints if there's enough of them.\n    ## For event_ids with too few, set all rows to aux=False. This handles these cases without any for loop logic.\n    ## MIN_PRIMARY_DATAPOINTS is a tuned value, based on optimizing the score on batches 101-110.\n    ## But the result was 'as low as possible'.\n    batch['primary_count'] = batch.groupby('event_id')['auxiliary'].transform('count') - batch.groupby('event_id')['auxiliary'].transform('sum')\n    batch.loc[batch.primary_count < MIN_PRIMARY_DATAPOINTS, 'auxiliary'] = False\n    batch = batch[batch.auxiliary == False]\n    print(batch.shape)\n\n    if USE_ENSEMBLE or ALGORITHM == 'center of charge':\n        df = center_of_charge(batch)\n        if USE_ENSEMBLE:\n            df1 = df\n    if USE_ENSEMBLE or ALGORITHM == 'least squares':\n        df = least_squares(batch, weighted=USE_WEIGHTED_LEAST_SQUARES)\n    if USE_ENSEMBLE:\n        df[['azimuth', 'zenith']] = average_angles([df1.azimuth.values, df.azimuth.values],\n                                                   [df1.zenith.values, df.zenith.values], \n                                                   weights=WEIGHTS)\n    return df","metadata":{"execution":{"iopub.status.busy":"2023-01-29T12:14:28.032354Z","iopub.execute_input":"2023-01-29T12:14:28.03351Z","iopub.status.idle":"2023-01-29T12:14:28.048426Z","shell.execute_reply.started":"2023-01-29T12:14:28.033433Z","shell.execute_reply":"2023-01-29T12:14:28.047005Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Section 2 - Find Best Points\nIn this section, I try to optimize which points are selected to be processed by the center of charge OR any other algorithm. This is just a baseline algorithm.\n\nThe algorithm with current parameters: Check the nearest 3 sensors above and 3 below, for any hits within 350ns of the hit. In current algorithm, it has to be a hit on the neighbor, a second hit on the same sensor within 375ns is not sufficient. For Deep Ice strings, check only the nearest sensor above and below within 180ns.\n\nIn early drafts, I tried calculating pure distance (which can be precomputed for a 5620x5620 sensor matrix). However, the early results showed an optimal distance well under the minimum xy distance separating strings. Which after reading a bunch of discussions and IceCube overviews, did NOT surprise me. So I simplified to simply using vertical neighbors and seeing how many is optimal. And likewise tuning how far apart to set the optimal max time delta. 1000ns is probably too long, if you are mainly using adjacent sensors that logically SHOULD both trigger for the high energy events you are looking for.\n\nOf course, future work can include a hugely stratified event characterization, so that we use different hyperparameters for 'best points' based on factors like total observed charge and/or number of rows, and other factors. Possibly stratified for 'edge' strings vs 'inner' strings. Another to-do that should see big gains is to use fuzzy logic, not binary logic. Convert this to a (sigmoidal?) likelihood algorithm and use the result as weights instead of just 'yes' vs 'no' binary weighting.\n\nDeep ice strings are NOT symmetric compared with the other strings! My 'center of charge' algorithm relies on symmetry. Current solution to that is to use different parameters that only take the highest quality points. It's more of a compromise than a solution to this particular issue. There might be some corner case math or other trick that could do a better job of preventing anti-symmetry bias from reducing the accuracy of the final prediction.","metadata":{}},{"cell_type":"code","source":"def proximity(e, col):\n    ## Since we explode the data to an NxN array, if N is too high it will take a long time and anyways we'll run out of memory.\n    ## TODO: see how high we can go before we run out of memory\n    ## TODO: consider subsampling instead of simply returning the baseline auxiliary setting?\n    ##       And/or splitting on string_id to get much smaller groups\n    if e.shape[0] > 2000:\n        return e[:, col['auxiliary']]\n\n    ## The magic None in the line below is a new thing I learned while working on this notebook.\n    ## It is shorthand for np.newaxis, and broadcasts the array to another dimension.\n    ## This allows us to create an NxN array for each event, so that we can\n    ## check if *any* other row in the event meets our proximity in time and height requirements.\n    ## For any matching pairs, then we set both rows as auxiliary = False. Rows without matches are auxiliary = True.\n    deltas = np.abs(e[:, [col['string_id'], col['depth_id'], col['time']]] - e[:, None, [col['string_id'], col['depth_id'], col['time']]])\n    dz = deltas[:, :, 1]\n\n    ## if same depth or different string id, ignore by setting dz > the max threshold used later.\n    dz[(dz == 0) | (deltas[:, :, 0] != 0)] = MAX_Z + MAX_DEEP_Z + 1\n\n    ## if sensor is not a deep ice sensor, and time > MAX_T, ignore\n    mask = (e[:, col['sensor_id']] < 4680)\n    mask = np.broadcast_to(mask, (mask.shape[0], mask.shape[0])).T\n    dz[mask & (deltas[:, :, 2] > MAX_T)] = MAX_Z + MAX_DEEP_Z + 1\n    ## if sensor IS a deep ice sensor, and time > MAX_DEEP_T, ignore\n    mask = (e[:, col['sensor_id']] >= 4680)\n    mask = np.broadcast_to(mask, (mask.shape[0], mask.shape[0])).T\n    dz[mask & (deltas[:, :, 2] > MAX_DEEP_T)] = MAX_Z + MAX_DEEP_Z + 1\n\n    ## Now take the min (best) result for each row com\n    dz = dz.min(axis=1)\n    ## If no matches, the default, then everything is aux=True\n    e[:, col['auxiliary']] = True\n    ## If not deep ice and distance less than threshold, or deep ice and distance less than other threshold, then we have a match!\n    e[((e[:, col['sensor_id']] < 4680) & (dz <= MAX_Z)) | ((e[:, col['sensor_id']] >= 4680) & (dz <= MAX_DEEP_Z)), col['auxiliary']] = False\n    ## Return only the data needed to speed up the np.concatenate called next.\n    return e[:, col['auxiliary']]","metadata":{"execution":{"iopub.status.busy":"2023-01-29T12:14:28.052145Z","iopub.execute_input":"2023-01-29T12:14:28.05277Z","iopub.status.idle":"2023-01-29T12:14:28.068531Z","shell.execute_reply.started":"2023-01-29T12:14:28.05273Z","shell.execute_reply":"2023-01-29T12:14:28.06721Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"## You can ignore this one unless interested in a deep dive on performance optimization\n## It is provided as a way of comparing the changes versus the pure numpy version.\n## Comments removed from this copy to conserve vertical space.\ndef proximity_unoptimized(df):\n    if df.shape[0] > 2000:\n        return df\n    deltas = np.abs(df[['string_id', 'depth_id', 'time']].values - df[['string_id', 'depth_id', 'time']].values[:, None, :])\n    mask = (df.sensor_id < 4680)\n    dz = deltas[:, :, 1]\n\n    dz[(dz == 0) | (deltas[:, :, 0] != 0)] = MAX_Z + MAX_DEEP_Z + 1\n\n    mask = (df.sensor_id < 4680)\n    mask = np.broadcast_to(mask, (mask.shape[0], mask.shape[0])).T\n    dz[mask & (deltas[:, :, 2] > MAX_T)] = MAX_Z + MAX_DEEP_Z + 1\n    mask = (df.sensor_id >= 4680)\n    mask = np.broadcast_to(mask, (mask.shape[0], mask.shape[0])).T\n    dz[mask & (deltas[:, :, 2] > MAX_DEEP_T)] = MAX_Z + MAX_DEEP_Z + 1\n\n    dz = dz.min(axis=1)\n    df.auxiliary = True\n    df.loc[((df.sensor_id < 4680) & (dz <= MAX_Z)) | ((df.sensor_id >= 4680) & (dz <= MAX_DEEP_Z)), 'auxiliary'] = False\n    return df\n","metadata":{"execution":{"iopub.status.busy":"2023-01-29T12:14:28.07072Z","iopub.execute_input":"2023-01-29T12:14:28.071138Z","iopub.status.idle":"2023-01-29T12:14:28.086754Z","shell.execute_reply.started":"2023-01-29T12:14:28.071102Z","shell.execute_reply":"2023-01-29T12:14:28.085423Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def find_best_points_pandas(batch):\n    if USE_UNOPTIMIZED:\n        cols = ['sensor_id', 'time', 'auxiliary', 'string_id', 'depth_id']\n        batch[cols] = batch[cols].groupby('event_id').progress_apply(proximity_unoptimized)\n        return batch\n\n    ## np.split used as a pure numpy equivalent of groupby\n    ## Note this version didn't minimize the size of the inputs, but does make sure all dtypes are the same for an efficient np array.\n    column_to_index = { k:v for v,k in enumerate(batch.columns)}\n    events = np.split(batch.values.astype('float32'), np.unique(batch.index.values, return_index=True)[1][1:])\n\n    ## Run each event sequentially in a list comprehension, then join back together with np.concatenate.\n    ## So far tried and failed to find a reasonable solution to avoid this groupby > apply > join loop.\n    ## Note that this line overrides the dtype of 'auxiliary' column to float32\n    batch.auxiliary = np.concatenate([proximity(e, column_to_index) for e in tqdm(events)])\n    return batch","metadata":{"execution":{"iopub.status.busy":"2023-01-29T12:14:28.088629Z","iopub.execute_input":"2023-01-29T12:14:28.089117Z","iopub.status.idle":"2023-01-29T12:14:28.099519Z","shell.execute_reply.started":"2023-01-29T12:14:28.089069Z","shell.execute_reply":"2023-01-29T12:14:28.097827Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sensor = pd.read_csv(f'{INPUT_DIR}/sensor_geometry.csv')\nsensor['string_id'] = sensor.sensor_id // 60\n# sensor['depth_id'] = sensor.sensor_id % 60\n\nsensor = sensor.sort_values('z', ascending=False)\nsensor['depth_id'] = pd.cut(sensor['z'], 60, labels=[i for i in range(60)])\nsensor = sensor.sort_values('sensor_id', ascending=True)\n\nsensor.to_csv('sensor_geometry_1.csv', index=False)\n\ndel sensor","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Run Sections 1 and 2 with pandas version","metadata":{}},{"cell_type":"code","source":"if not USE_POLARS:\n    print(MODE)\n    ## Note: this is so slow (until Polars). Never did check how much faster removing the unnecessary columns is.\n    ## I only actually use event_id, batch_id and azimuth, zenith labels.\n    ## It would also be very reasonable to create a dataset with 660 train_meta_{n}.parquet files,\n    ## to avoid the completely unnecessary loading costs of the huge 100M columns.\n    meta = pd.read_parquet(f'{INPUT_DIR}/{MODE}_meta.parquet')\n    print(meta.shape)\n\n#     print('load sensor data...')\n#     sensor = pd.read_csv(f'{INPUT_DIR}/sensor_geometry.csv')\n#     sensor = pd.read_csv('sensor_geometry_1.csv')\n    \n#     sensor['string_id'] = sensor.sensor_id // 60\n# #     sensor['depth_id'] = sensor.sensor_id % 60\n    \n#     sensor = sensor.sort_values('z', ascending=False)\n#     sensor['depth_id'] = pd.cut(sensor['z'], 60, labels=[i for i in range(60)])\n#     sensor = sensor.sort_values('sensor_id', ascending=True)\n    \n#     print(sensor.shape)\n\n    if MODE == 'train':\n        batch_id_start = TRAIN_BATCH_START\n        batch_id_end = batch_id_start + TRAIN_N_BATCHES\n    else:\n        batch_id_start = meta.batch_id.values[0]\n        batch_id_end = meta.batch_id.values[-1] + 1\n\n    sub = []\n    for batch_id in range(batch_id_start,batch_id_end):\n        print(batch_id)\n        t = time.time()\n        df = process_batch(batch_id, sensor, max_events=TRAIN_MAX_EVENTS)\n        print(f'Time: {time.time() - t:0.2f}s')\n        if MODE == 'test':\n            sub.append(df)\n        else:\n            if TRAIN_MAX_EVENTS is not None:\n                print(angular_dist_score(meta[meta.batch_id == batch_id].azimuth.values[:TRAIN_MAX_EVENTS],\n                                         meta[meta.batch_id == batch_id].zenith.values[:TRAIN_MAX_EVENTS],\n                                         df.azimuth.values, df.zenith.values))\n            else:\n                print(angular_dist_score(meta[meta.batch_id == batch_id].azimuth.values, meta[meta.batch_id == batch_id].zenith.values,\n                                         df.azimuth.values, df.zenith.values))\n\n    if MODE == 'test':\n        sub = pd.concat(sub, axis=0)\n        sub.to_csv('submission.csv', index=True)\n        print(sub)\n\n    print('done')","metadata":{"execution":{"iopub.status.busy":"2023-01-29T12:14:28.180055Z","iopub.execute_input":"2023-01-29T12:14:28.180997Z","iopub.status.idle":"2023-01-29T12:14:28.194854Z","shell.execute_reply.started":"2023-01-29T12:14:28.180942Z","shell.execute_reply":"2023-01-29T12:14:28.193369Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Section 3 - Polars\nLet's see what performance we get when using Polars instead of Pandas!","metadata":{}},{"cell_type":"code","source":"def time_weighted_centering(batch, charge_weighted=True):\n    ## Polars equivalent to groupby->transform is called 'over'. We again use this to get min and max without a for loop.\n    batch = batch.with_columns([pl.col('time').min().over('event_id').alias('ev_t_min'),\n                                pl.col('time').max().over('event_id').alias('ev_t_max')])\n    if charge_weighted:\n        batch = batch.with_columns((pl.col('charge') * (pl.col('time') - pl.col('ev_t_min'))\n                                    / (pl.col('ev_t_max') - pl.col('ev_t_min'))).alias('w1'))\n        batch = batch.with_columns((1 + pl.col('charge') - pl.col('w1')).alias('w0'))\n    else:\n        batch = batch.with_columns(((pl.col('time') - pl.col('ev_t_min'))\n                                    / (pl.col('ev_t_max') - pl.col('ev_t_min'))).alias('w1'))\n        batch = batch.with_columns((pl.lit(1) - pl.col('w1')).alias('w0'))\n\n    batch = batch.select(\n        [\n            pl.col('event_id'),\n            pl.col('w0'),\n            pl.col('w1'),\n            (pl.col('x') * pl.col('w0')).alias('wx0'),\n            (pl.col('y') * pl.col('w0')).alias('wy0'),\n            (pl.col('z') * pl.col('w0')).alias('wz0'),\n            (pl.col('x') * pl.col('w1')).alias('wx1'),\n            (pl.col('y') * pl.col('w1')).alias('wy1'),\n            (pl.col('z') * pl.col('w1')).alias('wz1'),\n        ]\n    ).collect().groupby('event_id', maintain_order=True).sum()\n\n    ## The direction the neutrino is traveling FROM is point0 - point1, instead of point1 - point0.\n    ## Counter-intuitive to me, but fortunately, easy to notice and correct if your score is > 1.57 instead of less.\n    batch_values = batch.select(\n        [\n            ((pl.col('wx0') / pl.col('w0')) - (pl.col('wx1') / pl.col('w1'))).alias('x'),\n            ((pl.col('wy0') / pl.col('w0')) - (pl.col('wy1') / pl.col('w1'))).alias('y'),\n            ((pl.col('wz0') / pl.col('w0')) - (pl.col('wz1') / pl.col('w1'))).alias('z'),\n        ]\n    ).to_numpy()\n    return angles_from_vectors(batch_values), batch","metadata":{"execution":{"iopub.status.busy":"2023-01-29T12:14:28.196702Z","iopub.execute_input":"2023-01-29T12:14:28.197595Z","iopub.status.idle":"2023-01-29T12:14:28.215244Z","shell.execute_reply.started":"2023-01-29T12:14:28.197546Z","shell.execute_reply":"2023-01-29T12:14:28.213727Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"## We use the same numpy proximity function, the only difference is we convert from and to a Polars df instead of a Pandas df.\ndef find_best_points_polars(batch):\n    ## Minimize the size of the inputs, and make sure all data types are the same for an efficient np array\n    df = batch.select([pl.col('sensor_id'), pl.col('time'), pl.col('auxiliary'), pl.col('string_id'), pl.col('depth_id')])\n    column_to_index = { k:v for v,k in enumerate(df.columns)}\n    batch_values = df.to_numpy().astype('float32')\n\n    ## Pure numpy equivalent of groupby\n    events = np.split(batch_values, np.unique(batch.select(pl.col('event_id')), return_index=True)[1][1:])\n\n    ## Run each event sequentially in a list comprehension, then join back together with np.concatenate. So far tried and failed to find a reasonable solution to avoid this groupby > apply > join loop.\n    batch = batch.with_columns(pl.Series(np.concatenate([proximity(e, column_to_index) for e in tqdm(events)]).astype('bool')).alias('auxiliary'))\n\n    return batch","metadata":{"execution":{"iopub.status.busy":"2023-01-29T12:14:28.217627Z","iopub.execute_input":"2023-01-29T12:14:28.218217Z","iopub.status.idle":"2023-01-29T12:14:28.231425Z","shell.execute_reply.started":"2023-01-29T12:14:28.218161Z","shell.execute_reply":"2023-01-29T12:14:28.229913Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Prep for running batches","metadata":{}},{"cell_type":"code","source":"if USE_POLARS:\n    print(MODE)\n    ## Scan parquet is part of Polars lazy evaluation, so no cost yet,\n    ## and it can figure out optimizations when we finally 'collect' it later.\n    meta = pl.scan_parquet(f'{INPUT_DIR}/{MODE}_meta.parquet')\n    if MODE == 'train':\n        batch_id_start = TRAIN_BATCH_START\n        batch_id_end = batch_id_start + TRAIN_N_BATCHES\n    else:\n        batch_id_start = meta.select(pl.col('batch_id')).collect()[0, 0]\n        batch_id_end = meta.select(pl.col('batch_id')).collect()[-1, 0] + 1\n\n    print(batch_id_start, batch_id_end)","metadata":{"execution":{"iopub.status.busy":"2023-01-29T12:24:48.014721Z","iopub.execute_input":"2023-01-29T12:24:48.015204Z","iopub.status.idle":"2023-01-29T12:24:48.031482Z","shell.execute_reply.started":"2023-01-29T12:24:48.01517Z","shell.execute_reply":"2023-01-29T12:24:48.030229Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Process each batch","metadata":{}},{"cell_type":"code","source":"print('load sensor data...')    \nsensor_1 = (pl.scan_csv('sensor_geometry_1.csv')\n            .with_columns([\n                pl.col('sensor_id').cast(pl.Int16),\n            ])\n         )\n\nsensor = (pl.scan_csv(f'{INPUT_DIR}/sensor_geometry.csv')\n            .with_columns([\n                pl.col('sensor_id').cast(pl.Int16),\n                (pl.col('sensor_id') // 60).alias('string_id'),\n                (pl.col('sensor_id') % 60).alias('depth_id'),                \n            ])\n         )","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if MODE == 'test':\n    if USE_POLARS:\n        sub = []\n        for batch_id in range(batch_id_start,batch_id_end):\n            print(batch_id)\n            t = time.time()\n            max_events=TRAIN_MAX_EVENTS\n            print('load batch...')\n            batch = pl.scan_parquet(f'{INPUT_DIR}/{MODE}/batch_{batch_id}.parquet')\n\n            ## Limit to max_events\n            if max_events is not None:\n                batch = batch.collect()\n                last_event_id = batch.select(pl.col('event_id')).unique()[TRAIN_MAX_EVENTS-1, 0]\n                batch = batch.lazy().filter(pl.col('event_id') <= last_event_id)\n\n            ## Merge in sensor x,y,z data\n            batch = batch.join(sensor, on='sensor_id', how='left').collect()\n\n            ## The logic for auxiliary = False is very basic, we can improve on it. Initial discussion here: \n            if FIND_BEST_POINTS:\n                batch = find_best_points_polars(batch)\n\n            ## Limit to primary (aux=False) datapoints if there's enough of them.\n            ## For event_ids with too few, set all rows to aux=False. This handles these cases without any for loop logic.\n            ## MIN_PRIMARY_DATAPOINTS is a tuned value, based on optimizing the score on batches 101-110.\n            ## But the result was 'as low as possible'.\n            batch = batch.lazy().with_columns((pl.col('auxiliary').count().over('event_id') - pl.col('auxiliary').sum().over('event_id')).alias('primary_count'))\n            batch = batch.filter((~pl.col('auxiliary')) | (pl.col('primary_count') < MIN_PRIMARY_DATAPOINTS))\n\n            preds, events = time_weighted_centering(batch)\n            if USE_ENSEMBLE:\n                preds1 = preds\n                preds, events = time_weighted_centering(batch, charge_weighted=False)\n                preds = average_angles([preds1[:, 0], preds[:, 0]], [preds1[:, 1], preds[:, 1]], weights=WEIGHTS)\n\n            if MODE == 'test':\n                sub.append(events.select([pl.col('event_id'), pl.Series(preds[:, 0]).alias('azimuth'),\n                                                             pl.Series(preds[:, 1]).alias('zenith')]))\n            else:\n                meta = meta.filter((pl.col('batch_id') >= batch_id_start) & (pl.col('batch_id') < batch_id_end))\n                if isinstance(meta, pl.LazyFrame):\n                    meta = meta.collect()\n                meta_values = meta.filter(pl.col('batch_id') == batch_id).select([pl.col('azimuth'), pl.col('zenith')]).to_numpy()\n                if TRAIN_MAX_EVENTS is not None:\n                    print(angular_dist_score(meta_values[:TRAIN_MAX_EVENTS, 0], meta_values[:TRAIN_MAX_EVENTS, 1], preds[:, 0], preds[:, 1]))\n                else:\n                    print(angular_dist_score(meta_values[:, 0], meta_values[:, 1], preds[:, 0], preds[:, 1]))\n            print(f'Time: {time.time() - t:0.2f}s')\n\n        if MODE == 'test':\n            sub = pl.concat(sub)\n            sub.write_csv('submission.csv')\n            print(sub)\n\n        print('done')","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}