{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a326f97f",
   "metadata": {},
   "source": [
    "# A groundwater model for Schoonhoven\n",
    "\n",
    "In this notebook we build a transient model for the area around Schoonhoven. Surface water is added to the model using a winter and a summer stage using the drain package. For the river Lek, we build a river package with a fixed stage of NAP+0.0 m."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "429d47d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import flopy\n",
    "import geopandas as gpd\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from shapely.geometry import LineString, Point\n",
    "\n",
    "import nlmod\n",
    "from nlmod.plot import DatasetCrossSection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01be41f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "nlmod.util.get_color_logger(\"INFO\")\n",
    "nlmod.show_versions()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26954c60",
   "metadata": {},
   "source": [
    "## Model settings\n",
    "We define some model settings, like the name, the directory of the model files, the model extent and the time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f8f9e3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"Schoonhoven\"\n",
    "model_ws = \"09_schoonhoven\"\n",
    "figdir, cachedir = nlmod.util.get_model_dirs(model_ws)\n",
    "extent = [116_500, 120_000, 439_000, 442_000]\n",
    "time = pd.date_range(\"2020\", \"2023\", freq=\"MS\")  # monthly timestep"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7c5efff",
   "metadata": {},
   "source": [
    "## Download data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9cf078d9",
   "metadata": {},
   "source": [
    "### layer 'waterdeel' from bgt\n",
    "The location of the surface water bodies is obtained from the GeoDataFrame that was created in the the surface water notebook. We saved this data as a geosjon file and load it here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a94094ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "fname_bgt = os.path.join(\"..\", \"data_sources\", \"02_surface_water\", \"cache\", \"bgt.gpkg\")\n",
    "if not os.path.isfile(fname_bgt):\n",
    "    raise (\n",
    "        Exception(\n",
    "            f\"{fname_bgt} not found. Please run notebook 02_surface_water.ipynb in the 'data_sources' directory first\"\n",
    "        )\n",
    "    )\n",
    "bgt = gpd.read_file(fname_bgt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1eb9d06",
   "metadata": {},
   "source": [
    "#### Plot summer stage of surface water bodies\n",
    "We can plot the summer stage. There are some surface water bodies without a summer stage, because the 'bronhouder' is not a water board. The main one is the river Lek, but there are also some surface water bodies without a summer stage in the north of the model area."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a328fb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "f, ax = nlmod.plot.get_map(extent)\n",
    "norm = matplotlib.colors.Normalize(vmin=-3, vmax=1)\n",
    "cmap = \"viridis\"\n",
    "bgt.plot(\"summer_stage\", ax=ax, norm=norm, cmap=cmap)\n",
    "nlmod.plot.colorbar_inside(norm=norm, cmap=cmap);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4aea6ed8",
   "metadata": {},
   "source": [
    "If no information about the stage is available, a constant stage is set to the minimal height of the digital terrain model (AHN) near the surface water body. We can plot these values as well:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1e552e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "f, ax = nlmod.plot.get_map(extent)\n",
    "bgt.plot(\"ahn_min\", ax=ax, norm=norm, cmap=cmap)\n",
    "nlmod.plot.colorbar_inside(norm=norm, cmap=cmap);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cefdaba6",
   "metadata": {},
   "source": [
    "### KNMI\n",
    "KNMI precipitation and evaporation data is used to estimate the recharge."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f829cfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# knmi data\n",
    "oc_knmi = nlmod.read.knmi.download_knmi(extent=extent, delr=100., delc=100., start='2000-1-1', end='2025-1-1')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbf531fa",
   "metadata": {},
   "source": [
    "### REGIS\n",
    "For the schematisation of the subsurface we use REGIS. Let's download this data for the required extent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22de3209",
   "metadata": {},
   "outputs": [],
   "source": [
    "layer_model = nlmod.read.regis.get_combined_layer_models(\n",
    "    extent,\n",
    "    use_regis=True,\n",
    "    use_geotop=False,\n",
    "    cachedir=cachedir,\n",
    "    cachename=\"layer_model.nc\",\n",
    ")\n",
    "layer_model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f64d15e",
   "metadata": {},
   "source": [
    "We then create a regular grid, add necessary variables and fill NaN's. For example, REGIS does not contain information about the hydraulic conductivity of the first layer ('HLc'). These NaN's are replaced by a default hydraulic conductivity (kh) of 1 m/d. This probably is not a good representation of the conductivity, but at least the model will run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ed4d540",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = nlmod.to_model_ds(layer_model, model_name, model_ws, delr=100.0, delc=100.0)\n",
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "363f0c3d",
   "metadata": {},
   "source": [
    "## Add grid refinement\n",
    "With the refine method, we can add grid refinement. The model will then use the disv-package instead of the dis-package. We can also test if the disv-package gives the same results as the dis-package by not specifying refinement_features: ds = nlmod.grid.refine(ds).\n",
    "\n",
    "This notebook can be run with or without running the cell below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed07f7df",
   "metadata": {},
   "outputs": [],
   "source": [
    "refinement_features = [(bgt[bgt[\"bronhouder\"] == \"L0002\"].dissolve().boundary, 2)]\n",
    "ds = nlmod.grid.refine(ds, refinement_features=refinement_features)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a4923c9",
   "metadata": {},
   "source": [
    "## Add information about time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d435d2b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = nlmod.time.set_ds_time(ds, time=time, start=3652)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3474fc22",
   "metadata": {},
   "source": [
    "## Add knmi recharge to the model dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00699ba5",
   "metadata": {},
   "outputs": [],
   "source": [
    "knmi_ds = nlmod.read.knmi.discretize_knmi(ds, oc_knmi, cachedir=cachedir, cachename=\"recharge\")\n",
    "ds.update(knmi_ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1447f11a",
   "metadata": {},
   "source": [
    "## Create a groundwater flow model\n",
    "Using the data from the xarray Dataset `ds` we generate a groundwater flow model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99dfcca7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create simulation\n",
    "sim = nlmod.sim.sim(ds)\n",
    "\n",
    "# create time discretisation\n",
    "tdis = nlmod.sim.tdis(ds, sim)\n",
    "\n",
    "# create ims\n",
    "ims = nlmod.sim.ims(sim)\n",
    "\n",
    "# create groundwater flow model\n",
    "gwf = nlmod.gwf.gwf(ds, sim)\n",
    "\n",
    "# Create discretization\n",
    "dis = nlmod.gwf.dis(ds, gwf)\n",
    "\n",
    "# create node property flow\n",
    "npf = nlmod.gwf.npf(ds, gwf, save_flows=True)\n",
    "\n",
    "# Create the initial conditions package\n",
    "ic = nlmod.gwf.ic(ds, gwf, starting_head=0.0)\n",
    "\n",
    "# Create the output control package\n",
    "oc = nlmod.gwf.oc(ds, gwf)\n",
    "\n",
    "# create storagee package\n",
    "sto = nlmod.gwf.sto(ds, gwf)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0d5d342",
   "metadata": {},
   "source": [
    "## Process surface water\n",
    "We intersect the surface water bodies with the grid, set a default bed resistance of 1 day, and seperate the large river 'Lek' form the other surface water bodies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d82c76c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "bed_resistance = 1.0\n",
    "\n",
    "bgt_grid = nlmod.grid.gdf_to_grid(bgt, ds).set_index(\"cellid\")\n",
    "\n",
    "bgt_grid[\"cond\"] = bgt_grid.area / bed_resistance\n",
    "\n",
    "# handle the lek as a river\n",
    "mask = bgt_grid[\"bronhouder\"] == \"L0002\"\n",
    "lek = bgt_grid[mask]\n",
    "bgt_grid = bgt_grid[~mask]\n",
    "\n",
    "# handle grote gracht and oude haven to model as a lake\n",
    "ids_grote_gracht = [\n",
    "    \"W0656.774b12049d9a4252bd61c4ea442b5158\",\n",
    "    \"W0656.59ab56cf0b2d4f15894c24369f0748df\",\n",
    "]\n",
    "ids_oude_haven = [\n",
    "    \"W0656.a6013e26cd9442de86eac2295eb0012b\",\n",
    "    \"W0656.2053970c192b4fe48bba882842e53eb5\",\n",
    "    \"W0656.540780b5c9944b51b53d8a98445b315a\",\n",
    "    \"W0656.a7c39fcaabe149c3b9eb4823f76db024\",\n",
    "    \"W0656.cb3c3a25de4141d18c573b561f02e84a\",\n",
    "]\n",
    "\n",
    "mask = bgt_grid[\"identificatie\"].isin(ids_grote_gracht + ids_oude_haven)\n",
    "lakes = bgt_grid[mask].copy()\n",
    "lakes[\"name\"] = \"\"\n",
    "lakes.loc[lakes[\"identificatie\"].isin(ids_grote_gracht), \"name\"] = \"grotegracht\"\n",
    "lakes.loc[lakes[\"identificatie\"].isin(ids_oude_haven), \"name\"] = \"oudehaven\"\n",
    "bgt_grid = bgt_grid[~mask]\n",
    "\n",
    "# cut rainfall and evaporation from model dataset\n",
    "lak_rainfall, lak_evaporation = nlmod.gwf.lake.copy_meteorological_data_from_ds(\n",
    "    lakes, ds, boundname_column=\"name\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd9adde9",
   "metadata": {},
   "source": [
    "### Lek as river\n",
    "Model the river Lek as a river with a fixed stage of 0.5 m NAP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "beff6f4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "lek[\"stage\"] = 0.0\n",
    "lek[\"rbot\"] = -3.0\n",
    "spd = nlmod.gwf.surface_water.build_spd(lek, \"RIV\", ds)\n",
    "riv = flopy.mf6.ModflowGwfriv(gwf, stress_period_data={0: spd})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8be299bd",
   "metadata": {},
   "source": [
    "### Other surface water as drains\n",
    "Model the other surface water using the drain package, with a summer stage and a  winter stage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e5e0a96",
   "metadata": {},
   "outputs": [],
   "source": [
    "drn = nlmod.gwf.surface_water.gdf_to_seasonal_pkg(bgt_grid, gwf, ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e894b678-42b7-469a-9b2f-675417d2168e",
   "metadata": {},
   "source": [
    "### Add lake\n",
    "\n",
    "Model de \"grote gracht\" and \"Oude Haven\" as lakes. Let the grote gracht overflow into the oude Haven."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1240a110-b71e-4afd-9c8f-0a1d868e8c2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# add general properties to the lake gdf\n",
    "summer_months = (4, 5, 6, 7, 8, 9)\n",
    "if pd.to_datetime(ds.time.start).month in summer_months:\n",
    "    lakes[\"strt\"] = lakes[\"summer_stage\"]\n",
    "else:\n",
    "    lakes[\"strt\"] = lakes[\"winter_stage\"]\n",
    "lakes[\"clake\"] = 100\n",
    "\n",
    "# add inflow to Oude Haven\n",
    "# ds['inflow_lake'] = xr.DataArray(100, dims=[\"time\"], coords=dict(time=ds.time))\n",
    "# lakes.loc[lakes['identificatie'].isin(ids_oude_haven), 'INFLOW'] = 'inflow_lake'\n",
    "\n",
    "# add outlet to Oude Haven, water flows from Oude Haven to Grote Gracht.\n",
    "lakes.loc[lakes[\"name\"] == \"oudehaven\", \"lakeout\"] = \"grotegracht\"\n",
    "lakes.loc[lakes[\"name\"] == \"oudehaven\", \"outlet_invert\"] = 1.0  # overstort hoogte\n",
    "\n",
    "# add lake to groundwaterflow model\n",
    "lak = nlmod.gwf.lake_from_gdf(\n",
    "    gwf,\n",
    "    lakes,\n",
    "    ds,\n",
    "    boundname_column=\"name\",\n",
    "    rainfall=lak_rainfall,\n",
    "    evaporation=lak_evaporation,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4cc10d1-d14c-4ddf-b32c-0444269ee38c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create recharge package\n",
    "rch = nlmod.gwf.rch(ds, gwf)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a7f416e",
   "metadata": {},
   "source": [
    "## Run the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c53cee0",
   "metadata": {},
   "outputs": [],
   "source": [
    "nlmod.sim.write_and_run(sim, ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "712f7f16",
   "metadata": {},
   "source": [
    "## Post-processing\n",
    "### Get the simulated head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "470f49b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "head = nlmod.gwf.get_heads_da(ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c4b8ddb",
   "metadata": {},
   "source": [
    "### Plot the average head in the first layer on a map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ff0c79d",
   "metadata": {},
   "outputs": [],
   "source": [
    "norm = matplotlib.colors.Normalize(-2.5, 0.0)\n",
    "\n",
    "\n",
    "pc = nlmod.plot.map_array(\n",
    "    head.sel(layer=\"HLc\").mean(\"time\"),\n",
    "    ds,\n",
    "    norm=norm,\n",
    "    colorbar=True,\n",
    "    colorbar_label=\"head [m NAP]\",\n",
    "    title=\"mean head\",\n",
    ")\n",
    "bgt.dissolve().plot(ax=pc.axes, edgecolor=\"k\", facecolor=\"none\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd9ddb76",
   "metadata": {},
   "source": [
    "### Plot the average head in a cross-section, from north to south"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59684fd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = 118228.0\n",
    "line = [(x, 439000), (x, 442000)]\n",
    "\n",
    "f, ax = plt.subplots(figsize=(10, 6), layout=\"constrained\")\n",
    "dcs = DatasetCrossSection(ds, line, ax=ax, zmin=-100.0, zmax=10.0)\n",
    "pc = dcs.plot_array(head.mean(\"time\"), norm=norm, head=head.mean(\"time\"))\n",
    "\n",
    "# add labels with layer names\n",
    "cbar = nlmod.plot.colorbar_inside(pc)\n",
    "dcs.plot_grid()\n",
    "dcs.plot_layers(colors=\"none\", min_label_area=1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01ddcb4f",
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "source": [
    "### Animate a cross-section with heads"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e195442",
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "# x = np.mean(extent[:2])\n",
    "# line = [(x, extent[2]), (x, extent[3])]\n",
    "\n",
    "# f, ax = plt.subplots(figsize=(10, 6))\n",
    "# norm\n",
    "# dcs = DatasetCrossSection(ds, line, ax=ax, zmin=-100.0, zmax=10.0)\n",
    "\n",
    "# # add labels with layer names\n",
    "# ax.set_xlabel(\"distance [m]\")\n",
    "# ax.set_ylabel(\"elevation [mNAP]\")\n",
    "\n",
    "# dcs.plot_grid(lw=0.25, edgecolor=\"k\", alpha=0.5, vertical=False)\n",
    "# dcs.plot_layers(alpha=0.0, min_label_area=5e4)\n",
    "# dcs.plot_surface(ds[\"top\"], lw=1.0, color=\"k\")\n",
    "\n",
    "# fname = os.path.join(ds.figdir, f\"anim_xsec_x{int(x)}_head.mp4\")\n",
    "# dcs.animate(\n",
    "#     head,\n",
    "#     cmap=\"Spectral_r\",\n",
    "#     head=head,\n",
    "#     plot_title=f\"doorsnede at x={int(x)}\",\n",
    "#     date_fmt=\"%Y-%m-%d\",\n",
    "#     fname=fname,\n",
    "# )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d543af4",
   "metadata": {},
   "source": [
    "### plot a time series at a certain location"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94b00624",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = 118228\n",
    "y = 439870\n",
    "head_point = nlmod.gwf.get_head_at_point(head, x=x, y=y, ds=ds)\n",
    "fig, ax = plt.subplots(1, 1, figsize=(10, 8))\n",
    "handles = head_point.plot.line(ax=ax, hue=\"layer\")\n",
    "ax.set_ylabel(\"head [m NAP]\");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a429d806-cf6d-4b3f-9d59-77f46fa66759",
   "metadata": {},
   "source": [
    "### plot the lake stages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d9f3bd7-a0c2-4aec-a52b-6278053f4fef",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv(os.path.join(model_ws, \"lak_STAGE.csv\"), index_col=0)\n",
    "df.index = ds.time.values\n",
    "ax = df.plot(figsize=(10, 3))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9355de12",
   "metadata": {
    "raw_mimetype": "text/x-python"
   },
   "source": [
    "## Compare with BRO measurements"
   ]
  },
  {
   "cell_type": "raw",
   "id": "9faabb6c-728a-477d-a659-2941ee684bbc",
   "metadata": {
    "raw_mimetype": "text/x-python",
    "vscode": {
     "languageId": "raw"
    }
   },
   "source": [
    "oc = nlmod.read.bro.download_bro(extent, cachedir=cachedir, cachename=\"bro\")\n",
    "oc_mod = nlmod.read.bro.add_modelled_head(oc, gwf, ds=ds)"
   ]
  },
  {
   "cell_type": "raw",
   "id": "1c736d07-5623-4df3-97be-b5ed14d301e1",
   "metadata": {
    "raw_mimetype": "text/x-python"
   },
   "source": [
    "# create interactive map\n",
    "oc_mod.plots.interactive_map(\"figures\", add_screen_to_legend=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25b68351",
   "metadata": {},
   "source": [
    "### Plot some properties of the first layer\n",
    "We can plot some properties of the first layer, called HLc. As REGIS does not contain data about hydraulic conductivities for this layer, default values of 1 m/d for kh and 0.1 m/d for hv are used, which can be seen in the graphs below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8196852",
   "metadata": {},
   "outputs": [],
   "source": [
    "layer = \"HLc\"\n",
    "\n",
    "f, axes = nlmod.plot.get_map(extent, nrows=2, ncols=2)\n",
    "variables = [\"top\", \"kh\", \"botm\", \"kv\"]\n",
    "for i, variable in enumerate(variables):\n",
    "    ax = axes.ravel()[i]\n",
    "    if variable == \"top\":\n",
    "        if layer == ds.layer[0]:\n",
    "            da = ds[\"top\"]\n",
    "        else:\n",
    "            da = ds[\"botm\"][np.where(ds.layer == layer)[0][0] - 1]\n",
    "    else:\n",
    "        da = ds[variable].sel(layer=layer)\n",
    "    pc = nlmod.plot.data_array(da, ds=ds, ax=ax)\n",
    "    nlmod.plot.colorbar_inside(pc, ax=ax)\n",
    "    ax.text(\n",
    "        0.5,\n",
    "        0.98,\n",
    "        f\"{variable} in layer {layer}\",\n",
    "        ha=\"center\",\n",
    "        va=\"top\",\n",
    "        transform=ax.transAxes,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "312f74ed-459e-4ea3-afb0-18da3eb5639f",
   "metadata": {},
   "source": [
    "## Add pathlines\n",
    "\n",
    "We create a modpath model which calculates the pathlines. We calculate the pathlines that start in the center of the modflow cells with a river boundary condition (the cells in the \"Lek\" river)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0682c7a-d6a2-49c2-b883-9b587f14e59b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a modpath model\n",
    "mpf = nlmod.modpath.mpf(gwf)\n",
    "\n",
    "# create the basic modpath package\n",
    "_mpfbas = nlmod.modpath.bas(mpf)\n",
    "\n",
    "# get the nodes from a package\n",
    "nodes = nlmod.modpath.package_to_nodes(gwf, \"RIV_0\", ibound=mpf.ib)\n",
    "\n",
    "# create a particle tracking group from cell centers\n",
    "pg = nlmod.modpath.pg_from_pd(nodes, localx=0.5, localy=0.5, localz=0.5)\n",
    "\n",
    "# create the modpath simulation file\n",
    "mpsim = nlmod.modpath.sim(mpf, pg, \"forward\", gwf=gwf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1d55129-bf37-4e81-ab15-8c318bd63ffb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# run modpath model\n",
    "nlmod.modpath.write_and_run(mpf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acad9afd-346d-43b4-92f2-cf9978a54083",
   "metadata": {},
   "outputs": [],
   "source": [
    "pdata = nlmod.modpath.load_pathline_data(mpf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25df4e09-f584-4e15-bf71-bcb4788ed38a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_segments(x, y, segments=None):\n",
    "    # split each flow path into multiple line segments\n",
    "    return [np.column_stack([x[i : i + 2], y[i : i + 2]]) for i in range(len(x) - 1)]\n",
    "\n",
    "\n",
    "def get_array(time, to_year=True):\n",
    "    # for each line-segment use the average time as the color\n",
    "    array = (time[:-1] + time[1:]) / 2\n",
    "    if to_year:\n",
    "        array = array / 365.25\n",
    "    return array\n",
    "\n",
    "\n",
    "cmap = plt.get_cmap(\"turbo\")\n",
    "norm = matplotlib.colors.BoundaryNorm(\n",
    "    [0, 1, 2, 5, 10, 25, 50, 100, 200, 500], cmap.N, extend=\"max\"\n",
    ")\n",
    "\n",
    "# get line segments and color values\n",
    "segments = []\n",
    "array = []\n",
    "for pid in np.unique(pdata[\"particleid\"]):\n",
    "    pf = pdata[pdata[\"particleid\"] == pid]\n",
    "    segments.extend(get_segments(pf[\"x\"], pf[\"y\"]))\n",
    "    array.extend(get_array(pf[\"time\"]))\n",
    "\n",
    "f, ax = nlmod.plot.get_map(extent)\n",
    "lc = matplotlib.collections.LineCollection(\n",
    "    segments, cmap=cmap, norm=norm, array=array, linewidth=1.0\n",
    ")\n",
    "line = ax.add_collection(lc)\n",
    "nlmod.plot.colorbar_inside(line, label=\"Travel time (years)\")\n",
    "\n",
    "bgt.dissolve().plot(ax=ax, edgecolor=\"k\", facecolor=\"none\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a3e8401-9504-47d9-99ef-f75d67182a9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = 118228.0\n",
    "line = LineString([(x, 439000), (x, 442000)])\n",
    "\n",
    "# get line segments and color values\n",
    "segments = []\n",
    "array = []\n",
    "for pid in np.unique(pdata[\"particleid\"]):\n",
    "    pf = pdata[pdata[\"particleid\"] == pid]\n",
    "    d = line.distance(Point(pf[\"x\"][0], pf[\"y\"][0]))\n",
    "    if d < 200.0:\n",
    "        x = [line.project(Point(x, y)) for x, y in zip(pf[\"x\"], pf[\"y\"])]\n",
    "        segments.extend(get_segments(x, pf[\"z\"]))\n",
    "        array.extend(get_array(pf[\"time\"]))\n",
    "\n",
    "f, ax = plt.subplots(figsize=(10, 6), layout=\"constrained\")\n",
    "ax.grid()\n",
    "dcs = DatasetCrossSection(ds, line, ax=ax, zmin=-100.0, zmax=10.0)\n",
    "lc = matplotlib.collections.LineCollection(\n",
    "    segments, cmap=cmap, norm=norm, array=array, linewidth=1.0\n",
    ")\n",
    "line = ax.add_collection(lc)\n",
    "nlmod.plot.colorbar_inside(line, label=\"Travel time (years)\")\n",
    "# add grid\n",
    "dcs.plot_grid()\n",
    "# add labels with layer names\n",
    "dcs.plot_layers(alpha=0.0, min_label_area=1000)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}