examples/notebooks/MODFLOW-API_extensions_objects.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5acdb8ec",
   "metadata": {},
   "source": [
    "# Interacting with MODFLOW-API Interface objects\n",
    "\n",
    "The purpose of this notebook is to show the MODFLOW-API interface objects and introduce the user to the data types and how to interact with the objects. \n",
    "\n",
    "**Note**: This notebook shows how to run a model using the modflow-api at the end of the notebook. However, the majority of the notebook is an illustration of how to access and work with the data types that are returned to a user defined callback function. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b0e6a93",
   "metadata": {},
   "outputs": [],
   "source": [
    "import modflowapi\n",
    "from modflowapi.extensions import ApiSimulation\n",
    "from pathlib import Path\n",
    "import platform"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e654316",
   "metadata": {},
   "source": [
    "Define the paths to the model and the Modflow shared library"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66c0f32c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "sim_ws = Path(\"../data/dis_model\")\n",
    "dll = \"./libmf6\"\n",
    "if platform.system().lower() == \"windows\":\n",
    "    ext = \".dll\"\n",
    "elif platform.system().lower() == \"linux\":\n",
    "    ext = \".so\"\n",
    "else:\n",
    "    ext = \".dylib\"\n",
    "    \n",
    "dll = Path(dll + ext)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d258b430",
   "metadata": {},
   "source": [
    "#### Initializing the API model object\n",
    "\n",
    "The modflow api allows users to initialize an object that can be used to interact with the model. This processes is done automatically with the `modflowapi.run_model` function call. We're going to initialize an object outside of that call as a demonstration of the interface data objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c18c3bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "mf6 = modflowapi.ModflowApi(dll, working_directory=sim_ws)\n",
    "mf6.initialize()\n",
    "\n",
    "# let's advance the model to the first timestep\n",
    "dt = mf6.get_time_step()\n",
    "mf6.prepare_time_step(dt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "424925fc",
   "metadata": {},
   "source": [
    "## The `ApiSimulation` object \n",
    "\n",
    "The `ApiSimulation` object is the top level container for the modflowapi interface classes. This container holds methods and other objects that allow the user to access boundary condition pointer data without assembling the specific memory addresses of the modflow data. \n",
    "\n",
    "Let's take a look at the `ApiSimulation` object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0e83b86",
   "metadata": {},
   "outputs": [],
   "source": [
    "sim = ApiSimulation.load(mf6)\n",
    "sim"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65030e7d",
   "metadata": {},
   "source": [
    "The simulation object allows the user to access models by name and has a number of handy properties and contains simulation level packages such as `sln`, `tdis`, `ats`, and `exchanges`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6930030",
   "metadata": {},
   "outputs": [],
   "source": [
    "mnames = sim.model_names\n",
    "mnames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "327a1c6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "kstp, kper = sim.kstp, sim.kper\n",
    "kstp, kper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93c4c417",
   "metadata": {},
   "outputs": [],
   "source": [
    "nstp = sim.nstp\n",
    "nstp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77c77460",
   "metadata": {},
   "outputs": [],
   "source": [
    "ims = sim.sln\n",
    "ims"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3805e031",
   "metadata": {},
   "outputs": [],
   "source": [
    "ims.dvclose"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "236d4a25",
   "metadata": {},
   "source": [
    "## The `ApiModel` object\n",
    "\n",
    "`ApiModel` objects are accessed from the `ApiSimulation` object and are a container for packages. These objects allow the user to view which packages are available and access those packages. \n",
    "\n",
    "The following cells show the main attributes and functions available on the `ApiModel` object"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e65a3f3",
   "metadata": {},
   "source": [
    "Model objects are accessible through the `get_model` function and as attributes on the sim object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "232c0660",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = sim.get_model('test_model')\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb9328af",
   "metadata": {},
   "outputs": [],
   "source": [
    "# approach 2\n",
    "model = sim.test_model\n",
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ddc33e1",
   "metadata": {},
   "source": [
    "There are also a number of other functions available including the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43aaed16",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07b0b3e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "088578c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.solution_id"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4eafd365",
   "metadata": {},
   "source": [
    "A list of all package names that are accesible is also available"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eee1f0a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.package_names"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "073bb268",
   "metadata": {},
   "source": [
    "## The `ApiPackage` object(s)\n",
    "\n",
    "Each package is contained in `ApiPackage` container. There are three types depending on the input data. We'll access and take a look at each of the types of `ApiPackage` containers."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79ac8066",
   "metadata": {},
   "source": [
    "Packages can be accessed from the `Model` object using `get_package()` or by attribute"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4914c509",
   "metadata": {},
   "outputs": [],
   "source": [
    "# example 1: get a package using get_package\n",
    "rch = model.get_package(\"rcha_0\")\n",
    "rch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbc0aadf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# example 2: get a package by package name attribute\n",
    "wel = model.wel_0\n",
    "wel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a705679",
   "metadata": {},
   "outputs": [],
   "source": [
    "# example 3: get all packages based on a package type\n",
    "rch_pkgs = model.rch\n",
    "rch_pkgs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d59405f",
   "metadata": {},
   "source": [
    "### `ListPackage` objects\n",
    "\n",
    "`ListPackage` objects are the primary object type of stress period data. The exception to this rule is the advanced packages which will be discussed later. \n",
    "\n",
    "`ListPackage` objects allow users to access stress period data as a numpy recarray or as a pandas dataframe."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbc7ed8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "recarray = rch.stress_period_data.values\n",
    "recarray[0:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6eefe129",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = rch.stress_period_data.dataframe\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62faee8b",
   "metadata": {},
   "source": [
    "### Updating values for `ListPackage` based data\n",
    "\n",
    "There are multiple ways to update values for `ListPackage` based data. The `.values` and `.dataframe` attributes can be used, or the object can be directly indexed if the user knows the underlying data. Here are some examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fecf116",
   "metadata": {},
   "outputs": [],
   "source": [
    "recarray['recharge'][0] *= 100\n",
    "rch.stress_period_data.values = recarray\n",
    "\n",
    "# show that values have been updated\n",
    "recarray = rch.stress_period_data.values\n",
    "recarray[0:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf87d47b",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = rch.stress_period_data.dataframe\n",
    "df.loc[1, \"recharge\"] = 10000\n",
    "rch.stress_period_data.dataframe = df\n",
    "\n",
    "# show that values have been updated\n",
    "df = rch.stress_period_data.dataframe\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd85083e",
   "metadata": {},
   "source": [
    "#### Interfacing directly with the `.stress_period_data` attribute\n",
    "\n",
    "The `.stress_period_data` attribute returns a container class that interacts with the internal modflow pointers. The data can be adjusted by interacting with `.stress_period_data` in the same fashion as changing data in a numpy recarray. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1d21b67",
   "metadata": {},
   "outputs": [],
   "source": [
    "rch.stress_period_data[\"recharge\"] *= 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a127471",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = rch.stress_period_data.dataframe\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85f6922c",
   "metadata": {},
   "source": [
    "#### Adding or removing a boundary condition\n",
    "In list packages the user can add and remove specific boundary conditions. Note: if a user adds a boundary condition, such as another well during a stress period, the total number of wells cannot be greater than the wel package's `maxbound` variable. Here's an example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4921081e",
   "metadata": {},
   "outputs": [],
   "source": [
    "wel = model.wel\n",
    "maxbound = wel.maxbound\n",
    "nbound = wel.nbound\n",
    "print(f\"{nbound=}\", f\"{maxbound=}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9639edaf",
   "metadata": {},
   "source": [
    "For the current stress period there are two active wells `nbound=2`, but there can be up to ten `maxbound=10`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9d7ba04",
   "metadata": {},
   "outputs": [],
   "source": [
    "recarray = wel.stress_period_data.values\n",
    "recarray"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56021e0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "recarray.resize((nbound + 1,), refcheck=False)\n",
    "recarray[-1] = ((0, 1, 5), -20., 0, 1)\n",
    "recarray"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1de51293",
   "metadata": {},
   "outputs": [],
   "source": [
    "wel.stress_period_data.values = recarray\n",
    "nbound = wel.nbound\n",
    "f\"{nbound=}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c21fb627",
   "metadata": {},
   "outputs": [],
   "source": [
    "wel.stress_period_data.dataframe"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42b72baa",
   "metadata": {},
   "source": [
    "### `ArrayPackage` objects\n",
    "\n",
    "The `ArrayPackage` class is used as a container for packages such as `DIS`, `NPF`, and `IC` that do not contain any sort of stress period data. These packages are used primarily to define model connectivity, initial conditions, and hydraulic parameters of the basin. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9939cae",
   "metadata": {},
   "outputs": [],
   "source": [
    "npf = model.npf\n",
    "npf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f55215ba",
   "metadata": {},
   "source": [
    "For an `ArrayPackage` type object, variable names can be viewed by calling the `.variable_names` property"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df4e7242",
   "metadata": {},
   "outputs": [],
   "source": [
    "npf.variable_names"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93c2006e",
   "metadata": {},
   "source": [
    "### Updating values for `ArrayPackage` objects\n",
    "\n",
    "Two methods are available for accessing and updating data in `ArrayPackage` objects. `get_array()` and `set_array()` methods can be used to get and set data. Arrays can also be accessed as attributes on the object."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f408dff",
   "metadata": {},
   "source": [
    "Using `get_array()` and `set_array()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d202974",
   "metadata": {},
   "outputs": [],
   "source": [
    "hk = npf.get_array(\"k11\")\n",
    "hk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17ce597b",
   "metadata": {},
   "outputs": [],
   "source": [
    "hk[0, 0:5, 0:5] = 50\n",
    "npf.set_array(\"k11\", hk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ed869a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# confirm that the data has been updated\n",
    "hk = npf.get_array(\"k11\")\n",
    "hk"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7bdbdb1e",
   "metadata": {},
   "source": [
    "Getting and setting data by attribute"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b589b5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# needs an update for inplace operations....\n",
    "npf.k33[0, 0:5, 0:5] = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38e79b05",
   "metadata": {},
   "outputs": [],
   "source": [
    "# confirm that the data has been updated\n",
    "npf.k33.values"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f245911d",
   "metadata": {},
   "source": [
    "## Accessing \"advanced variables\"\n",
    "\n",
    "Advanced variables in this context are variables that would not normally need to be accessed by the user, and in many cases changes to these variables would cause the Modflow simulation to do unexpected things. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff9a706b",
   "metadata": {},
   "source": [
    "For each package object a list of avanced variables can be returned by calling the `advanced_vars` attribute"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c943166",
   "metadata": {},
   "outputs": [],
   "source": [
    "wel = model.wel_0\n",
    "wel.advanced_vars"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcd0e792",
   "metadata": {},
   "source": [
    "The user can access and change these values, _at their own risk_, using the `.get_advanced_var()` and `.set_advanced_var()` methods. Data is returned to the user in the internal modflowapi structure. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a8ee821",
   "metadata": {},
   "outputs": [],
   "source": [
    "wel.get_advanced_var(\"ibcnum\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b034954d",
   "metadata": {},
   "source": [
    "### Advanced Packages\n",
    "\n",
    "Certain packages only support accessing data through the `.get_advanced_var()` and `.set_advanced_var()` methods. These packages, are sometimes refered to as \"advanced packages\" and include: BUY, CSUB, GNC, HFB, MAW, MVR, SFR, and UZF. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e2d960a",
   "metadata": {},
   "source": [
    "-------"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "853b6390",
   "metadata": {},
   "source": [
    "Let's close the existing modflowapi shared library object and look at an example of how this is all used in practice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6896d260",
   "metadata": {},
   "outputs": [],
   "source": [
    "mf6.finalize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27d3baa1",
   "metadata": {},
   "source": [
    "# Putting it all together and running a modflowapi simulation\n",
    "\n",
    "To run a simulation using the built in modflowapi runner the user needs to create a function that will receive callbacks at different steps in the simulation run. For the remainder of this notebook, we'll show how to create a callback function and use it with the `modflowapi.run_simulation()` method."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4404b2c2",
   "metadata": {},
   "source": [
    "## Create a callback function for adjusting model data\n",
    "\n",
    "The callback function allows users to wrap function that updates the modflow model at different steps. The `modflowapi.Callbacks` object allows users to find the particular solution step that they are currently in. `modflowapi.Callbacks` includes:\n",
    "\n",
    "   - `Callbacks.initalize`: the initialize callback sends loaded simulation data back to the user to make adjustments before the model begins solving. This callback only occurs once at the beginning of the MODFLOW6 simulation\n",
    "   - `Callbacks.stress_period_start`: the stress_period_start callback sends simulation data for each solution group to the user to make adjustments to stress packages at the beginning of each stress period.\n",
    "   - `Callbacks.stress_period_end`: the stress_period_end callback sends simulation data for each solution group to the user at the end of each stress period. This can be useful for writing custom output and coupling models\n",
    "   - `Callbacks.timestep_start`: the timestep_start callback sends simulation data for each solution group to the user to make adjustments to stress packages at the beginning of each timestep.\n",
    "   - `Callbacks.timestep_end`: the timestep_end callback sends simulation data for each solution group to the user at the end of each timestep. This can be useful for writing custom output and coupling models\n",
    "   - `Callbacks.iteration_start`: the iteration_start callback sends simulation data for each solution group to the user to make adjustments to stress packages at the beginning of each outer solution iteration.\n",
    "   - `Callbacks.iteration_end`: the iteration_end callback sends simulation data for each solution group to the user to make adjustments to stress packages and check values of stress packages at the end of each outer solution iteration.\n",
    "   - `Callbacks.finalize`: the finalize callback is useful for finalizing models coupled with the modflowapi.\n",
    "   \n",
    "The user can use any or all of these callbacks within their callback function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2783828",
   "metadata": {},
   "outputs": [],
   "source": [
    "from modflowapi import Callbacks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "395327da",
   "metadata": {},
   "outputs": [],
   "source": [
    "def callback_function(sim, callback_step):\n",
    "    \"\"\"\n",
    "    A demonstration function that dynamically adjusts recharge\n",
    "    and pumping in a modflow-6 model through the MODFLOW-API\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    sim : modflowapi.Simulation\n",
    "        A simulation object for the solution group that is \n",
    "        currently being solved\n",
    "    callback_step : enumeration\n",
    "        modflowapi.Callbacks enumeration object that indicates\n",
    "        the part of the solution modflow is currently in.\n",
    "    \"\"\"\n",
    "    ml = sim.test_model\n",
    "    if callback_step == Callbacks.initialize:\n",
    "        print(sim.models)\n",
    "    \n",
    "    if callback_step == Callbacks.stress_period_start:\n",
    "        # adjust recharge for stress periods 1 through 7\n",
    "        if sim.kper <= 6:\n",
    "            rcha = ml.rcha_0\n",
    "            spd = rcha.stress_period_data\n",
    "            print(f\"updating recharge: stress_period={ml.kper}\")\n",
    "            spd[\"recharge\"] += 0.40 * sim.kper\n",
    "        \n",
    "    \n",
    "    if callback_step == Callbacks.timestep_start:\n",
    "        print(f\"updating wel flux: stress_period={ml.kper}, timestep={ml.kstp}\")\n",
    "        ml.wel.stress_period_data[\"q\"] -= ml.kstp * 1.5\n",
    "    \n",
    "    if callback_step == Callbacks.iteration_start:\n",
    "        # we can implement complex solutions to boundary conditions here!\n",
    "        pass\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b751eb41",
   "metadata": {},
   "source": [
    "The callback function is then passed to `modflowapi.run_simulation`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0878e5b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "modflowapi.run_simulation(dll, sim_ws, callback_function, verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc1e31af",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}