modspa_samir.py

# -*- coding: UTF-8 -*-
# Python
"""
11-07-2023
@author: jeremy auclair

Usage of the SAMIR model in the Modspa framework. 

"""

import os  # os management
import xarray as xr  # to manage datasets
import rioxarray  # to better manage dataset projections
from numba import njit, types, boolean, uint8, uint16, int16, float32, int64, set_num_threads  # to compile functions in nopython mode for faster calculation
from numba.typed import List  #type: ignore, to type lists
import numpy as np  # for vectorized maths
import pandas as pd  # to manage dataframes
import netCDF4 as nc  # to efficiently read and write netCDF files
from tqdm import tqdm  # to show a progress bar
from psutil import virtual_memory  # to check available ram
from psutil import cpu_count  # to get number of physical cores available
from gc import collect  # to free up unused memory
from modspa_pixel.parameters.params_samir_class import samir_parameters  # to load SAMIR parameters
from modspa_pixel.source.code_toolbox import format_Byte_size  # to print memory requirements


def rasterize_samir_parameters(csv_param_file: str, land_cover_raster: str, irrigation_raster: str, init_RU_raster: str) -> tuple[np.ndarray, np.dtype]:
    """
    Creates a dictionnary containing raster parameter values and the scale factors from the csv parameter file
    and the land cover raster. For each parameter, the function loops on land cover classes to fill the raster.
    
    Before creating the dictionnary, it updates the parameter ``DataFrame`` and verifies the parameter range with
    the ``test_samir_parameter()`` function.

    Arguments
    =========

    1. csv_param_file: ``str``
        path to csv paramter file
    2. land_cover_raster: ``str``
        path to land cover netcdf raster
    3. irrigation_raster: ``str``
        path to irrigation mask raster
    4. init_RU_raster: ``str``
        path to initial soil water content. None if no input.

    Returns
    =======
    
    1. parameter_dict: ``Dict[str, np.ndarray]``
        the dictionnary containing all the rasterized Parameters
        and the scale factors

    """
    
    print('\nLoading SAMIR parameters')
    
    # Get path of min max csv file
    min_max_param_file = os.path.join(os.path.dirname(csv_param_file), 'params_samir_min_max.csv')

    # Load samir params into an object
    table_param = samir_parameters(csv_param_file)

    # Set general variables
    class_count = table_param.table.shape[1] - 2  # remove dtype and Default columns

    # Open land cover raster
    land_cover = xr.open_dataarray(land_cover_raster).to_numpy()
    shape = land_cover.shape

    # Modify parameter table based on its content
    table_param.test_samir_parameter(min_max_param_file)

    # If test_samir_parameter returns 0, parameters are incorrect
    if type(table_param.table) != pd.DataFrame:
        import sys  # to exit script
        print(f'\nModify the SAMIR parameter file: {csv_param_file}\n')
        sys.exit()
    
    # Loop through parameters to count them
    count = 0
    for parameter in table_param.table.index[1:]:
        if parameter == 'Irrig_auto' and table_param.table.at[parameter, 'load_raster'] == 1:
            pass
        else:
            count+=1
    
    parameter_array = np.zeros((count,) + shape, dtype = np.float32, order = 'C')
    param_list = []

    i = 0
    # Loop on samir parameters and create
    for parameter in table_param.table.index[1:]:
        
        if parameter == 'Irrig_auto' and table_param.table.at[parameter, 'load_raster'] == 1:
            
            # Load Irrig mask
            Irrig_auto = np.ascontiguousarray(xr.open_dataarray(irrigation_raster, decode_coords = 'all').astype(np.bool_).values, dtype = np.bool_)
            
            continue
        
        elif parameter == 'Init_RU' and table_param.table.at[parameter, 'load_raster'] == 1:
            
            param_list.append(parameter)
            
            # Load initial soil water content
            Init_RU = (xr.open_dataarray(init_RU_raster) / 1000).astype(np.float32).values
            i+=1
            
            continue
        
        # If scale_factor == 0, then the parameter does not have to be spatialized, it can stay scalar
        elif parameter == 'Kcmax':
            
            param_list.append(parameter)
            
            # Take Default parameter value
            parameter_array[i] = np.ones(shape = shape, dtype = np.float32) * np.float32(table_param.table.at[parameter, 'Default'])
            i+=1         
            
            continue
        
        # Check data type based on parameter
        if (parameter != 'Irrig_auto') and (parameter != 'Init_RU'):
            dtype = np.float32
            param_list.append(parameter)
            
        # Create 2D array from land cover
        value = land_cover.copy().astype(dtype)
        
        # Loop on classes to set parameter values for each class
        for class_val, class_name in zip(range(1, class_count + 1), table_param.table.columns[2:]):
            
            # Get parameter value
            param_value = table_param.table.at[parameter, class_name]

            # Load parameter value for each class
            value = np.where(land_cover == class_val, np.float32(param_value), value)
        
        # Assign parameter value
        if parameter == 'Irrig_auto':
            Irrig_auto = value.astype(np.bool_)
        elif parameter == 'Init_RU':
            Init_RU = value.astype(np.float32)
        else:
            parameter_array[i] = value.astype(dtype)
            i+=1
    
    # Assure parameter array is contiguous in memory
    parameter_array = np.ascontiguousarray(parameter_array, dtype = np.float32)
    
    # Change param_list type
    param_list = List(param_list)
    
    return parameter_array, Irrig_auto, Init_RU, param_list


