3D_gwModel.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3f1b65c6-90b5-42c3-96bc-30a6dee6dec7",
   "metadata": {},
   "source": [
    "# 3D Model Malta"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddb0b6ca-c5e8-4c20-9972-4cc9d66c7643",
   "metadata": {},
   "source": [
    "### 1) Import Python libraries and define input/output directories"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e14cc4ee-acf8-48dd-9648-17765174495d",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "We will use the Flopy library to build the individual SEAWAT models. Apart from that, only standard Python libraries such as Numpy and Pandas are used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "975e4dbe-5a86-43e1-9c61-d4f7c9e6bc84",
   "metadata": {},
   "outputs": [],
   "source": [
    "#   who needs warnings.. we are now on day 3 of Python coding we know what we are doing \n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7dd77342-2dfd-4559-849f-a996f4351241",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\User\\deltaforge\\lib\\site-packages\\xugrid\\ugrid\\snapping.py:30: DeprecationWarning: Shapely 2.0 is installed, but because PyGEOS is also installed, GeoPandas still uses PyGEOS by default. However, starting with version 0.14, the default will switch to Shapely. To force to use Shapely 2.0 now, you can either uninstall PyGEOS or set the environment variable USE_PYGEOS=0. You can do this before starting the Python process, or in your code before importing geopandas:\n",
      "\n",
      "import os\n",
      "os.environ['USE_PYGEOS'] = '0'\n",
      "import geopandas\n",
      "\n",
      "In the next release, GeoPandas will switch to using Shapely by default, even if PyGEOS is installed. If you only have PyGEOS installed to get speed-ups, this switch should be smooth. However, if you are using PyGEOS directly (calling PyGEOS functions on geometries from GeoPandas), this will then stop working and you are encouraged to migrate from PyGEOS to Shapely 2.0 (https://shapely.readthedocs.io/en/latest/migration_pygeos.html).\n",
      "  import geopandas as gpd\n"
     ]
    }
   ],
   "source": [
    "# import libraries\n",
    "import flopy\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import os\n",
    "import matplotlib\n",
    "import xarray as xr\n",
    "import PVGeo\n",
    "import pyvista as pv\n",
    "#import pyvistaqt\n",
    "import imod\n",
    "import flopy.utils.binaryfile as bf\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf549d30-a2d9-4d9b-b5ee-ba13a5a96ada",
   "metadata": {},
   "source": [
    "Define the paths below based on where you stored the input files and the seawat executable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ceced996-f974-4f19-a81f-71c05ff8b595",
   "metadata": {},
   "outputs": [],
   "source": [
    "# input and output directories\n",
    "input_dir = r'C:\\Users\\User\\WORK-PC-LOCAL\\AquaInfra-local\\model_inputs'\n",
    "out_dir = r'C:\\Users\\User\\WORK-PC-LOCAL\\AquaInfra-local\\model_outputs'\n",
    "seawat_exe_dir = r'C:\\Users\\User\\WORK-PC-LOCAL\\AquaInfra-local\\imod_files\\iMODexe'\n",
    "\n",
    "#   define the model name and model output directory\n",
    "model_name = '_Malta_case'\n",
    "model_dir = os.path.join(out_dir, model_name)\n",
    "\n",
    "# create the directories if they do not exist yet\n",
    "os.makedirs(input_dir, exist_ok = True)\n",
    "os.makedirs(out_dir, exist_ok = True)\n",
    "os.makedirs(model_dir, exist_ok = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f60dfac5-0807-4189-8035-f03c47430b27",
   "metadata": {},
   "source": [
    "### 2) Define grid dimensions and create the initial grid"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42aa2149-e38e-4a85-893d-51b46b713211",
   "metadata": {},
   "source": [
    "In the next steps we will create the 3D model grid. First we will define the grid cell dimensions (in meters) and then read in a CSV file that contains geological information and will be used to set up the 3D grid. This CSV file is created in Petrel as an ouput of a 3D geological model, but it could be any file with similar setup (e.g. Netcdf file). MODFLOW and SEAWAT basic model grid is called the IBOUND array - this is where we define active (1) and inactive (0) cells. In this tutorial we set all the cells based on their I J K indexes (column, row, layer) defined in the input CSV file. \n",
    "\n",
    "The input folder contains several CSV files, each for a different geological period - representing different depositional periods and stacking of lithological layers. The cmod_full.csv file represents the final geological environment (present time) and we will use it to set up our 3D IBOUND array. In such way we make sure that we have the full grid extent and can therefore accommodate all the geological layers. If we used the cmod_1.csv file that represents the geological environment during the fist geological period (i.e. the oldest geological layers) and its dimensions (at least in the vertical sense) might be smaller than the final geological environment. This could lead to a smaller IBOUND array where we could not fit the more recent geological layers. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "8b81ce16-65c7-4c05-95c3-30d8e59b52e5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Top elevation of the model grid      250.0 m bsl.\n",
      "Bottom elevation of the model grid   -410.0 m bsl.\n",
      "Total number of model layers         33\n",
      "Total number of model rows           255\n",
      "Total number of model columns        135\n"
     ]
    }
   ],
   "source": [
    "#   define grid dimensions in meters for row, col (dx and dy) and lay thickness (dz)\n",
    "dx = dy = 200\n",
    "dz = 20\n",
    "\n",
    "#   create the IBOUND array from the Petrel output, using the cmod_full, the fields list represents the column names of the CSV file\n",
    "fields = ['I', 'J', 'K', 'X', 'Y', 'Z', 'Lith', 'por','hyd']\n",
    "#   use Pandas to read the CSV file and round the Z values (elevation) and the V_shale column\n",
    "df_in = pd.read_csv(os.path.join(input_dir, 'malta_test'), header = 11, sep = \" \", index_col = False, names = fields)\n",
    "df_in['I'] = df_in['I'].round(0)\n",
    "df_in['J'] = df_in['J'].round(0)\n",
    "df_in['K'] = df_in['K'].round(0)\n",
    "df_in['Z'] = df_in['Z'].round(1)\n",
    "df_in['Lith'] = df_in['Lith'].round(1)\n",
    "df_in['por'] = df_in['por'].round(1)\n",
    "df_in['hyd'] = df_in['hyd'].round(1)\n",
    "\n",
    "#   make a 3D numpy array with the dimensions from the input file\n",
    "lay_arr = np.unique(df_in.K.values)  # unique layer numbers\n",
    "row_arr = np.unique(df_in.I.values)  # unique row numbers\n",
    "col_arr = np.unique(df_in.J.values)  # unique column numbers\n",
    "#   numpy.zeros function create a 3D array where all cells are set to a value = 0 \n",
    "full_ibound_arr = np.zeros((lay_arr.shape[0], row_arr.shape[0], col_arr.shape[0]))\n",
    "\n",
    "#   get unique z values \n",
    "z_vals = np.unique(df_in.Z.values)\n",
    "#   the max value will be the top elevation of the model\n",
    "top_elev = np.nanmax(z_vals)  \n",
    "#   have to do 2 * dz so it keeps the right end value\n",
    "bot_elev = np.arange(top_elev, np.nanmin(z_vals) - 2 * dz, -dz)[1:]  \n",
    "\n",
    "#   print the top, bottom elevations and the grid dimensions\n",
    "print(\"Top elevation of the model grid      \" + str(top_elev) + \" m bsl.\")\n",
    "print(\"Bottom elevation of the model grid   \" + str(bot_elev[-1]) + \" m bsl.\")\n",
    "print(\"Total number of model layers         \" + str(full_ibound_arr.shape[0]))\n",
    "print(\"Total number of model rows           \" + str(full_ibound_arr.shape[1]))\n",
    "print(\"Total number of model columns        \" + str(full_ibound_arr.shape[2]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49f25b78-fc02-4f31-9ceb-2853e1efbad1",
   "metadata": {},
   "source": [
    "### 3) Scenario definition "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3230f0ff-ecc4-478d-83e4-5a717f2186a6",
   "metadata": {},
   "source": [
    "Now that we have our main IBOUND grid dimensions and the array itself, we can think about what exactly we want our groundwater model to simulate. As we have multiple CSV files with geological environments at different times we can simulate two stress periods with different geologies and sea levels simulating a potential OFG deposition. Below we will define names of these stress periods, their length in years and corresponding sea levels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "9f5b0986-dd4e-4af7-a63d-994afc0e4d49",
   "metadata": {},
   "outputs": [],
   "source": [
    "#   define the model stress periods, names etc.\n",
    "sp_names = ['malta_test', 'malta_test','malta_test','malta_test','malta_test','malta_test','malta_test',\n",
    "            'malta_test','malta_test','malta_test','malta_test']  # just set it to match the csv names\n",
    "sp_time_dur = [2000, 2000, 2000, 2000, 2000, 2000, 2000, 2000, 2000, 2000, 2000]  # duriation of each stress period - in years\n",
    "sp_sea_level = [-117.5, -106, -75, -55.5, -50.5, -48.5, -42, -30, -18, -6, 0]  # sea level for each stress period\n",
    "\n",
    "#   we start with the first stress period, Python is 0 based so 0 represents the first index..\n",
    "a = 0\n",
    "\n",
    "#   create the model directory for the stress period model\n",
    "model_name_sp = sp_names[a]\n",
    "sp_dir = os.path.join(model_dir, model_name_sp)\n",
    "os.makedirs(sp_dir, exist_ok=True)\n",
    "sea_level = sp_sea_level[a]  # get the sea level for the stress period"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab2cc041-f7df-47ea-9e07-5f17b5927017",
   "metadata": {},
   "source": [
    "We will use the geology that was already exported in the CSV file - if you want to create a similar model to what is shown here you can use a similar procedure to build a groundwater model on top of your geological model. We also define a function that will reclassify the geological input based on the Vshale value into three lithological classes. You can of course increase the complexity if you wish but three different categories will be fine for our purposes here. Finally, we will assign a hydraulic conductivity value (in m/d) to all the classes. In this case we use a constant value for hydraulic conductivity which is a big simplification, in case you have a better geological model and hydraulic condcutivity estimation you can of course use that in the future. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec437239-cd6c-4e56-b78c-05721a811ffb",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "source": [
    "#   read in the csv file with active geology for the stress period\n",
    "df_sp_in = pd.read_csv(os.path.join(input_dir, model_name_sp + '.csv'), header = 11, sep = \" \", index_col = False, names = fields)\n",
    "\n",
    "#   function that defines lithology based on porosity limit values\n",
    "def get_lith_v2(data, res, non_res):\n",
    "    \"\"\"\n",
    "    :param data: dataframe input\n",
    "    :param res: lower porosity value treshold\n",
    "    :param non_res: upper porosity value treshold\n",
    "    \"\"\"\n",
    "    data['Code'] = pd.NaT\n",
    "    data['Lith'] = pd.NaT\n",
    "    #   sand is porosity higher than 0 but lower/equal than the res value\n",
    "    data.loc[(data[\"Vsh\"] <= res) & (data[\"Vsh\"] > 0), ['Lith', 'Code']] = ['Sand', int(1)]\n",
    "    #   shale has a high porosity equal or higher than non_res\n",
    "    data.loc[data.Vsh >= non_res, ['Lith', 'Code']] = ['Shale', int(3)]\n",
    "    #   mixed porosity is everything in between - so higher than res but lower than non_res\n",
    "    data.loc[(data[\"Vsh\"] > res) & (data[\"Vsh\"] < non_res), ['Lith', 'Code']] = ['Mixed', int(2)]\n",
    "    #   the rest will be assigned a 0 porosity\n",
    "    data.loc[data.Code.isnull(), ['Lith', 'Code']] = ['NaN', int(0)]\n",
    "    return data\n",
    "\n",
    "#   add the lithology codes to the dataframe\n",
    "df = get_lith_v2(df_sp_in, 0.3, 0.6)\n",
    "df['hyd_k'] = pd.NaT\n",
    "df['por'] = pd.NaT\n",
    "\n",
    "#   insert hydraulic conductivity and porosity values into the dataframe, using theoretical values for the hydraulic conductivity\n",
    "df.loc[df.Code == 1, ['hyd_k', 'por']] = [float(0.1), float(0.3)]\n",
    "df.loc[df.Code == 2, ['hyd_k', 'por']] = [float(0.01), float(0.45)]\n",
    "df.loc[df.Code == 3, ['hyd_k', 'por']] = [float(0.0001), float(0.60)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cad3df22-cd18-4928-a2c0-6cfb3ad5c594",
   "metadata": {},
   "source": [
    "### 4) Setting up the 3D model grid and geological conditions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd6d6e31-c82d-4804-bd48-53130c737fc6",
   "metadata": {},
   "source": [
    "Now we can create the IBOUND array, the Hydraulic conductivity (hk_arr) and the porosity arrasy (poro_arr) for the current stress periods. We will perform our first semi-illegal \"for loop\" in this tutorial and iterate through the rows of the CSV file to define active model cells and their hydraulic conductivity and porosity values - the rest of the IBOUND array (and the other arrays) will stay inactive."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "c5b484bd-f7c7-4143-a250-99a339a2a510",
   "metadata": {},
   "outputs": [],
   "source": [
    "#   now we can use this dataframe to fill in the ibound, poro and hk arrays as copy of the full ibound array\n",
    "ibound_arr = np.copy(full_ibound_arr)\n",
    "hk_arr = np.copy(full_ibound_arr)\n",
    "poro_arr = np.copy(full_ibound_arr)\n",
    "\n",
    "#   now gor row by row and if the hk_val is not nan or 0 then adapt the ibound (and other) array\n",
    "df_sel = df_in.loc[df_in['Lith'] > 0.]\n",
    "for row in df_sel.iterrows():\n",
    "    hk_arr[row[1]['K'].astype(int) - 1, row[1]['I'].astype(int) - 1, row[1]['J'].astype(int) - 1] = row[1]['hyd']\n",
    "    poro_arr[row[1]['K'].astype(int) - 1, row[1]['I'].astype(int) - 1, row[1]['J'].astype(int) - 1] = row[1]['por']\n",
    "    ibound_arr[row[1]['K'].astype(int) - 1, row[1]['I'].astype(int) - 1, row[1]['J'].astype(int) - 1] = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3e68597-842c-47d6-911c-4f7114b47936",
   "metadata": {},
   "source": [
    "Lets plot the 3D arrays and admire our effort so far! But first we need to create arrays with top and bottom elevation of each cell within the 3D grid. Since we work with a regular grid we can simply assign the same cell top and bottom elevation for each cell within a model layer. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "7d1056e1-7e90-4044-b021-5677a33dc4a1",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Created top and bottom elevation arrays with shape (layer, row, column) of 35 154 160\n"
     ]
    }
   ],
   "source": [
    "#   create a 3D array with the top and bottom cell elevation - we will use it for plotting the 3D array in the right scale (meters)\n",
    "#   first create 1D list with cell top and bottom elevation (same for each layer since we have a regular grid)\n",
    "top_arr_1d = np.arange(top_elev, bot_elev[-1], -dz)\n",
    "bot_arr_1d = np.arange(bot_elev[0], bot_elev[-1] - dz, -dz)\n",
    "\n",
    "#   now make 2D and 3D arrays that we will use for plotting\n",
    "arr_2d = np.ones((hk_arr.shape[1], hk_arr.shape[2]))\n",
    "top_arr_3d = np.ones((hk_arr.shape[0], hk_arr.shape[1], hk_arr.shape[2]))\n",
    "bot_arr_3d = np.ones((hk_arr.shape[0], hk_arr.shape[1], hk_arr.shape[2]))\n",
    "\n",
    "#   loop through the layers and assign the top/bottom elevation\n",
    "for lay in range(top_arr_1d.shape[0]):\n",
    "    top_arr_3d[lay, :, :] = arr_2d * top_arr_1d[lay]\n",
    "    bot_arr_3d[lay, :, :] = arr_2d * bot_arr_1d[lay]\n",
    "\n",
    "print(\"Created top and bottom elevation arrays with shape (layer, row, column) of \" + str(top_arr_3d.shape[0]) + ' '+ str(top_arr_3d.shape[1]) + ' ' + str(top_arr_3d.shape[2]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98eff264-9010-4cbe-8ce5-6f40893a374f",
   "metadata": {},
   "source": [
    "Now we can finally plot the 3D array of the hydraulic conductivity! First we will mask out the cells with hydraulic conductivity value of 0 (those are the inactive cells). In the next step, we will define a DataArray using the xarray Python library - you can imagine this as a netcdf file that would store hydraulic conductivity values for a 3D grid with 3 dimensions (layer, x, y). To plot the voxelized 3d grid we use a imod Python function - feel free to have a look into the source code if you feel extra nerdy! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "8b6657d9-0e65-47e7-b557-d7d1cbf86205",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\zamrs001\\AppData\\Local\\deltaforge\\envs\\3Dmodel_env\\lib\\site-packages\\ipywidgets\\widgets\\widget.py:528: DeprecationWarning: The `ipykernel.comm.Comm` class has been deprecated. Please use the `comm` module instead.For creating comms, use the function `from comm import create_comm`.\n",
      "  self.comm = Comm(**args)\n",
      "C:\\Users\\zamrs001\\AppData\\Local\\deltaforge\\envs\\3Dmodel_env\\lib\\site-packages\\ipywidgets\\widgets\\widget.py:528: DeprecationWarning: The `ipykernel.comm.Comm` class has been deprecated. Please use the `comm` module instead.For creating comms, use the function `from comm import create_comm`.\n",
      "  self.comm = Comm(**args)\n",
      "C:\\Users\\zamrs001\\AppData\\Local\\deltaforge\\envs\\3Dmodel_env\\lib\\site-packages\\ipywidgets\\widgets\\widget.py:528: DeprecationWarning: The `ipykernel.comm.Comm` class has been deprecated. Please use the `comm` module instead.For creating comms, use the function `from comm import create_comm`.\n",
      "  self.comm = Comm(**args)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "97ebb66046354ff6a58f43e7546b0b1c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Widget(value=\"<iframe src='http://localhost:55667/index.html?ui=P_0x28db9e9c910_1&reconnect=auto' style='width…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#   mask out the arrays where Hk = 0 m/d - those are the inactive cells\n",
    "hk_arr_plot = np.copy(hk_arr)\n",
    "hk_arr_plot[hk_arr_plot == 0] = np.nan\n",
    "\n",
    "#   create an xarray from the z_array - has to fit the IMOD standards\n",
    "da = xr.DataArray(data = hk_arr_plot,\n",
    "                  dims = [\"layer\", \"y\", \"x\"],\n",
    "                  coords = {\"top\" : ([\"layer\", \"y\", \"x\"], top_arr_3d),\n",
    "                            \"bottom\" : ([\"layer\", \"y\", \"x\"], bot_arr_3d),\n",
    "                            \"layer\" : np.arange(1, hk_arr_plot.shape[0] + 1),\n",
    "                            \"y\" : np.arange(1, 1 + hk_arr_plot.shape[1]) * dy,\n",
    "                            \"x\" : np.arange(1, 1 + hk_arr_plot.shape[2]) * dx})\n",
    "\n",
    "#   create the voxels and plot the 3D grid\n",
    "z_grid = imod.visualize.grid_3d(da, vertical_exaggeration=50, exterior_only=False, exterior_depth=1, return_index=False)\n",
    "z_grid.plot(show_grid = True, window_size=[1600, 800])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6c52b4c-db1f-4195-94fa-af4ed3dd304e",
   "metadata": {},
   "source": [
    "### 5) Initial conditions and stress period definition"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b7ac75f-73de-4bb3-87af-cb146d65f5a7",
   "metadata": {},
   "source": [
    "To build our model we will use the Flopy Python package which is the official USGS Python library to build MODFLOW and SEAWAT models. As a first step to build the model we will define the directory where model input files will be stored. We will first run one stress period that will be 9999 years long - which is the maximum stress period length allowed in imod-wq (please do not ask me why..). The last line in this cell creates a SEAWAT model object in Python - think of it as a person getting ready to leave the house in the morning. In the next steps we will help them to dress up by assigning clothes, watch or jewelery representing different MODFLOW/SEAWAT model properties.\n",
    "\n",
    "For more information you can check the official Flopy website - https://flopy.readthedocs.io/en/3.3.2/code.html we will use the MODFLOW 2005 and SEAWAT versions (NOT MODFLOW 6!), because imod-wq is tied to the MODFLOW 2005 version and MODFLOW 6 doesn't have parallel groundwater salinity module yet.."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c6623560-dedd-469a-919c-ff3198f0adc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "#   define the time discretization for the model run - we will run the model for 10 000 years\n",
    "time_start = 0\n",
    "time_end = 9999\n",
    "#   define the modelname and folder for the mini stress period\n",
    "mini_modelname = 'SP_' + str(time_start) + '_to_' + str(time_end)\n",
    "mini_sp_dir = os.path.join(sp_dir, mini_modelname)\n",
    "\n",
    "#   create the SEAWAT model object and start creating individual packages\n",
    "mswt = flopy.seawat.Seawat(mini_modelname, 'nam_swt', model_ws = mini_sp_dir, exe_name = seawat_exe_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac14a8d3-12ae-4928-ad1a-9af162f40ccd",
   "metadata": {},
   "source": [
    "Now that we have our 3D grid defined and filled it in with a geological information we can proceed to define the initial hydraulic head and salinity concentrations of our model. Ideally, we would have an estimated groundwater salinity and heads from a previous simulation or through an interpolation based on local/regional measurements. However, it is often the case that such information is not available so we need to make assumptions for the initial conditions. In this case we will first start with the whole domain being totally salinized and having groundwater heads equal to 0. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "477d2eb9-9097-4691-82fc-63984d1f3bff",
   "metadata": {},
   "outputs": [],
   "source": [
    "#   define starting concentrations and starting heads, we will start with all cells being saline (35 TDS mg/l) \n",
    "sconc_arr = ibound_arr * 35.\n",
    "strt_arr = ibound_arr * 0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8302e37d-0987-4890-96e6-e2cdf7c788ab",
   "metadata": {},
   "source": [
    "#### DIS package\n",
    "This is where we define the discretization (i.e. DIS) of the model domain - the number of layers, rows and columns as well as temporal discretization. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "2d3b30f1-bab2-4d2c-8c73-c55cfc34166d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#   make the DIS package, define the input first\n",
    "nlay = lay_arr.shape[0]\n",
    "nrow = row_arr.shape[0]\n",
    "ncol = col_arr.shape[0]\n",
    "delr = [dx] * ncol\n",
    "delc = [dy] * nrow\n",
    "perlen = [365.25 * 9999]\n",
    "nstp = [10]\n",
    "nper = len(perlen)\n",
    "dis = flopy.modflow.ModflowDis(mswt, nlay, nrow, ncol, nper = 1, delr = delr, delc = delc, top = top_elev,\n",
    "                               botm = bot_elev, perlen = perlen, nstp = nstp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1191d7d9-6273-496d-9aec-d4883b9dc581",
   "metadata": {},
   "source": [
    "#### BAS package\n",
    "Specify the active model cells in the domain (the IBOUND array) and the initial heads in these active cells."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "b5acb957-46f4-4b1c-b3c1-b504f4b3af7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#   next create the BAS package\n",
    "bas = flopy.modflow.ModflowBas(mswt, ibound = ibound_arr, strt = strt_arr)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c67fa37a-b9f8-4baa-8659-94b6ef4f102e",
   "metadata": {},
   "source": [
    "#### LPF package \n",
    "The layer property flow is used to define the hydraulic conductivities (the horizonatl and vertical conductivity) arrays. Here we use a uniform "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "c229e4fa-14d6-4dea-a5ee-40dfe637fd70",
   "metadata": {},
   "outputs": [],
   "source": [
    "#   in this setup we will use the BCF package\n",
    "vk_arr = hk_arr * 0.1\n",
    "lpf = flopy.modflow.ModflowLpf(mswt, laytyp = 0, hk = hk_arr, vka = vk_arr, ipakcb = 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae4b83b3-ce81-46f1-b247-03cd2c54d957",
   "metadata": {},
   "source": [
    "#### Boundary conditions - GHB, RCH, DRN\n",
    "Now it is time to define some simple initial boundary conditions\n",
    "\n",
    "- for the cells that are at the edge of our model domain we will define a GHB (general head boundary) condition which enables flow through these cells depending on a difference between the head in the adjacent cell and the GHB head elevation specified for the boundary cell. So if we specify the GHB head to be e.g. 10m the groundwater head in this cell will be kept 10m and if the groundwater head in the adjacent cells is lower then there will be inflow from the boundary. If the adjacent cells have a higher head than the GHB head elevation then there will be outflow through the boundary cell.\n",
    "  \n",
    "- for the cells where the top elevation is above the specified sea level we will apply freshwater recharge (RCH) of a constant value. In case you want to specify different recharge periods you need to setup multiple stress periods.\n",
    "  \n",
    "- to ensure that there is a simple surface drainage system in place (since we do not simulate rivers or lakes in this case) we implement a simple drainage network (DRN) by defining a drain elevation which limits the groundwater head elevation in this cell. This ensures that we do not inject fresh water into the model domain via RCH. \n",
    "\n",
    "In this process we create the ICBUND array where we specify boundary cells (= -1), this an input into the MODFLOW/SEAWAT model. Each of these boundaries also need to have a specified concentration, therefore we create a SSM list where each boundary/recharge cell will have a specified concentration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "70dc9e83-6994-470e-b52f-9a9cbe09b45b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#   create the icbund array\n",
    "icbund_arr = ibound_arr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "7a17c979-d688-40d4-8319-c3b05afb3f74",
   "metadata": {},
   "outputs": [],
   "source": [
    "#   create the GHB package input here, also start creating the SSM package\n",
    "itype = flopy.mt3d.Mt3dSsm.itype_dict()\n",
    "ghb_input_lst = []\n",
    "chb_input_lst = []\n",
    "ssmdata = []\n",
    "cond_const = 1000000.\n",
    "#   the inland part on the edges of the active model domain will be assigned the topographical head\n",
    "#   for each row check the first active cell - and if it is above sea level then assign fresh head\n",
    "col_idx = [0, ncol - 1]\n",
    "for i in range(nrow):\n",
    "    #   select the active cells only\n",
    "    for col in col_idx:\n",
    "        row_cells = [t for t in ibound_arr[:, i, col].tolist() if t == 1]\n",
    "        if len(row_cells) > 0:\n",
    "            lay_idx = ibound_arr[:, i, col].tolist().index(1)\n",
    "            if ibound_arr[lay_idx, i, col] == 1 and bot_elev[lay_idx] + dz >= sea_level:\n",
    "                #   check if the top elevation of the cell is above sea level (thats why we do bot_elev + dz)\n",
    "                for k in range(nlay):\n",
    "                    if ibound_arr[k, i, col] == 1:\n",
    "                        cond_val = hk_arr[k, i, col] * dz * 1000\n",
    "                        ghb_input_lst.append([k, i, col, bot_elev[k] + dz, cond_val])\n",
    "                        ssmdata.append([k, i, 0, 0.0, itype['GHB']])\n",
    "                        icbund_arr[k, i, col] = -1\n",
    "#   now do the same but for columns\n",
    "row_idx = [0, nrow - 1]\n",
    "for i in range(ncol):\n",
    "    #   select the active cells only\n",
    "    for row in row_idx:\n",
    "        col_cells = [t for t in ibound_arr[:, row, i].tolist() if t == 1]\n",
    "        if len(col_cells) > 0:\n",
    "            lay_idx = ibound_arr[:, row, i].tolist().index(1)\n",
    "            if ibound_arr[lay_idx, row, i] == 1 and bot_elev[lay_idx] + dz >= sea_level:\n",
    "                for k in range(nlay):\n",
    "                    if bot_elev[k] + dz >= sea_level:\n",
    "                        cond_val = hk_arr[k, row_idx, i] * dz * 1000\n",
    "                        ghb_input_lst.append([k, row_idx, i, bot_elev[k] + dz, cond_val])\n",
    "                        ssmdata.append([k, row_idx, i, 0.0, itype['GHB']])\n",
    "                        icbund_arr[k, row_idx, i] = -1\n",
    "#   now check for the offshore domain and set all the cells below sea level to saltwater concentration and\n",
    "#   head equal to sea level. Only in the top layer\n",
    "for i in range(nrow):\n",
    "    for j in range(ncol):\n",
    "        lay_cells = [t for t in ibound_arr[:, i, j].tolist() if t == 1]\n",
    "        if len(lay_cells) > 0:\n",
    "            lay_idx = ibound_arr[:, i, j].tolist().index(1)\n",
    "            if ibound_arr[lay_idx, i, j] == 1 and bot_elev[lay_idx] + dz < sea_level:\n",
    "                cond_val = (vk_arr[lay_idx, i, j] * dx * dy) / dz\n",
    "                ghb_input_lst.append([lay_idx, i, j, sea_level, cond_val])\n",
    "                ssmdata.append([lay_idx, i, j, 35.0, itype['GHB']])\n",
    "#   write the final output dictionary, inlcude each stress period\n",
    "ghb_arr_in = {}\n",
    "for d in range(len(perlen)):\n",
    "    ghb_arr_in[d] = ghb_input_lst\n",
    "ghb = flopy.modflow.ModflowGhb(mswt, ipakcb = 1, stress_period_data = ghb_arr_in)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "4ae1c63d-ca70-42ef-8f7a-f01c2b19f571",
   "metadata": {},
   "outputs": [],
   "source": [
    "#   the RCH package\n",
    "rch_val = 0.00025\n",
    "rch_arr = np.zeros((1, ibound_arr.shape[1], ibound_arr.shape[2]))\n",
    "#   only apply recharge to the cells above sea level, for each column find the first active layer\n",
    "for i in range(ibound_arr.shape[1]):\n",
    "    for j in range(ibound_arr.shape[2]):\n",
    "        lay_lst = [t for t in ibound_arr[:, i, j].tolist() if t == 1]\n",
    "        if len(lay_lst) > 0:\n",
    "            top_act_lay = ibound_arr[:, i, j].tolist().index(1)\n",
    "            top_act_elev = top_elev - top_act_lay * dz\n",
    "            #   if the top elevation of the cell is above sea level then assign recharge to it\n",
    "            if top_act_elev >= sea_level:\n",
    "                #print(top_act_elev)\n",
    "                rch_arr[0, i, j] = rch_val\n",
    "rch = flopy.modflow.ModflowRch(mswt, nrchop = 3, ipakcb = 1, rech = rch_arr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "e88755cf-7c99-4468-b7ef-08ecedaaa4ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "#   the DRN package is assigned only to cells that receive recharge - cells with elev above sea level\n",
    "drn_input_lst = []\n",
    "for i in range(ibound_arr.shape[1]):\n",
    "    for j in range(ibound_arr.shape[2]):\n",
    "        lay_lst = [t for t in ibound_arr[:, i, j].tolist() if t == 1]\n",
    "        if len(lay_lst) > 0:\n",
    "            top_act_lay = ibound_arr[:, i, j].tolist().index(1)\n",
    "            top_act_elev = top_elev - top_act_lay * dz\n",
    "            #   if the top elevation of the cell is above sea level then assign recharge to it\n",
    "            if top_act_elev >= sea_level:\n",
    "                cond_cell = (vk_arr[top_act_lay, i, j] * dx * dy) / dz\n",
    "                drn_input_lst.append([int(top_act_lay), i, j, top_act_elev, cond_cell])\n",
    "#   write the final output dictionary, inlcude each stress period\n",
    "if len(drn_input_lst) > 0:\n",
    "    drn_arr_in = {}\n",
    "    for c in range(len(perlen)):\n",
    "        drn_arr_in[c] = drn_input_lst\n",
    "    drn = flopy.modflow.ModflowDrn(mswt, ipakcb=1, stress_period_data=drn_arr_in)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bae5cd6-c074-4acc-802e-7e22e57a781e",
   "metadata": {},
   "source": [
    "#### OC package\n",
    "This step will help us define the frequency of model output to be stored (concentration, heads, flow vectors etc.). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "a774f0f6-1866-4b4e-8a4f-6f4f8112df5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#   write the OC package\n",
    "ihedfm = 1  # a code for the format in which heads will be printed.\n",
    "iddnfm = 0  # a code for the format in which drawdowns will be printed.\n",
    "extension = ['oc', 'hds', 'ddn', 'cbc']\n",
    "unitnumber = [14, 30, 52, 51]\n",
    "#   create the dictionary that defines how to write the output file\n",
    "spd = {(0, 0): ['SAVE HEAD', 'SAVE BUDGET', 'PRINT HEAD', 'PRINT BUDGET', 'SAVE HEADTEC', 'SAVE CONCTEC',\n",
    "                'SAVE VXTEC', 'SAVE VYTEC', 'SAVE VZTEC']}\n",
    "for t in range(0, nper):\n",
    "    per = t  # + 1\n",
    "    #   to save space on disk, every 10th timestep is saved\n",
    "    for g in range(1, nstp[t] + 1):\n",
    "        spd[(per, int(g))] = ['SAVE HEAD', 'SAVE BUDGET', 'PRINT HEAD', 'PRINT BUDGET', 'SAVE HEADTEC', 'SAVE CONCTEC',\n",
    "                              'SAVE VXTEC', 'SAVE VYTEC', 'SAVE VZTEC']\n",
    "oc = flopy.modflow.ModflowOc(mswt, ihedfm=ihedfm, stress_period_data=spd, unitnumber=unitnumber, compact=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c07876a3-6b7f-497d-9a50-30adbc3027b3",
   "metadata": {},
   "source": [
    "#### BTN package \n",
    "Package to define the initial concentration and porosity.. det0 = 0 means that SEAWAT/imod-wq will automatically compute the largest time step possible and therefore make sure the model is as fast as possible. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "f23b85f9-993b-43ac-aead-86366ecf30f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "#   the BTN package\n",
    "porosity = poro_arr\n",
    "dt0 = 365.25\n",
    "nprs = 1\n",
    "ifmtcn = 0\n",
    "chkmas = False\n",
    "nprmas = 10\n",
    "nprobs = 10\n",
    "timprs_lst = list(np.linspace(1., perlen[0], 10, endpoint=True, dtype=int))\n",
    "btn = flopy.mt3d.Mt3dBtn(mswt, nprs=nprs, timprs=timprs_lst, prsity=porosity, sconc=sconc_arr,\n",
    "                         ifmtcn=ifmtcn, chkmas=chkmas, nprobs=nprobs, nprmas=nprmas, dt0=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e863deb5-8b6c-4f61-8773-a2d4b858e4d6",
   "metadata": {},
   "source": [
    "#### ADV package \n",
    "Here we define the solver we will use in the model - mixelm (int) – MIXELM is an integer flag for the advection solution option.\n",
    "\n",
    "- MIXELM = 0, the standard finite-difference method with upstream or central-in-space weighting, depending on the value of NADVFD;\n",
    "- = 1, the forward-tracking method of characteristics (MOC);\n",
    "- = 2, the backward-tracking modified method of characteristics (MMOC);\n",
    "- = 3, the hybrid method of characteristics (HMOC) with MOC or MMOC automatically and dynamically selected;\n",
    "- = -1, the third-order\n",
    "\n",
    "We will use the simplest finite-difference method since it is the fastest - but it is also the most simple and probably wrong solver to use in general.."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "6ddc317f-69c9-4138-a55f-e2f0c0ee329c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#   write the ADV package\n",
    "adv = flopy.mt3d.Mt3dAdv(mswt, mixelm=0, mxpart=2000000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b38c7b2-5142-466d-b4b5-4610f513f7d2",
   "metadata": {},
   "source": [
    "#### DSP package\n",
    "This is where we define the dispersivity and diffusion coefficients.."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "e5deb4ad-b543-4796-888e-185d7964373c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#   write the DSP package\n",
    "dmcoef = 0.0000864  # effective molecular diffusion coefficient [M2/D]\n",
    "al = 1.\n",
    "trpt = 0.1\n",
    "trpv = 0.1\n",
    "dsp = flopy.mt3d.Mt3dDsp(mswt, al=al, trpt=trpt, trpv=trpv, dmcoef=dmcoef)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b871215e-1e6a-42e0-9992-10df079113fb",
   "metadata": {},
   "source": [
    "#### VDF package"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "51ccbd56-4e67-4600-b477-e4c6c80f57f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "#   write the VDF package\n",
    "iwtable = 0\n",
    "densemin = 1000.\n",
    "densemax = 1025.\n",
    "denseref = 1000.\n",
    "denseslp = 0.7143\n",
    "firstdt = 0.001\n",
    "vdf = flopy.seawat.SeawatVdf(mswt, iwtable=iwtable, densemin=densemin, densemax=densemax,\n",
    "                             denseref=denseref, denseslp=denseslp, firstdt=firstdt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b34c049-c532-4a41-af18-73101cf0139d",
   "metadata": {},
   "source": [
    "#### SSSM package \n",
    "Here we define the boundary conditions concentration values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "31dd1fff-63b5-4b4e-a161-80c663531154",
   "metadata": {},
   "outputs": [],
   "source": [
    "#   write the SSM package\n",
    "ssm_rch_in = np.copy(rch_arr) * 0.0\n",
    "ssmdata_dict = {0: ssmdata, 1: ssmdata}\n",
    "ssm = flopy.mt3d.Mt3dSsm(mswt, crch=ssm_rch_in, stress_period_data=ssmdata_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5aba9381-9a98-4dd7-90a7-9bef6a3450ed",
   "metadata": {},
   "source": [
    "#### Finally now we can use this one line to create all the model input files.\n",
    "After you run the line go to your model folder and you should see modflow files being created."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "1770f360-4ca8-4610-a5da-fdaacce76d05",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Util2d rech_1: locat is None, but model does not support free format and how is internal... resetting how = external\n"
     ]
    }
   ],
   "source": [
    "#   write packages and run model\n",
    "mswt.write_input()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56df49a0-5898-488d-a78f-81ca0454c256",
   "metadata": {},
   "source": [
    "#### Additionally, to use imod-wq, we also need to create files that will tell it what solvers to use and what grid to use to split the models into individual partitions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "4f5f407a-8471-49dd-9b47-261dca8e413e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#   write the ascii file with vertical sum of active cells in IBOUND\n",
    "ibound_arr_sum = np.sum(ibound_arr, axis=0, dtype=np.int32)\n",
    "ibound_arr_sum = ibound_arr_sum.astype(str)\n",
    "with open(os.path.join(mini_sp_dir, 'LOAD.ASC'), 'wb') as f:\n",
    "    f.write(ibound_arr_sum)\n",
    "\n",
    "#   create the pksf and pkst files - change it in case the grid discretization changes\n",
    "pksf_lines = ['ISOLVER 1', 'NPC 2', 'MXITER 200', 'RELAX .98', 'HCLOSEPKS 0.01', 'RCLOSEPKS 10000.0', 'PARTOPT 0',\n",
    "              'PARTDATA', 'external 40 1. (free) -1', 'GNCOL 274', 'GNROW 223', 'GDELR', '1000', 'GDELC', '1000',\n",
    "              'NOVLAPADV 2', 'END']\n",
    "pkst_lines = ['ISOLVER 2', 'NPC 2', 'MXITER 1000', 'INNERIT 50', 'RELAX .98', 'RCLOSEPKS 1.0E-05',\n",
    "              'HCLOSEPKS 1.0E+12', 'RELATIVE-L2NORM', 'END']\n",
    "# 'CCLOSEPKS=0.00001'\n",
    "\n",
    "with open(os.path.join(mini_sp_dir, mini_modelname + '.pksf'), 'w') as f:\n",
    "    for line in pksf_lines:\n",
    "        f.write(line)\n",
    "        f.write('\\n')\n",
    "\n",
    "with open(os.path.join(mini_sp_dir, mini_modelname + '.pkst'), 'w') as f:\n",
    "    for line in pkst_lines:\n",
    "        f.write(line)\n",
    "        f.write('\\n')\n",
    "\n",
    "#   open the nam_swt file and append these three lines\n",
    "nam_lines = ['PKSF          \t  27 ' + mini_modelname + '.pksf', 'PKST              35 ' +\n",
    "             mini_modelname + '.pkst', 'DATA 40 LOAD.ASC']\n",
    "\n",
    "with open(os.path.join(mini_sp_dir, mini_modelname + '.nam_swt'), 'a') as f:\n",
    "    for line in nam_lines:\n",
    "        f.write(line)\n",
    "        f.write('\\n')\n",
    "\n",
    "#   save the ibound_arr (if it doesnt exist)\n",
    "ibound_arr_dir = os.path.join(sp_dir, 'ibound_arr.npy')\n",
    "if not os.path.isfile(ibound_arr_dir):\n",
    "    np.save(ibound_arr_dir, ibound_arr)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf6177cd-55f1-4951-a6e4-5b758e5002ad",
   "metadata": {},
   "source": [
    "### 5) Running the SEAWAT/imod-wq model in Windows"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bb8de61-332b-4c7d-921b-0311826ea0d9",
   "metadata": {},
   "source": [
    "The SEAWAT executables - all is in _3Dmodel/imodwq folder.\n",
    "\n",
    "Place all the files either into the model directory, or make a C:Drive folder and put them there (e.g. C:\\SEAWAT). There are two batch files (run.bat and run_parallel.bat), these are for single-core or parallel runs respectively – therefore use accordingly. Note that the name file (ending with ‘.nam_swt’) in the batch file will need to be adjusted according to your model name.  Within the batch file, the location of the SEAWAT executables will also need to be updated if they are not in your model directory. \n",
    "\n",
    "Message Passing Interface (MPICH2) - MPICH is a tool used for distributed-memory applications used in parallel computing. It needs to be installed to run SEAWAT in parallel. The MPICH2 installation file is located in: _3Dmodel/imodwq folder.\n",
    "\n",
    "After installation, make sure that the batch file (run_parallel.bat) file you are using is linked to the right installation location.cation.\r\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc369320-3f32-487f-9da2-7ed49231a5d5",
   "metadata": {},
   "source": [
    "When the model run is finished you will see the following files created in your folder. As you can see we used 8 cores so the model domain was split into 8 partitions. Now we need to put the puzzle back together so we can examine the groundwater salinity for the whole model domain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "409d42cb-2a77-4392-9268-d198cf9734e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#   create the results directory if it doesnt exist\n",
    "main_res_dir = os.path.join(model_dir, 'results')\n",
    "os.makedirs(main_res_dir, exist_ok = True)\n",
    "\n",
    "#   create the main output folders\n",
    "netcdf_dir = os.path.join(main_res_dir, '_netcdf', model_name_sp)\n",
    "os.makedirs(netcdf_dir, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "id": "290d0456-2300-4bee-9b32-9bc93a902771",
   "metadata": {},
   "outputs": [],
   "source": [
    "#   get the nlay, nrow and ncol values\n",
    "nlay, nrow, ncol = ibound_arr.shape[0], ibound_arr.shape[1], ibound_arr.shape[2]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "369e48be-e1f9-4abc-83af-d059c7b8ba62",
   "metadata": {},
   "source": [
    "To put the puzzle back together we will loop through the individual partitions and glue them back together into the big 3D array of the whole model domain. Here we will open the \"list\" file (open it in your text editor) where the grid partitioning is specified."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "278071c5-0a1f-4079-b1ab-29f55b33f9f3",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'mini_sp_dir' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[17], line 6\u001b[0m\n\u001b[0;32m      4\u001b[0m str_start \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m p000 :   \u001b[39m\u001b[38;5;124m'\u001b[39m\n\u001b[0;32m      5\u001b[0m str_end \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m p007 :   \u001b[39m\u001b[38;5;124m'\u001b[39m\n\u001b[1;32m----> 6\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m os\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39misfile(os\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39mjoin(\u001b[43mmini_sp_dir\u001b[49m, mini_modelname \u001b[38;5;241m+\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m.list.p000\u001b[39m\u001b[38;5;124m'\u001b[39m)):\n\u001b[0;32m      7\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mopen\u001b[39m(os\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39mjoin(mini_sp_dir, mini_modelname \u001b[38;5;241m+\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m.list.p000\u001b[39m\u001b[38;5;124m'\u001b[39m), \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mr\u001b[39m\u001b[38;5;124m'\u001b[39m) \u001b[38;5;28;01mas\u001b[39;00m listfile:\n\u001b[0;32m      8\u001b[0m         lines \u001b[38;5;241m=\u001b[39m listfile\u001b[38;5;241m.\u001b[39mreadlines()\n",
      "\u001b[1;31mNameError\u001b[0m: name 'mini_sp_dir' is not defined"
     ]
    }
   ],
   "source": [
    "#   Extract the salinity concentrations and heads from the partition list and UCN files\n",
    "#   find the list.p0000 file\n",
    "partition_idxs_lst = []\n",
    "str_start = ' p000 :   '\n",
    "str_end = ' p007 :   '\n",
    "if os.path.isfile(os.path.join(mini_sp_dir, mini_modelname + '.list.p000')):\n",
    "    with open(os.path.join(mini_sp_dir, mini_modelname + '.list.p000'), 'r') as listfile:\n",
    "        lines = listfile.readlines()\n",
    "        #   loop through a list of core numbers\n",
    "        for i in range(24):\n",
    "            if i < 10:\n",
    "                str_i = '0' + str(i)\n",
    "            else:\n",
    "                str_i = str(i)\n",
    "            #   look for the line\n",
    "            for row in lines:\n",
    "                if row.find(' p0' + str_i + ' :   ') == 0:\n",
    "                    idx_lst = [int(i) for i in row.split('|')[1].split()]\n",
    "                    #   append to the list of partition_idxs_lst\n",
    "                    #   the format: partition, partition_str, row_start, row_end, col_start, col_end\n",
    "                    partition_idxs_lst.append([i, 'p0' + str_i, idx_lst[-2], idx_lst[-1],\n",
    "                                               idx_lst[-4], idx_lst[-3]])\n",
    "\n",
    "#    print the file location to open in your text editor\n",
    "print(\"Please open this file in your text editor!   \" + os.path.join(mini_sp_dir, mini_modelname + '.list.p000'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2eaf760-31e1-4cf9-8547-3e0442944733",
   "metadata": {},
   "source": [
    "Now we read the binary file with the model output, and print the time steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "bc7d2d70-e9f7-4696-a906-07fa3c911ca7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TIME STEPS  [1.0, 405790.0, 811590.0, 1217400.0, 1623200.0, 2029000.0, 2434800.0, 2840500.0, 3246300.0, 3652100.0]\n"
     ]
    }
   ],
   "source": [
    "#   read the first UCN file to get time steps\n",
    "ucnobj = bf.UcnFile(os.path.join(mini_sp_dir, 'MT3D001.UCN.p000'))\n",
    "time_steps = ucnobj.get_times()\n",
    "print(\"TIME STEPS \", time_steps)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a67f1cb-593a-4067-b842-4036f84903e2",
   "metadata": {},
   "source": [
    "Next, create the output arrays for each partition for the given time step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "f10e2c21-a2ce-44ab-878d-723b8c5130ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "#   only create the output for the last time step\n",
    "ts = time_steps[-1]\n",
    "\n",
    "#   make a 3D array with dimension of nlay (n partitions) and nrow and ncol\n",
    "partition_arr = np.zeros((len(partition_idxs_lst), nlay, nrow, ncol)) * np.nan\n",
    "partition_head_arr = np.zeros((len(partition_idxs_lst), nlay, nrow, ncol)) * np.nan\n",
    "#   now loop through the partition indxs list and assign the partition array to the corresponding vertical layer\n",
    "for k in range(len(partition_idxs_lst)):\n",
    "    #   read in the UCN file and get the right timestep\n",
    "    ucnobj = bf.UcnFile(os.path.join(mini_sp_dir, 'MT3D001.UCN.' + partition_idxs_lst[k][1]))\n",
    "    time_steps = ucnobj.get_times()\n",
    "    conc_arr = ucnobj.get_data(totim=ts).astype(dtype=np.float64)\n",
    "    conc_arr[conc_arr > 100.] = np.nan\n",
    "    partition_arr[k, :, partition_idxs_lst[k][2] - 1 : partition_idxs_lst[k][3],\n",
    "                  partition_idxs_lst[k][4] - 1 : partition_idxs_lst[k][5]] = conc_arr\n",
    "    nrow_part, ncol_part = conc_arr.shape[1], conc_arr.shape[2]\n",
    "\n",
    "    #   read the heads from the list file as well\n",
    "    #with open(os.path.join(mini_sp_dir, mini_modelname + '.list.p000'), 'r') as listfile:\n",
    "    with open(os.path.join(mini_sp_dir, mini_modelname + '.list.' + partition_idxs_lst[k][1]), 'r') as listfile:\n",
    "        lines = listfile.readlines()\n",
    "        #   create a head array for the partition and find the corresponding lines\n",
    "        head_arr = np.zeros((nlay, nrow_part, ncol_part)) * np.nan\n",
    "        #   loop through layers\n",
    "        for z in range(nlay):\n",
    "            #   now look for the string in the list file\n",
    "            lay_ts_str = 'HEAD IN LAYER' + str(z + 1).rjust(4) + ' AT END OF TIME STEP' + str(time_steps.index(ts) + 1).rjust(4)\n",
    "            #print(lay_ts_str)\n",
    "            lay_str_end = '1'\n",
    "            #   find the starting and end index\n",
    "            for row in lines:\n",
    "                if row.find(lay_ts_str) > 0:\n",
    "                    st_idx = lines.index(row)\n",
    "                    break\n",
    "            for row in lines[st_idx:]:\n",
    "                if row.find(lay_str_end) == 0:\n",
    "                    end_idx = st_idx + lines[st_idx:].index(row)\n",
    "                    break\n",
    "            #   loop through the selected lines and get all the row, col values into the head_arr\n",
    "            #   first, create a list of row indexes that we will then loop through to extract the data\n",
    "            row_idx_lst = []\n",
    "            lines_sel = lines[st_idx : end_idx]\n",
    "            for y in range(nrow_part):\n",
    "                for row in lines_sel:\n",
    "                    if row.find(str(y + 1).rjust(4)) == 0:\n",
    "                        row_idx_lst.append(lines_sel.index(row))\n",
    "            #   since it is a regular grid the number of rows in the text file will always be the same\n",
    "            #   so calculate the row step and then use it to get the actuall head values.. finally\n",
    "            row_step = row_idx_lst[1] - row_idx_lst[0]\n",
    "            for row in row_idx_lst:\n",
    "                lines_row = lines_sel[row : row + row_step]\n",
    "                heads_col_lst = []\n",
    "                for line in lines_row:\n",
    "                    if lines_sel.index(line) == row:\n",
    "                        #   append the values but skip the row index which is always the first string\n",
    "                        heads_col_lst.append([float(num) for num in line.split()][1:])\n",
    "                    else:\n",
    "                        heads_col_lst.append([float(num) for num in line.split()])\n",
    "                final_head_lst = [item for sublist in heads_col_lst for item in sublist]\n",
    "                #   insert the list into the head_arr\n",
    "                head_arr[z, row_idx_lst.index(row), :] = final_head_lst\n",
    "    #   add the partition head array into the overall head array\n",
    "    partition_head_arr[k, :, partition_idxs_lst[k][2] - 1 : partition_idxs_lst[k][3],\n",
    "                       partition_idxs_lst[k][4] - 1 : partition_idxs_lst[k][5]] = head_arr"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7382777d-1d4d-4849-845e-5e3e78d65e0e",
   "metadata": {},
   "source": [
    "There is always an overlay on the edges of each partition - in these overlapping cells we will take the mean concentration (and heads). So we average the arrays along an axis.."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "60b73779-e9a1-4261-89e8-2db5d417b16f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#   create a nanmean array from the partition_arr - along the axis = 0 (which is the number of partitions)\n",
    "final_arr = np.nanmean(partition_arr, axis=0)\n",
    "final_head_arr = np.nanmean(partition_head_arr, axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "271f6cf8-af10-4d98-bb8e-2b511df00cce",
   "metadata": {},
   "source": [
    "Create output netcdf and numpy dictionary files - these can be used in case you want to run a subsequent stress period. In such case you would use these files to load the initial heads and concentration for the next stress period."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "61e3a5ae-6132-4332-88a8-bf6423174e88",
   "metadata": {},
   "outputs": [],
   "source": [
    "#   create the output files\n",
    "sconc_dir = os.path.join(mini_sp_dir, '_sconc_arr.npy')\n",
    "strt_dir = os.path.join(mini_sp_dir, '_strt_arr.npy')\n",
    "np.save(sconc_dir, final_arr, allow_pickle=True)\n",
    "np.save(strt_dir, final_head_arr, allow_pickle=True)\n",
    "#   make the netcdf files\n",
    "tot_time_str = str(sp_time_dur[a] * a + time_end)\n",
    "conc_nc_dir = os.path.join(netcdf_dir, '_conc_time_' + tot_time_str + '_yrs.nc')\n",
    "head_nc_dir = os.path.join(netcdf_dir, '_head_time_' + tot_time_str + '_yrs.nc')\n",
    "#   create the netcdf files\n",
    "x_coord_lst = np.arange(dx / 2, final_arr.shape[2] * dx + dx / 2, dx).tolist()\n",
    "y_coord_lst = np.arange(dy / 2, final_arr.shape[1] * dy + dy / 2, dy).tolist()\n",
    "z_coord_lst = np.linspace(0, final_arr.shape[0], final_arr.shape[0]).tolist()\n",
    "conc_nc = xr.Dataset(data_vars={'salinity': (('z', 'y', 'x'), final_arr)},\n",
    "                     coords={'x': x_coord_lst,\n",
    "                             'y': y_coord_lst,\n",
    "                             'z': z_coord_lst,\n",
    "                             'time': sp_time_dur[a] * a + time_end})\n",
    "conc_nc.to_netcdf(conc_nc_dir)\n",
    "head_nc = xr.Dataset(data_vars={'gw_head': (('z', 'y', 'x'), final_head_arr)},\n",
    "                     coords={'x': x_coord_lst,\n",
    "                             'y': y_coord_lst,\n",
    "                             'z': z_coord_lst,\n",
    "                             'time': sp_time_dur[a] * a + time_end})\n",
    "head_nc.to_netcdf(head_nc_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d89d751-a0dc-44d2-98bc-e1a302f646c7",
   "metadata": {},
   "source": [
    "And finally, we will plot the concentration!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "id": "3cd0fd0a-b318-421c-95bd-b5dd9327dd31",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "503a8eaba36d4401b54b07bc3e302329",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Widget(value=\"<iframe src='http://localhost:52062/index.html?ui=P_0x27fa217f790_11&reconnect=auto' style='widt…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#   create an xarray from the z_array - has to fit the IMOD standards\n",
    "da = xr.DataArray(data = final_arr,\n",
    "                  dims = [\"layer\", \"y\", \"x\"],\n",
    "                  coords = {\"top\" : ([\"layer\", \"y\", \"x\"], top_arr_3d),\n",
    "                            \"bottom\" : ([\"layer\", \"y\", \"x\"], bot_arr_3d),\n",
    "                            \"layer\" : np.arange(1, hk_arr_plot.shape[0] + 1),\n",
    "                            \"y\" : np.arange(1, 1 + hk_arr_plot.shape[1]) * dy,\n",
    "                            \"x\" : np.arange(1, 1 + hk_arr_plot.shape[2]) * dx})\n",
    "\n",
    "#   define cmap\n",
    "cmap = plt.cm.jet  # define the colormap\n",
    "# extract all colors from the .jet map\n",
    "cmaplist = [cmap(i) for i in range(cmap.N)]\n",
    "# create the new map\n",
    "cmap = matplotlib.colors.LinearSegmentedColormap.from_list('Salinity cmap', cmaplist, cmap.N)\n",
    "# define the bins and normalize\n",
    "bounds = np.array([0., 0.1, 0.25, 0.5, 1, 2.5, 5, 10., 15., 25., 35.])\n",
    "norm = matplotlib.colors.BoundaryNorm(bounds, cmap.N)\n",
    "\n",
    "#   create the voxels and plot the 3D grid\n",
    "z_grid = imod.visualize.grid_3d(da, vertical_exaggeration=50, exterior_only=False, exterior_depth=1, return_index=False)\n",
    "z_grid.plot(show_grid = True, window_size=[1600, 800], cmap = cmap)"
   ]
  }
 ],
 "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}