Fastscape - Drainage area - Coarse grids

environment.yml

name: fastscape-demo
channels:
  - conda-forge
  - defaults
dependencies:
  - aiohttp=3.6.2=py38h0b31af3_0
  - ansiwrap=0.8.4=py_0
  - anyio=2.2.0=py38h50d1736_0
  - appdirs=1.4.4=pyh9f0ad1d_0
  - appnope=0.1.0=py38h32f6830_1001
  - argon2-cffi=20.1.0=py38h4d0b108_1
  - asciitree=0.3.3=py_2
  - async-timeout=3.0.1=py_1000
  - async_generator=1.10=py_0
  - atk=2.36.0=2
  - atk-1.0=2.36.0=haf1e3a3_2
  - attrs=20.2.0=pyh9f0ad1d_0
  - babel=2.9.0=pyhd3deb0d_0
  - backcall=0.2.0=pyh9f0ad1d_0
  - backports=1.0=py_2
  - backports.functools_lru_cache=1.6.1=py_0
  - black=20.8b1=py_1
  - bleach=3.2.1=pyh9f0ad1d_0
  - bokeh=2.2.1=py38h32f6830_0
  - brotlipy=0.7.0=py38h64e0658_1000
  - bzip2=1.0.8=hc929b4f_4
  - c-ares=1.17.1=hc929b4f_0
  - ca-certificates=2021.5.30=h033912b_0
  - cairo=1.16.0=ha8983da_1005
  - cartopy=0.18.0=py38h6c003aa_5
  - certifi=2021.5.30=py38h50d1736_0
  - cffi=1.14.3=py38hc4dd44e_0
  - cfgv=3.2.0=py_0
  - cftime=1.3.0=py38hfb243c8_0
  - chardet=3.0.4=py38h32f6830_1007
  - click=7.1.2=pyh9f0ad1d_0
  - cloudpickle=1.6.0=py_0
  - colorcet=2.0.1=py_0
  - cryptography=3.1.1=py38h52adbb4_0
  - curl=7.71.1=hcb81553_8
  - cycler=0.10.0=py_2
  - cytoolz=0.11.0=py38h4d0b108_0
  - dask=2021.6.0=pyhd8ed1ab_0
  - dask-core=2021.6.0=pyhd8ed1ab_0
  - dask-labextension=3.0.0=py_0
  - dataclasses=0.7=py38_0
  - datashader=0.11.1=pyh9f0ad1d_0
  - datashape=0.5.4=py_1
  - dbus=1.13.6=h2f22bb5_0
  - decorator=4.4.2=py_0
  - defusedxml=0.6.0=py_0
  - deprecation=2.1.0=pyh9f0ad1d_0
  - distlib=0.3.1=pyh9f0ad1d_0
  - distributed=2021.6.0=py38h50d1736_0
  - editdistance=0.5.3=py38h11c0d25_2
  - entrypoints=0.3=py38h32f6830_1001
  - expat=2.2.9=hb1e8313_2
  - fasteners=0.14.1=py_3
  - fastscapelib-f2py=2.8.2=py38hcf7560d_0
  - filelock=3.0.12=pyh9f0ad1d_0
  - font-ttf-dejavu-sans-mono=2.37=hab24e00_0
  - font-ttf-inconsolata=2.001=hab24e00_0
  - font-ttf-source-code-pro=2.030=hab24e00_0
  - font-ttf-ubuntu=0.83=hab24e00_0
  - fontconfig=2.13.1=h79c0d67_1002
  - fonts-conda-ecosystem=1=0
  - fonts-conda-forge=1=0
  - freetype=2.10.4=h4cff582_1
  - fribidi=1.0.10=h0b31af3_0
  - fsspec=0.8.3=py_0
  - gdk-pixbuf=2.38.2=h306395f_4
  - geos=3.8.1=h4a8c4bd_0
  - gettext=0.19.8.1=h46ab8bc_1002
  - giflib=5.2.1=h0b31af3_2
  - glib=2.66.1=h39b9ebd_0
  - gobject-introspection=1.66.1=py38h8ccf991_0
  - graphite2=1.3.13=h12caacf_1001
  - graphviz=2.42.3=h055b950_1
  - gtk2=2.24.32=hc8e9e3f_3
  - gts=0.7.6=h2684ab5_0
  - harfbuzz=2.7.2=h8810732_0
  - hdf4=4.2.13=h84186c3_1003
  - hdf5=1.10.6=nompi_hc457bb4_1110
  - heapdict=1.0.1=py_0
  - holoviews=1.14.2=pyhd8ed1ab_0
  - hvplot=0.7.0=pyhd3deb0d_0
  - icu=67.1=hb1e8313_0
  - identify=1.5.11=pyhd3deb0d_0
  - idna=2.10=pyh9f0ad1d_0
  - importlib-metadata=2.0.0=py38h32f6830_0
  - importlib_metadata=2.0.0=0
  - ipydatawidgets=4.2.0=pyhd3deb0d_0
  - ipyfastscape=0.2.0=pyhd8ed1ab_0
  - ipygany=0.5.0=pyhd8ed1ab_0
  - ipykernel=5.5.0=py38h9bb44b7_1
  - ipython=7.10.2=py38h5ca1d4c_0
  - ipython_genutils=0.2.0=py_1
  - ipywidgets=7.6.3=pyhd3deb0d_0
  - jedi=0.17.2=py38h32f6830_0
  - jinja2=2.11.2=pyh9f0ad1d_0
  - jpeg=9d=h0b31af3_0
  - json5=0.9.5=pyh9f0ad1d_0
  - jsonschema=3.2.0=py38h32f6830_1
  - jupyter=1.0.0=py_2
  - jupyter-packaging=0.9.2=pyhd8ed1ab_0
  - jupyter-server-proxy=1.5.0=py_0
  - jupyter_client=6.1.7=py_0
  - jupyter_console=6.2.0=py_0
  - jupyter_core=4.6.3=py38h32f6830_1
  - jupyter_server=1.6.2=py38h50d1736_0
  - jupyterlab=3.0.14=pyhd8ed1ab_0
  - jupyterlab_pygments=0.1.2=pyh9f0ad1d_0
  - jupyterlab_server=2.4.0=pyhd8ed1ab_0
  - jupyterlab_widgets=1.0.0=pyhd8ed1ab_1
  - kiwisolver=1.2.0=py38ha0d09dd_0
  - krb5=1.17.1=h75d18d8_3
  - lcms2=2.11=h174193d_0
  - libblas=3.8.0=17_openblas
  - libcblas=3.8.0=17_openblas
  - libclang=10.0.1=default_hf57f61e_1
  - libcurl=7.71.1=h9bf37e3_8
  - libcxx=10.0.1=h5f48129_0
  - libedit=3.1.20191231=hed1e85f_2
  - libev=4.33=haf1e3a3_1
  - libffi=3.2.1=hb1e8313_1007
  - libgfortran=4.0.0=h2d743fc_11
  - libgfortran4=7.5.0=h2d743fc_11
  - libiconv=1.16=haf1e3a3_0
  - liblapack=3.8.0=17_openblas
  - libllvm10=10.0.1=h009f743_3
  - libnetcdf=4.7.4=nompi_h9d8a93f_107
  - libnghttp2=1.41.0=h7580e61_2
  - libopenblas=0.3.10=openmp_h63d9170_4
  - libpng=1.6.37=hb0a8c7a_2
  - libpq=12.3=h489d428_0
  - libsodium=1.0.18=haf1e3a3_1
  - libssh2=1.9.0=h39bdce6_5
  - libtiff=4.1.0=h2ae36a8_6
  - libuv=1.40.0=haf1e3a3_0
  - libwebp=1.1.0=hd3bf737_4
  - libwebp-base=1.1.0=h0b31af3_3
  - libxml2=2.9.10=h7fdee97_2
  - llvm-openmp=10.0.1=h28b9765_0
  - llvmlite=0.34.0=py38h3707e27_1
  - locket=0.2.0=py_2
  - lz4-c=1.9.2=hb1e8313_3
  - markdown=3.2.2=py_0
  - markupsafe=1.1.1=py38h64e0658_1
  - matplotlib=3.3.4=py38h50d1736_0
  - matplotlib-base=3.3.4=py38h8b3ea08_0
  - mistune=0.8.4=py38h64e0658_1001
  - monotonic=1.5=py_0
  - msgpack-python=1.0.0=py38ha0d09dd_1
  - multidict=4.7.5=py38h64e0658_1
  - multipledispatch=0.6.0=py_0
  - mypy_extensions=0.4.3=py38h50d1736_2
  - mysql-common=8.0.21=2
  - mysql-libs=8.0.21=hfb8f7af_2
  - nbclassic=0.2.7=pyhd8ed1ab_0
  - nbclient=0.5.0=py_0
  - nbconvert=6.0.7=py38h32f6830_0
  - nbformat=5.0.7=py_0
  - ncurses=6.2=hb1e8313_1
  - nest-asyncio=1.4.1=py_0
  - netcdf4=1.5.5.1=nompi_py38h0bc7383_100
  - nodeenv=1.5.0=pyh9f0ad1d_0
  - nodejs=14.13.0=hdde0ff8_0
  - notebook=6.1.4=py38h32f6830_0
  - nspr=4.20=h0a44026_1000
  - nss=3.47=hcec2283_0
  - numba=0.51.2=py38h6be0db6_0
  - numcodecs=0.7.2=py38h11c0d25_0
  - numpy=1.19.1=py38h8ccc501_2
  - olefile=0.46=py_0
  - openssl=1.1.1k=h0d85af4_0
  - packaging=20.4=pyh9f0ad1d_0
  - pandas=1.1.2=py38h11c0d25_0
  - pandoc=2.10.1=haf1e3a3_0
  - pandocfilters=1.4.2=py_1
  - panel=0.9.7=py_0
  - pango=1.42.4=haa940fe_4
  - papermill=2.2.2=pyhd8ed1ab_0
  - param=1.9.3=py_0
  - parso=0.7.1=pyh9f0ad1d_0
  - partd=1.1.0=py_0
  - pathspec=0.8.1=pyhd3deb0d_0
  - pcre=8.44=h4a8c4bd_0
  - pexpect=4.8.0=py38h32f6830_1
  - pickleshare=0.7.5=py38h32f6830_1001
  - pillow=7.2.0=py38h83dc5e5_1
  - pip=20.2.3=py_0
  - pixman=0.38.0=h01d97ff_1003
  - pre-commit=2.9.3=py38h50d1736_0
  - proj=7.1.1=h45baca5_3
  - prometheus_client=0.8.0=pyh9f0ad1d_0
  - prompt-toolkit=3.0.7=py_0
  - prompt_toolkit=3.0.7=0
  - psutil=5.7.2=py38h4d0b108_0
  - ptyprocess=0.6.0=py_1001
  - pycparser=2.20=pyh9f0ad1d_2
  - pyct=0.4.6=py_0
  - pyct-core=0.4.6=py_0
  - pygments=2.7.1=py_0
  - pyopenssl=19.1.0=py38_0
  - pyparsing=2.4.7=pyh9f0ad1d_0
  - pyqt=5.12.3=py38hf180056_3
  - pyrsistent=0.17.3=py38h4d0b108_0
  - pyshp=2.1.3=pyh44b312d_0
  - pysocks=1.7.1=py38h32f6830_1
  - python=3.8.5=h0ed32c4_9_cpython
  - python-dateutil=2.8.1=py_0
  - python-graphviz=0.14.1=pyh9f0ad1d_0
  - python_abi=3.8=1_cp38
  - pytz=2020.1=pyh9f0ad1d_0
  - pyviz_comms=2.0.1=pyhd3deb0d_0
  - pyyaml=5.3.1=py38h64e0658_0
  - pyzmq=19.0.2=py38h2c785a9_1
  - qt=5.12.9=h717870c_0
  - qtconsole=4.7.7=pyh9f0ad1d_0
  - qtpy=1.9.0=py_0
  - readline=8.0=h0678c8f_2
  - regex=2020.11.13=py38h5406a74_0
  - requests=2.24.0=pyh9f0ad1d_0
  - scipy=1.5.2=py38h1402333_0
  - selenium=3.141.0=py38h5406a74_1002
  - send2trash=1.5.0=py_0
  - setuptools=49.6.0=py38h32f6830_1
  - shapely=1.7.1=py38h8918236_1
  - simpervisor=0.3=py_1
  - six=1.15.0=pyh9f0ad1d_0
  - sniffio=1.2.0=py38h50d1736_1
  - sortedcontainers=2.2.2=pyh9f0ad1d_0
  - sqlite=3.33.0=h960bd1c_0
  - tblib=1.6.0=py_0
  - tenacity=6.3.1=pyhd8ed1ab_0
  - terminado=0.9.1=py38h32f6830_0
  - testpath=0.4.4=py_0
  - textwrap3=0.9.2=py_0
  - tk=8.6.10=hbbe82c9_0
  - toml=0.10.2=pyhd8ed1ab_0
  - tomlkit=0.7.0=py38h50d1736_3
  - toolz=0.11.1=py_0
  - tornado=6.1=py38h5406a74_1
  - tqdm=4.50.0=pyh9f0ad1d_0
  - traitlets=5.0.4=py_1
  - traittypes=0.2.1=pyh9f0ad1d_2
  - typed-ast=1.4.2=py38h5406a74_0
  - typing_extensions=3.7.4.2=py_0
  - urllib3=1.25.10=py_0
  - virtualenv=20.2.2=py38h50d1736_0
  - wcwidth=0.2.5=pyh9f0ad1d_2
  - webencodings=0.5.1=py_1
  - wheel=0.35.1=pyh9f0ad1d_0
  - widgetsnbextension=3.5.1=py38h32f6830_1
  - xarray=0.18.2=pyhd8ed1ab_0
  - xarray-simlab=0.5.0=pyhd8ed1ab_0
  - xarray-spatial=0.1.2=pyhd3deb0d_0
  - xmovie=0.2.2=pyhd8ed1ab_0
  - xz=5.2.5=haf1e3a3_1
  - yaml=0.2.5=haf1e3a3_0
  - yarl=1.3.0=py38h0b31af3_1000
  - zarr=2.4.0=py_0
  - zeromq=4.3.3=hb1e8313_1
  - zict=2.0.0=py_0
  - zipp=3.3.0=py_0
  - zlib=1.2.11=h7795811_1009
  - zstd=1.4.5=h0384e3a_2
  - pip:
    - fastscape==0.1.0b2+1.g63b0f13
    - pyqt5-sip==4.19.18
    - pyqtchart==5.12
    - pyqtwebengine==5.12.1

