diff --git a/nes/interpolation/vertical_interpolation.py b/nes/interpolation/vertical_interpolation.py index 774edb76476b832ec81dbdf7fe424a5a9e7e1d64..cd31fd520cea1b479c780c42cc3d3c1e0ed62d09 100644 --- a/nes/interpolation/vertical_interpolation.py +++ b/nes/interpolation/vertical_interpolation.py @@ -62,9 +62,9 @@ def interpolate_vertical(self, new_levels, new_src_vertical=None, kind='linear', if info and self.master: print("\t{var} vertical interpolation".format(var=var_name)) sys.stdout.flush() - src_data = np.flip(var_info['data'], axis=1) - nt, nz, ny, nx = src_data.shape - dst_data = np.ma.masked_all((nt, nz_new, ny, nx)) + # src_data = np.flip(var_info['data'], axis=1) + nt, nz, ny, nx = var_info['data'].shape + dst_data = np.ma.masked_all((nt, nz_new, ny, nx), dtype=var_info['data'].dtype) for t in range(nt): if info and self.master: @@ -73,35 +73,30 @@ def interpolate_vertical(self, new_levels, new_src_vertical=None, kind='linear', for j in range(ny): for i in range(nx): if extrapolate is None: - fill_value = (np.float64(src_data[t, 0, j, i]), np.float64(src_data[t, -1, j, i])) + fill_value = (np.float64(var_info['data'][t, 0, j, i]), + np.float64(var_info['data'][t, -1, j, i])) else: fill_value = extrapolate try: - # f = 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) - # dst_data[t, :, j, i] = np.array(f(new_levels), dtype=np.float32) if 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=src_data.dtype) + np.array(var_info['data'][t, :, j, i], dtype=np.float64)), + dtype=var_info['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), + np.array(var_info['data'][t, :, j, i], dtype=np.float64), kind=kind, bounds_error=False, fill_value=fill_value)(new_levels), - dtype=src_data.dtype) + dtype=var_info['data'].dtype) except Exception as e: print("time lat lon", t, j, i) print("***********************") print("LEVELS", self.lev['data']) - print("VAR", src_data[t, :, j, i]) + print("VAR", var_info['data'][t, :, j, i]) print("+++++++++++++++++++++++") raise Exception(str(e)) self.variables[var_name]['data'] = copy(dst_data) @@ -112,9 +107,10 @@ def interpolate_vertical(self, new_levels, new_src_vertical=None, kind='linear', else: src_levels = self.variables[self.vertical_var_name]['data'] if self.vertical_var_name == 'layer_thickness': - src_levels = np.cumsum(np.flip(src_levels, axis=1), axis=1) + src_levels = np.flip(np.cumsum(np.flip(src_levels, axis=1), axis=1)) else: - src_levels = np.flip(src_levels, axis=1) + # src_levels = np.flip(src_levels, axis=1) + pass for var_name, var_info in self.variables.items(): if var_info['data'] is None: self.load(var_name) @@ -122,10 +118,9 @@ def interpolate_vertical(self, new_levels, new_src_vertical=None, kind='linear', if info and self.master: print("\t{var} vertical interpolation".format(var=var_name)) sys.stdout.flush() - src_data = np.flip(var_info['data'], axis=1) - nt, nz, ny, nx = src_data.shape + # src_data = np.flip(var_info['data'], axis=1) + nt, nz, ny, nx = var_info['data'].shape dst_data = np.empty((nt, nz_new, ny, nx), dtype=np.float32) - for t in range(nt): # if info and self.rank == self.size - 1: if self.info and self.master: @@ -133,18 +128,6 @@ def interpolate_vertical(self, new_levels, new_src_vertical=None, kind='linear', sys.stdout.flush() for j in range(ny): for i in range(nx): - # level_array = None - # nl = src_levels[t, 0, j, i] - src_levels[t, 1, j, i] - # if nl > 0: - # level_max = np.max(src_levels[t, :, j, i]) - # level_array = np.asarray(new_levels) - # level_array = level_array.astype('float32') - # level_array = np.where(level_array > level_max, level_max, level_array) - # if nl < 0: - # level_min = np.min(src_levels[t, :, j, i]) - # 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 @@ -156,7 +139,8 @@ def interpolate_vertical(self, new_levels, new_src_vertical=None, kind='linear', else: kind = kind # 'cubic' if extrapolate is None: - fill_value = (np.float64(src_data[t, 0, j, i]), np.float64(src_data[t, -1, j, i])) + fill_value = (np.float64(var_info['data'][t, 0, j, i]), + np.float64(var_info['data'][t, -1, j, i])) else: fill_value = extrapolate @@ -164,21 +148,21 @@ 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=src_data.dtype) + np.array(var_info['data'][t, :, j, i], dtype=np.float64)), + dtype=var_info['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), + np.array(var_info['data'][t, :, j, i], dtype=np.float64), kind=kind, bounds_error=False, fill_value=fill_value)(new_levels), - dtype=src_data.dtype) + dtype=var_info['data'].dtype) except Exception as e: print("time lat lon", t, j, i) print("***********************") print("LEVELS", np.array(src_levels[t, :, j, i], dtype=np.float64)) - print("DATA", np.array(src_data[t, :, j, i], dtype=np.float64)) + print("DATA", np.array(var_info['data'][t, :, j, i], dtype=np.float64)) print("METHOD", kind) print("FILL_VALUE", fill_value) print("+++++++++++++++++++++++") diff --git a/nes/nc_projections/default_nes.py b/nes/nc_projections/default_nes.py index 604acb10dae9eb5369acd5549753be7214d969db..1e824e0c900c00d4494dc16b1a541f492c17e1a9 100644 --- a/nes/nc_projections/default_nes.py +++ b/nes/nc_projections/default_nes.py @@ -1,6 +1,7 @@ #!/usr/bin/env python import sys +import gc import os import warnings import numpy as np @@ -287,8 +288,7 @@ class Nes(object): """ self.close() - for var_name, var_info in self.variables.items(): - del var_info['data'] + self.free_vars(list(self.variables.keys())) del self.variables try: del self.time @@ -309,6 +309,7 @@ class Nes(object): pass del self + gc.collect() return None @@ -531,11 +532,13 @@ class Nes(object): self.dataset = self.dataset.drop_vars(var_list) self.variables = self._get_lazy_variables() else: - for var_name in var_list: - if self.variables is not None: + if self.variables is not None: + for var_name in var_list: if var_name in self.variables: + if 'data' in self.variables[var_name].keys(): + del self.variables[var_name]['data'] del self.variables[var_name] - + gc.collect() return None def keep_vars(self, var_list): @@ -893,6 +896,42 @@ class Nes(object): return None + @staticmethod + def _get_axis_index_(axis): + if axis == 'T': + value = 0 + elif axis == 'Z': + value = 1 + elif axis == 'Y': + value = 2 + elif axis == 'X': + value = 3 + else: + raise ValueError("Unknown axis: {0}".format(axis)) + return value + + def sum_axis(self, axis='Z'): + if self.parallel_method == axis: + raise NotImplementedError("It is not possible to sum the axis with it is parallelized '{0}'".format( + self.parallel_method)) + + 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'].sum( + axis=self._get_axis_index_(axis), keepdims=True) + if axis == 'T': + self.variables[var_name]['cell_methods'] = "time: sum (interval: {0}hr)".format( + (self.time[-1] - self.time[0]).total_seconds() // 3600) + + if axis == 'T': + self.set_time_bnds([self.time[0], self.time[-1]]) + self.time = [self.time[0]] + self._time = [self._time[0]] + if axis == 'Z': + self.lev['data'] = [self.lev['data'][0]] + self._lev['data'] = [self._lev['data'][0]] + return None + # ================================================================================================================== # Reading # ================================================================================================================== @@ -2144,7 +2183,7 @@ class Nes(object): return None - def __to_grib2(self, path, grib_keys, grib_template_path, info=False): + def __to_grib2(self, path, grib_keys, grib_template_path, lat_flip=True, info=False): """ Private method to write output file with grib2 format. @@ -2196,36 +2235,43 @@ class Nes(object): # 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) + if value not in ['', 'None', None]: + try: + codes_set(clone_id, key, value) + except Exception as e: + print("Something went wrong while writing the Grib key '{0}': {1}".format(key, value)) + raise e + + # Time dependent keys + if 'dataTime' in grib_keys.keys() and grib_keys['dataTime'] in ['', 'None', None]: + codes_set(clone_id, 'dataTime', int(i_time * 100)) + if 'stepRange' in grib_keys.keys() and grib_keys['stepRange'] in ['', 'None', None]: + n_secs = (time - self._time[0]).total_seconds() + codes_set(clone_id, 'stepRange', int(n_secs // 3600)) + if 'forecastTime' in grib_keys.keys() and grib_keys['forecastTime'] in ['', 'None', None]: + n_secs = (time - self._time[0]).total_seconds() + codes_set(clone_id, 'forecastTime', int(n_secs)) # Level dependent keys - if 'typeOfFirstFixedSurface' in grib_keys.keys(): + if 'typeOfFirstFixedSurface' in grib_keys.keys() and \ + grib_keys['typeOfFirstFixedSurface'] in ['', 'None', None]: 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(): + if 'level' in grib_keys.keys() and grib_keys['level'] in ['', 'None', None]: 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) + if lat_flip: + newval = np.flipud(newval) + + # TODO Check default NaN Value + newval[np.isnan(newval)] = 0. + codes_set_values(clone_id, newval.ravel()) - # print('write') codes_write(clone_id, fout) codes_release(gid) fout.close() @@ -2233,7 +2279,7 @@ class Nes(object): return None - def to_grib2(self, path, grib_keys, grib_template_path, info=False): + def to_grib2(self, path, grib_keys, grib_template_path, lat_flip=True, info=False): """ Write output file with grib2 format. @@ -2256,9 +2302,9 @@ class Nes(object): 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) + new_nc.__to_grib2(path, grib_keys, grib_template_path, lat_flip=lat_flip, info=info) else: - self.__to_grib2(path, grib_keys, grib_template_path, info=info) + self.__to_grib2(path, grib_keys, grib_template_path, lat_flip=lat_flip, info=info) return None diff --git a/nes/nc_projections/latlon_nes.py b/nes/nc_projections/latlon_nes.py index e18d5f3625e7f2cb7ebda42618acaafd1ada3d6d..5fc8cb1cf418cfe0688a857e4d5e044a6ddb2136 100644 --- a/nes/nc_projections/latlon_nes.py +++ b/nes/nc_projections/latlon_nes.py @@ -221,7 +221,7 @@ class LatLonNes(Nes): return None - def to_grib2(self, path, grib_keys, grib_template_path, info=False): + def to_grib2(self, path, grib_keys, grib_template_path, lat_flip=False, info=False): """ Write output file with grib2 format. @@ -237,4 +237,4 @@ class LatLonNes(Nes): 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) + return super(LatLonNes, self).to_grib2(path, grib_keys, grib_template_path, lat_flip=lat_flip, info=info) diff --git a/nes/nc_projections/lcc_nes.py b/nes/nc_projections/lcc_nes.py index ef41fb3c5c2bc9c68c7f38f1ee0a24147ec784c7..9b87056100934f7fa305cc1f3889de327b67bdcd 100644 --- a/nes/nc_projections/lcc_nes.py +++ b/nes/nc_projections/lcc_nes.py @@ -426,7 +426,7 @@ class LCCNes(Nes): return None - def to_grib2(self, path, grib_keys, grib_template_path, info=False): + def to_grib2(self, path, grib_keys, grib_template_path, lat_flip=False, info=False): """ Write output file with grib2 format. diff --git a/nes/nc_projections/mercator_nes.py b/nes/nc_projections/mercator_nes.py index 3e10ef6d1f2e456eab677b4ed07173849020f07d..7da8a80cdcd6afefb0aeeff9fba2aa2619836adf 100644 --- a/nes/nc_projections/mercator_nes.py +++ b/nes/nc_projections/mercator_nes.py @@ -403,7 +403,7 @@ class MercatorNes(Nes): return None - def to_grib2(self, path, grib_keys, grib_template_path, info=False): + def to_grib2(self, path, grib_keys, grib_template_path, lat_flip=False, info=False): """ Write output file with grib2 format. diff --git a/nes/nc_projections/points_nes.py b/nes/nc_projections/points_nes.py index f42934c5862cdb2836521dbdfb5a6ac9495bd524..2595a0c6a07c7959adf57f81e942be498e082eee 100644 --- a/nes/nc_projections/points_nes.py +++ b/nes/nc_projections/points_nes.py @@ -619,7 +619,7 @@ class PointsNes(Nes): return points_nes_providentia - def to_grib2(self, path, grib_keys, grib_template_path, info=False): + def to_grib2(self, path, grib_keys, grib_template_path, lat_flip=False, info=False): """ Write output file with grib2 format. @@ -650,3 +650,13 @@ class PointsNes(Nes): self.shapefile = gdf return gdf + + @staticmethod + def _get_axis_index_(axis): + if axis == 'T': + value = 0 + elif axis == 'X': + value = 1 + else: + raise ValueError("Unknown axis: {0}".format(axis)) + return value diff --git a/nes/nc_projections/points_nes_ghost.py b/nes/nc_projections/points_nes_ghost.py index 4c66a0183e75c3176d599aa88e1d0cf09372c569..6cabd7409778bb15072afa391a5a235951707c17 100644 --- a/nes/nc_projections/points_nes_ghost.py +++ b/nes/nc_projections/points_nes_ghost.py @@ -751,3 +751,13 @@ class PointsNesGHOST(PointsNes): } return metadata_variables[GHOST_version] + + @staticmethod + def _get_axis_index_(axis): + if axis == 'T': + value = 1 + elif axis == 'X': + value = 0 + else: + raise ValueError("Unknown axis: {0}".format(axis)) + return value diff --git a/nes/nc_projections/points_nes_providentia.py b/nes/nc_projections/points_nes_providentia.py index 011b1dc71972777082945b8ef4312169e3114a4f..5bb6a5f98893c512706f4083691324c2accabefb 100644 --- a/nes/nc_projections/points_nes_providentia.py +++ b/nes/nc_projections/points_nes_providentia.py @@ -588,3 +588,13 @@ class PointsNesProvidentia(PointsNes): serial=True, info=info, chunking=chunking) return None + + @staticmethod + def _get_axis_index_(axis): + if axis == 'T': + value = 1 + elif axis == 'X': + value = 0 + else: + raise ValueError("Unknown axis: {0}".format(axis)) + return value diff --git a/nes/nc_projections/rotated_nes.py b/nes/nc_projections/rotated_nes.py index df82944b1841bb971d9f64d4c84d07d7d7bbcf0f..ec229b2dbb9a71d9c6048b1b2446261f17a583b3 100644 --- a/nes/nc_projections/rotated_nes.py +++ b/nes/nc_projections/rotated_nes.py @@ -453,7 +453,7 @@ class RotatedNes(Nes): return None - def to_grib2(self, path, grib_keys, grib_template_path, info=False): + def to_grib2(self, path, grib_keys, grib_template_path, lat_flip=False, info=False): """ Write output file with grib2 format. diff --git a/tests/scalability_test_nord3v2.bash b/tests/scalability_test_nord3v2.bash index c7256757d51f3597aa9da92d265745e507d8a1d3..b9c0540b121455552596d6efd0c8ec8226dd33e7 100644 --- a/tests/scalability_test_nord3v2.bash +++ b/tests/scalability_test_nord3v2.bash @@ -1,8 +1,8 @@ #!/bin/bash #EXPORTPATH="/esarchive/scratch/avilanova/software/NES" -EXPORTPATH="/gpfs/projects/bsc32/models/NES" -SRCPATH="/esarchive/scratch/avilanova/software/NES/tests" +EXPORTPATH="/gpfs/scratch/bsc32/bsc32538/NES_tests/NES" +SRCPATH="/gpfs/scratch/bsc32/bsc32538/NES_tests/NES/tests" EXE="2-nes_tests_by_projection.py" module purge @@ -18,7 +18,7 @@ module load eccodes-python/0.9.5-foss-2019b-Python-3.7.4 module load geopandas/0.8.1-foss-2019b-Python-3.7.4 module load Shapely/1.7.1-foss-2019b-Python-3.7.4 -for nprocs in 1 2 4 8 +for nprocs in 1 2 4 8 16 32 do JOB_ID=`sbatch --ntasks=${nprocs} --exclusive --job-name=nes_${nprocs} --output=./log_nord3v2_NES_${nprocs}_%J.out --error=./log_nord3v2_NES_${nprocs}_%J.err -D . --time=02:00:00 --wrap="export PYTHONPATH=${EXPORTPATH}:${PYTHONPATH}; cd ${SRCPATH}; mpirun --mca mpi_warn_on_fork 0 -np ${nprocs} python ${SRCPATH}/${EXE}"` done \ No newline at end of file