{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NumPy — Lab (solutions)\n",
    "\n",
    "**AI for Hydraulics · Part 1 · Week 1**\n",
    "\n",
    "In this lab you will practise NumPy on small hydraulic problems: channel geometry, Manning's equation, rating tables, rainfall statistics and a flood hydrograph.\n",
    "\n",
    "**How to work**\n",
    "- Run the cells from top to bottom (`Shift + Enter`).\n",
    "- Replace every `...` with your code.\n",
    "- Each exercise ends with a ✅ **Check** giving the values you should obtain.\n",
    "- No data file is needed: all data is defined or generated in the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "print(np.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 1 — Creating arrays\n",
    "\n",
    "**a)** Create an array `depths` of flow depths from 0.1 m to 3.0 m every 0.1 m (both included) with `np.linspace`. Print its `shape`, its `dtype`, and its first and last values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "depths = np.linspace(0.1, 3.0, 30)\n",
    "print(depths.shape, depths.dtype, depths[0], depths[-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**b)** Create an array `n_values` with the **same shape** as `depths`, filled with the Manning coefficient 0.03."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_values = np.full(depths.shape, 0.03)\n",
    "print(n_values.shape, n_values[:3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**c)** Using a random generator with seed `7`, draw 10 water temperatures from a normal distribution with mean 18 °C and standard deviation 2 °C. Print them rounded to 1 decimal, and their mean.\n",
    "\n",
    "✅ **Check:** `depths` has shape `(30,)` and dtype `float64`. The temperatures start with `18.  18.6 17.5`, with mean ≈ 17.6 °C."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = np.random.default_rng(7)\n",
    "temperatures = rng.normal(18, 2, size=10)\n",
    "print(temperatures.round(1), temperatures.mean().round(2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 2 — Selecting data\n",
    "\n",
    "`Q` contains the mean daily discharge (m³/s) of a river over one month (day 1 is `Q[0]`). A flood occurs around day 9."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Q = np.array([12.4, 11.8, 11.1, 10.6, 10.2, 9.9, 18.7, 35.2, 41.6, 30.3,\n",
    "              22.1, 17.5, 14.6, 12.9, 11.7, 10.8, 10.1, 9.5, 9.0, 8.6,\n",
    "              8.2, 7.9, 7.5, 7.2, 6.8, 6.5, 6.1, 5.8, 5.4, 4.9, 4.6])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**a)** Print the discharge of the first week and of the last 3 days."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "first_week = Q[:7]\n",
    "last_3_days = Q[-3:]\n",
    "print(first_week, last_3_days)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**b)** Using a boolean mask, print the discharges above 20 m³/s and the **day numbers** (1 to 31) on which they occur."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "flood_mask = Q > 20\n",
    "flood_values = Q[flood_mask]\n",
    "flood_days = np.flatnonzero(flood_mask) + 1\n",
    "print(flood_values, flood_days)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**c)** How many days are below the low-flow threshold of 5 m³/s?\n",
    "\n",
    "**d)** What is the mean discharge from day 10 to day 20 (both included)?\n",
    "\n",
    "✅ **Check:** flood days `[ 8  9 10 11]`, 2 low-flow days, mean ≈ 14.28 m³/s."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_low_flow = (Q < 5).sum()\n",
    "mean_10_20 = Q[9:20].mean()\n",
    "print(n_low_flow, round(mean_10_20, 2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 3 — Uniform flow in a trapezoidal channel\n",
    "\n",
    "An irrigation canal has a trapezoidal section with bottom width $b = 8$ m, side slope $m = 1.5$, Manning coefficient $n = 0.03$ and bed slope $S = 0.0005$.\n",
    "\n",
    "$$A = (b + m\\,y)\\,y \\qquad P = b + 2\\,y\\sqrt{1+m^2} \\qquad R_h = \\frac{A}{P} \\qquad Q = \\frac{1}{n} A R_h^{2/3} S^{1/2}$$\n",
    "\n",
    "**a)** For 40 depths `y` between 0.1 m and 4.0 m, compute `A`, `P`, `R` and the discharge `Qc` **without any loop**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "b, m, n, S, g = 8.0, 1.5, 0.03, 0.0005, 9.81\n",
    "y = np.linspace(0.1, 4.0, 40)\n",
    "\n",
    "A = (b + m * y) * y\n",
    "P = b + 2 * y * np.sqrt(1 + m**2)\n",
    "R = A / P\n",
    "Qc = (1 / n) * A * R ** (2 / 3) * np.sqrt(S)\n",
    "print(Qc[-1].round(2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**b)** Compute the mean velocity $V = Q/A$ and the Froude number $Fr = V / \\sqrt{g\\,D}$, where $D = A/T$ is the hydraulic depth and $T = b + 2\\,m\\,y$ the top width. Is the flow ever supercritical ($Fr > 1$) in this range?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "V = Qc / A\n",
    "T = b + 2 * m * y\n",
    "D = A / T\n",
    "Fr = V / np.sqrt(g * D)\n",
    "print(\"max V:\", V.max().round(2), \"m/s | max Fr:\", Fr.max().round(3),\n",
    "      \"| supercritical depths:\", (Fr > 1).sum())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**c)** The canal must carry a design discharge of 50 m³/s. Since `Qc` increases with `y`, use `np.interp` to estimate the **normal depth** $y_n$.\n",
    "\n",
    "✅ **Check:** Q at 4.0 m ≈ 76.83 m³/s, max Fr ≈ 0.262 (always subcritical), $y_n$ ≈ 3.20 m."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Q_design = 50.0\n",
    "y_normal = np.interp(Q_design, Qc, y)\n",
    "print(round(y_normal, 3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 4 — Broadcasting: a design table\n",
    "\n",
    "We now study the same canal for 4 possible bed slopes.\n",
    "\n",
    "**a)** Using broadcasting (`np.newaxis`), build `Q_table` of shape `(40, 4)`: one row per depth `y`, one column per slope."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "slopes = np.array([0.0002, 0.0005, 0.001, 0.002])\n",
    "\n",
    "Q_table = (1 / n) * A[:, np.newaxis] * R[:, np.newaxis] ** (2 / 3) * np.sqrt(slopes[np.newaxis, :])\n",
    "print(Q_table.shape)\n",
    "print(Q_table[-1].round(1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**b)** For each slope, find the **first depth** of the grid whose capacity is at least `Q_design`. Hint: `(Q_table >= Q_design).argmax(axis=0)` gives the row index of the first `True` in each column.\n",
    "\n",
    "**c)** ⚠️ Look carefully at the result for the gentlest slope (0.0002). Is it physically correct? Use `.any(axis=0)` to detect the columns where **no** depth reaches the design discharge, and replace those depths by `np.nan`.\n",
    "\n",
    "✅ **Check:** last row ≈ `[ 48.6  76.8 108.7 153.7]`. Correct first depths: `[nan 3.2 2.7 2.3]` m."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "enough = Q_table >= Q_design\n",
    "first_idx = enough.argmax(axis=0)\n",
    "print(\"naive:\", y[first_idx])          # 0.1 m for S = 0.0002 is wrong: argmax of all-False is 0\n",
    "\n",
    "y_first = np.where(enough.any(axis=0), y[first_idx], np.nan)\n",
    "print(\"corrected:\", y_first)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 5 — Monthly rainfall statistics\n",
    "\n",
    "The next cell generates 20 years (2001–2020) of synthetic monthly rainfall (mm) with a Mediterranean pattern (wet winters, dry summers). `rain` has one row per year and one column per month."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = np.random.default_rng(2026)\n",
    "seasonal = np.array([60, 50, 45, 35, 20, 8, 2, 4, 15, 35, 55, 65])  # typical monthly mean (mm)\n",
    "rain = rng.gamma(2.0, seasonal / 2.0, size=(20, 12)).round(1)\n",
    "\n",
    "years = np.arange(2001, 2021)\n",
    "months = np.array([\"Jan\", \"Feb\", \"Mar\", \"Apr\", \"May\", \"Jun\",\n",
    "                   \"Jul\", \"Aug\", \"Sep\", \"Oct\", \"Nov\", \"Dec\"])\n",
    "print(rain.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**a)** Compute the annual totals and the mean annual rainfall.\n",
    "\n",
    "**b)** Which year was the wettest, and which the driest? Give their totals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "annual = rain.sum(axis=1)\n",
    "mean_annual = annual.mean()\n",
    "print(\"mean annual rainfall:\", mean_annual.round(1), \"mm\")\n",
    "\n",
    "wettest_year = years[annual.argmax()]\n",
    "driest_year = years[annual.argmin()]\n",
    "print(\"wettest:\", wettest_year, annual.max().round(1), \"mm | driest:\", driest_year, annual.min().round(1), \"mm\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**c)** Compute the mean rainfall of each month and the name of the wettest month on average.\n",
    "\n",
    "**d)** Compute the anomaly matrix (each value minus the mean of its month) using broadcasting. Check that the mean anomaly of every month is (almost) zero."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "monthly_mean = rain.mean(axis=0)\n",
    "wettest_month = months[monthly_mean.argmax()]\n",
    "print(monthly_mean.round(1), wettest_month)\n",
    "\n",
    "anomaly = rain - monthly_mean\n",
    "print(anomaly.shape, np.allclose(anomaly.mean(axis=0), 0))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**e)** How many months in the whole record exceed 100 mm?\n",
    "\n",
    "**f)** For each year, compute the share (%) of the annual rainfall that falls in the wet season (October to March). What is the average share?\n",
    "\n",
    "✅ **Check:** mean annual ≈ 410.4 mm, wettest 2006 (696.6 mm), driest 2014 (228.6 mm), wettest month Dec, 15 months above 100 mm, wet-season share ≈ 78.6 %."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_above_100 = (rain > 100).sum()\n",
    "print(n_above_100)\n",
    "\n",
    "wet_season = [9, 10, 11, 0, 1, 2]   # column indices of Oct, Nov, Dec, Jan, Feb, Mar\n",
    "wet_share = rain[:, wet_season].sum(axis=1) / annual * 100\n",
    "print(wet_share.round(1), wet_share.mean().round(1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 6 — A flood hydrograph with missing data\n",
    "\n",
    "The next cell builds a 72-hour flood hydrograph (hourly discharge, m³/s) above a baseflow of 4 m³/s. The gauge failed at a few time steps: those values are `np.nan`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = np.arange(72.0)                        # time (h)\n",
    "Q_base, Q_peak, t_peak, k = 4.0, 60.0, 18.0, 3.0\n",
    "Qh = Q_base + Q_peak * (t / t_peak * np.exp(1 - t / t_peak)) ** k\n",
    "Qh = Qh.round(2)\n",
    "Qh[[5, 6, 20, 21, 22, 47, 60]] = np.nan    # sensor failures\n",
    "print(Qh[:8])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**a)** How many values are missing? What does `Qh.max()` return, and why?\n",
    "\n",
    "**b)** Fill the gaps by linear interpolation in time (`np.interp`), into a new array `Qf`. Do not modify `Qh`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_missing = np.isnan(Qh).sum()\n",
    "print(n_missing, Qh.max())     # max is nan: NaN propagates\n",
    "\n",
    "ok = ~np.isnan(Qh)\n",
    "Qf = Qh.copy()\n",
    "Qf[~ok] = np.interp(t[~ok], t[ok], Qh[ok])\n",
    "print(np.isnan(Qf).sum())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**c)** Find the peak discharge and the time to peak.\n",
    "\n",
    "**d)** Compute the total flood volume and the direct-runoff volume (above the 4 m³/s baseflow) with `np.trapezoid`, in hm³ (1 hm³ = 10⁶ m³). Careful with units: $t$ is in hours.\n",
    "\n",
    "**e)** What is the fastest rise of discharge in one hour, and when does it occur? How many hours is the discharge above 30 m³/s?\n",
    "\n",
    "✅ **Check:** 7 missing values, peak 64.0 m³/s at t = 18 h, total volume ≈ 6.79 hm³, direct runoff ≈ 5.77 hm³, fastest rise ≈ 5.82 m³/s per hour (between t = 7 h and 8 h), 27 hours above 30 m³/s."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Q_max = Qf.max()\n",
    "t_max = t[Qf.argmax()]\n",
    "print(\"peak:\", Q_max, \"m³/s at t =\", t_max, \"h\")\n",
    "\n",
    "V_total = np.trapezoid(Qf, t * 3600) / 1e6\n",
    "V_direct = np.trapezoid(Qf - Q_base, t * 3600) / 1e6\n",
    "print(\"volumes:\", round(V_total, 2), \"hm³ total,\", round(V_direct, 2), \"hm³ direct runoff\")\n",
    "\n",
    "rise = np.diff(Qf)\n",
    "max_rise = rise.max()\n",
    "print(\"fastest rise:\", max_rise.round(2), \"m³/s/h between t =\", t[rise.argmax()], \"and\", t[rise.argmax() + 1], \"h\")\n",
    "\n",
    "hours_above_30 = (Qf > 30).sum()\n",
    "print(hours_above_30, \"hours above 30 m³/s\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bonus — How much faster is vectorization?\n",
    "\n",
    "Compute the wetted area of the canal of Exercise 3 for one million random depths between 0.1 and 4 m, first with a Python loop (list comprehension over `y_big.tolist()`), then vectorized. Time both with `time.perf_counter()` and print the speed-up factor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "rng = np.random.default_rng(0)\n",
    "y_big = rng.uniform(0.1, 4.0, 1_000_000)\n",
    "\n",
    "start = time.perf_counter()\n",
    "A_loop = [(b + m * yi) * yi for yi in y_big.tolist()]\n",
    "t_loop = time.perf_counter() - start\n",
    "\n",
    "start = time.perf_counter()\n",
    "A_vec = (b + m * y_big) * y_big\n",
    "t_vec = time.perf_counter() - start\n",
    "\n",
    "print(f\"loop: {t_loop:.3f} s | vectorized: {t_vec:.4f} s | speed-up: x{t_loop / t_vec:.0f}\")\n",
    "print(np.allclose(A_loop, A_vec))"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}