From 617bb594554196c093fd15581a57fbf117e46aa9 Mon Sep 17 00:00:00 2001 From: Carles Tena Date: Fri, 29 Jul 2022 10:31:10 +0200 Subject: [PATCH 1/2] Grib2 write and some bugfix --- nes/create_nes.py | 17 +- nes/interpolation/horizontal_interpolation.py | 9 +- nes/interpolation/vertical_interpolation.py | 13 +- nes/nc_projections/default_nes.py | 174 +++++++++++++++++- nes/nc_projections/latlon_nes.py | 17 ++ nes/nc_projections/lcc_nes.py | 17 ++ nes/nc_projections/mercator_nes.py | 17 ++ nes/nc_projections/points_nes.py | 19 +- nes/nc_projections/rotated_nes.py | 17 ++ requirements.txt | 4 +- tests/scalability_test_nord3v2.bash | 4 +- tests/test_bash_mn4.cmd | 3 + tests/test_bash_nord3v2.cmd | 1 + 13 files changed, 296 insertions(+), 16 deletions(-) diff --git a/nes/create_nes.py b/nes/create_nes.py index 364101c..6c321a3 100644 --- a/nes/create_nes.py +++ b/nes/create_nes.py @@ -7,7 +7,7 @@ from .nc_projections import * def create_nes(comm=None, info=False, projection=None, parallel_method='Y', balanced=False, - strlen=75, times=None, **kwargs): + strlen=75, times=None, avoid_first_hours=0, avoid_last_hours=0, first_level=0, last_level=None, **kwargs): if comm is None: comm = MPI.COMM_WORLD @@ -55,23 +55,28 @@ def create_nes(comm=None, info=False, projection=None, parallel_method='Y', bala elif parallel_method == 'T': raise NotImplementedError("Parallel method T not implemented yet") nessy = PointsNes(comm=comm, dataset=None, xarray=False, info=info, parallel_method=parallel_method, - avoid_first_hours=0, avoid_last_hours=0, first_level=0, last_level=None, balanced=balanced, + avoid_first_hours=avoid_first_hours, avoid_last_hours=avoid_last_hours, + first_level=first_level, last_level=last_level, balanced=balanced, create_nes=True, strlen=strlen, times=times, **kwargs) elif projection == 'regular': nessy = LatLonNes(comm=comm, dataset=None, xarray=False, info=info, parallel_method=parallel_method, - avoid_first_hours=0, avoid_last_hours=0, first_level=0, last_level=None, balanced=balanced, + avoid_first_hours=avoid_first_hours, avoid_last_hours=avoid_last_hours, + first_level=first_level, last_level=last_level, balanced=balanced, create_nes=True, times=times, **kwargs) elif projection == 'rotated': nessy = RotatedNes(comm=comm, dataset=None, xarray=False, info=info, parallel_method=parallel_method, - avoid_first_hours=0, avoid_last_hours=0, first_level=0, last_level=None, balanced=balanced, + avoid_first_hours=avoid_first_hours, avoid_last_hours=avoid_last_hours, + first_level=first_level, last_level=last_level, balanced=balanced, create_nes=True, times=times, **kwargs) elif projection == 'lcc': nessy = LCCNes(comm=comm, dataset=None, xarray=False, info=info, parallel_method=parallel_method, - avoid_first_hours=0, avoid_last_hours=0, first_level=0, last_level=None, balanced=balanced, + avoid_first_hours=avoid_first_hours, avoid_last_hours=avoid_last_hours, + first_level=first_level, last_level=last_level, balanced=balanced, create_nes=True, times=times, **kwargs) elif projection == 'mercator': nessy = MercatorNes(comm=comm, dataset=None, xarray=False, info=info, parallel_method=parallel_method, - avoid_first_hours=0, avoid_last_hours=0, first_level=0, last_level=None, balanced=balanced, + avoid_first_hours=avoid_first_hours, avoid_last_hours=avoid_last_hours, + first_level=first_level, last_level=last_level, balanced=balanced, create_nes=True, times=times, **kwargs) else: raise NotImplementedError(projection) diff --git a/nes/interpolation/horizontal_interpolation.py b/nes/interpolation/horizontal_interpolation.py index 1bd41f4..e2aff54 100644 --- a/nes/interpolation/horizontal_interpolation.py +++ b/nes/interpolation/horizontal_interpolation.py @@ -84,7 +84,9 @@ def interpolate_horizontal(self, dst_grid, weight_matrix_path=None, kind='Neares final_dst.variables[var_name]['data'] = final_dst.variables[var_name]['data'].T # final_dst.variables[var_name]['dtype'] = final_dst.variables[var_name]['data'].dtype final_dst.variables[var_name]['dimensions'] = ('station', 'time') - print('final shape:', final_dst.variables[var_name]['data'].shape ) + print('final shape:', final_dst.variables[var_name]['data'].shape) + + final_dst.global_attrs = self.global_attrs return final_dst @@ -376,6 +378,11 @@ def create_nn_weight_matrix(self, dst_grid, n_neighbours=4, info=False): weight_matrix = dst_grid.copy() weight_matrix.time = [datetime(year=2000, month=1, day=1, hour=0, second=0, microsecond=0)] weight_matrix._time = [datetime(year=2000, month=1, day=1, hour=0, second=0, microsecond=0)] + weight_matrix.last_level = None + weight_matrix.first_level = 0 + weight_matrix.hours_start = 0 + weight_matrix.hours_end = 0 + weight_matrix.set_communicator(MPI.COMM_SELF) # take the reciprocals of the nearest neighbours distances inverse_dists = np.reciprocal(dists) diff --git a/nes/interpolation/vertical_interpolation.py b/nes/interpolation/vertical_interpolation.py index 3af5aaf..30eff96 100644 --- a/nes/interpolation/vertical_interpolation.py +++ b/nes/interpolation/vertical_interpolation.py @@ -85,14 +85,16 @@ def interpolate_vertical(self, new_levels, new_src_vertical=None, kind='linear', dst_data[t, :, j, i] = np.array( np.interp(new_levels, np.array(self.lev['data'], dtype=np.float64), - np.array(src_data[t, :, j, i], dtype=np.float64)), dtype=np.float32) + np.array(src_data[t, :, j, i], dtype=np.float64)), + dtype=src_data.dtype) else: dst_data[t, :, j, i] = np.array( interp1d(np.array(self.lev['data'], dtype=np.float64), np.array(src_data[t, :, j, i], dtype=np.float64), kind=kind, bounds_error=False, - fill_value=fill_value)(new_levels), dtype=np.float32) + fill_value=fill_value)(new_levels), + dtype=src_data.dtype) except Exception as e: print("time lat lon", t, j, i) print("***********************") @@ -141,7 +143,6 @@ def interpolate_vertical(self, new_levels, new_src_vertical=None, kind='linear', # level_array = np.asarray(new_levels) # level_array = level_array.astype('float32') # level_array = np.where(level_array < level_min, level_min, level_array) - curr_level_values = src_levels[t, :, j, i] try: # check if all values are identical or masked @@ -161,14 +162,16 @@ def interpolate_vertical(self, new_levels, new_src_vertical=None, kind='linear', dst_data[t, :, j, i] = np.array( np.interp(new_levels, np.array(src_levels[t, :, j, i], dtype=np.float64), - np.array(src_data[t, :, j, i], dtype=np.float64)), dtype=np.float32) + np.array(src_data[t, :, j, i], dtype=np.float64)), + dtype=src_data.dtype) else: dst_data[t, :, j, i] = np.array( interp1d(np.array(src_levels[t, :, j, i], dtype=np.float64), np.array(src_data[t, :, j, i], dtype=np.float64), kind=kind, bounds_error=False, - fill_value=fill_value)(new_levels), dtype=np.float32) + fill_value=fill_value)(new_levels), + dtype=src_data.dtype) except Exception as e: print("time lat lon", t, j, i) print("***********************") diff --git a/nes/nc_projections/default_nes.py b/nes/nc_projections/default_nes.py index fa819d0..b357c19 100644 --- a/nes/nc_projections/default_nes.py +++ b/nes/nc_projections/default_nes.py @@ -342,6 +342,12 @@ class Nes(object): return nessy + def get_full_times(self): + return self._time + + def get_full_levels(self): + return self._lev + def clear_communicator(self): """ Erase the communicator and the parallelization indexes. @@ -471,6 +477,49 @@ class Nes(object): return time_interval + def sel_time(self, time, copy=False): + """ + To select only one time step + + Parameters + ---------- + time : datetime.datetime + Time stamp to select + copy : bool + Indicates if you want a copy with the selected time step (True) or to modify te existen one (False) + + Returns + ------- + Nes + Nes object with the data (and metadata) of the selected time step + """ + if copy: + aux_nessy = self.copy(copy_vars=False) + aux_nessy.comm = self.comm + else: + aux_nessy = self + + aux_nessy.hours_start = 0 + aux_nessy.hours_end = 0 + + idx_time = aux_nessy.time.index(time) + + aux_nessy.time = [self.time[idx_time]] + aux_nessy._time = aux_nessy.time + for var_name, var_info in self.variables.items(): + if copy: + aux_nessy.variables[var_name] = {} + for att_name, att_value in var_info.items(): + if att_name == 'data': + if att_value is None: + raise ValueError("{} data not loaded".format(var_name)) + aux_nessy.variables[var_name][att_name] = att_value[[idx_time]] + else: + aux_nessy.variables[var_name][att_name] = att_value + else: + aux_nessy.variables[var_name]['data'] = aux_nessy.variables[var_name]['data'][[idx_time]] + return aux_nessy + # ================================================================================================================== # Statistics # ================================================================================================================== @@ -1231,6 +1280,12 @@ class Nes(object): return None + def to_dtype(self, data_type='float32'): + for var_name, var_info in self.variables.items(): + if var_info['data'] is not None: + self.variables[var_name]['data'] = self.variables[var_name]['data'].astype(data_type) + return None + def concatenate(self, aux_nessy): """ Concatenate different variables into the same nes object @@ -1253,7 +1308,9 @@ class Nes(object): new = True else: new = False - aux_nessy.load() + for var_name, var_info in aux_nessy.variables.items(): + if var_info['data'] is None: + aux_nessy.load(var_name) new_vars_added = [] for new_var_name, new_var_data in aux_nessy.variables.items(): @@ -1686,6 +1743,121 @@ class Nes(object): return None + def __to_grib2(self, path, grib_keys, grib_template_path, info=False): + """ + Private method to write output file with grib2 format. + + Parameters + ---------- + path : str + Path to the output file. + grib_keys : dict + Dictionary with the grib2 keys + grib_template_path : str + Path to the grib2 file to use as template + info : bool + Indicates if you want to print extra information during the process. + """ + from eccodes import codes_grib_new_from_file + from eccodes import codes_keys_iterator_new + from eccodes import codes_keys_iterator_next + from eccodes import codes_keys_iterator_get_name + from eccodes import codes_get_string + from eccodes import codes_keys_iterator_delete + from eccodes import codes_clone + from eccodes import codes_set + from eccodes import codes_set_values + from eccodes import codes_write + from eccodes import codes_release + + fout = open(path, 'wb') + + # read template + fin = open(grib_template_path, 'rb') + + gid = codes_grib_new_from_file(fin) + if gid is None: + sys.exit(1) + + iterid = codes_keys_iterator_new(gid, 'ls') + while codes_keys_iterator_next(iterid): + keyname = codes_keys_iterator_get_name(iterid) + keyval = codes_get_string(gid, keyname) + if info: + print("%s = %s" % (keyname, keyval)) + + codes_keys_iterator_delete(iterid) + for var_name, var_info in self.variables.items(): + for i_time, time in enumerate(self.time): + for i_lev, lev in enumerate(self.lev['data']): + clone_id = codes_clone(gid) + + # Adding grib2 keys to file + for key, value in grib_keys.items(): + if key not in ['typeOfFirstFixedSurface', 'level']: + if info: + print('key:', key, 'val:', value, 'type:', type(value)) + codes_set(clone_id, key, value) + # codes_set(clone_id, key, value) + + # Level dependent keys + if 'typeOfFirstFixedSurface' in grib_keys.keys(): + if float(lev) == 0: + codes_set(clone_id, 'typeOfFirstFixedSurface', 1) + # grib_keys['typeOfFirstFixedSurface'] = 1 + else: + codes_set(clone_id, 'typeOfFirstFixedSurface', 103) + # grib_keys['typeOfFirstFixedSurface'] = 103 + if 'level' in grib_keys.keys(): + codes_set(clone_id, 'level', float(lev)) + # grib_keys['level'] = float(lev) + + # # Adding grib2 keys to file + # for key, value in grib_keys.items(): + # print('key:', key, 'val:', value, 'type:', type(value)) + # codes_set(clone_id, key, value) + + # newval = vardata[step, nlev].round(int(keys['decimalPrecision'])) + newval = var_info['data'][i_time, i_lev] + newval = np.flipud(newval) + # newval = newval.reshape(newval.shape[-1], newval.shape[-2])[::-1, :] + # print(newval.dtype, newval) + codes_set_values(clone_id, newval.ravel()) + # print('write') + codes_write(clone_id, fout) + codes_release(gid) + fout.close() + fin.close() + return None + + def to_grib2(self, path, grib_keys, grib_template_path, info=False): + """ + Write output file with grib2 format. + + Parameters + ---------- + path : str + Path to the output file. + grib_keys : dict + Dictionary with the grib2 keys + grib_template_path : str + Path to the grib2 file to use as template + info : bool + Indicates if you want to print extra information during the process. + """ + + # if serial: + if self.parallel_method in ['X', 'Y'] and self.size > 1: + data = self._gather_data() + if self.master: + new_nc = self.copy(copy_vars=False) + new_nc.set_communicator(MPI.COMM_SELF) + new_nc.variables = data + new_nc.__to_grib2(path, grib_keys, grib_template_path, info=info) + else: + self.__to_grib2(path, grib_keys, grib_template_path, info=info) + return None + def __gather_data_py_object(self): """ Gather all the variable data into the MPI rank 0 to perform a serial write. diff --git a/nes/nc_projections/latlon_nes.py b/nes/nc_projections/latlon_nes.py index 77bb451..b24796e 100644 --- a/nes/nc_projections/latlon_nes.py +++ b/nes/nc_projections/latlon_nes.py @@ -158,3 +158,20 @@ class LatLonNes(Nes): mapping.inverse_flattening = 0 return None + + def to_grib2(self, path, grib_keys, grib_template_path, info=False): + """ + Write output file with grib2 format. + + Parameters + ---------- + path : str + Path to the output file. + grib_keys : dict + Dictionary with the grib2 keys + grib_template_path : str + Path to the grib2 file to use as template + info : bool + Indicates if you want to print extra information during the process. + """ + return super(LatLonNes, self).to_grib2(path, grib_keys, grib_template_path, info=info) diff --git a/nes/nc_projections/lcc_nes.py b/nes/nc_projections/lcc_nes.py index bee48a9..6bf09f4 100644 --- a/nes/nc_projections/lcc_nes.py +++ b/nes/nc_projections/lcc_nes.py @@ -270,3 +270,20 @@ class LCCNes(Nes): mapping.latitude_of_projection_origin = self.projection_data['latitude_of_projection_origin'] return None + + def to_grib2(self, path, grib_keys, grib_template_path, info=False): + """ + Write output file with grib2 format. + + Parameters + ---------- + path : str + Path to the output file. + grib_keys : dict + Dictionary with the grib2 keys + grib_template_path : str + Path to the grib2 file to use as template + info : bool + Indicates if you want to print extra information during the process. + """ + raise NotImplementedError("Grib2 format cannot be write in a Lambert Conformal Conic projection.") diff --git a/nes/nc_projections/mercator_nes.py b/nes/nc_projections/mercator_nes.py index 6e34d50..37dd6fd 100644 --- a/nes/nc_projections/mercator_nes.py +++ b/nes/nc_projections/mercator_nes.py @@ -265,3 +265,20 @@ class MercatorNes(Nes): mapping.longitude_of_projection_origin = self.projection_data['longitude_of_projection_origin'] return None + + def to_grib2(self, path, grib_keys, grib_template_path, info=False): + """ + Write output file with grib2 format. + + Parameters + ---------- + path : str + Path to the output file. + grib_keys : dict + Dictionary with the grib2 keys + grib_template_path : str + Path to the grib2 file to use as template + info : bool + Indicates if you want to print extra information during the process. + """ + raise NotImplementedError("Grib2 format cannot be write in a Mercator projection.") diff --git a/nes/nc_projections/points_nes.py b/nes/nc_projections/points_nes.py index 1f1ef99..82297e3 100644 --- a/nes/nc_projections/points_nes.py +++ b/nes/nc_projections/points_nes.py @@ -581,4 +581,21 @@ class PointsNes(Nes): lons.set_collective(True) lons[:] = self._lon['data'] - return None + return None + + def to_grib2(self, path, grib_keys, grib_template_path, info=False): + """ + Write output file with grib2 format. + + Parameters + ---------- + path : str + Path to the output file. + grib_keys : dict + Dictionary with the grib2 keys + grib_template_path : str + Path to the grib2 file to use as template + info : bool + Indicates if you want to print extra information during the process. + """ + raise NotImplementedError("Grib2 format cannot be write with point data.") diff --git a/nes/nc_projections/rotated_nes.py b/nes/nc_projections/rotated_nes.py index 935aa25..bdeb992 100644 --- a/nes/nc_projections/rotated_nes.py +++ b/nes/nc_projections/rotated_nes.py @@ -316,3 +316,20 @@ class RotatedNes(Nes): mapping.grid_north_pole_longitude = self.projection_data['grid_north_pole_longitude'] return None + + def to_grib2(self, path, grib_keys, grib_template_path, info=False): + """ + Write output file with grib2 format. + + Parameters + ---------- + path : str + Path to the output file. + grib_keys : dict + Dictionary with the grib2 keys + grib_template_path : str + Path to the grib2 file to use as template + info : bool + Indicates if you want to print extra information during the process. + """ + raise NotImplementedError("Grib2 format cannot be write in a Rotated pole projection.") diff --git a/requirements.txt b/requirements.txt index 177f449..74a35d2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,4 +9,6 @@ setuptools~=47.1.0 pytest~=6.2.5 Shapely~=1.8.0 scipy -filelock \ No newline at end of file +filelock +pyproj~=3.2.1 +eccodes-python~=0.9.5 \ No newline at end of file diff --git a/tests/scalability_test_nord3v2.bash b/tests/scalability_test_nord3v2.bash index bfde337..8c65cfa 100644 --- a/tests/scalability_test_nord3v2.bash +++ b/tests/scalability_test_nord3v2.bash @@ -12,7 +12,9 @@ module load xarray/0.17.0-foss-2019b-Python-3.7.4 module load pandas/1.2.4-foss-2019b-Python-3.7.4 module load mpi4py/3.0.3-foss-2019b-Python-3.7.4 module load filelock/3.7.1-foss-2019b-Python-3.7.4 -module load pyproj/2.5.0-foss-2019b-Python-3.7.4 +module load pyproj/2.5.0-foss-2019b-Python-3.7.4 +module load eccodes-python/0.9.5-foss-2019b-Python-3.7.4 + for nprocs in 1 2 4 8 do diff --git a/tests/test_bash_mn4.cmd b/tests/test_bash_mn4.cmd index 27c3bbd..3c92fcd 100644 --- a/tests/test_bash_mn4.cmd +++ b/tests/test_bash_mn4.cmd @@ -21,6 +21,9 @@ module load cfunits/1.8-foss-2019b-Python-3.7.4 module load xarray/0.17.0-foss-2019b-Python-3.7.4 module load OpenMPI/4.0.5-GCC-8.3.0-mn4 module load filelock/3.7.1-foss-2019b-Python-3.7.4 +module load eccodes-python/0.9.5-foss-2019b-Python-3.7.4 +module load pyproj/2.5.0-foss-2019b-Python-3.7.4 + export PYTHONPATH=/gpfs/scratch/bsc32/bsc32538/NES_tests/NES:${PYTHONPATH} cd /gpfs/scratch/bsc32/bsc32538/NES_tests/NES/tests diff --git a/tests/test_bash_nord3v2.cmd b/tests/test_bash_nord3v2.cmd index 6b97bb4..26ac901 100644 --- a/tests/test_bash_nord3v2.cmd +++ b/tests/test_bash_nord3v2.cmd @@ -21,6 +21,7 @@ module load xarray/0.17.0-foss-2019b-Python-3.7.4 module load pandas/1.2.4-foss-2019b-Python-3.7.4 module load mpi4py/3.0.3-foss-2019b-Python-3.7.4 module load filelock/3.7.1-foss-2019b-Python-3.7.4 +module load eccodes-python/0.9.5-foss-2019b-Python-3.7.4 module load pyproj/2.5.0-foss-2019b-Python-3.7.4 export PYTHONPATH=/esarchive/scratch/avilanova/software/NES:${PYTHONPATH} -- GitLab From 65ddde3ec609cc127472bc17f26e497cb6b09588 Mon Sep 17 00:00:00 2001 From: Carles Tena Date: Fri, 29 Jul 2022 12:14:29 +0200 Subject: [PATCH 2/2] Grib2 write --- Jupyter_notebooks/Jupyter_bash_nord3v2.cmd | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/Jupyter_notebooks/Jupyter_bash_nord3v2.cmd b/Jupyter_notebooks/Jupyter_bash_nord3v2.cmd index 450b725..da9fefb 100644 --- a/Jupyter_notebooks/Jupyter_bash_nord3v2.cmd +++ b/Jupyter_notebooks/Jupyter_bash_nord3v2.cmd @@ -30,10 +30,11 @@ module load xarray/0.19.0-foss-2019b-Python-3.7.4 module load cftime/1.0.1-foss-2019b-Python-3.7.4 module load cfunits/1.8-foss-2019b-Python-3.7.4 module load filelock/3.7.1-foss-2019b-Python-3.7.4 -module load pyproj/2.5.0-foss-2019b-Python-3.7.4 +module load pyproj/2.5.0-foss-2019b-Python-3.7.4 +module load eccodes-python/0.9.5-foss-2019b-Python-3.7.4 -# export PYTHONPATH=/gpfs/scratch/bsc32/bsc32538/NES_tests/NES:${PYTHONPATH} -export PYTHONPATH=/esarchive/scratch/avilanova/software/NES:${PYTHONPATH} +export PYTHONPATH=/gpfs/scratch/bsc32/bsc32538/NES_tests/NES:${PYTHONPATH} +# export PYTHONPATH=/esarchive/scratch/avilanova/software/NES:${PYTHONPATH} # DON'T USE ADDRESS BELOW. # DO USE TOKEN BELOW -- GitLab