plot_data_sideview_v3.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a8147690-cf50-452d-a62e-aa53ecd389fc",
   "metadata": {},
   "source": [
    "# Plotting function test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "34d81763-e3d9-4b7f-ae74-adf6ae6e698f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "from matplotlib import rc_context\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "from matplotlib.patches import Polygon\n",
    "import numpy as np\n",
    "from osgeo import gdal\n",
    "from pathlib import Path\n",
    "import rasterio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "68b51b5c-8cff-4114-9969-61854dfb8cd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "gdal.UseExceptions()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e40ca61-3dc5-4b1f-a9ce-fa0ce9f37cec",
   "metadata": {},
   "source": [
    "### Plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 178,
   "id": "40acdb7e-cc11-4152-b4d8-178c428c9ac4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the location as lon,lat\n",
    "# Bounding box: [-179.5,-50,5,85]\n",
    "lon,lat = (-108.64,34.52) # midwest desert area\n",
    "lon,lat = (-115.42, 51.09) # canmore\n",
    "lon,lat = (-102.51, 25.78) # mexico\n",
    "lon,lat = (-76.59,  18.11) # jamaica\n",
    "lon,lat = (-73.97,  40.78) # manhattan\n",
    "lon,lat = (-95.29,  39.35) # kansas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 179,
   "id": "42a10b8d-5c03-49c3-8dc7-badcb6c6059c",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_values = get_location_data_dict_from_geotiffs(file_paths,lon,lat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 180,
   "id": "f8fd6aab-7c8c-49ee-aefa-219cc05402b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "figure_path = Path('/Users/wmk934/data/perceptual_models/data/geotiff_same_resolution/res_0.01/figures')\n",
    "figure_name = f'res_0-01_lon-{lon}_lat-{lat}_X.png'\n",
    "figure_path.mkdir(exist_ok=True, parents=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 181,
   "id": "e47b746f-4da1-4a70-87ce-f89f873ff5ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_data_for_location(data_values, figure_path,figure_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34b6351c-8ae4-44b5-90bc-c5a74b4b1084",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### File specification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f03045dc-3526-48c3-9461-d46153ebf1b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Files\n",
    "base_path = Path('/Users/wmk934/data/perceptual_models/data/geotiff_same_resolution/res_0.01/')\n",
    "file_paths = [['cluster_num', base_path / 'clustering/res_0-01_mask-MERIT_datasets-29_reruns-10/res_0-01_datasets-29_clusters-15_reruns-10.tif'],              ['climate_fs', base_path / 'worldclim/wc2.1_30s_climate_index_fs_resampled.tif'],\n",
    "              ['climate_im', base_path / 'worldclim/wc2.1_30s_climate_index_im_resampled.tif'],\n",
    "              ['climate_imr', base_path / 'worldclim/wc2.1_30s_climate_index_imr_resampled.tif'],\n",
    "              ['lai_01', base_path / 'lai/2013_2023_01_MOD_Grid_MOD15A2H_Lai_500m_resampled.tif'],\n",
    "              ['lai_02', base_path / 'lai/2013_2023_02_MOD_Grid_MOD15A2H_Lai_500m_resampled.tif'],\n",
    "              ['lai_03', base_path / 'lai/2013_2023_03_MOD_Grid_MOD15A2H_Lai_500m_resampled.tif'],\n",
    "              ['lai_04', base_path / 'lai/2013_2023_04_MOD_Grid_MOD15A2H_Lai_500m_resampled.tif'],\n",
    "              ['lai_05', base_path / 'lai/2013_2023_05_MOD_Grid_MOD15A2H_Lai_500m_resampled.tif'],\n",
    "              ['lai_06', base_path / 'lai/2013_2023_06_MOD_Grid_MOD15A2H_Lai_500m_resampled.tif'],\n",
    "              ['lai_07', base_path / 'lai/2013_2023_07_MOD_Grid_MOD15A2H_Lai_500m_resampled.tif'],\n",
    "              ['lai_08', base_path / 'lai/2013_2023_08_MOD_Grid_MOD15A2H_Lai_500m_resampled.tif'],\n",
    "              ['lai_09', base_path / 'lai/2013_2023_09_MOD_Grid_MOD15A2H_Lai_500m_resampled.tif'],\n",
    "              ['lai_10', base_path / 'lai/2013_2023_10_MOD_Grid_MOD15A2H_Lai_500m_resampled.tif'],\n",
    "              ['lai_11', base_path / 'lai/2013_2023_11_MOD_Grid_MOD15A2H_Lai_500m_resampled.tif'],\n",
    "              ['lai_12', base_path / 'lai/2013_2023_12_MOD_Grid_MOD15A2H_Lai_500m_resampled.tif'],\n",
    "              ['veg_height', base_path / 'forest_height/forest_height_2020_resampled.tif'],\n",
    "              ['land_glclu', base_path / 'glclu2019/glclu2019_map_resampled.tif'],\n",
    "              ['land_igbp', base_path / 'modis_land/2013_2022_mode_MCD12Q1_LC_Type1_resampled.tif'],\n",
    "              ['agriculture', base_path / 'lgrip30/lgrip30_agriculture_resampled.tif'],\n",
    "              ['aspect', base_path / 'merit/merit_hydro_aspect_resampled.tif'],\n",
    "              ['slope', base_path / 'merit/merit_hydro_slope_resampled.tif'],\n",
    "              ['elev', base_path / 'merit/merit_hydro_elv_resampled.tif'],\n",
    "              ['clay1', base_path / 'soilgrids/clay_0-5cm_mean_resampled.tif'],\n",
    "              ['clay2', base_path / 'soilgrids/clay_5-15cm_mean_resampled.tif'],\n",
    "              ['clay3', base_path / 'soilgrids/clay_15-30cm_mean_resampled.tif'],\n",
    "              ['clay4', base_path / 'soilgrids/clay_30-60cm_mean_resampled.tif'],\n",
    "              ['clay5', base_path / 'soilgrids/clay_60-100cm_mean_resampled.tif'],\n",
    "              ['clay6', base_path / 'soilgrids/clay_100-200cm_mean_resampled.tif'],\n",
    "              ['sand1', base_path / 'soilgrids/sand_0-5cm_mean_resampled.tif'],\n",
    "              ['sand2', base_path / 'soilgrids/sand_5-15cm_mean_resampled.tif'],\n",
    "              ['sand3', base_path / 'soilgrids/sand_15-30cm_mean_resampled.tif'],\n",
    "              ['sand4', base_path / 'soilgrids/sand_30-60cm_mean_resampled.tif'],\n",
    "              ['sand5', base_path / 'soilgrids/sand_60-100cm_mean_resampled.tif'],\n",
    "              ['sand6', base_path / 'soilgrids/sand_100-200cm_mean_resampled.tif'],\n",
    "              ['silt1', base_path / 'soilgrids/silt_0-5cm_mean_resampled.tif'],\n",
    "              ['silt2', base_path / 'soilgrids/silt_5-15cm_mean_resampled.tif'],\n",
    "              ['silt3', base_path / 'soilgrids/silt_15-30cm_mean_resampled.tif'],\n",
    "              ['silt4', base_path / 'soilgrids/silt_30-60cm_mean_resampled.tif'],\n",
    "              ['silt5', base_path / 'soilgrids/silt_60-100cm_mean_resampled.tif'],\n",
    "              ['silt6', base_path / 'soilgrids/silt_100-200cm_mean_resampled.tif'],\n",
    "              ['soil_depth', base_path / 'pelletier/average_soil_and_sedimentary-deposit_thickness_resampled.tif']\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7557ba1a-69ee-45e9-8dfc-e01bc59fedc0",
   "metadata": {},
   "source": [
    "### Plotting functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "id": "ec86a76f-faf1-4828-ba0e-ab690a454c24",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_data_for_location(data_values, figure_path,figure_name):\n",
    "\n",
    "    # Make the 3 plots\n",
    "    plot_climate(data_values,  figure_path / figure_name.replace('X.png','climate.png'))\n",
    "    plot_lai(data_values,      figure_path / figure_name.replace('X.png','lai.png'))\n",
    "    plot_overview(data_values, figure_path / figure_name.replace('X.png','overview.png'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 155,
   "id": "e85271e3-063c-4fdf-9b92-3c6bdf764800",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_overview(data_values, save_file, xkcd=False):\n",
    "\n",
    "    def make_plot(data_values, save_file):\n",
    "    \n",
    "        # Data\n",
    "        vegetation_height    = data_values['veg_height'] # [m]\n",
    "        land_type1           = data_values['land_glclu'] # Categorical\n",
    "        land_type1_text      = glclu2019_dict[land_type1]\n",
    "        land_type2           = data_values['land_igbp'] # Categorical\n",
    "        land_type2_text      = igbp_dict[land_type2]\n",
    "        land_agriculture     = 1 # Categorical. 0: water, 1: non-cropland, 2: irrigated, 3: rainfed\n",
    "        land_agriculture_txt = lgrip30_dict[land_agriculture]\n",
    "        vegetation_lai       = np.array([data_values['lai_01'],data_values['lai_02'],data_values['lai_03'],\n",
    "                                         data_values['lai_04'],data_values['lai_05'],data_values['lai_06'],\n",
    "                                         data_values['lai_07'],data_values['lai_08'],data_values['lai_09'],\n",
    "                                         data_values['lai_10'],data_values['lai_11'],data_values['lai_12']])\n",
    "        topography_elevation = data_values['elev'] # [m.a.s.l.]\n",
    "        topography_slope     = math.tan(math.radians(data_values['slope'])) # [degrees] to [m m-1], derived\n",
    "        topography_aspect    = data_values['aspect'] # [degrees], derived\n",
    "        soil_depth           = data_values['soil_depth'] # [m]\n",
    "        soil_layer_top       = [0.00, 0.05, 0.15, 0.30, 0.60, 1.00] # Top of layer depths [m]\n",
    "        soil_layer_bot       = [0.05, 0.15, 0.30, 0.60, 1.00, 2.00] # Bottom of layer depths in [m]\n",
    "        soil_sand            = [data_values['sand1']/1000, data_values['sand2']/1000, data_values['sand3']/1000,\n",
    "                                data_values['sand4']/1000, data_values['sand5']/1000, data_values['sand6']/1000]\n",
    "        soil_silt            = [data_values['silt1']/1000, data_values['silt2']/1000, data_values['silt3']/1000,\n",
    "                                data_values['silt4']/1000, data_values['silt5']/1000, data_values['silt6']/1000]\n",
    "        soil_clay            = [data_values['clay1']/1000, data_values['clay2']/1000, data_values['clay3']/1000,\n",
    "                                data_values['clay4']/1000, data_values['clay5']/1000, data_values['clay6']/1000]\n",
    "        climate_aridity      = data_values['climate_im']\n",
    "        climate_seasonality  = data_values['climate_imr']\n",
    "        climate_snow         = data_values['climate_fs']\n",
    "        geology = 0 # Glhymps \n",
    "        lake = 0 # # Hydrolakes\n",
    "    \n",
    "        # General settings\n",
    "        text_settings = {'font.family': 'Courier New', 'font.size': 12}\n",
    "        \n",
    "        # Determine the spacing\n",
    "        bedrock_depth_factor = 1.05\n",
    "        bedrock_depth = soil_depth * bedrock_depth_factor\n",
    "        total_width  = 10 # [m]\n",
    "        total_height = vegetation_height + soil_depth + bedrock_depth # [m]\n",
    "        top = max(1,vegetation_height) # these max/min values create some space for the plot in case \n",
    "        bottom = min(-1,0-bedrock_depth) #   veg is small and/or soil is shallow\n",
    "        slope_l = 0.5*(topography_slope * total_width)\n",
    "        slope_r = -0.5*(topography_slope * total_width)\n",
    "        soil_depth_l = soil_depth + slope_l\n",
    "        soil_depth_r = soil_depth + slope_r\n",
    "                   \n",
    "        # General figure components\n",
    "        fig = plt.figure(figsize=(4,6))\n",
    "        ax = fig.add_axes((0.1, 0.1, 0.7, 0.7)) # left, bottom, width, height   \n",
    "        \n",
    "        # Specifics of the main figure\n",
    "        #ax.set_aspect('equal') \n",
    "        ax.spines[['top', 'right', 'bottom']].set_visible(False)\n",
    "        ax.set_xticks([])\n",
    "        ax.set_yticks([0-soil_depth,0,top])\n",
    "        ax.minorticks_on()\n",
    "        ax.tick_params(axis='y', which='minor', length=4)\n",
    "        ax.set_yticks(np.arange(0-soil_depth,top), minor=True)\n",
    "        plt.ylim([bottom, top])\n",
    "        \n",
    "        # Land use\n",
    "        # We'll primarily use IGBP categories here (land_type2)\n",
    "        land_patch = get_igbp_patch(land_type2, vegetation_height, slope_l, slope_r)\n",
    "        if land_patch:  # \"Barren\" returns a None type land_patch, which we skip with this IF-statement\n",
    "            for veg_patch in land_patch:\n",
    "                ax.add_patch(veg_patch)\n",
    "        \n",
    "        # Bedrock\n",
    "        n_bedrock_shading = 10\n",
    "        for x in np.linspace(0,total_width,n_bedrock_shading, endpoint=False):\n",
    "            plt.plot([x,x+total_width/n_bedrock_shading],\n",
    "                     [bottom,0-soil_depth], \n",
    "                     color=get_pm_plot_color('bedrock fracture'))\n",
    "        plt.plot([0,total_width],[0-soil_depth,0-soil_depth], color=get_pm_plot_color('bedrock boundary'))\n",
    "        \n",
    "        # Surface\n",
    "        #plt.plot([0,total_width],[slope_l,slope_r], color=get_pm_plot_color('surface'))\n",
    "        \n",
    "        # Soil\n",
    "        soil_text = 'Soil layer texture [%]\\ndepth [m]   sand silt clay'\n",
    "        for ix,(soil_t,soil_b) in enumerate(zip(soil_layer_top,soil_layer_bot)): # SOILGRIDS layer loop\n",
    "            if soil_depth >= soil_t: # the top of this soil layer is within estimated soil depth\n",
    "        \n",
    "                # Get the soil type in this layer\n",
    "                sand = soil_sand[ix]\n",
    "                silt = soil_silt[ix]\n",
    "                clay = soil_clay[ix]\n",
    "                _,soil_name = find_usda_soilclass(sand*100,silt*100,clay*100)\n",
    "                soil_color = get_pm_soil_color(soil_name)\n",
    "        \n",
    "                # Check where the bottom of the plotting area is\n",
    "                if soil_depth < soil_b: # Estimated depth ends before this soilgrids layer does\n",
    "                    soil_b = soil_depth # Therefore end the plotting layer at estimated soil depth\n",
    "        \n",
    "                # Handle the special case where estimated soil depth is deeper than deepest soil grid layer\n",
    "                if soil_b == 2: \n",
    "                    soil_b = soil_depth # extend this layer all the way to bedrock\n",
    "        \n",
    "                # Calculate at which fraction of total soil depth the layer interfaces are\n",
    "                soil_t_frac = soil_t / soil_depth\n",
    "                soil_b_frac = soil_b / soil_depth\n",
    "        \n",
    "                # Create a Polygon patch with the specified characteristics\n",
    "                x = [total_width, total_width, 0, 0, total_width]\n",
    "                y = [slope_r - soil_b_frac * soil_depth_r, slope_r - soil_t_frac * soil_depth_r,\n",
    "                     slope_l - soil_t_frac * soil_depth_l, slope_l - soil_b_frac * soil_depth_l,\n",
    "                     slope_r - soil_b_frac * soil_depth_r]\n",
    "                soil_patch = patches.Polygon(\n",
    "                    xy=list(zip(x, y)),  # Combine x and y coordinates\n",
    "                    closed=True,  # Connect the last point to the first to close the shape\n",
    "                    facecolor=soil_color,\n",
    "                    edgecolor='none'#,\n",
    "                    #linestyle='dotted',\n",
    "                    #linewidth=1.0  # Width of the outline\n",
    "                )  \n",
    "                ax.add_patch(soil_patch)\n",
    "                \n",
    "                # Append the annotation before we start to update soil_b\n",
    "                b_decimals = 3-len(str(int(soil_b)))\n",
    "                soil_text += f'\\n{soil_t:3.2f}-{soil_b:3.{b_decimals}f}   {sand:.2f} {silt:.2f} {clay:.2f}'\n",
    "           \n",
    "        # Text\n",
    "        clus_num = data_values['cluster_num']\n",
    "\n",
    "        if climate_aridity < 0:\n",
    "            clim_text = (f'Aridity       [-1,1]:{climate_aridity:3.2f}\\n'\n",
    "                         f'Seasonality    [0,2]: {climate_seasonality:3.2f}\\n'\n",
    "                         f'Fraction snow  [0,1]: {climate_snow:3.2f}')\n",
    "        else:\n",
    "            clim_text = (f'Aridity       [-1,1]: {climate_aridity:3.2f}\\n'\n",
    "                         f'Seasonality    [0,2]: {climate_seasonality:3.2f}\\n'\n",
    "                         f'Fraction snow  [0,1]: {climate_snow:3.2f}')\n",
    "        land_text = (f'Land class 1: {land_type2_text}\\n'\n",
    "                     f'Land class 2: {land_type1_text}\\n'\n",
    "                     f'Agriculture : {land_agriculture_txt}')\n",
    "        topo_text = (f'Elevation: {topography_elevation:4.1f} [m.a.s.l]\\n'\n",
    "                     f'Slope    : {topography_slope:4.3f} [m m-1]\\n'\n",
    "                     f'Aspect   : {topography_aspect:4.1f} [degrees]')\n",
    "        all_text = (f'Cluster {clus_num}\\n'\n",
    "                    f'Coordinates ({lon:5.2f},{lat:5.2f})\\n\\n'\n",
    "                    f'{clim_text}\\n\\n'\n",
    "                    f'{land_text}\\n\\n'\n",
    "                    f'{topo_text}\\n\\n'\n",
    "                    f'{soil_text}')\n",
    "        with rc_context(text_settings):\n",
    "            ax.text(total_width+1, 0-soil_depth, all_text, ha='left', va='bottom')\n",
    "        plt.savefig(save_file, dpi=300, bbox_inches='tight')\n",
    "        #plt.show()\n",
    "        plt.close()\n",
    "        return\n",
    "        \n",
    "    if xkcd:\n",
    "        with plt.xkcd(): \n",
    "            make_plot(data_values,save_file)\n",
    "    else: \n",
    "        make_plot(data_values,save_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 135,
   "id": "a765310f-cc0c-4784-b4c7-7b800ce6acec",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_climate(data_values, save_file, xkcd=False):\n",
    "\n",
    "    def make_plot(data_values, save_file):\n",
    "        c_im = data_values['climate_im']\n",
    "        c_imr = data_values['climate_imr']\n",
    "        c_fs = data_values['climate_fs']\n",
    "    \n",
    "        fig = plt.figure(figsize=(4,4))\n",
    "        ax = fig.add_axes((0.4, 0.4, 0.2, 0.2))\n",
    "        ax.spines[['top', 'right', 'left', 'bottom']].set_visible(False)\n",
    "        ax.set_ylim([0.5, 4])\n",
    "        ax.set_xlim([0, 2])\n",
    "        ax.set_xticks([0,2])\n",
    "        ax.set_yticks([1,2,3])\n",
    "        ax.set_xticklabels(['less','more'])\n",
    "        ax.set_yticklabels(['fraction snow','seasonality','aridity'])\n",
    "        ax.tick_params(axis='both', length=0)\n",
    "        \n",
    "        ax.plot([0,2],[1,1], color='blue', marker='D')\n",
    "        ax.plot([0,2],[2,2], color='green', marker='D')\n",
    "        ax.plot([0,2],[3,3], color='red', marker='D') \n",
    "        ax.plot([c_fs*2,c_fs*2],[1,1], color='black', marker='D') # scale [0,1] to [0,2]\n",
    "        ax.plot([c_imr,c_imr],[2,2],   color='black', marker='D') # no scaling, already on [0,2]\n",
    "        ax.plot([c_im+1,c_im+1],[3,3], color='black', marker='D') # scale [-1,1] to [0,2]\n",
    "        plt.savefig(save_file, dpi=300, bbox_inches='tight')\n",
    "        #plt.show()   \n",
    "        plt.close()\n",
    "        return\n",
    "\n",
    "    if xkcd:\n",
    "        with plt.xkcd(): \n",
    "            make_plot(data_values,save_file)\n",
    "    else: \n",
    "        make_plot(data_values,save_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 136,
   "id": "8c604d2a-e78f-4ff7-a47c-d368d2e99d56",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_lai(data_values, save_file, xkcd=False):\n",
    "\n",
    "    def make_plot(data_values,save_file):\n",
    "\n",
    "        vegetation_lai = np.array([data_values['lai_01'],data_values['lai_02'],data_values['lai_03'],\n",
    "                                   data_values['lai_04'],data_values['lai_05'],data_values['lai_06'],\n",
    "                                   data_values['lai_07'],data_values['lai_08'],data_values['lai_09'],\n",
    "                                   data_values['lai_10'],data_values['lai_11'],data_values['lai_12']])\n",
    "        \n",
    "        fig = plt.figure(figsize=(4,4))\n",
    "        ax = fig.add_axes((0.4, 0.4, 0.2, 0.2))\n",
    "        ax.spines[['top', 'right']].set_visible(False)\n",
    "        \n",
    "        ax.plot(vegetation_lai, color='green')\n",
    "        ax.set_ylabel('LAI [-]') #ax.set_ylabel('LAI [m2 m-2]')\n",
    "        ax.set_xticks([0,1,2,3,4,5,6,7,8,9,10,11])\n",
    "        ax.set_xticklabels(['J','','','','','','','','','','','D'])\n",
    "        plt.savefig(save_file, dpi=300, bbox_inches='tight')\n",
    "        #plt.show()\n",
    "        plt.close()\n",
    "        return\n",
    "\n",
    "    if xkcd:\n",
    "        with plt.xkcd(): \n",
    "            make_plot(data_values,save_file)\n",
    "    else: \n",
    "        make_plot(data_values,save_file)\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bf1d1ff-b1b1-48cc-b351-d4f32f833e6d",
   "metadata": {},
   "source": [
    "### Land class functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "bb74cce1-fb5e-45a9-aa19-7c36fbca75e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_rectangle(bottom_left, width, height):\n",
    "    vertices = [bottom_left,\n",
    "                (bottom_left[0] + width, bottom_left[1]),\n",
    "                (bottom_left[0] + width, bottom_left[1] + height),\n",
    "                (bottom_left[0], bottom_left[1] + height)]\n",
    "    return vertices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "9ca2d258-1600-408f-914d-7ceea109619e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_circle(center, radius, num_points=100):\n",
    "    angles = np.linspace(0, 2 * np.pi, num_points)\n",
    "    x = center[0] + radius * np.cos(angles)\n",
    "    y = center[1] + radius * np.sin(angles)\n",
    "    return np.column_stack((x, y))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "53cbd447-7318-4fd2-a39b-9fdb2a51fe04",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_ellipse(center, a_radius, b_radius, num_points=100):\n",
    "    angles = np.linspace(0, 2 * np.pi, num_points)\n",
    "    x = center[0] + a_radius * np.cos(angles)\n",
    "    y = center[1] + b_radius * np.sin(angles)\n",
    "    return np.column_stack((x, y))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "3a4dbe51-3194-4c12-9c42-d087a471f3d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_triangle(center, bottom, width, height):\n",
    "    vertices = [(center-0.5*width, bottom),\n",
    "                (center+0.5*width, bottom),\n",
    "                (center, height)]\n",
    "    return vertices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6e453b71-f335-4978-b579-8ba8491ff9cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_needle_tree(x,y,width,height):\n",
    "    tree_patches = [Polygon(create_rectangle((x-0.5*width,y),width,height*0.3), \n",
    "                            closed=True, \n",
    "                            edgecolor='black', \n",
    "                            facecolor=get_pm_plot_color('coniferous_trunk')),\n",
    "                    Polygon(create_triangle(x,height*0.2,height*0.25,height), \n",
    "                            closed=True, \n",
    "                            edgecolor='black', \n",
    "                            facecolor=get_pm_plot_color('coniferous_leafs'))]\n",
    "    return tree_patches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "36dc827b-b617-41ba-b46e-0694dbb96b53",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_leaf_tree(x,y,width,height):\n",
    "    tree_patches = [Polygon(create_rectangle((x-0.5*width,y),width,height*0.8), \n",
    "                            closed=True, \n",
    "                            edgecolor='black', \n",
    "                            facecolor=get_pm_plot_color('deciduous_trunk')),\n",
    "                    Polygon(create_circle((x,height*0.8), height*0.2), \n",
    "                            closed=True, \n",
    "                            edgecolor='black', \n",
    "                            facecolor=get_pm_plot_color('deciduous_leafs'))]\n",
    "    return tree_patches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "fc1b3030-d348-4fcb-a3ca-32c040787ac7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_shrub(x,y,width,height):\n",
    "    shrub_patches = Polygon(create_circle((x,height*0.1), height*0.9), \n",
    "                            closed=True, \n",
    "                            edgecolor='black', \n",
    "                            facecolor=get_pm_plot_color('shrub_leafs'))\n",
    "    return shrub_patches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "148114a3-95b3-47f6-88a2-12c5785e3eb8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_grass(lx,ly,rx,ry,height):\n",
    "    grass_patch = Polygon([(lx,ly),(rx,ry),(rx,ry+height),(lx,ly+height)],\n",
    "                          closed=True,\n",
    "                          edgecolor=get_pm_plot_color('grass_edge'),\n",
    "                          facecolor=get_pm_plot_color('grass'))\n",
    "    return grass_patch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "ab2912b4-4fd8-4034-a594-f73d032ea9c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_water(lx,ly,rx,ry,height=0):\n",
    "\n",
    "    # Check if we need to add some polygon height in case of small slopes\n",
    "    if (ly-ry) < 0.1: lx += 0.1\n",
    "    water_patch = Polygon([(lx,ly),(rx,ry),(rx,ly+height),(lx,ly+height)], # level\n",
    "                          closed=True,\n",
    "                          edgecolor=get_pm_plot_color('water_edge'),\n",
    "                          facecolor=get_pm_plot_color('water'),\n",
    "                          alpha=0.5)\n",
    "    return water_patch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "8946b6d7-ceb5-4422-b4c9-9a5ff7bbb060",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_snow(lx,ly,rx,ry,height=1):\n",
    "    snow_patch = Polygon([(lx,ly),(rx,ry),(rx,ly+height),(lx,ly+height)], # level\n",
    "                          closed=True,\n",
    "                          edgecolor='black',\n",
    "                          facecolor=get_pm_plot_color('snow'))\n",
    "    return snow_patch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "b7897bd2-8e46-40fa-b2e1-33c75d540fc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_urban(lx,ly,rx,ry,height=5):\n",
    "    urban_patch = [Polygon([(lx+2,ly),(rx-2,ry),(rx-2,ly+height),(lx+2,ly+height)],\n",
    "                           closed=True,\n",
    "                           edgecolor='black',\n",
    "                           facecolor=get_pm_plot_color('house')),\n",
    "                    Polygon([(lx,ly),(rx,ry),(rx,ry+0.1),(lx,ly+0.1)],\n",
    "                           closed=True,\n",
    "                           edgecolor='black',\n",
    "                           facecolor=get_pm_plot_color('road'))]\n",
    "    return urban_patch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "5bb1521f-b86d-43e8-bf9f-eb5bd500935c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_crop(x,y,height):\n",
    "    crop_patches = [Polygon([(x,y),(x,y+height*0.8)],\n",
    "                            closed=None,\n",
    "                            edgecolor=get_pm_plot_color('crop_edge')),\n",
    "                    Polygon(create_ellipse((x,y+height*0.8), height*0.1, height*0.2), \n",
    "                            closed=True, \n",
    "                            edgecolor=get_pm_plot_color('crop_edge'), \n",
    "                            facecolor=get_pm_plot_color('crop'))]\n",
    "    return crop_patches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 138,
   "id": "6ec4ab7c-37b3-4ff5-8fe7-68b6af5e7dc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_igbp_patch(land_type, veg_height, slope_l, slope_r):\n",
    "\n",
    "    # We know we have an x-dimension of size 10 to work with\n",
    "    # Y-dimension is dictated by veg_height\n",
    "    # We'll plot everything according to the slope\n",
    "\n",
    "    # find the y-coordinates\n",
    "    y = slope_l - (slope_l - slope_r) * np.arange(0,11)/10\n",
    "\n",
    "    # Needleleaf forest (>60% trees)\n",
    "    if (land_type == 1) or (land_type == 3):\n",
    "        land_patches = [patch for sublist in [make_needle_tree(2,y[2],0.5,veg_height), \n",
    "                                              make_needle_tree(8,y[8],0.5,veg_height), \n",
    "                                              make_needle_tree(5,y[5],0.5,veg_height)] \n",
    "                              for patch in sublist]        \n",
    "\n",
    "    # Broadleaf forest (>60% trees)\n",
    "    elif (land_type == 2) or (land_type == 4):\n",
    "        land_patches = [patch for sublist in [make_leaf_tree(2,y[2],0.5,veg_height), \n",
    "                                              make_leaf_tree(8,y[8],0.5,veg_height), \n",
    "                                              make_leaf_tree(5,y[6],0.5,veg_height)] \n",
    "                              for patch in sublist]\n",
    "\n",
    "    # Mixed forest (>60% trees)\n",
    "    elif land_type == 5:\n",
    "        land_patches = [patch for sublist in [make_needle_tree(2,y[2],0.5,veg_height), \n",
    "                                              make_needle_tree(8,y[8],0.5,veg_height), \n",
    "                                              make_leaf_tree(5,y[5],0.5,veg_height)] \n",
    "                              for patch in sublist]\n",
    "\n",
    "    # Closed shrubland (>60% shrubs)\n",
    "    elif land_type == 6:\n",
    "        land_patches = [make_shrub(2,y[2],0.4,veg_height), make_shrub(8,y[8],0.4,veg_height), \n",
    "                        make_shrub(4,y[4],0.4,veg_height), make_shrub(3,y[3],0.4,veg_height), \n",
    "                        make_shrub(5,y[5],0.4,veg_height), make_shrub(7,y[7],0.4,veg_height)]\n",
    "        \n",
    "    # Open shrubland (10-60% shrubs)\n",
    "    elif land_type == 7:\n",
    "        land_patches = [make_shrub(2,y[2],0.3,veg_height), make_shrub(8,y[8],0.3,veg_height), \n",
    "                        make_shrub(7,y[7],0.3,veg_height)]\n",
    "        \n",
    "    # Woody savanna (30-60% trees)\n",
    "    elif land_type == 8:\n",
    "        # randomly select 2m as the threshold for shrubs\n",
    "        if veg_height > 2: \n",
    "            shrub_h = 2\n",
    "            tree_h = veg_height\n",
    "        else:\n",
    "            shrub_h = veg_height\n",
    "            tree_h = 10 # doesn't matter, it'll get cut off on the plot anyway\n",
    "        land_patches = [make_leaf_tree(2,y[2],0.5,tree_h)[0], # trunk\n",
    "                        make_leaf_tree(2,y[2],0.5,tree_h)[1], # leafs\n",
    "                        make_shrub(3,y[3],0.4,shrub_h), make_shrub(7,y[7],0.4,shrub_h),\n",
    "                        make_leaf_tree(8,y[8],0.5,tree_h)[0], # trunk\n",
    "                        make_leaf_tree(8,y[8],0.5,tree_h)[1]] # leafs]\n",
    "        \n",
    "    # Savanna (10-30% trees)\n",
    "    elif land_type == 9:\n",
    "        # randomly select 2m as the threshold for shrubs\n",
    "        if veg_height > 2: \n",
    "            shrub_h = 2\n",
    "            tree_h = veg_height\n",
    "        else:\n",
    "            shrub_h = veg_height\n",
    "            tree_h = 10 # doesn't matter, it'll get cut off on the plot anyway\n",
    "        land_patches = [make_leaf_tree(2,y[2],0.5,tree_h)[0], # trunk\n",
    "                        make_leaf_tree(2,y[2],0.5,tree_h)[1], # leafs\n",
    "                        make_shrub(3,y[3],0.4,shrub_h), make_shrub(8,y[8],0.4,shrub_h),\n",
    "                        make_shrub(5,y[5],0.4,shrub_h),]\n",
    "        \n",
    "    # Grassland \n",
    "    elif land_type == 10:\n",
    "        if veg_height < 0.1: veg_height = 0.1 # so we at least have something\n",
    "        land_patches = [make_grass(0,slope_l,10,slope_r,veg_height)]\n",
    "        \n",
    "    # Permanent wetland\n",
    "    elif land_type == 11:\n",
    "        if veg_height < 0.2: veg_height = 0.2\n",
    "        land_patches = [make_leaf_tree(4,0,0.5,vh)[0], # trunk\n",
    "                        make_leaf_tree(4,0,0.5,vh)[1], # leafs\n",
    "                        make_water(0,sl,10,sr)]\n",
    "        \n",
    "    # Cropland\n",
    "    elif land_type == 12:\n",
    "        x_crops = np.arange(0.5,10,0.5) # Create narrower x-spacing\n",
    "        y_crops = slope_l - (slope_l - slope_r) * x_crops/10\n",
    "        land_patches = [patch for sublist in \n",
    "                        [make_crop(x_crop,y_crop,veg_height) for x_crop,y_crop in zip(x_crops,y_crops)]\n",
    "                              for patch in sublist]\n",
    "        \n",
    "    # Urban\n",
    "    elif land_type == 13:\n",
    "        land_patches = make_urban(0,sl,10,sr)\n",
    "        \n",
    "    # Cropland/natural mosaic\n",
    "    elif land_type == 14:\n",
    "        # We don't know what the satellite vegetation height is here, so specify something\n",
    "        tree_height = 10\n",
    "        shrub_height = 2\n",
    "        crop_height = 1\n",
    "        x_crops = np.arange(3.5,9.5,0.5) # Create narrower x-spacing\n",
    "        y_crops = slope_l - (slope_l - slope_r) * x_crops/10\n",
    "        tree = make_leaf_tree(3,y[3],0.5,tree_height)\n",
    "        crops = [patch for sublist in \n",
    "                        [make_crop(x_crop,y_crop,crop_height) for x_crop,y_crop in zip(x_crops,y_crops)]\n",
    "                              for patch in sublist]\n",
    "        land_patches = tree + crops + [make_shrub(2,y[2],0.4,shrub_height)]\n",
    "        \n",
    "    # Snow/ice\n",
    "    elif land_type == 15:\n",
    "        land_patches = [make_snow(0,sl,10,sr)]\n",
    "        \n",
    "    # Barren\n",
    "    elif land_type == 16:\n",
    "        land_patches = None\n",
    "\n",
    "    # Water\n",
    "    elif land_type == 17:\n",
    "        land_patches = [make_water(0,sl,10,sr,1)]\n",
    "        \n",
    "    return land_patches"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ce3007e-b514-49a9-abce-538bff0184af",
   "metadata": {},
   "source": [
    "### Plotting Support"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "00d91445-b94c-4dec-ab83-ac61ef43996e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_pm_plot_color(object):\n",
    "\n",
    "    '''Defines standard colors for objects in perceptual model plots'''\n",
    "\n",
    "    if object == 'bedrock boundary': color = [0.3, 0.3, 0.3] # Dark-ish grey\n",
    "    elif object == 'bedrock fracture': color = [0.7, 0.7, 0.7] # Light-ish grey\n",
    "    elif object == 'surface': color = 'saddlebrown' # Brown\n",
    "    elif object == 'soil':\n",
    "        print('WARNING: soil is defined based on sand/silt/clay percentages and requires get_pm_soil_color()')\n",
    "        color = None\n",
    "    elif object == 'deciduous_trunk': color = 'peru'\n",
    "    elif object == 'deciduous_leafs': color = 'forestgreen'\n",
    "    elif object == 'coniferous_trunk': color = 'saddlebrown'\n",
    "    elif object == 'coniferous_leafs': color = 'darkgreen'\n",
    "    elif object == 'shrub_leafs': color = 'green'\n",
    "    elif object == 'grass': color ='lightgreen'\n",
    "    elif object == 'grass_edge': color ='lightgreen'\n",
    "    elif object == 'water': color ='royalblue'\n",
    "    elif object == 'water_edge': color ='lightsteelblue'\n",
    "    elif object == 'snow': color = 'ghostwhite'\n",
    "    elif object == 'road': color = 'dimgrey'\n",
    "    elif object == 'house': color = 'darkgrey'\n",
    "    elif object == 'crop': color = 'gold'\n",
    "    elif object == 'crop_edge': color = 'goldenrod'\n",
    "        \n",
    "    return color"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "da1474bf-0bc1-494c-aa5a-f65059e556ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "lgrip30_dict = dict([(0, 'water'),\n",
    "                     (1, 'non-cropland'),\n",
    "                     (2, 'irrigated cropland'),\n",
    "                     (3, 'rain-fed cropland')])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "7dc09aca-9622-412f-80cb-f5baa47d53c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "glclu2019_dict = dict([(1,  'true desert'),\n",
    "                       (2,  'semi-arid'),\n",
    "                       (3,  'dense short vegetation'),\n",
    "                       (4,  'open tree cover'),\n",
    "                       (5,  'dense tree cover'),\n",
    "                       (6,  'tree cover gain'),\n",
    "                       (7,  'tree cover loss, not fire'),\n",
    "                       (8,  'salt pan'),\n",
    "                       (9,  'wetland sparse vegetation'),\n",
    "                       (10, 'wetland dense short vegetation'),\n",
    "                       (11, 'wetland open tree cover'),\n",
    "                       (12, 'wetland dense tree cover'),\n",
    "                       (13, 'wetland tree cover gain'),\n",
    "                       (14, 'wetland tree cover loss, not fire'),\n",
    "                       (15, 'ice'),\n",
    "                       (16, 'water'),\n",
    "                       (17, 'cropland'),\n",
    "                       (18, 'built-up'),\n",
    "                       (19, 'ocean'),\n",
    "                       (20, 'no data')])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "6707ba6b-e465-45c7-b572-278ca979d5ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "igbp_dict = dict([(1,  'evergreen needleleaf forest'),\n",
    "                  (2,  'evergreen broadleaf forest'),\n",
    "                  (3,  'deciduous needleleaf forest'),\n",
    "                  (4,  'deciduous broadleaf forest'),\n",
    "                  (5,  'mixed forest'),\n",
    "                  (6,  'closed shrublands'),\n",
    "                  (7,  'open shrublands'),\n",
    "                  (8,  'woody savannas'),\n",
    "                  (9,  'savannas'),\n",
    "                  (10, 'grasslands'),\n",
    "                  (11, 'permanent wetlands'),\n",
    "                  (12, 'croplands'),\n",
    "                  (13, 'urban and built-up land'),\n",
    "                  (14, 'cropland/natural mosaic'),\n",
    "                  (15, 'permanent snow/ice'),\n",
    "                  (16, 'barren'),\n",
    "                  (17, 'water'),\n",
    "                  (255,'unclassified')])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "3c832568-717d-4fa0-a167-8fad4b892b0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_usda_soilclass(sand,silt,clay,no_data_value=255):\n",
    "\n",
    "    # Adapted from: Knoben (2021)\n",
    "    # Knoben, W. J. M. (2021). Global USDA-NRCS soil texture class map, \n",
    "    # HydroShare, https://doi.org/10.4211/hs.1361509511e44adfba814f6950c6e742\n",
    "    \n",
    "    # Based on Benham et al., 2009 and matching the following soil class table:\n",
    "    # SUMMA-ROSETTA-STAS-RUC soil parameter table\n",
    "    # 1  'CLAY' \n",
    "    # 2  'CLAY LOAM'\n",
    "    # 3  'LOAM' \n",
    "    # 4  'LOAMY SAND'\n",
    "    # 5  'SAND'\n",
    "    # 6  'SANDY CLAY'\n",
    "    # 7  'SANDY CLAY LOAM'\n",
    "    # 8  'SANDY LOAM'\n",
    "    # 9  'SILT'\n",
    "    # 10 'SILTY CLAY'\n",
    "    # 11 'SILTY CLAY LOAM'\n",
    "    # 12 'SILT LOAM'\n",
    "\n",
    "    # Check inputs\n",
    "    if (sand<=1) & (silt<=1) & (clay<=1):\n",
    "        print(f'WARNING: find_usda_soilclass(): sand, silt and clay all < 1. Please provide inputs in % [0,100].')\n",
    "        return\n",
    "    \n",
    "    # Legend\n",
    "    soiltype = ['clay','clay loam','loam','loamy sand','sand','sandy clay','sandy clay loam',\n",
    "                'sandy loam','silt','silty clay','silty clay loam','silt loam']\n",
    "    \n",
    "    # Classify\n",
    "    if (clay >= 40) & (sand <= 45) & (silt < 40):\n",
    "        soilclass = 1\n",
    "    elif (clay >= 27) & (clay < 40) & (sand > 20) & (sand <= 45):\n",
    "        soilclass = 2\n",
    "    elif (clay >= 7) & (clay < 27) & (silt >= 28) & (silt < 50) & (sand < 52):\n",
    "        soilclass = 3\n",
    "    elif ((silt + 1.5 * clay) >= 15) & ((silt + 2* clay) < 30):\n",
    "        soilclass = 4\n",
    "    elif ((silt + 1.5 * clay) < 15):\n",
    "        soilclass = 5\n",
    "    elif (clay >= 35 )& (sand > 45):\n",
    "        soilclass = 6\n",
    "    elif (clay >= 20) & (clay < 35) & (silt < 28) & (sand > 45):\n",
    "        soilclass = 7\n",
    "    elif ((clay >= 7) & (clay < 20) & (sand > 52) & ((silt+2*clay) >= 30)) | ((clay < 7) & (silt < 50) & ((silt+2*clay >= 30))):\n",
    "        soilclass = 8\n",
    "    elif (silt >= 80) & (clay < 12):\n",
    "        soilclass = 9\n",
    "    elif (clay >= 40) & (silt >= 40):\n",
    "        soilclass = 10\n",
    "    elif (clay >= 27) & (clay < 40) & (sand <= 20):\n",
    "        soilclass = 11\n",
    "    elif ((silt >= 50) & (clay >= 12) & (clay < 27)) | ((silt >= 50) & (silt < 80) & (clay < 12)):\n",
    "        soilclass = 12\n",
    "    elif (sand == no_data_value) | (silt == no_data_value) | (clay == no_data_value):\n",
    "        soilclass = 0\n",
    "\n",
    "    soilname = 'undefined'\n",
    "    if soilclass > 0:\n",
    "        soilname = soiltype[soilclass-1] # -1 to match 0-based indexing\n",
    "    \n",
    "    return soilclass, soilname"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "f6a41a12-759d-467c-a75e-78ed17e08486",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_pm_soil_color(soil_name):\n",
    "\n",
    "    '''Creates a soil color for a give sand/silt/clay percentage'''\n",
    "\n",
    "    # Color source: https://kremesti.com/water/silt_clay_mud.htm (accessed 2023-11-27)\n",
    "    if soil_name == 'clay':\n",
    "        color = [97,62,38]\n",
    "    elif soil_name == 'clay loam':\n",
    "        color = [68,45,30]\n",
    "    elif soil_name == 'loam':\n",
    "        color = [71,66,49]\n",
    "    elif soil_name == 'loamy sand':\n",
    "        color = [66,61,41]\n",
    "    elif soil_name == 'sand':\n",
    "        color = [91,87,67]\n",
    "    elif soil_name == 'sandy clay':\n",
    "        color = [85,82,71]\n",
    "    elif soil_name == 'sandy clay loam':\n",
    "        color = [78,64,40]\n",
    "    elif soil_name == 'sandy loam':\n",
    "        color = [74,71,42]\n",
    "    elif soil_name == 'silt':\n",
    "        color = [48,48,48]\n",
    "    elif soil_name == 'silty clay':\n",
    "        color = [73,52,39]\n",
    "    elif soil_name == 'silty clay loam':\n",
    "        color = [61,52,47]\n",
    "    elif soil_name == 'silt loam':\n",
    "        color = [88,83,77]\n",
    "    else:\n",
    "        print(f'WARNING: get_pm_soil_color(): soil {soil_name} not defined.')\n",
    "        color = None\n",
    "    return [C/255 for C in color]    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac5fdb35-fbf9-4be0-8e93-1002f6ca2c63",
   "metadata": {},
   "source": [
    "### Data extraction functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "b2b4c049-27be-4952-b266-9396969b9c7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_geotiff_properties(file):\n",
    "\n",
    "    with rasterio.open(file) as src:\n",
    "        geoTransform = src.transform\n",
    "        pixelSize = [geoTransform[0],geoTransform[4]]\n",
    "        rasterSize = [src.width,src.height]\n",
    "        boundingBox = src.bounds\n",
    "    \n",
    "    return pixelSize, rasterSize, boundingBox"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "e23061aa-483e-473c-98e8-f14311e84bed",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_geotif_data_as_array(file, band=1):\n",
    "    ds = gdal.Open(file) # open the file\n",
    "    band = ds.GetRasterBand(band) # get the data band\n",
    "    data = band.ReadAsArray() # convert to numpy array for further manipulation   \n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "c56fc5f0-d257-409f-a4c9-7fa686af37f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_geotif_noData(src_file, band=1):\n",
    "    src_ds = gdal.Open(src_file)\n",
    "    src_band = src_ds.GetRasterBand(band)\n",
    "    no_data_value = src_band.GetNoDataValue()\n",
    "    src_ds = None\n",
    "    return no_data_value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "1d7d6345-5dfd-4510-9bd4-c2850eefb01f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def confirm_dimensions_match(files):\n",
    "    pixel_sizes = []\n",
    "    raster_sizes = []\n",
    "    bounding_boxes = []\n",
    "    \n",
    "    for sub_list in files:\n",
    "        file = sub_list[1]\n",
    "        pixel, raster, bb = get_geotiff_properties(file)\n",
    "        pixel_sizes.append(pixel)\n",
    "        raster_sizes.append(raster)\n",
    "        bounding_boxes.append(bb)\n",
    "\n",
    "    assert all([sub_list == pixel_sizes[0] for sub_list in pixel_sizes]), 'pixel sizes don''t match'\n",
    "    assert all([sub_list == raster_sizes[0] for sub_list in raster_sizes]), 'raster sizes don''t match'\n",
    "    assert all([sub_list == bounding_boxes[0] for sub_list in bounding_boxes]), 'bounding boxes don''t match'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "8c78434f-0fa3-4cc8-a727-4a0fc6ed203d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_pixel_value(tiff_path, pixel_x, pixel_y):\n",
    "\n",
    "    file = str(tiff_path)\n",
    "    \n",
    "    # Get data\n",
    "    with rasterio.open(file) as src:\n",
    "        pixel_value = src.read(1, window=((pixel_y-1, pixel_y), (pixel_x-1, pixel_x)))\n",
    "        pixel_value = pixel_value[0][0]\n",
    "\n",
    "    # Get scaling and offset\n",
    "    src = gdal.Open(file)\n",
    "    scale_factor = src.GetRasterBand(1).GetScale()\n",
    "    offset = src.GetRasterBand(1).GetOffset()\n",
    "    if scale_factor: pixel_value = pixel_value * scale_factor\n",
    "    if offset: pixel_value = pixel_value + offset\n",
    "    src = None\n",
    "    \n",
    "    return pixel_value # Extract the actual value from the array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "d3f0d2e6-dc21-4360-8c58-5ca3e64b4d42",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Main: called from get_location_data_dict_from_geotiffs()\n",
    "def convert_lat_lon_to_raster_x_y(file,lat,lon):\n",
    "    pixel, raster, bb = get_geotiff_properties(file)\n",
    "    assert is_point_inside_bbox((lon,lat),bb), f'({lon},{lat}) not in {bb}' # sub 1\n",
    "    x,y = lon_lat_to_pixel(lon,lat, pixel[0], raster, bb) # sub 2\n",
    "    return x,y\n",
    "\n",
    "# Sub 1\n",
    "def is_point_inside_bbox(point, bbox):\n",
    "    x, y = point\n",
    "    min_x, min_y, max_x, max_y = bbox.left, bbox.bottom, bbox.right, bbox.top\n",
    "    return min_x <= x <= max_x and min_y <= y <= max_y\n",
    "\n",
    "# Sub 2\n",
    "def lon_lat_to_pixel(lon, lat, pixel_size, raster_size, bounding_box):\n",
    "    # Calculate pixel coordinates\n",
    "    x = int((lon - bounding_box.left) / pixel_size)\n",
    "    y = int((bounding_box.top - lat) / pixel_size)\n",
    "\n",
    "    # Ensure the coordinates are within the raster bounds\n",
    "    x = max(0, min(x, raster_size[0] - 1))\n",
    "    y = max(0, min(y, raster_size[1] - 1))\n",
    "\n",
    "    return x,y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "a1b1c1e2-9c0e-4b1a-bb1a-6b3ce6a7b3ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_location_data_dict_from_geotiffs(files,lon,lat):\n",
    "\n",
    "    # Confirm the all rasters have the same dimension (function will error out if not)\n",
    "    confirm_dimensions_match(files)\n",
    "    \n",
    "    # Get the raster indices for the provided lat/lon coordinates\n",
    "    x,y = convert_lat_lon_to_raster_x_y(files[0][1],lat,lon)\n",
    "\n",
    "    # Get the data values we want\n",
    "    data_dict = {}\n",
    "    for sublist in files:\n",
    "        data_name,data_file = sublist\n",
    "        data_value = extract_pixel_value(data_file,x,y)\n",
    "        data_dict[data_name] = data_value    \n",
    "    \n",
    "    return data_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7407e158-c313-41f3-b456-f986e9557716",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "perc-model-env",
   "language": "python",
   "name": "perc-model-env"
  },
  "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}