{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Shear Extrapolation\n\nThis example shows how to perform shear extrapolation of a wind speed time series.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Simple shear extrapolation\nFirst, we need to prepare the sample data.\nLet's create some wind speed time series data with tree heights present.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import numpy as np\nimport pandas as pd\nimport xarray as xr\nimport windkit as wk\nimport matplotlib.pyplot as plt\n\nwind_speed = xr.DataArray(\n    np.array([[10.0, 12.0, 14.0], [10.5, 12.5, 14.5], [11.0, 13.0, 15.0]]).T,\n    dims=[\"time\", \"height\"],\n    coords={\n        \"time\": pd.date_range(\"2023-01-01\", periods=3, freq=\"h\"),\n        \"height\": [10.0, 30.0, 40.0],\n    },\n)\n\nprint(wind_speed)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's say we want to extrapolate this data to 100 m using a fixed shear exponent of 0.143.\nTo do that, we can use `windkit.shear_extrapolate`.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "wind_speed_new = wk.shear_extrapolate(wind_speed, 100, shear_exponent=0.143)\n\nprint(wind_speed_new)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "When more heights are present in the orignal wind speed DataArray,\nthe height closest to the requested height is used.\n\nMultiple target heights can be requested when the wind speed data has structured heights.:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "wind_speed_new = wk.shear_extrapolate(wind_speed, [50, 100, 200], shear_exponent=0.143)\n\nprint(wind_speed_new)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Varying shear\nThe shear exponent can be varied, including adding shear at different heights.\nWhen the shear exponent vary with the height, the shear nearest to the reference\nheights used, is taken, e.g., if the original wind speed is at 40.0 meters and the\nshear is at 30 and 60 meters, the shear at 30 is used.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "shear_exponent = xr.DataArray(\n    np.array([[0.113, 0.113, 0.113], [0.123, 0.123, 0.123], [0.133, 0.133, 0.133]]).T,\n    dims=[\"time\", \"height\"],\n    coords={\n        \"time\": pd.date_range(\"2023-01-01\", periods=3, freq=\"h\"),\n        \"height\": [15.0, 20.0, 35.0],\n    },\n)\n\nwind_speed_new = wk.shear_extrapolate(\n    wind_speed, [15, 25, 45], shear_exponent=shear_exponent\n)\n\nprint(wind_speed_new)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Unstructured data\nThe input wind speed can also be at unstructured heights, a-la datasets with \"point\"-structure\n, e.g., from a LiDAR. The shear extrapolation still works as expected.\nHowever, the only one height can be requested, or varying heights that match\nthe dimensions of the wind speed data.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "wind_speed = xr.DataArray(\n    np.array([[10.0, 10.0], [11.0, 11.0], [12.0, 12.0], [13.0, 13.0]]).T,\n    dims=(\"time\", \"point\"),\n    coords={\n        \"time\": pd.date_range(\"2023-01-01\", periods=2, freq=\"h\"),\n        \"height\": ((\"point\",), [10, 20.0, 30.0, 40.0]),\n        \"west_east\": ((\"point\",), [0, 0, 1, 2]),\n        \"south_north\": ((\"point\",), [-3, -3, -2, -1]),\n    },\n)\nwind_speed_new = wk.shear_extrapolate(wind_speed, 100.0, shear_exponent=0.143)\nprint(wind_speed_new)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "When the wind speeds are at unstructured heights, only one target height can be used, or\nvarying heights that match the dimensions of the wind speed data (e.g., different heights at different points)\nand similarly for the shear exponent. Let's say we want to extrapolate to different heights\nat different points, with a time-varying shear exponent. We can do that like this:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "height = xr.DataArray(\n    np.array([10, 50.0, 150.0, 250.0]),\n    dims=[\"point\"],\n    coords={\n        \"height\": ((\"point\",), [10, 50.0, 150.0, 250.0]),\n        \"west_east\": ((\"point\",), [0, 0, 1, 2]),\n        \"south_north\": ((\"point\",), [-3, -3, -2, -1]),\n    },\n)\n\nshear_exponent = xr.DataArray(\n    np.array([-0.113, 0.113]),\n    dims=[\"time\"],\n    coords={\n        \"time\": pd.date_range(\"2023-01-01\", periods=2, freq=\"h\"),\n    },\n)\n\nwind_speed_new = wk.shear_extrapolate(wind_speed, height, shear_exponent=shear_exponent)\nprint(wind_speed_new)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Custom shear exponent calculation\nThe shear exponent can also be calculated from the wind speed data itself,\nusing the `windkit.shear_exponent` function. This function computes the shear\nexponent from vertical wind speed profiles using finite differences in log-space.\nIf only a single height is present for a given horizontal location, the shear exponent\nis set to NaN for that location (so for points ``(1,-2)`` and ``(2,-1)`` in the below example).\nIn the example below, we first compute the shear exponent\nfrom the wind speed data, interpolate the shear to the original data structure and\nfinally use it for shear extrapolation to 100 m. Similar to before, the wind speed\ncan only be extrapolated to a single height when the data is at unstructured heights.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "wind_speed = xr.DataArray(\n    np.array([[10.0, 10.0], [11.0, 11.0], [12.0, 12.0], [13.0, 13.0]]).T,\n    dims=(\"time\", \"point\"),\n    coords={\n        \"time\": pd.date_range(\"2023-01-01\", periods=2, freq=\"h\"),\n        \"height\": ((\"point\",), [10, 40.0, 30.0, 40.0]),\n        \"west_east\": ((\"point\",), [0, 0, 1, 2]),\n        \"south_north\": ((\"point\",), [-3, -3, -2, -1]),\n    },\n)\nwind_speed = wk.spatial.set_crs(wind_speed, 4326)\nshear_exp = wk.shear_exponent(wind_speed)\nshear_exp = wk.spatial.interp_unstructured_like(shear_exp, wind_speed, method=\"nearest\")\nwind_speed_new = wk.shear_extrapolate(wind_speed, 100.0, shear_exponent=shear_exp)\nprint(wind_speed_new)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Finally, we can also extrapolate to varying heights when the shear exponent\nis calculated from the wind speed data itself. Here, we create a new dataset\nwith different heights at each point, and use the shear exponent calculated\nfrom the wind speed data for shear extrapolation.\nAgain, the shear exponent is interpolated to the new data structure before use.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "wind_speed = xr.DataArray(\n    np.array([[10.0, 10.0], [11.0, 11.0], [12.0, 12.0], [13.0, 13.0]]).T,\n    dims=(\"time\", \"point\"),\n    coords={\n        \"time\": pd.date_range(\"2023-01-01\", periods=2, freq=\"h\"),\n        \"height\": ((\"point\",), [10, 20.0, 30.0, 40.0]),\n        \"west_east\": ((\"point\",), [0, 0, 0, 0]),\n        \"south_north\": ((\"point\",), [-3, -3, -3, -3]),\n    },\n)\nwind_speed = wk.spatial.set_crs(wind_speed, 4326)\nwind_speed = wk.spatial.to_stacked_point(wind_speed)\nshear_exp = wk.shear_exponent(wind_speed)\nds = wk.spatial.create_stacked_point(\n    [0, 0, 0, 0], [-3, -3, -3, -3], [15, 25.0, 35.0, 45.0], 4326\n)\nwind_speed_new = wk.shear_extrapolate(wind_speed, ds.height, shear_exponent=shear_exp)\nprint(wind_speed_new)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Wind Veer Extrapolation\nThe wind direction can also be extrapolated using the `windkit.veer_extrapolate` function.\nThis requires a wind veer value, which can be constant or calculated from the data\nusing `windkit.wind_veer`.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "wind_direction = xr.DataArray(\n    np.array([[270.0, 280.0], [275.0, 285.0]]).T,\n    dims=[\"time\", \"height\"],\n    coords={\n        \"time\": pd.date_range(\"2023-01-01\", periods=2, freq=\"h\"),\n        \"height\": [10.0, 50.0],\n    },\n).expand_dims(stacked_point=[0])\nwind_direction[\"west_east\"] = (\"stacked_point\", [0])\nwind_direction[\"south_north\"] = (\"stacked_point\", [0])\n\n# Calculate wind veer\nveer = wk.wind_veer(wind_direction)\n\n# Extrapolate to new height\nwd_new = wk.veer_extrapolate(wind_direction, [75.0, 100.0], veer=veer)\n\n# Plot the results for the first time step\nplt.figure()\nwind_direction.isel(time=0).plot(y=\"height\", marker=\"o\", label=\"Input\")\nwd_new.isel(time=0).plot(y=\"height\", marker=\"x\", label=\"Extrapolated\")\nplt.legend()\nplt.title(\"Wind Direction Extrapolation\")"
      ]
    }
  ],
  "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.14.3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}