fan.ipynb

flow_routing.py

"""A set of utility functions for flow routing.

Note:

- single flow direction only (on a regular grid with 8 neighbors connectivity)
- no handling of closed depression areas

"""


import numba
import numpy as np


@numba.njit
def reset_receivers(receivers, nnodes):
    for inode in range(nnodes):
        receivers[inode] = inode


@numba.njit
def compute_receivers_d8(receivers, dist2receivers, elevation,
                         nx, ny, dx, dy):
    # queen (D8) neighbor lookup
    dr = np.array([-1, -1, -1, 0, 0, 1, 1, 1], dtype=np.intp)
    dc = np.array([-1, 0, 1, -1, 1, -1, 0, 1], dtype=np.intp)

    length = np.sqrt((dy * dr)**2 + (dx * dc)**2)

    tiny = np.finfo(elevation.dtype).tiny

    for r in range(1, ny - 1):
        for c in range(1, nx - 1):
            inode = r * nx + c
            slope_max = tiny

            for k in range(8):
                ineighbor = (r + dr[k]) * nx + (c + dc[k])
                slope = (elevation[inode] - elevation[ineighbor]) / length[k]

                if slope > slope_max:
                    slope_max = slope
                    receivers[inode] = ineighbor
                    dist2receivers[inode] = length[k]


