Newer
Older
Jeremy Auclair
committed
# -*- coding: UTF-8 -*-
# Python
"""
11-07-2023
Jeremy Auclair
committed
@author: jeremy auclair
Jeremy Auclair
committed
"""
Jeremy Auclair
committed
import os # os management
import xarray as xr # to manage datasets

Jérémy AUCLAIR
committed
import rioxarray # to better manage dataset projections

Jérémy AUCLAIR
committed
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
Jeremy Auclair
committed
from psutil import cpu_count # to get number of physical cores available
Jeremy Auclair
committed
from gc import collect # to free up unused memory
from modspa_pixel.parameters.params_samir_class import samir_parameters # to load SAMIR parameters
Jeremy Auclair
committed
from modspa_pixel.source.code_toolbox import format_Byte_size # to print memory requirements

Jérémy AUCLAIR
committed
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.
Jeremy Auclair
committed
Jeremy Auclair
committed
Arguments
=========
Jeremy Auclair
committed
path to csv paramter file
Jeremy Auclair
committed
path to land cover netcdf raster

Jérémy AUCLAIR
committed
3. irrigation_raster: ``str``

Jérémy AUCLAIR
committed
path to irrigation mask raster

Jérémy AUCLAIR
committed
4. init_RU_raster: ``str``
path to initial soil water content. None if no input.
Jeremy Auclair
committed
Jeremy Auclair
committed
Returns
=======

