Zarr_and_rechunker.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Install fsspec and gdrivefs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ! pip install fsspec==0.8.4\n",
    "# ! pip install 'git+git://github.com/tjcrone/gdrivefs.git@540a55fae06cbed4f7d9f8d540ac29fcb640c38a'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Drop variables I do not need"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ds_dropped = ds.drop({'AG1', 'AL', 'AL1', 'COD', 'CP', 'CWP', 'CZ', 'FRA', 'FRV', 'G11', 'G21', 'IWP', 'LHF', 'LWD', 'LWU', 'ME', 'QI', 'QQ', 'QQP', 'QQZ', 'QR', 'QS', 'QW', 'RF', 'RH', 'RU2', 'RZ', 'SAL', 'SF', 'SHF', 'SHSN0', 'SHSN3', 'SLO', 'SOL', 'SP', 'SRF', 'ST', 'ST2', 'SU', 'SW', 'SWD', 'TI1', 'TT', 'TTMAX', 'TTMIN', 'TTP', 'TTZ', 'UU', 'UUP', 'UUZ', 'UV', 'UVZ', 'VEG', 'VV', 'VVP', 'VVZ', 'WA1', 'WVP', 'ZN4', 'ZN5', 'ZN6', 'ZZ', 'ZZP'})\n",
    "# ds_dropped"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Rechunking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ds_dropped = ds_dropped.chunk({'TIME': 1})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create mapper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# mapper = gdfs.get_mapper('ICE.2000-2010.zarr')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Save data to Zarr file on Google Drive"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ds_dropped.to_zarr(mapper, consolidated=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Open data from Zarr file from Google Drive"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ds = xr.open_zarr(mapper, decode_coords=False) #, mask_and_scale=False) #, consolidated=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test opening tutorial data set, saving in Zarr and opening in Zarr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# gdfid_tutorial = '162zJ0ZyHaq9dkU4h3MdiZu1Vma8rPxkQ'\n",
    "# gdfs_tutorial = gdrivefs.GoogleDriveFileSystem(root_file_id = gdfid_tutorial, token = 'cache')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# mapper_tutorial = gdfs_tutorial.get_mapper('air_temperature.zarr')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ds_tutorial = xr.tutorial.open_dataset(\"air_temperature\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ds_tutorial.to_zarr(mapper_tutorial)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ds_tut_test = xr.open_zarr(mapper_tutorial)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ds_tut_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ds_tut_test.air[0,0,0].load()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test rechunker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "xr.set_options(display_style='text')\n",
    "import zarr\n",
    "import dask.array as dsa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre>&lt;xarray.Dataset&gt;\n",
       "Dimensions:  (lat: 25, lon: 53, time: 2920)\n",
       "Coordinates:\n",
       "  * lat      (lat) float32 75.0 72.5 70.0 67.5 65.0 ... 25.0 22.5 20.0 17.5 15.0\n",
       "  * lon      (lon) float32 200.0 202.5 205.0 207.5 ... 322.5 325.0 327.5 330.0\n",
       "  * time     (time) datetime64[ns] 2013-01-01 ... 2014-12-31T18:00:00\n",
       "Data variables:\n",
       "    air      (time, lat, lon) float32 dask.array&lt;chunksize=(100, 25, 53), meta=np.ndarray&gt;\n",
       "Attributes:\n",
       "    Conventions:  COARDS\n",
       "    title:        4x daily NMC reanalysis (1948)\n",
       "    description:  Data is from NMC initialized reanalysis\\n(4x/day).  These a...\n",
       "    platform:     Model\n",
       "    references:   http://www.esrl.noaa.gov/psd/data/gridded/data.ncep.reanaly...</pre>"
      ],
      "text/plain": [
       "<xarray.Dataset>\n",
       "Dimensions:  (lat: 25, lon: 53, time: 2920)\n",
       "Coordinates:\n",
       "  * lat      (lat) float32 75.0 72.5 70.0 67.5 65.0 ... 25.0 22.5 20.0 17.5 15.0\n",
       "  * lon      (lon) float32 200.0 202.5 205.0 207.5 ... 322.5 325.0 327.5 330.0\n",
       "  * time     (time) datetime64[ns] 2013-01-01 ... 2014-12-31T18:00:00\n",
       "Data variables:\n",
       "    air      (time, lat, lon) float32 dask.array<chunksize=(100, 25, 53), meta=np.ndarray>\n",
       "Attributes:\n",
       "    Conventions:  COARDS\n",
       "    title:        4x daily NMC reanalysis (1948)\n",
       "    description:  Data is from NMC initialized reanalysis\\n(4x/day).  These a...\n",
       "    platform:     Model\n",
       "    references:   http://www.esrl.noaa.gov/psd/data/gridded/data.ncep.reanaly..."
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds = xr.tutorial.open_dataset(\"air_temperature\")\n",
    "# create initial chunk structure\n",
    "ds = ds.chunk({'time': 100})\n",
    "ds.air.encoding = {} # helps when writing to zarr\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table>\n",
       "<tr>\n",
       "<td>\n",
       "<table>\n",
       "  <thead>\n",
       "    <tr><td> </td><th> Array </th><th> Chunk </th></tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr><th> Bytes </th><td> 15.48 MB </td> <td> 530.00 kB </td></tr>\n",
       "    <tr><th> Shape </th><td> (2920, 25, 53) </td> <td> (100, 25, 53) </td></tr>\n",
       "    <tr><th> Count </th><td> 31 Tasks </td><td> 30 Chunks </td></tr>\n",
       "    <tr><th> Type </th><td> float32 </td><td> numpy.ndarray </td></tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</td>\n",
       "<td>\n",
       "<svg width=\"159\" height=\"146\" style=\"stroke:rgb(0,0,0);stroke-width:1\" >\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"10\" y1=\"0\" x2=\"80\" y2=\"70\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"10\" y1=\"25\" x2=\"80\" y2=\"96\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Vertical lines -->\n",
       "  <line x1=\"10\" y1=\"0\" x2=\"10\" y2=\"25\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"12\" y1=\"2\" x2=\"12\" y2=\"27\" />\n",
       "  <line x1=\"14\" y1=\"4\" x2=\"14\" y2=\"30\" />\n",
       "  <line x1=\"17\" y1=\"7\" x2=\"17\" y2=\"32\" />\n",
       "  <line x1=\"19\" y1=\"9\" x2=\"19\" y2=\"35\" />\n",
       "  <line x1=\"22\" y1=\"12\" x2=\"22\" y2=\"37\" />\n",
       "  <line x1=\"24\" y1=\"14\" x2=\"24\" y2=\"39\" />\n",
       "  <line x1=\"26\" y1=\"16\" x2=\"26\" y2=\"42\" />\n",
       "  <line x1=\"29\" y1=\"19\" x2=\"29\" y2=\"44\" />\n",
       "  <line x1=\"31\" y1=\"21\" x2=\"31\" y2=\"47\" />\n",
       "  <line x1=\"34\" y1=\"24\" x2=\"34\" y2=\"49\" />\n",
       "  <line x1=\"36\" y1=\"26\" x2=\"36\" y2=\"52\" />\n",
       "  <line x1=\"39\" y1=\"29\" x2=\"39\" y2=\"54\" />\n",
       "  <line x1=\"41\" y1=\"31\" x2=\"41\" y2=\"56\" />\n",
       "  <line x1=\"43\" y1=\"33\" x2=\"43\" y2=\"59\" />\n",
       "  <line x1=\"46\" y1=\"36\" x2=\"46\" y2=\"61\" />\n",
       "  <line x1=\"48\" y1=\"38\" x2=\"48\" y2=\"64\" />\n",
       "  <line x1=\"51\" y1=\"41\" x2=\"51\" y2=\"66\" />\n",
       "  <line x1=\"53\" y1=\"43\" x2=\"53\" y2=\"68\" />\n",
       "  <line x1=\"55\" y1=\"45\" x2=\"55\" y2=\"71\" />\n",
       "  <line x1=\"58\" y1=\"48\" x2=\"58\" y2=\"73\" />\n",
       "  <line x1=\"60\" y1=\"50\" x2=\"60\" y2=\"76\" />\n",
       "  <line x1=\"63\" y1=\"53\" x2=\"63\" y2=\"78\" />\n",
       "  <line x1=\"65\" y1=\"55\" x2=\"65\" y2=\"81\" />\n",
       "  <line x1=\"68\" y1=\"58\" x2=\"68\" y2=\"83\" />\n",
       "  <line x1=\"70\" y1=\"60\" x2=\"70\" y2=\"85\" />\n",
       "  <line x1=\"72\" y1=\"62\" x2=\"72\" y2=\"88\" />\n",
       "  <line x1=\"75\" y1=\"65\" x2=\"75\" y2=\"90\" />\n",
       "  <line x1=\"77\" y1=\"67\" x2=\"77\" y2=\"93\" />\n",
       "  <line x1=\"80\" y1=\"70\" x2=\"80\" y2=\"95\" />\n",
       "  <line x1=\"80\" y1=\"70\" x2=\"80\" y2=\"96\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"10.000000,0.000000 80.588235,70.588235 80.588235,96.000852 10.000000,25.412617\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"10\" y1=\"0\" x2=\"38\" y2=\"0\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"12\" y1=\"2\" x2=\"40\" y2=\"2\" />\n",
       "  <line x1=\"14\" y1=\"4\" x2=\"43\" y2=\"4\" />\n",
       "  <line x1=\"17\" y1=\"7\" x2=\"45\" y2=\"7\" />\n",
       "  <line x1=\"19\" y1=\"9\" x2=\"48\" y2=\"9\" />\n",
       "  <line x1=\"22\" y1=\"12\" x2=\"50\" y2=\"12\" />\n",
       "  <line x1=\"24\" y1=\"14\" x2=\"52\" y2=\"14\" />\n",
       "  <line x1=\"26\" y1=\"16\" x2=\"55\" y2=\"16\" />\n",
       "  <line x1=\"29\" y1=\"19\" x2=\"57\" y2=\"19\" />\n",
       "  <line x1=\"31\" y1=\"21\" x2=\"60\" y2=\"21\" />\n",
       "  <line x1=\"34\" y1=\"24\" x2=\"62\" y2=\"24\" />\n",
       "  <line x1=\"36\" y1=\"26\" x2=\"65\" y2=\"26\" />\n",
       "  <line x1=\"39\" y1=\"29\" x2=\"67\" y2=\"29\" />\n",
       "  <line x1=\"41\" y1=\"31\" x2=\"69\" y2=\"31\" />\n",
       "  <line x1=\"43\" y1=\"33\" x2=\"72\" y2=\"33\" />\n",
       "  <line x1=\"46\" y1=\"36\" x2=\"74\" y2=\"36\" />\n",
       "  <line x1=\"48\" y1=\"38\" x2=\"77\" y2=\"38\" />\n",
       "  <line x1=\"51\" y1=\"41\" x2=\"79\" y2=\"41\" />\n",
       "  <line x1=\"53\" y1=\"43\" x2=\"82\" y2=\"43\" />\n",
       "  <line x1=\"55\" y1=\"45\" x2=\"84\" y2=\"45\" />\n",
       "  <line x1=\"58\" y1=\"48\" x2=\"86\" y2=\"48\" />\n",
       "  <line x1=\"60\" y1=\"50\" x2=\"89\" y2=\"50\" />\n",
       "  <line x1=\"63\" y1=\"53\" x2=\"91\" y2=\"53\" />\n",
       "  <line x1=\"65\" y1=\"55\" x2=\"94\" y2=\"55\" />\n",
       "  <line x1=\"68\" y1=\"58\" x2=\"96\" y2=\"58\" />\n",
       "  <line x1=\"70\" y1=\"60\" x2=\"98\" y2=\"60\" />\n",
       "  <line x1=\"72\" y1=\"62\" x2=\"101\" y2=\"62\" />\n",
       "  <line x1=\"75\" y1=\"65\" x2=\"103\" y2=\"65\" />\n",
       "  <line x1=\"77\" y1=\"67\" x2=\"106\" y2=\"67\" />\n",
       "  <line x1=\"80\" y1=\"70\" x2=\"108\" y2=\"70\" />\n",
       "  <line x1=\"80\" y1=\"70\" x2=\"109\" y2=\"70\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Vertical lines -->\n",
       "  <line x1=\"10\" y1=\"0\" x2=\"80\" y2=\"70\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"38\" y1=\"0\" x2=\"109\" y2=\"70\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"10.000000,0.000000 38.489733,0.000000 109.077968,70.588235 80.588235,70.588235\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"80\" y1=\"70\" x2=\"109\" y2=\"70\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"80\" y1=\"96\" x2=\"109\" y2=\"96\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Vertical lines -->\n",
       "  <line x1=\"80\" y1=\"70\" x2=\"80\" y2=\"96\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"109\" y1=\"70\" x2=\"109\" y2=\"96\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"80.588235,70.588235 109.077968,70.588235 109.077968,96.000852 80.588235,96.000852\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Text -->\n",
       "  <text x=\"94.833102\" y=\"116.000852\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" >53</text>\n",
       "  <text x=\"129.077968\" y=\"83.294544\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(0,129.077968,83.294544)\">25</text>\n",
       "  <text x=\"35.294118\" y=\"80.706734\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(45,35.294118,80.706734)\">2920</text>\n",
       "</svg>\n",
       "</td>\n",
       "</tr>\n",
       "</table>"
      ],
      "text/plain": [
       "dask.array<xarray-air, shape=(2920, 25, 53), dtype=float32, chunksize=(100, 25, 53), chunktype=numpy.ndarray>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds.air.data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "! rm -rf *.zarr # clean up any existing temporary data\n",
    "ds.to_zarr('air_temperature.zarr')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/\n",
      " ├── air (2920, 25, 53) float32\n",
      " ├── lat (25,) float32\n",
      " ├── lon (53,) float32\n",
      " └── time (2920,) float32\n"
     ]
    }
   ],
   "source": [
    "source_group = zarr.open('air_temperature.zarr')\n",
    "print(source_group.tree())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "zarr.hierarchy.Group"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(source_group)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table class=\"zarr-info\"><tbody><tr><th style=\"text-align: left\">Name</th><td style=\"text-align: left\">/air</td></tr><tr><th style=\"text-align: left\">Type</th><td style=\"text-align: left\">zarr.core.Array</td></tr><tr><th style=\"text-align: left\">Data type</th><td style=\"text-align: left\">float32</td></tr><tr><th style=\"text-align: left\">Shape</th><td style=\"text-align: left\">(2920, 25, 53)</td></tr><tr><th style=\"text-align: left\">Chunk shape</th><td style=\"text-align: left\">(100, 25, 53)</td></tr><tr><th style=\"text-align: left\">Order</th><td style=\"text-align: left\">C</td></tr><tr><th style=\"text-align: left\">Read-only</th><td style=\"text-align: left\">False</td></tr><tr><th style=\"text-align: left\">Compressor</th><td style=\"text-align: left\">Blosc(cname='lz4', clevel=5, shuffle=SHUFFLE, blocksize=0)</td></tr><tr><th style=\"text-align: left\">Store type</th><td style=\"text-align: left\">zarr.storage.DirectoryStore</td></tr><tr><th style=\"text-align: left\">No. bytes</th><td style=\"text-align: left\">15476000 (14.8M)</td></tr><tr><th style=\"text-align: left\">No. bytes stored</th><td style=\"text-align: left\">9209491 (8.8M)</td></tr><tr><th style=\"text-align: left\">Storage ratio</th><td style=\"text-align: left\">1.7</td></tr><tr><th style=\"text-align: left\">Chunks initialized</th><td style=\"text-align: left\">30/30</td></tr></tbody></table>"
      ],
      "text/plain": [
       "Name               : /air\n",
       "Type               : zarr.core.Array\n",
       "Data type          : float32\n",
       "Shape              : (2920, 25, 53)\n",
       "Chunk shape        : (100, 25, 53)\n",
       "Order              : C\n",
       "Read-only          : False\n",
       "Compressor         : Blosc(cname='lz4', clevel=5, shuffle=SHUFFLE, blocksize=0)\n",
       "Store type         : zarr.storage.DirectoryStore\n",
       "No. bytes          : 15476000 (14.8M)\n",
       "No. bytes stored   : 9209491 (8.8M)\n",
       "Storage ratio      : 1.7\n",
       "Chunks initialized : 30/30"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "source_array = source_group['air']\n",
    "source_array.info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<h2>Rechunked</h2>\n",
       "        <details>\n",
       "          <summary><b>Source</b></summary>\n",
       "          <p><code>&lt;zarr.core.Array (2920, 25, 53) float32&gt;</code></p>\n",
       "        </details>\n",
       "        <details>\n",
       "<summary><b>Intermediate</b></summary>\n",
       "<p><code>&lt;zarr.core.Array (2920, 25, 53) float32&gt;</code></p>\n",
       "</details>\n",
       "\n",
       "        <details>\n",
       "          <summary><b>Target</b></summary>\n",
       "          <p><code>&lt;zarr.core.Array (2920, 25, 53) float32&gt;</code></p>\n",
       "        </details>\n"
      ],
      "text/plain": [
       "<Rechunked>\n",
       "* Source      : <zarr.core.Array '/air' (2920, 25, 53) float32>\n",
       "* Intermediate: <zarr.core.Array (2920, 25, 53) float32>\n",
       "* Target      : <zarr.core.Array (2920, 25, 53) float32>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from rechunker import rechunk\n",
    "\n",
    "target_chunks_dict = {'time': 2920, 'lat': 25, 'lon': 1}\n",
    "max_mem = '1MB'\n",
    "\n",
    "target_store = 'air_rechunked.zarr'\n",
    "temp_store = 'air_rechunked-tmp.zarr'\n",
    "\n",
    "array_plan = rechunk(source_array, target_chunks_dict, max_mem, target_store, temp_store=temp_store)\n",
    "array_plan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[########################################] | 100% Completed |  0.7s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(2920, 25, 1)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from dask.diagnostics import ProgressBar\n",
    "\n",
    "with ProgressBar():\n",
    "    result = array_plan.execute()\n",
    "result.chunks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<zarr.core.Array (2920, 25, 53) float32>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "target_array = zarr.open('air_rechunked.zarr')\n",
    "target_array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table>\n",
       "<tr>\n",
       "<td>\n",
       "<table>\n",
       "  <thead>\n",
       "    <tr><td> </td><th> Array </th><th> Chunk </th></tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr><th> Bytes </th><td> 15.48 MB </td> <td> 292.00 kB </td></tr>\n",
       "    <tr><th> Shape </th><td> (2920, 25, 53) </td> <td> (2920, 25, 1) </td></tr>\n",
       "    <tr><th> Count </th><td> 54 Tasks </td><td> 53 Chunks </td></tr>\n",
       "    <tr><th> Type </th><td> float32 </td><td> numpy.ndarray </td></tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</td>\n",
       "<td>\n",
       "<svg width=\"159\" height=\"146\" style=\"stroke:rgb(0,0,0);stroke-width:1\" >\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"10\" y1=\"0\" x2=\"80\" y2=\"70\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"10\" y1=\"25\" x2=\"80\" y2=\"96\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Vertical lines -->\n",
       "  <line x1=\"10\" y1=\"0\" x2=\"10\" y2=\"25\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"80\" y1=\"70\" x2=\"80\" y2=\"96\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"10.000000,0.000000 80.588235,70.588235 80.588235,96.000852 10.000000,25.412617\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"10\" y1=\"0\" x2=\"38\" y2=\"0\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"80\" y1=\"70\" x2=\"109\" y2=\"70\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Vertical lines -->\n",
       "  <line x1=\"10\" y1=\"0\" x2=\"80\" y2=\"70\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"10\" y1=\"0\" x2=\"81\" y2=\"70\" />\n",
       "  <line x1=\"11\" y1=\"0\" x2=\"81\" y2=\"70\" />\n",
       "  <line x1=\"11\" y1=\"0\" x2=\"82\" y2=\"70\" />\n",
       "  <line x1=\"12\" y1=\"0\" x2=\"82\" y2=\"70\" />\n",
       "  <line x1=\"12\" y1=\"0\" x2=\"83\" y2=\"70\" />\n",
       "  <line x1=\"13\" y1=\"0\" x2=\"83\" y2=\"70\" />\n",
       "  <line x1=\"13\" y1=\"0\" x2=\"84\" y2=\"70\" />\n",
       "  <line x1=\"14\" y1=\"0\" x2=\"84\" y2=\"70\" />\n",
       "  <line x1=\"14\" y1=\"0\" x2=\"85\" y2=\"70\" />\n",
       "  <line x1=\"15\" y1=\"0\" x2=\"85\" y2=\"70\" />\n",
       "  <line x1=\"15\" y1=\"0\" x2=\"86\" y2=\"70\" />\n",
       "  <line x1=\"16\" y1=\"0\" x2=\"87\" y2=\"70\" />\n",
       "  <line x1=\"16\" y1=\"0\" x2=\"87\" y2=\"70\" />\n",
       "  <line x1=\"17\" y1=\"0\" x2=\"88\" y2=\"70\" />\n",
       "  <line x1=\"18\" y1=\"0\" x2=\"88\" y2=\"70\" />\n",
       "  <line x1=\"18\" y1=\"0\" x2=\"89\" y2=\"70\" />\n",
       "  <line x1=\"19\" y1=\"0\" x2=\"89\" y2=\"70\" />\n",
       "  <line x1=\"19\" y1=\"0\" x2=\"90\" y2=\"70\" />\n",
       "  <line x1=\"20\" y1=\"0\" x2=\"90\" y2=\"70\" />\n",
       "  <line x1=\"20\" y1=\"0\" x2=\"91\" y2=\"70\" />\n",
       "  <line x1=\"21\" y1=\"0\" x2=\"91\" y2=\"70\" />\n",
       "  <line x1=\"21\" y1=\"0\" x2=\"92\" y2=\"70\" />\n",
       "  <line x1=\"22\" y1=\"0\" x2=\"92\" y2=\"70\" />\n",
       "  <line x1=\"22\" y1=\"0\" x2=\"93\" y2=\"70\" />\n",
       "  <line x1=\"23\" y1=\"0\" x2=\"94\" y2=\"70\" />\n",
       "  <line x1=\"23\" y1=\"0\" x2=\"94\" y2=\"70\" />\n",
       "  <line x1=\"24\" y1=\"0\" x2=\"95\" y2=\"70\" />\n",
       "  <line x1=\"25\" y1=\"0\" x2=\"95\" y2=\"70\" />\n",
       "  <line x1=\"25\" y1=\"0\" x2=\"96\" y2=\"70\" />\n",
       "  <line x1=\"26\" y1=\"0\" x2=\"96\" y2=\"70\" />\n",
       "  <line x1=\"26\" y1=\"0\" x2=\"97\" y2=\"70\" />\n",
       "  <line x1=\"27\" y1=\"0\" x2=\"97\" y2=\"70\" />\n",
       "  <line x1=\"27\" y1=\"0\" x2=\"98\" y2=\"70\" />\n",
       "  <line x1=\"28\" y1=\"0\" x2=\"98\" y2=\"70\" />\n",
       "  <line x1=\"28\" y1=\"0\" x2=\"99\" y2=\"70\" />\n",
       "  <line x1=\"29\" y1=\"0\" x2=\"99\" y2=\"70\" />\n",
       "  <line x1=\"29\" y1=\"0\" x2=\"100\" y2=\"70\" />\n",
       "  <line x1=\"30\" y1=\"0\" x2=\"101\" y2=\"70\" />\n",
       "  <line x1=\"30\" y1=\"0\" x2=\"101\" y2=\"70\" />\n",
       "  <line x1=\"31\" y1=\"0\" x2=\"102\" y2=\"70\" />\n",
       "  <line x1=\"32\" y1=\"0\" x2=\"102\" y2=\"70\" />\n",
       "  <line x1=\"32\" y1=\"0\" x2=\"103\" y2=\"70\" />\n",
       "  <line x1=\"33\" y1=\"0\" x2=\"103\" y2=\"70\" />\n",
       "  <line x1=\"33\" y1=\"0\" x2=\"104\" y2=\"70\" />\n",
       "  <line x1=\"34\" y1=\"0\" x2=\"104\" y2=\"70\" />\n",
       "  <line x1=\"34\" y1=\"0\" x2=\"105\" y2=\"70\" />\n",
       "  <line x1=\"35\" y1=\"0\" x2=\"105\" y2=\"70\" />\n",
       "  <line x1=\"35\" y1=\"0\" x2=\"106\" y2=\"70\" />\n",
       "  <line x1=\"36\" y1=\"0\" x2=\"106\" y2=\"70\" />\n",
       "  <line x1=\"36\" y1=\"0\" x2=\"107\" y2=\"70\" />\n",
       "  <line x1=\"37\" y1=\"0\" x2=\"108\" y2=\"70\" />\n",
       "  <line x1=\"37\" y1=\"0\" x2=\"108\" y2=\"70\" />\n",
       "  <line x1=\"38\" y1=\"0\" x2=\"109\" y2=\"70\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"10.000000,0.000000 38.489733,0.000000 109.077968,70.588235 80.588235,70.588235\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"80\" y1=\"70\" x2=\"109\" y2=\"70\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"80\" y1=\"96\" x2=\"109\" y2=\"96\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Vertical lines -->\n",
       "  <line x1=\"80\" y1=\"70\" x2=\"80\" y2=\"96\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"81\" y1=\"70\" x2=\"81\" y2=\"96\" />\n",
       "  <line x1=\"81\" y1=\"70\" x2=\"81\" y2=\"96\" />\n",
       "  <line x1=\"82\" y1=\"70\" x2=\"82\" y2=\"96\" />\n",
       "  <line x1=\"82\" y1=\"70\" x2=\"82\" y2=\"96\" />\n",
       "  <line x1=\"83\" y1=\"70\" x2=\"83\" y2=\"96\" />\n",
       "  <line x1=\"83\" y1=\"70\" x2=\"83\" y2=\"96\" />\n",
       "  <line x1=\"84\" y1=\"70\" x2=\"84\" y2=\"96\" />\n",
       "  <line x1=\"84\" y1=\"70\" x2=\"84\" y2=\"96\" />\n",
       "  <line x1=\"85\" y1=\"70\" x2=\"85\" y2=\"96\" />\n",
       "  <line x1=\"85\" y1=\"70\" x2=\"85\" y2=\"96\" />\n",
       "  <line x1=\"86\" y1=\"70\" x2=\"86\" y2=\"96\" />\n",
       "  <line x1=\"87\" y1=\"70\" x2=\"87\" y2=\"96\" />\n",
       "  <line x1=\"87\" y1=\"70\" x2=\"87\" y2=\"96\" />\n",
       "  <line x1=\"88\" y1=\"70\" x2=\"88\" y2=\"96\" />\n",
       "  <line x1=\"88\" y1=\"70\" x2=\"88\" y2=\"96\" />\n",
       "  <line x1=\"89\" y1=\"70\" x2=\"89\" y2=\"96\" />\n",
       "  <line x1=\"89\" y1=\"70\" x2=\"89\" y2=\"96\" />\n",
       "  <line x1=\"90\" y1=\"70\" x2=\"90\" y2=\"96\" />\n",
       "  <line x1=\"90\" y1=\"70\" x2=\"90\" y2=\"96\" />\n",
       "  <line x1=\"91\" y1=\"70\" x2=\"91\" y2=\"96\" />\n",
       "  <line x1=\"91\" y1=\"70\" x2=\"91\" y2=\"96\" />\n",
       "  <line x1=\"92\" y1=\"70\" x2=\"92\" y2=\"96\" />\n",
       "  <line x1=\"92\" y1=\"70\" x2=\"92\" y2=\"96\" />\n",
       "  <line x1=\"93\" y1=\"70\" x2=\"93\" y2=\"96\" />\n",
       "  <line x1=\"94\" y1=\"70\" x2=\"94\" y2=\"96\" />\n",
       "  <line x1=\"94\" y1=\"70\" x2=\"94\" y2=\"96\" />\n",
       "  <line x1=\"95\" y1=\"70\" x2=\"95\" y2=\"96\" />\n",
       "  <line x1=\"95\" y1=\"70\" x2=\"95\" y2=\"96\" />\n",
       "  <line x1=\"96\" y1=\"70\" x2=\"96\" y2=\"96\" />\n",
       "  <line x1=\"96\" y1=\"70\" x2=\"96\" y2=\"96\" />\n",
       "  <line x1=\"97\" y1=\"70\" x2=\"97\" y2=\"96\" />\n",
       "  <line x1=\"97\" y1=\"70\" x2=\"97\" y2=\"96\" />\n",
       "  <line x1=\"98\" y1=\"70\" x2=\"98\" y2=\"96\" />\n",
       "  <line x1=\"98\" y1=\"70\" x2=\"98\" y2=\"96\" />\n",
       "  <line x1=\"99\" y1=\"70\" x2=\"99\" y2=\"96\" />\n",
       "  <line x1=\"99\" y1=\"70\" x2=\"99\" y2=\"96\" />\n",
       "  <line x1=\"100\" y1=\"70\" x2=\"100\" y2=\"96\" />\n",
       "  <line x1=\"101\" y1=\"70\" x2=\"101\" y2=\"96\" />\n",
       "  <line x1=\"101\" y1=\"70\" x2=\"101\" y2=\"96\" />\n",
       "  <line x1=\"102\" y1=\"70\" x2=\"102\" y2=\"96\" />\n",
       "  <line x1=\"102\" y1=\"70\" x2=\"102\" y2=\"96\" />\n",
       "  <line x1=\"103\" y1=\"70\" x2=\"103\" y2=\"96\" />\n",
       "  <line x1=\"103\" y1=\"70\" x2=\"103\" y2=\"96\" />\n",
       "  <line x1=\"104\" y1=\"70\" x2=\"104\" y2=\"96\" />\n",
       "  <line x1=\"104\" y1=\"70\" x2=\"104\" y2=\"96\" />\n",
       "  <line x1=\"105\" y1=\"70\" x2=\"105\" y2=\"96\" />\n",
       "  <line x1=\"105\" y1=\"70\" x2=\"105\" y2=\"96\" />\n",
       "  <line x1=\"106\" y1=\"70\" x2=\"106\" y2=\"96\" />\n",
       "  <line x1=\"106\" y1=\"70\" x2=\"106\" y2=\"96\" />\n",
       "  <line x1=\"107\" y1=\"70\" x2=\"107\" y2=\"96\" />\n",
       "  <line x1=\"108\" y1=\"70\" x2=\"108\" y2=\"96\" />\n",
       "  <line x1=\"108\" y1=\"70\" x2=\"108\" y2=\"96\" />\n",
       "  <line x1=\"109\" y1=\"70\" x2=\"109\" y2=\"96\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"80.588235,70.588235 109.077968,70.588235 109.077968,96.000852 80.588235,96.000852\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Text -->\n",
       "  <text x=\"94.833102\" y=\"116.000852\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" >53</text>\n",
       "  <text x=\"129.077968\" y=\"83.294544\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(0,129.077968,83.294544)\">25</text>\n",
       "  <text x=\"35.294118\" y=\"80.706734\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(45,35.294118,80.706734)\">2920</text>\n",
       "</svg>\n",
       "</td>\n",
       "</tr>\n",
       "</table>"
      ],
      "text/plain": [
       "dask.array<from-zarr, shape=(2920, 25, 53), dtype=float32, chunksize=(2920, 25, 1), chunktype=numpy.ndarray>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "target_array_dask = dsa.from_zarr('air_rechunked.zarr')\n",
    "target_array_dask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<h2>Rechunked</h2>\n",
       "        <details>\n",
       "          <summary><b>Source</b></summary>\n",
       "          <p><code>&lt;zarr.hierarchy.Group &#x27;/&#x27;&gt;</code></p>\n",
       "        </details>\n",
       "        <details>\n",
       "<summary><b>Intermediate</b></summary>\n",
       "<p><code>&lt;zarr.hierarchy.Group &#x27;/&#x27;&gt;</code></p>\n",
       "</details>\n",
       "\n",
       "        <details>\n",
       "          <summary><b>Target</b></summary>\n",
       "          <p><code>&lt;zarr.hierarchy.Group &#x27;/&#x27;&gt;</code></p>\n",
       "        </details>\n"
      ],
      "text/plain": [
       "<Rechunked>\n",
       "* Source      : <zarr.hierarchy.Group '/'>\n",
       "* Intermediate: <zarr.hierarchy.Group '/'>\n",
       "* Target      : <zarr.hierarchy.Group '/'>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "target_chunks = {\n",
    "    'air': {'time': 2920, 'lat': 25, 'lon': 1},\n",
    "    'time': None, # don't rechunk this array\n",
    "    'lon': None,\n",
    "    'lat': None,\n",
    "}\n",
    "max_mem = '1MB'\n",
    "\n",
    "target_store = 'group_rechunked.zarr'\n",
    "temp_store = 'group_rechunked-tmp.zarr'\n",
    "\n",
    "array_plan = rechunk(source_group, target_chunks, max_mem, target_store, temp_store=temp_store)\n",
    "array_plan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<zarr.hierarchy.Group '/'>"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "array_plan.execute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre>&lt;xarray.Dataset&gt;\n",
       "Dimensions:  (lat: 25, lon: 53, time: 2920)\n",
       "Coordinates:\n",
       "  * lat      (lat) float32 75.0 72.5 70.0 67.5 65.0 ... 25.0 22.5 20.0 17.5 15.0\n",
       "  * lon      (lon) float32 200.0 202.5 205.0 207.5 ... 322.5 325.0 327.5 330.0\n",
       "  * time     (time) datetime64[ns] 2013-01-01 ... 2014-12-31T18:00:00\n",
       "Data variables:\n",
       "    air      (time, lat, lon) float32 dask.array&lt;chunksize=(2920, 25, 1), meta=np.ndarray&gt;\n",
       "Attributes:\n",
       "    Conventions:  COARDS\n",
       "    description:  Data is from NMC initialized reanalysis\\n(4x/day).  These a...\n",
       "    platform:     Model\n",
       "    references:   http://www.esrl.noaa.gov/psd/data/gridded/data.ncep.reanaly...\n",
       "    title:        4x daily NMC reanalysis (1948)</pre>"
      ],
      "text/plain": [
       "<xarray.Dataset>\n",
       "Dimensions:  (lat: 25, lon: 53, time: 2920)\n",
       "Coordinates:\n",
       "  * lat      (lat) float32 75.0 72.5 70.0 67.5 65.0 ... 25.0 22.5 20.0 17.5 15.0\n",
       "  * lon      (lon) float32 200.0 202.5 205.0 207.5 ... 322.5 325.0 327.5 330.0\n",
       "  * time     (time) datetime64[ns] 2013-01-01 ... 2014-12-31T18:00:00\n",
       "Data variables:\n",
       "    air      (time, lat, lon) float32 dask.array<chunksize=(2920, 25, 1), meta=np.ndarray>\n",
       "Attributes:\n",
       "    Conventions:  COARDS\n",
       "    description:  Data is from NMC initialized reanalysis\\n(4x/day).  These a...\n",
       "    platform:     Model\n",
       "    references:   http://www.esrl.noaa.gov/psd/data/gridded/data.ncep.reanaly...\n",
       "    title:        4x daily NMC reanalysis (1948)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xr.open_zarr('group_rechunked.zarr')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}