{"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":"# Saving and Loading IceCube Data as TFRecord 🚀","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nimport random\nimport tensorflow as tf\nfrom contextlib import ExitStack\nimport os\nfrom datetime import datetime\nimport multiprocess","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"meta = pd.read_parquet('/kaggle/input/icecube-neutrinos-in-deep-ice/train_meta.parquet')\nsensor_geometry = pd.read_csv(\"/kaggle/input/icecube-neutrinos-in-deep-ice/sensor_geometry.csv\")\nbatch = pd.read_parquet(\"/kaggle/input/icecube-neutrinos-in-deep-ice/train/batch_1.parquet\")\nbatch.head()","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset Creation","metadata":{}},{"cell_type":"code","source":"def create_example_protobuff(event_pulses, azimuth, zenith):\n    # convert to binary string format for Example protobuf\n    event_data = tf.io.serialize_tensor(tf.cast(event_pulses, tf.float32))\n    return tf.train.Example(\n        features=tf.train.Features(\n            feature={\n                'event_pulses': tf.train.Feature(bytes_list=tf.train.BytesList(value=[event_data.numpy()])),\n                'azimuth': tf.train.Feature(float_list=tf.train.FloatList(value=[azimuth])),\n                'zenith': tf.train.Feature(float_list=tf.train.FloatList(value=[zenith])),\n            }\n        )\n    )\n\ndef write_tfrecords(filename, dataset):\n    with ExitStack() as stack:\n        writer = stack.enter_context(tf.io.TFRecordWriter(filename))\n\n        # create example protobuffs from instances\n        for event, az_zen_pair in dataset:\n            event = tf.Variable(event)\n            azimuth, zenith = az_zen_pair\n            example = create_example_protobuff(event, azimuth, zenith)\n            writer.write(example.SerializeToString())\n            \ndef save_to_tfrecord(X, Y, name):\n    dataset = zip(X, Y)\n    \n    # write Dataset to files\n    write_tfrecords(f\"tfrecords/{name}\", dataset)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"NUM_BATCHES_TO_CONVERT = 10\n\ntrain_batches = meta.batch_id.unique()\ntrain_meta_baseline = meta[meta.batch_id.isin(train_batches[:NUM_BATCHES_TO_CONVERT])]\nprint(len(train_meta_baseline))\ntrain_meta_baseline.head()","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Now we go through the parquets and create the tfrecords after each parquet file is loaded into memory","metadata":{}},{"cell_type":"code","source":"open_batch_dict = dict()\n\ndef get_event_data(event_idx, batch_meta_df):\n    # each meta row corresponds to an event\n    batch_id, first_pulse_index, last_pulse_index, azimuth, zenith = batch_meta_df.iloc[event_idx][[\"batch_id\", \"first_pulse_index\", \"last_pulse_index\", \"azimuth\", \"zenith\"]].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(f\"/kaggle/input/icecube-neutrinos-in-deep-ice/train/batch_{batch_id}.parquet\")})\n \n\n    train_batch = open_batch_dict[batch_id]\n    event_data = train_batch.iloc[first_pulse_index:last_pulse_index + 1]\n    \n    event_data_with_pos = pd.merge(event_data, sensor_geometry, on='sensor_id', how='left')\n    \n    # Add rank and sort by rank\n    # Find valid time window\n    t_valid_length = 6200\n    t_peak = event_data_with_pos[\"time\"][event_data_with_pos[\"charge\"].argmax()]\n    t_valid_min = t_peak - t_valid_length\n    t_valid_max = t_peak + t_valid_length\n    t_valid = (event_data_with_pos[\"time\"] > t_valid_min) * (event_data_with_pos[\"time\"] < t_valid_max)\n\n    # rank\n    event_data_with_pos[\"rank\"] = 2 * (1 - event_data_with_pos[\"auxiliary\"]) + (t_valid)\n\n    # sort by rank and charge (ascending order of importance)\n    event_data_with_pos = event_data_with_pos.sort_values(['rank', 'charge'], ascending=[True, True])\n\n    final_pulses = event_data_with_pos[['x', 'y', 'z', 'time', 'charge', 'auxiliary']]\n    x = final_pulses.to_numpy().astype(np.float16)\n    \n    # get label\n    y = [azimuth, zenith]\n    \n    return x, y\n\nfor batch_id in train_batches[:NUM_BATCHES_TO_CONVERT]:\n    batch_meta_df = train_meta_baseline.loc[train_meta_baseline.batch_id == batch_id]\n    X = []\n    Y = []\n    \n    def read_event(event_idx):\n        return get_event_data(event_idx, batch_meta_df)\n    \n    # Multiprocess Events\n    iterator = range(len(batch_meta_df))\n    with multiprocess.Pool() as pool:\n        for event_x, event_y in pool.map(read_event, iterator):\n            X.append(event_x)\n            Y.append(event_y)\n            \n            \n    save_to_tfrecord(X, Y, f\"batch_{batch_id}.tfrecord\")\n    print(f\"Batch {batch_id} done and saved\")\n    \n\n    ","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset Loading","metadata":{}},{"cell_type":"code","source":"def cartesian_to_spherical(x, y, z):\n    r = math.sqrt(x**2 + y**2 + z**2)\n    zenith = math.acos(z / r)\n    azimuth = math.atan2(y, x)\n    return azimuth, zenith\n\ndef adjust_spherical(azimuth, zenith):\n    if azimuth < 0:\n        azimuth += math.pi * 2\n    elif zenith < 0:\n        zenith += math.pi\n    return azimuth, zenith\n\ndef spherical_to_cartesian(azimuth, zenith):\n    x = np.cos(azimuth) * np.sin(zenith)\n    y = np.sin(azimuth) * np.sin(zenith)\n    z = np.cos(zenith)\n    return [x, y, z]\n\ndef spherical_to_cartesian_tf(azimuth, zenith):\n    x = tf.math.cos(azimuth) * tf.math.sin(zenith)\n    y = tf.math.sin(azimuth) * tf.math.sin(zenith)\n    z = tf.math.cos(zenith)\n    return [x, y, z]","metadata":{"_kg_hide-input":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class TFDatasetManager:\n    '''\n        Create TFRecordDataset from filepaths\n    '''\n    def __init__(self, filepaths, shuffle=True, shuffle_buffer_size=1000, ragged=False, cartesian_labels=True, batch_size=128, sort_by_time=False):\n        self.filepaths = filepaths\n        self.shuffle_buffer_size = shuffle_buffer_size\n        self.batch_size = batch_size\n        self.ragged = ragged\n        self.cartesian_labels = cartesian_labels\n        self.num_features = 6\n        self.sort_by_time = sort_by_time\n        self.shuffle = shuffle\n\n    # parse serialized Example protobuf\n    def preprocess(self, tfrecord):\n        # to parse we need the feature description of the protobuf\n        feature_description = {\n            'event_pulses': tf.io.RaggedFeature(value_key='event_pulses', dtype=tf.string),\n            'azimuth': tf.io.FixedLenFeature([], tf.float32, default_value=0),\n            'zenith': tf.io.FixedLenFeature([], tf.float32, default_value=0)\n        }\n        parsed_example = tf.io.parse_single_example(tfrecord, feature_description)\n        event_pulses = tf.io.parse_tensor(tf.squeeze(parsed_example['event_pulses']), out_type=tf.float32)    \n\n        if self.ragged == True:\n            event_pulses = tf.ensure_shape(event_pulses, (None, self.num_features))\n        else:\n            # Truncate or pad to size\n            end_size = max_pulse_count\n            current_size = tf.shape(event_pulses)[0]\n            if (current_size > end_size):\n#                 event_pulses = event_pulses[:end_size]\n                # these are ordered ascending by importance so we grab the last ones\n                event_pulses = event_pulses[-end_size:]\n                \n            elif (current_size < end_size):\n                diff = end_size - current_size\n                zeros = tf.zeros((diff, self.num_features))\n                event_pulses = tf.concat((event_pulses, zeros), axis=0)\n\n            event_pulses = tf.reshape(event_pulses, (end_size, self.num_features))\n            \n        if self.sort_by_time == True:\n            def sort_by_col(a, col):\n                return tf.gather(a, tf.nn.top_k( -a[:, col], k=a.get_shape()[0]).indices)\n            time_index = 3\n            event_pulses = sort_by_col(event_pulses, time_index)\n\n        azimuth = parsed_example['azimuth']\n        zenith = parsed_example['zenith']\n\n        if self.cartesian_labels == True:\n            xyz = spherical_to_cartesian_tf(azimuth, zenith)\n            return event_pulses, xyz\n        elif self.cartesian_labels == False:\n            return event_pulses, (azimuth, zenith)\n\n    def create_dataset(self):\n        # reading all filepaths in parallel\n        dataset = tf.data.TFRecordDataset(self.filepaths, num_parallel_reads=len(self.filepaths))\n\n        if self.shuffle == True:\n            dataset = dataset.shuffle(self.shuffle_buffer_size)\n        # parse serialized Dataset\n        dataset = dataset.map(self.preprocess, num_parallel_calls=tf.data.AUTOTUNE)\n\n        # batch size of -1 indicates that we want no batching. (In this case because we are using ragged tensors)\n        if self.batch_size != -1:\n            dataset = dataset.batch(self.batch_size, drop_remainder=True)\n        # be 1 batch ahead\n        return dataset.prefetch(tf.data.AUTOTUNE)\n    ","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"batch_size = 1024\n\nfilepaths = glob.glob(\"tfrecords/*\")\nneutrino_dataset = TFDatasetManager(filepaths, \n                                    shuffle=False,\n                                    batch_size=batch_size, \n                                    sort_by_time=True, \n                                    cartesian_labels=True, \n                                    ragged=False\n                                   )\n\nfull_set = neutrino_dataset.create_dataset()","metadata":{},"execution_count":null,"outputs":[]}]}