{"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":"model_names = [\n    \"PointPicker_mpc128bin16_LSTM128DENSE64_10epc1e-4_10epc1e-5\",\n    \"PointPicker_mpc128bin16_LSTM160DENSE0_10epc1e-4_10epc1e-5\",\n    \"PPAS_mpc128bin16_LSTM128DENSE64_10epc1e-4\",\n]\nmodel_shifted = [False, False, True]\nweights = [0.39, 0.22, 0.39]","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-03-07T12:57:16.407588Z","iopub.execute_input":"2023-03-07T12:57:16.408522Z","iopub.status.idle":"2023-03-07T12:57:16.414659Z","shell.execute_reply.started":"2023-03-07T12:57:16.408475Z","shell.execute_reply":"2023-03-07T12:57:16.413227Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Data I/O and preprocessing\nimport numpy as np\nimport pandas as pd\nimport pyarrow.parquet as pq\n\n# System\nimport time\nimport os\nimport gc\n\n# Graphic\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport plotly.express as px\nimport plotly.graph_objects as go\n\nfrom tqdm import tqdm\n\n# multiprocessing\nimport multiprocessing\n\n# TENSORFLOW\nimport tensorflow as tf","metadata":{"execution":{"iopub.status.busy":"2023-03-07T12:57:16.428127Z","iopub.execute_input":"2023-03-07T12:57:16.428907Z","iopub.status.idle":"2023-03-07T12:57:16.434636Z","shell.execute_reply.started":"2023-03-07T12:57:16.428856Z","shell.execute_reply":"2023-03-07T12:57:16.433775Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# directory\nhome_dir = \"/kaggle/input/icecube-neutrinos-in-deep-ice/\"\ntrain_format = home_dir + 'train/batch_{batch_id:d}.parquet'\ntest_format = home_dir + 'test/batch_{batch_id:d}.parquet'\n\nmodel_home = \"/kaggle/input/icecubemodels/\"\n\nweights = np.array(weights)","metadata":{"execution":{"iopub.status.busy":"2023-03-07T12:57:16.440875Z","iopub.execute_input":"2023-03-07T12:57:16.441526Z","iopub.status.idle":"2023-03-07T12:57:16.449354Z","shell.execute_reply.started":"2023-03-07T12:57:16.441474Z","shell.execute_reply":"2023-03-07T12:57:16.448442Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"models = list()\nfor model_name in model_names:\n    print(model_name)\n    \n    model_path = model_home + model_name\n    model = tf.keras.models.load_model(model_path)\n    model.summary()\n    \n    models.append(model)","metadata":{"execution":{"iopub.status.busy":"2023-03-07T12:57:16.456887Z","iopub.execute_input":"2023-03-07T12:57:16.457526Z","iopub.status.idle":"2023-03-07T12:57:35.196235Z","shell.execute_reply.started":"2023-03-07T12:57:16.457487Z","shell.execute_reply":"2023-03-07T12:57:35.193918Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"max_pulse_count = model.inputs[0].shape[1]\nn_features = model.inputs[0].shape[2]\noutput_bins = model.layers[-1].weights[0].shape[-1]\n\nbin_num = int(np.sqrt(output_bins))\n\nprint(\"     bin_num      : \", bin_num)\nprint(\"max_pulse_count: \", max_pulse_count)\nprint(\"       n_features       : \", n_features)","metadata":{"execution":{"iopub.status.busy":"2023-03-07T12:57:35.230622Z","iopub.execute_input":"2023-03-07T12:57:35.231111Z","iopub.status.idle":"2023-03-07T12:57:35.239066Z","shell.execute_reply.started":"2023-03-07T12:57:35.231069Z","shell.execute_reply":"2023-03-07T12:57:35.238079Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\n\n# sensor_geometry\nsensor_geometry_df = pd.read_csv(home_dir + \"sensor_geometry.csv\")\n\n# counts\ndoms_per_string = 60\nstring_num = 86\n\n# index\nouter_long_strings = np.concatenate([np.arange(0, 25), np.arange(27, 34), np.arange(37, 44), np.arange(46, 78)])\ninner_long_strings = np.array([25, 26, 34, 35, 36, 44, 45])\ninner_short_strings = np.array([78, 79, 80, 81, 82, 83, 84, 85])\n\n# known specs\nouter_xy_resolution = 125. / 2\ninner_xy_resolution = 70. / 2\nlong_z_resolution = 17. / 2\nshort_z_resolution = 7. / 2\n\n# evaluate error\nsensor_x = sensor_geometry_df.x\nsensor_y = sensor_geometry_df.y\nsensor_z = sensor_geometry_df.z\nsensor_r_err = np.ones(doms_per_string * string_num)\nsensor_z_err = np.ones(doms_per_string * string_num)\n\nfor string_id in outer_long_strings:\n    sensor_r_err[string_id * doms_per_string:(string_id + 1) * doms_per_string] *= outer_xy_resolution\nfor string_id in np.concatenate([inner_long_strings, inner_short_strings]):\n    sensor_r_err[string_id * doms_per_string:(string_id + 1) * doms_per_string] *= inner_xy_resolution\n\nfor string_id in outer_long_strings:\n    sensor_z_err[string_id * doms_per_string:(string_id + 1) * doms_per_string] *= long_z_resolution\nfor string_id in np.concatenate([inner_long_strings, inner_short_strings]):\n    for dom_id in range(doms_per_string):\n        z = sensor_z[string_id * doms_per_string + dom_id]\n        if (z < -156.) or (z > 95.5 and z < 191.5):\n            sensor_z_err[string_id * doms_per_string + dom_id] *= short_z_resolution\n# register\nsensor_geometry_df[\"r_err\"] = sensor_r_err\nsensor_geometry_df[\"z_err\"] = sensor_z_err","metadata":{"execution":{"iopub.status.busy":"2023-03-07T12:57:35.240306Z","iopub.execute_input":"2023-03-07T12:57:35.240753Z","iopub.status.idle":"2023-03-07T12:57:35.275055Z","shell.execute_reply.started":"2023-03-07T12:57:35.240667Z","shell.execute_reply":"2023-03-07T12:57:35.273567Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"theta = np.linspace(0, 2 * np.pi, 361)\n\nfig = plt.figure(figsize=(20, 5))\n\n# 3D plot\nax = fig.add_subplot(141, projection='3d')\n\ns = ax.scatter(sensor_x, sensor_y, sensor_z, s=0.5, c=np.arange(len(sensor_x)), alpha=0.5)\n\nax.set_xlabel(\"X [m]\")\nax.set_ylabel(\"Y [m]\")\nax.set_zlabel(\"Z [m]\")\n\nfig.colorbar(s, ax=ax)\n\n# X-Y plot\nax = fig.add_subplot(142)\n\nfor string_id in outer_long_strings:\n    x = sensor_x[string_id * doms_per_string]\n    y = sensor_y[string_id * doms_per_string]\n    r_err = sensor_r_err[string_id * doms_per_string]\n    scatter_outer_long = ax.scatter(x, y, color=\"blue\", label=\"outer long string\")\n    ax.plot(x + r_err * np.cos(theta), y + r_err * np.sin(theta), color=\"gray\", alpha=0.5)\n    \nfor string_id in inner_long_strings:\n    x = sensor_x[string_id * doms_per_string]\n    y = sensor_y[string_id * doms_per_string]\n    r_err = sensor_r_err[string_id * doms_per_string]\n    scatter_inner_long = ax.scatter(x, y, color=\"orange\", label=\"inner long string\")\n    ax.plot(x + r_err * np.cos(theta), y + r_err * np.sin(theta), color=\"gray\", alpha=0.5)\n    \nfor string_id in inner_short_strings:\n    x = sensor_x[string_id * doms_per_string]\n    y = sensor_y[string_id * doms_per_string]\n    r_err = sensor_r_err[string_id * doms_per_string]\n    scatter_inner_short = ax.scatter(x, y, color=\"red\", label=\"inner short string\")\n    ax.plot(x + r_err * np.cos(theta), y + r_err * np.sin(theta), color=\"gray\", alpha=0.5)\n\nax.set_xlabel(\"X [m]\")\nax.set_ylabel(\"Y [m]\")\nax.legend(handles=[scatter_outer_long, scatter_inner_long, scatter_inner_short])\n\n# X-Z plot\nax = fig.add_subplot(143)\n\nfor string_id in outer_long_strings:\n    x = sensor_x[string_id * doms_per_string:string_id * doms_per_string + doms_per_string]\n    z = sensor_z[string_id * doms_per_string:string_id * doms_per_string + doms_per_string]\n    scatter_outer_long = ax.scatter(x, z, s=0.5, color=\"blue\", label=\"outer long string\")\n    \nfor string_id in inner_long_strings:\n    x = sensor_x[string_id * doms_per_string:string_id * doms_per_string + doms_per_string]\n    z = sensor_z[string_id * doms_per_string:string_id * doms_per_string + doms_per_string]\n    scatter_inner_long = ax.scatter(x, z, s=0.5, color=\"orange\", label=\"inner long string\")\n\nfor string_id in inner_short_strings:\n    x = sensor_x[string_id * doms_per_string:string_id * doms_per_string + doms_per_string]\n    z = sensor_z[string_id * doms_per_string:string_id * doms_per_string + doms_per_string]\n    scatter_inner_short = ax.scatter(x, z, s=0.5, color=\"red\", label=\"inner short string\")\n\nax.set_xlabel(\"X [m]\")\nax.set_ylabel(\"Z [m]\")\nax.legend(handles=[scatter_outer_long, scatter_inner_long, scatter_inner_short])\n\n# X-Z plot zoom\nax = fig.add_subplot(144)\n\nfor string_id in outer_long_strings:\n    x = sensor_x[string_id * doms_per_string:string_id * doms_per_string + doms_per_string]\n    z = sensor_z[string_id * doms_per_string:string_id * doms_per_string + doms_per_string]\n    scatter_outer_long = ax.scatter(x, z, s=0.5, color=\"blue\", label=\"outer long string\")\n    \nfor string_id in inner_long_strings:\n    x = sensor_x[string_id * doms_per_string:string_id * doms_per_string + doms_per_string]\n    z = sensor_z[string_id * doms_per_string:string_id * doms_per_string + doms_per_string]\n    scatter_inner_long = ax.scatter(x, z, s=0.5, color=\"orange\", label=\"inner long string\")\n\nfor string_id in inner_short_strings:\n    x = sensor_x[string_id * doms_per_string:string_id * doms_per_string + doms_per_string]\n    z = sensor_z[string_id * doms_per_string:string_id * doms_per_string + doms_per_string]\n    scatter_inner_short = ax.scatter(x, z, s=0.5, color=\"red\", label=\"inner short string\")\n\nfor sensor_id in range(doms_per_string * string_num):\n    x = sensor_x[sensor_id]\n    z = sensor_z[sensor_id]\n    z_err = sensor_z_err[sensor_id]\n    if (x > -150 and x < 50) and (z > -200 and z < 250):\n        ax.plot(x + z_err * np.cos(theta), z + z_err * np.sin(theta), color=\"gray\", alpha=0.5)\n\nax.set_xlabel(\"X [m]\")\nax.set_ylabel(\"Z [m]\")\nax.set_xlim(-150, 50)\nax.set_ylim(-200, 250)\nax.legend(handles=[scatter_outer_long, scatter_inner_long, scatter_inner_short])\n\nfig.tight_layout()\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2023-03-07T12:57:35.278466Z","iopub.execute_input":"2023-03-07T12:57:35.278882Z","iopub.status.idle":"2023-03-07T12:57:43.701238Z","shell.execute_reply.started":"2023-03-07T12:57:35.278840Z","shell.execute_reply":"2023-03-07T12:57:43.699970Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# detector constants\nc_const = 0.299792458  # speed of light [m/ns]\n\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(\"t_valid_length: \", t_valid_length, \" ns\")","metadata":{"execution":{"iopub.status.busy":"2023-03-07T12:57:43.702935Z","iopub.execute_input":"2023-03-07T12:57:43.703393Z","iopub.status.idle":"2023-03-07T12:57:43.713366Z","shell.execute_reply.started":"2023-03-07T12:57:43.703352Z","shell.execute_reply":"2023-03-07T12:57:43.711985Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"azimuth_edges = np.linspace(0, 2 * np.pi, bin_num + 1)\nazimuth_shift = (azimuth_edges[1] - azimuth_edges[0]) / 2.\n\nzenith_edges_flat = np.linspace(0, np.pi, bin_num + 1)\nzenith_edges = list()\nzenith_edges.append(0)\nfor bin_idx in range(1, bin_num):\n    # cos(zen_before) - cos(zen_now) = 2 / bin_num\n    zen_now = np.arccos(np.cos(zenith_edges[-1]) - 2 / (bin_num))\n    zenith_edges.append(zen_now)\nzenith_edges.append(np.pi)\nzenith_edges = np.array(zenith_edges)","metadata":{"execution":{"iopub.status.busy":"2023-03-07T12:57:43.715287Z","iopub.execute_input":"2023-03-07T12:57:43.715801Z","iopub.status.idle":"2023-03-07T12:57:43.730230Z","shell.execute_reply.started":"2023-03-07T12:57:43.715726Z","shell.execute_reply":"2023-03-07T12:57:43.728723Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def y_to_onehot(batch_y):\n    # evaluate bin code\n    azimuth_code = (batch_y[:, 0] > azimuth_edges[1:].reshape((-1, 1))).sum(axis=0)\n    zenith_code = (batch_y[:, 1] > zenith_edges[1:].reshape((-1, 1))).sum(axis=0)\n    angle_code = bin_num * azimuth_code + zenith_code\n\n    # one-hot\n    batch_y_onehot = np.zeros((angle_code.size, bin_num * bin_num))\n    batch_y_onehot[np.arange(angle_code.size), angle_code] = 1\n    \n    return batch_y_onehot","metadata":{"execution":{"iopub.status.busy":"2023-03-07T12:57:43.731917Z","iopub.execute_input":"2023-03-07T12:57:43.732373Z","iopub.status.idle":"2023-03-07T12:57:43.742911Z","shell.execute_reply.started":"2023-03-07T12:57:43.732337Z","shell.execute_reply":"2023-03-07T12:57:43.741855Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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":{"execution":{"iopub.status.busy":"2023-03-07T12:57:43.744265Z","iopub.execute_input":"2023-03-07T12:57:43.745340Z","iopub.status.idle":"2023-03-07T12:57:43.759987Z","shell.execute_reply.started":"2023-03-07T12:57:43.745298Z","shell.execute_reply":"2023-03-07T12:57:43.758586Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def pred_to_angle(pred, epsilon=1e-8):\n    # convert prediction to vector\n    pred_vector = (pred.reshape((-1, bin_num * bin_num, 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\n\n\ndef pred_to_angle_azshift(pred, epsilon=1e-8):\n    # convert prediction to vector\n    pred_vector = (pred.reshape((-1, bin_num * bin_num, 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    # shift\n    azimuth -= azimuth_shift\n    azimuth[azimuth < 0] += 2 * np.pi\n    \n    # mask bad norm predictions as 0, 0\n    azimuth[mask] = 0.\n    zenith[mask] = 0.\n    \n    return azimuth, zenith","metadata":{"execution":{"iopub.status.busy":"2023-03-07T12:57:43.761602Z","iopub.execute_input":"2023-03-07T12:57:43.762532Z","iopub.status.idle":"2023-03-07T12:57:43.781939Z","shell.execute_reply.started":"2023-03-07T12:57:43.762491Z","shell.execute_reply":"2023-03-07T12:57:43.780191Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def pred_to_angle_argmax(pred):\n    # get the highest score codes\n    pred_code = pred.argmax(axis=1)\n    \n    # get the bin vector\n    pred_vector = angle_bin_vector_unit[pred_code, :]\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    return azimuth, zenith","metadata":{"execution":{"iopub.status.busy":"2023-03-07T12:57:43.787159Z","iopub.execute_input":"2023-03-07T12:57:43.787995Z","iopub.status.idle":"2023-03-07T12:57:43.800105Z","shell.execute_reply.started":"2023-03-07T12:57:43.787949Z","shell.execute_reply":"2023-03-07T12:57:43.798817Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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\n        sa = np.sin(az)\n        ca = np.cos(az)\n        sz = np.sin(zen)\n        cz = np.cos(zen)\n\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":{"execution":{"iopub.status.busy":"2023-03-07T12:57:43.801829Z","iopub.execute_input":"2023-03-07T12:57:43.802525Z","iopub.status.idle":"2023-03-07T12:57:43.818401Z","shell.execute_reply.started":"2023-03-07T12:57:43.802482Z","shell.execute_reply":"2023-03-07T12:57:43.817290Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"open_batch_dict = dict()\n\n\n# read single event from batch_meta_df\ndef read_event(event_idx, batch_meta_df, max_pulse_count, train=True):\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        if train:\n            open_batch_dict.update({batch_id: pd.read_parquet(train_format.format(batch_id=batch_id))})\n        else:\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 = [\n        (\"time\", \"float16\"),\n        (\"charge\", \"float16\"),\n        (\"auxiliary\", \"float16\"),\n        (\"x\", \"float16\"),\n        (\"y\", \"float16\"),\n        (\"z\", \"float16\"),\n        (\"r_err\", \"float16\"),\n        (\"z_err\", \"float16\"),\n        (\"rank\", \"short\"),\n    ]\n    event_x = np.zeros(last_pulse_index - first_pulse_index + 1, dtype)\n\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\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    event_x[\"r_err\"] = sensor_geometry_df.r_err[sensor_id].values\n    event_x[\"z_err\"] = sensor_geometry_df.z_err[sensor_id].values\n    \n    # For long event, pick-up\n    if len(event_x) > max_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\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[-max_pulse_count:]\n\n        # resort by time\n        event_x = np.sort(event_x, order=\"time\")\n\n    # for train data, give angles together\n    if train:\n        azimuth, zenith = batch_meta_df.iloc[event_idx][[\"azimuth\", \"zenith\"]].astype(\"float16\")\n        event_y = np.array([azimuth, zenith], dtype=\"float16\")\n        \n        return event_idx, len(event_x), event_x, event_y\n    \n    # for test data, just give feature \n    else:\n        return event_idx, len(event_x), event_x","metadata":{"execution":{"iopub.status.busy":"2023-03-07T12:57:43.820448Z","iopub.execute_input":"2023-03-07T12:57:43.821000Z","iopub.status.idle":"2023-03-07T12:57:43.839651Z","shell.execute_reply.started":"2023-03-07T12:57:43.820956Z","shell.execute_reply":"2023-03-07T12:57:43.838401Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_meta_df = pq.read_table(home_dir + \"test_meta.parquet\").to_pandas()\ntest_meta_df.head()","metadata":{"execution":{"iopub.status.busy":"2023-03-07T12:57:43.841332Z","iopub.execute_input":"2023-03-07T12:57:43.841751Z","iopub.status.idle":"2023-03-07T12:57:43.873047Z","shell.execute_reply.started":"2023-03-07T12:57:43.841684Z","shell.execute_reply":"2023-03-07T12:57:43.871735Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"batch_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\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":{"execution":{"iopub.status.busy":"2023-03-07T12:57:43.874557Z","iopub.execute_input":"2023-03-07T12:57:43.874939Z","iopub.status.idle":"2023-03-07T12:57:43.885642Z","shell.execute_reply.started":"2023-03-07T12:57:43.874901Z","shell.execute_reply":"2023-03-07T12:57:43.884307Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_batch_ids = test_meta_df.batch_id.unique()\n\ntest_event_id = list()\ntest_azimuth = list()\ntest_zenith = list()\n\nfor batch_id in test_batch_ids:\n    print(batch_id)\n    # READ ONE BATCH OF TEST DATA\n    # get batch meta data\n    batch_meta_df = test_meta_df_spliter(batch_id)\n\n    # register pulses\n    test_x = np.zeros((len(batch_meta_df), max_pulse_count, n_features), dtype=\"float16\")    \n    test_x[:, :, 2] = -1\n    \n\n    def read_event_local(event_idx):\n        return read_event(event_idx, batch_meta_df, max_pulse_count, train=False)\n\n    \n    # scan events\n    iterator = range(len(batch_meta_df))\n    with multiprocessing.Pool() as pool:\n        for event_idx, pulse_count, event_x in pool.map(read_event_local, iterator):\n            # feature\n            test_x[event_idx, :pulse_count, 0] = event_x[\"time\"]\n            test_x[event_idx, :pulse_count, 1] = event_x[\"charge\"]\n            test_x[event_idx, :pulse_count, 2] = event_x[\"auxiliary\"]\n            test_x[event_idx, :pulse_count, 3] = event_x[\"x\"]\n            test_x[event_idx, :pulse_count, 4] = event_x[\"y\"]\n            test_x[event_idx, :pulse_count, 5] = event_x[\"z\"]\n            test_x[event_idx, :pulse_count, 6] = event_x[\"r_err\"]\n            test_x[event_idx, :pulse_count, 7] = event_x[\"z_err\"]\n    \n    del batch_meta_df\n    \n    # CONVERT\n    test_x[:, :, 0] /= 1000  # time\n    test_x[:, :, 1] /= 300  # charge\n    test_x[:, :, 3:] /= 600  # space\n    \n    # PREDICT\n    pred_angles = list()\n    for model, shifted in zip(models, model_shifted):\n        pred_model = model.predict(test_x, verbose=0)\n        \n        if shifted:\n            az_model, zen_model = pred_to_angle_azshift(pred_model)\n        else:\n            az_model, zen_model = pred_to_angle(pred_model)\n    \n        pred_angles.append((az_model, zen_model))\n    \n    pred_azimuth, pred_zenith = weighted_vector_ensemble(pred_angles, weights)\n    \n    event_ids = test_meta_df.event_id[test_meta_df.batch_id == batch_id].values\n    \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.)","metadata":{"execution":{"iopub.status.busy":"2023-03-07T12:57:43.887497Z","iopub.execute_input":"2023-03-07T12:57:43.887898Z","iopub.status.idle":"2023-03-07T12:57:46.818089Z","shell.execute_reply.started":"2023-03-07T12:57:43.887858Z","shell.execute_reply":"2023-03-07T12:57:46.816881Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_result_dict = {\n    \"event_id\": test_event_id,\n    \"azimuth\": test_azimuth,\n    \"zenith\": test_zenith,\n}\n\ntest_result_df = pd.DataFrame(test_result_dict)\ntest_result_df = test_result_df.sort_values(by = [\"event_id\"])\n\ntest_result_df.to_csv(\"submission.csv\", index = False)\ntest_result_df.head()","metadata":{"execution":{"iopub.status.busy":"2023-03-07T12:57:46.819624Z","iopub.execute_input":"2023-03-07T12:57:46.820663Z","iopub.status.idle":"2023-03-07T12:57:46.838756Z","shell.execute_reply.started":"2023-03-07T12:57:46.820621Z","shell.execute_reply":"2023-03-07T12:57:46.837177Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}