diff --git a/Jupyter_notebooks/Jupyter_bash_nord3v2.cmd b/Jupyter_notebooks/Jupyter_bash_nord3v2.cmd index 450b725909ac10e129173fbfb3b2a4e7d3cff40a..da9fefbcf844a21b316e07af4ec8e5e6ec8adabf 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 diff --git a/nes/create_nes.py b/nes/create_nes.py index 364101c9ee28dfd39cb633f400f5169937bdb1bf..6c321a36a074179925118e9df300f39147db19de 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 1bd41f41f85ccee4313de18c173b283b1633e41a..e2aff54c1589fa1d7e1fedae54ca68dba7bb11cf 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 3af5aaf0c8fbd51f9815f27ec92baeae3cc9b03d..30eff96dd75ecd92ba2a2f6279a668ff04a72b8c 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 fa819d0672b926eb0b5e77dc1584016563b02d56..b357c1994beefc03430437f3d01416384841431e 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 77bb451f171d5e6a459e3d2749e67a637615a1af..b24796e66c648dbd8878bd73f6da69a86e5325a7 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 bee48a94a3fa36f3a2705be34d1f71713faab915..6bf09f4c6de858bfeb89cfc430b876577edf9626 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 6e34d5077f420d82db28f5864b2fff90a353cb00..37dd6fdb166091204a09b2799a4e33508c4805ec 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 b1421c7ec7aa522ba91dbd9c1567d60e2aa22cd3..3dd99455534488d692299f711afba6fcedb68d9e 100644 --- a/nes/nc_projections/points_nes.py +++ b/nes/nc_projections/points_nes.py @@ -615,4 +615,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 935aa252c7ceea693fe729d13308615beb73fa7d..bdeb992ff771233477826c41ac16d8dc0c6112a5 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 177f449ed12b759e767de38a4947fb11f2db1b1a..74a35d2b0177f006ea26eeb08b997d3cc0c4bc1f 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 bfde337f0556b5daa4d55ace8c12bcd0c6867586..8c65cfae453c0fb22901cd79bbbb1dd9842935a3 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 27c3bbde7d8886d8621e3f8b14ce3b2167b4680e..3c92fcd5c718f639d50357d320fb6990ae0cb19e 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 6b97bb43b5dfb0c0753b29c29ac8ba4236d42f6b..26ac9019b6d02e718c49d17b75ecb67b1c056a5b 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}