modspa_samir.py

            # Dimensions of ndvi dataset : (time, y, x)
            ds.variables['NDVI'].set_auto_mask(False)
            NDVI = ds.variables['NDVI'][i: i + time_slice, :, :]
        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'][i: i + time_slice, :, :]
            ds.variables['ET0'].set_auto_mask(False)
            ET0 = ds.variables['ET0'][i: i + time_slice, :, :]
    
    return NDVI, Rain, ET0


def write_outputs(save_path: str, DP: np.ndarray, SWCe: np.ndarray, SWCr: np.ndarray, E: np.ndarray, Tr: np.ndarray, Irrig: np.ndarray, additional_outputs: dict[str, int], additional_outputs_data: np.ndarray, i: int, time_slice: int, write_all = False) -> None:
    """
    Write outputs to netcdf file based on conditions of current loop.

    Arguments
    =========

    1. save_path: ``str``
        output netcdf save path
    2. DP: ``np.ndarray``
        deep percolaton ``np.ndarray``
    3. SWCe: ``np.ndarray``
        soil water content of evaporative layer ``np.ndarray``
    4. SWCr: ``np.ndarray``
        soil water content of root layer ``np.ndarray``
    5. E: ``np.ndarray``
        surface evaporation ``np.ndarray``
    6. Tr: ``np.ndarray``
        plant transpiration ``np.ndarray``
    7. Irrig: ``np.ndarray``
        simulated irrigation ``np.ndarray``
    8. additional_outputs: ``Dict[str, int]``
        dictionnary containing additional outputs and their scale factors
    9. additional_outputs_data: ``List[np.ndarray]``
        list of additional output ``np.ndarray``. Is ``None`` if no additional ouputs
    10. i: ``int``
        current loop counter
    11. time_slice: ``int``
        number of loaded time slices
    12. write_all: ``bool`` ``default = False``
        weather to write the whole output dataset

    Returns
    =======

    ``None``
    """
    
    
    # Write whole output dataset
    if write_all:
        
        with nc.Dataset(save_path, mode = 'a') as outputs:
            # Dimensions of output dataset : (time, y, x)
            # Deep percolation
            outputs.variables['DP'][:, :, :] = DP
            # Soil water content of the evaporative layer
            outputs.variables['SWCe'][:, :, :] = SWCe
            # Soil water content of the root layer
            outputs.variables['SWCr'][:, :, :] = SWCr
            # Evaporation
            outputs.variables['E'][:, :, :] = E
            # Transpiration
            outputs.variables['Tr'][:, :, :] = Tr
            # Irrigation
            outputs.variables['Irr'][:, :, :] = Irrig
            # Additionnal outputs
            if additional_outputs:
                k = 0
                for var in additional_outputs.keys():
                    outputs.variables[var][:, :, :] = additional_outputs_data[k,:,:,:]
                    k+=1
    
    else:
        
        # Write given number of time slices
        with nc.Dataset(save_path, mode = 'a') as outputs:
            # Dimensions of output dataset : (time, y, x)
            # Deep percolation
            outputs.variables['DP'][i - time_slice + 1: i + 1, :, :] = DP
            # Soil water content of the evaporative layer
            outputs.variables['SWCe'][i - time_slice + 1: i + 1, :, :] = SWCe
            # Soil water content of the root layer
            outputs.variables['SWCr'][i - time_slice + 1: i + 1, :, :] = SWCr
            # Evaporation
            outputs.variables['E'][i - time_slice + 1: i + 1, :, :] = E
            # Transpiration
            outputs.variables['Tr'][i - time_slice + 1: i + 1, :, :] = Tr
            # Irrigation
            outputs.variables['Irr'][i - time_slice + 1: i + 1, :, :] = Irrig
            # Additionnal outputs
            if additional_outputs:
                k=0
                for var in additional_outputs.keys():
                    outputs.variables[var][i - time_slice + 1: i + 1, :, :] = additional_outputs_data[k,:,:,:]
                    k+=1
        
    return None


