{"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":"import numpy as np\nimport pandas as pd\nimport polars as pl\nimport os\nimport time\nfrom tqdm.auto import tqdm\nimport numba as nb\n\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, TensorDataset, DataLoader\nimport math\n\nimport json\n\nfrom sklearn.ensemble import HistGradientBoostingClassifier\nimport joblib","metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","papermill":{"duration":3.040664,"end_time":"2023-04-08T09:58:26.942044","exception":false,"start_time":"2023-04-08T09:58:23.901380","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-19T11:41:26.607601Z","iopub.execute_input":"2023-04-19T11:41:26.608008Z","iopub.status.idle":"2023-04-19T11:41:26.617012Z","shell.execute_reply.started":"2023-04-19T11:41:26.607958Z","shell.execute_reply":"2023-04-19T11:41:26.615687Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def angular_dist_score(az_true, zen_true, az_pred, zen_pred, mean=True):\n    \"\"\" https://www.kaggle.com/code/sohier/mean-angular-error \"\"\"\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))) if mean else np.abs(np.arccos(scalar_prod))","metadata":{"papermill":{"duration":0.017828,"end_time":"2023-04-08T09:58:26.964704","exception":false,"start_time":"2023-04-08T09:58:26.946876","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-19T11:41:26.619333Z","iopub.execute_input":"2023-04-19T11:41:26.620296Z","iopub.status.idle":"2023-04-19T11:41:26.629726Z","shell.execute_reply.started":"2023-04-19T11:41:26.620254Z","shell.execute_reply":"2023-04-19T11:41:26.628683Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DATA_DIR = \"/kaggle/input/icecube-neutrinos-in-deep-ice/\"\nPREP_DIR = \"/kaggle/input/icecube-preprocessed-data/\"\nTRAIN_META_FORMAT = \"/kaggle/input/train-meta-parquet/train_meta_{:d}.parquet\"\nMODEL_DIR = \"/kaggle/input/icecube-models-final/\"","metadata":{"papermill":{"duration":0.012102,"end_time":"2023-04-08T09:58:26.981305","exception":false,"start_time":"2023-04-08T09:58:26.969203","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-19T11:41:26.631218Z","iopub.execute_input":"2023-04-19T11:41:26.632137Z","iopub.status.idle":"2023-04-19T11:41:26.646057Z","shell.execute_reply.started":"2023-04-19T11:41:26.632099Z","shell.execute_reply":"2023-04-19T11:41:26.645087Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"VALIDATE = False # For checking on validation data\n\nif not VALIDATE:\n    PARQUETS_DIR = os.path.join(DATA_DIR + 'test')\n    BATCH_LIST = list(sorted(os.listdir(PARQUETS_DIR)))\n    metadata = pl.read_parquet(f'{DATA_DIR}/test_meta.parquet')\n    CHECK_PREDICTION = False\nelse:\n    PARQUETS_DIR = os.path.join(DATA_DIR + 'train')\n    vbatches = [655]\n    BATCH_LIST = [f'batch_{vb}.parquet' for vb in vbatches]\n    META_FILES = [f'train_meta_{vb}.parquet' for vb in vbatches]\n    def read_metadata():\n        meta = []\n        for mf, vb in zip(META_FILES, vbatches):\n            bmeta = pl.read_parquet(TRAIN_META_FORMAT.format(vb))\n            # Polars isn't going to get much adoption with stupid syntax like this :facepalm:\n            bmeta = bmeta.with_columns(pl.lit(vb).alias('batch_id'))\n            meta.append(bmeta)\n        return pl.concat(meta)\n    metadata = read_metadata()\n    CHECK_PREDICTION = True\n    \nGEOMETRY = os.path.join(PREP_DIR, \"sensor_geometry_with_transparency.csv\")\ngeometry = pl.scan_csv(GEOMETRY).with_columns(\n                [pl.col('sensor_id').cast(pl.Int16)]\n            )\n    \nNUM_BINS = 128\nFEATURE_NAMES = ['time', 'charge', 'auxiliary', 'x', 'y', 'z', 'qe', 'scatter', 'absorp']\nCHARGE_IDX = FEATURE_NAMES.index('charge')\nTIME_IDX = FEATURE_NAMES.index('time')\nAUX_IDX = FEATURE_NAMES.index('auxiliary')\nN_FEATURES = len(FEATURE_NAMES)\nMAX_SEQUENCE_LENGTH = 256\nLONG_SEQ_MAX_LENGTH = 3072\nBATCH_SIZE = 500\nLONG_SEQ_BATCH_SIZE = 10\nULTTRA_LONG_RESAMPLE = 0\n\nMAX_EVENTS = 200_000 if VALIDATE else 1000000000000\n\ndevice = 'cuda' if torch.cuda.is_available() else 'cpu'\n\nwith open(os.path.join(PREP_DIR, f'angle_bins_{NUM_BINS}.json')) as fp:\n    bin_data = json.load(fp)\n    \nazimuth_bin_centers = torch.tensor(bin_data['azimuth_bin_centers']).type(torch.float32).to(device)\n\nzenith_centers = np.array(bin_data['zenith_bin_centers'])\nkernel_length = 15\nnum_zenith_padding = (kernel_length - 1)//2\npadded_zenith_bins = np.concatenate([\n        -zenith_centers[num_zenith_padding-1 : : -1],\n        zenith_centers,\n        2 * np.pi - zenith_centers[-1 : -num_zenith_padding-1 : -1],\n])\nzenith_bin_centers = torch.tensor(padded_zenith_bins).type(torch.float32).to(device)\nZENITH_NUM_BINS = len(zenith_bin_centers)\n\n# Model configs\n\n# 15 layer\nclass m3_15l_vmf:\n    name = 'v43vmf_ep349.ckpt'\n    n_embd = 512 + 128 + 128 + 14\n    middle_features = 8192\n    neck_features = 2048\n    bias = True\n    dropout = 0.0\n    unwanted_prefix = 'model'\n    batchnorm = True\n\nclass m3_15l_seq3072:\n    name = 'v99_v13_l3072_ep2.ckpt'\n    n_embd = [512]*15\n    n_heads = [8]*15\n    bias = False\n    dropout = 0.0\n    neck_dropout = 0.0\n    neck_features = 3072\n    unwanted_prefix = 'model'\n\nclass m3_15l:\n    name = 'v13r2_m3_s256_ep48.ckpt'\n    n_embd = [512]*15\n    n_heads = [8]*15\n    bias = False\n    dropout = 0.0\n    neck_dropout = 0.0\n    neck_features = 3072\n    unwanted_prefix = 'model'\n    long_seq_model_config = m3_15l_seq3072()\n    stack_model_config = m3_15l_vmf()\n    mixer_model = 'gboost_mixer_m315l_v13v43_sk1.0.2.pkl'\n    selector_model = 'gboost_selector_m315l_v13v43_sk1.0.2.pkl'\n\n# 18 layer\nclass m3_18l_vmf:\n    name = 'v49vmf_m3_18l_ep349.ckpt'\n    n_embd = 512 + 128 + 128 + 14\n    middle_features = 8192\n    neck_features = 2048\n    bias = True\n    dropout = 0.0\n    unwanted_prefix = 'model'\n    batchnorm = False\n\nclass m3_18l_seq3072:\n    name = 'v98_v97ep187_l3072_ep6.ckpt'\n    n_embd = [512]*18\n    n_heads = [8]*18\n    bias = False\n    dropout = 0.0\n    neck_dropout = 0.0\n    neck_features = 3072\n    unwanted_prefix = 'model'\n\nclass m3_18l:\n    name = 'v97r5_m3_s256_ep187.ckpt'\n    n_embd = [512]*18\n    n_heads = [8]*18\n    bias = False\n    dropout = 0.0\n    neck_dropout = 0.0\n    neck_features = 3072\n    unwanted_prefix = 'model'\n    long_seq_model_config = m3_18l_seq3072()\n    stack_model_config = m3_18l_vmf()\n    mixer_model = 'gboost_mixer_m318l_n2_v49_sk1.0.2.pkl'\n    selector_model = 'gboost_selector_m318l_n2_v49_sk1.0.2.pkl'\n\n# 18 layer different epoch\nclass m3_r6_vmf:\n    name = 'v3_vmfzn_m3r6_ep699.ckpt'\n    n_embd = 512 + 128 + 128 + 14\n    middle_features = 8192\n    neck_features = 2048\n    bias = True\n    dropout = 0.0\n    unwanted_prefix = 'model'\n    batchnorm = False\n\nclass m3_r6_seq3072:\n    name = 'v103_v99r6_l3072_ep7.ckpt'\n    n_embd = [512]*18\n    n_heads = [8]*18\n    bias = False\n    dropout = 0.0\n    neck_dropout = 0.0\n    neck_features = 3072\n    unwanted_prefix = 'model'\n\nclass m3_r6:\n    name = 'v99r6_m3_s256_ep26.ckpt'\n    n_embd = [512]*18\n    n_heads = [8]*18\n    bias = False\n    dropout = 0.0\n    neck_dropout = 0.0\n    neck_features = 3072\n    unwanted_prefix = 'model'\n    long_seq_model_config = m3_r6_seq3072()\n    stack_model_config = m3_r6_vmf()\n    mixer_model = 'gboost_mixer_v99r6_sk1.0.2.pkl'\n    selector_model = 'gboost_selector_v99r6_sk1.0.2.pkl'\n\n# MODEL_CONFIGS = [m3_18l()]\nMODEL_CONFIGS = [m3_18l(), m3_15l(), m3_r6()]\nMERGER_MODELS = ['gboost_merger_18ln2_15l.pkl', 'gboost_merger_15l_r6.pkl']","metadata":{"papermill":{"duration":3.225329,"end_time":"2023-04-08T09:58:30.211088","exception":false,"start_time":"2023-04-08T09:58:26.985759","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-19T11:41:26.648722Z","iopub.execute_input":"2023-04-19T11:41:26.649259Z","iopub.status.idle":"2023-04-19T11:41:26.687960Z","shell.execute_reply.started":"2023-04-19T11:41:26.649219Z","shell.execute_reply":"2023-04-19T11:41:26.686936Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def set_seed(value):\n    np.random.seed(value)\n\ndef sample_and_pad(data, pulse_indexes, padded_sequence_length):\n    data[:, CHARGE_IDX] = np.log10(data[:, CHARGE_IDX]) / 3.0\n    data[:, AUX_IDX] = data[:, AUX_IDX] - 0.5\n    data_x = np.zeros((len(pulse_indexes), padded_sequence_length, data.shape[-1]), dtype=np.float32)\n    sequence_lengths = np.zeros(len(pulse_indexes), dtype=np.int32)\n    for ii in range(len(pulse_indexes)):\n        event_data = data[pulse_indexes[ii, 0] : pulse_indexes[ii, 1] + 1]\n        if len(event_data) > padded_sequence_length:\n            naux_idx = np.where(event_data[:, AUX_IDX] == -0.5)[0]\n            aux_idx = np.where(event_data[:, AUX_IDX] == 0.5)[0]\n            if len(naux_idx) < padded_sequence_length:\n                max_length_possible = min(padded_sequence_length, len(event_data))\n                num_to_sample = max_length_possible - len(naux_idx)\n                aux_idx_sample = np.random.choice(aux_idx, size=num_to_sample, replace=False)\n                selected_idx = np.concatenate((naux_idx, aux_idx_sample))\n            else:\n                selected_idx = np.random.choice(naux_idx, size=padded_sequence_length, replace=False)\n            selected_idx = np.sort(selected_idx)\n            event_data = event_data[selected_idx]\n        event_data[:, TIME_IDX] = ( event_data[:, TIME_IDX] - event_data[:, TIME_IDX].min() ) / 3e4\n        assert np.all(np.isfinite(event_data))\n        data_x[ii, :len(event_data), :] = event_data\n        sequence_lengths[ii] = len(event_data)                       \n    return data_x, sequence_lengths","metadata":{"papermill":{"duration":5.646776,"end_time":"2023-04-08T09:58:35.862649","exception":false,"start_time":"2023-04-08T09:58:30.215873","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-19T11:41:26.725439Z","iopub.execute_input":"2023-04-19T11:41:26.725745Z","iopub.status.idle":"2023-04-19T11:41:26.737293Z","shell.execute_reply.started":"2023-04-19T11:41:26.725717Z","shell.execute_reply":"2023-04-19T11:41:26.736044Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def preprocess_data(bfile, seed, \n                    min_sequence_length, max_sequence_length, \n                    padded_sequence_length, print_time=False):\n    maybe_print = print if print_time else (lambda *args: '')\n    set_seed(seed)\n    maybe_print(\"Reading\", bfile)\n    start_time = time.perf_counter()\n    batch_id = int(bfile.split('.')[0].split('_')[-1])\n    batch = pl.scan_parquet(f'{PARQUETS_DIR}/{bfile}')\n    batch = batch.join(geometry, on='sensor_id', how='left')\n    batch_meta = metadata.filter(pl.col('batch_id') == batch_id)\n\n    data = batch.select(FEATURE_NAMES).collect().to_numpy()\n    pulse_indexes = batch_meta.select(['first_pulse_index', 'last_pulse_index']).to_numpy()[:MAX_EVENTS]\n    event_lengths = pulse_indexes[:, 1] - pulse_indexes[:, 0] + 1\n    filter_events = (event_lengths > min_sequence_length) & (event_lengths <= max_sequence_length)\n    event_to_keep = np.where(filter_events)[0]\n    pulse_indexes = pulse_indexes[event_to_keep]\n    maybe_print(\"Read and merge\", bfile, \"in\", time.perf_counter() - start_time, \"s\")\n    \n    start_time = time.perf_counter()\n    data_x, seq_lens = sample_and_pad(data, pulse_indexes, padded_sequence_length)\n    maybe_print(\"Processed\", bfile, \"in\", time.perf_counter() - start_time, \"s\")\n\n    return data_x, seq_lens, event_to_keep","metadata":{"papermill":{"duration":0.017797,"end_time":"2023-04-08T09:58:35.885279","exception":false,"start_time":"2023-04-08T09:58:35.867482","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-19T11:41:26.739972Z","iopub.execute_input":"2023-04-19T11:41:26.740422Z","iopub.status.idle":"2023-04-19T11:41:26.752482Z","shell.execute_reply.started":"2023-04-19T11:41:26.740365Z","shell.execute_reply":"2023-04-19T11:41:26.751358Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if VALIDATE:\n    def check_data_loading():\n        dx, dl, _ = preprocess_data(BATCH_LIST[0], seed=42, \n                                    min_sequence_length=0,\n                                    max_sequence_length=MAX_SEQUENCE_LENGTH,\n                                    padded_sequence_length=MAX_SEQUENCE_LENGTH,\n                                    print_time=True)\n        print(dx.shape)\n        dx, dl, _ = preprocess_data(BATCH_LIST[0], seed=42, \n                                    min_sequence_length=MAX_SEQUENCE_LENGTH,\n                                    max_sequence_length=np.inf,\n                                    padded_sequence_length=LONG_SEQ_MAX_LENGTH,\n                                    print_time=True)\n        print(dx.shape)\n\n    check_data_loading()","metadata":{"papermill":{"duration":0.014868,"end_time":"2023-04-08T09:58:35.904656","exception":false,"start_time":"2023-04-08T09:58:35.889788","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-19T11:41:26.754052Z","iopub.execute_input":"2023-04-19T11:41:26.754683Z","iopub.status.idle":"2023-04-19T11:41:26.766759Z","shell.execute_reply.started":"2023-04-19T11:41:26.754646Z","shell.execute_reply":"2023-04-19T11:41:26.765715Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class IceCubeDataset(Dataset):\n    def __init__(self, bfile, seed, min_sequence_lenth, max_sequence_length, padded_sequence_length):\n        super().__init__()\n        dx, sl, events = preprocess_data(bfile, seed, \n                                         min_sequence_lenth,\n                                         max_sequence_length, \n                                         padded_sequence_length)\n        self.x = torch.Tensor(dx)\n        self.l = torch.Tensor(sl)\n        \n        self.sort_idx = np.argsort(sl)\n        self.reverse_sort_idx = np.argsort(self.sort_idx)\n\n        self.events = events\n\n    def __len__(self):\n        return len(self.x)\n    \n    def __getitem__(self, index):\n        si = self.sort_idx[index] # Lazy packing, works well for batch size ~1000\n        return self.x[si], self.l[si]","metadata":{"papermill":{"duration":0.014778,"end_time":"2023-04-08T09:58:35.923842","exception":false,"start_time":"2023-04-08T09:58:35.909064","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-19T11:41:26.768316Z","iopub.execute_input":"2023-04-19T11:41:26.768949Z","iopub.status.idle":"2023-04-19T11:41:26.780725Z","shell.execute_reply.started":"2023-04-19T11:41:26.768910Z","shell.execute_reply":"2023-04-19T11:41:26.779725Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class LayerNorm(nn.Module):\n    \"\"\" LayerNorm but with an optional bias. PyTorch doesn't support simply bias=False \"\"\"\n\n    def __init__(self, ndim, bias=False):\n        super().__init__()\n        self.weight = nn.Parameter(torch.ones(ndim))\n        self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None\n\n    def forward(self, input):\n        return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)\n\ndef mlp(n_embd, bias=False, dropout=0.0, out_embd=None):\n    out_embd = n_embd if out_embd is None else out_embd\n    return nn.Sequential(\n        nn.Linear(n_embd, 4 * n_embd, bias=bias),\n        nn.GELU(approximate='tanh'),\n        nn.Linear(4 * n_embd, out_embd, bias=bias),\n        nn.Dropout(dropout)\n    )\n\n\nclass SelfAttention(nn.Module):\n    def __init__(self, prev_emdb, n_embd, n_heads, bias=False, dropout=0.0):\n        super().__init__()\n        self.prev_embd = prev_emdb\n        self.n_embd = n_embd\n        self.n_heads = n_heads\n        \n        self.c_attn = nn.Linear(prev_emdb, 3 * n_embd, bias=bias)\n        self.c_proj = nn.Linear(n_embd, n_embd, bias=bias)\n        self.dropout = dropout\n        self.resid_dropout = nn.Dropout(dropout)\n\n    def forward(self, x, attn_mask):\n        B, T, _ = x.shape\n        C = self.n_embd\n        q, k ,v  = self.c_attn(x).split(self.n_embd, dim=2)\n        k = k.view(B, T, self.n_heads, C // self.n_heads).transpose(1, 2) # (B, nh, T, hs)\n        q = q.view(B, T, self.n_heads, C // self.n_heads).transpose(1, 2) # (B, nh, T, hs)\n        v = v.view(B, T, self.n_heads, C // self.n_heads).transpose(1, 2) # (B, nh, T, hs)\n        if hasattr(F, 'scaled_dot_product_attention'):\n            y = F.scaled_dot_product_attention(q, k, v, \n                    attn_mask=attn_mask, dropout_p=self.dropout, is_causal=False)\n        else:\n            y = F._scaled_dot_product_attention(q, k, v, \n                    attn_mask=attn_mask, dropout_p=self.dropout, is_causal=False)[0]\n        y = y.transpose(1, 2).contiguous().view(B, T, C)\n        y = self.c_proj(y)\n        y = self.resid_dropout(y)\n        return y\n    \nclass AttentionBlock(nn.Module):\n    def __init__(self, prev_embd, n_embd, n_heads, bias=False, dropout=0.0):\n        super().__init__()\n        self.ln_1 = LayerNorm(prev_embd, bias)\n        assert n_embd % prev_embd == 0, f\"{prev_embd} {n_embd} should be divisble\"\n        self.attn = SelfAttention(prev_embd, n_embd, n_heads, bias, dropout)\n        self.ln_2 = LayerNorm(n_embd, bias)\n        self.mlp = mlp(n_embd, bias, dropout)\n\n    def forward(self, x, attn_mask):\n        x = x + self.attn(self.ln_1(x), attn_mask)\n        x = x + self.mlp(self.ln_2(x))\n        return x\n\nclass AttentionEncoder(nn.Module):\n    def __init__(self, config):\n        super().__init__()\n        dropout = config.dropout\n        bias = config.bias\n        \n        attn_layers = []\n        prev_embd = config.n_embd[0]\n        for n_embd, n_heads in zip(config.n_embd, config.n_heads):\n            attn_layers.append( AttentionBlock(prev_embd, n_embd, n_heads, bias, dropout) )\n            prev_embd = n_embd\n        self.attn = nn.ModuleList(attn_layers)\n    \n    def forward(self, x, attn_mask):\n        out = x\n        for attn_layer in self.attn:\n            out = attn_layer(out, attn_mask)\n        \n        return out\n\n    \nclass SequencePool(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.num_pools = 1\n    \n    def forward(self, x, sequence_lengths, padding_mask):\n        sumf = torch.sum(x * padding_mask.unsqueeze(2), dim=1) # Mask padded tokens\n        meanf = sumf / sequence_lengths.view(-1, 1) # Normalize avg pool values by seq length\n        out = meanf\n        return out\n    \nclass Neck(nn.Module):\n    def __init__(self, in_features, out_features, bias, dropout):\n        super().__init__()\n        self.mlp = nn.Sequential(\n            LayerNorm(in_features, bias=bias),\n            nn.Linear(in_features, 4 * in_features, bias=bias),\n            nn.GELU(approximate='tanh'),\n            nn.Linear(4 * in_features, out_features, bias=bias),\n            nn.Dropout(dropout)\n        )\n        self.n_repeats = out_features // in_features\n\n    def forward(self, x):\n        return x.repeat(1, self.n_repeats) + self.mlp(x)\n\n    \nclass MultiLabelClassifier(nn.Module):    \n    def __init__(self, n_features, max_block_size, num_classes, zenith_num_classes, config):\n        super().__init__()\n\n        self.inp = nn.Linear(n_features, config.n_embd[0])\n        self.drop_inputs = nn.Dropout(config.dropout)\n\n        self.encoder = AttentionEncoder(config)\n        \n        self.pool = SequencePool()\n\n        num_out_features = config.n_embd[-1] * self.pool.num_pools\n\n        self.neck_az = Neck(num_out_features, config.neck_features, config.bias, config.neck_dropout)\n        self.neck_zn = Neck(num_out_features, config.neck_features, config.bias, config.neck_dropout)\n        \n        self.azimuth = nn.Linear(config.neck_features, num_classes)\n        self.zenith = nn.Linear(config.neck_features, zenith_num_classes)\n\n    def get_masks(self, x, l):\n        key_padding_mask = torch.arange(x.shape[1]).view(1, -1).to(l.device) < l.view(-1, 1)\n        attn_mask = (key_padding_mask.unsqueeze(1) == key_padding_mask.unsqueeze(2)).unsqueeze(1)  # (B, 1, T, T)\n        return key_padding_mask, attn_mask\n\n    def forward(self, x):\n        inputs, seq_lengths = x\n        out = self.inp(inputs)\n        out = self.drop_inputs(out)\n        key_padding_mask, attn_mask = self.get_masks(inputs, seq_lengths)\n        out = self.encoder(out, attn_mask)\n        pool = self.pool(out, seq_lengths, key_padding_mask)\n    \n        az_out = self.azimuth(self.neck_az(pool))\n        zn_out = self.zenith(self.neck_zn(pool))\n        return az_out, zn_out, pool","metadata":{"papermill":{"duration":0.037914,"end_time":"2023-04-08T09:58:35.966121","exception":false,"start_time":"2023-04-08T09:58:35.928207","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-19T11:41:26.783938Z","iopub.execute_input":"2023-04-19T11:41:26.784406Z","iopub.status.idle":"2023-04-19T11:41:26.814907Z","shell.execute_reply.started":"2023-04-19T11:41:26.784377Z","shell.execute_reply":"2023-04-19T11:41:26.813891Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class StackNeck(nn.Module):\n    def __init__(self, in_features, middle_features, out_features, bias, dropout, batchnorm):\n        super().__init__()\n        if batchnorm:\n            self.mlp = nn.Sequential(\n                nn.Linear(in_features, middle_features, bias=bias),\n                nn.BatchNorm1d(middle_features),\n                nn.GELU(approximate='tanh'),\n                nn.Dropout(dropout),\n                nn.Linear(middle_features, out_features, bias=bias),\n                nn.BatchNorm1d(out_features),\n                nn.GELU(approximate='tanh'),\n                nn.Dropout(dropout),\n            )\n        else:\n            self.mlp = nn.Sequential(\n                nn.Linear(in_features, middle_features, bias=bias),\n                nn.LayerNorm(middle_features),\n                nn.GELU(approximate='tanh'),\n                nn.Dropout(dropout),\n                nn.Linear(middle_features, out_features, bias=bias),\n                nn.LayerNorm(out_features),\n                nn.GELU(approximate='tanh'),\n                nn.Dropout(dropout),\n            )\n\n    def forward(self, x):\n        return self.mlp(x)\n    \nclass StackModel(nn.Module):\n    def __init__(self, num_classes, zenith_num_classes, config):\n        super().__init__()\n\n        self.neck = StackNeck(config.n_embd, config.middle_features, config.neck_features, \n                              config.bias, config.dropout, config.batchnorm)\n        self.xyz = nn.Linear(config.neck_features, 3)\n\n    def forward(self, x):\n        return self.xyz(self.neck(x))","metadata":{"execution":{"iopub.status.busy":"2023-04-19T11:41:26.816395Z","iopub.execute_input":"2023-04-19T11:41:26.816749Z","iopub.status.idle":"2023-04-19T11:41:26.830151Z","shell.execute_reply.started":"2023-04-19T11:41:26.816712Z","shell.execute_reply":"2023-04-19T11:41:26.828952Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def prepare_submission(event_ids, azimuth, zenith, validate_ids=False):\n    if validate_ids:\n        sample_submission = pd.read_parquet(os.path.join(DATA_DIR, 'sample_submission.parquet'))\n        assert np.array_equal(event_ids, sample_submission.event_id.values)\n    zenith_clipped = torch.clip(zenith, 0.0, np.pi)\n    submission_df = pd.DataFrame(\n        {\n            'event_id': event_ids,\n            'azimuth': azimuth.cpu().numpy(),\n            'zenith': zenith_clipped.cpu().numpy(),\n        }\n    ).set_index('event_id')\n    return submission_df","metadata":{"papermill":{"duration":0.014677,"end_time":"2023-04-08T09:58:35.985181","exception":false,"start_time":"2023-04-08T09:58:35.970504","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-19T11:41:26.831615Z","iopub.execute_input":"2023-04-19T11:41:26.832192Z","iopub.status.idle":"2023-04-19T11:41:26.845574Z","shell.execute_reply.started":"2023-04-19T11:41:26.832155Z","shell.execute_reply":"2023-04-19T11:41:26.844091Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def load_model(model_config):\n    model = MultiLabelClassifier(n_features=N_FEATURES, \n                                max_block_size=MAX_SEQUENCE_LENGTH,\n                                num_classes=NUM_BINS,\n                                zenith_num_classes=ZENITH_NUM_BINS,\n                                config=model_config)\n    checkpoint_path = os.path.join(MODEL_DIR, model_config.name)\n    state_dict = torch.load(checkpoint_path)['state_dict']\n    old_keys = list(state_dict.keys())\n    # print(\"Removing prefix\", model_config.unwanted_prefix)\n    for key in old_keys:\n        if model_config.unwanted_prefix in key:\n            new_key = key.split(model_config.unwanted_prefix)[1][1:]\n            state_dict[new_key] = state_dict.pop(key)\n    model.load_state_dict(state_dict)\n    model.eval()\n    model.to(device)\n    return model\n\ndef load_stack_model(model_config):\n    stack_model = StackModel(num_classes=NUM_BINS, zenith_num_classes=ZENITH_NUM_BINS, config=model_config)\n    checkpoint_path = os.path.join(MODEL_DIR, model_config.name)\n    state_dict = torch.load(checkpoint_path)['state_dict']\n    old_keys = list(state_dict.keys())\n    for key in old_keys:\n        if model_config.unwanted_prefix in key:\n            new_key = key.split(model_config.unwanted_prefix)[1][1:]\n            state_dict[new_key] = state_dict.pop(key)\n    stack_model.load_state_dict(state_dict)\n    stack_model.eval()\n    stack_model.to(device)\n    return stack_model","metadata":{"papermill":{"duration":0.013723,"end_time":"2023-04-08T09:58:36.003268","exception":false,"start_time":"2023-04-08T09:58:35.989545","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-19T11:41:26.847106Z","iopub.execute_input":"2023-04-19T11:41:26.847724Z","iopub.status.idle":"2023-04-19T11:41:26.859979Z","shell.execute_reply.started":"2023-04-19T11:41:26.847679Z","shell.execute_reply":"2023-04-19T11:41:26.859243Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"### Classes to Angles\n\n# common\ndef simple_average(pred_mbs):\n    return torch.mean( torch.stack(pred_mbs, dim=2), dim=2)\n\ndef argmax_average(pred_mbs, centers):\n    pred_classes = []\n    for pred in pred_mbs:\n        pred_classes.append(centers[pred.argmax(axis=1)].unsqueeze(0))\n    angles = torch.mean(torch.cat(pred_classes), dim=0)\n    return angles\n\ndef simpleavg_argmax(pred_mbs, centers):\n    pred = simple_average(pred_mbs)\n    angle = centers[pred.argmax(dim=1)]\n    return angle\n\n# azimuth\nclass AzimuthXY:\n    def __init__(self):\n        self.azx = torch.cos(azimuth_bin_centers) \n        self.azy = torch.sin(azimuth_bin_centers)\n\n    def __call__(self, az_softmax):\n        self.azx = self.azx.to(az_softmax.device)\n        self.azy = self.azy.to(az_softmax.device)\n        return az_softmax * self.azx, az_softmax * self.azy\n\naz_xy = AzimuthXY()\n\n\ndef azimuth_vectorsum(azx, azy):\n    azmx, azmy = torch.sum(azx, dim=1), torch.sum(azy, dim=1)\n    azn = torch.sqrt(azmx**2 + azmy**2)\n    az_pred = ( torch.arccos(azmx / azn) * torch.sign(azmy) ) % (np.pi * 2)\n    return az_pred\n\ndef az_simpleavg_vectorsum(pred_mbs):\n    azsf = simple_average(pred_mbs)\n    azx, azy = az_xy(azsf)\n    az = azimuth_vectorsum(azx, azy)\n    return az\n\n# zenith\ndef zn_argmax_average(pred_mbs):\n    return argmax_average(pred_mbs, zenith_bin_centers)\n\n# predictions to angles\ndef ensemble_predictions(az_pred_mbs, zn_pred_mbs):\n    with torch.no_grad():\n        az_pred = az_simpleavg_vectorsum(az_pred_mbs)\n        zn_pred = zn_argmax_average(zn_pred_mbs)\n    return az_pred, zn_pred\n\ndef topksum(pred, centers, k=3):\n    zn_topk, zn_topk_idk = torch.topk(pred, k=k, dim=1)\n    zn_topk_smax = torch.softmax(zn_topk, axis=1)\n    angle = torch.sum(zn_topk_smax * centers[zn_topk_idk], dim=1)\n    return angle\n\n# single classifier to predictor\ndef discrete_to_angle(az_softmax,\n                      zn_pred,\n                      azimuth_bin_centers,\n                      zenith_bin_centers):\n    \n    azx, azy = torch.cos(azimuth_bin_centers), torch.sin(azimuth_bin_centers)\n    azmx, azmy = az_softmax @ azx, az_softmax @ azy\n    azn = torch.sqrt(azmx**2 + azmy**2)\n    az_pred_center = ( torch.arccos(azmx / azn) * torch.sign(azmy) ) % (np.pi * 2)\n    \n    zn_pred_center = zenith_bin_centers[zn_pred.argmax(1)]\n    \n    return az_pred_center, zn_pred_center\n\n# vmf predictions to angle\ndef xyz_to_angle(xyz):\n    z_normed = xyz[:, 2] / torch.norm(xyz ,dim=1)\n    x_normed = xyz[:, 0] / torch.norm(xyz[:, :2] ,dim=1)\n    azimuth = ( torch.arccos(x_normed) * torch.sign(xyz[:, 1]) ) % (np.pi * 2)\n    zenith = torch.arccos(z_normed)\n    return azimuth, zenith\n\ndef avg_az_np(az1, az2):\n    return np.arctan2(np.sin(az1) + np.sin(az2), np.cos(az1) + np.cos(az2)) % (np.pi * 2)\n\ndef avg_az(az1, az2):\n    return torch.atan2(torch.sin(az1) + torch.sin(az2), torch.cos(az1) + torch.cos(az2)) % (np.pi * 2)\n\ndef az_diff(az1, az2):\n    res = ( az1 - az2 ) % (np.pi * 2)\n    res[res > np.pi] = np.pi * 2 - res[res > np.pi]\n    return res","metadata":{"papermill":{"duration":0.026232,"end_time":"2023-04-08T09:58:36.033829","exception":false,"start_time":"2023-04-08T09:58:36.007597","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-19T11:41:26.861617Z","iopub.execute_input":"2023-04-19T11:41:26.862311Z","iopub.status.idle":"2023-04-19T11:41:26.882578Z","shell.execute_reply.started":"2023-04-19T11:41:26.862272Z","shell.execute_reply":"2023-04-19T11:41:26.881545Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def check_score(batch_id, az_pred, zn_pred, text, events_to_keep=None, convert_to_angle=True):\n    if not VALIDATE:\n        return\n    if events_to_keep is None:\n        events_to_keep = np.arange(len(az_pred), dtype=np.int32)\n    if convert_to_angle:\n        az, zn = ensemble_predictions([az_pred], [zn_pred])\n    else:\n        az, zn = az_pred, zn_pred\n    az_gt = metadata.filter(pl.col('batch_id') == batch_id).select('azimuth').to_numpy().squeeze()\n    zen_gt = metadata.filter(pl.col('batch_id') == batch_id).select('zenith').to_numpy().squeeze()\n    angular_dist = angular_dist_score(az_gt[events_to_keep], zen_gt[events_to_keep], \n                                      az.cpu().numpy(), zn.cpu().numpy())\n    print(f\"ang_dist {text}\", angular_dist, \"\\n\")\n    return angular_dist","metadata":{"papermill":{"duration":0.014168,"end_time":"2023-04-08T09:58:36.071813","exception":false,"start_time":"2023-04-08T09:58:36.057645","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-19T11:41:26.885671Z","iopub.execute_input":"2023-04-19T11:41:26.886190Z","iopub.status.idle":"2023-04-19T11:41:26.896649Z","shell.execute_reply.started":"2023-04-19T11:41:26.886138Z","shell.execute_reply":"2023-04-19T11:41:26.895670Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def model_load_check():\n    for model_config in MODEL_CONFIGS:\n        model = load_model(model_config)\n        stack_model = load_stack_model(model_config.stack_model_config)\n        del model, stack_model\nmodel_load_check()","metadata":{"papermill":{"duration":3.096961,"end_time":"2023-04-08T09:58:39.173127","exception":false,"start_time":"2023-04-08T09:58:36.076166","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-19T11:41:26.898311Z","iopub.execute_input":"2023-04-19T11:41:26.898659Z","iopub.status.idle":"2023-04-19T11:41:30.091320Z","shell.execute_reply.started":"2023-04-19T11:41:26.898623Z","shell.execute_reply":"2023-04-19T11:41:30.090291Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"FNAMES = ['time', 'charge', 'auxiliary', 'x', 'y', 'z', 'string']\n\nCHARGE_IDX = FNAMES.index('charge')\nTIME_IDX = FNAMES.index('time')\nAUX_IDX = FNAMES.index('auxiliary')\nX_IDX = FNAMES.index('x')\nY_IDX = FNAMES.index('y')\nZ_IDX = FNAMES.index('z')\nSTRING_IDX = FNAMES.index('string')\n\nTDIFF = 1000\nZDIFF = 60/500\n\ndef generate_features(data, pulse_indexes):\n    feats = np.zeros((len(pulse_indexes), 17))\n    for ii in tqdm(range(len(pulse_indexes))):\n        event_data = data[pulse_indexes[ii, 0] : pulse_indexes[ii, 1] + 1]\n        event_data[TIME_IDX] -= event_data[:, TIME_IDX].min()\n        naux_event_data = event_data[event_data[:, AUX_IDX] == 0]\n\n        event_length = len(event_data)\n        num_naux = len(naux_event_data)\n        if num_naux == 0:\n            continue\n        pct_naux = num_naux/event_length\n\n        total_charge = naux_event_data[:, CHARGE_IDX].sum()\n        charge_std = naux_event_data[:, CHARGE_IDX].std()\n\n        time_min = naux_event_data[:, TIME_IDX].min()\n        time_mean = naux_event_data[:, TIME_IDX].mean()\n        time_max = naux_event_data[:, TIME_IDX].max()\n        time_std = naux_event_data[:, TIME_IDX].std()\n        time_ratio = time_min / time_max\n\n        z_mean = naux_event_data[:, Z_IDX].mean()\n        z_min = naux_event_data[:, Z_IDX].min()\n        z_max = naux_event_data[:, Z_IDX].max()\n        z_std = naux_event_data[:, Z_IDX].std()\n\n        last_tval = naux_event_data[0, TIME_IDX].copy()\n        last_xyz = naux_event_data[0, [X_IDX, Y_IDX, Z_IDX]].copy()\n        last_string = naux_event_data[0, STRING_IDX].copy()\n        for pulse in naux_event_data:\n            tdiff = (pulse[TIME_IDX] - last_tval)\n            zdiff = (pulse[Z_IDX] - last_xyz[2])\n            dist = ((pulse[[X_IDX, Y_IDX, Z_IDX]] - last_xyz)**2).sum()\n            speed = dist / (tdiff + 1e-6)\n            same_pillar = ( ( tdiff < TDIFF ) &\n                            ( abs(zdiff) < ZDIFF) &\n                            ( pulse[STRING_IDX] == last_string ) \n            )\n            if same_pillar:\n                last_tval = pulse[TIME_IDX].copy()\n                last_xyz = pulse[[X_IDX, Y_IDX, Z_IDX]].copy()\n                last_string = pulse[STRING_IDX].copy()\n            else:\n                break\n        else:\n            tdiff = 0\n            zdiff = 0\n\n        feats[ii] = np.array([num_naux, pct_naux,\n            total_charge, charge_std,\n            time_min, time_mean, time_max, time_std,\n            z_mean, z_min, z_max, z_std, time_ratio,\n            tdiff, zdiff, dist, speed\n        ])\n    return feats\n\ndef generate_batch_features(bfile, events_to_keep):\n    batch_id = int(bfile.split('.')[0].split('_')[-1])\n    batch = pl.scan_parquet(f'{PARQUETS_DIR}/{bfile}')\n    batch = batch.join(geometry, on='sensor_id', how='left')\n    batch_meta = metadata.filter(pl.col('batch_id') == batch_id)\n\n    data = batch.select(FNAMES).collect().to_numpy()\n    pulse_indexes = batch_meta.select(['first_pulse_index', 'last_pulse_index']).to_numpy()[events_to_keep]\n    event_lengths = pulse_indexes[:, 1] - pulse_indexes[:, 0] + 1\n\n    feats = generate_features(data, pulse_indexes)\n    return event_lengths, feats\n\ndef generate_boosting_features(event_lengths, batch_features, \n                          az_base, zn_base, az_stack, zn_stack,\n                          prediction_data):\n    \n    az_base_prob = prediction_data[0].cpu().numpy()\n    zn_base_prob = prediction_data[1].cpu().numpy()\n    xyz = prediction_data[2].cpu().numpy()\n\n    stack_az = az_stack.cpu().numpy()\n    stack_zn = zn_stack.cpu().numpy()\n    az_cls = np.digitize(stack_az, bins=azimuth_bin_centers.cpu().numpy(), right=False) - 1\n    zn_cls = np.digitize(stack_zn, bins=zenith_bin_centers.cpu().numpy(), right=False) - 1 + num_zenith_padding\n    a0 = np.arange(len(az_cls))\n    pred_features = np.stack([\n        az_base.cpu().numpy(),\n        zn_base.cpu().numpy(),\n        az_base_prob.max(1),\n        zn_base_prob.max(1),\n        az_base_prob[a0, az_cls],\n        zn_base_prob[a0, zn_cls],\n        az_base_prob[a0, az_cls]/az_base_prob.max(1),\n        zn_base_prob[a0, zn_cls]/zn_base_prob.max(1),\n        stack_az,\n        stack_zn,\n        1/np.sqrt(np.linalg.norm(xyz, axis=1)),\n        1/np.sqrt(np.linalg.norm(xyz[:, :2], axis=1))\n    ]).T\n    adists = angular_dist_score(az_base.cpu().numpy(), \n                                zn_base.cpu().numpy(), \n                                az_stack.cpu().numpy(), \n                                zn_stack.cpu().numpy(), \n                                mean=False)[:, None]\n    all_feats = np.concatenate(\n            [event_lengths[:, None], pred_features, adists, batch_features], axis=1\n        )\n    return all_feats, pred_features","metadata":{"execution":{"iopub.status.busy":"2023-04-19T11:41:30.095366Z","iopub.execute_input":"2023-04-19T11:41:30.095916Z","iopub.status.idle":"2023-04-19T11:41:30.121477Z","shell.execute_reply.started":"2023-04-19T11:41:30.095879Z","shell.execute_reply":"2023-04-19T11:41:30.120360Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"@torch.no_grad()\ndef predict_model(dataset, batch_size, model, stack_model):\n\n    az_pred_batch, zn_pred_batch = [], []\n    xyz_pred_batch = []\n    dataloader = DataLoader(dataset, batch_size=batch_size)\n    for x, l in tqdm(dataloader):\n        b_max_len = int(l.max()) # Lazy packing, works well for batch size ~1000\n        azp, znp, pool = model((x[:, :b_max_len].to(device), l.to(device)))\n        xyz = stack_model(torch.cat([pool, azp, znp], dim=1))\n        az_pred_batch.append(azp)\n        zn_pred_batch.append(znp)\n        xyz_pred_batch.append(xyz)\n    \n    az_pred_batch = torch.cat(az_pred_batch, dim=0)[dataset.reverse_sort_idx]\n    zn_pred_batch = torch.cat(zn_pred_batch, dim=0)[dataset.reverse_sort_idx]\n    xyz_pred_batch = torch.cat(xyz_pred_batch, dim=0)[dataset.reverse_sort_idx]\n\n    az_pred_batch = torch.softmax(az_pred_batch, axis=1)\n    zn_pred_batch = torch.softmax(zn_pred_batch, axis=1)\n\n    return az_pred_batch, zn_pred_batch, xyz_pred_batch\n\n@torch.no_grad()\ndef predict_on_batch(model, stack_model, dataset, batch_size, mutlisample=0, dataset_fn=None):\n    if len(dataset) == 0:\n        return None, None, None, None, (None, None, None)\n    az_pred_batch, zn_pred_batch, xyz_pred_batch = predict_model(dataset, batch_size, model, stack_model)\n    \n    # Mutli sampling ultra long batches and average the logits\n    if mutlisample > 0:\n        az_pred_batch, zn_pred_batch, xyz_pred_batch = [az_pred_batch], [zn_pred_batch], [xyz_pred_batch]\n        for idx in range(mutlisample):\n            dataset = dataset_fn(seed=42 + idx + 1)\n            azp, znp, xyz = predict_model(dataset, batch_size, model, stack_model)\n            az_pred_batch.append(azp)\n            zn_pred_batch.append(znp)\n            xyz_pred_batch.append(xyz)\n        az_base, zn_base = ensemble_predictions(az_pred_batch, zn_pred_batch)\n        az_pred_batch = torch.stack(az_pred_batch).mean(0)\n        zn_pred_batch = torch.stack(zn_pred_batch).mean(0)\n        xyz_pred_batch = torch.stack(xyz_pred_batch).mean(0)\n    else:\n        az_base, zn_base = discrete_to_angle(az_pred_batch, zn_pred_batch, azimuth_bin_centers, zenith_bin_centers)\n    \n    az_stack, zn_stack = xyz_to_angle(xyz_pred_batch)\n\n    return az_base, zn_base, az_stack, zn_stack, (az_pred_batch, zn_pred_batch, xyz_pred_batch)","metadata":{"execution":{"iopub.status.busy":"2023-04-19T11:41:30.124243Z","iopub.execute_input":"2023-04-19T11:41:30.125024Z","iopub.status.idle":"2023-04-19T11:41:30.137672Z","shell.execute_reply.started":"2023-04-19T11:41:30.124967Z","shell.execute_reply":"2023-04-19T11:41:30.136692Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def join_predictions(preds, long_preds, ultralong_preds, ds, long_ds, ultralong_ds):\n    n_events = len( set(ds.events).union(long_ds.events).union(ultralong_ds.events) )\n    if len(preds.shape) == 1:\n        joint_preds = torch.zeros(n_events, dtype=preds.dtype, device=preds.device)\n    else:\n        joint_preds = torch.zeros(n_events, preds.shape[1], dtype=preds.dtype, device=preds.device)\n    joint_preds[torch.tensor(ds.events).type(torch.long)] = preds.clone()\n    joint_preds[torch.tensor(long_ds.events).type(torch.long)] = long_preds.clone()\n    if ultralong_preds is not None:\n        joint_preds[torch.tensor(ultralong_ds.events).type(torch.long)] = ultralong_preds.clone()\n    return joint_preds","metadata":{"execution":{"iopub.status.busy":"2023-04-19T11:41:30.141042Z","iopub.execute_input":"2023-04-19T11:41:30.141322Z","iopub.status.idle":"2023-04-19T11:41:30.152781Z","shell.execute_reply.started":"2023-04-19T11:41:30.141280Z","shell.execute_reply":"2023-04-19T11:41:30.151735Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"az_pred, zn_pred = [], []\nfor bfile in BATCH_LIST:\n    batch_id = int(bfile.split('.')[0].split('_')[-1])\n    dataset = IceCubeDataset(bfile, seed=42, \n                             min_sequence_lenth=0, \n                             max_sequence_length=MAX_SEQUENCE_LENGTH,\n                             padded_sequence_length=MAX_SEQUENCE_LENGTH,\n    )\n    long_seq_dataset = IceCubeDataset(bfile, seed=42,\n                             min_sequence_lenth=MAX_SEQUENCE_LENGTH,\n                             max_sequence_length=LONG_SEQ_MAX_LENGTH,\n                             padded_sequence_length=LONG_SEQ_MAX_LENGTH,\n    )\n    ul_ds_create_fn = lambda seed: IceCubeDataset(bfile, seed=seed, \n                             min_sequence_lenth=LONG_SEQ_MAX_LENGTH,\n                             max_sequence_length=np.inf,\n                             padded_sequence_length=LONG_SEQ_MAX_LENGTH,\n    )\n    ultralong_ds = ul_ds_create_fn(seed=42)\n    events_to_keep = list(set(dataset.events).union(long_seq_dataset.events).union(ultralong_ds.events))\n    event_lengths, batch_features = generate_batch_features(bfile, events_to_keep)\n    az_pred_mbs, zn_pred_mbs = [], []\n    model_pred_features = []\n    for model_config in MODEL_CONFIGS:\n        model = load_model(model_config)\n        stack_model = load_stack_model(model_config.stack_model_config)\n        long_seq_model = load_model(model_config.long_seq_model_config)\n\n        str_id = f'batch_{batch_id} model_{model_config.name}'\n\n        az_b, zn_b, az_s, zn_s, pdt = predict_on_batch(model, stack_model, dataset, BATCH_SIZE)\n        az_bls, zn_bls, az_sls, zn_sls, pdt_ls = predict_on_batch(long_seq_model, stack_model, \n                                                          long_seq_dataset, LONG_SEQ_BATCH_SIZE)\n        az_buls, zn_buls, az_suls, zn_suls, pdt_uls = predict_on_batch(long_seq_model, stack_model, \n                                                          ultralong_ds, LONG_SEQ_BATCH_SIZE,\n                                                          mutlisample=ULTTRA_LONG_RESAMPLE, \n                                                          dataset_fn=ul_ds_create_fn)\n        az_base = join_predictions(az_b, az_bls, az_buls, dataset, long_seq_dataset, ultralong_ds)\n        zn_base = join_predictions(zn_b, zn_bls, zn_buls, dataset, long_seq_dataset, ultralong_ds)\n        az_stack = join_predictions(az_s, az_sls, az_suls, dataset, long_seq_dataset, ultralong_ds)\n        zn_stack = join_predictions(zn_s, zn_sls, zn_suls, dataset, long_seq_dataset, ultralong_ds)\n        prediction_data = []\n        for p, p_ls, p_uls in zip(pdt, pdt_ls, pdt_uls):\n            prediction_data.append(\n                join_predictions(p, p_ls, p_uls, dataset, \n                                 long_seq_dataset, ultralong_ds)\n            )\n\n        ## combine classifier and vmf predictions\n        boosting_feats, pred_features = generate_boosting_features(event_lengths, batch_features, \n                                                       az_base, zn_base,\n                                                       az_stack, zn_stack,\n                                                       prediction_data\n                                                    )\n        model_pred_features.append(pred_features)\n        mixer_clf = joblib.load(f'{MODEL_DIR}/{model_config.mixer_model}')\n        mixer_preds = mixer_clf.predict(boosting_feats)\n        composite_zenith = zn_base.clone()\n        v = mixer_preds == 1\n        composite_zenith[v] = zn_stack[v].clone()\n        \n        selector_clf = joblib.load(f'{MODEL_DIR}/{model_config.selector_model}')\n        selector_preds = selector_clf.predict(boosting_feats)\n        v = selector_preds == 1\n        rep_azimuth = az_base.clone()\n        rep_zenith = composite_zenith.clone()\n        rep_azimuth[v] = ( az_base[v] + np.pi ) % (np.pi * 2) # reverse_azimuth\n        rep_zenith[v] = ( - composite_zenith[v] ) % np.pi # reverse_zenith\n\n        az_pred_mb, zn_pred_mb = rep_azimuth, rep_zenith\n        \n        az_pred_mbs.append(az_pred_mb)\n        zn_pred_mbs.append(zn_pred_mb)\n    \n\n    if len(MODEL_CONFIGS) == 2:\n        merger_model = joblib.load(f'{MODEL_DIR}/{MERGER_MODELS[0]}')\n        merger_feats = np.concatenate(\n            [event_lengths[:, None], model_pred_features[0], model_pred_features[1], batch_features],\n            axis=1\n        )\n        merger_preds = merger_model.predict(merger_feats)\n        az_pred_batch = az_pred_mbs[0].clone()\n        zn_pred_batch = zn_pred_mbs[0].clone()\n        v = merger_preds == 1\n        az_pred_batch[v] = az_pred_mbs[1][v].clone()\n        zn_pred_batch[v] = zn_pred_mbs[1][v].clone()\n    elif len(MODEL_CONFIGS) == 1:\n        az_pred_batch, zn_pred_batch = az_pred_mbs[0], zn_pred_mbs[0]\n    elif len(MODEL_CONFIGS) == 3:\n        # merge 18l and 15l\n        merger01_model = joblib.load(f'{MODEL_DIR}/{MERGER_MODELS[0]}')\n        merger01_feats = np.concatenate(\n            [event_lengths[:, None], model_pred_features[0], model_pred_features[1], batch_features],\n            axis=1\n        )\n        merger01_preds = merger01_model.predict(merger01_feats)\n        merge01_az = az_pred_mbs[0].clone()\n        merge01_zn = zn_pred_mbs[0].clone()\n        v = merger01_preds == 1\n        merge01_az[v] = az_pred_mbs[1][v].clone()\n        merge01_zn[v] = zn_pred_mbs[1][v].clone()\n\n        # merge 15l and r6\n        merger12_model = joblib.load(f'{MODEL_DIR}/{MERGER_MODELS[1]}')\n        merger12_feats = np.concatenate(\n            [event_lengths[:, None], model_pred_features[1], model_pred_features[2], batch_features],\n            axis=1\n        )\n        merger12_preds = merger12_model.predict(merger12_feats)\n        merge12_az = az_pred_mbs[1].clone()\n        merge12_zn = zn_pred_mbs[1].clone()\n        v = merger12_preds == 1\n        merge12_az[v] = az_pred_mbs[2][v].clone()\n        merge12_zn[v] = zn_pred_mbs[2][v].clone()\n\n        # average and vote on results of merge\n        vote_az = avg_az(merge01_az, merge12_az)\n        az_hdiff = torch.abs(az_diff(merge01_az, merge12_az)) > 0.4\n        vote_az[az_hdiff] = merge01_az[az_hdiff].clone()\n\n        vote_zn = ( merge01_zn + merge12_zn ) / 2\n        zn_hdiff = torch.abs(merge01_zn - merge12_zn) > 0.4\n        vote_zn[zn_hdiff] = merge01_zn[zn_hdiff].clone()\n\n        az_pred_batch, zn_pred_batch  = vote_az, vote_zn\n    else:\n        raise NotImplementedError\n\n    az_pred.append(az_pred_batch)\n    zn_pred.append(zn_pred_batch)\n\n    check_score(batch_id, az_pred_batch, zn_pred_batch,\n                events_to_keep=events_to_keep,\n                text=f'batch_{batch_id}', convert_to_angle=False)","metadata":{"papermill":{"duration":2.470144,"end_time":"2023-04-08T09:58:41.648497","exception":false,"start_time":"2023-04-08T09:58:39.178353","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-19T11:41:30.155973Z","iopub.execute_input":"2023-04-19T11:41:30.156297Z","iopub.status.idle":"2023-04-19T11:41:36.421053Z","shell.execute_reply.started":"2023-04-19T11:41:30.156270Z","shell.execute_reply":"2023-04-19T11:41:36.419950Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"azimuth = torch.cat(az_pred, axis=0)\nzenith = torch.cat(zn_pred, axis=0)\n\nevent_ids = metadata.select('event_id').to_numpy().squeeze()\nsubmission_df = prepare_submission(event_ids, azimuth, zenith, validate_ids=(not CHECK_PREDICTION))\nsubmission_df.to_csv('submission.csv')\nprint('Saved submission')","metadata":{"papermill":{"duration":0.105093,"end_time":"2023-04-08T09:58:41.758599","exception":false,"start_time":"2023-04-08T09:58:41.653506","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-19T11:41:36.423604Z","iopub.execute_input":"2023-04-19T11:41:36.423988Z","iopub.status.idle":"2023-04-19T11:41:36.519146Z","shell.execute_reply.started":"2023-04-19T11:41:36.423949Z","shell.execute_reply":"2023-04-19T11:41:36.517918Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if CHECK_PREDICTION:\n    submission = pd.read_csv('submission.csv')\n    az_pred, zen_pred = submission.azimuth.values, submission.zenith.values\n    az_gt = metadata['azimuth'].to_numpy()\n    zen_gt = metadata['zenith'].to_numpy()\n    angular_dist = angular_dist_score(az_gt, zen_gt, az_pred, zen_pred)\n    print(\"Angular Distance Score\", angular_dist)","metadata":{"papermill":{"duration":0.012776,"end_time":"2023-04-08T09:58:41.793330","exception":false,"start_time":"2023-04-08T09:58:41.780554","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2023-04-19T11:41:36.520705Z","iopub.execute_input":"2023-04-19T11:41:36.521077Z","iopub.status.idle":"2023-04-19T11:41:36.528150Z","shell.execute_reply.started":"2023-04-19T11:41:36.521037Z","shell.execute_reply":"2023-04-19T11:41:36.527058Z"},"trusted":true},"execution_count":null,"outputs":[]}]}