Jérémy AUCLAIR
committed
1. parameter_dict: ``Dict[str, np.ndarray]``
the dictionnary containing all the rasterized Parameters
and the scale factors
Jeremy Auclair
committed
"""
Jeremy Auclair
committed

Jérémy AUCLAIR
committed
print('\nLoading SAMIR parameters')
Jeremy Auclair
committed
# Get path of min max csv file

Jérémy AUCLAIR
committed
min_max_param_file = os.path.join(os.path.dirname(csv_param_file), 'params_samir_min_max.csv')
Jeremy Auclair
committed
# Load samir params into an object
table_param = samir_parameters(csv_param_file)
Jeremy Auclair
committed
# Set general variables
class_count = table_param.table.shape[1] - 2 # remove dtype and Default columns
Jeremy Auclair
committed
# Open land cover raster
land_cover = xr.open_dataarray(land_cover_raster).to_numpy()

Jérémy AUCLAIR
committed
shape = land_cover.shape
Jeremy Auclair
committed
# Modify parameter table based on its content
table_param.test_samir_parameter(min_max_param_file)

Jérémy AUCLAIR
committed
Jeremy Auclair
committed
# If test_samir_parameter returns 0, parameters are incorrect
Jeremy Auclair
committed
if type(table_param.table) != pd.DataFrame:
Jeremy Auclair
committed
import sys # to exit script
print(f'\nModify the SAMIR parameter file: {csv_param_file}\n')
sys.exit()

Jérémy AUCLAIR
committed

Jérémy AUCLAIR
committed
# 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
Jeremy Auclair
committed
for parameter in table_param.table.index[1:]:

Jérémy AUCLAIR
committed
if parameter == 'Irrig_auto' and table_param.table.at[parameter, 'load_raster'] == 1:
# Load Irrig mask

Jérémy AUCLAIR
committed
Irrig_auto = np.ascontiguousarray(xr.open_dataarray(irrigation_raster, decode_coords = 'all').astype(np.bool_).values, dtype = np.bool_)

Jérémy AUCLAIR
committed
continue

Jérémy AUCLAIR
committed
elif parameter == 'Init_RU' and table_param.table.at[parameter, 'load_raster'] == 1:

Jérémy AUCLAIR
committed

Jérémy AUCLAIR
committed
param_list.append(parameter)

Jérémy AUCLAIR
committed
# Load initial soil water content

Jérémy AUCLAIR
committed
Init_RU = (xr.open_dataarray(init_RU_raster) / 1000).astype(np.float32).values
i+=1

Jérémy AUCLAIR
committed
continue
Jeremy Auclair
committed
# If scale_factor == 0, then the parameter does not have to be spatialized, it can stay scalar

Jérémy AUCLAIR
committed
elif parameter == 'Kcmax':

Jérémy AUCLAIR
committed
param_list.append(parameter)
# Take Default parameter value

Jérémy AUCLAIR
committed
parameter_array[i] = np.ones(shape = shape, dtype = np.float32) * np.float32(table_param.table.at[parameter, 'Default'])
i+=1
continue

Jérémy AUCLAIR
committed
# Check data type based on parameter

Jérémy AUCLAIR
committed
if (parameter != 'Irrig_auto') and (parameter != 'Init_RU'):
dtype = np.float32
param_list.append(parameter)

Jérémy AUCLAIR
committed
# Create 2D array from land cover

Jérémy AUCLAIR
committed
value = land_cover.copy().astype(dtype)
Jeremy Auclair
committed
# 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]

Jérémy AUCLAIR
committed
# Load parameter value for each class
value = np.where(land_cover == class_val, np.float32(param_value), value)
Jeremy Auclair
committed
# Assign parameter value

Jérémy AUCLAIR
committed
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)
Jeremy Auclair
committed

Jérémy AUCLAIR
committed
return parameter_array, Irrig_auto, Init_RU, param_list

Jérémy AUCLAIR
committed
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:
Jeremy Auclair
committed
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.
Jeremy Auclair
committed
Arguments
=========
Jeremy Auclair
committed
1. ndvi_path: ``str``
path to ndvi cube
Jeremy Auclair
committed
2. dimensions: ``Dict[str, int]``
Jeremy Auclair
committed
frozen dictionnary containing the dimensions of the output dataset
Jeremy Auclair
committed
3. scaling_dict: ``Dict[str, int]`` ``default = {'E': 1000, 'Tr': 1000, 'SWCe': 1000, 'SWCr': 1000, 'DP': 100, 'Irr': 100}``
Jeremy Auclair
committed
scaling dictionnary for the nominal outputs
Jeremy Auclair
committed
4. additional_outputs: ``Dict[str, int]``
list of additional variable names to be saved
Jeremy Auclair
committed
Returns
=======
Jeremy Auclair
committed
model_outputs = xr.open_dataset(ndvi_path).drop_vars(['NDVI']).copy(deep = True)
Jeremy Auclair
committed
model_outputs = model_outputs.drop_sel(time = model_outputs.time)
Jeremy Auclair
committed
# Add scaling dictionnary to global attribute
Jeremy Auclair
committed
model_outputs.attrs['name'] = 'ModSpa Pixel SAMIR output'

Jérémy AUCLAIR
committed
model_outputs.attrs['description'] = 'Outputs of the ModSpa SAMIR (FAO-56) outputs at the pixel scale. Variables are upscaled to be stored as integers.'
Jeremy Auclair
committed
model_outputs.attrs['scaling'] = str(scaling_dict)
Jeremy Auclair
committed
model_outputs['E'] = (dimensions, np.zeros(tuple(dimensions[d] for d in list(dimensions)), dtype = np.uint16))

Jérémy AUCLAIR
committed
model_outputs['E'].attrs['unit'] = 'mm'
model_outputs['E'].attrs['standard_name'] = 'Evaporation'
model_outputs['E'].attrs['description'] = 'Accumulated daily evaporation in milimeters'
Jeremy Auclair
committed
model_outputs['E'].attrs['scale factor'] = scaling_dict['E']
Jeremy Auclair
committed
model_outputs['Tr'] = (dimensions, np.zeros(tuple(dimensions[d] for d in list(dimensions)), dtype = np.uint16))

Jérémy AUCLAIR
committed
model_outputs['Tr'].attrs['unit'] = 'mm'
model_outputs['Tr'].attrs['standard_name'] = 'Transpiration'
model_outputs['Tr'].attrs['description'] = 'Accumulated daily plant transpiration in milimeters'
Jeremy Auclair
committed
model_outputs['Tr'].attrs['scale factor'] = scaling_dict['Tr']
Jeremy Auclair
committed
model_outputs['SWCe'] = (dimensions, np.zeros(tuple(dimensions[d] for d in list(dimensions)), dtype = np.uint16))

Jérémy AUCLAIR
committed
model_outputs['SWCe'].attrs['unit'] = '%'
Jeremy Auclair
committed
model_outputs['SWCe'].attrs['standard_name'] = 'Soil Water Content of the evaporative layer'
model_outputs['SWCe'].attrs['description'] = 'Soil water content of the evaporative layer in milimeters'
Jeremy Auclair
committed
model_outputs['SWCe'].attrs['scale factor'] = scaling_dict['SWCe']
Jeremy Auclair
committed
model_outputs['SWCr'] = (dimensions, np.zeros(tuple(dimensions[d] for d in list(dimensions)), dtype = np.uint16))

Jérémy AUCLAIR
committed
model_outputs['SWCr'].attrs['unit'] = '%'
Jeremy Auclair
committed
model_outputs['SWCr'].attrs['standard_name'] = 'Soil Water Content of the root layer'
model_outputs['SWCr'].attrs['description'] = 'Soil water content of the root layer in milimeters'
Jeremy Auclair
committed
model_outputs['SWCr'].attrs['scale factor'] = scaling_dict['SWCr']
Jeremy Auclair
committed
model_outputs['Irr'] = (dimensions, np.zeros(tuple(dimensions[d] for d in list(dimensions)), dtype = np.uint16))

Jérémy AUCLAIR
committed
model_outputs['Irr'].attrs['unit'] = 'mm'
model_outputs['Irr'].attrs['standard_name'] = 'Irrigation'
model_outputs['Irr'].attrs['description'] = 'Simulated daily irrigation in milimeters'
Jeremy Auclair
committed
model_outputs['Irr'].attrs['scale factor'] = scaling_dict['Irr']
Jeremy Auclair
committed
model_outputs['DP'] = (dimensions, np.zeros(tuple(dimensions[d] for d in list(dimensions)), dtype = np.uint16))

Jérémy AUCLAIR
committed
model_outputs['DP'].attrs['unit'] = 'mm'
model_outputs['DP'].attrs['standard_name'] = 'Deep Percolation'
model_outputs['DP'].attrs['description'] = 'Simulated daily Deep Percolation in milimeters'
Jeremy Auclair
committed
model_outputs['DP'].attrs['scale factor'] = scaling_dict['DP']
Jeremy Auclair
committed
for var, scale in zip(additional_outputs.keys(), additional_outputs.values()):
Jeremy Auclair
committed
model_outputs[var] = (dimensions, np.zeros(tuple(dimensions[d] for d in list(dimensions)), dtype = np.int16))
Jeremy Auclair
committed
model_outputs[var].attrs['scale factor'] = scale

Jérémy AUCLAIR
committed
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
@njit(nogil = True, cache = True)
def find_index(numba_list: list, value: any) -> int:
"""
Equivalent .index() method of python lists adapted for numba lists.
Arguments
=========
1. numba_list: ``list``
input numba typed list
2. value: ``any``
value to find in list
Returns
=======
1. index: ``int``
index of value in list
"""
# Loop through list elements
for i in range(len(numba_list)):
if numba_list[i] == value:
# Return index of value in numba_list
return i
# Return -1 if element is not found
return -1
@njit((float32[:,:,:], types.ListType(types.unicode_type), uint8[:,:], float32[:,:], float32[:,:], float32[:,:], int64, int64), nogil = True, parallel = True, fastmath = True, cache = True)
def get_starting_conditions(parameter_array: np.ndarray, param_list: list[str], NDVI: np.ndarray, Init_RU: np.ndarray, Wfc: np.ndarray, Wwp: np.ndarray, y_size: int, x_size: int) -> tuple[np.ndarray]:
"""
Calculate SAMIR starting variables.
Arguments
=========
1. parameter_array: ``np.ndarray``
parameter array
2. param_list: ``list[str]``
list of param names
3. NDVI: ``np.ndarray``
NDVI for first date
4. Init_RU: ``np.ndarray``
initial RU
5. Wfc: ``np.ndarray``
field capacity array
6. Wwp: ``np.ndarray``
weilting point array
7. y_size: ``int``
y_size of arrays
8. x_size: ``int``
x_size of arrays
Returns
=======
1. output_arrays: ``tuple[np.ndarray]``
E0, Tr0, Zr0, Dei0, Dep0, Dr0, Dd0, Irrig_test
"""
E0 = np.zeros((y_size, x_size), dtype = np.uint16)
Tr0 = np.zeros((y_size, x_size), dtype = np.uint16)
Irrig_test = np.zeros((y_size, x_size), dtype = np.bool_)
Zr0 = np.maximum(parameter_array[find_index(param_list, 'minZr')] + (np.minimum(np.maximum(parameter_array[find_index(param_list, 'Fslope')] * (NDVI / np.float32(255)).astype(np.float32) + parameter_array[find_index(param_list, 'Foffset')], np.float32(0)), parameter_array[find_index(param_list, 'FCmax')]) / parameter_array[find_index(param_list, 'FCmax')]) * (parameter_array[find_index(param_list, 'maxZr')] - parameter_array[find_index(param_list, 'minZr')]), parameter_array[find_index(param_list, 'Ze')] + np.float32(0.001))
Dei0 = (Wfc - Wwp) * parameter_array[find_index(param_list, 'Ze')] * (np.float32(1) - Init_RU)
Dep0 = (Wfc - Wwp) * parameter_array[find_index(param_list, 'Ze')] * (np.float32(1) - Init_RU)
Dr0 = (Wfc - Wwp) * Zr0 * (np.float32(1) - Init_RU)
Dd0 = (Wfc - Wwp) * (parameter_array[find_index(param_list, 'Zsoil')] - Zr0) * (np.float32(1) - Init_RU)
return E0, Tr0, Zr0, Dei0, Dep0, Dr0, Dd0, Irrig_test
@njit((float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:]), nogil = True, parallel = True, fastmath = True, cache = 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:
Calculates the diffusion between the top soil layer and the root layer. Uses numba for faster and parallel calculation.
Jeremy Auclair
committed
Arguments
=========
water capacity of root layer
depletion of root layer
height of root layer
total available surface water = (Wfc-Wwp)*Ze
depletion of the evaporative layer
field capacity
height of evaporative layer (paramter)
diffusion coefficient between evaporative
and root layers (unitless, parameter)
Jeremy Auclair
committed
Returns
=======
the diffusion between the top soil layer and
the root layer
"""
# Temporary variables to make calculation easier to read
# TODO: check equation
tmp1 = (((TAW - Dr) / Zr - (RUE - De) / Ze) / Wfc) * DiffE
tmp2 = ((TAW * Ze) - (RUE - De - Dr) * Zr) / (Zr + Ze) - Dr
# Calculate diffusion according to SAMIR equation
# Return zero values where the 'DiffE' parameter is equal to 0
return np.where(DiffE == 0, np.float32(0), np.where(tmp1 < 0, np.maximum(tmp1, tmp2), np.minimum(tmp1, tmp2)))
@njit((float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:]), nogil = True, parallel = True, fastmath = True, cache = True)
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. Uses numba for faster and parallel calculation.
Jeremy Auclair
committed
Arguments
=========
water capacity of root layer
water capacity of deep layer
depletion of root layer
height of root layer
depletion of deep layer
field capacity
total height of soil (paramter)
Diffusion coefficient between root
and deep layers (unitless, parameter)
Jeremy Auclair
committed
Returns
=======
the diffusion between the root layer and the
deep layer
"""
# Temporary variables to make calculation easier to read
tmp1 = (((TDW - Dd) / (Zsoil - Zr) - (TAW - Dr) / Zr) / Wfc) * DiffR
tmp2 = ((TDW * Zr - (TAW - Dr - Dd) * (Zsoil - Zr)) / Zsoil) - Dd
# Calculate diffusion according to SAMIR equation
# Return zero values where the 'DiffR' parameter is equal to 0
return np.where(DiffR == 0, np.float32(0), np.where(tmp1 < 0, np.maximum(tmp1, tmp2), np.minimum(tmp1, tmp2)))
@njit((float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:]), nogil = True, parallel = True, fastmath = True, cache = True)
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.
Uses numba for faster and parallel calculation.
Jeremy Auclair
committed
Arguments
=========
water capacity of evaporative layer
depletion of the evaporative layer
(irrigation part)
depletion of the evaporative layer
(precipitation part)
soil fraction which is wetted by irrigation
and exposed to evaporation
soil fraction which is wetted by precipitation
and exposed to evaporation
Jeremy Auclair
committed
Returns
=======
weighting factor W
# Calculate the weighting factor to split the energy available for evaporation
# Equation: W = 1 / (1 + (fewp * (TEW - Dep) / fewi * (TEW - Dei)))
# Return W
# Equation: W = where(fewi * (TEW - Dei) > 0, W, 0)
return np.where(fewi * (TEW - Dei) > 0, 1 / (1 + (fewp * (TEW - Dep) / fewi * (TEW - Dei))), np.float32(0))
@njit((float32[:,:], float32[:,:], float32[:,:]), nogil = True, parallel = True, fastmath = True, cache = True)
def calculate_Kr(TEW: np.ndarray, De: np.ndarray, REW: np.ndarray) -> np.ndarray:
calculates of the reduction coefficient for evaporation dependent
on the amount of water in the soil using the FAO-56 method. Uses numba for faster and parallel calculation.
Jeremy Auclair
committed
Arguments
=========
water capacity of evaporative layer
depletion of evaporative layer
fixed readily evaporable water
Jeremy Auclair
committed
Returns
=======
Kr coefficient
"""
# Return Kr
return np.maximum(np.float32(0), np.minimum((TEW - De) / (TEW - REW), np.float32(1)))
@njit((float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:]), nogil = True, parallel = True, fastmath = True, cache = True)
def calculate_Ks(Dr: np.ndarray, TAW: np.ndarray, p: np.ndarray, E0: np.ndarray, Tr0: np.ndarray) -> np.ndarray:
Calculate Ks coefficient after day 1. Uses numba for faster and parallel calculation.
Arguments
=========
1. Dr: ``np.ndarray``
depletion of the root layer
water capacity of the root layer
fraction of TAW available for plant without inducing stress
surface evaporation of previous day
plant transpiration of previous day
Ks coefficient
"""
# Return Ks
# Equation: Ks = min((TAW - Dr) / (TAW * (1 - (p + 0.04 * (5 - (E0 + Tr0))))), 1)
return np.minimum((TAW - Dr) / (TAW * (np.float32(1) - (p + 0.04 * (np.float32(5) - (E0 + Tr0))))), np.float32(1)).astype(np.float32)

Jérémy AUCLAIR
committed
@njit((float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:]), nogil = True, parallel = True, fastmath = True, cache = True)
def calculate_Ke(W: np.ndarray, De: np.ndarray, TEW: np.ndarray, REW: np.ndarray, Kcmax: np.ndarray, Kcb: np.ndarray, few: np.ndarray) -> np.ndarray:
"""
Calculate the evaporation Ke coefficient.
Arguments
=========
1. W: ``np.ndarray``
weighting factor to split the energy available
for evaporation
2. De: ``np.ndarray``
Dei or Dep, depletion of the evaporative layer
3. TEW: ``np.ndarray``
water capacity of the evaporative layer
4. REW: ``np.ndarray``
fixed readily evaporable water
5. Kcmax: ``np.float32``
maximum possible evaporation in the atmosphere (scalar)
6. Kcb: ``np.ndarray``
crop coefficient
7. few: ``np.ndarray``
fewi or fewp, soil fraction which is wetted by
irrigation or precipitation and exposed to evaporation
Returns
=======
1. Ke: ``np.ndarray``
evaporation coefficient
"""
# Equation: Kei = np.minimum(W * Kri * (Kc_max - Kcb), fewi * Kc_max)
# Equation: Kep = np.minimum((1 - W) * Krp * (Kc_max - Kcb), fewp * Kc_max)
# Return Ke
return np.minimum(W * calculate_Kr(TEW, De, REW) * (Kcmax - Kcb), few * Kcmax)

Jérémy AUCLAIR
committed
@njit((float32[:,:], float32[:,:], float32[:,:], float32[:,:], boolean[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], boolean[:,:], float32[:,:], float32[:,:]), nogil = True, parallel = True, fastmath = True, cache = True)
def calculate_irrig(Dr: np.ndarray, TAW: np.ndarray, Rain: np.ndarray, Kcb: np.ndarray, Irrig_auto: np.ndarray, Lame_max: np.ndarray, Lame_min: np.ndarray, Kcb_min_start_Irrig: np.ndarray, frac_Kcb_stop_irrig: np.ndarray, Irrig_test: np.ndarray, frac_TAW: np.ndarray, Kcb_max_obs: np.ndarray) -> np.ndarray:
Calculate automatic irrigation after day one. Uses numba for faster and parallel calculation.
Arguments
=========
1. Dr: ``np.ndarray``
depletion of the root layer
water capacity of the root layer

Jérémy AUCLAIR
committed
3. Rain: ``np.ndarray``
precipitation of the current day

Jérémy AUCLAIR
committed
4. Kcb: ``np.ndarray``
crop coefficient value

Jérémy AUCLAIR
committed
5. Irrig_auto: ``np.ndarray`` ``boolean``
parameter for automatic irrigation

Jérémy AUCLAIR
committed
6. Lame_max: ``np.ndarray``
Jeremy Auclair
committed
maximum amount of possible irrigation in mm

Jérémy AUCLAIR
committed
7. Lame_min: ``np.ndarray``
Jeremy Auclair
committed
minimum amount of possible irrigation in mm

Jérémy AUCLAIR
committed
8. Kcb_min_start_Irrig: ``np.ndarray``
Jeremy Auclair
committed
minimum value of Kcb under which no irrigation is allowed

Jérémy AUCLAIR
committed
9. frac_Kcb_stop_irrig: ``np.ndarray``
Jeremy Auclair
committed
Kcb threshold to stop irrigation after reaching maximum

Jérémy AUCLAIR
committed
10. Irrig_test: ``np.ndarray`` ``boolean``
Jeremy Auclair
committed
boolean array that is true after the Kcb has reached
frac_Kcb_stop_irrig of its maximum

Jérémy AUCLAIR
committed
11. frac_TAW: ``np.ndarray``
Jeremy Auclair
committed
fraction of depleted TAW under which irrigation is triggered

Jérémy AUCLAIR
committed
12: Kcb_max_obs: ``np.ndarray``
Jeremy Auclair
committed
observed maximum value of Kcb
Returns
=======
1. Irrig: ``np.ndarray``
simulated irrigation for the current day
"""
Jeremy Auclair
committed
# First step: calculate irrigation to fill the depletion up to a maximum value of Lame_max
Jeremy Auclair
committed
Irrig = Irrig_auto * np.maximum(np.minimum(Dr - Rain, Lame_max), Lame_min)
Jeremy Auclair
committed
# Second step: only keep irrigation where the depletion is higher than a fraction (frac_TAW) of TAW and Kcb is higher than a minimum value
Irrig = np.where((Dr > frac_TAW * TAW) & (Kcb > Kcb_min_start_Irrig), Irrig, np.float32(0))
# Third step: if Kcb has already gone higher than a fraction (frac_Kcb_stop_irrig) of its maximum value,
# turn off irrigation when goes under that fraction of maximum Kcb

