{"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":"markdown","source":"# GraphNeT Baseline Submission\n<img style=\"float: right;\" src=\"https://raw.githubusercontent.com/graphnet-team/graphnet/main/assets/identity/graphnet-logo-and-wordmark.png\" width=\"600\" height=\"600\" />\n\nThis notebook submits predictions from the public pre-trained dynedge to the leaderboard. ","metadata":{}},{"cell_type":"code","source":"# Move software to working disk\n!rm  -r software\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/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":{"execution":{"iopub.status.busy":"2023-04-17T00:39:21.095176Z","iopub.execute_input":"2023-04-17T00:39:21.095648Z","iopub.status.idle":"2023-04-17T00:42:21.691426Z","shell.execute_reply.started":"2023-04-17T00:39:21.095616Z","shell.execute_reply":"2023-04-17T00:42:21.690185Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import pyarrow.parquet as pq\nimport sqlite3\nimport pandas as pd\nimport sqlalchemy\nfrom tqdm import tqdm\nimport os\nfrom typing import Any, Dict, List, Optional\nimport numpy as np\n# Import Modules\nimport gc\nimport multiprocessing\nimport time\nimport tensorflow as tf\n\nfrom graphnet.data.sqlite.sqlite_utilities import create_table\n\ndef load_input(meta_batch: pd.DataFrame, input_data_folder: str) -> pd.DataFrame:\n        \"\"\"\n        Will load the corresponding detector readings associated with the meta data batch.\n        \"\"\"\n        batch_id = pd.unique(meta_batch['batch_id'])\n\n        assert len(batch_id) == 1, \"contains multiple batch_ids. Did you set the batch_size correctly?\"\n        \n        detector_readings = pd.read_parquet(path = f'{input_data_folder}/batch_{batch_id[0]}.parquet')\n        sensor_positions = geometry_table.loc[detector_readings['sensor_id'], ['x', 'y', 'z']]\n        sensor_positions.index = detector_readings.index\n\n        for column in sensor_positions.columns:\n            if column not in detector_readings.columns:\n                detector_readings[column] = sensor_positions[column]\n\n        detector_readings['auxiliary'] = detector_readings['auxiliary'].replace({True: 1, False: 0})\n        return detector_readings.reset_index()\n\ndef add_to_table(database_path: str,\n                      df: pd.DataFrame,\n                      table_name:  str,\n                      is_primary_key: bool,\n                      ) -> None:\n    \"\"\"Writes meta data to sqlite table. \n\n    Args:\n        database_path (str): the path to the database file.\n        df (pd.DataFrame): the dataframe that is being written to table.\n        table_name (str, optional): The name of the meta table. Defaults to 'meta_table'.\n        is_primary_key(bool): Must be True if each row of df corresponds to a unique event_id. Defaults to False.\n    \"\"\"\n    try:\n        print(database_path)\n        create_table(   columns=  df.columns,\n                        database_path = database_path, \n                        table_name = table_name,\n                        integer_primary_key= is_primary_key,\n                        index_column = 'event_id')\n    except sqlite3.OperationalError as e:\n        if 'already exists' in str(e):\n            pass\n        else:\n            raise e\n    engine = sqlalchemy.create_engine(\"sqlite:///\" + database_path)\n    df.to_sql(table_name, con=engine, index=False, if_exists=\"append\", chunksize = 200000)\n    engine.dispose()\n    return\n\ndef convert_to_sqlite(meta_data_path: str,\n                      database_path: str,\n                      input_data_folder: str,\n                      batch_size: int = 200000) -> None:\n    \"\"\"Converts a selection of the Competition's parquet files to a single sqlite database.\n\n    Args:\n        meta_data_path (str): Path to the meta data file.\n        batch_size (int): the number of rows extracted from meta data file at a time. Keep low for memory efficiency.\n        database_path (str): path to database. E.g. '/my_folder/data/my_new_database.db'\n        input_data_folder (str): folder containing the parquet input files.\n        accepted_batch_ids (List[int]): The batch_ids you want converted. Defaults to None (all batches will be converted)\n    \"\"\"\n    meta_data_iter = pq.ParquetFile(meta_data_path).iter_batches(batch_size = batch_size)\n    \n    if not database_path.endswith('.db'):\n        database_path = database_path+'.db'\n        \n    converted_batches = [] \n    progress_bar = tqdm(total = None)\n    for meta_data_batch in meta_data_iter:\n        unique_batch_ids = pd.unique(meta_data_batch['event_id']).tolist()\n        meta_data_batch  = meta_data_batch.to_pandas()\n        add_to_table(database_path = database_path,\n                    df = meta_data_batch,\n                    table_name='meta_table',\n                    is_primary_key= True)\n        pulses = load_input(meta_batch=meta_data_batch, input_data_folder= input_data_folder)\n        del meta_data_batch # memory\n        add_to_table(database_path = database_path,\n                    df = pulses,\n                    table_name='pulse_table',\n                    is_primary_key= False)\n        del pulses # memory\n        progress_bar.update(1)\n    progress_bar.close()\n    del meta_data_iter # memory\n    print(f'Conversion Complete!. Database available at\\n {database_path}')","metadata":{"execution":{"iopub.status.busy":"2023-04-17T00:42:21.694102Z","iopub.execute_input":"2023-04-17T00:42:21.694447Z","iopub.status.idle":"2023-04-17T00:42:33.113662Z","shell.execute_reply.started":"2023-04-17T00:42:21.694415Z","shell.execute_reply":"2023-04-17T00:42:33.112574Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!rm '/kaggle/working/test_database.db'\ninput_data_folder = '/kaggle/input/icecube-neutrinos-in-deep-ice/test'\ngeometry_table = pd.read_csv('/kaggle/input/icecube-neutrinos-in-deep-ice/sensor_geometry.csv')\nmeta_data_path = '/kaggle/input/icecube-neutrinos-in-deep-ice/test_meta.parquet'\n\ndatabase_path = '/kaggle/working/test_database'\nconvert_to_sqlite(meta_data_path,\n                  database_path=database_path,\n                  input_data_folder=input_data_folder)","metadata":{"execution":{"iopub.status.busy":"2023-04-17T00:42:33.115563Z","iopub.execute_input":"2023-04-17T00:42:33.115939Z","iopub.status.idle":"2023-04-17T00:42:34.319255Z","shell.execute_reply.started":"2023-04-17T00:42:33.115902Z","shell.execute_reply":"2023-04-17T00:42:34.318071Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from pytorch_lightning.callbacks import EarlyStopping\nfrom torch.optim.adam import Adam\nfrom graphnet.data.constants import FEATURES, TRUTH\nfrom graphnet.models import StandardModel\nfrom graphnet.models.detector.icecube import IceCubeKaggle\nfrom graphnet.models.gnn import DynEdge\nfrom graphnet.models.graph_builders import KNNGraphBuilder\nfrom graphnet.models.task.reconstruction import DirectionReconstructionWithKappa, ZenithReconstructionWithKappa, AzimuthReconstructionWithKappa\nfrom graphnet.training.callbacks import ProgressBar, PiecewiseLinearLR\nfrom graphnet.training.loss_functions import VonMisesFisher3DLoss, VonMisesFisher2DLoss\nfrom graphnet.training.labels import Direction\nfrom graphnet.training.utils import make_dataloader\nfrom graphnet.utilities.logging import get_logger\nfrom pytorch_lightning import Trainer\nimport pandas as pd\n\nlogger = get_logger()\n\ndef build_model(config: Dict[str,Any], train_dataloader: Any) -> StandardModel:\n    \"\"\"Builds GNN from config\"\"\"\n    # Building model\n    detector = IceCubeKaggle(\n        graph_builder=KNNGraphBuilder(nb_nearest_neighbours=8),\n    )\n    gnn = DynEdge(\n        nb_inputs=detector.nb_outputs,\n        global_pooling_schemes=[\"min\", \"max\", \"mean\"],\n    )\n\n    if config[\"target\"] == 'direction':\n        task = DirectionReconstructionWithKappa(\n            hidden_size=gnn.nb_outputs,\n            target_labels=config[\"target\"],\n            loss_function=VonMisesFisher3DLoss(),\n        )\n        prediction_columns = [config[\"target\"] + \"_x\", \n                              config[\"target\"] + \"_y\", \n                              config[\"target\"] + \"_z\", \n                              config[\"target\"] + \"_kappa\" ]\n        additional_attributes = ['zenith', 'azimuth', 'event_id']\n\n    model = StandardModel(\n        detector=detector,\n        gnn=gnn,\n        tasks=[task],\n        optimizer_class=Adam,\n        optimizer_kwargs={\"lr\": 1e-03, \"eps\": 1e-03},\n        scheduler_class=PiecewiseLinearLR,\n        scheduler_kwargs={\n            \"milestones\": [\n                0,\n                len(train_dataloader) / 2,\n                len(train_dataloader) * config[\"fit\"][\"max_epochs\"],\n            ],\n            \"factors\": [1e-02, 1, 1e-02],\n        },\n        scheduler_config={\n            \"interval\": \"step\",\n        },\n    )\n    model.prediction_columns = prediction_columns\n    model.additional_attributes = additional_attributes\n    \n    return model\n\ndef load_pretrained_model(config: Dict[str,Any], state_dict_path: str = '/kaggle/input/dynedge-pretrained/dynedge_pretrained_batch_1_to_50/state_dict.pth') -> StandardModel:\n    train_dataloader, _ = make_dataloaders(config = config)\n    model = build_model(config = config, \n                        train_dataloader = train_dataloader)\n    #model._inference_trainer = Trainer(config['fit'])\n    model.load_state_dict(state_dict_path)\n    model.prediction_columns = [config[\"target\"] + \"_x\", \n                              config[\"target\"] + \"_y\", \n                              config[\"target\"] + \"_z\", \n                              config[\"target\"] + \"_kappa\" ]\n    model.additional_attributes = ['event_id'] #'zenith', 'azimuth',  not available in test data\n    return model\n\ndef make_dataloaders(config: Dict[str, Any]) -> List[Any]:\n    \"\"\"Constructs training and validation dataloaders for training with early stopping.\"\"\"\n    \n    train_dataloader = make_dataloader(db = config['path'],\n                                            selection = pd.read_csv(config['train_selection'])[config['index_column']].ravel().tolist() if config['train_selection'] else None,\n                                            pulsemaps = config['pulsemap'],\n                                            features = features,\n                                            truth = truth,\n                                            batch_size = config['batch_size'],\n                                            num_workers = config['num_workers'],\n                                            shuffle = True,\n                                            labels = {'direction': Direction()},\n                                            index_column = config['index_column'],\n                                            truth_table = config['truth_table'],\n                                            )\n    \n    validate_dataloader = make_dataloader(db = config['path'],\n                                            selection = pd.read_csv(config['validate_selection'])[config['index_column']].ravel().tolist() if config['validate_selection'] else None,\n                                            pulsemaps = config['pulsemap'],\n                                            features = features,\n                                            truth = truth,\n                                            batch_size = config['batch_size'],\n                                            num_workers = config['num_workers'],\n                                            shuffle = False,\n                                            labels = {'direction': Direction()},\n                                            index_column = config['index_column'],\n                                            truth_table = config['truth_table'],\n                                          \n                                            )\n    return train_dataloader, validate_dataloader\n\ndef inference(model, config: Dict[str, Any]) -> pd.DataFrame:\n    \"\"\"Applies model to the database specified in config['inference_database_path'] and saves results to disk.\"\"\"\n    # Make Dataloader\n    test_dataloader = make_dataloader(db = config['inference_database_path'],\n                                            selection = None, # Entire database\n                                            pulsemaps = config['pulsemap'],\n                                            features = features,\n                                            truth = truth,\n                                            batch_size = config['batch_size'],\n                                            num_workers = config['num_workers'],\n                                            shuffle = False,\n                                            labels = None, # Cannot make labels in test data\n                                            index_column = config['index_column'],\n                                            truth_table = config['truth_table'],\n                                            )\n    \n    # Get predictions\n    results = model.predict_as_dataframe(\n        gpus = [0],\n        dataloader = test_dataloader,\n        prediction_columns=model.prediction_columns,\n        additional_attributes=['event_id']\n    )\n    return results\n\n    ","metadata":{"execution":{"iopub.status.busy":"2023-04-17T00:42:34.322814Z","iopub.execute_input":"2023-04-17T00:42:34.32355Z","iopub.status.idle":"2023-04-17T00:42:35.547101Z","shell.execute_reply.started":"2023-04-17T00:42:34.323507Z","shell.execute_reply":"2023-04-17T00:42:35.546052Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def prepare_dataframe(df, angle_post_fix = '_reco', vec_post_fix = '') -> pd.DataFrame:\n    r = np.sqrt(df['direction_x'+ vec_post_fix]**2 + df['direction_y'+ vec_post_fix]**2 + df['direction_z' + vec_post_fix]**2)\n    df['zenith' + angle_post_fix] = np.arccos(df['direction_z'+ vec_post_fix]/r)\n    df['azimuth'+ angle_post_fix] = np.arctan2(df['direction_y'+ vec_post_fix],df['direction_x' + vec_post_fix]) #np.sign(results['true_y'])*np.arccos((results['true_x'])/(np.sqrt(results['true_x']**2 + results['true_y']**2)))\n    df['azimuth'+ angle_post_fix][df['azimuth'  + angle_post_fix]<0] = df['azimuth'  + angle_post_fix][df['azimuth'  +  angle_post_fix]<0] + 2*np.pi \n\n    drop_these_columns = []\n    for column in results.columns:\n        if column not in ['event_id', 'zenith', 'azimuth']:\n            drop_these_columns.append(column)\n    return df.drop(columns = drop_these_columns).iloc[:,[0,2,1]].set_index('event_id')\n    \n","metadata":{"execution":{"iopub.status.busy":"2023-04-17T00:42:35.550979Z","iopub.execute_input":"2023-04-17T00:42:35.551993Z","iopub.status.idle":"2023-04-17T00:42:35.567201Z","shell.execute_reply.started":"2023-04-17T00:42:35.551955Z","shell.execute_reply":"2023-04-17T00:42:35.56537Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Constants\nfeatures = FEATURES.KAGGLE\ntruth = TRUTH.KAGGLE\n\n# Configuration\nconfig = {\n        \"path\": '/kaggle/working/test_database.db',\n        \"inference_database_path\": '/kaggle/working/test_database.db',\n        \"pulsemap\": 'pulse_table',\n        \"truth_table\": 'meta_table',\n        \"features\": features,\n        \"truth\": truth,\n        \"index_column\": 'event_id',\n        \"run_name_tag\": 'graphnet_baseline_submission',\n        \"batch_size\": 200,\n        \"num_workers\": 2,\n        \"target\": 'direction',\n        \"early_stopping_patience\": 5,\n        \"fit\": {\n                \"max_epochs\": 50,\n                \"gpus\": [0],\n                \"distribution_strategy\": None,\n                },\n        'train_selection': None,\n        'validate_selection':  None,\n        'test_selection': None,\n        'base_dir': 'training'\n}\nmodel   = load_pretrained_model(config = config)\nsubmission_df0 = inference(model, config)\n","metadata":{"execution":{"iopub.status.busy":"2023-04-17T00:42:35.575196Z","iopub.execute_input":"2023-04-17T00:42:35.578398Z","iopub.status.idle":"2023-04-17T00:42:40.722035Z","shell.execute_reply.started":"2023-04-17T00:42:35.578355Z","shell.execute_reply":"2023-04-17T00:42:40.720391Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\n\n# Directories and constants\nhome_dir = \"/kaggle/input/icecube-neutrinos-in-deep-ice/\"\ntest_format = home_dir + 'test/batch_{batch_id:d}.parquet'\nmodel_home = \"/kaggle/input/lstmicecubesdata/\"\n\n# Model(s)\nmodel_names = [\"4347_MAE_1-02076_bin24_pp96_n6_batch2048_epoch29.h5\",\n               \"4347_MAE_1-02039_bin24_pp96_n6_batch2048_epoch25.h5\", \n               \"4346_MAE_1-02020_bin24_pp96_n6_batch2048_epoch27.h5\"]\nmodel_weights = np.array([0.30, \n                          0.30,\n                          0.40])\n\n","metadata":{"execution":{"iopub.status.busy":"2023-04-17T00:42:40.726732Z","iopub.execute_input":"2023-04-17T00:42:40.727079Z","iopub.status.idle":"2023-04-17T00:42:40.734435Z","shell.execute_reply.started":"2023-04-17T00:42:40.727036Z","shell.execute_reply":"2023-04-17T00:42:40.733307Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Load Models\nmodels = []\nfor model_name in model_names:\n    print(f'\\n========== Model File: {model_name}')\n    \n    # Load Model\n    model_path = model_home + model_name\n    model = tf.keras.models.load_model(model_path, compile = False)\n    models.append(model)      \n    \n    # Model summary\n    model.summary()\n    \n# Get Model Parameters\npulse_count = model.inputs[0].shape[1]\nfeature_count = model.inputs[0].shape[2]\noutput_bins = model.layers[-1].weights[0].shape[-1]\nbin_num = int(np.sqrt(output_bins))\n\n# Model Parameter Summary\nprint(\"\\n==== Model Parameters\")\nprint(f\"Bin Numbers: {bin_num}\")\nprint(f\"Maximum Pulse Count: {pulse_count}\")\nprint(f\"Features Count: {feature_count}\")","metadata":{"execution":{"iopub.status.busy":"2023-04-17T00:42:40.73582Z","iopub.execute_input":"2023-04-17T00:42:40.736591Z","iopub.status.idle":"2023-04-17T00:43:08.88239Z","shell.execute_reply.started":"2023-04-17T00:42:40.736552Z","shell.execute_reply":"2023-04-17T00:43:08.881233Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Load sensor_geometry\nsensor_geometry_df = pd.read_csv(home_dir + \"sensor_geometry.csv\")\n\n# Get Sensor Information\nsensor_x = sensor_geometry_df.x\nsensor_y = sensor_geometry_df.y\nsensor_z = sensor_geometry_df.z\n\n# Detector constants\nc_const = 0.299792458  # speed of light [m/ns]\n\n# Sensor Min / Max Coordinates\nx_min = sensor_x.min()\nx_max = sensor_x.max()\ny_min = sensor_y.min()\ny_max = sensor_y.max()\nz_min = sensor_z.min()\nz_max = sensor_z.max()\n\ndetector_length = np.sqrt((x_max - x_min)**2 + (y_max - y_min)**2 + (z_max - z_min)**2)\nt_valid_length = detector_length / c_const\n\nprint(f\"time valid length: {t_valid_length} ns\")\n# Create Azimuth Edges\nazimuth_edges = np.linspace(0, 2 * np.pi, bin_num + 1)\nprint(azimuth_edges)\n\n# Create Zenith Edges\nzenith_edges = []\nzenith_edges.append(0)\nfor bin_idx in range(1, bin_num):\n    zenith_edges.append(np.arccos(np.cos(zenith_edges[-1]) - 2 / (bin_num)))\nzenith_edges.append(np.pi)\nzenith_edges = np.array(zenith_edges)\nprint(zenith_edges)\nangle_bin_zenith0 = np.tile(zenith_edges[:-1], bin_num)\nangle_bin_zenith1 = np.tile(zenith_edges[1:], bin_num)\nangle_bin_azimuth0 = np.repeat(azimuth_edges[:-1], bin_num)\nangle_bin_azimuth1 = np.repeat(azimuth_edges[1:], bin_num)\n\nangle_bin_area = (angle_bin_azimuth1 - angle_bin_azimuth0) * (np.cos(angle_bin_zenith0) - np.cos(angle_bin_zenith1))\nangle_bin_vector_sum_x = (np.sin(angle_bin_azimuth1) - np.sin(angle_bin_azimuth0)) * ((angle_bin_zenith1 - angle_bin_zenith0) / 2 - (np.sin(2 * angle_bin_zenith1) - np.sin(2 * angle_bin_zenith0)) / 4)\nangle_bin_vector_sum_y = (np.cos(angle_bin_azimuth0) - np.cos(angle_bin_azimuth1)) * ((angle_bin_zenith1 - angle_bin_zenith0) / 2 - (np.sin(2 * angle_bin_zenith1) - np.sin(2 * angle_bin_zenith0)) / 4)\nangle_bin_vector_sum_z = (angle_bin_azimuth1 - angle_bin_azimuth0) * ((np.cos(2 * angle_bin_zenith0) - np.cos(2 * angle_bin_zenith1)) / 4)\n\nangle_bin_vector_mean_x = angle_bin_vector_sum_x / angle_bin_area\nangle_bin_vector_mean_y = angle_bin_vector_sum_y / angle_bin_area\nangle_bin_vector_mean_z = angle_bin_vector_sum_z / angle_bin_area\n\nangle_bin_vector = np.zeros((1, bin_num * bin_num, 3))\nangle_bin_vector[:, :, 0] = angle_bin_vector_mean_x\nangle_bin_vector[:, :, 1] = angle_bin_vector_mean_y\nangle_bin_vector[:, :, 2] = angle_bin_vector_mean_z\n\nangle_bin_vector_unit = angle_bin_vector[0].copy()\nangle_bin_vector_unit /= np.sqrt((angle_bin_vector_unit**2).sum(axis=1).reshape((-1, 1)))\ndef pred_to_angle(pred, epsilon = 1e-8):\n    # Convert prediction\n    pred_vector = (pred.reshape((-1, bin_num**2, 1)) * angle_bin_vector).sum(axis = 1)\n    \n    # Normalize\n    pred_vector_norm = np.sqrt((pred_vector**2).sum(axis = 1))\n    mask = pred_vector_norm < epsilon\n    pred_vector_norm[mask] = 1\n    \n    # Assign <1, 0, 0> to very small vectors (badly predicted)\n    pred_vector /= pred_vector_norm.reshape((-1, 1))\n    pred_vector[mask] = np.array([1., 0., 0.])\n    \n    # Convert to angle\n    azimuth = np.arctan2(pred_vector[:, 1], pred_vector[:, 0])\n    azimuth[azimuth < 0] += 2 * np.pi\n    zenith = np.arccos(pred_vector[:, 2])\n    \n    # Mask bad norm predictions as 0, 0\n    azimuth[mask] = 0.\n    zenith[mask] = 0.\n    \n    return azimuth, zenith\n\ndef weighted_vector_ensemble(angles, weight):\n    # Convert angle to vector\n    vec_models = list()\n    for angle in angles:\n        az, zen = angle\n        sa = np.sin(az)\n        ca = np.cos(az)\n        sz = np.sin(zen)\n        cz = np.cos(zen)\n        vec = np.stack([sz * ca, sz * sa, cz], axis=1)\n        vec_models.append(vec)\n    vec_models = np.array(vec_models)\n\n    # Weighted-mean\n    vec_mean = (weight.reshape((-1, 1, 1)) * vec_models).sum(axis=0) / weight.sum()\n    vec_mean /= np.sqrt((vec_mean**2).sum(axis=1)).reshape((-1, 1))\n\n    # Convert vector to angle\n    zenith = np.arccos(vec_mean[:, 2])\n    azimuth = np.arctan2(vec_mean[:, 1], vec_mean[:, 0])\n    azimuth[azimuth < 0] += 2 * np.pi\n    \n    return azimuth, zenith\n# Placeholder\nopen_batch_dict = dict()\n\n# Read single event from batch_meta_df\ndef read_event(event_idx, batch_meta_df, pulse_count):\n    # Read metadata\n    batch_id, first_pulse_index, last_pulse_index = batch_meta_df.iloc[event_idx][[\"batch_id\", \"first_pulse_index\", \"last_pulse_index\"]].astype(\"int\")\n\n    # close past batch df\n    if batch_id - 1 in open_batch_dict.keys():\n        del open_batch_dict[batch_id - 1]\n\n    # open current batch df\n    if batch_id not in open_batch_dict.keys():\n        open_batch_dict.update({batch_id: pd.read_parquet(test_format.format(batch_id=batch_id))})\n    \n    batch_df = open_batch_dict[batch_id]\n    \n    # Read event\n    event_feature = batch_df[first_pulse_index:last_pulse_index + 1]\n    sensor_id = event_feature.sensor_id\n    \n    # Merge features into single structured array\n    dtype = [(\"time\", \"float16\"),\n             (\"charge\", \"float16\"),\n             (\"auxiliary\", \"float16\"),\n             (\"x\", \"float16\"),\n             (\"y\", \"float16\"),\n             (\"z\", \"float16\"),\n             (\"rank\", \"short\")]    \n    \n    # Create event_x\n    event_x = np.zeros(last_pulse_index - first_pulse_index + 1, dtype)\n    event_x[\"time\"] = event_feature.time.values - event_feature.time.min()\n    event_x[\"charge\"] = event_feature.charge.values\n    event_x[\"auxiliary\"] = event_feature.auxiliary.values\n    event_x[\"x\"] = sensor_geometry_df.x[sensor_id].values\n    event_x[\"y\"] = sensor_geometry_df.y[sensor_id].values\n    event_x[\"z\"] = sensor_geometry_df.z[sensor_id].values\n\n    # For long event, pick-up\n    if len(event_x) > pulse_count:\n        # Find valid time window\n        t_peak = event_x[\"time\"][event_x[\"charge\"].argmax()]\n        t_valid_min = t_peak - t_valid_length\n        t_valid_max = t_peak + t_valid_length\n        t_valid = (event_x[\"time\"] > t_valid_min) * (event_x[\"time\"] < t_valid_max)\n\n        # Rank\n        event_x[\"rank\"] = 2 * (1 - event_x[\"auxiliary\"]) + (t_valid)\n\n        # Sort by Rank and Charge (important goes to backward)\n        event_x = np.sort(event_x, order = [\"rank\", \"charge\"])\n\n        # pick-up from backward\n        event_x = event_x[-pulse_count:]\n\n        # Sort events by time \n        event_x = np.sort(event_x, order = \"time\")\n\n    return event_idx, len(event_x), event_x\n","metadata":{"execution":{"iopub.status.busy":"2023-04-17T00:43:08.884099Z","iopub.execute_input":"2023-04-17T00:43:08.884644Z","iopub.status.idle":"2023-04-17T00:43:08.933193Z","shell.execute_reply.started":"2023-04-17T00:43:08.884603Z","shell.execute_reply":"2023-04-17T00:43:08.931826Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Read Test Meta data\ntest_meta_df = pq.read_table(home_dir + 'test_meta.parquet').to_pandas()\nbatch_counts = test_meta_df.batch_id.value_counts().sort_index()\n\nbatch_max_index = batch_counts.cumsum()\nbatch_max_index[test_meta_df.batch_id.min() - 1] = 0\nbatch_max_index = batch_max_index.sort_index()\n\n# Support Function\ndef test_meta_df_spliter(batch_id):\n    return test_meta_df.loc[batch_max_index[batch_id - 1]:batch_max_index[batch_id] - 1]","metadata":{"execution":{"iopub.status.busy":"2023-04-17T00:43:08.937462Z","iopub.execute_input":"2023-04-17T00:43:08.938681Z","iopub.status.idle":"2023-04-17T00:43:08.955602Z","shell.execute_reply.started":"2023-04-17T00:43:08.938642Z","shell.execute_reply":"2023-04-17T00:43:08.954862Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\n\n\n\ndef MeanAngErr(y_true, y_pred):\n    fcos=tf.math.scalar_mul(-0.99999,tf.keras.losses.cosine_similarity(y_true, y_pred))\n    return tf.reduce_mean(tf.math.acos(fcos), axis=-1)  # Note the `axis=-1`\ndef create_model():\n    \n    inputs = tf.keras.Input(shape=(1733,))\n    \n    intermediate = tf.keras.layers.Reshape((1733,1), input_shape=(1733,))(inputs)\n\n    inputs1 = tf.keras.layers.Cropping1D(cropping=(2,0))(intermediate)\n    inputs1 = tf.keras.layers.Reshape((1731,), input_shape=(1731,1))(inputs1)\n    \n    tensor_1 = tf.keras.layers.Dense(20, activation='tanh')(inputs)\n    tensor_1 = tf.keras.layers.Dense(1731, activation='tanh')(tensor_1)\n    #tensor_2 = tf.keras.layers.BatchNormalization()(tensor_2)\n    \n    #tensor_2 = tf.keras.layers.Dropout(0.5)(tensor_2)\n    \n    product = tf.keras.layers.Multiply()([tensor_1, inputs1])\n    \n    outputs = tf.keras.layers.Dense(3, activation = 'linear')(product)\n        \n        # Finalize Model\n    model = tf.keras.models.Model(inputs = inputs, outputs = outputs)\n\n        # Compile model\n    model.compile(loss = MeanAngErr,\n                      optimizer= tf.keras.optimizers.Adam())\n        \n        # Show Model Summary\n    model.summary()\n\n    return model\n\n\n\n\n\nmodelens=create_model()\nmodelens.load_weights(\"/kaggle/input/modelens3/model3(1).h5\")\n\ndef proc_X_models(X,submission_df0,testmeta):\n    ans=(np.concatenate((models[0].predict(X, verbose=0,batch_size=1024),models[1].predict(X, verbose=0,batch_size=1024),models[2].predict(X, verbose=0,batch_size=1024)),axis=1))\n    #e1=np.expand_dims(np.sum(np.log(ans[:,0:575]+1e-6)*ans[:,0:575],axis=1),axis=-1)\n    #e2=np.expand_dims(np.sum(np.log(ans[:,576:1151]+1e-6)*ans[:,576:1151],axis=1),axis=-1)\n    #e3=np.expand_dims(np.sum(np.log(ans[:,1152:1727]+1e-6)*ans[:,1152:1727],axis=1),axis=-1)\n    \n    submission_df0= submission_df0[['direction_x','direction_y','direction_z','direction_kappa']].to_numpy()\n    \n    pulses=testmeta[['last_pulse_index']].to_numpy()-testmeta[['first_pulse_index']].to_numpy()\n    \n    ans=np.concatenate((ans,submission_df0,pulses),axis=1)\n    \n    ans=modelens.predict(ans,batch_size=1024)\n    #print(ans)\n    norm=(np.maximum(np.sqrt((ans*ans)@np.array([1,1,1])),0.0000000000001))\n    \n    ans[:,0]=(ans[:,0]/norm)*0.9999999\n    ans[:,1]=(ans[:,1]/norm)*0.9999999\n    ans[:,2]=(ans[:,2]/norm)*0.9999999\n    #print(ans)\n    az=np.arctan2(ans[:,1],ans[:,0])*0.9999999\n    zen=np.arccos(ans[:,2])\n    az[az<0]+=2*np.pi-0.0000000000001\n    return az,zen","metadata":{"execution":{"iopub.status.busy":"2023-04-17T00:43:08.957354Z","iopub.execute_input":"2023-04-17T00:43:08.958071Z","iopub.status.idle":"2023-04-17T00:43:09.103172Z","shell.execute_reply.started":"2023-04-17T00:43:08.958032Z","shell.execute_reply":"2023-04-17T00:43:09.102243Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission_df0['event_id']=submission_df0['event_id'].astype(int)","metadata":{"execution":{"iopub.status.busy":"2023-04-17T00:43:09.10485Z","iopub.execute_input":"2023-04-17T00:43:09.105256Z","iopub.status.idle":"2023-04-17T00:43:09.11209Z","shell.execute_reply.started":"2023-04-17T00:43:09.105216Z","shell.execute_reply":"2023-04-17T00:43:09.110831Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Get Batch IDs\ntest_batch_ids = test_meta_df.batch_id.unique()\n\n# Submission Placeholders\ntest_event_id = []\ntest_azimuth = []\ntest_zenith = []\n\n# Batch Loop\nfor batch_id in test_batch_ids:\n    # Batch Meta DF\n    batch_meta_df = test_meta_df_spliter(batch_id)\n\n    # Set Pulses\n    test_x = np.zeros((len(batch_meta_df), pulse_count, feature_count), dtype = \"float16\")    \n    test_x[:, :, 2] = -1    \n\n    # Read Event Data\n    def read_event_local(event_idx):\n        return read_event(event_idx, batch_meta_df, pulse_count)\n    \n    # Multiprocess Events\n    iterator = range(len(batch_meta_df))\n    with multiprocessing.Pool() as pool:\n        for event_idx, pulsecount, event_x in pool.map(read_event_local, iterator):\n            # Features\n            test_x[event_idx, :pulsecount, 0] = event_x[\"time\"]\n            test_x[event_idx, :pulsecount, 1] = event_x[\"charge\"]\n            test_x[event_idx, :pulsecount, 2] = event_x[\"auxiliary\"]\n            test_x[event_idx, :pulsecount, 3] = event_x[\"x\"]\n            test_x[event_idx, :pulsecount, 4] = event_x[\"y\"]\n            test_x[event_idx, :pulsecount, 5] = event_x[\"z\"]\n    \n    \n    \n    # Normalize\n    test_x[:, :, 0] /= 1000  # time\n    test_x[:, :, 1] /= 300  # charge\n    test_x[:, :, 3:] /= 600  # space\n        \n    # Predict\n    pred_angles = []\n   \n    \n    # Get Event IDs\n    event_ids = test_meta_df[test_meta_df.batch_id == batch_id]\n    \n    graphpred=pd.merge(event_ids,submission_df0,on='event_id', how='left')\n\n    pred_azimuth, pred_zenith = proc_X_models(test_x,graphpred,event_ids)\n    gc.collect()\n    \n    del batch_meta_df\n    # Get Predicted Azimuth and Zenith\n    \n\n    event_ids=event_ids.event_id.values\n    \n    # Finalize \n    for event_id, azimuth, zenith in zip(event_ids, pred_azimuth, pred_zenith):\n        if np.isfinite(azimuth) and np.isfinite(zenith):\n            test_event_id.append(int(event_id))\n            test_azimuth.append(azimuth)\n            test_zenith.append(zenith)\n        else:\n            test_event_id.append(int(event_id))\n            test_azimuth.append(0.)\n            test_zenith.append(0.)\n# Create and Save Submission.csv\nsubmission_df1 = pd.DataFrame({\"event_id\": test_event_id,\n                              \"azimuth\": test_azimuth,\n                              \"zenith\": test_zenith})\nsubmission_df1 = submission_df1.sort_values(by = ['event_id'])\nsubmission_df1.to_csv(\"submission.csv\", index = False)","metadata":{"execution":{"iopub.status.busy":"2023-04-17T00:43:09.114243Z","iopub.execute_input":"2023-04-17T00:43:09.115096Z","iopub.status.idle":"2023-04-17T00:43:33.577052Z","shell.execute_reply.started":"2023-04-17T00:43:09.115056Z","shell.execute_reply":"2023-04-17T00:43:33.575088Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission_df1","metadata":{"execution":{"iopub.status.busy":"2023-04-17T00:43:33.582517Z","iopub.execute_input":"2023-04-17T00:43:33.583795Z","iopub.status.idle":"2023-04-17T00:43:33.606183Z","shell.execute_reply.started":"2023-04-17T00:43:33.583738Z","shell.execute_reply":"2023-04-17T00:43:33.604997Z"},"trusted":true},"execution_count":null,"outputs":[]}]}