diff --git a/docs/python/examples.ipynb b/docs/python/examples.ipynb
index 5256fc741..e7e3a67e0 100644
--- a/docs/python/examples.ipynb
+++ b/docs/python/examples.ipynb
@@ -2080,6 +2080,12 @@
    "source": [
     "datadir_basin = datadir / \"doc_example/results/basin.arrow\"\n",
     "df_basin = pd.read_feather(datadir_basin)\n",
+    "df_basin[\"vertical_flux\"] = (\n",
+    "    df_basin[\"precipitation\"]\n",
+    "    - df_basin[\"evaporation\"]\n",
+    "    + df_basin[\"drainage\"]\n",
+    "    + df_basin[\"infiltration\"]\n",
+    ")\n",
     "df_basin_wide = df_basin.pivot_table(\n",
     "    index=\"time\", columns=\"node_id\", values=[\"storage\", \"level\", \"vertical_flux\"]\n",
     ")\n",