Jérémy AUCLAIR
committed
return np.where((Irrig_test) & (Kcb < frac_Kcb_stop_irrig * Kcb_max_obs), np.float32(0), Irrig)
@njit((float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:]), nogil = True, parallel = True, fastmath = True, cache = True)
def calculate_Te(De: np.ndarray, Dr: np.ndarray, Ks: np.ndarray, Kcb: np.ndarray, Ze: np.ndarray, Zr: np.ndarray, TEW: np.ndarray, TAW: np.ndarray, ET0: np.ndarray) -> np.ndarray:
Calculate Te (root uptake of water) coefficient for current day. Uses numba for faster and parallel calculation.
Arguments
=========
1. De: ``np.ndarray``
Dei or Dep, depletion of the evaporative layer
depletion of the roor layer
stress coefficient
crop coefficient
height of the evaporative layer
heigh of the root layer
water capacity of the evaporative layer
water capacity of the root layer
reference evapotranspiration
Te coefficient
"""
# Equation: Kt = min( ((Ze / Zr)**0.6) * (1 - De/TEW) / max(1 - Dr / TAW, 0.01), 1)
# Equation: Te = Kt * Ks * Kcb * ET0

Jérémy AUCLAIR
committed
return (np.minimum(((Ze / Zr) ** 0.6) * (np.float32(1) - De / TEW) / np.maximum(np.float32(1) - Dr / TAW, 0.001), np.float32(1)) * Ks * Kcb * ET0).astype(np.float32)
@njit((float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:]), nogil = True, parallel = True, fastmath = True, cache = True)
def update_De_from_Diff(De : np.ndarray, few: np.ndarray, Ke: np.ndarray, Te: np.ndarray, Diff_re: np.ndarray, TEW: np.ndarray, ET0: np.ndarray) -> np.ndarray:
Last update step for Dei and Dep depletions. Uses numba for faster and parallel calculation.
Arguments
=========
1. De: ``np.ndarray``
Dei or Dep, depletion of the evaporative layer
fewi or fewp, soil fraction which is wetted by
irrigation or precipitation and exposed to evaporation
Kei or Kep, evaporation coefficient for soil fraction
irrigated or rainfed and exposed to evaporation
root uptake of water
dffusion between the root and evaporation layers
water capacity of the evaporative layer
reference evapotranspiration of the current day
updated Dei or Dep
# Update Dei and Dep depletions
# Equation: De = where(few > 0, min(max(De + ET0 * Ke / few + Te - Diff_re, 0), TEW), min(max(De + Te - Diff_re, 0), TEW))
return np.where(few > np.float32(0), np.minimum(np.maximum(De + ET0 * Ke / few + Te - Diff_re, np.float32(0)), TEW), np.minimum(np.maximum(De + Te - Diff_re, np.float32(0)), TEW))
@njit((float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:]), nogil = True, parallel = True, fastmath = True, cache = True)
def update_Dr_from_root(Wfc: np.ndarray, Wwp: np.ndarray, Zr: np.ndarray, Zsoil: np.ndarray, Dr0: np.ndarray, Dd0: np.ndarray, Zr0: np.ndarray) -> np.ndarray:
Return the updated depletion for the root layer. Uses numba for faster and parallel calculation.
Jeremy Auclair
committed
Arguments
=========
field capacity
field wilting point
root depth for current day
total soil depth (parameter)
depletion of the root layer for previous day
depletion of the deep laye for previous day
root layer height for previous day
Jeremy Auclair
committed
Returns
=======
updated depletion for the root layer
"""
# Temporary variables to make calculation easier to read
tmp1 = np.maximum(Dr0 + Dd0 * ((Wfc - Wwp) * (Zr - Zr0)) / ((Wfc - Wwp) * (Zsoil - Zr0)), np.float32(0))
tmp2 = np.maximum(Dr0 + Dr0 * ((Wfc - Wwp) * (Zr - Zr0)) / ((Wfc - Wwp) * Zr0), np.float32(0))
# Return updated Dr
return np.where(Zr > Zr0, tmp1, tmp2)
@njit((float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:]), nogil = True, parallel = True, fastmath = True, cache = True)
def update_Dd_from_root(Wfc: np.ndarray, Wwp: np.ndarray, Zr: np.ndarray, Zsoil: np.ndarray, Dr0: np.ndarray, Dd0: np.ndarray, Zr0: np.ndarray) -> np.ndarray:
Return the updated depletion for the deep layer. Uses numba for faster and parallel calculation.
Jeremy Auclair
committed
Arguments
=========
field capacity
field wilting point
root depth for current day
total soil depth (parameter)
depletion of the root layer for previous day
depletion of the deep laye for previous day
root layer height for previous day
Jeremy Auclair
committed
Returns
=======
updated depletion for the deep layer
"""
# Temporary variables to make calculation easier to read
tmp1 = np.maximum(Dd0 - Dd0 * ((Wfc - Wwp) * (Zr - Zr0)) / ((Wfc - Wwp) * (Zsoil - Zr0)), np.float32(0))
tmp2 = np.maximum(Dd0 - Dr0 * ((Wfc - Wwp) * (Zr - Zr0)) / ((Wfc - Wwp) * Zr0), np.float32(0))
# Return updated Dd
return np.where(Zr > Zr0, tmp1, tmp2)
@njit((float32[:,:], float32[:,:], float32[:,:], float32[:,:], float32[:,:]), nogil = True, parallel = True, fastmath = True, cache = True)
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
def calculate_SWCe(Dei: np.ndarray, Dep: np.ndarray, fewi: np.ndarray, fewp: np.ndarray, TEW: np.ndarray) -> np.ndarray:
"""
Calculate the soil water content of the evaporative layer.
Arguments
=========
1. Dei: ``np.ndarray``
depletion of the evaporative layer for the irrigated part
2. Dep: ``np.ndarray``
depletion of the evaporative layer for the rainfed part
3. fewi: ``np.ndarray``
soil fraction which is wetted by irrigation and exposed
to evaporation
4. fewp: ``np.ndarray``
soil fraction which is wetted by precipitation and exposed
to evaporation
5. TEW: ``np.ndarray``
water capacity of the evaporative layer
Returns
=======
1. SWCe: ``np.ndarray``
soil water content of the evaporative layer
"""
# Return SWCe

Jérémy AUCLAIR
committed
return np.where((fewi + fewp) > np.float32(0), np.float32(1) - ((Dei * fewi + Dep * fewp) / ((fewi + fewp) * TEW)), np.float32(1) - ((Dei + Dep) / (np.float32(2) * TEW)))
def calculate_memory_requirement(x_size: int, y_size: int, time_size: int, nb_inputs: int, nb_outputs: int, nb_variables: int, nb_params: int, nb_bytes: int) -> float:
Calculate memory requirement (GiB) of calculation if all datasets where loaded in memory.
Used to determine how to divide the datasets in times chunks for more efficient I/O
operations.
number of input variables
number of ouput variables
number of calculation variables
number of raster parameters
Jeremy Auclair
committed
8. nb_bytes: ``int``
number of bytes of datatype for inputs and outputs
Returns
=======
1. total_memory_requirement: ``float``
calculation memory requirement in GiB
"""
# Memory requirement of input datasets
input_memory_requirement = (x_size * y_size * time_size * nb_inputs * nb_bytes) / (1024**3)
Jeremy Auclair
committed
# Memory requirement of calculation variables

