{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# River runoff\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import os\nfrom datetime import datetime, timedelta\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom matplotlib.dates import DateFormatter\nimport opendrift\nfrom opendrift.models.oceandrift import OceanDrift\nfrom opendrift.readers import reader_oscillating\n\n\noutfile = 'runoff.nc'  # Raw simulation output\nhistogram_file = 'runoff_histogram.nc'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "First make a simulation with two seedings, marked by *origin_marker*\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "o = OceanDrift(loglevel=20)\no.set_config('drift:horizontal_diffusivity', 300)\no.set_config('general:coastline_action', 'previous')\nt1 = datetime.now()\nt2 = t1 + timedelta(hours=48)\nreader_x = reader_oscillating.Reader('x_sea_water_velocity', period_seconds=3600*24,\n                amplitude=1, zero_time=t1)\nreader_y = reader_oscillating.Reader('y_sea_water_velocity', period_seconds=3600*72,\n                amplitude=.5, zero_time=t2)\no.add_reader([reader_x, reader_y])\nnumber = 25000\no.seed_elements(time=[t1, t2], lon=9.017931, lat=58.562702, number=number,\n                origin_marker_name='River 1')  # River 1\no.seed_elements(time=[t1, t2], lon=8.824815, lat=58.425648, number=number,\n                origin_marker_name='River 2')  # River 2\nseed_times = o.elements_scheduled_time[0:number]\n\no.run(duration=timedelta(hours=48),\n      time_step=1800, time_step_output=3600, outfile=outfile)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Opening the output file lazily with Xarray.\nThis will work even if the file is too large to fit in memory, as it\nwill read and process data chuck-by-chunk directly from file using Dask.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "o = opendrift.open_xarray(outfile)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We want to extract timeseries of river water at the coordinates of a hypothetical measuring station\nas well as the amount of river water passing through two defined areas/regions\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "station_lon = 9.4\nstation_lat = 58.1\nbox1_lon = [8.4, 8.8]\nbox1_lat = [57.9, 58.1]\nbox2_lon = [9.5, 9.9]\nbox2_lat = [58.3, 58.5]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Animation of the spatial density of river runoff water.\nAlthough there are the same number of elements from each river, the density plots are\nweighted with the actual runoff at time of seeding. This weighting can be done/changed\nafterwards without needing to redo the simulation.\nThe calculated density fields will be stored/cached in the analysis file\nfor later re-use, as their calculation may be time consuming\nfor huge output files.\nNote that other analysis/plotting methods are not yet adapted\nto datasets opened lazily with open_xarray\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "runoff_river1 = np.abs(np.cos(np.arange(number)*2*np.pi/(number)))  # Impose a temporal variation of runoff\nrunoff_river2 = 10*runoff_river1  # Let river 2 have 10 times as large runoff as river 1\nrunoff = np.concatenate((runoff_river1, runoff_river2))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Calculate density with given pixel size, weighted by runoff amount per element\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "river_water = o.get_histogram(pixelsize_m=1500, weights=runoff, density=False)\n\nrw = river_water.sum(dim='origin_marker')  # For both rivers\n#rw = river_water.isel(origin_marker=1)    # For one of the rivers\nriver_water.name = 'River water [m3/cell]'\n\ntext = [{'s': o.origin_marker[0], 'x': 8.55, 'y': 58.56, 'fontsize': 20, 'color': 'g',\n         'backgroundcolor': 'white', 'bbox': dict(facecolor='white', alpha=0.8), 'zorder': 1000},\n        {'s': o.origin_marker[1], 'x': 8.35, 'y': 58.42, 'fontsize': 20, 'color': 'g',\n         'backgroundcolor': 'white', 'bbox': dict(facecolor='white', alpha=0.8), 'zorder': 1000},\n        {'s': '* Station', 'x': station_lon, 'y': station_lat, 'fontsize': 20, 'color': 'k',\n         'backgroundcolor': 'white', 'bbox': dict(facecolor='none', edgecolor='none', alpha=0.4), 'zorder': 1000}]\nbox = [{'lon': box1_lon, 'lat': box1_lat, 'text': 'Area 1', 'fc': 'none', 'alpha': 0.8, 'lw': 1, 'ec': 'k'},\n       {'lon': box2_lon, 'lat': box2_lat, 'text': 'Area 2', 'fc': 'none', 'alpha': 0.8, 'lw': 1, 'ec': 'k'}]\n\no.animation(background=rw.where(rw>0), bgalpha=1, fast=False,\n            show_elements=False, vmin=0, vmax=120, text=text, box=box)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<img src=\"file://gallery/animations/example_river_runoff_0.gif\">\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Plotting time series of river runoff, and corresponding water passing through the station and the two defined areas/boxes\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "fig, (ax1, ax2, ax3, ax4) = plt.subplots(4, 1)\n# Runoff\nax1.plot(seed_times, runoff_river1, label=o.origin_marker[0])\nax1.plot(seed_times, runoff_river2, label=o.origin_marker[1])\nax1.set_ylabel('Runoff  [m3/s]')\nax1.set_title('Runoff')\n# Area 1\nt1 = river_water.sel(lon_bin=slice(box1_lon[0], box1_lon[1]),\n                     lat_bin=slice(box1_lat[0], box1_lat[1]))\nt1 = t1.sum(('lon_bin', 'lat_bin'))\nt1.isel(origin_marker=0).plot(label=o.origin_marker[0], ax=ax2)\nt1.isel(origin_marker=1).plot(label=o.origin_marker[1], ax=ax2)\nt1.sum(dim='origin_marker').plot(label='Total', linestyle='--', ax=ax2)\nax2.set_title('Amount of water passing through Area 1')\n# Area 2\nt2 = river_water.sel(lon_bin=slice(box2_lon[0], box2_lon[1]),\n                     lat_bin=slice(box2_lat[0], box2_lat[1]))\nt2 = t2.sum(('lon_bin', 'lat_bin'))\nt2.isel(origin_marker=0).plot(label=o.origin_marker[0], ax=ax3)\nt2.isel(origin_marker=1).plot(label=o.origin_marker[1], ax=ax3)\nt2.sum(dim='origin_marker').plot(label='Total', linestyle='--', ax=ax3)\nax3.set_title('Amount of water passing through Area 2')\n# Extracting time series at the location of the station\nt = river_water.sel(lon_bin=station_lon, lat_bin=station_lat, method='nearest')\nt.isel(origin_marker=0).plot(label=o.origin_marker[0], ax=ax4)\nt.isel(origin_marker=1).plot(label=o.origin_marker[1], ax=ax4)\nt.sum(dim='origin_marker').plot(label='Total', linestyle='--', ax=ax4)\nax4.legend()\nax4.margins(x=0)\n\nfor ax in [ax1, ax2, ax3]:\n    ax.margins(x=0)\n    ax.legend()\n    #ax.set_xticks([])\n    ax.set_xlabel(None)\nax4.set_title('Density of water at Station')\n# TODO disabling due to recent problem with dateformatter\n#ax4.xaxis.set_major_formatter(DateFormatter(\"%d %b %H\"))\nplt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Finally, plot the spatial distribution of mean age of water \n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "num = o.get_histogram(pixelsize_m=1500, weights=None, density=False)\nnum.name = 'number'\nnum.to_netcdf(histogram_file)\n\nmas = o.get_histogram(pixelsize_m=1500, weights=o.ds['age_seconds'], density=False)\nmas = mas/3600  # in hours\nmas = mas/num  # per area\nmas.name='mean_age'\nmas.to_netcdf(histogram_file, 'a')\nmas = mas.mean(dim='time').sum(dim='origin_marker')  # Mean time of both rivers\n#mas = mas.mean(dim='time').isel(origin_marker=1)  # Mean age of a single river\nmas.name='Mean age of water [hours]'\n\no.plot(background=mas.where(mas>0), fast=True, show_elements=False)\n\n\n# Cleaning up\nos.remove(outfile)\nos.remove(histogram_file)"
      ]
    }
  ],
  "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.11.6"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}