@njit((uint8[:, :], uint16[:, :], uint16[:, :], float32[:, :], float32[:, :], float32[:, :], boolean[:, :], boolean[:, :], float32[:, :], float32[:, :], float32[:, :], uint16[:, :], uint16[:, :], float32[:, :], float32[:, :], int64, float32[:, :, :], types.ListType(types.unicode_type), types.ListType(types.unicode_type), int16[:], boolean, int16[:], types.ListType(types.unicode_type), int16[:], int16[:, :, :, :], int64), nogil = True, parallel = True, fastmath = True, cache = True)
def samir_daily(NDVI: np.ndarray, ET0: np.ndarray, Rain: np.ndarray, Wfc: np.ndarray, Wwp: np.ndarray, Kcb_max_obs: np.ndarray, Irrig_auto: np.ndarray, Irrig_test: np.ndarray, Dr0: np.ndarray, Dd0: np.ndarray, Zr0: np.ndarray, E0: np.ndarray, Tr0: np.ndarray, Dei0: np.ndarray, Dep0: np.ndarray, iday: int, params: np.ndarray, param_list: list[str], scaling_names: list[str], scaling_array: np.ndarray, write_additional_data: bool, additional_outputs: np.ndarray, additional_names: list[str], additional_idx: list[int], additional_outputs_data: np.ndarray, time_slice: int) -> tuple[np.ndarray]:
    """
    Run the SAMIR model on a single day. Inputs and outputs are `numpy.ndarray`.
    Calls functions compiled with numba for intermediary calculations.

    Arguments
    =========

    1. NDVI: ``np.ndarray``
        input NDVI
    2. ET0: ``np.ndarray``
        input ET0
    3. Rain: ``np.ndarray``
        input Rain
    4. Wfc: ``np.ndarray``
        field capacity
    5. Wwp: ``np.ndarray``
        field wilting point
    6. Kcb_max_obs: ``np.ndarray``
        observed maximum value of Kcb 
    7. Irrig_test: ``np.ndarray``
        boolean array that is true after the Kcb has reached 80% 
        of its maximum
    8. params: ``np.ndarray``
        dictionnary containing the rasterized
        samir parameters and their scale factors
    9. Dr0: ``np.ndarray``
        previous day root layer depletion
    10. Dd0: ``np.ndarray``
        previous day deep layer depletion
    11. Zr0: ``np.ndarray``
        previous day root depth
    12. E0: ``np.ndarray``
        previous day surface evaporation
    13. Tr0: ``np.ndarray``
        previous day plant transpiration
    14. Dei0: ``np.ndarray``
        previous day surface layer depletion
        for irrigation part
    15. Dep0: ``np.ndarray``
        previous day surface layer depletion
        for precipitation part
    16. iday: ``int``
        current loop counter
    17. scaling_dict: ``np.ndarray``
        scaling dictionnary for the nominal outputs
    18. additional_names: ``list[str]``
        list of names of additional outputs
    19. field_indices: ``list[int]``
        list of indices corresponding to additional names
    20. additional_outputs: ``np.ndarray``
        dictionary containing the names of additional
        output data and their integer scale
    21. additional_outputs_data: ``np.ndarray``
        list of numpy arrays containing the scaled values of
        the additional outputs
    22. time_slice: ``int``
        value of the current time slice, used to
        correctly write additional output data

    Returns
    =======

    1. current_day_outputs: `Tuple[np.ndarray]`
        multiple `numpy.ndarray` arrays are returned as a tuple for current day
    """
    
    # Create aliases
    DiffE = params[find_index(param_list, 'DiffE')]
    DiffR = params[find_index(param_list, 'DiffR')]
    FW = params[find_index(param_list, 'FW')]
    FCmax = params[find_index(param_list, 'FCmax')]
    Foffset = params[find_index(param_list, 'Foffset')]
    Fslope = params[find_index(param_list, 'Fslope')]
    frac_TAW = params[find_index(param_list, 'frac_TAW')]
    Init_RU = params[find_index(param_list, 'Init_RU')]
    Kcbmax = params[find_index(param_list, 'Kcbmax')]
    Kcmax = params[find_index(param_list, 'Kcmax')]
    Kslope = params[find_index(param_list, 'Kslope')]
    Koffset= params[find_index(param_list, 'Koffset')]
    Kcb_min_start_irrig = params[find_index(param_list, 'Kcb_min_start_irrig')]
    frac_Kcb_stop_irrig = params[find_index(param_list, 'frac_Kcb_stop_irrig')]
    Lame_max = params[find_index(param_list, 'Lame_max')]
    Lame_min = params[find_index(param_list, 'Lame_min')]
    REW = params[find_index(param_list, 'REW')]
    Ze = params[find_index(param_list, 'Ze')]
    Zsoil = params[find_index(param_list, 'Zsoil')]
    minZr = params[find_index(param_list, 'minZr')]
    maxZr = params[find_index(param_list, 'maxZr')]
    p = params[find_index(param_list, 'p')]
    
    # Scale input parameters
    NDVI = NDVI.astype(np.float32) / np.float32(255)
    Rain = Rain.astype(np.float32) / np.float32(100)
    ET0 = ET0.astype(np.float32) / np.float32(1000)
    E0 = E0.astype(np.float32) / scaling_array[find_index(scaling_names, 'E')]
    Tr0 = Tr0.astype(np.float32) / scaling_array[find_index(scaling_names, 'Tr')]

    # Update variables
    # Fraction cover
    # Equation: Fslope * NDVI + Foffset
    FCov = Fslope * NDVI + Foffset
    # Equation: min(max(FCov, 0), FCmax)
    FCov = np.minimum(np.maximum(FCov, np.float32(0)), FCmax)

    # Root depth upate
    # Equation: Zr = max(minZr + (FCov / FCmax) * (maxZr - minZr), Ze + 0.001)
    Zr = np.maximum(minZr + (FCov / FCmax) * (maxZr - minZr), Ze + np.float32(0.001))

    # Water capacities
    TAW = (Wfc - Wwp) * Zr
    TDW = (Wfc - Wwp) * (Zsoil - Zr)
    TEW = (Wfc - (Wwp / np.float32(2))) * Ze
    RUE = (Wfc - Wwp) * Ze

    # Update depletions from root increase
    Dei = Dei0
    Dep = Dep0
    Dr = update_Dr_from_root(Wfc, Wwp, Zr, Zsoil, Dr0, Dd0, Zr0)
    Dd = update_Dd_from_root(Wfc, Wwp, Zr, Zsoil, Dr0, Dd0, Zr0)
    
    # Kcb
    # Equation: Kslope * NDVI + Koffset
    Kcb = np.minimum(np.maximum(Kslope * NDVI + Koffset, np.float32(0)), Kcbmax)
    Irrig_test = np.where(np.invert(Irrig_test) & (Kcb > frac_Kcb_stop_irrig * Kcb_max_obs), True, Irrig_test)  # set Irrig_test to true when Kcb goes over Kcb_stop_irrig fraction of maximum and stay true

    # Irrigation 
    Irrig = calculate_irrig(Dr, TAW, Rain, Kcb, Irrig_auto, Lame_max, Lame_min, Kcb_min_start_irrig, frac_Kcb_stop_irrig, Irrig_test, frac_TAW, Kcb_max_obs)

    # Create temporary variable used multiple times
    temp = np.subtract(Dr, Rain + Irrig)
    
    # DP (Deep percolation)
    DP = - np.minimum(Dd + np.minimum(temp, np.float32(0)), np.float32(0))

    # Update depletions with Rainfall and/or irrigation
    # Dei and Dep
    # Equation: Dei = min(max(Dei - Rain - Irrig / FW, 0), TEW)
    Dei = np.minimum(np.maximum(Dei - Rain - (Irrig / FW), np.float32(0)), TEW)
    # Equation: Dep = min(max(Dep - Rain, 0), TEW)
    Dep = np.minimum(np.maximum(Dep - Rain, np.float32(0)), TEW)

    fewi = np.minimum(np.float32(1) - FCov, FW)
    fewp = np.float32(1) - FCov - fewi

    # De
    # Equation: De = (Dei * fewi + Dep * fewp) / (fewi + fewp)
    De = ((Dei * fewi) + (Dep * fewp)) / (fewi + fewp)
    # Equation: De = where(De.isfinite, De, Dei * FW + Dep * (1 - FW))
    De = np.where(np.isfinite(De), De, Dei * FW + Dep * (np.float32(1) - FW))

    # Update depletions from rain and irrigation
    Dr = np.minimum(np.maximum(temp, np.float32(0)), TAW)
    Dd = np.minimum(np.maximum(Dd + np.minimum(temp, np.float32(0)), np.float32(0)), TDW)
    temp = False # remove temp variable

    # Diffusion coefficients
    # Equation: check calculate_diff_re() and calculate_diff_dr functions
    Diff_rei = calculate_diff_re(TAW, Dr, Zr, RUE, Dei, Wfc, Ze, DiffE)
    Diff_rep = calculate_diff_re(TAW, Dr, Zr, RUE, Dep, Wfc, Ze, DiffE)
    Diff_dr = calculate_diff_dr(TAW, TDW, Dr, Zr, Dd, Wfc, Zsoil, DiffR)
    
    # Water Stress coefficient
    if iday == 0:
        Ks = np.minimum((TAW - Dr) / (TAW * (np.float32(1) - p)), np.float32(1))
    else:
        # When not first day
        Ks = calculate_Ks(Dr, TAW, p, E0, Tr0)

    # Reduction coefficient for evaporation
    W = calculate_W(TEW, Dei, Dep, fewi, fewp)

    # Equation: Kei = np.minimum(W * Kri * (Kc_max - Kcb), fewi * Kc_max)
    Kei = calculate_Ke(W, Dei, TEW, REW, Kcmax, Kcb, fewi)

    # Equation: Kep = np.minimum((1 - W) * Krp * (Kc_max - Kcb), fewp * Kc_max)
    Kep = calculate_Ke((np.float32(1)-W), Dep, TEW, REW, Kcmax, Kcb, fewp)

    # Prepare coefficients for evapotranspiration
    Tei = calculate_Te(Dei, Dr, Ks, Kcb, Ze, Zr, TEW, TAW, ET0)
    Tep = calculate_Te(Dep, Dr, Ks, Kcb, Ze, Zr, TEW, TAW, ET0)

    # Update depletions
    Dei = update_De_from_Diff(Dei, fewi, Kei, Tei, Diff_rei, TEW, ET0)
    Dep = update_De_from_Diff(Dep, fewp, Kep, Tep, Diff_rep, TEW, ET0)

    # Equation: De = (Dei * fewi + Dep * fewp) / (fewi + fewp)
    De = ((Dei * fewi) + (Dep * fewp)) / (fewi + fewp)
    # Equation: De = where(De.isfinite, De, Dei * FW + Dep * (1 - FW))
    De = np.where(np.isfinite(De), De, Dei * FW + Dep * (np.float32(1) - FW))

    # Evaporation
    E = np.maximum((Kei + Kep) * ET0, np.float32(0))

    # Transpiration
    Tr = Kcb * Ks * ET0

    # Update depletions (root and deep zones) at the end of the day
    Dr = np.minimum(np.maximum(Dr + E + Tr - Diff_dr, np.float32(0)), TAW)
    Dd = np.minimum(np.maximum(Dd + Diff_dr, np.float32(0)), TDW)

    # Soil water content of evaporative laye
    SWCe = calculate_SWCe(Dei, Dep, fewi, fewp, TEW)

    # Soil water content of root layer
    SWCr = np.float32(1) - Dr/TAW
    
    # Scale outputs before return
    DP = (DP * scaling_array[find_index(scaling_names, 'DP')]).astype(np.int16)
    Irrig = (Irrig * scaling_array[find_index(scaling_names, 'Irr')]).astype(np.int16)
    SWCe = (SWCe * scaling_array[find_index(scaling_names, 'SWCe')]).astype(np.int16)
    SWCr = (SWCr * scaling_array[find_index(scaling_names, 'SWCr')]).astype(np.int16)
    E = (E * scaling_array[find_index(scaling_names, 'E')]).astype(np.int16)
    Tr = (Tr * scaling_array[find_index(scaling_names, 'Tr')]).astype(np.uint16)
    
    # Collect additionnal outputs
    if write_additional_data:
        variable_mapping = {
            'FCov': FCov, 'Zr': Zr, 'TAW': TAW, 'TDW': TDW, 'TEW': TEW, 'RUE': RUE, 'Dei': Dei, 'Dep': Dep,
            'De': De, 'Dr': Dr, 'Dd': Dd, 'Kcb': Kcb, 'fewi': fewi, 'fewp': fewp, 'Diff_rei': Diff_rei,
            'Diff_rep': Diff_rep, 'Diff_dr': Diff_dr, 'Ks': Ks, 'W': W, 'Kei': Kei, 'Kep': Kep, 'Tei': Tei,
            'Tep': Tep}
        
        for var, idx in zip(additional_names, additional_idx):
            additional_outputs_data[idx, iday % time_slice, :, :] = np.round(variable_mapping[var] * additional_outputs[idx]).astype(np.int16)
    
    return DP, Dd, Dei, Dep, Dr, E, Irrig, Irrig_test, SWCe, SWCr, Tr, Zr


