Skip to content
Snippets Groups Projects
dev_samir_xarray.ipynb 42.1 KiB
Newer Older
{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "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": null,
23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453
   "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, max_GB: int = 2) -> 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",
    "    # Manage loading of data based on disk size of inputs\n",
    "    if ndvi_cube.nbytes < max_GB * (1024)**3:\n",
    "        ndvi_cube.load()\n",
    "        \n",
    "    if weather_cube.nbytes < max_GB * (1024)**3:\n",
    "        weather_cube.load()\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[i]})/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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = Client()\n",
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = '/mnt/e/DATA/DEV_inputs_test'\n",
    "\n",
    "ndvi_path = data_path + os.sep + 'ndvi.nc'\n",
    "weather_path = data_path + os.sep + 'weather.nc'\n",
    "land_cover_path = data_path + os.sep + 'land_cover.nc'\n",
    "json_config_file = '/home/auclairj/GIT/modspa-pixel/config/config_modspa.json'\n",
    "param_file = '/home/auclairj/GIT/modspa-pixel/parameters/csv_files/params_samir_test.csv'\n",
    "soil_path = data_path + os.sep + 'soil.nc'\n",
    "save_path = data_path + os.sep + 'outputs.nc'\n",
    "\n",
    "chunk_size = {'x': -1, 'y': -1, 'time': -1}\n",
    "\n",
    "\n",
    "run_samir(json_config_file, param_file, ndvi_path, weather_path, soil_path, land_cover_path, chunk_size, save_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "modspa_pixel",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}