{
    "cells": [
        { "cell_type": "markdown", "metadata": {}, "source": ["# Catchment (TriMesh) [Py]"] },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "A simple example simulating the evolution of a single synthetic catchment on a triangular (irregular) mesh.\n",
                "\n",
                "The local rate of elevation change, {math}`\\partial h/\\partial t`, is governed by bedrock river channel erosion (stream-power law).\n",
                "\n",
                "```{math}\n",
                ":label: eq_catchment\n",
                "\\frac{\\partial h}{\\partial t} = - K_f A^m (\\nabla h)^n\n",
                "```\n",
                "\n",
                "where {math}`A` is the drainage area (i.e., the total upslope area from which the water flow converges) and {math}`K_f`, {math}`m` and {math}`n` are parameters (uniform in space and time).\n",
                "\n",
                "The initial topography is a nearly flat plateau of a given elevation. The base level (catchment outlet) is lowered during the simulation at a constant rate."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import numpy as np\n",
                "import matplotlib.pyplot as plt\n",
                "import matplotlib.tri as mpltri\n",
                "import pygalmesh\n",
                "\n",
                "import fastscapelib as fs"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": { "tags": ["remove-cell"] },
            "outputs": [],
            "source": [
                "# Theme that looks reasonably fine on both dark/light modes\n",
                "import matplotlib\n",
                "matplotlib.style.use('Solarize_Light2')\n",
                "matplotlib.rcParams['axes.grid'] = False"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Setup the Mesh, Flow Graph and Eroders\n",
                "\n",
                "Import the coordinates of the catchment boundaries from a file and create a new irregular mesh (with quasi-uniform resolution) using a constrained triangulation algorithm available in the [pygalmesh](https://github.com/meshpro/pygalmesh) library."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "basin_poly = np.load(\"basin.npy\")\n",
                "n_poly = basin_poly.shape[0]\n",
                "constraints = list([list(s) for s in zip(range(n_poly), range(1, n_poly))])\n",
                "constraints[-1][1] = 0\n",
                "\n",
                "approx_resolution = 300.0\n",
                "\n",
                "mesh = pygalmesh.generate_2d(\n",
                "    basin_poly,\n",
                "    constraints,\n",
                "    max_edge_size=approx_resolution,\n",
                "    num_lloyd_steps=2,\n",
                ")\n",
                "\n",
                "points = mesh.points\n",
                "triangles = mesh.cells[0].data"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source":
                ["Determine the base level node, which corresponds to the catchment outlet (a boundary node)."]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "outlet_idx = np.argmax(points[:, 1])\n",
                "base_levels = {outlet_idx: fs.NodeStatus.FIXED_VALUE}"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source":
                ["Create a new {py:class}`~fastscapelib.TriMesh` object and pass the existing mesh data."]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": ["grid = fs.TriMesh(points, triangles, base_levels)"]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Create a {py:class}`~fastscapelib.FlowGraph` object with single direction flow routing and the resolution of closed depressions on the topographic surface. See {ref}`guide-flow-routing-strategies` for more examples on possible flow routing strategies.\n",
                "\n",
                "By default, base level nodes are set from fixed value boundary conditions (the outlet node set above in this example)."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source":
                ["flow_graph = fs.FlowGraph(grid, [fs.SingleFlowRouter(), fs.MSTSinkResolver()])"]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": ["Setup the bedrock channel eroder with a given set of parameter values."]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "spl_eroder = fs.SPLEroder(\n",
                "    flow_graph,\n",
                "    k_coef=2e-4,\n",
                "    area_exp=0.4,\n",
                "    slope_exp=1,\n",
                "    tolerance=1e-5,\n",
                ")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Setup Initial Conditions and External Forcing\n",
                "\n",
                "Create a flat (+ random perturbations) elevated surface topography as initial conditions. Also initialize the array for drainage area."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "init_elevation = np.random.uniform(0, 1, size=grid.shape) + 500.0\n",
                "elevation = init_elevation\n",
                "drainage_area = np.empty_like(elevation)"
            ]
        },
        { "cell_type": "markdown", "metadata": {}, "source": ["The base level lowering rate:"] },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": ["lowering_rate = 1e-4"]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": ["## Run the Model\n", "\n", "Run the model for one hundred time steps."]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "dt = 2e4\n",
                "nsteps = 100\n",
                "\n",
                "for step in range(100):\n",
                "    # base level (outlet) lowering \n",
                "    elevation[outlet_idx] -= lowering_rate * dt\n",
                "    \n",
                "    # flow routing\n",
                "    flow_graph.update_routes(elevation)\n",
                "    \n",
                "    # flow accumulation (drainage area)\n",
                "    flow_graph.accumulate(drainage_area, 1.0)\n",
                "    \n",
                "    # channel erosion (SPL)\n",
                "    spl_erosion = spl_eroder.erode(elevation, drainage_area, dt)\n",
                "    \n",
                "    # update topography\n",
                "    elevation -= spl_erosion"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": ["## Plot Outputs and Other Diagnostics"]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source":
                ["x, y = points.transpose()\n", "plot_tri = mpltri.Triangulation(x, y, triangles)"]
        },
        { "cell_type": "markdown", "metadata": {}, "source": ["- Topographic elevation"] },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "fig, ax = plt.subplots(figsize=(13, 13))\n",
                "ax.set_aspect('equal')\n",
                "cf = ax.tripcolor(plot_tri, elevation)\n",
                "plt.colorbar(cf)\n",
                "ax.triplot(plot_tri, linewidth=0.2, c=\"k\");"
            ]
        },
        { "cell_type": "markdown", "metadata": {}, "source": ["- Drainage area (log)"] },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "fig, ax = plt.subplots(figsize=(13, 13))\n",
                "ax.set_aspect('equal')\n",
                "cf = ax.tricontourf(plot_tri, np.log(drainage_area), cmap=plt.cm.Blues)\n",
                "plt.colorbar(cf);\n",
                "ax.triplot(plot_tri, linewidth=0.2, c=\"k\");"
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "Python [conda env:fastscapelib_dev]",
            "language": "python",
            "name": "conda-env-fastscapelib_dev-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.9.10"
        },
        "vscode": {
            "interpreter":
                { "hash": "828c8bd590fe165438ea9436ea3075ad1e985d4b674461d5f3ca65e5df0d3dd5" }
        }
    },
    "nbformat": 4,
    "nbformat_minor": 4
}