From dea394770198f065fd0e8815d0b58ab3f8444e77 Mon Sep 17 00:00:00 2001 From: Carles Tena Date: Tue, 27 Sep 2022 11:30:47 +0200 Subject: [PATCH 1/2] Improved vertical interpolation function --- nes/interpolation/vertical_interpolation.py | 60 ++++++++------------- 1 file changed, 22 insertions(+), 38 deletions(-) diff --git a/nes/interpolation/vertical_interpolation.py b/nes/interpolation/vertical_interpolation.py index 774edb7..cd31fd5 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("+++++++++++++++++++++++") -- GitLab From 735660b3fc9fac4ac8cc82e86439a1b1b7371e52 Mon Sep 17 00:00:00 2001 From: Carles Tena Date: Tue, 27 Sep 2022 11:32:46 +0200 Subject: [PATCH 2/2] - Improved memory usage while clean - Added sum_axis function - Modified grib2 function --- nes/nc_projections/default_nes.py | 100 ++++++++++++++----- nes/nc_projections/latlon_nes.py | 4 +- nes/nc_projections/lcc_nes.py | 2 +- nes/nc_projections/mercator_nes.py | 2 +- nes/nc_projections/points_nes.py | 12 ++- nes/nc_projections/points_nes_ghost.py | 10 ++ nes/nc_projections/points_nes_providentia.py | 10 ++ nes/nc_projections/rotated_nes.py | 2 +- tests/scalability_test_nord3v2.bash | 6 +- 9 files changed, 112 insertions(+), 36 deletions(-) diff --git a/nes/nc_projections/default_nes.py b/nes/nc_projections/default_nes.py index 604acb1..1e824e0 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 e18d5f3..5fc8cb1 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 ef41fb3..9b87056 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 3e10ef6..7da8a80 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 f42934c..2595a0c 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 4c66a01..6cabd74 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 011b1dc..5bb6a5f 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 df82944..ec229b2 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 c725675..b9c0540 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 -- GitLab