{"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":"markdown","source":"# 20th - Team Ice team Final Solution - Public LB 0.991, Private LB 0.992!","metadata":{}},{"cell_type":"markdown","source":"This notebook is the 20th place final submission for Team Ice team for Kaggle's IceCube - Neutrinos in Deep Ice competition. <b>Team members are junseonglee11 (@junseonglee11), Ayaan Jang(@ayaanjang).</b> This is an ensemble of the LSTM model of 6 (2 different versions).\n\n\n* TFRecord Dataset Notebook: https://www.kaggle.com/code/junseonglee11/icecube-data-to-tfrecord-v2-1\n* Train Notebook: https://www.kaggle.com/code/junseonglee11/20th-tensorflow-tfrecord-tpu-lstm-line-fit-train\n","metadata":{}},{"cell_type":"markdown","source":"# References\n### Robin smith's: notebooks:\n* https://www.kaggle.com/code/rsmits/tensorflow-lstm-model-inference\n* https://www.kaggle.com/code/rsmits/tensorflow-lstm-model-training-tpu\n* https://www.kaggle.com/code/rsmits/tensorflow-lstm-model-data-preprocessor/notebook \nI modified his notebook\n* a. Converted the dataset to TFRecords\n* b. Additional inputs and some preprocessing\n* c. Changed the model (more RNN layers, GRU --> LSTM, GELU activation, rectified adam optimizer)\n\n### Robert Hatch's notebook\n* https://www.kaggle.com/code/roberthatch/lb-1-183-lightning-fast-baseline-with-polars\n* It was crucial to improve our score. Used the results of this notebook as additional inputs in our model.\n\n* |TFRecord dataset generation: https://www.kaggle.com/code/junseonglee11/icecube-data-to-tfrecord-v2-1","metadata":{}},{"cell_type":"markdown","source":"In this inference notebook I show that with the correct data pre-processing, model architecture and training setup an LSTM model can achieve the same performance as the current top performing Graph based models.\n\nThis notebook is based on the notebook originally published with LSTM training and inference: [3 LSTMs; with Data Picking and Shifting](https://www.kaggle.com/code/seungmoklee/3-lstms-with-data-picking-and-shifting). So if you like my notebooks don't forget the work that it is based!\n\nThis Inference notebook is for the largest part the same. It doesn't use the shifting as applied in the original work, it is based on only using 96 pulses and only 6 features.\nIt does ensemble 3 models but they are from the same training run (just different epochs).\n\nCombining models from different model trainings with slightly different hyperparameters will very likely further increase the score. Increasing the number of bins and using more data for training is another way to increase the score of the models. Consider the Inference and Training notebooks a starting point to explore that yourself!\n\nThe training notebook for these models can be found [here](https://www.kaggle.com/code/rsmits/tensorflow-lstm-model-training-tpu). Note that I performed training on my local laptop (32GB RAM / NVidia 3070). For local training I loaded the files for 70 batches into RAM. The training notebook is equal to my local setup with the change that it loads multiple rounds of training data.\n\nIn the training notebook I will further explain what the differences are and how I improved the achieved score.\n\nI hope you enjoy this notebook and if you do please give it an upvote :-)\n\n!! Update in Latest Version: In the comments it was mentioned that using model files from a TPU training caused an error when trying to load the model. I've updated the load_model() with compile = False. This solves the error. Nothing else has changed in the notebook.","metadata":{"papermill":{"duration":0.005969,"end_time":"2023-04-01T07:34:19.174682","exception":false,"start_time":"2023-04-01T07:34:19.168713","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# Import Modules\nimport gc\nimport os\nimport multiprocessing\nimport time\nimport numpy as np\nimport pandas as pd\nimport pyarrow.parquet as pq\nimport tensorflow as tf\nfrom tqdm import tqdm","metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","papermill":{"duration":4.652047,"end_time":"2023-04-01T07:34:23.831899","exception":false,"start_time":"2023-04-01T07:34:19.179852","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-17T13:20:15.232838Z","iopub.execute_input":"2023-04-17T13:20:15.233308Z","iopub.status.idle":"2023-04-17T13:20:20.475848Z","shell.execute_reply.started":"2023-04-17T13:20:15.233231Z","shell.execute_reply":"2023-04-17T13:20:20.474804Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Directories and constants\nhome_dir = \"/kaggle/input/icecube-neutrinos-in-deep-ice/\"\ntest_format = home_dir + 'test/batch_{batch_id:d}.parquet'\n   \n# Model(s)\nmodel_names = [\n     # LB 0.993 (2) LSTM 7 layer\n    \"/kaggle/input/0419-lstm-1/epoch07-val_acc0.15238.h5\",    \n    \"/kaggle/input/0419-lstm-1/epoch01-val_acc0.15396.h5\",\n    \"/kaggle/input/0419-lstm-1/epoch02-val_acc0.15363.h5\",\n    \"/kaggle/input/0419-lstm-1/epoch03-val_acc0.15471.h5\",                        \n    \n    # LB 0.992 (2) LSTM 5 layer\n    \"/kaggle/input/0419-lstm-1/epoch03-val_acc0.15094.h5\",\n    \"/kaggle/input/0419-lstm-1/epoch07-val_acc0.15104.h5\",\n    \"/kaggle/input/0419-lstm-1/epoch08-val_acc0.15099.h5\",     \n    \"/kaggle/input/0419-lstm-1/epoch05-val_acc0.15092.h5\",                     \n    \n]\n\n# Weight earned through CV\nmodel_weights = np.array(\n    [0.05836763, 0.72707093, 0.71994111, 0.78264521, 0.14015217, 0.71328186, 0.47582921, 0.99825985]                                                                                                                                                                                                           \n)\n\n#0 9 features time, pos diff added\n#1 6 features pos --> pos - mean pos\n\ninput_norm_method = [\n    2,    \n    2,\n    2,    \n    2,\n    \n    2,\n    2,\n    2,\n    2,\n]","metadata":{"papermill":{"duration":0.017595,"end_time":"2023-04-01T07:34:23.855394","exception":false,"start_time":"2023-04-01T07:34:23.837799","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-17T13:20:20.481665Z","iopub.execute_input":"2023-04-17T13:20:20.484368Z","iopub.status.idle":"2023-04-17T13:20:20.496090Z","shell.execute_reply.started":"2023-04-17T13:20:20.484328Z","shell.execute_reply":"2023-04-17T13:20:20.494926Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Load Model(s)","metadata":{"papermill":{"duration":0.004622,"end_time":"2023-04-01T07:34:23.864792","exception":false,"start_time":"2023-04-01T07:34:23.86017","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# Load Models\nmodels = []\nfeature_counts = []\n\nfor model_name in model_names:\n    print(f'\\n========== Model File: {model_name}')\n    # Load Model\n    model_path = model_name\n    model = tf.keras.models.load_model(model_path, compile = False)\n    models.append(model)      \n    feature_count = model.inputs[0].shape[2]\n    feature_counts.append(feature_count)\n    # Model summary\n    model.summary()\n    \n# Get Model Parameters\npulse_count = model.inputs[0].shape[1]\n\n\n\n\noutput_bins = model.layers[-1].weights[0].shape[-1]\nbin_num = int(np.sqrt(output_bins))\n\n# Model Parameter Summary\nprint(\"\\n==== Model Parameters\")\nprint(f\"Bin Numbers: {bin_num}\")\nprint(f\"Maximum Pulse Count: {pulse_count}\")\nprint(f\"Features Count: {feature_count}\")","metadata":{"papermill":{"duration":13.962554,"end_time":"2023-04-01T07:34:37.832051","exception":false,"start_time":"2023-04-01T07:34:23.869497","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-17T13:20:20.501108Z","iopub.execute_input":"2023-04-17T13:20:20.504553Z","iopub.status.idle":"2023-04-17T13:20:22.217537Z","shell.execute_reply.started":"2023-04-17T13:20:20.504513Z","shell.execute_reply":"2023-04-17T13:20:22.215067Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Detector Information","metadata":{"papermill":{"duration":0.005398,"end_time":"2023-04-01T07:34:37.843359","exception":false,"start_time":"2023-04-01T07:34:37.837961","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# Load sensor_geometry\nsensor_geometry_df = pd.read_csv(home_dir + \"sensor_geometry.csv\")\n\n# Get Sensor Information\nsensor_x = sensor_geometry_df.x\nsensor_y = sensor_geometry_df.y\nsensor_z = sensor_geometry_df.z\n\n# Detector constants\nc_const = 0.299792458  # speed of light [m/ns]\n\n# Sensor Min / Max Coordinates\nx_min = sensor_x.min()\nx_max = sensor_x.max()\ny_min = sensor_y.min()\ny_max = sensor_y.max()\nz_min = sensor_z.min()\nz_max = sensor_z.max()\n\ndetector_length = np.sqrt((x_max - x_min)**2 + (y_max - y_min)**2 + (z_max - z_min)**2)\nt_valid_length = detector_length / c_const\n\nprint(f\"time valid length: {t_valid_length} ns\")","metadata":{"papermill":{"duration":0.032327,"end_time":"2023-04-01T07:34:37.881092","exception":false,"start_time":"2023-04-01T07:34:37.848765","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-17T13:20:22.218953Z","iopub.status.idle":"2023-04-17T13:20:22.219443Z","shell.execute_reply.started":"2023-04-17T13:20:22.219194Z","shell.execute_reply":"2023-04-17T13:20:22.219217Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Angle encoding edges\n\n- It is efficient to train the model by classification task, initially.\n- azimuth and zenith are independent\n- azimuth distribution is flat and zenith distribution is sinusoidal.\n  - Flat on the spherical surface\n  - $\\phi > \\pi$ events are a little bit rarer than $\\phi < \\pi$ events, (maybe) because of the neutrino attenuation by earth.\n- So, the uniform bin is used for azimuth, and $\\left| \\cos \\right|$ bin is used for zenith","metadata":{"papermill":{"duration":0.005324,"end_time":"2023-04-01T07:34:37.891946","exception":false,"start_time":"2023-04-01T07:34:37.886622","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# Create Azimuth Edges\nazimuth_edges = np.linspace(0, 2 * np.pi, bin_num + 1)\nprint(azimuth_edges)\n\n# Create Zenith Edges\nzenith_edges = []\nzenith_edges.append(0)\nfor bin_idx in range(1, bin_num):\n    zenith_edges.append(np.arccos(np.cos(zenith_edges[-1]) - 2 / (bin_num)))\nzenith_edges.append(np.pi)\nzenith_edges = np.array(zenith_edges)\nprint(zenith_edges)","metadata":{"papermill":{"duration":0.017351,"end_time":"2023-04-01T07:34:37.914765","exception":false,"start_time":"2023-04-01T07:34:37.897414","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-17T13:20:22.221042Z","iopub.status.idle":"2023-04-17T13:20:22.222030Z","shell.execute_reply.started":"2023-04-17T13:20:22.221749Z","shell.execute_reply":"2023-04-17T13:20:22.221782Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Define a function converts from prediction to angles\n\n- Calculation of the mean-vector in a bin $\\theta \\in ( \\theta_0, \\theta_1 )$ and $\\phi \\in ( \\phi_0, \\phi_1 )$\n  - $\\vec{r} \\left( \\theta, ~ \\phi \\right) = \\left< \\sin \\theta \\cos \\phi, ~ \\sin \\theta \\sin \\phi, ~ \\cos \\theta \\right>$\n  - $\\bar{\\vec{r}} = \\frac{ \\int_{\\theta_{0}}^{\\theta_{1}} \\int_{\\phi_0}^{\\phi_1} \\vec{r} \\left( \\theta, ~ \\phi \\right) \\sin \\theta \\,d\\phi \\,d\\theta }{ \\int_{\\theta_{0}}^{\\theta_{1}} \\int_{\\phi_0}^{\\phi_1} 1 \\sin \\theta \\,d\\phi \\,d\\theta }$\n  - $ \\int_{\\theta_{0}}^{\\theta_{1}} \\int_{\\phi_0}^{\\phi_1} 1 \\sin \\theta \\,d\\phi \\,d\\theta = \\left( \\phi_1 - \\phi_0 \\right) \\left( \\cos \\theta_0 - \\cos \\theta_1 \\right)$\n  - $\n\\int_{\\theta_{0}}^{\\theta_{1}} \\int_{\\phi_0}^{\\phi_1} {r}_{x} \\left( \\theta, ~ \\phi \\right) \\sin \\theta \\,d\\phi \\,d\\theta = \n\\int_{\\theta_{0}}^{\\theta_{1}} \\int_{\\phi_0}^{\\phi_1} \\sin^2 \\theta \\cos \\phi \\,d\\phi \\,d\\theta = \n\\left( \\sin \\phi_1 - \\sin \\phi_0 \\right) \\left( \\frac{\\theta_1 - \\theta_0}{2} - \\frac{\\sin 2 \\theta_1 - \\sin 2 \\theta_0}{4} \\right)\n$\n  - $\n\\int_{\\theta_{0}}^{\\theta_{1}} \\int_{\\phi_0}^{\\phi_1} {r}_{y} \\left( \\theta, ~ \\phi \\right) \\sin \\theta \\,d\\phi \\,d\\theta = \n\\int_{\\theta_{0}}^{\\theta_{1}} \\int_{\\phi_0}^{\\phi_1} \\sin^2 \\theta \\sin \\phi \\,d\\phi \\,d\\theta = \n\\left( \\cos \\phi_0 - \\cos \\phi_1 \\right) \\left( \\frac{\\theta_1 - \\theta_0}{2} - \\frac{\\sin 2 \\theta_1 - \\sin 2 \\theta_0}{4} \\right)\n$\n  - $\n\\int_{\\theta_{0}}^{\\theta_{1}} \\int_{\\phi_0}^{\\phi_1} {r}_{z} \\left( \\theta, ~ \\phi \\right) \\sin \\theta \\,d\\phi \\,d\\theta = \n\\int_{\\theta_{0}}^{\\theta_{1}} \\int_{\\phi_0}^{\\phi_1} \\sin \\theta \\cos \\theta \\,d\\phi \\,d\\theta = \n\\left( \\phi_1 - \\phi_0 \\right) \\left( \\frac{\\cos 2 \\theta_0 - \\cos 2 \\theta_1}{4} \\right)\n$","metadata":{"papermill":{"duration":0.00559,"end_time":"2023-04-01T07:34:37.925791","exception":false,"start_time":"2023-04-01T07:34:37.920201","status":"completed"},"tags":[]}},{"cell_type":"code","source":"angle_bin_zenith0 = np.tile(zenith_edges[:-1], bin_num)\nangle_bin_zenith1 = np.tile(zenith_edges[1:], bin_num)\nangle_bin_azimuth0 = np.repeat(azimuth_edges[:-1], bin_num)\nangle_bin_azimuth1 = np.repeat(azimuth_edges[1:], bin_num)\n\nangle_bin_area = (angle_bin_azimuth1 - angle_bin_azimuth0) * (np.cos(angle_bin_zenith0) - np.cos(angle_bin_zenith1))\nangle_bin_vector_sum_x = (np.sin(angle_bin_azimuth1) - np.sin(angle_bin_azimuth0)) * ((angle_bin_zenith1 - angle_bin_zenith0) / 2 - (np.sin(2 * angle_bin_zenith1) - np.sin(2 * angle_bin_zenith0)) / 4)\nangle_bin_vector_sum_y = (np.cos(angle_bin_azimuth0) - np.cos(angle_bin_azimuth1)) * ((angle_bin_zenith1 - angle_bin_zenith0) / 2 - (np.sin(2 * angle_bin_zenith1) - np.sin(2 * angle_bin_zenith0)) / 4)\nangle_bin_vector_sum_z = (angle_bin_azimuth1 - angle_bin_azimuth0) * ((np.cos(2 * angle_bin_zenith0) - np.cos(2 * angle_bin_zenith1)) / 4)\n\nangle_bin_vector_mean_x = angle_bin_vector_sum_x / angle_bin_area\nangle_bin_vector_mean_y = angle_bin_vector_sum_y / angle_bin_area\nangle_bin_vector_mean_z = angle_bin_vector_sum_z / angle_bin_area\n\nangle_bin_vector = np.zeros((1, bin_num * bin_num, 3))\nangle_bin_vector[:, :, 0] = angle_bin_vector_mean_x\nangle_bin_vector[:, :, 1] = angle_bin_vector_mean_y\nangle_bin_vector[:, :, 2] = angle_bin_vector_mean_z\n\nangle_bin_vector_unit = angle_bin_vector[0].copy()\nangle_bin_vector_unit /= np.sqrt((angle_bin_vector_unit**2).sum(axis=1).reshape((-1, 1)))","metadata":{"papermill":{"duration":0.019589,"end_time":"2023-04-01T07:34:37.950879","exception":false,"start_time":"2023-04-01T07:34:37.93129","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-17T13:20:22.223683Z","iopub.status.idle":"2023-04-17T13:20:22.224171Z","shell.execute_reply.started":"2023-04-17T13:20:22.223925Z","shell.execute_reply":"2023-04-17T13:20:22.223947Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def pred_to_angle(_pred, power = 1.35, epsilon = 1e-8):\n    # Convert prediction\n    pred = _pred.copy()\n    pred**= power    \n    \n    \n    pred_vector = (pred.reshape((-1, bin_num**2, 1)) * angle_bin_vector).sum(axis = 1)\n    \n    # Normalize\n    pred_vector_norm = np.sqrt((pred_vector**2).sum(axis = 1))\n    mask = pred_vector_norm < epsilon\n    pred_vector_norm[mask] = 1\n    \n    # Assign <1, 0, 0> to very small vectors (badly predicted)\n    pred_vector /= pred_vector_norm.reshape((-1, 1))\n    pred_vector[mask] = np.array([1., 0., 0.])\n    \n    # Convert to angle\n    azimuth = np.arctan2(pred_vector[:, 1], pred_vector[:, 0])\n    azimuth[azimuth < 0] += 2 * np.pi\n    zenith = np.arccos(pred_vector[:, 2])\n    \n    # Mask bad norm predictions as 0, 0\n    azimuth[mask] = 0.\n    zenith[mask] = 0.\n    \n    return azimuth, zenith","metadata":{"papermill":{"duration":0.016132,"end_time":"2023-04-01T07:34:37.972494","exception":false,"start_time":"2023-04-01T07:34:37.956362","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-17T13:20:22.226403Z","iopub.status.idle":"2023-04-17T13:20:22.227464Z","shell.execute_reply.started":"2023-04-17T13:20:22.227203Z","shell.execute_reply":"2023-04-17T13:20:22.227228Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Weighted-Vector Ensemble","metadata":{"papermill":{"duration":0.005411,"end_time":"2023-04-01T07:34:37.983522","exception":false,"start_time":"2023-04-01T07:34:37.978111","status":"completed"},"tags":[]}},{"cell_type":"code","source":"def weighted_vector_ensemble(angles, weight):\n    # Convert angle to vector\n    vec_models = list()\n    for angle in angles:\n        az, zen = angle\n        sa = np.sin(az)\n        ca = np.cos(az)\n        sz = np.sin(zen)\n        cz = np.cos(zen)\n        vec = np.stack([sz * ca, sz * sa, cz], axis=1)\n        vec_models.append(vec)\n    vec_models = np.array(vec_models)\n\n    # Weighted-mean\n    vec_mean = (weight.reshape((-1, 1, 1)) * vec_models).sum(axis=0) / weight.sum()\n    vec_mean /= np.sqrt((vec_mean**2).sum(axis=1)).reshape((-1, 1))\n\n    # Convert vector to angle\n    zenith = np.arccos(vec_mean[:, 2])\n    azimuth = np.arctan2(vec_mean[:, 1], vec_mean[:, 0])\n    azimuth[azimuth < 0] += 2 * np.pi\n    \n    return azimuth, zenith","metadata":{"papermill":{"duration":0.01725,"end_time":"2023-04-01T07:34:38.006221","exception":false,"start_time":"2023-04-01T07:34:37.988971","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-17T13:20:22.228712Z","iopub.status.idle":"2023-04-17T13:20:22.229562Z","shell.execute_reply.started":"2023-04-17T13:20:22.229310Z","shell.execute_reply":"2023-04-17T13:20:22.229333Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Single event reader function\n\n- Pick-up important data points first\n    - Rank 3 (First)\n        - not aux, in valid time window\n    - Rank 2\n        - not aux, out of valid time window\n    - Rank 1\n        - aux, in valid time window\n    - Rank 0 (Last)\n        - aux, out of valid time window\n    - In each ranks, take pulses from highest charge","metadata":{"papermill":{"duration":0.005282,"end_time":"2023-04-01T07:34:38.017208","exception":false,"start_time":"2023-04-01T07:34:38.011926","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# Placeholder\nopen_batch_dict = dict()\n\n# Read single event from batch_meta_df\ndef read_event(event_idx, batch_meta_df, pulse_count):\n    # Read metadata\n    batch_id, first_pulse_index, last_pulse_index = batch_meta_df.iloc[event_idx][[\"batch_id\", \"first_pulse_index\", \"last_pulse_index\"]].astype(\"int\")\n\n    # close past batch df\n    if batch_id - 1 in open_batch_dict.keys():\n        del open_batch_dict[batch_id - 1]\n\n    # open current batch df\n    if batch_id not in open_batch_dict.keys():\n        open_batch_dict.update({batch_id: pd.read_parquet(test_format.format(batch_id=batch_id))})\n    \n    batch_df = open_batch_dict[batch_id]\n    \n    # Read event\n    event_feature = batch_df[first_pulse_index:last_pulse_index + 1]\n    sensor_id = event_feature.sensor_id\n    \n    # Merge features into single structured array\n    dtype = [(\"time\", \"float16\"),\n             (\"charge\", \"float16\"),\n             (\"auxiliary\", \"float16\"),\n             (\"x\", \"float16\"),\n             (\"y\", \"float16\"),\n             (\"z\", \"float16\"),\n             (\"rank\", \"short\")]    \n    \n    # Create event_x\n    event_x = np.zeros(last_pulse_index - first_pulse_index + 1, dtype)\n    event_x[\"time\"] = event_feature.time.values - event_feature.time.min()\n    event_x[\"charge\"] = event_feature.charge.values\n    event_x[\"auxiliary\"] = event_feature.auxiliary.values\n    event_x[\"x\"] = sensor_geometry_df.x[sensor_id].values\n    event_x[\"y\"] = sensor_geometry_df.y[sensor_id].values\n    event_x[\"z\"] = sensor_geometry_df.z[sensor_id].values\n\n    # For long event, pick-up\n    if len(event_x) > pulse_count:\n        # Find valid time window\n        t_peak = event_x[\"time\"][event_x[\"charge\"].argmax()]\n        t_valid_min = t_peak - t_valid_length\n        t_valid_max = t_peak + t_valid_length\n        t_valid = (event_x[\"time\"] > t_valid_min) * (event_x[\"time\"] < t_valid_max)\n\n        # Rank\n        event_x[\"rank\"] = 2 * (1 - event_x[\"auxiliary\"]) + (t_valid)\n\n        # Sort by Rank and Charge (important goes to backward)\n        event_x = np.sort(event_x, order = [\"rank\", \"charge\"])\n\n        # pick-up from backward\n        event_x = event_x[-pulse_count:]\n\n        # Sort events by time \n        event_x = np.sort(event_x, order = \"time\")\n\n    return event_idx, len(event_x), event_x","metadata":{"papermill":{"duration":0.019846,"end_time":"2023-04-01T07:34:38.042527","exception":false,"start_time":"2023-04-01T07:34:38.022681","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-17T13:20:22.231290Z","iopub.status.idle":"2023-04-17T13:20:22.231947Z","shell.execute_reply.started":"2023-04-17T13:20:22.231679Z","shell.execute_reply":"2023-04-17T13:20:22.231704Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Test metadata","metadata":{"papermill":{"duration":0.005461,"end_time":"2023-04-01T07:34:38.053407","exception":false,"start_time":"2023-04-01T07:34:38.047946","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# Read Test Meta data\ntest_meta_df = pq.read_table(home_dir + 'test_meta.parquet').to_pandas()\nbatch_counts = test_meta_df.batch_id.value_counts().sort_index()\n\nbatch_max_index = batch_counts.cumsum()\nbatch_max_index[test_meta_df.batch_id.min() - 1] = 0\nbatch_max_index = batch_max_index.sort_index()\n\n# Support Function\ndef test_meta_df_spliter(batch_id):\n    return test_meta_df.loc[batch_max_index[batch_id - 1]:batch_max_index[batch_id] - 1]","metadata":{"papermill":{"duration":0.10286,"end_time":"2023-04-01T07:34:38.161687","exception":false,"start_time":"2023-04-01T07:34:38.058827","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-17T13:20:22.233937Z","iopub.status.idle":"2023-04-17T13:20:22.234416Z","shell.execute_reply.started":"2023-04-17T13:20:22.234171Z","shell.execute_reply":"2023-04-17T13:20:22.234194Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Read test data and predict batchwise","metadata":{"papermill":{"duration":0.005226,"end_time":"2023-04-01T07:34:38.172384","exception":false,"start_time":"2023-04-01T07:34:38.167158","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"# ⚡🧊⚡[LB 1.183] Polar Lightning\nI took this from the invaluable works of roberthatch  \nhttps://www.kaggle.com/code/roberthatch/lb-1-183-lightning-fast-baseline-with-polars/comments","metadata":{"execution":{"iopub.status.busy":"2023-04-14T23:06:02.159919Z","iopub.execute_input":"2023-04-14T23:06:02.160451Z","iopub.status.idle":"2023-04-14T23:06:02.167573Z","shell.execute_reply.started":"2023-04-14T23:06:02.160416Z","shell.execute_reply":"2023-04-14T23:06:02.166069Z"}}},{"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\nMIN_PRIMARY_DATAPOINTS = 2\nMAX_Z = 3\nMAX_DEEP_Z = 1\nMAX_T = 350\nMAX_DEEP_T = 180\nif USE_POLARS:\n    ## Only implmented in polars version.\n    AUX_FALSE_WEIGHT = 0.01\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\nWEIGHTS = [0.58, 0.42]\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\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\nif USE_ENSEMBLE:\n    USE_WEIGHTED_LEAST_SQUARES = False\n    \nif 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\n        \nimport numpy as np\nimport pandas as pd\nimport math\nimport time\nimport gc\nfrom tqdm.notebook import tqdm\ntqdm.pandas()\n\n## 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)))\n\n## 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)\n\n## 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))\n\ndef 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']])\n\ndef 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\n\ndef 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\n\ndef 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']]\n\n## 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\n\ndef 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\n\ndef 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((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\n\n\n## 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    ## Rename 'auxiliary' -> 'best', and flip the boolean value\n    batch = batch.with_columns(pl.Series(np.concatenate([proximity(e, column_to_index) for e in tqdm(events)]).astype('bool')).alias('best')).with_columns(pl.col('best').is_not())\n\n    return batch\n\n\n\nif USE_POLARS:\n    sub = []\n    #for batch_id in range(batch_id_start,batch_id_end):\ndef get_line_fit_angles(batch_id):\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\n    print('load sensor data...')\n    sensor = (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             )\n\n    print(sensor)\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.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)\n    \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    ## Use data point weights instead of filtering. Still set most points to 0 weight, unless they are the only data points available.\n    ## MIN_PRIMARY_DATAPOINTS and AUX_FALSE_WEIGHT are tuned values, based on optimizing the score on batches 101-110.\n    batch = batch.lazy().with_columns((pl.col('best').count().over('event_id') - pl.col('best').sum().over('event_id')).alias('best_count'))\n    batch = batch.lazy().with_columns((pl.col('auxiliary').count().over('event_id') - pl.col('auxiliary').sum().over('event_id')).alias('non_aux_count'))\n    batch = batch.with_columns(( pl.when( pl.col('best') | ((pl.col('best_count') < MIN_PRIMARY_DATAPOINTS) & (pl.col('non_aux_count') < MIN_PRIMARY_DATAPOINTS)) )\n                                            .then(1.0)\n                                            .otherwise(pl.when(pl.col('auxiliary').is_not())\n                                                         .then(AUX_FALSE_WEIGHT)\n                                                         .otherwise(0.0)\n                                                    ) \n                                      ).alias('trust'))\n    batch = batch.lazy().with_columns((pl.col('charge') * pl.col('trust')).alias('charge'))\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        ## Instead of charge_weighted=False, set charge*aux to just equal aux.\n        batch = batch.lazy().with_columns(pl.col('trust').alias('charge'))\n        preds, events = time_weighted_centering(batch)\n        preds = average_angles([preds1[:, 0], preds[:, 0]], [preds1[:, 1], preds[:, 1]], weights=WEIGHTS)\n\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    '''\n    print(f'Time: {time.time() - t:0.2f}s')\n    return preds\n\n#fitted_angle = get_line_fit_angles(1)\n","metadata":{"execution":{"iopub.status.busy":"2023-04-17T13:20:22.236682Z","iopub.status.idle":"2023-04-17T13:20:22.237178Z","shell.execute_reply.started":"2023-04-17T13:20:22.236924Z","shell.execute_reply":"2023-04-17T13:20:22.236945Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def normalize_data(_data, option = 0, batch_id = 0):\n    data = _data.copy()\n    if(option==0 or option==2):\n        data[:, :, 0] /= 1000   # time\n        data[:, :, 1] /= 300    # charge\n        data[:, :, 3:] /= 600   # space\n\n        #time to diff_time\n        data[:,:-1,0] = data[:,1:,0] - data[:,:-1,0]\n        data[:,-1,0] = 0    \n\n        #pseudo momentum (next time position - current position)\n        pseudo_momentum = data[:, :, 3:6].copy()\n        pseudo_momentum[:,:-1,:] = pseudo_momentum[:,1:,:] - data[:,:-1, 3:6]\n        pseudo_momentum[:,-1,:] = 0\n\n        for i in range(0, 3):\n            pseudo_momentum[:,:-1,i][data[:,:-1,0]<0]=0\n\n        data[:,:-1,0][data[:,:-1,0]<0] = 0\n\n        data[:,:,6:] = pseudo_momentum.copy()\n        del pseudo_momentum\n        gc.collect()\n        #data = np.append(data, pseudo_momentum, axis = 2)\n        if(option==2):\n            fitted_angle = get_line_fit_angles(batch_id)\n            data = {'event_pulses': data, 'fitted_targets': fitted_angle}\n    elif(option ==1):\n        # Normalize - Version 6\n        data[:, :, 0] /= 1000  # time\n        data[:, :, 1] /= 300  # charge\n        data[:, :, 3:] /= 600  # space\n\n        # distance feature\n        data[:, :, 6] = np.sqrt((data[:, :, 3] - data[:, :, 3].mean())**2 + (data[:, :, 4] - data[:, :, 4].mean())**2 + (data[:, :, 5] - data[:, :, 5].mean())**2)\n\n        data[:, :, 3] = data[:, :, 3] - data[:, :, 3].mean()\n        data[:, :, 4] = data[:, :, 4] - data[:, :, 4].mean()\n        data[:, :, 5] = data[:, :, 5] - data[:, :, 5].mean()    \n\n        # 충전과 거리 간의 상호작용 항 추가\n        data[:, :, 6] = data[:, :, 1] * data[:, :, 6]\n    return data","metadata":{"papermill":{"duration":0.019299,"end_time":"2023-04-01T07:34:38.197133","exception":false,"start_time":"2023-04-01T07:34:38.177834","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-17T13:20:22.238979Z","iopub.status.idle":"2023-04-17T13:20:22.239464Z","shell.execute_reply.started":"2023-04-17T13:20:22.239219Z","shell.execute_reply":"2023-04-17T13:20:22.239241Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Get Batch IDs\ntest_batch_ids = test_meta_df.batch_id.unique()\n\n# Submission Placeholders\ntest_event_id = []\ntest_azimuth = []\ntest_zenith = []\n\n# Batch Loop\nfor batch_id in test_batch_ids:\n    # Batch Meta DF\n    batch_meta_df = test_meta_df_spliter(batch_id)\n\n    # Set Pulses\n    test_x = np.zeros((len(batch_meta_df), pulse_count, 9), dtype = \"float16\")    \n    test_x[:, :, 2] = -1    \n\n    # Read Event Data\n    def read_event_local(event_idx):\n        return read_event(event_idx, batch_meta_df, pulse_count)\n    \n    # Multiprocess Events\n    iterator = range(len(batch_meta_df))\n    with multiprocessing.Pool() as pool:\n        for event_idx, pulsecount, event_x in pool.map(read_event_local, iterator):\n            # Features\n            test_x[event_idx, :pulsecount, 0] = event_x[\"time\"]\n            test_x[event_idx, :pulsecount, 1] = event_x[\"charge\"]\n            test_x[event_idx, :pulsecount, 2] = event_x[\"auxiliary\"]\n            test_x[event_idx, :pulsecount, 3] = event_x[\"x\"]\n            test_x[event_idx, :pulsecount, 4] = event_x[\"y\"]\n            test_x[event_idx, :pulsecount, 5] = event_x[\"z\"]\n    \n    del batch_meta_df\n    gc.collect()\n    \n\n    '''\n    test_x[:, :, 0] /= 1000  # time\n    test_x[:, :, 1] /= 300  # charge\n    test_x[:, :, 3:] /= 600  # space\n        \n    #230328 junseonglee11 time to time diff\n    test_x[:,:-1,0] = test_x[:,1:,0] - test_x[:,:-1,0]\n    test_x[:,-1,0] = 0    \n    test_x[:,:-1,0][test_x[:,:-1,0]<0] = 0\n    '''\n    # Predict\n    counter = 0\n    pred_angles = []\n    for model in models:\n        # Normalize\n        feature_count = feature_counts[counter]\n        \n        if(counter==0):\n            test_tmp_x = normalize_data(test_x, input_norm_method[counter], batch_id)                                          \n        elif(counter>0):\n            if(input_norm_method[counter-1]!=input_norm_method[counter]):\n                del test_tmp_x\n                gc.collect()\n                test_tmp_x = normalize_data(test_x, input_norm_method[counter], batch_id)     \n                \n                \n        try:\n            test_tmp_x = test_tmp_x[:,:,:feature_count]\n        except:\n            print('option 2 activated')\n            \n        \n        pred_model = model.predict(test_tmp_x, verbose=0)\n        if(input_norm_method[counter]==0 or input_norm_method[counter]==2):\n            az_model, zen_model = pred_to_angle(pred_model['encoded_angle'])\n        else:\n            az_model, zen_model = pred_to_angle(pred_model)\n        del pred_model\n        gc.collect()\n        pred_angles.append((az_model, zen_model))\n        counter +=1\n    \n    # Get Predicted Azimuth and Zenith\n    pred_azimuth, pred_zenith = weighted_vector_ensemble(pred_angles, model_weights)\n    del pred_angles\n    gc.collect()\n    \n    # Get Event IDs\n    event_ids = test_meta_df.event_id[test_meta_df.batch_id == batch_id].values\n    \n    # Finalize \n    for event_id, azimuth, zenith in zip(event_ids, pred_azimuth, pred_zenith):\n        if np.isfinite(azimuth) and np.isfinite(zenith):\n            test_event_id.append(int(event_id))\n            test_azimuth.append(azimuth)\n            test_zenith.append(zenith)\n        else:\n            test_event_id.append(int(event_id))\n            test_azimuth.append(0.)\n            test_zenith.append(0.)\n            \n    del test_x, test_tmp_x, pred_azimuth, pred_zenith \n    gc.collect()","metadata":{"papermill":{"duration":18.174466,"end_time":"2023-04-01T07:34:56.378588","exception":false,"start_time":"2023-04-01T07:34:38.204122","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-17T13:20:22.241383Z","iopub.status.idle":"2023-04-17T13:20:22.241883Z","shell.execute_reply.started":"2023-04-17T13:20:22.241623Z","shell.execute_reply":"2023-04-17T13:20:22.241645Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Create Submission","metadata":{"papermill":{"duration":0.005424,"end_time":"2023-04-01T07:34:56.389975","exception":false,"start_time":"2023-04-01T07:34:56.384551","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# Create and Save Submission.csv\nsubmission_df = pd.DataFrame({\"event_id\": test_event_id,\n                              \"azimuth\": test_azimuth,\n                              \"zenith\": test_zenith})\nsubmission_df = submission_df.sort_values(by = ['event_id'])\nsubmission_df.to_csv(\"submission.csv\", index = False)","metadata":{"papermill":{"duration":0.019622,"end_time":"2023-04-01T07:34:56.415111","exception":false,"start_time":"2023-04-01T07:34:56.395489","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-17T13:20:22.243781Z","iopub.status.idle":"2023-04-17T13:20:22.244260Z","shell.execute_reply.started":"2023-04-17T13:20:22.244016Z","shell.execute_reply":"2023-04-17T13:20:22.244038Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Summary\nsubmission_df.head()","metadata":{"papermill":{"duration":0.023362,"end_time":"2023-04-01T07:34:56.443982","exception":false,"start_time":"2023-04-01T07:34:56.42062","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-17T13:20:22.246138Z","iopub.status.idle":"2023-04-17T13:20:22.246611Z","shell.execute_reply.started":"2023-04-17T13:20:22.246368Z","shell.execute_reply":"2023-04-17T13:20:22.246390Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"papermill":{"duration":0.005447,"end_time":"2023-04-01T07:34:56.45505","exception":false,"start_time":"2023-04-01T07:34:56.449603","status":"completed"},"tags":[]},"execution_count":null,"outputs":[]}]}