modspa_samir.py

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

Usage of the SAMIR model in the Modspa framework. 

"""

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


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

    Arguments
    =========

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

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

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

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

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

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

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

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

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

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


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

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

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

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

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

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