def prepare_output_dataset(ndvi_path: str, dimensions: dict[str, int], scaling_dict: dict[str, int] = {'E': 1000, 'Tr': 1000, 'SWCe': 1000, 'SWCr': 1000, 'DP': 100, 'Irr': 100}, additional_outputs: dict[str, int] = None) -> xr.Dataset:
    """
    Creates the `xarray Dataset` containing the outputs of the SAMIR model that will be saved.
    Additional variables can be saved by adding their names to the `additional_outputs` list.

    Arguments
    =========
    
    1. ndvi_path: ``str``
        path to ndvi cube
    2. dimensions: ``Dict[str, int]``
        frozen dictionnary containing the dimensions of the output dataset
    3. scaling_dict: ``Dict[str, int]`` ``default = {'E': 1000, 'Tr': 1000, 'SWCe': 1000, 'SWCr': 1000, 'DP': 100, 'Irr': 100}``
        scaling dictionnary for the nominal outputs
    4. additional_outputs: ``Dict[str, int]``
        list of additional variable names to be saved

    Returns
    =======
    
    1. model_outputs: ``xr.Dataset``
        model outputs to be saved
    """

    # Evaporation and Transpiraion
    model_outputs = xr.open_dataset(ndvi_path).drop_vars(['NDVI']).copy(deep = True)
    model_outputs = model_outputs.drop_sel(time = model_outputs.time)
    
    # Add scaling dictionnary to global attribute
    model_outputs.attrs['name'] = 'ModSpa Pixel SAMIR output'
    model_outputs.attrs['description'] = 'Outputs of the ModSpa SAMIR (FAO-56) outputs at the pixel scale. Variables are upscaled to be stored as integers.'
    model_outputs.attrs['scaling'] = str(scaling_dict)

    model_outputs['E'] = (dimensions, np.zeros(tuple(dimensions[d] for d in list(dimensions)), dtype = np.uint16))
    model_outputs['E'].attrs['unit'] = 'mm'
    model_outputs['E'].attrs['standard_name'] = 'Evaporation'
    model_outputs['E'].attrs['description'] = 'Accumulated daily evaporation in milimeters'
    model_outputs['E'].attrs['scale factor'] = scaling_dict['E']

    model_outputs['Tr'] = (dimensions, np.zeros(tuple(dimensions[d] for d in list(dimensions)), dtype = np.uint16))
    model_outputs['Tr'].attrs['unit'] = 'mm'
    model_outputs['Tr'].attrs['standard_name'] = 'Transpiration'
    model_outputs['Tr'].attrs['description'] = 'Accumulated daily plant transpiration in milimeters'
    model_outputs['Tr'].attrs['scale factor'] = scaling_dict['Tr']

    # Soil Water Content
    model_outputs['SWCe'] = (dimensions, np.zeros(tuple(dimensions[d] for d in list(dimensions)), dtype = np.uint16))
    model_outputs['SWCe'].attrs['unit'] = '%'
    model_outputs['SWCe'].attrs['standard_name'] = 'Soil Water Content of the evaporative layer'
    model_outputs['SWCe'].attrs['description'] = 'Soil water content of the evaporative layer in milimeters'
    model_outputs['SWCe'].attrs['scale factor'] = scaling_dict['SWCe']

    model_outputs['SWCr'] = (dimensions, np.zeros(tuple(dimensions[d] for d in list(dimensions)), dtype = np.uint16))
    model_outputs['SWCr'].attrs['unit'] = '%'
    model_outputs['SWCr'].attrs['standard_name'] = 'Soil Water Content of the root layer'
    model_outputs['SWCr'].attrs['description'] = 'Soil water content of the root layer in milimeters'
    model_outputs['SWCr'].attrs['scale factor'] = scaling_dict['SWCr']

    # Irrigation
    model_outputs['Irr'] = (dimensions, np.zeros(tuple(dimensions[d] for d in list(dimensions)), dtype = np.uint16))
    model_outputs['Irr'].attrs['unit'] = 'mm'
    model_outputs['Irr'].attrs['standard_name'] = 'Irrigation'
    model_outputs['Irr'].attrs['description'] = 'Simulated daily irrigation in milimeters'
    model_outputs['Irr'].attrs['scale factor'] = scaling_dict['Irr']

    # Deep Percolation
    model_outputs['DP'] = (dimensions, np.zeros(tuple(dimensions[d] for d in list(dimensions)), dtype = np.uint16))
    model_outputs['DP'].attrs['unit'] = 'mm'
    model_outputs['DP'].attrs['standard_name'] = 'Deep Percolation'
    model_outputs['DP'].attrs['description'] = 'Simulated daily Deep Percolation in milimeters'
    model_outputs['DP'].attrs['scale factor'] = scaling_dict['DP']

    if additional_outputs:
        for var, scale in zip(additional_outputs.keys(), additional_outputs.values()):
            model_outputs[var] = (dimensions, np.zeros(tuple(dimensions[d] for d in list(dimensions)), dtype = np.int16))
            model_outputs[var].attrs['scale factor'] = scale

    return model_outputs


