{
    "cells": [
        {
            "cell_type": "markdown",
            "id": "89a6bd56-d67f-4161-a23e-b3a73f542f58",
            "metadata": {},
            "source": [
                "# Planetary [Py]\n",
                "\n",
                "A simple example simulating the evolution of a planet under the action of vertical uplift and bedrock channel erosion.\n",
                "\n",
                "This example is pretty similar to {doc}`mountain_py`, although solving equation {eq}`eq_mountain` without the hillslope linear diffusion term {math}`K_D \\nabla^2 h` on a spherical domain, leveraging Fastscapelib's support for the [HEALPix](https://healpix.sourceforge.io/) grid.\n",
                "\n",
                ":::{note}\n",
                "\n",
                "Fastscapelib support for HEALPix is only available in the Linux and MacOS packages published on conda-forge.\n",
                "\n",
                ":::\n"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "93184ad7-feb0-4570-a270-10c50e710511",
            "metadata": { "editable": true, "slideshow": { "slide_type": "" }, "tags": [] },
            "outputs": [],
            "source": [
                "import healpy as hp\n",
                "import numpy as np\n",
                "import matplotlib\n",
                "import matplotlib.pyplot as plt\n",
                "\n",
                "import fastscapelib as fs"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "db031bf3-34e0-4c50-b933-837817d8e123",
            "metadata": {},
            "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",
            "id": "0b9bda51-df1d-4216-b6e4-059cd92882cd",
            "metadata": {},
            "source": [
                "## Setup the Grid, Flow Graph and Eroders\n",
                "\n",
                "Create a {py:class}`~fastscapelib.HealpixGrid` with a given value of `nside`, which corresponds to a number of refinement levels from the initial tessellation of the sphere. From this value we can compute the grid resolution and total number of nodes."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "ddef7b12-b6f3-4eaa-86b0-9b25c06f5dcc",
            "metadata": {},
            "outputs": [],
            "source": [
                "nside = 100\n",
                "\n",
                "print(\"total nb. of nodes: \", hp.nside2npix(nside))\n",
                "print(\"approx. resolution (km): \", hp.nside2resol(nside) * 6.371e3)"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "4db64f42-da0f-4f90-9588-fcfacfe5815f",
            "metadata": {},
            "source":
                ["As \"boundary\" conditions, we set a strip of a given width along the equator that will represent the planet's ocean. This is done via the [healpy](https://healpy.readthedocs.io) library that allows handling HEALPix grids from within Python."]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "85499a73-9f37-4c49-8095-7614890eff65",
            "metadata": {},
            "outputs": [],
            "source": [
                "# set thin strips on each side from the equator (coastlines)\n",
                "# with \"FIXED VALUE\" node status\n",
                "nodes_status = np.zeros(hp.nside2npix(nside), dtype=np.uint8)\n",
                "ocean_idx = hp.query_strip(nside, np.deg2rad(85), np.deg2rad(91))\n",
                "nodes_status[ocean_idx] = fs.NodeStatus.FIXED_VALUE.value\n",
                "\n",
                "# set \"ghost\" nodes between the two thin strips (ocean)\n",
                "ocean_idx = hp.query_strip(nside, np.deg2rad(85.5), np.deg2rad(90.5))\n",
                "nodes_status[ocean_idx] = fs.NodeStatus.GHOST.value\n",
                "\n",
                "mask = nodes_status == fs.NodeStatus.GHOST.value"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "d47b1288-ef0b-49a2-bf9e-fa776f9862b2",
            "metadata": {},
            "source": ["Create the `grid` object with a radius equal to the Earth's radius."]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "0e82ce54-44a4-4146-9133-5f37c8ec09c3",
            "metadata": {},
            "outputs": [],
            "source": ["grid = fs.HealpixGrid(nside, nodes_status, 6.371e6)"]
        },
        {
            "cell_type": "markdown",
            "id": "bd6561c6-1a25-4630-a02a-f859a6df8c15",
            "metadata": {},
            "source":
                ["Create the flow graph and eroder objects. Set a mask for the flow graph from the \"ocean\" grid nodes."]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "ab3ad342-7df9-46c3-847f-ba54f0ade9dd",
            "metadata": {},
            "outputs": [],
            "source": [
                "flow_graph = fs.FlowGraph(\n",
                "    grid,\n",
                "    [fs.SingleFlowRouter(), fs.MSTSinkResolver()],\n",
                ")\n",
                "flow_graph.mask = mask\n",
                "\n",
                "spl_eroder = fs.SPLEroder(\n",
                "    flow_graph,\n",
                "    k_coef=1.2e-11,\n",
                "    area_exp=0.45,\n",
                "    slope_exp=1,\n",
                "    tolerance=1e-5,\n",
                ")\n",
                "\n",
                "urate = 9e-4\n",
                "\n",
                "dt = 1e4"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "5ad4566c-2dab-411e-8d25-26673ae90d16",
            "metadata": {},
            "source":
                ["## Create and Run a Simulation\n", "\n", "Create a simple simulation runner."]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "2b8a5fed-2c74-48af-adde-ad35672e6286",
            "metadata": {},
            "outputs": [],
            "source": [
                "def run_simulation(nsteps):\n",
                "    rng = np.random.Generator(np.random.PCG64(123456789))\n",
                "    init_elevation = rng.uniform(0, 1, size=grid.shape)\n",
                "    elevation = init_elevation.copy()\n",
                "    drainage_area = np.zeros_like(elevation)\n",
                "    uplift_rate = np.full_like(elevation, urate)\n",
                "    uplift_rate[grid.nodes_status() > fs.NodeStatus.CORE.value] = 0.\n",
                "\n",
                "    for step in range(nsteps):\n",
                "        uplifted_elevation = elevation + uplift_rate * dt\n",
                "        \n",
                "        # flow routing\n",
                "        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",
                "        \n",
                "        # update topography\n",
                "        elevation = uplifted_elevation - spl_erosion\n",
                "\n",
                "    return elevation, drainage_area"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "9a0de7a1-1795-4267-9f62-f27ebfa8bf48",
            "metadata": {},
            "source":
                ["Run the model for a few dozens of time steps (total simulation time: 400 000 years)."]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "a06e829a-74f4-4e33-b163-6ae119af8997",
            "metadata": {},
            "outputs": [],
            "source": ["elevation, drainage_area = run_simulation(40)"]
        },
        {
            "cell_type": "markdown",
            "id": "6d4ac708-300c-4826-9640-6c8d1549ee73",
            "metadata": {},
            "source": ["## Plot Outputs"]
        },
        {
            "cell_type": "markdown",
            "id": "2fae4cfa-1f6c-4dbd-802a-c690250e55a4",
            "metadata": {},
            "source": ["- Topographic elevation"]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "276dc362-e200-443e-be9e-5de53cb539ac",
            "metadata": {},
            "outputs": [],
            "source": ["hp.orthview(elevation, nest=False, cmap=plt.cm.copper)"]
        },
        {
            "cell_type": "markdown",
            "id": "cbb1fd5c-1205-4155-a51e-66f8c0ffa1a9",
            "metadata": {},
            "source": ["- Drainage area (log)"]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "79da3aea-e04c-48af-96f1-2e848a32edb0",
            "metadata": {},
            "outputs": [],
            "source": ["hp.orthview(np.log(drainage_area), nest=False);"]
        }
    ],
    "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": 5
}