cleanup after merge

test_samir.py

@@ -11,10 +11,8 @@ 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
@@ -72,8 +70,7 @@ def rasterize_samir_parameters(csv_param_file: str, empty_dataset: xr.Dataset, l
             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])
+        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:]):
@@ -81,18 +78,14 @@ def rasterize_samir_parameters(csv_param_file: str, empty_dataset: xr.Dataset, l
             # 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
             # TODO vr: formule trop longue, trouver un moyen de rendre lisible
-            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')
+            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]:
+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`
@@ -123,8 +116,7 @@ def setup_time_loop(calculation_variables_t1: List[str], calculation_variables_t
     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] = (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
@@ -134,8 +126,7 @@ def setup_time_loop(calculation_variables_t1: List[str], calculation_variables_t
     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] = (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
@@ -160,8 +151,7 @@ def prepare_outputs(empty_dataset: xr.Dataset, additional_outputs: List[str] = N
     # 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'] = (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'
@@ -175,8 +165,7 @@ def prepare_outputs(empty_dataset: xr.Dataset, additional_outputs: List[str] = N
     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'] = (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'
@@ -190,16 +179,14 @@ def prepare_outputs(empty_dataset: xr.Dataset, additional_outputs: List[str] = N
     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'] = (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'] = (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'
@@ -207,8 +194,7 @@ def prepare_outputs(empty_dataset: xr.Dataset, additional_outputs: List[str] = N
 
     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'))
+            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
 
@@ -247,8 +233,7 @@ def xr_minimum(ds: xr.DataArray, value: Union[xr.DataArray, float, int]) -> xr.D
     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,
-                      Wfc: np.ndarray, Ze: np.ndarray, DiffE: np.ndarray) -> np.ndarray:
+def calculate_diff_re(TAW: np.ndarray, Dr: np.ndarray, Zr: np.ndarray, RUE: np.ndarray, De: np.ndarray, Wfc: np.ndarray, Ze: np.ndarray, DiffE: np.ndarray) -> np.ndarray:
     """
     Calculates the diffusion between the top soil layer and the root layer.
 
@@ -288,8 +273,7 @@ def calculate_diff_re(TAW: np.ndarray, Dr: np.ndarray, Zr: np.ndarray, RUE: np.n
     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,
-                      Wfc: np.ndarray, Zsoil: np.ndarray, DiffR: np.ndarray) -> np.ndarray:
+def calculate_diff_dr(TAW: np.ndarray, TDW: np.ndarray, Dr: np.ndarray, Zr: np.ndarray, Dd: np.ndarray, Wfc: np.ndarray, Zsoil: np.ndarray, DiffR: np.ndarray) -> np.ndarray:
     """
     Calculates the diffusion between the root layer and the deep layer.
 
@@ -392,8 +376,7 @@ def calculate_Kr(TEW: np.ndarray, De: np.ndarray, REW: np.ndarray) -> np.ndarray
     return np.maximum(0, np.minimum(Kr, 1))
 
 
-def update_Dr_from_root(TAW: np.ndarray, Zr: np.ndarray, TAW0: np.ndarray, TDW0: np.ndarray,
-                        Dr0: np.ndarray, Dd0: np.ndarray, Zr0: np.ndarray) -> np.ndarray:
+def update_Dr_from_root(TAW: 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.
 
@@ -429,8 +412,7 @@ def update_Dr_from_root(TAW: np.ndarray, Zr: np.ndarray, TAW0: np.ndarray, TDW0:
     return np.where(Zr > Zr0, tmp1, tmp2)
 
 
-def update_Dd_from_root(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:
+def update_Dd_from_root(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 deep layer
 
@@ -489,22 +471,17 @@ def format_duration(timedelta: float) -> None:
     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')
+        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')
+        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')
+        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, Rain_path: str, ET0_path: str,
-              soil_params_path: str, land_cover_path: str, chunk_size: dict, save_path: str,
-              additional_outputs: List[str] = None, additional_outputs_scale: List[float] = None, max_ram_GB: int = 2) -> None:
+def run_samir(json_config_file: str, csv_param_file: str, ndvi_cube_path: str, Rain_path: str, ET0_path: str, soil_params_path: str, land_cover_path: str, chunk_size: dict, save_path: str, additional_outputs: List[str] = None, additional_outputs_scale: List[float] = None, max_ram_GB: int = 2) -> None:
 
     # Test inputs
     if len(additional_outputs) != len(additional_outputs_scale):
@@ -515,12 +492,12 @@ def run_samir(json_config_file: str, csv_param_file: str, ndvi_cube_path: str, R
     np.seterr(divide='ignore', invalid='ignore')
 
     # ============ General parameters ============#
-    # TODO vr: config_params jamais utilisé...?
+    # TODO vr: config_params jamais utilisé...?  ça viendra
     config_params = config(json_config_file)
-    # TODO vr: serait il possible de décrire le pourquoi de ses 2 listes?
-    calculation_variables_t2 = ['Diff_rei', 'Diff_rep', 'Diff_dr', 'Dd', 'De', 'Dei', 'Dep', 'DP', 'Dr', 'FCov', 'Irrig', 'Kcb', 'Kei',
-                                'Kep', 'Kri', 'Krp', 'Ks', 'Kti', 'Ktp', 'RUE', 'SWCe', 'SWCr', 'TAW', 'TDW', 'TEW', 'Tei', 'Tep', 'W',
-                                'Zr', 'fewi', 'fewp', 'p_cor']
+    
+    # calculation_variables_t2 contains the list of variables used for current day calculations
+    calculation_variables_t2 = ['Diff_rei', 'Diff_rep', 'Diff_dr', 'Dd', 'De', 'Dei', 'Dep', 'DP', 'Dr', 'FCov', 'Irrig', 'Kcb', 'Kei', 'Kep', 'Kri', 'Krp', 'Ks', 'Kti', 'Ktp', 'RUE', 'SWCe', 'SWCr', 'TAW', 'TDW', 'TEW', 'Tei', 'Tep', 'W', 'Zr', 'fewi', 'fewp', 'p_cor']
+    # calculation_variables_t1 contains the list of variables of the previous day necessary for current day calculations
     calculation_variables_t1 = ['Dr', 'Dd', 'TAW', 'TDW', 'Zr']
 
     # ============ Manage inputs ============#
@@ -649,8 +626,8 @@ def run_samir(json_config_file: str, csv_param_file: str, ndvi_cube_path: str, R
         # Dimensions of ndvi dataset : (time, x, y)
         NDVI = ds.variables['NDVI'][0, :, :] / 255
     with nc.Dataset(soil_params_path, mode='r') as ds:
-        Wfc = ds.variables['FC'][:, :]  # TODO vr: modifier FC en Wfc dans fichier sol
-        Wwp = ds.variables['WP'][:, :]  # TODO vr: modifier WP en Wwp dans fichier sol
+        Wfc = ds.variables['Wfc'][:, :]
+        Wwp = ds.variables['Wwp'][:, :]
         print('soil Wfc and Wwp:', Wfc[0, 0], Wwp[0, 0])
     with rio.open(Rain_path, mode='r') as ds:
         Rain = ds.read(1) / 1000
@@ -940,8 +917,8 @@ def run_samir(json_config_file: str, csv_param_file: str, ndvi_cube_path: str, R
 # %% MAIN
 if __name__ == '__main__':
 
-    # data_path = '/mnt/e/DATA/DEV_inputs_test'
-    data_path = './DEV_inputs_test'
+    data_path = '/mnt/e/DATA/DEV_inputs_test'
+    # data_path = './DEV_inputs_test'
 
     size = 10
 
@@ -965,10 +942,8 @@ if __name__ == '__main__':
 
     output_save_path = data_path + os.sep + 'pix_outputs_10.nc'
 
-    additional_outputs = ['Zr', 'Dei', 'Dep', 'Dr', 'Dd', 'Kei', 'Kep', 'Ks', 'W', 'Kcb', 'Kri', 'Krp', 'NDVI',
-                          'fewi', 'fewp', 'TDW', 'TAW', 'FCov', 'Tei', 'Tep', 'Diff_rei', 'Diff_rep', 'Diff_dr', 'Rain', 'p_cor']
-    additional_outputs_scale = [100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100,
-                                100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100]
+    additional_outputs = ['Zr', 'Dei', 'Dep', 'Dr', 'Dd', 'Kei', 'Kep', 'Ks', 'W', 'Kcb', 'Kri', 'Krp', 'NDVI', 'fewi', 'fewp', 'TDW', 'TAW', 'FCov', 'Tei', 'Tep', 'Diff_rei', 'Diff_rep', 'Diff_dr', 'Rain', 'p_cor']
+    additional_outputs_scale = [100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100, 100]
 
     chunk_size = {'x': 250, 'y': 250, 'time': -1}