@njit(nogil = True, cache = True)
def find_index(numba_list: list, value: any) -> int:
    """
    Equivalent .index() method of python lists adapted for numba lists.

    Arguments
    =========

    1. numba_list: ``list``
        input numba typed list
    2. value: ``any``
        value to find in list

    Returns
    =======

    1. index: ``int``
        index of value in list
    """
    
    # Loop through list elements
    for i in range(len(numba_list)):
        if numba_list[i] == value:
            # Return index of value in numba_list
            return i
    
    # Return -1 if element is not found
    return -1


@njit((float32[:,:,:], types.ListType(types.unicode_type), uint8[:,:], float32[:,:], float32[:,:], float32[:,:], int64, int64), nogil = True, parallel = True, fastmath = True, cache = True)
def get_starting_conditions(parameter_array: np.ndarray, param_list: list[str], NDVI: np.ndarray, Init_RU: np.ndarray, Wfc: np.ndarray, Wwp: np.ndarray, y_size: int, x_size: int) -> tuple[np.ndarray]:
    """
    Calculate SAMIR starting variables.

    Arguments
    =========

    1. parameter_array: ``np.ndarray``
        parameter array
    2. param_list: ``list[str]``
        list of param names
    3. NDVI: ``np.ndarray``
        NDVI for first date
    4. Init_RU: ``np.ndarray``
        initial RU
    5. Wfc: ``np.ndarray``
        field capacity array
    6. Wwp: ``np.ndarray``
        weilting point array
    7. y_size: ``int``
        y_size of arrays
    8. x_size: ``int``
        x_size of arrays

    Returns
    =======

    1. output_arrays: ``tuple[np.ndarray]``
        E0, Tr0, Zr0, Dei0, Dep0, Dr0, Dd0, Irrig_test
    """
    
    E0 = np.zeros((y_size, x_size), dtype = np.uint16)
    Tr0 = np.zeros((y_size, x_size), dtype = np.uint16)
    Irrig_test = np.zeros((y_size, x_size), dtype = np.bool_)
    Zr0 = np.maximum(parameter_array[find_index(param_list, 'minZr')] + (np.minimum(np.maximum(parameter_array[find_index(param_list, 'Fslope')] * (NDVI / np.float32(255)).astype(np.float32) + parameter_array[find_index(param_list, 'Foffset')], np.float32(0)), parameter_array[find_index(param_list, 'FCmax')]) / parameter_array[find_index(param_list, 'FCmax')]) * (parameter_array[find_index(param_list, 'maxZr')] - parameter_array[find_index(param_list, 'minZr')]), parameter_array[find_index(param_list, 'Ze')] + np.float32(0.001))
    Dei0 = (Wfc - Wwp) * parameter_array[find_index(param_list, 'Ze')] * (np.float32(1) - Init_RU)
    Dep0 = (Wfc - Wwp) * parameter_array[find_index(param_list, 'Ze')] * (np.float32(1) - Init_RU)
    Dr0 = (Wfc - Wwp) * Zr0 * (np.float32(1) - Init_RU)
    Dd0 = (Wfc - Wwp) * (parameter_array[find_index(param_list, 'Zsoil')] - Zr0) * (np.float32(1) - Init_RU)
    
    return E0, Tr0, Zr0, Dei0, Dep0, Dr0, Dd0, Irrig_test


@njit((float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:]), nogil = True, parallel = True, fastmath = True, cache = True)
def calculate_diff_re(TAW: np.ndarray, Dr: np.ndarray, Zr: np.ndarray, RUE: np.ndarray, De: np.ndarray, Wfc: np.ndarray, Ze: np.ndarray, DiffE: np.ndarray) -> np.ndarray:
    """
    Calculates the diffusion between the top soil layer and the root layer. Uses numba for faster and parallel calculation.

    Arguments
    =========
    
    1. TAW: ``np.ndarray``
        water capacity of root layer
    2. Dr: ``np.ndarray``
        depletion of root layer
    3. Zr: ``np.ndarray``
        height of root layer
    4. RUE: ``np.ndarray``
        total available surface water = (Wfc-Wwp)*Ze
    5. De: ``np.ndarray``
        depletion of the evaporative layer
    6. Wfc: ``np.ndarray``
        field capacity
    7. Ze: ``np.ndarray``
        height of evaporative layer (paramter)
    8. DiffE: ``np.ndarray``
        diffusion coefficient between evaporative
        and root layers (unitless, parameter)

    Returns
    =======
    
    1. diff_re: ``np.ndarray``
        the diffusion between the top soil layer and
        the root layer
    """

    # Temporary variables to make calculation easier to read
    # TODO: check equation
    tmp1 = (((TAW - Dr) / Zr - (RUE - De) / Ze) / Wfc) * DiffE
    tmp2 = ((TAW * Ze) - (RUE - De - Dr) * Zr) / (Zr + Ze) - Dr
    
    # Calculate diffusion according to SAMIR equation
    # Return zero values where the 'DiffE' parameter is equal to 0
    return np.where(DiffE == 0, np.float32(0), np.where(tmp1 < 0, np.maximum(tmp1, tmp2), np.minimum(tmp1, tmp2)))


