Skip to content
Snippets Groups Projects
modspa_samir.py 67.7 KiB
Newer Older
Jeremy Auclair's avatar
Jeremy Auclair committed
Usage of the SAMIR model in the Modspa framework. 

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
Jeremy Auclair's avatar
Jeremy Auclair committed
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.
Jeremy Auclair's avatar
Jeremy Auclair committed
    1. csv_param_file: ``str``
Jeremy Auclair's avatar
Jeremy Auclair committed
    2. land_cover_raster: ``str``
    4. init_RU_raster: ``str``
        path to initial soil water content. None if no input.
    1. parameter_dict: ``Dict[str, np.ndarray]``
        the dictionnary containing all the rasterized Parameters
        and the scale factors

    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)
    class_count = table_param.table.shape[1] - 2  # remove dtype and Default columns
    land_cover = xr.open_dataarray(land_cover_raster).to_numpy()
    # 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
    for parameter in table_param.table.index[1:]:
        if parameter == 'Irrig_auto' and table_param.table.at[parameter, 'load_raster'] == 1:
            Irrig_auto = np.ascontiguousarray(xr.open_dataarray(irrigation_raster, decode_coords = 'all').astype(np.bool_).values, dtype = np.bool_)
            
        elif parameter == 'Init_RU' and table_param.table.at[parameter, 'load_raster'] == 1:
            Init_RU = (xr.open_dataarray(init_RU_raster) / 1000).astype(np.float32).values
            i+=1
            
        # If scale_factor == 0, then the parameter does not have to be spatialized, it can stay scalar
            parameter_array[i] = np.ones(shape = shape, dtype = np.float32) * np.float32(table_param.table.at[parameter, 'Default'])
            i+=1         
        if (parameter != 'Irrig_auto') and (parameter != 'Init_RU'):
            dtype = np.float32
            param_list.append(parameter)
        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)
        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.

        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
        list of additional variable names to be saved

