diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index 7030919a..72dda3d5 100644 --- a/CONTRIBUTORS.md +++ b/CONTRIBUTORS.md @@ -13,6 +13,7 @@ List format (alphabetical order): Surname, Name. Employer/Affiliation * Harris, Lucas. GFDL. * Lee, Mi Kyung. GFDL. * Kung, Chris. NASA. +* Malatino, Frank. GFDL * McGibbon, Jeremy. Allen Institute for AI. * Niedermayr, Yannick. ETH. * Savarin, Ajda. University of Washington. 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 new file mode 100644 index 00000000..fb615b4e --- /dev/null +++ b/driver/examples/configs/test_external_C12_1x1.yaml @@ -0,0 +1,103 @@ +stencil_config: + compilation_config: + backend: numpy + rebuild: false + validate_args: true + format_source: false + device_sync: false +grid_config: + type: external + config: + grid_type: 0 + grid_file_path: "../../../test_input/C12.tile" + eta_file: '../../../tests/main/input/eta79.nc' +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..d07d142b --- /dev/null +++ b/driver/examples/configs/test_external_C12_2x2.yaml @@ -0,0 +1,103 @@ +stencil_config: + compilation_config: + backend: numpy + rebuild: false + validate_args: true + format_source: false + device_sync: false +grid_config: + type: external + config: + grid_type: 0 + grid_file_path: "../../../test_input/C12.tile" + eta_file: '../../../tests/main/input/eta79.nc' +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 d4e07d00..e1c3b057 100644 --- a/driver/pace/driver/grid.py +++ b/driver/pace/driver/grid.py @@ -3,11 +3,13 @@ from typing import ClassVar, Optional, Tuple import f90nml +import xarray as xr import pace.driver 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 @@ -65,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 @@ -198,6 +201,96 @@ def get_grid( return damping_coefficients, driver_grid_data, grid_data +@GridInitializerSelector.register("external") +@dataclasses.dataclass +class ExternalNetcdfGridConfig(GridInitializer): + """ + Configuration for grid initialized from external data. + Input is from tile files generated by FRE-NCtools methods + 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 + 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. 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 + + 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 + grid_file_path: str = "/input/tilefile" + eta_file: str = "None" + + def get_grid( + self, + quantity_factory: QuantityFactory, + communicator: Communicator, + ) -> Tuple[DampingCoefficients, DriverGridData, GridData]: + + 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( + communicator.rank, communicator.partitioner.total_ranks + ) + + 1 + ) + 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 + npx = ds.nxp.values.size + npy = ds.nyp.values.size + + subtile_slice_grid = communicator.partitioner.tile.subtile_slice( + rank=communicator.rank, + global_dims=[pace.util.Y_INTERFACE_DIM, pace.util.X_INTERFACE_DIM], + global_extent=(npy, npx), + overlap=True, + ) + + metric_terms = MetricTerms.from_external( + x=lon[subtile_slice_grid], + y=lat[subtile_slice_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) + 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, diff --git a/tests/main/driver/test_analytic_init.py b/tests/main/driver/test_analytic_init.py index 03dda5bc..97222db4 100644 --- a/tests/main/driver/test_analytic_init.py +++ b/tests/main/driver/test_analytic_init.py @@ -7,8 +7,13 @@ import pace.driver +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", + "../../../driver/examples/configs/analytic_test.yaml", ] @@ -20,7 +25,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/main/driver/test_example_configs.py b/tests/main/driver/test_example_configs.py index 5a9dcbcd..2e4b5d3e 100644 --- a/tests/main/driver/test_example_configs.py +++ b/tests/main/driver/test_example_configs.py @@ -29,6 +29,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/") diff --git a/tests/mpi_54rank/test_ext_grid/test_external_grid.py b/tests/mpi_54rank/test_ext_grid/test_external_grid.py new file mode 100644 index 00000000..03f50e65 --- /dev/null +++ b/tests/mpi_54rank/test_ext_grid/test_external_grid.py @@ -0,0 +1,200 @@ +import math +import os + +import numpy as np +import pytest +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__)) +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): + 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) + + +# TODO: Location of test configurations and data will be changed +# after refactor, need to update after + + +@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" + 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 + """ + + 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) + + with open( + os.path.join(TEST_CONFIGS_DIR, config_file_path), + "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) + + tile_num = get_tile_num(cube_comm) + 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 + 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 2x2:\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) 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 6181661f..2c556576 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, + extdgrid: bool = False, eta_file: str = "None", ): self._grid_type = grid_type @@ -412,11 +413,42 @@ def __init__( self._calculate_unit_vectors_lonlat = ( self._calculate_unit_vectors_lonlat_cube_sphere ) - self._init_dgrid() - self._init_agrid() + if extdgrid is False: + self._init_dgrid() + self._init_agrid() else: raise NotImplementedError(f"Unsupported grid_type = {grid_type}") + @classmethod + def from_external( + cls, + x, + y, + quantity_factory, + communicator, + grid_type, + eta_file: str = "None", + ) -> "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, + grid_type=grid_type, + extdgrid=True, + eta_file=eta_file, + ) + + rad_conv = PI / 180.0 + terms._grid_64.view[:, :, 0] = rad_conv * x + terms._grid_64.view[:, :, 1] = rad_conv * y + + terms._init_agrid() + + return terms + @classmethod def from_tile_sizing( cls,