@njit((float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:]), nogil = True, parallel = True, fastmath = True, cache = True)
def calculate_diff_dr(TAW: np.ndarray, TDW: np.ndarray, Dr: np.ndarray, Zr: np.ndarray, Dd: np.ndarray, Wfc: np.ndarray, Zsoil: np.ndarray, DiffR: np.ndarray) -> np.ndarray:
    """
    Calculates the diffusion between the root layer and the deep layer. Uses numba for faster and parallel calculation.

    Arguments
    =========

    1. TAW: ``np.ndarray``
        water capacity of root layer
    2. TDW: ``np.ndarray``
        water capacity of deep layer
    3. Dr: ``np.ndarray``
        depletion of root layer
    4. Zr: ``np.ndarray``
        height of root layer
    5. Dd: ``np.ndarray``
        depletion of deep layer
    6. Wfc: ``np.ndarray``
        field capacity
    7. Zsoil: ``np.ndarray``
        total height of soil (paramter)
    8. DiffR: ``np.ndarray``
        Diffusion coefficient between root
        and deep layers (unitless, parameter)

    Returns
    =======
    
    1. Diff_dr: ``np.ndarray``
        the diffusion between the root layer and the
        deep layer
    """

    # Temporary variables to make calculation easier to read
    tmp1 = (((TDW - Dd) / (Zsoil - Zr) - (TAW - Dr) / Zr) / Wfc) * DiffR
    tmp2 = ((TDW * Zr - (TAW - Dr - Dd) * (Zsoil - Zr)) / Zsoil) - Dd

    # Calculate diffusion according to SAMIR equation
    # Return zero values where the 'DiffR' parameter is equal to 0
    return np.where(DiffR == 0, np.float32(0), np.where(tmp1 < 0, np.maximum(tmp1, tmp2), np.minimum(tmp1, tmp2)))


@njit((float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:]), nogil = True, parallel = True, fastmath = True, cache = True)
def calculate_W(TEW: np.ndarray, Dei: np.ndarray, Dep: np.ndarray, fewi: np.ndarray, fewp: np.ndarray) -> np.ndarray:
    """
    Calculate W, the weighting factor to split the energy available
    for evaporation depending on the difference in water availability
    in the two evaporation components, ensuring that the larger and
    the wetter, the more the evaporation occurs from that component.
    Uses numba for faster and parallel calculation.

    Arguments
    =========

    1. TEW: ``np.ndarray``
        water capacity of evaporative layer
    2. Dei: ``np.ndarray``
        depletion of the evaporative layer
        (irrigation part)
    3. Dep: ``np.ndarray``
        depletion of the evaporative layer
        (precipitation part)
    4. fewi: ``np.ndarray``
        soil fraction which is wetted by irrigation
        and exposed to evaporation
    5. fewp: ``np.ndarray``
        soil fraction which is wetted by precipitation
        and exposed to evaporation

    Returns
    =======
    
    1. W: ``np.ndarray``
        weighting factor W
    """

    # Calculate the weighting factor to split the energy available for evaporation
    # Equation: W = 1 / (1 + (fewp * (TEW - Dep) / fewi * (TEW - Dei)))
    # Return W
    # Equation: W = where(fewi * (TEW - Dei) > 0, W, 0)
    return np.where(fewi * (TEW - Dei) > 0, 1 / (1 + (fewp * (TEW - Dep) / fewi * (TEW - Dei))), np.float32(0))


@njit((float32[:,:], float32[:,:], float32[:,:]), nogil = True, parallel = True, fastmath = True, cache = True)
def calculate_Kr(TEW: np.ndarray, De: np.ndarray, REW: np.ndarray) -> np.ndarray:
    """
    calculates of the reduction coefficient for evaporation dependent
    on the amount of water in the soil using the FAO-56 method. Uses numba for faster and parallel calculation.

    Arguments
    =========
    
    1. TEW: ``np.ndarray``
        water capacity of evaporative layer
    2. De: ``np.ndarray``
        depletion of evaporative layer
    3. REW: ``np.ndarray``
        fixed readily evaporable water

    Returns
    =======
    
    1. Kr: ``np.ndarray``
        Kr coefficient
    """

    # Return Kr
    return np.maximum(np.float32(0), np.minimum((TEW - De) / (TEW - REW), np.float32(1)))


@njit((float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:]), nogil = True, parallel = True, fastmath = True, cache = True)
def calculate_Ks(Dr: np.ndarray, TAW: np.ndarray, p: np.ndarray, E0: np.ndarray, Tr0: np.ndarray) -> np.ndarray:
    """
    Calculate Ks coefficient after day 1. Uses numba for faster and parallel calculation.

    Arguments
    =========

    1. Dr: ``np.ndarray``
        depletion of the root layer
    2. TAW: ``np.ndarray``
        water capacity of the root layer
    3. p: ``np.ndarray``
        fraction of TAW available for plant without inducing stress
    4. E0: ``np.ndarray``
        surface evaporation of previous day
    5. Tr0: ``np.ndarray``
        plant transpiration of previous day

    Returns
    =======

    1. Ks: ``np.ndarray``
        Ks coefficient
    """

    # Return Ks
    # Equation: Ks = min((TAW - Dr) / (TAW * (1 - (p + 0.04 * (5 - (E0 + Tr0))))), 1)
    return np.minimum((TAW - Dr) / (TAW * (np.float32(1) - (p + 0.04 * (np.float32(5) - (E0 + Tr0))))), np.float32(1)).astype(np.float32)


