POC.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e2aeb6f5-bbf7-4bb0-b5f6-7bce4ed4ad65",
   "metadata": {},
   "source": [
    "###### Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2602308d-9abc-4d8b-b301-056215b2bb49",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os,sys,py3dep,xarray,rasterio,folium,geopandas,math,numpy,datetime,shutil,pynhd,shapely,matplotlib,branca,shapely,whitebox_workflows,pygeohydro,multiprocessing"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aacb6515-5437-4486-975b-3edb2cd354dd",
   "metadata": {},
   "source": [
    "###### Import hydroframe library and register pin -- see https://hydroframe.org/hydrodata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a09ffde0-f144-4358-8d37-c99778b29ac6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import hf_hydrodata as hf\n",
    "hf.register_api_pin(\"evenson.grey@epa.gov\", \"2234\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1b96ebd-4c4f-4018-956e-5f7d7f06b9f2",
   "metadata": {},
   "source": [
    "###### Set run specific variables\n",
    "###### NOTE: pysdaDir must be set as path to pysda clone -- see https://github.com/ncss-tech/pysda.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f442f6f7-19c5-4077-9695-95ca86529175",
   "metadata": {},
   "outputs": [],
   "source": [
    "pysdaDir       = r'C:\\Users\\GEVENSON\\GitHub\\pysda'\n",
    "dirProject     = r'C:\\Users\\GEVENSON\\Projects\\nonperennial_stream_poc\\BUCKHORN_CREEK'\n",
    "DEMfname       = r'C:\\Users\\GEVENSON\\Projects\\nonperennial_stream_poc\\dem_preprocessing\\BUCKHORN_CREEK_2M_DEM.tiff'\n",
    "start_date     = \"2002-10-01\"\n",
    "end_date       = \"2006-09-30\"\n",
    "lat            = 37.44870490000000 # lat, lon info: https://www.waterqualitydata.us/provider/NWIS/USGS-KY/USGS-03278100/\n",
    "lon            = -83.1954512000000\n",
    "OVERWRITE_FLAG = False"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea3884f9-cc96-4ad9-9b76-b9b15bd717aa",
   "metadata": {},
   "source": [
    "###### Hard code grid info (see https://hf-hydrodata.readthedocs.io/en/latest/available_grids.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7c7ac425-8676-419a-8ac3-cd5de4182e51",
   "metadata": {},
   "outputs": [],
   "source": [
    "conus1_proj      = '+proj=lcc +lat_1=33 +lat_2=45 +lon_0=-96.0 +lat_0=39 +a=6378137.0 +b=6356752.31'\n",
    "conus1_spatext   = tuple([-121.47939483437318, 31.651836025255015, -76.09875469594509, 50.49802132270979])\n",
    "conus1_transform = rasterio.transform.Affine(1000.0,0.0,-1885055.4995,0.0,1000.0,-604957.0654)\n",
    "conus1_shape     = (1888,3342)\n",
    "conus2_proj      = '+proj=lcc +lat_1=30 +lat_2=60 +lon_0=-97.0 +lat_0=40.0000076294444 +a=6370000.0 +b=6370000'\n",
    "conus2_spatext   = tuple([-126.88755692881833, 21.8170599154073, -64.7677149695924, 53.20274381640737])\n",
    "conus2_transform = rasterio.transform.Affine(1000.0,0.0,-2208000.30881173,0.0,1000.0,-1668999.65483222)\n",
    "conus2_shape     = (3256,4442)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75ae2486-bc65-4766-8d17-a3ece05152b8",
   "metadata": {},
   "source": [
    "###### Hard code transmissivity decay factor from Table 1 of Zhang et al. (2016) https://doi.org/10.5194/bg-13-1387-2016"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "30bdbe01-cd51-4a6c-be24-653b1f9be89d",
   "metadata": {},
   "outputs": [],
   "source": [
    "transmissivity_decay_factor = {'clay heavy'     :3.2,\n",
    "                               'silty clay'     :3.1,\n",
    "                               'clay'           :2.8,\n",
    "                               'silty clay loam':2.9,\n",
    "                               'clay loam'      :2.7,\n",
    "                               'silt'           :3.4,\n",
    "                               'silt loam'      :2.6,\n",
    "                               'sandy clay'     :2.5,\n",
    "                               'loam'           :2.5,\n",
    "                               'sandy clay loam':2.4,\n",
    "                               'sandy loam'     :2.3,\n",
    "                               'loamy sand'     :2.2,\n",
    "                               'sand'           :2.1,\n",
    "                               'organic'        :2.5}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e84948e-0952-4f48-8043-4cbd653a48cc",
   "metadata": {},
   "source": [
    "###### Create directory architecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "18a8fd61-b944-4d9f-83c3-10d40a473883",
   "metadata": {},
   "outputs": [],
   "source": [
    "dirData               = os.path.join(dirProject,'Data')\n",
    "dirWTD                = os.path.join(dirData,'WTD')\n",
    "dirWTDParflow         = os.path.join(dirWTD,'Parflow')\n",
    "dirWTDParflowRaw      = os.path.join(dirWTDParflow,'raw')\n",
    "dirWTDParflowBilinear = os.path.join(dirWTDParflow,'resampled_bilinear')\n",
    "dirWTDParflowNearest  = os.path.join(dirWTDParflow,'resampled_nearest')\n",
    "dirWTDParflowCubic    = os.path.join(dirWTDParflow,'resampled_cubic')\n",
    "dirWTDFan             = os.path.join(dirWTD,'Fan')\n",
    "dirDEM                = os.path.join(dirData,'DEM')\n",
    "dirTWI                = os.path.join(dirData,'TWI')\n",
    "dirSoils              = os.path.join(dirData,'Soils')\n",
    "dirDomain             = os.path.join(dirData,'Domain')\n",
    "dirNHD                = os.path.join(dirData,'NHD')\n",
    "dirOutput             = os.path.join(dirProject,'Output')\n",
    "if not os.path.isdir(dirProject)           : os.mkdir(dirProject)\n",
    "if not os.path.isdir(dirData)              : os.mkdir(dirData)\n",
    "if not os.path.isdir(dirWTD)               : os.mkdir(dirWTD)\n",
    "if not os.path.isdir(dirWTDParflow)        : os.mkdir(dirWTDParflow)\n",
    "if not os.path.isdir(dirWTDParflowRaw)     : os.mkdir(dirWTDParflowRaw)\n",
    "if not os.path.isdir(dirWTDParflowBilinear): os.mkdir(dirWTDParflowBilinear)\n",
    "if not os.path.isdir(dirWTDParflowNearest) : os.mkdir(dirWTDParflowNearest)\n",
    "if not os.path.isdir(dirWTDParflowCubic)   : os.mkdir(dirWTDParflowCubic)\n",
    "if not os.path.isdir(dirWTDFan)            : os.mkdir(dirWTDFan)\n",
    "if not os.path.isdir(dirDEM)               : os.mkdir(dirDEM)\n",
    "if not os.path.isdir(dirTWI)               : os.mkdir(dirTWI)\n",
    "if not os.path.isdir(dirSoils)             : os.mkdir(dirSoils)\n",
    "if not os.path.isdir(dirDomain)            : os.mkdir(dirDomain)\n",
    "if not os.path.isdir(dirNHD)               : os.mkdir(dirNHD)\n",
    "if not os.path.isdir(dirOutput)            : os.mkdir(dirOutput)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6fe1c2a-0e30-4f4d-b6f2-6e87f8b68cb8",
   "metadata": {},
   "source": [
    "###### Set datetime dimension"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "77883ce4-157d-405e-9532-fba2dd70355a",
   "metadata": {},
   "outputs": [],
   "source": [
    "start_date_split = start_date.split('-')\n",
    "start_datetime = datetime.datetime(year=int(start_date_split[0]),month=int(start_date_split[1]),day=int(start_date_split[2]))\n",
    "end_date_split = end_date.split('-')\n",
    "end_datetime = datetime.datetime(year=int(end_date_split[0]),month=int(end_date_split[1]),day=int(end_date_split[2]))\n",
    "datetime_dim = list()\n",
    "idatetime = start_datetime\n",
    "while idatetime <= end_datetime:\n",
    "    datetime_dim.append(idatetime)\n",
    "    idatetime += datetime.timedelta(days=1)\n",
    "datetime_dim = numpy.array(datetime_dim)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae6daf7b-3c6c-4ece-8011-cb571500d17e",
   "metadata": {},
   "source": [
    "###### Identify HUC/domain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "11f5ca37-fb7a-4d78-bdf2-927bf5297c66",
   "metadata": {},
   "outputs": [],
   "source": [
    "huc_level = 'huc12' \n",
    "idp = shapely.geometry.Point(lon, lat)\n",
    "wbdbasins = pygeohydro.WBD(huc_level) \n",
    "wbdbasins  = wbdbasins.bygeom(idp)\n",
    "fname_domain = os.path.join(dirDomain,'domain.gpkg')\n",
    "domain = wbdbasins.dissolve()\n",
    "domain.to_file(fname_domain, driver=\"GPKG\")\n",
    "del wbdbasins, idp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90821e05-a07b-4ccc-ae41-48941316c9b8",
   "metadata": {},
   "source": [
    "###### Get domain lat/lon bounding box"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b088e655-bcc7-44f3-9a30-07dc3f1b25c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "buffered_domain = domain.to_crs(conus1_proj).buffer(distance=2000).to_crs(\"EPSG:4326\")\n",
    "lat_min = float(buffered_domain.bounds['miny'].iloc[0])\n",
    "lat_max = float(buffered_domain.bounds['maxy'].iloc[0])\n",
    "lon_min = float(buffered_domain.bounds['minx'].iloc[0])\n",
    "lon_max = float(buffered_domain.bounds['maxx'].iloc[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f243d704-8706-490e-966f-31413c505e04",
   "metadata": {},
   "source": [
    "###### Get domain bounding box for Parflow grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c85b99d4-a7c6-43d0-9a76-b2805526c374",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "conus1grid_minx, conus1grid_miny = hf.from_latlon(\"conus1\", lat_min, lon_min)\n",
    "conus1grid_maxx, conus1grid_maxy = hf.from_latlon(\"conus1\", lat_max, lon_max)\n",
    "conus1grid_minx, conus1grid_miny = math.floor(conus1grid_minx), math.floor(conus1grid_miny)\n",
    "conus1grid_maxx, conus1grid_maxy = math.ceil(conus1grid_maxx),  math.ceil(conus1grid_maxy)\n",
    "conus2grid_minx, conus2grid_miny = hf.from_latlon(\"conus2\", lat_min, lon_min)\n",
    "conus2grid_maxx, conus2grid_maxy = hf.from_latlon(\"conus2\", lat_max, lon_max)\n",
    "conus2grid_minx, conus2grid_miny = math.floor(conus2grid_minx), math.floor(conus2grid_miny)\n",
    "conus2grid_maxx, conus2grid_maxy = math.ceil(conus2grid_maxx),  math.ceil(conus2grid_maxy)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4829e33b-6cc3-4ab0-8b4f-8ee9c0a133f7",
   "metadata": {},
   "source": [
    "###### Reset lat/lon bounding box to include border Parflow grid cells"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c083f882-8b48-4a6d-a042-2da3df4e6851",
   "metadata": {},
   "source": [
    "latlon_bounds = hf.to_latlon(\"conus1\", *[conus1grid_minx, conus1grid_miny, conus1grid_maxx, conus1grid_maxy])\n",
    "lon_min = latlon_bounds[1]\n",
    "lat_min = latlon_bounds[0]\n",
    "lon_max = latlon_bounds[3]\n",
    "lat_max = latlon_bounds[2]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "957516e9-7960-43e8-8e3b-734ab7e37bf2",
   "metadata": {},
   "source": [
    "###### Download the Parflow water table depth data (if not already downloaded)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "eb5bc040-4826-4a13-bbc2-61e48a878923",
   "metadata": {},
   "outputs": [],
   "source": [
    "download_flag = False\n",
    "fnames_wtd = list()\n",
    "for idatetime in datetime_dim:\n",
    "    fname_wtd = os.path.join(dirWTDParflowRaw,'wtd_'+idatetime.strftime('%Y%m%d')+'.tiff')\n",
    "    if not os.path.isfile(fname_wtd):\n",
    "        download_flag = True\n",
    "        break\n",
    "if download_flag:\n",
    "    start_date_str = datetime_dim[0].strftime('%Y-%m-%d')\n",
    "    end_date_str = (datetime_dim[len(datetime_dim)-1]+datetime.timedelta(days=1)).strftime('%Y-%m-%d')\n",
    "    options_wtd = {\"dataset\": \"conus1_baseline_mod\", \"variable\": \"water_table_depth\", \"temporal_resolution\": \"daily\", \"start_time\": start_date_str, \"end_time\": end_date_str, \"grid_bounds\":[conus1grid_minx,conus1grid_miny,conus1grid_maxx,conus1grid_maxy]}  \n",
    "    hf_data = hf.get_gridded_data(options_wtd)\n",
    "    if hf_data.shape[0] != int((end_datetime - start_datetime).days) + 1: sys.exit('ERROR hydroframe returned data of unexpected time length.')\n",
    "    hf_conus1grid_temp = numpy.empty(conus1_shape)\n",
    "    for i in range(len(datetime_dim)):\n",
    "        idatetime = datetime_dim[i]\n",
    "        hf_conus1grid_temp[conus1grid_miny:conus1grid_maxy,conus1grid_minx:conus1grid_maxx] = hf_data[i,:,:]\n",
    "        memfile = rasterio.io.MemoryFile()\n",
    "        hf_conus1data = memfile.open(driver = \"GTiff\", height = hf_conus1grid_temp.shape[0], width = hf_conus1grid_temp.shape[1], crs=conus1_proj, transform = conus1_transform, nodata = numpy.nan, count = 1, dtype = numpy.float64)\n",
    "        hf_conus1data.write(hf_conus1grid_temp,1)\n",
    "        fname = os.path.join(dirWTDParflowRaw,'wtd_'+idatetime.strftime('%Y%m%d')+'.tiff')\n",
    "        wtd_data, wtd_transform = rasterio.mask.mask(hf_conus1data, buffered_domain.to_crs(hf_conus1data.crs), crop=True, all_touched=True, filled =True, nodata = numpy.nan)\n",
    "        wtd_meta = hf_conus1data.meta\n",
    "        wtd_meta.update({\"driver\": \"GTiff\",\"height\": wtd_data.shape[1],\"width\": wtd_data.shape[2],\"transform\": wtd_transform, \"nodata\" : numpy.nan})\n",
    "        with rasterio.open(fname,'w',**wtd_meta) as wtd_dataset:\n",
    "            wtd_dataset.write(wtd_data[0,:,:],1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6e3be77-1356-4259-a030-761109e8b4bd",
   "metadata": {},
   "source": [
    "###### Download the Fan (2013) equilibrium water table depth data (if not already downloaded)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "9bd54cf8-dd60-4a5f-a0ce-5a9f3695076d",
   "metadata": {},
   "outputs": [],
   "source": [
    "download_flag = False\n",
    "fnames_wtd = list()\n",
    "for idatetime in datetime_dim:\n",
    "    fname_wtd = os.path.join(dirWTDFan,'wtd_equilibrium.tiff')\n",
    "    if not os.path.isfile(fname_wtd):\n",
    "        download_flag = True\n",
    "if download_flag:\n",
    "    options_wtd = {\"dataset\": \"fan_2013\", \"variable\": \"water_table_depth\", \"grid_bounds\":[conus2grid_minx,conus2grid_miny,conus2grid_maxx,conus2grid_maxy]}  \n",
    "    hf_data = hf.get_gridded_data(options_wtd)\n",
    "    if hf_data.shape[0] != 1: sys.exit('ERROR hydroframe returned data of unexpected time length.')\n",
    "    hf_conus2grid_temp = numpy.empty(conus2_shape)\n",
    "    hf_conus2grid_temp[conus2grid_miny:conus2grid_maxy,conus2grid_minx:conus2grid_maxx] = hf_data[i,:,:]\n",
    "    memfile = rasterio.io.MemoryFile()\n",
    "    hf_conus2data = memfile.open(driver=\"GTiff\",height=hf_conus2grid_temp.shape[0],width=hf_conus2grid_temp.shape[1],crs=conus2_proj,transform=conus2_transform,nodata=numpy.nan,count=1,dtype=numpy.float64)\n",
    "    hf_conus1data.write(hf_conus1grid_temp,1)\n",
    "    fname = os.path.join(dirWTDParflowRaw,'wtd_'+idatetime.strftime('%Y%m%d')+'.tiff')\n",
    "    wtd_data, wtd_transform = rasterio.mask.mask(hf_conus1data, buffered_domain.to_crs(hf_conus1data.crs), crop=True, all_touched=True, filled =True, nodata = numpy.nan)\n",
    "    wtd_meta = hf_conus1data.meta\n",
    "    wtd_meta.update({\"driver\": \"GTiff\",\"height\": wtd_data.shape[1],\"width\": wtd_data.shape[2],\"transform\": wtd_transform, \"nodata\" : numpy.nan})\n",
    "    with rasterio.open(fname,'w',**wtd_meta) as wtd_dataset:\n",
    "        wtd_dataset.write(wtd_data[0,:,:],1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "804d5c6d-fbc3-4ec4-b2b4-85627ea4fa35",
   "metadata": {},
   "source": [
    "###### Increase resolution of ParFlow grid to 2 m (from 1000 m)\n",
    "###### Maybe just do this on the fly, in-memory, per time step, after putting the topographic info into this grid using one instance of the wtd grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "8b9781d1-515a-45e9-96fa-aeadda1f28b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#resample_methods = {'nearest':[rasterio.enums.Resampling.nearest,dirWTDParflowNearest],\n",
    "#                    'bilinear':[rasterio.enums.Resampling.bilinear,dirWTDParflowBilinear]}\n",
    "resample_methods = {'bilinear':[rasterio.enums.Resampling.bilinear,dirWTDParflowBilinear]}\n",
    "for resample_method_name, [resample_method, resample_dir] in resample_methods.items():\n",
    "    for idatetime in datetime_dim:\n",
    "        fname_wtd = os.path.join(dirWTDParflowRaw,'wtd_'+idatetime.strftime('%Y%m%d')+'.tiff')\n",
    "        fname_wtd_hr = os.path.join(resample_dir,'wtd_'+idatetime.strftime('%Y%m%d')+'_highresolution_'+resample_method_name+'.tiff')\n",
    "        if not os.path.isfile(fname_wtd_hr) or OVERWRITE_FLAG:\n",
    "            with rasterio.open(fname_wtd) as wtd_dataset:\n",
    "                wtd_data_hr = wtd_dataset.read(out_shape=(wtd_dataset.count,int(wtd_dataset.height * 500),int(wtd_dataset.width * 500)),resampling=resample_method)\n",
    "                wtd_transform_hr = wtd_dataset.transform * wtd_dataset.transform.scale((wtd_dataset.width / wtd_data_hr.shape[-1]),(wtd_dataset.height / wtd_data_hr.shape[-2]))\n",
    "                wtd_meta_hr = wtd_dataset.meta\n",
    "                wtd_meta_hr.update({\"driver\": \"GTiff\",\"height\": wtd_data_hr.shape[1],\"width\": wtd_data_hr.shape[2],\"transform\": wtd_transform_hr, \"nodata\" : numpy.nan})\n",
    "                with rasterio.open(fname_wtd_hr, \"w\", **wtd_meta_hr) as wtd_dataset_hr:\n",
    "                    wtd_dataset_hr.write(wtd_data_hr)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f6e324d-a199-4c81-a541-f45a534a3999",
   "metadata": {},
   "source": [
    "###### Create domain mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "8313e285-9d6e-4642-ab03-38721710a011",
   "metadata": {},
   "outputs": [],
   "source": [
    "fname_domain_mask = os.path.join(dirData,\"domain_mask.tiff\")\n",
    "if not os.path.isfile(fname_domain_mask) or OVERWRITE_FLAG:\n",
    "    fname_wtd_hr = os.path.join(dirWTD,'wtd_'+datetime_dim[0].strftime('%Y%m%d')+'_highresolution_bilinear_resample.tiff') # point to bilinear but all resampling methods should share the same grid / resolution\n",
    "    with rasterio.open(fname_wtd_hr,'r') as wtd_highres:\n",
    "        domain_data = wtd_highres.read(1)\n",
    "        domain_meta = wtd_highres.meta\n",
    "        domain_crs = wtd_highres.crs\n",
    "        domain_transform = wtd_highres.transform\n",
    "        domain_mask = rasterio.features.rasterize(shapes=domain.to_crs(domain_crs)['geometry'],out_shape=domain_data.shape,transform=domain_transform,fill=0,all_touched=True,dtype=rasterio.uint8,default_value=1)\n",
    "        domain_meta.update({\"driver\": \"GTiff\",\"height\": domain_data.shape[0],\"width\": domain_data.shape[1],\"transform\": domain_transform,\"dtype\": rasterio.uint8,\"nodata\":0})\n",
    "        with rasterio.open(fname_domain_mask, 'w', **domain_meta) as dst:\n",
    "            dst.write(domain_mask,indexes=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "899f5a08-dd0f-40c7-8d8c-49f19442fa48",
   "metadata": {},
   "source": [
    "###### Get soil texture data\n",
    "###### NOTE: uses pysda via https://github.com/ncss-tech/pysda.git -- pysdaDir must be set as path to pysda clone"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "961b5f8e-ce7a-43af-afca-b42d2f5f2c22",
   "metadata": {},
   "outputs": [],
   "source": [
    "fname_soil_texture = os.path.join(dirSoils,\"soil_texture.gpkg\")\n",
    "fname_soil_transmissivity_decay_factor = os.path.join(dirSoils,\"soil_transmissivity_decay_factor.tiff\")\n",
    "if not os.path.isfile(fname_soil_transmissivity_decay_factor) or OVERWRITE_FLAG:\n",
    "    if pysdaDir not in sys.path: sys.path.append(pysdaDir)\n",
    "    import sdapoly, sdaprop, sdainterp\n",
    "    soils_aoi = sdapoly.gdf(domain)\n",
    "    sandtotal_r=sdaprop.getprop(df=soils_aoi,column='mukey',method='dom_comp_num',top=0,bottom=400,prop='sandtotal_r',minmax=None,prnt=False,meta=False)\n",
    "    claytotal_r=sdaprop.getprop(df=soils_aoi,column='mukey',method='dom_comp_num',top=0,bottom=400,prop='claytotal_r',minmax=None,prnt=False,meta=False)\n",
    "    def calc_texture(row): \n",
    "        try: sand = float(row['sandtotal_r'])\n",
    "        except: return 'None'\n",
    "        try: clay = float(row['claytotal_r'])\n",
    "        except: return 'None'\n",
    "        return soiltexture.getTexture(sand,clay)\n",
    "    def calc_f(row):\n",
    "        if row['texture'] in transmissivity_decay_factor: return transmissivity_decay_factor[row['texture']]\n",
    "        else: \n",
    "            print('WARNING: Could not find transmissivity decay factor for soil texture ',row['texture'])\n",
    "            return numpy.mean(list(transmissivity_decay_factor.values()))\n",
    "    fname_wtd_hr = os.path.join(dirWTDParflowBilinear,'wtd_'+datetime_dim[0].strftime('%Y%m%d')+'_highresolution_bilinear.tiff') # point to bilinear but all resampling methods should share the same grid / resolution\n",
    "    with rasterio.open(fname_wtd_hr,'r') as wtd_highres:\n",
    "        dummy_meta = wtd_highres.meta\n",
    "        dummy_data = wtd_highres.read(1)\n",
    "        soils_aoi = soils_aoi.merge(pandas.merge(sandtotal_r,claytotal_r,on='mukey'),on='mukey')\n",
    "        soils_aoi['texture'] = soils_aoi.apply(calc_texture, axis=1)\n",
    "        soils_aoi['f'] = soils_aoi.apply(calc_f, axis=1)\n",
    "        soils_aoi = soils_aoi.to_crs(wtd_highres.crs)\n",
    "        soils_aoi.to_crs('EPSG:4326').to_file(fname_soil_texture, driver=\"GPKG\")\n",
    "        soils_shapes = ((geom,value) for geom, value in zip(soils_aoi.geometry, soils_aoi['f']))\n",
    "        texture_data = rasterio.features.rasterize(shapes=soils_shapes,out_shape=dummy_data.shape,transform=wtd_highres.transform,fill=numpy.nan,all_touched=True,dtype=rasterio.float32,default_value=numpy.nan)\n",
    "        dummy_meta.update({\"driver\": \"GTiff\",\"height\": dummy_data.shape[0],\"width\": dummy_data.shape[1],\"transform\": wtd_highres.transform,\"dtype\": rasterio.float32,\"nodata\":numpy.nan})\n",
    "        with rasterio.open(fname_soil_transmissivity_decay_factor, 'w', **dummy_meta) as dst:\n",
    "            dst.write(texture_data,indexes=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3237c950-0c0c-48d8-9d0b-a8cbf849711e",
   "metadata": {},
   "source": [
    "###### Get DEM\n",
    "###### NOTE: py3dep.check_3dep_availability may show availability of 1 meter resolution DEM but py3dep.get_dem does not work if called with resolution < 10, in my experience"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "8d2cebc9-9d32-4729-93e1-9a674aac0508",
   "metadata": {},
   "outputs": [],
   "source": [
    "fname_dem_original = os.path.join(dirDEM,\"dem_original.tif\")\n",
    "if not os.path.isfile(fname_dem_original) or OVERWRITE_FLAG:\n",
    "    if os.path.isfile(DEMfname):\n",
    "        shutil.copy2(DEMfname,fname_dem_original)\n",
    "    else:\n",
    "        #dtavailability = py3dep.check_3dep_availability([lon_min,lat_min,lon_max,lat_max]) # can see dem resolution availability here\n",
    "        dem_original = py3dep.get_dem(geometry=[lon_min,lat_min,lon_max,lat_max],resolution=10) # wpn't work with < 10 m\n",
    "        dem_original.rio.to_raster(fname_dem_original)\n",
    "        del dem_original"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4d6660a-d66a-4bfb-8c7b-3f110257d74f",
   "metadata": {},
   "source": [
    "###### Reproject DEM to match resampled (high resolution) Parflow grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "f91a9cb6-da08-44c1-bec6-27642a4d0b51",
   "metadata": {},
   "outputs": [],
   "source": [
    "fname_dem = os.path.join(dirDEM,\"dem_reprojected.tif\")\n",
    "if not os.path.isfile(fname_dem) or OVERWRITE_FLAG:\n",
    "    fname_wtd_hr = os.path.join(dirWTD,'wtd_'+datetime_dim[0].strftime('%Y%m%d')+'_highresolution_bilinear_resample.tiff') # point to bilinear but all resampling methods should share the same grid / resolution\n",
    "    with rasterio.open(fname_dem_original,'r') as dem_original and rasterio.open(fname_wtd_hr,'r') as dummy_hr_data:\n",
    "        dem_reprojected_data, dem_reprojected_transform = rasterio.warp.reproject(\n",
    "                source = dem_original.read(1),\n",
    "                destination = dummy_hr_data.read(1),\n",
    "                src_transform=dem_original.transform,\n",
    "                src_crs=dem_original.crs,\n",
    "                dst_transform=dummy_hr_data.transform,\n",
    "                dst_crs=dummy_hr_data.crs,\n",
    "                resampling=rasterio.enums.Resampling.bilinear)\n",
    "        dem_reprojected_data = dem_reprojected_data[0,:,:]\n",
    "        dem_reprojected_data = numpy.where(domain_mask==1,dem_reprojected_data,numpy.nan)\n",
    "        dem_reprojected_meta = dummy_hr_data.meta\n",
    "        dem_reprojected_meta.update({\"driver\": \"GTiff\",\"height\": dem_reprojected_data.shape[0],\"width\": dem_reprojected_data.shape[1],\"transform\": dem_reprojected_transform, \"nodata\":numpy.nan})\n",
    "        with rasterio.open(fname_dem, \"w\", **dem_reprojected_meta) as dem_reprojected:\n",
    "            dem_reprojected.write(dem_reprojected_data,1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3405d621-e5a9-44d6-9085-8706327136d4",
   "metadata": {},
   "source": [
    "###### Initialize whitebox environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "90dbebc5-3cea-4897-b146-d1d1fd3f8107",
   "metadata": {},
   "outputs": [],
   "source": [
    "wbe = whitebox_workflows.WbEnvironment()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb591d88-e137-4331-91a2-c108463bef8b",
   "metadata": {},
   "source": [
    "###### Breach the DEM (minimally invasive alternative to filling the DEM?) (see https://www.whiteboxgeo.com/manual/wbt_book/available_tools/hydrological_analysis.html#BreachDepressionsLeastCost)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "4ebe0b18-a722-4b82-8116-c7553906375a",
   "metadata": {},
   "outputs": [],
   "source": [
    "fname_dem_breached = os.path.join(dirDEM,\"dem_reprojected_breached.tif\")\n",
    "if not os.path.isfile(fname_dem_breached) or OVERWRITE_FLAG:\n",
    "    dem_wbe = wbe.read_raster(fname_dem)\n",
    "    dem_breached_wbe = wbe.breach_depressions_least_cost(dem=dem_wbe)\n",
    "    wbe.write_raster(dem_breached_wbe, fname_dem_breached, compress=False)\n",
    "    del dem_wbe, dem_breached_wbe"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "886463c1-b628-4787-b9cd-9892ee0ae2bf",
   "metadata": {},
   "source": [
    "###### Calculate flow accumulation / watershed area (required for TWI calculation) using d-infinity (https://www.whiteboxgeo.com/manual/wbt_book/available_tools/hydrological_analysis.html#DInfFlowAccumulation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "7702a2b0-3837-4386-bd15-2a61aee45537",
   "metadata": {},
   "outputs": [],
   "source": [
    "fname_acc = os.path.join(dirDEM,\"flow_acc.tif\")\n",
    "if not os.path.isfile(fname_acc) or OVERWRITE_FLAG:\n",
    "    dem_breached_wbe = wbe.read_raster(fname_dem_breached)\n",
    "    acc_wbe = wbe.dinf_flow_accum(dem=dem_breached_wbe,out_type='sca',log_transform=False)                                         \n",
    "    wbe.write_raster(acc_wbe, fname_acc, compress=False)\n",
    "    del dem_breached_wbe, acc_wbe"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28eca0c5-a5c3-4aca-a44a-39b95208846b",
   "metadata": {},
   "source": [
    "###### Calculate slope (required for TWI calculation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "4b718ee3-fb68-4520-8f8a-9644e73bcf34",
   "metadata": {},
   "outputs": [],
   "source": [
    "fname_slope = os.path.join(dirDEM,\"slope.tif\")\n",
    "if not os.path.isfile(fname_slope) or OVERWRITE_FLAG:\n",
    "    dem_wbe = wbe.read_raster(fname_dem_breached)\n",
    "    slp_wbe = wbe.slope(dem=dem_wbe)   \n",
    "    wbe.write_raster(slp_wbe, fname_slope, compress=False)\n",
    "    del dem_wbe, slp_wbe"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf090905-5ae7-4900-bda5-e7a646c02e45",
   "metadata": {},
   "source": [
    "###### Calculate TWI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "0d226f6d-c2b5-4e20-b01b-0bb06c7064d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "fname_twi = os.path.join(dirTWI,\"twi.tif\")\n",
    "if not os.path.isfile(fname_twi) or OVERWRITE_FLAG:\n",
    "    acc_wbe = wbe.read_raster(fname_acc)\n",
    "    slp_wbe = wbe.read_raster(fname_slope)\n",
    "    twi_wbe = wbe.wetness_index(specific_catchment_area=acc_wbe,slope=slp_wbe)\n",
    "    wbe.write_raster(twi_wbe, fname_twi, compress=False)\n",
    "    del acc_wbe, slp_wbe, twi_wbe"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f918bfa1-996d-4d26-a146-97fc9654d6c8",
   "metadata": {},
   "source": [
    "###### Upsample TWI to caculate mean TWI per Parflow grid cell \n",
    "###### TODO: change to moving window average TWI? Each pixel will have a mean TWI for it's unique window? Currently biased for pixels on border of Parflow grid window?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "50312f96-c751-4df2-bf33-19945d758fb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "fname_twi_upsample = os.path.join(dirTWI,\"twi_upsample.tif\")\n",
    "if not os.path.isfile(fname_twi_upsample) or OVERWRITE_FLAG:\n",
    "    with rasterio.open(fname_twi,'r') as twi_dataset:\n",
    "        twi_upsample_data      = twi_dataset.read(out_shape=(twi_dataset.count,int(twi_dataset.height / 100),int(twi_dataset.width / 100)),resampling=rasterio.enums.Resampling.average)\n",
    "        twi_upsample_transform = twi_dataset.transform * twi_dataset.transform.scale((twi_dataset.width / twi_upsample_data.shape[-1]),(twi_dataset.height / twi_upsample_data.shape[-2]))\n",
    "        twi_upsample_meta      = twi_dataset.meta\n",
    "        twi_upsample_meta.update({\"driver\": \"GTiff\",\"height\": twi_upsample_data.shape[1],\"width\": twi_upsample_data.shape[2],\"transform\": twi_upsample_transform})\n",
    "        with rasterio.open(fname_twi_upsample, \"w\", **twi_upsample_meta) as twi_upsample_dataset:\n",
    "            twi_upsample_dataset.write(twi_upsample_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55f3bd7b-8f0d-44af-ab4e-0891693008db",
   "metadata": {},
   "source": [
    "###### Downsample average Parflow grid cell TWI to high resolution ParFlow grid (to enable vectorized calculations of inundated area)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "e6e07adc-42d6-40b9-acf2-98bd4b7619cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "fname_twi_downsample = os.path.join(dirTWI,\"twi_downsample.tif\")\n",
    "if not os.path.isfile(fname_twi_downsample) or OVERWRITE_FLAG:\n",
    "    with rasterio.open(fname_twi_upsample,'r') as twi_upsample_dataset:\n",
    "        twi_downsample_data      = twi_upsample_dataset.read(out_shape=(twi_upsample_dataset.count,int(twi_upsample_dataset.height * 100),int(twi_upsample_dataset.width * 100)),resampling=rasterio.enums.Resampling.nearest)\n",
    "        twi_downsample_transform = twi_upsample_dataset.transform * twi_upsample_dataset.transform.scale((twi_upsample_dataset.width / twi_downsample_data.shape[-1]),(twi_upsample_dataset.height / twi_downsample_data.shape[-2]))\n",
    "        twi_downsample_meta      = twi_upsample_dataset.meta\n",
    "        twi_downsample_meta.update({\"driver\": \"GTiff\",\"height\": twi_downsample_data.shape[1],\"width\": twi_downsample_data.shape[2],\"transform\": twi_downsample_transform})\n",
    "        with rasterio.open(fname_twi_downsample, \"w\", **twi_downsample_meta) as twi_downsample_dataset:\n",
    "            twi_downsample_dataset.write(twi_downsample_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a4db4b6-14f7-4a1e-b82d-4ce7e429e71b",
   "metadata": {},
   "source": [
    "###### Define function to perform inundation calculation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "db08998a-841b-48af-8054-70c0efbcfae0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_indundation(wtd_mean,twi_local,twi_mean,f,domain_mask):\n",
    "    \"\"\"Calculate inundation using TOPMODEL-based equation - see Equation 3 of Zhang et al. (2016) (https://doi.org/10.5194/bg-13-1387-2016)\n",
    "    Arguments:\n",
    "        wtd_mean:    grid of mean water table depth values  (zeta sub m in equation 3 of Zhang et al.) \n",
    "        twi_local:   grid of local TWI values               (lambda sub l in equation 3 of Zhang et al.)\n",
    "        twi_mean:    grid of mean TWI values                (lamda sub m in equation 3 of Zhang et al.)\n",
    "        f:           grid of f parameter values             (f in equation 3 and table 1 of Zhang et al.)\n",
    "        domain_mask: grid of model domain                   (1=domain,0=not domain)\n",
    "    Returns:\n",
    "        wtd_local:   grid of local water table depth values (zeta sub l in equation 3 of Zhang et al.)\n",
    "    \"\"\"\n",
    "    wtd_mean  = wtd_mean*(-1)                                                             # values must be negative so multiply by -1\n",
    "    wtd_local = numpy.where(domain_mask==1,(1/f)*(twi_local-twi_mean)+wtd_mean,numpy.nan) # calculate local water depth using equation 3 from Zhang et al. (2016) where domain_mask=1 (i.e., within the model domain), otherise give a NaN value \n",
    "    wtd_local = numpy.where(wtd_local>=0,1,numpy.nan)                                     # give value of 1 where local water table depth is >= 0 (i.e. at or above the surface), otherwise give a NaN value\n",
    "    return wtd_local"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40639746-df94-4c32-aa4b-2c68b1c700cd",
   "metadata": {},
   "source": [
    "###### Do the inundated area calculations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "d85265ed-0e22-4975-bf23-11b7995550f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "twi_local = rasterio.open(fname_twi,'r').read(1)\n",
    "twi_mean = rasterio.open(fname_twi_downsample,'r').read(1)\n",
    "domain_mask = rasterio.open(fname_domain_mask,'r').read(1)\n",
    "transmissivity_decay_factor = rasterio.open(fname_soil_transmissivity_decay_factor,'r').read(1)\n",
    "dtp = {'bilinear':dirWTDParflowBilinear,'nearest':dirWTDParflowNearest}\n",
    "for idatetime in datetime_dim:\n",
    "    for method, dir_wtd in dtp.items():\n",
    "        fname_wtd_mean = os.path.join(dir_wtd,'wtd_'+idatetime.strftime('%Y%m%d')+'_highresolution_'+method+'.tiff')\n",
    "        fname_output = os.path.join(dirOutput,'inundatedarea_'+idatetime.strftime('%Y%m%d')+'_'+method+'.tiff')\n",
    "        if os.path.isfile(fname_wtd_mean):\n",
    "            with rasterio.open(fname_wtd_mean,'r') as wtd_mean_dataset:\n",
    "                wtd_mean = wtd_mean_dataset.read(1)\n",
    "                wtd_local = calculate_indundation(wtd_mean,twi_local,twi_mean,transmissivity_decay_factor,domain_mask)\n",
    "                with rasterio.open(fname_output, \"w\", **wtd_mean_dataset.meta) as wtd_local_dataset:\n",
    "                    wtd_local_dataset.write(wtd_local,1)\n",
    "del twi_local, twi_mean, domain_mask"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "936cdf1a-3ba9-4595-91a9-d224fcb15b0d",
   "metadata": {},
   "source": [
    "###### Create summary rasters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "f4036f4f-cfa4-41b2-ab28-1d373c1c5b84",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "domain_mask = rasterio.open(fname_domain_mask,'r').read(1)\n",
    "for method in ['bilinear','nearest','cubic']:\n",
    "    count = numpy.where(domain_mask==1,0,numpy.nan)\n",
    "    for idatetime in datetime_dim:\n",
    "        fname_wtd_local = os.path.join(dirOutput,'inundatedarea_'+idatetime.strftime('%Y%m%d')+'_'+method+'.tiff') \n",
    "        if os.path.isfile(fname_wtd_local):\n",
    "            inundation_data = rasterio.open(fname_wtd_local,'r').read(1)\n",
    "            count += numpy.where(inundation_data==1,1,0)\n",
    "    count = numpy.where(count==0,numpy.nan,count)\n",
    "    if numpy.any(count>0):\n",
    "        fname_output = os.path.join(dirOutput,'summary_grid_'+datetime_dim[0].strftime('%Y%m%d')+'_to_'+datetime_dim[len(datetime_dim)-1].strftime('%Y%m%d')+'_'+method+'.tiff')\n",
    "        with rasterio.open(fname_output, \"w\", **rasterio.open(fname_wtd_local,'r').meta) as summary_dataset:\n",
    "            summary_dataset.write(count,1)\n",
    "del domain_mask"
   ]
  }
 ],
 "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.12.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}