Chapter05_regridding_multiprocessing.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Regridding by multiprocessing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Based on [chapter04](https://github.com/zxdawn/GEOSChem-python-tutorial/blob/master/Chapter04_regridding_WRFChem_part1.ipynb), we can regrid many files at the same time by `multiprocessing`.\n",
    "\n",
    "Because there're so many files, it's better to comment the reminding of using `reuse_weights`.\n",
    "\n",
    "Edit `site-packages/xesmf/frontend.py` Line 189 - 204 like this:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```    \n",
    "    def _write_weight_file(self):\n",
    "\n",
    "        if os.path.exists(self.filename):\n",
    "            if self.reuse_weights:\n",
    "                #print('Reuse existing file: {}'.format(self.filename))\n",
    "                return  # do not compute it again, just read it\n",
    "            else:\n",
    "                #print('Overwrite existing file: {} \\n'.format(self.filename),\n",
    "                #      'You can set reuse_weights=True to save computing time.')\n",
    "                os.remove(self.filename)\n",
    "        #else:\n",
    "        #    print('Create weight file: {}'.format(self.filename))\n",
    "\n",
    "        regrid = esmf_regrid_build(self._grid_in, self._grid_out, self.method,\n",
    "                                   filename=self.filename)\n",
    "        esmf_regrid_finalize(regrid)  # only need weights, not regrid object\n",
    "```       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, fnmatch\n",
    "import multiprocessing\n",
    "from multiprocessing import Pool\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "import xesmf as xe\n",
    "from xgcm.autogenerate import generate_grid_ds\n",
    "from xgcm import Grid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Prepare global variables for multiprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set range of longitude/latitude and the rsolution of grid we want\n",
    "resolution = 0.1 # degree\n",
    "Lat_min = 20; Lat_max = 50\n",
    "Lon_min = -115; Lon_max = -65\n",
    "\n",
    "# The nested grid we want\n",
    "nested_grid = xe.util.grid_2d(Lon_min-resolution, Lon_max+resolution, resolution,  # longitude boundary range and resolution\n",
    "              Lat_min-resolution, Lat_max+resolution, resolution)  # latitude boundary range and resolution\n",
    "\n",
    "wrf_dir  = '/chenq3/zhangxin/BEHR/data/wrf_profiles/us/2014/all/' # the directory of wrfout* files\n",
    "save_dir = '/chenq3/zhangxin/BEHR/data/wrf_profiles/us/2014/regrid/'\n",
    "vnames   = ['IC_FLASHCOUNT','CG_FLASHCOUNT','no2','lno','lno2'] # extracted variables"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create several functions as [chapter04](https://github.com/zxdawn/GEOSChem-python-tutorial/blob/master/Chapter04_regridding_WRFChem_part1.ipynb):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_attrs(result):\n",
    "    result['no2'].attrs['description']  = 'NO2 mixing ratio'\n",
    "    result['lno'].attrs['description']  = 'LNO mixing ratio'\n",
    "    result['lno2'].attrs['description'] = 'LNO2 mixing ratio'\n",
    "    result['no2'].attrs['units']  = 'ppmv'\n",
    "    result['lno'].attrs['units']  = 'ppmv'\n",
    "    result['lno2'].attrs['units'] = 'ppmv'\n",
    "    result['IC_FLASHCOUNT'].attrs['description'] = 'Accumulated IC flash count'\n",
    "    result['CG_FLASHCOUNT'].attrs['description'] = 'Accumulated CG flash count'\n",
    "    result['TL_FLASHCOUNT'].attrs['description'] = 'Accumulated Total flash count'\n",
    "\n",
    "\n",
    "def bilinear_regridding(ds,nested_grid):\n",
    "    regridder_bilinear = xe.Regridder(ds, nested_grid, method='bilinear', reuse_weights=True)\n",
    "\n",
    "    bilinear_list = [] # an emtpy list to hold regridding result\n",
    "\n",
    "    for varname, dr in ds.data_vars.items():\n",
    "      # Olny pick variables we need\n",
    "      if varname in vnames:\n",
    "        dr_temp = regridder_bilinear(dr)  # temporary variable for the current tracer\n",
    "        bilinear_list.append(dr_temp)\n",
    "\n",
    "    bilinear_result = xr.merge(bilinear_list)  # merge a list of DataArray to a single Dataset\n",
    "    # NOTE: The next version of xESMF (v0.2) will be able to directly regrid a Dataset, \n",
    "    # so you will not need those additional code. But it is a nice coding exercise anyway.\n",
    "\n",
    "    bilinear_result['TL_FLASHCOUNT'] =  bilinear_result.CG_FLASHCOUNT + bilinear_result.IC_FLASHCOUNT #TL  = IC +CG\n",
    "\n",
    "    # Add attributes\n",
    "    add_attrs(bilinear_result)\n",
    "\n",
    "    return regridder_bilinear,bilinear_result\n",
    "\n",
    "\n",
    "def conservative_regridding(ds,nested_grid):\n",
    "    # To use conservative regridding, you need to specify the cell boundaries,\n",
    "    # since the boundary information is needed for calculating overlapping areas \n",
    "    # (thus ensures total mass conservation).\n",
    "\n",
    "    #Converting 2D coordinates from cell center to bound with xgcm.autogenerate\n",
    "    #https://gist.github.com/jbusecke/175d72d81e13f7f8d4dcf26aace511bd\n",
    "\n",
    "    # Generate outer dimensions from center input\n",
    "    ds_post = generate_grid_ds(ds, {'X':'west_east', 'Y':'south_north'}, position=('center', 'outer'))\n",
    "\n",
    "    # Generate the 2d coordinates for the cell boundaries\n",
    "    # with the \"standard\" xgcm interpolation for each logical dimension\n",
    "    grid_ds = Grid(ds_post, periodic=False)\n",
    "    bnd='extrapolate'\n",
    "    ds_post.coords['xb'] = grid_ds.interp(grid_ds.interp(ds_post['lon'], 'X', boundary=bnd, fill_value=np.nan),\n",
    "                                              'Y', boundary=bnd, fill_value=np.nan)\n",
    "    ds_post.coords['yb'] = grid_ds.interp(grid_ds.interp(ds_post['lat'], 'X', boundary=bnd, fill_value=np.nan),\n",
    "                                              'Y', boundary=bnd, fill_value=np.nan)\n",
    "\n",
    "    wrf_grid_with_bounds = {'lon': ds['lon'].values,\n",
    "                             'lat': ds['lat'].values,\n",
    "                             'lon_b': ds_post.xb.data,\n",
    "                             'lat_b': ds_post.yb.data,\n",
    "                           }\n",
    "\n",
    "    regridder_conserve = xe.Regridder(wrf_grid_with_bounds, nested_grid, method='conservative', reuse_weights=True)\n",
    "\n",
    "    conservative_list = []\n",
    "\n",
    "    for varname, dr in ds.data_vars.items():\n",
    "      if varname in vnames:\n",
    "        dr_temp = regridder_conserve(dr)\n",
    "        conservative_list.append(dr_temp)\n",
    "\n",
    "    conservative_result = xr.merge(conservative_list)\n",
    "\n",
    "    conservative_result['TL_FLASHCOUNT'] =  conservative_result.CG_FLASHCOUNT + conservative_result.IC_FLASHCOUNT\n",
    "\n",
    "    # Set attributes\n",
    "    add_attrs(conservative_result)\n",
    "\n",
    "    return regridder_conserve,conservative_result\n",
    "\n",
    "def regrid(filename):\n",
    "      # Read file, drop time dimension and rename lon/lat\n",
    "      ds = xr.open_dataset(wrf_dir+filename).isel(Time=0)\n",
    "      ds.rename({'XLONG':'lon','XLAT':'lat'}, inplace=True)\n",
    "\n",
    "      regridder_bilinear, bilinear_result     = bilinear_regridding(ds,nested_grid)\n",
    "      regridder_conserve, conservative_result = conservative_regridding(ds,nested_grid)\n",
    "\n",
    "      # Save to netcdf file\n",
    "      conservative_result.to_netcdf(save_dir+filename+'_conservative_regridding.nc')\n",
    "      bilinear_result.to_netcdf(save_dir+filename+'_bilinear_regridding.nc')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is the main function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def main():\n",
    "  filenames = [filename for filename in os.listdir(wrf_dir) if fnmatch.fnmatch(filename, 'wrfout*')]\n",
    "  \n",
    "  # Build Regridder\n",
    "  regrid(filenames[0])\n",
    "\n",
    "  # Multiprocessing using created regridder\n",
    "  p = Pool(multiprocessing.cpu_count())\n",
    "  p.map(regrid,filenames[1:])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's check the time used to finish the main function.\n",
    "Without mu;tiprocessing, it takes about `50 minutes` to deal with `1720` wrfout* files (12 km resolution, 430 * 345)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %%time\n",
    "# if __name__ == '__main__':\n",
    "#     main()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since the coumpting node has `24 cores`, I decide to run the script directly and print the time used here.\n",
    "\n",
    "But time used is almost as same as before ...\n",
    "\n",
    "I'm trying to figure out why ..."
   ]
  }
 ],
 "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}