preprocess_era5_to_compare_with_obs.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Preprocess ERA5 Reanalysis\n",
    "\n",
    "This notebook preprocesses ERA5 reanalysis so it's easier to compare against observations at the ENA site. This means:\n",
    "* renaming awkward variable names from the grib-to-netcdf conversion\n",
    "* calculating pressure at levels and interfaces\n",
    "* calculating geopotential height (needed because obs are in height coordinates)\n",
    "\n",
    "Contact: christopher.jones@pnnl.gov"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import xarray as xr\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# input_file_name = '/Users/jone003/tmp/EAGLES/ERA5_ENA_20170717-20170719_ml.nc'\n",
    "input_file_name = '/Users/jone003/tmp/EAGLES/ERA5_ENA_20180116-20180126_ml.nc'\n",
    "output_file_name = input_file_name.replace('_ml.nc', '_with_height_and_presl.nc')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load dataset and rename variables\n",
    "Note: ERA5 has already been converted from `grib` to `netcdf` format with `ncl_convert2nc`\n",
    "* Convert ERA5 mode levels to pressure and geopotential height\n",
    "* (optional) interpolate time and heights between ERA5 and obs for easier comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fix_dumb_names(ds, postfix='_P0_L105_GLL0', **kwargs):\n",
    "    \"\"\"removes suffixes from grib to netcdf conversion\n",
    "    optionally rename (orig: new) in kwargs\n",
    "    \"\"\"\n",
    "    to_rename = kwargs.copy()\n",
    "    for v in ds:\n",
    "        if postfix in v:\n",
    "            to_rename.update({v: v.replace(postfix, '')})\n",
    "            print(v, v.replace(postfix,''))\n",
    "    print(to_rename)\n",
    "    return ds.rename(to_rename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TMP_P0_L105_GLL0 TMP\n",
      "SPFH_P0_L105_GLL0 SPFH\n",
      "UGRD_P0_L105_GLL0 UGRD\n",
      "VGRD_P0_L105_GLL0 VGRD\n",
      "GP_P0_L105_GLL0 GP\n",
      "NLPRES_P0_L105_GLL0 NLPRES\n",
      "{'lat_0': 'lat', 'lon_0': 'lon', 'lv_HYBL0': 'lev', 'initial_time0_hours': 'time', 'TMP_P0_L105_GLL0': 'TMP', 'SPFH_P0_L105_GLL0': 'SPFH', 'UGRD_P0_L105_GLL0': 'UGRD', 'VGRD_P0_L105_GLL0': 'VGRD', 'GP_P0_L105_GLL0': 'GP', 'NLPRES_P0_L105_GLL0': 'NLPRES'}\n",
      "<xarray.Dataset>\n",
      "Dimensions:                (lat: 41, lev: 137, lon: 61, time: 264)\n",
      "Coordinates:\n",
      "  * time                   (time) datetime64[ns] 2018-01-16 ... 2018-01-26T23:00:00\n",
      "  * lat                    (lat) float32 45.0 44.75 44.5 ... 35.5 35.25 35.0\n",
      "  * lon                    (lon) float32 325.0 325.25 325.5 ... 339.75 340.0\n",
      "  * lev                    (lev) float32 1.0 2.0 3.0 4.0 ... 135.0 136.0 137.0\n",
      "Data variables:\n",
      "    TMP                    (time, lev, lat, lon) float32 ...\n",
      "    SPFH                   (time, lev, lat, lon) float32 ...\n",
      "    UGRD                   (time, lev, lat, lon) float32 ...\n",
      "    VGRD                   (time, lev, lat, lon) float32 ...\n",
      "    GP                     (time, lat, lon) float32 ...\n",
      "    NLPRES                 (time, lat, lon) float32 ...\n",
      "    initial_time0_encoded  (time) float64 ...\n",
      "    initial_time0          (time) |S18 ...\n",
      "Attributes:\n",
      "    creation_date:  Tue Feb 18 16:49:33 PST 2020\n",
      "    NCL_Version:    6.6.2\n",
      "    system:         Linux compy01.pnl.gov 3.10.0-957.5.1.el7.x86_64 #1 SMP Fr...\n",
      "    Conventions:    None\n",
      "    grib_source:    ERA5_ENA_20180116-20180126_ml.grb\n",
      "    title:          NCL: convert-GRIB-to-netCDF\n"
     ]
    }
   ],
   "source": [
    "ds_era = xr.open_dataset(input_file_name)\n",
    "ds_era = fix_dumb_names(ds_era, lat_0='lat', lon_0='lon', lv_HYBL0='lev', initial_time0_hours='time')\n",
    "print(ds_era)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convert from model level to pressure and geopotential height\n",
    "To convert from level to geopotential height, follow procedure outlined here\n",
    "         https://confluence.ecmwf.int/pages/viewpage.action?pageId=88261819\n",
    "\n",
    "The essence of the calculation is:\n",
    "    # compute the pressures (on half-levels)\n",
    "    ph_lev, ph_levplusone = get_ph_levs(values, lev)\n",
    " \n",
    "    if lev == 1:\n",
    "        dlog_p = np.log(ph_levplusone / 0.1)\n",
    "        alpha = np.log(2)\n",
    "    else:\n",
    "        dlog_p = np.log(ph_levplusone / ph_lev)\n",
    "        alpha = 1. - ((ph_lev / (ph_levplusone - ph_lev)) * dlog_p)\n",
    " \n",
    "    t_level = t_level * R_D\n",
    " \n",
    "    # z_f is the geopotential of this full level\n",
    "    # integrate from previous (lower) half-level z_h to the\n",
    "    # full level\n",
    "    z_f = z_h + (t_level * alpha)\n",
    " \n",
    "    # z_h is the geopotential of 'half-levels'\n",
    "    # integrate z_h to next half level\n",
    "    z_h = z_h + (t_level * dlog_p)\n",
    " \n",
    "    return z_h, z_f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# source: https://www.ecmwf.int/en/forecasts/documentation-and-support/137-model-levels\n",
    "model_level_table = \"\"\"\n",
    "0\t0.000000\t0.000000\t0.000000\t-  \t-\t-\t-\t-\n",
    "1\t2.000365\t0.000000\t0.0200\t0.0100\t79301.79\t80301.65\t198.05\t0.000018\n",
    "2\t3.102241\t0.000000\t0.0310\t0.0255\t73721.58\t74584.91\t209.21\t0.000042\n",
    "3\t4.666084\t0.000000\t0.0467\t0.0388\t71115.75\t71918.79\t214.42\t0.000063\n",
    "4\t6.827977\t0.000000\t0.0683\t0.0575\t68618.43\t69365.77\t221.32\t0.000090\n",
    "5\t9.746966\t0.000000\t0.0975\t0.0829\t66210.99\t66906.53\t228.06\t0.000127\n",
    "6\t13.605424\t0.000000\t0.1361\t0.1168\t63890.03\t64537.43\t234.56\t0.000173\n",
    "7\t18.608931\t0.000000\t0.1861\t0.1611\t61651.77\t62254.39\t240.83\t0.000233\n",
    "8\t24.985718\t0.000000\t0.2499\t0.2180\t59492.50\t60053.46\t246.87\t0.000308\n",
    "9\t32.985710\t0.000000\t0.3299\t0.2899\t57408.61\t57930.78\t252.71\t0.000400\n",
    "10\t42.879242\t0.000000\t0.4288\t0.3793\t55396.62\t55882.68\t258.34\t0.000512\n",
    "11\t54.955463\t0.000000\t0.5496\t0.4892\t53453.20\t53905.62\t263.78\t0.000646\n",
    "12\t69.520576\t0.000000\t0.6952\t0.6224\t51575.15\t51996.21\t269.04\t0.000806\n",
    "13\t86.895882\t0.000000\t0.8690\t0.7821\t49767.41\t50159.36\t270.65\t0.001007\n",
    "14\t107.415741\t0.000000\t1.0742\t0.9716\t48048.70\t48413.94\t270.65\t0.001251\n",
    "15\t131.425507\t0.000000\t1.3143\t1.1942\t46416.22\t46756.98\t269.02\t0.001546\n",
    "16\t159.279404\t0.000000\t1.5928\t1.4535\t44881.17\t45199.69\t264.72\t0.001913\n",
    "17\t191.338562\t0.000000\t1.9134\t1.7531\t43440.23\t43738.55\t260.68\t0.002343\n",
    "18\t227.968948\t0.000000\t2.2797\t2.0965\t42085.00\t42364.93\t256.89\t0.002843\n",
    "19\t269.539581\t0.000000\t2.6954\t2.4875\t40808.05\t41071.20\t253.31\t0.003421\n",
    "20\t316.420746\t0.000000\t3.1642\t2.9298\t39602.76\t39850.56\t249.94\t0.004084\n",
    "21\t368.982361\t0.000000\t3.6898\t3.4270\t38463.25\t38696.94\t246.75\t0.004838\n",
    "22\t427.592499\t0.000000\t4.2759\t3.9829\t37384.22\t37604.95\t243.73\t0.005693\n",
    "23\t492.616028\t0.000000\t4.9262\t4.6010\t36360.94\t36569.72\t240.86\t0.006655\n",
    "24\t564.413452\t0.000000\t5.6441\t5.2851\t35389.15\t35586.89\t238.14\t0.007731\n",
    "25\t643.339905\t0.000000\t6.4334\t6.0388\t34465.00\t34652.52\t235.55\t0.008931\n",
    "26\t729.744141\t0.000000\t7.2974\t6.8654\t33585.02\t33763.05\t233.09\t0.010261\n",
    "27\t823.967834\t0.000000\t8.2397\t7.7686\t32746.04\t32915.27\t230.74\t0.011729\n",
    "28\t926.344910\t0.000000\t9.2634\t8.7516\t31945.53\t32106.57\t228.60\t0.013337\n",
    "29\t1037.201172\t0.000000\t10.3720\t9.8177\t31177.59\t31330.96\t227.83\t0.015012\n",
    "30\t1156.853638\t0.000000\t11.5685\t10.9703\t30438.54\t30584.71\t227.09\t0.016829\n",
    "31\t1285.610352\t0.000000\t12.8561\t12.2123\t29726.69\t29866.09\t226.38\t0.018793\n",
    "32\t1423.770142\t0.000000\t14.2377\t13.5469\t29040.48\t29173.50\t225.69\t0.020910\n",
    "33\t1571.622925\t0.000000\t15.7162\t14.9770\t28378.46\t28505.47\t225.03\t0.023186\n",
    "34\t1729.448975\t0.000000\t17.2945\t16.5054\t27739.29\t27860.64\t224.39\t0.025624\n",
    "35\t1897.519287\t0.000000\t18.9752\t18.1348\t27121.74\t27237.73\t223.77\t0.028232\n",
    "36\t2076.095947\t0.000000\t20.7610\t19.8681\t26524.63\t26635.56\t223.17\t0.031013\n",
    "37\t2265.431641\t0.000000\t22.6543\t21.7076\t25946.90\t26053.04\t222.60\t0.033972\n",
    "38\t2465.770508\t0.000000\t24.6577\t23.6560\t25387.55\t25489.15\t222.04\t0.037115\n",
    "39\t2677.348145\t0.000000\t26.7735\t25.7156\t24845.63\t24942.93\t221.50\t0.040445\n",
    "40\t2900.391357\t0.000000\t29.0039\t27.8887\t24320.28\t24413.50\t220.97\t0.043967\n",
    "41\t3135.119385\t0.000000\t31.3512\t30.1776\t23810.67\t23900.02\t220.46\t0.047685\n",
    "42\t3381.743652\t0.000000\t33.8174\t32.5843\t23316.04\t23401.71\t219.97\t0.051604\n",
    "43\t3640.468262\t0.000000\t36.4047\t35.1111\t22835.68\t22917.85\t219.49\t0.055727\n",
    "44\t3911.490479\t0.000000\t39.1149\t37.7598\t22368.91\t22447.75\t219.02\t0.060059\n",
    "45\t4194.930664\t0.000000\t41.9493\t40.5321\t21915.16\t21990.82\t218.57\t0.064602\n",
    "46\t4490.817383\t0.000000\t44.9082\t43.4287\t21473.98\t21546.62\t218.12\t0.069359\n",
    "47\t4799.149414\t0.000000\t47.9915\t46.4498\t21045.00\t21114.77\t217.70\t0.074330\n",
    "48\t5119.895020\t0.000000\t51.1990\t49.5952\t20627.87\t20694.90\t217.28\t0.079516\n",
    "49\t5452.990723\t0.000000\t54.5299\t52.8644\t20222.24\t20286.66\t216.87\t0.084916\n",
    "50\t5798.344727\t0.000000\t57.9834\t56.2567\t19827.95\t19889.88\t216.65\t0.090458\n",
    "51\t6156.074219\t0.000000\t61.5607\t59.7721\t19443.55\t19503.09\t216.65\t0.096110\n",
    "52\t6526.946777\t0.000000\t65.2695\t63.4151\t19068.35\t19125.61\t216.65\t0.101968\n",
    "53\t6911.870605\t0.000000\t69.1187\t67.1941\t18701.27\t18756.34\t216.65\t0.108045\n",
    "54\t7311.869141\t0.000000\t73.1187\t71.1187\t18341.27\t18394.25\t216.65\t0.114355\n",
    "55\t7727.412109\t0.000007\t77.2810\t75.1999\t17987.41\t18038.35\t216.65\t0.120917\n",
    "56\t8159.354004\t0.000024\t81.6182\t79.4496\t17638.78\t17687.77\t216.65\t0.127751\n",
    "57\t8608.525391\t0.000059\t86.1450\t83.8816\t17294.53\t17341.62\t216.65\t0.134877\n",
    "58\t9076.400391\t0.000112\t90.8774\t88.5112\t16953.83\t16999.08\t216.65\t0.142321\n",
    "59\t9562.682617\t0.000199\t95.8280\t93.3527\t16616.09\t16659.55\t216.65\t0.150106\n",
    "60\t10065.978516\t0.000340\t101.0047\t98.4164\t16281.10\t16322.83\t216.65\t0.158248\n",
    "61\t10584.631836\t0.000562\t106.4153\t103.7100\t15948.85\t15988.88\t216.65\t0.166760\n",
    "62\t11116.662109\t0.000890\t112.0681\t109.2417\t15619.30\t15657.70\t216.65\t0.175655\n",
    "63\t11660.067383\t0.001353\t117.9714\t115.0198\t15292.44\t15329.24\t216.65\t0.184946\n",
    "64\t12211.547852\t0.001992\t124.1337\t121.0526\t14968.24\t15003.50\t216.65\t0.194646\n",
    "65\t12766.873047\t0.002857\t130.5637\t127.3487\t14646.68\t14680.44\t216.65\t0.204770\n",
    "66\t13324.668945\t0.003971\t137.2703\t133.9170\t14327.75\t14360.05\t216.65\t0.215331\n",
    "67\t13881.331055\t0.005378\t144.2624\t140.7663\t14011.41\t14042.30\t216.65\t0.226345\n",
    "68\t14432.139648\t0.007133\t151.5493\t147.9058\t13697.65\t13727.18\t216.65\t0.237825\n",
    "69\t14975.615234\t0.009261\t159.1403\t155.3448\t13386.45\t13414.65\t216.65\t0.249786\n",
    "70\t15508.256836\t0.011806\t167.0450\t163.0927\t13077.79\t13104.70\t216.65\t0.262244\n",
    "71\t16026.115234\t0.014816\t175.2731\t171.1591\t12771.64\t12797.30\t216.65\t0.275215\n",
    "72\t16527.322266\t0.018318\t183.8344\t179.5537\t12467.99\t12492.44\t216.65\t0.288713\n",
    "73\t17008.789063\t0.022355\t192.7389\t188.2867\t12166.81\t12190.10\t216.65\t0.302755\n",
    "74\t17467.613281\t0.026964\t201.9969\t197.3679\t11868.08\t11890.24\t216.65\t0.317357\n",
    "75\t17901.621094\t0.032176\t211.6186\t206.8078\t11571.79\t11592.86\t216.65\t0.332536\n",
    "76\t18308.433594\t0.038026\t221.6146\t216.6166\t11277.92\t11297.93\t216.65\t0.348308\n",
    "77\t18685.718750\t0.044548\t231.9954\t226.8050\t10986.70\t11005.69\t216.74\t0.364545\n",
    "78\t19031.289063\t0.051773\t242.7719\t237.3837\t10696.22\t10714.22\t218.62\t0.378253\n",
    "79\t19343.511719\t0.059728\t253.9549\t248.3634\t10405.61\t10422.64\t220.51\t0.392358\n",
    "80\t19620.042969\t0.068448\t265.5556\t259.7553\t10114.89\t10130.98\t222.40\t0.406868\n",
    "81\t19859.390625\t0.077958\t277.5852\t271.5704\t9824.08\t9839.26\t224.29\t0.421790\n",
    "82\t20059.931641\t0.088286\t290.0548\t283.8200\t9533.20\t9547.49\t226.18\t0.437130\n",
    "83\t20219.664063\t0.099462\t302.9762\t296.5155\t9242.26\t9255.70\t228.08\t0.452897\n",
    "84\t20337.863281\t0.111505\t316.3607\t309.6684\t8951.30\t8963.90\t229.97\t0.469097\n",
    "85\t20412.308594\t0.124448\t330.2202\t323.2904\t8660.32\t8672.11\t231.86\t0.485737\n",
    "86\t20442.078125\t0.138313\t344.5663\t337.3932\t8369.35\t8380.36\t233.75\t0.502825\n",
    "87\t20425.718750\t0.153125\t359.4111\t351.9887\t8078.41\t8088.67\t235.64\t0.520367\n",
    "88\t20361.816406\t0.168910\t374.7666\t367.0889\t7787.51\t7797.04\t237.53\t0.538370\n",
    "89\t20249.511719\t0.185689\t390.6450\t382.7058\t7496.68\t7505.51\t239.42\t0.556842\n",
    "90\t20087.085938\t0.203491\t407.0583\t398.8516\t7205.93\t7214.09\t241.31\t0.575790\n",
    "91\t19874.025391\t0.222333\t424.0190\t415.5387\t6915.29\t6922.80\t243.20\t0.595219\n",
    "92\t19608.572266\t0.242244\t441.5395\t432.7792\t6624.76\t6631.66\t245.09\t0.615138\n",
    "93\t19290.226563\t0.263242\t459.6321\t450.5858\t6334.38\t6340.68\t246.98\t0.635553\n",
    "94\t18917.460938\t0.285354\t478.3096\t468.9708\t6044.15\t6049.89\t248.86\t0.656471\n",
    "95\t18489.707031\t0.308598\t497.5845\t487.9470\t5754.10\t5759.30\t250.75\t0.677899\n",
    "96\t18006.925781\t0.332939\t517.4198\t507.5021\t5464.60\t5469.30\t252.63\t0.699815\n",
    "97\t17471.839844\t0.358254\t537.7195\t527.5696\t5176.77\t5180.98\t254.50\t0.722139\n",
    "98\t16888.687500\t0.384363\t558.3430\t548.0312\t4892.26\t4896.02\t256.35\t0.744735\n",
    "99\t16262.046875\t0.411125\t579.1926\t568.7678\t4612.58\t4615.92\t258.17\t0.767472\n",
    "100\t15596.695313\t0.438391\t600.1668\t589.6797\t4338.77\t4341.73\t259.95\t0.790242\n",
    "101\t14898.453125\t0.466003\t621.1624\t610.6646\t4071.80\t4074.41\t261.68\t0.812937\n",
    "102\t14173.324219\t0.493800\t642.0764\t631.6194\t3812.53\t3814.82\t263.37\t0.835453\n",
    "103\t13427.769531\t0.521619\t662.8084\t652.4424\t3561.70\t3563.69\t265.00\t0.857686\n",
    "104\t12668.257813\t0.549301\t683.2620\t673.0352\t3319.94\t3321.67\t266.57\t0.879541\n",
    "105\t11901.339844\t0.576692\t703.3467\t693.3043\t3087.75\t3089.25\t268.08\t0.900929\n",
    "106\t11133.304688\t0.603648\t722.9795\t713.1631\t2865.54\t2866.83\t269.52\t0.921768\n",
    "107\t10370.175781\t0.630036\t742.0855\t732.5325\t2653.58\t2654.69\t270.90\t0.941988\n",
    "108\t9617.515625\t0.655736\t760.5996\t751.3426\t2452.04\t2452.99\t272.21\t0.961527\n",
    "109\t8880.453125\t0.680643\t778.4661\t769.5329\t2260.99\t2261.80\t273.45\t0.980334\n",
    "110\t8163.375000\t0.704669\t795.6396\t787.0528\t2080.41\t2081.09\t274.63\t0.998368\n",
    "111\t7470.343750\t0.727739\t812.0847\t803.8622\t1910.19\t1910.76\t275.73\t1.015598\n",
    "112\t6804.421875\t0.749797\t827.7756\t819.9302\t1750.14\t1750.63\t276.77\t1.032005\n",
    "113\t6168.531250\t0.770798\t842.6959\t835.2358\t1600.04\t1600.44\t277.75\t1.047576\n",
    "114\t5564.382813\t0.790717\t856.8376\t849.7668\t1459.58\t1459.91\t278.66\t1.062310\n",
    "115\t4993.796875\t0.809536\t870.2004\t863.5190\t1328.43\t1328.70\t279.52\t1.076209\n",
    "116\t4457.375000\t0.827256\t882.7910\t876.4957\t1206.21\t1206.44\t280.31\t1.089286\n",
    "117\t3955.960938\t0.843881\t894.6222\t888.7066\t1092.54\t1092.73\t281.05\t1.101558\n",
    "118\t3489.234375\t0.859432\t905.7116\t900.1669\t987.00\t987.15\t281.73\t1.113047\n",
    "119\t3057.265625\t0.873929\t916.0815\t910.8965\t889.17\t889.29\t282.37\t1.123777\n",
    "120\t2659.140625\t0.887408\t925.7571\t920.9193\t798.62\t798.72\t282.96\t1.133779\n",
    "121\t2294.242188\t0.899900\t934.7666\t930.2618\t714.94\t715.02\t283.50\t1.143084\n",
    "122\t1961.500000\t0.911448\t943.1399\t938.9532\t637.70\t637.76\t284.00\t1.151724\n",
    "123\t1659.476563\t0.922096\t950.9082\t947.0240\t566.49\t566.54\t284.47\t1.159733\n",
    "124\t1387.546875\t0.931881\t958.1037\t954.5059\t500.91\t500.95\t284.89\t1.167147\n",
    "125\t1143.250000\t0.940860\t964.7584\t961.4311\t440.58\t440.61\t285.29\t1.173999\n",
    "126\t926.507813\t0.949064\t970.9046\t967.8315\t385.14\t385.16\t285.65\t1.180323\n",
    "127\t734.992188\t0.956550\t976.5737\t973.7392\t334.22\t334.24\t285.98\t1.186154\n",
    "128\t568.062500\t0.963352\t981.7968\t979.1852\t287.51\t287.52\t286.28\t1.191523\n",
    "129\t424.414063\t0.969513\t986.6036\t984.2002\t244.68\t244.69\t286.56\t1.196462\n",
    "130\t302.476563\t0.975078\t991.0230\t988.8133\t205.44\t205.44\t286.81\t1.201001\n",
    "131\t202.484375\t0.980072\t995.0824\t993.0527\t169.50\t169.51\t287.05\t1.205168\n",
    "132\t122.101563\t0.984542\t998.8081\t996.9452\t136.62\t136.62\t287.26\t1.208992\n",
    "133\t62.781250\t0.988500\t1002.2250\t1000.5165\t106.54\t106.54\t287.46\t1.212498\n",
    "134\t22.835938\t0.991984\t1005.3562\t1003.7906\t79.04\t79.04\t287.64\t1.215710\n",
    "135\t3.757813\t0.995003\t1008.2239\t1006.7900\t53.92\t53.92\t287.80\t1.218650\n",
    "136\t0.000000\t0.997630\t1010.8487\t1009.5363\t30.96\t30.96\t287.95\t1.221341\n",
    "137\t0.000000\t1.000000\t1013.2500\t1012.0494\t10.00\t10.00\t288.09\t1.223803\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# extract a and b coefficients to convert from model levels to pressure\n",
    "nlevel = []\n",
    "acoef = []\n",
    "bcoef = []\n",
    "for row in model_level_table.split('\\n'):\n",
    "    if row:\n",
    "        vals = row.split('\\t')[:3]\n",
    "        nlevel.append(int(vals[0]))\n",
    "        acoef.append(float(vals[1]))\n",
    "        bcoef.append(float(vals[2]))\n",
    "\n",
    "nlev = np.array(nlevel)\n",
    "ac = np.array(acoef)  # at half levels (interfaces)\n",
    "bc = np.array(bcoef)  # at half levels (interfaces)\n",
    "\n",
    "R_D = 287.06\n",
    "R_G = 9.80665\n",
    "\n",
    "da_hyba = xr.DataArray(data=ac, dims=['ilev'], coords=[nlev])\n",
    "da_hybb = xr.DataArray(data=bc, dims=['ilev'], coords=[nlev])\n",
    "presi = da_hyba + da_hybb * np.exp(ds_era['NLPRES'])  # pressure at interfaces\n",
    "# moist temperature:\n",
    "moist_temp = ds_era['TMP'] * (1 + 0.609133 * ds_era['SPFH'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_dlogp(presi):\n",
    "    \"\"\"calculate dlogp and alpha on interfaces (ilev)\"\"\"\n",
    "    axis = presi.dims.index('ilev')  # should be 0\n",
    "    if axis != 0:\n",
    "        return 'wrong axis!'\n",
    "    arr = presi.values\n",
    "    \n",
    "    ph_lev = arr[0:-1, ...]\n",
    "    ph_levplusone = arr[1:, ...]\n",
    "    \n",
    "    # rows 1:end\n",
    "    dlogp = np.log(ph_levplusone / ph_lev)\n",
    "    dlogp[0, ...] = np.log(arr[1, ...] / 0.1)\n",
    "    \n",
    "    alpha = 1 - (ph_lev / (ph_levplusone - ph_lev)) * dlogp\n",
    "    alpha[0, ...] = np.log(2)\n",
    "    \n",
    "    return dlogp, alpha\n",
    "\n",
    "def calc_geopotential(presi, moist_temp, z0):\n",
    "    \"\"\"Returns geopotential height at interfaces (gzi), model levels (gz), and height dataArray (z)\n",
    "    \"\"\"\n",
    "    R_D = 287.06\n",
    "    R_G = 9.80665\n",
    "    rhs = R_D * moist_temp.transpose('lev', 'time', 'lat', 'lon').values\n",
    "    \n",
    "    dlogp, alpha = calc_dlogp(presi)\n",
    "    gzi = np.zeros_like(presi.values)\n",
    "    gzi[0, ...] = z0.values.copy()\n",
    "    gzi[1:, ...] = gzi[0, ...] + np.cumsum((rhs * dlogp)[::-1, ...], axis=0)  # reversed ...\n",
    "    \n",
    "    gzi = gzi[::-1, ...]\n",
    "    gz = np.zeros_like(rhs)\n",
    "    gz = gzi[1:, ...] + (rhs * alpha)\n",
    "    \n",
    "    the_dims=['lev', 'time', 'lat', 'lon']\n",
    "    z = xr.DataArray(gz / (R_G * 1000), coords=[moist_temp[d] for d in the_dims], \n",
    "                     dims=the_dims,\n",
    "                     name='height',\n",
    "                     attrs={'units': 'km', 'long_name': 'height above sea level'})\n",
    "    return gzi, gz, z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/anaconda3/lib/python3.6/site-packages/ipykernel_launcher.py:12: RuntimeWarning: divide by zero encountered in true_divide\n",
      "  if sys.path[0] == '':\n"
     ]
    }
   ],
   "source": [
    "# calculate height and add to ds_era dataset\n",
    "gzi, gz, z = calc_geopotential(presi, moist_temp, ds_era['GP'])\n",
    "ds_era['height'] = z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<xarray.Dataset>\n",
       "Dimensions:                (lat: 41, lev: 137, lon: 61, time: 264)\n",
       "Coordinates:\n",
       "  * time                   (time) datetime64[ns] 2018-01-16 ... 2018-01-26T23:00:00\n",
       "  * lat                    (lat) float32 45.0 44.75 44.5 ... 35.5 35.25 35.0\n",
       "  * lon                    (lon) float32 325.0 325.25 325.5 ... 339.75 340.0\n",
       "  * lev                    (lev) float32 1.0 2.0 3.0 4.0 ... 135.0 136.0 137.0\n",
       "Data variables:\n",
       "    TMP                    (time, lev, lat, lon) float32 ...\n",
       "    SPFH                   (time, lev, lat, lon) float32 1.9870063e-06 ... 0.0055470048\n",
       "    UGRD                   (time, lev, lat, lon) float32 ...\n",
       "    VGRD                   (time, lev, lat, lon) float32 ...\n",
       "    GP                     (time, lat, lon) float32 -6.565628 ... -2.815628\n",
       "    NLPRES                 (time, lat, lon) float32 11.549789 ... 11.550724\n",
       "    initial_time0_encoded  (time) float64 ...\n",
       "    initial_time0          (time) |S18 ...\n",
       "    height                 (lev, time, lat, lon) float64 77.41 ... 0.009778\n",
       "    pressure               (lev, time, lat, lon) float64 1.0 1.0 ... 1.037e+05\n",
       "Attributes:\n",
       "    creation_date:  Tue Feb 18 16:49:33 PST 2020\n",
       "    NCL_Version:    6.6.2\n",
       "    system:         Linux compy01.pnl.gov 3.10.0-957.5.1.el7.x86_64 #1 SMP Fr...\n",
       "    Conventions:    None\n",
       "    grib_source:    ERA5_ENA_20180116-20180126_ml.grb\n",
       "    title:          NCL: convert-GRIB-to-netCDF"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# add pressure levels to the dataset\n",
    "the_dims = ['lev', 'time', 'lat', 'lon']\n",
    "\n",
    "# pressure levels are just taken as linear midpoint between interface pressures\n",
    "presl_vals = (presi[:-1, ...].values + presi[1:, ...].values) / 2\n",
    "presl = xr.DataArray(presl_vals, coords=[ds_era[d] for d in the_dims],\n",
    "                     dims=the_dims, name='pressure',\n",
    "                     attrs={'units': 'Pa', 'long_name': 'pressure on model levels'})\n",
    "ds_era['pressure'] = presl\n",
    "ds_era  # print to screen to verify contents"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optionally save the preprocessed file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "save_preprocessed_era_ds = True\n",
    "encoding = {'zlib': True, 'complevel': 6}\n",
    "if save_preprocessed_era_ds:\n",
    "    ds_era.to_netcdf(output_file_name,\n",
    "                     encoding={v: encoding for v in ds_era})"
   ]
  }
 ],
 "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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}