{"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":"!pip3 install -U polars\nimport polars as pl\nimport numpy as np\nimport pandas as pd\nimport sklearn.metrics as metrics\nfrom tqdm import tqdm\nimport polars as pl\n\n\ndata_path = \"/kaggle/input/icecube-neutrinos-in-deep-ice\"","metadata":{"execution":{"iopub.status.busy":"2023-02-08T15:35:02.732428Z","iopub.execute_input":"2023-02-08T15:35:02.732813Z","iopub.status.idle":"2023-02-08T15:35:18.217264Z","shell.execute_reply.started":"2023-02-08T15:35:02.732780Z","shell.execute_reply":"2023-02-08T15:35:18.216027Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def compute_angle(x, y, z):\n    \"\"\"\n    Compute Azimuth and Zenith from given XYZ.\n    Source: https://www.kaggle.com/code/gyozzza/icecube-3d-linear-regression\n\n    :param x:\n    :param y:\n    :param z:\n    :return:\n    \"\"\"\n    covx = np.cov(z, x)\n    mx = covx[0,1]/covx[0,0]\n    covy = np.cov(z, y)\n    my = covy[0,1]/covy[0,0]\n    zr = 1 # downward-going\n    #zr = -1 # upward-going\n    xr = mx*zr\n    yr = my*zr\n    r = np.sqrt(zr**2+xr**2+yr**2)\n\n    azimuth = np.arctan2(yr, xr)\n    if azimuth < 0: azimuth = 2*np.pi + azimuth\n\n    zenith = np.arccos(zr/r)\n\n    return azimuth, zenith\n","metadata":{"_uuid":"e4077ee6-8db2-439f-8b83-0aaa10255551","_cell_guid":"fb9e45f8-2957-48cc-a862-72070e54e192","collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2023-02-08T15:35:18.219603Z","iopub.execute_input":"2023-02-08T15:35:18.220067Z","iopub.status.idle":"2023-02-08T15:35:18.229433Z","shell.execute_reply.started":"2023-02-08T15:35:18.220018Z","shell.execute_reply":"2023-02-08T15:35:18.228007Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Dataset:\n    def __init__(self):\n        #self.sensor_geometry = pd.read_csv(f\"{data}/sensor_geometry.csv\", index_col='sensor_id')\n        #self.train_meta = pd.read_parquet(f'{data}/train_meta.parquet')\n        #self.submission_df = pd.read_parquet(f'{data}/sample_submission.parquet')\n        #self.submission_df = self.submission_df.set_index('event_id', drop=True)\n        self.sensor_geometry = pl.read_csv(f\"{data_path}/sensor_geometry.csv\")        \n        self.train_meta = pl.read_parquet(f'{data_path}/train_meta.parquet')\n        self.submission_df = pl.read_parquet(f'{data_path}/sample_submission.parquet')\n         \n        self.sensor_geometry = self.sensor_geometry.with_columns(pl.col('sensor_id').cast(pl.Int16))\n        \n\n\n    def __len__(self):\n        return len(self.train_meta.batch_id.unique())\n\n    def __getitem__(self, batch_id):\n        # item = self.train_meta.iloc[idx]\n        # batch_id = int(item[\"batch_id\"])\n        # event_id = int(item[\"event_id\"])\n        # first_pulse_index = int(item[\"first_pulse_index\"])\n        # last_pulse_index = int(item[\"last_pulse_index\"])\n        # azimuth = item[\"azimuth\"]\n        # zenith = item[\"zenith\"]\n\n        #batch_df = self.train_meta[self.train_meta.batch_id == batch_id].set_index('event_id', drop=True)\n        batch_df = self.train_meta.filter(pl.col(\"batch_id\") == batch_id)\n        \n        #batch_features = pd.read_parquet(f'{data}/train/batch_{batch_id}.parquet')\n        batch_features = pl.read_parquet(f'{data_path}/train/batch_{batch_id}.parquet')\n\n        batch_data = []\n        \n        #for event_id in tqdm(batch_df.index, desc=\"Loading\", leave=False):\n        unique_event_ids = batch_df.select('event_id').unique().to_series()\n        for event_id in tqdm(unique_event_ids, desc=\"Loading\", leave=False):\n            \n            # Get features for the current event by the first and last pulse index\n            #event_features = batch_features.iloc[batch_df.loc[event_id, 'first_pulse_index']:\n            #                                     batch_df.loc[event_id, 'last_pulse_index'] + 1]\n            event_features = batch_features[batch_df.filter(pl.col(\"event_id\") == event_id)[\"first_pulse_index\"][0] \n                                            : batch_df.filter(pl.col(\"event_id\") == event_id)[\"last_pulse_index\"][0]]\n                                                 \n            \n            \n            #position = self.sensor_geometry.loc[event_features.sensor_id].values\n            #time = event_features.time.values\n            #charge = event_features.charge.values\n            #auxiliary = event_features.auxiliary.values\n            \n            combined = event_features.join(self.sensor_geometry, on=\"sensor_id\", how = \"left\")\n            position = combined.select(['x','y','z'])\n            time = combined.select('time')\n            charge = combined.select('charge')\n            auxiliary = combined.select('auxiliary')\n\n            x = position[:, 0]\n            y = position[:, 1]\n            z = position[:, 2]\n            \n            #azimuth = self.train_meta.loc[event_id, \"azimuth\"]\n            #zenith = self.train_meta.loc[event_id, \"zenith\"]\n            azimuth = self.train_meta.filter(pl.col(\"event_id\") == event_id)[\"azimuth\"][0]\n            zenith = self.train_meta.filter(pl.col(\"event_id\") == event_id)[\"zenith\"][0]\n\n            batch_data.append([x.to_list(), y.to_list(), z.to_list(), azimuth, zenith])\n\n        return batch_data\n\n    def get_batch_ids(self):\n        return self.train_meta.select(\"batch_id\").unique()\n\n\ntraining_dataset = Dataset()","metadata":{"execution":{"iopub.status.busy":"2023-02-08T15:35:18.231259Z","iopub.execute_input":"2023-02-08T15:35:18.231723Z","iopub.status.idle":"2023-02-08T15:35:41.535276Z","shell.execute_reply.started":"2023-02-08T15:35:18.231661Z","shell.execute_reply":"2023-02-08T15:35:41.534193Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"maes = []\npbar = tqdm(training_dataset.get_batch_ids()[:5].to_series(), desc=\"Main\")\n\nb = None\n\nfor i in pbar:\n    batch = training_dataset[i]\n\n    buff = []\n    for data in tqdm(batch, desc=\"Metrics\", leave=False):\n        pred_azimuth, pred_zenith = compute_angle(data[0], data[1], data[2])\n        buff.append([data[3], data[4], pred_azimuth, pred_zenith])\n    b = buff\n    mae = metrics.mean_absolute_error((buff[0], buff[1]), (buff[2], buff[3]))\n    maes.append(mae)\n\n    pbar.set_postfix({\"MAE\": mae})\n\n\nprint(maes)\n","metadata":{"execution":{"iopub.status.busy":"2023-02-08T15:45:39.878760Z","iopub.execute_input":"2023-02-08T15:45:39.879251Z","iopub.status.idle":"2023-02-08T15:47:56.858396Z","shell.execute_reply.started":"2023-02-08T15:45:39.879209Z","shell.execute_reply":"2023-02-08T15:47:56.856922Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}