Create custom backends

_toc.yml

@@ -87,6 +87,7 @@ parts:
       - file: advanced/backends/backends.md
         sections:
           - file: advanced/backends/intro-backends.ipynb
+          - file: advanced/backends/create_custom_backend.ipynb
           - file: advanced/backends/1.Backend_without_Lazy_Loading.ipynb
           - file: advanced/backends/2.Backend_with_Lazy_Loading.ipynb
           - file: advanced/backends/rasterio_backend.ipynb

advanced/backends/create_custom_backend.ipynb

@@ -0,0 +1,315 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "0",
+   "metadata": {},
+   "source": [
+    "# Create your own Xarray Backend\n",
+    "In this lesson, we will learn about and how to create your own custom xarray backend\n",
+    "\n",
+    ":::{admonition} Learning Goals\n",
+    "- Learn about xarray's support for custom backends\n",
+    "- Learn how to use a custom `imageio` backend to open and manipulate GIFs\n",
+    "- Learn how to extend the `imagegio` backend to write out a GIF\n",
+    ":::\n",
+    "\n",
+    "## Why should you create your own Xarray backend?\n",
+    "\n",
+    "- Allows you to use xarray's interface\n",
+    "    - Attribute-like syntax, dict-like syntax and label based indexing\n",
+    "- You don’t need to integrate any code in Xarray\n",
+    "- Easy and fast!\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1",
+   "metadata": {},
+   "source": [
+    "## Setting up the BackendEntrypoint\n",
+    "\n",
+    "To set up a `BackendEntrypoint` we can implement a subclass of `BackendEntrypoint` and expose the `open_dataset` method to it. For this tutorial we have the `ImageIOBackend` already defined but we will extend the functionally to write GIFS."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "image_path = \"io.gif\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3",
+   "metadata": {},
+   "source": [
+    "We can write a simple image reader function that we can then plug into our `MyBackendEntrypoint` class. For this example we are going to use `imageio` an image reader and writer library. With `imageio` you can read an image file with `iio.imopen`\n",
+    "\n",
+    ":::{note}\n",
+    "The `ImageIOBackend` also defines a `ImageIOBackendArray` with basic indexing.\n",
+    ":"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import imageio as iio\n",
+    "from xarray.backends import BackendEntrypoint\n",
+    "\n",
+    "\n",
+    "def imageio_open(\n",
+    "    filename_or_obj,\n",
+    "):\n",
+    "    img = iio.imopen(filename_or_obj, io_mode=\"r\")\n",
+    "    return img.read()\n",
+    "\n",
+    "\n",
+    "imageio_open(image_path)\n",
+    "\n",
+    "\n",
+    "class MyImageReader(BackendEntrypoint):\n",
+    "    def open_dataset(\n",
+    "        self,\n",
+    "        filename_or_obj,\n",
+    "        *,\n",
+    "        drop_variables=None,\n",
+    "    ):\n",
+    "        return imageio_open(filename_or_obj)\n",
+    "\n",
+    "\n",
+    "type(MyImageReader().open_dataset(image_path))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5",
+   "metadata": {},
+   "source": [
+    "## Reading image data\n",
+    "\n",
+    "Lets use our `ImageIOBackend` to open a GIF.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import xarray as xr\n",
+    "from imageio_ import ImageIOBackend\n",
+    "\n",
+    "gif_ds = xr.open_dataset(image_path, engine=ImageIOBackend)\n",
+    "gif_ds"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7",
+   "metadata": {},
+   "source": [
+    "### Examining our image dataset\n",
+    "\n",
+    "Since our image is a `Dataset` object we can use xarray's interface for `Dataset` objects.\n",
+    "\n",
+    "Let's try listing all of the variables, dimensions and selecting data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "gif_ds.variables"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9",
+   "metadata": {},
+   "source": [
+    "We can list our dimensions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "10",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "gif_ds.dims"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "11",
+   "metadata": {},
+   "source": [
+    "Let's try getting our `DataArray`"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "12",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "gif_ds[\"data\"]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "13",
+   "metadata": {},
+   "source": [
+    "## GIF metadata\n",
+    "We can examine, update and add metadata to our `Dataset` object. \n",
+    "\n",
+    "We can examine the attributes in our GIF \"data\" variable with the `.attrs` method"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "14",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "gif_ds.data.attrs"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "15",
+   "metadata": {},
+   "source": [
+    "## Exercise"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "16",
+   "metadata": {},
+   "source": [
+    "::::{admonition} Exercise\n",
+    ":class: tip\n",
+    "\n",
+    "Can you add a new attribute to our GIF \"data\" variable called \"fps\". This is the frames per second we can write our GIF to. You can set it to any value.\n",
+    "\n",
+    ":::{admonition} Solution\n",
+    ":class: dropdown\n",
+    "\n",
+    "```python\n",
+    "gif_ds.data.attrs['fps'] = 100\n",
+    "```\n",
+    ":::\n",
+    "::::"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "17",
+   "metadata": {},
+   "source": [
+    "### GIF writer\n",
+    "We can extend our backend with an GIF writer. Here we use `matplotlib`'s animation functions and `PillowWriter`"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "18",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.animation as animation\n",
+    "import matplotlib.pyplot as plt\n",
+    "from imageio_ import ImageIOBackend\n",
+    "from matplotlib.animation import PillowWriter\n",
+    "\n",
+    "\n",
+    "class SimpleImageWriter(ImageIOBackend):\n",
+    "    def to_giff(\n",
+    "        self,\n",
+    "        dataset,\n",
+    "        variable,\n",
+    "        time_dim,\n",
+    "        out_filename,\n",
+    "        **kwargs,\n",
+    "    ):\n",
+    "        fig, ax = plt.subplots()\n",
+    "\n",
+    "        frames = []\n",
+    "        variable_da = dataset[variable]\n",
+    "\n",
+    "        for time in variable_da[time_dim]:\n",
+    "            variable_da = variable_da.transpose(\"time\", \"width\", \"height\", \"color\")\n",
+    "            to_plot = ax.pcolormesh(\n",
+    "                variable_da.height,\n",
+    "                variable_da.width,\n",
+    "                variable_da.sel(time=time),\n",
+    "                animated=True,\n",
+    "                shading=\"auto\",\n",
+    "                **kwargs,\n",
+    "            )\n",
+    "\n",
+    "            frames.append([to_plot])\n",
+    "\n",
+    "        try:\n",
+    "            writer = PillowWriter(fps=variable_da.attrs[\"fps\"], **kwargs)\n",
+    "        except KeyError:\n",
+    "            writer = PillowWriter(fps=50, **kwargs)\n",
+    "\n",
+    "        ani = animation.ArtistAnimation(fig, frames, blit=True, repeat=True)\n",
+    "        ani.save(filename=out_filename, writer=writer)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "19",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "img_writer = SimpleImageWriter()\n",
+    "img_writer.to_giff(gif_ds, \"data\", \"time\", \"io_writer.gif\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "20",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}