{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from glob import glob\n",
    "import xarray\n",
    "import numpy\n",
    "import dask\n",
    "from e3smplot.e3sm_utils import get_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = '/global/cfs/cdirs/e3sm/terai/SCREAM/DYAMOND2/Output/20201127/regridded'\n",
    "p_files = sorted(glob(f'{data_path}/*.eam.h4.*.nc'))\n",
    "t_files = sorted(glob(f'{data_path}/*.eam.h6.*.nc'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_p = xarray.open_mfdataset(p_files, drop_variables=('P3_input_dim', 'P3_output_dim'), chunks={'time': 1, 'lev': 1})\n",
    "ds_t = xarray.open_mfdataset(t_files, drop_variables=('P3_input_dim', 'P3_output_dim'), chunks={'time': 1, 'lev': 1})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = get_data(ds_p, 'PMID')\n",
    "t = get_data(ds_t, 'T')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "levels = (200, 300, 400, 500, 600, 700, 800, 850, 925, 1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "levels = xarray.DataArray(\n",
    "    numpy.array(levels), dims=('level'), attrs={'long_name': 'Pressure', 'units': 'hPa'}\n",
    ")\n",
    "t_p = xarray.DataArray(\n",
    "    dask.array.zeros([t.time.size, len(levels), t.lat.size, t.lon.size], chunks=(1, len(levels), t.lat.size, t.lon.size))\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "levels, *__ = xarray.broadcast(\n",
    "    levels, p.isel(lev=0)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [],
   "source": [
    "levels = levels.transpose('time', 'level', 'lat', 'lon')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('time', 'level', 'lat', 'lon')"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "levels.dims"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/global/homes/b/bhillma/.conda/envs/e3smplot/lib/python3.8/site-packages/dask/array/gufunc.py:430: PerformanceWarning: Increasing number of chunks by factor of 32\n",
      "  tmp = blockwise(  # First try to compute meta\n"
     ]
    }
   ],
   "source": [
    "t.data = t.data.rechunk((1, t.lev.size, t.lat.size, t.lon.size))\n",
    "p.data = p.data.rechunk((1, p.lev.size, p.lat.size, p.lon.size))\n",
    "t_p.data = xarray.apply_ufunc(\n",
    "    numpy.interp, levels, p, t,\n",
    "    exclude_dims=set(('lev',)),\n",
    "    input_core_dims=[[], ['lev',], ['lev',]],\n",
    "    dask='parallelized',\n",
    "    vectorize=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('done')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "e3smplot",
   "language": "python",
   "name": "e3smplot"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}