{"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":"# LSTM Preprocessing point picker\n\n- Important data point picker algorithm was included in preprocessing.","metadata":{}},{"cell_type":"markdown","source":"# Set-up\n- Import packages\n- Set hyperparameters","metadata":{}},{"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","metadata":{"execution":{"iopub.status.busy":"2023-02-01T12:38:33.162603Z","iopub.execute_input":"2023-02-01T12:38:33.163172Z","iopub.status.idle":"2023-02-01T12:38:35.079818Z","shell.execute_reply.started":"2023-02-01T12:38:33.163059Z","shell.execute_reply":"2023-02-01T12:38:35.078603Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Data setting\ntrain_batch_id_first = 101\ntrain_batch_id_last = 105\n\ntrain_batch_ids = range(train_batch_id_first, train_batch_id_last + 1)\n\nmax_pulse_count = 128\nbin_num = 8\nn_features = 9  # time, charge, aux, x, y, z, r_err, z_err, rank","metadata":{"execution":{"iopub.status.busy":"2023-02-01T12:38:35.081801Z","iopub.execute_input":"2023-02-01T12:38:35.082269Z","iopub.status.idle":"2023-02-01T12:38:35.088204Z","shell.execute_reply.started":"2023-02-01T12:38:35.082237Z","shell.execute_reply":"2023-02-01T12:38:35.087068Z"},"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\npoint_picker_format = './pointpicker_mpc128_n9_batch_{batch_id:d}.npz'","metadata":{"execution":{"iopub.status.busy":"2023-02-01T12:42:55.012876Z","iopub.execute_input":"2023-02-01T12:42:55.013436Z","iopub.status.idle":"2023-02-01T12:42:55.020475Z","shell.execute_reply.started":"2023-02-01T12:42:55.013395Z","shell.execute_reply":"2023-02-01T12:42:55.018871Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Set Detector Geometry","metadata":{}},{"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-02-01T12:38:55.819574Z","iopub.execute_input":"2023-02-01T12:38:55.820032Z","iopub.status.idle":"2023-02-01T12:38:55.866014Z","shell.execute_reply.started":"2023-02-01T12:38:55.819994Z","shell.execute_reply":"2023-02-01T12:38:55.865046Z"},"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-02-01T12:38:56.747611Z","iopub.execute_input":"2023-02-01T12:38:56.748039Z","iopub.status.idle":"2023-02-01T12:39:04.550275Z","shell.execute_reply.started":"2023-02-01T12:38:56.747995Z","shell.execute_reply":"2023-02-01T12:39:04.549118Z"},"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-02-01T12:39:04.552360Z","iopub.execute_input":"2023-02-01T12:39:04.553107Z","iopub.status.idle":"2023-02-01T12:39:04.563985Z","shell.execute_reply.started":"2023-02-01T12:39:04.553061Z","shell.execute_reply":"2023-02-01T12:39:04.562540Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Data I/O Helper","metadata":{}},{"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":{}},{"cell_type":"code","source":"# read single event from batch_meta_df\ndef read_event(event_idx, batch_meta_df, max_pulse_count, batch_df, 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    # 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-02-01T13:03:00.724553Z","iopub.execute_input":"2023-02-01T13:03:00.725004Z","iopub.status.idle":"2023-02-01T13:03:00.740672Z","shell.execute_reply.started":"2023-02-01T13:03:00.724966Z","shell.execute_reply":"2023-02-01T13:03:00.739549Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Data I/O (for CPU) & Normalization\n\n- Read data\n- Concatenate and convert\n- Normalize time, charge and position variables","metadata":{}},{"cell_type":"markdown","source":"## Read train metadata and define spliter (for CPU)","metadata":{}},{"cell_type":"code","source":"%%time\nprint(\"for CPU job\")\n\ntrain_meta_df = pd.read_parquet(home_dir + 'train_meta.parquet')\n\nbatch_counts = train_meta_df.batch_id.value_counts().sort_index()\n\nbatch_max_index = batch_counts.cumsum()\nbatch_max_index[train_meta_df.batch_id.min() - 1] = 0\nbatch_max_index = batch_max_index.sort_index()\n\n\ndef train_meta_df_spliter(batch_id):\n    return train_meta_df.loc[batch_max_index[batch_id - 1]:batch_max_index[batch_id] - 1]","metadata":{"execution":{"iopub.status.busy":"2023-02-01T13:03:03.140576Z","iopub.execute_input":"2023-02-01T13:03:03.142024Z","iopub.status.idle":"2023-02-01T13:03:17.383813Z","shell.execute_reply.started":"2023-02-01T13:03:03.141964Z","shell.execute_reply":"2023-02-01T13:03:17.382366Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Read training data\n\n- Takes 400 s / batch\n- About 900 MB / batch","metadata":{}},{"cell_type":"code","source":"%%time\n\nfor batch_id in train_batch_ids:\n    print(\"Reading batch \", batch_id, end=\"\")\n    # get batch meta data and data\n    batch_meta_df = train_meta_df_spliter(batch_id)\n    batch_df = pd.read_parquet(train_format.format(batch_id=batch_id))\n\n    # register pulses\n    batch_x = np.zeros((len(batch_meta_df), max_pulse_count, n_features), dtype=\"float16\")\n    batch_y = np.zeros((len(batch_meta_df), 2), dtype=\"float16\")\n    \n    batch_x[:, :, 2] = -1\n    \n\n    def read_event_local(event_idx):\n        return read_event(event_idx, batch_meta_df, max_pulse_count, batch_df, train=True)\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, event_y in pool.map(read_event_local, iterator):\n            batch_x[event_idx, :pulse_count, 0] = event_x[\"time\"]\n            batch_x[event_idx, :pulse_count, 1] = event_x[\"charge\"]\n            batch_x[event_idx, :pulse_count, 2] = event_x[\"auxiliary\"]\n            batch_x[event_idx, :pulse_count, 3] = event_x[\"x\"]\n            batch_x[event_idx, :pulse_count, 4] = event_x[\"y\"]\n            batch_x[event_idx, :pulse_count, 5] = event_x[\"z\"]\n            batch_x[event_idx, :pulse_count, 6] = event_x[\"r_err\"]\n            batch_x[event_idx, :pulse_count, 7] = event_x[\"z_err\"]\n            batch_x[event_idx, :pulse_count, 8] = event_x[\"rank\"]\n\n            batch_y[event_idx] = event_y\n\n    del batch_meta_df, batch_df\n    \n    # save\n    print(\" DONE! Saving...\")\n    np.savez(point_picker_format.format(batch_id=batch_id), x=batch_x, y=batch_y)","metadata":{"execution":{"iopub.status.busy":"2023-02-01T13:03:17.386251Z","iopub.execute_input":"2023-02-01T13:03:17.386959Z","iopub.status.idle":"2023-02-01T13:09:52.281667Z","shell.execute_reply.started":"2023-02-01T13:03:17.386913Z","shell.execute_reply":"2023-02-01T13:09:52.280722Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# END-OF-NOTE","metadata":{}}]}