{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Flow velocities and streamlines\n\nIn this section we will plot flow velocities and streamlines for some model results.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\nimport numpy as np\nimport xarray as xr"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We'll start with the usual imports\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import imod"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Load and unpack the data\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "ds_fluxes = imod.data.fluxes()\nds_fluxes = ds_fluxes.isel(time=-1)\n\nds_fluxes\n\nlower = ds_fluxes[\"bdgflf\"]\nright = ds_fluxes[\"bdgfrf\"]\nfront = ds_fluxes[\"bdgfff\"]\nheads = ds_fluxes[\"head\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Calculating flow velocity\n\nThe imod-python function imod.evaluate.flow_velocity() computes flow\nvelocities in m/d based on the budget results (bdgflf - flow lower face,\nbdgfrf - flow right face and bdgfff - flow front face). To apply this\nfunction, we first need to define a top_bot array.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "top_bot = xr.full_like(lower, 1.0)\ntop_bot[\"top\"] = top_bot[\"z\"] - 0.5 * top_bot[\"dz\"]\ntop_bot[\"bot\"] = top_bot[\"z\"] + 0.5 * top_bot[\"dz\"]\ntop_bot"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Next we'll calculate the velocities and plot a cross-section of the vertical\nvelocity at for y = 450050\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fig, ax = plt.subplots()\nvx, vy, vz = imod.evaluate.flow_velocity(\n    front=front, lower=lower, right=right, top_bot=top_bot, porosity=0.3\n)\nvz.sel(y=450050.0, method=\"nearest\").plot(cmap=\"RdYlBu\", yincrease=False)\n\nplt.title(\"Vz\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Quiver plot\n\nIt is also possible to make quiver plots for a cross section defined by two\npairs of coordinates. We will first define arrays indicating the start and end\nlocation of the cross section to be evaluated.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "start = np.array([97132.710, 457177.928])\nend = np.array([103736.517, 457215.557])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Using the function ``imod.evaluate.quiver_line()`` considering the starting and ending points\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "u, v = imod.evaluate.quiver_line(right, front, lower, start, end)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Adding top and bottom information to the previously obtained arrays\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "u[\"top\"] = u[\"z\"] + 0.5 * u[\"dz\"]\nu[\"bottom\"] = u[\"z\"] - 0.5 * u[\"dz\"]\n\nv[\"top\"] = v[\"z\"] + 0.5 * v[\"dz\"]\nv[\"bottom\"] = v[\"z\"] - 0.5 * v[\"dz\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Defining a cross section that shows the heads for the same location where the start and end points\nwere defined, to use as a background image for the plot\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "cross_section = imod.select.cross_section_line(heads, start, end)\ncross_section[\"top\"] = cross_section[\"z\"] - 0.5 * cross_section[\"dz\"]\ncross_section[\"bottom\"] = cross_section[\"z\"] + 0.5 * cross_section[\"dz\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Ploting the cross section\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "colors = \"magma\"\nlevels = np.arange(-8, 0.0, 0.5)\nskip = (slice(None, None, 2), slice(None, None, 2))\nfig, ax = plt.subplots()\nfig, ax = imod.visualize.cross_section(\n    cross_section, colors=colors, levels=levels, ax=ax, fig=fig\n)\nimod.visualize.quiver(u[skip], v[skip], ax, kwargs_quiver={\"color\": \"k\"})"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Streamline function\n\nThis function shows the streamlines for a line cross section through a 3D flow field.\nWe will use the previously created arrays that indicate the start and end location of the cross section and apply them in the function imod.evaluate.streamfunction_line()\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "streamfunction = imod.evaluate.streamfunction_line(right, front, start, end)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The previous array contains the streamfunction projected on the cross-section defined\nby provided linestring, with new dimension \u201cs\u201d along the cross-section.\nThe cellsizes along \u201cs\u201d are given in the \u201cds\u201d coordinate.\nThe streamline function can be plotted using imod.visualize.streamfunction(),\nbut first we need to define the 'top' and 'bottom' of the layers in the array\nNote: By default the steamlines are plotted in white.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "streamfunction[\"bottom\"] = streamfunction[\"z\"] + streamfunction[\"dz\"]\nstreamfunction[\"top\"] = streamfunction[\"z\"]\n\nfig, ax = plt.subplots()\nfig, ax = imod.visualize.cross_section(\n    cross_section, colors=colors, levels=levels, ax=ax, fig=fig\n)\nimod.visualize.streamfunction(streamfunction, ax=ax, n_streamlines=10)"
      ]
    }
  ],
  "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.10.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}