@njit((float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:]), nogil = True, parallel = True, fastmath = True, cache = True)
def calculate_Ke(W: np.ndarray, De: np.ndarray, TEW: np.ndarray, REW: np.ndarray, Kcmax: np.ndarray, Kcb: np.ndarray, few: np.ndarray) -> np.ndarray:
    """
    Calculate the evaporation Ke coefficient.

    Arguments
    =========

    1. W: ``np.ndarray``
        weighting factor to split the energy available
        for evaporation
    2. De: ``np.ndarray``
        Dei or Dep, depletion of the evaporative layer
    3. TEW: ``np.ndarray``
        water capacity of the evaporative layer
    4. REW: ``np.ndarray``
        fixed readily evaporable water
    5. Kcmax: ``np.float32``
        maximum possible evaporation in the atmosphere (scalar)
    6. Kcb: ``np.ndarray``
        crop coefficient
    7. few: ``np.ndarray``
        fewi or fewp, soil fraction which is wetted by
        irrigation or precipitation and exposed to evaporation

    Returns
    =======

    1. Ke: ``np.ndarray``
        evaporation coefficient
    """
    
    # Equation: Kei = np.minimum(W * Kri * (Kc_max - Kcb), fewi * Kc_max)
    # Equation: Kep = np.minimum((1 - W) * Krp * (Kc_max - Kcb), fewp * Kc_max)
    # Return Ke
    return np.minimum(W * calculate_Kr(TEW, De, REW) * (Kcmax - Kcb), few * Kcmax)


@njit((float32[:,:], float32[:,:], float32[:,:], float32[:,:], boolean[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], boolean[:,:], float32[:,:], float32[:,:]), nogil = True, parallel = True, fastmath = True, cache = True)
def calculate_irrig(Dr: np.ndarray, TAW: np.ndarray, Rain: np.ndarray, Kcb: np.ndarray, Irrig_auto: np.ndarray, Lame_max: np.ndarray, Lame_min: np.ndarray, Kcb_min_start_Irrig: np.ndarray, frac_Kcb_stop_irrig: np.ndarray, Irrig_test: np.ndarray, frac_TAW: np.ndarray, Kcb_max_obs: np.ndarray) -> np.ndarray:
    """
    Calculate automatic irrigation after day one. Uses numba for faster and parallel calculation.

    Arguments
    =========

    1. Dr: ``np.ndarray``
        depletion of the root layer
    2. TAW: ``np.ndarray``
        water capacity of the root layer
    3. Rain: ``np.ndarray``
        precipitation of the current day
    4. Kcb: ``np.ndarray``
        crop coefficient value
    5. Irrig_auto: ``np.ndarray`` ``boolean``
        parameter for automatic irrigation
    6. Lame_max: ``np.ndarray``
        maximum amount of possible irrigation in mm
    7. Lame_min: ``np.ndarray``
        minimum amount of possible irrigation in mm
    8. Kcb_min_start_Irrig: ``np.ndarray``
        minimum value of Kcb under which no irrigation is allowed
    9. frac_Kcb_stop_irrig: ``np.ndarray``
        Kcb threshold to stop irrigation after reaching maximum
    10. Irrig_test: ``np.ndarray`` ``boolean``
        boolean array that is true after the Kcb has reached
        frac_Kcb_stop_irrig of its maximum
    11. frac_TAW: ``np.ndarray``
        fraction of depleted TAW under which irrigation is triggered
    12: Kcb_max_obs: ``np.ndarray``
        observed maximum value of Kcb
        

    Returns
    =======

    1. Irrig: ``np.ndarray``
        simulated irrigation for the current day
    """
    
    # First step: calculate irrigation to fill the depletion up to a maximum value of Lame_max
    Irrig = Irrig_auto * np.maximum(np.minimum(Dr - Rain, Lame_max), Lame_min)
    
    # Second step: only keep irrigation where the depletion is higher than a fraction (frac_TAW) of TAW and Kcb is higher than a minimum value
    Irrig = np.where((Dr > frac_TAW * TAW) & (Kcb > Kcb_min_start_Irrig), Irrig, np.float32(0))
    
    # Third step: if Kcb has already gone higher than a fraction (frac_Kcb_stop_irrig) of its maximum value,
    # turn off irrigation when goes under that fraction of maximum Kcb
    return np.where((Irrig_test) & (Kcb < frac_Kcb_stop_irrig * Kcb_max_obs), np.float32(0), Irrig)
    

