{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Optical flow methods convergence\n\nIn this example we test the convergence of the optical flow methods available in\npySTEPS using idealized motion fields.\n\nTo test the convergence, using an example precipitation field we will:\n\n- Read precipitation field from a file\n- Morph the precipitation field using a given motion field (linear or rotor) to\n  generate a sequence of moving precipitation patterns.\n- Using the available optical flow methods, retrieve the motion field from the\n  precipitation time sequence (synthetic precipitation observations).\n\nLet's first load the libraries that we will use.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from datetime import datetime\nimport time\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nfrom matplotlib.cm import get_cmap\nfrom scipy.ndimage import uniform_filter\n\nimport pysteps as stp\nfrom pysteps import motion, io, rcparams\nfrom pysteps.motion.vet import morph\nfrom pysteps.visualization import plot_precip_field, quiver"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Load the reference precipitation data\n\nFirst, we will import a radar composite from the archive.\nYou need the pysteps-data archive downloaded and the pystepsrc file\nconfigured with the data_source paths pointing to data folders.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Selected case\ndate = datetime.strptime(\"201505151630\", \"%Y%m%d%H%M\")\ndata_source = rcparams.data_sources[\"mch\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Load the data from the archive\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "root_path = data_source[\"root_path\"]\npath_fmt = data_source[\"path_fmt\"]\nfn_pattern = data_source[\"fn_pattern\"]\nfn_ext = data_source[\"fn_ext\"]\nimporter_name = data_source[\"importer\"]\nimporter_kwargs = data_source[\"importer_kwargs\"]\n\n# Find the reference field in the archive\nfns = io.archive.find_by_date(\n    date, root_path, path_fmt, fn_pattern, fn_ext, timestep=5, num_prev_files=0\n)\n\n# Read the reference radar composite\nimporter = io.get_method(importer_name, \"importer\")\nreference_field, quality, metadata = io.read_timeseries(\n    fns, importer, **importer_kwargs\n)\n\ndel quality  # Not used\n\nreference_field = np.squeeze(reference_field)  # Remove time dimension"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Preprocess the data\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Convert to mm/h\nreference_field, metadata = stp.utils.to_rainrate(reference_field, metadata)\n\n# Mask invalid values\nreference_field = np.ma.masked_invalid(reference_field)\n\n# Plot the reference precipitation\nplot_precip_field(reference_field, title=\"Reference field\")\nplt.show()\n\n# Log-transform the data [dBR]\nreference_field, metadata = stp.utils.dB_transform(\n    reference_field, metadata, threshold=0.1, zerovalue=-15.0\n)\n\nprint(\"Precip. pattern shape: \" + str(reference_field.shape))\n\n# This suppress nan conversion warnings in plot functions\nreference_field.data[reference_field.mask] = np.nan"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Synthetic precipitation observations\n\nNow we need to create a series of precipitation fields by applying the ideal\nmotion field to the reference precipitation field \"n\" times.\n\nTo evaluate the accuracy of the computed_motion vectors, we will use\na relative RMSE measure.\nRelative MSE = <(expected_motion - computed_motion)^2> / <expected_motion^2>\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Relative RMSE = Rel_RMSE = sqrt(Relative MSE)\n#\n# - Rel_RMSE = 0%: no error\n# - Rel_RMSE = 100%: The retrieved motion field has an average error equal in\n#   magnitude to the motion field.\n#\n# Relative RMSE is computed over a region surrounding the precipitation\n# field, were there is enough information to retrieve the motion field.\n# The \"precipitation region\" includes the precipitation pattern plus a margin of\n# approximately 20 grid points."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's create a function to construct different motion fields.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def create_motion_field(input_precip, motion_type):\n    \"\"\"\n    Create idealized motion fields to be applied to the reference image.\n\n    Parameters\n    ----------\n\n    input_precip: numpy array (lat, lon)\n\n    motion_type: str\n        The supported motion fields are:\n\n            - linear_x: (u=2, v=0)\n            - linear_y: (u=0, v=2)\n            - rotor: rotor field\n\n    Returns\n    -------\n    ideal_motion : numpy array (u, v)\n    \"\"\"\n\n    # Create an imaginary grid on the image and create a motion field to be\n    # applied to the image.\n    ny, nx = input_precip.shape\n\n    x_pos = np.arange(nx)\n    y_pos = np.arange(ny)\n    x, y = np.meshgrid(x_pos, y_pos, indexing=\"ij\")\n\n    ideal_motion = np.zeros((2, nx, ny))\n\n    if motion_type == \"linear_x\":\n        ideal_motion[0, :] = 2  # Motion along x\n    elif motion_type == \"linear_y\":\n        ideal_motion[1, :] = 2  # Motion along y\n    elif motion_type == \"rotor\":\n        x_mean = x.mean()\n        y_mean = y.mean()\n        norm = np.sqrt(x * x + y * y)\n        mask = norm != 0\n        ideal_motion[0, mask] = 2 * (y - y_mean)[mask] / norm[mask]\n        ideal_motion[1, mask] = -2 * (x - x_mean)[mask] / norm[mask]\n    else:\n        raise ValueError(\"motion_type not supported.\")\n\n    # We need to swap the axes because the optical flow methods expect\n    # (lat, lon) or (y,x) indexing convention.\n    ideal_motion = ideal_motion.swapaxes(1, 2)\n    return ideal_motion"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's create another function that construct the temporal series of\nprecipitation observations.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def create_observations(input_precip, motion_type, num_times=9):\n    \"\"\"\n    Create synthetic precipitation observations by displacing the input field\n    using an ideal motion field.\n\n    Parameters\n    ----------\n\n    input_precip: numpy array (lat, lon)\n        Input precipitation field.\n\n    motion_type: str\n        The supported motion fields are:\n\n            - linear_x: (u=2, v=0)\n            - linear_y: (u=0, v=2)\n            - rotor: rotor field\n\n    num_times: int, optional\n        Length of the observations sequence.\n\n\n    Returns\n    -------\n    synthetic_observations: numpy array\n        Sequence of observations\n    \"\"\"\n\n    ideal_motion = create_motion_field(input_precip, motion_type)\n\n    # The morph function expects (lon, lat) or (x, y) dimensions.\n    # Hence, we need to swap the lat,lon axes.\n\n    # NOTE: The motion field passed to the morph function can't have any NaNs.\n    # Otherwise, it can result in a segmentation fault.\n    morphed_field, mask = morph(\n        input_precip.swapaxes(0, 1), ideal_motion.swapaxes(1, 2)\n    )\n\n    mask = np.array(mask, dtype=bool)\n\n    synthetic_observations = np.ma.MaskedArray(morphed_field, mask=mask)\n    synthetic_observations = synthetic_observations[np.newaxis, :]\n\n    for t in range(1, num_times):\n        morphed_field, mask = morph(\n            synthetic_observations[t - 1], ideal_motion.swapaxes(1, 2)\n        )\n        mask = np.array(mask, dtype=bool)\n\n        morphed_field = np.ma.MaskedArray(\n            morphed_field[np.newaxis, :], mask=mask[np.newaxis, :]\n        )\n\n        synthetic_observations = np.ma.concatenate(\n            [synthetic_observations, morphed_field], axis=0\n        )\n\n    # Swap  back to (lat, lon)\n    synthetic_observations = synthetic_observations.swapaxes(1, 2)\n\n    synthetic_observations = np.ma.masked_invalid(synthetic_observations)\n\n    synthetic_observations.data[np.ma.getmaskarray(synthetic_observations)] = 0\n\n    return ideal_motion, synthetic_observations\n\n\ndef plot_optflow_method_convergence(input_precip, optflow_method_name, motion_type):\n    \"\"\"\n    Test the convergence to the actual solution of the optical flow method used.\n\n    Parameters\n    ----------\n\n    input_precip: numpy array (lat, lon)\n        Input precipitation field.\n\n    optflow_method_name: str\n        Optical flow method name\n\n    motion_type: str\n        The supported motion fields are:\n\n            - linear_x: (u=2, v=0)\n            - linear_y: (u=0, v=2)\n            - rotor: rotor field\n    \"\"\"\n\n    if optflow_method_name.lower() != \"darts\":\n        num_times = 2\n    else:\n        num_times = 9\n\n    ideal_motion, precip_obs = create_observations(\n        input_precip, motion_type, num_times=num_times\n    )\n\n    oflow_method = motion.get_method(optflow_method_name)\n\n    elapsed_time = time.perf_counter()\n\n    computed_motion = oflow_method(precip_obs, verbose=False)\n\n    print(\n        f\"{optflow_method_name} computation time: \"\n        f\"{(time.perf_counter() - elapsed_time):.1f} [s]\"\n    )\n\n    precip_obs, _ = stp.utils.dB_transform(precip_obs, inverse=True)\n\n    precip_data = precip_obs.max(axis=0)\n    precip_data.data[precip_data.mask] = 0\n\n    precip_mask = (uniform_filter(precip_data, size=20) > 0.1) & ~precip_obs.mask.any(\n        axis=0\n    )\n\n    cmap = get_cmap(\"jet\").copy()\n    cmap.set_under(\"grey\", alpha=0.25)\n    cmap.set_over(\"none\")\n\n    # Compare retrieved motion field with the ideal one\n    plt.figure(figsize=(9, 4))\n    plt.subplot(1, 2, 1)\n    ax = plot_precip_field(precip_obs[0], title=\"Reference motion\")\n    quiver(ideal_motion, step=25, ax=ax)\n\n    plt.subplot(1, 2, 2)\n    ax = plot_precip_field(precip_obs[0], title=\"Retrieved motion\")\n    quiver(computed_motion, step=25, ax=ax)\n\n    # To evaluate the accuracy of the computed_motion vectors, we will use\n    # a relative RMSE measure.\n    # Relative MSE = < (expected_motion - computed_motion)^2 > / <expected_motion^2 >\n    # Relative RMSE = sqrt(Relative MSE)\n\n    mse = ((ideal_motion - computed_motion)[:, precip_mask] ** 2).mean()\n\n    rel_mse = mse / (ideal_motion[:, precip_mask] ** 2).mean()\n    plt.suptitle(\n        f\"{optflow_method_name} \" f\"Relative RMSE: {np.sqrt(rel_mse) * 100:.2f}%\"\n    )\n    plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Lucas-Kanade\n\n### Constant motion x-direction\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plot_optflow_method_convergence(reference_field, \"LucasKanade\", \"linear_x\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Constant motion y-direction\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plot_optflow_method_convergence(reference_field, \"LucasKanade\", \"linear_y\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Rotational motion\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plot_optflow_method_convergence(reference_field, \"LucasKanade\", \"rotor\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Variational Echo Tracking (VET)\n\n### Constant motion x-direction\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plot_optflow_method_convergence(reference_field, \"VET\", \"linear_x\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Constant motion y-direction\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plot_optflow_method_convergence(reference_field, \"VET\", \"linear_y\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Rotational motion\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plot_optflow_method_convergence(reference_field, \"VET\", \"rotor\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## DARTS\n\n### Constant motion x-direction\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plot_optflow_method_convergence(reference_field, \"DARTS\", \"linear_x\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Constant motion y-direction\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plot_optflow_method_convergence(reference_field, \"DARTS\", \"linear_y\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Rotational motion\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plot_optflow_method_convergence(reference_field, \"DARTS\", \"rotor\")\n\n# sphinx_gallery_thumbnail_number = 5"
      ]
    }
  ],
  "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.8.6"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}