{
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "# Inner Base Levels [Py]\n",
                "\n",
                "An example very similar to {doc}`mountain_py` but setting the base level nodes inside the raster grid instead of on the boundaries. The modeled domain is thus infinite (fully periodic), which is usefull for global (planetary) scale simulations."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "from random import random, uniform\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 201x201 nodes with a total length of 50 km in both y (rows) and x (columns).\n",
                "\n",
                "Set looped (reflective) boundary conditions at all border nodes. Also set fixed-value \"boundary\" for a given number of base level nodes randomly selected inside the domain."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "n_base_levels = 20\n",
                "base_level_row = np.random.uniform(1, 200, n_base_levels).astype(\"int\")\n",
                "base_level_col = np.random.uniform(1, 200, n_base_levels).astype(\"int\")\n",
                "\n",
                "base_levels = {\n",
                "    (i, j): fs.NodeStatus.FIXED_VALUE\n",
                "    for i, j in zip(base_level_row, base_level_col)\n",
                "}\n",
                "\n",
                "bs = fs.NodeStatus.LOOPED\n",
                "grid = fs.RasterGrid.from_length([201, 201], [5e4, 5e4], bs, 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 (random inner 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.where(\n",
                "    grid.nodes_status() == fs.NodeStatus.FIXED_VALUE.value, 0, 1e-3\n",
                ")"
            ]
        },
        { "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=(8, 8))\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=(8, 8))\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=(7, 7))\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.9.10"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 4
}