notebooks/ASDF_array_storage_intro.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "27a2e618-384e-4b9a-ab18-d46755e644be",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "# ASDF array storage"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "382ae69d-3ee2-4a58-bdd1-94a6caaa88a1",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": []
   },
   "source": [
    "One of the powerful features of asdf is the ability to read and write array data.\n",
    "\n",
    "These data can be small, large or 'big' each presenting a unique set of challenges.\n",
    "\n",
    "Being familiar with how array data is serialized, written to disk, deserialized and read can help to avoid performance issues, memory problems and puzzling bugs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe78aa88-c394-43f1-b549-6fe9ddd0614b",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "## ASDF file layout: the tree and blocks\n",
    "\n",
    "Asdf, the python implementation of the ASDF standard, supports reading and writing a wide variety of python objects to ASDF files. ASDF files are made of two major components:\n",
    "\n",
    "- a `tree` (stored as YAML)\n",
    "- binary `blocks` (typically one for each array in the tree)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d3aee0f-78a2-42d9-a934-a0ddf9dd1cf0",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "To begin exploring how asdf stores array data in blocks, let's start by making a simple ASDF file with a few arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "03ee0379-1cd7-4aed-a08d-a93ab7574925",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1mroot\u001b[0m (AsdfObject)\n",
      "\u001b[2m├─\u001b[0m\u001b[1marray0\u001b[0m (ndarray): shape=(42,), dtype=int64\n",
      "\u001b[2m├─\u001b[0m\u001b[1marray1\u001b[0m (ndarray): shape=(7,), dtype=int64\n",
      "\u001b[2m└─\u001b[0m\u001b[1marray2\u001b[0m (ndarray): shape=(26,), dtype=int64\n"
     ]
    }
   ],
   "source": [
    "import asdf\n",
    "import numpy as np\n",
    "\n",
    "# construct an AsdfFile object with 3 arrays\n",
    "example_af = asdf.AsdfFile({\n",
    "    'array0': np.arange(42),\n",
    "    'array1': np.arange(7),\n",
    "    'array2': np.arange(26),\n",
    "})\n",
    "\n",
    "# print out summary info for the tree\n",
    "example_af.info()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb30f484-46c7-43bc-abc3-6f877931ea4b",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "Next, let's define a few functions to help us look at what bytes are produced when asdf writes out an AsdfFile object to an ASDF file. To avoid writing files to disk (unless necessary) let's write a function to write an AsdfFile object to an `io.BytesIO` object (that can be used like a file) and apply it to our `example_af`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7916fe1a-6bfd-40aa-b763-1fca10d5a9ea",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "b'#ASDF 1.0.0\\n#ASDF_STANDARD 1.5.0\\n%YAML 1.1\\n%TAG ! tag:stsci.edu:asdf/\\n--- !core/asdf-1.1.0\\nasdf_library: !core/software-1.0.0 {author: The ASDF Developers, homepage: \\'http://github.com/asdf-format/asdf\\',\\n  name: asdf, version: 3.0.0.dev307+gb0c9a50}\\nhistory:\\n  extensions:\\n  - !core/extension_metadata-1.0.0\\n    extension_class: asdf.extension._manifest.ManifestExtension\\n    extension_uri: asdf://asdf-format.org/core/extensions/core-1.5.0\\n    software: !core/software-1.0.0 {name: asdf, version: 3.0.0.dev307+gb0c9a50}\\n  - !core/extension_metadata-1.0.0\\n    extension_class: asdf.extension.BuiltinExtension\\n    software: !core/software-1.0.0 {name: asdf, version: 3.0.0.dev307+gb0c9a50}\\narray0: !core/ndarray-1.0.0\\n  source: 0\\n  datatype: int64\\n  byteorder: little\\n  shape: [42]\\narray1: !core/ndarray-1.0.0\\n  source: 1\\n  datatype: int64\\n  byteorder: little\\n  shape: [7]\\narray2: !core/ndarray-1.0.0\\n  source: 2\\n  datatype: int64\\n  byteorder: little\\n  shape: [26]\\n...\\n\\xd3BLK\\x000\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x01P\\x00\\x00\\x00\\x00\\x00\\x00\\x01P\\x00\\x00\\x00\\x00\\x00\\x00\\x01P\\xa4=\\x1c\\xe7;q\\xf3\\xc0fC\\xe0kF\\xed\\x05\\xc8\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x02\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x03\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x04\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x05\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x06\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x07\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x08\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\t\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\n\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x0b\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x0c\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\r\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x0e\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x0f\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x10\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x11\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x12\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x13\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x14\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x15\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x16\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x17\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x18\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x19\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x1a\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x1b\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x1c\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x1d\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x1e\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x1f\\x00\\x00\\x00\\x00\\x00\\x00\\x00 \\x00\\x00\\x00\\x00\\x00\\x00\\x00!\\x00\\x00\\x00\\x00\\x00\\x00\\x00\"\\x00\\x00\\x00\\x00\\x00\\x00\\x00#\\x00\\x00\\x00\\x00\\x00\\x00\\x00$\\x00\\x00\\x00\\x00\\x00\\x00\\x00%\\x00\\x00\\x00\\x00\\x00\\x00\\x00&\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\'\\x00\\x00\\x00\\x00\\x00\\x00\\x00(\\x00\\x00\\x00\\x00\\x00\\x00\\x00)\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\xd3BLK\\x000\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x008\\x00\\x00\\x00\\x00\\x00\\x00\\x008\\x00\\x00\\x00\\x00\\x00\\x00\\x008\\x9f\\r\\x04\\xec\\x1aa\\xa3u\\xa4\\x16.P\\xc3\\x0eH\\xea\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x02\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x03\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x04\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x05\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x06\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\xd3BLK\\x000\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\xd0\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\xd0\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\xd0\\x0e\\xc9\\xeb\\xdc\\xc4-\\xe6\\t\\x12,:q\\x076\\xc0*\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x02\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x03\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x04\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x05\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x06\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x07\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x08\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\t\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\n\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x0b\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x0c\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\r\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x0e\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x0f\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x10\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x11\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x12\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x13\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x14\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x15\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x16\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x17\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x18\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x19\\x00\\x00\\x00\\x00\\x00\\x00\\x00#ASDF BLOCK INDEX\\n%YAML 1.1\\n---\\n- 968\\n- 1358\\n- 1468\\n...\\n'\n"
     ]
    }
   ],
   "source": [
    "import pathlib\n",
    "import tempfile\n",
    "\n",
    "_temporary_directory = tempfile.TemporaryDirectory()\n",
    "tmp_path = pathlib.Path(_temporary_directory.name)\n",
    "\n",
    "example_fn = tmp_path / 'example.asdf'\n",
    "example_af.write_to(example_fn)\n",
    "\n",
    "with open(example_fn, 'rb') as f:\n",
    "    print(f.read())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3ee24f5-bd06-4ee1-bcd7-f2d472134445",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "This wall of bytes isn't the easiest to view so let's define a few functions to split the ASDF file bytes into the ASDF tree and blocks. We can use the ASDF block 'magic' bytes `b'\\xd3BLK'` (which occur at the start of each ASDF block) to break up the file into:\n",
    "\n",
    "- tree (bytes before the first block)\n",
    "- blocks\n",
    "- block index (an optional YAML document at the end of the file that defines the byte offsets for each block)\n",
    "\n",
    "We will ignore the block index throughout this document (please see the [Low-Level file layout documentation for the asdf-standard](https://asdf-standard.readthedocs.io/en/latest/file_layout.html) for details about the block index and the ASDF file format)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "20343316-71ab-4c89-9cfd-9b1d91b9b601",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== ASDF Tree ===\n",
      "#ASDF 1.0.0\n",
      "#ASDF_STANDARD 1.5.0\n",
      "%YAML 1.1\n",
      "%TAG ! tag:stsci.edu:asdf/\n",
      "--- !core/asdf-1.1.0\n",
      "asdf_library: !core/software-1.0.0 {author: The ASDF Developers, homepage: 'http://github.com/asdf-format/asdf',\n",
      "  name: asdf, version: 3.0.0.dev307+gb0c9a50}\n",
      "history:\n",
      "  extensions:\n",
      "  - !core/extension_metadata-1.0.0\n",
      "    extension_class: asdf.extension._manifest.ManifestExtension\n",
      "    extension_uri: asdf://asdf-format.org/core/extensions/core-1.5.0\n",
      "    software: !core/software-1.0.0 {name: asdf, version: 3.0.0.dev307+gb0c9a50}\n",
      "  - !core/extension_metadata-1.0.0\n",
      "    extension_class: asdf.extension.BuiltinExtension\n",
      "    software: !core/software-1.0.0 {name: asdf, version: 3.0.0.dev307+gb0c9a50}\n",
      "array0: !core/ndarray-1.0.0\n",
      "  source: 0\n",
      "  datatype: int64\n",
      "  byteorder: little\n",
      "  shape: [42]\n",
      "array1: !core/ndarray-1.0.0\n",
      "  source: 1\n",
      "  datatype: int64\n",
      "  byteorder: little\n",
      "  shape: [7]\n",
      "array2: !core/ndarray-1.0.0\n",
      "  source: 2\n",
      "  datatype: int64\n",
      "  byteorder: little\n",
      "  shape: [26]\n",
      "...\n",
      "\n",
      "=== ASDF Blocks ===\n",
      "[b'\\x000\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x01P\\x00\\x00\\x00\\x00\\x00\\x00\\x01P\\x00\\x00\\x00\\x00\\x00\\x00\\x01P\\xa4=\\x1c\\xe7;q\\xf3\\xc0fC\\xe0kF\\xed\\x05\\xc8\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x02\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x03\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x04\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x05\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x06\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x07\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x08\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\t\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\n\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x0b\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x0c\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\r\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x0e\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x0f\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x10\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x11\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x12\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x13\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x14\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x15\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x16\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x17\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x18\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x19\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x1a\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x1b\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x1c\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x1d\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x1e\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x1f\\x00\\x00\\x00\\x00\\x00\\x00\\x00 \\x00\\x00\\x00\\x00\\x00\\x00\\x00!\\x00\\x00\\x00\\x00\\x00\\x00\\x00\"\\x00\\x00\\x00\\x00\\x00\\x00\\x00#\\x00\\x00\\x00\\x00\\x00\\x00\\x00$\\x00\\x00\\x00\\x00\\x00\\x00\\x00%\\x00\\x00\\x00\\x00\\x00\\x00\\x00&\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\'\\x00\\x00\\x00\\x00\\x00\\x00\\x00(\\x00\\x00\\x00\\x00\\x00\\x00\\x00)\\x00\\x00\\x00\\x00\\x00\\x00\\x00', b'\\x000\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x008\\x00\\x00\\x00\\x00\\x00\\x00\\x008\\x00\\x00\\x00\\x00\\x00\\x00\\x008\\x9f\\r\\x04\\xec\\x1aa\\xa3u\\xa4\\x16.P\\xc3\\x0eH\\xea\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x02\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x03\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x04\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x05\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x06\\x00\\x00\\x00\\x00\\x00\\x00\\x00', b'\\x000\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\xd0\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\xd0\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\xd0\\x0e\\xc9\\xeb\\xdc\\xc4-\\xe6\\t\\x12,:q\\x076\\xc0*\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x01\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x02\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x03\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x04\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x05\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x06\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x07\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x08\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\t\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\n\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x0b\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x0c\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\r\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x0e\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x0f\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x10\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x11\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x12\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x13\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x14\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x15\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x16\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x17\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x18\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x19\\x00\\x00\\x00\\x00\\x00\\x00\\x00']\n"
     ]
    }
   ],
   "source": [
    "def read_tree_and_blocks(filename):\n",
    "    \"\"\"\n",
    "    Split the bytes that make up an ASDF file and split these\n",
    "    bytes into:\n",
    "      - tree (decoded as ascii)\n",
    "      - a list of blocks (as bytes containing block headers and contents)\n",
    "    \"\"\"\n",
    "    with open(filename, 'rb') as f:\n",
    "        asdf_bytes = f.read()\n",
    "\n",
    "\n",
    "    # if these bytes don't contain any block 'magic' bytes\n",
    "    # (which mark the start of a block) then the written bytes\n",
    "    # contain no blocks (and only an ASDF tree)\n",
    "    if b'\\xd3BLK' not in asdf_bytes:\n",
    "        return asdf_bytes.decode('ascii'), []\n",
    "\n",
    "    # split the written bytes into the tree (which occurs first in the file)\n",
    "    # and any blocks based on the block 'magic' bytes\n",
    "    tree, *blocks = asdf_bytes.split(b'\\xd3BLK')\n",
    "    blocks[-1] = blocks[-1].split(b'#ASDF BLOCK INDEX')[0]\n",
    "    return tree.decode('ascii'), blocks\n",
    "\n",
    "asdf_tree, asdf_blocks = read_tree_and_blocks(example_fn)\n",
    "\n",
    "print(\"=== ASDF Tree ===\")\n",
    "print(asdf_tree)\n",
    "\n",
    "print(\"=== ASDF Blocks ===\")\n",
    "print(asdf_blocks)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdb9fda9-ff1d-4af8-b348-a11a91abad04",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "Since we're in Jupyter we can take this a step further and make the output syntax highlighted (for the ASDF tree YAML) and make Jupyter aware that we want to view AsdfFile objects using our new function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "babe6754-d5c3-45f4-95d9-745dacc66421",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "\n",
       "### ASDF Tree\n",
       "\n",
       "```yaml\n",
       "#ASDF 1.0.0\n",
       "#ASDF_STANDARD 1.5.0\n",
       "%YAML 1.1\n",
       "%TAG ! tag:stsci.edu:asdf/\n",
       "--- !core/asdf-1.1.0\n",
       "asdf_library: !core/software-1.0.0 {author: The ASDF Developers, homepage: 'http://github.com/asdf-format/asdf',\n",
       "  name: asdf, version: 3.0.0.dev307+gb0c9a50}\n",
       "history:\n",
       "  extensions:\n",
       "  - !core/extension_metadata-1.0.0\n",
       "    extension_class: asdf.extension._manifest.ManifestExtension\n",
       "    extension_uri: asdf://asdf-format.org/core/extensions/core-1.5.0\n",
       "    software: !core/software-1.0.0 {name: asdf, version: 3.0.0.dev307+gb0c9a50}\n",
       "  - !core/extension_metadata-1.0.0\n",
       "    extension_class: asdf.extension.BuiltinExtension\n",
       "    software: !core/software-1.0.0 {name: asdf, version: 3.0.0.dev307+gb0c9a50}\n",
       "array0: !core/ndarray-1.0.0\n",
       "  source: 0\n",
       "  datatype: int64\n",
       "  byteorder: little\n",
       "  shape: [42]\n",
       "array1: !core/ndarray-1.0.0\n",
       "  source: 1\n",
       "  datatype: int64\n",
       "  byteorder: little\n",
       "  shape: [7]\n",
       "array2: !core/ndarray-1.0.0\n",
       "  source: 2\n",
       "  datatype: int64\n",
       "  byteorder: little\n",
       "  shape: [26]\n",
       "...\n",
       "\n",
       "```\n",
       "\n",
       "### ASDF Blocks [N = 3]\n",
       "\tBlock0: b'\\x000\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00'...[386 B]\n",
       "\tBlock1: b'\\x000\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00'...[106 B]\n",
       "\tBlock2: b'\\x000\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00'...[258 B]\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": []
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython.display import display, Markdown\n",
    "\n",
    "\n",
    "def n_bytes_to_human(n_bytes):\n",
    "    \"\"\"Convert a number (of bytes) to a more 'human readable' format\n",
    "    by adding a SI prefix and unit\"\"\"\n",
    "    _units = ['B', 'KB', 'MB', 'GB', 'TB', 'PB']\n",
    "    i = 0\n",
    "    while n_bytes > 1024:\n",
    "        n_bytes /= 1024\n",
    "        i += 1\n",
    "    return f\"{n_bytes:0.0f} {_units[i]}\"\n",
    "\n",
    "def display_asdf_file(af):\n",
    "    \"\"\"\n",
    "    Use Jupyter Markdown formatting to display the ASDF\n",
    "    tree and blocks produced by writing an AsdfFile, af\n",
    "    \"\"\"\n",
    "    tmp_fn = tmp_path / 'temp.asdf'\n",
    "    af.write_to(tmp_fn)\n",
    "    tree, blocks = read_tree_and_blocks(tmp_fn)\n",
    "    # for the first 10 blocks, show the first 10 bytes and size of each block\n",
    "    block_text = \"\\n\".join([\n",
    "        f\"\\tBlock{i}: {blk[:10]}...[{n_bytes_to_human(len(blk))}]\"\n",
    "        for (i, blk) in enumerate(blocks[:10])\n",
    "    ])\n",
    "    # if there are more than 10 blocks, show a message that the remaining\n",
    "    # blocks were not displayed\n",
    "    if len(blocks) > 10:\n",
    "        block_text += f\"\\n... skipping display of {len(blocks)-10} blocks\"\n",
    "    # render the tree and block text as Markdown\n",
    "    display(Markdown(f\"\"\"\n",
    "### ASDF Tree\n",
    "\n",
    "```yaml\n",
    "{tree}\n",
    "```\n",
    "\n",
    "### ASDF Blocks [N = {len(blocks)}]\n",
    "{block_text}\n",
    "    \"\"\"))\n",
    "\n",
    "# register our 'show_file' function to display AsdfFile instances\n",
    "get_ipython().display_formatter.formatters['text/plain'].for_type(\n",
    "    asdf.AsdfFile, lambda n, p, cycle: display_asdf_file(n))\n",
    "\n",
    "# ending this cell with our AsdfFile object will cause Jupyter to use our new function\n",
    "example_af"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80bd3d06-972d-4ada-96c4-b1a0d1bf3d9b",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "source": [
    "Using our helper functions, we can observe that writing out our example AsdfFile produces an ASDF tree containing our 3 arrays, each with a `source` containing the index of the ASDF block that contains the corresponding array data. \n",
    "\n",
    "Each ASDF block consists of a header and a binary representation of the array data (more information can be found in the [Low-Level file layout documentation for the asdf-standard](https://asdf-standard.readthedocs.io/en/latest/file_layout.html)).\n",
    "\n",
    "Before returning to discussing asdf array storage let's define one more helper function to track the number of bytes read from a file (this will be useful when comparing some of the storage options)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "58668963-98d5-46e2-93ef-4d1b3c6d3abc",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class ReadMonitoringFile(asdf.generic_io.RandomAccessFile):\n",
    "    \"\"\"\n",
    "    Class to record the number of bytes read from an ASDF file\n",
    "    \"\"\"\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.n_bytes = 0\n",
    "\n",
    "    def read(self, *args, **kwargs):\n",
    "        bs = super().read(*args, **kwargs)\n",
    "        self.n_bytes += len(bs)\n",
    "        return bs\n",
    "\n",
    "import contextlib\n",
    "\n",
    "@contextlib.contextmanager\n",
    "def asdf_open_tracking_reads(fn):\n",
    "    with open(fn, 'rb') as f:\n",
    "        with ReadMonitoringFile(f, 'r') as mf:\n",
    "            with asdf.open(mf) as af:\n",
    "                yield af, mf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8873235e-3140-4c62-97ec-352f3d24753c",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "## Lazy loading blocks\n",
    "When the ASDF file is loaded, array data is (by default) \"lazy loaded\" meaning the ASDF block contents will not be read until the array is accessed.\n",
    "\n",
    "Note that this default behavior can be disabled by setting the keyword argument `lazy_load` to `False` in `asdf.open`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a7533b5c-e955-4c93-95cd-4a4729c28b5f",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "\n",
       "### ASDF Tree\n",
       "\n",
       "```yaml\n",
       "#ASDF 1.0.0\n",
       "#ASDF_STANDARD 1.5.0\n",
       "%YAML 1.1\n",
       "%TAG ! tag:stsci.edu:asdf/\n",
       "--- !core/asdf-1.1.0\n",
       "asdf_library: !core/software-1.0.0 {author: The ASDF Developers, homepage: 'http://github.com/asdf-format/asdf',\n",
       "  name: asdf, version: 3.0.0.dev307+gb0c9a50}\n",
       "history:\n",
       "  extensions:\n",
       "  - !core/extension_metadata-1.0.0\n",
       "    extension_class: asdf.extension._manifest.ManifestExtension\n",
       "    extension_uri: asdf://asdf-format.org/core/extensions/core-1.5.0\n",
       "    software: !core/software-1.0.0 {name: asdf, version: 3.0.0.dev307+gb0c9a50}\n",
       "  - !core/extension_metadata-1.0.0\n",
       "    extension_class: asdf.extension.BuiltinExtension\n",
       "    software: !core/software-1.0.0 {name: asdf, version: 3.0.0.dev307+gb0c9a50}\n",
       "array0: !core/ndarray-1.0.0\n",
       "  source: 0\n",
       "  datatype: int64\n",
       "  byteorder: little\n",
       "  shape: [42]\n",
       "array1: !core/ndarray-1.0.0\n",
       "  source: 1\n",
       "  datatype: int64\n",
       "  byteorder: little\n",
       "  shape: [7]\n",
       "array2: !core/ndarray-1.0.0\n",
       "  source: 2\n",
       "  datatype: int64\n",
       "  byteorder: little\n",
       "  shape: [26]\n",
       "...\n",
       "\n",
       "```\n",
       "\n",
       "### ASDF Blocks [N = 3]\n",
       "\tBlock0: b'\\x000\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00'...[386 B]\n",
       "\tBlock1: b'\\x000\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00'...[106 B]\n",
       "\tBlock2: b'\\x000\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00'...[258 B]\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": []
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# construct an AsdfFile object with 3 arrays\n",
    "example_af = asdf.AsdfFile({'array0': np.arange(42), 'array1': np.arange(7), 'array2': np.arange(26)})\n",
    "\n",
    "# write out our example file\n",
    "example_fn = tmp_path / 'example.asdf'\n",
    "example_af.write_to(example_fn)\n",
    "\n",
    "# and a reminder of what it contains\n",
    "example_af"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c086ca19-61bd-474b-9b3b-6290dc459358",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prior to access the array is unloaded\n",
      "\tArray = <array (unloaded) shape: [7] dtype: int64>\n",
      "\n",
      "Accessing the first element(0) of the array causes asdf to load the data\n",
      "\tArray = [0 1 2 3 4 5 6]\n"
     ]
    }
   ],
   "source": [
    "# open with lazy loading (enabled by default)\n",
    "with asdf.open(example_fn) as af:\n",
    "    print(\"Prior to access the array is unloaded\")\n",
    "    print(f\"\\tArray = {af['array1']}\")\n",
    "\n",
    "    first_element = af['array1'][0]\n",
    "    print(f\"\\nAccessing the first element({first_element}) of the array causes asdf to load the data\")\n",
    "    print(f\"\\tArray = {af['array1']}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1c8abfcf-e1a1-4059-8015-7d13def6d80e",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "With lazy_load=False the array data is immediately loaded\n",
      "\tArray = [0 1 2 3 4 5 6]\n"
     ]
    }
   ],
   "source": [
    "# open with lazy loading disabled\n",
    "with asdf.open(example_fn, lazy_load=False) as af:\n",
    "    print(\"With lazy_load=False the array data is immediately loaded\")\n",
    "    print(f\"\\tArray = {af['array1']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6042e232-435d-4617-ba43-21d44f83e640",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "## Memory mapping\n",
    "For applications where only a small portion of the array data is needed, loading the entire ASDF block when any portion of the array data is requested would be inefficient.\n",
    "\n",
    "For local files, asdf will memory map arrays allowing the operating system to read only the portion of the array that is requested.\n",
    "\n",
    "Note that array memory mapping can be disabled by setting the keyword argument `memmap` to `False` in `asdf.open`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b9559bb9-8432-48e3-8340-b98be480775c",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "With default memmap=True\n",
      "\tarray 0 base array type = <class 'numpy.memmap'>\n"
     ]
    }
   ],
   "source": [
    "# open our example file with memory mapping (enabled by default)\n",
    "with asdf.open(example_fn) as af:\n",
    "    print(\"With default memmap=True\") \n",
    "    # the array data has a base array that is of type numpy.memmap\n",
    "    print(f\"\\tarray 0 base array type = {type(af['array0'].base)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "05fb0a65-ab9f-45be-8a54-55f9de8f6093",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "With memmap=False\n",
      "\tarray 0 base array type = <class 'numpy.ndarray'>\n"
     ]
    }
   ],
   "source": [
    "# open our example file with memory mapping disabled\n",
    "with asdf.open(example_fn, memmap=False) as af:\n",
    "    print(\"With memmap=False\")\n",
    "    print(f\"\\tarray 0 base array type = {type(af['array0'].base)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "155c4eae-aa55-4068-a4f2-b98f49c549f4",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "## Chunked array storage\n",
    "\n",
    "Some applications benefit from breaking up large arrays into smaller chunks.\n",
    "\n",
    "Each chunk can be stored in a different ASDF block which can be loaded only when the data for that chunk is required.\n",
    "\n",
    "This can be useful when:\n",
    "  - memory mapping isn't possible (such as network file access) \n",
    "  - memory mapping is inefficient (macOS has poor memory mapping performance for arrays approaching 800 MB)\n",
    "  - paying per byte read (AWS S3)\n",
    "\n",
    "To support chunked storage, asdf uses [Zarr](https://zarr.readthedocs.io/en/stable/), a widely used, open source chunked array storage format."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee699b08-b399-4f4c-89fe-76340ac85a33",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "To compare the pros and cons of chunked array storage let's generate an AsdfFile containing a large non-chunked array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "fe9af82c-e3ef-4cfd-8472-5775656ad21f",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "large_array = np.ones((3, 4096, 4096), dtype='uint16')\n",
    "large_example_af = asdf.AsdfFile({'arr': large_array})\n",
    "\n",
    "# write it out to an ASDF file\n",
    "large_example_fn = tmp_path / 'large_example.asdf'\n",
    "large_example_af.write_to(large_example_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "8567de18-a760-46ad-95e1-e215b6fe64f4",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "\n",
       "### ASDF Tree\n",
       "\n",
       "```yaml\n",
       "#ASDF 1.0.0\n",
       "#ASDF_STANDARD 1.5.0\n",
       "%YAML 1.1\n",
       "%TAG ! tag:stsci.edu:asdf/\n",
       "--- !core/asdf-1.1.0\n",
       "asdf_library: !core/software-1.0.0 {author: The ASDF Developers, homepage: 'http://github.com/asdf-format/asdf',\n",
       "  name: asdf, version: 3.0.0.dev307+gb0c9a50}\n",
       "history:\n",
       "  extensions:\n",
       "  - !core/extension_metadata-1.0.0\n",
       "    extension_class: asdf.extension._manifest.ManifestExtension\n",
       "    extension_uri: asdf://asdf-format.org/core/extensions/core-1.5.0\n",
       "    software: !core/software-1.0.0 {name: asdf, version: 3.0.0.dev307+gb0c9a50}\n",
       "  - !core/extension_metadata-1.0.0\n",
       "    extension_class: asdf.extension.BuiltinExtension\n",
       "    software: !core/software-1.0.0 {name: asdf, version: 3.0.0.dev307+gb0c9a50}\n",
       "arr: !core/ndarray-1.0.0\n",
       "  source: 0\n",
       "  datatype: uint16\n",
       "  byteorder: little\n",
       "  shape: [3, 4096, 4096]\n",
       "...\n",
       "\n",
       "```\n",
       "\n",
       "### ASDF Blocks [N = 1]\n",
       "\tBlock0: b'\\x000\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00'...[96 MB]\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": []
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# our large array is stored in one ASDF block\n",
    "large_example_af"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5beba9c-ac4f-4e13-8a3d-f02be4e7d9ca",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "Storing this entire large array in one block means that any array access will result in reading the corresponding ASDF block."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "82b5aaf0-2006-4409-80c0-1cbd420c587f",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accessing one array element resulted in reading 96 MB\n"
     ]
    }
   ],
   "source": [
    "# use a helper function to monitor bytes read from the file\n",
    "with asdf_open_tracking_reads(large_example_fn) as (af, mon):\n",
    "    af['arr'][0, 0, 0]\n",
    "    print(f\"Accessing one array element resulted in reading {n_bytes_to_human(mon.n_bytes)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d1448f3-8d2a-40ff-a5aa-0cb2d06063c3",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "If we \"chunk\" this large array using Zarr we can store each chunk in a different ASDF block.\n",
    "\n",
    "When an array element is read from the Zarr array, asdf will only need to load the ASDF block that contains the corresponding chunk.\n",
    "\n",
    "This can drastically reduce the amount of data read from disk."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "338f24fd-e36d-48b6-8454-0a2363ff5631",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import zarr\n",
    "\n",
    "# make a Zarr array from the large NumPy array\n",
    "# Zarr, unless told otherwise, will pick a default chunk shape based on some\n",
    "# simple heuristics that aim to produce 1MB uncompressed chunks\n",
    "# we are setting the compressor to None to disable compression\n",
    "za = zarr.array(large_example_af['arr'], compressor=None)\n",
    "\n",
    "# add this Zarr array to an AsdfFile\n",
    "chunked_example_af = asdf.AsdfFile({'arr': za})\n",
    "\n",
    "# write it to an ASDF file for later use\n",
    "chunked_example_fn = tmp_path / 'chunked_example.asdf'\n",
    "chunked_example_af.write_to(chunked_example_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "98579d91-c7e6-4e82-b066-75e39c5a6a23",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "\n",
       "### ASDF Tree\n",
       "\n",
       "```yaml\n",
       "#ASDF 1.0.0\n",
       "#ASDF_STANDARD 1.5.0\n",
       "%YAML 1.1\n",
       "%TAG ! tag:stsci.edu:asdf/\n",
       "--- !core/asdf-1.1.0\n",
       "asdf_library: !core/software-1.0.0 {author: The ASDF Developers, homepage: 'http://github.com/asdf-format/asdf',\n",
       "  name: asdf, version: 3.0.0.dev307+gb0c9a50}\n",
       "history:\n",
       "  extensions:\n",
       "  - !core/extension_metadata-1.0.0\n",
       "    extension_class: asdf.extension._manifest.ManifestExtension\n",
       "    extension_uri: asdf://asdf-format.org/core/extensions/core-1.5.0\n",
       "    software: !core/software-1.0.0 {name: asdf, version: 3.0.0.dev307+gb0c9a50}\n",
       "  - !core/extension_metadata-1.0.0\n",
       "    extension_class: asdf_zarr.extensions.ZarrExtension\n",
       "    extension_uri: asdf://asdf-format.org/zarr/tags/zarr-1.0.0\n",
       "    software: !core/software-1.0.0 {name: asdf-zarr, version: 0.0.1}\n",
       "arr: !<asdf://asdf-format.org/zarr/tags/zarr-1.0.0>\n",
       "  .zarray:\n",
       "    chunks: [1, 512, 1024]\n",
       "    compressor: null\n",
       "    dtype: <u2\n",
       "    fill_value: 0\n",
       "    filters: null\n",
       "    order: C\n",
       "    shape: [3, 4096, 4096]\n",
       "    zarr_format: 2\n",
       "  chunk_block_map: 96\n",
       "...\n",
       "\n",
       "```\n",
       "\n",
       "### ASDF Blocks [N = 97]\n",
       "\tBlock0: b'\\x000\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00'...[1 MB]\n",
       "\tBlock1: b'\\x000\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00'...[1 MB]\n",
       "\tBlock2: b'\\x000\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00'...[1 MB]\n",
       "\tBlock3: b'\\x000\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00'...[1 MB]\n",
       "\tBlock4: b'\\x000\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00'...[1 MB]\n",
       "\tBlock5: b'\\x000\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00'...[1 MB]\n",
       "\tBlock6: b'\\x000\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00'...[1 MB]\n",
       "\tBlock7: b'\\x000\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00'...[1 MB]\n",
       "\tBlock8: b'\\x000\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00'...[1 MB]\n",
       "\tBlock9: b'\\x000\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00'...[1 MB]\n",
       "... skipping display of 87 blocks\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": []
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# print out the ASDF file produced with this Zarr array\n",
    "chunked_example_af"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "076cc006-3788-4ae6-a396-163d960d3feb",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "Zarr will default to chunks of ~1 MB.\n",
    "\n",
    "When indexed or sliced, the corresponding blocks will be read, combined and returned as a standard NumPy array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "0c3eaec9-58e9-4203-a248-a9c8761b99f9",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accessing one array element resulted in reading 1 MB\n",
      "\treturn type = <class 'numpy.uint16'>\n"
     ]
    }
   ],
   "source": [
    "with asdf_open_tracking_reads(chunked_example_fn) as (af, mon):\n",
    "    v = af['arr'][0, 0, 0]\n",
    "    print(f\"Accessing one array element resulted in reading {n_bytes_to_human(mon.n_bytes)}\")\n",
    "    print(f\"\\treturn type = {type(v)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df416519-290e-4c56-b154-eba70a1fff4c",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "## Applications for chunked array storage"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "270346bf-83b1-40c3-bde5-d24ed972d2e5",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "### Mosaic cutouts\n",
    "\n",
    "Taking \"cutouts\", accessing small sub-arrays, of a large mosaic is an obvious application where chunked array storage can be beneficial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "680f1984-81e6-4ab4-8816-d3d2572a4578",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# generate a large 2D montage (~380 MB)\n",
    "mosaic = zarr.array(np.zeros((10000, 10000), dtype='float32'), compressor=None)\n",
    "mosaic_af = asdf.AsdfFile({'mosaic': mosaic})\n",
    "\n",
    "# write it out to disk\n",
    "mosaic_example_fn = tmp_path / 'mosaic.asdf'\n",
    "mosaic_af.write_to(mosaic_example_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8789d83a-5c80-46b8-9564-e68d38aa2a03",
   "metadata": {},
   "source": [
    "Opening this large file and taking a small cutout involves very little data actually being read from disk."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "ee7f38ad-e085-4bf8-a5e6-b4c97b56ebf5",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accessing one cutout resulted in reading 2 MB\n"
     ]
    }
   ],
   "source": [
    "with asdf_open_tracking_reads(mosaic_example_fn) as (af, mon):\n",
    "    v = af['mosaic'][30:60, 90:120]\n",
    "    print(f\"Accessing one cutout resulted in reading {n_bytes_to_human(mon.n_bytes)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2543d0ff-ace3-4b6f-8f39-e815408f4bb9",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "### Spectral cubes\n",
    "\n",
    "High dimensional data (like spectral cubes) can also benefit from chunked array storage.\n",
    "\n",
    "For this application, selection of chunk shape can dramatically impact the performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "2d6ef452-c7ca-4f27-bcac-ee5fe4a36adc",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cube consumes 10 MB\n"
     ]
    }
   ],
   "source": [
    "# generate a fake spectral cube with 2 spatial and one spectral dimension\n",
    "shape = (32, 32, 2560)\n",
    "cube_array = np.zeros(shape, dtype='float32')\n",
    "print(f\"Cube consumes {n_bytes_to_human(cube_array.nbytes)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "049b6f3e-34c1-453e-9158-9c8ba2f8df39",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# chunk the spatial dimensions, preserve the spectral dimension\n",
    "chunks = (8, 8, 2560)  # array shape = 32, 32, 2560\n",
    "chunked_spatial_fn = tmp_path / \"chunked_spatial.asdf\"\n",
    "cube = zarr.array(cube_array, compressor=None, chunks=chunks)\n",
    "\n",
    "# write to an ASDF file\n",
    "cube_af = asdf.AsdfFile({'cube': cube})\n",
    "cube_af.write_to(chunked_spatial_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "240370ec-a21b-4d2b-ac02-128d85cd9144",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "read 10 MB\n"
     ]
    }
   ],
   "source": [
    "with asdf_open_tracking_reads(chunked_spatial_fn) as (af, mon):\n",
    "    # accessing all pixels at one wavelength we will end up loading all chunks\n",
    "    v = af['cube'][:, :, 0]\n",
    "    print(f\"read {n_bytes_to_human(mon.n_bytes)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "b5289475-447b-403a-ba91-2c694a8937b0",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "read 671 KB\n"
     ]
    }
   ],
   "source": [
    "with asdf_open_tracking_reads(chunked_spatial_fn) as (af, mon):\n",
    "    # accessing all wavelengths for one pixel results in loading a single chunk\n",
    "    v = af['cube'][0, 0, :]\n",
    "    print(f\"read {n_bytes_to_human(mon.n_bytes)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "0ebb172f-2124-4aed-82ed-f1b7122cfc1e",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# chunk the spectral dimension, preserve the spatial dimensions\n",
    "chunks = (32, 32, 256)  # array shape = 32, 32, 2560\n",
    "chunked_spectral_fn = tmp_path / \"chunked_spectral.asdf\"\n",
    "cube = zarr.array(cube_array, compressor=None, chunks=chunks)\n",
    "\n",
    "# write to an ASDF file\n",
    "cube_af = asdf.AsdfFile({'cube': cube})\n",
    "cube_af.write_to(chunked_spectral_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "8084e93e-5833-474c-b7c0-bfeecd82e5a6",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "read 1 MB\n"
     ]
    }
   ],
   "source": [
    "with asdf_open_tracking_reads(chunked_spectral_fn) as (af, mon):\n",
    "    # accessing all pixels for one wavelength results in loading one chunk\n",
    "    v = af['cube'][:, :, 0]\n",
    "    print(f\"read {n_bytes_to_human(mon.n_bytes)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "be232b9c-ed16-4cfd-aaae-8dffd2531512",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "read 10 MB\n"
     ]
    }
   ],
   "source": [
    "with asdf_open_tracking_reads(chunked_spectral_fn) as (af, mon):\n",
    "    # accessing all wavelengths for one pixel results in loading all chunks\n",
    "    v = af['cube'][0, 0, :]\n",
    "    print(f\"read {n_bytes_to_human(mon.n_bytes)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "cc5a5b83-fb02-45b2-ba8e-d0c81d0b87b2",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cube consumes 250 MB\n"
     ]
    }
   ],
   "source": [
    "# generate a larger spectral cube\n",
    "shape = (160, 160, 2560)\n",
    "large_cube_array = np.zeros(shape, dtype='float32')\n",
    "print(f\"Cube consumes {n_bytes_to_human(large_cube_array.nbytes)}\")\n",
    "\n",
    "# chunk the spectral dimension, preserve the spatial dimensions\n",
    "chunks = (32, 32, 256)  # array shape = 160, 160, 2560\n",
    "chunked_cube_fn = tmp_path / \"chunked_cube.asdf\"\n",
    "large_cube = zarr.array(large_cube_array, compressor=None, chunks=chunks)\n",
    "\n",
    "# write to an ASDF file\n",
    "large_cube_af = asdf.AsdfFile({'cube': large_cube})\n",
    "large_cube_af.write_to(chunked_cube_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "7a4cf6fc-e228-4c32-a9b9-300a81cebafd",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "read 25 MB\n"
     ]
    }
   ],
   "source": [
    "with asdf_open_tracking_reads(chunked_cube_fn) as (af, mon):\n",
    "    # accessing all pixels for one wavelength results in loading one chunk\n",
    "    v = af['cube'][:, :, 0]\n",
    "    print(f\"read {n_bytes_to_human(mon.n_bytes)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "2ec39a94-e972-4d45-b00e-fc40005c56cb",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "read 10 MB\n"
     ]
    }
   ],
   "source": [
    "with asdf_open_tracking_reads(chunked_cube_fn) as (af, mon):\n",
    "    # accessing all wavelengths for one pixel results in loading all chunks\n",
    "    v = af['cube'][0, 0, :]\n",
    "    print(f\"read {n_bytes_to_human(mon.n_bytes)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08acfebe-d63f-4e20-b6e2-e255fb256016",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "The selection of the chunk size and shape is critical.\n",
    "\n",
    "Chunking across the wrong dimensions or selecting a too small or too large chunk size can result in extra reads and poor performance.\n",
    "\n",
    "Think about how the data will be accessed.\n",
    "\n",
    "Treating the array as if it exists in RAM can result in poor performance and high storage access costs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3046489d-fcb8-4b02-8063-1b9dc06197d7",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": []
   },
   "source": [
    "## Questions?"
   ]
  }
 ],
 "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}