{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\nHail Nowcast (Hong Kong)\n===========================\nThis example demonstrates a rule-based\nhail nowcast for the next\n30 minutes, using data from Hong Kong.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Definitions\n--------------------------------------------------------\n\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Import all required modules and methods:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Python package to allow system command line functions\nimport os\n# Python package to manage warning message\nimport warnings\n# Python package for time calculations\nimport pandas as pd\n# Python package for numerical calculations\nimport numpy as np\n# Python package for xarrays to read and handle netcdf data\nimport xarray as xr\n# Python package for projection\nimport cartopy.crs as ccrs\n# Python package for land/sea features\nimport cartopy.feature as cfeature\n# Python package for reading map shape file\nimport cartopy.io.shapereader as shpreader\n# Python package for creating plots\nfrom matplotlib import pyplot as plt\n# Python package for ellipse patches\nfrom matplotlib.patches import Arc\n# Python package for colorbars\nfrom matplotlib.colors import BoundaryNorm, ListedColormap\n\n# swirlspy qpf function\nfrom swirlspy.qpf import rover\n# swirlspy data parser function\nfrom swirlspy.rad import read_iris_grid, calc_vil\n# swirlspy test data source locat utils function\nfrom swirlspy.qpe.utils import timestamps_ending, locate_file\n# swirlspy hail labeled and fitted function\nfrom swirlspy.object import get_labeled_frame, fit_ellipse\n# swirlspy standardize data function\nfrom swirlspy.utils import standardize_attr, FrameType\n# directory constants\nfrom swirlspy.tests.samples import DATA_DIR\nfrom swirlspy.tests.outputs import OUTPUT_DIR\n\nwarnings.filterwarnings(\"ignore\")\nplt.switch_backend('agg')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Define the working directories:\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "radar_dir = os.path.abspath(os.path.join(DATA_DIR, 'iris/3d'))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Define the basemap:\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# define the plot function\ndef plot_base(ax: plt.Axes, extents: list, crs: ccrs.Projection):\n    # fake the ocean color\n    ax.imshow(\n        np.tile(np.array([[[178, 208, 254]]], dtype=np.uint8), [2, 2, 1]),\n        origin='upper', transform=crs,\n        extent=extents, zorder=-1\n    )\n    # coastline, province and state, color\n    ax.add_feature(\n        map_with_province, facecolor=cfeature.COLORS['land'],\n        edgecolor='none', zorder=0\n    )\n    # overlay coastline, province and state without color\n    ax.add_feature(\n        map_with_province, facecolor='none', edgecolor='gray', linewidth=0.5\n    )\n    ax.set_title('')\n\n\n# Logging\nstart_time = pd.Timestamp.now()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Loading radar data\n-----------------------------------------------------------------------\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Specify the basetime\nbasetime = pd.Timestamp('201403301930')\n\n# Generate timestamps for the current and past 6 minute\n# radar scan\ntimestamps = timestamps_ending(\n    basetime,\n    duration=pd.Timedelta(6, 'm'),\n    exclude_end=False\n)\n\n# Locating the files\nlocated_files = []\nfor timestamp in timestamps:\n    located_files.append(locate_file(radar_dir, timestamp))\n\n# Reading the radar data\nreflectivity_list = []\nfor filename in located_files:\n    reflectivity = read_iris_grid(filename)\n    reflectivity_list.append(reflectivity)\n\n# Standardize reflectivity xarrays\nraw_frames = xr.concat(reflectivity_list,\n                       dim='time').sortby(['y'], ascending=False)\nstandardize_attr(raw_frames, frame_type=FrameType.dBZ)\n\ninitialising_time = pd.Timestamp.now()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Identify regions of interest for hail\n-----------------------------------------------------------------------\n\nHail has a chance of occurring when two phenomena co-happen:\n\n#. The 58 dBZ echo top exceeds 3 km.\n#. The Vertically Integrated Liquid (VIL) up to 2 km is less than 5mm.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# getting radar scan for basetime\nreflectivity = raw_frames.sel(time=basetime)\n\n# maximum reflectivity at elevations > 3km\nref_3km = reflectivity.sel(\n    height=slice(3000, None)\n).max(dim='height', keep_attrs=True)\n\n# 58 dBZ echo top exceeds 3km\ncond_1 = ref_3km > 58\n\n# vil up to 2km\nvil_2km = calc_vil(reflectivity.sel(height=slice(1000, 2000)))\n\n# VIL up 2km is less than 5mm\ncond_2 = vil_2km < 5\n\n# Region of interest for hail\nhail = xr.ufuncs.logical_and(cond_1, cond_2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The identified regions are then labeled and fitted with\nminimum enclosing ellipses.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# label image\n# image is binary so any threshold between 0 and 1 works\n# define minimum size as 4e6 m^2 or 4 km^2\nlabeled_hail, uids = get_labeled_frame(hail, 0.5, min_size=4e6)\n\n# fit ellipses to regions of interest\nellipse_list = []\nfor uid in uids:\n    ellipse, _ = fit_ellipse(labeled_hail == uid)\n    ellipse_list.append(ellipse)\n\nidentify_time = pd.Timestamp.now()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Extrapolation of hail region\n-----------------------------------------------------------------------\n\nFirst, we obtain the xarray.DataArray to generate\nthe motion field.\n\nIn this example, we generate the motion field from\ntwo consecutive 3km CAPPI radar scans closest to basetime.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Select radar data at 3km\nframes = raw_frames.sel(height=3000).drop('height')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Obtain the motion field by ROVER\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# ROVER\nmotion = rover(frames)\n\nmotion_time = pd.Timestamp.now()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Next, we extract the motion vector at the centroid of each ellipse,\nand calculate the displacement of the ellipse after 30 minutes.\nSince the distance is expressed in pixels, we need to convert\nthe distance of the motion vector to grid units (in this case, meters).\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# getting meters per pixel\n\narea_def = reflectivity.attrs['area_def']\nx_d = area_def.pixel_size_x\ny_d = area_def.pixel_size_y\n\n# time ratio\n# time ratio between nowcast interval and unit time\ntime_ratio = pd.Timedelta(30, 'm') / pd.Timedelta(6, 'm')\n\nellipse30_list = []\nfor ellipse in ellipse_list:\n    # getting motion in pixels/6 minutes\n    pu = motion[0].sel(x=ellipse['center'][0],\n                       y=ellipse['center'][1],\n                       method='nearest')\n    pv = motion[1].sel(x=ellipse['center'][0],\n                       y=ellipse['center'][1],\n                       method='nearest')\n    # converting to meters / 6 minutes\n    u = x_d * pu\n    v = y_d * pv\n    # get displacement in 30 minutes\n    dcenterx = u * time_ratio\n    dcentery = v * time_ratio\n    # get new position of ellipse after 30 minutes\n    x30 = ellipse['center'][0] + dcenterx\n    y30 = ellipse['center'][1] + dcentery\n    # get new ellipse, only the center is changed\n    ellipse30 = ellipse.copy()\n    ellipse30['center'] = (x30, y30)\n\n    ellipse30_list.append(ellipse30)\n\n\nextrapolate_time = pd.Timestamp.now()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Visualisation\n-----------------------------------------------------------------------\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Defining figure\nfig = plt.figure(figsize=(12, 6.5))\n# number of rows and columns in plot\nnrows = 2\nncols = 1\n# Defining the crs\ncrs = area_def.to_cartopy_crs()\n\n# Load the shape of Hong Kong\nmap_shape_file = os.path.abspath(\n    os.path.join(DATA_DIR, 'shape/hk_tms_aeqd.shp')\n)\n\n# coastline and province\nmap_with_province = cfeature.ShapelyFeature(\n    list(shpreader.Reader(map_shape_file).geometries()),\n    area_def.to_cartopy_crs()\n)\n\n# Defining extent\nx = labeled_hail.coords['x'].values\ny = labeled_hail.coords['y'].values\nx_d = x[1] - x[0]\ny_d = y[1] - y[0]\nextents = [x[0], y[0], x[-1], y[-1]]\n\n\n# Defining ellipse patches\ndef gen_patches(lst, ls='-'):\n    patch_list = []\n    for ellipse in lst:\n        patch = Arc(\n            ellipse['center'],\n            ellipse['b'] * 2,\n            ellipse['a'] * 2,\n            angle=ellipse['angle'],\n            ls=ls\n        )\n        patch_list.append(patch)\n    return patch_list\n\n\n# 1. Plotting maximum reflectivity above 3km\n# Define color scheme\n# Defining colour scale and  format.\nlevels = [\n    -32768, 10, 15, 20, 24, 28, 32, 34, 38, 41,\n    44, 47, 50, 53, 56, 58, 60, 62\n]\ncmap = ListedColormap([\n    '#FFFFFF00', '#08C5F5', '#0091F3', '#3898FF', '#008243', '#00A433',\n    '#00D100', '#01F508', '#77FF00', '#E0D100', '#FFDC01', '#EEB200',\n    '#F08100', '#F00101', '#E20200', '#B40466', '#ED02F0'\n])\n\nnorm = BoundaryNorm(levels, ncolors=cmap.N, clip=True)\n\nax = fig.add_subplot(ncols, nrows, 1, projection=crs)\nplot_base(ax, extents, crs)\nref_3km.where(ref_3km > levels[1]).plot(\n    ax=ax,\n    cmap=cmap,\n    norm=norm,\n    extend='max',\n    cbar_kwargs={\n        'ticks': levels[1:],\n        'format': '%.3g',\n        'fraction': 0.046,\n        'pad': 0.04\n    }\n)\npatches = gen_patches(ellipse_list)\nfor patch in patches:\n    ax.add_patch(patch)\npatches = gen_patches(ellipse30_list, ls='--')\nfor patch in patches:\n    ax.add_patch(patch)\nax.set_title('MAX_REF >= 3KM')\n\n# 2. Plotting VIL up to 2km\nlevels = [\n    0, 0.05,\n    0.1, 0.2, 0.4, 0.6, 0.8, 1,\n    2, 4, 6, 8, 15, 20,\n    25, 30, 32, 34\n]\ncmap = ListedColormap([\n    '#FFFFFF', '#08C5F5', '#0091F3', '#3898FF', '#008243', '#00A433',\n    '#00D100', '#01F508', '#77FF00', '#E0D100', '#FFDC01', '#EEB200',\n    '#F08100', '#F00101', '#E20200', '#B40466', '#ED02F0'\n])\n\nnorm = BoundaryNorm(levels, ncolors=cmap.N, clip=True)\nax = fig.add_subplot(ncols, nrows, 2, projection=crs)\nplot_base(ax, extents, crs)\n\nvil_2km.where(vil_2km > levels[1]).plot(\n    cmap=cmap,\n    norm=norm,\n    ax=ax,\n    extend='max',\n    cbar_kwargs={\n        'ticks': levels[1:],\n        'format': '%.3g',\n        'fraction': 0.046,\n        'pad': 0.04\n    }\n)\nax.set_title('VIL <= 2km')\npatches = gen_patches(ellipse_list)\nfor patch in patches:\n    ax.add_patch(patch)\npatches = gen_patches(ellipse30_list, ls='--')\nfor patch in patches:\n    ax.add_patch(patch)\n\nsuptitle1 = \"Hail Nowcast Tracks\"\nsuptitle2 = basetime.strftime('%Y-%m-%d')\nsuptitle3 = (f\"Based @ {basetime.strftime('%H:%MH')}\\n\"\n             f\"Valid @ {(basetime + pd.Timedelta(30, 'm')).strftime('%H:%MH')}\")\nfig.text(0., 0.93, suptitle1, va='top', ha='left', fontsize=16)\nfig.text(0.57, 0.93, suptitle2, va='top', ha='center', fontsize=16)\nfig.text(0.90, 0.93, suptitle3, va='top', ha='right', fontsize=16)\nplt.tight_layout()\nplt.savefig(\n    os.path.join(OUTPUT_DIR, \"hail.png\"),\n    dpi=300\n)\n\nvisualise_time = pd.Timestamp.now()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Checking run time of each component\n-----------------------------------------------------------------------\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "print(f\"Start time: {start_time}\")\nprint(f\"Initialising time: {initialising_time}\")\nprint(f\"Identify time: {identify_time}\")\nprint(f\"Motion field time: {motion_time}\")\nprint(f\"Extrapolate time: {extrapolate_time}\")\nprint(f\"Visualise time: {visualise_time}\")\n\nprint(f\"Time to initialise: {initialising_time - start_time}\")\nprint(f\"Time to identify hail regions: {identify_time - initialising_time}\")\nprint(f\"Time to generate motion field: {motion_time - identify_time}\")\nprint(f\"Time to extrapolate: {extrapolate_time - motion_time}\")\nprint(f\"Time to visualise: {visualise_time - extrapolate_time}\")\n\nprint(f\"Total: {visualise_time - start_time}\")"
      ]
    }
  ],
  "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
}