From e70894d77a45e61e2326ab2be41b4f3a434ecca2 Mon Sep 17 00:00:00 2001 From: Frank Malatino Date: Tue, 10 Oct 2023 10:39:53 -0400 Subject: [PATCH 01/31] Testing changes reflected across branches --- examples/build_scripts/build_gaea_c5.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/build_scripts/build_gaea_c5.sh b/examples/build_scripts/build_gaea_c5.sh index 94ad5611..4c920934 100644 --- a/examples/build_scripts/build_gaea_c5.sh +++ b/examples/build_scripts/build_gaea_c5.sh @@ -1,4 +1,4 @@ -#!/usr/bin/env bash +#!/usr/bin/env bash f # Example bash script to install Pace to run bare-metal on Gaea's c4 cluster From 2a8b97adbb74c506144f9ca5c59086f8401531df Mon Sep 17 00:00:00 2001 From: Frank Malatino Date: Tue, 10 Oct 2023 10:44:56 -0400 Subject: [PATCH 02/31] Undoing changes made in build_gaea_c5.sh --- examples/build_scripts/build_gaea_c5.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/build_scripts/build_gaea_c5.sh b/examples/build_scripts/build_gaea_c5.sh index 4c920934..94ad5611 100644 --- a/examples/build_scripts/build_gaea_c5.sh +++ b/examples/build_scripts/build_gaea_c5.sh @@ -1,4 +1,4 @@ -#!/usr/bin/env bash f +#!/usr/bin/env bash # Example bash script to install Pace to run bare-metal on Gaea's c4 cluster From e2cdd5c1d1020a2a00ed3777098668e756502b4e Mon Sep 17 00:00:00 2001 From: Frank Malatino Date: Tue, 10 Oct 2023 14:14:05 -0400 Subject: [PATCH 03/31] Testing vscode functionality, by adding a change to external_grid branch --- examples/build_scripts/build_gaea_c5.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/build_scripts/build_gaea_c5.sh b/examples/build_scripts/build_gaea_c5.sh index 94ad5611..4c920934 100644 --- a/examples/build_scripts/build_gaea_c5.sh +++ b/examples/build_scripts/build_gaea_c5.sh @@ -1,4 +1,4 @@ -#!/usr/bin/env bash +#!/usr/bin/env bash f # Example bash script to install Pace to run bare-metal on Gaea's c4 cluster From b4d9e5d80de5961811550fe4192bd1ab27cde3df Mon Sep 17 00:00:00 2001 From: Frank Malatino Date: Tue, 10 Oct 2023 14:15:18 -0400 Subject: [PATCH 04/31] Testing vscode functionality, by adding a change to external_grid branch --- examples/build_scripts/build_gaea_c5.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/build_scripts/build_gaea_c5.sh b/examples/build_scripts/build_gaea_c5.sh index 4c920934..94ad5611 100644 --- a/examples/build_scripts/build_gaea_c5.sh +++ b/examples/build_scripts/build_gaea_c5.sh @@ -1,4 +1,4 @@ -#!/usr/bin/env bash f +#!/usr/bin/env bash # Example bash script to install Pace to run bare-metal on Gaea's c4 cluster From 5b0991d3ee6ea5e30d52b34be8baf10a60204364 Mon Sep 17 00:00:00 2001 From: Fr4nk Date: Wed, 18 Oct 2023 22:39:15 -0400 Subject: [PATCH 05/31] Addition of from_generated method and calc_flag to util/pace/util/grid/generation.py --- util/pace/util/grid/generation.py | 41 +++++++++++++++++++++++++++++-- 1 file changed, 39 insertions(+), 2 deletions(-) diff --git a/util/pace/util/grid/generation.py b/util/pace/util/grid/generation.py index cf5e20af..c6af575d 100644 --- a/util/pace/util/grid/generation.py +++ b/util/pace/util/grid/generation.py @@ -225,6 +225,7 @@ def __init__( dx_const: float = 1000.0, dy_const: float = 1000.0, deglat: float = 15.0, + calc_flag: bool = False, ): assert grid_type < 3 self._grid_type = grid_type @@ -366,8 +367,44 @@ def __init__( self._vlon_64 = None self._vlat_64 = None - self._init_dgrid() - self._init_agrid() + if calc_flag is False: + self._init_dgrid() + self._init_agrid() + + @classmethod + def from_generated( + cls, + x, + y, + dx, + dy, + area, + quantity_factory: util.QuantityFactory, + communicator: util.Communicator, + grid_type: int = 0, + dx_const: float = 1000.0, + dy_const: float = 1000.0, + deglat: float = 15.0, + calc_flag: bool = True, + ) -> "MetricTerms": + mt_obj = MetricTerms( + quantity_factory=quantity_factory, + communicator=communicator, + grid_type=grid_type, + dx_const=dx_const, + dy_const=dy_const, + deglat=deglat, + calc_flag=calc_flag, + ) + + mt_obj.grid.data[:, :, 0] = x + mt_obj.grid.data[:, :, 1] = y + mt_obj._dx = dx + mt_obj._dy = dy + + mt_obj._init_agrid() + + return mt_obj @classmethod def from_tile_sizing( From 9001af1887281229249d464f2a9e435bf5aaca45 Mon Sep 17 00:00:00 2001 From: Frank Malatino Date: Thu, 19 Oct 2023 11:03:51 -0400 Subject: [PATCH 06/31] Added get_grid method for external grid data to driver/pace/driver/grid.py --- driver/pace/driver/grid.py | 70 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 70 insertions(+) diff --git a/driver/pace/driver/grid.py b/driver/pace/driver/grid.py index 1cf59d07..35d74b62 100644 --- a/driver/pace/driver/grid.py +++ b/driver/pace/driver/grid.py @@ -3,6 +3,7 @@ from typing import ClassVar, Optional, Tuple import f90nml +import numpy as np import pace.driver import pace.dsl @@ -196,6 +197,75 @@ def get_grid( return damping_coefficients, driver_grid_data, grid_data +@GridInitializerSelector.register("external") +@dataclasses.dataclass +class ExternalGridConfig(GridInitializer): + """ + Configuration for grid initialized from external data. + Input is from tile files generated by FRE-NCtools methods + """ + + restart_path: Optional[str] = None + grid_type: Optional[int] = 0 + dx_const: Optional[float] = 1000.0 + dy_const: Optional[float] = 1000.0 + deglat: Optional[float] = 15.0 + + # TODO: Add in xarray or Fortran read in method + x = np.zeros(9) + y = np.zeros(9) + dx = np.zeros(6) + dy = np.zeros(6) + area = np.zeros(4) + + def get_grid( + self, + quantity_factory: QuantityFactory, + communicator: Communicator, + ) -> Tuple[DampingCoefficients, DriverGridData, GridData]: + backend = quantity_factory.zeros( + dims=[pace.util.X_DIM, pace.util.Y_DIM], units="unknown" + ).gt4py_backend + + pace_log.info("Using external grid data") + + metric_terms = MetricTerms.from_generated( + x=self.x, + y=self.y, + dx=self.dx, + dy=self.dy, + area=self.area, + quantity_factory=quantity_factory, + communicator=communicator, + grid_type=self.grid_type, + dx_const=self.dx_const, + dy_const=self.dy_const, + deglat=self.deglat, + calc_flag=True, + ) + + horizontal_data = HorizontalGridData.new_from_metric_terms(metric_terms) + if self.restart_path is not None: + vertical_data = VerticalGridData.from_restart( + self.restart_path, quantity_factory=quantity_factory + ) + else: + vertical_data = VerticalGridData.new_from_metric_terms(metric_terms) + contravariant_data = ContravariantGridData.new_from_metric_terms(metric_terms) + angle_data = AngleGridData.new_from_metric_terms(metric_terms) + grid_data = GridData( + horizontal_data=horizontal_data, + vertical_data=vertical_data, + contravariant_data=contravariant_data, + angle_data=angle_data, + ) + + damping_coefficients = DampingCoefficients.new_from_metric_terms(metric_terms) + driver_grid_data = DriverGridData.new_from_metric_terms(metric_terms) + + return damping_coefficients, driver_grid_data, grid_data + + def _transform_horizontal_grid( metric_terms: MetricTerms, stretch_factor: float, From 65c0e14a5d6d4d323b7e2114bf3846c57ad03c2d Mon Sep 17 00:00:00 2001 From: Frank Malatino Date: Thu, 19 Oct 2023 15:40:34 -0400 Subject: [PATCH 07/31] Preliminary xarray netcdf read in method added to driver/pace/driver/grid.py --- driver/pace/driver/grid.py | 22 ++++++++++++---------- util/pace/util/grid/generation.py | 5 +++++ 2 files changed, 17 insertions(+), 10 deletions(-) diff --git a/driver/pace/driver/grid.py b/driver/pace/driver/grid.py index 35d74b62..6045678d 100644 --- a/driver/pace/driver/grid.py +++ b/driver/pace/driver/grid.py @@ -3,7 +3,7 @@ from typing import ClassVar, Optional, Tuple import f90nml -import numpy as np +import xarray as xr import pace.driver import pace.dsl @@ -203,6 +203,10 @@ class ExternalGridConfig(GridInitializer): """ Configuration for grid initialized from external data. Input is from tile files generated by FRE-NCtools methods + In its get_grid method it calls the MetricTerms class method + from_generated which generates an object of MetricTerms to be + used to generate the damping_coefficients, driver_grid_data, + and grid_data variables """ restart_path: Optional[str] = None @@ -211,21 +215,19 @@ class ExternalGridConfig(GridInitializer): dy_const: Optional[float] = 1000.0 deglat: Optional[float] = 15.0 - # TODO: Add in xarray or Fortran read in method - x = np.zeros(9) - y = np.zeros(9) - dx = np.zeros(6) - dy = np.zeros(6) - area = np.zeros(4) + # TODO: Probably still need to change this, but a good placeholder + ds = xr.open_dataset("../../../input/file.nc") + x = ds.x.values + y = ds.y.values + dx = ds.dx.values + dy = ds.dy.values + area = ds.area.values def get_grid( self, quantity_factory: QuantityFactory, communicator: Communicator, ) -> Tuple[DampingCoefficients, DriverGridData, GridData]: - backend = quantity_factory.zeros( - dims=[pace.util.X_DIM, pace.util.Y_DIM], units="unknown" - ).gt4py_backend pace_log.info("Using external grid data") diff --git a/util/pace/util/grid/generation.py b/util/pace/util/grid/generation.py index c6af575d..df49e740 100644 --- a/util/pace/util/grid/generation.py +++ b/util/pace/util/grid/generation.py @@ -371,6 +371,11 @@ def __init__( self._init_dgrid() self._init_agrid() + # from_generated class method + # Generates a metric terms object, which will perform + # the same function as the __init__ method for MetricTerms + # initializing the dgrid by input from data contained in an + # externally generated tile file @classmethod def from_generated( cls, From dedd7116bc90dd1cf540af24e7d284db8208425e Mon Sep 17 00:00:00 2001 From: Frank Malatino Date: Thu, 26 Oct 2023 15:10:47 -0400 Subject: [PATCH 08/31] Updating util/pace/util/grid/generation.py from_generated method --- driver/pace/driver/grid.py | 5 +++-- util/pace/util/grid/generation.py | 13 ++++++------- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/driver/pace/driver/grid.py b/driver/pace/driver/grid.py index 6045678d..e520c31b 100644 --- a/driver/pace/driver/grid.py +++ b/driver/pace/driver/grid.py @@ -223,6 +223,8 @@ class ExternalGridConfig(GridInitializer): dy = ds.dy.values area = ds.area.values + # TODO: Area read in? + def get_grid( self, quantity_factory: QuantityFactory, @@ -236,14 +238,13 @@ def get_grid( y=self.y, dx=self.dx, dy=self.dy, - area=self.area, quantity_factory=quantity_factory, communicator=communicator, grid_type=self.grid_type, dx_const=self.dx_const, dy_const=self.dy_const, deglat=self.deglat, - calc_flag=True, + extdgrid=True, ) horizontal_data = HorizontalGridData.new_from_metric_terms(metric_terms) diff --git a/util/pace/util/grid/generation.py b/util/pace/util/grid/generation.py index df49e740..253294fd 100644 --- a/util/pace/util/grid/generation.py +++ b/util/pace/util/grid/generation.py @@ -225,7 +225,7 @@ def __init__( dx_const: float = 1000.0, dy_const: float = 1000.0, deglat: float = 15.0, - calc_flag: bool = False, + extdgrid: bool = False, ): assert grid_type < 3 self._grid_type = grid_type @@ -367,7 +367,7 @@ def __init__( self._vlon_64 = None self._vlat_64 = None - if calc_flag is False: + if extdgrid is False: self._init_dgrid() self._init_agrid() @@ -383,14 +383,13 @@ def from_generated( y, dx, dy, - area, quantity_factory: util.QuantityFactory, communicator: util.Communicator, grid_type: int = 0, dx_const: float = 1000.0, dy_const: float = 1000.0, deglat: float = 15.0, - calc_flag: bool = True, + extdgrid: bool = True, ) -> "MetricTerms": mt_obj = MetricTerms( quantity_factory=quantity_factory, @@ -399,11 +398,11 @@ def from_generated( dx_const=dx_const, dy_const=dy_const, deglat=deglat, - calc_flag=calc_flag, + extdgrid=extdgrid, ) - mt_obj.grid.data[:, :, 0] = x - mt_obj.grid.data[:, :, 1] = y + mt_obj.grid.data[:, :, 0] = x * (PI / 180) + mt_obj.grid.data[:, :, 1] = y * (PI / 180) mt_obj._dx = dx mt_obj._dy = dy From 17e281b2fdfd26d4e837dc3231b2f04ccfd3a0e5 Mon Sep 17 00:00:00 2001 From: Fr4nk Date: Tue, 31 Oct 2023 11:21:38 -0400 Subject: [PATCH 09/31] Addition of external grid data read in methods for initialization of grid. Current method uses xarray to interact with netcdf tile files. Values for longitutde, latitude, number of points in x an y, grid edge distances read in. --- CONTRIBUTORS.md | 1 + driver/pace/driver/grid.py | 31 +++++++++++++++++++------------ util/pace/util/grid/generation.py | 5 +++-- util/pace/util/partitioner.py | 2 +- 4 files changed, 24 insertions(+), 15 deletions(-) diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index 0ca00ff5..496b3cca 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -12,6 +12,7 @@ List format (alphabetical order): Surname, Name. Employer/Affiliation * George, Rhea. Allen Institute for AI. * Harris, Lucas. GFDL. * Kung, Chris. NASA. +* Malatino, Frank. GFDL * McGibbon, Jeremy. Allen Institute for AI. * Niedermayr, Yannick. ETH. * Savarin, Ajda. University of Washington. diff --git a/driver/pace/driver/grid.py b/driver/pace/driver/grid.py index e520c31b..737b5981 100644 --- a/driver/pace/driver/grid.py +++ b/driver/pace/driver/grid.py @@ -215,14 +215,6 @@ class ExternalGridConfig(GridInitializer): dy_const: Optional[float] = 1000.0 deglat: Optional[float] = 15.0 - # TODO: Probably still need to change this, but a good placeholder - ds = xr.open_dataset("../../../input/file.nc") - x = ds.x.values - y = ds.y.values - dx = ds.dx.values - dy = ds.dy.values - area = ds.area.values - # TODO: Area read in? def get_grid( @@ -233,11 +225,26 @@ def get_grid( pace_log.info("Using external grid data") + ds = xr.open_dataset("../../../input/file.nc") + lon = ds.x.values + lat = ds.y.values + dx = ds.dx.values + dy = ds.dy.values + npx = ds.nxp.values.size + npy = ds.nyp.values.size + + subtile_slice = communicator.partitioner.subtile_slice( + rank=communicator.rank, + global_dims=[pace.util.X_DIM, pace.util.Y_DIM], + global_extent=(npx, npy), + overlap=True, + ) + metric_terms = MetricTerms.from_generated( - x=self.x, - y=self.y, - dx=self.dx, - dy=self.dy, + x=lon[subtile_slice], + y=lat[subtile_slice], + dx=dx, + dy=dy, quantity_factory=quantity_factory, communicator=communicator, grid_type=self.grid_type, diff --git a/util/pace/util/grid/generation.py b/util/pace/util/grid/generation.py index d0adee31..548d21ad 100644 --- a/util/pace/util/grid/generation.py +++ b/util/pace/util/grid/generation.py @@ -446,8 +446,9 @@ def from_generated( extdgrid=extdgrid, ) - terms.grid.data[:, :, 0] = x * (PI / 180) - terms.grid.data[:, :, 1] = y * (PI / 180) + rad_conv = PI / 180.0 + terms.grid.data[:, :, 0] = np.dot(rad_conv, x) + terms.grid.data[:, :, 1] = np.dot(rad_conv, y) terms._dx = dx terms._dy = dy diff --git a/util/pace/util/partitioner.py b/util/pace/util/partitioner.py index 0e59ddfa..c2560159 100644 --- a/util/pace/util/partitioner.py +++ b/util/pace/util/partitioner.py @@ -755,7 +755,7 @@ def subtile_index( ) -> Tuple[int, int]: within_tile_rank = rank % ranks_per_tile j = within_tile_rank // layout[1] - i = within_tile_rank % layout[1] + i = within_tile_rank % layout[0] return j, i From 3a627aa5b2485712a3662cc038515c46d71fd6ba Mon Sep 17 00:00:00 2001 From: Frank Malatino Date: Tue, 7 Nov 2023 10:53:30 -0500 Subject: [PATCH 10/31] driver/examples/configs/test_external_C12_1x1.yaml --- .../configs/test_external_C12_1x1.yaml | 105 ++++++++++++++++++ .../configs/test_external_C12_2x2.yaml | 105 ++++++++++++++++++ driver/pace/driver/grid.py | 55 +++++---- util/pace/util/grid/generation.py | 12 +- 4 files changed, 247 insertions(+), 30 deletions(-) create mode 100644 driver/examples/configs/test_external_C12_1x1.yaml create mode 100644 driver/examples/configs/test_external_C12_2x2.yaml diff --git a/driver/examples/configs/test_external_C12_1x1.yaml b/driver/examples/configs/test_external_C12_1x1.yaml new file mode 100644 index 00000000..96c0c247 --- /dev/null +++ b/driver/examples/configs/test_external_C12_1x1.yaml @@ -0,0 +1,105 @@ +stencil_config: + compilation_config: + backend: numpy + rebuild: false + validate_args: true + format_source: false + device_sync: false +grid_config: + type: external + config: + grid_type: 4 + dx_const: 3000.0 + dy_const: 3000.0 + deglat: 10.0 + grid_file_path: /input/C12_grid.tile +initialization: + type: analytic + config: + case: baroclinic +performance_config: + collect_performance: true + experiment_name: c12_baroclinic +nx_tile: 12 +nz: 79 +dt_atmos: 225 +minutes: 15 +layout: + - 1 + - 1 +diagnostics_config: + path: output + output_format: netcdf + names: + - u + - v + - ua + - va + - pt + - delp + - qvapor + - qliquid + - qice + - qrain + - qsnow + - qgraupel + z_select: + - level: 65 + names: + - pt +dycore_config: + a_imp: 1.0 + beta: 0. + consv_te: 0. + d2_bg: 0. + d2_bg_k1: 0.2 + d2_bg_k2: 0.1 + d4_bg: 0.15 + d_con: 1.0 + d_ext: 0.0 + dddmp: 0.5 + delt_max: 0.002 + do_sat_adj: true + do_vort_damp: true + fill: true + hord_dp: 6 + hord_mt: 6 + hord_tm: 6 + hord_tr: 8 + hord_vt: 6 + hydrostatic: false + k_split: 1 + ke_bg: 0. + kord_mt: 9 + kord_tm: -9 + kord_tr: 9 + kord_wz: 9 + n_split: 1 + nord: 3 + nwat: 6 + p_fac: 0.05 + rf_cutoff: 3000. + rf_fast: true + tau: 10. + vtdm4: 0.06 + z_tracer: true + do_qa: true + tau_i2s: 1000. + tau_g2v: 1200. + ql_gen: 0.001 + ql_mlt: 0.002 + qs_mlt: 0.000001 + qi_lim: 1.0 + dw_ocean: 0.1 + dw_land: 0.15 + icloud_f: 0 + tau_l2v: 300. + tau_v2l: 90. + fv_sg_adj: 0 + n_sponge: 48 + u_max: 355.0 + +physics_config: + hydrostatic: false + nwat: 6 + do_qa: true diff --git a/driver/examples/configs/test_external_C12_2x2.yaml b/driver/examples/configs/test_external_C12_2x2.yaml new file mode 100644 index 00000000..8c3c4acf --- /dev/null +++ b/driver/examples/configs/test_external_C12_2x2.yaml @@ -0,0 +1,105 @@ +stencil_config: + compilation_config: + backend: numpy + rebuild: false + validate_args: true + format_source: false + device_sync: false +grid_config: + type: external + config: + grid_type: 4 + dx_const: 3000.0 + dy_const: 3000.0 + deglat: 10.0 + grid_file_path: /input/C12_grid.tile +initialization: + type: analytic + config: + case: baroclinic +performance_config: + collect_performance: true + experiment_name: c12_baroclinic +nx_tile: 12 +nz: 79 +dt_atmos: 225 +minutes: 15 +layout: + - 2 + - 2 +diagnostics_config: + path: output + output_format: netcdf + names: + - u + - v + - ua + - va + - pt + - delp + - qvapor + - qliquid + - qice + - qrain + - qsnow + - qgraupel + z_select: + - level: 65 + names: + - pt +dycore_config: + a_imp: 1.0 + beta: 0. + consv_te: 0. + d2_bg: 0. + d2_bg_k1: 0.2 + d2_bg_k2: 0.1 + d4_bg: 0.15 + d_con: 1.0 + d_ext: 0.0 + dddmp: 0.5 + delt_max: 0.002 + do_sat_adj: true + do_vort_damp: true + fill: true + hord_dp: 6 + hord_mt: 6 + hord_tm: 6 + hord_tr: 8 + hord_vt: 6 + hydrostatic: false + k_split: 1 + ke_bg: 0. + kord_mt: 9 + kord_tm: -9 + kord_tr: 9 + kord_wz: 9 + n_split: 1 + nord: 3 + nwat: 6 + p_fac: 0.05 + rf_cutoff: 3000. + rf_fast: true + tau: 10. + vtdm4: 0.06 + z_tracer: true + do_qa: true + tau_i2s: 1000. + tau_g2v: 1200. + ql_gen: 0.001 + ql_mlt: 0.002 + qs_mlt: 0.000001 + qi_lim: 1.0 + dw_ocean: 0.1 + dw_land: 0.15 + icloud_f: 0 + tau_l2v: 300. + tau_v2l: 90. + fv_sg_adj: 0 + n_sponge: 48 + u_max: 355.0 + +physics_config: + hydrostatic: false + nwat: 6 + do_qa: true diff --git a/driver/pace/driver/grid.py b/driver/pace/driver/grid.py index 737b5981..f8bf9d74 100644 --- a/driver/pace/driver/grid.py +++ b/driver/pace/driver/grid.py @@ -9,6 +9,7 @@ import pace.dsl import pace.physics import pace.stencils +import pace.util import pace.util.grid from pace.stencils.testing import TranslateGrid from pace.util import Communicator, QuantityFactory @@ -66,7 +67,8 @@ def get_grid( communicator: Communicator, ) -> Tuple[DampingCoefficients, DriverGridData, GridData]: return self.config.get_grid( - quantity_factory=quantity_factory, communicator=communicator + quantity_factory=quantity_factory, + communicator=communicator, ) @classmethod @@ -209,11 +211,8 @@ class ExternalGridConfig(GridInitializer): and grid_data variables """ - restart_path: Optional[str] = None grid_type: Optional[int] = 0 - dx_const: Optional[float] = 1000.0 - dy_const: Optional[float] = 1000.0 - deglat: Optional[float] = 15.0 + grid_file_path: str = "/input/tilefile" # TODO: Area read in? @@ -225,7 +224,15 @@ def get_grid( pace_log.info("Using external grid data") - ds = xr.open_dataset("../../../input/file.nc") + if self.grid_type <= 3: + tile_num = pace.util.get_tile_index( + communicator.rank, communicator.partitioner.total_ranks + ) + tile_file = self.grid_file_path + str(tile_num) + ".nc" + else: + tile_file = self.grid_file_path + + ds = xr.open_dataset(tile_file) lon = ds.x.values lat = ds.y.values dx = ds.dx.values @@ -233,34 +240,40 @@ def get_grid( npx = ds.nxp.values.size npy = ds.nyp.values.size - subtile_slice = communicator.partitioner.subtile_slice( + subtile_slice_grid = communicator.partitioner.subtile_slice( rank=communicator.rank, - global_dims=[pace.util.X_DIM, pace.util.Y_DIM], + global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_INTERFACE_DIM], global_extent=(npx, npy), overlap=True, ) + subtile_slice_dx = communicator.partitioner.subtile_slice( + rank=communicator.rank, + global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_DIM], + global_extent=(npx - 1, npy), + overlap=True, + ) + + subtile_slice_dy = communicator.partitioner.subtile_slice( + rank=communicator.rank, + global_dims=[pace.util.X_DIM, pace.util.Y_INTERFACE_DIM], + global_extent=(npx, npy - 1), + overlap=True, + ) + metric_terms = MetricTerms.from_generated( - x=lon[subtile_slice], - y=lat[subtile_slice], - dx=dx, - dy=dy, + x=lon[subtile_slice_grid], + y=lat[subtile_slice_grid], + dx=dx[subtile_slice_dx], + dy=dy[subtile_slice_dy], quantity_factory=quantity_factory, communicator=communicator, grid_type=self.grid_type, - dx_const=self.dx_const, - dy_const=self.dy_const, - deglat=self.deglat, extdgrid=True, ) horizontal_data = HorizontalGridData.new_from_metric_terms(metric_terms) - if self.restart_path is not None: - vertical_data = VerticalGridData.from_restart( - self.restart_path, quantity_factory=quantity_factory - ) - else: - vertical_data = VerticalGridData.new_from_metric_terms(metric_terms) + vertical_data = VerticalGridData.new_from_metric_terms(metric_terms) contravariant_data = ContravariantGridData.new_from_metric_terms(metric_terms) angle_data = AngleGridData.new_from_metric_terms(metric_terms) grid_data = GridData( diff --git a/util/pace/util/grid/generation.py b/util/pace/util/grid/generation.py index 548d21ad..56435134 100644 --- a/util/pace/util/grid/generation.py +++ b/util/pace/util/grid/generation.py @@ -428,21 +428,15 @@ def from_generated( y, dx, dy, - quantity_factory: util.QuantityFactory, - communicator: util.Communicator, - grid_type: int = 0, - dx_const: float = 1000.0, - dy_const: float = 1000.0, - deglat: float = 15.0, + quantity_factory, + communicator, + grid_type, extdgrid: bool = True, ) -> "MetricTerms": terms = MetricTerms( quantity_factory=quantity_factory, communicator=communicator, grid_type=grid_type, - dx_const=dx_const, - dy_const=dy_const, - deglat=deglat, extdgrid=extdgrid, ) From 396025c072602b5f4fbf8bb0549c6ae5eee2e268 Mon Sep 17 00:00:00 2001 From: Frank Malatino Date: Tue, 7 Nov 2023 16:04:09 -0500 Subject: [PATCH 11/31] Preliminary unit test for external grid data read in --- tests/mpi_54rank/test_external_grid_init.py | 471 ++++++++++++++++++++ 1 file changed, 471 insertions(+) create mode 100644 tests/mpi_54rank/test_external_grid_init.py diff --git a/tests/mpi_54rank/test_external_grid_init.py b/tests/mpi_54rank/test_external_grid_init.py new file mode 100644 index 00000000..9e266046 --- /dev/null +++ b/tests/mpi_54rank/test_external_grid_init.py @@ -0,0 +1,471 @@ +import numpy as np +import xarray as xr + +import pace.driver +import pace.util +from pace.util.grid import MetricTerms +from pace.util.mpi import MPIComm + + +def get_quantity_factory(layout, nx_tile, ny_tile, nz): + nx = nx_tile // layout[0] + ny = ny_tile // layout[1] + return pace.util.QuantityFactory( + sizer=pace.util.SubtileGridSizer( + nx=nx, ny=ny, nz=nz, n_halo=3, extra_dim_lengths={} + ), + numpy=np, + ) + + +def get_cube_comm(layout, comm: MPIComm): + return pace.util.CubedSphereCommunicator( + comm=comm, + partitioner=pace.util.CubedSpherePartitioner( + pace.util.TilePartitioner(layout=layout) + ), + ) + + +def test_extgrid_equals_generated_1x1_6(): + nx_tile, ny_tile, nz = 6, 6, 5 + comm_1by1 = MPIComm() + cube_comm = get_cube_comm(layout=(1, 1), comm=comm_1by1) + metric_terms_gen = MetricTerms( + quantity_factory=get_quantity_factory( + layout=(1, 1), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz + ), + communicator=cube_comm, + ) + + tile_file = "/input/tilefile" + ds = xr.open_dataset(tile_file) + lon = ds.x.values + lat = ds.y.values + dx = ds.dx.values + dy = ds.dy.values + npx = ds.nxp.values.size + npy = ds.nyp.values.size + + subtile_slice_grid = cube_comm.partitioner.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_INTERFACE_DIM], + global_extent=(npx, npy), + overlap=True, + ) + + subtile_slice_dx = cube_comm.partitioner.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_DIM], + global_extent=(npx - 1, npy), + overlap=True, + ) + + subtile_slice_dy = cube_comm.partitioner.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.X_DIM, pace.util.Y_INTERFACE_DIM], + global_extent=(npx, npy - 1), + overlap=True, + ) + + metric_terms_ext = MetricTerms.from_generated( + x=lon[subtile_slice_grid], + y=lat[subtile_slice_grid], + dx=dx[subtile_slice_dx], + dy=dy[subtile_slice_dy], + quantity_factory=get_quantity_factory( + layout=(1, 1), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz + ), + communicator=cube_comm, + grid_type=0, + extdgrid=True, + ) + + errors = [] + + if metric_terms_gen.grid.data[:, :, 0] != metric_terms_ext.grid.data[:, :, 0]: + errors.append("lon data mismatch between generated and external grid data") + + if metric_terms_gen.grid.data[:, :, 1] != metric_terms_ext.grid.data[:, :, 1]: + errors.append("lat data mismatch between generated and external grid data") + + if metric_terms_gen._dx != metric_terms_ext._dx: + errors.append("dx data mismatch between generated and external grid data") + + if metric_terms_gen._dy != metric_terms_ext._dy: + errors.append("dy data mismatch between generated and external grid data") + + if metric_terms_gen._area != metric_terms_ext._area: + errors.append("area data mismatch between generated and external grid data") + + assert not errors, "errors occured in 1x1_6:\n{}".format("\n".join(errors)) + + +def test_extgrid_equals_generated_1x1_24(): + nx_tile, ny_tile, nz = 24, 24, 5 + comm_1by1 = MPIComm() + cube_comm = get_cube_comm(layout=(1, 1), comm=comm_1by1) + metric_terms_gen = MetricTerms( + quantity_factory=get_quantity_factory( + layout=(1, 1), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz + ), + communicator=cube_comm, + ) + + tile_file = "/input/tilefile" + ds = xr.open_dataset(tile_file) + lon = ds.x.values + lat = ds.y.values + dx = ds.dx.values + dy = ds.dy.values + npx = ds.nxp.values.size + npy = ds.nyp.values.size + + subtile_slice_grid = cube_comm.partitioner.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_INTERFACE_DIM], + global_extent=(npx, npy), + overlap=True, + ) + + subtile_slice_dx = cube_comm.partitioner.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_DIM], + global_extent=(npx - 1, npy), + overlap=True, + ) + + subtile_slice_dy = cube_comm.partitioner.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.X_DIM, pace.util.Y_INTERFACE_DIM], + global_extent=(npx, npy - 1), + overlap=True, + ) + + metric_terms_ext = MetricTerms.from_generated( + x=lon[subtile_slice_grid], + y=lat[subtile_slice_grid], + dx=dx[subtile_slice_dx], + dy=dy[subtile_slice_dy], + quantity_factory=get_quantity_factory( + layout=(1, 1), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz + ), + communicator=cube_comm, + grid_type=0, + extdgrid=True, + ) + + errors = [] + + if metric_terms_gen.grid.data[:, :, 0] != metric_terms_ext.grid.data[:, :, 0]: + errors.append("lon data mismatch between generated and external grid data") + + if metric_terms_gen.grid.data[:, :, 1] != metric_terms_ext.grid.data[:, :, 1]: + errors.append("lat data mismatch between generated and external grid data") + + if metric_terms_gen._dx != metric_terms_ext._dx: + errors.append("dx data mismatch between generated and external grid data") + + if metric_terms_gen._dy != metric_terms_ext._dy: + errors.append("dy data mismatch between generated and external grid data") + + if metric_terms_gen._area != metric_terms_ext._area: + errors.append("area data mismatch between generated and external grid data") + + assert not errors, "errors occured in 1x1_24:\n{}".format("\n".join(errors)) + + +def test_extgrid_equals_generated_1x1_54(): + nx_tile, ny_tile, nz = 54, 54, 5 + comm_1by1 = MPIComm() + cube_comm = get_cube_comm(layout=(1, 1), comm=comm_1by1) + metric_terms_gen = MetricTerms( + quantity_factory=get_quantity_factory( + layout=(1, 1), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz + ), + communicator=cube_comm, + ) + + tile_file = "/input/tilefile" + ds = xr.open_dataset(tile_file) + lon = ds.x.values + lat = ds.y.values + dx = ds.dx.values + dy = ds.dy.values + npx = ds.nxp.values.size + npy = ds.nyp.values.size + + subtile_slice_grid = cube_comm.partitioner.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_INTERFACE_DIM], + global_extent=(npx, npy), + overlap=True, + ) + + subtile_slice_dx = cube_comm.partitioner.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_DIM], + global_extent=(npx - 1, npy), + overlap=True, + ) + + subtile_slice_dy = cube_comm.partitioner.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.X_DIM, pace.util.Y_INTERFACE_DIM], + global_extent=(npx, npy - 1), + overlap=True, + ) + + metric_terms_ext = MetricTerms.from_generated( + x=lon[subtile_slice_grid], + y=lat[subtile_slice_grid], + dx=dx[subtile_slice_dx], + dy=dy[subtile_slice_dy], + quantity_factory=get_quantity_factory( + layout=(1, 1), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz + ), + communicator=cube_comm, + grid_type=0, + extdgrid=True, + ) + + errors = [] + + if metric_terms_gen.grid.data[:, :, 0] != metric_terms_ext.grid.data[:, :, 0]: + errors.append("lon data mismatch between generated and external grid data") + + if metric_terms_gen.grid.data[:, :, 1] != metric_terms_ext.grid.data[:, :, 1]: + errors.append("lat data mismatch between generated and external grid data") + + if metric_terms_gen._dx != metric_terms_ext._dx: + errors.append("dx data mismatch between generated and external grid data") + + if metric_terms_gen._dy != metric_terms_ext._dy: + errors.append("dy data mismatch between generated and external grid data") + + if metric_terms_gen._area != metric_terms_ext._area: + errors.append("area data mismatch between generated and external grid data") + + assert not errors, "errors occured in 1x1_54:\n{}".format("\n".join(errors)) + + +def test_extgrid_equals_generated_2x2_6(): + nx_tile, ny_tile, nz = 6, 6, 5 + comm_1by1 = MPIComm() + cube_comm = get_cube_comm(layout=(2, 2), comm=comm_1by1) + metric_terms_gen = MetricTerms( + quantity_factory=get_quantity_factory( + layout=(2, 2), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz + ), + communicator=cube_comm, + ) + + tile_file = "/input/tilefile" + ds = xr.open_dataset(tile_file) + lon = ds.x.values + lat = ds.y.values + dx = ds.dx.values + dy = ds.dy.values + npx = ds.nxp.values.size + npy = ds.nyp.values.size + + subtile_slice_grid = cube_comm.partitioner.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_INTERFACE_DIM], + global_extent=(npx, npy), + overlap=True, + ) + + subtile_slice_dx = cube_comm.partitioner.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_DIM], + global_extent=(npx - 1, npy), + overlap=True, + ) + + subtile_slice_dy = cube_comm.partitioner.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.X_DIM, pace.util.Y_INTERFACE_DIM], + global_extent=(npx, npy - 1), + overlap=True, + ) + + metric_terms_ext = MetricTerms.from_generated( + x=lon[subtile_slice_grid], + y=lat[subtile_slice_grid], + dx=dx[subtile_slice_dx], + dy=dy[subtile_slice_dy], + quantity_factory=get_quantity_factory( + layout=(1, 1), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz + ), + communicator=cube_comm, + grid_type=0, + extdgrid=True, + ) + + errors = [] + + if metric_terms_gen.grid.data[:, :, 0] != metric_terms_ext.grid.data[:, :, 0]: + errors.append("lon data mismatch between generated and external grid data") + + if metric_terms_gen.grid.data[:, :, 1] != metric_terms_ext.grid.data[:, :, 1]: + errors.append("lat data mismatch between generated and external grid data") + + if metric_terms_gen._dx != metric_terms_ext._dx: + errors.append("dx data mismatch between generated and external grid data") + + if metric_terms_gen._dy != metric_terms_ext._dy: + errors.append("dy data mismatch between generated and external grid data") + + if metric_terms_gen._area != metric_terms_ext._area: + errors.append("area data mismatch between generated and external grid data") + + assert not errors, "errors occured in 2x2_6:\n{}".format("\n".join(errors)) + + +def test_extgrid_equals_generated_2x2_24(): + nx_tile, ny_tile, nz = 24, 24, 5 + comm_1by1 = MPIComm() + cube_comm = get_cube_comm(layout=(2, 2), comm=comm_1by1) + metric_terms_gen = MetricTerms( + quantity_factory=get_quantity_factory( + layout=(2, 2), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz + ), + communicator=cube_comm, + ) + + tile_file = "/input/tilefile" + ds = xr.open_dataset(tile_file) + lon = ds.x.values + lat = ds.y.values + dx = ds.dx.values + dy = ds.dy.values + npx = ds.nxp.values.size + npy = ds.nyp.values.size + + subtile_slice_grid = cube_comm.partitioner.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_INTERFACE_DIM], + global_extent=(npx, npy), + overlap=True, + ) + + subtile_slice_dx = cube_comm.partitioner.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_DIM], + global_extent=(npx - 1, npy), + overlap=True, + ) + + subtile_slice_dy = cube_comm.partitioner.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.X_DIM, pace.util.Y_INTERFACE_DIM], + global_extent=(npx, npy - 1), + overlap=True, + ) + + metric_terms_ext = MetricTerms.from_generated( + x=lon[subtile_slice_grid], + y=lat[subtile_slice_grid], + dx=dx[subtile_slice_dx], + dy=dy[subtile_slice_dy], + quantity_factory=get_quantity_factory( + layout=(1, 1), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz + ), + communicator=cube_comm, + grid_type=0, + extdgrid=True, + ) + + errors = [] + + if metric_terms_gen.grid.data[:, :, 0] != metric_terms_ext.grid.data[:, :, 0]: + errors.append("lon data mismatch between generated and external grid data") + + if metric_terms_gen.grid.data[:, :, 1] != metric_terms_ext.grid.data[:, :, 1]: + errors.append("lat data mismatch between generated and external grid data") + + if metric_terms_gen._dx != metric_terms_ext._dx: + errors.append("dx data mismatch between generated and external grid data") + + if metric_terms_gen._dy != metric_terms_ext._dy: + errors.append("dy data mismatch between generated and external grid data") + + if metric_terms_gen._area != metric_terms_ext._area: + errors.append("area data mismatch between generated and external grid data") + + assert not errors, "errors occured in 2x2_24:\n{}".format("\n".join(errors)) + + +def test_extgrid_equals_generated_2x2_54(): + nx_tile, ny_tile, nz = 54, 54, 5 + comm_1by1 = MPIComm() + cube_comm = get_cube_comm(layout=(2, 2), comm=comm_1by1) + metric_terms_gen = MetricTerms( + quantity_factory=get_quantity_factory( + layout=(2, 2), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz + ), + communicator=cube_comm, + ) + + tile_file = "/input/tilefile" + ds = xr.open_dataset(tile_file) + lon = ds.x.values + lat = ds.y.values + dx = ds.dx.values + dy = ds.dy.values + npx = ds.nxp.values.size + npy = ds.nyp.values.size + + subtile_slice_grid = cube_comm.partitioner.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_INTERFACE_DIM], + global_extent=(npx, npy), + overlap=True, + ) + + subtile_slice_dx = cube_comm.partitioner.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_DIM], + global_extent=(npx - 1, npy), + overlap=True, + ) + + subtile_slice_dy = cube_comm.partitioner.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.X_DIM, pace.util.Y_INTERFACE_DIM], + global_extent=(npx, npy - 1), + overlap=True, + ) + + metric_terms_ext = MetricTerms.from_generated( + x=lon[subtile_slice_grid], + y=lat[subtile_slice_grid], + dx=dx[subtile_slice_dx], + dy=dy[subtile_slice_dy], + quantity_factory=get_quantity_factory( + layout=(1, 1), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz + ), + communicator=cube_comm, + grid_type=0, + extdgrid=True, + ) + + errors = [] + + if metric_terms_gen.grid.data[:, :, 0] != metric_terms_ext.grid.data[:, :, 0]: + errors.append("lon data mismatch between generated and external grid data") + + if metric_terms_gen.grid.data[:, :, 1] != metric_terms_ext.grid.data[:, :, 1]: + errors.append("lat data mismatch between generated and external grid data") + + if metric_terms_gen._dx != metric_terms_ext._dx: + errors.append("dx data mismatch between generated and external grid data") + + if metric_terms_gen._dy != metric_terms_ext._dy: + errors.append("dy data mismatch between generated and external grid data") + + if metric_terms_gen._area != metric_terms_ext._area: + errors.append("area data mismatch between generated and external grid data") + + assert not errors, "errors occured in 2x2_54:\n{}".format("\n".join(errors)) From a63be710bf9e0eff36e20b5a84078a00f6473a91 Mon Sep 17 00:00:00 2001 From: Frank Malatino Date: Mon, 27 Nov 2023 21:58:00 -0500 Subject: [PATCH 12/31] Current state of unit tests as of 27 Nov 2023 --- driver/pace/driver/grid.py | 10 +- .../mpi_54rank/test_external_grid_6rank1x1.py | 116 +++++ .../mpi_54rank/test_external_grid_6rank2x2.py | 116 +++++ tests/mpi_54rank/test_external_grid_init.py | 471 ------------------ tests/mpi_54rank/test_grid_init.py | 4 +- util/pace/util/grid/generation.py | 36 +- 6 files changed, 271 insertions(+), 482 deletions(-) create mode 100644 tests/mpi_54rank/test_external_grid_6rank1x1.py create mode 100644 tests/mpi_54rank/test_external_grid_6rank2x2.py delete mode 100644 tests/mpi_54rank/test_external_grid_init.py diff --git a/driver/pace/driver/grid.py b/driver/pace/driver/grid.py index f8bf9d74..69cfe48b 100644 --- a/driver/pace/driver/grid.py +++ b/driver/pace/driver/grid.py @@ -249,19 +249,19 @@ def get_grid( subtile_slice_dx = communicator.partitioner.subtile_slice( rank=communicator.rank, - global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_DIM], - global_extent=(npx - 1, npy), + global_dims=[pace.util.X_DIM, pace.util.Y_INTERFACE_DIM], + global_extent=(npx, npy), overlap=True, ) subtile_slice_dy = communicator.partitioner.subtile_slice( rank=communicator.rank, - global_dims=[pace.util.X_DIM, pace.util.Y_INTERFACE_DIM], - global_extent=(npx, npy - 1), + global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_DIM], + global_extent=(npx, npy), overlap=True, ) - metric_terms = MetricTerms.from_generated( + metric_terms = MetricTerms.from_external( x=lon[subtile_slice_grid], y=lat[subtile_slice_grid], dx=dx[subtile_slice_dx], diff --git a/tests/mpi_54rank/test_external_grid_6rank1x1.py b/tests/mpi_54rank/test_external_grid_6rank1x1.py new file mode 100644 index 00000000..166570ab --- /dev/null +++ b/tests/mpi_54rank/test_external_grid_6rank1x1.py @@ -0,0 +1,116 @@ +import numpy as np +import xarray as xr + +import pace.driver +import pace.util +from pace.util.grid import MetricTerms +from pace.util.mpi import MPIComm + + +def get_quantity_factory(layout, nx_tile, ny_tile, nz): + nx = nx_tile // layout[0] + ny = ny_tile // layout[1] + return pace.util.QuantityFactory( + sizer=pace.util.SubtileGridSizer( + nx=nx, ny=ny, nz=nz, n_halo=3, extra_dim_lengths={} + ), + numpy=np, + ) + + +def get_cube_comm(layout, comm: MPIComm): + return pace.util.CubedSphereCommunicator( + comm=comm, + partitioner=pace.util.CubedSpherePartitioner( + pace.util.TilePartitioner(layout=layout) + ), + ) + + +def get_tile_num(comm: MPIComm): + return pace.util.get_tile_index(comm.rank, comm.partitioner.total_ranks) + + +def test_extgrid_equals_generated_1x1_6(): + nx_tile, ny_tile, nz = 6, 6, 5 + comm_1by1 = MPIComm() + cube_comm = get_cube_comm(layout=(1, 1), comm=comm_1by1) + metric_terms_gen = MetricTerms( + quantity_factory=get_quantity_factory( + layout=(1, 1), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz + ), + communicator=cube_comm, + ) + + metric_terms_gen._dx, metric_terms_gen._dy = metric_terms_gen._compute_dxdy() + + tile_num = get_tile_num(cube_comm) + 1 + tile_file = "../../test_input/ext_test_6.tile" + str(tile_num) + ".nc" + ds = xr.open_dataset(tile_file) + lon = ds.x.values + lat = ds.y.values + dx = ds.dx.values + dy = ds.dy.values + npx = ds.nxp.values.size + npy = ds.nyp.values.size + + subtile_slice_grid = cube_comm.partitioner.tile.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.Y_INTERFACE_DIM, pace.util.X_INTERFACE_DIM], + global_extent=(npx, npy), + overlap=True, + ) + + subtile_slice_dx = cube_comm.partitioner.tile.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.Y_INTERFACE_DIM, pace.util.X_DIM], + global_extent=(npx, npy), + overlap=True, + ) + + subtile_slice_dy = cube_comm.partitioner.tile.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.Y_DIM, pace.util.X_INTERFACE_DIM], + global_extent=(npx, npy), + overlap=True, + ) + + metric_terms_ext = MetricTerms.from_external( + x=lon[subtile_slice_grid], + y=lat[subtile_slice_grid], + dx=dx[subtile_slice_dx], + dy=dy[subtile_slice_dy], + quantity_factory=get_quantity_factory( + layout=(1, 1), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz + ), + communicator=cube_comm, + grid_type=0, + extdgrid=True, + ) + + print(metric_terms_ext._dx) + + errors = [] + + if ( + metric_terms_gen.grid.data[:, :, 0].any() + != metric_terms_ext.grid.data[:, :, 0].any() + ): + errors.append("lon data mismatch between generated and external grid data") + + if ( + metric_terms_gen.grid.data[:, :, 1].any() + != metric_terms_ext.grid.data[:, :, 1].any() + ): + errors.append("lat data mismatch between generated and external grid data") + + if metric_terms_gen._dx.view[:, :].any() != metric_terms_ext._dx.view[:, :].any(): + errors.append("dx data mismatch between generated and external grid data") + + if metric_terms_gen._dy.view[:, :].any() != metric_terms_ext._dy.view[:, :].any(): + errors.append("dy data mismatch between generated and external grid data") + + # if metric_terms_gen._area.data[:,:].any() != metric_terms_ext._area.any(): + # errors.append("area data mismatch between generated and external grid data") + + assert not errors, "errors occured in 1x1_6:\n{}".format("\n".join(errors)) diff --git a/tests/mpi_54rank/test_external_grid_6rank2x2.py b/tests/mpi_54rank/test_external_grid_6rank2x2.py new file mode 100644 index 00000000..6d5fba03 --- /dev/null +++ b/tests/mpi_54rank/test_external_grid_6rank2x2.py @@ -0,0 +1,116 @@ +import numpy as np +import xarray as xr + +import pace.driver +import pace.util +from pace.util.grid import MetricTerms +from pace.util.mpi import MPIComm + + +def get_quantity_factory(layout, nx_tile, ny_tile, nz): + nx = nx_tile // layout[0] + ny = ny_tile // layout[1] + return pace.util.QuantityFactory( + sizer=pace.util.SubtileGridSizer( + nx=nx, ny=ny, nz=nz, n_halo=3, extra_dim_lengths={} + ), + numpy=np, + ) + + +def get_cube_comm(layout, comm: MPIComm): + return pace.util.CubedSphereCommunicator( + comm=comm, + partitioner=pace.util.CubedSpherePartitioner( + pace.util.TilePartitioner(layout=layout) + ), + ) + + +def get_tile_num(comm: MPIComm): + return pace.util.get_tile_index(comm.rank, comm.partitioner.total_ranks) + + +def test_extgrid_equals_generated_2x2_6(): + nx_tile, ny_tile, nz = 7, 7, 5 + comm_2by2 = MPIComm() + cube_comm = get_cube_comm(layout=(2, 2), comm=comm_2by2) + metric_terms_gen = MetricTerms( + quantity_factory=get_quantity_factory( + layout=(2, 2), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz + ), + communicator=cube_comm, + ) + + metric_terms_gen._dx, metric_terms_gen._dy = metric_terms_gen._compute_dxdy() + + tile_num = get_tile_num(cube_comm) + 1 + tile_file = "../../test_input/ext_test_6.tile" + str(tile_num) + ".nc" + ds = xr.open_dataset(tile_file) + lon = ds.x.values + lat = ds.y.values + dx = ds.dx.values + dy = ds.dy.values + npx = ds.nxp.values.size + npy = ds.nyp.values.size + + subtile_slice_grid = cube_comm.partitioner.tile.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.Y_INTERFACE_DIM, pace.util.X_INTERFACE_DIM], + global_extent=(npx, npy), + overlap=True, + ) + + subtile_slice_dx = cube_comm.partitioner.tile.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.Y_INTERFACE_DIM, pace.util.X_DIM], + global_extent=(npx, npy), + overlap=True, + ) + + subtile_slice_dy = cube_comm.partitioner.tile.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.Y_DIM, pace.util.X_INTERFACE_DIM], + global_extent=(npx, npy), + overlap=True, + ) + + metric_terms_ext = MetricTerms.from_external( + x=lon[subtile_slice_grid], + y=lat[subtile_slice_grid], + dx=dx[subtile_slice_dx], + dy=dy[subtile_slice_dy], + quantity_factory=get_quantity_factory( + layout=(2, 2), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz + ), + communicator=cube_comm, + grid_type=0, + extdgrid=True, + ) + + print(metric_terms_ext._dx) + + errors = [] + + if ( + metric_terms_gen.grid.data[:, :, 0].any() + != metric_terms_ext.grid.data[:, :, 0].any() + ): + errors.append("lon data mismatch between generated and external grid data") + + if ( + metric_terms_gen.grid.data[:, :, 1].any() + != metric_terms_ext.grid.data[:, :, 1].any() + ): + errors.append("lat data mismatch between generated and external grid data") + + if metric_terms_gen._dx.view[:, :].any() != metric_terms_ext._dx.view[:, :].any(): + errors.append("dx data mismatch between generated and external grid data") + + if metric_terms_gen._dy.view[:, :].any() != metric_terms_ext._dy.view[:, :].any(): + errors.append("dy data mismatch between generated and external grid data") + + # if metric_terms_gen._area.data[:,:].any() != metric_terms_ext._area.any(): + # errors.append("area data mismatch between generated and external grid data") + + assert not errors, "errors occured in 1x1_6:\n{}".format("\n".join(errors)) diff --git a/tests/mpi_54rank/test_external_grid_init.py b/tests/mpi_54rank/test_external_grid_init.py deleted file mode 100644 index 9e266046..00000000 --- a/tests/mpi_54rank/test_external_grid_init.py +++ /dev/null @@ -1,471 +0,0 @@ -import numpy as np -import xarray as xr - -import pace.driver -import pace.util -from pace.util.grid import MetricTerms -from pace.util.mpi import MPIComm - - -def get_quantity_factory(layout, nx_tile, ny_tile, nz): - nx = nx_tile // layout[0] - ny = ny_tile // layout[1] - return pace.util.QuantityFactory( - sizer=pace.util.SubtileGridSizer( - nx=nx, ny=ny, nz=nz, n_halo=3, extra_dim_lengths={} - ), - numpy=np, - ) - - -def get_cube_comm(layout, comm: MPIComm): - return pace.util.CubedSphereCommunicator( - comm=comm, - partitioner=pace.util.CubedSpherePartitioner( - pace.util.TilePartitioner(layout=layout) - ), - ) - - -def test_extgrid_equals_generated_1x1_6(): - nx_tile, ny_tile, nz = 6, 6, 5 - comm_1by1 = MPIComm() - cube_comm = get_cube_comm(layout=(1, 1), comm=comm_1by1) - metric_terms_gen = MetricTerms( - quantity_factory=get_quantity_factory( - layout=(1, 1), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz - ), - communicator=cube_comm, - ) - - tile_file = "/input/tilefile" - ds = xr.open_dataset(tile_file) - lon = ds.x.values - lat = ds.y.values - dx = ds.dx.values - dy = ds.dy.values - npx = ds.nxp.values.size - npy = ds.nyp.values.size - - subtile_slice_grid = cube_comm.partitioner.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_INTERFACE_DIM], - global_extent=(npx, npy), - overlap=True, - ) - - subtile_slice_dx = cube_comm.partitioner.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_DIM], - global_extent=(npx - 1, npy), - overlap=True, - ) - - subtile_slice_dy = cube_comm.partitioner.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.X_DIM, pace.util.Y_INTERFACE_DIM], - global_extent=(npx, npy - 1), - overlap=True, - ) - - metric_terms_ext = MetricTerms.from_generated( - x=lon[subtile_slice_grid], - y=lat[subtile_slice_grid], - dx=dx[subtile_slice_dx], - dy=dy[subtile_slice_dy], - quantity_factory=get_quantity_factory( - layout=(1, 1), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz - ), - communicator=cube_comm, - grid_type=0, - extdgrid=True, - ) - - errors = [] - - if metric_terms_gen.grid.data[:, :, 0] != metric_terms_ext.grid.data[:, :, 0]: - errors.append("lon data mismatch between generated and external grid data") - - if metric_terms_gen.grid.data[:, :, 1] != metric_terms_ext.grid.data[:, :, 1]: - errors.append("lat data mismatch between generated and external grid data") - - if metric_terms_gen._dx != metric_terms_ext._dx: - errors.append("dx data mismatch between generated and external grid data") - - if metric_terms_gen._dy != metric_terms_ext._dy: - errors.append("dy data mismatch between generated and external grid data") - - if metric_terms_gen._area != metric_terms_ext._area: - errors.append("area data mismatch between generated and external grid data") - - assert not errors, "errors occured in 1x1_6:\n{}".format("\n".join(errors)) - - -def test_extgrid_equals_generated_1x1_24(): - nx_tile, ny_tile, nz = 24, 24, 5 - comm_1by1 = MPIComm() - cube_comm = get_cube_comm(layout=(1, 1), comm=comm_1by1) - metric_terms_gen = MetricTerms( - quantity_factory=get_quantity_factory( - layout=(1, 1), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz - ), - communicator=cube_comm, - ) - - tile_file = "/input/tilefile" - ds = xr.open_dataset(tile_file) - lon = ds.x.values - lat = ds.y.values - dx = ds.dx.values - dy = ds.dy.values - npx = ds.nxp.values.size - npy = ds.nyp.values.size - - subtile_slice_grid = cube_comm.partitioner.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_INTERFACE_DIM], - global_extent=(npx, npy), - overlap=True, - ) - - subtile_slice_dx = cube_comm.partitioner.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_DIM], - global_extent=(npx - 1, npy), - overlap=True, - ) - - subtile_slice_dy = cube_comm.partitioner.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.X_DIM, pace.util.Y_INTERFACE_DIM], - global_extent=(npx, npy - 1), - overlap=True, - ) - - metric_terms_ext = MetricTerms.from_generated( - x=lon[subtile_slice_grid], - y=lat[subtile_slice_grid], - dx=dx[subtile_slice_dx], - dy=dy[subtile_slice_dy], - quantity_factory=get_quantity_factory( - layout=(1, 1), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz - ), - communicator=cube_comm, - grid_type=0, - extdgrid=True, - ) - - errors = [] - - if metric_terms_gen.grid.data[:, :, 0] != metric_terms_ext.grid.data[:, :, 0]: - errors.append("lon data mismatch between generated and external grid data") - - if metric_terms_gen.grid.data[:, :, 1] != metric_terms_ext.grid.data[:, :, 1]: - errors.append("lat data mismatch between generated and external grid data") - - if metric_terms_gen._dx != metric_terms_ext._dx: - errors.append("dx data mismatch between generated and external grid data") - - if metric_terms_gen._dy != metric_terms_ext._dy: - errors.append("dy data mismatch between generated and external grid data") - - if metric_terms_gen._area != metric_terms_ext._area: - errors.append("area data mismatch between generated and external grid data") - - assert not errors, "errors occured in 1x1_24:\n{}".format("\n".join(errors)) - - -def test_extgrid_equals_generated_1x1_54(): - nx_tile, ny_tile, nz = 54, 54, 5 - comm_1by1 = MPIComm() - cube_comm = get_cube_comm(layout=(1, 1), comm=comm_1by1) - metric_terms_gen = MetricTerms( - quantity_factory=get_quantity_factory( - layout=(1, 1), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz - ), - communicator=cube_comm, - ) - - tile_file = "/input/tilefile" - ds = xr.open_dataset(tile_file) - lon = ds.x.values - lat = ds.y.values - dx = ds.dx.values - dy = ds.dy.values - npx = ds.nxp.values.size - npy = ds.nyp.values.size - - subtile_slice_grid = cube_comm.partitioner.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_INTERFACE_DIM], - global_extent=(npx, npy), - overlap=True, - ) - - subtile_slice_dx = cube_comm.partitioner.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_DIM], - global_extent=(npx - 1, npy), - overlap=True, - ) - - subtile_slice_dy = cube_comm.partitioner.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.X_DIM, pace.util.Y_INTERFACE_DIM], - global_extent=(npx, npy - 1), - overlap=True, - ) - - metric_terms_ext = MetricTerms.from_generated( - x=lon[subtile_slice_grid], - y=lat[subtile_slice_grid], - dx=dx[subtile_slice_dx], - dy=dy[subtile_slice_dy], - quantity_factory=get_quantity_factory( - layout=(1, 1), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz - ), - communicator=cube_comm, - grid_type=0, - extdgrid=True, - ) - - errors = [] - - if metric_terms_gen.grid.data[:, :, 0] != metric_terms_ext.grid.data[:, :, 0]: - errors.append("lon data mismatch between generated and external grid data") - - if metric_terms_gen.grid.data[:, :, 1] != metric_terms_ext.grid.data[:, :, 1]: - errors.append("lat data mismatch between generated and external grid data") - - if metric_terms_gen._dx != metric_terms_ext._dx: - errors.append("dx data mismatch between generated and external grid data") - - if metric_terms_gen._dy != metric_terms_ext._dy: - errors.append("dy data mismatch between generated and external grid data") - - if metric_terms_gen._area != metric_terms_ext._area: - errors.append("area data mismatch between generated and external grid data") - - assert not errors, "errors occured in 1x1_54:\n{}".format("\n".join(errors)) - - -def test_extgrid_equals_generated_2x2_6(): - nx_tile, ny_tile, nz = 6, 6, 5 - comm_1by1 = MPIComm() - cube_comm = get_cube_comm(layout=(2, 2), comm=comm_1by1) - metric_terms_gen = MetricTerms( - quantity_factory=get_quantity_factory( - layout=(2, 2), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz - ), - communicator=cube_comm, - ) - - tile_file = "/input/tilefile" - ds = xr.open_dataset(tile_file) - lon = ds.x.values - lat = ds.y.values - dx = ds.dx.values - dy = ds.dy.values - npx = ds.nxp.values.size - npy = ds.nyp.values.size - - subtile_slice_grid = cube_comm.partitioner.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_INTERFACE_DIM], - global_extent=(npx, npy), - overlap=True, - ) - - subtile_slice_dx = cube_comm.partitioner.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_DIM], - global_extent=(npx - 1, npy), - overlap=True, - ) - - subtile_slice_dy = cube_comm.partitioner.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.X_DIM, pace.util.Y_INTERFACE_DIM], - global_extent=(npx, npy - 1), - overlap=True, - ) - - metric_terms_ext = MetricTerms.from_generated( - x=lon[subtile_slice_grid], - y=lat[subtile_slice_grid], - dx=dx[subtile_slice_dx], - dy=dy[subtile_slice_dy], - quantity_factory=get_quantity_factory( - layout=(1, 1), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz - ), - communicator=cube_comm, - grid_type=0, - extdgrid=True, - ) - - errors = [] - - if metric_terms_gen.grid.data[:, :, 0] != metric_terms_ext.grid.data[:, :, 0]: - errors.append("lon data mismatch between generated and external grid data") - - if metric_terms_gen.grid.data[:, :, 1] != metric_terms_ext.grid.data[:, :, 1]: - errors.append("lat data mismatch between generated and external grid data") - - if metric_terms_gen._dx != metric_terms_ext._dx: - errors.append("dx data mismatch between generated and external grid data") - - if metric_terms_gen._dy != metric_terms_ext._dy: - errors.append("dy data mismatch between generated and external grid data") - - if metric_terms_gen._area != metric_terms_ext._area: - errors.append("area data mismatch between generated and external grid data") - - assert not errors, "errors occured in 2x2_6:\n{}".format("\n".join(errors)) - - -def test_extgrid_equals_generated_2x2_24(): - nx_tile, ny_tile, nz = 24, 24, 5 - comm_1by1 = MPIComm() - cube_comm = get_cube_comm(layout=(2, 2), comm=comm_1by1) - metric_terms_gen = MetricTerms( - quantity_factory=get_quantity_factory( - layout=(2, 2), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz - ), - communicator=cube_comm, - ) - - tile_file = "/input/tilefile" - ds = xr.open_dataset(tile_file) - lon = ds.x.values - lat = ds.y.values - dx = ds.dx.values - dy = ds.dy.values - npx = ds.nxp.values.size - npy = ds.nyp.values.size - - subtile_slice_grid = cube_comm.partitioner.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_INTERFACE_DIM], - global_extent=(npx, npy), - overlap=True, - ) - - subtile_slice_dx = cube_comm.partitioner.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_DIM], - global_extent=(npx - 1, npy), - overlap=True, - ) - - subtile_slice_dy = cube_comm.partitioner.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.X_DIM, pace.util.Y_INTERFACE_DIM], - global_extent=(npx, npy - 1), - overlap=True, - ) - - metric_terms_ext = MetricTerms.from_generated( - x=lon[subtile_slice_grid], - y=lat[subtile_slice_grid], - dx=dx[subtile_slice_dx], - dy=dy[subtile_slice_dy], - quantity_factory=get_quantity_factory( - layout=(1, 1), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz - ), - communicator=cube_comm, - grid_type=0, - extdgrid=True, - ) - - errors = [] - - if metric_terms_gen.grid.data[:, :, 0] != metric_terms_ext.grid.data[:, :, 0]: - errors.append("lon data mismatch between generated and external grid data") - - if metric_terms_gen.grid.data[:, :, 1] != metric_terms_ext.grid.data[:, :, 1]: - errors.append("lat data mismatch between generated and external grid data") - - if metric_terms_gen._dx != metric_terms_ext._dx: - errors.append("dx data mismatch between generated and external grid data") - - if metric_terms_gen._dy != metric_terms_ext._dy: - errors.append("dy data mismatch between generated and external grid data") - - if metric_terms_gen._area != metric_terms_ext._area: - errors.append("area data mismatch between generated and external grid data") - - assert not errors, "errors occured in 2x2_24:\n{}".format("\n".join(errors)) - - -def test_extgrid_equals_generated_2x2_54(): - nx_tile, ny_tile, nz = 54, 54, 5 - comm_1by1 = MPIComm() - cube_comm = get_cube_comm(layout=(2, 2), comm=comm_1by1) - metric_terms_gen = MetricTerms( - quantity_factory=get_quantity_factory( - layout=(2, 2), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz - ), - communicator=cube_comm, - ) - - tile_file = "/input/tilefile" - ds = xr.open_dataset(tile_file) - lon = ds.x.values - lat = ds.y.values - dx = ds.dx.values - dy = ds.dy.values - npx = ds.nxp.values.size - npy = ds.nyp.values.size - - subtile_slice_grid = cube_comm.partitioner.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_INTERFACE_DIM], - global_extent=(npx, npy), - overlap=True, - ) - - subtile_slice_dx = cube_comm.partitioner.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_DIM], - global_extent=(npx - 1, npy), - overlap=True, - ) - - subtile_slice_dy = cube_comm.partitioner.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.X_DIM, pace.util.Y_INTERFACE_DIM], - global_extent=(npx, npy - 1), - overlap=True, - ) - - metric_terms_ext = MetricTerms.from_generated( - x=lon[subtile_slice_grid], - y=lat[subtile_slice_grid], - dx=dx[subtile_slice_dx], - dy=dy[subtile_slice_dy], - quantity_factory=get_quantity_factory( - layout=(1, 1), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz - ), - communicator=cube_comm, - grid_type=0, - extdgrid=True, - ) - - errors = [] - - if metric_terms_gen.grid.data[:, :, 0] != metric_terms_ext.grid.data[:, :, 0]: - errors.append("lon data mismatch between generated and external grid data") - - if metric_terms_gen.grid.data[:, :, 1] != metric_terms_ext.grid.data[:, :, 1]: - errors.append("lat data mismatch between generated and external grid data") - - if metric_terms_gen._dx != metric_terms_ext._dx: - errors.append("dx data mismatch between generated and external grid data") - - if metric_terms_gen._dy != metric_terms_ext._dy: - errors.append("dy data mismatch between generated and external grid data") - - if metric_terms_gen._area != metric_terms_ext._area: - errors.append("area data mismatch between generated and external grid data") - - assert not errors, "errors occured in 2x2_54:\n{}".format("\n".join(errors)) diff --git a/tests/mpi_54rank/test_grid_init.py b/tests/mpi_54rank/test_grid_init.py index 9404a3e0..855c539e 100644 --- a/tests/mpi_54rank/test_grid_init.py +++ b/tests/mpi_54rank/test_grid_init.py @@ -4,7 +4,9 @@ import pace.fv3core import pace.util -from pace.fv3core.initialization import init_baroclinic_state +from pace.fv3core.initialization.test_cases.initialize_baroclinic import ( + init_baroclinic_state, +) from pace.util.grid import MetricTerms from pace.util.mpi import MPIComm from pace.util.quantity import Quantity diff --git a/util/pace/util/grid/generation.py b/util/pace/util/grid/generation.py index 56435134..b494466c 100644 --- a/util/pace/util/grid/generation.py +++ b/util/pace/util/grid/generation.py @@ -422,7 +422,7 @@ def __init__( # initializing the dgrid by input from data contained in an # externally generated tile file @classmethod - def from_generated( + def from_external( cls, x, y, @@ -440,11 +440,37 @@ def from_generated( extdgrid=extdgrid, ) + # TODO: Using quantity factory to create quantities + # for dx and dy, quantity_factory.from_array() + rad_conv = PI / 180.0 - terms.grid.data[:, :, 0] = np.dot(rad_conv, x) - terms.grid.data[:, :, 1] = np.dot(rad_conv, y) - terms._dx = dx - terms._dy = dy + terms._grid_64.view[:, :, 0] = np.dot(rad_conv, x) + terms._grid_64.view[:, :, 1] = np.dot(rad_conv, y) + + dx_64 = terms.quantity_factory.zeros( + [util.Y_INTERFACE_DIM, util.X_DIM], + "m", + dtype=np.float64, + allow_mismatch_float_precision=True, + ) + + dx_64.view[:, :] = dx + terms._dx_64 = dx_64 + + dy_64 = terms.quantity_factory.zeros( + [util.Y_DIM, util.X_INTERFACE_DIM], + "m", + dtype=np.float64, + allow_mismatch_float_precision=True, + ) + dy_64.view[:, :] = dy + terms._dy_64 = dy_64 + + terms._comm.halo_update(terms._grid_64, n_points=terms._halo) + terms._comm.vector_halo_update(terms._dx_64, terms._dy_64, n_points=terms._halo) + + # terms._dx = quantity_cast_to_model_float(terms.quantity_factory, terms._dx_64) + # terms._dy = quantity_cast_to_model_float(terms.quantity_factory, terms._dy_64) terms._init_agrid() From d95e1fd18c08e76d5c61eaf54ac92aa51bbc6e5a Mon Sep 17 00:00:00 2001 From: Frank Malatino Date: Mon, 4 Dec 2023 12:30:05 -0500 Subject: [PATCH 13/31] External grid method and unit tests added --- .../configs/test_external_C12_1x1.yaml | 7 +- .../configs/test_external_C12_2x2.yaml | 7 +- driver/pace/driver/grid.py | 32 ++++++--- ..._6rank1x1.py => test_external_grid_1x1.py} | 66 ++++++++++++------- ..._6rank2x2.py => test_external_grid_2x2.py} | 66 ++++++++++++------- util/pace/util/grid/generation.py | 27 +++++--- util/pace/util/partitioner.py | 2 +- 7 files changed, 129 insertions(+), 78 deletions(-) rename tests/mpi_54rank/{test_external_grid_6rank1x1.py => test_external_grid_1x1.py} (62%) rename tests/mpi_54rank/{test_external_grid_6rank2x2.py => test_external_grid_2x2.py} (62%) diff --git a/driver/examples/configs/test_external_C12_1x1.yaml b/driver/examples/configs/test_external_C12_1x1.yaml index 96c0c247..269921ac 100644 --- a/driver/examples/configs/test_external_C12_1x1.yaml +++ b/driver/examples/configs/test_external_C12_1x1.yaml @@ -8,11 +8,8 @@ stencil_config: grid_config: type: external config: - grid_type: 4 - dx_const: 3000.0 - dy_const: 3000.0 - deglat: 10.0 - grid_file_path: /input/C12_grid.tile + grid_type: 0 + grid_file_path: "../../test_input/C12.tile" initialization: type: analytic config: diff --git a/driver/examples/configs/test_external_C12_2x2.yaml b/driver/examples/configs/test_external_C12_2x2.yaml index 8c3c4acf..b12166b8 100644 --- a/driver/examples/configs/test_external_C12_2x2.yaml +++ b/driver/examples/configs/test_external_C12_2x2.yaml @@ -8,11 +8,8 @@ stencil_config: grid_config: type: external config: - grid_type: 4 - dx_const: 3000.0 - dy_const: 3000.0 - deglat: 10.0 - grid_file_path: /input/C12_grid.tile + grid_type: 0 + grid_file_path: "../../test_input/C12.tile" initialization: type: analytic config: diff --git a/driver/pace/driver/grid.py b/driver/pace/driver/grid.py index 69cfe48b..06bc9195 100644 --- a/driver/pace/driver/grid.py +++ b/driver/pace/driver/grid.py @@ -225,8 +225,11 @@ def get_grid( pace_log.info("Using external grid data") if self.grid_type <= 3: - tile_num = pace.util.get_tile_index( - communicator.rank, communicator.partitioner.total_ranks + tile_num = ( + pace.util.get_tile_index( + communicator.rank, communicator.partitioner.total_ranks + ) + + 1 ) tile_file = self.grid_file_path + str(tile_num) + ".nc" else: @@ -237,27 +240,37 @@ def get_grid( lat = ds.y.values dx = ds.dx.values dy = ds.dy.values + area = ds.area.values + nx = ds.nx.values.size + ny = ds.ny.values.size npx = ds.nxp.values.size npy = ds.nyp.values.size - subtile_slice_grid = communicator.partitioner.subtile_slice( + subtile_slice_grid = communicator.partitioner.tile.subtile_slice( rank=communicator.rank, global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_INTERFACE_DIM], global_extent=(npx, npy), overlap=True, ) - subtile_slice_dx = communicator.partitioner.subtile_slice( + subtile_slice_dx = communicator.partitioner.tile.subtile_slice( rank=communicator.rank, - global_dims=[pace.util.X_DIM, pace.util.Y_INTERFACE_DIM], - global_extent=(npx, npy), + global_dims=[pace.util.Y_INTERFACE_DIM, pace.util.X_DIM], + global_extent=(npy, nx), overlap=True, ) - subtile_slice_dy = communicator.partitioner.subtile_slice( + subtile_slice_dy = communicator.partitioner.tile.subtile_slice( rank=communicator.rank, - global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_DIM], - global_extent=(npx, npy), + global_dims=[pace.util.Y_DIM, pace.util.X_INTERFACE_DIM], + global_extent=(ny, npx), + overlap=True, + ) + + subtile_slice_area = communicator.partitioner.tile.subtile_slice( + rank=communicator.rank, + global_dims=[pace.util.Y_DIM, pace.util.X_DIM], + global_extent=(ny, nx), overlap=True, ) @@ -266,6 +279,7 @@ def get_grid( y=lat[subtile_slice_grid], dx=dx[subtile_slice_dx], dy=dy[subtile_slice_dy], + area=area[subtile_slice_area], quantity_factory=quantity_factory, communicator=communicator, grid_type=self.grid_type, diff --git a/tests/mpi_54rank/test_external_grid_6rank1x1.py b/tests/mpi_54rank/test_external_grid_1x1.py similarity index 62% rename from tests/mpi_54rank/test_external_grid_6rank1x1.py rename to tests/mpi_54rank/test_external_grid_1x1.py index 166570ab..d069e1a2 100644 --- a/tests/mpi_54rank/test_external_grid_6rank1x1.py +++ b/tests/mpi_54rank/test_external_grid_1x1.py @@ -1,8 +1,9 @@ import numpy as np import xarray as xr +import yaml -import pace.driver import pace.util +from pace.driver import Driver, DriverConfig from pace.util.grid import MetricTerms from pace.util.mpi import MPIComm @@ -31,26 +32,28 @@ def get_tile_num(comm: MPIComm): return pace.util.get_tile_index(comm.rank, comm.partitioner.total_ranks) -def test_extgrid_equals_generated_1x1_6(): - nx_tile, ny_tile, nz = 6, 6, 5 +def test_extgrid_equals_generated_1x1(): + + with open("../../driver/examples/configs/test_external_C12_1x1.yaml", "r") as ext_f: + ext_config = yaml.safe_load(ext_f) + ext_driver_config = DriverConfig.from_dict(ext_config) + + ext_driver = Driver(ext_driver_config) + + nx_tile, ny_tile, nz = 12, 12, 5 comm_1by1 = MPIComm() cube_comm = get_cube_comm(layout=(1, 1), comm=comm_1by1) - metric_terms_gen = MetricTerms( - quantity_factory=get_quantity_factory( - layout=(1, 1), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz - ), - communicator=cube_comm, - ) - - metric_terms_gen._dx, metric_terms_gen._dy = metric_terms_gen._compute_dxdy() tile_num = get_tile_num(cube_comm) + 1 - tile_file = "../../test_input/ext_test_6.tile" + str(tile_num) + ".nc" + tile_file = "../../test_input/C12.tile" + str(tile_num) + ".nc" ds = xr.open_dataset(tile_file) lon = ds.x.values lat = ds.y.values dx = ds.dx.values dy = ds.dy.values + area = ds.area.values + nx = ds.nx.values.size + ny = ds.ny.values.size npx = ds.nxp.values.size npy = ds.nyp.values.size @@ -64,14 +67,21 @@ def test_extgrid_equals_generated_1x1_6(): subtile_slice_dx = cube_comm.partitioner.tile.subtile_slice( rank=cube_comm.rank, global_dims=[pace.util.Y_INTERFACE_DIM, pace.util.X_DIM], - global_extent=(npx, npy), + global_extent=(npy, nx), overlap=True, ) subtile_slice_dy = cube_comm.partitioner.tile.subtile_slice( rank=cube_comm.rank, global_dims=[pace.util.Y_DIM, pace.util.X_INTERFACE_DIM], - global_extent=(npx, npy), + global_extent=(ny, npx), + overlap=True, + ) + + subtile_slice_area = cube_comm.partitioner.tile.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.Y_DIM, pace.util.X_DIM], + global_extent=(ny, nx), overlap=True, ) @@ -80,6 +90,7 @@ def test_extgrid_equals_generated_1x1_6(): y=lat[subtile_slice_grid], dx=dx[subtile_slice_dx], dy=dy[subtile_slice_dy], + area=area[subtile_slice_area], quantity_factory=get_quantity_factory( layout=(1, 1), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz ), @@ -88,29 +99,36 @@ def test_extgrid_equals_generated_1x1_6(): extdgrid=True, ) - print(metric_terms_ext._dx) - errors = [] if ( - metric_terms_gen.grid.data[:, :, 0].any() + ext_driver.state.grid_data.lon.data.any() != metric_terms_ext.grid.data[:, :, 0].any() ): errors.append("lon data mismatch between generated and external grid data") if ( - metric_terms_gen.grid.data[:, :, 1].any() + ext_driver.state.grid_data.lat.data.any() != metric_terms_ext.grid.data[:, :, 1].any() ): - errors.append("lat data mismatch between generated and external grid data") + errors.append("lon data mismatch between generated and external grid data") - if metric_terms_gen._dx.view[:, :].any() != metric_terms_ext._dx.view[:, :].any(): + if ( + ext_driver.state.grid_data.dx.data.any() + != metric_terms_ext._dx.view[:, :].any() + ): errors.append("dx data mismatch between generated and external grid data") - if metric_terms_gen._dy.view[:, :].any() != metric_terms_ext._dy.view[:, :].any(): + if ( + ext_driver.state.grid_data.dy.data.any() + != metric_terms_ext._dy.view[:, :].any() + ): errors.append("dy data mismatch between generated and external grid data") - # if metric_terms_gen._area.data[:,:].any() != metric_terms_ext._area.any(): - # errors.append("area data mismatch between generated and external grid data") + if ( + ext_driver.state.grid_data.area.data.any() + != metric_terms_ext._area.view[:, :].any() + ): + errors.append("area data mismatch between generated and external grid data") - assert not errors, "errors occured in 1x1_6:\n{}".format("\n".join(errors)) + assert not errors, "errors occured in 1x1:\n{}".format("\n".join(errors)) diff --git a/tests/mpi_54rank/test_external_grid_6rank2x2.py b/tests/mpi_54rank/test_external_grid_2x2.py similarity index 62% rename from tests/mpi_54rank/test_external_grid_6rank2x2.py rename to tests/mpi_54rank/test_external_grid_2x2.py index 6d5fba03..98c8095f 100644 --- a/tests/mpi_54rank/test_external_grid_6rank2x2.py +++ b/tests/mpi_54rank/test_external_grid_2x2.py @@ -1,8 +1,9 @@ import numpy as np import xarray as xr +import yaml -import pace.driver import pace.util +from pace.driver import Driver, DriverConfig from pace.util.grid import MetricTerms from pace.util.mpi import MPIComm @@ -31,26 +32,28 @@ def get_tile_num(comm: MPIComm): return pace.util.get_tile_index(comm.rank, comm.partitioner.total_ranks) -def test_extgrid_equals_generated_2x2_6(): - nx_tile, ny_tile, nz = 7, 7, 5 +def test_extgrid_equals_generated_2x2(): + + with open("../../driver/examples/configs/test_external_C12_2x2.yaml", "r") as ext_f: + ext_config = yaml.safe_load(ext_f) + ext_driver_config = DriverConfig.from_dict(ext_config) + + ext_driver = Driver(ext_driver_config) + + nx_tile, ny_tile, nz = 12, 12, 5 comm_2by2 = MPIComm() cube_comm = get_cube_comm(layout=(2, 2), comm=comm_2by2) - metric_terms_gen = MetricTerms( - quantity_factory=get_quantity_factory( - layout=(2, 2), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz - ), - communicator=cube_comm, - ) - - metric_terms_gen._dx, metric_terms_gen._dy = metric_terms_gen._compute_dxdy() tile_num = get_tile_num(cube_comm) + 1 - tile_file = "../../test_input/ext_test_6.tile" + str(tile_num) + ".nc" + tile_file = "../../test_input/C12.tile" + str(tile_num) + ".nc" ds = xr.open_dataset(tile_file) lon = ds.x.values lat = ds.y.values dx = ds.dx.values dy = ds.dy.values + area = ds.area.values + nx = ds.nx.values.size + ny = ds.ny.values.size npx = ds.nxp.values.size npy = ds.nyp.values.size @@ -64,14 +67,21 @@ def test_extgrid_equals_generated_2x2_6(): subtile_slice_dx = cube_comm.partitioner.tile.subtile_slice( rank=cube_comm.rank, global_dims=[pace.util.Y_INTERFACE_DIM, pace.util.X_DIM], - global_extent=(npx, npy), + global_extent=(npy, nx), overlap=True, ) subtile_slice_dy = cube_comm.partitioner.tile.subtile_slice( rank=cube_comm.rank, global_dims=[pace.util.Y_DIM, pace.util.X_INTERFACE_DIM], - global_extent=(npx, npy), + global_extent=(ny, npx), + overlap=True, + ) + + subtile_slice_area = cube_comm.partitioner.tile.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.Y_DIM, pace.util.X_DIM], + global_extent=(ny, nx), overlap=True, ) @@ -80,6 +90,7 @@ def test_extgrid_equals_generated_2x2_6(): y=lat[subtile_slice_grid], dx=dx[subtile_slice_dx], dy=dy[subtile_slice_dy], + area=area[subtile_slice_area], quantity_factory=get_quantity_factory( layout=(2, 2), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz ), @@ -88,29 +99,36 @@ def test_extgrid_equals_generated_2x2_6(): extdgrid=True, ) - print(metric_terms_ext._dx) - errors = [] if ( - metric_terms_gen.grid.data[:, :, 0].any() + ext_driver.state.grid_data.lon.data.any() != metric_terms_ext.grid.data[:, :, 0].any() ): errors.append("lon data mismatch between generated and external grid data") if ( - metric_terms_gen.grid.data[:, :, 1].any() + ext_driver.state.grid_data.lat.data.any() != metric_terms_ext.grid.data[:, :, 1].any() ): - errors.append("lat data mismatch between generated and external grid data") + errors.append("lon data mismatch between generated and external grid data") - if metric_terms_gen._dx.view[:, :].any() != metric_terms_ext._dx.view[:, :].any(): + if ( + ext_driver.state.grid_data.dx.data.any() + != metric_terms_ext._dx.view[:, :].any() + ): errors.append("dx data mismatch between generated and external grid data") - if metric_terms_gen._dy.view[:, :].any() != metric_terms_ext._dy.view[:, :].any(): + if ( + ext_driver.state.grid_data.dy.data.any() + != metric_terms_ext._dy.view[:, :].any() + ): errors.append("dy data mismatch between generated and external grid data") - # if metric_terms_gen._area.data[:,:].any() != metric_terms_ext._area.any(): - # errors.append("area data mismatch between generated and external grid data") + if ( + ext_driver.state.grid_data.area.data.any() + != metric_terms_ext._area.view[:, :].any() + ): + errors.append("area data mismatch between generated and external grid data") - assert not errors, "errors occured in 1x1_6:\n{}".format("\n".join(errors)) + assert not errors, "errors occured in 2x2:\n{}".format("\n".join(errors)) diff --git a/util/pace/util/grid/generation.py b/util/pace/util/grid/generation.py index b494466c..14203dd4 100644 --- a/util/pace/util/grid/generation.py +++ b/util/pace/util/grid/generation.py @@ -1,7 +1,7 @@ import dataclasses import functools import warnings -from typing import Tuple +from typing import Optional, Tuple import numpy as np @@ -275,13 +275,13 @@ def __init__( # for the selected floating point precision self._agrid = None self._np = self._grid_64.np - self._dx = None - self._dy = None + self._dx: Optional[util.Quantity] = None + self._dy: Optional[util.Quantity] = None self._dx_agrid = None self._dy_agrid = None self._dx_center = None self._dy_center = None - self._area = None + self._area: Optional[util.Quantity] = None self._area_c = None self._ks = None self._ak = None @@ -428,6 +428,7 @@ def from_external( y, dx, dy, + area, quantity_factory, communicator, grid_type, @@ -440,9 +441,6 @@ def from_external( extdgrid=extdgrid, ) - # TODO: Using quantity factory to create quantities - # for dx and dy, quantity_factory.from_array() - rad_conv = PI / 180.0 terms._grid_64.view[:, :, 0] = np.dot(rad_conv, x) terms._grid_64.view[:, :, 1] = np.dot(rad_conv, y) @@ -453,7 +451,6 @@ def from_external( dtype=np.float64, allow_mismatch_float_precision=True, ) - dx_64.view[:, :] = dx terms._dx_64 = dx_64 @@ -469,8 +466,18 @@ def from_external( terms._comm.halo_update(terms._grid_64, n_points=terms._halo) terms._comm.vector_halo_update(terms._dx_64, terms._dy_64, n_points=terms._halo) - # terms._dx = quantity_cast_to_model_float(terms.quantity_factory, terms._dx_64) - # terms._dy = quantity_cast_to_model_float(terms.quantity_factory, terms._dy_64) + terms._dx = quantity_cast_to_model_float(terms.quantity_factory, terms._dx_64) + terms._dy = quantity_cast_to_model_float(terms.quantity_factory, terms._dy_64) + + area_64 = terms.quantity_factory.zeros( + [util.X_DIM, util.Y_DIM], + "m^2", + dtype=np.float64, + allow_mismatch_float_precision=True, + ) + area_64.view[:, :] = area + + terms._area = quantity_cast_to_model_float(terms.quantity_factory, area_64) terms._init_agrid() diff --git a/util/pace/util/partitioner.py b/util/pace/util/partitioner.py index c2560159..0e59ddfa 100644 --- a/util/pace/util/partitioner.py +++ b/util/pace/util/partitioner.py @@ -755,7 +755,7 @@ def subtile_index( ) -> Tuple[int, int]: within_tile_rank = rank % ranks_per_tile j = within_tile_rank // layout[1] - i = within_tile_rank % layout[0] + i = within_tile_rank % layout[1] return j, i From a536d956d47b5553df06bfd8d2c155028c15f229 Mon Sep 17 00:00:00 2001 From: Frank Malatino Date: Fri, 8 Dec 2023 13:58:30 -0500 Subject: [PATCH 14/31] Re-excluding external grid data yamls from test_example_configs.py --- tests/main/driver/test_example_configs.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/main/driver/test_example_configs.py b/tests/main/driver/test_example_configs.py index 1fc5dec1..d4e75355 100644 --- a/tests/main/driver/test_example_configs.py +++ b/tests/main/driver/test_example_configs.py @@ -28,6 +28,8 @@ "baroclinic_c12_orch_cpu.yaml", "tropical_read_restart_fortran.yml", "tropicalcyclone_c128.yaml", + "test_external_C12_1x1.yaml", + "test_external_C12_2x2.yaml", ] JENKINS_CONFIGS_DIR = os.path.join(dirname, "../../../.jenkins/driver_configs/") From 6b9ff5e2f854f8145de6fb9ba3cef6ccab36300b Mon Sep 17 00:00:00 2001 From: fmalatino <142349306+fmalatino@users.noreply.github.com> Date: Sat, 9 Dec 2023 09:37:34 -0500 Subject: [PATCH 15/31] Update driver/pace/driver/grid.py Co-authored-by: Florian Deconinck --- driver/pace/driver/grid.py | 1 + 1 file changed, 1 insertion(+) diff --git a/driver/pace/driver/grid.py b/driver/pace/driver/grid.py index 06bc9195..7ebff00f 100644 --- a/driver/pace/driver/grid.py +++ b/driver/pace/driver/grid.py @@ -224,6 +224,7 @@ def get_grid( pace_log.info("Using external grid data") + # ToDo: refactor when grid_type is an enum if self.grid_type <= 3: tile_num = ( pace.util.get_tile_index( From 8e0b82938f8c480da0b6cc73e5fe76825738f7c7 Mon Sep 17 00:00:00 2001 From: Fr4nk Date: Mon, 11 Dec 2023 13:54:14 -0500 Subject: [PATCH 16/31] Changed name of grid initializer function to match NetCDF dependency and class descriptor --- driver/pace/driver/grid.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/driver/pace/driver/grid.py b/driver/pace/driver/grid.py index 7ebff00f..cc5621bc 100644 --- a/driver/pace/driver/grid.py +++ b/driver/pace/driver/grid.py @@ -201,14 +201,14 @@ def get_grid( @GridInitializerSelector.register("external") @dataclasses.dataclass -class ExternalGridConfig(GridInitializer): +class ExternalNetcdfGridConfig(GridInitializer): """ Configuration for grid initialized from external data. Input is from tile files generated by FRE-NCtools methods - In its get_grid method it calls the MetricTerms class method - from_generated which generates an object of MetricTerms to be - used to generate the damping_coefficients, driver_grid_data, - and grid_data variables + The ExternalNetcdfGridConfig get_grid member method calls + the MetricTerms class method from_generated which generates + an object of MetricTerms to be used to generate the + damping_coefficients, driver_grid_data, and grid_data variables """ grid_type: Optional[int] = 0 From caa2d43f980bf81fa94f45cda0951bc424f1fe73 Mon Sep 17 00:00:00 2001 From: fmalatino <142349306+fmalatino@users.noreply.github.com> Date: Mon, 11 Dec 2023 15:53:05 -0500 Subject: [PATCH 17/31] Update util/pace/util/grid/generation.py Moved position of doc string for "from_external" MetricTerms class method Co-authored-by: Oliver Elbert --- util/pace/util/grid/generation.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/util/pace/util/grid/generation.py b/util/pace/util/grid/generation.py index 14203dd4..377cfbe4 100644 --- a/util/pace/util/grid/generation.py +++ b/util/pace/util/grid/generation.py @@ -416,11 +416,6 @@ def __init__( else: raise NotImplementedError(f"Unsupported grid_type = {grid_type}") - # from_generated class method - # Generates a metric terms object, which will perform - # the same function as the __init__ method for MetricTerms - # initializing the dgrid by input from data contained in an - # externally generated tile file @classmethod def from_external( cls, @@ -434,6 +429,10 @@ def from_external( grid_type, extdgrid: bool = True, ) -> "MetricTerms": + """ + Generates a metric terms object, using input from data contained in an + externally generated tile file + """ terms = MetricTerms( quantity_factory=quantity_factory, communicator=communicator, From 6690e6821050ca716099d46e7f2dea0ca2e1fdb5 Mon Sep 17 00:00:00 2001 From: Frank Malatino Date: Mon, 11 Dec 2023 22:40:17 -0500 Subject: [PATCH 18/31] Fixed indentation error in generation.py from suggestion in PR 42 --- util/pace/util/grid/generation.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/util/pace/util/grid/generation.py b/util/pace/util/grid/generation.py index 377cfbe4..304de1ff 100644 --- a/util/pace/util/grid/generation.py +++ b/util/pace/util/grid/generation.py @@ -429,10 +429,10 @@ def from_external( grid_type, extdgrid: bool = True, ) -> "MetricTerms": - """ - Generates a metric terms object, using input from data contained in an - externally generated tile file - """ + """ + Generates a metric terms object, using input from data contained in an + externally generated tile file + """ terms = MetricTerms( quantity_factory=quantity_factory, communicator=communicator, From ada497b8b0b1428b503a2ca63f8eb64c1223a35b Mon Sep 17 00:00:00 2001 From: Frank Malatino Date: Tue, 12 Dec 2023 10:28:45 -0500 Subject: [PATCH 19/31] Removal of TODO comment in grid.py, changes to method of file accessing in test_analytic_init, test_external_grid_* --- driver/pace/driver/grid.py | 2 -- tests/main/driver/test_analytic_init.py | 6 ++++-- tests/mpi_54rank/test_external_grid_1x1.py | 12 ++++++++++-- tests/mpi_54rank/test_external_grid_2x2.py | 12 ++++++++++-- 4 files changed, 24 insertions(+), 8 deletions(-) diff --git a/driver/pace/driver/grid.py b/driver/pace/driver/grid.py index cc5621bc..e73a849c 100644 --- a/driver/pace/driver/grid.py +++ b/driver/pace/driver/grid.py @@ -214,8 +214,6 @@ class ExternalNetcdfGridConfig(GridInitializer): grid_type: Optional[int] = 0 grid_file_path: str = "/input/tilefile" - # TODO: Area read in? - def get_grid( self, quantity_factory: QuantityFactory, diff --git a/tests/main/driver/test_analytic_init.py b/tests/main/driver/test_analytic_init.py index 03dda5bc..d13c8af2 100644 --- a/tests/main/driver/test_analytic_init.py +++ b/tests/main/driver/test_analytic_init.py @@ -7,8 +7,10 @@ import pace.driver +DIR = os.path.dirname(os.path.abspath(__file__)) + TESTED_CONFIGS: List[str] = [ - "driver/examples/configs/analytic_test.yaml", + "../../../driver/examples/configs/analytic_test.yaml", ] @@ -20,7 +22,7 @@ ) def test_analytic_init_config(tested_configs: List[str]): for config_file in tested_configs: - with open(os.path.abspath(config_file), "r") as f: + with open(os.path.join(DIR, config_file), "r") as f: config = yaml.safe_load(f) driver_config = pace.driver.DriverConfig.from_dict(config) assert driver_config.initialization.type == "analytic" diff --git a/tests/mpi_54rank/test_external_grid_1x1.py b/tests/mpi_54rank/test_external_grid_1x1.py index d069e1a2..af82a186 100644 --- a/tests/mpi_54rank/test_external_grid_1x1.py +++ b/tests/mpi_54rank/test_external_grid_1x1.py @@ -1,3 +1,5 @@ +import os + import numpy as np import xarray as xr import yaml @@ -8,6 +10,9 @@ from pace.util.mpi import MPIComm +DIR = os.path.dirname(os.path.abspath(__file__)) + + def get_quantity_factory(layout, nx_tile, ny_tile, nz): nx = nx_tile // layout[0] ny = ny_tile // layout[1] @@ -34,7 +39,10 @@ def get_tile_num(comm: MPIComm): def test_extgrid_equals_generated_1x1(): - with open("../../driver/examples/configs/test_external_C12_1x1.yaml", "r") as ext_f: + with open( + os.path.join(DIR, "../../driver/examples/configs/test_external_C12_2x2.yaml"), + "r", + ) as ext_f: ext_config = yaml.safe_load(ext_f) ext_driver_config = DriverConfig.from_dict(ext_config) @@ -46,7 +54,7 @@ def test_extgrid_equals_generated_1x1(): tile_num = get_tile_num(cube_comm) + 1 tile_file = "../../test_input/C12.tile" + str(tile_num) + ".nc" - ds = xr.open_dataset(tile_file) + ds = xr.open_dataset(os.path.join(DIR, tile_file)) lon = ds.x.values lat = ds.y.values dx = ds.dx.values diff --git a/tests/mpi_54rank/test_external_grid_2x2.py b/tests/mpi_54rank/test_external_grid_2x2.py index 98c8095f..5838cc5d 100644 --- a/tests/mpi_54rank/test_external_grid_2x2.py +++ b/tests/mpi_54rank/test_external_grid_2x2.py @@ -1,3 +1,5 @@ +import os + import numpy as np import xarray as xr import yaml @@ -8,6 +10,9 @@ from pace.util.mpi import MPIComm +DIR = os.path.dirname(os.path.abspath(__file__)) + + def get_quantity_factory(layout, nx_tile, ny_tile, nz): nx = nx_tile // layout[0] ny = ny_tile // layout[1] @@ -34,7 +39,10 @@ def get_tile_num(comm: MPIComm): def test_extgrid_equals_generated_2x2(): - with open("../../driver/examples/configs/test_external_C12_2x2.yaml", "r") as ext_f: + with open( + os.path.join(DIR, "../../driver/examples/configs/test_external_C12_2x2.yaml"), + "r", + ) as ext_f: ext_config = yaml.safe_load(ext_f) ext_driver_config = DriverConfig.from_dict(ext_config) @@ -46,7 +54,7 @@ def test_extgrid_equals_generated_2x2(): tile_num = get_tile_num(cube_comm) + 1 tile_file = "../../test_input/C12.tile" + str(tile_num) + ".nc" - ds = xr.open_dataset(tile_file) + ds = xr.open_dataset(os.path.join(DIR, tile_file)) lon = ds.x.values lat = ds.y.values dx = ds.dx.values From 6c1acfb6ec6e71d0bad90ae76e90152181a40a28 Mon Sep 17 00:00:00 2001 From: Frank Malatino Date: Tue, 12 Dec 2023 11:24:31 -0500 Subject: [PATCH 20/31] Changed grid data read-in unit tests to compare data directly from file to driver grid data generated from yaml --- tests/mpi_54rank/test_external_grid_1x1.py | 42 ++++------------------ tests/mpi_54rank/test_external_grid_2x2.py | 40 +++------------------ 2 files changed, 11 insertions(+), 71 deletions(-) diff --git a/tests/mpi_54rank/test_external_grid_1x1.py b/tests/mpi_54rank/test_external_grid_1x1.py index af82a186..a3e9432f 100644 --- a/tests/mpi_54rank/test_external_grid_1x1.py +++ b/tests/mpi_54rank/test_external_grid_1x1.py @@ -6,7 +6,6 @@ import pace.util from pace.driver import Driver, DriverConfig -from pace.util.grid import MetricTerms from pace.util.mpi import MPIComm @@ -40,7 +39,7 @@ def get_tile_num(comm: MPIComm): def test_extgrid_equals_generated_1x1(): with open( - os.path.join(DIR, "../../driver/examples/configs/test_external_C12_2x2.yaml"), + os.path.join(DIR, "../../driver/examples/configs/test_external_C12_1x1.yaml"), "r", ) as ext_f: ext_config = yaml.safe_load(ext_f) @@ -93,50 +92,21 @@ def test_extgrid_equals_generated_1x1(): overlap=True, ) - metric_terms_ext = MetricTerms.from_external( - x=lon[subtile_slice_grid], - y=lat[subtile_slice_grid], - dx=dx[subtile_slice_dx], - dy=dy[subtile_slice_dy], - area=area[subtile_slice_area], - quantity_factory=get_quantity_factory( - layout=(1, 1), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz - ), - communicator=cube_comm, - grid_type=0, - extdgrid=True, - ) - errors = [] - if ( - ext_driver.state.grid_data.lon.data.any() - != metric_terms_ext.grid.data[:, :, 0].any() - ): + if ext_driver.state.grid_data.lon.data.any() != lon[subtile_slice_grid].any(): errors.append("lon data mismatch between generated and external grid data") - if ( - ext_driver.state.grid_data.lat.data.any() - != metric_terms_ext.grid.data[:, :, 1].any() - ): + if ext_driver.state.grid_data.lat.data.any() != lat[subtile_slice_grid].any(): errors.append("lon data mismatch between generated and external grid data") - if ( - ext_driver.state.grid_data.dx.data.any() - != metric_terms_ext._dx.view[:, :].any() - ): + if ext_driver.state.grid_data.dx.data.any() != dx[subtile_slice_dx].any(): errors.append("dx data mismatch between generated and external grid data") - if ( - ext_driver.state.grid_data.dy.data.any() - != metric_terms_ext._dy.view[:, :].any() - ): + if ext_driver.state.grid_data.dy.data.any() != dy[subtile_slice_dy].any(): errors.append("dy data mismatch between generated and external grid data") - if ( - ext_driver.state.grid_data.area.data.any() - != metric_terms_ext._area.view[:, :].any() - ): + if ext_driver.state.grid_data.area.data.any() != area[subtile_slice_area].any(): errors.append("area data mismatch between generated and external grid data") assert not errors, "errors occured in 1x1:\n{}".format("\n".join(errors)) diff --git a/tests/mpi_54rank/test_external_grid_2x2.py b/tests/mpi_54rank/test_external_grid_2x2.py index 5838cc5d..049c4557 100644 --- a/tests/mpi_54rank/test_external_grid_2x2.py +++ b/tests/mpi_54rank/test_external_grid_2x2.py @@ -6,7 +6,6 @@ import pace.util from pace.driver import Driver, DriverConfig -from pace.util.grid import MetricTerms from pace.util.mpi import MPIComm @@ -93,50 +92,21 @@ def test_extgrid_equals_generated_2x2(): overlap=True, ) - metric_terms_ext = MetricTerms.from_external( - x=lon[subtile_slice_grid], - y=lat[subtile_slice_grid], - dx=dx[subtile_slice_dx], - dy=dy[subtile_slice_dy], - area=area[subtile_slice_area], - quantity_factory=get_quantity_factory( - layout=(2, 2), nx_tile=nx_tile, ny_tile=ny_tile, nz=nz - ), - communicator=cube_comm, - grid_type=0, - extdgrid=True, - ) - errors = [] - if ( - ext_driver.state.grid_data.lon.data.any() - != metric_terms_ext.grid.data[:, :, 0].any() - ): + if ext_driver.state.grid_data.lon.data.any() != lon[subtile_slice_grid].any(): errors.append("lon data mismatch between generated and external grid data") - if ( - ext_driver.state.grid_data.lat.data.any() - != metric_terms_ext.grid.data[:, :, 1].any() - ): + if ext_driver.state.grid_data.lat.data.any() != lat[subtile_slice_grid].any(): errors.append("lon data mismatch between generated and external grid data") - if ( - ext_driver.state.grid_data.dx.data.any() - != metric_terms_ext._dx.view[:, :].any() - ): + if ext_driver.state.grid_data.dx.data.any() != dx[subtile_slice_dx].any(): errors.append("dx data mismatch between generated and external grid data") - if ( - ext_driver.state.grid_data.dy.data.any() - != metric_terms_ext._dy.view[:, :].any() - ): + if ext_driver.state.grid_data.dy.data.any() != dy[subtile_slice_dy].any(): errors.append("dy data mismatch between generated and external grid data") - if ( - ext_driver.state.grid_data.area.data.any() - != metric_terms_ext._area.view[:, :].any() - ): + if ext_driver.state.grid_data.area.data.any() != area[subtile_slice_area].any(): errors.append("area data mismatch between generated and external grid data") assert not errors, "errors occured in 2x2:\n{}".format("\n".join(errors)) From d48f2babc993084d1c030b07df40f927b6721dfe Mon Sep 17 00:00:00 2001 From: Frank Malatino Date: Tue, 12 Dec 2023 14:18:58 -0500 Subject: [PATCH 21/31] Change to reading in lon and lat, other metric terms calculated as needed --- driver/pace/driver/grid.py | 29 ----------------- tests/mpi_54rank/test_external_grid_1x1.py | 1 - tests/mpi_54rank/test_external_grid_2x2.py | 1 - util/pace/util/grid/generation.py | 37 ---------------------- 4 files changed, 68 deletions(-) diff --git a/driver/pace/driver/grid.py b/driver/pace/driver/grid.py index e73a849c..79cda613 100644 --- a/driver/pace/driver/grid.py +++ b/driver/pace/driver/grid.py @@ -237,11 +237,6 @@ def get_grid( ds = xr.open_dataset(tile_file) lon = ds.x.values lat = ds.y.values - dx = ds.dx.values - dy = ds.dy.values - area = ds.area.values - nx = ds.nx.values.size - ny = ds.ny.values.size npx = ds.nxp.values.size npy = ds.nyp.values.size @@ -252,33 +247,9 @@ def get_grid( overlap=True, ) - subtile_slice_dx = communicator.partitioner.tile.subtile_slice( - rank=communicator.rank, - global_dims=[pace.util.Y_INTERFACE_DIM, pace.util.X_DIM], - global_extent=(npy, nx), - overlap=True, - ) - - subtile_slice_dy = communicator.partitioner.tile.subtile_slice( - rank=communicator.rank, - global_dims=[pace.util.Y_DIM, pace.util.X_INTERFACE_DIM], - global_extent=(ny, npx), - overlap=True, - ) - - subtile_slice_area = communicator.partitioner.tile.subtile_slice( - rank=communicator.rank, - global_dims=[pace.util.Y_DIM, pace.util.X_DIM], - global_extent=(ny, nx), - overlap=True, - ) - metric_terms = MetricTerms.from_external( x=lon[subtile_slice_grid], y=lat[subtile_slice_grid], - dx=dx[subtile_slice_dx], - dy=dy[subtile_slice_dy], - area=area[subtile_slice_area], quantity_factory=quantity_factory, communicator=communicator, grid_type=self.grid_type, diff --git a/tests/mpi_54rank/test_external_grid_1x1.py b/tests/mpi_54rank/test_external_grid_1x1.py index a3e9432f..54959336 100644 --- a/tests/mpi_54rank/test_external_grid_1x1.py +++ b/tests/mpi_54rank/test_external_grid_1x1.py @@ -47,7 +47,6 @@ def test_extgrid_equals_generated_1x1(): ext_driver = Driver(ext_driver_config) - nx_tile, ny_tile, nz = 12, 12, 5 comm_1by1 = MPIComm() cube_comm = get_cube_comm(layout=(1, 1), comm=comm_1by1) diff --git a/tests/mpi_54rank/test_external_grid_2x2.py b/tests/mpi_54rank/test_external_grid_2x2.py index 049c4557..a926b8c2 100644 --- a/tests/mpi_54rank/test_external_grid_2x2.py +++ b/tests/mpi_54rank/test_external_grid_2x2.py @@ -47,7 +47,6 @@ def test_extgrid_equals_generated_2x2(): ext_driver = Driver(ext_driver_config) - nx_tile, ny_tile, nz = 12, 12, 5 comm_2by2 = MPIComm() cube_comm = get_cube_comm(layout=(2, 2), comm=comm_2by2) diff --git a/util/pace/util/grid/generation.py b/util/pace/util/grid/generation.py index 304de1ff..f56c63dc 100644 --- a/util/pace/util/grid/generation.py +++ b/util/pace/util/grid/generation.py @@ -421,9 +421,6 @@ def from_external( cls, x, y, - dx, - dy, - area, quantity_factory, communicator, grid_type, @@ -444,40 +441,6 @@ def from_external( terms._grid_64.view[:, :, 0] = np.dot(rad_conv, x) terms._grid_64.view[:, :, 1] = np.dot(rad_conv, y) - dx_64 = terms.quantity_factory.zeros( - [util.Y_INTERFACE_DIM, util.X_DIM], - "m", - dtype=np.float64, - allow_mismatch_float_precision=True, - ) - dx_64.view[:, :] = dx - terms._dx_64 = dx_64 - - dy_64 = terms.quantity_factory.zeros( - [util.Y_DIM, util.X_INTERFACE_DIM], - "m", - dtype=np.float64, - allow_mismatch_float_precision=True, - ) - dy_64.view[:, :] = dy - terms._dy_64 = dy_64 - - terms._comm.halo_update(terms._grid_64, n_points=terms._halo) - terms._comm.vector_halo_update(terms._dx_64, terms._dy_64, n_points=terms._halo) - - terms._dx = quantity_cast_to_model_float(terms.quantity_factory, terms._dx_64) - terms._dy = quantity_cast_to_model_float(terms.quantity_factory, terms._dy_64) - - area_64 = terms.quantity_factory.zeros( - [util.X_DIM, util.Y_DIM], - "m^2", - dtype=np.float64, - allow_mismatch_float_precision=True, - ) - area_64.view[:, :] = area - - terms._area = quantity_cast_to_model_float(terms.quantity_factory, area_64) - terms._init_agrid() return terms From 136db1a400f7fb6cd556af04b474fa86090451c3 Mon Sep 17 00:00:00 2001 From: Frank Malatino Date: Wed, 13 Dec 2023 22:54:55 -0500 Subject: [PATCH 22/31] Removed read-in of dx, dy, and area. Changed unit tests to compare calculated area to 'ideal' surface area as given by selected constants type. --- driver/pace/driver/grid.py | 3 + tests/mpi_54rank/test_external_grid_1x1.py | 97 +++------------------- tests/mpi_54rank/test_external_grid_2x2.py | 95 +++------------------ 3 files changed, 24 insertions(+), 171 deletions(-) diff --git a/driver/pace/driver/grid.py b/driver/pace/driver/grid.py index 79cda613..8a1e7a84 100644 --- a/driver/pace/driver/grid.py +++ b/driver/pace/driver/grid.py @@ -209,6 +209,9 @@ class ExternalNetcdfGridConfig(GridInitializer): the MetricTerms class method from_generated which generates an object of MetricTerms to be used to generate the damping_coefficients, driver_grid_data, and grid_data variables + It is imperative that the user ensures that the constants used + to generate the input grid data matches that of the constants + used in Pace to develop accurate metric terms. """ grid_type: Optional[int] = 0 diff --git a/tests/mpi_54rank/test_external_grid_1x1.py b/tests/mpi_54rank/test_external_grid_1x1.py index 54959336..a40d0fde 100644 --- a/tests/mpi_54rank/test_external_grid_1x1.py +++ b/tests/mpi_54rank/test_external_grid_1x1.py @@ -1,42 +1,18 @@ +import math import os import numpy as np -import xarray as xr import yaml -import pace.util from pace.driver import Driver, DriverConfig +from pace.util.constants import PI, RADIUS from pace.util.mpi import MPIComm DIR = os.path.dirname(os.path.abspath(__file__)) -def get_quantity_factory(layout, nx_tile, ny_tile, nz): - nx = nx_tile // layout[0] - ny = ny_tile // layout[1] - return pace.util.QuantityFactory( - sizer=pace.util.SubtileGridSizer( - nx=nx, ny=ny, nz=nz, n_halo=3, extra_dim_lengths={} - ), - numpy=np, - ) - - -def get_cube_comm(layout, comm: MPIComm): - return pace.util.CubedSphereCommunicator( - comm=comm, - partitioner=pace.util.CubedSpherePartitioner( - pace.util.TilePartitioner(layout=layout) - ), - ) - - -def get_tile_num(comm: MPIComm): - return pace.util.get_tile_index(comm.rank, comm.partitioner.total_ranks) - - -def test_extgrid_equals_generated_1x1(): +def test_extgrid_equals_generated_2x2(): with open( os.path.join(DIR, "../../driver/examples/configs/test_external_C12_1x1.yaml"), @@ -47,65 +23,14 @@ def test_extgrid_equals_generated_1x1(): ext_driver = Driver(ext_driver_config) - comm_1by1 = MPIComm() - cube_comm = get_cube_comm(layout=(1, 1), comm=comm_1by1) - - tile_num = get_tile_num(cube_comm) + 1 - tile_file = "../../test_input/C12.tile" + str(tile_num) + ".nc" - ds = xr.open_dataset(os.path.join(DIR, tile_file)) - lon = ds.x.values - lat = ds.y.values - dx = ds.dx.values - dy = ds.dy.values - area = ds.area.values - nx = ds.nx.values.size - ny = ds.ny.values.size - npx = ds.nxp.values.size - npy = ds.nyp.values.size - - subtile_slice_grid = cube_comm.partitioner.tile.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.Y_INTERFACE_DIM, pace.util.X_INTERFACE_DIM], - global_extent=(npx, npy), - overlap=True, - ) - - subtile_slice_dx = cube_comm.partitioner.tile.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.Y_INTERFACE_DIM, pace.util.X_DIM], - global_extent=(npy, nx), - overlap=True, - ) - - subtile_slice_dy = cube_comm.partitioner.tile.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.Y_DIM, pace.util.X_INTERFACE_DIM], - global_extent=(ny, npx), - overlap=True, - ) - - subtile_slice_area = cube_comm.partitioner.tile.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.Y_DIM, pace.util.X_DIM], - global_extent=(ny, nx), - overlap=True, - ) - - errors = [] - - if ext_driver.state.grid_data.lon.data.any() != lon[subtile_slice_grid].any(): - errors.append("lon data mismatch between generated and external grid data") - - if ext_driver.state.grid_data.lat.data.any() != lat[subtile_slice_grid].any(): - errors.append("lon data mismatch between generated and external grid data") - - if ext_driver.state.grid_data.dx.data.any() != dx[subtile_slice_dx].any(): - errors.append("dx data mismatch between generated and external grid data") + surface_area_true = 4 * PI * (RADIUS ** 2) - if ext_driver.state.grid_data.dy.data.any() != dy[subtile_slice_dy].any(): - errors.append("dy data mismatch between generated and external grid data") + mpicomm = MPIComm() - if ext_driver.state.grid_data.area.data.any() != area[subtile_slice_area].any(): - errors.append("area data mismatch between generated and external grid data") + tmp = [math.fsum(row) for row in ext_driver.state.grid_data.area.view[:, :]] + rank_area_sum = math.fsum(tmp) + tile_area = mpicomm._comm.gather(rank_area_sum, root=0) - assert not errors, "errors occured in 1x1:\n{}".format("\n".join(errors)) + if mpicomm.Get_rank() == 0: + total_area = math.fsum(tile_area) + assert np.isclose(total_area, surface_area_true) diff --git a/tests/mpi_54rank/test_external_grid_2x2.py b/tests/mpi_54rank/test_external_grid_2x2.py index a926b8c2..311998eb 100644 --- a/tests/mpi_54rank/test_external_grid_2x2.py +++ b/tests/mpi_54rank/test_external_grid_2x2.py @@ -1,41 +1,17 @@ +import math import os import numpy as np -import xarray as xr import yaml -import pace.util from pace.driver import Driver, DriverConfig +from pace.util.constants import PI, RADIUS from pace.util.mpi import MPIComm DIR = os.path.dirname(os.path.abspath(__file__)) -def get_quantity_factory(layout, nx_tile, ny_tile, nz): - nx = nx_tile // layout[0] - ny = ny_tile // layout[1] - return pace.util.QuantityFactory( - sizer=pace.util.SubtileGridSizer( - nx=nx, ny=ny, nz=nz, n_halo=3, extra_dim_lengths={} - ), - numpy=np, - ) - - -def get_cube_comm(layout, comm: MPIComm): - return pace.util.CubedSphereCommunicator( - comm=comm, - partitioner=pace.util.CubedSpherePartitioner( - pace.util.TilePartitioner(layout=layout) - ), - ) - - -def get_tile_num(comm: MPIComm): - return pace.util.get_tile_index(comm.rank, comm.partitioner.total_ranks) - - def test_extgrid_equals_generated_2x2(): with open( @@ -47,65 +23,14 @@ def test_extgrid_equals_generated_2x2(): ext_driver = Driver(ext_driver_config) - comm_2by2 = MPIComm() - cube_comm = get_cube_comm(layout=(2, 2), comm=comm_2by2) - - tile_num = get_tile_num(cube_comm) + 1 - tile_file = "../../test_input/C12.tile" + str(tile_num) + ".nc" - ds = xr.open_dataset(os.path.join(DIR, tile_file)) - lon = ds.x.values - lat = ds.y.values - dx = ds.dx.values - dy = ds.dy.values - area = ds.area.values - nx = ds.nx.values.size - ny = ds.ny.values.size - npx = ds.nxp.values.size - npy = ds.nyp.values.size - - subtile_slice_grid = cube_comm.partitioner.tile.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.Y_INTERFACE_DIM, pace.util.X_INTERFACE_DIM], - global_extent=(npx, npy), - overlap=True, - ) - - subtile_slice_dx = cube_comm.partitioner.tile.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.Y_INTERFACE_DIM, pace.util.X_DIM], - global_extent=(npy, nx), - overlap=True, - ) - - subtile_slice_dy = cube_comm.partitioner.tile.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.Y_DIM, pace.util.X_INTERFACE_DIM], - global_extent=(ny, npx), - overlap=True, - ) - - subtile_slice_area = cube_comm.partitioner.tile.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.Y_DIM, pace.util.X_DIM], - global_extent=(ny, nx), - overlap=True, - ) - - errors = [] - - if ext_driver.state.grid_data.lon.data.any() != lon[subtile_slice_grid].any(): - errors.append("lon data mismatch between generated and external grid data") - - if ext_driver.state.grid_data.lat.data.any() != lat[subtile_slice_grid].any(): - errors.append("lon data mismatch between generated and external grid data") - - if ext_driver.state.grid_data.dx.data.any() != dx[subtile_slice_dx].any(): - errors.append("dx data mismatch between generated and external grid data") + surface_area_true = 4 * PI * (RADIUS ** 2) - if ext_driver.state.grid_data.dy.data.any() != dy[subtile_slice_dy].any(): - errors.append("dy data mismatch between generated and external grid data") + mpicomm = MPIComm() - if ext_driver.state.grid_data.area.data.any() != area[subtile_slice_area].any(): - errors.append("area data mismatch between generated and external grid data") + tmp = [math.fsum(row) for row in ext_driver.state.grid_data.area.view[:, :]] + rank_area_sum = math.fsum(tmp) + tile_area = mpicomm._comm.gather(rank_area_sum, root=0) - assert not errors, "errors occured in 2x2:\n{}".format("\n".join(errors)) + if mpicomm.Get_rank() == 0: + total_area = math.fsum(tile_area) + assert np.isclose(total_area, surface_area_true) From 0be34399f50f6e125e3b7a07e4ab63c6ce232f7d Mon Sep 17 00:00:00 2001 From: fmalatino <142349306+fmalatino@users.noreply.github.com> Date: Thu, 14 Dec 2023 11:00:30 -0500 Subject: [PATCH 23/31] Update tests/mpi_54rank/test_external_grid_1x1.py Incorrect name of test in test_external_grid_1x1.py changed to match file name Co-authored-by: Oliver Elbert --- tests/mpi_54rank/test_external_grid_1x1.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/mpi_54rank/test_external_grid_1x1.py b/tests/mpi_54rank/test_external_grid_1x1.py index a40d0fde..616f865f 100644 --- a/tests/mpi_54rank/test_external_grid_1x1.py +++ b/tests/mpi_54rank/test_external_grid_1x1.py @@ -12,7 +12,7 @@ DIR = os.path.dirname(os.path.abspath(__file__)) -def test_extgrid_equals_generated_2x2(): +def test_extgrid_equals_generated_1x1(): with open( os.path.join(DIR, "../../driver/examples/configs/test_external_C12_1x1.yaml"), From 2d8eb41ad1867348003d7ea85ce5001f09822d14 Mon Sep 17 00:00:00 2001 From: Frank Malatino Date: Thu, 14 Dec 2023 16:21:35 -0500 Subject: [PATCH 24/31] Added comparisons for read-in vs generated by driver lon, lat, dx, dy, and area data to unit tests --- driver/pace/driver/grid.py | 11 ++- tests/mpi_54rank/test_external_grid_1x1.py | 91 +++++++++++++++++++++ tests/mpi_54rank/test_external_grid_2x2.py | 94 ++++++++++++++++++++++ util/pace/util/grid/generation.py | 15 ++-- 4 files changed, 197 insertions(+), 14 deletions(-) diff --git a/driver/pace/driver/grid.py b/driver/pace/driver/grid.py index 8a1e7a84..83d68888 100644 --- a/driver/pace/driver/grid.py +++ b/driver/pace/driver/grid.py @@ -209,9 +209,9 @@ class ExternalNetcdfGridConfig(GridInitializer): the MetricTerms class method from_generated which generates an object of MetricTerms to be used to generate the damping_coefficients, driver_grid_data, and grid_data variables - It is imperative that the user ensures that the constants used - to generate the input grid data matches that of the constants - used in Pace to develop accurate metric terms. + We do not read in the dx, dy, or area values as there may be + inconsistencies in the constants used during calculation of the + input data and the model use. """ grid_type: Optional[int] = 0 @@ -245,8 +245,8 @@ def get_grid( subtile_slice_grid = communicator.partitioner.tile.subtile_slice( rank=communicator.rank, - global_dims=[pace.util.X_INTERFACE_DIM, pace.util.Y_INTERFACE_DIM], - global_extent=(npx, npy), + global_dims=[pace.util.Y_INTERFACE_DIM, pace.util.X_INTERFACE_DIM], + global_extent=(npy, npx), overlap=True, ) @@ -256,7 +256,6 @@ def get_grid( quantity_factory=quantity_factory, communicator=communicator, grid_type=self.grid_type, - extdgrid=True, ) horizontal_data = HorizontalGridData.new_from_metric_terms(metric_terms) diff --git a/tests/mpi_54rank/test_external_grid_1x1.py b/tests/mpi_54rank/test_external_grid_1x1.py index 616f865f..6a9efacd 100644 --- a/tests/mpi_54rank/test_external_grid_1x1.py +++ b/tests/mpi_54rank/test_external_grid_1x1.py @@ -2,8 +2,10 @@ import os import numpy as np +import xarray as xr import yaml +import pace.util from pace.driver import Driver, DriverConfig from pace.util.constants import PI, RADIUS from pace.util.mpi import MPIComm @@ -12,6 +14,19 @@ DIR = os.path.dirname(os.path.abspath(__file__)) +def get_cube_comm(layout, comm: MPIComm): + return pace.util.CubedSphereCommunicator( + comm=comm, + partitioner=pace.util.CubedSpherePartitioner( + pace.util.TilePartitioner(layout=layout) + ), + ) + + +def get_tile_num(comm: MPIComm): + return pace.util.get_tile_index(comm.rank, comm.partitioner.total_ranks) + + def test_extgrid_equals_generated_1x1(): with open( @@ -23,6 +38,82 @@ def test_extgrid_equals_generated_1x1(): ext_driver = Driver(ext_driver_config) + comm_1by1 = MPIComm() + cube_comm = get_cube_comm(layout=(1, 1), comm=comm_1by1) + + tile_num = get_tile_num(cube_comm) + 1 + tile_file = "../../test_input/C12.tile" + str(tile_num) + ".nc" + ds = xr.open_dataset(os.path.join(DIR, tile_file)) + lon = ds.x.values + lat = ds.y.values + dx = ds.dx.values + dy = ds.dy.values + area = ds.area.values + nx = ds.nx.values.size + ny = ds.ny.values.size + npx = ds.nxp.values.size + npy = ds.nyp.values.size + + subtile_slice_grid = cube_comm.partitioner.tile.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.Y_INTERFACE_DIM, pace.util.X_INTERFACE_DIM], + global_extent=(npy, npx), + overlap=True, + ) + + subtile_slice_dx = cube_comm.partitioner.tile.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.Y_INTERFACE_DIM, pace.util.X_DIM], + global_extent=(npy, nx), + overlap=True, + ) + + subtile_slice_dy = cube_comm.partitioner.tile.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.Y_DIM, pace.util.X_INTERFACE_DIM], + global_extent=(ny, npx), + overlap=True, + ) + + subtile_slice_area = cube_comm.partitioner.tile.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.Y_DIM, pace.util.X_DIM], + global_extent=(ny, nx), + overlap=True, + ) + + lon_rad = lon * (PI / 180) + lat_rad = lat * (PI / 180) + + errors = [] + + if not np.isclose( + ext_driver.state.grid_data.lon.view[:, :], lon_rad[subtile_slice_grid] + ).all(): + errors.append("Lon data mismatch") + + if not np.isclose( + ext_driver.state.grid_data.lat.view[:, :], lat_rad[subtile_slice_grid] + ).all(): + errors.append("Lat data mismatch") + + if not np.isclose( + ext_driver.state.grid_data.dy.view[:, :], dx[subtile_slice_dx] + ).all(): + errors.append("dx data mismatch") + + if not np.isclose( + ext_driver.state.grid_data.dx.view[:, :], dy[subtile_slice_dy] + ).all(): + errors.append("dy data mismatch") + + if not np.isclose( + ext_driver.state.grid_data.area.view[:, :], area[subtile_slice_area] + ).all(): + errors.append("area data mismatch") + + assert not errors, "errors occured in 2x2:\n{}".format("\n".join(errors)) + surface_area_true = 4 * PI * (RADIUS ** 2) mpicomm = MPIComm() diff --git a/tests/mpi_54rank/test_external_grid_2x2.py b/tests/mpi_54rank/test_external_grid_2x2.py index 311998eb..3b201d35 100644 --- a/tests/mpi_54rank/test_external_grid_2x2.py +++ b/tests/mpi_54rank/test_external_grid_2x2.py @@ -2,8 +2,10 @@ import os import numpy as np +import xarray as xr import yaml +import pace.util from pace.driver import Driver, DriverConfig from pace.util.constants import PI, RADIUS from pace.util.mpi import MPIComm @@ -12,6 +14,19 @@ DIR = os.path.dirname(os.path.abspath(__file__)) +def get_cube_comm(layout, comm: MPIComm): + return pace.util.CubedSphereCommunicator( + comm=comm, + partitioner=pace.util.CubedSpherePartitioner( + pace.util.TilePartitioner(layout=layout) + ), + ) + + +def get_tile_num(comm: MPIComm): + return pace.util.get_tile_index(comm.rank, comm.partitioner.total_ranks) + + def test_extgrid_equals_generated_2x2(): with open( @@ -23,6 +38,82 @@ def test_extgrid_equals_generated_2x2(): ext_driver = Driver(ext_driver_config) + comm_2by2 = MPIComm() + cube_comm = get_cube_comm(layout=(2, 2), comm=comm_2by2) + + tile_num = get_tile_num(cube_comm) + 1 + tile_file = "../../test_input/C12.tile" + str(tile_num) + ".nc" + ds = xr.open_dataset(os.path.join(DIR, tile_file)) + lon = ds.x.values + lat = ds.y.values + dx = ds.dx.values + dy = ds.dy.values + area = ds.area.values + nx = ds.nx.values.size + ny = ds.ny.values.size + npx = ds.nxp.values.size + npy = ds.nyp.values.size + + subtile_slice_grid = cube_comm.partitioner.tile.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.Y_INTERFACE_DIM, pace.util.X_INTERFACE_DIM], + global_extent=(npy, npx), + overlap=True, + ) + + subtile_slice_dx = cube_comm.partitioner.tile.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.Y_INTERFACE_DIM, pace.util.X_DIM], + global_extent=(npy, nx), + overlap=True, + ) + + subtile_slice_dy = cube_comm.partitioner.tile.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.Y_DIM, pace.util.X_INTERFACE_DIM], + global_extent=(ny, npx), + overlap=True, + ) + + subtile_slice_area = cube_comm.partitioner.tile.subtile_slice( + rank=cube_comm.rank, + global_dims=[pace.util.Y_DIM, pace.util.X_DIM], + global_extent=(ny, nx), + overlap=True, + ) + + lon_rad = lon * (PI / 180) + lat_rad = lat * (PI / 180) + + errors = [] + + if not np.isclose( + ext_driver.state.grid_data.lon.view[:, :], lon_rad[subtile_slice_grid] + ).all(): + errors.append("Lon data mismatch") + + if not np.isclose( + ext_driver.state.grid_data.lat.view[:, :], lat_rad[subtile_slice_grid] + ).all(): + errors.append("Lat data mismatch") + + if not np.isclose( + ext_driver.state.grid_data.dy.view[:, :], dx[subtile_slice_dx] + ).all(): + errors.append("dx data mismatch") + + if not np.isclose( + ext_driver.state.grid_data.dx.view[:, :], dy[subtile_slice_dy] + ).all(): + errors.append("dy data mismatch") + + if not np.isclose( + ext_driver.state.grid_data.area.view[:, :], area[subtile_slice_area] + ).all(): + errors.append("area data mismatch") + + assert not errors, "errors occured in 2x2:\n{}".format("\n".join(errors)) + surface_area_true = 4 * PI * (RADIUS ** 2) mpicomm = MPIComm() @@ -32,5 +123,8 @@ def test_extgrid_equals_generated_2x2(): tile_area = mpicomm._comm.gather(rank_area_sum, root=0) if mpicomm.Get_rank() == 0: + # print(ext_driver.state.grid_data.dy.view[:, :]) + # print("") + # print(dx[subtile_slice_dx]) total_area = math.fsum(tile_area) assert np.isclose(total_area, surface_area_true) diff --git a/util/pace/util/grid/generation.py b/util/pace/util/grid/generation.py index f56c63dc..3e81a056 100644 --- a/util/pace/util/grid/generation.py +++ b/util/pace/util/grid/generation.py @@ -1,7 +1,7 @@ import dataclasses import functools import warnings -from typing import Optional, Tuple +from typing import Tuple import numpy as np @@ -275,13 +275,13 @@ def __init__( # for the selected floating point precision self._agrid = None self._np = self._grid_64.np - self._dx: Optional[util.Quantity] = None - self._dy: Optional[util.Quantity] = None + self._dx = None + self._dy = None self._dx_agrid = None self._dy_agrid = None self._dx_center = None self._dy_center = None - self._area: Optional[util.Quantity] = None + self._area = None self._area_c = None self._ks = None self._ak = None @@ -424,7 +424,6 @@ def from_external( quantity_factory, communicator, grid_type, - extdgrid: bool = True, ) -> "MetricTerms": """ Generates a metric terms object, using input from data contained in an @@ -434,12 +433,12 @@ def from_external( quantity_factory=quantity_factory, communicator=communicator, grid_type=grid_type, - extdgrid=extdgrid, + extdgrid=True, ) rad_conv = PI / 180.0 - terms._grid_64.view[:, :, 0] = np.dot(rad_conv, x) - terms._grid_64.view[:, :, 1] = np.dot(rad_conv, y) + terms._grid_64.view[:, :, 0] = rad_conv * x + terms._grid_64.view[:, :, 1] = rad_conv * y terms._init_agrid() From b7f379f2b0544d25abb6d82891e78f0b57ba00cd Mon Sep 17 00:00:00 2001 From: Frank Malatino Date: Mon, 18 Dec 2023 10:41:16 -0500 Subject: [PATCH 25/31] Added relative error calculations to unit tests for external grid data read-in --- tests/mpi_54rank/test_external_grid_1x1.py | 29 +++++++++++++++++++++- tests/mpi_54rank/test_external_grid_2x2.py | 26 +++++++++++++++++++ 2 files changed, 54 insertions(+), 1 deletion(-) diff --git a/tests/mpi_54rank/test_external_grid_1x1.py b/tests/mpi_54rank/test_external_grid_1x1.py index 6a9efacd..b6d951cf 100644 --- a/tests/mpi_54rank/test_external_grid_1x1.py +++ b/tests/mpi_54rank/test_external_grid_1x1.py @@ -86,33 +86,59 @@ def test_extgrid_equals_generated_1x1(): lat_rad = lat * (PI / 180) errors = [] + diffs = [] if not np.isclose( ext_driver.state.grid_data.lon.view[:, :], lon_rad[subtile_slice_grid] ).all(): errors.append("Lon data mismatch") + diff_lon = np.amax( + ext_driver.state.grid_data.lon.view[:, :] - lon_rad[subtile_slice_grid] + ) / np.amax(lon_rad[subtile_slice_grid]) + diffs.append(f"Lon maximum relative error = {diff_lon}") + if not np.isclose( ext_driver.state.grid_data.lat.view[:, :], lat_rad[subtile_slice_grid] ).all(): errors.append("Lat data mismatch") + diff_lat = np.amax( + ext_driver.state.grid_data.lat.view[:, :] - lat_rad[subtile_slice_grid] + ) / np.amax(lat_rad[subtile_slice_grid]) + diffs.append(f"Lat maximum relative error = {diff_lat}") + if not np.isclose( ext_driver.state.grid_data.dy.view[:, :], dx[subtile_slice_dx] ).all(): errors.append("dx data mismatch") + diff_dx = np.amax( + ext_driver.state.grid_data.dy.view[:, :] - dx[subtile_slice_dx] + ) / np.amax(dx[subtile_slice_dx]) + diffs.append(f"dx maximum relative error = {diff_dx}") + if not np.isclose( ext_driver.state.grid_data.dx.view[:, :], dy[subtile_slice_dy] ).all(): errors.append("dy data mismatch") + diff_dy = np.amax( + ext_driver.state.grid_data.dx.view[:, :] - dy[subtile_slice_dy] + ) / np.amax(dy[subtile_slice_dy]) + diffs.append(f"dy maximum relative error = {diff_dy}") + if not np.isclose( ext_driver.state.grid_data.area.view[:, :], area[subtile_slice_area] ).all(): errors.append("area data mismatch") - assert not errors, "errors occured in 2x2:\n{}".format("\n".join(errors)) + diff_area = np.amax( + ext_driver.state.grid_data.area.view[:, :] - area[subtile_slice_area] + ) / np.amax(area[subtile_slice_area]) + diffs.append(f"Area maximum relative error = {diff_area}") + + assert not errors, "errors occured in 1x1:\n{}".format("\n".join(errors)) surface_area_true = 4 * PI * (RADIUS ** 2) @@ -124,4 +150,5 @@ def test_extgrid_equals_generated_1x1(): if mpicomm.Get_rank() == 0: total_area = math.fsum(tile_area) + print(diffs) assert np.isclose(total_area, surface_area_true) diff --git a/tests/mpi_54rank/test_external_grid_2x2.py b/tests/mpi_54rank/test_external_grid_2x2.py index 3b201d35..8ea8f1ff 100644 --- a/tests/mpi_54rank/test_external_grid_2x2.py +++ b/tests/mpi_54rank/test_external_grid_2x2.py @@ -86,32 +86,58 @@ def test_extgrid_equals_generated_2x2(): lat_rad = lat * (PI / 180) errors = [] + diffs = [] if not np.isclose( ext_driver.state.grid_data.lon.view[:, :], lon_rad[subtile_slice_grid] ).all(): errors.append("Lon data mismatch") + diff_lon = np.amax( + ext_driver.state.grid_data.lon.view[:, :] - lon_rad[subtile_slice_grid] + ) / np.amax(lon_rad[subtile_slice_grid]) + diffs.append(f"Lon maximum relative error = {diff_lon}") + if not np.isclose( ext_driver.state.grid_data.lat.view[:, :], lat_rad[subtile_slice_grid] ).all(): errors.append("Lat data mismatch") + diff_lat = np.amax( + ext_driver.state.grid_data.lat.view[:, :] - lat_rad[subtile_slice_grid] + ) / np.amax(lat_rad[subtile_slice_grid]) + diffs.append(f"Lat maximum relative error = {diff_lat}") + if not np.isclose( ext_driver.state.grid_data.dy.view[:, :], dx[subtile_slice_dx] ).all(): errors.append("dx data mismatch") + diff_dx = np.amax( + ext_driver.state.grid_data.dy.view[:, :] - dx[subtile_slice_dx] + ) / np.amax(dx[subtile_slice_dx]) + diffs.append(f"dx maximum relative error = {diff_dx}") + if not np.isclose( ext_driver.state.grid_data.dx.view[:, :], dy[subtile_slice_dy] ).all(): errors.append("dy data mismatch") + diff_dy = np.amax( + ext_driver.state.grid_data.dx.view[:, :] - dy[subtile_slice_dy] + ) / np.amax(dy[subtile_slice_dy]) + diffs.append(f"dy maximum relative error = {diff_dy}") + if not np.isclose( ext_driver.state.grid_data.area.view[:, :], area[subtile_slice_area] ).all(): errors.append("area data mismatch") + diff_area = np.amax( + ext_driver.state.grid_data.area.view[:, :] - area[subtile_slice_area] + ) / np.amax(area[subtile_slice_area]) + diffs.append(f"Area maximum relative error = {diff_area}") + assert not errors, "errors occured in 2x2:\n{}".format("\n".join(errors)) surface_area_true = 4 * PI * (RADIUS ** 2) From 264549ae8f6a3e3d8746ed065cbea4d38d9e3396 Mon Sep 17 00:00:00 2001 From: Frank Malatino Date: Wed, 3 Jan 2024 20:11:49 -0500 Subject: [PATCH 26/31] External grid data read in tests changed: relative errors printed by each rank and get_tile_number replacing get_tile_index --- tests/mpi_54rank/test_external_grid_1x1.py | 7 ++++--- tests/mpi_54rank/test_external_grid_2x2.py | 6 ++++-- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/tests/mpi_54rank/test_external_grid_1x1.py b/tests/mpi_54rank/test_external_grid_1x1.py index b6d951cf..6ca43828 100644 --- a/tests/mpi_54rank/test_external_grid_1x1.py +++ b/tests/mpi_54rank/test_external_grid_1x1.py @@ -24,7 +24,7 @@ def get_cube_comm(layout, comm: MPIComm): def get_tile_num(comm: MPIComm): - return pace.util.get_tile_index(comm.rank, comm.partitioner.total_ranks) + return pace.util.get_tile_number(comm.rank, comm.partitioner.total_ranks) def test_extgrid_equals_generated_1x1(): @@ -41,7 +41,7 @@ def test_extgrid_equals_generated_1x1(): comm_1by1 = MPIComm() cube_comm = get_cube_comm(layout=(1, 1), comm=comm_1by1) - tile_num = get_tile_num(cube_comm) + 1 + tile_num = get_tile_num(cube_comm) tile_file = "../../test_input/C12.tile" + str(tile_num) + ".nc" ds = xr.open_dataset(os.path.join(DIR, tile_file)) lon = ds.x.values @@ -138,6 +138,8 @@ def test_extgrid_equals_generated_1x1(): ) / np.amax(area[subtile_slice_area]) diffs.append(f"Area maximum relative error = {diff_area}") + print(diffs) + assert not errors, "errors occured in 1x1:\n{}".format("\n".join(errors)) surface_area_true = 4 * PI * (RADIUS ** 2) @@ -150,5 +152,4 @@ def test_extgrid_equals_generated_1x1(): if mpicomm.Get_rank() == 0: total_area = math.fsum(tile_area) - print(diffs) assert np.isclose(total_area, surface_area_true) diff --git a/tests/mpi_54rank/test_external_grid_2x2.py b/tests/mpi_54rank/test_external_grid_2x2.py index 8ea8f1ff..467aaf7c 100644 --- a/tests/mpi_54rank/test_external_grid_2x2.py +++ b/tests/mpi_54rank/test_external_grid_2x2.py @@ -24,7 +24,7 @@ def get_cube_comm(layout, comm: MPIComm): def get_tile_num(comm: MPIComm): - return pace.util.get_tile_index(comm.rank, comm.partitioner.total_ranks) + return pace.util.get_tile_number(comm.rank, comm.partitioner.total_ranks) def test_extgrid_equals_generated_2x2(): @@ -41,7 +41,7 @@ def test_extgrid_equals_generated_2x2(): comm_2by2 = MPIComm() cube_comm = get_cube_comm(layout=(2, 2), comm=comm_2by2) - tile_num = get_tile_num(cube_comm) + 1 + tile_num = get_tile_num(cube_comm) tile_file = "../../test_input/C12.tile" + str(tile_num) + ".nc" ds = xr.open_dataset(os.path.join(DIR, tile_file)) lon = ds.x.values @@ -138,6 +138,8 @@ def test_extgrid_equals_generated_2x2(): ) / np.amax(area[subtile_slice_area]) diffs.append(f"Area maximum relative error = {diff_area}") + print(diffs) + assert not errors, "errors occured in 2x2:\n{}".format("\n".join(errors)) surface_area_true = 4 * PI * (RADIUS ** 2) From 75a3e1d277ca79513a53b8fdb1796925d46dd62a Mon Sep 17 00:00:00 2001 From: fmalatino <142349306+fmalatino@users.noreply.github.com> Date: Tue, 23 Jan 2024 20:10:41 -0500 Subject: [PATCH 27/31] Removing commented out sections in test_external_grid_2x2.py Co-authored-by: Oliver Elbert --- tests/mpi_54rank/test_external_grid_2x2.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/mpi_54rank/test_external_grid_2x2.py b/tests/mpi_54rank/test_external_grid_2x2.py index 467aaf7c..bd71ea20 100644 --- a/tests/mpi_54rank/test_external_grid_2x2.py +++ b/tests/mpi_54rank/test_external_grid_2x2.py @@ -151,8 +151,5 @@ def test_extgrid_equals_generated_2x2(): tile_area = mpicomm._comm.gather(rank_area_sum, root=0) if mpicomm.Get_rank() == 0: - # print(ext_driver.state.grid_data.dy.view[:, :]) - # print("") - # print(dx[subtile_slice_dx]) total_area = math.fsum(tile_area) assert np.isclose(total_area, surface_area_true) From f0a4a4dcf3e06f234228c74530e6a5adac9624bc Mon Sep 17 00:00:00 2001 From: Frank Malatino Date: Thu, 25 Jan 2024 11:30:56 -0500 Subject: [PATCH 28/31] Updated external grid data read-in to take configuration and input data locations from command line, updated test description, and added documentation on grid construction to external grid data configuration selection dataclass. --- constraints.txt | 4 +- .../configs/test_external_C12_1x1.yaml | 2 +- .../configs/test_external_C12_2x2.yaml | 2 +- driver/pace/driver/grid.py | 18 +- tests/main/driver/test_analytic_init.py | 3 + tests/mpi_54rank/test_ext_grid/conftest.py | 13 ++ .../test_external_grid.py} | 47 +++++- tests/mpi_54rank/test_external_grid_1x1.py | 155 ------------------ 8 files changed, 78 insertions(+), 166 deletions(-) create mode 100644 tests/mpi_54rank/test_ext_grid/conftest.py rename tests/mpi_54rank/{test_external_grid_2x2.py => test_ext_grid/test_external_grid.py} (70%) delete mode 100644 tests/mpi_54rank/test_external_grid_1x1.py diff --git a/constraints.txt b/constraints.txt index 73f71625..02b6b94f 100644 --- a/constraints.txt +++ b/constraints.txt @@ -81,7 +81,7 @@ coverage==5.5 # pytest-cov cytoolz==0.12.1 # via gt4py -dace==0.14.4 +dace==0.15 # via # -r requirements_dev.txt # pace-dsl @@ -184,7 +184,7 @@ googleapis-common-protos==1.53.0 # via google-api-core gprof2dot==2021.2.21 # via pytest-profiling -gridtools-cpp==2.3.0 +gridtools-cpp==2.3.1 # via gt4py h5netcdf==0.11.0 # via -r util/requirements.txt diff --git a/driver/examples/configs/test_external_C12_1x1.yaml b/driver/examples/configs/test_external_C12_1x1.yaml index 269921ac..7e721889 100644 --- a/driver/examples/configs/test_external_C12_1x1.yaml +++ b/driver/examples/configs/test_external_C12_1x1.yaml @@ -9,7 +9,7 @@ grid_config: type: external config: grid_type: 0 - grid_file_path: "../../test_input/C12.tile" + grid_file_path: "../../../test_input/C12.tile" initialization: type: analytic config: diff --git a/driver/examples/configs/test_external_C12_2x2.yaml b/driver/examples/configs/test_external_C12_2x2.yaml index b12166b8..74e24222 100644 --- a/driver/examples/configs/test_external_C12_2x2.yaml +++ b/driver/examples/configs/test_external_C12_2x2.yaml @@ -9,7 +9,7 @@ grid_config: type: external config: grid_type: 0 - grid_file_path: "../../test_input/C12.tile" + grid_file_path: "../../../test_input/C12.tile" initialization: type: analytic config: diff --git a/driver/pace/driver/grid.py b/driver/pace/driver/grid.py index 83d68888..79c00e9d 100644 --- a/driver/pace/driver/grid.py +++ b/driver/pace/driver/grid.py @@ -211,7 +211,23 @@ class ExternalNetcdfGridConfig(GridInitializer): damping_coefficients, driver_grid_data, and grid_data variables We do not read in the dx, dy, or area values as there may be inconsistencies in the constants used during calculation of the - input data and the model use. + input data and the model use. An example of two adjacent finite + volume cells in the supergrid should appear like: + + X----X----X----X----X + | | | + X X X X X + | | | + X----X----X----X----X + + Grid construction follows a gnomonic coordinate system. An + inscribed cube is projected onto the surface of the sphere of + the Earth. This creates six, non-overlapping faces/tiles that + contain the surface area of the sphere. For more information + on the NOAA-GFDL grid construction methods, please visit: + https://repository.library.noaa.gov/view/noaa/30725 + which is a pdf describing the FV3 dynamical core, by + Harris, L.; Chen, X.; Putnam, W.; et. al. """ grid_type: Optional[int] = 0 diff --git a/tests/main/driver/test_analytic_init.py b/tests/main/driver/test_analytic_init.py index d13c8af2..97222db4 100644 --- a/tests/main/driver/test_analytic_init.py +++ b/tests/main/driver/test_analytic_init.py @@ -9,6 +9,9 @@ DIR = os.path.dirname(os.path.abspath(__file__)) +# TODO: Location of test configurations will be changed after refactor, +# need to update after + TESTED_CONFIGS: List[str] = [ "../../../driver/examples/configs/analytic_test.yaml", ] diff --git a/tests/mpi_54rank/test_ext_grid/conftest.py b/tests/mpi_54rank/test_ext_grid/conftest.py new file mode 100644 index 00000000..3929a5f8 --- /dev/null +++ b/tests/mpi_54rank/test_ext_grid/conftest.py @@ -0,0 +1,13 @@ +def pytest_addoption(parser): + parser.addoption( + "--config_file_path", + action="store", + default="config.yaml", + help="Configuration for testing.", + ) + parser.addoption( + "--tile_file_base_path", + action="store", + default="tile.nc", + help="Tile file base name to be read in", + ) diff --git a/tests/mpi_54rank/test_external_grid_2x2.py b/tests/mpi_54rank/test_ext_grid/test_external_grid.py similarity index 70% rename from tests/mpi_54rank/test_external_grid_2x2.py rename to tests/mpi_54rank/test_ext_grid/test_external_grid.py index bd71ea20..7400f87c 100644 --- a/tests/mpi_54rank/test_external_grid_2x2.py +++ b/tests/mpi_54rank/test_ext_grid/test_external_grid.py @@ -27,10 +27,43 @@ def get_tile_num(comm: MPIComm): return pace.util.get_tile_number(comm.rank, comm.partitioner.total_ranks) -def test_extgrid_equals_generated_2x2(): +# TODO: Location of test configurations and data will be changed +# after refactor, need to update after + + +def test_extgrid_equals_generated(request): + """ + Test for matching values between what exists in external grid + file, and what is generated by Pace when using an "external" + grid configuration. The "external" grid configuration takes + d-grid values as input and Pace then calculates the remaining + metric terms from this data. + + External grid files are expected to be of Netcdf type + (for reference see method of generating grid data in + FRE-NCtools) and contain values for the longitudinal, + latitude, distance between longitude and latitude points, + and cell areas. + + On a run, the user should specify the correct number of ranks, + location of configuration yaml, and base name for the tile files + containing the grid data. + + ON PASS: Input and generated data will match, with zero, or + insignificant relative difference. Each rank will print a + list of maximum relative differences for each metric term + tested. + + ON FAIL: There will be a mismatch present between input and + generated grid data values. Mismatch can be determined from the + printed list of errors and the list of maximum relative differences + """ + + config_file_path = request.config.getoption("--config_file_path") + tile_file_base_path = request.config.getoption("--tile_file_base_path") with open( - os.path.join(DIR, "../../driver/examples/configs/test_external_C12_2x2.yaml"), + os.path.join(DIR, "../../../", config_file_path), "r", ) as ext_f: ext_config = yaml.safe_load(ext_f) @@ -38,12 +71,14 @@ def test_extgrid_equals_generated_2x2(): ext_driver = Driver(ext_driver_config) - comm_2by2 = MPIComm() - cube_comm = get_cube_comm(layout=(2, 2), comm=comm_2by2) + comm = MPIComm() + size = comm.Get_size() + side = int(math.sqrt(size // 6)) + cube_comm = get_cube_comm(layout=(side, side), comm=comm) tile_num = get_tile_num(cube_comm) - tile_file = "../../test_input/C12.tile" + str(tile_num) + ".nc" - ds = xr.open_dataset(os.path.join(DIR, tile_file)) + tile_file = tile_file_base_path + str(tile_num) + ".nc" + ds = xr.open_dataset(os.path.join(DIR, "../../../", tile_file)) lon = ds.x.values lat = ds.y.values dx = ds.dx.values diff --git a/tests/mpi_54rank/test_external_grid_1x1.py b/tests/mpi_54rank/test_external_grid_1x1.py deleted file mode 100644 index 6ca43828..00000000 --- a/tests/mpi_54rank/test_external_grid_1x1.py +++ /dev/null @@ -1,155 +0,0 @@ -import math -import os - -import numpy as np -import xarray as xr -import yaml - -import pace.util -from pace.driver import Driver, DriverConfig -from pace.util.constants import PI, RADIUS -from pace.util.mpi import MPIComm - - -DIR = os.path.dirname(os.path.abspath(__file__)) - - -def get_cube_comm(layout, comm: MPIComm): - return pace.util.CubedSphereCommunicator( - comm=comm, - partitioner=pace.util.CubedSpherePartitioner( - pace.util.TilePartitioner(layout=layout) - ), - ) - - -def get_tile_num(comm: MPIComm): - return pace.util.get_tile_number(comm.rank, comm.partitioner.total_ranks) - - -def test_extgrid_equals_generated_1x1(): - - with open( - os.path.join(DIR, "../../driver/examples/configs/test_external_C12_1x1.yaml"), - "r", - ) as ext_f: - ext_config = yaml.safe_load(ext_f) - ext_driver_config = DriverConfig.from_dict(ext_config) - - ext_driver = Driver(ext_driver_config) - - comm_1by1 = MPIComm() - cube_comm = get_cube_comm(layout=(1, 1), comm=comm_1by1) - - tile_num = get_tile_num(cube_comm) - tile_file = "../../test_input/C12.tile" + str(tile_num) + ".nc" - ds = xr.open_dataset(os.path.join(DIR, tile_file)) - lon = ds.x.values - lat = ds.y.values - dx = ds.dx.values - dy = ds.dy.values - area = ds.area.values - nx = ds.nx.values.size - ny = ds.ny.values.size - npx = ds.nxp.values.size - npy = ds.nyp.values.size - - subtile_slice_grid = cube_comm.partitioner.tile.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.Y_INTERFACE_DIM, pace.util.X_INTERFACE_DIM], - global_extent=(npy, npx), - overlap=True, - ) - - subtile_slice_dx = cube_comm.partitioner.tile.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.Y_INTERFACE_DIM, pace.util.X_DIM], - global_extent=(npy, nx), - overlap=True, - ) - - subtile_slice_dy = cube_comm.partitioner.tile.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.Y_DIM, pace.util.X_INTERFACE_DIM], - global_extent=(ny, npx), - overlap=True, - ) - - subtile_slice_area = cube_comm.partitioner.tile.subtile_slice( - rank=cube_comm.rank, - global_dims=[pace.util.Y_DIM, pace.util.X_DIM], - global_extent=(ny, nx), - overlap=True, - ) - - lon_rad = lon * (PI / 180) - lat_rad = lat * (PI / 180) - - errors = [] - diffs = [] - - if not np.isclose( - ext_driver.state.grid_data.lon.view[:, :], lon_rad[subtile_slice_grid] - ).all(): - errors.append("Lon data mismatch") - - diff_lon = np.amax( - ext_driver.state.grid_data.lon.view[:, :] - lon_rad[subtile_slice_grid] - ) / np.amax(lon_rad[subtile_slice_grid]) - diffs.append(f"Lon maximum relative error = {diff_lon}") - - if not np.isclose( - ext_driver.state.grid_data.lat.view[:, :], lat_rad[subtile_slice_grid] - ).all(): - errors.append("Lat data mismatch") - - diff_lat = np.amax( - ext_driver.state.grid_data.lat.view[:, :] - lat_rad[subtile_slice_grid] - ) / np.amax(lat_rad[subtile_slice_grid]) - diffs.append(f"Lat maximum relative error = {diff_lat}") - - if not np.isclose( - ext_driver.state.grid_data.dy.view[:, :], dx[subtile_slice_dx] - ).all(): - errors.append("dx data mismatch") - - diff_dx = np.amax( - ext_driver.state.grid_data.dy.view[:, :] - dx[subtile_slice_dx] - ) / np.amax(dx[subtile_slice_dx]) - diffs.append(f"dx maximum relative error = {diff_dx}") - - if not np.isclose( - ext_driver.state.grid_data.dx.view[:, :], dy[subtile_slice_dy] - ).all(): - errors.append("dy data mismatch") - - diff_dy = np.amax( - ext_driver.state.grid_data.dx.view[:, :] - dy[subtile_slice_dy] - ) / np.amax(dy[subtile_slice_dy]) - diffs.append(f"dy maximum relative error = {diff_dy}") - - if not np.isclose( - ext_driver.state.grid_data.area.view[:, :], area[subtile_slice_area] - ).all(): - errors.append("area data mismatch") - - diff_area = np.amax( - ext_driver.state.grid_data.area.view[:, :] - area[subtile_slice_area] - ) / np.amax(area[subtile_slice_area]) - diffs.append(f"Area maximum relative error = {diff_area}") - - print(diffs) - - assert not errors, "errors occured in 1x1:\n{}".format("\n".join(errors)) - - surface_area_true = 4 * PI * (RADIUS ** 2) - - mpicomm = MPIComm() - - tmp = [math.fsum(row) for row in ext_driver.state.grid_data.area.view[:, :]] - rank_area_sum = math.fsum(tmp) - tile_area = mpicomm._comm.gather(rank_area_sum, root=0) - - if mpicomm.Get_rank() == 0: - total_area = math.fsum(tile_area) - assert np.isclose(total_area, surface_area_true) From 1b7ca1a26495f655d7fb70a0bb8393630c6d6e1c Mon Sep 17 00:00:00 2001 From: Frank Malatino Date: Thu, 25 Jan 2024 16:08:07 -0500 Subject: [PATCH 29/31] Updated documentation in grid.py --- driver/pace/driver/grid.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/driver/pace/driver/grid.py b/driver/pace/driver/grid.py index 79c00e9d..240f2559 100644 --- a/driver/pace/driver/grid.py +++ b/driver/pace/driver/grid.py @@ -220,14 +220,10 @@ class ExternalNetcdfGridConfig(GridInitializer): | | | X----X----X----X----X - Grid construction follows a gnomonic coordinate system. An - inscribed cube is projected onto the surface of the sphere of - the Earth. This creates six, non-overlapping faces/tiles that - contain the surface area of the sphere. For more information - on the NOAA-GFDL grid construction methods, please visit: - https://repository.library.noaa.gov/view/noaa/30725 - which is a pdf describing the FV3 dynamical core, by - Harris, L.; Chen, X.; Putnam, W.; et. al. + The grid data must define the verticies, centroids, and mid-points + on edge of the cells contained in the computation. For more information + on grid discretization for the FV3 dynamical core please visit: + https://www.gfdl.noaa.gov/fv3/ """ grid_type: Optional[int] = 0 From 39467a9181b70bd1b6836e4d9bd6c4c3d7eaea44 Mon Sep 17 00:00:00 2001 From: Frank Malatino Date: Fri, 26 Jan 2024 14:02:21 -0500 Subject: [PATCH 30/31] Updated external grid data read in unit test to use parametrize functionality of pytest --- tests/mpi_54rank/test_ext_grid/conftest.py | 13 ------- .../test_ext_grid/test_external_grid.py | 36 +++++++++++++------ 2 files changed, 25 insertions(+), 24 deletions(-) delete mode 100644 tests/mpi_54rank/test_ext_grid/conftest.py diff --git a/tests/mpi_54rank/test_ext_grid/conftest.py b/tests/mpi_54rank/test_ext_grid/conftest.py deleted file mode 100644 index 3929a5f8..00000000 --- a/tests/mpi_54rank/test_ext_grid/conftest.py +++ /dev/null @@ -1,13 +0,0 @@ -def pytest_addoption(parser): - parser.addoption( - "--config_file_path", - action="store", - default="config.yaml", - help="Configuration for testing.", - ) - parser.addoption( - "--tile_file_base_path", - action="store", - default="tile.nc", - help="Tile file base name to be read in", - ) diff --git a/tests/mpi_54rank/test_ext_grid/test_external_grid.py b/tests/mpi_54rank/test_ext_grid/test_external_grid.py index 7400f87c..7cef4657 100644 --- a/tests/mpi_54rank/test_ext_grid/test_external_grid.py +++ b/tests/mpi_54rank/test_ext_grid/test_external_grid.py @@ -2,6 +2,7 @@ import os import numpy as np +import pytest import xarray as xr import yaml @@ -12,6 +13,15 @@ DIR = os.path.dirname(os.path.abspath(__file__)) +TEST_CONFIGS_DIR = os.path.join(DIR, "../../../driver/examples/configs/") +TEST_DATA_DIR = os.path.join(DIR, "../../../test_input/") + +TEST_CONFIG_FILE_RANKS = [ + ("test_external_C12_1x1.yaml", 6), + ("test_external_C12_2x2.yaml", 24), +] + +TILE_FILE_BASE_NAME = "C12.tile" def get_cube_comm(layout, comm: MPIComm): @@ -31,7 +41,8 @@ def get_tile_num(comm: MPIComm): # after refactor, need to update after -def test_extgrid_equals_generated(request): +@pytest.mark.parametrize("config_file_path, ranks", TEST_CONFIG_FILE_RANKS) +def test_extgrid_equals_generated(config_file_path: str, ranks: int): """ Test for matching values between what exists in external grid file, and what is generated by Pace when using an "external" @@ -59,11 +70,19 @@ def test_extgrid_equals_generated(request): printed list of errors and the list of maximum relative differences """ - config_file_path = request.config.getoption("--config_file_path") - tile_file_base_path = request.config.getoption("--tile_file_base_path") + comm = MPIComm() + size = comm.Get_size() + if size != ranks: + pytest.skip("Number of ranks does not match amount needed for layout") + side = int(math.sqrt(size // 6)) + cube_comm = get_cube_comm(layout=(side, side), comm=comm) + + print("THIS") + print(TEST_CONFIGS_DIR) + print("THAT") with open( - os.path.join(DIR, "../../../", config_file_path), + os.path.join(TEST_CONFIGS_DIR, config_file_path), "r", ) as ext_f: ext_config = yaml.safe_load(ext_f) @@ -71,14 +90,9 @@ def test_extgrid_equals_generated(request): ext_driver = Driver(ext_driver_config) - comm = MPIComm() - size = comm.Get_size() - side = int(math.sqrt(size // 6)) - cube_comm = get_cube_comm(layout=(side, side), comm=comm) - tile_num = get_tile_num(cube_comm) - tile_file = tile_file_base_path + str(tile_num) + ".nc" - ds = xr.open_dataset(os.path.join(DIR, "../../../", tile_file)) + tile_file = TILE_FILE_BASE_NAME + str(tile_num) + ".nc" + ds = xr.open_dataset(os.path.join(TEST_DATA_DIR, tile_file)) lon = ds.x.values lat = ds.y.values dx = ds.dx.values From f37c4fde065ac8b72b7972a47834e6e8087d8eb2 Mon Sep 17 00:00:00 2001 From: Frank Malatino Date: Tue, 30 Jan 2024 12:26:45 -0500 Subject: [PATCH 31/31] Ammended files to reference changes to PR 36 --- driver/examples/configs/test_external_C12_1x1.yaml | 1 + driver/examples/configs/test_external_C12_2x2.yaml | 1 + driver/pace/driver/grid.py | 2 ++ tests/mpi_54rank/test_ext_grid/test_external_grid.py | 4 ---- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/driver/examples/configs/test_external_C12_1x1.yaml b/driver/examples/configs/test_external_C12_1x1.yaml index 7e721889..fb615b4e 100644 --- a/driver/examples/configs/test_external_C12_1x1.yaml +++ b/driver/examples/configs/test_external_C12_1x1.yaml @@ -10,6 +10,7 @@ grid_config: config: grid_type: 0 grid_file_path: "../../../test_input/C12.tile" + eta_file: '../../../tests/main/input/eta79.nc' initialization: type: analytic config: diff --git a/driver/examples/configs/test_external_C12_2x2.yaml b/driver/examples/configs/test_external_C12_2x2.yaml index 74e24222..d07d142b 100644 --- a/driver/examples/configs/test_external_C12_2x2.yaml +++ b/driver/examples/configs/test_external_C12_2x2.yaml @@ -10,6 +10,7 @@ grid_config: config: grid_type: 0 grid_file_path: "../../../test_input/C12.tile" + eta_file: '../../../tests/main/input/eta79.nc' initialization: type: analytic config: diff --git a/driver/pace/driver/grid.py b/driver/pace/driver/grid.py index f15a7d5b..e1c3b057 100644 --- a/driver/pace/driver/grid.py +++ b/driver/pace/driver/grid.py @@ -230,6 +230,7 @@ class ExternalNetcdfGridConfig(GridInitializer): grid_type: Optional[int] = 0 grid_file_path: str = "/input/tilefile" + eta_file: str = "None" def get_grid( self, @@ -270,6 +271,7 @@ def get_grid( quantity_factory=quantity_factory, communicator=communicator, grid_type=self.grid_type, + eta_file=self.eta_file, ) horizontal_data = HorizontalGridData.new_from_metric_terms(metric_terms) diff --git a/tests/mpi_54rank/test_ext_grid/test_external_grid.py b/tests/mpi_54rank/test_ext_grid/test_external_grid.py index 7cef4657..03f50e65 100644 --- a/tests/mpi_54rank/test_ext_grid/test_external_grid.py +++ b/tests/mpi_54rank/test_ext_grid/test_external_grid.py @@ -77,10 +77,6 @@ def test_extgrid_equals_generated(config_file_path: str, ranks: int): side = int(math.sqrt(size // 6)) cube_comm = get_cube_comm(layout=(side, side), comm=comm) - print("THIS") - print(TEST_CONFIGS_DIR) - print("THAT") - with open( os.path.join(TEST_CONFIGS_DIR, config_file_path), "r",