Merge remote-tracking branch 'pace/main' into read_akbk

.jenkins/driver_configs/baroclinic_c192_54ranks.yaml

@@ -6,7 +6,9 @@ stencil_config:
     format_source: false
     device_sync: true
 initialization:
-  type: baroclinic
+  type: analytic
+  config:
+    case: baroclinic
 diagnostics_config:
   path: "output.zarr"
   names:

.jenkins/driver_configs/baroclinic_c192_6ranks.yaml

@@ -6,7 +6,9 @@ stencil_config:
     format_source: false
     device_sync: true
 initialization:
-  type: baroclinic
+  type: analytic
+  config:
+    case: baroclinic
 diagnostics_config:
   path: "output.zarr"
   names:

.jenkins/driver_configs/baroclinic_c48_6ranks_dycore_only.yaml

@@ -9,7 +9,9 @@ stencil_config:
     device_sync: false
     run_mode: Run
 initialization:
-  type: baroclinic
+  type: analytic
+  config:
+    case: baroclinicc
 performance_config:
   collect_performance: false
 nx_tile: 48

CONTRIBUTORS.md

@@ -2,6 +2,7 @@
 
 List format (alphabetical order):  Surname, Name. Employer/Affiliation
 
+* Abbott, Tristan. GFDL.
 * Cheeseman, Mark. Vulcan Inc.
 * Dahm, Johann. Allen Institute for AI.
 * Davis, Eddie. Allen Institute for AI.

README.md

@@ -70,7 +70,7 @@ After the run completes, you will see an output direcotry `output.zarr`. An exam
 ### Environment variable configuration
 
 - `PACE_CONSTANTS`: Pace is bundled with various constants (see _util/pace/util/constants.py_).
-  - `FV3DYCORE` NOAA's FV3 dynamical core constants (original port)
+  - `GFDL` NOAA's FV3 dynamical core constants (original port)
   - `GFS` Constant as defined in NOAA GFS
   - `GEOS`  Constant as defined in GEOS v13
 - `PACE_FLOAT_PRECISION`: default precision of the field & scalars in the numerics. Default to 64.

constraints.txt

@@ -188,7 +188,7 @@ gridtools-cpp==2.3.0
     # via gt4py
 h5netcdf==0.11.0
     # via -r util/requirements.txt
-h5py==2.10.0
+h5py==3.9.0
     # via
     #   -r util/requirements.txt
     #   h5netcdf
@@ -277,7 +277,7 @@ nest-asyncio==1.5.6
     #   ipykernel
     #   jupyter-client
     #   nbclient
-netcdf4==1.5.7
+netcdf4==1.6.4
     # via
     #   -r requirements_dev.txt
     #   pace-driver

driver/examples/configs/analytic_test.yaml

@@ -0,0 +1,101 @@
+stencil_config:
+  compilation_config:
+    backend: numpy
+    rebuild: false
+    validate_args: true
+    format_source: false
+    device_sync: false
+initialization:
+  type: analytic
+  config:
+    case: baroclinic
+performance_config:
+  collect_performance: true
+  experiment_name: c12_baroclinic
+comm_config:
+  type: null_comm
+  config:
+    rank: 0
+    total_ranks: 6
+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
+
+physics_config:
+  hydrostatic: false
+  nwat: 6
+  do_qa: true

driver/examples/configs/baroclinic_c12.yaml

@@ -6,7 +6,9 @@ stencil_config:
     format_source: false
     device_sync: false
 initialization:
-  type: baroclinic
+  type: analytic
+  config:
+    case: baroclinic
 performance_config:
   collect_performance: true
   experiment_name: c12_baroclinic

driver/examples/configs/baroclinic_c12_comm_read.yaml

@@ -6,7 +6,9 @@ stencil_config:
     format_source: false
     device_sync: false
 initialization:
-  type: baroclinic
+  type: analytic
+  config:
+    case: baroclinic
 performance_config:
   collect_performance: false
   experiment_name: c12_baroclinic

driver/examples/configs/baroclinic_c12_comm_write.yaml

@@ -6,7 +6,9 @@ stencil_config:
     format_source: false
     device_sync: false
 initialization:
-  type: baroclinic
+  type: analytic
+  config:
+    case: baroclinic
 performance_config:
   collect_performance: false
   experiment_name: c12_baroclinic

driver/examples/configs/baroclinic_c12_dp.yaml

@@ -13,7 +13,9 @@ grid_config:
     dy_const: 3000.0
     deglat: 10.0
 initialization:
-  type: baroclinic
+  type: analytic
+  config:
+    case: baroclinic
 performance_config:
   collect_performance: true
   experiment_name: c12_baroclinic

driver/examples/configs/baroclinic_c12_null_comm.yaml

@@ -6,7 +6,9 @@ stencil_config:
     format_source: false
     device_sync: false
 initialization:
-  type: baroclinic
+  type: analytic
+  config:
+    case: baroclinic
 performance_config:
   collect_performance: false
   experiment_name: c12_baroclinic

driver/examples/configs/baroclinic_c12_orch_cpu.yaml

@@ -6,7 +6,9 @@ stencil_config:
     format_source: false
     device_sync: false
 initialization:
-  type: baroclinic
+  type: analytic
+  config:
+    case: baroclinic
 performance_config:
   collect_performance: false
 nx_tile: 12

driver/examples/configs/baroclinic_c12_write_restart.yaml

@@ -6,7 +6,9 @@ stencil_config:
     format_source: false
     device_sync: false
 initialization:
-  type: baroclinic
+  type: analytic
+  config:
+    case: baroclinic
 performance_config:
   collect_performance: false
   experiment_name: c12_baroclinic