@njit((float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:]), nogil = True, parallel = True, fastmath = True, cache = True)
def calculate_Te(De: np.ndarray, Dr: np.ndarray, Ks: np.ndarray, Kcb: np.ndarray, Ze: np.ndarray, Zr: np.ndarray, TEW: np.ndarray, TAW: np.ndarray, ET0: np.ndarray) -> np.ndarray:
    """
    Calculate Te (root uptake of water) coefficient for current day. Uses numba for faster and parallel calculation.

    Arguments
    =========

    1. De: ``np.ndarray``
        Dei or Dep, depletion of the evaporative layer
    2. Dr: ``np.ndarray``
        depletion of the roor layer
    3. Ks: ``np.ndarray``
        stress coefficient
    4. Kcb: ``np.ndarray``
        crop coefficient
    5. Ze: ``np.ndarray``
        height of the evaporative layer
    6. Zr: ``np.ndarray``
        heigh of the root layer
    7. TEW: ``np.ndarray``
        water capacity of the evaporative layer
    8. TAW: ``np.ndarray``
        water capacity of the root layer
    9. ET0: ``np.ndarray``
        reference evapotranspiration

    Returns
    =======

    1. Te: ``np.ndarray``
        Te coefficient
    """
    
    # Equation: Kt = min( ((Ze / Zr)**0.6) * (1 - De/TEW) / max(1 - Dr / TAW, 0.01), 1)
    # Equation: Te = Kt * Ks * Kcb * ET0
    return (np.minimum(((Ze / Zr) ** 0.6) * (np.float32(1) - De / TEW) / np.maximum(np.float32(1) - Dr / TAW, 0.001), np.float32(1)) * Ks * Kcb * ET0).astype(np.float32)


@njit((float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:]), nogil = True, parallel = True, fastmath = True, cache = True)
def update_De_from_Diff(De : np.ndarray, few: np.ndarray, Ke: np.ndarray, Te: np.ndarray, Diff_re: np.ndarray, TEW: np.ndarray, ET0: np.ndarray) -> np.ndarray:
    """
    Last update step for Dei and Dep depletions. Uses numba for faster and parallel calculation.

    Arguments
    =========

    1. De: ``np.ndarray``
        Dei or Dep, depletion of the evaporative layer
    2. few: ``np.ndarray``
        fewi or fewp, soil fraction which is wetted by
        irrigation or precipitation and exposed to evaporation
    3. Ke: ``np.ndarray``
        Kei or Kep, evaporation coefficient for soil fraction
        irrigated or rainfed and exposed to evaporation
    4. Te: ``np.ndarray``
        root uptake of water
    5. Diff_re: ``np.ndarray``
        dffusion between the root and evaporation layers
    6. TEW: ``np.ndarray``
        water capacity of the evaporative layer
    7. ET0: ``np.ndarray``
        reference evapotranspiration of the current day

    Returns
    =======

    De: ``np.ndarray``
        updated Dei or Dep
    """
    
    # Update Dei and Dep depletions
    # Equation: De = where(few > 0, min(max(De + ET0 * Ke / few + Te - Diff_re, 0), TEW), min(max(De + Te - Diff_re, 0), TEW))
    return np.where(few > np.float32(0), np.minimum(np.maximum(De + ET0 * Ke / few + Te - Diff_re, np.float32(0)), TEW), np.minimum(np.maximum(De + Te - Diff_re, np.float32(0)), TEW))
        

@njit((float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:]), nogil = True, parallel = True, fastmath = True, cache = True)
def update_Dr_from_root(Wfc: np.ndarray, Wwp: np.ndarray, Zr: np.ndarray, Zsoil: np.ndarray, Dr0: np.ndarray, Dd0: np.ndarray, Zr0: np.ndarray) -> np.ndarray:
    """
    Return the updated depletion for the root layer. Uses numba for faster and parallel calculation.

    Arguments
    =========

    1. Wfc: ``np.ndarray``
        field capacity
    2. Wwp: ``np.ndarray``
        field wilting point
    3. Zr: ``np.ndarray``
        root depth for current day
    4. Zsoil: ``np.ndarray``
        total soil depth (parameter)
    5. Dr0: ``np.ndarray``
        depletion of the root layer for previous day
    6. Dd0: ``np.ndarray``
        depletion of the deep laye for previous day
    7. Zr0: ``np.ndarray``
        root layer height for previous day

    Returns
    =======
    
    1. output: ``np.ndarray``
        updated depletion for the root layer
    """

    # Temporary variables to make calculation easier to read
    tmp1 = np.maximum(Dr0 + Dd0 * ((Wfc - Wwp) * (Zr - Zr0)) / ((Wfc - Wwp) * (Zsoil - Zr0)), np.float32(0))
    tmp2 = np.maximum(Dr0 + Dr0 * ((Wfc - Wwp) * (Zr - Zr0)) / ((Wfc - Wwp) * Zr0), np.float32(0))

    # Return updated Dr
    return np.where(Zr > Zr0, tmp1, tmp2)


