dev_samir_xarray.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "from dask.distributed import Client\n",
    "import os\n",
    "import numpy as np\n",
    "from typing import List, Tuple, Union\n",
    "import warnings\n",
    "import gc\n",
    "from parameters.params_samir_class import samir_parameters\n",
    "from config.config import config\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def rasterize_samir_parameters(csv_param_file: str, empty_dataset: xr.Dataset, land_cover_raster: str, chunk_size: dict) -> Tuple[xr.Dataset, dict]:\n",
    "    \"\"\"\n",
    "    Creates a raster `xarray` dataset from the csv parameter file, the land cover raster and an empty dataset\n",
    "    that contains the right structure (emptied ndvi dataset for example). For each parameter, the function loops\n",
    "    on land cover classes to fill the raster.\n",
    "\n",
    "    ## Arguments\n",
    "    1. csv_param_file: `str`\n",
    "        path to csv paramter file\n",
    "    2. empty_dataset: `xr.Dataset`\n",
    "        empty dataset that contains the right structure (emptied ndvi dataset for example).\n",
    "    3. land_cover_raster: `str`\n",
    "        path to land cover netcdf raster\n",
    "    4. chunk_size: `dict`\n",
    "        chunk_size for dask computation\n",
    "\n",
    "    ## Returns\n",
    "    1. parameter_dataset: `xr.Dataset`\n",
    "        the dataset containing all the rasterized Parameters\n",
    "    2. scale_factor: `dict`\n",
    "        dictionnary containing the scale factors for each parameter\n",
    "    \"\"\"\n",
    "    \n",
    "    # Load samir params into an object\n",
    "    table_param = samir_parameters(csv_param_file)\n",
    "    \n",
    "    # Set general variables\n",
    "    class_count = table_param.table.shape[1] - 2  # remove dtype and default columns\n",
    "    \n",
    "    # Open land cover raster\n",
    "    land_cover = xr.open_dataarray(land_cover_raster, chunks = chunk_size)\n",
    "    \n",
    "    # Create dataset\n",
    "    parameter_dataset = empty_dataset.copy(deep = True)\n",
    "    \n",
    "    # Create dictionnary containing the scale factors\n",
    "    scale_factor = {}\n",
    "    \n",
    "    # Loop on samir parameters and create \n",
    "    for parameter in table_param.table.index[1:]:\n",
    "        \n",
    "        # Create new variable and set attributes\n",
    "        parameter_dataset[parameter] = land_cover.copy(deep = True).astype('f4')\n",
    "        parameter_dataset[parameter].attrs['name'] = parameter\n",
    "        parameter_dataset[parameter].attrs['description'] = 'cf SAMIR Doc for detail'\n",
    "        parameter_dataset[parameter].attrs['scale factor'] = str(table_param.table.loc[table_param.table.index == parameter]['scale_factor'].values[0])\n",
    "        \n",
    "        # Assigne value in dictionnary\n",
    "        scale_factor[parameter] = 1/int(table_param.table.loc[table_param.table.index == parameter]['scale_factor'].values[0])\n",
    "        \n",
    "        # Loop on classes to set parameter values for each class\n",
    "        for class_val, class_name in zip(range(1, class_count + 1), table_param.table.columns[2:]):\n",
    "            \n",
    "            # Parameter values are multiplied by the scale factor in order to store all values as int16 types\n",
    "            # These values are then rounded to make sure there isn't any decimal point issues when casting the values to int16\n",
    "            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')\n",
    "    \n",
    "    # Return dataset converted to 'int16' data type to reduce memory usage\n",
    "    # and scale_factor dictionnary for later conversion\n",
    "    return parameter_dataset, scale_factor\n",
    "\n",
    "\n",
    "def setup_time_loop(calculation_variables_t1: List[str], calculation_variables_t2: List[str], empty_dataset: xr.Dataset) -> Tuple[xr.Dataset, xr.Dataset]:\n",
    "    \"\"\"\n",
    "    Creates two temporary `xarray Datasets` that will be used in the SAMIR time loop.\n",
    "    `variables_t1` corresponds to the variables for the previous day and `variables_t2`\n",
    "    corresponds to the variables for the current day. After each loop, `variables_t1`\n",
    "    takes the value of `variables_t2` for the corresponding variables.\n",
    "\n",
    "    ## Arguments\n",
    "    1. calculation_variables_t1: `List[str]`\n",
    "        list of strings containing the variable names\n",
    "        for the previous day dataset\n",
    "    2. calculation_variables_t2: `List[str]`\n",
    "        list of strings containing the variable names\n",
    "        for the current day dataset\n",
    "    3. empty_dataset: `xr.Dataset`\n",
    "        empty dataset that contains the right structure\n",
    "\n",
    "    ## Returns\n",
    "    1. variables_t1: `xr.Dataset`\n",
    "        output dataset for previous day\n",
    "    2. variables_t2: `xr.Dataset`\n",
    "        output dataset for current day\n",
    "    \"\"\"\n",
    "    \n",
    "    # Create new dataset\n",
    "    variables_t1 = empty_dataset.copy(deep = True)\n",
    "    \n",
    "    # Create empty DataArray for each variable\n",
    "    for variable in calculation_variables_t1:\n",
    "        \n",
    "        # Assign new empty DataArray\n",
    "        variables_t1[variable] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype = 'float32'))\n",
    "        variables_t1[variable].attrs['name'] = variable  # set name in attributes\n",
    "    \n",
    "    # Create new dataset\n",
    "    variables_t2 = empty_dataset.copy(deep = True)\n",
    "    \n",
    "    # Create empty DataArray for each variable\n",
    "    for variable in calculation_variables_t2:\n",
    "        \n",
    "        # Assign new empty DataArray\n",
    "        variables_t2[variable] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype = 'float32'))\n",
    "        variables_t2[variable].attrs['name'] = variable  # set name in attributes\n",
    "    \n",
    "    return variables_t1, variables_t2\n",
    "\n",
    "\n",
    "def prepare_outputs(empty_dataset: xr.Dataset, additional_outputs: List[str] = None) -> xr.Dataset:\n",
    "    \"\"\"\n",
    "    Creates the `xarray Dataset` containing the outputs of the SAMIR model that will be saved.\n",
    "    Additional variables can be saved by adding their names to the `additional_outputs` list.\n",
    "\n",
    "    ## Arguments\n",
    "    1. empty_dataset: `xr.Dataset`\n",
    "        empty dataset that contains the right structure\n",
    "    2. additional_outputs: `List[str]`\n",
    "        list of additional variable names to be saved\n",
    "\n",
    "    ## Returns\n",
    "    1. model_outputs: `xr.Dataset`\n",
    "        model outputs to be saved\n",
    "    \"\"\"\n",
    "    \n",
    "    # Evaporation and Transpiraion\n",
    "    model_outputs = empty_dataset.copy(deep = True)\n",
    "    \n",
    "    model_outputs['E'] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype = 'int16'))\n",
    "    model_outputs['E'].attrs['units'] = 'mm'\n",
    "    model_outputs['E'].attrs['standard_name'] = 'Evaporation'\n",
    "    model_outputs['E'].attrs['description'] = 'Accumulated daily evaporation in milimeters'\n",
    "    model_outputs['E'].attrs['scale factor'] = '1'\n",
    "    \n",
    "    model_outputs['Tr'] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype = 'int16'))\n",
    "    model_outputs['Tr'].attrs['units'] = 'mm'\n",
    "    model_outputs['Tr'].attrs['standard_name'] = 'Transpiration'\n",
    "    model_outputs['Tr'].attrs['description'] = 'Accumulated daily plant transpiration in milimeters'\n",
    "    model_outputs['Tr'].attrs['scale factor'] = '1'\n",
    "    \n",
    "    # Soil Water Content\n",
    "    model_outputs['SWCe'] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype = 'int16'))\n",
    "    model_outputs['SWCe'].attrs['units'] = 'mm'\n",
    "    model_outputs['SWCe'].attrs['standard_name'] = 'Soil Water Content of the evaporative zone'\n",
    "    model_outputs['SWCe'].attrs['description'] = 'Soil water content of the evaporative zone in milimeters'\n",
    "    model_outputs['SWCe'].attrs['scale factor'] = '1'\n",
    "    \n",
    "    model_outputs['SWCr'] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype = 'int16'))\n",
    "    model_outputs['SWCr'].attrs['units'] = 'mm'\n",
    "    model_outputs['SWCr'].attrs['standard_name'] = 'Soil Water Content of the root zone'\n",
    "    model_outputs['SWCr'].attrs['description'] = 'Soil water content of the root zone in milimeters'\n",
    "    model_outputs['SWCr'].attrs['scale factor'] = '1'\n",
    "    \n",
    "    # Irrigation\n",
    "    model_outputs['Irr'] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype = 'int16'))\n",
    "    model_outputs['Irr'].attrs['units'] = 'mm'\n",
    "    model_outputs['Irr'].attrs['standard_name'] = 'Irrigation'\n",
    "    model_outputs['Irr'].attrs['description'] = 'Simulated daily irrigation in milimeters'\n",
    "    model_outputs['Irr'].attrs['scale factor'] = '1'\n",
    "    \n",
    "    # Deep Percolation\n",
    "    model_outputs['DP'] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype = 'int16'))\n",
    "    model_outputs['DP'].attrs['units'] = 'mm'\n",
    "    model_outputs['DP'].attrs['standard_name'] = 'Deep Percolation'\n",
    "    model_outputs['DP'].attrs['description'] = 'Simulated daily Deep Percolation in milimeters'\n",
    "    model_outputs['DP'].attrs['scale factor'] = '1'\n",
    "    \n",
    "    if additional_outputs:\n",
    "        for var in additional_outputs:\n",
    "            model_outputs[var] = (empty_dataset.dims, np.zeros(tuple(empty_dataset.dims[d] for d in list(empty_dataset.dims)), dtype = 'int16'))\n",
    "    \n",
    "    return model_outputs\n",
    "\n",