{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# QPE (Hong Kong)\n\nThis example demonstrates how to perform QPE,\nusing raingauge and radar data from Hong Kong.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Definitions\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import os\nimport time\nimport tarfile\nimport numpy as np\nimport pandas as pd\nimport copy\nimport xarray\nimport scipy\nimport pyproj\nimport matplotlib\nfrom matplotlib.colors import BoundaryNorm\nimport matplotlib.pyplot as plt\nfrom PIL import Image\nfrom pyresample import utils\nfrom cartopy.io import shapereader\nfrom sklearn.gaussian_process import kernels\n\nfrom swirlspy.obs import Rain\nfrom swirlspy.rad.iris import read_iris_grid\nfrom swirlspy.core.resample import grid_resample\nfrom swirlspy.qpe.rfmap import rg_interpolate, comp_qpe, show_raingauge\nfrom swirlspy.qpe.utils import timestamps_ending, locate_file, dbz2rr, rr2rf, \\\n    temporal_interp, acc\n\nplt.switch_backend('agg')\nTHIS_DIR = os.getcwd()\nos.chdir(THIS_DIR)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Initialising\n\nThis section demonstrates extracting raingauge and radar data.\n\nStep 1: Defining an end-time for accumulating rainfall.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "acctime = pd.Timestamp('20190420150000').floor('min')\nacctime_str = acctime.strftime('%Y%m%d%H%M')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Step 2: Setting up directories for raingauge and radar files.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "rad_dir = THIS_DIR + '/../tests/samples/iris/'\nrg_dir = THIS_DIR + '/../tests/samples/rfmap/'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Step 3: Generating timestamps and pattern for both radar and\nraingauge files.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Timestamps of raingauges\nrg_timestrings = timestamps_ending(\n    acctime + pd.Timedelta(minutes=5),\n    duration=pd.Timedelta(hours=1),\n    interval=pd.Timedelta(minutes=5)\n)\n\n# Raingauge pattern\nrg_pattern = ['rf5m_20'+ts for ts in rg_timestrings]\n\n# Finding time nearest radar file\n# to accumulation end time\nminute = acctime.minute\nnearest_6min = acctime.minute // 6 * 6\n\nnearest_rad_timestamp = pd.Timestamp(\n    acctime_str[:-2]+f'{nearest_6min:02}'\n)\n\nrad_timestrings = timestamps_ending(\n    nearest_rad_timestamp,\n    duration=pd.Timedelta(hours=1),\n    interval=pd.Timedelta(minutes=6)\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Step 4: Extracting raingauge and radar files from\ntheir respective directories.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "located_rg_files = []\nfor pat in rg_pattern:\n    located_rg_files.append(locate_file(rg_dir, pat))\n\nlocated_radar_files = []\nfor ts in rad_timestrings:\n    located_radar_files.append(locate_file(rad_dir, ts))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Step 5: Read data from raingauge files into a Rain object.\nCoordinates are geodetic, following that in the files.\nThere is some noise in the raingauge, so known problematic stations\nare filtered away.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "rg_object_geodetic = Rain(\n    located_rg_files,\n    'geodetic',\n    duration=pd.Timedelta(minutes=5),\n    NAN=[3276.7, 32767]\n)\n\nbad_stations = ['N25', 'SSP', 'D25', 'TWN', 'TMS', 'N14']\nrg_object_geodetic = rg_object_geodetic.remove_bad_stations(\n    bad_stations\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Step 6: Read radar files into xarray.DataArrays.\nThe data in the files are already in Cartesian Coordinates,\nin the Centered Azimuthal Projection.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "reflec_list = []\nfor file in located_radar_files:\n    reflec = read_iris_grid(\n        file\n    )\n\n    reflec_list.append(reflec)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Data Reprojection\n\nThis section demonstrates the reprojection of extracted raingauge\nand radar data to a user-defined grid.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Step 1: Define the target grid as a pyresample AreaDefinition.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Defining target grid\narea_id = \"hk1980_grid\"\ndescription = (\"A grid centered about Hong Kong with a resolution \"\n               \"880m in the x-direction and 660m in the y-direction \"\n               \"HK1980 easting/northing coordinates\")\nproj_id = 'hk1980'\nprojection = ('+proj=tmerc +lat_0=22.31213333333334 '\n              '+lon_0=114.1785555555556 +k=1 +x_0=836694.05 '\n              '+y_0=819069.8 +ellps=intl +towgs84=-162.619,-276.959,'\n              '-161.764,0.067753,-2.24365,-1.15883,-1.09425 +units=m '\n              '+no_defs')\n\nx_size = 1000\ny_size = 1000\n\narea_extent = (792000, 796000, 880000, 862000)\n\narea_def = utils.get_area_def(\n    area_id, description, proj_id, projection, x_size, y_size, area_extent\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Step 2: Convert coordinates of raingauge object to desired projection.\nIn this example, the desired projection is HK1980.\nThis can be achieved by the .reproject() method of the Rain object.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "inProj = pyproj.Proj(init=\"epsg:4326\")\noutProj = pyproj.Proj(area_def.proj_str)\n\nrg_object = rg_object_geodetic.reproject(inProj, outProj, \"HK1980\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Step 3: Regrid radar reflectivity from Centered Azimuthal\nProjection to HK1980.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "reproj_reflec_list = []\nfor reflec in reflec_list:\n    reproj_reflec = grid_resample(\n        reflec,\n        reflec.attrs['area_def'],\n        area_def, coord_label=['easting', 'northing']\n    )\n    reproj_reflec_list.append(reproj_reflec)\n\nreflectivity = xarray.concat(reproj_reflec_list, dim='time')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Accumulating and interpolating rainfall\n\nInterpolate rainfall recorded by raingauges into the user-defined grid\nand accumulate radar rainfall over an hour\nafter making the necessary\nadjustments.\n\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Step 1: Interpolate Rain object to user-defined grid. In this example,\nordinary kriging is used.\n\nUsing kriging may require further customisation of certain\nparameters.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Perform some primitive quality control\nupperQ = np.quantile(rg_object.rainfall, .75)\nlowerQ = np.quantile(rg_object.rainfall, .25)\niqr = upperQ - lowerQ\n\nnoisePos = np.logical_or(rg_object.rainfall > upperQ + 1.5*iqr,\n                         rg_object.rainfall < lowerQ - 1.5*iqr)\n\nalpha = np.where(noisePos, 1e4, 1e-10)\n\nkernel = kernels.Matern()\n\ninterpolated_rg = rg_interpolate(\n    rg_object, area_def, 'ordinary_kriging',\n    coord_label=['easting', 'northing'],\n    kernel=kernel,\n    alpha=alpha,\n    n_restarts_optimizer=20\n    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Step 2: Convert to radar reflectivity to rainrates,\ninterpolate rainrates to times of raingauges,\nand convert to rainfalls accumulated every 5 minutes.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "rainrates = dbz2rr(reflectivity, a=58.53, b=1.56)\n# Convert time coordinates of rainrates from start time\n# to end time\nrainrates_time_endtime = rainrates.copy()\nrainrates_time_endtime.coords['time'] = \\\n    [\n        pd.Timestamp(t) + pd.Timedelta(minutes=6)\n        for t in rainrates.coords['time'].values\n    ]\n\nrainrates_5min = temporal_interp(\n    rainrates_time_endtime,\n    rg_object.start_time + pd.Timedelta(minutes=5),\n    rg_object.end_time,\n    intvl=pd.Timedelta(minutes=5),\n    interp_type='quadratic'\n    )\nrainfalls = rr2rf(rainrates_5min, scan_duration=5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Step 3: Accumulate rainfall over an hour.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "acc_rf = acc(\n    rainfalls,\n    rg_object.end_time,\n    acc_period=pd.Timedelta(minutes=60)\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Compositing rainfall\n\nPerform compositing on radar and raingauge derived rainfall\nto obtain a composite QPE.\n\nSome parameter tuning may be required to make the observations\nfit better with each other.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "comprf = comp_qpe(\n    area_def,\n    rg_object=rg_object,\n    rg_interp=interpolated_rg,\n    rad_rf=acc_rf,\n    rg_radius=5000,\n    max_truth={'rg': 1., 'radar': 0.1}\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Plotting\n\nPlot composited radar and raingauge rainfall.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Plotting function for neatness.\n\n\ndef plot_map(\n    da, rg_object, acctime, area_def,\n    based='raingauge and radar',\n    savepath='',\n):\n\n    \"\"\"\n    A custom function for plotting a map.\n\n    Parameters\n    --------------\n    da: xarray.DataArray\n        Contains data to be plotted.\n    rg_object: Rain\n        Contains raingauge data.\n    acctime: pd.Timestamp\n        Contains the endtime of the accumulation\n        period.\n    area_def: pyresample.geometry.AreaDefinition\n        AreaDefinition of the grid.\n    based: str\n        Type of data plotted in the map.\n    savepath: str\n        Path to save the image to.\n\n    \"\"\"\n    # Defining the colour scheme\n    levels = [\n        0, 0.5, 2, 5, 10, 20,\n        30, 40, 50, 70, 100, 150,\n        200, 300, 400, 500, 600, 700\n    ]\n\n    cmap = matplotlib.colors.ListedColormap([\n        '#ffffff', '#9bf7f7', '#00ffff', '#00d5cc', '#00bd3d', '#2fd646',\n        '#9de843', '#ffdd41', '#ffac33', '#ff621e', '#d23211', '#9d0063',\n        '#e300ae', '#ff00ce', '#ff57da', '#ff8de6', '#ffe4fd'\n    ])\n\n    norm = BoundaryNorm(levels, ncolors=cmap.N, clip=True)\n\n    # Plotting axes\n    plt.figure(figsize=(28, 21))\n    crs = area_def.to_cartopy_crs()\n    ax = plt.axes(projection=crs)\n\n    # Plot data\n    quadmesh = da.plot(\n        cmap=cmap,\n        norm=norm,\n        extend='max',\n        cbar_kwargs={'ticks': levels, 'format': '%.3g'}\n    )\n\n    # Adjusting size of colorbar\n    cb = quadmesh.colorbar\n    cb.ax.set_ylabel(\n        da.attrs['long_name']+'['+da.attrs['units']+']',\n        fontsize=28\n    )\n\n    cb.ax.tick_params(labelsize=24)\n\n    # Setting labels\n    ax.xaxis.set_visible(True)\n    ax.yaxis.set_visible(True)\n\n    for tick in ax.xaxis.get_major_ticks():\n        tick.label.set_fontsize(24)\n    for tick in ax.yaxis.get_major_ticks():\n        tick.label.set_fontsize(24)\n\n    ax.xaxis.label.set_size(28)\n    ax.yaxis.label.set_size(28)\n\n    # Coastlines\n    shpfile = THIS_DIR + \\\n        '/../tests/samples/gadm36_HKG_shp/gadm36_HKG_0_hk1980.shp'\n    shp = shapereader.Reader(shpfile)\n    for record, geometry in zip(shp.records(), shp.geometries()):\n        ax.add_geometries([geometry], crs, facecolor='None', edgecolor='black')\n\n    # Show title\n    plt.title(\n        (f\"Last Hour Rainfall at {acctime.strftime('%H:%MH %d-%b-%Y')}\\n\"\n         f\"(based on {based} data)\"),\n        fontsize=32\n    )\n\n    plt.savefig(savepath, dpi=300)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Plotting maps\n\n# Raingauge only\nplot_map(\n    interpolated_rg, rg_object,\n    acctime, area_def,\n    based='raingauge',\n    savepath=THIS_DIR+f'/../tests/outputs/raingauge_{acctime_str}.png'\n)\n\n# Radar only\nplot_map(\n    acc_rf, rg_object,\n    acctime, area_def,\n    based='radar',\n    savepath=THIS_DIR+f'/../tests/outputs/radar_{acctime_str}.png'\n)\n\n# Composite raingauge and radar\nplot_map(\n    comprf, rg_object,\n    acctime, area_def,\n    savepath=THIS_DIR+f'/../tests/outputs/comp_{acctime_str}.png'\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.6.15"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}