{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\nConvert to RGB format (vector data)\n===========================\nThis example demonstrates convertion from wind data to uv rgb format (with HKO earth format).\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 image generation\nfrom PIL import Image\n# Python package for image preview\nimport matplotlib.pyplot as plt\n\n# swirlspy data convertion function\nfrom swirlspy.utils.conversion import to_rgb_data, to_hko_earth_format\n# directory constants\nfrom swirlspy.tests.samples import DATA_DIR\nfrom swirlspy.tests.outputs import OUTPUT_DIR\n\nwarnings.filterwarnings(\"ignore\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Define the working directories:\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "data_dir = os.path.join(DATA_DIR, 'gts')\n\n# Logging\nstart_time = pd.Timestamp.now()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Loading wind data\n-----------------------------------------------------------------------\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Specify the basetime\nbasetime = pd.Timestamp('202107180000')\n\n# Reading the wind data\nds = xr.open_dataset(\n    os.path.join(\n        data_dir,\n        f\"romas_plv_{basetime.strftime('%Y%m%d%H%M')}.nc\"\n    )\n)\n\n# pick pressure level, e.g. 200 in this case\nds_200 = ds.sel(isobaricInhPa=200)\nu = ds_200['u'].data\nv = ds_200['v'].data\ny = ds['latitude'].data\nx = ds['longitude'].data\n\n# Close dataset\nds.close()\n\ninitialising_time = pd.Timestamp.now()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Extract all necessary values for rgb image\n-----------------------------------------------------------------------\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# prepare placeholder for necessary values\nd = np.empty((len(y), len(x), 2))\n\nd[:] = 0\nd[:, :, 0] = u\nd[:, :, 1] = v\n\n# calculate x, y step size\ny_step = (y[-1] - y[0]) / len(y)\nx_step = (x[-1] - x[0]) / len(x)\n\n# build image data\ndata = xr.DataArray(\n    d,\n    dims=['y', 'x', 'wind'],\n    coords={\n        'wind': ['u', 'v'],\n        'y': y,\n        'x': x\n    },\n    attrs={\n        'y_step': y_step,\n        'x_step': x_step\n    }\n)\n\npreparation_time = pd.Timestamp.now()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Convert image data into rgb format\n-----------------------------------------------------------------------\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# this step is not necessary, depends on any meta data and preprocess is required by your platform\nearth_data = to_hko_earth_format(data, coords_dp=3, values_dp=0, uv_mode=True, y_step=y_step, x_step=x_step)\nrgb = to_rgb_data(earth_data, uv_mode=True)\n\nconvertion_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": [
        "path = os.path.join(OUTPUT_DIR, \"rgb_wind.png\")\nwith Image.fromarray(rgb, 'RGBA') as img:\n    img.save(path, 'png')\n\n# preview\nwith Image.open(path) as image:\n    plt.axis('off')\n    plt.imshow(image)\n    plt.plot()\n    plt.show()\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\"Preparation time: {preparation_time}\")\nprint(f\"Convertion time: {convertion_time}\")\nprint(f\"Visualise time: {visualise_time}\")\n\nprint(f\"Time to initialise: {initialising_time - start_time}\")\nprint(f\"Time to prepare information: {preparation_time - initialising_time}\")\nprint(f\"Time to convertion: {convertion_time - preparation_time}\")\nprint(f\"Time to visualise: {visualise_time - convertion_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
}