{"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":"# Relation-Shape-RNN, is it the edge of RNN based model?\n\n- Why RNN-based?\n    - As I have shared my code in the *early sharing prize* period, very light LSTM model showed some performance.\n        - [3 LSTMs; with Data Picking and Shifting](https://www.kaggle.com/code/seungmoklee/3-lstms-with-data-picking-and-shifting)\n    - And @rsmits even improved the performance of RNN model (actually, GRU) further. His score was 1.017 in this [note](https://www.kaggle.com/code/rsmits/tensorflow-lstm-model-inference).\n    - These may imply there could be some potential on RNN-based method. So I decided to bet my effort on it.\n        - It is much easier to understand than GNN. Don't need other packages like `~_geometry` or `~_knn~`.\n- Relation-Shape?\n    - When the host opened their baseline, I've searched for some more researches about point clouds, and found this one.\n        - [Relation-Shape Convolutional Neural Network for Point Cloud Analysis](https://arxiv.org/abs/1904.07601)\n    - The key point was to train the `local shape` features. They gathered some nearest neighbors using knn, evaluated some local shape features like $d\\vec{x}$ or euclidean distance, and applied cnn on those features.\n    - They claimed their performance was better than DGCNN.\n- How to train Relation-Shape in RNN?\n    - I've interpreted our dataset as 4-dimensional point cloud.\n        - $x^{\\mu} = \\left< t, ~ x, ~ y, ~ z \\right>$\n    - My guess (and maybe the most important reason for the low performance of my model) was that the *nearest neighbor* in our data can be approximated as *time-nearest neighbor*.\n    - So, I have implemented the relation-shape learning as 1D-Conv network, where the points are sorted in the time order.\n","metadata":{}},{"cell_type":"markdown","source":"# Model Description\n\n- Input Features\n    - 6 features are used as input.\n        - $\\vec{f} \\equiv $ $<$ $t$, $a$(uxiliary), $c$(harge), $x$, $y$, $z$ $>$\n        - As @rsmits has claimed in his/her [note](https://www.kaggle.com/code/rsmits/tensorflow-lstm-model-inference), the resolution term was discarded in this version.\n- Two RSBlocks\n    - I have implemented `RSBlock` class. It consists of H (low level feature) layer, M (mapping) layer, FR (feature raising) layer, A (aggregation) layer and CR (channel raising) layer.\n    - H, Low Level Features for Relation-Shape\n        - $\\rm \\left(B_{atch}, N_{umber-of-pulses}, F_{eatures} \\right) \\rightarrow \\left( B, N, K_{nearest-neighbors}, H_{low-level-features} \\right)$\n        - For each points $\\vec{f_{0}}$ and its time-nearest points $\\vec{f_{k}}$, evaluate low level feautres $\\vec{h_{k, 0}} = \\left< t_{k} - t_{0}, x_{k} - x_{0}, \\rm{distance}, \\cdots \\right>$.\n    - M, Mapping Layer\n        - $\\rm \\left(B, N, K, F_{in} \\right) \\rightarrow \\left( B, N, K, F_{out} \\right)$\n        - From $\\rm F_{in}$ number of low-level features, generated $\\rm F_{out}$ number of higher-level features.\n        - Conv2d with $1 \\times 1$ window was used.\n    - FR, Feature Raising Layer\n        - $\\rm \\left(B, N, F_{eatures} \\right) \\rightarrow \\left( B, N, F_{FR} \\right)$\n        - For the feature vector of the point (not neighbor), $\\vec{f_{0}}$, increase its features into higher level.\n        - Conv1d with window length of 1 was used.\n    - A, Aggregation Layer\n        - So, now we have two tensors.\n            - $FR_{n, f}$ with shape of $\\rm \\left( B, N, F_{FR} \\right)$ from the point, and $M_{nkf}$ with shape of $\\rm \\left( B, N, K, F_{out} \\right)$ from its neighbors.\n            - $\\rm F_{FR}$ and $\\rm F_{out}$ must be same, \n        - Calculate tensor $T$ as component-wise product, $T_{n, k, f} = FR_{n, f} \\cdot M_{n, k, f}$.\n        - Aggregate features from its neighbors. $A_{n, f} = \\rm{max}_{among ~ \\it k} \\it \\left( T_{n, k, f} \\right)$.\n        - Now, we have $A_{n, f}$ tensor with shape of $\\rm \\left(B, N, F_{FR} \\right)$\n    - CR, Channel Raising Layer\n        - $\\rm \\left(B, N, F_{FR} \\right) \\rightarrow \\left( B, N, F_{CR} \\right)$\n        - Raise number of channels using 1D Conv with the window length of 1.\n- GABlock with RNN\n    - $\\rm \\left(B, N, F_{CR} \\right) \\rightarrow \\left( B, 3 \\right)$\n    - Globally aggregate relation-shape pulse features to predict the angle.\n    - As @rsmits suggested, 3 GRU layers were adopted here.\n    - Attached one hidden Dense layer after the series of GRUs, and one output Dense layer for prediction.\n- You can find all the implementation in the `Layer.py` file in the [zb-icecubemodels](https://www.kaggle.com/datasets/seungmoklee/zb-icecubemodels) dataset.\n","metadata":{}},{"cell_type":"markdown","source":"# Loss Function\n\n- I've implemented and used `VonMisesFisher3DLoss` which was used by the host, in tensorflow.\n- You can find the implementation in `Loss.py` in the [zb-icecubemodels](https://www.kaggle.com/datasets/seungmoklee/zb-icecubemodels) dataset.","metadata":{}},{"cell_type":"markdown","source":"# Import and Setting","metadata":{}},{"cell_type":"code","source":"!mkdir ./Packages\n!cp /kaggle/input/zb-icecubemodels/*.py ./Packages\n!ls ./Packages/","metadata":{"execution":{"iopub.status.busy":"2023-04-19T16:12:59.526405Z","iopub.execute_input":"2023-04-19T16:12:59.526915Z","iopub.status.idle":"2023-04-19T16:13:02.634070Z","shell.execute_reply.started":"2023-04-19T16:12:59.526851Z","shell.execute_reply":"2023-04-19T16:13:02.632712Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# %% Imports\nimport os\nimport random\nimport time\nimport gc\n\nimport matplotlib as mpl\nimport matplotlib.pyplot as plt\nfrom tqdm import tqdm\nimport numpy as np\nimport pandas as pd\nimport pyarrow.parquet as pq\nimport pyarrow as pa\nimport tensorflow as tf\n\nfrom Packages.Layers import GABlockResRNN, RSBlock\nfrom Packages.Losses import VonMisesFisher3DLoss\nfrom Packages.Metrics import AngularDistScore, angular_dist_score\nfrom Packages.Utils import GpuMemoryManagement\n\nfrom typing import List","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-04-19T16:13:02.636940Z","iopub.execute_input":"2023-04-19T16:13:02.637417Z","iopub.status.idle":"2023-04-19T16:13:12.109112Z","shell.execute_reply.started":"2023-04-19T16:13:02.637366Z","shell.execute_reply":"2023-04-19T16:13:12.107830Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Sensor Geometry","metadata":{}},{"cell_type":"code","source":"# sensor_geometry\nsensor_geometry_df = pd.read_csv(\"/kaggle/input/icecube-neutrinos-in-deep-ice/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\n# r-error\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\n# z-error\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        else:\n            sensor_z_err[string_id * doms_per_string + dom_id] = long_z_resolution\n# register\nsensor_geometry_df[\"r_err\"] = sensor_r_err\nsensor_geometry_df[\"z_err\"] = sensor_z_err","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2023-04-19T16:05:56.925329Z","iopub.execute_input":"2023-04-19T16:05:56.925843Z","iopub.status.idle":"2023-04-19T16:05:56.959592Z","shell.execute_reply.started":"2023-04-19T16:05:56.925799Z","shell.execute_reply":"2023-04-19T16:05:56.958669Z"},"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":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2023-04-19T16:05:56.962074Z","iopub.execute_input":"2023-04-19T16:05:56.962712Z","iopub.status.idle":"2023-04-19T16:05:56.972752Z","shell.execute_reply.started":"2023-04-19T16:05:56.962673Z","shell.execute_reply":"2023-04-19T16:05:56.971439Z"},"jupyter":{"source_hidden":true},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Generator From Parquet","metadata":{}},{"cell_type":"code","source":"def generator_from_parquet(\n    meta_data_path,\n    data_path_header,\n    max_pulse_count: int=128,\n    max_batch_id: int=9999,\n):\n    # type conversion for tf\n    if type(meta_data_path) == bytes:\n        meta_data_path = str(meta_data_path, 'utf-8')\n    if type(data_path_header) == bytes:\n        data_path_header = str(data_path_header, 'utf-8')\n        \n    print(meta_data_path)\n    print(data_path_header)\n    \n    # read one batch\n    meta_data = pq.ParquetFile(meta_data_path)\n    for meta_batch in meta_data.iter_batches(batch_size=200000):\n        # batch meta data\n        batch_id = meta_batch[\"batch_id\"][0].as_py()\n        event_ids = meta_batch[\"event_id\"].to_numpy()\n        \n        if batch_id > max_batch_id:\n            print(\"Reached max_batch_id!\")\n            return\n        \n        # batch data\n        data_batch = pq.read_table(data_path_header + f\"{batch_id:d}.parquet\")\n        sensor_id = data_batch[\"sensor_id\"].combine_chunks().to_numpy()\n        time = data_batch[\"time\"].combine_chunks().to_numpy()\n        charge = data_batch[\"charge\"].combine_chunks().to_numpy()\n        auxiliary = data_batch[\"auxiliary\"].combine_chunks().to_numpy(False)\n        pulse_index = np.append(meta_batch[\"first_pulse_index\"].to_numpy(), [data_batch.num_rows])\n        \n        # read each event\n        for event_id, first_idx, last_idx in zip(event_ids, pulse_index[:-1], pulse_index[1:]):\n            event_time = time[first_idx:last_idx]\n            event_time = event_time - event_time.min()\n            event_charge = charge[first_idx:last_idx]\n            event_auxiliary = auxiliary[first_idx:last_idx]\n            event_x = sensor_geometry_df.x[sensor_id[first_idx:last_idx]].values\n            event_y = sensor_geometry_df.y[sensor_id[first_idx:last_idx]].values\n            event_z = sensor_geometry_df.z[sensor_id[first_idx:last_idx]].values\n            \n            dtype = [\n                (\"time\", \"float16\"),\n                (\"charge\", \"float16\"),\n                (\"auxiliary\", \"float16\"),\n                (\"x\", \"float16\"),\n                (\"y\", \"float16\"),\n                (\"z\", \"float16\"),\n                (\"rank\", \"short\")\n            ]\n            event_features = np.zeros(last_idx - first_idx, dtype)\n            event_features[\"time\"] = event_time\n            event_features[\"charge\"] = event_charge\n            event_features[\"auxiliary\"] = event_auxiliary\n            event_features[\"x\"] = event_x\n            event_features[\"y\"] = event_y\n            event_features[\"z\"] = event_z\n\n            # point picker\n            if len(event_x) > max_pulse_count:\n                # find valid time window\n                t_peak = event_features[\"time\"][event_features[\"charge\"].argmax()]\n                t_valid_min = t_peak - t_valid_length\n                t_valid_max = t_peak + t_valid_length\n                \n                # rank\n                t_valid = (event_features[\"time\"] > t_valid_min) * (event_features[\"time\"] < t_valid_max)\n                event_features[\"rank\"] = 2 * (1 - event_features[\"auxiliary\"]) + (t_valid)\n                \n                # sort by rank and charge\n                event_features = np.sort(event_features, order=[\"rank\", \"charge\"])\n                \n                # pick-up from backward\n                event_features = event_features[-max_pulse_count:]\n                \n                # resort by time\n                event_features = np.sort(event_features, order=\"time\")\n            \n            pulse_count = min(len(event_x), max_pulse_count)\n            \n            # yield\n            features = np.zeros((max_pulse_count, 6), dtype=\"float32\")\n            features[:pulse_count, 0] = (event_features[\"time\"] - event_features[\"time\"].min()).astype(\"float32\") / 1000.\n            features[:pulse_count, 1] = event_features[\"charge\"].astype(\"float32\") / 300.\n            features[:pulse_count, 2] = event_features[\"auxiliary\"].astype(\"float32\") * 1.\n            features[:pulse_count, 3] = event_features[\"x\"].astype(\"float32\") / 600.\n            features[:pulse_count, 4] = event_features[\"y\"].astype(\"float32\") / 600.\n            features[:pulse_count, 5] = event_features[\"z\"].astype(\"float32\") / 600.\n\n            features[:, :][features[:, 1] == 0.0] = 0.0  # for masking layer\n            \n            yield event_id, features\n        \n        # memory management\n        del auxiliary, charge, time, sensor_id, data_batch, pulse_index, event_id\n        _=gc.collect()","metadata":{"execution":{"iopub.status.busy":"2023-04-19T16:05:56.974971Z","iopub.execute_input":"2023-04-19T16:05:56.975713Z","iopub.status.idle":"2023-04-19T16:05:56.997732Z","shell.execute_reply.started":"2023-04-19T16:05:56.975672Z","shell.execute_reply":"2023-04-19T16:05:56.996612Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Load Model","metadata":{}},{"cell_type":"code","source":"model_dir = \"/kaggle/input/zb-icecubemodels/\"\n\nmodel_da_name = \"RSGRUDA_s3260114_s595_s756_s639_s585_s403_s158_s263_s756\"\nmodel_da = tf.keras.models.load_model(\n    model_dir + model_da_name + \".h5\",\n    custom_objects={\n        \"VonMisesFisher3DLoss\": VonMisesFisher3DLoss,\n        \"AngularDistScore\": AngularDistScore,\n        \"RSBlock\": RSBlock,\n        \"GABlockResRNN\": GABlockResRNN,\n    },\n)\nmodel_da.summary()","metadata":{"execution":{"iopub.status.busy":"2023-04-19T16:05:57.022719Z","iopub.execute_input":"2023-04-19T16:05:57.023562Z","iopub.status.idle":"2023-04-19T16:06:07.356153Z","shell.execute_reply.started":"2023-04-19T16:05:57.023526Z","shell.execute_reply":"2023-04-19T16:06:07.355355Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Main","metadata":{}},{"cell_type":"code","source":"max_batch_id = 99999\nbatch_size = 128\n\nif max_batch_id < 9999:\n    print(\"TESTING!!!\")\n    split_name = \"train\"\nelse:\n    split_name = \"test\"\n\nhome_dir = \"/kaggle/input/icecube-neutrinos-in-deep-ice/\"\nmeta_data_path = home_dir + split_name + '_meta.parquet'\ndata_dir = home_dir + split_name + \"/\"\ndata_path_header = data_dir + \"batch_\"\n\nmax_pulse_count = 128\nmax_pulse_count_lite = 80\n\nAUTOTUNE = tf.data.experimental.AUTOTUNE","metadata":{"execution":{"iopub.status.busy":"2023-04-19T16:13:12.111253Z","iopub.execute_input":"2023-04-19T16:13:12.111966Z","iopub.status.idle":"2023-04-19T16:13:12.123344Z","shell.execute_reply.started":"2023-04-19T16:13:12.111933Z","shell.execute_reply":"2023-04-19T16:13:12.121204Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_da_ds = tf.data.Dataset.from_generator(\n    generator_from_parquet,\n    args=[meta_data_path, data_path_header, max_pulse_count, max_batch_id],\n    output_types=(tf.int32, tf.float32),\n    output_shapes=((None), (max_pulse_count, 6)),\n)\ntest_da_ds = test_da_ds.batch(batch_size).prefetch(AUTOTUNE)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Predict and Save","metadata":{}},{"cell_type":"code","source":"with open('pred_da.csv', 'w') as pred_csv:\n    pred_csv.write('event_id,pred_x,pred_y,pred_z\\n')\n\nwith open('pred_lite.csv', 'w') as pred_csv:\n    pred_csv.write('event_id,pred_x,pred_y,pred_z,pulse\\n')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\n\nsteps_per_batch = np.ceil(200000 // batch_size)\nprint(f\"One batch has {int(steps_per_batch):d} steps\")\n\nstep = 0\nfor batch_data in test_da_ds:\n    step += 1\n    print(f\"{step:7d}\", end=\"\")\n    if (step % 20 == 0) or (step % steps_per_batch == 0):\n        print()\n        \n    batch_event_id, batch_features = batch_data\n#     batch_features[:, :, :][batch_features[:, :, 1] == 0.0] = 0.0  # for masking\n    \n    test_pred = model_da.predict_on_batch(batch_features)\n    \n    with open('pred_da.csv', 'a') as pred_csv:\n        for event_id, pred_x, pred_y, pred_z in zip(batch_event_id, test_pred[:, 0], test_pred[:, 1], test_pred[:, 2]):\n            pred_csv.write(f\"{event_id:d},{pred_x:f},{pred_y:f},{pred_z:f}\\n\")\n    \n    del batch_event_id, batch_features, test_pred\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Memory Management\n\n- Discard models from GPU and RAM","metadata":{}},{"cell_type":"code","source":"del model_da\n\ntf.keras.backend.clear_session()\n\nprint(\"gc collect: \", gc.collect())\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Convert to angle","metadata":{}},{"cell_type":"code","source":"# read predictions\npred_da = pd.read_csv(\"pred_da.csv\")\n\nevent_id = pred_da.event_id.values\n\npred_da_x = pred_da.pred_x.values\npred_da_y = pred_da.pred_y.values\npred_da_z = pred_da.pred_z.values","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_pred = np.zeros((len(pred_da_x), 3))\n\ntest_pred[:, 0] = pred_da_x\ntest_pred[:, 1] = pred_da_y\ntest_pred[:, 2] = pred_da_z\n\n# convert to angle\nkappa_pred = np.linalg.norm(test_pred, axis=1)\nvec_x_pred = test_pred[:, 0] / kappa_pred\nvec_y_pred = test_pred[:, 1] / kappa_pred\nvec_z_pred = test_pred[:, 2] / kappa_pred\naz_pred = np.arctan2(vec_y_pred, vec_x_pred)\naz_pred = np.where(az_pred < 0, az_pred + 2*np.pi, az_pred)\nzen_pred = np.arccos(vec_z_pred)\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Submit","metadata":{}},{"cell_type":"code","source":"with open('submission.csv', 'w') as submission:\n    submission.write('event_id,azimuth,zenith\\n')\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with open('submission.csv', 'a') as submission:\n    for event, azimuth, zenith in zip(event_id, az_pred, zen_pred):\n        submission.write(f\"{event:d},{azimuth:f},{zenith:f}\\n\")\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# END-OF-NOTE","metadata":{}}]}