{"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":"# Fast preprocessing with Polars and Numba\n\nIf you're using a LSTM like method for this competition such as [LSTM Baseline by ZhaounBooty](https://www.kaggle.com/code/seungmoklee/lstm-preprocessing-point-picker). Every time you want to try out a different preprocessing idea, you have to regenerate a large set of files. This is both time consuming and takes up a lot of space if you're using your own machine.\n\nThis notebook contains an optimized preprocessing setup based on Polars and Numba. With it I'm able to quickly try out different preprocessing ideas without having to pre-cache the data. \n\n### Timings for preprocessing one batch: \n- Kaggle - 7 seconds\n- i5 12400 Processor - 1.5 seconds (probably because Kaggle disk read is slow)\n\n### Key Insights\n1. Use `polars` for loading parquet files, way faster than `pandas`\n2. For padding, pre-assign the `np.zeros` array instead of using `np.pad`\n3. For loop based operations operations, such as sampling points or filtering based on particular data, use `numba`. \n\nNumba compatible code is slightly annoying to develop. Removing numba raises the per batch time to around 3 seconds on my computer, which is still acceptable for quickly trying out a new idea. Here I provide a decent numba example.","metadata":{}},{"cell_type":"code","source":"import os, time\nimport polars as pl\nimport numpy as np\nimport numba as nb","metadata":{"execution":{"iopub.status.busy":"2023-03-30T21:27:48.582386Z","iopub.execute_input":"2023-03-30T21:27:48.583407Z","iopub.status.idle":"2023-03-30T21:27:50.218854Z","shell.execute_reply.started":"2023-03-30T21:27:48.583353Z","shell.execute_reply":"2023-03-30T21:27:50.217248Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DATA_DIR = '/kaggle/input/icecube-neutrinos-in-deep-ice/'","metadata":{"execution":{"iopub.status.busy":"2023-03-30T21:27:50.221364Z","iopub.execute_input":"2023-03-30T21:27:50.222216Z","iopub.status.idle":"2023-03-30T21:27:50.228786Z","shell.execute_reply.started":"2023-03-30T21:27:50.222158Z","shell.execute_reply":"2023-03-30T21:27:50.227586Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Load geometry and set parameters\n\nYou can also use a custom sensor geomtry file with more preprocessed columns\n\nThe main parameters I set are the features to use and the maximum length of the sampled event sequence.","metadata":{}},{"cell_type":"code","source":"GEOMETRY = os.path.join(DATA_DIR, \"sensor_geometry.csv\")\ngeometry = pl.scan_csv(GEOMETRY).with_columns(\n                [pl.col('sensor_id').cast(pl.Int16)]\n            )","metadata":{"execution":{"iopub.status.busy":"2023-03-30T21:27:50.230755Z","iopub.execute_input":"2023-03-30T21:27:50.231219Z","iopub.status.idle":"2023-03-30T21:27:50.294202Z","shell.execute_reply.started":"2023-03-30T21:27:50.231172Z","shell.execute_reply":"2023-03-30T21:27:50.293123Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"FEATURE_NAMES = ['time', 'charge', 'auxiliary', 'x', 'y', 'z']\nCHARGE_IDX = FEATURE_NAMES.index('charge')\nTIME_IDX = FEATURE_NAMES.index('time')\nAUX_IDX = FEATURE_NAMES.index('auxiliary')\nN_FEATURES = len(FEATURE_NAMES)\n\n# This affect the RAM usage - For 64 the usage is around 2 GB for a full batch\nMAX_SEQUENCE_LENGTH = 64","metadata":{"execution":{"iopub.status.busy":"2023-03-30T21:27:50.297906Z","iopub.execute_input":"2023-03-30T21:27:50.298865Z","iopub.status.idle":"2023-03-30T21:27:50.306298Z","shell.execute_reply.started":"2023-03-30T21:27:50.298807Z","shell.execute_reply":"2023-03-30T21:27:50.304878Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Preprocessing with Numba\n\nAt first I wanted to find a way to vectorize a sampling operation on groups. But that up a big unreadable mess and hard to modify. This code is much simpler, easier to read and modify.\n\n### Gist of the code\n1. Prepare an zeros array of shape `[num_events, max_sequence_length, num_features]` at the beginning, when an event is processed just place it into the array instead of `np.pad`\n2. Loop through the events using the pulse indexess.\n3. If the`~auxiliary` can fill the sequence, sample from those, otherwise include the `auxiliary` indexes. \n4. Sort the indexes so that they are ordered by time.\n\nI didn't find a big difference in performance when using this vs the [Point Picker method](https://www.kaggle.com/code/seungmoklee/lstm-preprocessing-point-picker). Ofcourse, the point of this numba code is so that you can implement your custom preprocessing.\n\n### Working with Numba\nIf numba compile gives an error, its takes some time to fix. Unfortunately the error messages are not very beginner friendly. Here are 2 tips:\n\n1. Be careful about the input and output types.\n2. Comment out all the lines and keep adding them one by one to find out which line is offending the compiler.\n\nJust remove the decorators to switch to numpy/python code.\n","metadata":{}},{"cell_type":"code","source":"@nb.njit\ndef set_seed(value):\n    # To seed numba sampling, its important to set the seed with this decorated function \n    np.random.seed(value)\n\n@nb.jit( nb.float32[:,:,:](nb.float64[:,:], nb.int64[:,:]) )\ndef sample_and_pad(data, pulse_indexes):\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), MAX_SEQUENCE_LENGTH, data.shape[-1]), dtype=np.float32)\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) > MAX_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) < MAX_SEQUENCE_LENGTH:\n                max_length_possible = min(MAX_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=MAX_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        data_x[ii, :len(event_data), :] = event_data                  \n    return data_x","metadata":{"execution":{"iopub.status.busy":"2023-03-30T21:27:50.308198Z","iopub.execute_input":"2023-03-30T21:27:50.309450Z","iopub.status.idle":"2023-03-30T21:27:58.374065Z","shell.execute_reply.started":"2023-03-30T21:27:50.309397Z","shell.execute_reply":"2023-03-30T21:27:58.372759Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Load batch data using polars","metadata":{}},{"cell_type":"code","source":"batch_id = 100","metadata":{"execution":{"iopub.status.busy":"2023-03-30T21:27:58.375359Z","iopub.execute_input":"2023-03-30T21:27:58.375706Z","iopub.status.idle":"2023-03-30T21:27:58.381751Z","shell.execute_reply.started":"2023-03-30T21:27:58.375674Z","shell.execute_reply":"2023-03-30T21:27:58.380475Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# This is perhaps the slowest part, but luckily you don't need to run it over and over.\n# Alternative options: \n# 1. Load the full metadata parquet into memory - This is fine if you have enough memory\n# 2. Split the metadata into individual batches and save them on disk - Only takes around 3 GB\n\nstart_time = time.perf_counter()\nMETADATA = os.path.join(DATA_DIR, \"train_meta.parquet\")\nmetadata = pl.scan_parquet(METADATA)\nbatch_meta = metadata.filter(pl.col('batch_id') == batch_id)\npulse_indexes = batch_meta.select(['first_pulse_index', 'last_pulse_index']).collect().to_numpy()\ndata_y = batch_meta.select(['azimuth', 'zenith']).collect().to_numpy()\nprint(\"Processed batch in\", time.perf_counter() - start_time, \"s\")","metadata":{"execution":{"iopub.status.busy":"2023-03-30T21:27:58.383175Z","iopub.execute_input":"2023-03-30T21:27:58.383772Z","iopub.status.idle":"2023-03-30T21:28:12.017904Z","shell.execute_reply.started":"2023-03-30T21:27:58.383736Z","shell.execute_reply":"2023-03-30T21:28:12.016774Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Here I use the split batches shared by SolverWorld\n# https://www.kaggle.com/datasets/solverworld/train-meta-parquet\n\nstart_time = time.perf_counter()\nsplit_meta_file = f'/kaggle/input/train-meta-parquet/train_meta_{batch_id}.parquet'\nbatch_meta = pl.read_parquet(split_meta_file)\npulse_indexes = batch_meta.select(['first_pulse_index', 'last_pulse_index']).to_numpy()\ndata_y = batch_meta.select(['azimuth', 'zenith']).to_numpy()\n# Even this load time is likely affected by Kaggle disk fetch time\nprint(\"Processed batch in\", time.perf_counter() - start_time, \"s\")","metadata":{"execution":{"iopub.status.busy":"2023-03-30T21:28:12.021768Z","iopub.execute_input":"2023-03-30T21:28:12.022107Z","iopub.status.idle":"2023-03-30T21:28:12.558733Z","shell.execute_reply.started":"2023-03-30T21:28:12.022075Z","shell.execute_reply":"2023-03-30T21:28:12.557752Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Polars data load\n\nThis is fairly simple, but I was surprised how much faster polars is for this job.","metadata":{}},{"cell_type":"code","source":"start_time = time.perf_counter()\n\nbatch_file = os.path.join(DATA_DIR, 'train', f'batch_{batch_id}.parquet')\nbatch = pl.scan_parquet(batch_file)\nbatch = batch.join(geometry, on='sensor_id', how='left')\n\n# Load is slow probably because of Kaggle disk load time\n# Much much faster (<0.5 seconds) on my consumer pc\ndata = batch.select(FEATURE_NAMES).collect().to_numpy() \n\nprint(\"Loaded batch in\", time.perf_counter() - start_time, \"s\")\nstart_time = time.perf_counter()\n\n# main preprocessing on numba\nset_seed(42)\ndata_x = sample_and_pad(data, pulse_indexes)\n\nprint(\"Processed batch in\", time.perf_counter() - start_time, \"s\")","metadata":{"execution":{"iopub.status.busy":"2023-03-30T21:28:23.483543Z","iopub.execute_input":"2023-03-30T21:28:23.483943Z","iopub.status.idle":"2023-03-30T21:28:29.831702Z","shell.execute_reply.started":"2023-03-30T21:28:23.483902Z","shell.execute_reply":"2023-03-30T21:28:29.830531Z"},"trusted":true},"execution_count":null,"outputs":[]}]}