Jérémy AUCLAIR
committed
calculation_memory_requirement = (x_size * y_size * (nb_params * 4 + nb_variables * 4)) / (1024**3) # calculation done in float32, params in float32
# Memory requirement of output datasets
output_memory_requirement = (x_size * y_size * time_size * nb_outputs * nb_bytes) / (1024**3)
# Total memory requirement

Jérémy AUCLAIR
committed
# TODO : add adjustment factor elsewhere
Jeremy Auclair
committed
total_memory_requirement = (input_memory_requirement + calculation_memory_requirement + output_memory_requirement) * 1.05 # 5% adjustment factor
return total_memory_requirement
Jeremy Auclair
committed

Jérémy AUCLAIR
committed
def calculate_time_slices_to_load(x_size: int, y_size: int, time_size: int, nb_inputs: int, nb_outputs: int, nb_variables: int, nb_params: int, nb_bytes: int, available_ram: int) -> tuple[int, int, bool]:
Calculate how many time slices to load in memory (for input and output data)
based on available ram and calculation requirements.
Jeremy Auclair
committed
1. x_size: ``int``
x size of dataset (pixels)
2. y_size: ``int``
y size of dataset (pixels)
3. time_size: ``int``
number of time bands (dates)
4. nb_inputs: ``int``
number of input variables
5. nb_outputs: ``int``
number of ouput variables
6. nb_variables: ``int``
number of calculation variables
7. nb_params: ``int``
number of raster parameters
8. nb_bytes: ``int``
number of bytes of datatype for inputs and outputs
Jeremy Auclair
committed
9. available_ram: ``int``
available ram for computation in GiB
number of times slices to load for
input and output data
remainder of euclidian division for
the number of time slices to load
(last block of data to load)
used to know wheather data has been loaded
when the whole time values fit in memory, not
used otherwise
"""
Jeremy Auclair
committed
# Test if dataset is not too big
Jeremy Auclair
committed
if calculate_memory_requirement(x_size, y_size, 1, nb_inputs, nb_outputs, nb_variables, nb_params, nb_bytes) > available_ram:
Jeremy Auclair
committed
import sys
print('Calculation area is too large for available memory.')
sys.exit()
Jeremy Auclair
committed
# Boolean to get out of while loop
not_finished = True
Jeremy Auclair
committed
# Divisor starts at one
Jeremy Auclair
committed
divisor = 1
Jeremy Auclair
committed
Jeremy Auclair
committed
# Increase divisor by one at every loop
while not_finished and divisor < time_size / 2:
Jeremy Auclair
committed
memory_requirement = calculate_memory_requirement(x_size, y_size, time_size // divisor, nb_inputs, nb_outputs, nb_variables, nb_params, nb_bytes)
Jeremy Auclair
committed
# Test if memory requirement divided by divisor fits in the available ram
Jeremy Auclair
committed
if memory_requirement < available_ram:
Jeremy Auclair
committed
# If all the inputs and outputs can be loaded
if divisor == 1:
Jeremy Auclair
committed
time_slice = time_size # whole dataset is loaded
remainder_to_load = None # no remainder to load
already_loaded = False # this boolean is set to true once the whole inputs have been loaded
not_finished = False # boolean set to false to get out of while loop
return time_slice, remainder_to_load, already_loaded
Jeremy Auclair
committed
# Otherwise return the number of time bands that can be loaded
else:
Jeremy Auclair
committed
time_slice = time_size // divisor
Jeremy Auclair
committed
remainder_to_load = time_size % time_slice # remainder, used to load the last block of inputs
already_loaded = None # this boolean is set to None if the whole dataset can't be loaded
not_finished = False # boolean set to false to get out of while loop
return time_slice, remainder_to_load, already_loaded
Jeremy Auclair
committed
Jeremy Auclair
committed
divisor += 1 # increment divisor
# If dataset is to big, load only one time slice per loop
time_slice = 1
remainder_to_load = 1 # in order to correctly save last date
Jeremy Auclair
committed
already_loaded = None # this boolean is set to None if the whole dataset can't be loaded
return time_slice, remainder_to_load, already_loaded

Jérémy AUCLAIR
committed
def get_empty_arrays(x_size: int, y_size: int, time_size: int, nb_arrays: int, array: bool = False) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
Short function to make the `run_samir()` function easier to read.
Generates a varying number (`nb_arrays`) of empty numpy arrays of shape
`(x_size, y_size, time_size)` in list or tuple mode. Used to
store variable values before writing in in the output file.
x size of dataset
y size of dataset
number of time bands
number of arrays to generate

Jérémy AUCLAIR
committed
5. array: ``bool`` ``default = False``
weather to return a tuple or a list
Jeremy Auclair
committed
output: ``Tuple[np.ndarray * nb_arrays]`` or ``List[np.ndarray * nb_arrays]``
output empty arrays
"""
# Return empty arrays into a list

