{"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-02-03T21:40:52.413871Z","iopub.execute_input":"2023-02-03T21:40:52.414292Z","iopub.status.idle":"2023-02-03T21:45:44.289823Z","shell.execute_reply.started":"2023-02-03T21:40:52.414208Z","shell.execute_reply":"2023-02-03T21:45:44.288643Z"},"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\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-02-03T21:45:44.29244Z","iopub.execute_input":"2023-02-03T21:45:44.292765Z","iopub.status.idle":"2023-02-03T21:45:48.645758Z","shell.execute_reply.started":"2023-02-03T21:45:44.292734Z","shell.execute_reply":"2023-02-03T21:45:48.64477Z"},"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-02-03T21:45:48.647333Z","iopub.execute_input":"2023-02-03T21:45:48.648102Z","iopub.status.idle":"2023-02-03T21:45:49.804156Z","shell.execute_reply.started":"2023-02-03T21:45:48.648066Z","shell.execute_reply":"2023-02-03T21:45:49.803137Z"},"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\": 5e-04, \"eps\": 5e-04},\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 build_model1(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=12),\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\": 5e-04, \"eps\": 5e-04},\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\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 load_pretrained_model2(config: Dict[str,Any], state_dict_path: str = '/kaggle/input/icecubemodel/model(1).pth') -> StandardModel:\n    train_dataloader, _ = make_dataloaders(config = config)\n    model = build_model1(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\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(models, 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    results1 = models[0].predict_as_dataframe(\n        gpus = [0],\n        dataloader = test_dataloader,\n        prediction_columns=models[0].prediction_columns,\n        additional_attributes=['event_id']\n    )\n    \n    print(results1)\n    \n    results2 = models[1].predict_as_dataframe(\n        gpus = [0],\n        dataloader = test_dataloader,\n        prediction_columns=models[1].prediction_columns,\n        additional_attributes=['event_id']\n    )\n\n    results1[\"direction_x\"] = 0.9 * results1[\"direction_x\"] + 0.1 * results2[\"direction_x\"]\n    results1[\"direction_y\"] = 0.9 * results1[\"direction_y\"] + 0.1 * results2[\"direction_y\"]\n    results1[\"direction_z\"] = 0.9 * results1[\"direction_z\"] + 0.1 * results2[\"direction_z\"]\n    results1[\"direction_kappa\"] = 0.9 * results1[\"direction_kappa\"] + 0.1 * results2[\"direction_kappa\"]\n\n    return results1\n\n    ","metadata":{"execution":{"iopub.status.busy":"2023-02-03T21:45:49.807506Z","iopub.execute_input":"2023-02-03T21:45:49.80781Z","iopub.status.idle":"2023-02-03T21:45:50.639506Z","shell.execute_reply.started":"2023-02-03T21:45:49.807781Z","shell.execute_reply":"2023-02-03T21:45:50.63852Z"},"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-02-03T21:45:50.641102Z","iopub.execute_input":"2023-02-03T21:45:50.64148Z","iopub.status.idle":"2023-02-03T21:45:50.652383Z","shell.execute_reply.started":"2023-02-03T21:45:50.641416Z","shell.execute_reply":"2023-02-03T21:45:50.650447Z"},"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\": 250,\n        \"num_workers\": 2,\n        \"target\": 'direction',\n        \"early_stopping_patience\": 5,\n        \"fit\": {\n                \"max_epochs\": 70,\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}\nmodels = [load_pretrained_model(config = config), load_pretrained_model2(config = config)]\n\nresults = inference(models, config)","metadata":{"execution":{"iopub.status.busy":"2023-02-03T21:45:50.65428Z","iopub.execute_input":"2023-02-03T21:45:50.654802Z","iopub.status.idle":"2023-02-03T21:45:54.087245Z","shell.execute_reply.started":"2023-02-03T21:45:50.654763Z","shell.execute_reply":"2023-02-03T21:45:54.085557Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"results","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission_df = prepare_dataframe(results, angle_post_fix = '')","metadata":{"execution":{"iopub.status.busy":"2023-02-03T21:45:54.09177Z","iopub.execute_input":"2023-02-03T21:45:54.092091Z","iopub.status.idle":"2023-02-03T21:45:54.110743Z","shell.execute_reply.started":"2023-02-03T21:45:54.092059Z","shell.execute_reply":"2023-02-03T21:45:54.109404Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission_df.to_csv('submission.csv')","metadata":{"execution":{"iopub.status.busy":"2023-02-03T21:45:54.11238Z","iopub.execute_input":"2023-02-03T21:45:54.11308Z","iopub.status.idle":"2023-02-03T21:45:54.126265Z","shell.execute_reply.started":"2023-02-03T21:45:54.113046Z","shell.execute_reply":"2023-02-03T21:45:54.125506Z"},"trusted":true},"execution_count":null,"outputs":[]}]}