ReshapeDataset.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How can Datasets be reshaped?\n",
    "\n",
    "Question: Is there a way to reshape a Dataset (given that the user provides shape compatible dimensions)?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'0.14.1'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import xarray as xr\n",
    "import numpy as np\n",
    "xr.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre>&lt;xarray.Dataset&gt;\n",
       "Dimensions:  (draw: 2, row: 24)\n",
       "Coordinates:\n",
       "  * draw     (draw) int32 0 1\n",
       "  * row      (row) int32 0 1 2 3 4 5 6 7 8 9 ... 14 15 16 17 18 19 20 21 22 23\n",
       "Data variables:\n",
       "    obs      (draw, row) int32 0 1 2 3 4 5 6 7 8 ... 39 40 41 42 43 44 45 46 47</pre>"
      ],
      "text/plain": [
       "<xarray.Dataset>\n",
       "Dimensions:  (draw: 2, row: 24)\n",
       "Coordinates:\n",
       "  * draw     (draw) int32 0 1\n",
       "  * row      (row) int32 0 1 2 3 4 5 6 7 8 9 ... 14 15 16 17 18 19 20 21 22 23\n",
       "Data variables:\n",
       "    obs      (draw, row) int32 0 1 2 3 4 5 6 7 8 ... 39 40 41 42 43 44 45 46 47"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = np.arange(48).reshape((2,-1))\n",
    "ds = xr.Dataset({\"obs\": ([\"draw\", \"row\"], data)}, coords={\"draw\": [0,1], \"row\": range(24)})\n",
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We want to convert the `row` dimension to 3 dimensions of sizes `(3, 2, 4)` respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre>&lt;xarray.Dataset&gt;\n",
       "Dimensions:  (dim: 24, draw: 2)\n",
       "Coordinates:\n",
       "  * draw     (draw) int32 0 1\n",
       "  * dim      (dim) MultiIndex\n",
       "  - dim_0    (dim) object &#x27;a&#x27; &#x27;a&#x27; &#x27;a&#x27; &#x27;a&#x27; &#x27;a&#x27; &#x27;a&#x27; ... &#x27;c&#x27; &#x27;c&#x27; &#x27;c&#x27; &#x27;c&#x27; &#x27;c&#x27; &#x27;c&#x27;\n",
       "  - dim_1    (dim) int64 0 0 0 0 1 1 1 1 0 0 0 0 1 1 1 1 0 0 0 0 1 1 1 1\n",
       "  - dim_2    (dim) int64 0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3\n",
       "Data variables:\n",
       "    obs      (draw, dim) int32 0 1 2 3 4 5 6 7 8 ... 39 40 41 42 43 44 45 46 47</pre>"
      ],
      "text/plain": [
       "<xarray.Dataset>\n",
       "Dimensions:  (dim: 24, draw: 2)\n",
       "Coordinates:\n",
       "  * draw     (draw) int32 0 1\n",
       "  * dim      (dim) MultiIndex\n",
       "  - dim_0    (dim) object 'a' 'a' 'a' 'a' 'a' 'a' ... 'c' 'c' 'c' 'c' 'c' 'c'\n",
       "  - dim_1    (dim) int64 0 0 0 0 1 1 1 1 0 0 0 0 1 1 1 1 0 0 0 0 1 1 1 1\n",
       "  - dim_2    (dim) int64 0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3 0 1 2 3\n",
       "Data variables:\n",
       "    obs      (draw, dim) int32 0 1 2 3 4 5 6 7 8 ... 39 40 41 42 43 44 45 46 47"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "aux_ds = ds.assign_coords(\n",
    "    dim_0=[\"a\", \"b\", \"c\"], dim_1=[0,1], dim_2=range(4)\n",
    ").stack(\n",
    "    dim=(\"dim_0\", \"dim_1\", \"dim_2\")\n",
    ").reset_index(\n",
    "    \"row\", drop=True\n",
    ")\n",
    "multiindex_ds = aux_ds.rename(row=\"dim\")\n",
    "multiindex_ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre>&lt;xarray.Dataset&gt;\n",
       "Dimensions:  (dim_0: 3, dim_1: 2, dim_2: 4, draw: 2)\n",
       "Coordinates:\n",
       "  * draw     (draw) int32 0 1\n",
       "  * dim_0    (dim_0) object &#x27;a&#x27; &#x27;b&#x27; &#x27;c&#x27;\n",
       "  * dim_1    (dim_1) int64 0 1\n",
       "  * dim_2    (dim_2) int64 0 1 2 3\n",
       "Data variables:\n",
       "    obs      (draw, dim_0, dim_1, dim_2) int32 0 1 2 3 4 5 ... 42 43 44 45 46 47</pre>"
      ],
      "text/plain": [
       "<xarray.Dataset>\n",
       "Dimensions:  (dim_0: 3, dim_1: 2, dim_2: 4, draw: 2)\n",
       "Coordinates:\n",
       "  * draw     (draw) int32 0 1\n",
       "  * dim_0    (dim_0) object 'a' 'b' 'c'\n",
       "  * dim_1    (dim_1) int64 0 1\n",
       "  * dim_2    (dim_2) int64 0 1 2 3\n",
       "Data variables:\n",
       "    obs      (draw, dim_0, dim_1, dim_2) int32 0 1 2 3 4 5 ... 42 43 44 45 46 47"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reshaped_ds = multiindex_ds.unstack(\"dim\")\n",
    "reshaped_ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "----\n",
    "**Note**: if instead of using `rename` on `aux_df` we were to use `set_index`, which could seem more natural, we would not get the same result, the multiindex object is converted into a single index object with tuples as coordinate values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre>&lt;xarray.Dataset&gt;\n",
       "Dimensions:  (draw: 2, row: 24)\n",
       "Coordinates:\n",
       "  * draw     (draw) int32 0 1\n",
       "  * row      (row) object (&#x27;a&#x27;, 0, 0) (&#x27;a&#x27;, 0, 1) ... (&#x27;c&#x27;, 1, 2) (&#x27;c&#x27;, 1, 3)\n",
       "Data variables:\n",
       "    obs      (draw, row) int32 0 1 2 3 4 5 6 7 8 ... 39 40 41 42 43 44 45 46 47</pre>"
      ],
      "text/plain": [
       "<xarray.Dataset>\n",
       "Dimensions:  (draw: 2, row: 24)\n",
       "Coordinates:\n",
       "  * draw     (draw) int32 0 1\n",
       "  * row      (row) object ('a', 0, 0) ('a', 0, 1) ... ('c', 1, 2) ('c', 1, 3)\n",
       "Data variables:\n",
       "    obs      (draw, row) int32 0 1 2 3 4 5 6 7 8 ... 39 40 41 42 43 44 45 46 47"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "aux_ds.set_index(row=\"dim\")"
   ]
  },
  {
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}