{"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":"!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/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/pytorch_lightning-1.9.0-py3-none-any.whl\n!pip install /kaggle/working/software/dependencies/torch_geometric-2.0.4.tar.gz","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# -----------------------------------------------------------------------------\n# Imports\n# -----------------------------------------------------------------------------\n\nimport os, random, copy, yaml, pickle\nfrom time import time, sleep\nfrom tqdm import tqdm\nfrom math import floor\n\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport pyarrow.parquet as pq\n\n# torch imports\nimport torch\nimport torch.nn as nn\nimport torch.nn.init as init\nimport torch.nn.functional as F\nfrom torch import Tensor\nfrom torch.utils.data import random_split, IterableDataset\n\nimport torch_geometric\nfrom torch_geometric.data import Data, Batch\nfrom torch_geometric.utils import homophily\nfrom torch_geometric.loader import DataLoader\nfrom torch_geometric.nn import EdgeConv, knn_graph\nfrom torch_scatter import \\\n    scatter_add, scatter_mean, scatter_max, scatter_min\n\nimport pytorch_lightning as pl","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# -----------------------------------------------------------------------------\n# Helper function\n# -----------------------------------------------------------------------------\n\nimport logging\nimport logging.config\nimport os, re, gc, psutil\n\ndef get_logger(name, msg):\n    \"\"\"\n    :param name: string\n    :param msg: DEBUG, INFO, WARNING, ERROR\n    :return: Logger() instance\n    \"\"\"\n    level = {\"DEBUG\": logging.DEBUG,\n             \"INFO\": logging.INFO,\n             \"WARNING\": logging.WARNING,\n             \"ERROR\": logging.ERROR}\n    logging.basicConfig(level=level[msg],\n                        format=\"== %(name)s == %(asctime)s %(levelname)s:\\t%(message)s\",\n                        datefmt=\"%H:%M:%S\")\n    logger = logging.getLogger(name)\n\n    return logger\n\n\ndef walk_dir(dirname, batch_ids):\n    files = dict()\n    pattern = r\"_(\\d+)\\.parquet\"\n\n    if batch_ids is None:\n        batch_ids = list()\n        for base, _, names in os.walk(dirname):\n            for name in names:\n                match = re.findall(pattern, name)\n                batch_ids.append(int(match[0]))\n                files[int(match[0])] = os.path.join(base, name)\n        return files, batch_ids\n    \n    for base, _, names in os.walk(dirname):\n        selected_files = dict()\n        for name in names:\n            match = re.findall(pattern, name)\n            if int(match[0]) in batch_ids:\n                selected_files[int(match[0])] = os.path.join(base, name)\n        files.update(selected_files)\n    return files, batch_ids\n\n\ndef memory_check(logger, msg=\"\"):\n    gc.collect()\n    logger.debug(f\"memory usage {psutil.virtual_memory().used / 1024**3:.2f} \"\n                f\"of {psutil.virtual_memory().total / 1024**3:.2f} GB {msg}\")\n","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# -----------------------------------------------------------------------------\n# Basic settings\n# -----------------------------------------------------------------------------\n\n# basic\nBATCH_SIZE = 400\nBATCHES_TEST = None #####################\nEVENTS_PER_FILE = 200_000\n\n# paths\nBASE_PATH = \"/kaggle/\" ######################\nPATH = os.path.join(BASE_PATH, \"input\", \"icecube-neutrinos-in-deep-ice\")\nMODEL_PATH = os.path.join(BASE_PATH, \"input\", \"ice-cube-model\")\nOUTPUT_PATH = os.path.join(BASE_PATH, \"working\")\nTEST_PATH = os.path.join(PATH, \"test\")\nMETA_PATH = os.path.join(OUTPUT_PATH, \"test_meta\")\n\n# files\nFILES_TEST, BATCHES_TEST = walk_dir(TEST_PATH, BATCHES_TEST)\nFILE_META = os.path.join(PATH, \"test_meta.parquet\")\nFILE_SENSOR_GEO = os.path.join(PATH, \"sensor_geometry.csv\")\nFILE_GNN = os.path.join(MODEL_PATH, \"finetuned.ckpt\")\nFILE_BDT = os.path.join(MODEL_PATH, \"BDT_clf.Deep.0414.sklearn\")\nCUT_BDT = 0.00015","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# -----------------------------------------------------------------------------\n# Split meta file\n# -----------------------------------------------------------------------------\n\nif not os.path.exists(META_PATH): \n    os.mkdir(META_PATH)\n\nmeta_test = pd.read_parquet(FILE_META)\n\nfor i, df in meta_test.groupby(\"batch_id\"):\n    print(f\"processing {i} -> {df.shape}\")\n    splitted = os.path.join(META_PATH, f\"meta_{i}.parquet\")\n    df.to_parquet(splitted)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# -----------------------------------------------------------------------------\n# Some Logging\n# -----------------------------------------------------------------------------\nLOGGER = get_logger(\"IceCube\", \"DEBUG\")\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nLOGGER.info(f\"using {DEVICE}\")\nLOGGER.info(f\"{len(FILES_TEST)} files for testing\")\nmemory_check(LOGGER)\n\nLOGGER.info(f\"GNN model:{FILE_GNN}\")\nLOGGER.info(f\"BDT model:{FILE_BDT}\")","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# -----------------------------------------------------------------------------\n# Dataset\n# -----------------------------------------------------------------------------\n\n# sensor geometry\ndef prepare_sensors(scale=None):\n    sensors = pd.read_csv(FILE_SENSOR_GEO).astype({\n        \"sensor_id\": np.int16,\n        \"x\": np.float32,\n        \"y\": np.float32,\n        \"z\": np.float32,\n    })\n\n    if scale is not None and isinstance(scale, float):\n        sensors[\"x\"] *= scale\n        sensors[\"y\"] *= scale\n        sensors[\"z\"] *= scale\n\n    return sensors\n\n\ndef angle_to_xyz(angles_b):\n    az, zen = angles_b.t()\n    x = torch.cos(az) * torch.sin(zen)\n    y = torch.sin(az) * torch.sin(zen)\n    z = torch.cos(zen)\n    return torch.stack([x, y, z], dim=1)\n\n\ndef xyz_to_angle(xyz_b):\n    x, y, z = xyz_b.t()\n    az = torch.arccos(x / torch.sqrt(x**2 + y**2)) * torch.sign(y)\n    zen = torch.arccos(z / torch.sqrt(x**2 + y**2 + z**2))\n    return torch.stack([az, zen], dim=1)\n\n\ndef angular_error(xyz_pred_b, xyz_true_b):\n    return torch.arccos(torch.clip_(torch.sum(xyz_pred_b * xyz_true_b, dim=1), -1, 1))\n\n\ndef angles2vector(df):\n    df[\"nx\"] = np.sin(df.zenith) * np.cos(df.azimuth)\n    df[\"ny\"] = np.sin(df.zenith) * np.sin(df.azimuth)\n    df[\"nz\"] = np.cos(df.zenith) \n    return df\n\n\ndef vector2angles(n, eps=1e-8):\n    n = n / (np.linalg.norm(n, axis=1, keepdims=True) + eps)    \n    azimuth = np.arctan2( n[:,1],  n[:,0])    \n    azimuth[azimuth < 0] += 2*np.pi\n    zenith = np.arccos( n[:,2].clip(-1,1) )                                \n    return np.concatenate([azimuth[:, np.newaxis], zenith[:, np.newaxis]], axis=1)\n\n\ndef series2tensor(series, set_device=None):\n    ret = torch.from_numpy(series.values).float()\n    if set_device is not None:\n        return ret.to(DEVICE)\n    return ret\n\n\n\ndef solve_linear(xw, yw, zw, xxw, yyw, xyw, yzw, zxw):\n    A = torch.tensor([\n        [xxw, xyw, xw],\n        [xyw, yyw, yw],\n        [xw,  yw,  1 ],\n    ])\n    b = torch.tensor([\n        zxw, yzw, zw\n    ])\n    try:\n        coeff = torch.linalg.solve(A, b)\n        return coeff\n    except Exception:\n        LOGGER.debug(\"linear system not solvable\")\n        return torch.zeros((3, ))\n\n\ndef feature_extraction(df, fun=None, eps=1e-8):                                           # list of variables\n    # sort by time\n    df.sort_values([\"time\"], inplace=True)\n\n    t = series2tensor(df.time)\n    c = series2tensor(df.charge)\n    x = series2tensor(df.x)\n    y = series2tensor(df.y)\n    z = series2tensor(df.z)\n\n    hits = t.numel()                                                                            # hits\n\n    # weighted values\n    Sx = torch.sum(x); Sxx = torch.sum(x*x); Sxy = torch.sum(x*y)\n    Sy = torch.sum(y); Syy = torch.sum(y*y); Syz = torch.sum(y*z)\n    Sz = torch.sum(z); Szx = torch.sum(z*x)\n\n    # error of plane estimate\n    coeff = solve_linear(Sx, Sy, Sz, Sxx, Syy, Sxy, Syz, Szx)\n    error = torch.sum((z - coeff[0] * x - coeff[1] * y - coeff[2]))\n    error = torch.square(error * 1e3)                                                           # error\n\n    # plane norm vector\n    norm_vec = torch.tensor([coeff[0], coeff[1], -1], dtype=torch.float)\n    norm_vec /= torch.sqrt(coeff[0]**2 + coeff[1]**2 + 1)                                       # norm_vec -> (3, )\n\n    # delta t -> median time\n    dt = torch.quantile(t, torch.tensor([0.15, 0.50, 0.85], dtype=torch.float))                 # dt -> (3, )\n\n    # charge centre (vector)\n    sumq = torch.sum(c)                                                                         # sumq\n    meanq = sumq / hits                                                                         # meanq\n    qv = torch.tensor([torch.sum(x*c), torch.sum(y*c), torch.sum(z*c)], dtype=torch.float)\n    qv /= sumq                                                                                  # qv -> (3, )\n\n    # bright sensor ratio\n    bratio = c[c > 5 * meanq].numel() / hits                                                    # bratio\n\n    # grouping by time (remember to sort by time)\n    n_groups = 4                                                                                # xyzt -> (16, )\n\n    if hits > n_groups:\n        sec_len = floor(hits / n_groups)\n        remain_len = hits - (n_groups - 1) * sec_len\n        xyzt = series2tensor(df[[\"x\", \"y\", \"z\", \"time\"]])\n        xyzt = torch.split(xyzt, [sec_len, sec_len, sec_len, remain_len])\n        xyzt = torch.concat([xx.mean(axis=0) for xx in xyzt])\n    else:\n        xyzt = torch.zeros(n_groups * 4)\n        _xxxx = list()\n        for i in range(hits):\n            _xxxx.append(x[i]); _xxxx.append(y[i]); _xxxx.append(z[i]); _xxxx.append(t[i])\n        xyzt[: hits * 4] = torch.tensor(_xxxx, dtype=torch.float)\n\n    # unique xyz\n    unique = torch.tensor([_x.unique().numel() for _x in [x, y, z]], dtype=torch.float)         # unique -> (3, )\n\n    # global features\n    glob_feat = torch.tensor([hits, error, sumq, meanq, bratio, ], dtype=torch.float)\n\n    return torch.concat([norm_vec, dt, qv, xyzt, unique, glob_feat]).unsqueeze(0)\n\n\ndef prepare_feature(df):\n    df = df.reset_index(drop=True)\n    # remove auxiliary\n    df = df[~df.auxiliary]\n    df.x *= 1e-3; df.y *= 1e-3; df.z *= 1e-3\n    df.time -= np.min(df.time)\n    return df[[\"time\", \"charge\", \"x\", \"y\", \"z\"]]\n\n\n# Dataset\nclass IceCube(IterableDataset):\n    def __init__(\n        self, parquet_dir, meta_dir, chunk_ids,\n        batch_size=200, max_pulses=200, extra=False\n    ):\n        self.parquet_dir = parquet_dir\n        self.meta_dir = meta_dir\n        self.chunk_ids = chunk_ids\n        self.batch_size = batch_size\n        self.max_pulses = max_pulses\n        self.extra = extra\n\n    def __iter__(self):\n        # Handle num_workers > 1 and multi-gpu\n        is_dist = torch.distributed.is_initialized()\n        world_size = torch.distributed.get_world_size() if is_dist else 1\n        rank_id = torch.distributed.get_rank() if is_dist else 0\n\n        info = torch.utils.data.get_worker_info()\n        num_worker = info.num_workers if info else 1\n        worker_id = info.id if info else 0\n\n        num_replica = world_size * num_worker\n        offset = rank_id * num_worker + worker_id\n        chunk_ids = self.chunk_ids[offset::num_replica]\n\n        # Sensor data\n        sensor = prepare_sensors()\n\n        # Read each chunk and meta iteratively into memory and build mini-batch\n        for c, chunk_id in enumerate(chunk_ids):\n            data = pd.read_parquet(os.path.join(self.parquet_dir, f\"batch_{chunk_id}.parquet\"))\n            meta = pd.read_parquet(os.path.join(self.meta_dir, f\"meta_{chunk_id}.parquet\"))\n\n            eids = meta[\"event_id\"].tolist()\n            eids_batches = [\n                eids[i : i + self.batch_size]\n                for i in range(0, len(eids), self.batch_size)\n            ]\n\n            for batch_eids in eids_batches:\n                batch = []\n\n                # For each sample, extract features\n                for eid in batch_eids:\n                    df = data.loc[eid]\n                    df = pd.merge(df, sensor, on=\"sensor_id\")\n                    # sampling of pulses if number exceeds maximum\n                    if len(df) > self.max_pulses:\n                        df_pass = df[~df.auxiliary]\n                        df_fail = df[df.auxiliary]\n                        if len(df_pass) >= self.max_pulses:\n                            df = df_pass.sample(self.max_pulses)\n                        else:\n                            df_fail = df_fail.sample(self.max_pulses - len(df_pass))\n                            df = pd.concat([df_fail, df_pass])\n\n                    df.sort_values([\"time\"], inplace=True)\n\n                    t = series2tensor(df.time)\n                    c = series2tensor(df.charge)\n                    a = series2tensor(df.auxiliary)\n                    x = series2tensor(df.x)\n                    y = series2tensor(df.y)\n                    z = series2tensor(df.z)\n                        \n                    feat = torch.stack([x, y, z, t, c, a], dim=1)\n\n                    batch_data = Data(x=feat, n_pulses=len(feat), eid=torch.tensor([eid]).long())\n\n                    if self.extra:\n                        feats = feature_extraction(prepare_feature(df))\n                        setattr(batch_data, \"extra_feat\", feats)\n\n                    batch.append(batch_data)\n\n                yield Batch.from_data_list(batch)\n\n            del data\n            del meta\n            gc.collect()\n","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# -----------------------------------------------------------------------------\n# GraphNet GNN model\n# -----------------------------------------------------------------------------\n\nclass MLP(nn.Sequential):\n    def __init__(self, feats):\n        layers = []\n        for i in range(1, len(feats)):\n            layers.append(nn.Linear(feats[i - 1], feats[i]))\n            layers.append(nn.LeakyReLU())\n        super().__init__(*layers)\n\n\nclass Model(pl.LightningModule):\n    def __init__(\n        self, max_lr=1e-3, \n        num_warmup_step=1_000,\n        remaining_step=1_000,\n    ):\n        super().__init__()\n        self.save_hyperparameters()\n        self.conv0 = EdgeConv(MLP([17 * 2, 128, 256]), aggr=\"add\")\n        self.conv1 = EdgeConv(MLP([512, 336, 256]), aggr=\"add\")\n        self.conv2 = EdgeConv(MLP([512, 336, 256]), aggr=\"add\")\n        self.conv3 = EdgeConv(MLP([512, 336, 256]), aggr=\"add\")\n        self.post = MLP([1024 + 17, 336, 256])\n        self.readout = MLP([768, 128])\n        self.pred = nn.Linear(128, 3)\n\n    def forward(self, data: Batch):\n        vert_feat = data.x\n        batch = data.batch\n\n        # x, y, z, t, c, a\n        # 0  1  2  3  4  5\n        vert_feat[:, 0] /= 500.0  # x\n        vert_feat[:, 1] /= 500.0  # y\n        vert_feat[:, 2] /= 500.0  # z\n        vert_feat[:, 3] = (vert_feat[:, 3] - 1.0e04) / 3.0e4  # time\n        vert_feat[:, 4] = torch.log10(vert_feat[:, 4]) / 3.0  # charge\n\n        edge_index = knn_graph(vert_feat[:, :3], 8, batch)\n\n        # Construct global features\n        hx = homophily(edge_index, vert_feat[:, 0], batch).reshape(-1, 1)\n        hy = homophily(edge_index, vert_feat[:, 1], batch).reshape(-1, 1)\n        hz = homophily(edge_index, vert_feat[:, 2], batch).reshape(-1, 1)\n        ht = homophily(edge_index, vert_feat[:, 3], batch).reshape(-1, 1)\n        means = scatter_mean(vert_feat, batch, dim=0)\n        n_p = torch.log10(data.n_pulses).reshape(-1, 1)\n        global_feats = torch.cat([means, hx, hy, hz, ht, n_p], dim=1)  # [B, 11]\n\n        # Distribute global_feats to each vertex\n        _, cnts = torch.unique_consecutive(batch, return_counts=True)\n        global_feats = torch.repeat_interleave(global_feats, cnts, dim=0)\n        vert_feat = torch.cat((vert_feat, global_feats), dim=1)\n\n        # Convolutions\n        feats = [vert_feat]\n        # Conv 0\n        vert_feat = self.conv0(vert_feat, edge_index)\n        feats.append(vert_feat)\n        # Conv 1\n        edge_index = knn_graph(vert_feat[:, :3], k=8, batch=batch)\n        vert_feat = self.conv1(vert_feat, edge_index)\n        feats.append(vert_feat)\n        # Conv 2\n        edge_index = knn_graph(vert_feat[:, :3], k=8, batch=batch)\n        vert_feat = self.conv2(vert_feat, edge_index)\n        feats.append(vert_feat)\n        # Conv 3\n        edge_index = knn_graph(vert_feat[:, :3], k=8, batch=batch)\n        vert_feat = self.conv3(vert_feat, edge_index)\n        feats.append(vert_feat)\n\n        # Postprocessing\n        post_inp = torch.cat(feats, dim=1)\n        post_out = self.post(post_inp)\n\n        # Readout\n        readout_inp = torch.cat(\n            [\n                scatter_min(post_out, batch, dim=0)[0],\n                scatter_max(post_out, batch, dim=0)[0],\n                scatter_mean(post_out, batch, dim=0),\n            ],\n            dim=1,\n        )\n        readout_out = self.readout(readout_inp)\n\n        # Predict\n        pred = self.pred(readout_out)\n        kappa = pred.norm(dim=1, p=2) + 1e-8\n        pred_x = pred[:, 0] / kappa\n        pred_y = pred[:, 1] / kappa\n        pred_z = pred[:, 2] / kappa\n        pred = torch.stack([pred_x, pred_y, pred_z, kappa], dim=1)\n\n        return pred\n\n    def train_or_valid_step(self, data, prefix):\n        pred_xyzk = self.forward(data)  # [B, 4]\n        true_xyz = data.gt.view(-1, 3)  # [B, 3]\n        loss = VonMisesFisher3DLoss()(pred_xyzk, true_xyz).mean()\n        error = angular_error(pred_xyzk[:, :3], true_xyz).mean()\n        self.log(f\"loss-{prefix}\", loss, batch_size=len(true_xyz), \n            on_epoch=True, prog_bar=True, logger=True, sync_dist=True)\n        self.log(f\"error-{prefix}\", error, batch_size=len(true_xyz), \n            on_epoch=True, prog_bar=True, logger=True, sync_dist=True)\n        return loss\n\n    def training_step(self, data, _):\n        return self.train_or_valid_step(data, \"train\")\n\n    def validation_step(self, data, _):\n        return self.train_or_valid_step(data, \"valid\")\n    \n    def configure_optimizers(self):\n        optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.max_lr)\n        scheduler = torch.optim.lr_scheduler.SequentialLR(\n            optimizer,\n            schedulers=[\n                torch.optim.lr_scheduler.LinearLR(\n                    optimizer, 1e-2, 1, self.hparams.num_warmup_step\n                ),\n                torch.optim.lr_scheduler.LinearLR(\n                    optimizer, 1, 1e-3, self.hparams.remaining_step\n                ),\n            ],\n            milestones=[self.hparams.num_warmup_step],\n        )\n        return {\n            \"optimizer\": optimizer,\n            \"lr_scheduler\": {\n                \"scheduler\": scheduler,\n                \"interval\": \"step\",\n            },\n        }\n","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# -----------------------------------------------------------------------------\n# GNN prediction\n# -----------------------------------------------------------------------------\n\ndef predict_gnn(model):\n\n    test_set = IceCube(TEST_PATH, META_PATH, BATCHES_TEST, batch_size=400, extra=True)\n    test_loader = DataLoader(test_set, batch_size=1)\n\n    pred = None\n    eid = None\n    \n    with torch.no_grad():\n        for i, data in tqdm(enumerate(test_loader)):\n            pred_xyzk = model(data.to(DEVICE))\n            angles = np.concatenate([\n                # +-------------+------------------------------------+----------------------+\n                # | x, y, z, kp |           azimuth, zenith          |    extra features    |\n                # +-------------+------------------------------------+----------------------+\n                pred_xyzk.cpu(), xyz_to_angle(pred_xyzk[:, :3]).cpu(), data.extra_feat.cpu()\n                # +-------------+------------------------------------+----------------------+\n            ], axis=1)\n            pred = angles if pred is None else np.concatenate([pred, angles])\n            eid = data.eid.cpu() if eid is None else np.concatenate([eid, data.eid.cpu()])\n\n    col_xyzk = [\"x\", \"y\", \"z\", \"kappa\"]\n    col_angles = [\"azimuth\", \"zenith\"]\n    col_norm_vec = [\"ex\", \"ey\", \"ez\"]\n    col_dt = [\"dt_15\", \"dt_50\", \"dt_85\"]\n    col_qv = [\"qx\", \"qy\", \"qz\"]\n    col_xyzt = [\n        \"x0\", \"y0\", \"z0\", \"t0\",\n        \"x1\", \"y1\", \"z1\", \"t1\",\n        \"x2\", \"y2\", \"z2\", \"t2\",\n        \"x3\", \"y3\", \"z3\", \"t3\", ]\n    col_unique = [\"uniq_x\", \"uniq_y\", \"uniq_z\"]\n    col_glob_feat = [\"hits\", \"error\", \"sumq\", \"meanq\", \"bratio\"]\n    col_extra = col_norm_vec + col_dt + col_qv + \\\n        col_xyzt + col_unique + col_glob_feat\n    \n    res = pd.DataFrame(pred, columns=col_xyzk+col_angles+col_extra)\n    \n    res[\"azimuth\"] = np.remainder(res[\"azimuth\"], 2 * np.pi)\n    res[\"zenith\"] = np.remainder(res[\"zenith\"], 2 * np.pi)\n    res[\"event_id\"] = eid\n\n    return res\n\n\nmodel = Model.load_from_checkpoint(FILE_GNN)\nLOGGER.info(f\"loaded {FILE_GNN}\")\n\nmodel.eval()\nmodel.freeze()\nmodel.to(DEVICE)\n\nreco_df = predict_gnn(model)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# -----------------------------------------------------------------------------\n# GNN + Plane projection prediction\n# -----------------------------------------------------------------------------\n\nn_hat = reco_df[[\"x\", \"y\", \"z\"]].to_numpy()\ne = reco_df[[\"ex\", \"ey\", \"ez\"]].to_numpy()\nxe = np.sum(n_hat * e, axis=1, keepdims=True)\n\nproj = n_hat - xe * e\nproj /= (np.linalg.norm(proj, axis=1, keepdims=True) + 1e-8)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# -----------------------------------------------------------------------------\n# GNN + Plane projection + BDT prediction\n# -----------------------------------------------------------------------------\n\n# reco_df inputs\nreco_df[\"error\"] = np.log10(reco_df[\"error\"] + 1e-6)\nreco_df[\"sumq\"] = np.log10(reco_df[\"sumq\"] + 1e-3)\nreco_df[\"dt_50\"] = np.log10(reco_df[\"dt_50\"] + 1e-3)\nreco_df[\"dt_85\"] = np.log10(reco_df[\"dt_85\"] + 1e-3)\nreco_df[\"kappa\"] = np.log10(reco_df[\"kappa\"] + 1e-3)\ncolumns = [\"kappa\", \"zenith\", \"error\", \"sumq\", \"qz\", \"dt_50\", \"dt_85\", \"ez\"]\nreco = reco_df[columns].to_numpy()\nxe = np.arccos(xe)\n\n# trajectory display\ncol_xyzt = [\n    \"x0\", \"y0\", \"z0\", \"t0\",\n    \"x1\", \"y1\", \"z1\", \"t1\",\n    \"x2\", \"y2\", \"z2\", \"t2\",\n    \"x3\", \"y3\", \"z3\", \"t3\", ]\ntraj = reco_df[col_xyzt].values\ntraj = traj.reshape(-1, 4, 4)\n\nv1 = 1e3 * (traj[:, 1, :3] - traj[:, 0, :3]) / (traj[:, 1, 3] - traj[:, 0, 3] + 1)[:, np.newaxis]\nv2 = 1e3 * (traj[:, 2, :3] - traj[:, 1, :3]) / (traj[:, 2, 3] - traj[:, 1, 3] + 1)[:, np.newaxis]\nv3 = 1e3 * (traj[:, 3, :3] - traj[:, 2, :3]) / (traj[:, 3, 3] - traj[:, 2, 3] + 1)[:, np.newaxis]\n\nv1scale = np.linalg.norm(v1, axis=1, keepdims=True) + 1e-1\nv2scale = np.linalg.norm(v2, axis=1, keepdims=True) + 1e-1\nv3scale = np.linalg.norm(v3, axis=1, keepdims=True) + 1e-1\n\nev1 = np.sum(-v1 * e / v1scale, axis=1, keepdims=True)\nev2 = np.sum(-v2 * e / v2scale, axis=1, keepdims=True)\nev3 = np.sum(-v3 * e / v3scale, axis=1, keepdims=True)\n\nev1 = np.arccos(ev1)\nev2 = np.arccos(ev2)\nev3 = np.arccos(ev3)\n\nvv12 = np.sum(v1 * v2 / v1scale / v2scale, axis=1, keepdims=True)\nvv23 = np.sum(v2 * v3 / v2scale / v3scale, axis=1, keepdims=True)\nvv31 = np.sum(v3 * v1 / v3scale / v1scale, axis=1, keepdims=True)\n\nvavg = np.log10(np.mean((v1scale, v2scale, v3scale), axis=0))\nevvv = np.mean((ev1, ev2, ev3), axis=0)\nvvvv = np.mean((vv12, vv23, vv31), axis=0)\n\npos = np.mean(traj[:, :, :3], axis=1)\nxyzq = reco_df[[\"qx\", \"qy\", \"qz\"]].to_numpy()\ndistq = pos - xyzq\ndistq = np.linalg.norm(distq, axis=1, keepdims=True) + 1e-3\n\n# load the model and predict\nLOGGER.info(\"Loading BDT model...\")\nclf = pickle.load(open(FILE_BDT, 'rb'))\nLOGGER.info(\"Predicting...\")\nX = np.concatenate([reco, xe, ev1, ev2, ev3, vavg, evvv, vvvv, distq], axis=1)\nX[np.isnan(X)] = 0\ny_hat = clf.decision_function(X) > CUT_BDT\ny_hat = y_hat[:, np.newaxis]\n\ngnn = reco_df[[\"azimuth\", \"zenith\"]].to_numpy()\nfit = vector2angles(proj)\nbdt = (gnn * ~y_hat) + fit * y_hat\n\nLOGGER.info(f\"GNN prediction\\n{gnn}\")\nLOGGER.info(f\"Plane prediction\\n{fit}\")\nLOGGER.info(f\"BDT prediction\\n{bdt}\")\n\nsubmit_df = pd.DataFrame(bdt, columns=[\"azimuth\", \"zenith\"])\n\nsubmit_df[\"event_id\"] = reco_df.event_id.values\nsubmit_df = submit_df.set_index(\"event_id\")\nsubmit_df = submit_df.sort_values([\"event_id\"]) # sort by event_id\nsubmit_df.to_csv(os.path.join(OUTPUT_PATH, \"submission.csv\"))\nprint(submit_df)","metadata":{},"execution_count":null,"outputs":[]}]}