{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Higher Dimensional Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pathlib \n",
    "from collections import defaultdict\n",
    "\n",
    "import h5py\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import LogNorm,Normalize\n",
    "import matplotlib.cm as cm\n",
    "\n",
    "from bluesky_tutorial_utils import nexus, fetch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Un-comment this for better-looking plots on high-resolution screens.\n",
    "# %config InlineBackend.figure_format = 'retina'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fetch.rsoxs_simulation_data();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Working with Multiple xArrays: Gathering Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build Index Table (Table of Contents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_pandas_index(nxs_path):\n",
    "    nxs_path = pathlib.Path(nxs_path)\n",
    "    nxs_files = list(nxs_path.glob('*nxs'))\n",
    "    \n",
    "    #progress = ipywidgets.IntProgress(0,0,len(nxs_files))\n",
    "    #display(progress)\n",
    "    \n",
    "    index_table = []\n",
    "    for i,nxs_file in enumerate(nxs_files):\n",
    "        #progress.value = i\n",
    "        with h5py.File(nxs_file,'r') as nxs:\n",
    "            notes = nxs[u'entry/instrument/simulation_engine/notes']\n",
    "            config =  {k:v[()] for k,v in notes.items()}\n",
    "            config['nxs'] = nxs_file\n",
    "            index_table.append(config)\n",
    "    return pd.DataFrame(index_table)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "toc = build_pandas_index('./rsoxs_simulation_data//')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "toc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "toc.describe().loc[['count','min','max']]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Select subset of data From Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sdf = toc.query('Radius==40.0 & EndAngle==360.0 & PhysSize==5 & NumX==512')\n",
    "sdf = sdf.sort_values('Energy')\n",
    "sdf.describe().loc[['count','min','max']]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gather Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gather(df):\n",
    "    coords = defaultdict(list)\n",
    "    data_arrays = []\n",
    "    for row_index,row in df.iterrows():\n",
    "        da_img = nexus.read_singleimg_nxs(row['nxs'])\n",
    "        data_arrays.append(da_img)\n",
    "        \n",
    "        for col_index,value in row.iteritems():\n",
    "            if col_index=='nxs':\n",
    "                continue\n",
    "            coords[col_index].append(value)\n",
    "    return data_arrays,coords"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_arrays,coords = gather(sdf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_arrays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "coords"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_arrays[1].plot(norm=LogNorm(1e-9,1),aspect=1.2,size=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multiple xArrays: simple xr.concat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sdf = toc.query('Radius==40.0 & EndAngle==360.0 & PhysSize==5 & NumX==512')\n",
    "sdf = sdf.sort_values('Energy')\n",
    "display(sdf.describe().loc[['count','min','max']])\n",
    "\n",
    "data_arrays,coords = gather(sdf)\n",
    "\n",
    "da = xr.concat(data_arrays,dim='Energy')\n",
    "da"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "da = da.assign_coords(Energy=sdf.Energy.values)\n",
    "da"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "da.sel(Qy=0,method='nearest').plot(norm=LogNorm(1e-9,1),yscale='log')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# da.plot(col='Energy',col_wrap=3,norm=LogNorm(1e-9,1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building xArrays: Multi-Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sdf = toc.query('EndAngle==360.0 & PhysSize==5 & NumX==512')\n",
    "sdf = sdf.sort_values(['Energy','Radius'])\n",
    "display(sdf.describe().loc[['count','min','max']])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "raises-exception"
    ]
   },
   "outputs": [],
   "source": [
    "data_arrays,coords = gather(sdf)\n",
    "\n",
    "da = xr.concat(data_arrays,dim=['Energy','Radius'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "hmmm...that didn't work...\n",
    "\n",
    "\n",
    "Let's try a multi-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "keys =  ['Energy','Radius']\n",
    "tuples = [(i,j) for i,j in sdf[keys].values]\n",
    "index = pd.MultiIndex.from_tuples(tuples,names=keys)\n",
    "index.name = 'system'\n",
    "da = xr.concat(data_arrays,dim=index)\n",
    "da"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "da.sel(Energy=285.,method='nearest')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "da.sel(Energy=285.,Qy=0,method='nearest').plot.line(x='Qx',yscale='log',xscale='log')#(norm=LogNorm(1e-9,1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "da.sel(Energy=285.,method='nearest').plot(col='Radius',col_wrap=3,norm=LogNorm(1e-9,1))"
   ]
  }
 ],
 "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}