{"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":"# Move software to working disk\n!scp -r /kaggle/input/graphnet-and-dependencies/software .\n\n# Install dependencies\n!pip install /kaggle/working/software/dependencies/torch-1.11.0+cu115-cp37-cp37m-linux_x86_64.whl\n!pip install /kaggle/input/icecube-submission-dynemb/torchvision-0.12.0cu115-cp37-cp37m-linux_x86_64.whl\n!pip install /kaggle/working/software/dependencies/torch_cluster-1.6.0-cp37-cp37m-linux_x86_64.whl\n!pip install /kaggle/working/software/dependencies/torch_scatter-2.0.9-cp37-cp37m-linux_x86_64.whl\n!pip install /kaggle/working/software/dependencies/torch_sparse-0.6.13-cp37-cp37m-linux_x86_64.whl\n!pip install /kaggle/working/software/dependencies/torch_geometric-2.0.4.tar.gz\n\n# Install GraphNeT\n!cd software/graphnet; pip install --no-index --find-links=\"/kaggle/working/software/dependencies\" -e .[torch]\n\n# Append to PATH\nimport sys\nsys.path.append('/kaggle/working/software/graphnet/src')","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import glob\nfrom typing import List, Tuple, Union, Optional\nimport joblib\nimport numpy as np\nimport pandas as pd\nfrom pytorch_lightning import LightningModule\nimport torch\nimport torch.nn as nn\nfrom torch import Tensor, LongTensor\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch_geometric.data import Data, Batch\nfrom torch_scatter import scatter_max, scatter_mean, scatter_min, scatter_sum\nfrom tqdm import tqdm\nfrom graphnet.models.components.layers import DynEdgeConv\nfrom graphnet.models.detector.detector import Detector\nfrom graphnet.models.gnn import DynEdge\nfrom graphnet.models.gnn.gnn import GNN\nfrom graphnet.models.graph_builders import KNNGraphBuilder\nfrom graphnet.models.task.reconstruction import DirectionReconstructionWithKappa\nfrom graphnet.models.utils import calculate_xyzt_homophily\nfrom graphnet.training.loss_functions import VonMisesFisher3DLoss\n\n# Define constants\nCOORDS = joblib.load('/kaggle/input/icecube-submission-dynemb/coords.np')\nGLOBAL_POOLINGS = {\"min\": scatter_min, \"max\": scatter_max, \"sum\": scatter_sum, \"mean\": scatter_mean}","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class LitIceCubeKaggle(Detector):\n    \"\"\"`Detector` class for Kaggle Competition.\"\"\"\n    \n    # Implementing abstract class attribute\n    features = ['x', 'y', 'z', 'time', 'charge', 'auxiliary', 'sensor_id']\n    \n    def __init__(self, graph_builder, scalers: List[dict]=None):\n        \"\"\"Construct `Detector`.\"\"\"\n        # Base class constructor\n        super().__init__(graph_builder, scalers)\n\n    @property\n    def nb_outputs(self) -> int:\n        \"\"\"Return number of output features.\n        This the default, but may be overridden by specific inheriting classes.\"\"\"\n        return self.nb_inputs\n\n    def _forward(self, data):\n        \"\"\"Ingest data, build graph, and preprocess features.\n\n        Args:\n            data: Input graph data.\n\n        Returns:\n            Connected and preprocessed graph data.\"\"\"\n        \n        # Check(s)\n        self._validate_features(data)\n\n        # Preprocessing\n        data.x[:, 0] /= 500.0  # x\n        data.x[:, 1] /= 500.0  # y\n        data.x[:, 2] /= 500.0  # z\n        data.x[:, 3] = (data.x[:, 3] - 1.0e04) / 3.0e4  # time\n        data.x[:, 4] = torch.log10(data.x[:, 4]) / 3.0  # charge\n        return data\n\nclass LitDynEdge(GNN):\n    \"\"\"DynEmb (v5) model class.\"\"\"\n    def __init__(\n        self,\n        nb_inputs: int,\n        *,\n        nb_neighbours: int=8,\n        features_subset: Optional[Union[List[int], slice]]=None,\n        dynedge_layer_sizes: Optional[List[Tuple[int, ...]]]=None,\n        post_processing_layer_sizes: Optional[List[int]]=None,\n        readout_layer_sizes: Optional[List[int]]=None,\n        global_pooling_schemes: Optional[Union[str, List[str]]]=None,\n        add_global_variables_after_pooling: bool=False):\n        \"\"\"Construct `DynEdge`.\n\n        Args:\n            nb_inputs: Number of input features on each node.\n            nb_neighbours: Number of neighbours to used in the k-nearest\n                neighbour clustering which is performed after each (dynamical)\n                edge convolution.\n            features_subset: The subset of latent features on each node that\n                are used as metric dimensions when performing the k-nearest\n                neighbours clustering. Defaults to [0,1,2].\n            dynedge_layer_sizes: The layer sizes, or latent feature dimenions,\n                used in the `DynEdgeConv` layer. Each entry in\n                `dynedge_layer_sizes` corresponds to a single `DynEdgeConv`\n                layer; the integers in the corresponding tuple corresponds to\n                the layer sizes in the multi-layer perceptron (MLP) that is\n                applied within each `DynEdgeConv` layer. That is, a list of\n                size-two tuples means that all `DynEdgeConv` layers contain a\n                two-layer MLP.\n                Defaults to [(128, 256), (336, 256), (336, 256), (336, 256)].\n            post_processing_layer_sizes: Hidden layer sizes in the MLP\n                following the skip-concatenation of the outputs of each\n                `DynEdgeConv` layer. Defaults to [336, 256].\n            readout_layer_sizes: Hidden layer sizes in the MLP following the\n                post-processing _and_ optional global pooling. As this is the\n                last layer(s) in the model, the last layer in the read-out\n                yields the output of the `DynEdge` model. Defaults to [128,].\n            global_pooling_schemes: The list global pooling schemes to use.\n                Options are: \"min\", \"max\", \"mean\", and \"sum\".\n            add_global_variables_after_pooling: Whether to add global variables\n                after global pooling. The alternative is to  added (distribute)\n                them to the individual nodes before any convolutional\n                operations.\n        \"\"\"\n        # Latent feature subset for computing nearest neighbours in DynEdge.\n        if features_subset is None:\n            features_subset = slice(0, 5)\n\n        # DynEdge layer sizes\n        if dynedge_layer_sizes is None:\n            dynedge_layer_sizes = [\n                (\n                    128,\n                    256,\n                ),\n                (\n                    336,\n                    256,\n                ),\n                (\n                    336,\n                    256,\n                ),\n                (\n                    336,\n                    256,\n                ),\n            ]\n\n        assert isinstance(dynedge_layer_sizes, list)\n        assert len(dynedge_layer_sizes)\n        assert all(isinstance(sizes, tuple) for sizes in dynedge_layer_sizes)\n        assert all(len(sizes) > 0 for sizes in dynedge_layer_sizes)\n        assert all(\n            all(size > 0 for size in sizes) for sizes in dynedge_layer_sizes\n        )\n\n        self._dynedge_layer_sizes = dynedge_layer_sizes\n\n        # Post-processing layer sizes\n        if post_processing_layer_sizes is None:\n            post_processing_layer_sizes = [\n                336,\n                256,\n            ]\n\n        assert isinstance(post_processing_layer_sizes, list)\n        assert len(post_processing_layer_sizes)\n        assert all(size > 0 for size in post_processing_layer_sizes)\n\n        self._post_processing_layer_sizes = post_processing_layer_sizes\n\n        # Read-out layer sizes\n        if readout_layer_sizes is None:\n            readout_layer_sizes = [\n                128,\n            ]\n\n        assert isinstance(readout_layer_sizes, list)\n        assert len(readout_layer_sizes)\n        assert all(size > 0 for size in readout_layer_sizes)\n\n        self._readout_layer_sizes = readout_layer_sizes\n\n        # Global pooling scheme(s)\n        if isinstance(global_pooling_schemes, str):\n            global_pooling_schemes = [global_pooling_schemes]\n\n        if isinstance(global_pooling_schemes, list):\n            for pooling_scheme in global_pooling_schemes:\n                assert (\n                    pooling_scheme in GLOBAL_POOLINGS\n                ), f\"Global pooling scheme {pooling_scheme} not supported.\"\n        else:\n            assert global_pooling_schemes is None\n\n        self._global_pooling_schemes = global_pooling_schemes\n\n        if add_global_variables_after_pooling:\n            assert self._global_pooling_schemes, (\n                \"No global pooling schemes were request, so cannot add global\"\n                \" variables after pooling.\"\n            )\n        self._add_global_variables_after_pooling = (add_global_variables_after_pooling)\n\n        # Base class constructor\n        super().__init__(nb_inputs, self._readout_layer_sizes[-1])\n\n        # Switch to SiLU function\n        self._activation = torch.nn.SiLU()\n        \n        self._nb_inputs = nb_inputs - 1 + 3\n        self._nb_global_variables = 5 + nb_inputs - 1 + 3\n        self._nb_neighbours = nb_neighbours\n        self._features_subset = features_subset\n        \n        # get embeddings\n        self.id_embedder = nn.Embedding.from_pretrained(torch.Tensor(COORDS), freeze=False)\n        self._construct_layers()\n\n    def _construct_layers(self) -> None:\n        \"\"\"Construct layers (torch.nn.Modules).\"\"\"\n        # Convolutional operations\n        nb_input_features = self._nb_inputs\n        if not self._add_global_variables_after_pooling:\n            nb_input_features += self._nb_global_variables\n\n        self._conv_layers = torch.nn.ModuleList()\n        \n        nb_latent_features = nb_input_features \n        for sizes in self._dynedge_layer_sizes:\n            layers = []\n            layer_sizes = [nb_latent_features] + list(sizes)\n            for ix, (nb_in, nb_out) in enumerate(\n                zip(layer_sizes[:-1], layer_sizes[1:])\n            ):\n                if ix == 0:\n                    nb_in *= 2\n                layers.append(torch.nn.BatchNorm1d(nb_in))\n                layers.append(torch.nn.utils.weight_norm(torch.nn.Linear(nb_in, nb_out)))\n                layers.append(self._activation)\n            \n            conv_layer = DynEdgeConv(\n                torch.nn.Sequential(*layers),\n                aggr='max',\n                nb_neighbors=self._nb_neighbours,\n                features_subset=self._features_subset)\n            self._conv_layers.append(conv_layer)\n\n            nb_latent_features = nb_out\n\n        # Post-processing operations\n        nb_latent_features = (sum(sizes[-1] for sizes in self._dynedge_layer_sizes)+ nb_input_features)\n        \n        post_processing_layers = []\n        layer_sizes = [nb_latent_features] + list(self._post_processing_layer_sizes)\n        for nb_in, nb_out in zip(layer_sizes[:-1], layer_sizes[1:]):\n            post_processing_layers.append(torch.nn.utils.weight_norm(torch.nn.Linear(nb_in, nb_out)))\n            post_processing_layers.append(self._activation)\n\n        self._post_processing = torch.nn.Sequential(*post_processing_layers)\n\n        # Read-out operations\n        nb_poolings = (\n            len(self._global_pooling_schemes)\n            if self._global_pooling_schemes\n            else 1)\n        \n        nb_latent_features = nb_out * nb_poolings\n        if self._add_global_variables_after_pooling:\n            nb_latent_features += self._nb_global_variables\n\n        readout_layers = []\n        layer_sizes = [nb_latent_features] + list(self._readout_layer_sizes)\n        for nb_in, nb_out in zip(layer_sizes[:-1], layer_sizes[1:]):\n            readout_layers.append(torch.nn.utils.weight_norm(torch.nn.Linear(nb_in, nb_out)))\n            readout_layers.append(self._activation)\n\n        self._readout = torch.nn.Sequential(*readout_layers)\n\n    def _global_pooling(self, x: Tensor, batch: LongTensor) -> Tensor:\n        \"\"\"Perform global pooling.\"\"\"\n        assert self._global_pooling_schemes\n        pooled = []\n        for pooling_scheme in self._global_pooling_schemes:\n            pooling_fn = GLOBAL_POOLINGS[pooling_scheme]\n            pooled_x = pooling_fn(x, index=batch, dim=0)\n            if isinstance(pooled_x, tuple) and len(pooled_x) == 2:\n                pooled_x, _ = pooled_x\n            pooled.append(pooled_x)\n\n        return torch.cat(pooled, dim=1)\n\n    def _calculate_global_variables(self,\n                                    x: Tensor,\n                                    edge_index: LongTensor,\n                                    batch: LongTensor,\n                                    *additional_attributes: Tensor\n                                    ) -> Tensor:\n        \"\"\"Calculate global variables.\"\"\"\n        \n        # Calculate homophily (scalar variables)\n        h_x, h_y, h_z, h_t = calculate_xyzt_homophily(x, edge_index, batch)\n\n        # Calculate mean features\n        global_means = scatter_mean(x, batch, dim=0)\n\n        # Add global variables\n        global_variables = torch.cat(\n            [\n                global_means,\n                h_x,\n                h_y,\n                h_z,\n                h_t,\n            ]\n            + [attr.unsqueeze(dim=1) for attr in additional_attributes],\n            dim=1,\n        )\n\n        return global_variables\n\n    def forward(self, data: Data) -> Tensor:\n        \"\"\"Apply learnable forward pass.\"\"\"\n        # Convenience variables\n        x, edge_index, batch = data.x, data.edge_index, data.batch\n        embeds = self.id_embedder(x[:, 6].long())\n        x = torch.cat([x[:, :6], embeds, x[:, 7:]], dim=1)\n\n        global_variables = self._calculate_global_variables(\n            x,\n            edge_index,\n            batch,\n            torch.log10(data.n_pulses),\n        )\n\n        # Distribute global variables out to each node\n        if not self._add_global_variables_after_pooling:\n            distribute = (\n                batch.unsqueeze(dim=1) == torch.unique(batch).unsqueeze(dim=0)\n            ).type(torch.float)\n\n            global_variables_distributed = torch.sum(\n                distribute.unsqueeze(dim=2)\n                * global_variables.unsqueeze(dim=0),\n                dim=1,\n            )\n\n            x = torch.cat((x, global_variables_distributed), dim=1)\n\n        # DynEdge-convolutions\n        skip_connections = [x]\n        for conv_layer in self._conv_layers:\n            x, edge_index = conv_layer(x, edge_index, batch)\n            skip_connections.append(x)\n\n        # Skip-cat\n        x = torch.cat(skip_connections, dim=1)\n\n        # Post-processing\n        x = self._post_processing(x)\n\n        # (Optional) Global pooling\n        if self._global_pooling_schemes:\n            x = self._global_pooling(x, batch=batch)\n            if self._add_global_variables_after_pooling:\n                x = torch.cat(\n                    [\n                        x,\n                        global_variables,\n                    ],\n                    dim=1,\n                )\n        # Read-out\n        x = self._readout(x)\n        return x\n\nclass LitDirectionReconstructionWithKappa(DirectionReconstructionWithKappa):\n    \"\"\"Reconstructs direction with kappa from the 3D-vMF distribution.\"\"\"\n    # Requires three features: untransformed points in (x,y,z)-space.\n    nb_inputs = 3\n    def __init__(\n        self,\n        *,\n        hidden_size: int,\n        target_labels,\n        loss_function,\n        transform_prediction_and_target=None,\n        transform_target=None,\n        transform_inference=None,\n        transform_support=None,\n        loss_weight=None):\n      \n        # Base class constructor\n        super().__init__(hidden_size=hidden_size, target_labels=target_labels, loss_function=loss_function)\n        self._affine = torch.nn.utils.weight_norm(torch.nn.Linear(hidden_size, self.nb_inputs))\n\n        # init bias\n        bias = torch.Tensor((0.0000043, 0.0001597, 0.0305320))\n        self._affine.bias.data = bias\n        \n        shape = self._affine.weight.data.shape\n        self._affine.weight.data = torch.zeros(shape[0], shape[1], requires_grad=True)\n\nclass LitModel(LightningModule):\n    \"\"\"Lightning wrapper for DynEmb model.\"\"\"\n    def __init__(self):\n        super().__init__()\n        self._detector = LitIceCubeKaggle(graph_builder=KNNGraphBuilder(nb_nearest_neighbours=9))\n        self._gnn = LitDynEdge(nb_inputs=self._detector.nb_outputs,\n                               dynedge_layer_sizes=[(128, 256), (512, 256), (512, 256), (512, 256)],\n                               post_processing_layer_sizes=[512, 256],\n                               readout_layer_sizes=[256],\n                               global_pooling_schemes=[\"min\", \"max\", \"mean\", 'sum'])\n        self._tasks = torch.nn.ModuleList([LitDirectionReconstructionWithKappa(hidden_size=self._gnn.nb_outputs,\n                                                                               target_labels='direction',\n                                                                               loss_function=VonMisesFisher3DLoss())])\n\n    def forward(self, data: Data) -> List[Union[Tensor, Data]]:\n        \"\"\"Forward pass, chaining model components.\"\"\"\n        data = self._detector(data)\n        x = self._gnn(data)\n        preds = self._tasks[0](x)\n        return preds\n\nclass BatchDataset(Dataset):\n    # Fast graph dataset.\n    def __init__(self, batch_id, additional_features=['sensor_id']):\n        super().__init__()\n\n        self._df = pd.read_parquet(f'/kaggle/input/icecube-neutrinos-in-deep-ice/test/batch_{batch_id}.parquet')\n        self._df['charge'] = self._df['charge'].astype(np.float32)\n        self._df['auxiliary'] = self._df['auxiliary'].astype(np.uint8)\n        self._df = self._df.reset_index().merge(get_sensor_geometry(), on='sensor_id', how='left').set_index('event_id')\n        self._features = ['x', 'y', 'z', 'time', 'charge', 'auxiliary'] + additional_features\n        self._indexes = self._df.index.unique()\n\n    def __len__(self):\n        return len(self._indexes)\n\n    def __getitem__(self, idx):\n\n        event_id = self._indexes[idx]\n        x = torch.tensor(self._df.loc[event_id][self._features].values, dtype=torch.float32)\n\n        graph = Data(x=x,\n                     features=self._features,\n                     event_id=event_id,\n                     n_pulses=torch.tensor(x.shape[0], dtype=torch.int32),\n                     edge_index=None,\n                     )\n\n        return graph\n\ndef get_sensor_geometry():\n    \"\"\"Read and process sensor geometry file.\"\"\"\n    sensor_geometry = pd.read_csv('/kaggle/input/icecube-neutrinos-in-deep-ice/sensor_geometry.csv')\n    sensor_geometry['sensor_id'] = sensor_geometry['sensor_id'].astype(np.uint16)\n    sensor_geometry['x'] = sensor_geometry['x'].astype(np.float32)\n    sensor_geometry['y'] = sensor_geometry['y'].astype(np.float32)\n    sensor_geometry['z'] = sensor_geometry['z'].astype(np.float32)\n    sensor_geometry.set_index('sensor_id', inplace=True)\n    return sensor_geometry\n\ndef inference(dataloader, model):\n    \"\"\"perform inference for given model and data.\"\"\"\n    results = []\n    event_ids = []\n\n    with torch.no_grad():\n        model = model.to('cuda')\n        model.eval()\n\n        for batch in tqdm(dataloader):\n            batch = batch.to('cuda')\n            predict = model(batch)\n            results += predict.detach().cpu().numpy().tolist()\n            event_ids += batch.event_id\n\n    results = pd.DataFrame(results, columns=['direction_x', 'direction_y', 'direction_z', 'direction_kappa'])\n    event_ids = pd.DataFrame(event_ids, columns=['event_id'])\n    results = pd.concat([results, event_ids], axis=1).reset_index(drop=True)\n    return results\n\ndef prepare_submission(df):\n    \"\"\"Get zenith and azimuth.\"\"\"\n    df['zenith'] = np.arccos(df['direction_z'].values)\n    df['azimuth'] = np.arctan2(df['direction_y'].values, df['direction_x'].values)\n    df['azimuth'][df['azimuth'] < 0] = df['azimuth'][df['azimuth'] < 0] + 2*np.pi\n    return df[['event_id', 'azimuth', 'zenith']].set_index('event_id')","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# define model and load parameters\nmodel = LitModel()\nmodel.load_state_dict(torch.load('/kaggle/input/icecube-submission-dynemb/checkpoint-dynemb.ckpt')['state_dict'])\n\n# do inference\ninference_batch_ids = []\nfor batch_path in glob.glob('/kaggle/input/icecube-neutrinos-in-deep-ice/test/batch_*'):\n    batch_name = batch_path.split('/')[-1]\n    batch_name = batch_name.replace('batch_', '')\n    batch_name = batch_name.replace('.parquet', '')\n    inference_batch_ids.append(int(batch_name))\ninference_batch_ids = sorted(inference_batch_ids)\n\nresults = []\nfor i in inference_batch_ids:\n    dataset = BatchDataset(batch_id=i)\n    dataloader = DataLoader(dataset,\n                            batch_size=32,\n                            num_workers=2,\n                            shuffle=False,\n                            drop_last=False,\n                            collate_fn=lambda g: Batch.from_data_list(g))\n    results.append(inference(dataloader, model))\nresults = pd.concat(results).reset_index(drop=True)\n\n# form submission\nsubmission = prepare_submission(results)\nsubmission.to_csv('/kaggle/working/submission.csv')","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}