{"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":"#This block is to install everything not on kaggle or graphnet\n\n\n# Install dependencies\n!pip install /kaggle/input/graphnet-and-dependencies/software/dependencies/torch-1.11.0+cu115-cp37-cp37m-linux_x86_64.whl\n!pip install /kaggle/input/graphnet-and-dependencies/software/dependencies/torch_cluster-1.6.0-cp37-cp37m-linux_x86_64.whl\n!pip install /kaggle/input/graphnet-and-dependencies/software/dependencies/torch_scatter-2.0.9-cp37-cp37m-linux_x86_64.whl\n!pip install /kaggle/input/graphnet-and-dependencies/software/dependencies/torch_sparse-0.6.13-cp37-cp37m-linux_x86_64.whl\n!pip install /kaggle/input/graphnet-and-dependencies/software/dependencies/torch_geometric-2.0.4.tar.gz\n!pip install /kaggle/input/graphnet-and-dependencies/software/dependencies/torch_geometric-2.0.4.tar.gz\n\n#import getpass\nfrom pathlib import Path\nfrom typing import Any, Callable, List, Optional, Sequence, Tuple, Union\n\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom scipy.interpolate import interp1d\nfrom sklearn.preprocessing import RobustScaler\nfrom torch import LongTensor, Tensor\n\nfrom torch_geometric.loader import DataLoader\nfrom torch_geometric.data import Data, Dataset\nfrom torch_geometric.data import Data\nfrom torch_geometric.nn import EdgeConv\nfrom torch_geometric.nn.pool import knn_graph\nfrom torch_geometric.typing import Adj\nfrom torch_scatter import scatter_max, scatter_mean, scatter_min, scatter_sum\nimport gc\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2023-03-26T01:09:20.496900Z","iopub.execute_input":"2023-03-26T01:09:20.497534Z","iopub.status.idle":"2023-03-26T01:11:08.275101Z","shell.execute_reply.started":"2023-03-26T01:09:20.497493Z","shell.execute_reply":"2023-03-26T01:11:08.272006Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from pathlib import Path\nfrom typing import Any, Callable, List, Optional, Sequence, Tuple, Union\n\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nfrom scipy.interpolate import interp1d\nfrom sklearn.preprocessing import RobustScaler\nfrom torch import LongTensor, Tensor\n\nfrom torch_geometric.loader import DataLoader\nfrom torch_geometric.data import Data, Dataset\nfrom torch_geometric.data import Data\nfrom torch_geometric.nn import EdgeConv\nfrom torch_geometric.nn.pool import knn_graph\nfrom torch_geometric.typing import Adj\nfrom torch_scatter import scatter_max, scatter_mean, scatter_min, scatter_sum\nimport gc","metadata":{"execution":{"iopub.status.busy":"2023-03-24T13:07:04.804205Z","iopub.execute_input":"2023-03-24T13:07:04.806576Z","iopub.status.idle":"2023-03-24T13:07:04.820054Z","shell.execute_reply.started":"2023-03-24T13:07:04.806481Z","shell.execute_reply":"2023-03-24T13:07:04.818954Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def ice_transparency(data_path, datum=1950):\n    # Data from page 31 of https://arxiv.org/pdf/1301.5361.pdf\n    # Datum is from footnote 8 of page 29\n    df = pd.read_csv(data_path, delim_whitespace=True)\n    df[\"z\"] = df[\"depth\"] - datum\n    df[\"z_norm\"] = df[\"z\"] / 500\n    df[[\"scattering_len_norm\", \"absorption_len_norm\"]] = RobustScaler().fit_transform(\n        df[[\"scattering_len\", \"absorption_len\"]]\n    )\n\n    # These are both roughly equivalent after scaling\n    f_scattering = interp1d(df[\"z_norm\"], df[\"scattering_len_norm\"])\n    f_absorption = interp1d(df[\"z_norm\"], df[\"absorption_len_norm\"])\n    return f_scattering, f_absorption\n\n    \n# preprocessing.py\ndef prepare_sensors():\n    sensors = pd.read_csv(PATH_DATASET + \"sensor_geometry.csv\",).astype(\n        {\n            \"sensor_id\": np.int16,\n            \"x\": np.float32,\n            \"y\": np.float32,\n            \"z\": np.float32,\n        }\n    )\n    sensors[\"string\"] = 0\n    sensors[\"qe\"] = 1\n\n    for i in range(len(sensors) // 60):\n        start, end = i * 60, (i * 60) + 60\n        sensors.loc[start:end, \"string\"] = i\n\n        # High Quantum Efficiency in the lower 50 DOMs - https://arxiv.org/pdf/2209.03042.pdf (Figure 1)\n        if i in range(78, 86):\n            start_veto, end_veto = i * 60, (i * 60) + 10\n            start_core, end_core = end_veto + 1, (i * 60) + 60\n            sensors.loc[start_core:end_core, \"qe\"] = 1.35\n\n    # https://github.com/graphnet-team/graphnet/blob/b2bad25528652587ab0cdb7cf2335ee254cfa2db/src/graphnet/models/detector/icecube.py#L33-L41\n    # Assume that \"rde\" (relative dom efficiency) is equivalent to QE\n    sensors[\"x\"] /= 500\n    sensors[\"y\"] /= 500\n    sensors[\"z\"] /= 500\n    sensors[\"qe\"] -= 1.25\n    sensors[\"qe\"] /= 0.25\n\n    return sensors\n\n# Convert angular to xyz\ndef angle_to_xyz(data):\n    #Data is [n,2] size\n    az = data[:, 0]\n    zen = data[:, 1]\n\n    # pre-compute all sine and cosine values\n    sa = torch.sin(az)\n    ca = torch.cos(az)\n    sz = torch.sin(zen)\n    cz = torch.cos(zen)\n    \n    #return to [N,3]\n    return torch.stack([sz*ca, sz*sa, cz], dim = 1)\nclass IceCubeSubmissionDataset(Dataset): #pytorch geometric dataset, not regular pytorch dataset\n    def __init__(\n        self,\n        batch_id,\n        #event_ids,\n        sensor_df,\n        mode=\"train\",\n        pulse_limit=400,\n        transform=None,\n        pre_transform=None,\n        pre_filter=None,\n        auxiliary = True,\n        min_pulse = -1, #subset the events based on the number of pulses\n        max_pulse = -1,\n\n    ):\n        super().__init__(transform, pre_transform, pre_filter)\n        self.meta_batch  = meta[meta[\"batch_id\"]==batch_id].copy() #meta set of the batches\n        #subset based on pulses\n        if min_pulse > -1:\n            self.meta_batch[\"pulses\"]= self.meta_batch[\"last_pulse_index\"]-self.meta_batch[\"first_pulse_index\"]+1\n            self.meta_batch = self.meta_batch[(self.meta_batch[\"pulses\"]>=min_pulse)&\n                                             (self.meta_batch[\"pulses\"]<max_pulse)]\n            \n        self.event_ids = self.meta_batch[\"event_id\"].tolist() #event ids in a batch\n        #self.event_ids = event_ids\n        self.batch_df = pd.read_parquet(PATH_DATASET + f\"/{mode}/batch_{batch_id}.parquet\")\n        self.auxiliary = auxiliary\n        self.sensor_df = sensor_df\n        self.pulse_limit = pulse_limit\n        self.f_scattering, self.f_absorption = ice_transparency(TRANSPARENCY_PATH)\n        \n\n        self.batch_df[\"time\"] = (self.batch_df[\"time\"] - 1.0e04) / 3.0e4\n        self.batch_df[\"charge\"] = np.log10(self.batch_df[\"charge\"]) / 3.0\n        self.batch_df[\"auxiliary\"] = self.batch_df[\"auxiliary\"].astype(int) - 0.5 # true ->0.5, false ->-0.5\n\n    def len(self):\n        return len(self.event_ids)\n\n    def get(self, idx, train = True):\n        #train = True also retrieves actuals\n        \n        ####indexing directly using meta file for fastering search###\n        event_meta = self.meta_batch.iloc[[idx]]\n        event = self.batch_df.iloc[int(event_meta[\"first_pulse_index\"]):int(event_meta[\"last_pulse_index\"])+1]\n        pulses = (event_meta[\"last_pulse_index\"]-event_meta[\"first_pulse_index\"]+1).values\n        ###indexing using the event_ids slower ###\n#         event_id = self.event_ids[idx]\n#         event = self.batch_df.loc[event_id]\n\n        event = pd.merge(event, self.sensor_df, on=\"sensor_id\")\n\n        #filter out auxiliary data\n        if not self.auxiliary:\n            event = event[event[\"auxiliary\"] < 0]\n            x = event[[\"x\", \"y\", \"z\", \"time\", \"charge\", \"qe\"]].values\n        else:\n            x = event[[\"x\", \"y\", \"z\", \"time\", \"charge\", \"qe\", \"auxiliary\"]].values \n        x = torch.tensor(x, dtype=torch.float32)\n        data = Data(x=x, n_pulses=torch.tensor(x.shape[0], dtype=torch.int32),\n                   real_pulses = torch.tensor(x.shape[0], dtype=torch.int32))\n\n        # Add ice transparency data\n        z = data.x[:, 2].numpy()\n        scattering = torch.tensor(self.f_scattering(z), dtype=torch.float32).view(-1, 1)\n        # absorption = torch.tensor(self.f_absorption(z), dtype=torch.float32).view(-1, 1)\n\n        data.x = torch.cat([data.x, scattering], dim=1)\n\n        #####################################################################################\n        # Downsample the large events, we could find better ways to preprocess this shit\n        #####################################################################################\n        if data.n_pulses > self.pulse_limit:\n            #sort by auxiliary\n            data.x = data.x[data.x[:,6]<0]\n            \n            #sort by time < 6200 ns, barely noticeable effect\n            #data.x = data.x[data.x[:,3]<(data.x[:,3].min()+6300/3.0e4)]\n            \n            data.n_pulses = torch.tensor(data.x.shape[0], dtype=torch.int32)\n            if data.n_pulses > self.pulse_limit:\n                #data.x = data.x[:300]\n                data.x = data.x[np.random.choice(data.n_pulses, self.pulse_limit)]\n                data.n_pulses = torch.tensor(self.pulse_limit, dtype=torch.int32)\n            \n            \n        #get the actuals when training or validating\n        if train:\n            data.y = angle_to_xyz(self.get_actuals(idx))\n        return data\n    \n    def get_actuals(self, idx):\n        return torch.tensor([self.meta_batch.iloc[idx][[\"azimuth\",\"zenith\"]]], dtype=torch.float32)    ","metadata":{"execution":{"iopub.status.busy":"2023-03-24T13:07:07.815177Z","iopub.execute_input":"2023-03-24T13:07:07.815628Z","iopub.status.idle":"2023-03-24T13:07:07.850228Z","shell.execute_reply.started":"2023-03-24T13:07:07.815583Z","shell.execute_reply":"2023-03-24T13:07:07.849194Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#run this block to initiatlize meta file and input data\nfrom pyarrow.parquet import ParquetFile\nimport pyarrow as pa \n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nfrom scipy.interpolate import interp1d\n\nPATH_DATASET = f\"/kaggle/input/icecube-neutrinos-in-deep-ice/\"\nTRANSPARENCY_PATH = f\"/kaggle/input/icecubetransparency/ice_transparency.txt\"\n\n#prepare sensor file, with QE data.\nsensors = prepare_sensors()\n\n#read meta file\n_dtype = {\n    \"batch_id\": \"int16\",\n    \"event_id\": \"int64\",\n    \"first_pulse_index\": \"int64\",\n    \"last_pulse_index\": \"int64\",\n    \"azimuth\": \"float64\",\n    \"zenith\": \"float64\"\n}\n\n# meta = pd.read_parquet(\n#     PATH_DATASET + f\"train_meta.parquet\", columns=[\"batch_id\", \"event_id\",\"azimuth\", \"zenith\"]\n# ).astype(_dtype)\n\npf = ParquetFile(PATH_DATASET + f\"train_meta.parquet\")\nmeta = next(pf.iter_batches(batch_size = 55000000)) \nmeta = pa.Table.from_batches([meta]).to_pandas().astype(_dtype)\n\n\n","metadata":{"execution":{"iopub.status.busy":"2023-03-24T13:07:15.681653Z","iopub.execute_input":"2023-03-24T13:07:15.682816Z","iopub.status.idle":"2023-03-24T13:07:38.652370Z","shell.execute_reply.started":"2023-03-24T13:07:15.682755Z","shell.execute_reply":"2023-03-24T13:07:38.651228Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"meta","metadata":{"execution":{"iopub.status.busy":"2023-03-24T13:07:52.930490Z","iopub.execute_input":"2023-03-24T13:07:52.930946Z","iopub.status.idle":"2023-03-24T13:07:52.966462Z","shell.execute_reply.started":"2023-03-24T13:07:52.930906Z","shell.execute_reply":"2023-03-24T13:07:52.965109Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import csv\nimport os\n\n\nfor batch_id in range(190,200):\n    print(f\"Processing batch: {batch_id}\")\n    dataset = IceCubeSubmissionDataset(batch_id, sensors)\n    \n    gc.collect()\n    new_meta1 = np.zeros([200000,2], dtype = np.int32) #first, last pulse\n    new_meta2 = np.zeros([200000,3]) #x,y,z\n\n    indexer = 0\n    with open('temp.csv', 'w') as f:\n        for i in range(200000):\n            event = dataset[i]\n            start_index = indexer\n            indexer += event.n_pulses.item()\n            end_index = indexer-1\n\n\n            #new_batch_df = np.append(new_batch_df, event.x.detach().numpy(), axis = 0)\n            ########\n            rows = [\"{},{},{},{},{},{},{},{}\".format(n1,n2,n3,n4,n5,n6,n7,n8) for n1,n2,n3,n4,n5,n6,n7,n8\n            in event.x.detach().numpy()]\n            text = \"\\n\".join(rows)+\"\\n\"\n            f.write(text)\n            ################\n\n            new_meta1[i] = np.array([start_index, end_index])\n            new_meta2[i] = event.y.detach().numpy()\n\n            if i% 20000 ==0:\n                print(i, \" complete\")\n    new_meta = pd.DataFrame(new_meta1, columns = ([\"first_pulse_index\",\"last_pulse_index\"]))\n    new_meta[[\"x\",'y','z']]= new_meta2\n    new_meta.to_parquet(f'{batch_id}_meta')\n\n    _dtype2 = {\n        \"x\": \"float32\",\n        \"y\": \"float32\",\n        \"z\": \"float32\",\n        \"time\": \"float32\",\n        \"charge\": \"float32\",\n        \"qe\": \"float32\",\n        \"auxiliary\": \"float32\",\n        \"scattering\": \"float32\",\n    }\n    new_todf = pd.read_csv('temp.csv', header = None, names = [\"x\", \"y\", \"z\", \"time\", \"charge\", \"qe\", \"auxiliary\", \"scattering\"],\n                        dtype = _dtype2)\n    new_todf.to_parquet(f'{batch_id}_processed', index = False)\n    os.remove('temp.csv') ","metadata":{"execution":{"iopub.status.busy":"2023-03-24T13:08:51.085146Z","iopub.execute_input":"2023-03-24T13:08:51.085642Z","iopub.status.idle":"2023-03-24T13:32:47.618462Z","shell.execute_reply.started":"2023-03-24T13:08:51.085597Z","shell.execute_reply":"2023-03-24T13:32:47.616392Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# new_meta = pd.DataFrame(new_meta1, columns = ([\"first_pulse_index\",\"last_pulse_index\"]))\n# new_meta[[\"x\",'y','z']]= new_meta2\n# new_meta.to_parquet(f'{batch_id}_meta')\n\n# _dtype2 = {\n#     \"x\": \"float32\",\n#     \"y\": \"float32\",\n#     \"z\": \"float32\",\n#     \"time\": \"float32\",\n#     \"charge\": \"float32\",\n#     \"qe\": \"float32\",\n#     \"auxiliary\": \"float32\",\n#     \"scattering\": \"float32\",\n# }\n# new_todf = pd.read_csv('temp.csv', header = None, names = [\"x\", \"y\", \"z\", \"time\", \"charge\", \"qe\", \"auxiliary\", \"scattering\"],\n#                     dtype = _dtype2)\n# new_todf.to_parquet(f'{batch_id}_processed', index = False)","metadata":{"execution":{"iopub.status.busy":"2023-03-22T17:08:48.901105Z","iopub.execute_input":"2023-03-22T17:08:48.90315Z","iopub.status.idle":"2023-03-22T17:08:49.068323Z","shell.execute_reply.started":"2023-03-22T17:08:48.903072Z","shell.execute_reply":"2023-03-22T17:08:49.066587Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class IceCubeSubmissionDataset2(Dataset): #reading pre-processed batch version of original\n    def __init__(\n        self,\n        batch_id,\n        sensor_df,\n        datapath = \"\",\n        mode=\"train\",\n        pulse_limit=300,\n        transform=None,\n        pre_transform=None,\n        pre_filter=None,\n        auxiliary = True,\n        min_pulse = -1, #subset the events based on the number of pulses\n        max_pulse = -1,\n\n    ):\n        super().__init__(transform, pre_transform, pre_filter)\n        self.meta_batch  = pd.read_parquet(datapath+f'{batch_id}_meta') #meta set of the batches\n        \n        #subset based on pulses\n        if min_pulse > -1:\n            self.meta_batch[\"pulses\"]= self.meta_batch[\"last_pulse_index\"]-self.meta_batch[\"first_pulse_index\"]+1\n            self.meta_batch = self.meta_batch[(self.meta_batch[\"pulses\"]>=min_pulse)&\n                                             (self.meta_batch[\"pulses\"]<max_pulse)]\n        #self.event_ids = event_ids\n        self.batch_df = pd.read_parquet(datapath+f'{batch_id}_processed')\n    def len(self):\n        return len(self.meta_batch)\n\n    def get(self, idx, train = True):\n        #train = True also retrieves actuals\n        \n        ####indexing directly using meta file for fastering search###\n        event_meta = self.meta_batch.iloc[[idx]]\n        event = self.batch_df.iloc[int(event_meta[\"first_pulse_index\"]):int(event_meta[\"last_pulse_index\"])+1]\n        pulses = (event_meta[\"last_pulse_index\"]-event_meta[\"first_pulse_index\"]+1).values\n        ###indexing using the event_ids slower ###\n#         event_id = self.event_ids[idx]\n#         event = self.batch_df.loc[event_id]\n\n        #x = event[[\"x\", \"y\", \"z\", \"time\", \"charge\", \"qe\", \"auxiliary\"]].values \n        x = torch.tensor(event.values, dtype=torch.float32)\n        data = Data(x=x, n_pulses=torch.tensor(x.shape[0], dtype=torch.int32))\n\n        #get the actuals when training or validating\n        if train:\n            data.y = self.get_actuals(idx)\n        return data\n    \n    def get_actuals(self, idx):\n        return torch.tensor([self.meta_batch.iloc[idx][[\"x\",\"y\",\"z\"]]], dtype=torch.float32)    ","metadata":{"execution":{"iopub.status.busy":"2023-03-22T17:30:23.903337Z","iopub.execute_input":"2023-03-22T17:30:23.903811Z","iopub.status.idle":"2023-03-22T17:30:23.920995Z","shell.execute_reply.started":"2023-03-22T17:30:23.90377Z","shell.execute_reply":"2023-03-22T17:30:23.919558Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#df2 = IceCubeSubmissionDataset2(1, sensors,\"\")","metadata":{"execution":{"iopub.status.busy":"2023-03-22T17:30:50.940072Z","iopub.execute_input":"2023-03-22T17:30:50.940614Z","iopub.status.idle":"2023-03-22T17:30:51.774103Z","shell.execute_reply.started":"2023-03-22T17:30:50.940566Z","shell.execute_reply":"2023-03-22T17:30:51.771929Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# t1 = time()\n# for i in range(10000):\n#     df2[i]\n# print(time()-t1)","metadata":{"execution":{"iopub.status.busy":"2023-03-22T17:36:17.81255Z","iopub.execute_input":"2023-03-22T17:36:17.812951Z","iopub.status.idle":"2023-03-22T17:36:29.383784Z","shell.execute_reply.started":"2023-03-22T17:36:17.812916Z","shell.execute_reply":"2023-03-22T17:36:29.382154Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# t1 = time()\n# for i in range(10000):\n#     dataset[i]\n# print(time()-t1)","metadata":{"execution":{"iopub.status.busy":"2023-03-22T17:36:29.386329Z","iopub.execute_input":"2023-03-22T17:36:29.386766Z","iopub.status.idle":"2023-03-22T17:37:19.175859Z","shell.execute_reply.started":"2023-03-22T17:36:29.386728Z","shell.execute_reply":"2023-03-22T17:37:19.174328Z"},"trusted":true},"execution_count":null,"outputs":[]}]}