{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\nProbabilistic Lightning Density Nowcast (Hong Kong)\n========================================================\nThis example demonstrates how to perform\nprobabilistic lightning nowcast by advecting\nlightning density, 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 text formatting\nimport textwrap\n# Python package for projection description\nimport pyproj\nfrom pyresample import get_area_def\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 output grid format\nfrom matplotlib.gridspec import GridSpec\n# Python package for colorbars\nfrom matplotlib.colors import BoundaryNorm, ListedColormap\nfrom matplotlib.cm import ScalarMappable\n\n# Python com-swirls package to calculate motion field (rover) and semi-lagrangian advection\nfrom swirlspy.qpf import rover, sla\n# swirlspy data parser function\nfrom swirlspy.rad.iris import read_iris_grid\n# swirlspy test data source locat utils function\nfrom swirlspy.qpe.utils import timestamps_ending, locate_file\n# swirlspy regrid function\nfrom swirlspy.core.resample import grid_resample\n# swirlspy lightning function\nfrom swirlspy.obs import Lightning\nfrom swirlspy.ltg.map import gen_ltg_density\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": [
        "Initialising\n-------------------------------------------------------------------------\n\nDefine the target grid:\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Defining target grid\narea_id = \"hk1980_250km\"\ndescription = (\"A 1km resolution rectangular grid \"\n               \"centred at HKO and extending to 250 km \"\n               \"in each direction in 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 = 500\ny_size = 500\n\narea_extent = (587000, 569000, 1087000, 1069000)\n\narea_def = get_area_def(\n    area_id, description, proj_id, projection, x_size, y_size, area_extent\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Define basemap:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Load the shape of Hong Kong\nmap_shape_file = os.path.abspath(os.path.join(\n    DATA_DIR,\n    'shape/hk_hk1980'\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\n# 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('')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Defining a basetime, nowcast interval and nowcast duration:\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "basetime = pd.Timestamp('201904201200').floor('min')\nbasetime_str = basetime.strftime('%Y%m%d%H%M')\nnowcast_interval = pd.Timedelta(30, 'm')\nnowcast_duration = pd.Timedelta(2, 'h')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Loading Lightning and Radar Data\n--------------------------------------------------------------------------\n\nExtract Lightning Location Information System (LLIS) files to read.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# The interval between LLIS files\nllis_file_interval = pd.Timedelta(1, 'm')\n\n# Finding the string representation\n# of timestamps marking the relevant LLIS files\ntimestamps = timestamps_ending(\n    basetime,\n    duration=nowcast_interval,\n    interval=llis_file_interval,\n    format='%Y%m%d%H%M'\n)\n\n# Locating files\nllis_dir = os.path.join(DATA_DIR, 'llis')\n\nlocated_files = []\nfor timestamp in timestamps:\n    located_file = locate_file(llis_dir, timestamp)\n    located_files.append(located_file)\n    if located_file is not None:\n        print(f\"LLIS file for {timestamp} has been located\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Read data from files into Lightning objects.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Initialising Lightning objects\n# Coordinates are geodetic\nltg_object_geodetic = Lightning(\n    located_files,\n    'geodetic',\n    basetime - nowcast_interval,\n    basetime\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Locating IRIS REF radar files needed to generate motion field.\nThe IRIS REF files are marked by start time.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# The interval between the radar files\nrad_interval = pd.Timedelta(6, 'm')\n\n# Finding the string representation\n# of the timestamps\n# marking the required files\n# In this case, the most recent radar files\n# separated by `nowcast_interval` is required\n\nrad_timestamps = timestamps_ending(\n    basetime,\n    duration=nowcast_interval + rad_interval,\n    interval=rad_interval\n)\n\nrad_timestamps = [rad_timestamps[0], rad_timestamps[-1]]\n\nrad_dir = os.path.join(DATA_DIR, 'iris')\nlocated_rad_files = []\nfor timestamp in rad_timestamps:\n    located_file = locate_file(rad_dir, timestamp)\n    located_rad_files.append(located_file)\n    if located_file is not None:\n        print(f\"IRIS REF file for {timestamp} has been located\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Read from iris radar files using read_iris_grid().\nThe data is in the Azimuthal Equidistant Projection.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "reflec_unproj_list = []\nfor file in located_rad_files:\n    reflec = read_iris_grid(file)\n    reflec_unproj_list.append(reflec)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Data Reprojection\n---------------------------------------------------------------------------\n\nSince the coordinates of the Lightning objects are geodetic, a conversion\nfrom geodetic coordinates to that of the user-defined projection system is\nrequired.\nThis can be achieved by `swirlspy.obs.Lightning.reproject()`\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# pyproj representation of the map projection of the input data\ninProj = pyproj.Proj(init='epsg:4326')\n# pyproj representation of the intended map projection\noutProj = pyproj.Proj(area_def.proj_str)\n\n# reproject\nltg_object = ltg_object_geodetic.reproject(inProj, outProj, 'HK1980')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Radar data also needs to be reprojected from Azimuthal Equidistant\nto HK1980. This is achieved by the `swirlspy.core.resample` function.\nFollowing reprojection, the individual xarray.DataArrays can be concatenated\nto a single xarray.DataArrays.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "reflec_list = []\nfor z in reflec_unproj_list:\n    reflec = grid_resample(\n        z,\n        z.attrs['area_def'],\n        area_def,\n        coord_label=['easting', 'northing']\n    )\n    reflec_list.append(reflec)\n\nreflectivity = xr.concat(reflec_list, dim='time')\nstandardize_attr(reflectivity, frame_type=FrameType.dBZ)\n\nprint(reflectivity)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Generating lightning density field\n------------------------------------------------------------------\n\nGenerate a lightning density field.\n\n#. Count the number of lightning strikes in\n   a circular search area of 8km surrounding a pixel over\n   the past 30 minutes. Do this for all the pixels.\n#. This value is then spatially and temporally averaged to\n   give a lightning density, an indicator of lightning activity given as the\n   number of strikes/km\\ :sup:`2`/hr.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Obtaining lightning density from lightning objects\nltg_density = gen_ltg_density(\n    ltg_object,\n    area_def,\n    ltg_object.start_time,\n    ltg_object.end_time,\n    tolerance=8000,\n    coord_label=['easting', 'northing']\n)\n\nprint(ltg_density)\n\n# Duplicating ltg_density along the time dimension\n# so that it can be advected.\nltg_density1 = ltg_density.copy()\nltg_density1.coords['time'] = [basetime - nowcast_interval]\nltg_density_concat = xr.concat(\n    [ltg_density1, ltg_density],\n    dim='time'\n)\n# Identifying zero value\nltg_density_concat.attrs['zero_value'] = np.nan"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Generating motion fields and extrapolating\n--------------------------------------------------------------\n\n#. Generate motion field using ROVER.\n#. Extrapolate lightning density using the motion field\n   by Semi-Lagrangian Advection.\n#. Repeat using different parameters to form an ensemble.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Different ROVER parameters\nparams = pd.DataFrame()\nparams['smallest_scale'] = [1, 1, 1, 2]\nparams['largest_scale'] = [7, 7, 7, 7]\nparams['rho'] = [9, 9, 9, 9]\nparams['alpha'] = [2000, 2000, 10000, 2000]\nparams['sigma'] = [1.5, 2.5, 1.5, 2.5]\nparams['tracking_interval'] = [30, 30, 30, 30]\n\nxtrp_frames = []\nfor i in range(len(params)):\n    # Generating motion field using ROVER\n    motion = rover(\n        reflectivity,\n        start_level=params['smallest_scale'][i],\n        max_level=params['largest_scale'][i],\n        rho=params['rho'][i],\n        alpha=params['alpha'][i],\n        sigma=params['sigma'][i]\n    )\n    # SLA\n    steps = int(nowcast_duration.total_seconds()\n                // nowcast_interval.total_seconds())\n    ltg_density_xtrp = sla(ltg_density_concat, motion, nowcast_steps=steps)\n    xtrp_frames.append(ltg_density_xtrp)\n\n\n# Create coordinates for `ensemble` dimension\nens_coords = [\"Member-{:02}\".format(i) for i in range(len(xtrp_frames))]\n\n# Concat ltg_frames along the ensemble dimension\nltg_frames = xr.concat(xtrp_frames, pd.Index(ens_coords, name='ensemble'))\n\nprint(ltg_frames)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Now, let us visualise the extrapolated lightning density fields\nfor the next 30 minutes.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "levels = [\n    0, 0.125, 0.25, 0.5, 1,\n    2, 4, 8, 16, 32, 999\n]\ncmap = ListedColormap([\n    '#FFFFFF', '#5050FF', '#00FFFF', '#00C030', '#A0FF80',\n    '#FFFF40', '#FF9830', '#FF0000', '#FF80FF', '#800080'\n])\n\nnorm = BoundaryNorm(levels, ncolors=cmap.N, clip=True)\n\nmappable = ScalarMappable(cmap=cmap, norm=norm)\nmappable.set_array([])\n\n# Defining the crs\ncrs = area_def.to_cartopy_crs()\n\n# Defining area\nx = ltg_frames.coords['easting'].values\ny = ltg_frames.coords['northing'].values\nextents = [\n    np.min(x), np.max(x),\n    np.min(y), np.max(y)\n]\n\nqx = motion.coords['easting'].values[::20]\nqy = motion.coords['northing'].values[::20]\nqu = motion.values[0, ::20, ::20]\nqv = motion.values[1, ::20, ::20]\n\nfig = plt.figure(figsize=(8, 8), frameon=False)\ngs = GridSpec(\n    2, 2, figure=fig, wspace=0.03, hspace=-0.25, top=0.95,\n    bottom=0.05, left=0.17, right=0.845\n)\n\n# The time to visualise\nvis_time = basetime + nowcast_interval\n\nfor i, mem in enumerate(ltg_frames.ensemble.values):\n    row = i // 2\n    col = i % 2\n    ax = fig.add_subplot(gs[row, col], projection=crs)\n\n    # plot base map\n    plot_base(ax, extents, crs)\n\n    # plot lightning\n    frame = ltg_frames.sel(ensemble=mem, time=vis_time)\n    frame.where(frame > levels[1]).plot(\n        cmap=cmap,\n        norm=norm,\n        add_colorbar=False\n    )\n\n    # plot motion vector\n    ax.quiver(\n        qx, qy, qu, qv, pivot='mid'\n    )\n\n    ax.set_title(mem.split(\"-\")[1], fontsize=10)\n\n    ax.text(\n        extents[0],\n        extents[3],\n        textwrap.dedent(\n            \"\"\"\n            Lightning Density\n            Based @ {baseTime}\n            \"\"\"\n        ).format(\n            baseTime=basetime.strftime('%H:%MH')\n        ).strip(),\n        fontsize=10,\n        va='bottom',\n        ha='left',\n        linespacing=1\n    )\n    ax.text(\n        extents[1],\n        extents[3],\n        textwrap.dedent(\n            \"\"\"\n            {validDate}\n            Valid @ {validTime}\n            \"\"\"\n        ).format(\n            validDate=basetime.strftime('%Y-%m-%d'),\n            validTime=vis_time.strftime('%H:%MH')\n        ).strip(),\n        fontsize=10,\n        va='bottom',\n        ha='right',\n        linespacing=1\n    )\n\ncbar_ax = fig.add_axes([0.875, 0.125, 0.03, 0.75])\ncbar = fig.colorbar(\n    mappable, cax=cbar_ax, ticks=levels[:-1], extend='max', format='%.3g')\ncbar.ax.set_ylabel(ltg_frames.attrs['long_name'], rotation=90)\n\nfig.savefig(\n    os.path.join(OUTPUT_DIR, \"ltg-hk.png\"),\n    bbox_inches='tight'\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Deriving lightning probabilities\n------------------------------------------------------------------------------\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Finally, it is possible to derive\n# the probability of lightning exceeding threshold from our ensemble forecast.\n#\n\n# Compute exceedance probabililities for a 1 strikes/km^2/h threshold\nP = (ltg_frames > 1).sum(dim='ensemble') / len(ltg_frames.ensemble) * 100\n\n# Define cmap, norm\nlevels = [\n    0, 5, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100\n]\ncmap = plt.get_cmap(\"OrRd\", len(levels))\nnorm = BoundaryNorm(levels, ncolors=cmap.N, clip=True)\n\n# Plotting the Exceedance Probability for the first 30 minutes nowcast\nfig = plt.figure(figsize=(15, 10))\nax = plt.axes(projection=area_def.to_cartopy_crs())\n\nplot_base(ax, extents, area_def.to_cartopy_crs())\n\nP.where(P > levels[0]).sel(time=vis_time).plot(\n    cmap=cmap,\n    norm=norm\n)\n\nax.text(\n    extents[0],\n    extents[3],\n    textwrap.dedent(\n        \"\"\"\n        Probability exceeding 1 strikes/$km^2$/hr\n        Based @ {baseTime}\n        \"\"\"\n    ).format(\n        baseTime=basetime.strftime('%H:%MH')\n    ).strip(),\n    fontsize=20,\n    va='bottom',\n    ha='left',\n    linespacing=1\n)\nax.text(\n    extents[1],\n    extents[3],\n    textwrap.dedent(\n        \"\"\"\n        {validDate}\n        Valid @ {validTime}\n        \"\"\"\n    ).format(\n        validDate=basetime.strftime('%Y-%m-%d'),\n        validTime=vis_time.strftime('%H:%MH')\n    ).strip(),\n    fontsize=20,\n    va='bottom',\n    ha='right',\n    linespacing=1\n)\n\nax.set_title(\"\")\n\nfig.savefig(\n    os.path.join(OUTPUT_DIR, \"ltg-hk-p.png\"),\n    bbox_inches='tight'\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
}