driver/examples/configs/tropicalcyclone_c128.yaml

@@ -8,7 +8,9 @@ stencil_config:
     format_source: false
     device_sync: false
 initialization:
-  type: tropicalcyclone
+  type: analytic
+  config:
+    case: tropicalcyclone
 performance_config:
   performance_mode: true
   experiment_name: c128_tropical

driver/pace/driver/__init__.py

@@ -9,7 +9,7 @@
 from .diagnostics import Diagnostics, DiagnosticsConfig
 from .driver import Driver, DriverConfig, RestartConfig
 from .grid import GeneratedGridConfig, SerialboxGridConfig
-from .initialization import BaroclinicInit, PredefinedStateInit, RestartInit
+from .initialization import AnalyticInit, PredefinedStateInit, RestartInit
 from .performance import PerformanceConfig
 from .registry import Registry
 from .state import DriverState, TendencyState

driver/pace/driver/diagnostics.py

@@ -10,7 +10,7 @@
 import pace.util
 import pace.util.grid
 from pace.dsl.dace.orchestration import dace_inhibitor
-from pace.fv3core.initialization.dycore_state import DycoreState
+from pace.fv3core.dycore_state import DycoreState
 from pace.util.constants import RGRAV
 
 from .state import DriverState

driver/pace/driver/driver.py

@@ -24,7 +24,11 @@
 
 # TODO: move update_atmos_state into pace.driver
 from pace.stencils import update_atmos_state
-from pace.util.communicator import CubedSphereCommunicator
+from pace.util.communicator import (
+    Communicator,
+    CubedSphereCommunicator,
+    TileCommunicator,
+)
 from pace.util.logging import pace_log
 
 from . import diagnostics
@@ -90,6 +94,7 @@ class DriverConfig:
     nz: int
     layout: Tuple[int, int]
     dt_atmos: float
+    grid_type: Optional[int] = 0
     grid_config: GridInitializerSelector = dataclasses.field(
         default_factory=lambda: GridInitializerSelector(
             type="generated", config=GeneratedGridConfig()
@@ -158,7 +163,7 @@ def apply_tendencies(self) -> bool:
 
     def get_grid(
         self,
-        communicator: pace.util.CubedSphereCommunicator,
+        communicator: pace.util.Communicator,
         quantity_factory: Optional[pace.util.QuantityFactory] = None,
     ) -> Tuple[
         pace.util.grid.DampingCoefficients,
@@ -187,7 +192,7 @@ def get_grid(
 
     def get_driver_state(
         self,
-        communicator: pace.util.CubedSphereCommunicator,
+        communicator: pace.util.Communicator,
         damping_coefficients: pace.util.grid.DampingCoefficients,
         driver_grid_data: pace.util.grid.DriverGridData,
         grid_data: pace.util.grid.GridData,
@@ -213,7 +218,7 @@ def get_driver_state(
             if stencil_factory is None:
                 grid_indexing = (
                     pace.dsl.stencil.GridIndexing.from_sizer_and_communicator(
-                        sizer=sizer, cube=communicator
+                        sizer=sizer, comm=communicator
                     )
                 )
                 stencil_factory = pace.dsl.StencilFactory(
@@ -322,6 +327,9 @@ def write_for_restart(
         config_dict["initialization"]["type"] = "restart"
         config_dict["initialization"]["config"]["start_time"] = time
         config_dict["initialization"]["config"]["path"] = restart_path
+        # restart config doesn't have 'case'
+        if "case" in config_dict["initialization"]["config"].keys():
+            del config_dict["initialization"]["config"]["case"]
         with open(f"{restart_path}/restart.yaml", "w") as file:
             yaml.safe_dump(config_dict, file)
 
@@ -404,11 +412,19 @@ def __init__(
                 if self.config.performance_config.collect_communication
                 else None
             )
-            communicator = CubedSphereCommunicator.from_layout(
-                comm=self.comm,
-                layout=self.config.layout,
-                timer=comm_timer,
-            )
+            communicator: Communicator
+            if self.config.grid_type <= 3:
+                communicator = CubedSphereCommunicator.from_layout(
+                    comm=self.comm,
+                    layout=self.config.layout,
+                    timer=comm_timer,
+                )
+            else:
+                communicator = TileCommunicator.from_layout(
+                    comm=self.comm,
+                    layout=self.config.layout,
+                    timer=comm_timer,
+                )
             self._update_driver_config_with_communicator(communicator)
 
             if self.config.stencil_config.compilation_config.run_mode == RunMode.Build:
@@ -544,7 +560,7 @@ def exit_instead_of_build(self):
         pace_log.info("initialization of the object done")
 
     def _update_driver_config_with_communicator(
-        self, communicator: CubedSphereCommunicator
+        self, communicator: Communicator
     ) -> None:
         dace_config = DaceConfig(
             communicator=communicator,
@@ -707,7 +723,7 @@ def log_subtile_location(partitioner: pace.util.TilePartitioner, rank: int):
 
 def _setup_factories(
     config: DriverConfig,
-    communicator: pace.util.CubedSphereCommunicator,
+    communicator: pace.util.Communicator,
     stencil_compare_comm,
 ) -> Tuple[pace.util.QuantityFactory, pace.dsl.StencilFactory]:
     """
@@ -735,7 +751,7 @@ def _setup_factories(
     )
 
     grid_indexing = pace.dsl.stencil.GridIndexing.from_sizer_and_communicator(
-        sizer=sizer, cube=communicator
+        sizer=sizer, comm=communicator
     )
     quantity_factory = pace.util.QuantityFactory.from_backend(
         sizer, backend=config.stencil_config.compilation_config.backend