{"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":"This notebook allows to convert specified train data to SQLite databases. You can use `MAX_EVENTS` variable to exclude events with large number of pulses (set to None if not required).\n\nBased on https://www.kaggle.com/code/rasmusrse/graphnet-example","metadata":{}},{"cell_type":"code","source":"# Move software to working disk\n!rm  -r software\n!scp -r /kaggle/input/graphnet-and-dependencies-cpu/software .\n\n# Install dependencies\n!pip install /kaggle/working/software/dependencies/torch-1.11.0+cpu-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":{"scrolled":true,"execution":{"iopub.status.busy":"2023-04-10T18:06:24.114255Z","iopub.execute_input":"2023-04-10T18:06:24.114875Z","iopub.status.idle":"2023-04-10T18:08:40.023978Z","shell.execute_reply.started":"2023-04-10T18:06:24.114735Z","shell.execute_reply":"2023-04-10T18:08:40.022349Z"},"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, max_events: int = None) -> 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        if max_events is not None:\n            detector_readings['n_events'] = -1\n            detector_readings['n_events'] = detector_readings.groupby(detector_readings.index)['n_events'].count()\n            detector_readings = detector_readings[detector_readings['n_events'] < max_events+1]\n            detector_readings = detector_readings.drop('n_events', axis=1)\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        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,\n                      batch_ids: Optional[List[int]] = None,\n                      max_events: int = None) -> 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        batch_ids (List[int]): The batch_ids you want converted. Defaults to None (all batches will be converted)\n    \"\"\"\n    if batch_ids is None:\n        batch_ids = np.arange(1,661,1).to_list()\n    else:\n        assert isinstance(batch_ids,list), \"Variable 'batch_ids' must be list.\"\n    if not database_path.endswith('.db'):\n        database_path = database_path+'.db'\n    meta_data_iter = pq.ParquetFile(meta_data_path).iter_batches(batch_size = batch_size)\n    batch_id = 1\n    converted_batches = []\n    progress_bar = tqdm(total = len(batch_ids))\n    for meta_data_batch in meta_data_iter:\n        if batch_id in batch_ids:\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, max_events=max_events)\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            converted_batches.append(batch_id)\n        batch_id +=1\n        if len(batch_ids) == len(converted_batches):\n            break\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-10T18:08:40.028679Z","iopub.execute_input":"2023-04-10T18:08:40.029123Z","iopub.status.idle":"2023-04-10T18:08:44.856181Z","shell.execute_reply.started":"2023-04-10T18:08:40.029086Z","shell.execute_reply":"2023-04-10T18:08:44.85515Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DATA_DIR = '/kaggle/input/icecube-neutrinos-in-deep-ice'\nOUT_DIR = '/kaggle/working'\n\nSTART_BATCH = 52\nEND_BATCH = 59\nMAX_EVENTS = 200","metadata":{"execution":{"iopub.status.busy":"2023-04-10T18:08:44.857788Z","iopub.execute_input":"2023-04-10T18:08:44.859091Z","iopub.status.idle":"2023-04-10T18:08:44.86411Z","shell.execute_reply.started":"2023-04-10T18:08:44.859047Z","shell.execute_reply":"2023-04-10T18:08:44.86277Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from pathlib import Path\nfrom tqdm.auto import tqdm\nimport pandas as pd\n\ninput_data_folder = str(Path(DATA_DIR)/'train')\ngeometry_table = pd.read_csv(Path(DATA_DIR)/'sensor_geometry.csv')\nmeta_data_path = str(Path(DATA_DIR)/'train_meta.parquet')\n\nfor batch_id in tqdm(range(START_BATCH, END_BATCH+1)):\n    database_path = str(Path(OUT_DIR)/f'batch_{batch_id}')\n    convert_to_sqlite(meta_data_path,\n                      database_path=database_path,\n                      input_data_folder=input_data_folder,\n                      batch_ids = [batch_id],\n                      max_events = MAX_EVENTS)","metadata":{"execution":{"iopub.status.busy":"2023-04-10T18:08:44.865942Z","iopub.execute_input":"2023-04-10T18:08:44.866766Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!rm -rf ./software","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}