xarray_rechunker_example.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from rechunker import api\n",
    "import xarray as xr\n",
    "xr.set_options(display_style='text');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre>&lt;xarray.Dataset&gt;\n",
       "Dimensions:     (eta_rho: 191, ocean_time: 2, s_rho: 30, xi_rho: 371)\n",
       "Coordinates:\n",
       "    Cs_r        (s_rho) float64 dask.array&lt;chunksize=(30,), meta=np.ndarray&gt;\n",
       "    lon_rho     (eta_rho, xi_rho) float64 dask.array&lt;chunksize=(191, 371), meta=np.ndarray&gt;\n",
       "    hc          float64 ...\n",
       "    h           (eta_rho, xi_rho) float64 dask.array&lt;chunksize=(191, 371), meta=np.ndarray&gt;\n",
       "    lat_rho     (eta_rho, xi_rho) float64 dask.array&lt;chunksize=(191, 371), meta=np.ndarray&gt;\n",
       "    Vtransform  int32 ...\n",
       "  * ocean_time  (ocean_time) datetime64[ns] 2001-08-01 2001-08-08\n",
       "  * s_rho       (s_rho) float64 -0.9833 -0.95 -0.9167 ... -0.05 -0.01667\n",
       "Dimensions without coordinates: eta_rho, xi_rho\n",
       "Data variables:\n",
       "    salt        (ocean_time, s_rho, eta_rho, xi_rho) float32 dask.array&lt;chunksize=(1, 30, 191, 371), meta=np.ndarray&gt;\n",
       "    zeta        (ocean_time, eta_rho, xi_rho) float32 dask.array&lt;chunksize=(1, 191, 371), meta=np.ndarray&gt;</pre>"
      ],
      "text/plain": [
       "<xarray.Dataset>\n",
       "Dimensions:     (eta_rho: 191, ocean_time: 2, s_rho: 30, xi_rho: 371)\n",
       "Coordinates:\n",
       "    Cs_r        (s_rho) float64 dask.array<chunksize=(30,), meta=np.ndarray>\n",
       "    lon_rho     (eta_rho, xi_rho) float64 dask.array<chunksize=(191, 371), meta=np.ndarray>\n",
       "    hc          float64 ...\n",
       "    h           (eta_rho, xi_rho) float64 dask.array<chunksize=(191, 371), meta=np.ndarray>\n",
       "    lat_rho     (eta_rho, xi_rho) float64 dask.array<chunksize=(191, 371), meta=np.ndarray>\n",
       "    Vtransform  int32 ...\n",
       "  * ocean_time  (ocean_time) datetime64[ns] 2001-08-01 2001-08-08\n",
       "  * s_rho       (s_rho) float64 -0.9833 -0.95 -0.9167 ... -0.05 -0.01667\n",
       "Dimensions without coordinates: eta_rho, xi_rho\n",
       "Data variables:\n",
       "    salt        (ocean_time, s_rho, eta_rho, xi_rho) float32 dask.array<chunksize=(1, 30, 191, 371), meta=np.ndarray>\n",
       "    zeta        (ocean_time, eta_rho, xi_rho) float32 dask.array<chunksize=(1, 191, 371), meta=np.ndarray>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds = xr.tutorial.open_dataset('ROMS_example.nc', chunks={'ocean_time': 1})\n",
    "ds.attrs = {}\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'salt': (1, 30, 191, 371), 'zeta': (1, 191, 371)}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "{v: ds[v].data.chunksize for v in ds}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm -rf /tmp/{temp,target}.zarr\n",
    "api.rechunk(\n",
    "    ds, \n",
    "    target_chunks=dict(\n",
    "        salt=(2, 10, 10, 10)\n",
    "    ),\n",
    "    max_mem='500MB',\n",
    "    target_store='/tmp/target.zarr',\n",
    "    temp_store='/tmp/temp.zarr'\n",
    ").execute();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'salt': (2, 10, 10, 10), 'zeta': (1, 191, 371)}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds = xr.open_zarr('/tmp/target.zarr')\n",
    "{v: ds[v].data.chunksize for v in ds}"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gwas",
   "language": "python",
   "name": "gwas"
  },
  "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.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}