{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Particle Tracking with MODPATH\n",
    "\n",
    "This notebook shows how to create a particle tracking model using MODPATH.\n",
    "    \n",
    "## To-Do\n",
    "- make the examples from a package and from a model layer faster\n",
    "- update toc    \n",
    "- add cross section"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "import flopy\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "\n",
    "import nlmod"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nlmod.util.get_color_logger(\"INFO\")\n",
    "nlmod.show_versions()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Groundwater Flow Model\n",
    "\n",
    "We use the groundwater flow model from the [03_local_grid_refinement notebook](03_local_grid_refinement.ipynb). Make sure to run that notebook before you run this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load lgr model dataset\n",
    "model_ws = os.path.join(\"..\", \"examples\", \"03_local_grid_refinement\")\n",
    "model_name = \"IJm_planeten\"\n",
    "\n",
    "ds = xr.open_dataset(os.path.join(model_ws, f\"{model_name}.nc\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load simulation and groundwaterflow model\n",
    "# set exe_name to point to mf6 version in nlmod bin directory\n",
    "exe_name = os.path.join(os.path.dirname(nlmod.__file__), \"bin\", \"mf6\")\n",
    "if sys.platform.startswith(\"win\"):\n",
    "    exe_name += \".exe\"\n",
    "\n",
    "sim = flopy.mf6.MFSimulation.load(\"mfsim.nam\", sim_ws=model_ws, exe_name=exe_name)\n",
    "gwf = sim.get_model(model_name=model_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Modpath"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Backward tracking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# list with xy coordinates to start particle tracking from\n",
    "xy_start = [(101500, 496500), (101500, 499100)]\n",
    "\n",
    "# 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",
    "# find the nodes for given xy\n",
    "nodes = nlmod.modpath.xy_to_nodes(xy_start, mpf, ds, layer=5)\n",
    "\n",
    "# create a particle tracking group at the cell faces\n",
    "pg = nlmod.modpath.pg_from_fdt(nodes)\n",
    "\n",
    "# create the modpath simulation file\n",
    "mpsim = nlmod.modpath.sim(mpf, pg, \"backward\", gwf=gwf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run modpath model\n",
    "nlmod.modpath.write_and_run(mpf, script_path=\"10_modpath.ipynb\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pdata = nlmod.modpath.load_pathline_data(mpf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f, ax = plt.subplots(nrows=1, ncols=1, figsize=(10, 10))\n",
    "ax.set_aspect(\"equal\")\n",
    "ax = nlmod.plot.modelgrid(ds, ax=ax)\n",
    "\n",
    "for pid in np.unique(pdata[\"particleid\"]):\n",
    "    pf = pdata[pdata[\"particleid\"] == pid]\n",
    "    ax.plot(pf[\"x\"], pf[\"y\"], color=\"k\", linewidth=0.5)\n",
    "ax.plot(pf[\"x\"], pf[\"y\"], color=\"k\", linewidth=0.5, label=\"pathline\")\n",
    "\n",
    "cids = [nlmod.grid.get_icell2d_from_xy(xy[0], xy[1], ds) for xy in xy_start]\n",
    "ax.plot(\n",
    "    ds.x[cids],\n",
    "    ds.y[cids],\n",
    "    label=\"start of backwards tracking\",\n",
    "    ls=\"\",\n",
    "    marker=\"o\",\n",
    "    color=\"red\",\n",
    ")\n",
    "ax.set_title(\"pathlines\")\n",
    "ax.legend(loc=\"upper right\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f, ax = plt.subplots(nrows=1, ncols=1, figsize=(12, 4))\n",
    "\n",
    "for _, pid in enumerate(np.unique(pdata[\"particleid\"])):\n",
    "    pf = pdata[pdata[\"particleid\"] == pid]\n",
    "    x0, y0, z0 = pf[[\"x\", \"y\", \"z\"]][0]\n",
    "    distance = np.sqrt((pf[\"x\"] - x0) ** 2 + (pf[\"y\"] - y0) ** 2 + (pf[\"z\"] - z0) ** 2)\n",
    "    ax.plot(pf[\"time\"] / 365.25, distance, label=pid)\n",
    "\n",
    "ax.set_ylabel(\"distance [m]\")\n",
    "ax.set_xlabel(\"time [year]\")\n",
    "ax.set_title(\"distance travelled per particle\")\n",
    "ax.grid()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Forward tracking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# list with xy coordinates to start particle tracking from\n",
    "xy_start = [(101500, 496500), (101500, 499100)]\n",
    "\n",
    "# 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",
    "# find the nodes for given xy\n",
    "nodes = nlmod.modpath.xy_to_nodes(xy_start, mpf, ds, layer=5)\n",
    "\n",
    "# create a particle tracking group at the cell faces\n",
    "# pg = nlmod.modpath.pg_from_pd(nodes, localx=0.5, localy=0.5, localz=1.0)\n",
    "pg = nlmod.modpath.pg_from_fdt(nodes)\n",
    "\n",
    "# create the modpath simulation file\n",
    "mpsim = nlmod.modpath.sim(mpf, pg, \"forward\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run modpath model\n",
    "nlmod.modpath.write_and_run(mpf, script_path=\"10_modpath.ipynb\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pdata = nlmod.modpath.load_pathline_data(mpf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f, axl = plt.subplots(nrows=1, ncols=3, figsize=(30, 10))\n",
    "for i, ax in enumerate(axl):\n",
    "    ax.set_aspect(\"equal\")\n",
    "    ax = nlmod.plot.modelgrid(ds, ax=ax)\n",
    "\n",
    "    for pid in np.unique(pdata[\"particleid\"]):\n",
    "        pf = pdata[pdata[\"particleid\"] == pid]\n",
    "        ax.plot(pf[\"x\"], pf[\"y\"], color=\"k\", linewidth=0.5)\n",
    "    ax.plot(pf[\"x\"], pf[\"y\"], color=\"k\", linewidth=0.5, label=\"pathline\")\n",
    "\n",
    "    cids = [nlmod.grid.get_icell2d_from_xy(xy[0], xy[1], ds) for xy in xy_start]\n",
    "    ax.plot(\n",
    "        ds.x[cids],\n",
    "        ds.y[cids],\n",
    "        label=\"start of forward tracking\",\n",
    "        ls=\"\",\n",
    "        marker=\"o\",\n",
    "        color=\"red\",\n",
    "    )\n",
    "    ax.set_title(\"pathlines\")\n",
    "    ax.legend(loc=\"upper right\")\n",
    "\n",
    "    if i == 1:\n",
    "        ax.set_xlim(101200, 101700)\n",
    "        ax.set_ylim(498700, 499300)\n",
    "    elif i == 2:\n",
    "        ax.set_xlim(101200, 101700)\n",
    "        ax.set_ylim(496300, 496700)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f, ax = plt.subplots(nrows=1, ncols=1, figsize=(12, 4))\n",
    "\n",
    "for _, pid in enumerate(np.unique(pdata[\"particleid\"])):\n",
    "    pf = pdata[pdata[\"particleid\"] == pid]\n",
    "    x0, y0, z0 = pf[[\"x\", \"y\", \"z\"]][0]\n",
    "    distance = np.sqrt((pf[\"x\"] - x0) ** 2 + (pf[\"y\"] - y0) ** 2 + (pf[\"z\"] - z0) ** 2)\n",
    "    ax.plot(pf[\"time\"] / 365.25, distance, label=pid)\n",
    "\n",
    "ax.set_ylabel(\"distance [m]\")\n",
    "ax.set_xlabel(\"time [year]\")\n",
    "ax.set_title(\"distance travelled per particle\")\n",
    "ax.grid()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Backward tracking from general head boundaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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, \"GHB\", 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, \"backward\", gwf=gwf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run modpath model\n",
    "nlmod.modpath.write_and_run(mpf, script_path=\"10_modpath.ipynb\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pdata = nlmod.modpath.load_pathline_data(mpf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f, axl = plt.subplots(nrows=1, ncols=3, figsize=(30, 10))\n",
    "for i, ax in enumerate(axl):\n",
    "    ax.set_aspect(\"equal\")\n",
    "    ax = nlmod.plot.modelgrid(ds, ax=ax)\n",
    "\n",
    "    for pid in np.unique(pdata[\"particleid\"]):\n",
    "        pf = pdata[pdata[\"particleid\"] == pid]\n",
    "        ax.plot(pf[\"x\"], pf[\"y\"], color=\"k\", linewidth=0.5)\n",
    "    ax.plot(pf[\"x\"], pf[\"y\"], color=\"k\", linewidth=0.5, label=\"pathline\")\n",
    "\n",
    "    if i > 0:\n",
    "        cids = np.where((ds[\"rws_oppwater_cond\"] != 0).values)[0]\n",
    "        ax.plot(\n",
    "            ds.x[cids],\n",
    "            ds.y[cids],\n",
    "            label=\"start of backwards tracking\",\n",
    "            ls=\"\",\n",
    "            marker=\"o\",\n",
    "            color=\"red\",\n",
    "        )\n",
    "    ax.set_title(\"pathlines\")\n",
    "    ax.legend(loc=\"upper right\")\n",
    "\n",
    "    if i == 1:\n",
    "        ax.set_xlim(101000, 102000)\n",
    "        ax.set_ylim(498300, 499300)\n",
    "    elif i == 2:\n",
    "        ax.set_xlim(101000, 102000)\n",
    "        ax.set_ylim(496300, 497300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f, ax = plt.subplots(nrows=1, ncols=1, figsize=(12, 4))\n",
    "\n",
    "for i, pid in enumerate(np.unique(pdata[\"particleid\"])):\n",
    "    pf = pdata[pdata[\"particleid\"] == pid]\n",
    "    x0, y0, z0 = pf[[\"x\", \"y\", \"z\"]][0]\n",
    "    distance = np.sqrt((pf[\"x\"] - x0) ** 2 + (pf[\"y\"] - y0) ** 2 + (pf[\"z\"] - z0) ** 2)\n",
    "    ax.plot(pf[\"time\"] / 365.25, distance, label=pid)\n",
    "\n",
    "ax.set_xlim(0, 5000)\n",
    "ax.set_ylabel(\"distance [m]\")\n",
    "ax.set_xlabel(\"time [year]\")\n",
    "ax.set_title(\"distance travelled per particle\")\n",
    "ax.grid()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Forward tracking from each cell in the top layer\n",
    "\n",
    "Stop after 10 years."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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 nodes of all cells in the top modellayer\n",
    "nodes = nlmod.modpath.layer_to_nodes(mpf, 0)\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, stoptime=10 * 365)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run modpath model\n",
    "nlmod.modpath.write_and_run(mpf, script_path=\"10_modpath.ipynb\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pdata = nlmod.modpath.load_pathline_data(mpf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f, axl = plt.subplots(nrows=1, ncols=3, figsize=(30, 10))\n",
    "for i, ax in enumerate(axl):\n",
    "    ax.set_aspect(\"equal\")\n",
    "    ax = nlmod.plot.modelgrid(ds, ax=ax)\n",
    "\n",
    "    for pid in np.unique(pdata[\"particleid\"]):\n",
    "        pf = pdata[pdata[\"particleid\"] == pid]\n",
    "        ax.plot(pf[\"x\"], pf[\"y\"], color=\"k\", linewidth=0.5)\n",
    "    ax.plot(pf[\"x\"], pf[\"y\"], color=\"k\", linewidth=0.5, label=\"pathline\")\n",
    "\n",
    "    if i > 0:\n",
    "        ax.plot(\n",
    "            ds.x.values,\n",
    "            ds.y.values,\n",
    "            label=\"start of forward tracking\",\n",
    "            ls=\"\",\n",
    "            marker=\"o\",\n",
    "            color=\"red\",\n",
    "        )\n",
    "    ax.set_title(\"pathlines\")\n",
    "    ax.legend(loc=\"upper right\")\n",
    "\n",
    "    if i == 1:\n",
    "        ax.set_xlim(101000, 102000)\n",
    "        ax.set_ylim(498300, 499300)\n",
    "    elif i == 2:\n",
    "        ax.set_xlim(101000, 102000)\n",
    "        ax.set_ylim(496300, 497300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f, ax = plt.subplots(nrows=1, ncols=1, figsize=(12, 4))\n",
    "\n",
    "for i, pid in enumerate(np.unique(pdata[\"particleid\"])):\n",
    "    pf = pdata[pdata[\"particleid\"] == pid]\n",
    "    x0, y0, z0 = pf[[\"x\", \"y\", \"z\"]][0]\n",
    "    distance = np.sqrt((pf[\"x\"] - x0) ** 2 + (pf[\"y\"] - y0) ** 2 + (pf[\"z\"] - z0) ** 2)\n",
    "    ax.plot(pf[\"time\"] / 365.25, distance, label=pid)\n",
    "\n",
    "ax.set_xlim(0, 11)\n",
    "ax.set_ylabel(\"distance [m]\")\n",
    "ax.set_xlabel(\"time [year]\")\n",
    "ax.set_title(\"distance travelled per particle\")\n",
    "ax.grid()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}