{
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "# Mountain [Py]\n",
                "\n",
                "A simple example simulating the formation of an orogenic belt (on a 2D raster grid) from a nearly-flat surface under the action of block-uplift vs. bedrock channel and hillslope erosion. See also {doc}`mountain_cpp`.\n",
                "\n",
                "The local rate of elevation change, {math}`\\partial h/\\partial t`, is determined by the balance between uplift (uniform in space and time) {math}`U` and erosion (stream-power law for river erosion and linear diffusion for hillslope erosion).\n",
                "\n",
                "```{math}\n",
                ":label: eq_mountain\n",
                "\\frac{\\partial h}{\\partial t} = U - K_f A^m (\\nabla h)^n + K_D \\nabla^2 h\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`, {math}`n` and {math}`K_D` are parameters (uniform in space and time).\n",
                "\n",
                "The initial topography is nearly flat with small random perturbations and elevation remains fixed at the boundaries of the spatial domain."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "from random import random\n",
                "\n",
                "import fastscapelib as fs\n",
                "import numpy as np\n",
                "import matplotlib\n",
                "import matplotlib.pyplot as plt"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": { "tags": ["remove-cell"] },
            "outputs": [],
            "source": [
                "# Theme that looks reasonably fine on both dark/light modes\n",
                "matplotlib.style.use('Solarize_Light2')\n",
                "matplotlib.rcParams['axes.grid'] = False"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Setup the Grid, Flow Graph and Eroders\n",
                "\n",
                "Create a {py:class}`~fastscapelib.RasterGrid` of 201x301 nodes with a total length of 50 km in y (rows) and 75 km in x (columns).\n",
                "\n",
                "Set fixed value boundary conditions at all border nodes."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source":
                ["grid = fs.RasterGrid.from_length([201, 301], [5e4, 7.5e4], fs.NodeStatus.FIXED_VALUE)"]
        },
        {
            "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 (all border nodes 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 eroder classes (bedrock channel + hillslope) 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",
                ")\n",
                "\n",
                "diffusion_eroder = fs.DiffusionADIEroder(grid, 0.01)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": ["## Setup Initial Conditions and External Forcing"]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source":
                ["Create a flat (+ random perturbations) surface topography as initial conditions. Also initialize the array for drainage area."]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "rng = np.random.Generator(np.random.PCG64(1234))\n",
                "\n",
                "init_elevation = rng.uniform(0, 1, size=grid.shape)\n",
                "\n",
                "elevation = init_elevation\n",
                "drainage_area = np.empty_like(elevation)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source":
                ["Set upflit rate as uniform (fixed value) within the domain and to zero at all grid boundaries."]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "uplift_rate = np.full_like(elevation, 1e-3)\n",
                "uplift_rate[[0, -1], :] = 0.\n",
                "uplift_rate[:, [0, -1]] = 0."
            ]
        },
        { "cell_type": "markdown", "metadata": {}, "source": ["## Run the Model"] },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source":
                ["Run the model for a few dozens of time steps (total simulation time: 1M years)."]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "dt = 2e4\n",
                "nsteps = 50\n",
                "\n",
                "for step in range(nsteps):\n",
                "    # uplift (no uplift at fixed elevation boundaries)\n",
                "    uplifted_elevation = elevation + dt * uplift_rate\n",
                "    \n",
                "    # flow routing\n",
                "    filled_elevation = flow_graph.update_routes(uplifted_elevation)\n",
                "    \n",
                "    # flow accumulation (drainage area)\n",
                "    flow_graph.accumulate(drainage_area, 1.0)\n",
                "    \n",
                "    # apply channel erosion then hillslope diffusion\n",
                "    spl_erosion = spl_eroder.erode(uplifted_elevation, drainage_area, dt)\n",
                "    diff_erosion = diffusion_eroder.erode(uplifted_elevation - spl_erosion, dt)\n",
                "    \n",
                "    # update topography\n",
                "    elevation = uplifted_elevation - spl_erosion - diff_erosion\n"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": ["## Plot Outputs and Other Diagnostics\n"]
        },
        { "cell_type": "markdown", "metadata": {}, "source": ["- Topographic elevation"] },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "fig, ax = plt.subplots(figsize=(12, 6))\n",
                "plt.imshow(elevation)\n",
                "plt.colorbar();"
            ]
        },
        { "cell_type": "markdown", "metadata": {}, "source": ["- Drainage area (log)"] },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "fig, ax = plt.subplots(figsize=(12, 6))\n",
                "plt.imshow(np.log(drainage_area), cmap=plt.cm.Blues)\n",
                "plt.colorbar();"
            ]
        },
        { "cell_type": "markdown", "metadata": {}, "source": ["- Drainage basins"] },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "colors = [(1,1,1)] + [(random(),random(),random()) for i in range(255)]\n",
                "rnd_cm = matplotlib.colors.LinearSegmentedColormap.from_list('new_map', colors, N=256)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "fig, ax = plt.subplots(figsize=(12, 6))\n",
                "plt.imshow(flow_graph.basins(), cmap=rnd_cm);"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": []
        }
    ],
    "metadata": {
        "kernelspec":
            { "display_name": "Python 3 (ipykernel)", "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.12.3"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 4
}