def run_samir(csv_param_file: str, ndvi_cube_path: str, weather_path: str, soil_params_path: str, land_cover_path: str, irrigation_raster: str, init_RU_path: str, Kcb_max_obs_path: str, save_path: str, scaling_dict: dict[str, int] = {'E': 1000, 'Tr': 1000, 'SWCe': 1000, 'SWCr': 1000, 'DP': 100, 'Irr': 100}, additional_outputs: dict[str, int] = None, available_ram: int = 8, available_cpu: int = 4, compress_outputs: bool = False) -> None:
    """
    Run the *SAMIR* model on the prepared inputs. Calls the ``samir_daily()`` in a time loop.
    Maximizes memory usage with given limits to run faster.

    Arguments
    =========

    1. csv_param_file: ``str``
        SAMIR csv parameter file
    2. ndvi_cube_path: ``str``
        path to ndvi cube
    3. weather_path: ``str``
        path to weather cube
    4. soil_params_path: ``str``
        path to soil dataset
    5. land_cover_path: ``str``
        path to land cover raster
    6. irrigation_raster: ``str``
        path to netCDF file containing an irrigation mask
        for input parameters
    7. init_RU_path: ``str``
        path to netCDF file containing initial soil water content raster
    8. Kcb_max_obs_path: ``str``
        path to netCDF containing a single
        array of maximum observed Kcb values
    8. save_path: ``str``
        path to save the model outputs
    9. scaling_dict: ``Dict[str, int]`` ``default = {'E': 1000, 'Tr': 1000, 'SWCe': 1000, 'SWCr': 1000, 'DP': 100, 'Irr': 100}``
        scaling dictionnary for the nominal outputs
    10. additional_outputs: ``Dict[str, int]`` ``default = None``
        dictionnary containing the names and scale
        factors of potential additional outouts
    11. available_ram: ``int`` ``default = 8``
        available RAM in **GiB** for the model
    12. available_cpu: ``int`` ``default = 4``
        number of available **physical** CPU cores
    12. compress_outputs: ``bool`` ``default = False``
        choice to compress output file (takes longer)

    Returns
    =======

    ``None``
    """

    # Turn off numpy warings
    np.seterr(divide='ignore', invalid='ignore')
    
    # Test if memory requirement is not loo large
    if np.ceil(virtual_memory().available / (1024**3)) < available_ram:
        print('\nRequested', available_ram, 'GiB of memory when available memory is approximately', round(virtual_memory().available / (1024**3), 1), 'GiB.\n\nExiting script.\n')
        return None

    # Set maximum number of usable CPU cores
    # Get number of CPU cores and limit max value (working on a cluster requires os.sched_getaffinity to get true number of available CPUs, 
    # this is not true on a "personnal" computer, hence the use of the min function)
    try:
        nb_threads = min([available_cpu * 2, cpu_count(logical = True), len(os.sched_getaffinity(0))])
    except:
        nb_threads = min([available_cpu * 2, cpu_count(logical = True)])  # os.sched_getaffinity won't work on windows
    set_num_threads(nb_threads)

    # ============ Manage inputs ============ #
    # NDVI (to have a correct empty dataset structure)
    ndvi_cube = xr.open_dataset(ndvi_cube_path)
    ndvi_cube_empty = ndvi_cube.drop_sel(time = ndvi_cube.time)  # create dataset with a time dimension of length 0

    # SAMIR Parameters
    parameter_array, Irrig_auto, Init_RU, param_list = rasterize_samir_parameters(csv_param_file, land_cover_path, irrigation_raster, init_RU_path)
    
    # ============ Get size of dataset ============ #
    x_size = ndvi_cube.sizes['x']
    y_size = ndvi_cube.sizes['y']
    time_size = ndvi_cube.sizes['time']
    dimensions = ndvi_cube_empty.sizes  # to create empty output dataset
    dates = ndvi_cube.time
    ndvi_cube.close()
    
    # ============ Scaling factors ============ #
    # Standard scaling dict
    scaling_array = []
    for var in scaling_dict.keys():
        scaling_array.append(scaling_dict[var])
    scaling_array = np.ascontiguousarray(np.array(scaling_array, dtype = np.int16))
    scaling_names = List(list(scaling_dict.keys()))
    
    # Manage additional output parameters
    if additional_outputs is not None:
        # Additional output scaling dict
        additional_scaling_array = []
        for var in additional_outputs.keys():
            additional_scaling_array.append(additional_outputs[var])
        additional_scaling_array = np.array(additional_scaling_array, dtype = np.int16)
        additional_names = list(additional_outputs.keys())
        field_name_to_index = {name: idx for idx, name in enumerate(additional_outputs)}
        additional_idx = np.ascontiguousarray(np.array([field_name_to_index[var] for var in additional_names], dtype = np.int16))
        additional_names = List(additional_names)
        write_additional_data = True
    else:
        # Create empty variables for correct numba parsing
        additional_scaling_array = np.array([0], dtype = np.int16)
        additional_idx = np.array([0], dtype = np.int16)
        additional_names = List(['0'])
        additional_outputs_data = np.empty(shape = (1,1,1,1), dtype = np.int16)
        write_additional_data = False
    
    # ============ Memory handling ============ #
    # Check how much memory the calculation would take if all the inputs would be loaded in memory
    # Unit is GiB
    # Datatype of variables is float32 for calculation
    nb_bytes = 2  # uint16 or int16
    nb_inputs = 3  # NDVI, Rain, ET0
    if additional_outputs:
        nb_outputs = 6 + len(additional_outputs)  # DP, E, Irr, SWCe, SWCr, Tr
    else:
        nb_outputs = 6  # DP, E, Irr, SWCe, SWCr, Tr
    nb_variables = 38  # calculation variables (e.g:  Dd, Diff_rei, FCov, etc.)
    nb_params = parameter_array.shape[0]
    
    # Get total memory requirement
    total_memory_requirement = calculate_memory_requirement(x_size, y_size, time_size, nb_inputs, nb_outputs, nb_variables, nb_params, nb_bytes)
    
    # Determine how many time slices can be loaded in memory at once
    # This will allow faster I/O operations and a faster runtime
    time_slice, remainder_to_load, already_loaded = calculate_time_slices_to_load(x_size, y_size, time_size, nb_inputs, nb_outputs, nb_variables, nb_params, nb_bytes, available_ram)
    
    # Get actual memory requirement
    actual_memory_requirement = calculate_memory_requirement(x_size, y_size, time_slice, nb_inputs, nb_outputs, nb_variables, nb_params, nb_bytes)
    
    print_size1, print_unit1 = format_Byte_size(total_memory_requirement, decimals = 1)
    print_size2, print_unit2, = format_Byte_size(actual_memory_requirement, decimals = 1)
    print('\nApproximate memory requirement of calculation:', print_size1, print_unit1 + ', available memory:', available_ram, 'GiB\n\nLoading blocks of', time_slice, 'time bands, actual memory requirement:', print_size2, print_unit2, '\n')
    
    # ============ Prepare outputs ============ #
    model_outputs = prepare_output_dataset(ndvi_cube_path, dimensions, scaling_dict, additional_outputs = additional_outputs)

    # Create encoding dictionnary
    encoding_dict = {}
    for variable in list(model_outputs.keys()):
        # Write encoding dict
        encod = {}
        if variable not in scaling_dict.keys():  # additional outputs are not in the scaling dict
            encod['dtype'] = 'i2'
        else:
            encod['dtype'] = 'u2'
        # TODO: figure out optimal file chunk size
        encod['chunksizes'] = (1, y_size, x_size)
        if compress_outputs:
            encod['zlib'] = True
            encod['complevel'] = 1
        encoding_dict[variable] = encod

    # Save empty output
    print('Writing empty output dataset\n')
    model_outputs.to_netcdf(save_path, encoding = encoding_dict, unlimited_dims = 'time')  # add time as an unlimited dimension, allows to append data along the time dimension
    model_outputs.close()

    # ============ Prepare time iterations ============#

    # input soil data
    with nc.Dataset(soil_params_path, mode = 'r') as ds:
        ds.variables['Wfc'].set_auto_mask(False)
        Wfc = ds.variables['Wfc'][:, :]
        ds.variables['Wwp'].set_auto_mask(False)
        Wwp = ds.variables['Wwp'][:, :]
    
    # Observed Kcb max data
    with nc.Dataset(Kcb_max_obs_path, mode = 'r') as ds:
        ds.variables['Kcb_max_obs'].set_auto_mask(False)
        Kcb_max_obs = ds.variables['Kcb_max_obs'][:, :]

    # ============ Time loop ============ #    
    # Create progress bar
    progress_bar = tqdm(total = len(dates), desc = '', unit = ' days')

    for i in range(0, len(dates)):

        # ============ Load input data and prepare output data ============ #
        # Based on available memory and previous calculation of time slices to load,
        # load a given number of time slices
        if time_slice == time_size and not already_loaded:  # if whole dataset fits in memory and it has not already been loaded
            
            # Update progress bar
            progress_bar.set_description_str(desc = 'Reading inputs')
            
            NDVI, Rain, ET0 = read_inputs(ndvi_cube_path, weather_path, i, time_slice, load_all = True)
            already_loaded = True
            write_remainder = False
            
            # Prepare output data
            # Standard outputs
            DP, E, Irrig, SWCe, SWCr, Tr = get_empty_arrays(x_size, y_size, time_slice, 6)
            # Additionnal outputs
            if additional_outputs:
                additional_outputs_data = get_empty_arrays(x_size, y_size, time_slice, len(additional_outputs.keys()), array = True)
                
            # Update progress bar
            progress_bar.set_description_str(desc = 'Running model')
            
        elif i % time_slice == 0:  # load a time slice every time i is divisible by the size of the time slice
            if i + time_slice <= time_size:  # if the time slice does not gow over the dataset size
                
                # Update progress bar
                progress_bar.set_description_str(desc = 'Reading inputs')
                
                NDVI, Rain, ET0 = read_inputs(ndvi_cube_path, weather_path, i, time_slice)
                write_remainder = False
                
                # Prepare output data
                DP, E, Irrig, SWCe, SWCr, Tr = get_empty_arrays(x_size, y_size, time_slice, 6)
                # Additionnal outputs
                if additional_outputs:
                    additional_outputs_data = get_empty_arrays(x_size, y_size, time_slice, len(additional_outputs.keys()), array = True)
                
                # Update progress bar
                progress_bar.set_description_str(desc = 'Running model')

            else:  # load the remainder when the time slice would go over the dataset size
                
                # Update progress bar
                progress_bar.set_description_str(desc = 'Reading inputs')
                
                NDVI, Rain, ET0 = read_inputs(ndvi_cube_path, weather_path, i, remainder_to_load)
                write_remainder = True
                
                # Prepare output data
                DP, E, Irrig, SWCe, SWCr, Tr = get_empty_arrays(x_size, y_size, remainder_to_load, 6)
                # Additionnal outputs
                if additional_outputs:
                    additional_outputs_data = get_empty_arrays(x_size, y_size, remainder_to_load, len(additional_outputs.keys()), array = True)
                
                # Update progress bar
                progress_bar.set_description_str(desc = 'Running model')

        if i == 0:
            E0, Tr0, Zr0, Dei0, Dep0, Dr0, Dd0, Irrig_test = get_starting_conditions(parameter_array, param_list, NDVI[0], Init_RU, Wfc, Wwp, y_size, x_size)

        # Run SAMIR daily
        DP[i % time_slice], Dd0, Dei0, Dep0, Dr0, E[i % time_slice], Irrig[i % time_slice], Irrig_test, SWCe[i % time_slice], SWCr[i % time_slice], Tr[i % time_slice], Zr0 = samir_daily(NDVI[i % time_slice], ET0[i % time_slice], Rain[i % time_slice], Wfc, Wwp, Kcb_max_obs, Irrig_auto, Irrig_test, Dr0, Dd0, Zr0, E0, Tr0, Dei0, Dep0, i, parameter_array, param_list, scaling_names, scaling_array, write_additional_data, additional_scaling_array, additional_names, additional_idx, additional_outputs_data, time_slice)
        
        # Update previous day values
        E0, Tr0  = E[i % time_slice], Tr[i % time_slice]
        
        # ============ Write outputs ============ #
        # Based on available memory and previous calculation of time slices to load,
        # write a given number of time slices
        if time_slice == time_size and i == time_size - 1:  # if whole dataset fits in memory, write data at the end of the loop
            
            # Remove unecessary data
            del NDVI, Rain, ET0
            collect()  # free up unused memory
            
            # Update progress bar
            progress_bar.set_description_str(desc = 'Writing outputs')
            
            write_outputs(save_path, DP, SWCe, SWCr, E, Tr, Irrig, additional_outputs, additional_outputs_data, i, time_slice, write_all = True)
            
            # Remove written data
            del DP, SWCe, SWCr, E, Tr, Irrig, additional_outputs_data
            additional_outputs_data = None
            collect()  # free up unused memory
            
        elif (i % time_slice == time_slice - 1) and (remainder_to_load != None):  # write a time slice every time i is divisible by the size of the time slice
            
            # Remove unecessary data
            del NDVI, Rain, ET0
            collect()  # free up unused memory
            
            # Update progress bar
            progress_bar.set_description_str(desc = 'Writing outputs')
            
            write_outputs(save_path, DP, SWCe, SWCr, E, Tr, Irrig, additional_outputs, additional_outputs_data, i, time_slice)
            
            # Remove written data
            del DP, SWCe, SWCr, E, Tr, Irrig, additional_outputs_data
            additional_outputs_data = None
            collect()  # free up unused memory

        elif i == time_size - 1 and write_remainder:  # write the remainder when the time slice would go over the dataset size
            
            # Remove unecessary data
            del NDVI, Rain, ET0
            collect()  # free up unused memory
            
            # Update progress bar
            progress_bar.set_description_str(desc = 'Writing outputs')
            
            write_outputs(save_path, DP, SWCe, SWCr, E, Tr, Irrig, additional_outputs, additional_outputs_data, i, remainder_to_load)

            # Remove written data
            del DP, SWCe, SWCr, E, Tr, Irrig, additional_outputs_data
            additional_outputs_data = None
            collect()  # free up unused memory
        
        # Update progress bar
        progress_bar.update()

    # Close progress bar
    progress_bar.close()
    
    print('\nWritting output dataset:', save_path)
    
    # Reopen output model file to reinsert dates  
    with nc.Dataset(save_path, mode = 'a') as model_outputs:
        model_outputs.variables['time'].units = f'days since {np.datetime_as_string(dates[0], unit = "D")} 00:00:00'  # set correct unit
        model_outputs.variables['time'][:] = np.arange(0, len(dates))  # save dates as integers representing the number of days since the first day
        model_outputs.sync() # flush data to disk

    return None