dev_samir_xarray.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "from dask.distributed import Client\n",
    "import os\n",
    "import numpy as np\n",
    "from typing import List, Tuple, Union\n",
    "import warnings\n",
    "import gc\n",
    "from parameters.params_samir_class import samir_parameters\n",
    "from config.config import config\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def rasterize_samir_parameters(csv_param_file: str, empty_dataset: xr.Dataset, land_cover_raster: str, chunk_size: dict) -> Tuple[xr.Dataset, dict]:\n",
    "    \"\"\"\n",
    "    Creates a raster `xarray` dataset from the csv parameter file, the land cover raster and an empty dataset\n",
    "    that contains the right structure (emptied ndvi dataset for example). For each parameter, the function loops\n",
    "    on land cover classes to fill the raster.\n",
    "\n",
    "    ## Arguments\n",
    "    1. csv_param_file: `str`\n",
    "        path to csv paramter file\n",
    "    2. empty_dataset: `xr.Dataset`\n",
    "        empty dataset that contains the right structure (emptied ndvi dataset for example).\n",
    "    3. land_cover_raster: `str`\n",
    "        path to land cover netcdf raster\n",
    "    4. chunk_size: `dict`\n",
    "        chunk_size for dask computation\n",
    "\n",
    "    ## Returns\n",
    "    1. parameter_dataset: `xr.Dataset`\n",
    "        the dataset containing all the rasterized Parameters\n",
    "    2. scale_factor: `dict`\n",
    "        dictionnary containing the scale factors for each parameter\n",
    "    \"\"\"\n",
    "    \n",
    "    # Load samir params into an object\n",
    "    table_param = samir_parameters(csv_param_file)\n",
    "    \n",
    "    # Set general variables\n",
    "    class_count = table_param.table.shape[1] - 2  # remove dtype and default columns\n",
    "    \n",
    "    # Open land cover raster\n",
    "    land_cover = xr.open_dataarray(land_cover_raster, chunks = chunk_size)\n",
    "    \n",
    "    # Create dataset\n",
    "    parameter_dataset = empty_dataset.copy(deep = True)\n",
    "    \n",
    "    # Create dictionnary containing the scale factors\n",
    "    scale_factor = {}\n",
    "    \n",
    "    # Loop on samir parameters and create \n",
    "    for parameter in table_param.table.index[1:]:\n",
    "        \n",
    "        # Create new variable and set attributes\n",
    "        parameter_dataset[parameter] = land_cover.copy(deep = True).astype('f4')\n",
    "        parameter_dataset[parameter].attrs['name'] = parameter\n",
    "        parameter_dataset[parameter].attrs['description'] = 'cf SAMIR Doc for detail'\n",
    "        parameter_dataset[parameter].attrs['scale factor'] = str(table_param.table.loc[table_param.table.index == parameter]['scale_factor'].values[0])\n",
    "        \n",
    "        # Assigne value in dictionnary\n",
    "        scale_factor[parameter] = 1/int(table_param.table.loc[table_param.table.index == parameter]['scale_factor'].values[0])\n",
    "        \n",
    "        # Loop on classes to set parameter values for each class\n",
    "        for class_val, class_name in zip(range(1, class_count + 1), table_param.table.columns[2:]):\n",
    "            \n",
    "            # Parameter values are multiplied by the scale factor in order to store all values as int16 types\n",
    "            # These values are then rounded to make sure there isn't any decimal point issues when casting the values to int16\n",
    "            parameter_dataset[parameter].values = np.where(parameter_dataset[parameter].values == class_val, round(table_param.table.loc[table_param.table.index == parameter][class_name].values[0]*table_param.table.loc[table_param.table.index == parameter]['scale_factor'].values[0]), parameter_dataset[parameter].values).astype('f4')\n",
    "    \n",
    "    # Return dataset converted to 'int16' data type to reduce memory usage\n",
    "    # and scale_factor dictionnary for later conversion\n",
    "    return parameter_dataset, scale_factor\n",
    "\n",
    "\n",
    "def setup_time_loop(calculation_variables_t1: List[str], calculation_variables_t2: List[str], empty_dataset: xr.Dataset) -> Tuple[xr.Dataset, xr.Dataset]:\n",
    "    \"\"\"\n",
    "    Creates two temporary `xarray Datasets` that will be used in the SAMIR time loop.\n",
    "    `variables_t1` corresponds to the variables for the previous day and `variables_t2`\n",
    "    corresponds to the variables for the current day. After each loop, `variables_t1`\n",
    "    takes the value of `variables_t2` for the corresponding variables.\n",
    "\n",
    "    ## Arguments\n",
    "    1. calculation_variables_t1: `List[str]`\n",
    "        list of strings containing the variable names\n",
    "        for the previous day dataset\n",
    "    2. calculation_variables_t2: `List[str]`\n",
    "        list of strings containing the variable names\n",
    "        for the current day dataset\n",
    "    3. empty_dataset: `xr.Dataset`\n",
    "        empty dataset that contains the right structure\n",
    "\n",
    "    ## Returns\n",
    "    1. variables_t1: `xr.Dataset`\n",
    "        output dataset for previous day\n",
    "    2. variables_t2: `xr.Dataset`\n",
    "        output dataset for current day\n",
    "    \"\"\"\n",
    "    \n",
    "    # Create new dataset\n",
    "    variables_t1 = empty_dataset.copy(deep = True)\n",
    "    \n",
    "    # Create empty DataArray for each variable\n",
    "    for variable in calculation_variables_t1:\n",
    "        \n",
    "        # Assign new empty DataArray\n",
    "        variables_t1[variable] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype = 'float32'))\n",
    "        variables_t1[variable].attrs['name'] = variable  # set name in attributes\n",
    "    \n",
    "    # Create new dataset\n",
    "    variables_t2 = empty_dataset.copy(deep = True)\n",
    "    \n",
    "    # Create empty DataArray for each variable\n",
    "    for variable in calculation_variables_t2:\n",
    "        \n",
    "        # Assign new empty DataArray\n",
    "        variables_t2[variable] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype = 'float32'))\n",
    "        variables_t2[variable].attrs['name'] = variable  # set name in attributes\n",
    "    \n",
    "    return variables_t1, variables_t2\n",
    "\n",
    "\n",
    "def prepare_outputs(empty_dataset: xr.Dataset, additional_outputs: List[str] = None) -> xr.Dataset:\n",
    "    \"\"\"\n",
    "    Creates the `xarray Dataset` containing the outputs of the SAMIR model that will be saved.\n",
    "    Additional variables can be saved by adding their names to the `additional_outputs` list.\n",
    "\n",
    "    ## Arguments\n",
    "    1. empty_dataset: `xr.Dataset`\n",
    "        empty dataset that contains the right structure\n",
    "    2. additional_outputs: `List[str]`\n",
    "        list of additional variable names to be saved\n",
    "\n",
    "    ## Returns\n",
    "    1. model_outputs: `xr.Dataset`\n",
    "        model outputs to be saved\n",
    "    \"\"\"\n",
    "    \n",
    "    # Evaporation and Transpiraion\n",
    "    model_outputs = empty_dataset.copy(deep = True)\n",
    "    \n",
    "    model_outputs['E'] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype = 'int16'))\n",
    "    model_outputs['E'].attrs['units'] = 'mm'\n",
    "    model_outputs['E'].attrs['standard_name'] = 'Evaporation'\n",
    "    model_outputs['E'].attrs['description'] = 'Accumulated daily evaporation in milimeters'\n",
    "    model_outputs['E'].attrs['scale factor'] = '1'\n",
    "    \n",
    "    model_outputs['Tr'] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype = 'int16'))\n",
    "    model_outputs['Tr'].attrs['units'] = 'mm'\n",
    "    model_outputs['Tr'].attrs['standard_name'] = 'Transpiration'\n",
    "    model_outputs['Tr'].attrs['description'] = 'Accumulated daily plant transpiration in milimeters'\n",
    "    model_outputs['Tr'].attrs['scale factor'] = '1'\n",
    "    \n",
    "    # Soil Water Content\n",
    "    model_outputs['SWCe'] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype = 'int16'))\n",
    "    model_outputs['SWCe'].attrs['units'] = 'mm'\n",
    "    model_outputs['SWCe'].attrs['standard_name'] = 'Soil Water Content of the evaporative zone'\n",
    "    model_outputs['SWCe'].attrs['description'] = 'Soil water content of the evaporative zone in milimeters'\n",
    "    model_outputs['SWCe'].attrs['scale factor'] = '1'\n",
    "    \n",
    "    model_outputs['SWCr'] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype = 'int16'))\n",
    "    model_outputs['SWCr'].attrs['units'] = 'mm'\n",
    "    model_outputs['SWCr'].attrs['standard_name'] = 'Soil Water Content of the root zone'\n",
    "    model_outputs['SWCr'].attrs['description'] = 'Soil water content of the root zone in milimeters'\n",
    "    model_outputs['SWCr'].attrs['scale factor'] = '1'\n",
    "    \n",
    "    # Irrigation\n",
    "    model_outputs['Irr'] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype = 'int16'))\n",
    "    model_outputs['Irr'].attrs['units'] = 'mm'\n",
    "    model_outputs['Irr'].attrs['standard_name'] = 'Irrigation'\n",
    "    model_outputs['Irr'].attrs['description'] = 'Simulated daily irrigation in milimeters'\n",
    "    model_outputs['Irr'].attrs['scale factor'] = '1'\n",
    "    \n",
    "    # Deep Percolation\n",
    "    model_outputs['DP'] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype = 'int16'))\n",
    "    model_outputs['DP'].attrs['units'] = 'mm'\n",
    "    model_outputs['DP'].attrs['standard_name'] = 'Deep Percolation'\n",
    "    model_outputs['DP'].attrs['description'] = 'Simulated daily Deep Percolation in milimeters'\n",
    "    model_outputs['DP'].attrs['scale factor'] = '1'\n",
    "    \n",
    "    if additional_outputs:\n",
    "        for var in additional_outputs:\n",
    "            model_outputs[var] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype = 'int16'))\n",
    "    \n",
    "    return model_outputs\n",
    "\n",
    "\n",
    "def xr_maximum(ds: xr.DataArray, value: Union[xr.DataArray, float, int]) -> xr.DataArray:\n",
    "    \"\"\"\n",
    "    Equivalent of `numpy.maximum(ds, value)` for xarray DataArrays\n",
    "\n",
    "    ## Arguments\n",
    "    1. ds: `xr.DataArray`\n",
    "        datarray to compare\n",
    "    2. value: `Union[xr.DataArray, float, int]`\n",
    "        value (scalar or dataarray) to compare\n",
    "\n",
    "    ## Returns\n",
    "    1. output: `xr.DataArray`\n",
    "        resulting dataarray with maximum value element-wise\n",
    "    \"\"\"\n",
    "    return xr.where(ds <= value, value, ds, keep_attrs = True)\n",
    "\n",
    "\n",
    "def xr_minimum(ds: xr.DataArray, value: Union[xr.DataArray, float, int]) -> xr.DataArray:\n",
    "    \"\"\"\n",
    "    Equivalent of `numpy.minimum(ds, value)` for xarray DataArrays\n",
    "\n",
    "    ## Arguments\n",
    "    1. ds: `xr.DataArray`\n",
    "        datarray to compare\n",
    "    2. value: `Union[xr.DataArray, float, int]`\n",
    "        value (scalar or dataarray) to compare\n",
    "\n",
    "    ## Returns\n",
    "    1. output: `xr.DataArray`\n",
    "        resulting dataarray with minimum value element-wise\n",
    "    \"\"\"\n",
    "    return xr.where(ds >= value, value, ds, keep_attrs = True)\n",
    "\n",
    "\n",
    "def calculate_diff_re(TAW: xr.DataArray, Dr: xr.DataArray, Zr: xr.DataArray, RUE: xr.DataArray, De: xr.DataArray, FCov: xr.DataArray, Ze_: xr.DataArray, DiffE_: xr.DataArray, scale_dict: dict) -> xr.DataArray:\n",
    "    \"\"\"\n",
    "    Calculates the diffusion between the top soil layer and the root layer.\n",
    "\n",
    "    ## Arguments\n",
    "    1. TAW: `xr.DataArray`\n",
    "        water capacity of root layer\n",
    "    2. Dr: `xr.DataArray`\n",
    "        depletion of root layer\n",
    "    3. Zr: `xr.DataArray`\n",
    "        height of root layer\n",
    "    4. RUE: `xr.DataArray`\n",
    "        total available surface water\n",
    "    5. De: `xr.DataArray`\n",
    "        depletion of the evaporative layer\n",
    "    6. FCov: `xr.DataArray`\n",
    "        fraction cover of plants\n",
    "    7. Ze_: `xr.DataArray`\n",
    "        height of evaporative layer (paramter)\n",
    "    8. DiffE_: `xr.DataArray`\n",
    "        diffusion coefficient between evaporative\n",
    "        and root layers (unitless, parameter)\n",
    "    9. scale_dict: `dict`\n",
    "        dictionnary containing the scale factors for\n",
    "        the rasterized parameters\n",
    "\n",
    "    ## Returns\n",
    "    1. diff_re: `xr.Dataset`\n",
    "        the diffusion between the top soil layer and\n",
    "        the root layer\n",
    "    \"\"\"\n",
    "    \n",
    "    # Temporary variables to make calculation easier to read\n",
    "    tmp1 = (((TAW - Dr) / Zr - (RUE - De) / (scale_dict['Ze'] * Ze_)) / FCov) * (scale_dict['DiffE'] * DiffE_)\n",
    "    tmp2 = ((TAW * scale_dict['Ze'] * Ze_) - (RUE - De - Dr) * Zr) / (Zr + scale_dict['Ze'] * Ze_) - Dr\n",
    "    \n",
    "    # Calculate diffusion according to SAMIR equation\n",
    "    diff_re = xr.where(tmp1 < 0, xr_maximum(tmp1, tmp2), xr_minimum(tmp1, tmp2))\n",
    "\n",
    "    # Return zero values where the 'DiffE' parameter is equal to 0\n",
    "    return xr.where(DiffE_ == 0, 0, diff_re)\n",
    "\n",
    "\n",
    "def calculate_diff_dr(TAW: xr.DataArray, TDW: xr.DataArray, Dr: xr.DataArray, Zr: xr.DataArray, Dd: xr.DataArray, FCov: xr.DataArray, Zsoil_: xr.DataArray, DiffR_: xr.DataArray, scale_dict: dict) -> xr.DataArray:\n",
    "    \"\"\"\n",
    "    Calculates the diffusion between the root layer and the deep layer.\n",
    "\n",
    "    ## Arguments\n",
    "    1. TAW: `xr.DataArray`\n",
    "        water capacity of root layer\n",
    "    2. TDW: `xr.DataArray`\n",
    "        water capacity of deep layer\n",
    "    3. Dr: `xr.DataArray`\n",
    "        depletion of root layer\n",
    "    4. Zr: `xr.DataArray`\n",
    "        height of root layer\n",
    "    5. Dd: `xr.DataArray`\n",
    "        depletion of deep layer\n",
    "    6. FCov: `xr.DataArray`\n",
    "        fraction cover of plants\n",
    "    7. Zsoil_: `xr.DataArray`\n",
    "        total height of soil (paramter)\n",
    "    8. DiffR_: `xr.DataArray`\n",
    "        Diffusion coefficient between root\n",
    "        and deep layers (unitless, parameter)\n",
    "    9. scale_dict: `dict`\n",
    "        dictionnary containing the scale factors for\n",
    "        the rasterized parameters\n",
    "\n",
    "    ## Returns\n",
    "    1. diff_dr: `xr.Dataset`\n",
    "        the diffusion between the root layer and the\n",
    "        deep layer\n",
    "    \"\"\"\n",
    "    \n",
    "    # Temporary variables to make calculation easier to read\n",
    "    tmp1 = (((TDW - Dd) / (scale_dict['Zsoil'] * Zsoil_ - Zr) - (TAW - Dr) / Zr) / FCov) * scale_dict['DiffR'] * DiffR_\n",
    "    tmp2 = (TDW *Zr - (TAW - Dr - Dd) * (scale_dict['Zsoil'] * Zsoil_ - Zr)) / (scale_dict['Zsoil'] * Zsoil_) - Dd\n",
    "    \n",
    "    # Calculate diffusion according to SAMIR equation\n",
    "    diff_dr = xr.where(tmp1 < 0, xr_maximum(tmp1, tmp2), xr_minimum(tmp1, tmp2))\n",
    "    \n",
    "    # Return zero values where the 'DiffR' parameter is equal to 0\n",
    "    return xr.where(DiffR_ == 0, 0, diff_dr)\n",
    "\n",
    "\n",
    "def calculate_W(TEW: xr.DataArray, Dei: xr.DataArray, Dep: xr.DataArray, fewi: xr.DataArray, fewp: xr.DataArray) -> xr.DataArray:\n",
    "    \"\"\"\n",
    "    Calculate W, the weighting factor to split the energy available\n",
    "    for evaporation depending on the difference in water availability\n",
    "    in the two evaporation components, ensuring that the larger and\n",
    "    the wetter, the more the evaporation occurs from that component\n",
    "\n",
    "    ## Arguments\n",
    "    1. TEW: `xr.DataArray`\n",
    "        water capacity of evaporative layer\n",
    "    2. Dei: `xr.DataArray`\n",
    "        depletion of the evaporative layer\n",
    "        (irrigation part)\n",
    "    3. Dep: `xr.DataArray`\n",
    "        depletion of the evaporative layer\n",
    "        (precipitation part)\n",
    "    4. fewi: `xr.DataArray`\n",
    "        soil fraction which is wetted by irrigation\n",
    "        and exposed to evaporation\n",
    "    5. fewp: `xr.DataArray`\n",
    "        soil fraction which is wetted by precipitation\n",
    "        and exposed to evaporation\n",
    "\n",
    "    ## Returns\n",
    "    1. W: `xr.DataArray`\n",
    "        weighting factor W\n",
    "    \"\"\"\n",
    "    \n",
    "    # Temporary variables to make calculation easier to read\n",
    "    tmp = fewi * (TEW - Dei)\n",
    "    \n",
    "    # Calculate the weighting factor to split the energy available for evaporation\n",
    "    W = 1 / (1 + (fewp * (TEW - Dep) / tmp ))\n",
    "\n",
    "    # Return W \n",
    "    return xr.where(tmp > 0, W, 0)\n",
    "\n",
    "\n",
    "def calculate_Kr(TEW: xr.DataArray, De: xr.DataArray, REW_: xr.DataArray, scale_dict: dict) -> xr.DataArray:\n",
    "    \"\"\"\n",
    "    calculates of the reduction coefficient for evaporation dependent \n",
    "    on the amount of water in the soil using the FAO-56 method\n",
    "\n",
    "    ## Arguments\n",
    "    1. TEW: `xr.DataArray`\n",
    "        water capacity of evaporative layer\n",
    "    2. De: `xr.DataArray`\n",
    "        depletion of evaporative layer\n",
    "    3. REW_: `xr.DataArray`\n",
    "        readily evaporable water\n",
    "    4. scale_dict: `dict`\n",
    "        dictionnary containing the scale factors for\n",
    "        the rasterized parameters\n",
    "\n",
    "    ## Returns\n",
    "    1. Kr: `xr.DataArray`\n",
    "        Kr coefficient\n",
    "    \"\"\"\n",
    "    \n",
    "    # Formula for calculating Kr\n",
    "    Kr = (TEW - De) / (TEW - scale_dict['REW'] * REW_)\n",
    "    \n",
    "    # Return Kr\n",
    "    return xr_maximum(0, xr_minimum(Kr, 1))\n",
    "\n",
    "\n",
    "def update_Dr(TAW: xr.DataArray, TDW: xr.DataArray, Zr: xr.DataArray, TAW0: xr.DataArray, TDW0: xr.DataArray, Dr0: xr.DataArray, Dd0: xr.DataArray, Zr0: xr.DataArray) -> xr.DataArray:\n",
    "    \"\"\"\n",
    "    Return the updated depletion for the root layer\n",
    "\n",
    "    ## Arguments\n",
    "    1. TAW: `xr.DataArray`\n",
    "        water capacity of root layer for current day\n",
    "    2. TDW: `xr.DataArray`\n",
    "        water capacity of deep layer for current day\n",
    "    3. Zr: `xr.DataArray`\n",
    "        root layer height for current day\n",
    "    4. TAW0: `xr.DataArray`\n",
    "        water capacity of root layer for previous day\n",
    "    5. TDW0: `xr.DataArray`\n",
    "        water capacity of deep layer for previous day\n",
    "    6. Dr0: `xr.DataArray`\n",
    "        depletion of the root layer for previous day\n",
    "    7. Dd0: `xr.DataArray`\n",
    "        depletion of the deep laye for previous day\n",
    "    8. Zr0: `xr.DataArray`\n",
    "        root layer height for previous day\n",
    "\n",
    "    ## Returns\n",
    "    1. output: `xr.DataArray`\n",
    "        updated depletion for the root layer\n",
    "    \"\"\"\n",
    "    \n",
    "    # Temporary variables to make calculation easier to read\n",
    "    tmp1 = xr_maximum(Dr0 + Dd0 * (TAW - TAW0) / TDW0, 0)\n",
    "    tmp2 = xr_minimum(Dr0 + Dd0 * (TAW - TAW0) / TDW0, TDW)\n",
    "\n",
    "    # Return updated Dr\n",
    "    return xr.where(Zr > Zr0, tmp1, tmp2)\n",
    "\n",
    "\n",
    "def update_Dd(TAW: xr.DataArray, TDW: xr.DataArray, Zr: xr.DataArray, TAW0: xr.DataArray, TDW0: xr.DataArray, Dd0: xr.DataArray, Zr0: xr.DataArray) -> xr.DataArray:\n",
    "    \"\"\"\n",
    "    Return the updated depletion for the deep layer\n",
    "\n",
    "    ## Arguments\n",
    "    1. TAW: `xr.DataArray`\n",
    "        water capacity of root layer for current day\n",
    "    2. TDW: `xr.DataArray`\n",
    "        water capacity of deep layer for current day\n",
    "    3. TAW0: `xr.DataArray`\n",
    "        water capacity of root layer for previous day\n",
    "    5. TDW0: `xr.DataArray`\n",
    "        water capacity of deep layer for previous day\n",
    "    6. Dd0: `xr.DataArray`\n",
    "        depletion of the deep laye for previous day\n",
    "    7. Zr0: `xr.DataArray`\n",
    "        root layer height for previous day\n",
    "\n",
    "    ## Returns\n",
    "    1. output: `xr.DataArray`\n",
    "        updated depletion for the deep layer\n",
    "    \"\"\"\n",
    "    \n",
    "    # Temporary variables to make calculation easier to read\n",
    "    tmp1 = xr_maximum(Dd0 - Dd0 * (TAW - TAW0) / TDW0, 0)\n",
    "    tmp2 = xr_minimum(Dd0 - Dd0 * (TAW - TAW0) / TDW0, TDW)\n",
    "    \n",
    "    # Return updated Dd\n",
    "    return xr.where(Zr > Zr0, tmp1, tmp2)\n",
    "\n",
    "\n",
    "def run_samir(json_config_file: str, csv_param_file: str, ndvi_cube_path: str, weather_cube_path: str, soil_params_path: str, land_cover_path: str, chunk_size: dict, save_path: str) -> None:\n",
    "    \n",
    "    # warnings.simplefilter(\"error\", category = RuntimeWarning())\n",
    "    warnings.filterwarnings(\"ignore\", message=\"invalid value encountered in cast\")\n",
    "    warnings.filterwarnings(\"ignore\", message=\"invalid value encountered in divide\")\n",
    "    np.errstate(all = 'raise')\n",
    "    gc.disable()\n",
    "    \n",
    "    #============ General parameters ============#\n",
    "    config_params = config(json_config_file)\n",
    "    calculation_variables_t2 = ['diff_rei', 'diff_rep', 'diff_dr' , 'Dd', 'De', 'Dei', 'Dep', 'Dr', 'FCov', 'Irrig', 'Kcb', 'Kei', 'Kep', 'Ks', 'Kti', 'Ktp', 'RUE', 'TAW', 'TDW', 'TEW', 'Tei', 'Tep', 'W', 'Zr', 'fewi', 'fewp']\n",
    "    calculation_variables_t1 = ['Dr', 'Dd', 'TAW', 'TDW', 'Zr']\n",
    "    \n",
    "    #============ Manage inputs ============#\n",
    "    # NDVI\n",
    "    ndvi_cube = xr.open_dataset(ndvi_cube_path, chunks = chunk_size).astype('f4')\n",
    "    \n",
    "    # # Create a daily DateTimeIndex for the desired date range\n",
    "    # daily_index = pd.date_range(start = config_params.start_date, end = config_params.end_date, freq = 'D')\n",
    "\n",
    "    # # Resample the dataset to a daily frequency and reindex with the new DateTimeIndex\n",
    "    # ndvi_cube = ndvi_cube.resample(time = '1D').asfreq().reindex(time = daily_index)\n",
    "\n",
    "    # # Interpolate the dataset along the time dimension to fill nan values\n",
    "    # ndvi_cube = ndvi_cube.interpolate_na(dim = 'time', method = 'linear', fill_value = 'extrapolate').astype('u1')\n",
    "    \n",
    "    # Weather\n",
    "    weather_cube = xr.open_dataset(weather_cube_path, chunks = chunk_size).astype('f4')\n",
    "    \n",
    "    # Soil\n",
    "    soil_params = xr.open_dataset(soil_params_path, chunks = chunk_size).astype('f4')\n",
    "    \n",
    "    # SAMIR Parameters\n",
    "    param_dataset, scale_factor = rasterize_samir_parameters(csv_param_file, ndvi_cube.drop_vars(['ndvi', 'time']), land_cover_path, chunk_size = chunk_size)\n",
    "    \n",
    "    # SAMIR Variables\n",
    "    variables_t1, variables_t2 = setup_time_loop(calculation_variables_t1, calculation_variables_t2, ndvi_cube.drop_vars(['ndvi', 'time']))\n",
    "\n",
    "    #============ Prepare outputs ============#\n",
    "    model_outputs = prepare_outputs(ndvi_cube.drop_vars(['ndvi']))\n",
    "    \n",
    "    #============ Prepare time iterations ============#\n",
    "    dates = ndvi_cube.time.values\n",
    "    \n",
    "    #============ Create aliases for better readability ============#\n",
    "    \n",
    "    # Variables for current day\n",
    "    diff_rei = variables_t2['diff_rei']\n",
    "    diff_rep = variables_t2['diff_rep']\n",
    "    diff_dr = variables_t2['diff_dr']\n",
    "    Dd = variables_t2['Dd']\n",
    "    De = variables_t2['De']\n",
    "    Dei = variables_t2['Dei']\n",
    "    Dep = variables_t2['Dep']\n",
    "    Dr = variables_t2['Dr']\n",
    "    FCov = variables_t2['FCov']\n",
    "    Irrig = variables_t2['Irrig']\n",
    "    Kcb = variables_t2['Kcb']\n",
    "    Kei = variables_t2['Kei']\n",
    "    Kep = variables_t2['Kep']\n",
    "    Ks = variables_t2['Ks']\n",
    "    Kti = variables_t2['Kti']\n",
    "    Ktp = variables_t2['Ktp']\n",
    "    RUE = variables_t2['RUE']\n",
    "    TAW = variables_t2['TAW']\n",
    "    TDW = variables_t2['TDW']\n",
    "    TEW = variables_t2['TEW']\n",
    "    Tei = variables_t2['Tei']\n",
    "    Tep = variables_t2['Tep']\n",
    "    Zr = variables_t2['Zr']\n",
    "    W = variables_t2['W']\n",
    "    fewi = variables_t2['fewi']\n",
    "    fewp = variables_t2['fewp']\n",
    "    \n",
    "    # Variables for previous day\n",
    "    TAW0 = variables_t1['TAW']\n",
    "    TDW0 = variables_t1['TDW']\n",
    "    Dr0 = variables_t1['Dr']\n",
    "    Dd0 = variables_t1['Dd']\n",
    "    Zr0 = variables_t1['Zr']\n",
    "    \n",
    "    # Parameters\n",
    "    # Parameters have an underscore (_) behind their name for recognition \n",
    "    DiffE_ = param_dataset['DiffE']\n",
    "    DiffR_ = param_dataset['DiffR']\n",
    "    FW_ = param_dataset['FW']\n",
    "    Fc_stop_ = param_dataset['Fc_stop']\n",
    "    FmaxFC_ = param_dataset['FmaxFC']\n",
    "    Foffset_ = param_dataset['Foffset']\n",
    "    Fslope_ = param_dataset['Fslope']\n",
    "    Init_RU_ = param_dataset['Init_RU']\n",
    "    Irrig_auto_ = param_dataset['Irrig_auto']\n",
    "    Kcmax_ = param_dataset['Kcmax']\n",
    "    KmaxKcb_ = param_dataset['KmaxKcb']\n",
    "    Koffset_ = param_dataset['Koffset']\n",
    "    Kslope_ = param_dataset['Kslope']\n",
    "    Lame_max_ = param_dataset['Lame_max']\n",
    "    REW_ = param_dataset['REW']\n",
    "    Ze_ = param_dataset['Ze']\n",
    "    Zsoil_ = param_dataset['Zsoil']\n",
    "    maxZr_ = param_dataset['maxZr']\n",
    "    minZr_ = param_dataset['minZr']\n",
    "    p_ = param_dataset['p']\n",
    "    \n",
    "    # scale factors\n",
    "    # Scale factors have the following name scheme : s_ + parameter_name\n",
    "    s_DiffE = scale_factor['DiffE']\n",
    "    s_DiffR = scale_factor['DiffR']\n",
    "    s_FW = scale_factor['FW']\n",
    "    s_Fc_stop = scale_factor['Fc_stop']\n",
    "    s_FmaxFC = scale_factor['FmaxFC']\n",
    "    s_Foffset = scale_factor['Foffset']\n",
    "    s_Fslope = scale_factor['Fslope']\n",
    "    s_Init_RU = scale_factor['Init_RU']\n",
    "    # s_Irrig_auto = scale_factor['Irrig_auto']\n",
    "    s_Kcmax = scale_factor['Kcmax']\n",
    "    s_KmaxKcb = scale_factor['KmaxKcb']\n",
    "    s_Koffset = scale_factor['Koffset']\n",
    "    s_Kslope = scale_factor['Kslope']\n",
    "    s_Lame_max = scale_factor['Lame_max']\n",
    "    s_REW = scale_factor['REW']\n",
    "    s_Ze = scale_factor['Ze']\n",
    "    s_Zsoil = scale_factor['Zsoil']\n",
    "    s_maxZr = scale_factor['maxZr']\n",
    "    s_minZr = scale_factor['minZr']\n",
    "    s_p = scale_factor['p']\n",
    "    \n",
    "    #============ First day initialization ============#\n",
    "    # Fraction cover\n",
    "    FCov = s_Fslope * Fslope_ * (ndvi_cube['ndvi'].sel({'time': dates[0]})/255) + s_Foffset * Foffset_\n",
    "    FCov = xr_minimum(xr_maximum(FCov, 0), s_Fc_stop * Fc_stop_)\n",
    "    \n",
    "    # Root depth upate\n",
    "    Zr = s_minZr * minZr_ + (FCov / (s_FmaxFC * FmaxFC_)) * s_maxZr * (maxZr_ - minZr_)\n",
    "    \n",
    "    # Water capacities\n",
    "    TEW = (soil_params['FC'] - soil_params['WP']/2) * s_Ze * Ze_\n",
    "    RUE = (soil_params['FC'] - soil_params['WP']) * s_Ze * Ze_\n",
    "    TAW = (soil_params['FC'] - soil_params['WP']) * Zr\n",
    "    TDW = (soil_params['FC'] - soil_params['WP']) * (s_Zsoil * Zsoil_ - Zr)  # Zd = Zsoil - Zr\n",
    "    \n",
    "    # Depletions\n",
    "    Dei = RUE * (1 - s_Init_RU * Init_RU_)\n",
    "    Dep = RUE * (1 - s_Init_RU * Init_RU_)\n",
    "    Dr = TAW * (1 - s_Init_RU * Init_RU_)\n",
    "    Dd = TDW * (1 - s_Init_RU * Init_RU_)\n",
    "    \n",
    "    # Irrigation   ==============!!!!!\n",
    "    Irrig = xr_minimum(xr_maximum(Dr - weather_cube['tp'].sel({'time': dates[0]}) / 1000, 0), s_Lame_max * Lame_max_) * Irrig_auto_\n",
    "    Irrig = xr.where(Dr > TAW * s_p * p_, Irrig, 0)\n",
    "    \n",
    "    # Kcb\n",
    "    Kcb = xr_minimum(s_Kslope * Kslope_ * (ndvi_cube['ndvi'].sel({'time': dates[0]}) / 255) + s_Koffset * Koffset_, s_KmaxKcb * KmaxKcb_)\n",
    "    \n",
    "    # Update depletions with rainfall and/or irrigation    \n",
    "    ## DP\n",
    "    model_outputs['DP'].loc[{'time': dates[0]}] = -xr_minimum(Dd + xr_minimum(Dr - (weather_cube['tp'].sel({'time': dates[0]}) / 1000) - Irrig, 0), 0)\n",
    "    \n",
    "    ## De\n",
    "    Dei = xr_minimum(xr_maximum(Dei - (weather_cube['tp'].sel({'time': dates[0]}) / 1000) - Irrig / (s_FW * FW_ / 100), 0), TEW)\n",
    "    Dep = xr_minimum(xr_maximum(Dep - (weather_cube['tp'].sel({'time': dates[0]}) / 1000), 0), TEW)\n",
    "    \n",
    "    fewi = xr_minimum(1 - FCov, (s_FW * FW_ / 100))\n",
    "    fewp = 1 - FCov - fewi\n",
    "    \n",
    "    De = (Dei * fewi + Dep * fewp) / (fewi + fewp)\n",
    "    # variables_t1['De'] = xr.where(variables_t1['De'].isfinite(), variables_t1['De'], variables_t1['Dei'] * (s_FW * FW_ / 100) + variables_t1['Dep'] * (1 - (s_FW * FW_ / 100)))  #================= find replacement for .isfinite() method !!!!!!!!!\n",
    "\n",
    "    ## Dr\n",
    "    Dr = xr_minimum(xr_maximum(Dr - (weather_cube['tp'].sel({'time': dates[0]}) / 1000) - Irrig, 0), TAW)\n",
    "    \n",
    "    ## Dd\n",
    "    Dd = xr_minimum(xr_maximum(Dd + xr_minimum(Dr - (weather_cube['tp'].sel({'time': dates[0]}) / 1000) - Irrig, 0), 0), TDW)\n",
    "    \n",
    "    # Diffusion coefficients\n",
    "    diff_rei = calculate_diff_re(TAW, Dr, Zr, RUE, Dei, FCov, Ze_, DiffE_, scale_factor)\n",
    "    diff_rep = calculate_diff_re(TAW, Dr, Zr, RUE, Dep, FCov, Ze_, DiffE_, scale_factor)\n",
    "    diff_dr = calculate_diff_dr(TAW, TDW, Dr, Zr, Dd, FCov, Zsoil_, DiffR_, scale_factor)    \n",
    "    \n",
    "    # Weighing factor W\n",
    "    W = calculate_W(TEW, Dei, Dep, fewi, fewp)\n",
    "    \n",
    "    # Soil water contents\n",
    "    model_outputs['SWCe'].loc[{'time': dates[0]}] = 1 - De/TEW\n",
    "    model_outputs['SWCr'].loc[{'time': dates[0]}] = 1 - Dr/TAW\n",
    "    \n",
    "    # Water Stress coefficient\n",
    "    Ks = xr_minimum((TAW - Dr) / (TAW * (1 - s_p * p_)), 1)\n",
    "    \n",
    "    # Reduction coefficient for evaporation\n",
    "    Kei = xr_minimum(W * calculate_Kr(TEW, Dei, REW_, scale_factor) * (s_Kcmax * Kcmax_ - Kcb), fewi * s_Kcmax * Kcmax_)\n",
    "    Kep = xr_minimum((1 - W) * calculate_Kr(TEW, Dep, REW_, scale_factor) * (s_Kcmax * Kcmax_ - Kcb), fewp * s_Kcmax * Kcmax_)\n",
    "    \n",
    "    # Prepare coefficients for evapotranspiration\n",
    "    Kti = xr_minimum(((s_Ze * Ze_ / Zr)**6) * (1 - Dei / TEW) / xr_maximum(1 - Dr / TAW, 0.001), 1)\n",
    "    Ktp = xr_minimum(((s_Ze * Ze_ / Zr)**6) * (1 - Dep / TEW) / xr_maximum(1 - Dr / TAW, 0.001), 1)\n",
    "    Tei = Kti * Ks * Kcb * weather_cube['ET0'].sel({'time': dates[0]}) / 1000\n",
    "    Tep = Ktp * Ks * Kcb * weather_cube['ET0'].sel({'time': dates[0]}) / 1000\n",
    "    \n",
    "    # Update depletions\n",
    "    Dei = xr.where(fewi > 0, xr_minimum(xr_maximum(Dei + (weather_cube['ET0'].sel({'time': dates[0]}) / 1000) * Kei / fewi + Tei - diff_rei, 0), TEW), xr_minimum(xr_maximum(Dei + Tei - diff_rei, 0), TEW))\n",
    "    Dep = xr.where(fewp > 0, xr_minimum(xr_maximum(Dep + (weather_cube['ET0'].sel({'time': dates[0]}) / 1000) * Kep / fewp + Tep - diff_rep, 0), TEW), xr_minimum(xr_maximum(Dep + Tep - diff_rep, 0), TEW))\n",
    "    \n",
    "    De = (Dei * fewi + Dep * fewp) / (fewi + fewp)\n",
    "    # De = xr.where(De.isfinite(), De, Dei * (s_FW * FW_ / 100) + Dep * (1 - (s_FW * FW_ / 100)))  #================= find replacement for .isfinite() method !!!!!!!!!\n",
    "    \n",
    "    # Evaporation\n",
    "    model_outputs['E'].loc[{'time': dates[0]}]  = xr_maximum((Kei + Kep) * weather_cube['ET0'].sel({'time': dates[0]}) / 1000, 0)\n",
    "    \n",
    "    # Transpiration\n",
    "    model_outputs['Tr'].loc[{'time': dates[0]}]  = Kcb * Ks * weather_cube['ET0'].sel({'time': dates[0]}) / 1000\n",
    "    \n",
    "    # Potential evapotranspiration and evaporative fraction ??\n",
    "    \n",
    "    # Update depletions (root and deep zones) at the end of the day\n",
    "    Dr = xr_minimum(xr_maximum(Dr + model_outputs['E'].loc[{'time': dates[0]}] + model_outputs['Tr'].loc[{'time': dates[0]}] - diff_dr, 0), TAW)\n",
    "    Dd = xr_minimum(xr_maximum(Dd + diff_dr, 0), TDW)\n",
    "    \n",
    "    # Write outputs\n",
    "    model_outputs['Irr'].loc[{'time': dates[0]}] = Irrig\n",
    "    \n",
    "    # Update variable_t1 values\n",
    "    for variable in calculation_variables_t1:\n",
    "        variables_t1[variable] = variables_t2[variable].copy(deep = True)\n",
    "    \n",
    "    #============ Time loop ============#\n",
    "    for i in range(1, len(dates)):\n",
    "        \n",
    "        # Update variables\n",
    "        ## Fraction cover\n",
    "        FCov = s_Fslope * Fslope_ * (ndvi_cube['ndvi'].sel({'time': dates[0]})/255) + s_Foffset * Foffset_\n",
    "        FCov = xr_minimum(xr_maximum(FCov, 0), s_Fc_stop * Fc_stop_)\n",
    "        \n",
    "        ## Root depth upate\n",
    "        Zr = s_minZr * minZr_ + (FCov / (s_FmaxFC * FmaxFC_)) * s_maxZr * (maxZr_ - minZr_)\n",
    "        \n",
    "        # Water capacities\n",
    "        TAW = (soil_params['FC'] - soil_params['WP']) * Zr\n",
    "        TDW = (soil_params['FC'] - soil_params['WP']) * (s_Zsoil * Zsoil_ - Zr)\n",
    "        \n",
    "        # Update depletions\n",
    "        Dr = update_Dr(TAW, TDW, Zr, TAW0, TDW0, Dr0, Dd0, Zr0)\n",
    "        Dd = update_Dd(TAW, TDW, Zr, TAW0, TDW0, Dd0, Zr0)\n",
    "        \n",
    "        # Update param p\n",
    "        p_ = (xr_minimum(xr_maximum(s_p * p_ + 0.04 * (5 - weather_cube['ET0'].sel({'time': dates[i-1]}) / 1000), 0.1), 0.8) * (1 / s_p)).round(0).astype('i2')\n",
    "        \n",
    "        # Irrigation   ==============!!!!!\n",
    "        Irrig = xr_minimum(xr_maximum(Dr - weather_cube['tp'].sel({'time': dates[i]}) / 1000, 0), s_Lame_max * Lame_max_) * Irrig_auto_\n",
    "        Irrig = xr.where(Dr > TAW * s_p * p_, Irrig, 0)\n",
    "    \n",
    "        # Kcb\n",
    "        Kcb = xr_minimum(s_Kslope * Kslope_ * (ndvi_cube['ndvi'].sel({'time': dates[i]}) / 255) + s_Koffset * Koffset_, s_KmaxKcb * KmaxKcb_)\n",
    "        \n",
    "        # Update depletions with rainfall and/or irrigation    \n",
    "        ## DP\n",
    "        model_outputs['DP'].loc[{'time': dates[i]}] = -xr_minimum(Dd + xr_minimum(Dr - (weather_cube['tp'].sel({'time': dates[i]}) / 1000) - Irrig, 0), 0)\n",
    "        \n",
    "        ## De\n",
    "        Dei = xr_minimum(xr_maximum(Dei - (weather_cube['tp'].sel({'time': dates[i]}) / 1000) - Irrig / (s_FW * FW_ / 100), 0), TEW)\n",
    "        Dep = xr_minimum(xr_maximum(Dep - (weather_cube['tp'].sel({'time': dates[i]}) / 1000), 0), TEW)\n",
    "        \n",
    "        fewi = xr_minimum(1 - FCov, (s_FW * FW_ / 100))\n",
    "        fewp = 1 - FCov - fewi\n",
    "        \n",
    "        De = (Dei * fewi + Dep * fewp) / (fewi + fewp)\n",
    "        # variables_t1['De'] = xr.where(variables_t1['De'].isfinite(), variables_t1['De'], variables_t1['Dei'] * (s_FW * FW_ / 100) + variables_t1['Dep'] * (1 - (s_FW * FW_ / 100)))  #================= find replacement for .isfinite() method !!!!!!!!!\n",
    "\n",
    "        ## Dr\n",
    "        Dr = xr_minimum(xr_maximum(Dr - (weather_cube['tp'].sel({'time': dates[i]}) / 1000) - Irrig, 0), TAW)\n",
    "        \n",
    "        ## Dd\n",
    "        Dd = xr_minimum(xr_maximum(Dd + xr_minimum(Dr - (weather_cube['tp'].sel({'time': dates[i]}) / 1000) - Irrig, 0), 0), TDW)\n",
    "        \n",
    "        # Diffusion coefficients\n",
    "        diff_rei = calculate_diff_re(TAW, Dr, Zr, RUE, Dei, FCov, Ze_, DiffE_, scale_factor)\n",
    "        diff_rep = calculate_diff_re(TAW, Dr, Zr, RUE, Dep, FCov, Ze_, DiffE_, scale_factor)\n",
    "        diff_dr = calculate_diff_dr(TAW, TDW, Dr, Zr, Dd, FCov, Zsoil_, DiffR_, scale_factor) \n",
    "        \n",
    "        # Weighing factor W\n",
    "        W = calculate_W(TEW, Dei, Dep, fewi, fewp)\n",
    "        \n",
    "        # Soil water contents\n",
    "        model_outputs['SWCe'].loc[{'time': dates[i]}] = 1 - De/TEW\n",
    "        model_outputs['SWCr'].loc[{'time': dates[i]}] = 1 - Dr/TAW\n",
    "        \n",
    "        # Water Stress coefficient\n",
    "        Ks = xr_minimum((TAW - Dr) / (TAW * (1 - s_p * p_)), 1)\n",
    "        \n",
    "        # Reduction coefficient for evaporation\n",
    "        Kei = xr_minimum(W * calculate_Kr(TEW, Dei, REW_, scale_factor) * (s_Kcmax * Kcmax_ - Kcb), fewi * s_Kcmax * Kcmax_)\n",
    "        Kep = xr_minimum((1 - W) * calculate_Kr(TEW, Dei, REW_, scale_factor) * (s_Kcmax * Kcmax_ - Kcb), fewp * s_Kcmax * Kcmax_)\n",
    "        \n",
    "        # Prepare coefficients for evapotranspiration\n",
    "        Kti = xr_minimum(((s_Ze * Ze_ / Zr)**6) * (1 - Dei / TEW) / xr_maximum(1 - Dr / TAW, 0.001), 1)\n",
    "        Ktp = xr_minimum(((s_Ze * Ze_ / Zr)**6) * (1 - Dep / TEW) / xr_maximum(1 - Dr / TAW, 0.001), 1)\n",
    "        Tei = Kti * Ks * Kcb * weather_cube['ET0'].sel({'time': dates[i]}) / 1000\n",
    "        Tep = Ktp * Ks * Kcb * weather_cube['ET0'].sel({'time': dates[i]}) / 1000\n",
    "        \n",
    "        # Update depletions\n",
    "        Dei = xr.where(fewi > 0, xr_minimum(xr_maximum(Dei + (weather_cube['ET0'].sel({'time': dates[i]}) / 1000) * Kei / fewi + Tei - diff_rei, 0), TEW), xr_minimum(xr_maximum(Dei + Tei - diff_rei, 0), TEW))\n",
    "        Dep = xr.where(fewp > 0, xr_minimum(xr_maximum(Dep + (weather_cube['ET0'].sel({'time': dates[i]}) / 1000) * Kep / fewp + Tep - diff_rep, 0), TEW), xr_minimum(xr_maximum(Dep + Tep - diff_rep, 0), TEW))\n",
    "        \n",
    "        De = (Dei * fewi + Dep * fewp) / (fewi + fewp)\n",
    "        # De = xr.where(De.isfinite(), De, Dei * (s_FW * FW_ / 100) + Dep * (1 - (s_FW * FW_ / 100)))  #================= find replacement for .isfinite() method !!!!!!!!!\n",
    "        \n",
    "        # Evaporation\n",
    "        model_outputs['E'].loc[{'time': dates[i]}]  = xr_maximum((Kei + Kep) * weather_cube['ET0'].sel({'time': dates[i]}) / 1000, 0)\n",
    "        \n",
    "        # Transpiration\n",
    "        model_outputs['Tr'].loc[{'time': dates[i]}]  = Kcb * Ks * weather_cube['ET0'].sel({'time': dates[i]}) / 1000\n",
    "        \n",
    "        # Potential evapotranspiration and evaporative fraction ??\n",
    "        \n",
    "        # Update depletions (root and deep zones) at the end of the day\n",
    "        Dr = xr_minimum(xr_maximum(Dr + model_outputs['E'].loc[{'time': dates[i]}] + model_outputs['Tr'].loc[{'time': dates[i]}] - diff_dr, 0), TAW)\n",
    "        Dd = xr_minimum(xr_maximum(Dd + diff_dr, 0), TDW)\n",
    "        \n",
    "        # Write outputs\n",
    "        model_outputs['Irr'].loc[{'time': dates[i]}] = Irrig\n",
    "        \n",
    "        # Update variable_t1 values\n",
    "        for variable in calculation_variables_t1:\n",
    "            variables_t1[variable] = variables_t2[variable].copy(deep = True)\n",
    "        \n",
    "        print('day ', i+1, '/', len(dates), '   ', end = '\\r')\n",
    "    \n",
    "    # Scale the model_outputs variable to save in int16 format\n",
    "    model_outputs = model_outputs * 1000\n",
    "    \n",
    "    # Write encoding dict\n",
    "    encoding_dict = {}\n",
    "    for variable in list(model_outputs.keys()):\n",
    "        encod = {}\n",
    "        encod['dtype'] = 'i2'\n",
    "        encoding_dict[variable] = encod\n",
    "    \n",
    "    # Save model outputs to netcdf\n",
    "    model_outputs.to_netcdf(save_path, encoding = encoding_dict)\n",
    "    \n",
    "    return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/auclairj/anaconda3/envs/modspa_pixel/lib/python3.10/site-packages/distributed/node.py:182: UserWarning: Port 8787 is already in use.\n",
      "Perhaps you already have a cluster running?\n",
      "Hosting the HTTP server on port 37667 instead\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "    <div style=\"width: 24px; height: 24px; background-color: #e1e1e1; border: 3px solid #9D9D9D; border-radius: 5px; position: absolute;\"> </div>\n",
       "    <div style=\"margin-left: 48px;\">\n",
       "        <h3 style=\"margin-bottom: 0px;\">Client</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">Client-29d78c7f-2653-11ee-9cc7-00155d33b451</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "\n",
       "        <tr>\n",
       "        \n",
       "            <td style=\"text-align: left;\"><strong>Connection method:</strong> Cluster object</td>\n",
       "            <td style=\"text-align: left;\"><strong>Cluster type:</strong> distributed.LocalCluster</td>\n",
       "        \n",
       "        </tr>\n",
       "\n",
       "        \n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard: </strong> <a href=\"http://127.0.0.1:37667/status\" target=\"_blank\">http://127.0.0.1:37667/status</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\"></td>\n",
       "            </tr>\n",
       "        \n",
       "\n",
       "        </table>\n",
       "\n",
       "        \n",
       "\n",
       "        \n",
       "            <details>\n",
       "            <summary style=\"margin-bottom: 20px;\"><h3 style=\"display: inline;\">Cluster Info</h3></summary>\n",
       "            <div class=\"jp-RenderedHTMLCommon jp-RenderedHTML jp-mod-trusted jp-OutputArea-output\">\n",
       "    <div style=\"width: 24px; height: 24px; background-color: #e1e1e1; border: 3px solid #9D9D9D; border-radius: 5px; position: absolute;\">\n",
       "    </div>\n",
       "    <div style=\"margin-left: 48px;\">\n",
       "        <h3 style=\"margin-bottom: 0px; margin-top: 0px;\">LocalCluster</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">c5e1ba98</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard:</strong> <a href=\"http://127.0.0.1:37667/status\" target=\"_blank\">http://127.0.0.1:37667/status</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Workers:</strong> 4\n",
       "                </td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total threads:</strong> 8\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total memory:</strong> 23.47 GiB\n",
       "                </td>\n",
       "            </tr>\n",
       "            \n",
       "            <tr>\n",
       "    <td style=\"text-align: left;\"><strong>Status:</strong> running</td>\n",
       "    <td style=\"text-align: left;\"><strong>Using processes:</strong> True</td>\n",
       "</tr>\n",
       "\n",
       "            \n",
       "        </table>\n",
       "\n",
       "        <details>\n",
       "            <summary style=\"margin-bottom: 20px;\">\n",
       "                <h3 style=\"display: inline;\">Scheduler Info</h3>\n",
       "            </summary>\n",
       "\n",
       "            <div style=\"\">\n",
       "    <div>\n",
       "        <div style=\"width: 24px; height: 24px; background-color: #FFF7E5; border: 3px solid #FF6132; border-radius: 5px; position: absolute;\"> </div>\n",
       "        <div style=\"margin-left: 48px;\">\n",
       "            <h3 style=\"margin-bottom: 0px;\">Scheduler</h3>\n",
       "            <p style=\"color: #9D9D9D; margin-bottom: 0px;\">Scheduler-205d3daa-9675-4977-8d69-da12e45dc32c</p>\n",
       "            <table style=\"width: 100%; text-align: left;\">\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Comm:</strong> tcp://127.0.0.1:42111\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Workers:</strong> 4\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Dashboard:</strong> <a href=\"http://127.0.0.1:37667/status\" target=\"_blank\">http://127.0.0.1:37667/status</a>\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total threads:</strong> 8\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Started:</strong> Just now\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total memory:</strong> 23.47 GiB\n",
       "                    </td>\n",
       "                </tr>\n",
       "            </table>\n",
       "        </div>\n",
       "    </div>\n",
       "\n",
       "    <details style=\"margin-left: 48px;\">\n",
       "        <summary style=\"margin-bottom: 20px;\">\n",
       "            <h3 style=\"display: inline;\">Workers</h3>\n",
       "        </summary>\n",
       "\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 0</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:43845\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 2\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://127.0.0.1:44805/status\" target=\"_blank\">http://127.0.0.1:44805/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 5.87 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:46421\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /tmp/dask-scratch-space/worker-t_o8kxq0\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: 1</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tcp://127.0.0.1:34535\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 2\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://127.0.0.1:33207/status\" target=\"_blank\">http://127.0.0.1:33207/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 5.87 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tcp://127.0.0.1:36817\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /tmp/dask-scratch-space/worker-ekfvxctp\n",