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 dask.array as da\n",
    "import os\n",
    "import numpy as np\n",
    "from typing import List, Tuple, Union\n",
    "import warnings\n",
    "import pandas as pd\n",
    "import netCDF4 as nc\n",
    "from parameters.params_samir_class import samir_parameters\n",
    "from config.config import config\n",
    "from time import time"
   ]
  },
  {
   "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",