test_samir.py

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

Test file
"""

if __name__ == '__main__':
    
    
    import xarray as xr
    from dask.distributed import Client
    import os
    import numpy as np
    import pandas as pd
    import rasterio as rio
    from typing import List, Tuple, Union
    import warnings
    import netCDF4 as nc
    from tqdm import tqdm
    from parameters.params_samir_class import samir_parameters
    from config.config import config
    from time import time
    import webbrowser  # to open dask dashboard


    def rasterize_samir_parameters(csv_param_file: str, empty_dataset: xr.Dataset, land_cover_raster: str, chunk_size: dict) -> Tuple[xr.Dataset, dict]:
        """
        Creates a raster `xarray` dataset from the csv parameter file, the land cover raster and an empty dataset
        that contains the right structure (emptied ndvi dataset for example). For each parameter, the function loops
        on land cover classes to fill the raster.

        ## Arguments
        1. csv_param_file: `str`
            path to csv paramter file
        2. empty_dataset: `xr.Dataset`
            empty dataset that contains the right structure (emptied ndvi dataset for example).
        3. land_cover_raster: `str`
            path to land cover netcdf raster
        4. chunk_size: `dict`
            chunk_size for dask computation

        ## Returns
        1. parameter_dataset: `xr.Dataset`
            the dataset containing all the rasterized Parameters
        2. scale_factor: `dict`
            dictionnary containing the scale factors for each parameter
        """
        
        # 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, chunks = chunk_size)
        
        # Create dataset
        parameter_dataset = empty_dataset.copy(deep = True)
        
        # Create dictionnary containing the scale factors
        scale_factor = {}
        
        # Loop on samir parameters and create 
        for parameter in table_param.table.index[1:]:
            
            # Create new variable and set attributes
            parameter_dataset[parameter] = land_cover.copy(deep = True).astype('f4')
            parameter_dataset[parameter].attrs['name'] = parameter
            parameter_dataset[parameter].attrs['description'] = 'cf SAMIR Doc for detail'
            parameter_dataset[parameter].attrs['scale factor'] = str(table_param.table.loc[table_param.table.index == parameter]['scale_factor'].values[0])
            
            # Assigne value in dictionnary
            scale_factor[parameter] = 1/int(table_param.table.loc[table_param.table.index == parameter]['scale_factor'].values[0])
            
            # 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:]):
                
                # Parameter values are multiplied by the scale factor in order to store all values as int16 types
                # These values are then rounded to make sure there isn't any decimal point issues when casting the values to int16
                parameter_dataset[parameter].values = np.where(parameter_dataset[parameter].values == class_val, round(table_param.table.loc[table_param.table.index == parameter][class_name].values[0]*table_param.table.loc[table_param.table.index == parameter]['scale_factor'].values[0]), parameter_dataset[parameter].values).astype('f4')
        
        # Return dataset converted to 'int16' data type to reduce memory usage
        # and scale_factor dictionnary for later conversion
        return parameter_dataset, scale_factor


    def setup_time_loop(calculation_variables_t1: List[str], calculation_variables_t2: List[str], empty_dataset: xr.Dataset) -> Tuple[xr.Dataset, xr.Dataset]:
        """
        Creates two temporary `xarray Datasets` that will be used in the SAMIR time loop.
        `variables_t1` corresponds to the variables for the previous day and `variables_t2`
        corresponds to the variables for the current day. After each loop, `variables_t1`
        takes the value of `variables_t2` for the corresponding variables.

        ## Arguments
        1. calculation_variables_t1: `List[str]`
            list of strings containing the variable names
            for the previous day dataset
        2. calculation_variables_t2: `List[str]`
            list of strings containing the variable names
            for the current day dataset
        3. empty_dataset: `xr.Dataset`
            empty dataset that contains the right structure

        ## Returns
        1. variables_t1: `xr.Dataset`
            output dataset for previous day
        2. variables_t2: `xr.Dataset`
            output dataset for current day
        """
        
        # Create new dataset
        variables_t1 = empty_dataset.copy(deep = True)
        
        # Create empty DataArray for each variable
        for variable in calculation_variables_t1:
            
            # Assign new empty DataArray
            variables_t1[variable] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype = 'float32'))
            variables_t1[variable].attrs['name'] = variable  # set name in attributes
        
        # Create new dataset
        variables_t2 = empty_dataset.copy(deep = True)
        
        # Create empty DataArray for each variable
        for variable in calculation_variables_t2:
            
            # Assign new empty DataArray
            variables_t2[variable] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype = 'float32'))
            variables_t2[variable].attrs['name'] = variable  # set name in attributes
        
        return variables_t1, variables_t2


    def prepare_outputs(empty_dataset: xr.Dataset, additional_outputs: List[str] = 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. empty_dataset: `xr.Dataset`
            empty dataset that contains the right structure
        2. additional_outputs: `List[str]`
            list of additional variable names to be saved

        ## Returns
        1. model_outputs: `xr.Dataset`
            model outputs to be saved
        """
        
        # Evaporation and Transpiraion
        model_outputs = empty_dataset.copy(deep = True)
        
        model_outputs['E'] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype = 'int16'))
        model_outputs['E'].attrs['units'] = 'mm'
        model_outputs['E'].attrs['standard_name'] = 'Evaporation'
        model_outputs['E'].attrs['description'] = 'Accumulated daily evaporation in milimeters'
        model_outputs['E'].attrs['scale factor'] = '1000'
        
        model_outputs['Tr'] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype = 'int16'))
        model_outputs['Tr'].attrs['units'] = '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'] = '1000'
        
        # Soil Water Content
        model_outputs['SWCe'] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype = 'int16'))
        model_outputs['SWCe'].attrs['units'] = 'mm'
        model_outputs['SWCe'].attrs['standard_name'] = 'Soil Water Content of the evaporative zone'
        model_outputs['SWCe'].attrs['description'] = 'Soil water content of the evaporative zone in milimeters'
        model_outputs['SWCe'].attrs['scale factor'] = '1000'
        
        model_outputs['SWCr'] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype = 'int16'))
        model_outputs['SWCr'].attrs['units'] = 'mm'
        model_outputs['SWCr'].attrs['standard_name'] = 'Soil Water Content of the root zone'
        model_outputs['SWCr'].attrs['description'] = 'Soil water content of the root zone in milimeters'
        model_outputs['SWCr'].attrs['scale factor'] = '1000'
        
        # Irrigation
        model_outputs['Irr'] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype = 'int16'))
        model_outputs['Irr'].attrs['units'] = 'mm'
        model_outputs['Irr'].attrs['standard_name'] = 'Irrigation'
        model_outputs['Irr'].attrs['description'] = 'Simulated daily irrigation in milimeters'
        model_outputs['Irr'].attrs['scale factor'] = '1000'
        
        # Deep Percolation
        model_outputs['DP'] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype = 'int16'))
        model_outputs['DP'].attrs['units'] = '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'] = '1000'
        
        if additional_outputs:
            for var in additional_outputs:
                model_outputs[var] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype = 'int16'))
        
        return model_outputs


    def xr_maximum(ds: xr.DataArray, value: Union[xr.DataArray, float, int]) -> xr.DataArray:
        """
        Equivalent of `numpy.maximum(ds, value)` for xarray DataArrays

        ## Arguments
        1. ds: `xr.DataArray`
            datarray to compare
        2. value: `Union[xr.DataArray, float, int]`
            value (scalar or dataarray) to compare

        ## Returns
        1. output: `xr.DataArray`
            resulting dataarray with maximum value element-wise
        """
        return xr.where(ds <= value, value, ds, keep_attrs = True)


    def xr_minimum(ds: xr.DataArray, value: Union[xr.DataArray, float, int]) -> xr.DataArray:
        """
        Equivalent of `numpy.minimum(ds, value)` for xarray DataArrays

        ## Arguments
        1. ds: `xr.DataArray`
            datarray to compare
        2. value: `Union[xr.DataArray, float, int]`
            value (scalar or dataarray) to compare

        ## Returns
        1. output: `xr.DataArray`
            resulting dataarray with minimum value element-wise
        """
        return xr.where(ds >= value, value, ds, keep_attrs = True)


    def calculate_diff_re(TAW: np.ndarray, Dr: np.ndarray, Zr: np.ndarray, RUE: np.ndarray, De: np.ndarray, FCov: np.ndarray, Ze_: np.ndarray, DiffE_: np.ndarray, scale_dict: dict) -> np.ndarray:
        """
        Calculates the diffusion between the top soil layer and the root layer.

        ## 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
        5. De: `np.ndarray`
            depletion of the evaporative layer
        6. FCov: `np.ndarray`
            fraction cover of plants
        7. Ze_: `np.ndarray`
            height of evaporative layer (paramter)
        8. DiffE_: `np.ndarray`
            diffusion coefficient between evaporative
            and root layers (unitless, parameter)
        9. scale_dict: `dict`
            dictionnary containing the scale factors for
            the rasterized parameters

        ## 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
        tmp1 = (((TAW - Dr) / Zr - (RUE - De) / (scale_dict['Ze'] * Ze_)) / FCov) * (scale_dict['DiffE'] * DiffE_)
        tmp2 = ((TAW * scale_dict['Ze'] * Ze_) - (RUE - De - Dr) * Zr) / (Zr + scale_dict['Ze'] * Ze_) - Dr
        
        # Calculate diffusion according to SAMIR equation
        diff_re = np.where(tmp1 < 0, np.maximum(tmp1, tmp2), np.minimum(tmp1, tmp2))

        # Return zero values where the 'DiffE' parameter is equal to 0
        return np.where(DiffE_ == 0, 0, diff_re)


    def calculate_diff_dr(TAW: np.ndarray, TDW: np.ndarray, Dr: np.ndarray, Zr: np.ndarray, Dd: np.ndarray, FCov: np.ndarray, Zsoil_: np.ndarray, DiffR_: np.ndarray, scale_dict: dict) -> np.ndarray:
        """
        Calculates the diffusion between the root layer and the deep layer.

        ## 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. FCov: `np.ndarray`
            fraction cover of plants
        7. Zsoil_: `np.ndarray`
            total height of soil (paramter)
        8. DiffR_: `np.ndarray`
            Diffusion coefficient between root
            and deep layers (unitless, parameter)
        9. scale_dict: `dict`
            dictionnary containing the scale factors for
            the rasterized parameters

        ## 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) / (scale_dict['Zsoil'] * Zsoil_ - Zr) - (TAW - Dr) / Zr) / FCov) * scale_dict['DiffR'] * DiffR_
        tmp2 = (TDW *Zr - (TAW - Dr - Dd) * (scale_dict['Zsoil'] * Zsoil_ - Zr)) / (scale_dict['Zsoil'] * Zsoil_) - Dd
        
        # Calculate diffusion according to SAMIR equation
        diff_dr = np.where(tmp1 < 0, np.maximum(tmp1, tmp2), np.minimum(tmp1, tmp2))
        
        # Return zero values where the 'DiffR' parameter is equal to 0
        return np.where(DiffR_ == 0, 0, diff_dr)


    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

        ## 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
        """
        
        # Temporary variables to make calculation easier to read
        tmp = fewi * (TEW - Dei)
        
        # Calculate the weighting factor to split the energy available for evaporation
        W = 1 / (1 + (fewp * (TEW - Dep) / tmp ))

        # Return W 
        return np.where(tmp > 0, W, 0)


    def calculate_Kr(TEW: np.ndarray, De: np.ndarray, REW_: np.ndarray, scale_dict: dict) -> np.ndarray:
        """
        calculates of the reduction coefficient for evaporation dependent 
        on the amount of water in the soil using the FAO-56 method

        ## Arguments
        1. TEW: `np.ndarray`
            water capacity of evaporative layer
        2. De: `np.ndarray`
            depletion of evaporative layer
        3. REW_: `np.ndarray`
            readily evaporable water
        4. scale_dict: `dict`
            dictionnary containing the scale factors for
            the rasterized parameters

        ## Returns
        1. Kr: `np.ndarray`
            Kr coefficient
        """
        
        # Formula for calculating Kr
        Kr = (TEW - De) / (TEW - scale_dict['REW'] * REW_)
        
        # Return Kr
        return np.maximum(0, np.minimum(Kr, 1))


    def update_Dr(TAW: np.ndarray, TDW: np.ndarray, Zr: np.ndarray, TAW0: np.ndarray, TDW0: np.ndarray, Dr0: np.ndarray, Dd0: np.ndarray, Zr0: np.ndarray) -> np.ndarray:
        """
        Return the updated depletion for the root layer

        ## Arguments
        1. TAW: `np.ndarray`
            water capacity of root layer for current day
        2. TDW: `np.ndarray`
            water capacity of deep layer for current day
        3. Zr: `np.ndarray`
            root layer height for current day
        4. TAW0: `np.ndarray`
            water capacity of root layer for previous day
        5. TDW0: `np.ndarray`
            water capacity of deep layer for previous day
        6. Dr0: `np.ndarray`
            depletion of the root layer for previous day
        7. Dd0: `np.ndarray`
            depletion of the deep laye for previous day
        8. 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 * (TAW - TAW0) / TDW0, 0)
        tmp2 = np.minimum(Dr0 + Dd0 * (TAW - TAW0) / TDW0, TDW)

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


    def update_Dd(TAW: np.ndarray, TDW: np.ndarray, Zr: np.ndarray, TAW0: np.ndarray, TDW0: np.ndarray, Dd0: np.ndarray, Zr0: np.ndarray) -> np.ndarray:
        """
        Return the updated depletion for the deep layer

        ## Arguments
        1. TAW: `np.ndarray`
            water capacity of root layer for current day
        2. TDW: `np.ndarray`
            water capacity of deep layer for current day
        3. TAW0: `np.ndarray`
            water capacity of root layer for previous day
        5. TDW0: `np.ndarray`
            water capacity of deep 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 * (TAW - TAW0) / TDW0, 0)
        tmp2 = np.minimum(Dd0 - Dd0 * (TAW - TAW0) / TDW0, TDW)
        
        # Return updated Dd
        return np.where(Zr > Zr0, tmp1, tmp2)


    def format_duration(timedelta: float) -> None:
            """
            Print formatted timedelta in human readable format
            (days, hours, minutes, seconds, microseconds, milliseconds, nanoseconds).

            ## Arguments
            timedelta: `float`
                time value in seconds to format

            ## Returns
            `None`
            """
            
            if timedelta < 0.9e-6:
                print(round(timedelta*1e9, 1), 'ns')
            elif timedelta < 0.9e-3:
                print(round(timedelta*1e6, 1), 'µs')
            elif timedelta < 0.9:
                print(round(timedelta*1e3, 1), 'ms')
            elif timedelta < 60:
                print(round(timedelta, 1), 's')
            elif timedelta < 3.6e3:
                print(round(timedelta//60), 'm', round(timedelta%60, 1),  's')
            elif timedelta < 24*3.6e3:
                print(round((timedelta/3.6e3)//1), 'h', round((timedelta/3.6e3)%1*60//1), 'm', round((timedelta/3.6e3)%1*60%1*60, 1), 's' ) 
            elif timedelta < 48*3.6e3:
                print(round((timedelta/(24*3.6e3))//1), 'day,', round(((timedelta/(24*3.6e3))%1*24)//1), 'h,', round((timedelta/(24*3.6e3))%1*24%1*60//1), 'm,',  round((timedelta/(24*3.6e3))%1*24%1*60%1*60, 1), 's')
            else:
                print(round((timedelta/(24*3.6e3))//1), 'days,', round(((timedelta/(24*3.6e3))%1*24)//1), 'h,', round((timedelta/(24*3.6e3))%1*24%1*60//1), 'm,',  round((timedelta/(24*3.6e3))%1*24%1*60%1*60, 1), 's')
            
            return None


    @profile  # type: ignore
    def run_samir(json_config_file: str, csv_param_file: str, ndvi_cube_path: str, precip_cube_path: str, ET0_cube_path: str, soil_params_path: str, land_cover_path: str, chunk_size: dict, save_path: str, max_GB: int = 2) -> None:
        
        # warnings.simplefilter("error", category = RuntimeWarning())
        warnings.filterwarnings("ignore", message="invalid value encountered in cast")
        warnings.filterwarnings("ignore", message="invalid value encountered in divide")
        np.errstate(all = 'ignore')
        
        #============ General parameters ============#
        config_params = config(json_config_file)
        calculation_variables_t2 = ['diff_rei', 'diff_rep', 'diff_dr' , 'Dd', 'De', 'Dei', 'Dep', 'DP', 'Dr', 'FCov', 'Irrig', 'Kcb', 'Kei', 'Kep', 'Ks', 'Kti', 'Ktp', 'RUE', 'SWCe', 'SWCr', 'TAW', 'TDW', 'TEW', 'Tei', 'Tep', 'W', 'Zr', 'fewi', 'fewp']
        calculation_variables_t1 = ['Dr', 'Dd', 'TAW', 'TDW', 'Zr']
        
        #============ Manage inputs ============#
        # NDVI
        ndvi_cube = xr.open_mfdataset(ndvi_cube_path, chunks = chunk_size, parallel = True)
        
        # Weather
        # ## Open geotiff cubes and rename variables and coordinates
        # prec_cube = xr.open_mfdataset(precip_cube_path, chunks = chunk_size, parallel = True).astype('u2').rename({'band': 'time', 'band_data': 'prec'})
        # ET0_cube = xr.open_mfdataset(ET0_cube_path, chunks = chunk_size, parallel = True).astype('u2').rename({'band': 'time', 'band_data': 'ET0'})
        
        # ## Reset times values 
        # prec_cube['time'] = pd.date_range(start = config_params.start_date, end = config_params.end_date, freq = 'D')
        # ET0_cube['time'] = pd.date_range(start = config_params.start_date, end = config_params.end_date, freq = 'D')
        
        # ## Remove unwanted attributes
        # del prec_cube.prec.attrs['AREA_OR_POINT'], ET0_cube.ET0.attrs['AREA_OR_POINT']
        
        # # Soil
        # soil_params = xr.open_mfdataset(soil_params_path, chunks = chunk_size, parallel = True).astype('f4')
        
        # SAMIR Parameters
        param_dataset, scale_factor = rasterize_samir_parameters(csv_param_file, ndvi_cube.drop_vars(['ndvi', 'time']), land_cover_path, chunk_size = chunk_size)
        
        # SAMIR Variables
        variables_t1, variables_t2 = setup_time_loop(calculation_variables_t1, calculation_variables_t2, ndvi_cube.drop_vars(['ndvi', 'time']))
        
        # # Manage loading of data based on disk size of inputs
        # if ndvi_cube.nbytes < max_GB * (1024)**3:
        #     ndvi_cube.load()
            
        # if weather_cube.nbytes < max_GB * (1024)**3:
        #     weather_cube.load()

        #============ Prepare outputs ============#
        model_outputs = prepare_outputs(ndvi_cube.drop_vars(['ndvi']))
        
        # Create encoding dictionnary
        for variable in list(model_outputs.keys()):
            # Write encoding dict
            encoding_dict = {}
            encod = {}
            encod['dtype'] = 'i2'
            encoding_dict[variable] = encod
            
        # Save empty output
        model_outputs.to_netcdf(save_path, encoding = encoding_dict)
        model_outputs.close()
        
        #============ Prepare time iterations ============#
        dates = ndvi_cube.time.values
        ndvi_cube.close()
        
        #============ Create aliases for better readability ============#
        
        # Variables for current day
        # var = da.from_array(dataarray, chunks = (5, 5))
        diff_rei = variables_t2.diff_rei.to_numpy()
        diff_rep = variables_t2.diff_rep.to_numpy()
        diff_dr = variables_t2.diff_dr.to_numpy()
        Dd = variables_t2.Dd.to_numpy()
        De = variables_t2.De.to_numpy()
        Dei = variables_t2.Dei.to_numpy()
        Dep = variables_t2.Dep.to_numpy()
        DP = variables_t2.DP.to_numpy()
        Dr = variables_t2.Dr.to_numpy()
        FCov = variables_t2.FCov.to_numpy()
        Irrig = variables_t2.Irrig.to_numpy()
        Kcb = variables_t2.Kcb.to_numpy()
        Kei = variables_t2.Kei.to_numpy()
        Kep = variables_t2.Kep.to_numpy()
        Ks = variables_t2.Ks.to_numpy()
        Kti = variables_t2.Kti.to_numpy()
        Ktp = variables_t2.Ktp.to_numpy()
        RUE = variables_t2.RUE.to_numpy()
        SWCe = variables_t2.SWCe.to_numpy()
        SWCr = variables_t2.SWCr.to_numpy()
        TAW = variables_t2.TAW.to_numpy()
        TDW = variables_t2.TDW.to_numpy()
        TEW = variables_t2.TEW.to_numpy()
        Tei = variables_t2.Tei.to_numpy()
        Tep = variables_t2.Tep.to_numpy()
        Zr = variables_t2.Zr.to_numpy()
        W = variables_t2.W.to_numpy()
        fewi = variables_t2.fewi.to_numpy()
        fewp = variables_t2.fewp.to_numpy()
        
        # Variables for previous day
        TAW0 = variables_t1.TAW.to_numpy()
        TDW0 = variables_t1.TDW.to_numpy()
        Dr0 = variables_t1.Dr.to_numpy()
        Dd0 = variables_t1.Dd.to_numpy()
        Zr0 = variables_t1.Zr.to_numpy()
        
        # Parameters
        # Parameters have an underscore (_) behind their name for recognition 
        DiffE_ = param_dataset.DiffE.to_numpy()
        DiffR_ = param_dataset.DiffR.to_numpy()
        FW_ = param_dataset.FW.to_numpy()
        Fc_stop_ = param_dataset.Fc_stop.to_numpy()
        FmaxFC_ = param_dataset.FmaxFC.to_numpy()
        Foffset_ = param_dataset.Foffset.to_numpy()
        Fslope_ = param_dataset.Fslope.to_numpy()
        Init_RU_ = param_dataset.Init_RU.to_numpy()
        Irrig_auto_ = param_dataset.Irrig_auto.to_numpy()
        Kcmax_ = param_dataset.Kcmax.to_numpy()
        KmaxKcb_ = param_dataset.KmaxKcb.to_numpy()
        Koffset_ = param_dataset.Koffset.to_numpy()
        Kslope_ = param_dataset.Kslope.to_numpy()
        Lame_max_ = param_dataset.Lame_max.to_numpy()
        REW_ = param_dataset.REW.to_numpy()
        Ze_ = param_dataset.Ze.to_numpy()
        Zsoil_ = param_dataset.Zsoil.to_numpy()
        maxZr_ = param_dataset.maxZr.to_numpy()
        minZr_ = param_dataset.minZr.to_numpy()
        p_ = param_dataset.p.to_numpy()
        
        # scale factors
        # Scale factors have the following name scheme : s_ + parameter_name
        s_FW = scale_factor['FW']
        s_Fc_stop = scale_factor['Fc_stop']
        s_FmaxFC = scale_factor['FmaxFC']
        s_Foffset = scale_factor['Foffset']
        s_Fslope = scale_factor['Fslope']
        s_Init_RU = scale_factor['Init_RU']
        # s_Irrig_auto = scale_factor['Irrig_auto']
        s_Kcmax = scale_factor['Kcmax']
        s_KmaxKcb = scale_factor['KmaxKcb']
        s_Koffset = scale_factor['Koffset']
        s_Kslope = scale_factor['Kslope']
        s_Lame_max = scale_factor['Lame_max']
        s_REW = scale_factor['REW']
        s_Ze = scale_factor['Ze']
        s_Zsoil = scale_factor['Zsoil']
        s_maxZr = scale_factor['maxZr']
        s_minZr = scale_factor['minZr']
        s_p = scale_factor['p']
        
        # input data
        with nc.Dataset(ndvi_cube_path, mode = 'r') as ds:
            # Dimensions of ndvi dataset : (time, x, y)
            ndvi = ds.variables['ndvi'][0,:,:] / 255
        with nc.Dataset(soil_params_path, mode = 'r') as ds:
            FC = ds.variables['FC'][:,:]
            WP = ds.variables['WP'][:,:]
        with rio.open(precip_cube_path, mode = 'r') as ds:
            prec = ds.read(1) / 1000
        with rio.open(ET0_cube_path, mode = 'r') as ds:
            ET0 = ds.read(1) / 1000
        
        # ndvi = ndvi_cube.ndvi.sel({'time': dates[0]}).to_numpy() / 255
        # prec = prec_cube.prec.sel({'time': dates[0]}).to_numpy() / 1000
        # ET0 = ET0_cube.ET0.sel({'time': dates[0]}).to_numpy() / 1000
        # FC = soil_params.FC.to_numpy()
        # WP = soil_params.WP.to_numpy()
        # ndvi_cube.close()
        # prec_cube.close()
        # ET0_cube.close()
        # soil_params.close()
        
        # Create progress bar
        progress_bar = tqdm(total = len(dates), desc = 'Running model', unit = ' days')

        #============ First day initialization ============#
        # Fraction cover
        FCov = s_Fslope * Fslope_ * ndvi + s_Foffset * Foffset_
        FCov = np.minimum(np.maximum(FCov, 0), s_Fc_stop * Fc_stop_)
        
        # Root depth upate
        Zr = s_minZr * minZr_ + (FCov / (s_FmaxFC * FmaxFC_)) * s_maxZr * (maxZr_ - minZr_)
        
        # Water capacities
        TEW = (FC - WP/2) * s_Ze * Ze_
        RUE = (FC - WP) * s_Ze * Ze_
        TAW = (FC - WP) * Zr
        TDW = (FC - WP) * (s_Zsoil * Zsoil_ - Zr)  # Zd = Zsoil - Zr
        
        # Depletions
        Dei = RUE * (1 - s_Init_RU * Init_RU_)
        Dep = RUE * (1 - s_Init_RU * Init_RU_)
        Dr = TAW * (1 - s_Init_RU * Init_RU_)
        Dd = TDW * (1 - s_Init_RU * Init_RU_)
        
        # Irrigation  TODO : find correct method for irrigation
        Irrig = np.minimum(np.maximum(Dr - prec, 0), s_Lame_max * Lame_max_) * Irrig_auto_
        Irrig = np.where(Dr > TAW * s_p * p_, Irrig, 0)
        
        # Kcb
        Kcb = np.minimum(s_Kslope * Kslope_ * ndvi + s_Koffset * Koffset_, s_KmaxKcb * KmaxKcb_)
        
        # Update depletions with rainfall and/or irrigation
        
        ## DP  
        DP = - np.minimum(Dd + np.minimum(Dr - prec - Irrig, 0), 0)
        
        ## De
        Dei = np.minimum(np.maximum(Dei - prec - Irrig / (s_FW * FW_ / 100), 0), TEW)
        Dep = np.minimum(np.maximum(Dep - prec, 0), TEW)
        
        fewi = np.minimum(1 - FCov, (s_FW * FW_ / 100))
        fewp = 1 - FCov - fewi
        
        De = np.divide((Dei * fewi + Dep * fewp), (fewi + fewp))
        De = np.where(np.isfinite(De), De, Dei * (s_FW * FW_ / 100) + Dep * (1 - (s_FW * FW_ / 100)))

        ## Dr
        Dr = np.minimum(np.maximum(Dr - prec - Irrig, 0), TAW)
        
        ## Dd
        Dd = np.minimum(np.maximum(Dd + np.minimum(Dr - prec - Irrig, 0), 0), TDW)
        
        # Diffusion coefficients
        diff_rei = calculate_diff_re(TAW, Dr, Zr, RUE, Dei, FCov, Ze_, DiffE_, scale_factor)
        diff_rep = calculate_diff_re(TAW, Dr, Zr, RUE, Dep, FCov, Ze_, DiffE_, scale_factor)
        diff_dr = calculate_diff_dr(TAW, TDW, Dr, Zr, Dd, FCov, Zsoil_, DiffR_, scale_factor)    
        
        # Weighing factor W
        W = calculate_W(TEW, Dei, Dep, fewi, fewp)
        
        # Soil water content of evaporative layer
        SWCe = 1 - De/TEW
        # Soil water content of root layer
        SWCr = 1 - Dr/TAW
        
        # Water Stress coefficient
        Ks = np.minimum((TAW - Dr) / (TAW * (1 - s_p * p_)), 1)
        
        # Reduction coefficient for evaporation
        Kei = np.minimum(W * calculate_Kr(TEW, Dei, REW_, scale_factor) * (s_Kcmax * Kcmax_ - Kcb), fewi * s_Kcmax * Kcmax_)
        Kep = np.minimum((1 - W) * calculate_Kr(TEW, Dep, REW_, scale_factor) * (s_Kcmax * Kcmax_ - Kcb), fewp * s_Kcmax * Kcmax_)
        
        # Prepare coefficients for evapotranspiration
        Kti = np.minimum(((s_Ze * Ze_ / Zr)**6) * (1 - Dei / TEW) / np.maximum(1 - Dr / TAW, 0.001), 1)
        Ktp = np.minimum(((s_Ze * Ze_ / Zr)**6) * (1 - Dep / TEW) / np.maximum(1 - Dr / TAW, 0.001), 1)
        Tei = Kti * Ks * Kcb * ET0
        Tep = Ktp * Ks * Kcb * ET0
        
        # Update depletions
        Dei = np.where(fewi > 0, np.minimum(np.maximum(Dei + ET0 * Kei / fewi + Tei - diff_rei, 0), TEW), np.minimum(np.maximum(Dei + Tei - diff_rei, 0), TEW))
        Dep = np.where(fewp > 0, np.minimum(np.maximum(Dep + ET0 * Kep / fewp + Tep - diff_rep, 0), TEW), np.minimum(np.maximum(Dep + Tep - diff_rep, 0), TEW))
        
        De = (Dei * fewi + Dep * fewp) / (fewi + fewp)
        De = np.where(np.isfinite(De), De, Dei * (s_FW * FW_ / 100) + Dep * (1 - (s_FW * FW_ / 100)))
        
        # Evaporation
        E = np.maximum((Kei + Kep) * ET0, 0)
        
        # Transpiration
        Tr = Kcb * Ks * ET0
        
        # Write outputs
        with nc.Dataset(save_path, mode='r+') as outputs:
            # Dimensions of output dataset : (x, y, time)
            # Deep percolation
            outputs.variables['DP'][:,:,0] = np.round(DP * 1000).astype('int16')
            # Soil water content of the evaporative layer
            outputs.variables['SWCe'][:,:,0] = np.round(SWCe * 1000).astype('int16')
            # Soil water content of the root layer
            outputs.variables['SWCr'][:,:,0] = np.round(SWCr * 1000).astype('int16')
            # Evaporation
            outputs.variables['E'][:,:,0] = np.round(E * 1000).astype('int16')
            # Transpiration
            outputs.variables['Tr'][:,:,0] = np.round(Tr * 1000).astype('int16')
            # Irrigation
            outputs.variables['Irr'][:,:,0] = np.round(Irrig * 1000).astype('int16')
        
        # Potential evapotranspiration and evaporative fraction ??
        
        # Update depletions (root and deep zones) at the end of the day
        Dr = np.minimum(np.maximum(Dr + E + Tr - diff_dr, 0), TAW)
        Dd = np.minimum(np.maximum(Dd + diff_dr, 0), TDW)
        del E, Tr
        
        # Update previous day values
        TAW0 = TAW
        TDW0 = TDW
        Dr0 = Dr
        Dd0 = Dd
        Zr0 = Zr
        
        # Update progress bar
        progress_bar.update()
        
        #============ Time loop ============#
        for i in range(1, len(dates)):
            
            # Reset input aliases
            # input data          
            with nc.Dataset(ndvi_cube_path, mode = 'r') as ds:
                # Dimensions of ndvi dataset : (time, x, y)
                ndvi = ds.variables['ndvi'][i,:,:] / 255
            with rio.open(precip_cube_path, mode = 'r') as ds:
                prec = ds.read(i+1) / 1000
            with rio.open(ET0_cube_path, mode = 'r') as ds:
                ET0 = ds.read(i+1) / 1000
                ET0_previous = ds.read(i) / 1000
        
            # Update variables
            ## Fraction cover
            FCov = s_Fslope * Fslope_ * ndvi + s_Foffset * Foffset_
            FCov = np.minimum(np.maximum(FCov, 0), s_Fc_stop * Fc_stop_)
            
            ## Root depth upate
            Zr = s_minZr * minZr_ + (FCov / (s_FmaxFC * FmaxFC_)) * s_maxZr * (maxZr_ - minZr_)
            
            # Water capacities
            TAW = (FC - WP) * Zr
            TDW = (FC - WP) * (s_Zsoil * Zsoil_ - Zr)
            
            # Update depletions
            Dr = update_Dr(TAW, TDW, Zr, TAW0, TDW0, Dr0, Dd0, Zr0)
            Dd = update_Dd(TAW, TDW, Zr, TAW0, TDW0, Dd0, Zr0)
            
            # Update param p
            p_ = np.round((np.minimum(np.maximum(s_p * p_ + 0.04 * (5 - ET0_previous), 0.1), 0.8) * (1 / s_p))).astype('i2')
            
            # Irrigation   TODO : find correct method for irrigation
            Irrig = np.minimum(np.maximum(Dr - prec, 0), s_Lame_max * Lame_max_) * Irrig_auto_
            Irrig = np.where(Dr > TAW * s_p * p_, Irrig, 0)
        
            # Kcb
            Kcb = np.minimum(s_Kslope * Kslope_ * ndvi + s_Koffset * Koffset_, s_KmaxKcb * KmaxKcb_)
            
            # DP (Deep percolation)
            DP = - np.minimum(Dd + np.minimum(Dr - prec - Irrig, 0), 0)
            
            # Update depletions with rainfall and/or irrigation  
            
            ## De
            Dei = np.minimum(np.maximum(Dei - prec - Irrig / (s_FW * FW_ / 100), 0), TEW)
            Dep = np.minimum(np.maximum(Dep - prec, 0), TEW)
            
            fewi = np.minimum(1 - FCov, (s_FW * FW_ / 100))
            fewp = 1 - FCov - fewi
            
            De = (Dei * fewi + Dep * fewp) / (fewi + fewp)
            De = np.where(np.isfinite(De), De, Dei * (s_FW * FW_ / 100) + Dep * (1 - (s_FW * FW_ / 100)))

            ## Dr
            Dr = np.minimum(np.maximum(Dr - prec - Irrig, 0), TAW)
            
            ## Dd
            Dd = np.minimum(np.maximum(Dd + np.minimum(Dr - prec - Irrig, 0), 0), TDW)
            
            # Diffusion coefficients
            diff_rei = calculate_diff_re(TAW, Dr, Zr, RUE, Dei, FCov, Ze_, DiffE_, scale_factor)
            diff_rep = calculate_diff_re(TAW, Dr, Zr, RUE, Dep, FCov, Ze_, DiffE_, scale_factor)
            diff_dr = calculate_diff_dr(TAW, TDW, Dr, Zr, Dd, FCov, Zsoil_, DiffR_, scale_factor) 
            
            # Weighing factor W
            W = calculate_W(TEW, Dei, Dep, fewi, fewp)
            
            # Soil water content of evaporative layer
            SWCe = 1 - De/TEW
            # Soil water content of root layer
            SWCr = 1 - Dr/TAW
            
            # Water Stress coefficient
            Ks = np.minimum((TAW - Dr) / (TAW * (1 - s_p * p_)), 1)
            
            # Reduction coefficient for evaporation
            Kei = np.minimum(W * calculate_Kr(TEW, Dei, REW_, scale_factor) * (s_Kcmax * Kcmax_ - Kcb), fewi * s_Kcmax * Kcmax_)
            Kep = np.minimum((1 - W) * calculate_Kr(TEW, Dei, REW_, scale_factor) * (s_Kcmax * Kcmax_ - Kcb), fewp * s_Kcmax * Kcmax_)
            
            # Prepare coefficients for evapotranspiration
            Kti = np.minimum(((s_Ze * Ze_ / Zr)**6) * (1 - Dei / TEW) / np.maximum(1 - Dr / TAW, 0.001), 1)
            Ktp = np.minimum(((s_Ze * Ze_ / Zr)**6) * (1 - Dep / TEW) / np.maximum(1 - Dr / TAW, 0.001), 1)
            Tei = Kti * Ks * Kcb * ET0
            Tep = Ktp * Ks * Kcb * ET0
            
            # Update depletions
            Dei = np.where(fewi > 0, np.minimum(np.maximum(Dei + ET0 * Kei / fewi + Tei - diff_rei, 0), TEW), np.minimum(np.maximum(Dei + Tei - diff_rei, 0), TEW))
            Dep = np.where(fewp > 0, np.minimum(np.maximum(Dep + ET0 * Kep / fewp + Tep - diff_rep, 0), TEW), np.minimum(np.maximum(Dep + Tep - diff_rep, 0), TEW))
            
            De = (Dei * fewi + Dep * fewp) / (fewi + fewp)
            De = np.where(np.isfinite(De), De, Dei * (s_FW * FW_ / 100) + Dep * (1 - (s_FW * FW_ / 100)))
            
            # Evaporation
            E = np.maximum((Kei + Kep) * ET0, 0)
            
            # Transpiration
            Tr = Kcb * Ks * ET0
            
            # Write outputs
            with nc.Dataset(save_path, mode='r+') as outputs:
                # Dimensions of output dataset : (x, y, time)
                # Deep percolation
                outputs.variables['DP'][:,:,i] = np.round(DP * 1000).astype('int16')
                # Soil water content of the evaporative layer
                outputs.variables['SWCe'][:,:,i] = np.round(SWCe * 1000).astype('int16')
                # Soil water content of the root layer
                outputs.variables['SWCr'][:,:,i] = np.round(SWCr * 1000).astype('int16')
                # Evaporation
                outputs.variables['E'][:,:,i] = np.round(E * 1000).astype('int16')
                # Transpiration
                outputs.variables['Tr'][:,:,i] = np.round(Tr * 1000).astype('int16')
                # Irrigation
                outputs.variables['Irr'][:,:,i] = np.round(Irrig * 1000).astype('int16')
            
            # Potential evapotranspiration and evaporative fraction ??
            
            # Update depletions (root and deep zones) at the end of the day
            Dr = np.minimum(np.maximum(Dr + E + Tr - diff_dr, 0), TAW)
            Dd = np.minimum(np.maximum(Dd + diff_dr, 0), TDW)
            del E, Tr
        
            # Update previous day values
            TAW0 = TAW
            TDW0 = TDW
            Dr0 = Dr
            Dd0 = Dd
            Zr0 = Zr
            
            # Update progress bar
            progress_bar.update()
        
        # Close progress bar
        progress_bar.close()
        
        return None


    data_path = '/mnt/e/DATA/DEV_inputs_test'

    size = 100

    ndvi_path = data_path + os.sep + 'ndvi_' + str(size) + '.nc'
    prec_path = data_path + os.sep + 'rain_' + str(size) + '.tif'
    ET0_path = data_path + os.sep + 'ET0_' + str(size) + '.tif'
    land_cover_path = data_path + os.sep + 'land_cover_' + str(size) + '.nc'
    json_config_file = '/home/auclairj/GIT/modspa-pixel/config/config_modspa.json'
    param_file = '/home/auclairj/GIT/modspa-pixel/parameters/csv_files/params_samir_test.csv'
    soil_path = data_path + os.sep + 'soil_' + str(size) + '.nc'
    save_path = data_path + os.sep + 'outputs_' + str(size) + '.nc'

    chunk_size = {'x': 250, 'y': 250, 'time': -1}

    t = time()

    client = Client()
    # webbrowser.open('http://127.0.0.1:8787/status', new=2, autoraise=True)

    run_samir(json_config_file, param_file, ndvi_path, prec_path, ET0_path, soil_path, land_cover_path, chunk_size, save_path)

    format_duration(time() - t)

    client.close()