ExampleVariableSPL.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Simple example run for the use of the VariableSPL process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import xsimlab as xs\n",
    "import xarray as xr\n",
    "%load_ext xsimlab.ipython\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We import the basic model from FastScape and modify it by replacing the SPL process by the VariableSPL process and dropping the diffusion process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fastscape.models import basic_model\n",
    "from VariableSPL import VariableSPL\n",
    "spl_model = basic_model.update_processes({'spl': VariableSPL}).drop_processes(('diffusion'))\n",
    "\n",
    "spl_model.visualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We create a local cluster to run a large number of model runs in bacth mode"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Basic geometry, time (clock) and uplift function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xl = 10e3\n",
    "yl = 10e3\n",
    "nx = 101\n",
    "ny = 101\n",
    "tim = xr.DataArray(np.linspace(0,5e7,101), dims='tstep')\n",
    "u = xr.DataArray(np.where(tim<2e7, 1e-3, 0), dims='tstep')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We create a simple xarray input dataset for the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %create_setup spl_model -d\n",
    "import xsimlab as xs\n",
    "\n",
    "ds_in = xs.create_setup(\n",
    "    model=spl_model,\n",
    "    clocks={'tstep':tim,\n",
    "           'time':[2e7,tim[-1]]},\n",
    "    master_clock='tstep',\n",
    "    input_vars={\n",
    "        'grid__shape': [ny,nx],\n",
    "        'grid__length': [yl,xl],\n",
    "        'boundary__status': ['looped','looped','fixed_value','fixed_value'],\n",
    "        'uplift__rate': u,\n",
    "        'init_topography__seed': 123456789,\n",
    "        'spl__k_coef': ('n', [1e-5,1e-7,1e-9,1e-11,1e-13]),\n",
    "        'spl__area_exp': ('n', [0.45,0.9,1.35,1.8,2.25]),\n",
    "        'spl__slope_exp': ('n', [1,2,3,4,5]),\n",
    "        'spl__recession_exp': ('b', [1,2]),\n",
    "        'spl__slope_exp_critical_stream_flow': ('n', [3/8,6/8,9/8,12/8,15/8]),#3/2,#1,\n",
    "        'spl__daily_stream_power_exp': 0.75,\n",
    "        'spl__threshold_erate': ('threshold', [1e-5,3e-5,1e-4,3e-4,1e-3,3e-3,1e-2]),\n",
    "        'spl__PET': 5,\n",
    "        'spl__soil_moisture_capacity': 1,\n",
    "        'spl__storm_depth': 10,\n",
    "        'spl__daily_rainfall_frequency': 0.6,\n",
    "        'spl__storm_size': 1e12,\n",
    "    },\n",
    "    output_vars={'topography__elevation': 'tstep',\n",
    "                'flow__receivers': 'time',\n",
    "                'flow__nb_donors': 'time',\n",
    "                'flow__donors': 'time',\n",
    "                'drainage__area': 'time',\n",
    "                'spl__streamflow_variability': 'time',\n",
    "                'spl__response_time': 'time',\n",
    "                'spl__streamflow_ratio': 'time'}\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We run the batch of models and save the resuls in a zarr archive on disk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import zarr\n",
    "zgroup = zarr.group(\"Multiple_thresholds.zarr\", overwrite=True)\n",
    "with spl_model:\n",
    "    ds_out = (ds_in.stack(batch=['threshold','n','b'])\n",
    "              .xsimlab.run(store=zgroup, batch_dim='batch', parallel=True, scheduler='processes')\n",
    "              .unstack('batch')\n",
    "              .assign_coords({'threshold':ds_in.spl__threshold_erate/ds_in.uplift__rate.max()})\n",
    "              .assign_coords({'n':ds_in.spl__slope_exp})\n",
    "              .assign_coords({'b':ds_in.spl__recession_exp}))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We show the result of one model run as the maximum and final topography"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_out.topography__elevation.sel(tstep=(2e7,5e7)).sel(b=1,n=1).isel(threshold=-1).plot(col='tstep')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We reproduce one of the figures of the manuscript bu computing the ratio of mean topography at the end of the post-orogenic phase to the same thing at the end of the orogenic phase, for various values of the threshold (normalized by the uplift rate) and the parameters n and b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(6,4))\n",
    "((ds_out.topography__elevation.isel(tstep=-1).mean(('x','y'))\n",
    "  /ds_out.topography__elevation.sel(tstep=2e7).mean(('x','y')))\n",
    " .isel(b=0).plot.line(x='threshold',xscale='log',marker='o',ax=ax))\n",
    "plt.gca().set_prop_cycle(None)\n",
    "((ds_out.topography__elevation.isel(tstep=-1).mean(('x','y'))\n",
    "  /ds_out.topography__elevation.sel(tstep=2e7).mean(('x','y')))\n",
    " .isel(b=1).plot.line(x='threshold',xscale='log',marker='+',ax=ax))\n",
    "plt.gca().set_prop_cycle(None)\n",
    "((ds_out.topography__elevation.isel(tstep=-1).mean(('x','y'))\n",
    "  /ds_out.topography__elevation.sel(tstep=2e7).mean(('x','y')))\n",
    " .isel(b=1).plot.line(x='threshold',xscale='log',ax=ax))\n",
    "ax.set_title('o b=1; + b=2')\n",
    "ax.set_xlabel(r'$\\epsilon_c/U$');\n",
    "ax.set_ylabel(r'$R_\\infty$');\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We show the distribution of response time, tau, and srteam flow variability for all models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig,ax = plt.subplots(ncols=2, figsize=(12,4))\n",
    "(ds_out.spl__response_time\n",
    "             .isel(threshold=0)\n",
    "             .isel(n=0)\n",
    "             .isel(b=0)\n",
    "             .isel(time=-1)\n",
    "             .plot(ax=ax[0])\n",
    "            )\n",
    "(ds_out.spl__streamflow_variability\n",
    "             .isel(threshold=0)\n",
    "             .isel(n=0)\n",
    "             .isel(b=0)\n",
    "             .isel(time=-1)\n",
    "             .plot(ax=ax[1])\n",
    "            )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:fastscape]",
   "language": "python",
   "name": "conda-env-fastscape-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.10.8"
  },
  "toc-autonumbering": true,
  "toc-showcode": false,
  "toc-showmarkdowntxt": false,
  "toc-showtags": false
 },
 "nbformat": 4,
 "nbformat_minor": 4
}