{"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":"# Animated events with matplotlib","metadata":{}},{"cell_type":"code","source":"%matplotlib widget\n\nimport numpy as np\nimport pandas as pd\nimport pyarrow.parquet as pq\nimport matplotlib.pyplot as plt\nimport matplotlib.animation as animation\n\nfrom mpl_toolkits.mplot3d import Axes3D","metadata":{"execution":{"iopub.status.busy":"2023-01-28T18:39:54.807835Z","iopub.execute_input":"2023-01-28T18:39:54.808440Z","iopub.status.idle":"2023-01-28T18:39:54.977554Z","shell.execute_reply.started":"2023-01-28T18:39:54.808312Z","shell.execute_reply":"2023-01-28T18:39:54.976605Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Reading datasets","metadata":{}},{"cell_type":"code","source":"train_df = pd.read_parquet(f\"/kaggle/input/icecube-neutrinos-in-deep-ice/train_meta.parquet\")","metadata":{"execution":{"iopub.status.busy":"2023-01-28T18:39:54.979354Z","iopub.execute_input":"2023-01-28T18:39:54.979869Z","iopub.status.idle":"2023-01-28T18:40:41.760487Z","shell.execute_reply.started":"2023-01-28T18:39:54.979814Z","shell.execute_reply":"2023-01-28T18:40:41.759124Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_df.head()","metadata":{"execution":{"iopub.status.busy":"2023-01-28T18:40:41.761863Z","iopub.execute_input":"2023-01-28T18:40:41.762819Z","iopub.status.idle":"2023-01-28T18:40:41.785091Z","shell.execute_reply.started":"2023-01-28T18:40:41.762779Z","shell.execute_reply":"2023-01-28T18:40:41.783509Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"batch_df = pd.read_parquet(f'/kaggle/input/icecube-neutrinos-in-deep-ice/train/batch_1.parquet')\nbatch_df.reset_index(inplace=True)","metadata":{"execution":{"iopub.status.busy":"2023-01-28T18:40:41.788360Z","iopub.execute_input":"2023-01-28T18:40:41.789334Z","iopub.status.idle":"2023-01-28T18:40:45.665764Z","shell.execute_reply.started":"2023-01-28T18:40:41.789250Z","shell.execute_reply":"2023-01-28T18:40:45.664271Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sensors = pd.read_csv(f'/kaggle/input/icecube-neutrinos-in-deep-ice/sensor_geometry.csv')\nsensors.head()","metadata":{"execution":{"iopub.status.busy":"2023-01-28T18:40:45.667540Z","iopub.execute_input":"2023-01-28T18:40:45.668087Z","iopub.status.idle":"2023-01-28T18:40:45.694440Z","shell.execute_reply.started":"2023-01-28T18:40:45.668035Z","shell.execute_reply":"2023-01-28T18:40:45.693046Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Select and load an event","metadata":{}},{"cell_type":"code","source":"eid = 1128590","metadata":{"execution":{"iopub.status.busy":"2023-01-28T18:40:45.696129Z","iopub.execute_input":"2023-01-28T18:40:45.696555Z","iopub.status.idle":"2023-01-28T18:40:45.701993Z","shell.execute_reply.started":"2023-01-28T18:40:45.696517Z","shell.execute_reply":"2023-01-28T18:40:45.700851Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"event_df = batch_df[batch_df[\"event_id\"] == eid].reset_index(drop=True)\n\nlabel = train_df[train_df[\"event_id\"] == eid][[\"azimuth\", \"zenith\"]].values[0]\n\nazimuth = label[0]\nzenith = label[1]\n\nlabel_x = np.cos(azimuth) * np.sin(zenith)\nlabel_y = np.sin(azimuth) * np.sin(zenith)\nlabel_z = np.cos(zenith)\n\nlabel_line = [\n    [-label_x * 500, label_x * 500],\n    [-label_y * 500, label_y * 500], \n    [-label_z * 500, label_z * 500]\n]\n\nevent_df[\"charge\"] = np.log(1 + event_df[\"charge\"])\nevent_df[\"time_ns\"] = pd.to_datetime(event_df[\"time\"], unit='ns')\nevent_df = event_df.set_index(\"time_ns\")","metadata":{"execution":{"iopub.status.busy":"2023-01-28T18:41:02.205430Z","iopub.execute_input":"2023-01-28T18:41:02.206320Z","iopub.status.idle":"2023-01-28T18:41:04.428946Z","shell.execute_reply.started":"2023-01-28T18:41:02.206270Z","shell.execute_reply":"2023-01-28T18:41:04.427306Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"w = event_df[\"charge\"].rolling(\"6000N\", min_periods=1).sum().argmax()\nwt = event_df.iloc[w][\"time\"]\n\nevent_df = event_df.loc[pd.Timestamp(wt - 6000, unit=\"ns\"):pd.Timestamp(wt, unit=\"ns\")]\nevent_df[\"time\"] = event_df[\"time\"] - event_df[\"time\"].min()\nevent_df[\"time_ns\"] = pd.to_datetime(event_df[\"time\"], unit='ns')\nevent_df = event_df.set_index(\"time_ns\")","metadata":{"execution":{"iopub.status.busy":"2023-01-28T18:41:04.431143Z","iopub.execute_input":"2023-01-28T18:41:04.431613Z","iopub.status.idle":"2023-01-28T18:41:04.452960Z","shell.execute_reply.started":"2023-01-28T18:41:04.431557Z","shell.execute_reply":"2023-01-28T18:41:04.451501Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"event_df.head()","metadata":{"execution":{"iopub.status.busy":"2023-01-28T18:41:06.930823Z","iopub.execute_input":"2023-01-28T18:41:06.931305Z","iopub.status.idle":"2023-01-28T18:41:06.946798Z","shell.execute_reply.started":"2023-01-28T18:41:06.931269Z","shell.execute_reply":"2023-01-28T18:41:06.945278Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Animation with matplotlib","metadata":{}},{"cell_type":"code","source":"MAX_SIZE = 1\nNUM_FRAMES = 100\nCOLORMAP = plt.cm.viridis\n\nLAST_TIME = event_df.iloc[-1][\"time\"]","metadata":{"execution":{"iopub.status.busy":"2023-01-28T18:41:13.135328Z","iopub.execute_input":"2023-01-28T18:41:13.135867Z","iopub.status.idle":"2023-01-28T18:41:13.143461Z","shell.execute_reply.started":"2023-01-28T18:41:13.135820Z","shell.execute_reply":"2023-01-28T18:41:13.142433Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"first_time = event_df.groupby(by=\"sensor_id\").min()[[\"time\"]].reset_index()\nfirst_time[\"time\"] = first_time[\"time\"] / first_time[\"time\"].max()\n\nevent_sensors = sensors[sensors[\"sensor_id\"].isin(event_df[\"sensor_id\"].unique())].copy()\nevent_sensors.sort_values(by=\"sensor_id\", inplace=True)\nevent_sensors = event_sensors.merge(first_time, on=\"sensor_id\", how=\"left\")\n\nevent_df[\"charge\"] = (event_df[\"charge\"] / event_df[\"charge\"].max()) * MAX_SIZE","metadata":{"execution":{"iopub.status.busy":"2023-01-28T18:41:13.595914Z","iopub.execute_input":"2023-01-28T18:41:13.596337Z","iopub.status.idle":"2023-01-28T18:41:13.627209Z","shell.execute_reply.started":"2023-01-28T18:41:13.596304Z","shell.execute_reply":"2023-01-28T18:41:13.626098Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"FRAMES = []\n\ndef next_step(frame_num):\n    \n    i = (frame_num * LAST_TIME) / NUM_FRAMES\n    \n    step = event_df.loc[\n        event_df.index < pd.Timestamp(i, unit='ns')\n    ].groupby(by=\"sensor_id\").sum()[[\"charge\"]].reset_index()\n\n    results = event_sensors.copy()\n    results = results.merge(step, on=\"sensor_id\", how=\"left\")\n    results.fillna(0, inplace=True)\n    \n    return results\n\ndef update(num):\n    \n    data = FRAMES[num]\n    \n    for i, row in data.iterrows():\n        data_points[sensor_map[row[\"sensor_id\"]]].set(\n            markersize=row[\"charge\"],\n        )\n        \n    return data_points\n\nfor i in range(NUM_FRAMES):\n    \n    new_frame = next_step(i)\n    \n    if i > 0:\n        old_frame = next_step(i-1)\n        old_frame.rename(columns={'charge': \"old_charge\"}, inplace=True)\n        new_frame = new_frame.merge(old_frame[[\"sensor_id\", \"old_charge\"]], on=[\"sensor_id\"], how=\"left\")\n        new_frame[\"delta_charge\"] = new_frame[\"charge\"] - new_frame[\"old_charge\"]\n        new_frame = new_frame[new_frame[\"delta_charge\"] > 0].copy()\n        new_frame.drop(columns=[\"old_charge\", \"delta_charge\"], inplace=True)\n    \n    FRAMES.append(new_frame.copy())\n    ","metadata":{"execution":{"iopub.status.busy":"2023-01-28T18:41:15.540351Z","iopub.execute_input":"2023-01-28T18:41:15.540814Z","iopub.status.idle":"2023-01-28T18:41:17.450135Z","shell.execute_reply.started":"2023-01-28T18:41:15.540776Z","shell.execute_reply":"2023-01-28T18:41:17.448990Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"fig = plt.figure(figsize=(9, 6))\nax = fig.add_subplot(111, projection='3d')\ntitle = ax.set_title('Event #{0}'.format(eid))\n\ndata_points = []\ndata = FRAMES[0]\nsensor_map = {}\n\nfor i, row in data.iterrows():\n    sensor_map[row[\"sensor_id\"]] = i\n    data_points.append(\n        ax.plot(\n            row[\"x\"], row[\"y\"], row[\"z\"],\n            markerfacecolor=COLORMAP(row[\"time\"]),\n            markeredgecolor=COLORMAP(row[\"time\"]),\n            marker='o',\n            markersize=0,\n            alpha=1.0\n        )[0]\n    )\n\nax.set_xlim3d(-800, 800)\nax.set_ylim3d(-800, 800)\nax.set_zlim3d(-800, 800)\nax.plot(xs=label_line[0], ys=label_line[1], zs=label_line[2], color=\"r\")\n\nanim = animation.FuncAnimation(fig, update, frames=NUM_FRAMES, interval=15, repeat=True, blit=True)\n\n# anim.save(f'event_{eid}.gif', writer='imagemagick', fps=24)\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2023-01-28T18:41:50.562172Z","iopub.execute_input":"2023-01-28T18:41:50.562887Z","iopub.status.idle":"2023-01-28T18:41:51.293410Z","shell.execute_reply.started":"2023-01-28T18:41:50.562844Z","shell.execute_reply":"2023-01-28T18:41:51.292165Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"![](https://i.imgur.com/kcsHHOk.gif)","metadata":{}},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}