Jérémy AUCLAIR
committed
if array:
# return [np.empty((time_size, y_size, x_size), dtype = np.int16) for k in range(nb_arrays)]
return np.empty((nb_arrays, time_size, y_size, x_size), dtype = np.int16, order = 'C')
# Return empty arrays into a tuple
Jeremy Auclair
committed
return tuple([np.empty((time_size, y_size, x_size), dtype = np.uint16) for k in range(nb_arrays)])
# @profile # type: ignore

Jérémy AUCLAIR
committed
def read_inputs(ndvi_cube_path: str, weather_path: str, i: int, time_slice: int, load_all: bool = False) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
Jeremy Auclair
committed
Read input data into numpy arrays based on loop conditions.
Arguments
=========
1. ndvi_cube_path: ``str``
path of input NDVI file
path of input weather file
current loop counter
number of time slices to load
boolean to load the whole datasets
Returns
=======
1. NDVI: ``np.ndarray``
Jeremy Auclair
committed
integer NDVI values into numpy array
Jeremy Auclair
committed
integer Rain values into numpy array
Jeremy Auclair
committed
integer ET0 values into numpy array
"""
# Load whole dataset
if load_all:
with nc.Dataset(ndvi_cube_path, mode='r') as ds:
Jeremy Auclair
committed
# Dimensions of ndvi dataset : (time, y, x)
Jeremy Auclair
committed
ds.variables['NDVI'].set_auto_mask(False)
NDVI = ds.variables['NDVI'][:, :, :]
with nc.Dataset(weather_path, mode='r') as ds:
Jeremy Auclair
committed
# Dimensions of ndvi dataset : (time, y, x)
Jeremy Auclair
committed
ds.variables['Rain'].set_auto_mask(False)
Rain = ds.variables['Rain'][:, :, :]
ds.variables['ET0'].set_auto_mask(False)
ET0 = ds.variables['ET0'][:, :, :]
# Load given number of time slices
else:
with nc.Dataset(ndvi_cube_path, mode='r') as ds: