From 83c5c8f2580056ae83da15e27f16d6b3729091ae Mon Sep 17 00:00:00 2001 From: Ryan Grout Date: Fri, 20 Nov 2020 10:50:35 -0600 Subject: [PATCH 1/4] Don't construct lat/long arrays before scattering. --- core/regrid.py | 36 ++++++++++++++++++++++-------------- 1 file changed, 22 insertions(+), 14 deletions(-) diff --git a/core/regrid.py b/core/regrid.py index 659c5f4..31b6057 100755 --- a/core/regrid.py +++ b/core/regrid.py @@ -2368,19 +2368,19 @@ def calculate_weights(id_tmp, force_count, input_forcings, config_options, mpi_c lon_tmp = None if mpi_config.rank == 0: # Process lat/lon values from the GFS grid. - if len(id_tmp.variables['latitude'].shape) == 3: + if id_tmp.variables['latitude'].ndim == 3: # We have 2D grids already in place. lat_tmp = id_tmp.variables['latitude'][0, :, :] lon_tmp = id_tmp.variables['longitude'][0, :, :] - elif len(id_tmp.variables['longitude'].shape) == 2: + elif id_tmp.variables['longitude'].ndim == 2: # We have 2D grids already in place. lat_tmp = id_tmp.variables['latitude'][:, :] lon_tmp = id_tmp.variables['longitude'][:, :] - elif len(id_tmp.variables['latitude'].shape) == 1: + elif id_tmp.variables['latitude'].ndim == 1: # We have 1D lat/lons we need to translate into # 2D grids. - lat_tmp = np.repeat(id_tmp.variables['latitude'][:][:, np.newaxis], input_forcings.nx_global, axis=1) - lon_tmp = np.tile(id_tmp.variables['longitude'][:], (input_forcings.ny_global, 1)) + lat_tmp = id_tmp.variables['latitude'][:] + lon_tmp = id_tmp.variables['longitude'][:] err_handler.check_program_status(config_options, mpi_config) # Scatter global GFS latitude grid to processors.. @@ -2412,8 +2412,12 @@ def calculate_weights(id_tmp, force_count, input_forcings, config_options, mpi_c err_handler.log_critical(config_options, mpi_config) err_handler.check_program_status(config_options, mpi_config) - input_forcings.esmf_lats[:, :] = var_sub_lat_tmp - input_forcings.esmf_lons[:, :] = var_sub_lon_tmp + if var_sub_lat_tmp.ndim == 1 and var_sub_lon_tmp.ndim == 1: + input_forcings.esmf_lats[:] = np.atleast_2d(var_sub_lat_tmp).T + input_forcings.esmf_lons[:] = np.atleast_2d(var_sub_lon_tmp) + else: + input_forcings.esmf_lats[:, :] = var_sub_lat_tmp + input_forcings.esmf_lons[:, :] = var_sub_lon_tmp del var_sub_lat_tmp del var_sub_lon_tmp del lat_tmp @@ -2596,19 +2600,19 @@ def calculate_supp_pcp_weights(supplemental_precip, id_tmp, tmp_file, config_opt lat_tmp = lon_tmp = None if mpi_config.rank == 0: # Process lat/lon values from the GFS grid. - if len(id_tmp.variables['latitude'].shape) == 3: + if id_tmp.variables['latitude'].ndim == 3: # We have 2D grids already in place. lat_tmp = id_tmp.variables['latitude'][0, :, :] lon_tmp = id_tmp.variables['longitude'][0, :, :] - elif len(id_tmp.variables['longitude'].shape) == 2: + elif id_tmp.variables['longitude'].ndim == 2: # We have 2D grids already in place. lat_tmp = id_tmp.variables['latitude'][:, :] lon_tmp = id_tmp.variables['longitude'][:, :] - elif len(id_tmp.variables['latitude'].shape) == 1: + elif id_tmp.variables['latitude'].ndim == 1: # We have 1D lat/lons we need to translate into # 2D grids. - lat_tmp = np.repeat(id_tmp.variables['latitude'][:][:, np.newaxis], supplemental_precip.nx_global, axis=1) - lon_tmp = np.tile(id_tmp.variables['longitude'][:], (supplemental_precip.ny_global, 1)) + lat_tmp = id_tmp.variables['latitude'][:] + lon_tmp = id_tmp.variables['longitude'][:] # mpi_config.comm.barrier() # Scatter global GFS latitude grid to processors.. @@ -2642,8 +2646,12 @@ def calculate_supp_pcp_weights(supplemental_precip, id_tmp, tmp_file, config_opt err_handler.err_out(config_options) # mpi_config.comm.barrier() - supplemental_precip.esmf_lats[:, :] = var_sub_lat_tmp - supplemental_precip.esmf_lons[:, :] = var_sub_lon_tmp + if var_sub_lat_tmp.ndim == 1 and var_sub_lon_tmp.ndim == 1: + supplemental_precip.esmf_lats[:] = np.atleast_2d(var_sub_lat_tmp).T + supplemental_precip.esmf_lons[:] = np.atleast_2d(var_sub_lon_tmp) + else: + supplemental_precip.esmf_lats[:, :] = var_sub_lat_tmp + supplemental_precip.esmf_lons[:, :] = var_sub_lon_tmp del var_sub_lat_tmp del var_sub_lon_tmp del lat_tmp From 239c5e1aa9d94908fe75c46ab9f502ab44da2fe8 Mon Sep 17 00:00:00 2001 From: Ryan Grout Date: Fri, 20 Nov 2020 11:59:07 -0600 Subject: [PATCH 2/4] Improve dimension checking/handling. --- core/regrid.py | 55 ++++++++++++++++++++++++++++++-------------------- 1 file changed, 33 insertions(+), 22 deletions(-) diff --git a/core/regrid.py b/core/regrid.py index 31b6057..d9a09cd 100755 --- a/core/regrid.py +++ b/core/regrid.py @@ -2368,19 +2368,24 @@ def calculate_weights(id_tmp, force_count, input_forcings, config_options, mpi_c lon_tmp = None if mpi_config.rank == 0: # Process lat/lon values from the GFS grid. - if id_tmp.variables['latitude'].ndim == 3: + _lat = id_tmp.variables['latitude'].ndim + _lon = id_tmp.variables['longitude'].ndim + if _lat != _lon: + config_options.errMsg = f"Latitude and Longitude dimensions must match ({_lat} != {_lon})" + err_handler.log_error(config_options, mpi_config) + err_handler.check_program_status(config_options, mpi_config) + + if _lat == 3: # We have 2D grids already in place. - lat_tmp = id_tmp.variables['latitude'][0, :, :] - lon_tmp = id_tmp.variables['longitude'][0, :, :] - elif id_tmp.variables['longitude'].ndim == 2: + lat_tmp = id_tmp.variables['latitude'][0] + lon_tmp = id_tmp.variables['longitude'][0] + elif 0 < _lat < 3: # We have 2D grids already in place. - lat_tmp = id_tmp.variables['latitude'][:, :] - lon_tmp = id_tmp.variables['longitude'][:, :] - elif id_tmp.variables['latitude'].ndim == 1: - # We have 1D lat/lons we need to translate into - # 2D grids. lat_tmp = id_tmp.variables['latitude'][:] lon_tmp = id_tmp.variables['longitude'][:] + else: + config_options.errMsg = f"Latitude/Longitude arrays must be 1D, 2D, or 3D (got: {_lat})" + err_handler.log_error(config_options, mpi_config) err_handler.check_program_status(config_options, mpi_config) # Scatter global GFS latitude grid to processors.. @@ -2416,8 +2421,8 @@ def calculate_weights(id_tmp, force_count, input_forcings, config_options, mpi_c input_forcings.esmf_lats[:] = np.atleast_2d(var_sub_lat_tmp).T input_forcings.esmf_lons[:] = np.atleast_2d(var_sub_lon_tmp) else: - input_forcings.esmf_lats[:, :] = var_sub_lat_tmp - input_forcings.esmf_lons[:, :] = var_sub_lon_tmp + input_forcings.esmf_lats[:] = var_sub_lat_tmp + input_forcings.esmf_lons[:] = var_sub_lon_tmp del var_sub_lat_tmp del var_sub_lon_tmp del lat_tmp @@ -2600,19 +2605,25 @@ def calculate_supp_pcp_weights(supplemental_precip, id_tmp, tmp_file, config_opt lat_tmp = lon_tmp = None if mpi_config.rank == 0: # Process lat/lon values from the GFS grid. - if id_tmp.variables['latitude'].ndim == 3: + _lat = id_tmp.variables['latitude'].ndim + _lon = id_tmp.variables['longitude'].ndim + if _lat != _lon: + config_options.errMsg = f"Latitude and Longitude dimensions must match ({_lat} != {_lon})" + err_handler.log_error(config_options, mpi_config) + err_handler.check_program_status(config_options, mpi_config) + + if _lat == 3: # We have 2D grids already in place. - lat_tmp = id_tmp.variables['latitude'][0, :, :] - lon_tmp = id_tmp.variables['longitude'][0, :, :] - elif id_tmp.variables['longitude'].ndim == 2: + lat_tmp = id_tmp.variables['latitude'][0] + lon_tmp = id_tmp.variables['longitude'][0] + elif 0 < _lat < 3: # We have 2D grids already in place. - lat_tmp = id_tmp.variables['latitude'][:, :] - lon_tmp = id_tmp.variables['longitude'][:, :] - elif id_tmp.variables['latitude'].ndim == 1: - # We have 1D lat/lons we need to translate into - # 2D grids. lat_tmp = id_tmp.variables['latitude'][:] lon_tmp = id_tmp.variables['longitude'][:] + else: + config_options.errMsg = f"Latitude/Longitude arrays must be 1D, 2D, or 3D (got: {_lat})" + err_handler.log_error(config_options, mpi_config) + err_handler.check_program_status(config_options, mpi_config) # mpi_config.comm.barrier() # Scatter global GFS latitude grid to processors.. @@ -2650,8 +2661,8 @@ def calculate_supp_pcp_weights(supplemental_precip, id_tmp, tmp_file, config_opt supplemental_precip.esmf_lats[:] = np.atleast_2d(var_sub_lat_tmp).T supplemental_precip.esmf_lons[:] = np.atleast_2d(var_sub_lon_tmp) else: - supplemental_precip.esmf_lats[:, :] = var_sub_lat_tmp - supplemental_precip.esmf_lons[:, :] = var_sub_lon_tmp + supplemental_precip.esmf_lats[:] = var_sub_lat_tmp + supplemental_precip.esmf_lons[:] = var_sub_lon_tmp del var_sub_lat_tmp del var_sub_lon_tmp del lat_tmp From aa4715b60058c1fa96ba0961f0b99bd35e60d083 Mon Sep 17 00:00:00 2001 From: Ryan Grout Date: Fri, 20 Nov 2020 12:02:05 -0600 Subject: [PATCH 3/4] Create gmask more efficiently. --- core/regrid.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/core/regrid.py b/core/regrid.py index d9a09cd..d101f5a 100755 --- a/core/regrid.py +++ b/core/regrid.py @@ -2353,11 +2353,9 @@ def calculate_weights(id_tmp, force_count, input_forcings, config_options, mpi_c border) err_handler.log_msg(config_options, mpi_config) - gmask = np.ones([input_forcings.ny_global, input_forcings.nx_global]) - gmask[:+border, :] = 0. # top edge - gmask[-border:, :] = 0. # bottom edge - gmask[:, :+border] = 0. # left edge - gmask[:, -border:] = 0. # right edge + gmask = np.zeros([input_forcings.ny_global, input_forcings.nx_global]) + # Set center block to 1 + gmask[border:-border, border:-border] = 1 mask[:, :] = mpi_config.scatter_array(input_forcings, gmask, config_options) err_handler.check_program_status(config_options, mpi_config) From cce26ba2a53e81eeb410060a77f1a02f2afc6719 Mon Sep 17 00:00:00 2001 From: Ryan Grout Date: Fri, 20 Nov 2020 13:16:19 -0600 Subject: [PATCH 4/4] Minor typo fixes. --- core/regrid.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/core/regrid.py b/core/regrid.py index d101f5a..10a6479 100755 --- a/core/regrid.py +++ b/core/regrid.py @@ -2369,16 +2369,15 @@ def calculate_weights(id_tmp, force_count, input_forcings, config_options, mpi_c _lat = id_tmp.variables['latitude'].ndim _lon = id_tmp.variables['longitude'].ndim if _lat != _lon: - config_options.errMsg = f"Latitude and Longitude dimensions must match ({_lat} != {_lon})" + config_options.errMsg = f"Latitude and Longitude number of dimensions must match ({_lat} != {_lon})" err_handler.log_error(config_options, mpi_config) err_handler.check_program_status(config_options, mpi_config) if _lat == 3: - # We have 2D grids already in place. + # Extract 2D slice from 3D array lat_tmp = id_tmp.variables['latitude'][0] lon_tmp = id_tmp.variables['longitude'][0] elif 0 < _lat < 3: - # We have 2D grids already in place. lat_tmp = id_tmp.variables['latitude'][:] lon_tmp = id_tmp.variables['longitude'][:] else: @@ -2606,16 +2605,15 @@ def calculate_supp_pcp_weights(supplemental_precip, id_tmp, tmp_file, config_opt _lat = id_tmp.variables['latitude'].ndim _lon = id_tmp.variables['longitude'].ndim if _lat != _lon: - config_options.errMsg = f"Latitude and Longitude dimensions must match ({_lat} != {_lon})" + config_options.errMsg = f"Latitude and Longitude number of dimensions must match ({_lat} != {_lon})" err_handler.log_error(config_options, mpi_config) err_handler.check_program_status(config_options, mpi_config) if _lat == 3: - # We have 2D grids already in place. + # Extract 2D slice from 3D array lat_tmp = id_tmp.variables['latitude'][0] lon_tmp = id_tmp.variables['longitude'][0] elif 0 < _lat < 3: - # We have 2D grids already in place. lat_tmp = id_tmp.variables['latitude'][:] lon_tmp = id_tmp.variables['longitude'][:] else: