D-grid data input method (#42)

* Testing changes reflected across branches

* Undoing changes made in build_gaea_c5.sh

* Testing vscode functionality, by adding a change to external_grid branch

* Testing vscode functionality, by adding a change to external_grid branch

* Addition of from_generated method and calc_flag to util/pace/util/grid/generation.py

* Added get_grid method for external grid data to driver/pace/driver/grid.py

* Preliminary xarray netcdf read in method added to driver/pace/driver/grid.py

* Updating util/pace/util/grid/generation.py from_generated method

* 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.

* driver/examples/configs/test_external_C12_1x1.yaml

* Preliminary unit test for external grid data read in

* Current state of unit tests as of 27 Nov 2023

* External grid method and unit tests added

* Re-excluding external grid data yamls from test_example_configs.py

* Update driver/pace/driver/grid.py

Co-authored-by: Florian Deconinck <[email protected]>

* Changed name of grid initializer function to match NetCDF dependency and class descriptor

* Update util/pace/util/grid/generation.py

Moved position of doc string for "from_external" MetricTerms class method

Co-authored-by: Oliver Elbert <[email protected]>

* Fixed indentation error in generation.py from suggestion in PR 42

* Removal of TODO comment in grid.py, changes to method of file accessing in test_analytic_init, test_external_grid_*

* Changed grid data read-in unit tests to compare data directly from file to driver grid data generated from yaml

* Change to reading in lon and lat, other metric terms calculated as needed

* 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.

* 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 <[email protected]>

* Added comparisons for read-in vs generated by driver lon, lat, dx, dy, and area data to unit tests

* Added relative error calculations to unit tests for external grid data read-in

* External grid data read in tests changed: relative errors printed by each rank and get_tile_number replacing get_tile_index

* Removing commented out sections in test_external_grid_2x2.py

Co-authored-by: Oliver Elbert <[email protected]>

* 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.

* Updated documentation in grid.py

* Updated external grid data read in unit test to use parametrize functionality of pytest

* Ammended files to reference changes to PR 36

---------

Co-authored-by: Frank Malatino <[email protected]>
Co-authored-by: Florian Deconinck <[email protected]>
Co-authored-by: Oliver Elbert <[email protected]>

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.

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

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

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

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,

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"

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/")