{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from wrf.extension import _slp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from wrf import getvar\n",
    "from netCDF4 import Dataset as nc\n",
    "ncfile = nc(\"/Users/ladwig/Documents/wrf_files/wrf_vortex_multi/wrfout_d01_2005-08-28_00:00:00\")\n",
    "b = getvar([ncfile,ncfile,ncfile], \"slp\", timeidx=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "b = getvar([ncfile,ncfile,ncfile], \"td\", None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(b)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "b = getvar([ncfile,ncfile,ncfile], \"tk\", None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "b = getvar([ncfile,ncfile,ncfile], \"rh\", None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print (b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# 500 MB Heights\n",
    "from wrf import getvar, interplevel\n",
    "\n",
    "z = getvar([ncfile,ncfile,ncfile], \"z\", timeidx=None)\n",
    "p = getvar([ncfile,ncfile,ncfile], \"pressure\", timeidx=None)\n",
    "ht_500mb = interplevel(z, p, 500)\n",
    "\n",
    "print(ht_500mb)\n",
    "del ht_500mb, z, p"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Pressure using pivot and angle\n",
    "from wrf import getvar, vertcross\n",
    "\n",
    "z = getvar(ncfile, \"z\", timeidx=None)\n",
    "p = getvar(ncfile, \"pressure\", timeidx=None)\n",
    "pivot_point = (z.shape[-2] / 2, z.shape[-1] / 2) \n",
    "angle = 90.0\n",
    "\n",
    "p_vert = vertcross(p, z, pivot_point=pivot_point, angle=angle)\n",
    "print(p_vert)\n",
    "del p_vert\n",
    "\n",
    "# Pressure using start_point and end_point\n",
    "start_point = (z.shape[-2]/2, 0)\n",
    "end_point = (z.shape[-2]/2, -1)\n",
    "\n",
    "p_vert = vertcross(p, z, start_point=start_point, end_point=end_point)\n",
    "print(p_vert)\n",
    "del p_vert, p, z\n"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# T2 using pivot and angle\n",
    "from wrf import interpline, getvar\n",
    "\n",
    "t2 = getvar([ncfile,ncfile,ncfile], \"T2\", timeidx=None)\n",
    "pivot_point = (t2.shape[-2] / 2, t2.shape[-1] / 2) \n",
    "angle = 90.0\n",
    "\n",
    "t2_line = interpline(t2, pivot_point=pivot_point, angle=angle)\n",
    "print(t2_line)\n",
    "\n",
    "del t2_line\n",
    "\n",
    "# T2 using start_point and end_point\n",
    "start_point = (t2.shape[-2]/2, 0)\n",
    "end_point = (t2.shape[-2]/2, -1)\n",
    "\n",
    "t2_line = interpline(t2, start_point=start_point, end_point=end_point)\n",
    "print(t2_line)\n",
    "\n",
    "del t2_line, t2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from wrf import getvar\n",
    "from netCDF4 import Dataset as nc\n",
    "lambertnc = nc(\"/Users/ladwig/Documents/wrf_files/wrfout_d01_2010-06-13_21:00:00\")\n",
    "uvmet = getvar([lambertnc,lambertnc], \"uvmet\", timeidx=None)\n",
    "uvmet10 = getvar([lambertnc,lambertnc], \"uvmet10\", timeidx=None)\n",
    "uv_wspd_wdir = getvar([lambertnc,lambertnc], \"wspd_wdir_uvmet\", timeidx=None)\n",
    "uv_wspd_wdir10 = getvar([lambertnc,lambertnc], \"wspd_wdir_uvmet10\", timeidx=None)\n",
    "wspd_wdir = getvar([lambertnc,lambertnc], \"wspd_wdir\", timeidx=None)\n",
    "wspd_wdir10 = getvar([lambertnc,lambertnc], \"wspd_wdir10\", timeidx=None)\n",
    "print (uvmet)\n",
    "print (uvmet10)\n",
    "print (uv_wspd_wdir)\n",
    "print (uv_wspd_wdir10)\n",
    "print (wspd_wdir)\n",
    "print (wspd_wdir10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from wrf import (ALL_TIMES, npvalues, Constants, getvar, extract_vars, destagger, \n",
    "                 interp1d, interp2dxy, interpz3d, \n",
    "                 slp, tk, td, rh, uvmet, smooth2d)\n",
    "from netCDF4 import Dataset as nc\n",
    "\n",
    "wrfnc = nc(\"/Users/ladwig/Documents/wrf_files/wrf_vortex_multi/wrfout_d01_2005-08-28_00:00:00\")\n",
    "timeidx = ALL_TIMES\n",
    "method = \"cat\"\n",
    "squeeze = True\n",
    "cache = None\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "\n",
    "varnames=(\"T\", \"P\", \"PB\", \"QVAPOR\", \"PH\", \"PHB\")\n",
    "ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache,\n",
    "                          meta=True)\n",
    "\n",
    "t = ncvars[\"T\"]\n",
    "p = ncvars[\"P\"]\n",
    "pb = ncvars[\"PB\"]\n",
    "ph = ncvars[\"PH\"]\n",
    "phb = ncvars[\"PHB\"]\n",
    "\n",
    "\n",
    "ncvars = extract_vars(wrfnc, timeidx, (\"QVAPOR\",), method, squeeze, cache,\n",
    "                          meta=False)\n",
    "qvapor = ncvars[\"QVAPOR\"]\n",
    "\n",
    "    \n",
    "full_t = t + Constants.T_BASE\n",
    "full_p = p + pb\n",
    "qvapor[qvapor < 0] = 0.\n",
    "    \n",
    "full_ph = (ph + phb) / Constants.G\n",
    "    \n",
    "destag_ph = destagger(full_ph, -3)\n",
    "    \n",
    "_tk = tk(full_p, full_t)\n",
    "_slp = slp(destag_ph, _tk, full_p, qvapor)\n",
    "_td = td(full_p, qvapor)\n",
    "_smooth2d = smooth2d(npvalues(_slp), 3)\n",
    "\n",
    "_uvmet = getvar(wrfnc, \"uvmet\", timeidx=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print (_tk)\n",
    "print (\"\\n\")\n",
    "print (_slp)\n",
    "print (\"\\n\")\n",
    "print (_td)\n",
    "print (\"\\n\")\n",
    "print (_smooth2d)\n",
    "print (_uvmet)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from wrf import (ALL_TIMES, npvalues, Constants, getvar, extract_vars, destagger, \n",
    "                 interp1d, interp2dxy, interpz3d, \n",
    "                 slp, tk, td, rh, uvmet, smooth2d, extract_global_attrs)\n",
    "from math import fabs, log, tan, sin, cos\n",
    "from wrf.util import either\n",
    "from netCDF4 import Dataset as nc\n",
    "\n",
    "wrfnc = nc(\"/Users/ladwig/Documents/wrf_files/wrfout_d01_2010-06-13_21:00:00\")\n",
    "wrfnc = [wrfnc, wrfnc, wrfnc]\n",
    "ten_m = True\n",
    "method = \"cat\"\n",
    "squeeze=True\n",
    "cache=None\n",
    "timeidx = ALL_TIMES\n",
    "\n",
    "if not ten_m:\n",
    "    varname = either(\"U\", \"UU\")(wrfnc)\n",
    "    u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache,\n",
    "                          meta=True)\n",
    "    \n",
    "    #renamed = u_vars[varname].rename({\"west_east_stag\" : \"test\"})\n",
    "    u = destagger(u_vars[varname], -1, meta=True)\n",
    "    #print (renamed)\n",
    "    #u = destagger(renamed, -1, meta=True)\n",
    "    #print (u)\n",
    "\n",
    "    varname = either(\"V\", \"VV\")(wrfnc)\n",
    "    v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache,\n",
    "                          meta=True)\n",
    "    v = destagger(v_vars[varname], -2)\n",
    "else:\n",
    "    varname = either(\"U10\", \"UU\")(wrfnc)\n",
    "    u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache,\n",
    "                          meta=True)\n",
    "    u = (u_vars[varname] if varname == \"U10\" else \n",
    "         destagger(u_vars[varname][...,0,:,:], -1)) \n",
    "\n",
    "    varname = either(\"V10\", \"VV\")(wrfnc)\n",
    "    v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache,\n",
    "                          meta=True)\n",
    "    v = (v_vars[varname] if varname == \"V10\" else \n",
    "         destagger(v_vars[varname][...,0,:,:], -2))\n",
    "\n",
    "map_proj_attrs = extract_global_attrs(wrfnc, attrs=\"MAP_PROJ\")\n",
    "map_proj = map_proj_attrs[\"MAP_PROJ\"]\n",
    "\n",
    "# 1 - Lambert\n",
    "# 2 - Polar Stereographic\n",
    "# 3 - Mercator\n",
    "# 6 - Lat/Lon\n",
    "# Note:  NCL has no code to handle other projections (0,4,5) as they \n",
    "# don't appear to be supported any longer\n",
    "\n",
    "if map_proj in (0,3,6):\n",
    "    # No rotation needed for Mercator and Lat/Lon, but still need\n",
    "    # u,v aggregated in to one array\n",
    "\n",
    "    end_idx = -3 if not ten_m else -2\n",
    "    resdim = (2,) + u.shape[0:end_idx] + u.shape[end_idx:]\n",
    "\n",
    "    # Make a new output array for the result\n",
    "    res = np.empty(resdim, u.dtype)\n",
    "\n",
    "    # For 2D array, this makes (0,...,:,:) and (1,...,:,:)\n",
    "    # For 3D array, this makes (0,...,:,:,:) and (1,...,:,:,:)\n",
    "    idx0 = (0,) + (Ellipsis,) + (slice(None),)*(-end_idx)\n",
    "    idx1 = (1,) + (Ellipsis,) + (slice(None),)*(-end_idx)\n",
    "\n",
    "    res[idx0] = u[:]\n",
    "    res[idx1] = v[:]\n",
    "\n",
    "    result = res\n",
    "elif map_proj in (1,2):\n",
    "    lat_attrs = extract_global_attrs(wrfnc, attrs=(\"TRUELAT1\",\n",
    "                                                   \"TRUELAT2\"))\n",
    "    radians_per_degree = Constants.PI/180.0\n",
    "    # Rotation needed for Lambert and Polar Stereographic\n",
    "    true_lat1 = lat_attrs[\"TRUELAT1\"]\n",
    "    true_lat2 = lat_attrs[\"TRUELAT2\"]\n",
    "\n",
    "    try:\n",
    "        lon_attrs = extract_global_attrs(wrfnc, attrs=\"STAND_LON\")\n",
    "    except AttributeError:\n",
    "        try:\n",
    "            cen_lon_attrs = extract_global_attrs(wrfnc, attrs=\"CEN_LON\")\n",
    "        except AttributeError:\n",
    "            raise RuntimeError(\"longitude attributes not found in NetCDF\")\n",
    "        else:\n",
    "            cen_lon = cen_lon_attrs[\"CEN_LON\"]\n",
    "    else:\n",
    "        cen_lon = lon_attrs[\"STAND_LON\"]\n",
    "\n",
    "\n",
    "    varname = either(\"XLAT_M\", \"XLAT\")(wrfnc)\n",
    "    xlat_var = extract_vars(wrfnc, timeidx, varname, \n",
    "                            method, squeeze, cache, meta=True)\n",
    "    lat = xlat_var[varname]\n",
    "\n",
    "    varname = either(\"XLONG_M\", \"XLONG\")(wrfnc)\n",
    "    xlon_var = extract_vars(wrfnc, timeidx, varname, \n",
    "                            method, squeeze, cache, meta=False)\n",
    "    lon = xlon_var[varname]\n",
    "\n",
    "    if map_proj == 1:\n",
    "        if((fabs(true_lat1 - true_lat2) > 0.1) and\n",
    "                (fabs(true_lat2 - 90.) > 0.1)): \n",
    "            cone = (log(cos(true_lat1*radians_per_degree)) \n",
    "                - log(cos(true_lat2*radians_per_degree)))\n",
    "            cone = (cone / \n",
    "                    (log(tan((45.-fabs(true_lat1/2.))*radians_per_degree)) \n",
    "                - log(tan((45.-fabs(true_lat2/2.))*radians_per_degree)))) \n",
    "        else:\n",
    "            cone = sin(fabs(true_lat1)*radians_per_degree)\n",
    "    else:\n",
    "        cone = 1\n",
    "    \n",
    "    result = uvmet(u, v, lat, lon, cen_lon, cone, meta=True)\n",
    "    \n",
    "print (result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from wrf import (ALL_TIMES, npvalues, Constants, getvar, extract_vars, destagger, \n",
    "                 interp1d, interp2dxy, interpz3d, \n",
    "                 slp, tk, td, rh, uvmet, smooth2d, extract_global_attrs, xy)\n",
    "from math import fabs, log, tan, sin, cos\n",
    "from wrf.util import either\n",
    "from netCDF4 import Dataset as nc\n",
    "\n",
    "timeidx = None\n",
    "wrfnc = nc(\"/Users/ladwig/Documents/wrf_files/wrfout_d01_2010-06-13_21:00:00\")\n",
    "\n",
    "wrfnc = [wrfnc, wrfnc]\n",
    "field3d = getvar(wrfnc, \"P\", timeidx, meta=True)\n",
    "\n",
    "z = getvar(wrfnc, \"z\", timeidx)\n",
    "\n",
    "interp = interpz3d(field3d, z, 500, missingval=Constants.DEFAULT_FILL, meta=True)\n",
    "print (interp)\n",
    "print (\"\\n\")\n",
    "\n",
    "# 2dxy test\n",
    "pivot_point = (field3d.shape[-1] / 2, field3d.shape[-2] / 2 ) \n",
    "angle = 90.0\n",
    "\n",
    "start_point = (0, z.shape[-2]/2)\n",
    "end_point = (-1, z.shape[-2]/2)\n",
    "\n",
    "#start_point = (0, 0)\n",
    "#end_point = (-1, -1)\n",
    "\n",
    "_xy = xy(field3d, pivot_point=pivot_point, angle=angle)\n",
    "print(_xy)\n",
    "print (\"\\n\")\n",
    "\n",
    "_xy = xy(field3d, start_point=start_point, end_point=end_point)\n",
    "print(_xy)\n",
    "print (\"\\n\")\n",
    "\n",
    "interpxy = interp2dxy(field3d, _xy)\n",
    "print (interpxy)\n",
    "print (\"\\n\")\n",
    "\n",
    "# 1D test\n",
    "only_height = [0]*field3d.ndim\n",
    "only_height[-3] = slice(None)\n",
    "only_height = (Ellipsis,) + tuple(only_height[-3:])\n",
    "\n",
    "v_in = field3d[only_height]\n",
    "print(\"v_in\", v_in)\n",
    "z_in = z[only_height]\n",
    "z_out = np.asarray([100.,200.,300.,500.,1000.,5000.], field3d.dtype)\n",
    "\n",
    "int1d = interp1d(v_in, z_in, z_out, missingval=Constants.DEFAULT_FILL, meta=True)\n",
    "print(int1d)\n",
    "print (\"\\n\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "\n",
    "import numpy as np\n",
    "\n",
    "from wrf import (ALL_TIMES, getvar, cape_2d, cape_3d, extract_vars, tk, destagger, ConversionFactors, Constants)\n",
    "from netCDF4 import Dataset as nc\n",
    "\n",
    "timeidx = ALL_TIMES\n",
    "method = \"cat\"\n",
    "squeeze=True\n",
    "cache=None\n",
    "timeidx = None\n",
    "\n",
    "wrfnc = nc(\"/Users/ladwig/Documents/wrf_files/wrfout_d01_2010-06-13_21:00:00\")\n",
    "\n",
    "wrfnc = [wrfnc, wrfnc]\n",
    "\n",
    "varnames = (\"T\", \"P\", \"PB\", \"QVAPOR\", \"PH\", \"PHB\", \"HGT\", \"PSFC\")\n",
    "ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache,\n",
    "                          meta=True)\n",
    "t = ncvars[\"T\"]\n",
    "p = ncvars[\"P\"]\n",
    "pb = ncvars[\"PB\"]\n",
    "qvapor = ncvars[\"QVAPOR\"]\n",
    "ph = ncvars[\"PH\"]\n",
    "phb = ncvars[\"PHB\"]\n",
    "ter = ncvars[\"HGT\"]\n",
    "psfc = ncvars[\"PSFC\"]\n",
    "    \n",
    "full_t = t + Constants.T_BASE\n",
    "full_p = p + pb\n",
    "tkel = tk(full_p, full_t)\n",
    "    \n",
    "geopt = ph + phb\n",
    "geopt_unstag = destagger(geopt, -3)\n",
    "height = geopt_unstag/Constants.G\n",
    "    \n",
    "# Convert pressure to hPa\n",
    "p_hpa = ConversionFactors.PA_TO_HPA * full_p\n",
    "psfc_hpa = ConversionFactors.PA_TO_HPA * psfc \n",
    "\n",
    "ter_follow = 1\n",
    "\n",
    "cape2d = cape_2d(p_hpa, tkel, qvapor, height, ter, psfc_hpa, ter_follow, \n",
    "            missing=Constants.DEFAULT_FILL, meta=True)\n",
    "\n",
    "cape3d = cape_3d(p_hpa, tkel, qvapor, height, ter, psfc_hpa, ter_follow, \n",
    "            missing=Constants.DEFAULT_FILL, meta=True)\n",
    "\n",
    "print (cape2d)\n",
    "print (cape3d)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}