@numba.njit
def compute_donors(ndonors, donors, receivers, nnodes):
    ndonors[:] = 0

    for inode in range(nnodes):
        if receivers[inode] != inode:
            irec = receivers[inode]
            donors[irec, ndonors[irec]] = inode
            ndonors[irec] += 1


@numba.njit
def _add2stack(inode, ndonors, donors, stack, istack):
    for k in range(ndonors[inode]):
        idonor = donors[inode, k]
        stack[istack] = idonor
        istack += 1
        istack = _add2stack(idonor, ndonors, donors, stack, istack)

    return istack


@numba.njit
def compute_stack(stack, ndonors, donors, receivers, nnodes):
    istack = 0

    for inode in range(nnodes):
        if receivers[inode] == inode:
            stack[istack] = inode
            istack += 1
            istack = _add2stack(inode, ndonors, donors,
                                stack, istack)


@numba.njit
def propagate_area(area, stack, receiver):
    for inode in stack[-1::-1]:
        if receiver[inode] != inode:
            area[receiver[inode]] += area[inode]

flow_routing_coarse_grid.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import xarray as xr\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import fastscapelib_fortran as fs\n",
    "from ipyfastscape import TopoViz3d"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Open and visualize model outputs (see \"fan.ipynb\" notebook)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = xr.open_dataset(\"fan_out.nc\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "app = TopoViz3d(ds, time_dim=\"out\")\n",
    "app.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A routine to re-compute drainage area from model outputs.\n",
    "\n",
    "It reuses the [fastscapelib-fortran](https://github.com/fastscape-lem/fastscapelib-fortran) library which implements multiple flow directions (MFD) routing with closed depression resolving. Caveats: not possible at the moment to reuse it with custom flow partioning routines (e.g., computed from a trained neural network)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_flowacc(ds, time_pos=-1):\n",
    "    \"\"\"Re-computes drainage area from fastscape outputs.\"\"\"\n",
    "    topo = ds.topography__elevation.isel(out=time_pos)\n",
    "    shape_xy = np.flip(topo.shape)\n",
    "    # assuming origin (0, 0)\n",
    "    length_xy = (topo.x.max(), topo.y.max())\n",
    "    \n",
    "    fs.fastscape_init()\n",
    "    fs.fastscape_set_nx_ny(*shape_xy)\n",
    "    fs.fastscape_setup()\n",
    "    fs.fastscape_set_xl_yl(*length_xy)\n",
    "    fs.fastscape_set_erosional_parameters(\n",
    "        1e-4, 1e-4, 0.5, 1.0, 1e-4, 1e-4, 0., 0., float(ds.flow__slope_exp)\n",
    "    )\n",
    "    # assuming reflective top/bottom + fixed left borders\n",
    "    fs.fastscape_set_bc(1)\n",
    "    fs.fastscape_init_h(topo.values.ravel())\n",
    "    \n",
    "    fs.flowrouting()\n",
    "    fs.flowaccumulation()\n",
    "    \n",
    "    a = fs.fastscapecontext.a.copy()\n",
    "    a2d = a.reshape(topo.shape)\n",
    "    \n",
    "    fs.fastscape_destroy()\n",
    "    \n",
    "    return xr.DataArray(\n",
    "        a2d, dims=(\"y\", \"x\"), coords={\"x\": topo.x, \"y\": topo.y}, name=\"drainage__area\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compare drainage area:\n",
    "\n",
    "- Top: computed from elevation values resampled on the coarse grid\n",
    "- Bottom: computed from elevation values on the native resolution grid then resampled on the coarse grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# coarsening factor from model native resolution grid to the resampled (coarsen) grid\n",
    "coarsen = 5\n",
    "\n",
    "area_native = compute_flowacc(ds)\n",
    "area_resampled = area_native.coarsen(\n",
    "    {\"x\": coarsen, \"y\": coarsen}, boundary=\"trim\"\n",
    ").sum()\n",
    "\n",
    "\n",
    "ds_resampled = ds.coarsen(\n",
    "    {\"x\": coarsen, \"y\": coarsen}, boundary=\"trim\"\n",
    ").mean()\n",
    "area_resampled_dem = compute_flowacc(ds_resampled)\n",
    "\n",
    "area = xr.concat([area_resampled_dem, area_resampled], \"how\")\n",
    "area[\"how\"] = [\"resampled_dem\", \"resampled\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.log(area).plot.pcolormesh(row=\"how\", x=\"x\", y=\"y\", figsize=(16, 13), vmin=13, vmax=21);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Show drainage area computed directly on the native resolution model grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.log(area_native).plot.pcolormesh(x=\"x\", y=\"y\", figsize=(19.3, 8), vmin=13, vmax=21);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.log(area_native).plot.hist(bins=20);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.log(area_resampled).plot.hist(bins=20);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.log(area_resampled_dem).plot.hist(bins=20);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:fastscape-demo]",
   "language": "python",
   "name": "conda-env-fastscape-demo-py"
  },
  "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}