{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1fc61c8f",
   "metadata": {},
   "source": [
    "# GRLP river networks: a 5-segment example\n",
    "\n",
    "The single-profile tutorial ([`example_1d`](example_1d.ipynb)) evolved one river long profile. GRLP also solves whole **river networks** — in fact *every* GRLP solution is a network solution, and a single profile is just the trivial one-edge case.\n",
    "\n",
    "In a network, **segments** (reaches) meet at **confluences**, where two rules hold:\n",
    "\n",
    "1. **Sediment and water sum**: a segment's supply is the sum of its upstream tributaries'.\n",
    "2. **Elevation is continuous** across the junction.\n",
    "\n",
    "The solver assembles **one global system** over the whole network by walking the topology, then solves it just as for a single segment. Here we build a small 5-segment network — two confluences — by hand, evolve it to steady state, and plot it. Compare each step with the single-profile tutorial.\n",
    "\n",
    "Reference: [McNab et al. (2025, ESurf)](https://doi.org/10.5194/esurf-13-1059-2025) develop and apply the network-scale theory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9722533a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "import grlp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2553298",
   "metadata": {},
   "source": [
    "## 1. Wire up the topology\n",
    "\n",
    "A network is defined by how its segments connect. Each segment lists its **upstream** and **downstream** neighbours by ID. **Channel heads** have no upstream segments; the **outlet** has no downstream segment.\n",
    "\n",
    "Our little network (heads `0`, `1`, `3`):\n",
    "\n",
    "    head 0 ─┐\n",
    "            ├─ segment 2 ─┐\n",
    "    head 1 ─┘             ├─ segment 4 ──▶ outlet (base level)\n",
    "    head 3 ──────────────┘"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d91bad2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Each segment's upstream / downstream neighbours, by ID.\n",
    "# heads 0, 1 join at 2;  segment 2 and head 3 join at 4;  4 is the outlet.\n",
    "upstream_segment_IDs   = [[], [], [0, 1], [], [2, 3]]\n",
    "downstream_segment_IDs = [[2], [2], [4], [4], []]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd8c7103",
   "metadata": {},
   "source": [
    "## 2. Give each segment geometry and discharge\n",
    "\n",
    "Every segment carries its own arrays: node positions `x` (downstream distance), initial bed `z`, discharge `Q`, and valley width `B`. **Discharge grows downstream** as tributaries join — here the two headwater segments each carry 5 m³/s, so their trunk (segment 2) carries 10, and adding head 3 (10 m³/s) makes the outlet trunk (segment 4) carry 20.\n",
    "\n",
    "One rule to remember: a **confluence segment needs at least 3 nodes** (the conservative junction cell reaches two nodes in from the boundary). Valley width is uniform here to keep the example simple."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "089b50cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "numel = [5, 7, 4, 8, 6]        # nodes per segment; confluence segments need >= 3\n",
    "\n",
    "# Downstream distance [m] of each node, per segment (a small dendritic layout)\n",
    "x = [ 1000. * np.array([2, 4, 6.5, 9, 10]),\n",
    "      1000. * np.array([0, 1, 2, 3, 6, 8, 10.5]),\n",
    "      1000. * np.array([12, 15, 18, 20]),\n",
    "      1000. * np.array([2, 6, 8, 12, 14, 16, 18, 20]),\n",
    "      1000. * np.array([23, 24, 27, 29, 29.5, 30]) ]\n",
    "\n",
    "# Discharge [m3/s], constant within each segment and summing at confluences:\n",
    "#   heads 0,1 (5 + 5) -> trunk 2 (10);  trunk 2 + head 3 (10 + 10) -> outlet 4 (20)\n",
    "Q_in = [5., 5., 10., 10., 20.]\n",
    "\n",
    "z = [np.zeros(n) for n in numel]                  # start flat\n",
    "Q = [q * np.ones(n) for q, n in zip(Q_in, numel)]\n",
    "B = [100. * np.ones(n) for n in numel]            # uniform valley width [m]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b768856a",
   "metadata": {},
   "source": [
    "## 3. Boundary conditions\n",
    "\n",
    "Same idea as the single profile, applied at the network's edges:\n",
    "\n",
    "* **Upstream** — each channel head gets an equilibrium slope `S0` (here 1.5%), which sets that head's sediment supply (exactly as `set_S0` does for one profile).\n",
    "* **Downstream** — the outlet has a **base level**, the point `(x_bl, z_bl)` (recall `set_bl` from the 1-D tutorial). The network places it on its single river mouth."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90fe3483",
   "metadata": {},
   "outputs": [],
   "source": [
    "S0 = [0.015, 0.015, 0.015]   # equilibrium slope at each channel head (0, 1, 3)\n",
    "x_bl = 1000. * 32            # base-level position [m]\n",
    "z_bl = 0.                    # base-level elevation [m]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb8a132a",
   "metadata": {},
   "source": [
    "## 4. Build and evolve\n",
    "\n",
    "Instantiate a `Network`, pass everything through `initialize()`, then evolve to steady state — the same kind of `evolve_*` call as a single profile, now over the whole graph. (`get_z_lengths()` tells the solver how many nodes each segment has before the first solve.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7f95189",
   "metadata": {},
   "outputs": [],
   "source": [
    "net = grlp.Network()\n",
    "net.initialize(\n",
    "    config_file=None,\n",
    "    x_bl=x_bl, z_bl=z_bl,\n",
    "    S0=S0, Q_s_0=None,\n",
    "    upstream_segment_IDs=upstream_segment_IDs,\n",
    "    downstream_segment_IDs=downstream_segment_IDs,\n",
    "    x=x, z=z, Q=Q, B=B,\n",
    "    overwrite=False,\n",
    ")\n",
    "net.set_niter(3)\n",
    "net.get_z_lengths()\n",
    "net.evolve_threshold_width_river_network(nt=100, dt=1e11)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57387931",
   "metadata": {},
   "source": [
    "## 5. Plot the network\n",
    "\n",
    "Each segment is drawn in its own colour; the thin black lines are the **connectors** — a segment's downstream end joining its neighbour across each confluence, and the outlet joining base level. The profiles are concave-up, descend to base level, and steepen in the smaller-discharge headwaters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f93fc48",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 6))\n",
    "for lp in net.segments:\n",
    "    # connector: downstream end -> neighbour across the junction\n",
    "    # (or, at the outlet, down to base level)\n",
    "    if len(lp.downstream_segment_IDs) > 0:\n",
    "        ds = net.segments[lp.downstream_segment_IDs[0]]\n",
    "        plt.plot([lp.x[-1] / 1000., ds.x[0] / 1000.],\n",
    "                 [lp.z[-1], ds.z[0]], 'k-', lw=1, alpha=0.6)\n",
    "    else:\n",
    "        plt.plot([lp.x[-1] / 1000., lp.x_ghost_downstream / 1000.],\n",
    "                 [lp.z[-1], lp.z_bl], 'k-', lw=1, alpha=0.6)\n",
    "    plt.plot(lp.x / 1000., lp.z, '-', lw=3, label='segment %d' % lp.ID)\n",
    "plt.xlabel('Downstream distance [km]', fontsize=14)\n",
    "plt.ylabel('Elevation [m]', fontsize=14)\n",
    "plt.title('5-segment network at steady state')\n",
    "plt.legend()\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24a82ac7",
   "metadata": {},
   "source": [
    "## Next steps\n",
    "\n",
    "You built a network by hand. To generate networks **at scale** — random topologies with no DEM required, plus network-structure analysis (Hack's law, Strahler orders, Horton ratios) — see the random-network tutorial ([`example_random_network`](example_random_network.ipynb)).\n",
    "\n",
    "Full documentation and API reference: [grlp.readthedocs.io](https://grlp.readthedocs.io)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}