{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Mountain (Flow Kernel Eroder) [Py]\n", "\n", "This example performs the same simulation than in {doc}`mountain_py`, although here the Stream-Power bedrock channel erosion process has been re-implemented using Fastscapelib's {ref}`flow kernels `." ] }, { "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\n", "import numba as nb" ] }, { "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": [ "## Re-Implement the Stream-Power Law as a Flow Kernel Eroder\n", "\n", "The class below re-implements the basic Stream-Power Law (SPL) for simulating bedrock channel erosion, using Fastscapelib's {ref}`flow kernels `. It inherits from the {py:class}`~fastscapelib.FlowKernelEroder` base class and implements the required methods (parameter and input specs, kernel function and kernel application direction).\n", "\n", "This eroder class provides the exact same interface than {py:class}`~fastscapelib.SPLEroder` and it should also provide the same results. It is used below as a drop-in replacement." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "class SPLFlowKernelEroder(fs.FlowKernelEroder):\n", " \"\"\"Stream-Power Law (SPL) berock channel erosion.\n", " \n", " SPL is implemented here as a Fastscapelib flow kernel eroder.\n", " The core erosion logic is implemented in pure-Python and\n", " accelerated with Numba.\n", " \n", " \"\"\"\n", " def __init__(\n", " self,\n", " flow_graph: fs.FlowGraph,\n", " k_coef: float,\n", " area_exp: float,\n", " slope_exp: float,\n", " tolerance: float = 1e-3,\n", " ):\n", " super().__init__(flow_graph)\n", "\n", " self.kernel_data.bind(\n", " k_coef=k_coef, area_exp=area_exp, slope_exp=slope_exp, tolerance=tolerance\n", " )\n", "\n", " @staticmethod\n", " def param_spec():\n", " \"\"\"Returns a dictionary with parameter names\n", " and their (numba) value type.\n", " \"\"\"\n", " return {\n", " \"k_coef\": nb.float64,\n", " \"area_exp\": nb.float64,\n", " \"slope_exp\": nb.float64,\n", " \"tolerance\": nb.float64,\n", " }\n", "\n", " @staticmethod\n", " def input_spec():\n", " \"\"\"Returns a dictionary with input variable names\n", " and their (numba) value type.\n", " \"\"\"\n", " return {\n", " \"elevation\": nb.float64[::1],\n", " \"drainage_area\": nb.float64[::1],\n", " \"dt\": nb.float64,\n", " }\n", "\n", " @staticmethod\n", " def kernel_apply_dir() -> fs.FlowGraphTraversalDir:\n", " \"\"\"Returns the kernel application direction and order.\"\"\"\n", " return fs.FlowGraphTraversalDir.BREADTH_UPSTREAM\n", "\n", " @staticmethod\n", " def kernel_func(node):\n", " \"\"\"The eroder flow kernel function.\"\"\"\n", " r_count = node.receivers.count\n", "\n", " # skip base level node (erosion = 0)\n", " if r_count == 1 and node.receivers.distance[0] == 0.0:\n", " return\n", "\n", " # determine whether the current node is in a closed depression\n", " # of the topography (lake) and skip it if that's the case\n", " # (erosion = 0)\n", " elevation_flooded = np.finfo(np.double).max\n", "\n", " for r in range(r_count):\n", " irec_elevation_next = (\n", " node.receivers.elevation[r] - node.receivers.erosion[r]\n", " )\n", "\n", " if irec_elevation_next < elevation_flooded:\n", " elevation_flooded = irec_elevation_next\n", "\n", " if node.elevation <= elevation_flooded:\n", " return\n", "\n", " # compute new elevation at the current node by applying SPL\n", " eq_num = node.elevation\n", " eq_den = 1.0\n", "\n", " for r in range(r_count):\n", " irec_elevation = node.receivers.elevation[r]\n", " irec_elevation_next = irec_elevation - node.receivers.erosion[r]\n", "\n", " if irec_elevation > node.elevation:\n", " continue\n", "\n", " irec_weight = node.receivers.weight[r]\n", " irec_distance = node.receivers.distance[r]\n", "\n", " factor = (\n", " node.k_coef\n", " * node.dt\n", " * np.power(node.drainage_area * irec_weight, node.area_exp)\n", " )\n", " factor /= irec_distance\n", " eq_num += factor * irec_elevation_next\n", " eq_den += factor\n", "\n", " elevation_updated = eq_num / eq_den\n", "\n", " # prevent the creation of new closed depressions in the\n", " # topography\n", " if elevation_updated < elevation_flooded:\n", " elevation_updated = elevation_flooded + np.finfo(np.double).tiny\n", "\n", " # compute erosion\n", " node.erosion = node.elevation - elevation_updated\n", "\n", " def erode(\n", " self, elevation: np.ndarray, drainage_area: np.ndarray, dt: float\n", " ) -> np.ndarray:\n", " \"\"\"Compute SPL erosion for one time step.\"\"\"\n", " return super().erode(elevation=elevation, drainage_area=drainage_area, dt=dt)\n" ] }, { "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 = SPLFlowKernelEroder(\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.13.5" }, "pixi-kernel": { "environment": "doc" } }, "nbformat": 4, "nbformat_minor": 4 }