transport_flux_calculation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c032281",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\"\"\" This code generates transport flux for different wind components i.e. c,w and wspd\n",
    "\t=>> written by Mukesh Rai,, 2021/11/24 \"\"\"\n",
    "\n",
    "#=== load library\n",
    "from netCDF4 import Dataset\n",
    "from mpl_toolkits.basemap import Basemap\n",
    "import numpy as np\n",
    "import wrf\n",
    "from wrf import (getvar, to_np, get_cartopy,get_basemap, latlon_coords, vertcross,ALL_TIMES,cartopy_xlim, cartopy_ylim, interpline, CoordPair)\n",
    "from matplotlib.cm import get_cmap\n",
    "import matplotlib.pyplot as plt\n",
    "import cmaps\n",
    "from pylab import *\n",
    "import scipy as sp\n",
    "import scipy.ndimage\n",
    "import matplotlib.gridspec as gridspec\n",
    "import cartopy.crs as crs\n",
    "from cartopy.feature import NaturalEarthFeature\n",
    "from matplotlib.patches import FancyArrowPatch\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "from matplotlib import rcParams\n",
    "\n",
    "\"\"\"Control font\"\"\"\n",
    "matplotlib.rcParams['font.family'] = \"sans-serif\"\n",
    "matplotlib.rcParams['font.sans-serif'] = \"Times New Roman\"\n",
    "\n",
    "#===== load files =========================================\n",
    "data = [Dataset(\"wrfout_d01_2017-01-01_00_00_00\"),\n",
    "       Dataset(\"wrfout_d01_2017-02-01_00_00_00\"),\n",
    "       Dataset(\"wrfout_d01_2017-12-01_00_00_00\"),\n",
    "       Dataset(\"wrfout_d01_2017-03-01_00_00_00.nc\"),\n",
    "       Dataset(\"wrfout_d01_2017-04-01_00_00_00.nc\"),\n",
    "       Dataset(\"wrfout_d01_2017-05-01_00_00_00.nc\"),\n",
    "       Dataset(\"wrfout_d01_2017-06-01_00_00_00\"),\n",
    "       Dataset(\"wrfout_d01_2017-07-01_00_00_00\"),\n",
    "       Dataset(\"wrfout_d01_2017-08-01_00_00_00\"),\n",
    "       Dataset(\"wrfout_d01_2017-09-01_00_00_00\"),\n",
    "       Dataset(\"wrfout_d01_2017-10-01_00_00_00\"),\n",
    "       Dataset(\"wrfout_d01_2017-11-01_00_00_00\")]\n",
    "\n",
    "#====== get variables =========================================\n",
    "b1  = wrf.getvar(data, 'BC1', timeidx=ALL_TIMES)[:,:,:]    #bc1\n",
    "b2  = wrf.getvar(data, 'BC2', timeidx=ALL_TIMES)[:,:,:]    #bc2\n",
    "v   = wrf.getvar(data, 'va',  timeidx=150)[:,:,:]  # meridional wind\n",
    "w   = wrf.getvar(data, 'wa',  timeidx=200)[:,:,:]  # vertical wind\n",
    "ws  = wrf.getvar(data, 'wspd',timeidx=150)[:,:,:]  # wind speed\n",
    "\n",
    "#===== time dimesion average out ====================\n",
    "bc1 = np.average(b1,axis=0)\n",
    "bc2 = np.average(b2,axis=0)\n",
    "bc  = bc1+bc2 \n",
    "\n",
    "\"\"\"=== BC Transport Flux ==>> (tF)======\"\"\"\n",
    "tf1 = bc * v  # TF toward TP \n",
    "tf2 = bc * w  # TF  vertical\n",
    "tf3 = bc * ws # TG due to wind speed\n",
    "\n",
    "\"\"\"======== Now interpolation \"\"\"\n",
    "z1 = 5**(tf1/5.)\n",
    "z2 = 5**(tf2/5.)\n",
    "z3 = 5**(tf3/5.)\n",
    "\n",
    "\"\"\"========  Define the cross section start and end points=========\"\"\"\n",
    "\n",
    "cross_start = CoordPair(lat=22, lon=70)\n",
    "cross_end   = CoordPair(lat=36, lon=119)\n",
    "\n",
    "\"\"\"===== Get terrain height ============================\"\"\"\n",
    "ter      = getvar(data, \"ter\", units=\"km\")\n",
    "\"\"\"============= Get model height for mass grid ==========================================\"\"\"\n",
    "ht  = getvar(data, \"height_agl\", units=\"km\")[:,:,:]\n",
    "\"\"\"====  Get the terrain heights along the cross section line =============================\"\"\"\n",
    "ter_line = interpline(ter, wrfin=data, start_point=cross_start,end_point=cross_end)\n",
    "\"\"\"Compute the vertical cross-section interpolation.\n",
    "   Also, include the lat/lon points along the cross-section in the metadata by setting latlon to True\"\"\"\n",
    "z1_cross = vertcross(z1, ht, wrfin=data,start_point=cross_start,end_point=cross_end,latlon=True, meta=True)\n",
    "z2_cross = vertcross(z2, ht, wrfin=data,start_point=cross_start,end_point=cross_end,latlon=True, meta=True)\n",
    "z3_cross = vertcross(z3, ht, wrfin=data,start_point=cross_start,end_point=cross_end,latlon=True, meta=True)\n",
    "\n",
    "#== for wind component =====================================================================================\n",
    "v1_cross  = vertcross(v, ht, wrfin=data,start_point=cross_start,end_point=cross_end,latlon=True, meta=True)\n",
    "w1_cross  = vertcross(w, ht, wrfin=data,start_point=cross_start,end_point=cross_end,latlon=True, meta=True)\n",
    "ws1_cross = vertcross(ws, ht, wrfin=data,start_point=cross_start,end_point=cross_end,latlon=True, meta=True)\n",
    "\n",
    "\"\"\"===== Convert back to dBz after interpolation ==============================\"\"\"\n",
    "z1_cross  = 5 * np.log10(z1_cross)\n",
    "z2_cross  = 5 * np.log10(z2_cross)\n",
    "z3_cross  = 5 * np.log10(z3_cross)\n",
    "\n",
    "#=== wind compoent ====================\n",
    "v1_cross   = 5 * np.log10(v1_cross)\n",
    "w1_cross   = 5 * np.log10(w1_cross)\n",
    "ws1_cross  = 5 * np.log10(ws1_cross)\n",
    "\n",
    "\n",
    "\"\"\"Make a copy of the z cross data. Let's use regular numpy arrays for this\"\"\"\n",
    "xs1 = np.arange(0, z1_cross.shape[-1], 1)\n",
    "ys1 = to_np(z1_cross.coords[\"vertical\"])\n",
    "xs2 = np.arange(0, z2_cross.shape[-1], 1)\n",
    "ys2 = to_np(z2_cross.coords[\"vertical\"])\n",
    "xs3 = np.arange(0, z3_cross.shape[-1], 1)\n",
    "ys3 = to_np(z3_cross.coords[\"vertical\"])\n",
    "\n",
    "#== For wind =================================\n",
    "xs4 = np.arange(0, v1_cross.shape[-1], 1)\n",
    "ys4 = to_np(v1_cross.coords[\"vertical\"])\n",
    "xs5 = np.arange(0, w1_cross.shape[-1], 1)\n",
    "ys5 = to_np(w1_cross.coords[\"vertical\"])\n",
    "xs6 = np.arange(0, ws1_cross.shape[-1], 1)\n",
    "ys6 = to_np(ws1_cross.coords[\"vertical\"])\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\"\"\"Make a copy of the z cross data. Let's use regular numpy arrays for this\"\"\"\n",
    "z1_cross_filled = np.ma.copy(to_np(z1_cross))\n",
    "z2_cross_filled = np.ma.copy(to_np(z2_cross))\n",
    "z3_cross_filled = np.ma.copy(to_np(z3_cross))\n",
    "\n",
    "#====wind ====================================\n",
    "v1_cross_filled  = np.ma.copy(to_np(v1_cross))\n",
    "w1_cross_filled  = np.ma.copy(to_np(w1_cross))\n",
    "ws1_cross_filled = np.ma.copy(to_np(ws1_cross))\n",
    "\n",
    "\n",
    "\"\"\"=========================================================================\n",
    "   For each cross section column, find the first index with non-missing\n",
    "    values and copy these to the missing elements below\n",
    "   =========================================================================\"\"\"\n",
    "for i in range(z1_cross_filled.shape[-1]):\n",
    "    column_vals = z1_cross_filled[:,i]\n",
    "    first_idx = int(np.transpose((column_vals > -200).nonzero())[0])\n",
    "    z1_cross_filled[0:first_idx, i] = z1_cross_filled[first_idx, i]\n",
    "\n",
    "for i in range(z2_cross_filled.shape[-1]):\n",
    "    column_vals = z2_cross_filled[:,i]\n",
    "    first_idx = int(np.transpose((column_vals > -200).nonzero())[0])\n",
    "    z2_cross_filled[0:first_idx, i] = z2_cross_filled[first_idx, i]\n",
    "\n",
    "for i in range(z3_cross_filled.shape[-1]):\n",
    "    column_vals = z3_cross_filled[:,i]\n",
    "    first_idx = int(np.transpose((column_vals > -200).nonzero())[0])\n",
    "    z3_cross_filled[0:first_idx, i] = z3_cross_filled[first_idx, i]\n",
    "\n",
    "\n",
    "\n",
    "\"\"\"======Plot data ==================================================\"\"\"\n",
    "\n",
    "fig, (ax1,ax2,ax3) = plt.subplots(1, 3,figsize=(9,3),dpi=200)\n",
    "\n",
    "\"\"\"===== BC ==========================================================\"\"\"\n",
    "p1 = ax1.contourf(xs1,ys1,to_np(z1_cross_filled),50,cmap=cmaps.MPL_RdBu_r)      #v\n",
    "p2 = ax2.contourf(xs2,ys2,to_np(z2_cross_filled),50,cmap=cmaps.BlueWhiteOrangeRed)    #w\n",
    "p3 = ax3.contourf(xs3,ys3,to_np(z3_cross_filled),50,cmap=cmaps.WhiteBlueGreenYellowRed)  #ws\n",
    "\n",
    "#=== wind contour plot ========\n",
    "#===== v component ==================\n",
    "p4 = ax1.contour(xs4,ys4,to_np(v1_cross_filled), cmap='hsv',levels=15,linestyles=\"dashed\",linewidths=0.3)\n",
    "#p4 = ax1.contour(xs4,ys4,to_np(v1_cross_filled), cmap='tab20c',levels=5,linestyles=\"dashed\",linewidths=0.6)\n",
    "ax1.clabel(p4, inline=True, fontsize=5,inline_spacing=5,fmt='%1.0f')\n",
    "#===== w component ==================\n",
    "#p5 = ax2.contour(xs5,ys5,to_np(w1_cross_filled), cmap=cmaps.cmocean_curl,levels=5,linestyles=\"dashed\",linewidths=0.6)\n",
    "p5 = ax2.contour(xs5,ys5,to_np(w1_cross_filled), cmap='hsv',levels=6, linestyles=\"dashed\",linewidths=0.3)\n",
    "ax1.clabel(p5, inline=True, fontsize=5,inline_spacing=5,fmt='%1.0f')\n",
    "#===== v component ==================\n",
    "#p6 = ax3.contour(xs6,ys6,to_np(ws1_cross_filled), cmap=cmaps.cmocean_curl,levels=5,linestyles=\"dashed\",linewidths=0.5)\n",
    "p6 = ax3.contour(xs6,ys6,to_np(ws1_cross_filled), cmap='gnuplot',levels=13,linestyles=\"dashed\",linewidths=0.3)\n",
    "ax1.clabel(p6, inline=True, fontsize=5,inline_spacing=5,fmt='%1.0f')\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\"\"\"===== Add colorbar==================================================\"\"\"\n",
    "divider = make_axes_locatable(ax1)\n",
    "cax = divider.append_axes(\"right\", size=\"5%\", pad=0.05)\n",
    "#b1 = plt.colorbar(p1, cax=cax,ticks=[-9,-7.5,-6,-4.5,-3,-1.5,0,1.5,3,4.5,6,7.5])\n",
    "b1 = plt.colorbar(p1, cax=cax,ticks=[-7,-6.5,-6,-5.5,-5,-4.5,-4,-3.5,-3,-2.5,-2,-1.5,-1,-0.5,0,0.5,1,1.5,2,2.5,3,3.5,4])\n",
    "\n",
    "b1.ax.tick_params(labelsize=5)\n",
    "#================================================\n",
    "divider = make_axes_locatable(ax2)\n",
    "cax = divider.append_axes(\"right\", size=\"5%\", pad=0.05)\n",
    "b2 = plt.colorbar(p2, cax=cax,ticks=[-0.055,-0.05,-0.045,-0.04,-0.035,-0.03,-0.025,-0.02,-0.015,-0.010,-0.005,0,0.005,0.010,0.015,0.02,0.025,0.03,0.035,0.04,0.045,0.05,0.055,0.060])\n",
    "b2.ax.tick_params(labelsize=5)\n",
    "#===============================================\n",
    "divider = make_axes_locatable(ax3)\n",
    "cax = divider.append_axes(\"right\", size=\"5%\", pad=0.05)\n",
    "#b3 = plt.colorbar(p3, cax=cax,ticks=[0,2,4,6,8,10,12,14,16,18])\n",
    "b3 = plt.colorbar(p3, cax=cax,ticks=[0,1,2,3,4,5,6,7,8,9,10,11])\n",
    "b3.ax.tick_params(labelsize=5)\n",
    "\n",
    "\"\"\"=========  Fill in the mountain area =========================================\"\"\"\n",
    "ax1.fill_between(xs1, 0, to_np(ter_line),facecolor=\"#4D4D4D\",zorder=5)\n",
    "ax2.fill_between(xs2, 0, to_np(ter_line),facecolor=\"#4D4D4D\",zorder=5)\n",
    "ax3.fill_between(xs3, 0, to_np(ter_line),facecolor=\"#4D4D4D\",zorder=5)\n",
    "\n",
    "\n",
    "\"\"\"=====  Set the x-ticks to use latitude and longitude labels===\"\"\"\n",
    "coord_pairs = to_np(z1_cross.coords[\"xy_loc\"])\n",
    "x_ticks     = np.arange(coord_pairs.shape[0])\n",
    "x_labels    = [pair.latlon_str(fmt=\"{:.1f}, {:.1f}\")\n",
    "              for pair in to_np(coord_pairs)]\n",
    "\"\"\"====== Set the desired number of x ticks below===== \"\"\"\n",
    "num_ticks = 5\n",
    "thin = int((len(x_ticks) / num_ticks) + .8)\n",
    "\n",
    "\"\"\"===== Set xtick label =============================================\"\"\"\n",
    "ax1.set_xticklabels(x_labels[::thin],rotation=25, fontsize=7)\n",
    "ax2.set_xticklabels(x_labels[::thin],rotation=25, fontsize=7)\n",
    "ax3.set_xticklabels(x_labels[::thin],rotation=25, fontsize=7)\n",
    "\"\"\"====  Set ytick label =============================================\"\"\"\n",
    "ax1.yaxis.set_tick_params(labelsize=7)\n",
    "ax2.yaxis.set_tick_params(labelsize=7)\n",
    "ax3.yaxis.set_tick_params(labelsize=7)\n",
    "\n",
    "\"\"\" Add  text \"\"\"\n",
    "ax1.text(-35,10,'Elevation [Km]', va='center', rotation='vertical',fontsize=7)\n",
    "ax1.text(70,18,\"[a]\",fontsize=7)\n",
    "ax2.text(70,18,\"[b]\",fontsize=7)\n",
    "ax3.text(70,18,\"[c]\",fontsize=7)\n",
    "\n",
    "ax1.text(145,-0.8,\"$μg m ^{-2} s ^{-1}$\",fontsize=5)\n",
    "ax2.text(145,-0.8,\"$μg m ^{-2} s ^{-1}$\",fontsize=5)\n",
    "ax3.text(145,-0.8,\"$μg m ^{-2} s ^{-1}$\",fontsize=5)\n",
    "\n",
    "\n",
    "\n",
    "\"\"\" Adjust\"\"\"\n",
    "fig.subplots_adjust(top=0.988,\n",
    "                        bottom=0.124,\n",
    "                        left=0.06,\n",
    "                        right=0.96,\n",
    "                        hspace=0.2,\n",
    "                        wspace=0.4)\n",
    "\n",
    "#plt.savefig(\"/mnt/g/2nd_Paper/tf_2000dpi.png\",dpi=2000)\n",
    "plt.show()\n"
   ]
  }
 ],
 "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.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}