Awkward_Target_for_Kaitai_Struct.ipynb

{"cells":[{"cell_type":"markdown","metadata":{},"source":["#  Awkward Target for Kaitai Struct\n","Data formats for scientific data often differ across experiments due the hardware design and availability constraints. To interact with these unique data formats, researchers have to develop, document and maintain specific analysis software which are often tightly coupled with a particular data format.\n","\n","This proliferation of custom data formats has been a prominent challenge for the Nuclear and High Energy Physics (NHEP) community. Within the Large Hadron Collider (LHC) experiments, this problem has largely been mitigated with the widespread adoption of `ROOT`.\n","\n","However, not all experiments in the NHEP community use `ROOT` for their data formats. Experiments such as [Cryogenic Dark Matter Search (CDMS)](https://supercdms.slac.stanford.edu/overview) continue to use custom data formats to meet specific research needs. Therefore, simplifying the process of converting a unique data format to analysis code while requiring no changes to existing data formats still holds immense value for the broader NHEP community."]},{"cell_type":"markdown","metadata":{},"source":["<p>\n","    <div>\n","        <img src=\"img/Kaitai%20logo.png\" width=\"35\" style=\"float: left; margin-right: 10px;\"/>\n","        <h2>Kaitai Struct YAML</h2>\n","    </div>\n","</p>\n","\n","[Kaitai Struct](https://kaitai.io/) is a declarative language that uses a YAML-like description of a binary data structure to generate code, in any supported language, for reading a raw data file.\n","\n","The main idea is that a particular format is described in `Kaitai Struct YAML (KSY)` language and then can be compiled with [Kaitai Struct Compiler (ksc)](https://github.com/kaitai-io/kaitai_struct_compiler) into source files in one of the supported programming languages. These modules include a generated code for a parser that can read the described data structure from a file or stream and can be included as a library in any external code."]},{"cell_type":"markdown","metadata":{},"source":["<p style=\"text-align: center;\">\n","    <img src=\"img/Kaitai%20workflow.png\" width=\"900\">\n","</p>"]},{"cell_type":"markdown","metadata":{},"source":["### Kaitai Example\n","\n","Take an example of a simple binary data structure that describes an animal species, age, and weight. You want to analyze this animal data in Python. This data structure can be described in `KSY` format. This file can then be passed to the Kaitai Struct Compiler which generates a Python script that can be imported into any other Python code.\n"]},{"cell_type":"markdown","metadata":{},"source":["<p style=\"float: right;\">\n","    <img src=\"img/animal.png\" width=\"600\" style=\"margin-bottom: 10px;\"> <br>\n","    <img src=\"img/animal%20data%20format.png\" width=\"600\" style=\"margin-top: 10px; margin-bottom: 10px;\"> <br>\n","</p>\n","\n","```yaml\n","meta:\n","  id: animal\n","  endian: le\n","  license: CC0-1.0\n","  ks-version: 0.8\n","\n","seq:\n","\n","  - id: entry\n","    type: animal_entry\n","    repeat: eos\n","\n","types:\n","\n","  animal_entry:\n","    seq:\n","      - id: str_len\n","        type: u1\n","\n","      - id: species\n","        type: str\n","        size: str_len\n","        encoding: UTF-8\n","\n","      - id: age\n","        type: u1\n","\n","      - id: weight\n","        type: u2\n","```"]},{"cell_type":"markdown","metadata":{},"source":["This is how it looks on the [Kaitai WebIDE](https://ide.kaitai.io/).\n","<p style=\"text-align: center;\">\n","    <img src=\"img/Kaitai%20WebIDE.png\" width=\"1000\">\n","</p>"]},{"cell_type":"markdown","metadata":{},"source":["#### Generate the code with Kaitai Struct Compiler\n","<p style=\"text-align: center;\">\n","    <img src=\"img/animal%20python.png\" width=\"550\" style=\"margin-top: 10px;\">\n","</p>"]},{"cell_type":"markdown","metadata":{},"source":["#### Generated Python script `animal.py`\n","\n","```python\n","import kaitaistruct\n","from kaitaistruct import KaitaiStruct, KaitaiStream, BytesIO\n","\n","\n","if getattr(kaitaistruct, 'API_VERSION', (0, 9)) < (0, 9):\n","    raise Exception(\"Incompatible Kaitai Struct Python API: 0.9 or later is required, but you have %s\" % (kaitaistruct.__version__))\n","\n","class Animal(KaitaiStruct):\n","    def __init__(self, _io, _parent=None, _root=None):\n","        self._io = _io\n","        self._parent = _parent\n","        self._root = _root if _root else self\n","        self._read()\n","\n","    def _read(self):\n","        self.entry = []\n","        i = 0\n","        while not self._io.is_eof():\n","            self.entry.append(Animal.AnimalEntry(self._io, self, self._root))\n","            i += 1\n","\n","\n","    class AnimalEntry(KaitaiStruct):\n","        def __init__(self, _io, _parent=None, _root=None):\n","            self._io = _io\n","            self._parent = _parent\n","            self._root = _root if _root else self\n","            self._read()\n","\n","        def _read(self):\n","            self.str_len = self._io.read_u1()\n","            self.species = (self._io.read_bytes(self.str_len)).decode(u\"UTF-8\")\n","            self.age = self._io.read_u1()\n","            self.weight = self._io.read_u2le()\n","```"]},{"cell_type":"markdown","metadata":{},"source":["#### Output\n","```python\n","Species: cat\n","Age: 5\n","Weight: 12\n","Species: dog\n","Age: 3\n","Weight: 43\n","Species: turtle\n","Age: 10\n","Weight: 5\n","```"]},{"cell_type":"markdown","metadata":{},"source":["<div>\n","    <img src=\"img/Awkward%20logo.png\" width=\"200\" style=\"align=left;\"/>\n","</div>\n","\n","When dealing with large files and complicated data structures of scientific data which consists of nested arrays, even the most efficient Python code can be quite time and resource-heavy."]},{"cell_type":"markdown","metadata":{},"source":["<p style=\"text-align: center;\">\n","    <strong>MIDAS Event Format</strong> <br> <br>\n","    <img src=\"img/MIDAS.jpg\" width=\"800\">\n","</p>"]},{"cell_type":"markdown","metadata":{},"source":["[Awkward Array](https://awkward-array.org/doc/main/) is a Scikit-HEP library that uses NumPy-like idioms to create arrays of nested records, variable-length lists, mixed types, and missing data. It offers a dynamic and memory efficient approach to represent complex data structures into Numpy-like arrays."]},{"cell_type":"code","execution_count":15,"metadata":{},"outputs":[{"data":{"text/html":["<pre>[[{x: 1.1, y: [1]}, {x: 2.2, y: [...]}, {x: 3.3, y: [1, 2, 3]}],\n"," [],\n"," [{x: 4.4, y: [1, 2, 3, 4]}, {x: 5.5, y: [1, ..., 5]}]]\n","----------------------------------------------------------------\n","type: 3 * var * {\n","    x: float64,\n","    y: var * int64\n","}</pre>"],"text/plain":["<Array [[{x: 1.1, y: [1]}, ..., {...}], ...] type='3 * var * {x: float64, y...'>"]},"execution_count":15,"metadata":{},"output_type":"execute_result"}],"source":["import numpy as np\n","import awkward as ak\n","\n","array = ak.Array([\n","    [{\"x\": 1.1, \"y\": [1]}, {\"x\": 2.2, \"y\": [1, 2]}, {\"x\": 3.3, \"y\": [1, 2, 3]}],\n","    [],\n","    [{\"x\": 4.4, \"y\": [1, 2, 3, 4]}, {\"x\": 5.5, \"y\": [1, 2, 3, 4, 5]}]\n","])\n","\n","array"]},{"cell_type":"markdown","metadata":{},"source":["Therefore, we propose adding Awkward Array as one of the target languages for Kaitai  Struct using the header-only `LayoutBuilder`.\n","\n","Refer to [Ianna Osborne's talk]((https://github.com/ianna/PyHEP-Users-Workshop-2023/blob/main/slides/Awkward_LayoutBuilder.ipynb)) and the [User Guide](https://awkward-array.org/doc/main/user-guide/how-to-use-header-only-layoutbuilder.html?highlight=layout%20builder) to know more about `LayoutBuilder`."]},{"cell_type":"markdown","metadata":{},"source":["## Awkward Target for Kaitai: How it works?\n","\n","Users can simply describe their data format in the `.ksy` format just once. Then this KSY file can be converted into a compiled Python module which takes the raw data and converts it into Awkward Arrays."]},{"cell_type":"markdown","metadata":{},"source":["<p style=\"text-align: center;\">\n","    <img src=\"img/Awkward%20Kaitai%20workflow.png\" width=\"1100\">\n","</p>"]},{"cell_type":"markdown","metadata":{},"source":["### Example"]},{"cell_type":"markdown","metadata":{},"source":["Take an example of this simple `fake.ksy` file - \n","\n","```yaml\n","meta:\n","  id: fake\n","  file-extension: raw\n","  endian: le\n","\n","seq:\n","  - id: points\n","    type: point\n","    repeat: eos\n","\n","types:\n","  point:\n","    seq:\n","      - id: x\n","        type: u4\n","      - id: y\n","        type: u4\n","      - id: z\n","        type: u4\n","```"]},{"cell_type":"markdown","metadata":{},"source":["### Steps"]},{"cell_type":"markdown","metadata":{},"source":["#### 1. First, Clone the [*kaitai_awkward_runtime*](https://github.com/ManasviGoyal/kaitai_awkward_runtime.git) and change the directory"]},{"cell_type":"markdown","metadata":{"tags":[],"trusted":true},"source":["```bash\n","git clone --recursive https://github.com/ManasviGoyal/kaitai_awkward_runtime.git\n","cd kaitai_awkward_runtime\n","````"]},{"cell_type":"markdown","metadata":{},"source":["#### 2. Then, Install the Python module via *pip*"]},{"cell_type":"markdown","metadata":{"tags":[],"trusted":true},"source":["```bash\n","pip install . --config-settings 'cmake.define.KSY=schemas/fake.ksy'\n","```"]},{"cell_type":"markdown","metadata":{},"source":["<p style=\"text-align: center;\">\n","    <img src=\"img/fake%20awkward.png\" width=\"500\">\n","</p>"]},{"cell_type":"markdown","metadata":{},"source":["#### Scala Code snippet: For generating C++ string for Builder Declaration\n","``` scala\n","case class RecordBuilder(\n","  fields: ListBuffer[String],\n","  contents: ListBuffer[LayoutBuilder]\n",") extends LayoutBuilder {\n","\n","  override def printStructure(indent: Int, builderName: String): String = {\n","    UserDefinedMap(builderName)\n","    val fieldStrings = fields.zip(contents).zipWithIndex.map { case ((field, content), i) =>\n","      s\"${\"\\t\" * (indent + 2)}RecordField<Field_${builderName}::$field, ${content.printStructure(indent + 1, s\"${field}\")}>\"\n","    }\n","    s\"${if (indent == 0) \"\\nusing \" + builderName.capitalize + \"BuilderType =\\n\\t\" else \"\"}\" +\n","    s\"RecordBuilder<\\n${fieldStrings.mkString(\",\\n\")}\\n${\"\\t\" * (indent)}\\t>\"\n","  }\n","\n","  def UserDefinedMap(builderName: String): Unit = {\n","    val mapStrings = fields.zipWithIndex.map { case (field, i) =>\n","      s\"\"\"{Field_${builderName}::$field, \"$field\"}\"\"\"\n","    }\n","    outHdrAwkward.puts(s\"enum Field_${builderName} : std::size_t {${fields.mkString(\", \")}};\")\n","    outSrcAwkward.puts\n","    outSrcAwkward.puts(s\"UserDefinedMap ${builderName}_fields_map({\\n\\t${mapStrings.mkString(\",\\n\\t\")}});\")\n","  }\n","}\n","```"]},{"cell_type":"markdown","metadata":{},"source":["#### Generated C++ libraries with awkward code"]},{"cell_type":"markdown","metadata":{},"source":["<p style=\"text-align: center;\">\n","    <strong><h4>fake.h</h4></strong>\n","</p>\n","\n","```C++\n","#ifndef FAKE_H_\n","#define FAKE_H_\n","\n","// This is a generated file! Please edit source .ksy file and use kaitai-struct-compiler to rebuild\n","\n","#include \"kaitai/kaitaistruct.h\"\n","#include <stdint.h>\n","#include \"awkward/LayoutBuilder.h\"\n","#include <vector>\n","\n","using UserDefinedMap = std::map<std::size_t, std::string>;\n","template<class... BUILDERS>\n","using RecordBuilder = awkward::LayoutBuilder::Record<UserDefinedMap, BUILDERS...>;\n","template<std::size_t field_name, class BUILDER>\n","using RecordField = awkward::LayoutBuilder::Field<field_name, BUILDER>;\n","template<class PRIMITIVE, class BUILDER>\n","using ListOffsetBuilder = awkward::LayoutBuilder::ListOffset<PRIMITIVE, BUILDER>;\n","template<class PRIMITIVE>\n","using NumpyBuilder = awkward::LayoutBuilder::Numpy<PRIMITIVE>;\n","\n","enum Field_fake : std::size_t {points};\n","enum Field_points : std::size_t {x, y, z};\n","\n","using FakeBuilderType =\n","\tRecordBuilder<\n","\t\tRecordField<Field_fake::points, ListOffsetBuilder<int64_t, RecordBuilder<\n","\t\t\tRecordField<Field_points::x, NumpyBuilder<uint32_t>>,\n","\t\t\tRecordField<Field_points::y, NumpyBuilder<uint32_t>>,\n","\t\t\tRecordField<Field_points::z, NumpyBuilder<uint32_t>>\n","\t\t>>>\n","\t>;\n","using PointsBuilderType = decltype(std::declval<FakeBuilderType>().content<Field_fake::points>().content());\n","\n","\n","#if KAITAI_STRUCT_VERSION < 9000L\n","#error \"Incompatible Kaitai Struct C++/STL API: version 0.9 or later is required\"\n","#endif\n","\n","class fake_t : public kaitai::kstruct {\n","\n","public:\n","    class point_t;\n","\n","    fake_t(kaitai::kstream* p__io, kaitai::kstruct* p__parent = 0, fake_t* p__root = 0);\n","\n","private:\n","    void _read();\n","    void _clean_up();\n","\n","public:\n","    ~fake_t();\n","\n","    class point_t : public kaitai::kstruct {\n","\n","    public:\n","\n","        point_t(kaitai::kstream* p__io, fake_t* p__parent = 0, fake_t* p__root = 0);\n","\n","    private:\n","        void _read();\n","        void _clean_up();\n","\n","    public:\n","        ~point_t();\n","\n","    private:\n","        uint32_t m_x;\n","        uint32_t m_y;\n","        uint32_t m_z;\n","        fake_t* m__root;\n","        fake_t* m__parent;\n","\n","    public:\n","        uint32_t x() const { return m_x; }\n","        uint32_t y() const { return m_y; }\n","        uint32_t z() const { return m_z; }\n","        fake_t* _root() const { return m__root; }\n","        fake_t* _parent() const { return m__parent; }\n","        PointsBuilderType points_builder;\n","    };\n","\n","private:\n","    std::vector<point_t*>* m_points;\n","    fake_t* m__root;\n","    kaitai::kstruct* m__parent;\n","\n","public:\n","    std::vector<point_t*>* points() const { return m_points; }\n","    fake_t* _root() const { return m__root; }\n","    kaitai::kstruct* _parent() const { return m__parent; }\n","    FakeBuilderType fake_builder;\n","};\n","\n","#endif  // FAKE_H_\n","```"]},{"cell_type":"markdown","metadata":{},"source":["<p style=\"text-align: center;\">\n","    <strong><h4>fake.cpp</h4></strong>\n","</p>\n","\n","```C++\n","// This is a generated file! Please edit source .ksy file and use kaitai-struct-compiler to rebuild\n","\n","#include \"fake.h\"\n","\n","UserDefinedMap fake_fields_map({\n","\t{Field_fake::points, \"points\"}});\n","\n","UserDefinedMap points_fields_map({\n","\t{Field_points::x, \"x\"},\n","\t{Field_points::y, \"y\"},\n","\t{Field_points::z, \"z\"}});\n","\n","fake_t::fake_t(kaitai::kstream* p__io, kaitai::kstruct* p__parent, fake_t* p__root) : kaitai::kstruct(p__io) {\n","    m__parent = p__parent;\n","    m__root = this;\n","    m_points = 0;\n","\n","    fake_builder.set_fields(fake_fields_map);\n","\n","    try {\n","        _read();\n","    } catch(...) {\n","        _clean_up();\n","        throw;\n","    }\n","}\n","\n","void fake_t::_read() {\n","    m_points = new std::vector<point_t*>();\n","    {\n","        int i = 0;\n","        auto& sub_points_builder = fake_builder.content<Field_fake::points>();\n","        while (!m__io->is_eof()) {\n","            sub_points_builder.begin_list();\n","            m_points->push_back(new point_t(m__io, this, m__root));\n","            sub_points_builder.end_list();\n","            i++;\n","        }\n","    }\n","}\n","\n","fake_t::~fake_t() {\n","    _clean_up();\n","}\n","\n","void fake_t::_clean_up() {\n","    if (m_points) {\n","        for (std::vector<point_t*>::iterator it = m_points->begin(); it != m_points->end(); ++it) {\n","            delete *it;\n","        }\n","        delete m_points; m_points = 0;\n","    }\n","}\n","\n","fake_t::point_t::point_t(kaitai::kstream* p__io, fake_t* p__parent, fake_t* p__root) : kaitai::kstruct(p__io),\n","\tpoints_builder(p__parent->fake_builder.content<Field_fake::points>().content()) {\n","    m__parent = p__parent;\n","    m__root = p__root;\n","\n","    points_builder.set_fields(points_fields_map);\n","\n","    try {\n","        _read();\n","    } catch(...) {\n","        _clean_up();\n","        throw;\n","    }\n","}\n","\n","void fake_t::point_t::_read() {\n","    m_x = m__io->read_u4le();\n","    auto& x_builder = points_builder.content<Field_points::x>();\n","    x_builder.append(m_x);\n","    m_y = m__io->read_u4le();\n","    auto& y_builder = points_builder.content<Field_points::y>();\n","    y_builder.append(m_y);\n","    m_z = m__io->read_u4le();\n","    auto& z_builder = points_builder.content<Field_points::z>();\n","    z_builder.append(m_z);\n","}\n","\n","fake_t::point_t::~point_t() {\n","    _clean_up();\n","}\n","\n","void fake_t::point_t::_clean_up() {\n","}\n","```"]},{"cell_type":"markdown","metadata":{},"source":["3. Finally, you can print the returned `ak.Array`\n","```python\n","import kaitai_awkward_runtime\n","awkward_arrays = kaitai_awkward_runtime.load(\"data/fake.raw\")\n","print(awkward_arrays)\n","\n","[{points: [{x: 1, y: 2, z: 3}]}, {...}, ..., {points: [{x: 11, y: 22, ...}]}]\n","```"]},{"cell_type":"markdown","metadata":{},"source":["### Future Plans\n","- Improving the User Interface by making it simpler and more customizable\n","- Adding support for `IndexedOption` Arrays"]},{"cell_type":"markdown","metadata":{},"source":["For any questions or suggestions, feel free to contact:\n","- Manasvi Goyal - `mg.manasvi@gmail.com`\n","- Amy Roberts - `amy.roberts@ucdenver.edu`"]},{"cell_type":"markdown","metadata":{},"source":["### Acknowledgements\n","\n","<div>\n","    <img src=\"img/NSF%20logo.png\" width=\"60\" style=\"float: left;\">\n","    Support for this work was provided by NSF cooperative agreements <a href=\"https://www.nsf.gov/awardsearch/showAward?AWD_ID=1836650\">OAC-1836650</a> and <a href=\"https://www.nsf.gov/awardsearch/showAward?AWD_ID=2323298\">PHY-2323298</a> (IRIS-HEP) and grants <a href=\"https://www.nsf.gov/awardsearch/showAward?AWD_ID=2104003\">OAC-2104003</a> (PONDD) and <a href=\"https://www.nsf.gov/awardsearch/showAward?AWD_ID=2103945\">OAC-2103945</a> (Awkward Array).\n","</div>"]}],"metadata":{"kernelspec":{"display_name":"Python 3 (ipykernel)","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.11.5"}},"nbformat":4,"nbformat_minor":4}