D-grid data input method

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.

driver/examples/configs/test_external_C12_1x1.yaml

@@ -0,0 +1,102 @@
+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"
+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,102 @@
+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"
+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
@@ -196,6 +199,110 @@ 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
+    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
+    """
+
+    grid_type: Optional[int] = 0
+    grid_file_path: str = "/input/tilefile"
+
+    # TODO: Area read in?
+
+    def get_grid(
+        self,
+        quantity_factory: QuantityFactory,
+        communicator: Communicator,
+    ) -> Tuple[DampingCoefficients, DriverGridData, GridData]:
+
+        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
+                )
+                + 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
+        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.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.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,
+            extdgrid=True,
+        )
+
+        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_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/")