examples/pstfrom_mf6_ppu.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setting up a PEST interface from MODFLOW6 using the `PstFrom` class with `PyPestUtils` for advanced pilot point parameterization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import pyemu\n",
    "import flopy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append(os.path.join(\"..\",\"..\",\"pypestutils\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pypestutils as ppu"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An existing MODFLOW6 model is in the directory `freyberg_mf6`.  Lets check it out:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "org_model_ws = os.path.join('freyberg_mf6')\n",
    "os.listdir(org_model_ws)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can see that all the input array and list data for this model have been written \"externally\" - this is key to using the `PstFrom` class. \n",
    "\n",
    "Let's quickly viz the model top just to remind us of what we are dealing with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "id_arr = np.loadtxt(os.path.join(org_model_ws,\"freyberg6.dis_idomain_layer3.txt\"))\n",
    "top_arr = np.loadtxt(os.path.join(org_model_ws,\"freyberg6.dis_top.txt\"))\n",
    "top_arr[id_arr==0] = np.nan\n",
    "plt.imshow(top_arr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's copy those files to a temporary location just to make sure we don't goof up those original files:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tmp_model_ws = \"temp_pst_from_ppu\"\n",
    "if os.path.exists(tmp_model_ws):\n",
    "    shutil.rmtree(tmp_model_ws)\n",
    "shutil.copytree(org_model_ws,tmp_model_ws)\n",
    "os.listdir(tmp_model_ws)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we need just a tiny bit of info about the spatial discretization of the model - this is needed to work out separation distances between parameters for build a geostatistical prior covariance matrix later.\n",
    "\n",
    "Here we will load the flopy sim and model instance just to help us define some quantities later - flopy is not required to use the `PstFrom` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim = flopy.mf6.MFSimulation.load(sim_ws=tmp_model_ws)\n",
    "m = sim.get_model(\"freyberg6\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we use the simple `SpatialReference` pyemu implements to help us spatially locate parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sr = pyemu.helpers.SpatialReference.from_namfile(\n",
    "        os.path.join(tmp_model_ws, \"freyberg6.nam\"),\n",
    "        delr=m.dis.delr.array, delc=m.dis.delc.array)\n",
    "sr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can instantiate a `PstFrom` class instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "template_ws = \"freyberg6_template\"\n",
    "pf = pyemu.utils.PstFrom(original_d=tmp_model_ws, new_d=template_ws,\n",
    "                 remove_existing=True,\n",
    "                 longnames=True, spatial_reference=sr,\n",
    "                 zero_based=False,start_datetime=\"1-1-2018\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Observations\n",
    "\n",
    "So now that we have a `PstFrom` instance, but its just an empty container at this point, so we need to add some PEST interface \"observations\" and \"parameters\".  Let's start with observations using MODFLOW6 head.  These are stored in `heads.csv`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv(os.path.join(tmp_model_ws,\"heads.csv\"),index_col=0)\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The main entry point for adding observations is (surprise) `PstFrom.add_observations()`.  This method works on the list-type observation output file.  We need to tell it what column is the index column (can be string if there is a header or int if no header) and then what columns contain quantities we want to monitor (e.g. \"observe\") in the control file - in this case we want to monitor all columns except the index column:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hds_df = pf.add_observations(\"heads.csv\",insfile=\"heads.csv.ins\",index_cols=\"time\",\n",
    "                    use_cols=list(df.columns.values),prefix=\"hds\",)\n",
    "hds_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that it returned a dataframe with lots of useful info: the observation names that were formed (`obsnme`), the values that were read from `heads.csv` (`obsval`) and also some generic weights and group names.  At this point, no control file has been created, we have simply prepared to add this observations to the control file later.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "[f for f in os.listdir(template_ws) if f.endswith(\".ins\")]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Nice!  We also have a PEST-style instruction file for those obs.\n",
    "\n",
    "Now lets do the same for SFR observations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv(os.path.join(tmp_model_ws, \"sfr.csv\"), index_col=0)\n",
    "sfr_df = pf.add_observations(\"sfr.csv\", insfile=\"sfr.csv.ins\", index_cols=\"time\", use_cols=list(df.columns.values))\n",
    "sfr_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sweet as!  Now that we have some observations, let's add parameters!\n",
    "\n",
    "## Pilot points and `PyPestUtils`\n",
    "\n",
    "This notebook is mostly meant to demonstrate some advanced pilot point parameterization that is possible with `PyPestUtils`, so we will only focus on HK and VK pilot point parameters.  This is just to keep the example short.  In practice, please please please parameterize boundary conditions too!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v = pyemu.geostats.ExpVario(contribution=1.0,a=5000,bearing=0,anisotropy=1)\n",
    "pp_gs = pyemu.geostats.GeoStruct(variograms=v, transform='log')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pp_gs.plot()\n",
    "print(\"spatial variogram\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's get the idomain array to use as a zone array - this keeps us from setting up parameters in inactive model cells:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ib = m.dis.idomain[0].array"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Find HK files for the upper and lower model layers (assuming model layer 2 is a semi-confining unit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hk_arr_files = [f for f in os.listdir(tmp_model_ws) if \"npf_k_\" in f and f.endswith(\".txt\") and \"layer2\" not in f]\n",
    "hk_arr_files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr_file = \"freyberg6.npf_k_layer1.txt\"\n",
    "tag = arr_file.split('.')[1].replace(\"_\",\"-\")\n",
    "pf.add_parameters(filenames=arr_file,par_type=\"pilotpoints\",\n",
    "                   par_name_base=tag,pargp=tag,zone_array=ib,\n",
    "                   upper_bound=10.,lower_bound=0.1,ult_ubound=100,ult_lbound=0.01,\n",
    "                   pp_options={\"pp_space\":3},geostruct=pp_gs)\n",
    "#let's also add the resulting hk array that modflow sees as observations\n",
    "# so we can make easy plots later...\n",
    "pf.add_observations(arr_file,prefix=tag,\n",
    "                    obsgp=tag,zone_array=ib)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you are familiar with how `PstFrom` has worked historically, we handed off the process to solve for the factor file (which requires solving the kriging equations for each active node) to a pure python (well, with pandas and numpy).  This was ok for toy models, but hella slow for big ugly models.  If you look at the log entries above, you should see that the instead, `PstFrom` successfully handed off the solve to `PyPestUtils`, which is exponentially faster for big models.  sweet ez! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tpl_files = [f for f in os.listdir(template_ws) if f.endswith(\".tpl\")]\n",
    "tpl_files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(template_ws,tpl_files[0]),'r') as f:\n",
    "    for _ in range(2):\n",
    "        print(f.readline().strip())\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "So those might look like pretty redic parameter names, but they contain heaps of metadata to help you post process things later..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So those are you standard pilot points for HK in layer 1 - same as it ever was..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Geostatistical hyper-parameters\n",
    "\n",
    "For the HK layer 1 pilot points, we used a standard geostatistical structure - the ever popular exponential variogram.  But what if the properties that define that variogram were themselves uncertain?  Like what is the anisotropy ellipse varied in space across the model domain?  What does this imply?  Well, technically speaking, those variogram properties can be conceptualized as \"hyper parameters\" in that they influence the underlying parameters (in this case, the pilot points) in hierarchical sense.  That is, the bearing of the anisotropy of the variogram changes, then the resulting interpolation from the pilot points to grid changes.  But where it gets really deep is that we need to define correlation structures for these spatially varying hyper pars, so they themselves have plausible spatial patterns...Seen that movie inception?!\n",
    "\n",
    "In `PyPestUtils`, we can supply the pilot-point-to-grid interpolation process with arrays of hyper-parameter values, one array for each variogram property.  The result of this hyper parameter mess is referred to as a non-stationary spatial parameterization.  buckle up...\n",
    "\n",
    "First let's define some additional geostatistical structures:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "value_v = pyemu.geostats.ExpVario(contribution=1, a=5000, anisotropy=2, bearing=0.0)\n",
    "value_gs = pyemu.geostats.GeoStruct(variograms=value_v)\n",
    "bearing_v = pyemu.geostats.ExpVario(contribution=1,a=5000,anisotropy=2,bearing=90.0)\n",
    "bearing_gs = pyemu.geostats.GeoStruct(variograms=bearing_v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr_file = \"freyberg6.npf_k_layer3.txt\"\n",
    "tag = arr_file.split('.')[1].replace(\"_\",\"-\")\n",
    "pf.add_parameters(filenames=arr_file,par_type=\"pilotpoints\",\n",
    "                   par_name_base=tag,pargp=tag,zone_array=ib,\n",
    "                   upper_bound=10.,lower_bound=0.1,ult_ubound=100,ult_lbound=0.01,\n",
    "                 pp_options={\"pp_space\":3,\"prep_hyperpars\":True},geostruct=value_gs,\n",
    "                 apply_order=2)\n",
    "pf.add_observations(arr_file,prefix=tag,\n",
    "                    obsgp=tag,zone_array=ib)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hyperpar_files = [f for f in os.listdir(pf.new_d) if tag in f]\n",
    "hyperpar_files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "when we supplied the \"prep_hyperpars\" as `True` above, that triggered `PstFrom` to do something different.  Instead of solving for the pilot point kriging factors as before, now, we have array-based files for the geostatistical hyper parameters, as well as some additional quantities we need to \"apply\" these hyper parameter at runtime.  This is a key difference:  When the pilot point variogram is changing for each model run, we need to re-solve for the kriging factors for each model run...\n",
    "\n",
    "We snuck in something else too - see that `apply_order` argument?  That is how we can control what order of files being processed by the run-time multiplier parameter function.  Since we are going to parameterize the hyper parameters and there is an implicit order between these hyper parameters and the underlying pilot points, we need to make sure the hyper parameters are processed first.  \n",
    "\n",
    "Lets setup some hyper parameters for estimation.  We will use a constant for the anisotropy ratio, but use pilot points for the bearing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "afile = 'npf-k-layer3.aniso.dat'\n",
    "tag = afile.split('.')[0].replace(\"_\",\"-\")+\"-aniso\"\n",
    "pf.add_parameters(afile,par_type=\"constant\",par_name_base=tag,\n",
    "                  pargp=tag,lower_bound=-1.0,upper_bound=1.0,\n",
    "                  apply_order=1,\n",
    "                  par_style=\"a\",transform=\"none\",initial_value=0.0)\n",
    "pf.add_observations(afile, prefix=tag, obsgp=tag)\n",
    "bfile = 'npf-k-layer3.bearing.dat'\n",
    "tag = bfile.split('.')[0].replace(\"_\",\"-\")+\"-bearing\"\n",
    "pf.add_parameters(bfile, par_type=\"pilotpoints\", par_name_base=tag,\n",
    "                  pargp=tag, pp_space=6,lower_bound=-45,upper_bound=45,\n",
    "                  par_style=\"a\",transform=\"none\",\n",
    "                  pp_options={\"try_use_ppu\":True},\n",
    "                  apply_order=1,geostruct=bearing_gs)\n",
    "pf.add_observations(bfile, prefix=tag, obsgp=tag)                "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that the `apply_order` for these hyper pars is 1 so that any processing for these quantities happens before the actual underlying pilot points are processed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## \"These go to 11\" - amp'ing things up with categorization\n",
    "\n",
    "Sometimes, the world we want to simulate might be better represented as categorical instead continuous.  That is, rather than smoothly varying property fields, we want fields that are either a high value or a low value (please dont ask for more than 2 categories!).  In this case, depending on how you plan to assimilate data (that is, what inversion algorithm you are planning to you), we can accommodate this preference for categorical fields.  \n",
    "\n",
    "This is pretty advanced and also dense.  There is another example notebook the describes the categorization process in detail.  Here we will just blast thru it....\n",
    "\n",
    "lets setup non-stationary categorical parameterization for the VK of layer 2 (the semi confining unit).  We can conceptualize this as a semi-confining unit that has \"windows\" in it that connects the two aquifers.  Where there is not a window, the Vk will be very low, where there is a window, the VK will be much higher. Let's also assume the windows in the confining unit where created when a stream eroded thru it, so the shape of these windows will be higher-order (not derived from a standard geostatistical 2-point process), but rather from connected features.\n",
    "\n",
    "In what follows, we setup this complex parameterization.  We also add lots of aux observations to lets plot and viz the steps in this parameterization process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr_file = \"freyberg6.npf_k33_layer2.txt\"\n",
    "print(arr_file)\n",
    "k = int(arr_file.split(\".\")[1][-1]) - 1\n",
    "pth_arr_file = os.path.join(pf.new_d,arr_file)\n",
    "arr = np.loadtxt(pth_arr_file)\n",
    "cat_dict = {1:[0.4,arr.mean()],2:[0.6,arr.mean()]}\n",
    "\n",
    "#this is where we initialize the categorization process - it will operate on the \n",
    "# layer 2 VK array just before MODFLOW runs\n",
    "thresharr,threshcsv = pyemu.helpers.setup_threshold_pars(pth_arr_file,cat_dict=cat_dict,\n",
    "                                                         testing_workspace=pf.new_d,inact_arr=ib)\n",
    "\n",
    "# the corresponding apply function\n",
    "pf.pre_py_cmds.append(\"pyemu.helpers.apply_threshold_pars('{0}')\".format(os.path.split(threshcsv)[1]))\n",
    "prefix = arr_file.split('.')[1].replace(\"_\",\"-\")\n",
    "\n",
    "pth_arr_file = os.path.join(pf.new_d,arr_file)\n",
    "arr = np.loadtxt(pth_arr_file)\n",
    "\n",
    "tag = arr_file.split('.')[1].replace(\"_\",\"-\") + \"_pp\"\n",
    "prefix = arr_file.split('.')[1].replace(\"_\",\"-\")\n",
    "#setup pilot points with hyper pars for the thresholding array (the array that will drive the \n",
    "# categorization process).  Notice the apply_order arg being used \n",
    "pf.add_parameters(filenames=os.path.split(thresharr)[1],par_type=\"pilotpoints\",transform=\"none\",\n",
    "                  par_name_base=tag+\"-threshpp_k:{0}\".format(k),\n",
    "                  pargp=tag + \"-threshpp_k:{0}\".format(k),\n",
    "                  lower_bound=0.0,upper_bound=2.0,par_style=\"m\",\n",
    "                  pp_options={\"try_use_ppu\":False,\"prep_hyperpars\":True,\"pp_space\":5},\n",
    "                  apply_order=2,geostruct=value_gs\n",
    "                  )\n",
    "\n",
    "tag = arr_file.split('.')[1].replace(\"_\",\"-\")\n",
    "# a constant parameter for the anisotropy of the thresholding array\n",
    "# Notice the apply_order arg being used\n",
    "tfiles = [f for f in os.listdir(pf.new_d) if tag in f]\n",
    "afile = [f for f in tfiles if \"aniso\" in f][0]\n",
    "pf.add_parameters(afile,par_type=\"constant\",par_name_base=tag+\"-aniso\",\n",
    "                  pargp=tag+\"-aniso\",lower_bound=-1.0,upper_bound=1.0,\n",
    "                  apply_order=1,\n",
    "                  par_style=\"a\",transform=\"none\",initial_value=0.0)\n",
    "# obs for the anisotropy field\n",
    "pf.add_observations(afile, prefix=tag+\"-aniso\", obsgp=tag+\"-aniso\")\n",
    "\n",
    "# pilot points for the bearing array of the geostructure of the thresholding array\n",
    "# Notice the apply_order arg being used\n",
    "bfile = [f for f in tfiles if \"bearing\" in f][0]\n",
    "pf.add_parameters(bfile, par_type=\"pilotpoints\", par_name_base=tag + \"-bearing\",\n",
    "                  pargp=tag + \"-bearing\", pp_space=6,lower_bound=-45,upper_bound=45,\n",
    "                  par_style=\"a\",transform=\"none\",\n",
    "                  pp_options={\"try_use_ppu\":True},\n",
    "                  apply_order=1,geostruct=bearing_gs)\n",
    "# obs for the bearing array\n",
    "pf.add_observations(bfile, prefix=tag + \"-bearing\", obsgp=tag + \"-bearing\")                \n",
    "\n",
    "# list style parameters for the quantities used in the categorization process\n",
    "# We will manipulate these initial values and bounds later\n",
    "pf.add_parameters(filenames=os.path.split(threshcsv)[1], par_type=\"grid\",index_cols=[\"threshcat\"],\n",
    "                  use_cols=[\"threshproportion\",\"threshfill\"],\n",
    "                  par_name_base=[prefix+\"threshproportion_k:{0}\".format(k),prefix+\"threshfill_k:{0}\".format(k)],\n",
    "                  pargp=[prefix+\"threshproportion_k:{0}\".format(k),prefix+\"threshfill_k:{0}\".format(k)],\n",
    "                  lower_bound=[0.1,0.1],upper_bound=[10.0,10.0],transform=\"none\",par_style='d')\n",
    "\n",
    "# obs of the resulting Vk array that MODFLOW uses\n",
    "pf.add_observations(arr_file,prefix=tag,\n",
    "                    obsgp=tag,zone_array=ib)\n",
    "\n",
    "# observations of the categorized array\n",
    "pf.add_observations(arr_file+\".threshcat.dat\", prefix=\"tcatarr-\" + prefix+\"_k:{0}\".format(k),\n",
    "                    obsgp=\"tcatarr-\" + prefix+\"_k:{0}\".format(k),zone_array=ib)\n",
    "\n",
    "# observations of the thresholding array\n",
    "pf.add_observations(arr_file + \".thresharr.dat\",\n",
    "                    prefix=tag+'-thresharr',\n",
    "                    obsgp=tag+'-thresharr', zone_array=ib)\n",
    "\n",
    "# observations of the results of the thresholding process\n",
    "df = pd.read_csv(threshcsv.replace(\".csv\",\"_results.csv\"),index_col=0)\n",
    "pf.add_observations(os.path.split(threshcsv)[1].replace(\".csv\",\"_results.csv\"),index_cols=\"threshcat\",use_cols=df.columns.tolist(),prefix=prefix+\"-results_k:{0}\".format(k),\n",
    "                    obsgp=prefix+\"-results_k:{0}\".format(k),ofile_sep=\",\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### build the control file, pest interface files, and forward run script\n",
    "At this point, we have some parameters and some observations, so we can create a control file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pf.mod_sys_cmds.append(\"mf6\")\n",
    "pf.pre_py_cmds.insert(0,\"import sys\")\n",
    "pf.pre_py_cmds.insert(1,\"sys.path.append(os.path.join('..','..','..','pypestutils'))\")\n",
    "pst = pf.build_pst()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = [print(line.rstrip()) for line in open(os.path.join(template_ws,\"forward_run.py\"))]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting initial parameter bounds and values\n",
    "\n",
    "Here is some gory detail regarding defining the hyper parameters for both layer 3 HK and layer 2 VK..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#set the initial and bounds for the fill values\n",
    "par = pst.parameter_data\n",
    "\n",
    "apar = par.loc[par.pname.str.contains(\"aniso\"),:]\n",
    "bpar = par.loc[par.pname.str.contains(\"bearing\"), :]\n",
    "par.loc[apar.parnme.str.contains(\"layer3\").index,\"parval1\"] = 3\n",
    "par.loc[apar.parnme.str.contains(\"layer3\").index,\"parlbnd\"] = 1\n",
    "par.loc[apar.parnme.str.contains(\"layer3\").index,\"parubnd\"] = 5\n",
    "\n",
    "par.loc[apar.parnme.str.contains(\"layer2\").index,\"parval1\"] = 2\n",
    "par.loc[apar.parnme.str.contains(\"layer2\").index,\"parlbnd\"] = 0\n",
    "par.loc[apar.parnme.str.contains(\"layer2\").index,\"parubnd\"] = 4\n",
    "\n",
    "par.loc[bpar.parnme.str.contains(\"layer3\").index,\"parval1\"] = 0\n",
    "par.loc[bpar.parnme.str.contains(\"layer3\").index,\"parlbnd\"] = -90\n",
    "par.loc[bpar.parnme.str.contains(\"layer3\").index,\"parubnd\"] = 90\n",
    "\n",
    "par.loc[bpar.parnme.str.contains(\"layer2\").index,\"parval1\"] = 0\n",
    "par.loc[bpar.parnme.str.contains(\"layer2\").index,\"parlbnd\"] = -90\n",
    "par.loc[bpar.parnme.str.contains(\"layer2\").index,\"parubnd\"] = 90\n",
    "\n",
    "cat1par = par.loc[par.apply(lambda x: x.threshcat==\"0\" and x.usecol==\"threshfill\",axis=1),\"parnme\"]\n",
    "cat2par = par.loc[par.apply(lambda x: x.threshcat == \"1\" and x.usecol == \"threshfill\", axis=1), \"parnme\"]\n",
    "assert cat1par.shape[0] == 1\n",
    "assert cat2par.shape[0] == 1\n",
    "\n",
    "cat1parvk = [p for p in cat1par if \"k:1\" in p]\n",
    "cat2parvk = [p for p in cat2par if \"k:1\" in p]\n",
    "for lst in [cat2parvk,cat1parvk]:\n",
    "    assert len(lst) > 0\n",
    "\n",
    "#these are the values that will fill the two categories of VK - \n",
    "# one is low (clay) and one is high (sand - the windows)\n",
    "par.loc[cat1parvk, \"parval1\"] = 0.0001\n",
    "par.loc[cat1parvk, \"parubnd\"] = 0.01\n",
    "par.loc[cat1parvk, \"parlbnd\"] = 0.000001\n",
    "par.loc[cat1parvk, \"partrans\"] = \"log\"\n",
    "par.loc[cat2parvk, \"parval1\"] = 0.1\n",
    "par.loc[cat2parvk, \"parubnd\"] = 1\n",
    "par.loc[cat2parvk, \"parlbnd\"] = 0.01\n",
    "par.loc[cat2parvk, \"partrans\"] = \"log\"\n",
    "\n",
    "\n",
    "cat1par = par.loc[par.apply(lambda x: x.threshcat == \"0\" and x.usecol == \"threshproportion\", axis=1), \"parnme\"]\n",
    "cat2par = par.loc[par.apply(lambda x: x.threshcat == \"1\" and x.usecol == \"threshproportion\", axis=1), \"parnme\"]\n",
    "\n",
    "assert cat1par.shape[0] == 1\n",
    "assert cat2par.shape[0] == 1\n",
    "\n",
    "#these are the proportions of clay and sand in the resulting categorical array\n",
    "#really under the hood, only the first one is used, so we can fix the other.\n",
    "par.loc[cat1par, \"parval1\"] = 0.95\n",
    "par.loc[cat1par, \"parubnd\"] = 1.0\n",
    "par.loc[cat1par, \"parlbnd\"] = 0.9\n",
    "par.loc[cat1par,\"partrans\"] = \"none\"\n",
    "\n",
    "# since the apply method only looks that first proportion, we can just fix this one\n",
    "par.loc[cat2par, \"parval1\"] = 1\n",
    "par.loc[cat2par, \"parubnd\"] = 1\n",
    "par.loc[cat2par, \"parlbnd\"] = 1\n",
    "par.loc[cat2par,\"partrans\"] = \"fixed\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generating a prior parameter ensemble, then run and viz a real"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(122341)\n",
    "pe = pf.draw(num_reals=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pe.to_csv(os.path.join(template_ws,\"prior.csv\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "real = 0\n",
    "pst_name = \"real_{0}.pst\".format(real)\n",
    "pst.parameter_data.loc[pst.adj_par_names,\"parval1\"] = pe.loc[real,pst.adj_par_names].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pst.control_data.noptmax = 0\n",
    "pst.write(os.path.join(pf.new_d,pst_name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyemu.os_utils.run(\"pestpp-ies {0}\".format(pst_name),cwd=pf.new_d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pst.set_res(os.path.join(pf.new_d,pst_name.replace(\".pst\",\".base.rei\")))\n",
    "res = pst.res\n",
    "obs = pst.observation_data\n",
    "grps = [o for o in obs.obgnme.unique() if o.startswith(\"npf\") and \"result\" not in o and \"aniso\" not in o]\n",
    "grps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gobs = obs.loc[obs.obgnme.isin(grps),:].copy()\n",
    "gobs[\"i\"] = gobs.i.astype(int)\n",
    "gobs[\"j\"] = gobs.j.astype(int)\n",
    "gobs[\"k\"] = gobs.obgnme.apply(lambda x: int(x.split('-')[2].replace(\"layer\",\"\")) - 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "uk = gobs.k.unique()\n",
    "uk.sort()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for k in uk:\n",
    "    kobs = gobs.loc[gobs.k==k,:]\n",
    "    ug = kobs.obgnme.unique()\n",
    "    ug.sort()\n",
    "    fig,axes = plt.subplots(1,4,figsize=(20,6))\n",
    "    axes = np.atleast_1d(axes)\n",
    "    for ax in axes:\n",
    "        ax.set_frame_on(False)\n",
    "        ax.set_yticks([])\n",
    "        ax.set_xticks([])\n",
    "    for g,ax in zip(ug,axes):\n",
    "        gkobs = kobs.loc[kobs.obgnme==g,:]\n",
    "        \n",
    "        arr = np.zeros_like(top_arr)\n",
    "        arr[gkobs.i,gkobs.j] = res.loc[gkobs.obsnme,\"modelled\"].values\n",
    "        ax.set_aspect(\"equal\")\n",
    "        label = \"\"\n",
    "        if \"bearing\" not in g and \"aniso\" not in g:\n",
    "            arr = np.log10(arr)\n",
    "            label = \"$log_{10}$\"\n",
    "        cb = ax.imshow(arr)\n",
    "        plt.colorbar(cb,ax=ax,label=label)\n",
    "        ax.set_title(\"layer: {0} group: {1}\".format(k+1,g),loc=\"left\",fontsize=15)\n",
    "        \n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "    plt.close(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Stunning isn't it?!  There is clearly a lot subjectivity in the form of defining the prior for the hyper parameters required to use these non-stationary geostats, but they do afford more opportunities to express (stochastic) expert knowledge.  To be honest, there was a lot of experimenting with this notebook to get these figures to look this way - playing with variograms and parameter initial values and bounds a lot.  You encouraged to do the same!  scroll back up, change things, and \"restart kernel and run all\" - this will help build some better intution, promise...."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}