@njit((float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:]), nogil = True, parallel = True, fastmath = True, cache = True)
def update_Dd_from_root(Wfc: np.ndarray, Wwp: np.ndarray, Zr: np.ndarray, Zsoil: np.ndarray, Dr0: np.ndarray, Dd0: np.ndarray, Zr0: np.ndarray) -> np.ndarray:
    """
    Return the updated depletion for the deep layer. Uses numba for faster and parallel calculation.

    Arguments
    =========

    1. Wfc: ``np.ndarray``
        field capacity
    2. Wwp: ``np.ndarray``
        field wilting point
    3. Zr: ``np.ndarray``
        root depth for current day
    4. Zsoil: ``np.ndarray``
        total soil depth (parameter)
    5. Dr0: ``np.ndarray``
        depletion of the root layer for previous day
    6. Dd0: ``np.ndarray``
        depletion of the deep laye for previous day
    7. Zr0: ``np.ndarray``
        root layer height for previous day

    Returns
    =======
    
    1. output: ``np.ndarray``
        updated depletion for the deep layer
    """

    # Temporary variables to make calculation easier to read
    tmp1 = np.maximum(Dd0 - Dd0 * ((Wfc - Wwp) * (Zr - Zr0)) / ((Wfc - Wwp) * (Zsoil - Zr0)), np.float32(0))
    tmp2 = np.maximum(Dd0 - Dr0 * ((Wfc - Wwp) * (Zr - Zr0)) / ((Wfc - Wwp) * Zr0), np.float32(0))

    # Return updated Dd
    return np.where(Zr > Zr0, tmp1, tmp2)


@njit((float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:]), nogil = True, parallel = True, fastmath = True, cache = True)
def calculate_SWCe(Dei: np.ndarray, Dep: np.ndarray, fewi: np.ndarray, fewp: np.ndarray, TEW: np.ndarray) -> np.ndarray:
    """
    Calculate the soil water content of the evaporative layer.

    Arguments
    =========

    1. Dei: ``np.ndarray``
        depletion of the evaporative layer for the irrigated part
    2. Dep: ``np.ndarray``
        depletion of the evaporative layer for the rainfed part
    3. fewi: ``np.ndarray``
        soil fraction which is wetted by irrigation and exposed
        to evaporation
    4. fewp: ``np.ndarray``
        soil fraction which is wetted by precipitation and exposed
        to evaporation
    5. TEW: ``np.ndarray``
        water capacity of the evaporative layer

    Returns
    =======

    1. SWCe: ``np.ndarray``
        soil water content of the evaporative layer
    """
    
    # Return SWCe
    return np.where((fewi + fewp) > np.float32(0), np.float32(1) - ((Dei * fewi + Dep * fewp) / ((fewi + fewp) * TEW)), np.float32(1) - ((Dei + Dep) / (np.float32(2) * TEW)))


def calculate_memory_requirement(x_size: int, y_size: int, time_size: int, nb_inputs: int, nb_outputs: int, nb_variables: int, nb_params: int, nb_bytes: int) -> float:
    """
    Calculate memory requirement (GiB) of calculation if all datasets where loaded in memory.
    Used to determine how to divide the datasets in times chunks for more efficient I/O
    operations.

    Arguments
    =========

    1. x_size: ``int``
        x size of dataset (pixels)
    2. y_size: ``int``
        y size of dataset (pixels)
    3. time_size: ``int``
        number of time bands (dates)
    4. nb_inputs: ``int``
        number of input variables
    5. nb_outputs: ``int``
        number of ouput variables
    6. nb_variables: ``int``
        number of calculation variables
    7. nb_params: ``int``
        number of raster parameters
    8. nb_bytes: ``int``
        number of bytes of datatype for inputs and outputs

    Returns
    =======

    1. total_memory_requirement: ``float``
        calculation memory requirement in GiB
    """
    
    # Memory requirement of input datasets
    input_memory_requirement = (x_size * y_size * time_size * nb_inputs * nb_bytes) / (1024**3)
    
    # Memory requirement of calculation variables
    calculation_memory_requirement = (x_size * y_size * (nb_params * 4 + nb_variables * 4)) / (1024**3)  # calculation done in float32, params in float32
    
    # Memory requirement of output datasets
    output_memory_requirement = (x_size * y_size * time_size * nb_outputs * nb_bytes) / (1024**3)
    
    # Total memory requirement
    # TODO : add adjustment factor elsewhere
    total_memory_requirement = (input_memory_requirement + calculation_memory_requirement + output_memory_requirement) * 1.05  # 5% adjustment factor

    return total_memory_requirement