Jeremy Auclair's avatar
Jeremy Auclair committed
    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)
    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['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['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['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['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['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['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.
Jeremy Auclair's avatar
Jeremy Auclair committed
    1. TAW: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    2. Dr: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    3. Zr: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    4. RUE: ``np.ndarray``
        total available surface water = (Wfc-Wwp)*Ze
Jeremy Auclair's avatar
Jeremy Auclair committed
    5. De: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    6. Wfc: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    7. Ze: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    8. DiffE: ``np.ndarray``
        diffusion coefficient between evaporative
        and root layers (unitless, parameter)
Jeremy Auclair's avatar
Jeremy Auclair committed
    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.

Jeremy Auclair's avatar
Jeremy Auclair committed
    1. TAW: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    2. TDW: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    3. Dr: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    4. Zr: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    5. Dd: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    6. Wfc: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    7. Zsoil: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    8. DiffR: ``np.ndarray``
        Diffusion coefficient between root
        and deep layers (unitless, parameter)

Jeremy Auclair's avatar
Jeremy Auclair committed
    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.
Jeremy Auclair's avatar
Jeremy Auclair committed
    1. TEW: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    2. Dei: ``np.ndarray``
        depletion of the evaporative layer
        (irrigation part)
Jeremy Auclair's avatar
Jeremy Auclair committed
    3. Dep: ``np.ndarray``
        depletion of the evaporative layer
        (precipitation part)
Jeremy Auclair's avatar
Jeremy Auclair committed
    4. fewi: ``np.ndarray``
        soil fraction which is wetted by irrigation
        and exposed to evaporation
Jeremy Auclair's avatar
Jeremy Auclair committed
    5. fewp: ``np.ndarray``
        soil fraction which is wetted by precipitation
        and exposed to evaporation
Jeremy Auclair's avatar
Jeremy Auclair committed
    1. W: ``np.ndarray``

    # Calculate the weighting factor to split the energy available for evaporation
    # Equation: W = 1 / (1 + (fewp * (TEW - Dep) / fewi * (TEW - Dei)))
    # 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.

Jeremy Auclair's avatar
Jeremy Auclair committed
    1. TEW: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    2. De: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    3. REW: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    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.
Jeremy Auclair's avatar
Jeremy Auclair committed
    Arguments
    =========

    1. Dr: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    2. TAW: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    3. p: ``np.ndarray``
        fraction of TAW available for plant without inducing stress
Jeremy Auclair's avatar
Jeremy Auclair committed
    4. E0: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    5. Tr0: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    Returns
    =======

    1. Ks: ``np.ndarray``
    # 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 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.

Jeremy Auclair's avatar
Jeremy Auclair committed
    Arguments
    =========

    1. Dr: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    2. TAW: ``np.ndarray``
    5. Irrig_auto: ``np.ndarray`` ``boolean``
        maximum amount of possible irrigation in mm
        minimum value of Kcb under which no irrigation is allowed
        Kcb threshold to stop irrigation after reaching maximum
        boolean array that is true after the Kcb has reached
        frac_Kcb_stop_irrig of its maximum
        fraction of depleted TAW under which irrigation is triggered
Jeremy Auclair's avatar
Jeremy Auclair committed
    Returns
    =======

    1. Irrig: ``np.ndarray``
    # 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.

Jeremy Auclair's avatar
Jeremy Auclair committed
    Arguments
    =========

    1. De: ``np.ndarray``
        Dei or Dep, depletion of the evaporative layer
Jeremy Auclair's avatar
Jeremy Auclair committed
    2. Dr: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    3. Ks: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    4. Kcb: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    5. Ze: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    6. Zr: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    7. TEW: ``np.ndarray``
        water capacity of the evaporative layer
Jeremy Auclair's avatar
Jeremy Auclair committed
    8. TAW: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    9. ET0: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    Returns
    =======

    1. Te: ``np.ndarray``
    # 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.

Jeremy Auclair's avatar
Jeremy Auclair committed
    Arguments
    =========

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

Jeremy Auclair's avatar
Jeremy Auclair committed
    Returns
    =======

    De: ``np.ndarray``
    # 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.

Jeremy Auclair's avatar
Jeremy Auclair committed
    1. Wfc: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    2. Wwp: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    3. Zr: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    4. Zsoil: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    5. Dr0: ``np.ndarray``
        depletion of the root layer for previous day
Jeremy Auclair's avatar
Jeremy Auclair committed
    6. Dd0: ``np.ndarray``
        depletion of the deep laye for previous day
Jeremy Auclair's avatar
Jeremy Auclair committed
    7. Zr0: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    1. output: ``np.ndarray``
    # 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.

Jeremy Auclair's avatar
Jeremy Auclair committed
    1. Wfc: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    2. Wwp: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    3. Zr: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    4. Zsoil: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    5. Dr0: ``np.ndarray``
        depletion of the root layer for previous day
Jeremy Auclair's avatar
Jeremy Auclair committed
    6. Dd0: ``np.ndarray``
        depletion of the deep laye for previous day
Jeremy Auclair's avatar
Jeremy Auclair committed
    7. Zr0: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    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.

Jeremy Auclair's avatar
Jeremy Auclair committed
    Arguments
    =========

    1. x_size: ``int``
        x size of dataset (pixels)
Jeremy Auclair's avatar
Jeremy Auclair committed
    2. y_size: ``int``
        y size of dataset (pixels)
Jeremy Auclair's avatar
Jeremy Auclair committed
    3. time_size: ``int``
        number of time bands (dates)
Jeremy Auclair's avatar
Jeremy Auclair committed
    4. nb_inputs: ``int``
Jeremy Auclair's avatar
Jeremy Auclair committed
    5. nb_outputs: ``int``
Jeremy Auclair's avatar
Jeremy Auclair committed
    6. nb_variables: ``int``
Jeremy Auclair's avatar
Jeremy Auclair committed
    7. nb_params: ``int``
        number of bytes of datatype for inputs and outputs
Jeremy Auclair's avatar
Jeremy Auclair committed
    Returns
    =======

    1. total_memory_requirement: ``float``
    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
    output_memory_requirement = (x_size * y_size * time_size * nb_outputs * nb_bytes) / (1024**3)
    total_memory_requirement = (input_memory_requirement + calculation_memory_requirement + output_memory_requirement) * 1.05  # 5% adjustment factor
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.

Jeremy Auclair's avatar
Jeremy Auclair committed
    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
Jeremy Auclair's avatar
Jeremy Auclair committed
    Returns
    =======

    1. time_slice: ``int``
        number of times slices to load for
        input and output data
Jeremy Auclair's avatar
Jeremy Auclair committed
    2. remainder_to_load: ``int``
        remainder of euclidian division for
        the number of time slices to load
        (last block of data to load)
Jeremy Auclair's avatar
Jeremy Auclair committed
    3. already_loaded: ``bool``
        used to know wheather data has been loaded
        when the whole time values fit in memory, not
        used otherwise
    """
    
    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
    # 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 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
                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
    
    # 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.

Jeremy Auclair's avatar
Jeremy Auclair committed
    Arguments
    =========

    1. x_size: ``int``
Jeremy Auclair's avatar
Jeremy Auclair committed
    2. y_size: ``int``
Jeremy Auclair's avatar
Jeremy Auclair committed
    3. time_size: ``int``
Jeremy Auclair's avatar
Jeremy Auclair committed
    4. nb_arrays: ``int``
Jeremy Auclair's avatar
Jeremy Auclair committed
    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 tuple([np.empty((time_size, y_size, x_size), dtype = np.uint16) for k in range(nb_arrays)])
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.
Jeremy Auclair's avatar
Jeremy Auclair committed
    Arguments
    =========

    1. ndvi_cube_path: ``str``
Jeremy Auclair's avatar
Jeremy Auclair committed
    2. weather_path: ``str``
Jeremy Auclair's avatar
Jeremy Auclair committed
    3. i: ``int``
Jeremy Auclair's avatar
Jeremy Auclair committed
    4. time_slice: ``int``
Jeremy Auclair's avatar
Jeremy Auclair committed
    5. load_all: ``bool`` ``default = False``
Jeremy Auclair's avatar
Jeremy Auclair committed
    Returns
    =======

    1. NDVI: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    2. Rain: ``np.ndarray``
Jeremy Auclair's avatar
Jeremy Auclair committed
    3. ET0: ``np.ndarray``
        
    """
    
    # 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: