{"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":"!mkdir -p /tmp/pip/cache/\n!cp /kaggle/input/wpca-01-py3-whl/wpca-0.1-py3-none-any.whl /tmp/pip/cache/\n!pip install --no-index --find-links /tmp/pip/cache/ wpca","metadata":{"execution":{"iopub.status.busy":"2023-04-17T12:48:36.065051Z","iopub.execute_input":"2023-04-17T12:48:36.065443Z","iopub.status.idle":"2023-04-17T12:48:51.219616Z","shell.execute_reply.started":"2023-04-17T12:48:36.065409Z","shell.execute_reply":"2023-04-17T12:48:51.217967Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport pytorch_lightning as pl\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nimport pandas as pd\nimport numpy as np\nfrom sklearn.model_selection import train_test_split\nfrom wpca import WPCA\nfrom sklearn.decomposition import PCA\nimport gc\nimport polars","metadata":{"execution":{"iopub.status.busy":"2023-04-17T12:48:51.225007Z","iopub.execute_input":"2023-04-17T12:48:51.225538Z","iopub.status.idle":"2023-04-17T12:49:00.277885Z","shell.execute_reply.started":"2023-04-17T12:48:51.225477Z","shell.execute_reply":"2023-04-17T12:49:00.276729Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def normalize_targets(y):\n    y_min = np.array([0.0, 0.0], dtype=np.float32)  # azimuth: [0, 2*pi], zenith: [0, pi]\n    y_max = np.array([2 * np.pi, np.pi], dtype=np.float32)\n    return (y - y_min) / (y_max - y_min)\n\ndef denormalize_targets(y_norm):\n    y_min = np.array([0.0, 0.0], dtype=np.float32)\n    y_max = np.array([2 * np.pi, np.pi], dtype=np.float32)\n    return y_norm * (y_max - y_min) + y_min\n\ndef normalize_features(df, features):\n    df = df.copy()\n    cols_min = df[features].min()\n    cols_max = df[features].max()\n    df[features] = (df[features] - cols_min) / (cols_max - cols_min)\n    return df","metadata":{"execution":{"iopub.status.busy":"2023-04-17T12:49:00.279500Z","iopub.execute_input":"2023-04-17T12:49:00.279889Z","iopub.status.idle":"2023-04-17T12:49:00.287784Z","shell.execute_reply.started":"2023-04-17T12:49:00.279854Z","shell.execute_reply":"2023-04-17T12:49:00.286773Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# TO TEST DATASET\ndata_dir = '/kaggle/input/icecube-neutrinos-in-deep-ice/'\nmetadata_file = f'{data_dir}/train_meta.parquet'\nbatch_files = [f'{data_dir}train/batch_{i}.parquet' for i in range(1, 2)] # You can set up to 660 files\n\nsensor_geometry = pd.read_csv(os.path.join(data_dir, 'sensor_geometry.csv'))\nmetadata = pd.read_parquet(metadata_file)","metadata":{"execution":{"iopub.status.busy":"2023-04-17T12:49:00.289419Z","iopub.execute_input":"2023-04-17T12:49:00.290020Z","iopub.status.idle":"2023-04-17T12:49:41.160097Z","shell.execute_reply.started":"2023-04-17T12:49:00.289976Z","shell.execute_reply":"2023-04-17T12:49:41.159069Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def generating_some_features(batch_file, metadata, sensor_geometry):\n    \n    \n    def join_tables_all(df_meta, df_batch, df_sensor):\n        return df_meta.join(df_batch, on='event_id').join(df_sensor, on='sensor_id').with_columns([\n            (polars.col('time') - polars.col('time').min()).over('event_id')\n        ])\n\n    def generate_features_grouped(dataf):\n        return dataf.groupby('event_id').agg([\n        polars.col('x').mean().alias('x_mean'),\n        polars.col('x').median().alias('x_median'),\n        polars.col('y').mean().alias('y_mean'),\n        polars.col('y').median().alias('y_median'),\n        polars.col('z').mean().alias('z_mean'),\n        polars.col('z').median().alias('z_median'),    \n        polars.col('time').mean().alias('event_mean_time'),\n        polars.col('time').max().alias('event_max_time'),\n        polars.col('charge').min().alias('event_min_charge'),\n        polars.col('charge').mean().alias('event_mean_charge'),\n        polars.col('charge').max().alias('event_max_charge'),\n        polars.col('charge').count().alias('overall_count'),\n        polars.col('auxiliary').sum().alias('overall_aux_sum'),\n        polars.col('charge').sum().alias('sum_charge'),\n        (polars.col('auxiliary').sum() / polars.col('auxiliary').count()).alias('aux_ratio'),\n        polars.col('sensor_id').n_unique().alias('sensor_count'),\n    ])\n    \n    def add_ranks(dataf):\n        return dataf.with_columns(\n[\n    polars.col('time').rank('ordinal').over('event_id').alias('time_rank_asc'),\n    polars.col('time').rank('ordinal', descending=True).over('event_id').alias('time_rank_des'),\n    polars.col('charge').rank('ordinal').over('event_id').alias('charge_rank_asc'),\n    polars.col('charge').rank('ordinal').over('event_id').alias('charge_rank_des')\n])\n\n    def make_geometrical_features(dataf):\n        geometrical_features = dataf.select('event_id').unique()\n        for direction in ['time_rank_asc','time_rank_des', 'charge_rank_asc', 'charge_rank_des']:\n            for direction_axis in ['x', 'y', 'z']:    \n                temp_col_1 = dataf.filter(polars.col(direction) == 1).select([\n                    polars.col('event_id'),\n                    polars.col(direction_axis).over('event_id')\n                ]).with_columns([\n                    polars.col(direction_axis).alias(direction_axis+'_'+direction+'_1')]\n                ).select(polars.col('event_id'), polars.col(direction_axis+'_'+direction+'_1'))\n\n                temp_col_2 = dataf.filter(polars.col(\"time_rank_asc\") == 2).select([\n                    polars.col('event_id'),\n                    polars.col(direction_axis).over('event_id')\n                ]).with_columns([\n                    polars.col(direction_axis).alias(direction_axis+'_'+direction+'_2')]\n                ).select(polars.col('event_id'), polars.col(direction_axis+'_'+direction+'_2'))\n\n                temp_col_3 = dataf.filter(polars.col(\"time_rank_asc\") == 3).select([\n                    polars.col('event_id'),\n                    polars.col(direction_axis).over('event_id')\n                ]).with_columns([\n                    polars.col(direction_axis).alias(direction_axis+'_'+direction+'_3')]\n                ).select(polars.col('event_id'), polars.col(direction_axis+'_'+direction+'_3'))\n\n                geometrical_features = geometrical_features.join(temp_col_1, on='event_id', how='left'\n                               ).join(temp_col_2, on='event_id', how='left'\n                               ).join(temp_col_3, on='event_id', how='left'\n                               )\n        return geometrical_features.fill_null(1000)\n\n    train_batch = polars.scan_parquet(batch_file).lazy()\n    df_train_meta = polars.DataFrame(metadata).lazy()\n    df_sensor_geometry = polars.DataFrame(sensor_geometry).with_columns(polars.col('sensor_id').cast(polars.Int16)).lazy()\n    \n        #Not accounting for aux\n    features_grouped_metrics = df_train_meta.pipe(join_tables_all, train_batch, df_sensor_geometry\n                      ).pipe(generate_features_grouped).collect()\n\n    geometrical_features = df_train_meta.pipe(join_tables_all, train_batch, df_sensor_geometry\n                      ).pipe(add_ranks\n                      ).collect().pipe(make_geometrical_features)\n\n    temp_1 = features_grouped_metrics.join(geometrical_features, on='event_id', how='left')\n\n\n    #AUX = FALSE\n\n    features_grouped_metrics = df_train_meta.pipe(join_tables_all, train_batch, df_sensor_geometry\n                      ).filter(polars.col('auxiliary') == False).pipe(generate_features_grouped).collect()\n\n    geometrical_features = df_train_meta.pipe(join_tables_all, train_batch, df_sensor_geometry\n                      ).filter(polars.col('auxiliary') == False).pipe(add_ranks\n                      ).collect().pipe(make_geometrical_features)\n\n    temp_2 = features_grouped_metrics.join(geometrical_features, on='event_id', how='left')\n    \n    temp_3 = temp_1.join(temp_2, on = 'event_id', how='left').fill_null(0)\n    del temp_1, temp_2, features_grouped_metrics, geometrical_features\n    \n    temp_3 = temp_3.to_pandas().set_index('event_id')\n    \n    temp_3 = (temp_3-temp_3.mean())/temp_3.std()\n    \n    return temp_3","metadata":{"execution":{"iopub.status.busy":"2023-04-17T12:49:41.163062Z","iopub.execute_input":"2023-04-17T12:49:41.163704Z","iopub.status.idle":"2023-04-17T12:49:41.190102Z","shell.execute_reply.started":"2023-04-17T12:49:41.163667Z","shell.execute_reply":"2023-04-17T12:49:41.188516Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class NeutrinoDataset(Dataset):\n    def __init__(self, metadata, sensor_geometry, batch_file, mode=\"train\"):\n        self.metadata = metadata\n        self.sensor_geometry = normalize_features(sensor_geometry, ['x', 'y', 'z'])       \n        self.batch_data = pd.read_parquet(batch_file)       \n        self.mode = mode\n        self.get_time_valid_length()\n\n        self.batch_features = generating_some_features(batch_file, metadata, sensor_geometry)\n        \n    def __len__(self):\n        return len(self.metadata)\n\n    \n    def __getitem__(self, idx):\n        event = self.metadata.iloc[idx]\n        event_data = self.batch_data.loc[event.event_id]\n        if len(event_data) > 1010:\n            event_data = event_data.sample(1000)\n        features = self.combine_features(event_data, idx)\n        \n        if self.mode == \"train\":\n            y = np.array([event['azimuth'], event['zenith']], dtype=np.float32)\n            y_norm = normalize_targets(y)\n            return features, y_norm\n        else:\n            return event.event_id, features\n        \n\n    # Return dataframe with PCA and WPCA predicted angles (using all and non-aux pulses) for single event\n    def process_event_pca(self, event_pulses):\n        event_pulses = event_pulses.reset_index(drop=True)\n        event_pulses['auxiliary'] = event_pulses['auxiliary'] * 1\n        event_pulses[\"time_from_start\"] = event_pulses.time.values - event_pulses.time.min()\n        time_peak_charge = event_pulses.loc[event_pulses['charge'].argmax(), 'time_from_start']\n        time_valid_min = time_peak_charge - self.time_valid_length\n        time_valid_max = time_peak_charge + self.time_valid_length\n        time_valid = ((event_pulses['time_from_start'] > time_valid_min) * \n                      (event_pulses['time_from_start'] < time_valid_max))\n        event_pulses['pulse_rank'] = 2 * (1 - event_pulses['auxiliary'].values) + (time_valid)\n        event_pulses = event_pulses.drop('time_from_start', axis=1)\n\n        event_pulses[['x', 'y', 'z']] = self.sensor_geometry.loc[event_pulses.sensor_id].reset_index()[['x', 'y', 'z']]\n        \n        event_pulses.x = event_pulses.x - event_pulses.x.agg('mean')\n        event_pulses.y = event_pulses.y - event_pulses.y.agg('mean')\n        event_pulses.z = event_pulses.z - event_pulses.z.agg('mean')\n        \n\n        result = []\n        result.extend(self.get_angles(event_pulses, 'PCA', drop_aux=True))\n        result.extend(self.get_angles(event_pulses, 'PCA', drop_aux=False))\n        result.extend(self.get_angles(event_pulses, 'WPCA', drop_aux=True))\n        result.extend(self.get_angles(event_pulses, 'WPCA', drop_aux=False))\n        result.extend(self.get_angles(event_pulses, 'WPCA_new_weights', drop_aux=True))\n        result.extend(self.get_angles(event_pulses, 'WPCA_new_weights', drop_aux=False))\n\n        return np.array(result)\n    \n    \n    def get_time_valid_length(self):\n        c_const = 0.299792458  # speed of light [m/ns]\n\n        x_min, x_max = self.sensor_geometry.x.agg([min, max]).values\n        y_min, y_max = self.sensor_geometry.y.agg([min, max]).values\n        z_min, z_max = self.sensor_geometry.z.agg([min, max]).values\n\n        detector_length = np.sqrt((x_max - x_min)**2 + (y_max - y_min)**2 + (z_max - z_min)**2)\n        self.time_valid_length = detector_length / c_const\n    \n    \n    def get_weight_all(self, pulse_rank):\n        if pulse_rank == 2:\n            return 1\n        if pulse_rank == 3:\n            return 5\n        else:\n            return 0.001\n    \n    \n    # Calculate target angles for single event\n    def get_angles(self, event_pulses, ThisPCA='WPCA', drop_aux=False):\n        num_of_nonaux_pulses = event_pulses.auxiliary.count() - event_pulses.auxiliary.sum()\n        \n        # Computing the weighted PCA\n        if ThisPCA == 'WPCA':\n            ThisPCA = WPCA\n            #if it is necessary to throw out auxiliary pulses\n            if drop_aux and num_of_nonaux_pulses > 5:\n                event_pulses = event_pulses.loc[event_pulses.auxiliary == 0].copy()\n                weights = np.stack([event_pulses.pulse_rank, event_pulses.pulse_rank, \\\n                                   event_pulses.pulse_rank], axis=1)\n            else:\n                weights = np.stack([event_pulses.pulse_rank, event_pulses.pulse_rank, \\\n                                   event_pulses.pulse_rank], axis=1)\n            kwds = {'weights': weights}\n            \n        elif ThisPCA == 'WPCA_new_weights':\n            ThisPCA = WPCA\n            event_pulses['weights_all'] = event_pulses['pulse_rank'].apply(self.get_weight_all)\n            #if it is necessary to throw out auxiliary pulses\n            if drop_aux and num_of_nonaux_pulses > 5:\n                event_pulses = event_pulses.loc[event_pulses.auxiliary == 0].copy()\n                weights = np.stack([event_pulses.weights_all, event_pulses.weights_all, \\\n                                   event_pulses.weights_all], axis=1)\n            else:\n                weights = np.stack([event_pulses.weights_all, event_pulses.weights_all, \\\n                                   event_pulses.weights_all], axis=1)\n            kwds = {'weights': weights}\n\n        # Computing the stardard PCA\n        else:\n            ThisPCA = PCA\n            kwds = {}\n            if drop_aux and num_of_nonaux_pulses > 5:\n                event_pulses = event_pulses.loc[event_pulses.auxiliary == 0].copy()\n\n        # Computing the PCA vectors\n        X = event_pulses.loc[:, ['x', 'y', 'z']]\n        pca = ThisPCA().fit(X, **kwds)\n            \n        if len(pca.components_) > 2:\n            v_1, v_2, v_3, v_4, v_5, v_6, v_7, v_8, v_9 =  map(float, np.array(pca.components_[:3]).reshape(-1, 1))\n        else:\n            v_1, v_2, v_3, v_4, v_5, v_6, v_7, v_8, v_9 =  np.array([0., 0., 0., 0., 0., 0., 0., 0., 0.])\n                \n        # Calculating the first principal component vector\n        vector = pca.components_[0]\n\n        # Calculating a distance of a point from a plane\n        event_pulses['distance'] = (event_pulses.x * vector[0] +\n                                        event_pulses.y * vector[1] +\n                                        event_pulses.z * vector[2]) / np.linalg.norm(vector)\n\n        # Flip the vector direction if it points away from the neutrino origin\n        if (event_pulses.loc[event_pulses.distance > 0].time.mean() >\n                event_pulses.loc[event_pulses.distance < 0].time.mean()):\n            vector = -1 * np.array(vector)\n\n        vector = np.clip(vector, -1, 1)\n\n        # Calculating the angles\n        zenith = np.arccos(vector[2])\n        azimuth = np.arctan2(vector[1], vector[0])\n        if azimuth < 0:\n            azimuth = 2 * np.pi + azimuth\n\n        azimuth = azimuth / (2 * np.pi)\n        zenith = zenith / np.pi\n\n        return float(zenith), float(azimuth), v_1, v_2, v_3, v_4, v_5, v_6, v_7, v_8, v_9\n\n    \n    def combine_features(self, event_data, idx):\n        pca_features = self.process_event_pca(event_data)\n        batch_features = self.batch_features.loc[self.metadata.iloc[idx]['event_id']].fillna(0).values\n\n        combined_features = np.hstack([\n            pca_features,\n            batch_features\n        ])\n        assert np.all(~np.isnan(combined_features))\n        return combined_features.astype(np.float32)\n","metadata":{"execution":{"iopub.status.busy":"2023-04-17T12:49:41.192239Z","iopub.execute_input":"2023-04-17T12:49:41.192702Z","iopub.status.idle":"2023-04-17T12:49:41.230666Z","shell.execute_reply.started":"2023-04-17T12:49:41.192641Z","shell.execute_reply":"2023-04-17T12:49:41.228486Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def angular_dist_score(az_true, zen_true, az_pred, zen_pred):\n    '''\n    calculate the MAE of the angular distance between two directions.\n    The two vectors are first converted to cartesian unit vectors,\n    and then their scalar product is computed, which is equal to\n    the cosine of the angle between the two vectors. The inverse \n    cosine (arccos) thereof is then the angle between the two input vectors\n    \n    Parameters:\n    -----------\n    \n    az_true : float (or array thereof)\n        true azimuth value(s) in radian\n    zen_true : float (or array thereof)\n        true zenith value(s) in radian\n    az_pred : float (or array thereof)\n        predicted azimuth value(s) in radian\n    zen_pred : float (or array thereof)\n        predicted zenith value(s) in radian\n    \n    Returns:\n    --------\n    \n    dist : float\n        mean over the angular distance(s) in radian\n    '''\n    \n    if not (np.all(np.isfinite(az_true)) and\n            np.all(np.isfinite(zen_true)) and\n            np.all(np.isfinite(az_pred)) and\n            np.all(np.isfinite(zen_pred))):\n        raise ValueError(\"All arguments must be finite\")\n    \n    # pre-compute all sine and cosine values\n    sa1 = np.sin(az_true)\n    ca1 = np.cos(az_true)\n    sz1 = np.sin(zen_true)\n    cz1 = np.cos(zen_true)\n    \n    sa2 = np.sin(az_pred)\n    ca2 = np.cos(az_pred)\n    sz2 = np.sin(zen_pred)\n    cz2 = np.cos(zen_pred)\n    \n    # scalar product of the two cartesian vectors (x = sz*ca, y = sz*sa, z = cz)\n    scalar_prod = sz1*sz2*(ca1*ca2 + sa1*sa2) + (cz1*cz2)\n    \n    # scalar product of two unit vectors is always between -1 and 1, this is against nummerical instability\n    # that might otherwise occure from the finite precision of the sine and cosine functions\n    scalar_prod =  np.clip(scalar_prod, -1, 1)\n    \n    # convert back to an angle (in radian)\n    return np.average(np.abs(np.arccos(scalar_prod)))","metadata":{"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2023-04-17T12:49:41.231957Z","iopub.execute_input":"2023-04-17T12:49:41.232304Z","iopub.status.idle":"2023-04-17T12:49:41.244724Z","shell.execute_reply.started":"2023-04-17T12:49:41.232267Z","shell.execute_reply":"2023-04-17T12:49:41.243793Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport pytorch_lightning as pl\n\nclass NeutrinoModel(pl.LightningModule):\n    def __init__(self, input_dims, num_layers=2, hidden_units=None, learning_rate=0.001):\n        super().__init__()\n        self.save_hyperparameters()\n\n        layers = []\n        layers.append(nn.BatchNorm1d(input_dims))\n\n        if hidden_units is None:\n            hidden_units = [input_dims*2] * (num_layers - 1)\n\n        for i in range(num_layers):\n            in_dim = input_dims if i == 0 else hidden_units\n            out_dim = 2 if i == num_layers - 1 else hidden_units\n            layers.append(nn.Linear(in_dim, out_dim))\n\n            if i < num_layers - 1:\n                layers.append(nn.ReLU())\n            else:\n                layers.append(nn.Sigmoid())\n\n        self.layers = nn.Sequential(*layers)\n\n    def forward(self, x):\n        return self.layers(x)\n\n    def training_step(self, batch, batch_idx):\n        x, y = batch\n        y_hat = self(x)\n        loss = F.l1_loss(y_hat, y, reduction='mean')\n        self.log(\"train_loss\", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)\n        return loss\n\n    def validation_step(self, batch, batch_idx):\n        x, y_norm = batch\n        y_hat_norm = self(x)\n        loss = F.l1_loss(y_hat_norm, y_norm, reduction='mean')\n\n        # Denormalize target values and predictions\n        y = denormalize_targets(y_norm.cpu().numpy())\n        y_hat = denormalize_targets(y_hat_norm.detach().cpu().numpy())\n\n        # Calculate angular distance score\n        az_true, zen_true = y[:, 0], y[:, 1]\n        az_pred, zen_pred = y_hat[:, 0], y_hat[:, 1]\n        ang_dist = angular_dist_score(az_true, zen_true, az_pred, zen_pred)\n\n        self.log(\"val_loss\", loss, on_step=False, on_epoch=True, prog_bar=True, logger=True)\n        self.log(\"angular_dist_score\", ang_dist, on_step=False, on_epoch=True, prog_bar=True, logger=True)\n\n        return loss\n\n    def configure_optimizers(self):\n        return torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)\n","metadata":{"execution":{"iopub.status.busy":"2023-04-17T12:49:41.246299Z","iopub.execute_input":"2023-04-17T12:49:41.247079Z","iopub.status.idle":"2023-04-17T12:49:41.263407Z","shell.execute_reply.started":"2023-04-17T12:49:41.247014Z","shell.execute_reply":"2023-04-17T12:49:41.262417Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_one_epoch_on_batch(model, metadata, sensor_geometry, batch_file, val_split=0.2):\n    batch_id = int(batch_file.split('_')[-1].split('.')[0])\n    batch_metadata = metadata[metadata['batch_id'] == batch_id]\n    \n    train_metadata, val_metadata = train_test_split(batch_metadata, test_size=val_split, random_state=42)\n    \n    train_dataset = NeutrinoDataset(train_metadata, sensor_geometry, batch_file, mode=\"train\")\n    train_dataloader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=4)\n    \n    val_dataset = NeutrinoDataset(val_metadata, sensor_geometry, batch_file, mode=\"train\")\n    val_dataloader = DataLoader(val_dataset, batch_size=64, shuffle=False, num_workers=4)\n\n    trainer = pl.Trainer(max_epochs=1)\n    trainer.fit(model, train_dataloader, val_dataloader)\n    \n    del train_dataset, train_dataloader, val_dataset, val_dataloader\n","metadata":{"execution":{"iopub.status.busy":"2023-04-17T12:49:41.265128Z","iopub.execute_input":"2023-04-17T12:49:41.265867Z","iopub.status.idle":"2023-04-17T12:49:41.281026Z","shell.execute_reply.started":"2023-04-17T12:49:41.265821Z","shell.execute_reply":"2023-04-17T12:49:41.279877Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"data_dir = '/kaggle/input/icecube-neutrinos-in-deep-ice/'\nmetadata_file = f'{data_dir}/train_meta.parquet'\nbatch_files = [f'{data_dir}train/batch_{i}.parquet' for i in range(1, 6)] # You can set up to 660 files\n\nsensor_geometry = pd.read_csv(os.path.join(data_dir, 'sensor_geometry.csv'))\nmetadata = pd.read_parquet(metadata_file)","metadata":{"execution":{"iopub.status.busy":"2023-04-17T12:49:41.282504Z","iopub.execute_input":"2023-04-17T12:49:41.283655Z","iopub.status.idle":"2023-04-17T12:49:52.637451Z","shell.execute_reply.started":"2023-04-17T12:49:41.283606Z","shell.execute_reply":"2023-04-17T12:49:52.636289Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"config = {\n    \"input_dims\": 170,\n    \"num_layers\": 4,\n    \"hidden_units\": 512,\n    \"learning_rate\": 0.001,\n}\n\nmodel = NeutrinoModel(\n    input_dims=config['input_dims'],\n    num_layers=config['num_layers'],\n    hidden_units=config['hidden_units'],\n    learning_rate=config['learning_rate']\n)","metadata":{"execution":{"iopub.status.busy":"2023-04-17T12:49:52.639233Z","iopub.execute_input":"2023-04-17T12:49:52.639555Z","iopub.status.idle":"2023-04-17T12:49:52.695013Z","shell.execute_reply.started":"2023-04-17T12:49:52.639524Z","shell.execute_reply":"2023-04-17T12:49:52.693541Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for epoch in range(1): # You can set any number of epochs but first use all of 660 files\n    for batch_file in batch_files:\n        train_one_epoch_on_batch(model, metadata, sensor_geometry, batch_file)","metadata":{"execution":{"iopub.status.busy":"2023-04-17T12:49:52.696438Z","iopub.execute_input":"2023-04-17T12:49:52.696809Z","iopub.status.idle":"2023-04-17T14:46:08.766622Z","shell.execute_reply.started":"2023-04-17T12:49:52.696774Z","shell.execute_reply":"2023-04-17T14:46:08.764537Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def predict(model, metadata, sensor_geometry, batch_files, output_file=\"submission.csv\"):\n    model.eval()\n    \n    with open(output_file, 'w') as f:\n        f.write('event_id,azimuth,zenith\\n')\n\n    for batch_file in batch_files:\n        batch_id = int(batch_file.split('_')[-1].split('.')[0])\n        batch_metadata = metadata[metadata['batch_id'] == batch_id]\n        test_dataset = NeutrinoDataset(batch_metadata, sensor_geometry, batch_file, mode=\"test\")\n        test_dataloader = DataLoader(test_dataset, batch_size=64, num_workers=4)\n        \n        for event_ids, features in test_dataloader:\n            with torch.no_grad():\n                pred_norm = model(features)\n                pred = denormalize_targets(pred_norm.detach().numpy())\n                \n                with open(output_file, 'a') as f:\n                    for event_id, prediction in zip(event_ids, pred):\n                        f.write(f'{event_id},{prediction[0]},{prediction[1]}\\n')","metadata":{"execution":{"iopub.status.busy":"2023-04-17T14:46:08.771469Z","iopub.execute_input":"2023-04-17T14:46:08.772069Z","iopub.status.idle":"2023-04-17T14:46:08.791471Z","shell.execute_reply.started":"2023-04-17T14:46:08.771976Z","shell.execute_reply":"2023-04-17T14:46:08.790064Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_metadata_file = f'{data_dir}/test_meta.parquet'\ntest_metadata = pd.read_parquet(test_metadata_file)\n\n\ntest_dir = '/kaggle/input/icecube-neutrinos-in-deep-ice/test/'\ntest_batch_files = [f'{test_dir}/{file}' for file in os.listdir(test_dir)]\n\npredict(model, test_metadata, sensor_geometry, test_batch_files, \"submission.csv\")","metadata":{"execution":{"iopub.status.busy":"2023-04-17T14:46:08.800732Z","iopub.execute_input":"2023-04-17T14:46:08.801196Z","iopub.status.idle":"2023-04-17T14:46:10.220509Z","shell.execute_reply.started":"2023-04-17T14:46:08.801155Z","shell.execute_reply":"2023-04-17T14:46:10.218256Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pd.read_csv('submission.csv')","metadata":{"execution":{"iopub.status.busy":"2023-04-17T14:46:10.223139Z","iopub.execute_input":"2023-04-17T14:46:10.223646Z","iopub.status.idle":"2023-04-17T14:46:10.276356Z","shell.execute_reply.started":"2023-04-17T14:46:10.223558Z","shell.execute_reply":"2023-04-17T14:46:10.274605Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}