def calculate_time_slices_to_load(x_size: int, y_size: int, time_size: int, nb_inputs: int, nb_outputs: int, nb_variables: int, nb_params: int, nb_bytes: int, available_ram: int) -> tuple[int, int, bool]:
    """
    Calculate how many time slices to load in memory (for input and output data)
    based on available ram and calculation requirements.

    Arguments
    =========

    1. x_size: ``int``
        x size of dataset (pixels)
    2. y_size: ``int``
        y size of dataset (pixels)
    3. time_size: ``int``
        number of time bands (dates)
    4. nb_inputs: ``int``
        number of input variables
    5. nb_outputs: ``int``
        number of ouput variables
    6. nb_variables: ``int``
        number of calculation variables
    7. nb_params: ``int``
        number of raster parameters
    8. nb_bytes: ``int``
        number of bytes of datatype for inputs and outputs
    9. available_ram: ``int``
        available ram for computation in GiB

    Returns
    =======

    1. time_slice: ``int``
        number of times slices to load for
        input and output data
    2. remainder_to_load: ``int``
        remainder of euclidian division for
        the number of time slices to load
        (last block of data to load)
    3. already_loaded: ``bool``
        used to know wheather data has been loaded
        when the whole time values fit in memory, not
        used otherwise
    """
    
    # Test if dataset is not too big
    if calculate_memory_requirement(x_size, y_size, 1, nb_inputs, nb_outputs, nb_variables, nb_params, nb_bytes) > available_ram:
        import sys
        print('Calculation area is too large for available memory.')
        sys.exit()
    
    # Boolean to get out of while loop
    not_finished = True
    
    # Divisor starts at one
    divisor = 1
    
    # Increase divisor by one at every loop
    while not_finished and divisor < time_size / 2:
        
        memory_requirement = calculate_memory_requirement(x_size, y_size, time_size // divisor, nb_inputs, nb_outputs, nb_variables, nb_params, nb_bytes)
        
        # Test if memory requirement divided by divisor fits in the available ram
        if memory_requirement < available_ram:
            
            # If all the inputs and outputs can be loaded
            if divisor == 1:
                
                time_slice = time_size  # whole dataset is loaded
                remainder_to_load = None  # no remainder to load
                already_loaded = False  # this boolean is set to true once the whole inputs have been loaded
                not_finished = False  # boolean set to false to get out of while loop
                
                return time_slice, remainder_to_load, already_loaded
            
            # Otherwise return the number of time  bands that can be loaded
            else:
                
                time_slice = time_size // divisor
                remainder_to_load = time_size % time_slice  # remainder, used to load the last block of inputs
                already_loaded = None  # this boolean is set to None if the whole dataset can't be loaded
                not_finished = False  # boolean set to false to get out of while loop
                
                return time_slice, remainder_to_load, already_loaded

        divisor += 1  # increment divisor
    
    # If dataset is to big, load only one time slice per loop
    time_slice = 1
    remainder_to_load = 1  # in order to correctly save last date
    already_loaded = None  # this boolean is set to None if the whole dataset can't be loaded
        
    return time_slice, remainder_to_load, already_loaded


def get_empty_arrays(x_size: int, y_size: int, time_size: int, nb_arrays: int, array: bool = False) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """
    Short function to make the `run_samir()` function easier to read.
    Generates a varying number (`nb_arrays`) of empty numpy arrays of shape 
    `(x_size, y_size, time_size)` in list or tuple mode. Used to
    store variable values before writing in in the output file.

    Arguments
    =========

    1. x_size: ``int``
        x size of dataset
    2. y_size: ``int``
        y size of dataset
    3. time_size: ``int``
        number of time bands
    4. nb_arrays: ``int``
        number of arrays to generate
    5. array: ``bool`` ``default = False``
        weather to return a tuple or a list

    Returns
    =======

    output: ``Tuple[np.ndarray * nb_arrays]`` or ``List[np.ndarray * nb_arrays]``
        output empty arrays
    """
    
    # Return empty arrays into a list
    if array:
        # return [np.empty((time_size, y_size, x_size), dtype = np.int16) for k in range(nb_arrays)]
        return np.empty((nb_arrays, time_size, y_size, x_size), dtype = np.int16, order = 'C')
    
    # Return empty arrays into a tuple
    return tuple([np.empty((time_size, y_size, x_size), dtype = np.uint16) for k in range(nb_arrays)])


# @profile  # type: ignore
def read_inputs(ndvi_cube_path: str, weather_path: str, i: int, time_slice: int, load_all: bool = False) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Read input data into numpy arrays based on loop conditions.

    Arguments
    =========

    1. ndvi_cube_path: ``str``
        path of input NDVI file
    2. weather_path: ``str``
        path of input weather file
    3. i: ``int``
        current loop counter
    4. time_slice: ``int``
        number of time slices to load
    5. load_all: ``bool`` ``default = False``
        boolean to load the whole datasets

    Returns
    =======

    1. NDVI: ``np.ndarray``
        integer NDVI values into numpy array
    2. Rain: ``np.ndarray``
        integer Rain values into numpy array
    3. ET0: ``np.ndarray``
        integer ET0 values into numpy array
        
    """
    
    # Load whole dataset
    if load_all:
        
        with nc.Dataset(ndvi_cube_path, mode='r') as ds:
            # Dimensions of ndvi dataset : (time, y, x)
            ds.variables['NDVI'].set_auto_mask(False)
            NDVI = ds.variables['NDVI'][:, :, :]
        with nc.Dataset(weather_path, mode='r') as ds:
            # Dimensions of ndvi dataset : (time, y, x)
            ds.variables['Rain'].set_auto_mask(False)
            Rain = ds.variables['Rain'][:, :, :]
            ds.variables['ET0'].set_auto_mask(False)
            ET0 = ds.variables['ET0'][:, :, :]
    
    # Load given number of time slices
    else:
        
        with nc.Dataset(ndvi_cube_path, mode='r') as ds: