python/notebooks/mandelbrot.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup\n",
    "\n",
    "Before start, please make sure you have all the dependencies installed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install --quiet vizarr zarr numba"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import zarr\n",
    "import numcodecs\n",
    "import numba\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Zoomable Mandelbrot Set\n",
    "\n",
    "This notebook a contains `vizarr` example visualizing a generic multiscale zarr. The first cell contains code to create the underlying generatvie zarr store. It dynamically creates \"chunks\" at different zoom levels and associated array metadata."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def create_meta_store(levels: int, tilesize: int, compressor, dtype: str):\n",
    "    store = dict()\n",
    "    zarr.storage.init_group(store)\n",
    "    \n",
    "    datasets = [{\"path\": str(i)} for i in range(levels)]\n",
    "    root_attrs = {\"multiscales\": [{\"datasets\": datasets, \"version\": \"0.1\"}]}\n",
    "    store[\".zattrs\"] = zarr.util.json_dumps(root_attrs)\n",
    "    \n",
    "    base_width = tilesize * 2 ** levels\n",
    "    for level in range(levels):\n",
    "        width = int(base_width / 2 ** level)\n",
    "        zarr.storage.init_array(\n",
    "            store,\n",
    "            path=str(level),\n",
    "            shape=(width, width),\n",
    "            chunks=(tilesize, tilesize),\n",
    "            dtype=dtype,\n",
    "            compressor=compressor,\n",
    "        )\n",
    "    return store\n",
    "\n",
    "\n",
    "@numba.njit\n",
    "def mandelbrot(out, from_x, from_y, to_x, to_y, grid_size, maxiter):\n",
    "    step_x = (to_x - from_x) / grid_size\n",
    "    step_y = (to_y - from_y) / grid_size\n",
    "    creal = from_x\n",
    "    cimag = from_y\n",
    "    for i in range(grid_size):\n",
    "        cimag = from_y\n",
    "        for j in range(grid_size):\n",
    "            nreal = real = imag = n = 0\n",
    "            for _ in range(maxiter):\n",
    "                nreal = real * real - imag * imag + creal\n",
    "                imag = 2 * real * imag + cimag\n",
    "                real = nreal\n",
    "                if real * real + imag * imag > 4.0:\n",
    "                    break\n",
    "                n += 1\n",
    "            out[j * grid_size + i] = n\n",
    "            cimag += step_y\n",
    "        creal += step_x\n",
    "    return out\n",
    "\n",
    "\n",
    "@numba.njit\n",
    "def tile_bounds(level, x, y, max_level, min_coord=-2.5, max_coord=2.5):\n",
    "    max_width = max_coord - min_coord\n",
    "    tile_width = max_width / 2 ** (max_level - level)\n",
    "    from_x = min_coord + x * tile_width\n",
    "    to_x = min_coord + (x + 1) * tile_width\n",
    "\n",
    "    from_y = min_coord + y * tile_width\n",
    "    to_y = min_coord + (y + 1) * tile_width\n",
    "\n",
    "    return from_x, from_y, to_x, to_y\n",
    "\n",
    "\n",
    "class MandlebrotStore(zarr.storage.BaseStore):\n",
    "    \n",
    "    def __init__(self, levels, tilesize, maxiter=255, compressor=None):\n",
    "        self.levels = levels\n",
    "        self.tilesize = tilesize\n",
    "        self.compressor = compressor\n",
    "        self.dtype = np.dtype(np.uint8 if maxiter < 256 else np.uint16)\n",
    "        self.maxiter = maxiter\n",
    "        self._store = create_meta_store(levels, tilesize, compressor, self.dtype)\n",
    "\n",
    "    def __getitem__(self, key):\n",
    "        if key in self._store:\n",
    "            return self._store[key]\n",
    "\n",
    "        try:\n",
    "            # Try parsing pyramidal coords\n",
    "            level, chunk_key = key.split(\"/\")\n",
    "            level = int(level)\n",
    "            y, x = map(int, chunk_key.split(\".\"))\n",
    "        except:\n",
    "            raise KeyError\n",
    "\n",
    "        from_x, from_y, to_x, to_y = tile_bounds(level, x, y, self.levels)\n",
    "        out = np.zeros(self.tilesize * self.tilesize, dtype=self.dtype)\n",
    "        tile = mandelbrot(out, from_x, from_y, to_x, to_y, self.tilesize, self.maxiter)\n",
    "        tile = tile.reshape(self.tilesize, self.tilesize)\n",
    "\n",
    "        if self.compressor:\n",
    "            return self.compressor.encode(tile)\n",
    "\n",
    "        return tile.tobytes()\n",
    "\n",
    "    def __iter__(self):\n",
    "        return iter(self._store)\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self._store)\n",
    "\n",
    "    def __delitem__(self, key):\n",
    "        raise NotImplementedError(\"read-only store\")\n",
    "        \n",
    "    def __setitem__(self, key):\n",
    "        raise NotImplementedError(\"read-only store\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Running vizarr\n",
    "\n",
    "Initalize the store implemented above, and open as a `zarr.Group` for vizarr. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the store\n",
    "store = MandlebrotStore(levels=50, tilesize=512, compressor=numcodecs.Blosc())\n",
    "# Wrap in a cache so that tiles don't need to be computed as often\n",
    "store = zarr.LRUStoreCache(store, max_size=1e9)\n",
    "\n",
    "# This store implements the 'multiscales' zarr specfiication which is recognized by vizarr\n",
    "grp = zarr.open(store, mode=\"r\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import vizarr\n",
    "\n",
    "viewer = vizarr.Viewer()\n",
    "viewer.add_image(source=grp, name=\"mandelbrot\")\n",
    "viewer"
   ]
  }
 ],
 "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}